Rename TileStrategy.get_block_index to CompileEnvironment.get_block_id (#151)

Author: jansel

helion/_compiler/compile_environment.py

@@ -74,11 +74,9 @@ def __init__(self, device: torch.device, settings: Settings) -> None:
         self.loop_dependency_checker = LoopDependencyChecker()
 
     def add_kernel_tensor_size(self, sizes: Sequence[int | torch.SymInt]) -> None:
-        from .tile_strategy import TileStrategy
-
         for size in sizes:
             if isinstance(size, torch.SymInt):
-                block_idx = TileStrategy.get_block_index(size)
+                block_idx = self.get_block_id(size)
                 if block_idx is None:
                     value = self.shape_env.replace(size._sympy_())
                     if value.free_symbols:
@@ -315,6 +313,33 @@ def has_current() -> bool:
         except NoCurrentEnvironment:
             return False
 
+    def get_block_id(self, size: int | torch.SymInt | sympy.Expr) -> int | None:
+        """
+        Get the block ID associated with a given size expression.
+
+        This method determines if a size expression corresponds to a registered block size
+        in the current compilation environment. It looks up the origin information of
+        symbolic expressions to find their associated block IDs.
+
+        Args:
+            size: The size expression to check. Can be an integer, torch.SymInt, or sympy.Expr.
+
+        Returns:
+            The block ID if the size corresponds to a registered block size, None otherwise.
+        """
+        if isinstance(size, torch.SymInt):
+            return self.get_block_id(size._sympy_())
+        if isinstance(size, sympy.Symbol):
+            from .host_function import HostFunction
+
+            origin_info = HostFunction.current().expr_to_origin.get(size)
+            if origin_info is not None and isinstance(
+                origin_info.origin,
+                BlockSizeOrigin,
+            ):
+                return origin_info.origin.block_id
+        return None
+
 
 class NoCurrentEnvironment(RuntimeError):
     pass

helion/_compiler/device_function.py

@@ -29,7 +29,6 @@
 from .host_function import HostFunction
 from .host_function import NoCurrentFunction
 from .output_header import reserved_names
-from .tile_strategy import TileStrategy
 from .variable_origin import BlockSizeOrigin
 from .variable_origin import Origin
 from .variable_origin import TensorSizeOrigin
@@ -209,7 +208,7 @@ def user_sympy_expr(self, expr: sympy.Expr) -> str:
         replacements = {}
         for sym in sorted(expr.free_symbols, key=lambda s: s.name):
             assert isinstance(sym, sympy.Symbol)
-            block_idx = TileStrategy.get_block_index(sym)
+            block_idx = CompileEnvironment.current().get_block_id(sym)
             if block_idx is not None:
                 replacements[sym] = self.tile_strategy.user_size(block_idx)
         if replacements:

helion/_compiler/indexing_strategy.py

@@ -14,7 +14,6 @@
 from .compile_environment import CompileEnvironment
 from .host_function import HostFunction
 from .tile_strategy import DeviceLoopState
-from .tile_strategy import TileStrategy
 from .variable_origin import BlockSizeOrigin
 
 if TYPE_CHECKING:
@@ -207,6 +206,7 @@ def compute_shape(
         assert isinstance(index, (list, tuple)), index
         input_size = collections.deque(tensor.size())
         output_size = []
+        env = CompileEnvironment.current()
         for k in index:
             if k is None:
                 output_size.append(1)
@@ -218,11 +218,7 @@ def compute_shape(
                 if isinstance(symbol, sympy.Symbol):
                     origin = HostFunction.current().expr_to_origin.get(symbol)
                     if origin and isinstance(origin.origin, BlockSizeOrigin):
-                        if (
-                            CompileEnvironment.current()
-                            .block_sizes[origin.origin.block_id]
-                            .is_grid()
-                        ):
+                        if env.block_sizes[origin.origin.block_id].is_grid():
                             pass
                         elif tensor.size(tensor.ndim - len(input_size) - 1) != 1:
                             output_size.append(k)
@@ -231,9 +227,7 @@ def compute_shape(
             elif isinstance(k, slice) and str(k) == "slice(None, None, None)":
                 size = input_size.popleft()
                 if size != 1:
-                    rdim = CompileEnvironment.current().allocate_reduction_dimension(
-                        size
-                    )
+                    rdim = env.allocate_reduction_dimension(size)
                     output_size.append(rdim.var)
                 else:
                     output_size.append(1)
@@ -308,9 +302,7 @@ def create(
                 assert len(ast_index) == len(index)
                 index_var = state.codegen.lift(ast_index[n], prefix="index").id
                 index_values.append(f"({index_var}){expand}")
-                if (
-                    block_idx := TileStrategy.get_block_index(output_size[output_idx])
-                ) is not None:
+                if (block_idx := env.get_block_id(output_size[output_idx])) is not None:
                     if mask := state.codegen.mask_var(block_idx):
                         mask_values.setdefault(f"({mask}){expand}")
                 output_idx += 1
@@ -325,7 +317,7 @@ def create(
                 index_values.append(index_var)
                 output_idx += k.ndim
                 for n, s in enumerate(output_size):
-                    if (block_idx := TileStrategy.get_block_index(s)) is not None and (
+                    if (block_idx := env.get_block_id(s)) is not None and (
                         mask := state.codegen.mask_var(block_idx)
                     ):
                         mask_values.setdefault(

helion/_compiler/inductor_lowering.py

@@ -51,7 +51,6 @@
 from .node_masking import getitem_masked_value
 from .node_masking import inductor_masked_value
 from .node_masking import mask_node_inputs
-from .tile_strategy import TileStrategy
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -429,7 +428,7 @@ def __init__(
         reduction_var = reduction_ranges[0]
         assert isinstance(reduction_var, sympy.Symbol)
 
-        block_index = TileStrategy.get_block_index(reduction_var)
+        block_index = CompileEnvironment.current().get_block_id(reduction_var)
         assert block_index is not None
         self.block_index: int = block_index
 
@@ -497,7 +496,7 @@ def codegen(self, ctx: GraphInterpreter, node: torch.fx.Node) -> object:
         (dim,) = [
             i
             for i, v in enumerate(repr_input.shape)
-            if TileStrategy.get_block_index(v) == self.block_index
+            if CompileEnvironment.current().get_block_id(v) == self.block_index
         ]
 
         return strategy.codegen_reduction(
@@ -764,7 +763,7 @@ def apply_dot_requirements(
     a, b, c = min_dot_size(lproxy.device, lproxy.dtype, rproxy.dtype)
     env = CompileEnvironment.current()
     for shape, min_size in [(n, a), (k, b), (m, c)]:
-        block_idx = TileStrategy.get_block_index(shape)
+        block_idx = CompileEnvironment.current().get_block_id(shape)
         if block_idx is not None:
             env.block_sizes[block_idx].update_min_block(min_size, allow_flattened=True)
     # inputs to the dot operation must be zero-masked

helion/_compiler/reduction_strategy.py

@@ -177,8 +177,6 @@ def codegen_preamble(self, state: CodegenState) -> None:
                 f"{mask_var} = {index_var} < {self.fn.sympy_expr(numel)}"
             )
         # Extract end_var_name from the numel expression
-        env = CompileEnvironment.current()
-        numel = env.block_sizes[self.block_index].numel
         end_var_name = {self.block_index: self.fn.sympy_expr(numel)}
         state.codegen.set_active_loops(
             PersistentReductionState(self, end_var_name=end_var_name)

helion/_compiler/roll_reduction.py

@@ -6,10 +6,10 @@
 import torch
 from torch.fx import map_arg
 
+from helion._compiler.compile_environment import CompileEnvironment
 from helion._compiler.inductor_lowering import APIFuncLowering
 from helion._compiler.inductor_lowering import ReductionLowering
 from helion._compiler.inductor_lowering import aten_lowering_dispatch
-from helion._compiler.tile_strategy import TileStrategy
 from helion.language._tracing_ops import _for_loop
 from helion.language._tracing_ops import _get_symnode
 from helion.language._tracing_ops import _host_tensor
@@ -102,7 +102,7 @@ def should_go_in_inner_graph(self, node: torch.fx.Node) -> bool:
         num_rdims = 0
         if isinstance(val, torch.Tensor):
             for size in val.size():
-                block_idx = TileStrategy.get_block_index(size)
+                block_idx = CompileEnvironment.current().get_block_id(size)
                 num_rdims += block_idx == self.rdim.block_id
             if num_rdims > 1:
                 raise NotImplementedError(
@@ -113,7 +113,7 @@ def should_go_in_inner_graph(self, node: torch.fx.Node) -> bool:
             for item in val:
                 if isinstance(item, torch.Tensor):
                     for size in item.size():
-                        block_idx = TileStrategy.get_block_index(size)
+                        block_idx = CompileEnvironment.current().get_block_id(size)
                         num_rdims += block_idx == self.rdim.block_id
             if num_rdims > 1:
                 raise NotImplementedError(
@@ -263,7 +263,7 @@ def is_matmul_with_rdim(node: torch.fx.Node) -> bool:
                 val = input_node.meta.get("val", None)
                 if isinstance(val, torch.Tensor):
                     for size in val.size():
-                        block_idx = TileStrategy.get_block_index(size)
+                        block_idx = CompileEnvironment.current().get_block_id(size)
                         if block_idx == self.rdim.block_id:
                             return True
             return False

helion/_compiler/tile_dispatch.py

@@ -125,7 +125,7 @@ def codegen_device_loop(
     def _compact_shape(self, shapes: ShapeLike) -> list[CompactedShape]:
         compacted_shapes = []
         for idx, shape in enumerate(shapes):
-            block_idx = TileStrategy.get_block_index(shape)
+            block_idx = CompileEnvironment.current().get_block_id(shape)
             if block_idx is None:
                 compacted_shapes.append(
                     CompactedShape(self.strategies[0].fn.literal_expr(shape), [idx], [])

helion/_compiler/tile_strategy.py

@@ -26,7 +26,6 @@
 from .program_id import ProgramIDs
 from .program_id import SharedProgramID
 from .program_id import VirtualProgramIDs
-from .variable_origin import BlockSizeOrigin
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -119,19 +118,6 @@ def codegen_preamble(self, state: CodegenState) -> None:
     def compact_shape(self, shapes: list[CompactedShape]) -> list[CompactedShape]:
         raise NotImplementedError
 
-    @classmethod
-    def get_block_index(cls, size: int | torch.SymInt | sympy.Expr) -> int | None:
-        if isinstance(size, torch.SymInt):
-            return cls.get_block_index(size._sympy_())
-        if isinstance(size, sympy.Symbol):
-            origin_info = HostFunction.current().expr_to_origin.get(size)
-            if origin_info is not None and isinstance(
-                origin_info.origin,
-                BlockSizeOrigin,
-            ):
-                return origin_info.origin.block_id
-        return None
-
 
 class BlockSizeTileStrategy(TileStrategy):
     def __init__(
@@ -260,7 +246,8 @@ def codegen_grid(self, state: CodegenState) -> DeviceGridState:
         block_size_var, offsets_var, total_numel, statements = self._codegen_common(
             state
         )
-        dtype = CompileEnvironment.current().triton_index_type()
+        env = CompileEnvironment.current()
+        dtype = env.triton_index_type()
         state.add_statement(
             f"{offsets_var} = tl.program_id(0) * ({block_size_var}) + tl.arange(0, {block_size_var}).to({dtype})"
         )
@@ -276,7 +263,7 @@ def codegen_grid(self) -> ast.AST:
 
         end_var_name = {}
         for block_id in self.block_ids:
-            end_bound = CompileEnvironment.current().block_sizes[block_id].numel
+            end_bound = env.block_sizes[block_id].numel
             end_var_name[block_id] = state.device_function.sympy_expr(end_bound)
         return DeviceGridState(self, end_var_name=end_var_name)
 
@@ -313,12 +300,13 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
 
     @classmethod
     def update_allow_flattened(cls, shape: Sequence[sympy.Expr]) -> None:
+        env = CompileEnvironment.current()
         used_indices = {}
         for i, x in enumerate(shape):
-            block_idx = cls.get_block_index(x)
+            block_idx = env.get_block_id(x)
             if block_idx is not None:
                 used_indices[block_idx] = i
-        flatten_loops = CompileEnvironment.current().config_spec.flatten_loops
+        flatten_loops = env.config_spec.flatten_loops
         for spec in [*flatten_loops]:
             block_ids = spec.block_ids
             if not (
@@ -405,7 +393,7 @@ def codegen_grid(self, state: CodegenState) -> DeviceGridState:
                 )
             else:
                 block_size_var = "1"
-                dtype = CompileEnvironment.current().triton_index_type()
+                dtype = env.triton_index_type()
                 state.add_statement(f"{offset_var} = {pid_var}")
                 state.add_statement(
                     f"{index_var} = {offset_var} + tl.zeros([1], {dtype})"
@@ -428,7 +416,7 @@ def codegen_grid(self, state: CodegenState) -> DeviceGridState:
         # Extract end_var_name from end bound expressions
         end_var_name = {}
         for block_id in self.block_ids:
-            end_bound = CompileEnvironment.current().block_sizes[block_id].numel
+            end_bound = env.block_sizes[block_id].numel
             end_var_name[block_id] = state.device_function.sympy_expr(end_bound)
         return DeviceGridState(self, end_var_name=end_var_name)
 

helion/language/_tracing_ops.py

@@ -13,7 +13,6 @@
 from .._compiler.compile_environment import CompileEnvironment
 from .._compiler.host_function import HostFunction
 from .._compiler.tile_index_proxy import TileIndexProxy
-from .._compiler.tile_strategy import TileStrategy
 from ..exc import NotInsideKernel
 from . import _decorators
 
@@ -39,7 +38,7 @@ def _get_symnode(debug_name: str) -> int:
 def _(state: CodegenState) -> ast.AST:
     val = state.fx_node.meta["val"]
     assert isinstance(val, (torch.SymInt, torch.SymFloat, torch.SymBool)), val
-    if (block_idx := TileStrategy.get_block_index(val)) is not None:
+    if (block_idx := CompileEnvironment.current().get_block_id(val)) is not None:
         if state.device_function.block_size_var(block_idx) is None:
             # this should be unused
             return expr_from_string("block_size_var_optimized_away")
@@ -261,7 +260,7 @@ def _(state: CodegenState) -> ast.AST:
     mask_exprs = []
     input_sizes = [*tensor.size()]
     for dim, size in enumerate(input_sizes):
-        if (index := TileStrategy.get_block_index(size)) is not None and (
+        if (index := CompileEnvironment.current().get_block_id(size)) is not None and (
             mask_var := state.codegen.mask_var(index)
         ) is not None:
             expand = state.tile_strategy.expand_str(input_sizes, dim)

helion/language/loops.py

@@ -211,9 +211,9 @@ def _(
         elif isinstance(bs, int):
             results.append(TileIndexType.allocate_fixed(size, bs, origin))
         elif isinstance(bs, torch.SymInt):
-            from helion._compiler.tile_strategy import TileStrategy
+            from helion._compiler.compile_environment import CompileEnvironment
 
-            index = TileStrategy.get_block_index(bs)
+            index = CompileEnvironment.current().get_block_id(bs)
             if index is None:
                 results.append(TileIndexType.allocate_fixed(size, bs, origin))
             else: