Merge branch 'boyanl/debug-artifacts' into 'main'

blinxt · blinxt · commit 09be17039532 · 2025-11-24T13:32:46.000-08:00
dump debug artifacts when TileCompilerError is raised

See merge request dl/tileir/cutile-python!32
diff --git a/src/cuda/tile/_compile.py b/src/cuda/tile/_compile.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from dataclasses import dataclass
+import datetime
 import functools
 from functools import cache
 import logging
@@ -15,12 +16,17 @@
 import threading
 import traceback
 from typing import Callable, Optional
+import zipfile
 
 from cuda.tile._ast2ir import get_function_ir
 from cuda.tile._cext import get_compute_capability, TileContext, default_tile_context
 from cuda.tile._compiler_options import CompilerOptions
 from cuda.tile._const_utils import get_constant_annotations
-from cuda.tile._exception import TileCompilerError, TileCompilerTimeoutError
+from cuda.tile._exception import (
+    TileCompilerError,
+    TileCompilerExecutionError,
+    TileCompilerTimeoutError,
+)
 from cuda.tile._ir import ir
 from cuda.tile._passes.code_motion import hoist_loop_invariants
 from cuda.tile._passes.loop_split import split_loops
@@ -36,6 +42,7 @@
 from cuda.tile._passes.dce import dead_code_elimination_pass
 from cuda.tile._passes.token_order import token_order_pass
 from cuda.tile._ir2bytecode import generate_bytecode_for_kernel
+from cuda.tile._version import __version__ as cutile_version
 import cuda.tile._bytecode as bc
 
 
@@ -101,6 +108,38 @@ def _log_mlir(bytecode_buf):
     print(f"Lowering\n==== TILEIR MLIR module ====\n\n{text}", file=sys.stderr)
 
 
+def _compiler_crash_dump(func_ir,
+                         bytecode_generator,
+                         error_msg,
+                         compiler_flags,
+                         compiler_version):
+    debug_info = (
+        f"error:\n{error_msg}\n\n"
+        f"compiler flags:\n{compiler_flags}\n\n"
+        f"compiler version:\n{compiler_version or 'Unkown'}\n\n"
+        f"cutile version:\n{cutile_version}\n"
+    )
+
+    # Anonymize debug attributes in the bytecode
+    bytecode_buf = bytearray()
+    with bc.write_bytecode(num_functions=1, buf=bytecode_buf) as writer:
+        bytecode_generator(writer, anonymize_debug_attr=True)
+
+    artifacts = {
+        f"{func_ir.qualname}.bytecode": bytes(bytecode_buf),
+        f"{func_ir.qualname}.cutileir": f"{func_ir.to_string(include_loc=False)}\n",
+        "debug_info.txt": debug_info,
+    }
+
+    timestamp = datetime.datetime.now().timestamp()
+    zip_filename = os.path.abspath(f"crash_dump_{func_ir.qualname}_{timestamp}.zip")
+    print(f"Dumping crash artifacts to {zip_filename}\n", file=sys.stderr)
+
+    with zipfile.ZipFile(zip_filename, "w") as z:
+        for filename, content in artifacts.items():
+            z.writestr(filename, content)
+
+
 @global_compiler_lock
 def compile_tile(pyfunc,
                  args,
@@ -115,9 +154,12 @@ def compile_tile(pyfunc,
 
     sm_arch = get_sm_arch()
 
+    bytecode_generator = functools.partial(generate_bytecode_for_kernel,
+                                           func_ir, compiler_options, sm_arch)
+
     bytecode_buf = bytearray()
     with bc.write_bytecode(num_functions=1, buf=bytecode_buf) as writer:
-        generate_bytecode_for_kernel(func_ir, compiler_options, sm_arch, writer)
+        bytecode_generator(writer, anonymize_debug_attr=False)
 
     if 'TILEIR' in context.config.log_keys:
         _log_mlir(bytecode_buf)
@@ -150,14 +192,21 @@ def compile_tile(pyfunc,
             print("Can't print MLIR because the internal extension is missing", file=sys.stderr)
 
     # Compile MLIR module and generate cubin
-    with tempfile.NamedTemporaryFile(suffix='.mlirbc', prefix=func_ir.qualname,
+    with tempfile.NamedTemporaryFile(suffix='.bytecode', prefix=func_ir.qualname,
                                      dir=context.config.temp_dir, delete=False) as f:
         f.write(bytecode_buf)
         f.flush()
-        cubin_file = compile_cubin(f.name,
-                                   compiler_options,
-                                   sm_arch,
-                                   timeout_sec=context.config.compiler_timeout_sec)
+
+        try:
+            cubin_file = compile_cubin(f.name, compiler_options, sm_arch,
+                                       timeout_sec=context.config.compiler_timeout_sec)
+        except TileCompilerError as e:
+            if context.config.enable_crash_dump:
+                _compiler_crash_dump(func_ir, bytecode_generator, e.message,
+                                     e.compiler_flags, e.compiler_version)
+
+            raise e
+
     return TileLibrary(func_ir.qualname, cubin_file, bytecode_buf, func_ir)
 
 
@@ -223,6 +272,15 @@ def _find_compiler_bin() -> tuple[str, str, str]:
                             f"make sure it is available in $PATH or ${cuda_home_var}/bin")
 
 
+def _try_get_compiler_version(compiler_bin) -> Optional[str]:
+    try:
+        res = subprocess.run([str(compiler_bin), "--version"],
+                             check=True, capture_output=True, text=True)
+        return res.stdout
+    except Exception:
+        return None
+
+
 @cache
 def get_sm_arch() -> str:
     major, minor = get_compute_capability()
@@ -237,30 +295,37 @@ def compile_cubin(
     compiler_bin, bin_path, ld_path = _find_compiler_bin()
     fname_cubin = Path(fname_bytecode).with_suffix(".cubin")
     compiler_hints = compiler_options.specialize_for_target(sm_arch)
+
     command = [
         str(compiler_bin),
         str(fname_bytecode),
+        "-o",
+        str(fname_cubin),
+    ]
+
+    flags = [
         "--gpu-name",
         sm_arch,
         f"-O{compiler_hints.opt_level}",
-        "-o",
-        str(fname_cubin),
+        "--lineinfo"
     ]
-    # compile with line info
-    command.append("--lineinfo")
-    logger.debug(f"Invoke tile compiler: {' '.join(command)}\n"
+
+    logger.debug(f"Invoke tile compiler: {' '.join(command + flags)}\n"
                  f"LD_LIBRARY_PATH:{ld_path}\n"
                  f"PATH:{bin_path}")
     try:
         env = os.environ.copy()
         env['LD_LIBRARY_PATH'] = ld_path
         env['PATH'] = bin_path
-        subprocess.run(command, env=env, check=True, capture_output=True, timeout=timeout_sec)
+        subprocess.run(command + flags, env=env, check=True, capture_output=True,
+                       timeout=timeout_sec)
     except subprocess.CalledProcessError as e:
-        raise TileCompilerError(e.returncode, e.stderr.decode())
+        raise TileCompilerExecutionError(e.returncode, e.stderr.decode(), ' '.join(flags),
+                                         _try_get_compiler_version(compiler_bin))
     except subprocess.TimeoutExpired:
         message = (f"`tileiras` compiler exceeded timeout {timeout_sec}s. "
                    "Using a smaller tile size may reduce compilation time.")
-        raise TileCompilerTimeoutError(message)
+        raise TileCompilerTimeoutError(message, ' '.join(flags),
+                                       _try_get_compiler_version(compiler_bin))
 
     return fname_cubin
diff --git a/src/cuda/tile/_context.py b/src/cuda/tile/_context.py
@@ -15,13 +15,15 @@ class TileContextConfig:
     temp_dir: str
     log_keys: list[str]
     compiler_timeout_sec: Optional[int]
+    enable_crash_dump: bool
 
 
 def init_context_config_from_env():
     config = TileContextConfig(
             temp_dir=get_temp_dir_from_env(),
             log_keys=get_log_keys_from_env(),
-            compiler_timeout_sec=get_compile_timeout_from_env()
+            compiler_timeout_sec=get_compile_timeout_from_env(),
+            enable_crash_dump=get_enable_crash_dump_from_env()
             )
     return config
 
@@ -61,3 +63,9 @@ def get_temp_dir_from_env() -> str:
         dir = tempfile.mkdtemp()
         atexit.register(_clean_tmp_dir, dir)
     return dir
+
+
+def get_enable_crash_dump_from_env() -> bool:
+    key = "CUDA_TILE_ENABLE_CRASH_DUMP"
+    env = os.environ.get(key, "0").lower()
+    return env in ("1", "true", "yes", "on")
diff --git a/src/cuda/tile/_debug.py b/src/cuda/tile/_debug.py
@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 import os
 
 # Comma separated string for selective debug logging
@@ -36,3 +37,6 @@ def parse_cuda_tile_logs_env():
 
 CUDA_TILE_TESTING_DISABLE_TOKEN_ORDER = (
     os.environ.get("CUDA_TILE_TESTING_DISABLE_TOKEN_ORDER", "0") == "1")
+
+
+logger = logging.getLogger(__name__)
diff --git a/src/cuda/tile/_exception.py b/src/cuda/tile/_exception.py
@@ -156,12 +156,32 @@ def _parse_tileir_stderr(stderr: str) -> tuple[str, Optional[Loc]]:
 
 
 class TileCompilerError(TileInternalError):
-    def __init__(self, return_code: int, stderr: str):
+    def __init__(self,
+                 message: str,
+                 loc: Loc,
+                 compiler_flags: str,
+                 compiler_version: Optional[str]):
+        super().__init__(message, loc)
+        self.compiler_flags = compiler_flags
+        self.compiler_version = compiler_version
+
+
+class TileCompilerExecutionError(TileCompilerError):
+    def __init__(self,
+                 return_code: int,
+                 stderr: str,
+                 compiler_flags: str,
+                 compiler_version: Optional[str]):
         message, loc = _parse_tileir_stderr(stderr)
         if loc is None:
             loc = _unknown_loc
-        super().__init__(f"Return code {return_code}\n{message}", loc)
+        super().__init__(f"Return code {return_code}\n{message}", loc,
+                         compiler_flags, compiler_version)
 
 
-class TileCompilerTimeoutError(TileInternalError):
-    pass
+class TileCompilerTimeoutError(TileCompilerError):
+    def __init__(self,
+                 message: str,
+                 compiler_flags: str,
+                 compiler_version: Optional[str]):
+        super().__init__(message, _unknown_loc, compiler_flags, compiler_version)
diff --git a/src/cuda/tile/_ir2bytecode.py b/src/cuda/tile/_ir2bytecode.py
@@ -466,10 +466,11 @@ def lower_scan(ctx: "BytecodeContext", x: bc.Value, input_ty: Type,
 
 
 class DebugAttrMap:
-    def __init__(self, debug_attr_table: bc.DebugAttrTable, linkage_name: str):
+    def __init__(self, debug_attr_table: bc.DebugAttrTable, linkage_name: str, anonymize: bool):
         self._subprogram_cache = {}
         self._debug_attr_table = debug_attr_table
         self._linkage_name = linkage_name
+        self._anonymize = anonymize
 
     def get_subprogram(self, pyfunc) -> bc.DebugAttrId:
         try:
@@ -495,6 +496,9 @@ def get_subprogram(self, pyfunc) -> bc.DebugAttrId:
         return ret
 
     def get_debugattr(self, loc: Loc) -> bc.DebugAttrId:
+        if self._anonymize:
+            return bc.MISSING_DEBUG_ATTR_ID
+
         subprogram = self.get_subprogram(loc.function)
         attr = self._debug_attr_table.loc(subprogram, loc.filename, loc.line, loc.col)
         if loc.call_site is not None:
@@ -809,7 +813,8 @@ def generate_bytecode_for_block(ctx: BytecodeContext, block: Block):
 def generate_bytecode_for_kernel(func_ir: Function,
                                  compiler_options: CompilerOptions,
                                  sm_arch: str,
-                                 writer: bc.BytecodeWriter):
+                                 writer: bc.BytecodeWriter,
+                                 anonymize_debug_attr: bool):
     target_options = compiler_options.specialize_for_target(sm_arch)
     entry_hints = bc.EntryHints(num_cta_in_cga=target_options.num_ctas,
                                 occupancy=target_options.occupancy)
@@ -822,7 +827,7 @@ def generate_bytecode_for_kernel(func_ir: Function,
         param_type_ids.extend(typeid_tuple(writer.type_table, ty))
     param_offsets.append(len(param_type_ids))
 
-    debug_attr_map = DebugAttrMap(writer.debug_attr_table, func_ir.qualname)
+    debug_attr_map = DebugAttrMap(writer.debug_attr_table, func_ir.qualname, anonymize_debug_attr)
     func_debug_attr = debug_attr_map.get_debugattr(func_ir.loc)
 
     with writer.function(name=func_ir.qualname,
