Always set triton allocator (#416)

Author: jansel

helion/_compiler/device_function.py

@@ -419,7 +419,7 @@ def sorted_args(self) -> list[Argument]:
 
     def codegen_function_def(self) -> list[ast.stmt]:
         prefix = []
-        if self._tensor_descriptor_args:
+        if CompileEnvironment.current().settings.set_triton_allocator:
             prefix.append(
                 statement_from_string("helion.runtime.set_triton_allocator()")
             )

helion/runtime/__init__.py

@@ -19,15 +19,8 @@ def _alloc_fn(size: int, alignment: int, stream: int | None) -> torch.Tensor:
 
 @functools.cache
 def set_triton_allocator() -> None:
-    try:
-        from triton.runtime._allocation import NullAllocator
-        from triton.runtime._allocation import _allocator
-
-        if not isinstance(_allocator, NullAllocator):
-            return
-    except ImportError:
-        pass
-    triton.set_allocator(_alloc_fn)
+    if hasattr(triton, "set_allocator"):
+        triton.set_allocator(_alloc_fn)
 
 
 def get_num_sm(device: torch.device) -> int:

helion/runtime/settings.py

@@ -95,6 +95,7 @@ class _Settings:
         RefMode.EAGER if os.environ.get("HELION_INTERPRET", "") == "1" else RefMode.OFF
     )
     autotuner_fn: AutotunerFunction = default_autotuner_fn
+    set_triton_allocator: bool = True
 
 
 class Settings(_Settings):
@@ -117,6 +118,7 @@ class Settings(_Settings):
         "allow_warp_specialize": "If True, allow warp specialization for tl.range calls on CUDA devices.",
         "ref_mode": "Reference mode for kernel execution. Can be RefMode.OFF or RefMode.EAGER.",
         "autotuner_fn": "Function to create an autotuner",
+        "set_triton_allocator": "If True, insert helion.runtime.set_triton_allocator() call in generated code. Default is True.",
     }
     assert __slots__.keys() == {field.name for field in dataclasses.fields(_Settings)}