Fix bug with computations based on hl.register_block_size (#157)

Author: jansel

helion/_compiler/device_function.py

@@ -152,6 +152,7 @@ def __init__(self, name: str, config: Config) -> None:
         self._variable_renames: dict[str, list[str]] = {}
         self.dce_vars: list[str] = []
         self.block_size_var_cache: dict[tuple[int, ...], str] = {}
+        self.expr_to_var_name: dict[sympy.Expr, str] = {}
 
         from .indexing_strategy import IndexingStrategy
         from .tile_dispatch import TileStrategyDispatch
@@ -175,17 +176,26 @@ def set_pid(self, pid: SharedProgramID | ProgramIDs) -> None:
         self.pid = pid
 
     def sympy_expr(self, expr: sympy.Expr) -> str:
-        expr_to_origin = HostFunction.current().expr_to_origin
         expr = CompileEnvironment.current().shape_env.simplify(expr)
         if not expr.free_symbols:
             return texpr(expr)
+        if expr in self.expr_to_var_name:
+            return self.expr_to_var_name[expr]
+        expr_to_origin = HostFunction.current().expr_to_origin
         if expr in expr_to_origin:
             return self._lift_sympy_arg(expr)
         replacements = {}
         for sym in sorted(expr.free_symbols, key=lambda x: x.name):
             assert isinstance(sym, sympy.Symbol)
-            assert sym in expr_to_origin, f"no origin found for {sym.name}"
-            replacements[sym] = sympy.Symbol(self._lift_sympy_arg(sym), integer=True)
+            if sym in self.expr_to_var_name:
+                replacements[sym] = sympy.Symbol(
+                    self.expr_to_var_name[sym], integer=True
+                )
+            else:
+                assert sym in expr_to_origin, f"no origin found for {sym.name}"
+                replacements[sym] = sympy.Symbol(
+                    self._lift_sympy_arg(sym), integer=True
+                )
         return texpr(expr.xreplace(replacements))
 
     def _lift_sympy_arg(self, expr: sympy.Expr) -> str:

helion/_compiler/generate_ast.py

@@ -80,11 +80,15 @@ def set_statements(self, new_statements: list[ast.AST] | None) -> Iterator[None]
         if new_statements is None:
             yield
         else:
+            expr_to_var_name = self.device_function.expr_to_var_name
+            # We don't want to reuse vars assigned in a nested scope, so copy it
+            self.device_function.expr_to_var_name = expr_to_var_name.copy()
             self.statements_stack.append(new_statements)
             try:
                 yield
             finally:
                 self.statements_stack.pop()
+                self.device_function.expr_to_var_name = expr_to_var_name
 
     @contextlib.contextmanager
     def set_on_device(self) -> Iterator[None]:

helion/_compiler/inductor_lowering.py

@@ -927,13 +927,26 @@ def run_node(self, n: Node) -> object:
                     return result
                 assert isinstance(result, ast.expr)
                 if len(n.users) > 0:
-                    if isinstance(result, (ast.Name, ast.Constant)):
-                        return result
-                    name = self.cg.device_function.new_var(n.name)
-                    self.cg.add_statement(
-                        statement_from_string(f"{name} = result", result=result)
-                    )
-                    return create(ast.Name, id=name, ctx=ast.Load())
+                    if not isinstance(result, (ast.Name, ast.Constant)):
+                        name = self.cg.device_function.new_var(n.name)
+                        self.cg.add_statement(
+                            statement_from_string(f"{name} = result", result=result)
+                        )
+                        result = create(ast.Name, id=name, ctx=ast.Load())
+                    if (
+                        isinstance(val := n.meta["val"], torch.SymInt)
+                        and len((expr := val._sympy_()).free_symbols) > 0
+                    ):
+                        # Keep track of what variable symints are stored in to support DeviceFunction.sympy_expr()
+                        expr = CompileEnvironment.current().shape_env.simplify(expr)
+                        if isinstance(result, ast.Name):
+                            self.cg.device_function.expr_to_var_name[expr] = result.id
+                        else:
+                            assert isinstance(result, ast.Constant)
+                            self.cg.device_function.expr_to_var_name[expr] = repr(
+                                result.value
+                            )
+                    return result
                 if not isinstance(result, (ast.Name, ast.Constant)):
                     self.cg.add_statement(create(ast.Expr, value=result))
                 return None
@@ -997,3 +1010,6 @@ def config(self) -> Config:
 
     def add_statement(self, statement: ast.AST | str) -> None:
         return self.codegen.add_statement(statement)
+
+    def sympy_expr(self, expr: sympy.Expr) -> str:
+        return self.codegen.device_function.sympy_expr(expr)

helion/_compiler/program_id.py

@@ -28,7 +28,7 @@ def host_cdiv(self) -> str:
         return f"triton.cdiv({numel_str}, {self.block_size_var})"
 
     def device_cdiv(self, state: CodegenState) -> str:
-        numel_str = state.device_function.sympy_expr(self.numel)
+        numel_str = state.sympy_expr(self.numel)
         if self.block_size_var == "1":
             return numel_str
         return f"tl.cdiv({numel_str}, {self.block_size_var})"

helion/_compiler/reduction_strategy.py

@@ -227,7 +227,6 @@ def __init__(
     def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
         env = CompileEnvironment.current()
         block_index = self.block_index
-        device_function = state.device_function
         numel = env.block_sizes[block_index].numel
         offset_var = self.offset_var(block_index)
         index_var = self.index_var(block_index)
@@ -245,21 +244,21 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
         if (mask_var := self._mask_var) is not None:
             body.append(
                 statement_from_string(
-                    f"{mask_var} = {index_var} < {device_function.sympy_expr(numel)}"
+                    f"{mask_var} = {index_var} < {state.sympy_expr(numel)}"
                 )
             )
         for_node = create(
             ast.For,
             target=create(ast.Name, id=offset_var, ctx=ast.Store()),
             iter=expr_from_string(
-                f"range(0, ({device_function.sympy_expr(numel)}), {block_size_var})"
+                f"range(0, ({state.sympy_expr(numel)}), {block_size_var})"
             ),
             body=body,
             orelse=[],
             type_comment=None,
         )
         # Extract end_var_name from the actual numel expression used in the range()
-        end_var_name = {block_index: device_function.sympy_expr(numel)}
+        end_var_name = {block_index: state.sympy_expr(numel)}
         return DeviceLoopState(
             self,
             for_node=for_node,

helion/_compiler/tile_strategy.py

@@ -212,7 +212,6 @@ def _codegen_common(
         block_ids = self.block_ids
         env = CompileEnvironment.current()
         total_numel = sympy.S.One
-        device_function = state.device_function
         offsets_var = self.new_var("offsets", dce=True)
         block_size_var = self.block_size_var(-1)
         statements = []
@@ -227,17 +226,17 @@ def _codegen_common(
             block_index_var = self.index_var(block_idx)
             expr = offsets_var
             if total_numel != sympy.S.One:
-                expr = f"({expr}) // ({device_function.sympy_expr(total_numel)})"
+                expr = f"({expr}) // ({state.sympy_expr(total_numel)})"
             if i + 1 < len(block_ids):
-                expr = f"({expr}) % ({device_function.sympy_expr(numel)})"
+                expr = f"({expr}) % ({state.sympy_expr(numel)})"
             statements.append(statement_from_string(f"{block_index_var} = {expr}"))
             total_numel = total_numel * numel
 
         mask_var = self.mask_var(-1)
         if mask_var is not None:
             statements.append(
                 statement_from_string(
-                    f"{mask_var} = {offsets_var} < ({device_function.sympy_expr(total_numel)})"
+                    f"{mask_var} = {offsets_var} < ({state.sympy_expr(total_numel)})"
                 )
             )
         return block_size_var, offsets_var, total_numel, statements
@@ -264,7 +263,7 @@ def codegen_grid(self) -> ast.AST:
         end_var_name = {}
         for block_id in self.block_ids:
             end_bound = env.block_sizes[block_id].numel
-            end_var_name[block_id] = state.device_function.sympy_expr(end_bound)
+            end_var_name[block_id] = state.sympy_expr(end_bound)
         return DeviceGridState(self, end_var_name=end_var_name)
 
     def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
@@ -277,7 +276,7 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
             ast.For,
             target=create(ast.Name, id=lid, ctx=ast.Store()),
             iter=expr_from_string(
-                f"range(tl.cdiv({state.device_function.sympy_expr(total_numel)}, {block_size_var}))"
+                f"range(tl.cdiv({state.sympy_expr(total_numel)}, {block_size_var}))"
             ),
             body=(
                 body := [
@@ -417,7 +416,7 @@ def codegen_grid(self, state: CodegenState) -> DeviceGridState:
         end_var_name = {}
         for block_id in self.block_ids:
             end_bound = env.block_sizes[block_id].numel
-            end_var_name[block_id] = state.device_function.sympy_expr(end_bound)
+            end_var_name[block_id] = state.sympy_expr(end_bound)
         return DeviceGridState(self, end_var_name=end_var_name)
 
     def select_pid_strategy(self) -> ProgramIDs:

helion/_compiler/variable_origin.py

@@ -237,6 +237,9 @@ def host_str(self) -> str:
             return "1"
         return var
 
+    def suggest_var_name(self) -> str:
+        return f"block_size_{self.block_id}"
+
 
 @dataclasses.dataclass
 class ReductionDimensionOrigin(Origin):

helion/language/_tracing_ops.py

@@ -43,7 +43,7 @@ def _(state: CodegenState) -> ast.AST:
             # this should be unused
             return expr_from_string("block_size_var_optimized_away")
     return state.codegen.lift(
-        expr_from_string(state.device_function.sympy_expr(val._sympy_())),
+        expr_from_string(state.sympy_expr(val._sympy_())),
         dce=True,
         prefix="symnode",
     )

helion/language/loops.py

@@ -8,6 +8,7 @@
 from typing import overload
 
 import torch
+from torch._inductor.codegen.simd import constant_repr
 from torch._inductor.runtime.runtime_utils import next_power_of_2
 from torch._inductor.runtime.triton_heuristics import get_max_y_grid
 
@@ -374,7 +375,7 @@ def _(state: CodegenState) -> ast.AST:
 
 
 @_decorators.api(is_device_only=False, cache_type=True, tiles_as_sizes=True)
-def register_block_size(min_or_max: int, max_or_none: int | None = None, /) -> Tile:
+def register_block_size(min_or_max: int, max_or_none: int | None = None, /) -> int:
     """
     Explicitly register a block size that should be autotuned and can be used for
     allocations and inside hl.tile(..., block_size=...).
@@ -396,6 +397,8 @@ def register_block_size(min_or_max: int, max_or_none: int | None = None, /) -> T
 def _(
     min_or_max: TypeInfo, max_or_none: TypeInfo | None = None, /, *, origin: Origin
 ) -> TypeInfo:
+    from .._compiler.type_propagation import SymIntType
+
     min_type, max_type = _normalize_begin_end(min_or_max, max_or_none, origin=origin)
     min_proxy = _to_proxy(min_type)
     max_proxy = _to_proxy(max_type)
@@ -412,13 +415,43 @@ def _(
     loop_spec = env.config_spec.block_sizes.block_id_lookup(result.block_id)
     loop_spec.min_size = assert_integer_power_of_two(max(1, min_proxy))
     loop_spec.max_size = next_power_of_2(env.size_hint(max_proxy))
-    return result
+    block_id = result.block_id
+    return SymIntType(origin, env.block_sizes[block_id].var)
+
+
+def _block_id_from_state(state: CodegenState) -> int:
+    """Extract the block_id from the current state for nodes hl.register_block_size."""
+    from .._compiler.type_propagation import SymIntType
+
+    env = CompileEnvironment.current()
+    if state.fx_node is not None:
+        val = state.fx_node.meta["val"]
+        assert isinstance(val, SymIntType)
+        block_id = env.get_block_id(val.value)
+        assert block_id is not None
+        return block_id
+    current_node = ExtendedAST.current()[-1]
+    type_info = current_node._type_info
+    assert isinstance(type_info, SymIntType)
+    block_id = env.get_block_id(type_info.value)
+    assert block_id is not None
+    return block_id
+
+
+@_decorators.codegen(register_block_size)
+def _(state: CodegenState) -> ast.AST:
+    env = CompileEnvironment.current()
+    block_size = env.config_spec.block_sizes.config_get(
+        state.config.block_sizes, _block_id_from_state(state)
+    )
+    assert block_size is not None
+    return expr_from_string(constant_repr(block_size))
 
 
 @_decorators.api(is_device_only=False, cache_type=True, tiles_as_sizes=True)
 def register_reduction_dim(
     size: int,
-) -> torch.SymInt:
+) -> int:
     """
     Explicitly register a reduction dimension that should be used for reduction operations.
 
@@ -432,21 +465,10 @@ def register_reduction_dim(
     raise exc.NotInsideKernel
 
 
-@_decorators.register_fake(register_reduction_dim)
-def _(size: int) -> torch.SymInt:
-    """Fake implementation that returns the registered reduction dimension size(s)"""
-    from .._compiler.compile_environment import CompileEnvironment
-
-    env = CompileEnvironment.current()
-
-    rdim = env.allocate_reduction_dimension(size)
-    return rdim.var
-
-
 @_decorators.type_propagation(register_reduction_dim)
 def _(sizes: TypeInfo, *, origin: Origin) -> TypeInfo:
     from .._compiler.compile_environment import CompileEnvironment
-    from .._compiler.type_propagation import ReductionDimType
+    from .._compiler.type_propagation import SymIntType
 
     try:
         proxy_sizes = sizes.proxy()
@@ -464,16 +486,16 @@ def _(sizes: TypeInfo, *, origin: Origin) -> TypeInfo:
     env = CompileEnvironment.current()
 
     rdim = env.allocate_reduction_dimension(proxy_sizes)
-    return ReductionDimType(origin, rdim.block_id)
+    return SymIntType(origin, rdim.var)
 
 
 @_decorators.codegen(register_reduction_dim)
 def _(state: CodegenState) -> ast.AST:
     """Generate code for register_reduction_dim - return the size expression"""
-    from .._compiler.type_propagation import ReductionDimType
+    from .._compiler.type_propagation import SymIntType
 
     current_node = ExtendedAST.current()[-1]
     type_info = current_node._type_info
 
-    assert isinstance(type_info, ReductionDimType)
+    assert isinstance(type_info, SymIntType)
     return current_node.args[0]  # pyre-ignore[16]

helion/language/tiles.py

@@ -58,7 +58,7 @@ def tile_begin(tile: Tile) -> int:
 
 
 @_decorators.register_fake(tile_begin)
-def _(state: CodegenState) -> torch.SymInt:
+def _(tile: torch.SymInt) -> torch.SymInt:
     return CompileEnvironment.current().create_unbacked_symint()
 
 
@@ -93,7 +93,7 @@ def tile_end(tile: Tile) -> int:
 
 
 @_decorators.register_fake(tile_end)
-def _(state: CodegenState) -> torch.SymInt:
+def _(tile: torch.SymInt) -> torch.SymInt:
     return CompileEnvironment.current().create_unbacked_symint()