update

Author: yf225

helion/_compiler/compile_environment.py

@@ -69,8 +69,6 @@ def __init__(self, device: torch.device, settings: Settings) -> None:
             collections.Counter()
         )
 
-        self.symint_to_grid_index_type: dict[str, GridIndexType] = {}
-
     def add_kernel_tensor_size(self, sizes: Sequence[int | torch.SymInt]) -> None:
         self.kernel_tensor_sizes[(*map(_to_sympy, sizes),)] += 1
 

helion/_compiler/indexing_strategy.py

@@ -11,9 +11,10 @@
 
 from .. import exc
 from .ast_extension import expr_from_string
-from .compile_environment import CompileEnvironment
+from .compile_environment import CompileEnvironment, GridBlockSizeSource
 from .host_function import HostFunction
 from .tile_strategy import TileStrategy
+from .tile_strategy import GridTileStrategy
 from .variable_origin import BlockSizeOrigin
 
 if TYPE_CHECKING:
@@ -169,10 +170,7 @@ def compute_shape(
                 if isinstance(symbol, sympy.Symbol):
                     origin = HostFunction.current().symbol_to_origin.get(symbol.name)
                     if origin and isinstance(origin.origin, BlockSizeOrigin):
-                        if (
-                            str(k)
-                            in CompileEnvironment.current().symint_to_grid_index_type
-                        ):
+                        if isinstance(CompileEnvironment.current().block_sizes[origin.origin.block_size_idx].block_size_source, GridBlockSizeSource):
                             pass
                         elif tensor.size(tensor.ndim - len(input_size) - 1) != 1:
                             output_size.append(k)
@@ -217,7 +215,7 @@ def create(
                 if isinstance(symbol, sympy.Symbol):
                     origin = HostFunction.current().symbol_to_origin.get(symbol.name)
                 if origin and isinstance(origin.origin, BlockSizeOrigin):
-                    if str(k) in CompileEnvironment.current().symint_to_grid_index_type:
+                    if isinstance(CompileEnvironment.current().block_sizes[origin.origin.block_size_idx].block_size_source, GridBlockSizeSource):
                         first_non_grid_index = n + 1
                         expand = tile_strategy.expand_str(output_size, output_idx)
                     else:

helion/_compiler/tile_dispatch.py

@@ -61,7 +61,6 @@ def _add_loop_strategy(
         env = CompileEnvironment.current()
         block_size_infos = [env.block_sizes[i] for i in block_indices]
         loop_order = block_size_infos[0].get_order(config, len(block_size_infos))
-        print(f"block_size_infos: {block_size_infos}, block_size_infos[0]: {block_size_infos[0]}")
         if isinstance(block_size_infos[0].block_size_source, GridBlockSizeSource):
             strategy: TileStrategy = GridTileStrategy(
                 fn,

helion/_compiler/tile_strategy.py

@@ -145,10 +145,6 @@ def __init__(
         )
         self.block_size = block_size
         self.loop_order = loop_order
-        if str(block_size) == "[1, 1]":
-            import traceback
-            traceback.print_stack()
-        print(f"block_size: {block_size}, loop_order: {loop_order}, block_indices: {block_indices}")
 
     def _reorder(self, block_indices: list[_T]) -> list[_T]:
         if len(block_indices) <= 1:
@@ -363,7 +359,6 @@ def __init__(
         loop_order: list[int],
         l2_grouping: int,
     ) -> None:
-        print(f"will create new NDTileStrategy")
         assert isinstance(block_size, list)
         super().__init__(fn, block_indices, block_size, loop_order)
         self.mask_vars: dict[int, str | None] = {}

helion/_compiler/type_propagation.py

@@ -984,7 +984,6 @@ def __init__(self, origin: Origin, block_size_idx: int) -> None:
 
         env = CompileEnvironment.current()
         super().__init__(origin, env.block_sizes[block_size_idx].var)
-        env.symint_to_grid_index_type[str(self.value)] = self
         self.block_size_idx = block_size_idx
 
     def __str__(self) -> str:  # pragma: no cover – debug helper