Support in-tile reshapes for symbolic block sizes

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

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

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 {}))

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

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,13 +2016,33 @@ 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()
         }
         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,