Rename block_size_idx to block_id (#134)

Some minor refactoring to make naming more consistent

Author: jansel

helion/_compiler/compile_environment.py

@@ -101,7 +101,7 @@ def allocate_block_size(
         idx = len(self.block_sizes)
         self.block_sizes.append(
             info := BlockSizeInfo(
-                block_size_idx=idx,
+                block_id=idx,
                 size=size,
                 var=self.create_block_var(
                     f"block_size_{idx}" if not reduction else f"rdim_{idx}",
@@ -310,7 +310,7 @@ class BlockSizeInfo:
     Used to track the block size for a given dimension.
     """
 
-    block_size_idx: int
+    block_id: int
     size: torch.SymInt | int | AutoSize | None
     var: torch.SymInt
     reduction: bool
@@ -348,7 +348,7 @@ def mark_alternate_size(self, size: torch.SymInt | int | None) -> None:
                 with contextlib.suppress(KeyError):
                     # update the size hint now that we know the size
                     env.config_spec.block_sizes.block_id_lookup(
-                        self.block_size_idx
+                        self.block_id
                     ).update_hint(env.size_hint(size))
         elif size is None or self.size is None or self.size != size:
             self.size = None
@@ -357,7 +357,7 @@ def symbol(self) -> sympy.Symbol:
         return self.var._sympy_()
 
     def from_config(self, config: Config) -> int | torch.SymInt | None:
-        return self.block_size_source.from_config(config, self.block_size_idx)
+        return self.block_size_source.from_config(config, self.block_id)
 
     def from_config_assert(self, config: Config) -> int | torch.SymInt:
         val = self.from_config(config)
@@ -366,19 +366,17 @@ def from_config_assert(self, config: Config) -> int | torch.SymInt:
 
     def is_flattened(self, config: Config) -> bool:
         spec = CompileEnvironment.current().config_spec
-        return spec.flatten_loops.config_get(
-            config.flatten_loops, self.block_size_idx, False
-        )
+        return spec.flatten_loops.config_get(config.flatten_loops, self.block_id, False)
 
     def is_grid(self) -> bool:
         return self.block_size_source.is_grid()
 
     def update_min_block(self, value: int, *, allow_flattened: bool = True) -> None:
         spec = CompileEnvironment.current().config_spec
         if not allow_flattened:
-            spec.flatten_loops.disable_block_id(self.block_size_idx)
+            spec.flatten_loops.disable_block_id(self.block_id)
         with contextlib.suppress(KeyError):
-            spec.block_sizes.block_id_lookup(self.block_size_idx).update_min(value)
+            spec.block_sizes.block_id_lookup(self.block_id).update_min(value)
 
 
 class BlockSizeSource:

helion/_compiler/device_function.py

@@ -158,8 +158,8 @@ def __init__(self, name: str, config: Config) -> None:
         self.tile_strategy: TileStrategyDispatch = TileStrategyDispatch(self, config)
         self.indexing_strategy: IndexingStrategy = IndexingStrategy.select(config)
 
-    def block_size_var(self, block_size_idx: int) -> str | None:
-        return self.block_size_var_cache.get((block_size_idx,))
+    def block_size_var(self, block_id: int) -> str | None:
+        return self.block_size_var_cache.get((block_id,))
 
     def merge_variable_names(self, a: str, b: str) -> None:
         name_group = [
@@ -197,7 +197,7 @@ def _lift_sympy_arg(self, expr: sympy.Expr) -> str:
             )
             return arg.name
         if isinstance(origin.origin, BlockSizeOrigin):
-            result = self.block_size_var(origin.origin.block_size_idx)
+            result = self.block_size_var(origin.origin.block_id)
             assert result is not None
             return result
         return self.expr_arg(expr, origin.origin).name

helion/_compiler/device_ir.py

@@ -321,7 +321,7 @@ def build_rolled_reductions(self) -> None:
                     new_graph, type(graph_info), **graph_info.kwargs()
                 )
                 reduction_info = RolledReductionInfo(
-                    rolled_block_indices=[rdim.block_size_idx],
+                    rolled_block_indices=[rdim.block_id],
                     original_graph_id=graph_id,
                     new_graph_id=new_graph_id,
                     used_rdim=len(roller.graphs_added) > 0,
@@ -335,7 +335,7 @@ def build_rolled_reductions(self) -> None:
                 # TODO(jansel): we should add support for rolling multiple dims at once
                 env.config_spec.reduction_loops.append(
                     ReductionLoopSpec(
-                        block_id=rdim.block_size_idx,
+                        block_id=rdim.block_id,
                         size_hint=rdim.size_hint(),
                     )
                 )
@@ -545,7 +545,7 @@ def run_subgraph(*args: object) -> list[object]:
                 graph_idx = self.device_ir.add_graph(
                     graph,
                     ForLoopGraphInfo,
-                    block_indices=[x.block_size_idx for x in iter_vars],
+                    block_indices=[x.block_id for x in iter_vars],
                     node_args=inputs.get_node_args(tracer),
                 )
                 args = (
@@ -826,9 +826,9 @@ def visit_For(self, node: ast.For) -> None:
             assert isinstance(iter_type, IterType)
             inner = iter_type.inner
             if isinstance(inner, SequenceType):
-                block_indices = [x.block_size_idx for x in inner.unpack()]
+                block_indices = [x.block_id for x in inner.unpack()]
             else:
-                block_indices = [inner.block_size_idx]
+                block_indices = [inner.block_id]
             self.device_ir.grid_block_indices.append(block_indices)
         else:
             self.generic_visit(node)

helion/_compiler/generate_ast.py

@@ -212,7 +212,7 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
         if self.on_device:
             pass
         elif isinstance(type_info := node._type_info, TileIndexType):
-            block_info = env.block_sizes[type_info.block_size_idx]
+            block_info = env.block_sizes[type_info.block_id]
             return expr_from_string(
                 self.host_function.literal_expr(
                     block_info.from_config(self.device_function.config)
@@ -221,7 +221,7 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
         elif isinstance(type_info, SequenceType):
             values = type_info.unpack()
             if all(isinstance(x, TileIndexType) for x in values):
-                block_infos = [env.block_sizes[x.block_size_idx] for x in values]
+                block_infos = [env.block_sizes[x.block_id] for x in values]
                 return expr_from_string(
                     self.host_function.literal_expr(
                         [

helion/_compiler/indexing_strategy.py

@@ -220,7 +220,7 @@ def compute_shape(
                     if origin and isinstance(origin.origin, BlockSizeOrigin):
                         if (
                             CompileEnvironment.current()
-                            .block_sizes[origin.origin.block_size_idx]
+                            .block_sizes[origin.origin.block_id]
                             .is_grid()
                         ):
                             pass
@@ -272,15 +272,15 @@ def create(
                 if isinstance(symbol, sympy.Symbol):
                     origin = HostFunction.current().expr_to_origin.get(symbol)
                 if origin and isinstance(origin.origin, BlockSizeOrigin):
-                    index_var = state.codegen.index_var(origin.origin.block_size_idx)
-                    if env.block_sizes[origin.origin.block_size_idx].is_grid():
+                    index_var = state.codegen.index_var(origin.origin.block_id)
+                    if env.block_sizes[origin.origin.block_id].is_grid():
                         index_values.append(index_var)
                         continue
                     expand = tile_strategy.expand_str(output_size, output_idx)
                     i = len(index_values)
                     index_values.append(f"({index_var}){expand}")
                     if (
-                        mask := state.codegen.mask_var(origin.origin.block_size_idx)
+                        mask := state.codegen.mask_var(origin.origin.block_id)
                     ) and fake_value.size(i) != 1:
                         mask_values.setdefault(f"({mask}){expand}")
                     output_idx += 1
@@ -293,7 +293,7 @@ def create(
                 size = fake_value.size(len(index_values))
                 if size != 1:
                     rdim = env.allocate_reduction_dimension(size)
-                    block_idx = rdim.block_size_idx
+                    block_idx = rdim.block_id
                     index_var = state.codegen.index_var(block_idx)
                     index_values.append(f"({index_var}){expand}")
                     if mask := state.codegen.mask_var(block_idx):
@@ -451,7 +451,7 @@ def is_supported(
                 if isinstance(symbol, sympy.Symbol):
                     origin = HostFunction.current().expr_to_origin.get(symbol)
                 if origin and isinstance(origin.origin, BlockSizeOrigin):
-                    block_index = origin.origin.block_size_idx
+                    block_index = origin.origin.block_id
                     try:
                         state.codegen.offset_var(block_index)
                     except NotImplementedError:
@@ -507,7 +507,7 @@ def create(
                 if origin and isinstance(origin.origin, BlockSizeOrigin):
                     if fake_value.size(len(res.offsets)) != 1:
                         res.offsets.append(
-                            state.codegen.offset_var(origin.origin.block_size_idx)
+                            state.codegen.offset_var(origin.origin.block_id)
                         )
                         res.block_shape.append(k)
                     else:
@@ -521,7 +521,7 @@ def create(
                 if size != 1:
                     env = CompileEnvironment.current()
                     rdim = env.allocate_reduction_dimension(size)
-                    res.offsets.append(state.codegen.offset_var(rdim.block_size_idx))
+                    res.offsets.append(state.codegen.offset_var(rdim.block_id))
                     res.block_shape.append(rdim.var)
                 else:
                     res.offsets.append("0")

helion/_compiler/roll_reduction.py

@@ -60,7 +60,7 @@ def is_reduction(self, node: torch.fx.Node) -> bool:
         return (
             node.op == "call_function"
             and isinstance(lowering := node.meta["lowering"], ReductionLowering)
-            and lowering.block_index == self.rdim.block_size_idx
+            and lowering.block_index == self.rdim.block_id
         )
 
     def should_go_in_inner_graph(self, node: torch.fx.Node) -> bool:
@@ -103,7 +103,7 @@ def should_go_in_inner_graph(self, node: torch.fx.Node) -> bool:
         if isinstance(val, torch.Tensor):
             for size in val.size():
                 block_idx = TileStrategy.get_block_index(size)
-                num_rdims += block_idx == self.rdim.block_size_idx
+                num_rdims += block_idx == self.rdim.block_id
             if num_rdims > 1:
                 raise NotImplementedError(
                     "multiple reduction dims of same size not supported"
@@ -121,7 +121,7 @@ def size_node(self, meta: dict[str, object]) -> torch.fx.Node:
             return self._size_node
         self._size_node = node = self.outer_graph.call_function(
             _get_symnode,
-            (f"rdim{self.rdim.block_size_idx}",),
+            (f"rdim{self.rdim.block_id}",),
             {},
         )
         node.meta.update(meta)
@@ -143,7 +143,7 @@ def start_new_graph(self) -> None:
         graph.output([*outputs.values()])
         graph_id = self.device_ir.add_reduction_loop_graph(
             graph,
-            block_index=self.rdim.block_size_idx,
+            block_index=self.rdim.block_id,
             node_args=self.inner_args,
         )
         self.graphs_added.append(graph_id)

helion/_compiler/tile_dispatch.py

@@ -97,7 +97,7 @@ def _add_loop_strategy(
 
     def _add_reduction_strategies(self, fn: DeviceFunction, config: Config) -> None:
         env = CompileEnvironment.current()
-        rdims = [bs.block_size_idx for bs in env.block_sizes if bs.reduction]
+        rdims = [bs.block_id for bs in env.block_sizes if bs.reduction]
         for block_id in rdims:
             reduction_loop = env.config_spec.reduction_loops.config_get(
                 config.reduction_loops, block_id, None

helion/_compiler/tile_strategy.py

@@ -127,7 +127,7 @@ def get_block_index(cls, size: int | torch.SymInt | sympy.Expr) -> int | None:
                 origin_info.origin,
                 BlockSizeOrigin,
             ):
-                return origin_info.origin.block_size_idx
+                return origin_info.origin.block_id
         return None
 
 

helion/_compiler/type_propagation.py

@@ -468,7 +468,7 @@ def _device_indexing_size(self, key: TypeInfo) -> list[int | torch.SymInt]:
                     output_sizes.append(1)
             elif isinstance(k, TileIndexType):
                 inputs_consumed += 1
-                output_sizes.append(env.block_sizes[k.block_size_idx].var)
+                output_sizes.append(env.block_sizes[k.block_id].var)
             elif isinstance(k, TypeNotAllowedOnDevice):
                 raise exc.TypePropagationError(k)
             elif isinstance(k, TensorType) and k.fake_value.ndim == 1:
@@ -944,22 +944,22 @@ def _get_hint(numel: int | torch.SymInt | AutoSize | None) -> int:
 
 
 class TileIndexType(TypeInfo):
-    block_size_idx: int
+    block_id: int
 
     def __str__(self) -> str:
-        return f"{type(self).__name__}({self.block_size_idx})"
+        return f"{type(self).__name__}({self.block_id})"
 
-    def __init__(self, origin: Origin, block_size_idx: int) -> None:
+    def __init__(self, origin: Origin, block_id: int) -> None:
         super().__init__(origin)
-        self.block_size_idx = block_size_idx
+        self.block_id = block_id
 
     def proxy(self) -> object:
         with proxy_tensor.disable_proxy_modes_tracing():
             fake_mode = torch._C._unset_dispatch_mode(
                 torch._C._TorchDispatchModeKey.FAKE
             )
             try:
-                return TileIndexProxy(self.block_size_idx)
+                return TileIndexProxy(self.block_id)
             finally:
                 assert fake_mode is not None
                 torch._C._set_dispatch_mode(fake_mode)
@@ -992,10 +992,10 @@ def allocate_fixed(
 
     def merge(self, other: TypeInfo) -> TypeInfo:
         if isinstance(other, TileIndexType):
-            if self.block_size_idx == other.block_size_idx:
+            if self.block_id == other.block_id:
                 return self
             return UnknownType(
-                debug_msg=f"TileIndexType mismatch in control flow: {self.block_size_idx} and {other.block_size_idx}",
+                debug_msg=f"TileIndexType mismatch in control flow: {self.block_id} and {other.block_id}",
                 origin=other.origin,
             )
         return super().merge(other)
@@ -1007,17 +1007,17 @@ def propagate_attribute(self, attr: str, origin: AttributeOrigin) -> TypeInfo:
 
 
 class GridIndexType(SymIntType):
-    block_size_idx: int
+    block_id: int
 
-    def __init__(self, origin: Origin, block_size_idx: int) -> None:
+    def __init__(self, origin: Origin, block_id: int) -> None:
         from .._compiler.compile_environment import CompileEnvironment
 
         env = CompileEnvironment.current()
-        super().__init__(origin, env.block_sizes[block_size_idx].var)
-        self.block_size_idx = block_size_idx
+        super().__init__(origin, env.block_sizes[block_id].var)
+        self.block_id = block_id
 
     def __str__(self) -> str:  # pragma: no cover – debug helper
-        return f"{type(self).__name__}({self.block_size_idx})"
+        return f"{type(self).__name__}({self.block_id})"
 
     @staticmethod
     def allocate(numel: int | torch.SymInt, origin: Origin) -> GridIndexType:
@@ -1030,10 +1030,10 @@ def allocate(numel: int | torch.SymInt, origin: Origin) -> GridIndexType:
 
     def merge(self, other: TypeInfo) -> TypeInfo:  # type: ignore[override]
         if isinstance(other, GridIndexType):
-            if self.block_size_idx == other.block_size_idx:
+            if self.block_id == other.block_id:
                 return self
             return UnknownType(
-                debug_msg=f"GridIndexType mismatch in control flow: {self.block_size_idx} vs {other.block_size_idx}",
+                debug_msg=f"GridIndexType mismatch in control flow: {self.block_id} vs {other.block_id}",
                 origin=other.origin,
             )
         return super().merge(other)

helion/_compiler/variable_origin.py

@@ -221,7 +221,7 @@ class DeviceOrigin(Origin):
 
 @dataclasses.dataclass
 class BlockSizeOrigin(Origin):
-    block_size_idx: int
+    block_id: int
 
     def host_str(self) -> str:
         """
@@ -232,7 +232,7 @@ def host_str(self) -> str:
         from .device_function import DeviceFunction
 
         # Look up the block size variable name; if not set (e.g., size==1), use literal 1
-        var = DeviceFunction.current().block_size_var(self.block_size_idx)
+        var = DeviceFunction.current().block_size_var(self.block_id)
         if var is None:
             return "1"
         return var