Rename TileIndexProxy to hl.Tile (#171)

This makes the docs more readable when they are auto-generated from
docstrings.

Author: jansel

helion/_compiler/device_ir.py

@@ -44,8 +44,6 @@
 from .node_masking import remove_unnecessary_masking
 from .roll_reduction import ReductionRoller
 from .source_location import current_location
-from .tile_index_proxy import CheckForIndexCalls
-from .tile_index_proxy import TileIndexProxy
 from .type_propagation import CallableType
 from .type_propagation import GridIndexType
 from .type_propagation import IterType
@@ -58,6 +56,8 @@
 from .type_propagation import _eval_binary
 from .type_propagation import _eval_compare
 from .type_propagation import _eval_unary
+from helion.language.tile_proxy import Tile
+from helion.language.tile_proxy import _CheckForIndexCalls
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -83,7 +83,7 @@ def _get_proxy_slot(
         default: object = proxy_tensor.no_default,
         transform: Callable[[object], object] = lambda x: x,
     ) -> object:
-        if isinstance(obj, torch.Tensor) and not isinstance(obj, TileIndexProxy):
+        if isinstance(obj, torch.Tensor) and not isinstance(obj, Tile):
             tracker = tracer.tensor_tracker
             if obj not in tracker:
                 origin = HostFunction.current().tensor_to_origin[obj]
@@ -473,7 +473,7 @@ def disable_tracing() -> Iterator[proxy_tensor.PythonKeyTracer]:
 
     @staticmethod
     def should_become_arg(value: object) -> bool:
-        if isinstance(value, (TileIndexProxy, torch.SymInt)):
+        if isinstance(value, (Tile, torch.SymInt)):
             return False
         if isinstance(value, torch.Tensor):
             if (
@@ -584,7 +584,7 @@ def run_subgraph(*args: object) -> list[object]:
                     tracer=tracer,
                 )
             for name, value in outputs.unflatten().items():
-                if isinstance(value, TileIndexProxy):
+                if isinstance(value, Tile):
                     continue
                 if name in self.scope:
                     try:
@@ -803,7 +803,7 @@ def visit_Call(self, node: ast.Call) -> object:
             func = self.visit(node.func)
 
         # pyre-ignore[6]
-        return CheckForIndexCalls.retry_call(func, args, kwargs)
+        return _CheckForIndexCalls.retry_call(func, args, kwargs)
 
     def visit_Attribute(self, node: ast.Attribute) -> object:
         return getattr(self.visit(node.value), node.attr)
@@ -825,7 +825,7 @@ def __init__(self, values: dict[str, object]) -> None:
         self.tensor_indices = [
             i
             for i, v in enumerate(self.flat_values)
-            if isinstance(v, torch.Tensor) and not isinstance(v, TileIndexProxy)
+            if isinstance(v, torch.Tensor) and not isinstance(v, Tile)
         ]
 
     def unflatten(self) -> dict[str, object]:

helion/_compiler/type_propagation.py

@@ -39,8 +39,6 @@
 from .output_header import library_imports
 from .source_location import SourceLocation
 from .source_location import current_location
-from .tile_index_proxy import CheckForIndexCalls
-from .tile_index_proxy import TileIndexProxy
 from .variable_origin import ArgumentOrigin
 from .variable_origin import AttributeOrigin
 from .variable_origin import BuiltinOrigin
@@ -51,6 +49,8 @@
 from .variable_origin import SourceOrigin
 from .variable_origin import TensorSizeOrigin
 import helion
+from helion.language.tile_proxy import Tile
+from helion.language.tile_proxy import _CheckForIndexCalls
 
 # pyre-ignore-all-errors[8,15,58]: visit_* overrides
 if TYPE_CHECKING:
@@ -623,7 +623,7 @@ def propagate_call(
         try:
             fn = getattr(self.tensor.fake_value, attr)
             output_type = TypeInfo.from_example(
-                CheckForIndexCalls.retry_call(fn, proxy_args, proxy_kwargs), origin
+                _CheckForIndexCalls.retry_call(fn, proxy_args, proxy_kwargs), origin
             )
         except exc.Base:
             raise
@@ -775,7 +775,9 @@ def to_proxy(arg: TypeInfo) -> object:
         try:
             with patch.object(torch.SymInt, "__index__", _raise_shape_specializing):
                 output_type = TypeInfo.from_example(
-                    CheckForIndexCalls.retry_call(self.value, proxy_args, proxy_kwargs),
+                    _CheckForIndexCalls.retry_call(
+                        self.value, proxy_args, proxy_kwargs
+                    ),
                     origin,
                 )
             output_type.tree_map(warn_wrong_device)
@@ -1006,7 +1008,7 @@ def proxy(self) -> object:
                 torch._C._TorchDispatchModeKey.FAKE
             )
             try:
-                return TileIndexProxy(self.block_id)
+                return Tile(self.block_id)
             finally:
                 assert fake_mode is not None
                 torch._C._set_dispatch_mode(fake_mode)
@@ -1048,7 +1050,7 @@ def merge(self, other: TypeInfo) -> TypeInfo:
         return super().merge(other)
 
     def propagate_attribute(self, attr: str, origin: AttributeOrigin) -> TypeInfo:
-        if isinstance(getattr(TileIndexProxy, attr, None), property):
+        if isinstance(getattr(Tile, attr, None), property):
             return TypeInfo.from_example(getattr(self.proxy(), attr), origin)
         return super().propagate_attribute(attr, origin)
 

helion/language/__init__.py

@@ -5,16 +5,16 @@
 from .creation_ops import full as full
 from .creation_ops import zeros as zeros
 from .device_print import device_print as device_print
-from .loops import Tile as Tile
 from .loops import grid as grid
 from .loops import tile as tile
 from .memory_ops import atomic_add as atomic_add
 from .memory_ops import load as load
 from .memory_ops import store as store
-from .tiles import tile_begin as tile_begin
-from .tiles import tile_block_size as tile_block_size
-from .tiles import tile_end as tile_end
-from .tiles import tile_index as tile_index
+from .tile_ops import tile_begin as tile_begin
+from .tile_ops import tile_block_size as tile_block_size
+from .tile_ops import tile_end as tile_end
+from .tile_ops import tile_index as tile_index
+from .tile_proxy import Tile as Tile
 from .tunable_ops import register_block_size as register_block_size
 from .tunable_ops import register_reduction_dim as register_reduction_dim
 from .tunable_ops import register_tunable as register_tunable

helion/language/_decorators.py

@@ -111,9 +111,9 @@ def unpack(x: object) -> object:
 
 
 def tiles_as_sizes_prepare_args(*args: object) -> tuple[object, ...]:
-    from helion._compiler.tile_index_proxy import TileIndexProxy
+    from helion.language.tile_proxy import Tile
 
-    return TileIndexProxy.tiles_to_sizes(args)
+    return Tile._tiles_to_sizes(args)
 
 
 def no_op_prepare_args(*args: object) -> tuple[object, ...]:
@@ -273,12 +273,12 @@ def _default_type_function(
     def type_prop_with_fake_fn(
         *args: object, origin: Origin, **kwargs: object
     ) -> TypeInfo:
-        from .._compiler.tile_index_proxy import TileIndexProxy
         from .._compiler.type_propagation import TypeInfo
+        from helion.language.tile_proxy import Tile
 
         args, kwargs = tree_map_only(TypeInfo, _to_proxy, (args, kwargs))
         if tiles_as_sizes:
-            args, kwargs = TileIndexProxy.tiles_to_sizes((args, kwargs))
+            args, kwargs = Tile._tiles_to_sizes((args, kwargs))
         return TypeInfo.from_example(fake_fn(*args, **kwargs), origin)
 
     return type_prop_with_fake_fn

helion/language/_tracing_ops.py

@@ -12,9 +12,9 @@
 from .._compiler.ast_extension import expr_from_string
 from .._compiler.compile_environment import CompileEnvironment
 from .._compiler.host_function import HostFunction
-from .._compiler.tile_index_proxy import TileIndexProxy
 from ..exc import NotInsideKernel
 from . import _decorators
+from helion.language.tile_proxy import Tile
 
 if TYPE_CHECKING:
     from .._compiler.inductor_lowering import CodegenState
@@ -96,9 +96,9 @@ def _phi(lhs: object, rhs: object) -> object:
 
 @_decorators.register_fake(_phi)
 def _(lhs: object, rhs: object) -> object:
-    if isinstance(lhs, TileIndexProxy):
-        assert isinstance(rhs, TileIndexProxy)
-        assert lhs.block_size_index == rhs.block_size_index
+    if isinstance(lhs, Tile):
+        assert isinstance(rhs, Tile)
+        assert lhs.block_id == rhs.block_id
         return lhs
     assert isinstance(lhs, torch.Tensor), lhs
     assert isinstance(rhs, torch.Tensor), rhs

helion/language/loops.py

@@ -15,7 +15,6 @@
 from .._compiler.ast_extension import LoopType
 from .._compiler.ast_extension import expr_from_string
 from .._compiler.compile_environment import CompileEnvironment
-from .._compiler.tile_index_proxy import TileIndexProxy
 from .._compiler.type_propagation import GridIndexType
 from .._compiler.type_propagation import IterType
 from .._compiler.type_propagation import Origin
@@ -28,15 +27,15 @@
 from ..autotuner.config_spec import L2GroupingSpec
 from ..autotuner.config_spec import LoopOrderSpec
 from . import _decorators
+from helion.language.tile_proxy import Tile
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     from .._compiler.inductor_lowering import CodegenState
 
 
-__all__ = ["Tile", "grid", "tile"]
-Tile = TileIndexProxy
+__all__ = ["grid", "tile"]
 
 
 @overload
@@ -186,7 +185,7 @@ def _(
     proxy_end = _to_proxy(end)
     _check_matching(proxy_begin, proxy_end)
     if _not_none(block_size):
-        proxy_block_size = TileIndexProxy.tiles_to_sizes(_to_proxy(block_size))
+        proxy_block_size = Tile._tiles_to_sizes(_to_proxy(block_size))
         _check_matching(proxy_end, proxy_block_size)
     else:
         proxy_block_size = begin.tree_map(lambda n: None)

helion/language/memory_ops.py

@@ -44,11 +44,11 @@ def _(
     value: torch.Tensor | torch.SymInt | float,
     extra_mask: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, list[object], torch.Tensor | torch.SymInt | int | float]:
-    from helion._compiler.tile_index_proxy import TileIndexProxy
+    from helion.language.tile_proxy import Tile
 
     if hasattr(value, "dtype") and value.dtype != tensor.dtype:
         value = value.to(tensor.dtype)
-    index = TileIndexProxy.tiles_to_sizes(index)
+    index = Tile._tiles_to_sizes(index)
     return (tensor, index, value, extra_mask)
 
 
@@ -147,16 +147,16 @@ def _(
     value: torch.Tensor | float,
     sem: str = "relaxed",
 ) -> tuple[torch.Tensor, object, torch.Tensor | float | int, str]:
-    from helion._compiler.tile_index_proxy import TileIndexProxy
+    from helion.language.tile_proxy import Tile
 
     valid_sems = {"relaxed", "acquire", "release", "acq_rel"}
     if sem not in valid_sems:
         raise ValueError(
             f"Invalid memory semantic '{sem}'. Must be one of {valid_sems}."
         )
 
-    index = TileIndexProxy.prepare_index(index)
-    index = TileIndexProxy.tiles_to_sizes(index)
+    index = Tile._prepare_index(index)
+    index = Tile._tiles_to_sizes(index)
 
     return (target, index, value, sem)