Add shader dump when compile triton (#4482)

Add dump information for debug  when autotune the Triton kernel:
1. Print the Triton cache director when dump the autotuner information
with `TRITON_PRINT_AUTOTUNING`
2. Add a new knob `TRITON_INTEL_ENABLE_IGC_SHADER_DUMP` to dump the
shader information in IGC when compiling the Triton kernel.

Signed-off-by: Lu,Chengjun <[email protected]>

Author: chengjunlu

python/triton/compiler/compiler.py

@@ -320,6 +320,8 @@ def compile(src, target=None, options=None):
         **options.__dict__,
         **env_vars,
     }
+
+    metadata["cache_dir"] = fn_cache_manager.cache_dir
     metadata["triton_version"] = __version__
     # run compilation pipeline  and populate metadata
     stages = dict()

python/triton/knobs.py

@@ -489,6 +489,7 @@ class nvidia_knobs(base_knobs):
 class intel_knobs(base_knobs):
     spirv_dis: env_intel_tool = env_intel_tool("spirv-dis")
 
+    dump_shader_info: env_bool = env_bool("TRITON_INTEL_ENABLE_IGC_SHADER_DUMP", False)
     gen_native_code: env_bool = env_bool("TRITON_XPU_GEN_NATIVE_CODE", False)
     tile_load_ll: env_bool = env_bool("TRITON_XPU_ENABLE_TILE_LOAD_LINEAR_LAYOUT", True)
     advanced_path: env_bool = env_bool("TRITON_INTEL_ADVANCED_PATH", False)

third_party/intel/backend/compiler.py

@@ -55,7 +55,8 @@ def __post_init__(self):
         object.__setattr__(self, 'extern_libs', tuple(extern_libs.items()))
         if self.num_warps <= 0 or (self.num_warps & (self.num_warps - 1)) != 0:
             raise AssertionError("num_warps must be a power of 2")
-        self.generate_native_code = knobs.intel.gen_native_code or self.generate_native_code
+        self.generate_native_code = (knobs.intel.gen_native_code
+                                     or knobs.intel.dump_shader_info) or self.generate_native_code
 
     def hash(self):
         key = '_'.join([f'{name}-{val}' for name, val in self.__dict__.items()])
@@ -390,6 +391,13 @@ def make_spv(src, metadata, options):
         if knobs.intel.disable_igc_opt:
             metadata["build_flags"] += " -cl-opt-disable"
 
+        shader_dump_opt = ""
+        if knobs.intel.dump_shader_info:
+            # The IGC (Intel Graphic Compiler) only parses the options at first time in JIT-ing the binary per process.
+            # Have to use the `ocloc` to generate the binary in sub-process to work around the limitation.
+            assert options.generate_native_code, "Only support native code generation with shader dump"
+            shader_dump_opt = f" -igc_opts ',DumpToCustomDir={metadata['cache_dir']},ShaderDumpEnable=1'"
+
         metadata["generate_native_code"] = options.generate_native_code
 
         if options.generate_native_code:
@@ -401,7 +409,7 @@ def make_spv(src, metadata, options):
 
                 ocloc_cmd = [
                     'ocloc', 'compile', '-file', fsrc.name, '-o', fbin, '-spirv_input', '-device', 'pvc', '-options',
-                    metadata["build_flags"]
+                    metadata["build_flags"] + shader_dump_opt
                 ]
 
                 try:
@@ -417,7 +425,7 @@ def make_spv(src, metadata, options):
                                 """
                                 metadata["build_flags"] += " -cl-intel-256-GRF-per-thread"
                                 # re-run with new build flags
-                                ocloc_cmd[-1] = metadata["build_flags"]
+                                ocloc_cmd[-1] = metadata["build_flags"] + shader_dump_opt
                                 subprocess.run(ocloc_cmd, check=True, close_fds=False, stdout=flog,
                                                stderr=subprocess.STDOUT)
                         os.remove(flog.name)