Make imports relative (#310)

jansel · web-flow · commit 366a3b39b9a4 · 2025-07-14T11:36:34.000-07:00
diff --git a/examples/matmul_layernorm.py b/examples/matmul_layernorm.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import torch
+import torch.nn.functional as F
 
 import helion
 from helion._testing import run_example
@@ -37,8 +38,6 @@ def matmul_layernorm(
 def matmul_layernorm_pytorch(
     x: torch.Tensor, y: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
 ) -> torch.Tensor:
-    import torch.nn.functional as F
-
     matmul_out = torch.matmul(x, y)
 
     ln_out = F.layer_norm(
diff --git a/helion/__init__.py b/helion/__init__.py
@@ -11,8 +11,8 @@
 from .runtime import Kernel
 from .runtime import kernel
 from .runtime import kernel as jit  # alias
-from helion.runtime.settings import Settings
-from helion.runtime.settings import set_default_settings
+from .runtime.settings import Settings
+from .runtime.settings import set_default_settings
 
 __all__ = [
     "Config",
diff --git a/helion/_compiler/device_ir.py b/helion/_compiler/device_ir.py
@@ -32,6 +32,8 @@
 from ..language._decorators import args_to_proxies
 from ..language._decorators import get_device_func_replacement
 from ..language._tracing_ops import _new_var
+from ..language.tile_proxy import Tile
+from ..language.tile_proxy import _CheckForIndexCalls
 from .ast_extension import ExtendedAST
 from .ast_extension import LoopType
 from .ast_extension import NodeVisitor
@@ -58,8 +60,6 @@
 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
diff --git a/helion/_compiler/lift_closures.py b/helion/_compiler/lift_closures.py
@@ -5,10 +5,10 @@
 
 from torch._dynamo.utils import make_cell
 
-from helion import exc
-from helion._compiler.host_function import HostFunction
-from helion._compiler.variable_origin import ClosureOrigin
-from helion._compiler.variable_origin import Origin
+from .. import exc
+from .host_function import HostFunction
+from .variable_origin import ClosureOrigin
+from .variable_origin import Origin
 
 
 class CaptureGlobals(dict[str, object]):
diff --git a/helion/_compiler/node_masking.py b/helion/_compiler/node_masking.py
@@ -16,13 +16,13 @@
 from torch.fx.experimental import proxy_tensor
 from torch.utils._sympy.value_ranges import ValueRanges
 
-from helion.language._tracing_ops import _for_loop
-from helion.language._tracing_ops import _if
-from helion.language._tracing_ops import _mask_to
-from helion.language._tracing_ops import _phi
+from ..language._tracing_ops import _for_loop
+from ..language._tracing_ops import _if
+from ..language._tracing_ops import _mask_to
+from ..language._tracing_ops import _phi
 
 if TYPE_CHECKING:
-    from helion._compiler.inductor_lowering import InductorLowering
+    from .inductor_lowering import InductorLowering
 
     ValueRangesAny = ValueRanges[Any]
 
@@ -49,7 +49,7 @@ def apply_masking(
             if user.args[1] == other:
                 assert user.args[0] is node
                 return user  # reuse existing mask_to node
-    from helion._compiler.inductor_lowering import APIFuncLowering
+    from .inductor_lowering import APIFuncLowering
 
     # If we reach here, we need to create a new mask_to node
     with node.graph.inserting_before(base_node):
@@ -79,9 +79,9 @@ def cached_masked_value(
         return node.meta["masked_value"]
 
     if node.op == "placeholder":
-        from helion._compiler.device_ir import DeviceIR
-        from helion._compiler.device_ir import ForLoopGraphInfo
-        from helion._compiler.device_ir import NodeArgsGraphInfo
+        from .device_ir import DeviceIR
+        from .device_ir import ForLoopGraphInfo
+        from .device_ir import NodeArgsGraphInfo
 
         """
         We are inside a for loop or an if statement, which is represented as a subgraph.
@@ -128,7 +128,7 @@ def getitem_masked_value(
     Retrieve the masked value for a node that is a getitem operation.
     This handles loop outputs, since the `_for` node has multiple outputs.
     """
-    from helion._compiler.device_ir import DeviceIR
+    from .device_ir import DeviceIR
 
     assert not getitem_node.kwargs, "getitem kwargs not supported"
     node, index = getitem_node.args
diff --git a/helion/_compiler/roll_reduction.py b/helion/_compiler/roll_reduction.py
@@ -6,21 +6,21 @@
 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.language._tracing_ops import _for_loop
-from helion.language._tracing_ops import _get_symnode
-from helion.language._tracing_ops import _host_tensor
-from helion.language._tracing_ops import _if
-from helion.language.memory_ops import store
-from helion.language.reduce_ops import _reduce
+from ..language._tracing_ops import _for_loop
+from ..language._tracing_ops import _get_symnode
+from ..language._tracing_ops import _host_tensor
+from ..language._tracing_ops import _if
+from ..language.memory_ops import store
+from ..language.reduce_ops import _reduce
+from .compile_environment import CompileEnvironment
+from .inductor_lowering import APIFuncLowering
+from .inductor_lowering import ReductionLowering
+from .inductor_lowering import aten_lowering_dispatch
 
 if TYPE_CHECKING:
-    from helion._compiler.compile_environment import BlockSizeInfo
-    from helion._compiler.device_ir import DeviceIR
-    from helion._compiler.device_ir import RolledReductionInfo
+    from .compile_environment import BlockSizeInfo
+    from .device_ir import DeviceIR
+    from .device_ir import RolledReductionInfo
 
 _duplicate_ops: tuple[object, ...] = (
     _host_tensor,
diff --git a/helion/_compiler/tile_dispatch.py b/helion/_compiler/tile_dispatch.py
@@ -4,28 +4,28 @@
 import operator
 from typing import TYPE_CHECKING
 
-from helion._compiler.compile_environment import CompileEnvironment
-from helion._compiler.device_function import DeviceFunction
-from helion._compiler.device_ir import ForLoopGraphInfo
-from helion._compiler.device_ir import ReductionLoopGraphInfo
-from helion._compiler.host_function import HostFunction
-from helion._compiler.reduction_strategy import LoopedReductionStrategy
-from helion._compiler.reduction_strategy import PersistentReductionStrategy
-from helion._compiler.reduction_strategy import ReductionStrategy
-from helion._compiler.tile_strategy import CompactedShape
-from helion._compiler.tile_strategy import DeviceLoopState
-from helion._compiler.tile_strategy import FlattenedTileStrategy
-from helion._compiler.tile_strategy import NDTileStrategy
-from helion._compiler.tile_strategy import TileStrategy
+from .compile_environment import CompileEnvironment
+from .device_function import DeviceFunction
+from .device_ir import ForLoopGraphInfo
+from .device_ir import ReductionLoopGraphInfo
+from .host_function import HostFunction
+from .reduction_strategy import LoopedReductionStrategy
+from .reduction_strategy import PersistentReductionStrategy
+from .reduction_strategy import ReductionStrategy
+from .tile_strategy import CompactedShape
+from .tile_strategy import DeviceLoopState
+from .tile_strategy import FlattenedTileStrategy
+from .tile_strategy import NDTileStrategy
+from .tile_strategy import TileStrategy
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     import sympy
     import torch
 
-    from helion import Config
-    from helion._compiler.inductor_lowering import CodegenState
+    from .. import Config
+    from .inductor_lowering import CodegenState
 
     SymIntLike = torch.SymInt | int
     ShapeLike = Sequence[SymIntLike]
diff --git a/helion/_compiler/type_propagation.py b/helion/_compiler/type_propagation.py
@@ -27,6 +27,8 @@
 from ..autotuner.config_spec import BlockSizeSpec
 from ..language._decorators import get_device_func_replacement
 from ..language._decorators import is_api_func
+from ..language.tile_proxy import Tile
+from ..language.tile_proxy import _CheckForIndexCalls
 from .ast_extension import ExtendedAST
 from .ast_extension import LoopType
 from .ast_extension import create
@@ -50,8 +52,6 @@
 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
 
 if TYPE_CHECKING:
     from collections.abc import Callable
diff --git a/helion/autotuner/config_generation.py b/helion/autotuner/config_generation.py
@@ -7,14 +7,14 @@
 from typing import TYPE_CHECKING
 from typing import cast
 
-from helion._compat import warps_to_threads
-from helion.autotuner.config_fragment import Category
-from helion.autotuner.config_fragment import ConfigSpecFragment
-from helion.autotuner.config_fragment import PowerOfTwoFragment
+from .._compat import warps_to_threads
+from .config_fragment import Category
+from .config_fragment import ConfigSpecFragment
+from .config_fragment import PowerOfTwoFragment
 
 if TYPE_CHECKING:
-    from helion import Config
-    from helion.autotuner import ConfigSpec
+    from .. import Config
+    from . import ConfigSpec
 
 FlatConfig = list[object]
 
diff --git a/helion/autotuner/config_spec.py b/helion/autotuner/config_spec.py
@@ -28,8 +28,8 @@
     from collections.abc import Callable
     from collections.abc import Sequence
 
-    from helion.runtime.config import IndexingLiteral
-    from helion.runtime.config import PidTypeLiteral
+    from ..runtime.config import IndexingLiteral
+    from ..runtime.config import PidTypeLiteral
 
 DEFAULT_NUM_WARPS = 4
 DEFAULT_NUM_STAGES = 3
diff --git a/helion/autotuner/random_search.py b/helion/autotuner/random_search.py
@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from .config_generation import ConfigGeneration
-from helion.autotuner.finite_search import FiniteSearch
+from .finite_search import FiniteSearch
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
diff --git a/helion/language/_decorators.py b/helion/language/_decorators.py
@@ -17,14 +17,14 @@
 from torch.utils._pytree import tree_map_only
 from torch.utils._thunk import Thunk
 
-from helion import exc
+from .. import exc
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from helion._compiler.inductor_lowering import CodegenState
-    from helion._compiler.type_propagation import TypeInfo
-    from helion._compiler.variable_origin import Origin
+    from .._compiler.inductor_lowering import CodegenState
+    from .._compiler.type_propagation import TypeInfo
+    from .._compiler.variable_origin import Origin
 
     _T = TypeVar("_T")
     _C = TypeVar("_C", bound=Callable[..., object])
@@ -112,7 +112,7 @@ def unpack(x: object) -> object:
 
 
 def tiles_as_sizes_prepare_args(*args: object) -> tuple[object, ...]:
-    from helion.language.tile_proxy import Tile
+    from .tile_proxy import Tile
 
     return Tile._tiles_to_sizes(args)
 
@@ -140,7 +140,7 @@ def wrapper(*args: object, **kwargs: object) -> object:
 
             mode = proxy_tensor.get_proxy_mode()
             if mode is None:
-                from helion._compiler.compile_environment import CompileEnvironment
+                from .._compiler.compile_environment import CompileEnvironment
 
                 if CompileEnvironment.has_current():
                     assert api._fake_fn is not None
@@ -296,7 +296,7 @@ def type_prop_with_fake_fn(
         *args: object, origin: Origin, **kwargs: object
     ) -> TypeInfo:
         from .._compiler.type_propagation import TypeInfo
-        from helion.language.tile_proxy import Tile
+        from .tile_proxy import Tile
 
         args, kwargs = tree_map_only(TypeInfo, _to_proxy, (args, kwargs))
         if tiles_as_sizes:
diff --git a/helion/language/_tracing_ops.py b/helion/language/_tracing_ops.py
@@ -17,7 +17,7 @@
 from .._compiler.host_function import HostFunction
 from ..exc import NotInsideKernel
 from . import _decorators
-from helion.language.tile_proxy import Tile
+from .tile_proxy import Tile
 
 if TYPE_CHECKING:
     from .._compiler.inductor_lowering import CodegenState
diff --git a/helion/language/loops.py b/helion/language/loops.py
@@ -40,7 +40,7 @@
 from ..autotuner.config_spec import RangeWarpSpecializeSpec
 from ..autotuner.config_spec import StaticRangeSpec
 from . import _decorators
-from helion.language.tile_proxy import Tile
+from .tile_proxy import Tile
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -335,7 +335,7 @@ def _(
         elif isinstance(bs, int):
             results.append(TileIndexType.allocate(size, origin, bs))
         elif isinstance(bs, torch.SymInt):
-            from helion._compiler.compile_environment import CompileEnvironment
+            from .._compiler.compile_environment import CompileEnvironment
 
             index = CompileEnvironment.current().get_block_id(bs)
             if index is None:
diff --git a/helion/language/memory_ops.py b/helion/language/memory_ops.py
@@ -52,7 +52,7 @@ def _(
 ) -> tuple[
     torch.Tensor, list[object], torch.Tensor | torch.SymInt | float, torch.Tensor | None
 ]:
-    from helion.language.tile_proxy import Tile
+    from .tile_proxy import Tile
 
     if isinstance(value, torch.Tensor) and value.dtype != tensor.dtype:
         value = value.to(tensor.dtype)
@@ -189,7 +189,7 @@ def _(
     value: torch.Tensor | float,
     sem: str = "relaxed",
 ) -> tuple[torch.Tensor, object, torch.Tensor | float | int, str]:
-    from helion.language.tile_proxy import Tile
+    from .tile_proxy import Tile
 
     valid_sems = {"relaxed", "acquire", "release", "acq_rel"}
     if sem not in valid_sems:
diff --git a/helion/language/signal_wait.py b/helion/language/signal_wait.py
@@ -57,7 +57,7 @@ def _(
     scope: str = "gpu",
     skip_sync: bool = False,
 ) -> tuple[torch.Tensor, object, int, int | None, str, str, str, bool]:
-    from helion.language.tile_proxy import Tile
+    from .tile_proxy import Tile
 
     valid_ops = {"ld", "atomic_cas"}
     valid_sems = {"relaxed", "acquire", "acq_rel"}
@@ -194,7 +194,7 @@ def _(
     scope: str = "gpu",
     skip_sync: bool = False,
 ) -> tuple[torch.Tensor, object, int, int | None, str, str, str, bool]:
-    from helion.language.tile_proxy import Tile
+    from .tile_proxy import Tile
 
     valid_ops = {"atomic_add", "atomic_xchg", "atomic_cas"}
     valid_sems = {"relaxed", "release", "acq_rel"}
diff --git a/helion/language/tile_proxy.py b/helion/language/tile_proxy.py
@@ -11,8 +11,8 @@
 import torch
 from torch.utils._pytree import tree_map_only
 
-from helion import exc
-from helion._compiler.compile_environment import CompileEnvironment
+from .. import exc
+from .._compiler.compile_environment import CompileEnvironment
 
 if TYPE_CHECKING:
     from collections.abc import Callable
diff --git a/helion/runtime/config.py b/helion/runtime/config.py
@@ -7,8 +7,8 @@
 from typing import Literal
 from typing import cast
 
-from helion.autotuner.config_spec import DEFAULT_NUM_STAGES
-from helion.autotuner.config_spec import DEFAULT_NUM_WARPS
+from ..autotuner.config_spec import DEFAULT_NUM_STAGES
+from ..autotuner.config_spec import DEFAULT_NUM_WARPS
 
 IndexingLiteral = Literal["pointer", "tensor_descriptor", "block_ptr"]
 PidTypeLiteral = Literal["flat", "xyz", "persistent_blocked", "persistent_interleaved"]
