[FRONTEND] Add sass to asm dict with lazy evaluation (#2309)

python/triton/compiler/compiler.py

@@ -30,6 +30,14 @@ from .utils import (InfoFromBackendForTensorMap, TensorMapManager,
                     get_ids_of_tensormaps, parse_tma_info)
 
 
+class LazyDict(dict):
+    def __getitem__(self, key):
+        val = dict.__getitem__(self, key)
+        if callable(val):
+            return val()
+        return val
+
+
 def inline_triton_ir(mod):
     pm = ir.pass_manager(mod.context)
     pm.enable_debug()
@@ -489,7 +497,7 @@ def compile(fn, **kwargs):
     metadata["device_type"] = device_type
 
     first_stage = list(stages.keys()).index(ext)
-    asm = dict()
+    asm = LazyDict()
     module = fn
     # run compilation pipeline  and populate metadata
     for ir_name, (parse, compile_kernel) in list(stages.items())[first_stage:]:
@@ -523,11 +531,7 @@ def compile(fn, **kwargs):
 
         if ir_name == "cubin":
             asm[ir_name] = next_module
-            sass_ir = "sass"
-            sass_fname = f"{name}.{sass_ir}"
-            asm[sass_ir] = get_sass(next_module)
-            metadata_group[sass_fname] = fn_cache_manager.put(asm[sass_ir], sass_fname)
-
+            asm["sass"] = lambda: get_sass(next_module)
         elif ir_name == "amdgcn":
             asm[ir_name] = str(next_module[0])
         else:

python/triton/tools/disasm.py

@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import functools
 import os
 import re
 import subprocess
@@ -64,6 +65,7 @@ def processSassLines(fline, sline, labels):
     return (f'{ctrl}', f'{asm}')
 
 
+@functools.lru_cache()
 def get_sass(cubin_asm, fun=None):
     fd, path = tempfile.mkstemp()
     try: