Suggest HELION_AUTOTUNE_PRECOMPILE=spawn when IMA happens (#984)

Author: jansel

helion/autotuner/base_search.py

@@ -47,6 +47,7 @@
 from .logger import LambdaLogger
 from .logger import classify_triton_exception
 from .logger import format_triton_compile_failure
+from .logger import match_unrecoverable_runtime_error
 from .progress_bar import iter_with_progress
 
 if TYPE_CHECKING:
@@ -300,6 +301,14 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
                 self.best_perf_so_far = res
             return res
         except Exception as e:
+            if match_unrecoverable_runtime_error(e):
+                raise exc.TritonUnrecoverableRuntimeError(
+                    reason=str(e),
+                    decorator=self.kernel.format_kernel_decorator(
+                        config, self.settings
+                    ),
+                    error=f"{type(e).__qualname__}: {e}",
+                ) from e
             action = classify_triton_exception(e)
             if self.settings.autotune_ignore_errors:
                 pass

helion/autotuner/logger.py

@@ -121,15 +121,30 @@ def format_triton_compile_failure(
                 "triton.compiler.errors.CompilationError",  # Triton CompilationError
                 "out of resource: shared memory",  # Triton shared memory OOM
                 "ZE_RESULT_ERROR_INVALID_KERNEL_NAME",  # Level Zero compile failed
-                "an illegal memory access was encountered",  # workaround triton bugs
-                "misaligned address",  # workaround triton bugs
-                "unspecified launch failure",  # workaround ptxas bugs
                 "exceeds triton maximum tensor numel",  # needs smaller config
             ],
         )
     )
 )
 
+_UNRECOVERABLE_RUNTIME_ERROR_RE: re.Pattern[str] = re.compile(
+    "|".join(
+        map(
+            re.escape,
+            [
+                "illegal memory access",
+                "misaligned address",
+                "unspecified launch failure",
+            ],
+        )
+    ),
+    re.IGNORECASE,
+)
+
+
+def match_unrecoverable_runtime_error(err: BaseException) -> bool:
+    return bool(_UNRECOVERABLE_RUNTIME_ERROR_RE.search(str(err)))
+
 
 def classify_triton_exception(err: BaseException) -> Literal["raise", "warn", "debug"]:
     """
@@ -150,6 +165,6 @@ def classify_triton_exception(err: BaseException) -> Literal["raise", "warn", "d
     msg = str(err)
     if "PassManager::run failed" in msg:
         return "warn"
-    if _EXPECTED_TRITON_ERRORS_RE.search(msg):
+    if _EXPECTED_TRITON_ERRORS_RE.search(msg) or match_unrecoverable_runtime_error(err):
         return "debug"
     return "raise"

helion/exc.py

@@ -345,6 +345,15 @@ class TritonError(BaseError):
 Set autotune_ignore_errors=True or HELION_AUTOTUNE_IGNORE_ERRORS=1 to ignore Triton errors in autotuning."""
 
 
+class TritonUnrecoverableRuntimeError(BaseError):
+    message = """\
+An unrecoverable Triton runtime error occurred: {reason}.
+This likely indicates a bug in Triton and cannot be recovered from.
+{decorator}
+Original error: {error}
+Set HELION_AUTOTUNE_PRECOMPILE="spawn" to isolate these errors in a subprocess so tuning can continue."""
+
+
 class BaseWarning(_FixedMessage):
     message = "A warning occurred."