diff --git a/helion/_compiler/compile_environment.py b/helion/_compiler/compile_environment.py index d45c55695..2f7ad7d2f 100644 --- a/helion/_compiler/compile_environment.py +++ b/helion/_compiler/compile_environment.py @@ -563,23 +563,52 @@ def _extract_tensor_numel_constraints(self) -> None: # VMEM byte budget is enforced separately at runtime. return None + uses_triton_codegen = self.codegen_name == "triton" + cs_block_sizes = self.config_spec.block_sizes + config_block_ids = set(cs_block_sizes.valid_block_ids()) block_sym_to_id: dict[sympy.Symbol, int] = {} - for bs in self.block_sizes: - block_sym_to_id[bs.symbol()] = bs.block_id + block_sym_to_info: dict[sympy.Symbol, BlockSizeInfo] = {} + if uses_triton_codegen: + for info in self.block_sizes: + symbol = info.symbol() + # Reused reduction dimensions deliberately share the original + # tile symbol. Preserve the first origin rather than allowing a + # later fixed/reduction alias to erase its tunable provenance. + block_sym_to_info.setdefault(symbol, info) + if info.block_id in config_block_ids: + block_sym_to_id.setdefault(symbol, info.block_id) + else: + for bs in self.block_sizes: + block_sym_to_id[bs.symbol()] = bs.block_id seen_exprs: set[str] = set() - cs_block_sizes = self.config_spec.block_sizes for shape in self.kernel_tensor_sizes: if not shape: continue numel_expr = sympy.Mul(*shape) if len(shape) > 1 else shape[0] + if uses_triton_codegen: + substitutions: dict[sympy.Basic, sympy.Basic] = {} + for symbol in numel_expr.free_symbols: + if ( + not isinstance(symbol, sympy.Symbol) + or symbol in block_sym_to_id + or (info := block_sym_to_info.get(symbol)) is None + ): + continue + extent = self._search_invariant_extent_for_numel_constraint(info) + if extent is not None: + substitutions[symbol] = sympy.Integer(extent) + if substitutions: + numel_expr = numel_expr.xreplace(substitutions) + all_free = numel_expr.free_symbols involved_syms = all_free & block_sym_to_id.keys() if not involved_syms: continue # Skip expressions with non-block-size free symbols (e.g., # runtime tensor dimensions) — they can't be evaluated at - # config generation time. + # config generation time. A rollable reduction remains symbolic + # here and is skipped by the same rule. if all_free - block_sym_to_id.keys(): log.debug( "skipping numel constraint for shape %s: expression has " @@ -622,6 +651,34 @@ def _extract_tensor_numel_constraints(self) -> None: ) ) + def _search_invariant_extent_for_numel_constraint( + self, block_size: BlockSizeInfo + ) -> int | None: + """Return an extent fixed across the generated Triton search choices.""" + source = block_size.block_size_source + if isinstance(source, FixedBlockSizeSource): + value = source.value + elif isinstance(source, ReductionLoopBlockSizeSource): + reduction_loops = self.config_spec.reduction_loops + if block_size.block_id in reduction_loops.valid_block_ids(): + loop_spec = reduction_loops.block_id_lookup(block_size.block_id) + if loop_spec._flat_fragment(self.config_spec).low < loop_spec.size_hint: + return None + value = block_size.size + else: + return None + if not isinstance(value, (int, torch.SymInt)): + return None + expr = self.specialize_expr(self.shape_env.replace(_to_sympy(value))) + if expr.free_symbols or not expr.is_Integer: + return None + extent = int(expr) + if isinstance(source, ReductionLoopBlockSizeSource): + # Every fragment-generated choice is persistent, so codegen uses + # the full backend-rounded reduction dimension. + extent = self.backend.static_rdim_size(extent) + return extent + def _disable_range_num_stages_for_aliasing(self) -> None: """ Disable range_num_stages choices if any kernel argument name is both read and written. diff --git a/test/test_tensor_numel_constraints.py b/test/test_tensor_numel_constraints.py index 11070242d..62076f5fb 100644 --- a/test/test_tensor_numel_constraints.py +++ b/test/test_tensor_numel_constraints.py @@ -1,15 +1,22 @@ from __future__ import annotations import logging +from types import SimpleNamespace +from typing import cast import unittest import sympy +from helion._compiler.backend import TileIRBackend from helion._compiler.backend import TritonBackend +from helion._compiler.compile_environment import CompileEnvironment +from helion._compiler.compile_environment import FixedBlockSizeSource +from helion._compiler.compile_environment import ReductionLoopBlockSizeSource from helion.autotuner.config_generation import TRITON_MAX_TENSOR_NUMEL from helion.autotuner.config_generation import ConfigGeneration from helion.autotuner.config_spec import BlockSizeSpec from helion.autotuner.config_spec import ConfigSpec +from helion.autotuner.config_spec import ReductionLoopSpec from helion.autotuner.config_spec import TensorNumelConstraint @@ -31,8 +38,6 @@ def _make_constraint( def _check_numel_constraints(gen: ConfigGeneration, flat_config: list[object]) -> bool: """Return True if all tensor numel constraints are satisfied.""" - from typing import cast - for constraint in gen.config_spec.tensor_numel_constraints: args = [ cast("int", flat_config[gen.block_size_indices[i]]) @@ -280,26 +285,182 @@ def test_random_flat_satisfies_constraints(self) -> None: ) -class TestMixedSymbolExtraction(unittest.TestCase): - """Constraints with non-block-size free symbols must be skipped.""" +class TestNumelConstraintExtraction(unittest.TestCase): + """Extract constraints while preserving only config-time symbols.""" - def test_mixed_symbol_shape_skipped(self) -> None: - """A tensor shape like block_size_0 * runtime_dim should not produce - a constraint, because lambdify cannot substitute runtime_dim.""" - from types import SimpleNamespace + @staticmethod + def _mock_env( + backend: TritonBackend | None = None, + ) -> CompileEnvironment: + env = object.__new__(CompileEnvironment) + env._backend = backend or TritonBackend() + env.shape_env = SimpleNamespace(replace=lambda expr: expr) + env.specialized_vars = set() + env.config_spec = ConfigSpec(backend=env._backend) + return env + + def test_config_invariant_reduction_extents_are_substituted(self) -> None: + """Persistent-only reductions contribute to the per-tile tensor numel.""" + b0, b1, r0, r1, r2 = sympy.symbols("b0 b1 r0 r1 r2") + env = self._mock_env() + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=0, size_hint=128)) + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=1, size_hint=128)) + env.config_spec.reduction_loops.append( + ReductionLoopSpec(block_id=2, size_hint=4) + ) + env.block_sizes = [ + SimpleNamespace(symbol=lambda: b0, block_id=0), + SimpleNamespace(symbol=lambda: b1, block_id=1), + SimpleNamespace( + symbol=lambda: r0, + block_id=2, + size=4, + block_size_source=ReductionLoopBlockSizeSource(0), + ), + SimpleNamespace( + symbol=lambda: r1, + block_id=3, + size=32, + block_size_source=ReductionLoopBlockSizeSource(1), + ), + SimpleNamespace( + symbol=lambda: r2, + block_id=4, + size=128, + block_size_source=ReductionLoopBlockSizeSource(2), + ), + ] + env.kernel_tensor_sizes = [(b0, r0, r1, b1, r2)] + + env._extract_tensor_numel_constraints() - from helion._compiler.compile_environment import CompileEnvironment + self.assertEqual(len(env.config_spec.tensor_numel_constraints), 1) + constraint = env.config_spec.tensor_numel_constraints[0] + self.assertEqual(constraint.block_indices, (0, 1)) + self.assertTrue(constraint.check_fn(8, 8)) + self.assertFalse(constraint.check_fn(16, 8)) + gen = ConfigGeneration(env.config_spec) + flat = gen.default_flat() + config_block_indices, _ = gen._key_to_flat_indices["block_sizes"] + self.assertEqual([flat[i] for i in config_block_indices], [8, 8]) + + def test_tileir_substitutes_persistent_reduction_extent(self) -> None: + block, reduction = sympy.symbols("u0 r0") + env = self._mock_env(TileIRBackend()) + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=0, size_hint=128)) + env.config_spec.reduction_loops.append( + ReductionLoopSpec(block_id=1, size_hint=5) + ) + env.block_sizes = [ + SimpleNamespace(symbol=lambda: block, block_id=0), + SimpleNamespace( + symbol=lambda: reduction, + block_id=1, + size=5, + block_size_source=ReductionLoopBlockSizeSource(0), + ), + ] + env.kernel_tensor_sizes = [(block, reduction)] + + env._extract_tensor_numel_constraints() + + self.assertEqual(len(env.config_spec.tensor_numel_constraints), 1) + constraint = env.config_spec.tensor_numel_constraints[0] + self.assertEqual(constraint.expr_str, "8*u0 <= 1048576") + self.assertTrue(constraint.check_fn(131072)) + self.assertFalse(constraint.check_fn(262144)) + + def test_rollable_reduction_skips_constraint(self) -> None: + """A genuinely configurable reduction cannot use a block-only check.""" + b0, b1, r0 = sympy.symbols("b0 b1 r0") + env = self._mock_env() + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=0, size_hint=128)) + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=1, size_hint=128)) + env.config_spec.reduction_loops.append( + ReductionLoopSpec(block_id=2, size_hint=16384) + ) + env.block_sizes = [ + SimpleNamespace(symbol=lambda: b0, block_id=0), + SimpleNamespace(symbol=lambda: b1, block_id=1), + SimpleNamespace( + symbol=lambda: r0, + block_id=2, + size=16384, + block_size_source=ReductionLoopBlockSizeSource(0), + ), + ] + env.kernel_tensor_sizes = [(b0, b1, r0)] + env._extract_tensor_numel_constraints() + + self.assertEqual(env.config_spec.tensor_numel_constraints, []) + + def test_reused_reduction_symbol_preserves_tunable_origin(self) -> None: + """A reduction alias must not overwrite its tunable tile symbol.""" + block = sympy.Symbol("block") + env = self._mock_env() + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=0, size_hint=128)) + env.block_sizes = [ + SimpleNamespace(symbol=lambda: block, block_id=0), + SimpleNamespace( + symbol=lambda: block, + block_id=1, + size=16384, + block_size_source=ReductionLoopBlockSizeSource(0), + ), + ] + env.kernel_tensor_sizes = [(block, sympy.Integer(16384))] + + env._extract_tensor_numel_constraints() + + self.assertEqual(len(env.config_spec.tensor_numel_constraints), 1) + constraint = env.config_spec.tensor_numel_constraints[0] + self.assertEqual(constraint.block_indices, (0,)) + self.assertTrue(constraint.check_fn(64)) + self.assertFalse(constraint.check_fn(128)) + + def test_fixed_block_extent_is_substituted(self) -> None: + block, fixed = sympy.symbols("block fixed") + env = self._mock_env() + env.config_spec.block_sizes.append(BlockSizeSpec(block_id=0, size_hint=128)) + env.block_sizes = [ + SimpleNamespace(symbol=lambda: block, block_id=0), + SimpleNamespace( + symbol=lambda: fixed, + block_id=1, + size=128, + block_size_source=FixedBlockSizeSource(128), + ), + ] + env.kernel_tensor_sizes = [(block, fixed, sympy.Integer(128))] + + env._extract_tensor_numel_constraints() + + self.assertEqual(len(env.config_spec.tensor_numel_constraints), 1) + constraint = env.config_spec.tensor_numel_constraints[0] + self.assertTrue(constraint.check_fn(64)) + self.assertFalse(constraint.check_fn(128)) + + def test_mixed_symbol_shape_skipped(self) -> None: + """A tensor shape like block_size_0 * runtime_dim should not produce + a constraint because it cannot be evaluated during config generation.""" b0 = sympy.Symbol("u0", integer=True) runtime_dim = sympy.Symbol("u8", integer=True) # Minimal mock of CompileEnvironment for _extract_tensor_numel_constraints env = object.__new__(CompileEnvironment) - env._backend = SimpleNamespace(max_tensor_numel=TRITON_MAX_TENSOR_NUMEL) + env._backend = SimpleNamespace( + name="triton", + codegen_name="triton", + max_tensor_numel=TRITON_MAX_TENSOR_NUMEL, + ) + env.shape_env = SimpleNamespace(replace=lambda expr: expr) + env.specialized_vars = set() env.block_sizes = [SimpleNamespace(symbol=lambda: b0, block_id=0)] env.config_spec = SimpleNamespace( block_sizes=SimpleNamespace( block_id_to_index=lambda bid: bid, + valid_block_ids=lambda: [0], ), tensor_numel_constraints=[], ) @@ -324,10 +485,6 @@ class TestBackendMaxTensorNumel(unittest.TestCase): def test_backend_with_no_cap_extracts_no_constraints(self) -> None: """When backend.max_tensor_numel is None, _extract_tensor_numel_constraints emits no constraints regardless of how large the tile would be.""" - from types import SimpleNamespace - - from helion._compiler.compile_environment import CompileEnvironment - b0 = sympy.Symbol("u0", integer=True) env = object.__new__(CompileEnvironment) env._backend = SimpleNamespace(max_tensor_numel=None)