From 511373abbf6c601a37fddf90732a5e6cf9904b4c Mon Sep 17 00:00:00 2001 From: minjang Date: Mon, 10 Jun 2024 15:25:59 -0700 Subject: [PATCH 1/2] [CPU] Dump human-readable asm code in TRITON_CACHE_DIR --- python/triton/compiler/compiler.py | 7 +++++++ third_party/cpu/backend/compiler.py | 3 --- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python/triton/compiler/compiler.py b/python/triton/compiler/compiler.py index 6eb7aa292b4f..f301e6a0167b 100644 --- a/python/triton/compiler/compiler.py +++ b/python/triton/compiler/compiler.py @@ -293,6 +293,13 @@ def compile(src, target=None, options=None): ttgir_full_name = fn_cache_manager.get_file(ir_filename) next_module = parse(ttgir_full_name, ext, context) print(f"re-parse ttgir with {ttgir_full_name}") + # It's ugly, but a quick hack to generate human-readable asm. + if ext == "bc" and backend.target.backend == "cpu": + from triton._C.libtriton import llvm + + asm_filename = f"{src.name}.asm" + asm = llvm.translate_to_host_asm(module, options.enable_fp_fusion) + metadata_group[asm_filename] = fn_cache_manager.put(asm, asm_filename) module = next_module # write-back metadata metadata_group[metadata_filename] = fn_cache_manager.put(json.dumps(metadata, default=vars), metadata_filename, diff --git a/third_party/cpu/backend/compiler.py b/third_party/cpu/backend/compiler.py index d48fbf3a96bf..86b59690487b 100644 --- a/third_party/cpu/backend/compiler.py +++ b/third_party/cpu/backend/compiler.py @@ -140,9 +140,6 @@ 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)) ret = llvm.translate_to_bc(src) return ret From a57e2f7793571224273305920919343a6d216870 Mon Sep 17 00:00:00 2001 From: minjang Date: Mon, 10 Jun 2024 20:36:34 -0700 Subject: [PATCH 2/2] Don't touch the main compiler.py --- python/triton/compiler/compiler.py | 7 ------- third_party/cpu/backend/compiler.py | 8 +++++++- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/python/triton/compiler/compiler.py b/python/triton/compiler/compiler.py index f301e6a0167b..6eb7aa292b4f 100644 --- a/python/triton/compiler/compiler.py +++ b/python/triton/compiler/compiler.py @@ -293,13 +293,6 @@ def compile(src, target=None, options=None): ttgir_full_name = fn_cache_manager.get_file(ir_filename) next_module = parse(ttgir_full_name, ext, context) print(f"re-parse ttgir with {ttgir_full_name}") - # It's ugly, but a quick hack to generate human-readable asm. - if ext == "bc" and backend.target.backend == "cpu": - from triton._C.libtriton import llvm - - asm_filename = f"{src.name}.asm" - asm = llvm.translate_to_host_asm(module, options.enable_fp_fusion) - metadata_group[asm_filename] = fn_cache_manager.put(asm, asm_filename) module = next_module # write-back metadata metadata_group[metadata_filename] = fn_cache_manager.put(json.dumps(metadata, default=vars), metadata_filename, diff --git a/third_party/cpu/backend/compiler.py b/third_party/cpu/backend/compiler.py index 86b59690487b..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 @@ -140,6 +139,13 @@ def make_llir(src, metadata, options): @staticmethod def make_bc(src, metadata, options): + if os.environ.get("TRITON_CPU_ASM_DUMP", "0") == "1": + 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