Add ptxas_config autotuning option

stack-info: PR: #793, branch: jansel/stack/158

Author: jansel

helion/_compat.py

@@ -88,3 +88,12 @@ def warps_to_threads(num_warps: int) -> int:
         )
         return num_warps * (props.warp_size or 32)
     return num_warps * 32
+
+
+def supports_ptxas(device: torch.device) -> bool:
+    """Return True if PTXAS controls are supported for the given device."""
+    if device.type != "cuda":
+        return False
+    if torch.version.hip is not None:
+        return False
+    return supports_tensor_descriptor()

helion/_compiler/compile_environment.py

@@ -20,6 +20,7 @@
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
 from .. import exc
+from .._compat import supports_ptxas
 from ..language.constexpr import ConstExpr
 from .loop_dependency_checker import LoopDependencyChecker
 from .source_location import SourceLocation
@@ -90,6 +91,7 @@ def __init__(self, device: torch.device, settings: Settings) -> None:
         self.block_sizes: list[BlockSizeInfo] = []
         self.debug_shape_renames: dict[sympy.Expr, sympy.Expr] = {}
         self.config_spec = ConfigSpec()
+        self.config_spec.ptxas_supported = supports_ptxas(device)
         self.kernel_tensor_sizes: dict[tuple[sympy.Expr, ...], int] = (
             collections.Counter()
         )

helion/_compiler/device_function.py

@@ -573,6 +573,12 @@ def codegen_function_call(self) -> ast.AST:
                 f"num_stages={self.config.num_stages}",
             ]
         )
+        ptxas_config = self.config.ptxas_config
+        if ptxas_config:
+            from ..runtime.ptxas_configs import get_ptxas_option
+
+            ptx_option = get_ptxas_option(ptxas_config)
+            args.append(f"ptx_options={ptx_option!r}")
         pid = self.pid
         assert pid is not None
         # TODO(jansel): we should run CSE this statement

helion/autotuner/config_spec.py

@@ -50,6 +50,7 @@
         "num_stages",
         "pid_type",
         "indexing",
+        "ptxas_config",
     ]
 )
 VALID_PID_TYPES = ("flat", "xyz", "persistent_blocked", "persistent_interleaved")
@@ -97,6 +98,7 @@ class ConfigSpec:
         default_factory=functools.partial(tuple, VALID_PID_TYPES)
     )
     grid_block_ids: list[int] = dataclasses.field(default_factory=list)
+    ptxas_supported: bool = True
 
     @staticmethod
     def _valid_indexing_types() -> tuple[IndexingLiteral, ...]:
@@ -226,6 +228,11 @@ def normalize(self, config: helion.Config | dict[str, object]) -> None:
             else:
                 config[name] = values[0]
 
+        if self.ptxas_supported:
+            value = config.get("ptxas_config") or 0
+            if not isinstance(value, int):
+                raise InvalidConfig(f"ptxas_config must be integer, got {value!r}")
+
         # Set default values for grid indices when pid_type is not persistent
         pid_type = config["pid_type"]
         if pid_type in ("flat", "xyz") and self.grid_block_ids:
@@ -267,6 +274,10 @@ def flat_config(self, fn: Callable[[ConfigSpecFragment], object]) -> helion.Conf
             "indexing": fn(EnumFragment(self._valid_indexing_types())),
             "pid_type": fn(EnumFragment(self.allowed_pid_types)),
         }
+        if self.ptxas_supported:
+            from ..runtime.ptxas_configs import search_ptxas_configs
+
+            config["ptxas_config"] = fn(EnumFragment((0, *search_ptxas_configs())))
         # Add tunable parameters
         for key, fragment in self.user_defined_tunables.items():
             config[key] = fn(fragment)

helion/runtime/__init__.py

@@ -61,8 +61,15 @@ def default_launcher(
     *args: object,
     num_warps: int,
     num_stages: int,
+    ptx_options: str | None = None,
 ) -> object:
     """Default launcher function that executes the kernel immediately."""
-    return triton_kernel.run(
-        *args, grid=grid, warmup=False, num_warps=num_warps, num_stages=num_stages
-    )
+    run_kwargs = {
+        "grid": grid,
+        "warmup": False,
+        "num_warps": num_warps,
+        "num_stages": num_stages,
+    }
+    if ptx_options:
+        run_kwargs["ptx_options"] = ptx_options
+    return triton_kernel.run(*args, **run_kwargs)

