Rename block_indices to block_ids (#135)

Some minor refactoring to make naming more consistent

Author: jansel

helion/_compiler/device_ir.py

@@ -188,7 +188,7 @@ def kwargs(self) -> dict[str, object]:
 
 @dataclasses.dataclass
 class ForLoopGraphInfo(NodeArgsGraphInfo):
-    block_indices: list[int]
+    block_ids: list[int]
 
     @property
     def name(self) -> str:
@@ -197,7 +197,7 @@ def name(self) -> str:
     def kwargs(self) -> dict[str, object]:
         return {
             **super().kwargs(),
-            "block_indices": [*self.block_indices],
+            "block_ids": [*self.block_ids],
         }
 
     def codegen(self, state: CodegenState) -> list[object]:
@@ -206,7 +206,7 @@ def codegen(self, state: CodegenState) -> list[object]:
         assert all(isinstance(x, ast.AST) for x in args)
         with state.codegen.add_device_loop(
             state.device_function.tile_strategy.codegen_device_loop(
-                state, self.block_indices
+                state, self.block_ids
             )
         ):
             return codegen_call_with_graph(
@@ -238,7 +238,7 @@ def codegen(self, state: CodegenState) -> list[object]:
 
 
 class RolledReductionInfo(NamedTuple):
-    rolled_block_indices: list[int]
+    rolled_block_ids: list[int]
     original_graph_id: int
     new_graph_id: int | None
     used_rdim: bool
@@ -251,7 +251,7 @@ def __init__(self) -> None:
         self.graphs: list[GraphInfo] = []
         self.root_id: int | None = None
         self.rolled_reductions: list[RolledReductionInfo] = []
-        self.grid_block_indices: list[list[int]] = []
+        self.grid_block_ids: list[list[int]] = []
 
     def get_root(self, config: Config) -> torch.fx.Graph:
         """ " If we are using a rolled reduction, return the rolled reduction graph otherwise
@@ -296,7 +296,7 @@ def add_reduction_loop_graph(
         return self.add_graph(
             graph,
             graph_info_cls=ReductionLoopGraphInfo,
-            block_indices=[block_index],
+            block_ids=[block_index],
             node_args=node_args,
         )
 
@@ -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_id],
+                    rolled_block_ids=[rdim.block_id],
                     original_graph_id=graph_id,
                     new_graph_id=new_graph_id,
                     used_rdim=len(roller.graphs_added) > 0,
@@ -545,7 +545,7 @@ def run_subgraph(*args: object) -> list[object]:
                 graph_idx = self.device_ir.add_graph(
                     graph,
                     ForLoopGraphInfo,
-                    block_indices=[x.block_id for x in iter_vars],
+                    block_ids=[x.block_id for x in iter_vars],
                     node_args=inputs.get_node_args(tracer),
                 )
                 args = (
@@ -826,10 +826,10 @@ 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_id for x in inner.unpack()]
+                block_ids = [x.block_id for x in inner.unpack()]
             else:
-                block_indices = [inner.block_id]
-            self.device_ir.grid_block_indices.append(block_indices)
+                block_ids = [inner.block_id]
+            self.device_ir.grid_block_ids.append(block_ids)
         else:
             self.generic_visit(node)
 

helion/_compiler/generate_ast.py

@@ -97,22 +97,22 @@ def set_on_device(self) -> Iterator[None]:
     @contextlib.contextmanager
     def add_device_loop(self, device_loop: DeviceLoopState) -> Iterator[None]:
         with self.set_statements(device_loop.inner_statements):
-            for idx in device_loop.block_indices:
+            for idx in device_loop.block_ids:
                 active_loops = self.active_device_loops[idx]
                 active_loops.append(device_loop)
                 if len(active_loops) > 1:
                     raise exc.NestedDeviceLoopsConflict
             try:
                 yield
             finally:
-                for idx in device_loop.block_indices:
+                for idx in device_loop.block_ids:
                     self.active_device_loops[idx].pop()
         self.statements_stack[-1].extend(device_loop.outer_prefix)
         self.add_statement(device_loop.for_node)
         self.statements_stack[-1].extend(device_loop.outer_suffix)
 
     def set_active_loops(self, device_grid: DeviceLoopOrGridState) -> None:
-        for idx in device_grid.block_indices:
+        for idx in device_grid.block_ids:
             self.active_device_loops[idx] = [device_grid]
 
     def generic_visit(self, node: ast.AST) -> ast.AST:

helion/_compiler/reduction_strategy.py

@@ -39,7 +39,7 @@ def __init__(
     ) -> None:
         super().__init__(
             fn=fn,
-            block_indices=[block_index],
+            block_ids=[block_index],
         )
         self._mask_var = mask_var
         if block_size_var is not None:
@@ -51,7 +51,7 @@ def mask_var(self, block_idx: int) -> str | None:
 
     @property
     def block_index(self) -> int:
-        return self.block_indices[0]
+        return self.block_ids[0]
 
     def user_size(self, block_index: int) -> sympy.Expr:
         return CompileEnvironment.current().block_sizes[block_index].numel

helion/_compiler/tile_dispatch.py

@@ -41,37 +41,37 @@ def __init__(
     ) -> None:
         super().__init__()
         self.strategies: list[TileStrategy] = []
-        self.block_indices_to_strategy: dict[tuple[int, ...], TileStrategy] = {}
+        self.block_id_to_strategy: dict[tuple[int, ...], TileStrategy] = {}
         self._add_loop_strategies(fn, config)
         self._add_reduction_strategies(fn, config)
 
     def _add_loop_strategies(self, fn: DeviceFunction, config: Config) -> None:
         device_ir = HostFunction.current().device_ir
-        for block_indices in device_ir.grid_block_indices:
-            self._add_loop_strategy(block_indices, fn, config)
+        for block_ids in device_ir.grid_block_ids:
+            self._add_loop_strategy(block_ids, fn, config)
         for graph in device_ir.graphs:
             if isinstance(graph, ForLoopGraphInfo) and not isinstance(
                 graph, ReductionLoopGraphInfo
             ):
-                block_indices = [*graph.block_indices]
-                self._add_loop_strategy(block_indices, fn, config)
+                block_ids = [*graph.block_ids]
+                self._add_loop_strategy(block_ids, fn, config)
 
     def _add_loop_strategy(
-        self, block_indices: list[int], fn: DeviceFunction, config: Config
+        self, block_ids: list[int], fn: DeviceFunction, config: Config
     ) -> None:
         env = CompileEnvironment.current()
-        block_size_infos = [env.block_sizes[i] for i in block_indices]
+        block_size_infos = [env.block_sizes[i] for i in block_ids]
         loop_order = env.config_spec.loop_orders.config_get(
-            config.loop_orders, block_indices[0]
-        ) or [*range(len(block_indices))]
+            config.loop_orders, block_ids[0]
+        ) or [*range(len(block_ids))]
         l2_grouping = env.config_spec.l2_groupings.config_get(
-            config.l2_groupings, block_indices[0], 1
+            config.l2_groupings, block_ids[0], 1
         )
 
         if block_size_infos[0].is_grid():
             strategy: TileStrategy = NDGridTileStrategy(
                 fn,
-                block_indices,
+                block_ids,
                 loop_order=loop_order,
             )
         elif block_size_infos[0].is_flattened(config):
@@ -80,20 +80,20 @@ def _add_loop_strategy(
             )
             strategy: TileStrategy = FlattenedTileStrategy(
                 fn,
-                block_indices,
+                block_ids,
                 block_size=block_size,
                 loop_order=loop_order,
             )
         else:
             strategy = NDTileStrategy(
                 fn,
-                block_indices,
+                block_ids,
                 block_size=[bs.from_config_assert(config) for bs in block_size_infos],
                 loop_order=loop_order,
                 l2_grouping=l2_grouping,
             )
         self.strategies.append(strategy)
-        self.block_indices_to_strategy[tuple(block_indices)] = strategy
+        self.block_id_to_strategy[tuple(block_ids)] = strategy
 
     def _add_reduction_strategies(self, fn: DeviceFunction, config: Config) -> None:
         env = CompileEnvironment.current()
@@ -107,20 +107,20 @@ def _add_reduction_strategies(self, fn: DeviceFunction, config: Config) -> None:
             else:
                 strategy = LoopedReductionStrategy(fn, block_id, reduction_loop)
             self.strategies.append(strategy)
-            self.block_indices_to_strategy[(block_id,)] = strategy
+            self.block_id_to_strategy[(block_id,)] = strategy
 
-    def codegen_grid(self, state: CodegenState, block_indices: list[int]) -> None:
-        strategy = self.block_indices_to_strategy[tuple(block_indices)]
+    def codegen_grid(self, state: CodegenState, block_ids: list[int]) -> None:
+        strategy = self.block_id_to_strategy[tuple(block_ids)]
         strategy.codegen_grid(state)
         for other_strategy in self.strategies:
             if other_strategy is not strategy:
                 other_strategy.codegen_preamble(state)
         state.codegen.set_active_loops(DeviceGridState(strategy))
 
     def codegen_device_loop(
-        self, state: CodegenState, block_indices: list[int]
+        self, state: CodegenState, block_ids: list[int]
     ) -> DeviceLoopState:
-        strategy = self.block_indices_to_strategy[tuple(block_indices)]
+        strategy = self.block_id_to_strategy[tuple(block_ids)]
         return strategy.codegen_device_loop(state)
 
     def _compact_shape(self, shapes: ShapeLike) -> list[CompactedShape]:
@@ -161,14 +161,14 @@ def expand_str(self, shape: ShapeLike, i: int) -> str:
         return f"[{', '.join(result)}]"
 
     def get_reduction_strategy(self, block_idx: int) -> ReductionStrategy:
-        strategy = self.block_indices_to_strategy[(block_idx,)]
+        strategy = self.block_id_to_strategy[(block_idx,)]
         assert isinstance(strategy, ReductionStrategy)
         return strategy
 
     def user_size(self, block_index: int) -> sympy.Expr:
         """The user-visible size of the block index."""
         # This only does something special for reduction loops, only need to check for 1D loop
-        strategy = self.block_indices_to_strategy.get((block_index,))
+        strategy = self.block_id_to_strategy.get((block_index,))
         if strategy is None:
             return CompileEnvironment.current().block_sizes[block_index].symbol()
         return strategy.user_size(block_index)