diff --git a/helion/_compiler/compile_environment.py b/helion/_compiler/compile_environment.py index 525924514..44d5aa2c6 100644 --- a/helion/_compiler/compile_environment.py +++ b/helion/_compiler/compile_environment.py @@ -3,6 +3,7 @@ import collections import contextlib import dataclasses +import itertools import logging import sys import threading @@ -140,7 +141,7 @@ def tensor_descriptor_layout_signature_from_strides( def _make_numel_check( - symbols: list[sympy.Basic], expr: sympy.Basic + symbols: typing.Sequence[sympy.Basic], expr: sympy.Basic ) -> typing.Callable[..., bool]: """Evaluate a sympy constraint with concrete block-size values.""" @@ -158,6 +159,7 @@ def check(*args: int) -> bool: from torch._guards import Source from .. import Config + from ..autotuner.config_spec import BlockSizeConstraint from ..runtime.settings import Settings from .backend import Backend from .pallas.compact_worklist import CompactWorklistPlan @@ -1059,6 +1061,133 @@ def known_equal(self, a: int | torch.SymInt, b: int | torch.SymInt) -> bool: return bool(res) return a == b + def prepare_tile_reshape( + self, + input_shape: typing.Sequence[int | torch.SymInt], + output_shape: typing.Sequence[int | torch.SymInt], + ) -> None: + """Record reshape guards involving direct tunable block-size symbols.""" + if any(isinstance(dim, int) and dim == -1 for dim in output_shape): + # PyTorch already handles inferred dimensions. Supporting them here + # requires normalizing -1 to an exact symbolic quotient first. + return + + input_numel = sympy.prod(_to_sympy(dim) for dim in input_shape) + output_numel = sympy.prod(_to_sympy(dim) for dim in output_shape) + constraint_expr = sympy.simplify(sympy.Eq(input_numel, output_numel)) + if constraint_expr is sympy.true or constraint_expr is sympy.false: + return + + from ..autotuner.config_spec import BlockSizeConstraint + + block_ids: set[int] = set() + config_block_ids = set(self.config_spec.block_sizes.valid_block_ids()) + for symbol in constraint_expr.free_symbols: + if not isinstance(symbol, sympy.Symbol): + return + block_id = self.get_block_id(symbol) + # Accept only direct tunable symbols so this path does not need + # block-size alias substitution or canonicalization. + if ( + block_id is None + or block_id not in config_block_ids + or self.block_sizes[block_id].symbol() != symbol + ): + return + block_ids.add(block_id) + ordered_block_ids = tuple(sorted(block_ids)) + symbols = [ + self.block_sizes[block_id].symbol() for block_id in ordered_block_ids + ] + constraint = BlockSizeConstraint( + check_fn=_make_numel_check(symbols, constraint_expr), + block_ids=ordered_block_ids, + expr_str=str(constraint_expr), + ) + + constraints = [*self.config_spec.block_size_constraints] + constraint_exists = any( + existing.block_ids == constraint.block_ids + and existing.expr_str == constraint.expr_str + for existing in constraints + ) + if not constraint_exists: + constraints.append(constraint) + + assignment = self._find_block_constraint_assignment( + constraints, + preferred=shape_env_var_hints(self.shape_env), + ) + if assignment is None: + raise exc.InvalidConfig( + "No power-of-two block size assignment satisfies reshape " + f"constraint: {constraint.expr_str}" + ) + + if not constraint_exists: + self.config_spec.block_size_constraints.append(constraint) + + hints = shape_env_var_hints(self.shape_env) + for block_id, value in assignment.items(): + info = self.block_sizes[block_id] + hints[info.symbol()] = sympy.Integer(value) + self.config_spec.block_sizes.block_id_lookup( + block_id + ).autotuner_default = value + + def _find_block_constraint_assignment( + self, + constraints: typing.Sequence[BlockSizeConstraint], + *, + preferred: typing.Mapping[sympy.Symbol, sympy.Integer], + max_attempts: int = 65536, + ) -> dict[int, int] | None: + """Find concrete power-of-two block sizes satisfying all shape guards.""" + block_ids = tuple( + sorted( + { + block_id + for constraint in constraints + for block_id in constraint.block_ids + } + ) + ) + if not block_ids: + return {} + + domains: list[list[int]] = [] + for block_id in block_ids: + spec = self.config_spec.block_sizes.block_id_lookup(block_id) + low = max(spec.min_size, 1) + high = spec.max_size + values = [ + 1 << exponent + for exponent in range(low.bit_length() - 1, high.bit_length()) + if (1 << exponent) >= low + ] + preferred_value = int( + preferred.get(self.block_sizes[block_id].symbol(), values[0]) + ) + values.sort( + key=lambda value: abs( + (value.bit_length() - 1) - (preferred_value.bit_length() - 1) + ) + ) + domains.append(values) + + for attempt, values in enumerate(itertools.product(*domains)): + if attempt >= max_attempts: + break + assignment = dict(zip(block_ids, values, strict=True)) + if all( + constraint.check_fn( + *(assignment[block_id] for block_id in constraint.block_ids) + ) + for constraint in constraints + ): + return assignment + return None + def known_multiple(self, a: sympy.Expr, b: int | torch.SymInt) -> bool: if isinstance(a, (int, sympy.Integer)) and isinstance(b, int): return (int(a) % b) == 0 diff --git a/helion/_compiler/kernel_compiler.py b/helion/_compiler/kernel_compiler.py index 8768ab2e9..f9dc7a928 100644 --- a/helion/_compiler/kernel_compiler.py +++ b/helion/_compiler/kernel_compiler.py @@ -15,6 +15,7 @@ from .host_function import HostFunction from .host_function import KernelDefinition from .tensor_utils import patch_tensor_factories +from .tile_shape_constraints import TileShapeConstraintMode if TYPE_CHECKING: import types @@ -172,6 +173,7 @@ def _compilation_context(self) -> typing.Generator[None, None, None]: # suppress_guards() prevents this by skipping guard # installation (including replacements) in evaluate_expr. HostFunction._suppress_guards_if_profiler_enabled(self.env), + TileShapeConstraintMode(), torch.device(self.env.device), ): yield diff --git a/helion/_compiler/tile_shape_constraints.py b/helion/_compiler/tile_shape_constraints.py new file mode 100644 index 000000000..49b124dee --- /dev/null +++ b/helion/_compiler/tile_shape_constraints.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from .compile_environment import CompileEnvironment + +if TYPE_CHECKING: + from collections.abc import Callable + + +_RESHAPE_OPS = { + torch.ops.aten.reshape.default, + torch.ops.aten.view.default, +} + + +class TileShapeConstraintMode(TorchDispatchMode): + """Record block-size constraints before FakeTensor validates shape ops.""" + + def __torch_dispatch__( + self, + func: Callable[..., object], + types: tuple[type, ...], + args: tuple[object, ...] = (), + kwargs: dict[str, object] | None = None, + ) -> object: + if func in _RESHAPE_OPS and len(args) >= 2: + tensor, shape = args[:2] + if isinstance(tensor, torch.Tensor) and isinstance(shape, (list, tuple)): + if all(isinstance(dim, (int, torch.SymInt)) for dim in shape): + CompileEnvironment.current().prepare_tile_reshape( + tensor.shape, + shape, + ) + return func(*args, **(kwargs or {})) diff --git a/helion/autotuner/config_generation.py b/helion/autotuner/config_generation.py index 9617e9482..1feeb9e03 100644 --- a/helion/autotuner/config_generation.py +++ b/helion/autotuner/config_generation.py @@ -632,7 +632,10 @@ def random_population( except InvalidConfig: attempts += 1 # Retry to fill the population to the requested size - while len(result) < n and attempts < 64: + max_attempts = ( + max(64, n * 256) if self.config_spec.block_size_constraints else 64 + ) + while len(result) < n and attempts < max_attempts: with contextlib.suppress(InvalidConfig): result.append(self.unflatten(self.random_flat())) attempts += 1 diff --git a/helion/autotuner/config_spec.py b/helion/autotuner/config_spec.py index 765831d5a..fe291fe16 100644 --- a/helion/autotuner/config_spec.py +++ b/helion/autotuner/config_spec.py @@ -96,6 +96,14 @@ class TensorNumelConstraint(NamedTuple): expr_str: str +class BlockSizeConstraint(NamedTuple): + """Compile-time condition that concrete block sizes must satisfy.""" + + check_fn: Callable[..., bool] + block_ids: tuple[int, ...] + expr_str: str + + class MatmulFact(NamedTuple): """Shape facts recorded when matmul requirements are applied.""" @@ -605,6 +613,7 @@ def __init__( self.max_num_sm_multiplier: int = MAX_NUM_SM_MULTIPLIER self.grid_block_ids: list[int] = [] self.tensor_numel_constraints: list[TensorNumelConstraint] = [] + self.block_size_constraints: list[BlockSizeConstraint] = [] self.load_eviction_policies = ListOf( EnumFragment(choices=get_valid_eviction_policies(self.backend_name)), length=0, @@ -2007,6 +2016,8 @@ def normalize( for key in _CUTE_IMPLICIT_DEFAULT_KEYS - provided_keys - preserve_keys: config.pop(key, None) + self._validate_block_size_constraints(config) + # Allow tunable parameter keys in addition to backend-supported keys. allowed_keys = self.supported_config_keys() | { *self.user_defined_tunables.keys() @@ -2014,6 +2025,24 @@ def normalize( if invalid_keys := ({*config} - allowed_keys): raise InvalidConfig(f"Invalid config keys {sorted(invalid_keys)!r}") + def _validate_block_size_constraints(self, config: dict[str, object]) -> None: + block_sizes = config.get("block_sizes") + if not isinstance(block_sizes, list): + return + for constraint in self.block_size_constraints: + values = [ + self.block_sizes.config_get(block_sizes, block_id) + for block_id in constraint.block_ids + ] + if any(type(value) is not int for value in values): + raise InvalidConfig( + f"Unable to evaluate block size constraint: {constraint.expr_str}" + ) + if not constraint.check_fn(*cast("list[int]", values)): + raise InvalidConfig( + f"block_sizes violate reshape constraint: {constraint.expr_str}" + ) + def raise_grid_block_minimums(self) -> None: """Raise min_size for grid block dimensions based on problem size. @@ -2097,10 +2126,17 @@ def default_config(self) -> helion.Config: # A promoted seed only specifies the knobs it cares about (e.g. block_sizes); layer it over # the full base defaults so every other key — including user register_tunable defaults — is # preserved rather than dropped. - merged = dict(self._base_default_config().config) + base = self._base_default_config() + merged = dict(base.config) merged.update(self.compiler_default_config.config) config = helion.Config.from_dict(merged) self._shrink_for_numel_constraints(config) + try: + self._validate_block_size_constraints(config.config) + except InvalidConfig: + config.config["block_sizes"] = [*base.block_sizes] + self._shrink_for_numel_constraints(config) + self._validate_block_size_constraints(config.config) return config def _shrink_for_numel_constraints(self, config: helion.Config) -> None: @@ -2448,6 +2484,7 @@ def __init__( bounded_hint = max(bounded_hint, 1) self.min_size: int = min_size self.autotuner_min: int = min_size + self.autotuner_default: int | None = None # Largest power-of-two block that fits inside the dimension. allow_overshoot # may raise max_size above this for matmul dims, but the default block size # stays clamped to dim_max_size (see _fragment). @@ -2529,11 +2566,15 @@ def _fragment(self, base: ConfigSpec) -> BlockSizeFragment: else: default = 1 low = min(max(self.min_size, self.autotuner_min), self.max_size) + if self.autotuner_default is not None: + low = min(low, self.autotuner_default) # Clamp the default within the dimension so allow_overshoot only widens # the autotuner *search*, never the default (non-autotuned) block size. # Needed for matmul dims smaller than the heuristic default (e.g. M<16), # where the default would otherwise overshoot to a masked tile. default = min(default, self.dim_max_size) + if self.autotuner_default is not None: + default = self.autotuner_default return BlockSizeFragment( low, self.max_size, diff --git a/test/test_indexing.py b/test/test_indexing.py index ab90c54ea..efc359870 100644 --- a/test/test_indexing.py +++ b/test/test_indexing.py @@ -21,6 +21,8 @@ from helion._testing import _get_backend from helion._testing import code_and_output from helion._testing import onlyBackends +from helion._testing import skipIfCudaCapabilityLessThan +from helion._testing import skipIfCute from helion._testing import skipIfLowVRAM from helion._testing import skipIfNormalMode from helion._testing import skipIfRefEager @@ -29,6 +31,7 @@ from helion._testing import skipUnlessTensorDescriptor from helion._testing import xfailIfCute from helion._testing import xfailIfPallas +from helion.autotuner.config_generation import ConfigGeneration import helion.language as hl _LARGE_BF16_SHAPE = (51200, 51200) @@ -1354,6 +1357,127 @@ def size1_reshape_kernel( size1_reshape_kernel(x2, out2) torch.testing.assert_close(out2, x2) + @skipIfCudaCapabilityLessThan((9, 0), reason="FP8 requires CUDA capability >= 9.0") + @skipIfCute("CuTe does not support scalar float32 to float8 conversion") + @skipIfRefEager("Test validates compiler block-size constraints") + def test_tile_reshape_for_grouped_quantization(self): + group_size = 128 + + def quantize(block: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + fp8_max = torch.finfo(torch.float8_e4m3fn).max + *lead, last = block.shape + grouped = block.reshape(*lead, last // 128, 128) + amax = grouped.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + scale = amax.to(torch.float32) / fp8_max + quantized = (grouped.to(torch.float32) / scale).to(torch.float8_e4m3fn) + return quantized.reshape(*lead, last), scale.squeeze(-1) + + @helion.kernel(config=helion.Config(block_sizes=[32, group_size], num_warps=4)) + def grouped_quantize(x: torch.Tensor, recipe): + m, n = x.shape + out = torch.empty_like(x, dtype=torch.float8_e4m3fn) + scales = torch.empty((m, n // 128), dtype=torch.float32, device=x.device) + for tile_m, tile_n in hl.tile([m, n]): + out_local, scale_local = recipe(x[tile_m, tile_n]) + out[tile_m, tile_n] = out_local + scales[tile_m, tile_n.begin // 128] = scale_local.squeeze(-1) + return out, scales + + x = torch.randn(64, 256, dtype=torch.bfloat16, device=DEVICE) + bound = grouped_quantize.bind((x, quantize)) + self.assertEqual(bound.config_spec.block_sizes[1].min_size, 1) + self.assertEqual(bound.config_spec.block_sizes[1].autotuner_default, group_size) + self.assertGreaterEqual(len(bound.config_spec.block_size_constraints), 1) + self.assertEqual(bound.config_spec.default_config().block_sizes[1], group_size) + generated_configs = ConfigGeneration(bound.config_spec).random_population(10) + self.assertEqual(len(generated_configs), 10) + self.assertTrue( + all(config.block_sizes[1] == group_size for config in generated_configs) + ) + for invalid_n in (64, 256): + with self.assertRaisesRegex( + helion.exc.InvalidConfig, + "block_sizes violate reshape constraint", + ): + bound.config_spec.normalize( + helion.Config(block_sizes=[32, invalid_n], num_warps=4) + ) + actual_quantized, actual_scales = bound(x, quantize) + expected_quantized, expected_scales = quantize(x) + + torch.testing.assert_close(actual_quantized.float(), expected_quantized.float()) + torch.testing.assert_close(actual_scales, expected_scales) + + @skipIfRefEager("Test validates compiler block-size constraints") + def test_tile_reshape_additive_constraint(self): + @helion.kernel(config=helion.Config(block_sizes=[2, 2])) + def additive_reshape(x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + for tile_m, tile_n in hl.tile(x.shape): + block = x[tile_m, tile_n] + m, n = block.shape + flat = block.reshape(m + n) + out[tile_m, tile_n] = flat.reshape(m, n) + return out + + x = torch.randn(4, 4, device=DEVICE) + bound = additive_reshape.bind((x,)) + self.assertEqual(bound.config_spec.default_config().block_sizes, [2, 2]) + with self.assertRaisesRegex( + helion.exc.InvalidConfig, + "block_sizes violate reshape constraint", + ): + bound.config_spec.normalize(helion.Config(block_sizes=[4, 4])) + torch.testing.assert_close(bound(x), x) + + @skipIfRefEager("Test validates compiler block-size constraints") + def test_tile_reshape_modulo_constraint(self): + @helion.kernel(config=helion.Config(block_sizes=[4, 128])) + def modulo_reshape(x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + for tile_m, tile_n in hl.tile(x.shape): + block = x[tile_m, tile_n] + m, n = block.shape + reshaped = block.reshape(m, n - n % 128) + out[tile_m, tile_n] = reshaped + return out + + x = torch.randn(8, 256, device=DEVICE) + bound = modulo_reshape.bind((x,)) + self.assertEqual( + bound.config_spec.default_config().block_sizes[1] % 128, + 0, + ) + generated_configs = ConfigGeneration(bound.config_spec).random_population(10) + self.assertEqual(len(generated_configs), 10) + self.assertTrue( + all(config.block_sizes[1] % 128 == 0 for config in generated_configs) + ) + with self.assertRaisesRegex( + helion.exc.InvalidConfig, + "block_sizes violate reshape constraint", + ): + bound.config_spec.normalize(helion.Config(block_sizes=[4, 64])) + torch.testing.assert_close(bound(x), x) + + @skipIfRefEager("Test validates compiler block-size constraints") + def test_tile_reshape_unsatisfiable_constraint(self): + @helion.kernel + def invalid_reshape(x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + for tile_m, tile_n in hl.tile(x.shape): + block = x[tile_m, tile_n] + m, n = block.shape + out[tile_m, tile_n] = block.reshape(m + 1, n) + return out + + x = torch.randn(4, 4, device=DEVICE) + with self.assertRaisesRegex( + helion.exc.InvalidConfig, + "No power-of-two block size assignment satisfies reshape constraint", + ): + invalid_reshape.bind((x,)) + def test_size1_dimension_variable_tile_range(self): """Test tile indexing on size-1 dimensions with variable tile ranges.