helion/runtime/config.py

@@ -38,6 +38,7 @@ def __init__(
         num_stages: int | None = None,
         pid_type: PidTypeLiteral | None = None,
         indexing: IndexingLiteral | None = None,
+        ptxas_config: int | None = None,
         # For user-defined properties
         **kwargs: object,
     ) -> None:
@@ -78,6 +79,7 @@ def __init__(
             "num_stages": num_stages,
             "indexing": indexing,
             "pid_type": pid_type,
+            "ptxas_config": ptxas_config,
         }
         for key, value in core_props.items():
             if value is not None:
@@ -169,6 +171,10 @@ def pid_type(self) -> PidTypeLiteral:
     def range_unroll_factors(self) -> list[int]:
         return cast("list[int]", self.config.get("range_unroll_factors", []))
 
+    @property
+    def ptxas_config(self) -> int:
+        return cast("int", self.config.get("ptxas_config", 0))
+
     @property
     def range_warp_specializes(self) -> list[bool | None]:
         return cast("list[bool | None]", self.config.get("range_warp_specializes", []))

helion/runtime/ptxas_configs/__init__.py

@@ -0,0 +1,43 @@
+"""Utilities for working with packaged PTXAS control files."""
+
+from __future__ import annotations
+
+from functools import cache
+from pathlib import Path
+
+_CONFIG_FILES: dict[int, str] = {
+    1: "spiffy-bee-104.bin",
+}
+
+
+def _config_root() -> Path:
+    return Path(__file__).resolve().parent
+
+
+@cache
+def search_ptxas_configs() -> tuple[int, ...]:
+    """Return the sorted tuple of available PTXAS config IDs."""
+
+    return tuple(sorted(_CONFIG_FILES))
+
+
+def _config_file_path(config_id: int) -> str:
+    """Return the absolute path to the PTXAS control file for ``config_id``."""
+
+    try:
+        filename = _CONFIG_FILES[config_id]
+    except KeyError as exc:  # pragma: no cover - defensive
+        raise ValueError(f"Unknown PTXAS config id: {config_id}") from exc
+    resolved = (_config_root() / filename).resolve()
+    if not resolved.is_file():
+        raise FileNotFoundError(f"Missing PTXAS config file: {resolved}")
+    return str(resolved)
+
+
+@cache
+def get_ptxas_option(config_value: int) -> str | None:
+    """Translate a config enum value into a PTXAS option string."""
+
+    if config_value == 0:
+        return None
+    return f"--apply-controls {_config_file_path(config_value)}"

helion/runtime/ptxas_configs/spiffy-bee-104.bin

@@ -85,6 +85,7 @@ packages = ["helion"]
 include = [
   "helion/**/*.py",
   "helion/**/*.pyi",
+  "helion/runtime/ptxas_configs/*.bin",
   "LICENSE",
 ]
 exclude = [

pyproject.toml

@@ -0,0 +1,27 @@
+This file is automatically generated by assertExpectedJournal calls in test_ptxas_config.py.
+Update expected outputs by running tests with the EXPECTTEST_ACCEPT=1 environment variable set.
+
+--- assertExpectedJournal(TestPtxasConfig.test_ptxas_config_apply_controls_flag)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+@triton.jit
+def _helion__copy_kernel(x_flat, out_flat, x_size_0, out_flat_stride_0, x_flat_stride_0, _BLOCK_SIZE_0: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0 * _BLOCK_SIZE_0
+    indices_0 = (offset_0 + tl.arange(0, _BLOCK_SIZE_0)).to(tl.int32)
+    mask_0 = indices_0 < x_size_0
+    load = tl.load(x_flat + indices_0 * x_flat_stride_0, mask_0, other=0)
+    tl.store(out_flat + indices_0 * out_flat_stride_0, load, mask_0)
+
+def _copy_kernel(x: torch.Tensor, *, _launcher=_default_launcher):
+    out = torch.empty_like(x)
+    x_flat = x.view(-1)
+    out_flat = out.view(-1)
+    _BLOCK_SIZE_0 = 32
+    _launcher(_helion__copy_kernel, (triton.cdiv(x.size(0), _BLOCK_SIZE_0),), x_flat, out_flat, x.size(0), out_flat.stride(0), x_flat.stride(0), _BLOCK_SIZE_0, num_warps=4, num_stages=3, ptx_options='--apply-controls <path>')
+    return out