from __future__ import annotations
import ast
from collections.abc import Sequence
import contextlib
import dataclasses
import functools
import hashlib
import inspect
import itertools
import logging
import operator
import os
import re
import sys
import textwrap
import threading
import types
from typing import TYPE_CHECKING
from typing import Any
from typing import Callable
from typing import Generic
from typing import Hashable
from typing import Literal
from typing import NamedTuple
from typing import TypeVar
from typing import cast
from typing import overload
from typing_extensions import Protocol
import weakref
import sympy
import torch
from torch._dynamo.source import GetItemSource
from torch._dynamo.source import LocalSource
from torch._dynamo.source import TensorProperty
from torch._dynamo.source import TensorPropertySource
from torch._inductor.codecache import PyCodeCache
from torch._inductor.codecache import compiled_fx_graph_hash
from torch._subclasses import FakeTensor
from torch._subclasses.fake_tensor import unset_fake_temporarily
import torch.distributed as dist
from torch.utils._pytree import tree_map_only
from torch.utils.weak import WeakIdKeyDictionary
from .. import exc
from .._compat import shape_env_size_hint
from .._compat import target_device_capability
from .._compile_time import measure
from .._compiler.ast_extension import unparse
from .._compiler.autotuner_heuristics import compiler_seed_configs
from .._compiler.compile_environment import CompileEnvironment
from .._compiler.compile_environment import TensorDescriptorLayoutGuard
from .._compiler.compile_environment import _is_supported_tensor_input_source
from .._compiler.compile_environment import _symint_free_symbols
from .._compiler.compile_environment import (
tensor_descriptor_layout_signature_from_strides,
)
from .._compiler.generate_ast import generate_ast
from .._compiler.inductor_lowering_extra import patch_inductor_lowerings
from .._compiler.kernel_compiler import KernelCompiler
from .._compiler.output_header import assert_no_conflicts
from .._compiler.variable_origin import ArgumentOrigin
from .._dist_utils import _find_process_group_name
from .._dist_utils import check_config_consistancy as dist_check_config_consistancy
from .._dist_utils import kernel_declares_process_group
from .._dist_utils import kernel_uses_symm_mem
from .._logging import LazyString
from .._utils import counters
from ..autotuner.base_search import _AutotunableKernel
from ..language.constexpr import ConstExpr
from .config import Config
from .ref_mode import RefModeContext
from .ref_mode import is_ref_mode_enabled
from .settings import Settings
if TYPE_CHECKING:
from collections.abc import Generator
from collections.abc import Hashable
from torch._guards import Source
from .._compiler.host_function import HostFunction
from ..autotuner import ConfigSpec
from ..autotuner.base_cache import BoundKernelInMemoryCacheKey
ConfigLike = Config | dict[str, object]
log: logging.Logger = logging.getLogger(__name__)
def _indexing_config_uses_tensor_descriptor(indexing: object, index: int) -> bool:
if indexing == "tensor_descriptor":
return True
if isinstance(indexing, list):
return index < len(indexing) and indexing[index] == "tensor_descriptor"
return False
def _td_layout_guard_active_for_config(
guard: TensorDescriptorLayoutGuard, config: Config
) -> bool:
return any(
_indexing_config_uses_tensor_descriptor(config.indexing, index)
for index in guard.memory_op_indices
) or any(
_indexing_config_uses_tensor_descriptor(config.atomic_indexing, index)
for index in guard.atomic_op_indices
)
_R = TypeVar("_R")
CompiledConfig = Callable[..., _R]
# Opt-in: auto-capture Pallas kernels under torch.compile (see
# pallas._tpu_compile_capture).
# Off by default so the eager dispatch path is unchanged.
_TPU_COMPILE_CAPTURE = os.environ.get("HELION_TPU_COMPILE_CAPTURE", "0") == "1"
# Cache for GraphModule hashes
_graph_module_hash_cache: WeakIdKeyDictionary = WeakIdKeyDictionary()
_INT32_INDEX_LIMIT = torch.iinfo(torch.int32).max
_CUSTOM_KEY_UNSET = object()
_HostSemanticInputNormalization = tuple[tuple[int, str, int, int], ...]
class _FastDispatchEntry(NamedTuple):
"""A dispatch key together with the state used to build it."""
key: tuple[Hashable, ...]
extra_guards: tuple[tuple[Callable[[Sequence[object]], Hashable], Hashable], ...]
specialization_generation: int
class _SpecializationAlias:
"""An omitted-default signature backed by a normalized signature."""
__slots__ = ("schemas", "canonical_signature", "trailing_defaults")
def __init__(
self,
schemas: dict[Hashable, list[Callable[[Sequence[object]], Hashable]]],
canonical_signature: tuple[Hashable, ...],
trailing_defaults: tuple[object, ...],
) -> None:
self.schemas = schemas
self.canonical_signature = canonical_signature
self.trailing_defaults = trailing_defaults
def __call__(self, values: Sequence[object]) -> Hashable:
normalized = (*values, *self.trailing_defaults)
return tuple(
extractor(normalized)
for extractor in self.schemas[self.canonical_signature]
)
def _make_prepared_arg_guard(
kernel: Kernel,
args: tuple[object, ...],
) -> Callable[[tuple[object, ...]], bool]:
namespace: dict[str, object] = {}
checks = [f"len(args) == {len(args)}"]
for index, arg in enumerate(args):
prefix = f"args[{index}]"
namespace[f"type_{index}"] = type(arg)
checks.append(f"type({prefix}) is type_{index}")
if type(arg) in (torch.Tensor, torch.nn.Parameter):
assert isinstance(arg, torch.Tensor)
namespace[f"dtype_{index}"] = arg.dtype
namespace[f"shape_{index}"] = arg.shape
namespace[f"stride_{index}"] = arg.stride()
namespace[f"device_{index}"] = arg.device
checks.extend(
(
f"{prefix}.dtype is dtype_{index}",
f"{prefix}.shape == shape_{index}",
f"{prefix}.stride() == stride_{index}",
f"{prefix}.device == device_{index}",
)
)
static_indices = getattr(arg, "_dynamo_static_indices", None)
if static_indices is None:
checks.append(
f"getattr({prefix}, '_dynamo_static_indices', None) is None"
)
else:
namespace[f"static_indices_{index}"] = frozenset(static_indices)
checks.extend(
(
f"getattr({prefix}, '_dynamo_static_indices', None) is not None",
f"frozenset({prefix}._dynamo_static_indices) == static_indices_{index}",
)
)
elif type(arg) in (bool, int, float):
# Ordinary numeric arguments are runtime values in Helion's base
# specialization. A constexpr annotation is the one exception;
# hl.specialize() constraints are covered by ``_extra_guards``.
if kernel._annotations[index] is ConstExpr:
if type(arg) is float and arg != arg:
# Python's equality/hash behavior does not provide a stable
# equivalence class for constexpr NaNs. Keep those calls on
# the regular specialization path rather than weakening its
# semantics in the prepared guard.
raise TypeError("constexpr NaN cannot use a prepared call")
namespace[f"value_{index}"] = arg.hex() if type(arg) is float else arg
checks.append(
f"{prefix}.hex() == value_{index}"
if type(arg) is float
else f"{prefix} == value_{index}"
)
elif type(arg) in (str, type(None), torch.dtype, torch.device):
namespace[f"value_{index}"] = arg
checks.append(f"{prefix} == value_{index}")
else:
raise TypeError(f"unsupported prepared-call argument: {type(arg)!r}")
return eval(f"lambda args: {' and '.join(checks)}", namespace)
class _PreparedCall:
"""Monomorphic eager-call guard for a compiled ``BoundKernel``.
``Kernel._dispatch_cache`` remains the general multi-specialization cache.
This object only makes its most recently used entry cheap to revisit: it
compares argument metadata without rebuilding and hashing a nested cache
key, then calls the already-compiled host wrapper directly.
"""
__slots__ = (
"_dist_initialized",
"_extra_guards",
"_is_distributed",
"_matches_args",
"bound",
)
def __init__(
self,
bound: BoundKernel,
args: tuple[object, ...],
*,
dist_initialized: bool,
extra_guards: tuple[
tuple[Callable[[Sequence[object]], Hashable], Hashable], ...
],
is_distributed: bool,
) -> None:
self._matches_args = _make_prepared_arg_guard(bound.kernel, args)
self._dist_initialized = dist_initialized
self._is_distributed = is_distributed
self._extra_guards = extra_guards
self.bound = bound
@classmethod
def build(
cls,
kernel: Kernel,
bound: BoundKernel,
args: tuple[object, ...],
*,
dist_initialized: bool,
extra_guards: tuple[
tuple[Callable[[Sequence[object]], Hashable], Hashable], ...
],
is_distributed: bool,
) -> _PreparedCall | None:
if (
not bound.env.backend.supports_eager_prepared_call
or kernel._key_fn is not None
):
return None
# The caller already obtained a non-None ``_fast_dispatch_key``, which
# proves every argument has an exact supported type.
try:
return cls(
bound,
args,
dist_initialized=dist_initialized,
extra_guards=extra_guards,
is_distributed=is_distributed,
)
except Exception:
# Preparation is optional and must not introduce a new user-visible
# failure after a kernel has run successfully.
return None
def matches(self, kernel: Kernel, args: tuple[object, ...]) -> bool:
try:
if not self._matches_args(args):
return False
dist_initialized = dist.is_initialized()
# ``kernel_uses_symm_mem`` and declared distributed intent are both
# false before process-group initialization. If that state changes,
# reject the prepared call before rechecking it.
if dist_initialized != self._dist_initialized or (
dist_initialized
and kernel._compute_is_distributed(
args, dist_initialized=dist_initialized
)
!= self._is_distributed
):
return False
for extractor, expected in self._extra_guards:
if extractor(args) != expected:
return False
return True
except Exception:
# Guard evaluation is an optional fast path. Falling through lets
# the normal dispatch machinery preserve its own error semantics.
return False
def _current_device_index(device_type: str) -> int:
device_module = getattr(torch, device_type, None)
current_device = getattr(device_module, "current_device", None)
if callable(current_device):
return cast("int", current_device())
accelerator = getattr(torch, "accelerator", None)
current_accelerator = getattr(accelerator, "current_accelerator", None)
current_device_index = getattr(accelerator, "current_device_index", None)
if callable(current_accelerator) and callable(current_device_index):
accelerator_device = cast("torch.device", current_accelerator())
if accelerator_device.type == device_type:
return cast("int", current_device_index())
raise exc.InvalidAPIUsage(
f"autotune_multi requires a current indexed accelerator, got {device_type!r}"
)
def _canonicalize_multi_shape_device(device: torch.device) -> torch.device:
if device.type in ("cpu", "meta", "mps"):
raise exc.InvalidAPIUsage(
f"autotune_multi requires an indexed accelerator device, got {device}"
)
if device.index is not None:
return device
return torch.device(device.type, _current_device_index(device.type))
def _has_unspecialized_numeric_value(value: object) -> bool:
"""Return whether normal specialization records only a numeric value's type."""
if isinstance(value, ConstExpr):
return False
if type(value) in (bool, int, float):
return True
if dataclasses.is_dataclass(value) and not isinstance(value, type):
return any(
_has_unspecialized_numeric_value(getattr(value, field.name))
for field in dataclasses.fields(value)
)
if isinstance(value, dict):
return any(_has_unspecialized_numeric_value(item) for item in value.values())
if isinstance(value, (list, tuple)):
return any(_has_unspecialized_numeric_value(item) for item in value)
return False
def _resolve_index_dtype(
settings: Settings,
args: Sequence[object] | tuple[object, ...],
) -> torch.dtype:
if (index_dtype := settings.index_dtype) is not None:
limit = torch.iinfo(index_dtype).max
else:
limit = _INT32_INDEX_LIMIT
over_limit = False
def _check(tensor: torch.Tensor) -> None:
nonlocal over_limit
if over_limit:
return
try:
over_limit = bool(tensor.numel() > limit)
except RuntimeError: # unbacked SymInt
if index_dtype is None:
over_limit = True
tree_map_only(torch.Tensor, _check, args)
# pyrefly: ignore [unbound-name]
if index_dtype is None: # Auto-select when not provided
return torch.int64 if over_limit else torch.int32
if over_limit:
# pyrefly: ignore [unbound-name]
raise exc.InputTensorNumelExceedsIndexType(index_dtype=index_dtype)
# pyrefly: ignore [unbound-name]
return index_dtype
def _device_specialization_key(
args: Sequence[object],
) -> tuple[str | None, tuple[int, int] | None]:
"""Return the recursive argument device key used by bound-kernel caching.
`_find_device` intentionally searches tensors and bare `torch.device`
objects inside supported containers to match binding device selection.
Capability splits mixed-sm cache entries. Device index remains excluded so
same-capability devices share bound kernels, matching prior behavior.
"""
try:
device = _find_device(tuple(args))
except exc.NoTensorArgs:
return None, None
return device.type, target_device_capability(device)
@dataclasses.dataclass
class OutputCodeOptions:
"""Options for :meth:`BoundKernel.to_code`.
Passing ``options=None`` (the default) keeps ``to_code``'s original behavior.
Attributes:
allow_helion_deps: When ``False``, emit a self-contained module that does
not import ``helion`` at runtime -- the dependency-free launcher is
inlined (and any in-kernel runtime helpers are embedded) so the only
deps are ``torch`` + the backend DSL.
jax_fn: Pallas only. When ``True``, emit a module whose entrypoint operates
on ``jax.Array`` inputs instead of TorchTPU tensors. Orthogonal to
``allow_helion_deps``: combine with ``allow_helion_deps=False`` for a
pure-JAX module (launch core inlined), or leave ``allow_helion_deps=True``
to import the launch core from helion.
"""
allow_helion_deps: bool = True
jax_fn: bool = False
[docs]
class Kernel(Generic[_R]):
[docs]
def __init__(
self,
fn: Callable[..., _R],
*,
configs: Sequence[ConfigLike] | None = None,
settings: Settings | None,
key: Callable[..., Hashable] | None = None,
) -> None:
"""
Initialize the Kernel object. This is typically called from the `@helion.kernel` decorator.
Args:
fn: The function to be compiled as a Helion kernel.
configs: A list of configurations to use for the kernel.
settings: The settings to be used by the Kernel. If None, a new `Settings()` instance is created.
key: Optional callable that returns an extra hashable component for specialization.
"""
super().__init__()
assert isinstance(fn, types.FunctionType)
assert_no_conflicts(fn)
self.name: str = fn.__name__
# pyrefly: ignore [read-only]
self.fn: types.FunctionType = fn
self.signature: inspect.Signature = inspect.signature(fn)
self.settings: Settings = settings or Settings()
self._key_fn: Callable[..., Hashable] | None = key
# Whether the kernel declares distributed intent via an hl.ProcessGroupName
# argument. Computed once so the per-call is_distributed check stays cheap
# on the dispatch hot path (avoids re-running inspect.signature). See #3024.
self._declares_process_group: bool = kernel_declares_process_group(fn)
self.configs: list[Config] = [
# pyrefly: ignore [bad-argument-type]
Config(**config) if isinstance(config, dict) else config
for config in configs or []
]
self._bind_lock = threading.RLock()
self._specialize_extra_lock = threading.Lock()
self._bound_kernels: dict[BoundKernelInMemoryCacheKey, BoundKernel] = {}
# Fast dispatch cache: maps a cheap, fine-grained argument key (exact
# dtype/shape/stride/device per tensor) directly to a BoundKernel,
# skipping the full specialization-key machinery on repeat calls.
self._dispatch_cache: dict[Hashable, BoundKernel] = {}
self._prepared_call: _PreparedCall | None = None
self._specialize_extra: dict[
Hashable, list[Callable[[Sequence[object]], Hashable]]
] = {}
self._specialization_aliases: dict[
tuple[Hashable, ...], _SpecializationAlias
] = {}
self._specialization_generation = 0
self._reset_generation = 0
self._has_specialization_extras = False
self._cute_grouped_static_tail_extra_descriptors: dict[
Hashable, set[Hashable]
] = {}
if any(
param.kind
in (
inspect.Parameter.VAR_POSITIONAL,
inspect.Parameter.VAR_KEYWORD,
inspect.Parameter.KEYWORD_ONLY,
)
for param in self.signature.parameters.values()
):
raise TypeError(
f"Kernel({self.name}) cannot have *args, **kwargs, or keyword-only arguments"
)
self._annotations: list[object] = []
for param in self.signature.parameters.values():
ann = param.annotation
if isinstance(ann, str) and re.search(r"constexpr", ann, re.IGNORECASE):
self._annotations.append(ConstExpr)
else:
self._annotations.append(ann)
# Cache the number of parameters to avoid accessing self.signature.parameters
# during torch.compile tracing.
self._num_params: int = len(self.signature.parameters)
# Expose function attributes for compatibility with torch.library.custom_op
# These are set as instance attributes to allow the Kernel to be used
# as if it were a regular function for introspection purposes
functools.update_wrapper(self, fn)
# Manually add function-specific attributes not copied by update_wrapper
self.__globals__ = fn.__globals__
self.__code__ = fn.__code__
self.__defaults__ = fn.__defaults__
self.__kwdefaults__ = fn.__kwdefaults__
# Opt-in: register the torch.compile(backend="tpu") capture op now, so an
# annotated/functional/benchmark-free Pallas kernel is captured with zero
# warm-up (like a hand-written custom_op). None if it must use first-call
# registration (unannotated, mutating, or autotuning among configs=[...]).
self._capture_op: Callable[..., Any] | None = None
if self.settings.backend == "pallas" and _TPU_COMPILE_CAPTURE:
from .pallas._tpu_compile_capture import register_decoration_op
self._capture_op = register_decoration_op(self)
[docs]
@functools.cache # noqa: B019
def kernel_source(self) -> str:
"""
Return the kernel's source text.
This is the stable identifier across processes/runs, suitable for
grouping telemetry rows by kernel during analysis.
"""
return inspect.getsource(self.fn)
def _get_bound_kernel_cache_key(
self, args: tuple[object, ...], signature: tuple[Hashable, ...]
) -> BoundKernelInMemoryCacheKey | None:
from ..autotuner.base_cache import BoundKernelInMemoryCacheKey
extra_results = self._stable_specialization_extra_results(args, signature)
if extra_results is None:
return None
return BoundKernelInMemoryCacheKey(signature, extra_results)
def _stable_specialization_extra_results(
self,
args: Sequence[object],
signature: tuple[Hashable, ...],
) -> tuple[Hashable, ...] | None:
# Evaluate outside the schema lock because some backend extractors may
# synchronize device data. Schema changes are finite and retry here.
while True:
with self._specialize_extra_lock:
extra_fns = self._specialize_extra.get(signature)
generation = self._specialization_generation
extra_results = (
None if extra_fns is None else tuple(fn(args) for fn in extra_fns)
)
with self._specialize_extra_lock:
if (
self._specialization_generation == generation
and self._specialize_extra.get(signature) is extra_fns
):
return extra_results
def _create_bound_kernel_cache_key(
self,
bound_kernel: BoundKernel,
args: tuple[object, ...],
signature: tuple[Hashable, ...],
*,
extra_fns: list[Callable[[Sequence[object]], Hashable]] | None = None,
) -> BoundKernelInMemoryCacheKey:
from ..autotuner.base_cache import BoundKernelInMemoryCacheKey
if extra_fns is None:
extra_fns = bound_kernel._specialize_extra()
if not bound_kernel._cache_managed:
extra_results = tuple(s(args) for s in extra_fns)
return BoundKernelInMemoryCacheKey(signature, extra_results)
# Autotune cache keys can be generated outside eager binding, so
# schema synchronization must not wait for the eager bind lock or hold
# the schema lock while backend extractors run.
while True:
with self._specialize_extra_lock:
if bound_kernel._reset_generation != self._reset_generation:
active_extra_fns = extra_fns
generation = None
else:
published_extra_fns = self._specialize_extra.get(signature)
if published_extra_fns is None:
self._specialize_extra[signature] = extra_fns
if extra_fns:
self._has_specialization_extras = True
else:
# Late specialization can extend this schema after the
# bound is constructed. The published list is the
# authoritative schema.
extra_fns = published_extra_fns
generation = self._specialization_generation
active_extra_fns = extra_fns
extra_results = tuple(s(args) for s in active_extra_fns)
cache_key = BoundKernelInMemoryCacheKey(signature, extra_results)
if generation is None:
return cache_key
with self._specialize_extra_lock:
if bound_kernel._reset_generation != self._reset_generation:
return cache_key
if (
self._specialization_generation == generation
and self._specialize_extra.get(signature) is active_extra_fns
):
return cache_key
def _extend_bound_kernel_specializations(
self,
bound_kernel: BoundKernel,
signature: tuple[Hashable, ...],
extractors: list[Callable[[Sequence[object]], Hashable]],
args: Sequence[object],
) -> bool:
if not extractors:
return False
from ..autotuner.base_cache import BoundKernelInMemoryCacheKey
with self._bind_lock:
if bound_kernel._reset_generation != self._reset_generation:
return False
full_args = tuple(args)
aliases = {
alias_signature: alias
for alias_signature, alias in self._specialization_aliases.items()
if alias.canonical_signature == signature
}
# Keep cache-key generation from snapshotting the old schema while
# this extension is being validated and published.
with self._specialize_extra_lock:
updated_extractors = [
*self._specialize_extra.get(signature, []),
*extractors,
]
with unset_fake_temporarily():
current_results = tuple(
extractor(full_args) for extractor in updated_extractors
)
updated_cache_key = BoundKernelInMemoryCacheKey(
signature, current_results
)
hash(updated_cache_key)
for alias_signature, alias in aliases.items():
alias_arg_count = len(full_args) - len(alias.trailing_defaults)
alias_args = full_args[:alias_arg_count]
alias_normalized_args = (
*alias_args,
*alias.trailing_defaults,
)
alias_results = tuple(
extractor(alias_normalized_args)
for extractor in updated_extractors
)
hash(
BoundKernelInMemoryCacheKey(
alias_signature,
(alias_results,),
)
)
# Publish the extended specialization only after every extractor
# succeeds. Otherwise a failed extension could leave a prepared
# call guarding an already-mutated extractor list.
self._has_specialization_extras = True
self._specialize_extra[signature] = updated_extractors
for alias_signature, alias in aliases.items():
self._specialize_extra[alias_signature] = [alias]
self._specialization_generation += 1
affected_signatures = {signature, *aliases}
stale_keys = [
key
for key in self._bound_kernels
if key.specialization_key in affected_signatures
]
for stale_key in stale_keys:
self._bound_kernels.pop(stale_key)
for fast_key, cached_bound in list(self._dispatch_cache.items()):
if cached_bound._base_spec_key == signature:
self._dispatch_cache.pop(fast_key)
self._prepared_call = None
self._bound_kernels[updated_cache_key] = bound_kernel
return True
def _compute_is_distributed(
self,
args: Sequence[object],
*,
dist_initialized: bool | None = None,
) -> bool:
"""Whether this call should compile as a distributed kernel.
True when the arguments carry symmetric-memory tensors, or the author
declared distributed intent (``settings.distributed`` or an
``hl.ProcessGroupName`` argument) inside an initialized process group.
This is folded into the specialization key so a symmetric-memory call
and an identically-shaped ordinary call never share a compiled kernel.
See GitHub issue #3024.
"""
if dist_initialized is None:
dist_initialized = dist.is_initialized()
return kernel_uses_symm_mem(tuple(args), dist_initialized=dist_initialized) or (
dist_initialized
and (self.settings.distributed or self._declares_process_group)
)
def _fast_dispatch_key(
self,
args: tuple[object, ...],
*,
is_distributed: bool | None = None,
signature: tuple[Hashable, ...] | None = None,
_extra_guards: list[tuple[Callable[[Sequence[object]], Hashable], Hashable]]
| None = None,
) -> tuple[Hashable, ...] | None:
"""Build the exact eager dispatch key, optionally collecting its guards."""
key: list[Hashable] = []
has_tensor = False
for a in args:
t = type(a)
if t is torch.Tensor or t is torch.nn.Parameter:
tensor = cast("torch.Tensor", a)
has_tensor = True
si = getattr(tensor, "_dynamo_static_indices", None)
key.append(
(
tensor.dtype,
tensor.shape,
tensor.stride(),
tensor.device,
None if si is None else frozenset(si),
)
)
elif t is int or t is float or t is bool or t is str:
key.append((t, a))
elif a is None:
key.append(None)
elif t is torch.dtype or t is torch.device:
key.append(a)
else:
return None
if not has_tensor:
return None
if is_distributed is None:
is_distributed = self._compute_is_distributed(args)
key.append(is_distributed)
if signature is None and self._has_specialization_extras:
signature = self._base_specialization_key(
args, is_distributed=is_distributed
)
if self._key_fn is not None:
key.append(self._key_fn(*args) if signature is None else signature[-1])
if signature is not None:
extra_fns = self._specialize_extra.get(signature)
if extra_fns:
extra_results: list[Hashable] = []
for extractor in extra_fns:
result = extractor(args)
extra_results.append(result)
if _extra_guards is not None:
_extra_guards.append((extractor, result))
key.append(tuple(extra_results))
return tuple(key)
def _fast_dispatch_key_and_guards(
self,
args: tuple[object, ...],
*,
is_distributed: bool | None = None,
signature: tuple[Hashable, ...] | None = None,
) -> _FastDispatchEntry | None:
"""
Build a cheap dispatch key for the fast-path cache in ``__call__``.
The key records exact per-argument metadata (dtype/shape/stride/device
for tensors, type and value for scalars), which is strictly finer than
the full specialization key: any two argument lists that produce
different full keys also produce different fast keys. That makes it
safe to map a fast key directly to the BoundKernel that a full
``bind()`` resolved for the same arguments.
If a base signature has extra specialization extractors, their results
are appended to preserve the same no-collision invariant for
value-based specializations.
Returns None when an argument type is not handled (tensor subclasses,
containers, ...), or when there is no tensor argument to pin down the
device; callers must then take the regular ``bind()`` path.
"""
specialization_generation = self._specialization_generation
extra_guards: list[tuple[Callable[[Sequence[object]], Hashable], Hashable]] = []
key = self._fast_dispatch_key(
args,
is_distributed=is_distributed,
signature=signature,
_extra_guards=extra_guards,
)
if key is None:
return None
return _FastDispatchEntry(
key,
tuple(extra_guards),
specialization_generation,
)
def _prepare_dispatch_entry(
self,
args: tuple[object, ...],
bound: BoundKernel[_R],
fast_entry: _FastDispatchEntry,
) -> tuple[_PreparedCall | None, bool] | None:
"""Validate and construct eager fast paths from one runtime snapshot."""
try:
if fast_entry.specialization_generation != self._specialization_generation:
return None
if bound._reset_generation != self._reset_generation:
return None
dist_initialized = dist.is_initialized()
is_distributed = self._compute_is_distributed(
args, dist_initialized=dist_initialized
)
if fast_entry.key[len(args)] != is_distributed:
return None
if self._key_fn is None:
raw_signature = self._base_specialization_key(
args, is_distributed=is_distributed
)
else:
# Reuse the custom key captured by the dispatch lookup; user
# key functions must execute exactly once per launch.
raw_signature = self._base_specialization_key(
args,
is_distributed=is_distributed,
custom_key=fast_entry.key[len(args) + 1],
)
alias = self._specialization_aliases.get(raw_signature)
signature = raw_signature if alias is None else alias.canonical_signature
if signature != bound._base_spec_key:
return None
if self._key_fn is None:
prepared = _PreparedCall.build(
self,
bound,
args,
dist_initialized=dist_initialized,
extra_guards=fast_entry.extra_guards,
is_distributed=is_distributed,
)
keyed_direct_dispatch = False
else:
prepared = None
keyed_direct_dispatch = (
not self.settings.distributed
and not self._declares_process_group
and not bound._env._is_distributed
)
# Process-group initialization is external to ``_bind_lock``. Do
# not publish a key assembled across a state transition.
if dist.is_initialized() != dist_initialized:
return None
return prepared, keyed_direct_dispatch
except Exception:
# Fast-path publication is optional and must not add a failure
# after the compiled kernel has already run successfully.
return None
[docs]
def bind(self, args: tuple[object, ...]) -> BoundKernel[_R]:
"""
Bind the given arguments to the Kernel and return a BoundKernel object.
Args:
args: The arguments to bind to the Kernel.
Returns:
BoundKernel: A BoundKernel object with the given arguments bound.
"""
# Dynamo executes bind while capturing the call but cannot trace an RLock.
# Capture only needs an independent host function and compile environment.
if torch.compiler.is_compiling():
return self._bind_isolated(args)
with self._bind_lock:
return self._bind(args)
def _validate_bind_args(
self, args: tuple[object, ...] | list[object]
) -> tuple[object, ...]:
if not isinstance(args, tuple):
assert isinstance(args, list), "args must be a tuple or list"
args = tuple(args)
if len(args) > self._num_params:
raise TypeError(
f"Too many arguments passed to the kernel, expected: {self._num_params} got: {len(args)}."
)
return args
def _bind_isolated(self, args: tuple[object, ...]) -> BoundKernel[_R]:
"""Construct a canonical bound without reading or publishing shared caches."""
args = self._validate_bind_args(args)
args = self.normalize_args(*args)
dist_initialized = dist.is_initialized()
is_distributed = self._compute_is_distributed(
args, dist_initialized=dist_initialized
)
signature = self._base_specialization_key(args, is_distributed=is_distributed)
return BoundKernel(
self,
args,
base_spec_key=signature,
is_distributed=is_distributed,
cache_managed=False,
)
def _bind(self, args: tuple[object, ...]) -> BoundKernel[_R]:
with measure("Kernel.bind"):
args = self._validate_bind_args(args)
dist_initialized = dist.is_initialized()
is_distributed = self._compute_is_distributed(
args, dist_initialized=dist_initialized
)
signature = self._base_specialization_key(
args, is_distributed=is_distributed
)
cache_key = self._get_bound_kernel_cache_key(args, signature)
bound_kernel = (
None if cache_key is None else self._bound_kernels.get(cache_key, None)
)
if bound_kernel is None:
normalized_args: tuple[object, ...] = self.normalize_args(*args)
extra_fns: list[Callable[[Sequence[object]], Hashable]] | None = None
if len(normalized_args) != len(args):
# we had default args that needed to be applied
bound_kernel = self._bind(normalized_args)
canonical_signature = bound_kernel._base_spec_key
alias = self._specialization_aliases.get(signature)
if alias is None:
trailing_defaults = normalized_args[len(args) :]
alias = _SpecializationAlias(
self._specialize_extra,
canonical_signature,
trailing_defaults,
)
self._specialization_aliases[signature] = alias
extra_fns = (
[alias]
if self._specialize_extra.get(canonical_signature)
else []
)
else:
bound_kernel = BoundKernel(
self,
args,
base_spec_key=signature,
is_distributed=is_distributed,
)
if cache_key is None:
cache_key = self._create_bound_kernel_cache_key(
bound_kernel,
args,
signature,
extra_fns=extra_fns,
)
self._bound_kernels[cache_key] = bound_kernel
return bound_kernel
def _base_specialization_key(
self,
args: Sequence[object],
*,
is_distributed: bool | None = None,
custom_key: object = _CUSTOM_KEY_UNSET,
) -> tuple[Hashable, ...]:
"""
Generate the base specialization key from input argument metadata only,
using the per-type extractor functions defined in `_specialization_extractors`,
without any extras discovered during compilation. Used internally for
_specialize_extra lookups.
"""
result: list[Hashable] = []
assert len(args) <= len(self._annotations)
for value, annotation in zip(args, self._annotations, strict=False):
if isinstance(value, ConstExpr):
result.append(value.value)
elif annotation is ConstExpr:
result.append(value)
else:
result.append(self._specialization_key(value))
device_type, device_capability = _device_specialization_key(args)
if is_distributed is None:
is_distributed = self._compute_is_distributed(args)
if self._key_fn is not None:
if custom_key is _CUSTOM_KEY_UNSET:
custom_key = self._key_fn(*args)
return (
*result,
device_type,
device_capability,
is_distributed,
cast("Hashable", custom_key),
)
return (*result, device_type, device_capability, is_distributed)
[docs]
def specialization_key(self, args: Sequence[object]) -> tuple[Hashable, ...]:
"""
Generate the full specialization key for the given arguments, including
any additional specialization constraints discovered during compilation
(e.g. from hl.specialize() calls).
Before the first compilation, these extras are not yet known and the
key may be incomplete.
Args:
args: The arguments to generate a specialization key for.
Returns:
Hashable: A hashable key representing the specialization of the arguments.
"""
base = self._base_specialization_key(args)
extra_results = self._stable_specialization_extra_results(args, base)
return base if extra_results is None else (*base, *extra_results)
def _specialization_key(self, obj: object) -> Hashable:
"""
Helper used to generate a specialization key for the given object.
This method determines a unique key for the object based on its type
and the corresponding extractor function defined in `_specialization_extractors`.
Args:
obj: The argument to generate a specialization key for.
Returns:
Hashable: A hashable key representing the specialization of the object.
"""
extractor = _specialization_extractors.get(type(obj))
if extractor is None:
if isinstance(obj, torch.fx.GraphModule):
# GraphModule subclasses need special handling
extractor = _specialization_extractors[torch.fx.GraphModule]
elif isinstance(obj, torch.Tensor):
# torch.Tensor subclasses (e.g. the JAX-export adapter)
# share the standard tensor specialization key. Use the
# SymInt-safe extractor: unlike exact ``torch.Tensor``,
# subclasses may carry symbolic sizes/strides.
extractor = _specialization_extractors["tensor_subclass"]
elif isinstance(obj, tuple) and hasattr(obj, "_fields"):
# this is a namedtuple
extractor = _specialization_extractors["namedtuple"]
elif dataclasses.is_dataclass(obj):
extractor = _specialization_extractors["dataclass"]
else:
raise TypeError(f"unsupported argument type: {type(obj).__name__}")
return extractor(self, obj)
[docs]
def normalize_args(self, *args: object, **kwargs: object) -> tuple[object, ...]:
"""
Normalize the given arguments and keyword arguments according to the function signature.
Args:
args: The positional arguments to normalize.
kwargs: The keyword arguments to normalize.
Returns:
tuple[object, ...]: A tuple of normalized positional arguments.
"""
bound_args = self.signature.bind(*args, **kwargs)
bound_args.apply_defaults()
return tuple(bound_args.args)
[docs]
def autotune(
self,
args: Sequence[object],
*,
force: bool = True,
**options: object,
) -> Config:
"""
Perform autotuning to find the optimal configuration for the kernel. This uses the
default setting, you can call helion.autotune.* directly for more customization.
If config= or configs= is provided to helion.kernel(), the search will be restricted to
the provided configs. Use force=True to ignore the provided configs.
Mutates (the bound version of) self so that `__call__` will run the best config found.
Args:
args: Example arguments used for benchmarking during autotuning.
force: If True, force full autotuning even if a config is provided.
options: Additional keyword options forwarded to the autotuner.
Returns:
Config: The best configuration found during autotuning.
"""
args = self.normalize_args(*args)
return self.bind(args).autotune(args, force=force, **options)
[docs]
def autotune_multi(
self,
arg_sets: Sequence[Sequence[object]],
*,
aggregation: Literal["geomean", "max"] = "geomean",
relative_to: Literal["default", "baseline"] | None = None,
cache_tag: str | None = None,
force: bool = True,
**options: object,
) -> Config:
"""Find one config using an objective measured across several inputs.
The first argument set anchors config generation. Each candidate is measured
on every set, then reduced with a geometric mean or maximum. ``relative_to``
optionally optimizes per-shape latency relative to each shape's default config
or custom baseline. Only the supplied bound specializations are configured.
Every argument set must bind normally to the same exact, current accelerator
device. Distributed processes are not supported. A non-empty ``cache_tag`` is
required for custom callbacks, dynamic-shape tuning, and runtime numeric
arguments; callers own tag invalidation in those cases.
Args:
arg_sets: Non-empty sequence of representative kernel argument sequences.
aggregation: Joint objective, either ``"geomean"`` or ``"max"``.
relative_to: Optional ``"default"`` or ``"baseline"`` normalization.
cache_tag: User-managed cache discriminator for dynamic shapes, runtime
numeric arguments, or callbacks.
force: If true, ignore pinned configs and cache reads during the search.
options: Additional options forwarded to the registered autotuner.
Returns:
The config selected for all supplied specializations.
"""
from ..autotuner.benchmark_provider import LocalBenchmarkProvider
from ..autotuner.benchmark_provider import _has_valid_multi_shape_measurement
from ..autotuner.benchmark_provider import _materialize_multi_shape_config
from ..autotuner.benchmark_provider import _MultiShapeAutotuneArgs
from .settings import default_autotuner_fn
if aggregation not in ("geomean", "max"):
raise exc.InvalidAPIUsage(
"autotune_multi aggregation must be 'geomean' or 'max'"
)
if relative_to not in (None, "default", "baseline"):
raise exc.InvalidAPIUsage(
"autotune_multi relative_to must be None, 'default', or 'baseline'"
)
if cache_tag is not None and (not isinstance(cache_tag, str) or not cache_tag):
raise exc.InvalidAPIUsage(
"autotune_multi cache_tag must be a non-empty string"
)
if (
not isinstance(arg_sets, Sequence)
or isinstance(arg_sets, (str, bytes))
or not arg_sets
):
raise exc.InvalidAPIUsage(
"autotune_multi arg_sets must be a non-empty sequence"
)
if dist.is_initialized():
raise exc.InvalidAPIUsage(
"autotune_multi does not support an initialized distributed process group"
)
if self.settings.backend not in {"triton", "tileir", "cute", "pallas"}:
raise exc.InvalidAPIUsage(
f"autotune_multi does not support backend {self.settings.backend!r}"
)
if self.settings.autotuner_fn is not default_autotuner_fn:
raise exc.InvalidAPIUsage(
"autotune_multi requires the default registered autotuner"
)
if self.settings.autotune_cache == "AOTAutotuneCache":
raise exc.InvalidAPIUsage(
"autotune_multi does not support AOTAutotuneCache"
)
if self.settings.autotune_cache not in {
"LocalAutotuneCache",
"StrictLocalAutotuneCache",
"RemoteAutotuneCache",
"StrictRemoteAutotuneCache",
}:
raise exc.InvalidAPIUsage(
"autotune_multi requires a built-in local or remote best-config cache"
)
if "benchmark_provider_cls" in options:
if options.pop("benchmark_provider_cls") is not LocalBenchmarkProvider:
raise exc.InvalidAPIUsage(
"autotune_multi does not support a custom benchmark provider"
)
if relative_to == "baseline" and self.settings.autotune_baseline_fn is None:
raise exc.InvalidAPIUsage(
"autotune_multi relative_to='baseline' requires autotune_baseline_fn"
)
if self.settings.autotune_benchmark_fn is not None:
raise exc.InvalidAPIUsage(
"autotune_multi does not support autotune_benchmark_fn"
)
custom_callbacks = (
self.settings.autotune_baseline_fn,
self.settings.autotune_baseline_accuracy_check_fn,
self.settings.autotune_config_filter,
)
if cache_tag is None and any(fn is not None for fn in custom_callbacks):
raise exc.InvalidAPIUsage(
"autotune_multi requires cache_tag when a custom baseline, "
"accuracy check, or config filter is configured"
)
if cache_tag is None and not self.settings.static_shapes:
raise exc.InvalidAPIUsage(
"autotune_multi requires cache_tag when static_shapes=False"
)
normalized_arg_sets: list[tuple[object, ...]] = []
for case_index, arg_set in enumerate(arg_sets):
if not isinstance(arg_set, Sequence) or isinstance(arg_set, (str, bytes)):
raise exc.InvalidAPIUsage(
f"autotune_multi arg_sets[{case_index}] must be a sequence"
)
try:
normalized_arg_sets.append(self.normalize_args(*arg_set))
except TypeError as error:
raise exc.InvalidAPIUsage(
f"autotune_multi arg_sets[{case_index}] does not match the "
f"kernel signature: {error}"
) from error
if cache_tag is None:
for case_index, normalized_args in enumerate(normalized_arg_sets):
if any(
annotation is not ConstExpr
and _has_unspecialized_numeric_value(value)
for value, annotation in zip(
normalized_args, self._annotations, strict=True
)
):
raise exc.InvalidAPIUsage(
"autotune_multi requires cache_tag when an argument set "
"contains a runtime numeric value; "
f"arg_sets[{case_index}] does"
)
cases: list[tuple[BoundKernel[_R], tuple[object, ...]]] = []
case_keys: list[BoundKernelInMemoryCacheKey] = []
for case_index, normalized_args in enumerate(normalized_arg_sets):
try:
bound_kernel = self.bind(normalized_args)
except exc.NoTensorArgs as error:
raise exc.InvalidAPIUsage(
"autotune_multi requires each argument set to have a device "
f"discoverable by normal kernel binding; arg_sets[{case_index}] did not"
) from error
signature = self._base_specialization_key(normalized_args)
case_key = self._get_bound_kernel_cache_key(normalized_args, signature)
assert case_key is not None
cases.append((bound_kernel, normalized_args))
case_keys.append(case_key)
anchor = cases[0][0]
canonical_device = _canonicalize_multi_shape_device(anchor.env.device)
current_index = _current_device_index(canonical_device.type)
if canonical_device.index != current_index:
raise exc.InvalidAPIUsage(
"autotune_multi requires the indexed accelerator to be current: "
f"got {canonical_device}, current index is {current_index}"
)
unique_bound_kernels: list[BoundKernel[_R]] = []
seen_bound_kernel_ids: set[int] = set()
for bound_kernel, _ in cases:
if id(bound_kernel) not in seen_bound_kernel_ids:
seen_bound_kernel_ids.add(id(bound_kernel))
unique_bound_kernels.append(bound_kernel)
anchor_backend = anchor.env.backend.name
anchor_capability = target_device_capability(anchor.env.device)
advanced_controls_files = self.settings.autotune_search_acf or None
anchor_fingerprint = anchor.config_spec.structural_fingerprint(
advanced_controls_files=advanced_controls_files
)
for case_index, (bound_kernel, _) in enumerate(cases):
if bound_kernel.env.process_group_name is not None:
raise exc.InvalidAPIUsage(
"autotune_multi does not support a bound kernel with a process group"
)
bound_device = _canonicalize_multi_shape_device(bound_kernel.env.device)
if bound_device != canonical_device:
raise exc.InvalidAPIUsage(
"autotune_multi bound a different device for "
f"arg_sets[{case_index}]: {bound_device} != {canonical_device}"
)
if bound_kernel.env.backend.name != anchor_backend:
raise exc.InvalidAPIUsage(
"autotune_multi requires every case to use the same backend"
)
if target_device_capability(bound_kernel.env.device) != anchor_capability:
raise exc.InvalidAPIUsage(
"autotune_multi requires every case to have the same device capability"
)
fingerprint = bound_kernel.config_spec.structural_fingerprint(
advanced_controls_files=advanced_controls_files
)
if fingerprint != anchor_fingerprint:
raise exc.InvalidAPIUsage(
"autotune_multi requires structurally compatible ConfigSpec "
f"instances; arg_sets[{case_index}] is incompatible with the anchor"
)
multi_args = _MultiShapeAutotuneArgs(
cases=tuple(cases),
aggregation=aggregation,
relative_to=relative_to,
cache_tag=cache_tag,
workload_key=(
"multi_shape:v1",
aggregation,
relative_to,
cache_tag,
tuple(
(case_key.specialization_key, case_key.extra_results)
for case_key in case_keys
),
),
reference_latencies=None,
)
ephemeral = anchor.env.backend.make_ephemeral_cache()
ctx = ephemeral if ephemeral is not None else contextlib.nullcontext()
with ctx:
config = anchor.env.backend.autotune(
anchor,
cast("Sequence[object]", multi_args),
force=force,
**options,
)
if multi_args.search_started and not multi_args.found_valid_config:
raise exc.NoConfigFound
if multi_args.search_started and not _has_valid_multi_shape_measurement(
multi_args,
anchor.config_spec,
config,
):
raise exc.NoConfigFound
winner = _materialize_multi_shape_config(anchor.config_spec, config)
if ephemeral is not None:
for bound_kernel in unique_bound_kernels:
bound_kernel.env.backend.finalize_ephemeral_cache(bound_kernel, winner)
for bound_kernel in unique_bound_kernels:
bound_kernel.compile_config(winner)
for bound_kernel in unique_bound_kernels:
bound_kernel.set_config(winner)
return winner
[docs]
def __call__(self, *args: object, **kwargs: object) -> _R:
"""
Call the Kernel with the given arguments and keyword arguments.
Args:
args: The positional arguments to pass to the Kernel.
kwargs: The keyword arguments to pass to the Kernel.
Returns:
_R: The result of the Kernel function call.
"""
if kwargs:
args = self.normalize_args(*args, **kwargs)
is_compiling = torch.compiler.is_compiling()
if (
not is_compiling
and (prepared := self._prepared_call) is not None
and prepared.matches(self, args)
and prepared.bound._run is not None
):
return prepared.bound._run(*args)
if self._dispatch_cache:
# Fast path: repeat call with argument metadata seen before. The
# cache is only populated by calls that already took the slow
# path below, so hitting it cannot skip autotuning/compilation.
# (The cache stays empty under TPU compile capture, so this
# cannot bypass auto_capture_call below.)
fast_entry: _FastDispatchEntry | None = None
if (
self._key_fn is not None
and not self.settings.distributed
and not self._declares_process_group
):
# Keyed kernels cannot use a prepared call. Build their legacy
# exact key without allocating guard metadata on every launch.
specialization_generation = self._specialization_generation
dist_initialized = dist.is_initialized()
is_distributed = self._compute_is_distributed(
args, dist_initialized=dist_initialized
)
fast_key = self._fast_dispatch_key(args, is_distributed=is_distributed)
if fast_key is not None:
bound = self._dispatch_cache.get(fast_key)
if bound is not None and bound._run is not None:
run = bound._run
if is_compiling:
return run(*args)
if (
bound._dispatch_generation == specialization_generation
and bound._run is run
and self._specialization_generation
== specialization_generation
and dist.is_initialized() == dist_initialized
):
return run(*args)
fast_entry = _FastDispatchEntry(
fast_key,
(),
specialization_generation,
)
if fast_entry is None:
fast_entry = self._fast_dispatch_key_and_guards(args)
if fast_entry is not None:
fast_key = fast_entry.key
bound = self._dispatch_cache.get(fast_key)
if bound is not None and bound._run is not None:
run = bound._run
if is_compiling:
return run(*args)
# A compilation can discover a late specialization while a
# different thread is reading this cache. Revalidate fast-
# path state while the mapping is still current; the same
# lock protects specialization-driven invalidation.
with self._bind_lock:
if (
self._dispatch_cache.get(fast_key) is bound
and bound._run is run
):
entry = self._prepare_dispatch_entry(
args,
bound,
fast_entry,
)
if entry is not None:
self._prepared_call = entry[0]
if entry[1]:
bound._dispatch_generation = (
fast_entry.specialization_generation
)
else:
run = None
else:
run = None
if run is not None:
return run(*args)
if self.settings.backend == "pallas" and _TPU_COMPILE_CAPTURE:
# Local import: _tpu_compile_capture pulls in the dynamo HOP machinery,
# not ready when kernel.py first loads during ``import helion``.
from .pallas._tpu_compile_capture import RUN_NORMAL
from .pallas._tpu_compile_capture import auto_capture_call
result = auto_capture_call(self, args)
if result is not RUN_NORMAL:
return cast("_R", result)
return self.bind(args)(*args)
bound = self.bind(args)
result = bound(*args)
if is_compiling:
return result
# Avoid a second bind for signatures the fast paths cannot support.
# This cheap key check also captures a custom key exactly once.
fast_entry = self._fast_dispatch_key_and_guards(args)
if fast_entry is not None:
fast_key = fast_entry.key
# Resolve any late specialization and publish both fast paths as
# one atomic update with respect to discovery and reset.
with self._bind_lock:
if self._has_specialization_extras:
bound = self._bind(args)
if bound._run is not None:
entry = self._prepare_dispatch_entry(
args,
bound,
fast_entry,
)
if entry is not None:
self._dispatch_cache[fast_key] = bound
self._prepared_call = entry[0]
if entry[1]:
bound._dispatch_generation = (
fast_entry.specialization_generation
)
return result
[docs]
def reset(self) -> None:
"""
Clears the cache of bound kernels, meaning subsequent calls will
recompile and re-autotune.
"""
with self._bind_lock:
self._bound_kernels.clear()
self._dispatch_cache.clear()
self._prepared_call = None
# Specialization extractors are discovered by tracing the host
# function and can change after an explicit reset. Keeping the old
# schema could hide newly discovered hl.specialize() calls.
with self._specialize_extra_lock:
# Replace rather than clear: in-flight omitted-default aliases
# retain the old schema mapping while new work starts fresh.
self._specialize_extra = {}
self._specialization_aliases = {}
self._cute_grouped_static_tail_extra_descriptors = {}
self._has_specialization_extras = False
self._specialization_generation += 1
self._reset_generation += 1
@property
def jax_fn(self) -> Callable[..., Any]:
"""A pure-JAX callable view of this Helion kernel.
Pallas-backend kernels can be called directly with JAX arrays
or tracers (i.e. inside ``jax.jit``) by going through this
property. The kernel's compile/specialize path runs the first
time the callable is invoked; subsequent calls reuse the
cached compilation just like ``__call__``.
This only supports kernels compiled for the Pallas backend.
"""
from .pallas.jax_export import make_jax_fn
cached = getattr(self, "_jax_fn_callable", None)
if cached is not None:
return cast("Callable[..., Any]", cached)
rv = make_jax_fn(self)
# pyrefly: ignore [unsupported-attribute-set]
self._jax_fn_callable = rv
return rv
class BoundKernel(_AutotunableKernel, Generic[_R]):
def __init__(
self,
kernel: Kernel[_R],
args: tuple[object, ...],
*,
base_spec_key: tuple[Hashable, ...] | None = None,
is_distributed: bool | None = None,
cache_managed: bool = True,
) -> None:
"""
Initialize a BoundKernel object.
This constructor sets up the environment, compiles the kernel function, and prepares
the arguments for execution.
Args:
kernel: The Kernel object to bind.
args: A tuple of arguments to bind to the kernel.
cache_managed: Whether this bound participates in the kernel's shared
specialization and bound caches.
"""
super().__init__()
self.kernel = kernel
self._reset_generation = kernel._reset_generation
# Extending this bound's schema evicts all of its dispatch mappings.
self._dispatch_generation: int | None = None
self._cache_managed = cache_managed
if is_distributed is None:
dist_initialized = dist.is_initialized()
is_distributed = kernel._compute_is_distributed(
args, dist_initialized=dist_initialized
)
# Base specialization key from the REAL args (the same key bind() uses to
# distinguish BoundKernels). Reused in compile_config as the PyCodeCache
# ``extra`` discriminator for backends that cache shape-specific module state
# (Backend.requires_shape_specialized_module). Must be computed on the real
# args: fake_args turn scalar int/float/bool into SymInt/SymFloat/SymBool,
# which the specialization extractors reject.
self._base_spec_key = (
kernel._base_specialization_key(args, is_distributed=is_distributed)
if base_spec_key is None
else base_spec_key
)
self._run: Callable[..., _R] | None = None
self._config: Config | None = None
self._compile_cache: dict[Config, CompiledConfig] = {}
self._cache_path_map: dict[Config, str | None] = {}
self._host_semantic_fingerprints: dict[
tuple[_HostSemanticInputNormalization, tuple[object, ...]], str
] = {}
# Direct to_code() has no call arguments, so keep this bound kernel's
# construction-time tensor values as its stable weak fallback.
self._runtime_tensor_refs_by_name = {
name: weakref.ref(value)
for name, value in zip(self.kernel.signature.parameters, args, strict=False)
if isinstance(value, torch.Tensor)
}
self._first_compile_lock = threading.RLock()
self._backward_compiled: (
tuple[Kernel[object], str, BoundKernel[object]] | None
) = None
# Distributed detection lives on Kernel so the same value feeds both the
# specialization cache key and CompileEnvironment gating; that keeps a
# symmetric-memory call from ever aliasing a non-distributed compiled
# kernel of the same shape. See GitHub issue #3024.
self._env = CompileEnvironment(
_find_device(args),
self.kernel.settings,
index_dtype=_resolve_index_dtype(self.kernel.settings, args),
is_distributed=is_distributed,
)
if is_ref_mode_enabled(self.kernel.settings):
self.fake_args = [] # type: ignore[assignment]
self.host_function = None # type: ignore[assignment]
return
with self.env:
self._env.process_group_name = _find_process_group_name(
kernel.fn, args, is_distributed
)
assert len(args) == len(self.kernel.signature.parameters)
self.fake_args: list[object] = []
constexpr_args = {}
for name, arg, annotation in zip(
self.kernel.signature.parameters,
args,
self.kernel._annotations,
strict=False,
):
if isinstance(arg, ConstExpr):
assert not isinstance(arg.value, torch.Tensor), (
"ConstExpr cannot be a tensor"
)
self.fake_args.append(arg.value)
constexpr_args[name] = arg.value
elif annotation is ConstExpr:
assert not isinstance(arg, torch.Tensor), (
"ConstExpr cannot be a tensor"
)
self.fake_args.append(arg)
constexpr_args[name] = arg
else:
self.fake_args.append(self.env.to_fake(arg, ArgumentOrigin(name)))
self._apply_mark_static(args)
with (
_maybe_skip_dtype_check_in_meta_registrations(),
patch_inductor_lowerings(),
measure("BoundKernel.create_host_function"),
):
try:
compiler = KernelCompiler(self.env)
self.host_function: HostFunction = compiler.compile(
self.kernel.fn,
self.fake_args,
constexpr_args,
)
except Exception:
config = self.env.config_spec.default_config()
self.maybe_log_repro(log.warning, args, config=config)
raise
self.env.restrict_pid_types_for_persistent(args)
self.env.config_spec.configure_epilogue_subtile_autotune(args)
self.env.config_spec.compiler_seed_configs = compiler_seed_configs(
self.env,
self.host_function.device_ir,
)
# Post-compile FX-graph scan to detect kernels
# whose tcgen05 matmul is followed by an
# aux-fused store
# (``out[tile] = (acc + residual[tile]).to(...)``
# and variants — see
# ``host_function_has_tcgen05_aux_kernel_pattern``
# for the accepted shapes). When detected, the
# autotune surface widens to admit
# ``tcgen05_strategy=ROLE_LOCAL_WITH_SCHEDULER``
# + ``tcgen05_warp_spec_c_input_warps=1`` so the
# productive C-input warp lift is reachable from
# the normal autotune path. For pure-matmul
# kernels the detector returns False and the
# autotune surface keeps the narrow
# ``MONOLITHIC + c_input_warps=0`` shape so
# autotune cannot sample the strictly-worse
# inert C-input warp configuration.
# The exact-shape detector is narrower: it gates the
# ``tcgen05_aux_load_mode=tma`` seed/search axis.
from .._compiler.cute.aux_tensor import (
host_function_has_tcgen05_aux_kernel_pattern,
)
from .._compiler.cute.aux_tensor import (
host_function_has_tcgen05_exact_shape_aux_kernel_pattern,
)
from .._compiler.cute.aux_tensor import (
host_function_matmul_has_non_tcgen05_operand,
)
self.env.config_spec.cute_tcgen05_aux_kernel_detected = (
host_function_has_tcgen05_aux_kernel_pattern(self.host_function)
)
self.env.config_spec.cute_tcgen05_exact_shape_aux_kernel_detected = (
host_function_has_tcgen05_exact_shape_aux_kernel_pattern(
self.host_function
)
)
self.env.config_spec.cute_tcgen05_matmul_has_non_tcgen05_operand = (
host_function_matmul_has_non_tcgen05_operand(self.host_function)
)
if not self.env.settings.disable_autotuner_heuristics:
for seed_config in self.env.config_spec.autotune_seed_configs():
if (
seed_config
not in self.env.config_spec.compiler_seed_configs
):
self.env.config_spec.compiler_seed_configs.append(
seed_config
)
def _apply_mark_static(self, args: tuple[object, ...]) -> None:
"""
Apply torch._dynamo.mark_static() markings from input tensors.
This reads _dynamo_static_indices from each tensor argument and marks
the corresponding dimensions as specialized (constant) in the kernel.
"""
for arg, fake_arg in zip(args, self.fake_args, strict=True):
if isinstance(arg, torch.Tensor) and isinstance(fake_arg, torch.Tensor):
for dim in getattr(arg, "_dynamo_static_indices", ()):
size = fake_arg.size(dim)
if isinstance(size, torch.SymInt):
self.env.specialized_vars.update(_symint_free_symbols(size))
@property
def env(self) -> CompileEnvironment: # pyrefly: ignore[bad-override]
return self._env
@property
def settings(self) -> Settings:
"""
Retrieve the settings associated with the kernel.
Returns:
Settings: The settings of the kernel.
"""
return self.kernel.settings
@property
def config_spec(self) -> ConfigSpec:
"""
Retrieve the configuration specification for the kernel.
Returns:
ConfigSpec: The configuration specification.
"""
return self.env.config_spec
@property
def configs(self) -> list[Config]:
"""Return the kernel's configured configs (alias for `self.kernel.configs`)."""
return self.kernel.configs
def _normalize_config(self, config: ConfigLike) -> Config:
if isinstance(config, Config):
return config
# pyrefly: ignore [bad-argument-type]
return Config(**config)
def _normalized_config_copy(self, config: ConfigLike) -> Config:
normalized = self._normalize_config(config)
normalized = Config(**normalized.config) # pyrefly: ignore[bad-argument-type]
self.env.config_spec.normalize(normalized)
return normalized
def format_kernel_decorator(self, config: Config, settings: Settings) -> str:
"""Return the @helion.kernel decorator snippet capturing configs and settings that influence Triton code generation."""
parts = [
f"config={config.__repr__()}",
f"static_shapes={settings.static_shapes}",
]
if settings.index_dtype is not None:
parts.append(f"index_dtype={settings.index_dtype}")
return f"@helion.kernel({', '.join(parts)})"
def to_code(
self,
config: ConfigLike | None = None,
*,
options: OutputCodeOptions | None = None,
emit_repro_caller: bool = False,
output_origin_lines: bool | None = None,
) -> str:
"""
Generate backend-specific code for the kernel based on the given configuration.
Args:
config: The configuration to use for code generation.
options: Optional :class:`~helion.runtime.precompile.OutputCodeOptions`.
With ``allow_helion_deps=False`` the returned module is
self-contained (no ``helion`` import at runtime); ``jax_fn=True``
(Pallas only) emits a pure-JAX module operating on ``jax.Array``s.
``None`` keeps the default behavior.
emit_repro_caller: Emits a main function to call the kernel with example inputs.
Returns:
str: The generated code as a string.
"""
if config is None:
config = self._require_implicit_config()
with self.env, measure("BoundKernel.to_code"):
# Work on a copy so the caller's Config is not mutated with defaults
# specific to this BoundKernel's config_spec.
config = self._normalized_config_copy(config)
with (
self._runtime_arg_values_for_codegen(),
measure("BoundKernel.generate_ast"),
):
# pyrefly: ignore [bad-argument-type]
root = generate_ast(self.host_function, config, emit_repro_caller)
self._register_cute_grouped_static_tail_specializations()
if output_origin_lines is None:
output_origin_lines = self.settings.output_origin_lines
import_lines: list[str] = []
body_start = 0
for i, stmt in enumerate(root.body):
if isinstance(stmt, (ast.Import, ast.ImportFrom)):
if not (
isinstance(stmt, ast.ImportFrom) and stmt.module == "__future__"
):
import_lines.append(ast.unparse(stmt))
continue
body_start = i
break
else:
body_start = len(root.body)
body_root = ast.Module(body=root.body[body_start:], type_ignores=[])
ast.fix_missing_locations(body_root)
# One optional AST processing step, then the single unparse. Both rewrites run
# after generate_ast and outside the fake-tensor env above: jax_fn's launch
# capture runs the compiled kernel on *real* tensors (which specializes
# fake_args, so codegen must already be done); dep-free is pure-AST and
# unaffected by placement. jax_fn is checked first -- it spans both dep modes.
if options is not None and options.jax_fn:
from .._compiler.output_code_utils import build_jax_fn_module
from .._compiler.output_code_utils import capture_jax_launch_metadata
jax_meta = capture_jax_launch_metadata(self, config)
body_root = build_jax_fn_module(
self, options, import_lines, body_root, jax_meta
)
elif options is not None and not options.allow_helion_deps:
from .._compiler.output_code_utils import build_dependency_free_code
body_root = build_dependency_free_code(
self, options, import_lines, body_root
)
with measure("BoundKernel.unparse"):
body = unparse(body_root, output_origin_lines=output_origin_lines)
imports = "\n".join(import_lines)
if imports:
return f"from __future__ import annotations\n\n{imports}\n\n{body}"
return f"from __future__ import annotations\n\n{body}"
def to_triton_code(
self,
config: ConfigLike | None = None,
*,
emit_repro_caller: bool = False,
output_origin_lines: bool | None = None,
) -> str:
"""Backward-compatible alias for :meth:`to_code`."""
return self.to_code(
config,
emit_repro_caller=emit_repro_caller,
output_origin_lines=output_origin_lines,
)
def compile_config(
self, config: ConfigLike | None = None, *, allow_print: bool = True
) -> CompiledConfig:
"""
Compile the kernel for a specific configuration.
Args:
config: The configuration to compile the kernel with.
allow_print: Set to suppress printing the output code when autotuning.
Returns:
CompiledConfig: A callable object representing the compiled kernel.
"""
if config is None:
config = self._require_implicit_config()
requested_config = self._normalize_config(config)
config = self._normalized_config_copy(requested_config)
dist_check_config_consistancy(
config, process_group_name=self._env.process_group_name
)
if (rv := self._compile_cache.get(config)) is not None:
return rv
device_index = (
self._env.device.index if self._env.device.index is not None else 0
)
self.env.backend.setup_compile_cache_dir(device_index)
try:
triton_code = self.to_triton_code(
config, emit_repro_caller=self.settings.print_output_code
)
# static_shapes=True keys a distinct BoundKernel per input shape (see
# _tensor_key), but PyCodeCache keys compiled modules by SOURCE TEXT, so
# two shapes that emit byte-identical source (e.g. compact_worklist, whose
# token dim enters only via runtime offsets + a data-dependent loop) would
# share one module. For backends whose generated module caches shape-
# specific state (Backend.requires_shape_specialized_module -- e.g. Pallas,
# whose output-meta descriptor / launcher cache / ds-pad decision /
# signature lock are all monomorphic) that shared module returns the first
# shape's cached output extent for both. Fold the specialization key into
# the PyCodeCache key (via ``extra``, leaving the source untouched) so each
# specialization gets its own module. Reusing the same key that
# distinguishes BoundKernels (``self._base_spec_key``, computed from the
# real args in __init__) guarantees distinct BoundKernel => distinct module
# and already covers shape/dtype/stride recursively through nested
# container args.
cache_extra = ""
if (
self.settings.static_shapes
and self.env.backend.requires_shape_specialized_module
):
cache_extra = repr(self._base_spec_key)
with measure("BoundKernel.PyCodeCache.load"):
module = PyCodeCache.load(triton_code, extra=cache_extra)
self.env.backend.annotate_compiled_module(
module, triton_code, self.kernel.name
)
except Exception:
log.warning(
"Helion compiler triton codegen error for %s",
self.format_kernel_decorator(requested_config, self.settings),
exc_info=True,
)
self.maybe_log_repro(log.warning, self.fake_args, config=requested_config)
raise
if allow_print:
log.info("Output code written to: %s", module.__file__)
log.debug("Debug string: \n%s", LazyString(lambda: self._debug_str()))
# for distributed kernel, print rank1 code since rank0
# code can skip some offset computation.
if (
not dist.is_initialized() or dist.get_rank() == 1
) and self.settings.print_output_code:
log.info("Output code: \n%s", triton_code)
print(f"# Output code written to: {module.__file__}", file=sys.stderr)
print(triton_code, file=sys.stderr)
rv = getattr(module, self.kernel.name)
self._compile_cache[config] = rv
self._cache_path_map[config] = module.__file__
return rv
def bench_compile_config(
self,
config: Config | dict[str, object] | None = None,
*,
allow_print: bool = True,
) -> Callable[..., object]:
return self.compile_config(config, allow_print=allow_print)
def extra_cache_key(self) -> str:
"""Return extra data folded into the disk-cache key.
Returns ``""`` by default, leaving the cache key unchanged.
"""
return ""
def supports_subprocess_benchmark(self) -> bool:
return True
def is_cacheable(self) -> bool:
return True
def get_cached_path(self, config: ConfigLike | None = None) -> str | None:
"""
Get the file path of the generated Triton code for a specific configuration.
Args:
config: The configuration to get the file path for.
Returns:
str | None: The file path of the generated Triton code, or None if not found.
"""
if config is None:
config = self._require_implicit_config()
requested_config = self._normalize_config(config)
if requested_config in self._cache_path_map:
return self._cache_path_map[requested_config]
try:
config = self._normalized_config_copy(requested_config)
except exc.InvalidConfig:
return None
return self._cache_path_map.get(config, None)
def _debug_str(self) -> str:
"""
Generate a debug string for the kernel.
Returns:
str: A string containing debug information about the kernel.
"""
if self.host_function is None:
# In ref mode, host_function is not created
return f"<BoundKernel {self.kernel.fn.__name__} in ref mode>"
with self.env:
return self.host_function.debug_str()
def _get_host_semantic_input_normalization(
self,
) -> _HostSemanticInputNormalization:
"""Return capture-time equivalence classes for dynamic input symbols."""
if self.settings.static_shapes:
return ()
canonical_by_expr: dict[sympy.Expr, int] = {}
normalization: list[tuple[int, str, int, int]] = []
with self.env:
tensor_args = self._host_semantic_input_tensors()
for tensor_index, arg in enumerate(tensor_args):
for property_name, values in (
("size", arg.size()),
("stride", arg.stride()),
):
for index, value in enumerate(values):
if not isinstance(value, torch.SymInt):
continue
expr = self.env.shape_env.simplify(value._sympy_())
if isinstance(expr, sympy.Integer):
continue
canonical = canonical_by_expr.get(expr)
if canonical is None:
canonical = len(canonical_by_expr)
canonical_by_expr[expr] = canonical
normalization.append(
(tensor_index, property_name, index, canonical)
)
return tuple(normalization)
def _host_semantic_debug_renames(
self,
normalization: _HostSemanticInputNormalization,
) -> tuple[dict[sympy.Basic, sympy.Basic], bool]:
tensor_args = self._host_semantic_input_tensors()
renames: dict[sympy.Basic, sympy.Basic] = {}
class_by_expr: dict[sympy.Basic, int] = {}
compatible = True
for tensor_index, property_name, index, canonical in normalization:
# A missing or concretized property is left unnormalized. If it is
# semantically used, the rendered host/device trace still differs.
if tensor_index >= len(tensor_args):
continue
tensor = tensor_args[tensor_index]
values = tensor.size() if property_name == "size" else tensor.stride()
if index >= len(values):
continue
value = values[index]
if not isinstance(value, torch.SymInt):
continue
expr = self.env.shape_env.simplify(value._sympy_())
existing = class_by_expr.get(expr)
if existing is not None and existing != canonical:
compatible = False
continue
class_by_expr[expr] = canonical
renames[expr] = sympy.Symbol(
f"<helion_input_symbol_{canonical}>",
integer=True,
)
return renames, compatible
def _host_semantic_input_tensors(self) -> list[torch.Tensor]:
return [
tensor
for tensor, source in self.env.input_sources.items()
if _is_supported_tensor_input_source(source)
]
def _host_semantic_external_tensor_keys(self) -> tuple[object, ...]:
return tuple(
(
type(source).__name__,
tuple(self._semantic_dimension_key(dim) for dim in tensor.shape),
tuple(self._semantic_dimension_key(dim) for dim in tensor.stride()),
str(tensor.dtype),
str(tensor.device),
str(tensor.layout),
tensor.requires_grad,
)
for tensor, source in self.env.input_sources.items()
if not _is_supported_tensor_input_source(source)
)
def _get_host_semantic_fingerprint(
self,
outputs: Sequence[object],
*,
input_normalization: _HostSemanticInputNormalization | None = None,
) -> str:
assert self.host_function is not None
if input_normalization is None:
input_normalization = self._get_host_semantic_input_normalization()
output_keys: list[object] = []
for output in outputs:
if isinstance(output, torch.Tensor):
if self.settings.static_shapes:
shape = tuple(
self._semantic_dimension_key(dim) for dim in output.shape
)
else:
shape = ("rank", output.ndim)
output_keys.append(
(
"tensor",
shape,
str(output.dtype),
str(output.device),
str(output.layout),
)
)
elif isinstance(output, (torch.SymInt, torch.SymFloat, torch.SymBool)):
output_keys.append((type(output).__name__,))
elif output is None or type(output) in (bool, float, int, str):
output_type = type(output)
output_keys.append(
(
f"{output_type.__module__}.{output_type.__qualname__}",
repr(output),
)
)
else:
raise TypeError(
f"Unsupported Helion host output in semantic fingerprint: "
f"{type(output).__name__}"
)
output_key = tuple(output_keys)
cache_key = (input_normalization, output_key)
fingerprint = self._host_semantic_fingerprints.get(cache_key)
if fingerprint is not None:
return fingerprint
normalization_compatible = True
with self.env:
if input_normalization:
renames, normalization_compatible = self._host_semantic_debug_renames(
input_normalization
)
with self.env.use_debug_shape_renames(renames):
host_source = self.host_function.semantic_debug_str()
else:
host_source = self.host_function.semantic_debug_str()
payload = (
input_normalization,
normalization_compatible,
self.env.backend_name,
str(self.env.index_dtype),
self._host_semantic_external_tensor_keys(),
host_source,
output_key,
)
fingerprint = hashlib.sha256(repr(payload).encode()).hexdigest()
self._host_semantic_fingerprints[cache_key] = fingerprint
return fingerprint
def _semantic_dimension_key(self, dim: object) -> object:
if isinstance(dim, torch.SymInt) and dim.node.shape_env is self.env.shape_env:
return shape_env_size_hint(self.env.shape_env, dim.node.expr)
return dim
def autotune(
self,
args: Sequence[object],
*,
force: bool = True,
**kwargs: object,
) -> Config:
"""
Perform autotuning to find the optimal configuration for the kernel. This uses the
default setting, you can call helion.autotune.* directly for more customization.
If config= or configs= is provided to helion.kernel(), the search will be restricted to
the provided configs. Use force=True to ignore the provided configs.
Mutates self so that `__call__` will run the best config found.
Args:
args: Example arguments used for benchmarking during autotuning.
force: If True, force full autotuning even if a config is provided.
kwargs: Additional keyword options forwarded to the autotuner.
Returns:
Config: The best configuration found during autotuning.
"""
ephemeral = self.env.backend.make_ephemeral_cache()
ctx = ephemeral if ephemeral is not None else contextlib.nullcontext()
with ctx:
config = self.env.backend.autotune(self, args, force=force, **kwargs)
if ephemeral is not None:
self.env.backend.finalize_ephemeral_cache(self, config)
self.set_config(config)
return config
def set_config(self, config: ConfigLike) -> None:
"""
Set the configuration for the kernel and compile it.
Mutates self so that `__call__` will run the provided config.
Args:
config: The configuration to set.
"""
config = self._normalize_config(config)
self._run = self.compile_config(config)
self._config = config
counters["best_config_decorator"][
self.format_kernel_decorator(config, self.settings)
] = 1
def _specialize_extra(self) -> list[Callable[[Sequence[object]], Hashable]]:
"""
Returns a list of functions that will be called to generate extra specialization keys.
This is used to specialize on the values hl.specialize()'ed arguments.
Returns:
list[Callable[[Sequence[object]], Hashable]]: A list of functions that generate extra specialization keys.
"""
if (
not self.env.specialized_vars
and not self.env.specialized_strides
and not self.env.tensor_descriptor_layout_guards
):
return []
def make_extractor(v: Source) -> Callable[[Sequence[object]], Hashable]:
if isinstance(v, TensorPropertySource):
index = v.idx
assert index is not None
inner = make_extractor(v.base)
if v.prop == TensorProperty.SIZE:
def size_extractor(
args: Sequence[object],
_inner: Callable[[Sequence[object]], Hashable] = inner,
_index: int = index,
) -> Hashable:
result = _inner(args)
# Handle list of tensors: return tuple of sizes for all tensors
if isinstance(result, (list, tuple)):
return tuple(
cast("torch.Tensor", t).size(_index) for t in result
)
return cast("torch.Tensor", result).size(_index)
return size_extractor
if v.prop == TensorProperty.STRIDE:
def stride_extractor(
args: Sequence[object],
_inner: Callable[[Sequence[object]], Hashable] = inner,
_index: int = index,
) -> Hashable:
result = _inner(args)
# Handle list of tensors: return tuple of strides for all tensors
if isinstance(result, (list, tuple)):
return tuple(
cast("torch.Tensor", t).stride(_index) for t in result
)
return cast("torch.Tensor", result).stride(_index)
return stride_extractor
raise exc.SpecializeArgType(v)
if isinstance(v, GetItemSource):
if not isinstance(v.index, (int, str)) or v.index_is_slice:
raise exc.SpecializeArgType(v)
inner = make_extractor(v.base)
def getitem_extractor(
args: Sequence[object],
_inner: Callable[[Sequence[object]], Hashable] = inner,
_index: int | str = v.index,
) -> Hashable:
result = _inner(args)
if isinstance(result, dict):
return cast("Hashable", result[_index])
if isinstance(_index, str):
return cast("Hashable", getattr(result, _index))
return cast("Sequence[Hashable]", result)[_index]
return getitem_extractor
if isinstance(v, LocalSource):
index = arg_name_to_index[v.local_name]
return operator.itemgetter(index)
raise exc.SpecializeArgType(v)
arg_name_to_index: dict[str, int] = {
n: i for i, n in enumerate(self.kernel.signature.parameters.keys())
}
extractors: list[Callable[[Sequence[object]], Hashable]] = []
extracted_strides: set[TensorPropertySource] = set()
for v in sorted(self.env.specialized_vars, key=lambda v: v.name):
source = self.env.shape_env.var_to_sources[v][0]
extractors.append(make_extractor(source))
if (
isinstance(source, TensorPropertySource)
and source.prop == TensorProperty.STRIDE
and source.idx is not None
):
extracted_strides.add(source)
for source in sorted(self.env.specialized_strides, key=repr):
if source in extracted_strides:
continue
extractors.append(make_extractor(source))
implicit_config = self._fixed_config_for_td_layout_guards()
for source, guard in sorted(
self.env.tensor_descriptor_layout_guards.items(),
key=lambda item: repr(item[0]),
):
if implicit_config is not None and not _td_layout_guard_active_for_config(
guard, implicit_config
):
continue
extract_tensor = make_extractor(source)
def td_layout_extractor(
args: Sequence[object],
_extract_tensor: Callable[
[Sequence[object]], Hashable
] = extract_tensor,
_ndim: int = guard.ndim,
_element_size: int = guard.element_size,
) -> Hashable:
tensor = cast("torch.Tensor", _extract_tensor(args))
if tensor.ndim != _ndim:
return ("ndim", tensor.ndim)
return tensor_descriptor_layout_signature_from_strides(
tensor.stride(),
_element_size,
)
extractors.append(td_layout_extractor)
return extractors
@contextlib.contextmanager
def _runtime_arg_values_for_codegen(self) -> Generator[None, None, None]:
values: dict[str, object] = self.env.runtime_arg_values_by_name
if not values:
values = {
name: value
for name, ref in self._runtime_tensor_refs_by_name.items()
if (value := ref()) is not None
}
with self.env.use_runtime_arg_values(values):
yield
def _register_cute_grouped_static_tail_specializations(self) -> None:
if self.kernel.settings.backend != "cute" or not self._cache_managed:
return
signature = self._base_spec_key
descriptors = _cute_grouped_static_tail_extra_descriptors(
self.env.cute_resolved_wrapper_plans
)
if not descriptors:
return
# Serialize descriptor discovery with schema extension so concurrent
# code generation cannot publish duplicate equivalent extractors.
with self.kernel._bind_lock:
if self._reset_generation != self.kernel._reset_generation:
return
seen = self.kernel._cute_grouped_static_tail_extra_descriptors.setdefault(
signature,
set(),
)
new_descriptors: list[Hashable] = []
new_extractors: list[Callable[[Sequence[object]], Hashable]] = []
for descriptor in descriptors:
if descriptor in seen:
continue
new_descriptors.append(descriptor)
new_extractors.append(
_make_cute_grouped_static_tail_extractor(descriptor)
)
runtime_args = tuple(
self.env.runtime_arg_values_by_name.get(name)
for name in self.kernel.signature.parameters
)
if self.kernel._extend_bound_kernel_specializations(
self,
signature,
new_extractors,
runtime_args,
):
seen.update(new_descriptors)
def _fixed_config_for_td_layout_guards(self) -> Config | None:
"""Return the fixed config if TD layout guards can be filtered safely."""
if self._config is not None:
return self._config
if self.kernel.settings.autotune_effort == "none" and (
len(self.kernel.configs) == 0 or self.settings.force_autotune
):
return self.config_spec.default_config()
if self.settings.force_autotune:
return None
if len(self.kernel.configs) == 1:
return self.kernel.configs[0]
return None
def _user_provided_config(self) -> Config | None:
"""Return a config if the user explicitly provided one, else None.
Checks the kernel's config list and settings to determine if
a config can be resolved without autotuning.
"""
configs = self.kernel.configs
if self.kernel.settings.autotune_effort == "none" and (
len(configs) == 0 or self.settings.force_autotune
):
config = self.config_spec.default_config()
if not is_ref_mode_enabled(self.kernel.settings):
kernel_decorator = self.format_kernel_decorator(config, self.settings)
print(
f"Using implicit config: {kernel_decorator}",
file=sys.stderr,
)
return config
if self.settings.force_autotune:
return None
if len(configs) == 1:
return configs[0]
return None
def _implicit_config(self) -> Config | None:
"""
Returns a single config that is implicitly used by this kernel, if any.
"""
if self._config is not None:
return self._config
return self._user_provided_config()
def _require_implicit_config(self) -> Config:
"""
Returns the implicit config for this kernel, or raises an error if no implicit config is available.
"""
if (config := self._implicit_config()) is None:
raise RuntimeError("no config provided and no implicit config available")
return config
def ensure_config_exists(self, args: Sequence[object]) -> None:
"""
Ensure a config is available, triggering autotuning if needed.
If an implicit config is available (from configs list or default), it will be used.
Otherwise, autotuning will be triggered with the provided args.
"""
if self._config is not None:
return # Already have a config
if (config := self._implicit_config()) is not None:
with measure("BoundKernel.set_config"):
self.set_config(config)
else:
with measure("BoundKernel.autotune"):
self.autotune(args, force=False)
# pyrefly: ignore [bad-return]
def run_ref(self, *args: object) -> _R:
# Unwrap ConstExpr arguments
clean_args = []
for arg in args:
if isinstance(arg, ConstExpr):
clean_args.append(arg.value)
else:
clean_args.append(arg)
# Pass the config to RefModeContext
with RefModeContext(self.env, self._config):
result = self.kernel.fn(*clean_args)
return cast("_R", result)
def __call__(self, *args: object) -> _R:
"""
Execute the kernel with the given arguments.
Args:
args: The arguments to pass to the kernel.
Returns:
_R: The result of the kernel execution.
"""
if self._cache_managed and self.kernel._has_specialization_extras:
rebound = self.kernel.bind(args)
if rebound is not self:
return rebound(*args)
if self._run is None:
with self._first_compile_lock:
# Another caller may have discovered a late specialization while
# this caller waited for the first compile.
if self._cache_managed and self.kernel._has_specialization_extras:
rebound = self.kernel.bind(args)
if rebound is not self:
return rebound(*args)
if self._run is None:
if is_ref_mode_enabled(self.kernel.settings):
if (config := self._implicit_config()) is not None:
self._config = config
return self.run_ref(*args)
runtime_args: dict[str, object] = {
name: value
for name, value in zip(
self.kernel.signature.parameters, args, strict=False
)
if isinstance(value, torch.Tensor)
}
with self.env.use_runtime_arg_values(runtime_args):
self.ensure_config_exists(args)
assert self._run is not None
self.maybe_log_repro(log.warning, args)
return self._run(*args)
def backend_cache_key(self, config: ConfigLike | None = None) -> str | None:
"""
Return the backend cache key for the compiled kernel.
For the Triton backend, this is the base32 encoding of the SHA-256
hash that Triton uses to cache compiled GPU binaries under
``TRITON_CACHE_DIR/<key>/``. For the CuTe backend, it is the base32
encoding of the SHA-256 hash of the compiled IR module, which names the
``CUTE_DSL_CACHE_DIR/<key>.mlir`` artifact.
Args:
config: The configuration to look up. Defaults to the implicit config.
Returns:
str | None: The cache key, or None if the kernel hasn't been
JIT-compiled yet or the backend doesn't support cache keys.
"""
if config is None:
config = self._require_implicit_config()
config = self._normalized_config_copy(config)
compiled_fn = self._compile_cache.get(config)
if compiled_fn is None:
return None
return self.env.backend.compiled_cache_key(self, compiled_fn)
def maybe_log_repro(
self,
log_func: Callable[[str], None],
args: Sequence[object],
config: Config | None = None,
) -> None:
if not self.settings.print_repro:
return
effective_config = config or self._config
assert effective_config is not None
# Get kernel source
try:
raw_source = inspect.getsource(self.kernel.fn)
source_lines = textwrap.dedent(raw_source).splitlines()
# Skip decorator lines (including multi-line decorators)
start_idx = 0
while start_idx < len(source_lines) and not source_lines[
start_idx
].lstrip().startswith("def "):
start_idx += 1
kernel_body = "\n".join(source_lines[start_idx:])
except (OSError, TypeError):
kernel_body = f"# Source unavailable for {self.kernel.fn.__module__}.{self.kernel.fn.__qualname__}"
# Format decorator
decorator = self.format_kernel_decorator(effective_config, self.settings)
# Build output
output_lines = [
"# === HELION KERNEL REPRO ===",
"import helion",
"import helion.language as hl",
"import torch",
"from torch._dynamo.testing import rand_strided",
"",
decorator,
kernel_body,
]
# Generate caller function
if args:
def _render_input_arg_assignment(name: str, value: object) -> list[str]:
if isinstance(value, torch.Tensor):
shape = tuple(int(d) for d in value.shape)
stride = tuple(int(s) for s in value.stride())
device = str(value.device)
dtype = str(value.dtype)
lines = [
f"{name} = rand_strided({shape!r}, {stride!r}, dtype={dtype}, device={device!r})"
]
if value.requires_grad:
lines.append(f"{name}.requires_grad_(True)")
return lines
return [f"{name} = {value!r}"]
sig_param_names = list(self.kernel.signature.parameters.keys())
assert len(args) == len(sig_param_names)
output_lines.extend(["", "def helion_repro_caller():"])
output_lines.append(" torch.manual_seed(0)")
arg_names: list[str] = []
for i, value in enumerate(args):
var_name = sig_param_names[i]
arg_names.append(var_name)
# Add assignment lines with indentation
for line in _render_input_arg_assignment(var_name, value):
output_lines.append(f" {line}")
# Add return statement
call_args = ", ".join(arg_names)
output_lines.append(f" return {self.kernel.name}({call_args})")
output_lines.extend(["", "helion_repro_caller()"])
output_lines.append("# === END HELION KERNEL REPRO ===")
repro_text = "\n" + "\n".join(output_lines)
log_func(repro_text)
class _KernelDecorator(Protocol):
def __call__(
self,
fn: Callable[..., _R],
) -> Kernel[_R]: ...
@overload
def kernel(
fn: Callable[..., _R],
*,
config: ConfigLike | None = None,
configs: Sequence[ConfigLike] | None = None,
key: Callable[..., Hashable] | None = None,
**settings: object,
) -> Kernel[_R]: ...
@overload
def kernel(
fn: None = None,
*,
config: ConfigLike | None = None,
configs: Sequence[ConfigLike] | None = None,
key: Callable[..., Hashable] | None = None,
**settings: object,
) -> _KernelDecorator: ...
[docs]
def kernel(
fn: Callable[..., _R] | None = None,
*,
config: ConfigLike | None = None,
configs: Sequence[ConfigLike] | None = None,
key: Callable[..., Hashable] | None = None,
**settings: object,
) -> Kernel[_R] | _KernelDecorator:
"""
Decorator to create a Kernel object from a Python function.
Args:
fn: The function to be wrapped by the Kernel. If None, a decorator is returned.
config: A single configuration to use for the kernel. Refer to the
``helion.Config`` class for details.
configs: A list of configurations to use for the kernel. Can only specify
one of config or configs. Refer to the ``helion.Config`` class for
details.
key: Optional callable returning a hashable that augments the specialization key.
settings: Keyword arguments representing settings for the Kernel.
Can also use settings=Settings(...) to pass a Settings object
directly. Refer to the ``helion.Settings`` class for available
options.
Returns:
object: A Kernel object or a decorator that returns a Kernel object.
See Also:
- :class:`~helion.Settings`: Controls compilation behavior and debugging options
- :class:`~helion.Config`: Controls GPU execution parameters and optimization strategies
"""
if config is not None:
assert not configs, "Cannot specify both config and configs"
configs = [config]
elif configs is None:
configs = []
if settings_obj := settings.get("settings"):
assert len(settings) == 1, "settings must be the only keyword argument"
assert isinstance(settings_obj, Settings), "settings must be a Settings object"
else:
settings_obj = Settings(**settings)
if fn is None:
return functools.partial(
kernel,
configs=configs,
settings=settings_obj,
key=key,
)
return Kernel(
fn,
configs=configs,
settings=settings_obj,
key=key,
)
def _hashable_dim(s: int | torch.SymInt) -> Hashable:
if isinstance(s, torch.SymInt):
return (id(s.node.shape_env), s.node.expr)
return s
def _safe_bucket_dim(s: int | torch.SymInt) -> Hashable:
if isinstance(s, torch.SymInt):
return (id(s.node.shape_env), s.node.expr)
# Dynamic-shape kernels should not get separate bound kernels for sizes
# 0 or 1. Keep 2 as the canonical "dynamic dimension" bucket that was
# already used for all concrete sizes >= 2.
return 2
_EMPTY_FROZENSET: frozenset[int] = frozenset()
def _bucketed_size(obj: torch.Tensor) -> tuple[Hashable, ...]:
sz = obj.size()
n = len(sz)
if n == 1:
return (_safe_bucket_dim(sz[0]),)
if n == 2:
return (_safe_bucket_dim(sz[0]), _safe_bucket_dim(sz[1]))
if n == 3:
return (
_safe_bucket_dim(sz[0]),
_safe_bucket_dim(sz[1]),
_safe_bucket_dim(sz[2]),
)
return tuple(_safe_bucket_dim(s) for s in sz)
def _hashable_dims(dims: Sequence[int | torch.SymInt]) -> tuple[Hashable, ...]:
n = len(dims)
if n == 1:
return (_hashable_dim(dims[0]),)
if n == 2:
return (_hashable_dim(dims[0]), _hashable_dim(dims[1]))
if n == 3:
return (_hashable_dim(dims[0]), _hashable_dim(dims[1]), _hashable_dim(dims[2]))
return tuple(_hashable_dim(s) for s in dims)
def _concrete_tensor_key(fn: Kernel, obj: torch.Tensor) -> Hashable:
# Fast extractor for plain ``torch.Tensor`` / ``torch.nn.Parameter``:
# exact-type dispatch guarantees concrete int sizes/strides, so
# ``torch.Size`` and the stride tuple can be used directly (both are
# tuple subclasses that hash/compare identically to plain int tuples).
# The ``_hashable_dims`` wrap in ``_tensor_key`` exists only to
# normalize SymInts, which appear on FakeTensors during tracing.
si = getattr(obj, "_dynamo_static_indices", None)
static_indices = frozenset(si) if si is not None else _EMPTY_FROZENSET
if fn.settings.static_shapes:
return (obj.dtype, obj.size(), obj.stride(), static_indices)
bucketed = _bucketed_size(obj)
if fn.settings.index_dtype is None:
try:
needs_int64 = bool(obj.numel() > _INT32_INDEX_LIMIT)
except RuntimeError:
needs_int64 = True # unbacked SymInt
return (
obj.dtype,
bucketed,
needs_int64,
static_indices,
)
return (
obj.dtype,
bucketed,
static_indices,
)
def _tensor_key(fn: Kernel, obj: torch.Tensor) -> Hashable:
si = getattr(obj, "_dynamo_static_indices", None)
static_indices = frozenset(si) if si is not None else _EMPTY_FROZENSET
if fn.settings.static_shapes:
return (
obj.dtype,
_hashable_dims(obj.size()),
_hashable_dims(obj.stride()),
static_indices,
)
bucketed = _bucketed_size(obj)
if fn.settings.index_dtype is None:
try:
needs_int64 = bool(obj.numel() > _INT32_INDEX_LIMIT)
except RuntimeError:
needs_int64 = True # unbacked SymInt
return (
obj.dtype,
bucketed,
needs_int64,
static_indices,
)
return (
obj.dtype,
bucketed,
static_indices,
)
def _sequence_key(fn: Kernel, obj: Sequence) -> Hashable:
return type(obj), tuple([fn._specialization_key(item) for item in obj])
def _mapping_key(
fn: Kernel, obj: dict[str | int, object], real_type: type[object]
) -> Hashable:
return real_type, tuple(
sorted((k, fn._specialization_key(v)) for k, v in obj.items())
)
def _number_key(fn: Kernel, n: float | bool) -> object:
return type(n)
def _function_key(fn: Kernel, obj: types.FunctionType) -> object:
if obj.__closure__:
closures = [
fn._specialization_key(cell.cell_contents) for cell in obj.__closure__
]
return (obj.__code__, *closures)
return obj.__code__
def _cute_grouped_layout_has_m_tail(
layout_values: tuple[int, ...],
*,
bm: int,
group_count: int,
) -> bool | None:
cursor = 0
has_m_tail = False
for expected_group in range(group_count):
if expected_group > 0:
next_m_boundary = ((cursor + bm - 1) // bm) * bm
while (
cursor < len(layout_values)
and cursor < next_m_boundary
and layout_values[cursor] < 0
):
cursor += 1
if cursor != next_m_boundary or (
cursor < len(layout_values) and layout_values[cursor] < 0
):
return None
if cursor >= len(layout_values) or layout_values[cursor] != expected_group:
return None
start = cursor
while cursor < len(layout_values) and layout_values[cursor] == expected_group:
cursor += 1
actual_m = cursor - start
if start % bm != 0:
return None
has_m_tail = has_m_tail or actual_m % bm != 0
if cursor != len(layout_values):
if all(value < 0 for value in layout_values[cursor:]):
cursor = len(layout_values)
if cursor != len(layout_values):
return None
return has_m_tail
def _cute_grouped_static_tail_extra_descriptors(
plans: Sequence[dict[str, object]],
) -> tuple[Hashable, ...]:
descriptors: list[Hashable] = []
for plan in plans:
if plan.get("kind") != "tcgen05_grouped_static_persistent" or bool(
plan.get("worklist_metadata")
):
continue
if not (
isinstance(plan.get("grouped_static_has_m_tail"), bool)
or isinstance(plan.get("grouped_static_has_n_tail"), bool)
):
continue
layout_idx = plan.get("layout_bind_idx")
group_count = plan.get("group_count")
bm = plan.get("bm")
if not (
isinstance(layout_idx, int)
and isinstance(group_count, int)
and isinstance(bm, int)
):
continue
n_sizes_idx = plan.get("n_sizes_bind_idx")
bn = plan.get("bn")
descriptors.append(
(
"cute_grouped_static_tail",
layout_idx,
group_count,
bm,
n_sizes_idx if isinstance(n_sizes_idx, int) else None,
bn if isinstance(bn, int) else None,
)
)
return tuple(sorted(descriptors, key=repr))
def _cute_int_1d_tensor_values(
value: object,
cache: WeakIdKeyDictionary,
) -> tuple[int, ...] | None:
if not (
isinstance(value, torch.Tensor)
and value.ndim == 1
and value.dtype in (torch.int32, torch.int64)
):
return None
if torch.is_inference(value):
return tuple(int(v) for v in value.detach().cpu().tolist())
signature = (
int(value._version),
int(value.data_ptr()),
tuple(value.shape),
tuple(value.stride()),
value.dtype,
)
try:
cached_signature, cached_values = cache[value]
except KeyError:
pass
else:
if cached_signature == signature:
return cached_values
values = tuple(int(v) for v in value.detach().cpu().tolist())
cache[value] = (signature, values)
return values
def _make_cute_grouped_static_tail_extractor(
descriptor: Hashable,
) -> Callable[[Sequence[object]], Hashable]:
(
_label,
layout_idx,
group_count,
bm,
n_sizes_idx,
bn,
) = cast("tuple[object, int, int, int, int | None, int | None]", descriptor)
tensor_values_cache: WeakIdKeyDictionary = WeakIdKeyDictionary()
def cute_grouped_static_tail_extractor(
args: Sequence[object],
*,
_descriptor: Hashable = descriptor,
_layout_idx: int = layout_idx,
_group_count: int = group_count,
_bm: int = bm,
_n_sizes_idx: int | None = n_sizes_idx,
_bn: int | None = bn,
) -> Hashable:
layout_has_m_tail: bool | None = None
if _layout_idx < len(args):
layout_values = _cute_int_1d_tensor_values(
args[_layout_idx],
tensor_values_cache,
)
if layout_values is not None:
layout_has_m_tail = _cute_grouped_layout_has_m_tail(
layout_values,
bm=_bm,
group_count=_group_count,
)
n_sizes_has_n_tail: bool | None = None
if _n_sizes_idx is not None and _bn is not None and _n_sizes_idx < len(args):
n_sizes_values = _cute_int_1d_tensor_values(
args[_n_sizes_idx],
tensor_values_cache,
)
if n_sizes_values is not None and len(n_sizes_values) == _group_count:
n_sizes_has_n_tail = any(
group_n % _bn != 0 for group_n in n_sizes_values
)
return (_descriptor, layout_has_m_tail, n_sizes_has_n_tail)
return cute_grouped_static_tail_extractor
def _graph_module_key(fn: Kernel, obj: torch.fx.GraphModule) -> Hashable:
"""Generate a specialization key for GraphModule arguments."""
# Check if already cached
if obj in _graph_module_hash_cache:
return _graph_module_hash_cache[obj]
# Check for unsupported operations
unsupported_ops = {
node.op
for node in itertools.chain(
obj.graph.find_nodes(op="call_module"),
obj.graph.find_nodes(op="get_attr"),
)
}
if unsupported_ops:
raise exc.GraphModuleUnsupportedOps(", ".join(sorted(unsupported_ops)))
_graph_module_hash_cache[obj] = rv = str(compiled_fx_graph_hash(obj, [], {}, []))
return rv
_specialization_extractors: dict[
type[object] | str,
Callable[[Kernel, object], Hashable],
# pyrefly: ignore [bad-assignment]
] = {
# Exact-type dispatch (see ``_specialization_key``): plain tensors and
# Parameters always have concrete int sizes/strides and take the fast
# extractor. Subclasses (FakeTensor below, or anything hitting the
# ``isinstance`` fallback) go through SymInt-safe ``_tensor_key``.
torch.Tensor: _concrete_tensor_key,
torch.nn.Parameter: _concrete_tensor_key,
FakeTensor: _tensor_key,
# SymInt-safe extractor for torch.Tensor subclasses reached via the
# isinstance fallback in ``_specialization_key`` (string key so the
# fallback stays loosely typed, like "namedtuple" / "dataclass").
"tensor_subclass": _tensor_key,
torch.dtype: lambda fn, x: x,
torch.device: lambda fn, x: x,
int: _number_key,
float: _number_key,
bool: _number_key,
str: lambda fn, x: x,
list: _sequence_key,
tuple: _sequence_key,
# pyrefly: ignore [bad-argument-type]
dict: lambda fn, x: _mapping_key(fn, x, type(x)),
# pyrefly: ignore [missing-attribute]
"namedtuple": lambda fn, x: _mapping_key(fn, x._asdict(), type(x)),
# pyrefly: ignore [no-matching-overload, bad-argument-type]
"dataclass": lambda fn, x: _mapping_key(fn, dataclasses.asdict(x), type(x)),
types.FunctionType: _function_key,
types.BuiltinFunctionType: lambda fn, x: x,
torch.fx.GraphModule: _graph_module_key,
# pyrefly: ignore [missing-attribute]
ConstExpr: lambda fn, x: x.value,
type(None): lambda fn, x: None,
}
def _find_device(args: tuple[object, ...]) -> torch.device:
"""
Extract the device from the arguments.
Args:
args: The arguments to extract the device from.
Returns:
torch.device: The extracted device
"""
for arg in args:
if isinstance(arg, torch.device):
return arg
if isinstance(arg, torch.Tensor):
return arg.device
if isinstance(arg, ConstExpr):
continue # a constexpr-wrapped value carries no device; skip it
if isinstance(arg, (tuple, list)):
for item in arg:
try:
return _find_device((item,))
except exc.NoTensorArgs:
pass
elif isinstance(arg, dict):
for item in arg.values():
try:
return _find_device((item,))
except exc.NoTensorArgs:
pass
raise exc.NoTensorArgs
def _maybe_skip_dtype_check_in_meta_registrations() -> (
contextlib.AbstractContextManager[None, None]
):
# pyrefly: ignore [implicit-import]
if hasattr(torch.fx.experimental._config, "skip_dtype_check_in_meta_registrations"):
# pyrefly: ignore [implicit-import, missing-attribute]
return torch.fx.experimental._config.patch(
skip_dtype_check_in_meta_registrations=True
)
return contextlib.nullcontext()
