From 38ee9e5331b64652a6d87b34b37f61e5ae746503 Mon Sep 17 00:00:00 2001 From: Minjang Kim Date: Tue, 11 Jun 2024 08:40:15 -0700 Subject: [PATCH] [CPU] Dump human-readable asm code in TRITON_CACHE_DIR (#19) * [CPU] Dump human-readable asm code in TRITON_CACHE_DIR * Don't touch the main compiler.py --- third_party/cpu/backend/compiler.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/third_party/cpu/backend/compiler.py b/third_party/cpu/backend/compiler.py index d48fbf3a96bf..c3f11334750a 100644 --- a/third_party/cpu/backend/compiler.py +++ b/third_party/cpu/backend/compiler.py @@ -1,7 +1,6 @@ import functools import hashlib import os -import re from dataclasses import dataclass from typing import Any, Tuple @@ -141,8 +140,12 @@ def make_llir(src, metadata, options): @staticmethod def make_bc(src, metadata, options): if os.environ.get("TRITON_CPU_ASM_DUMP", "0") == "1": - print("********** Module ASM **********") - print(llvm.translate_to_host_asm(src, options.enable_fp_fusion)) + from triton.runtime.cache import get_cache_manager + + asm = llvm.translate_to_host_asm(src, options.enable_fp_fusion) + fn_cache_manager = get_cache_manager(metadata['hash']) + fn_cache_manager.put(asm, f"{metadata['name']}.asm") + ret = llvm.translate_to_bc(src) return ret