Import updates

stack-info: PR: #953, branch: jansel/stack/198

Author: jansel

helion/_compiler/reduction_strategy.py

@@ -26,12 +26,12 @@
 from .tile_strategy import PersistentReductionState
 from .tile_strategy import TileStrategy
 
-ARG_REDUCE_MAP = {"argmax": ("max", "maximum"), "argmin": ("min", "minimum")}
-
 if TYPE_CHECKING:
     from .device_function import DeviceFunction
     from .inductor_lowering import CodegenState
 
+ARG_REDUCE_MAP = {"argmax": ("max", "maximum"), "argmin": ("min", "minimum")}
+
 
 class ReductionStrategy(TileStrategy):
     def __init__(

helion/_compiler/type_propagation.py

@@ -22,6 +22,7 @@
 from torch.utils._pytree import tree_map_only
 
 from .. import exc
+from .. import language as language_module
 from ..autotuner.config_fragment import ConfigSpecFragment
 from ..autotuner.config_spec import BlockSizeSpec
 from ..language._decorators import get_device_func_replacement
@@ -54,7 +55,6 @@
 from .variable_origin import Origin
 from .variable_origin import SourceOrigin
 from .variable_origin import TensorSizeOrigin
-import helion
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -97,7 +97,7 @@ def _get(self, name: str) -> TypeInfo:
             else:
                 raise exc.UndefinedVariable(name) from None
         else:
-            if value is helion.language:
+            if value is language_module:
                 origin = GlobalOrigin(name="hl", function=self.function)
                 return TypeInfo.from_example(value, origin)
             if name in library_imports:

helion/_testing.py

@@ -20,13 +20,13 @@
 from torch.utils._pytree import tree_map
 import triton
 
+from ._compat import get_tensor_descriptor_fn_name
 from ._utils import counters
+from .autotuner.benchmarking import compute_repeat
+from .autotuner.benchmarking import interleaved_bench
 from .runtime.config import Config
-import helion
-from helion._compat import get_tensor_descriptor_fn_name
-from helion.autotuner.benchmarking import compute_repeat
-from helion.autotuner.benchmarking import interleaved_bench
-from helion.runtime.ref_mode import is_ref_mode_enabled
+from .runtime.ref_mode import is_ref_mode_enabled
+from .runtime.settings import RefMode
 
 if TYPE_CHECKING:
     import types
@@ -93,7 +93,7 @@ def track_run_ref_calls() -> Generator[list[int], None, None]:
     Yields:
         A list that will contain the count of run_ref calls.
     """
-    from helion.runtime.kernel import BoundKernel
+    from .runtime.kernel import BoundKernel
 
     original_run_ref = BoundKernel.run_ref
     run_ref_count = [0]
@@ -112,7 +112,7 @@ def tracked_run_ref(self: BoundKernel, *args: object) -> object:
 
 @contextlib.contextmanager
 def assert_helion_ref_mode(
-    ref_mode: helion.RefMode = helion.RefMode.OFF,
+    ref_mode: RefMode = RefMode.OFF,
 ) -> Generator[None, None, None]:
     """Context manager that asserts Helion compilation behavior based on RefMode.
 
@@ -122,12 +122,12 @@ def assert_helion_ref_mode(
     with track_run_ref_calls() as run_ref_count:
         yield
 
-        if ref_mode == helion.RefMode.OFF:
+        if ref_mode == RefMode.OFF:
             # In normal mode (RefMode.OFF), run_ref should not be called
             assert run_ref_count[0] == 0, (
                 f"Expected run_ref to not be called in normal mode (RefMode.OFF), but got: run_ref={run_ref_count[0]}"
             )
-        elif ref_mode == helion.RefMode.EAGER:
+        elif ref_mode == RefMode.EAGER:
             # In ref eager mode (RefMode.EAGER), run_ref should be called
             assert run_ref_count[0] > 0, (
                 f"Expected run_ref to be called in ref eager mode (RefMode.EAGER), but got: run_ref={run_ref_count[0]}"
@@ -137,11 +137,11 @@ def assert_helion_ref_mode(
 
 
 assert_helion_compilation = functools.partial(
-    assert_helion_ref_mode, ref_mode=helion.RefMode.OFF
+    assert_helion_ref_mode, ref_mode=RefMode.OFF
 )
 
 assert_ref_eager_mode = functools.partial(
-    assert_helion_ref_mode, ref_mode=helion.RefMode.EAGER
+    assert_helion_ref_mode, ref_mode=RefMode.EAGER
 )
 
 

helion/_utils.py

@@ -35,7 +35,7 @@ def convert_size_arg(size: object) -> object:
     - Other values -> unchanged
     """
     # Import here to avoid circular dependency
-    from helion.language.ref_tile import RefTile
+    from .language.ref_tile import RefTile
 
     if isinstance(size, (list, tuple)):
         return [convert_size_arg(item) for item in size]
@@ -54,7 +54,7 @@ def convert_tile_indices_to_slices(index: object) -> object:
         Index with RefTile objects replaced by their slice objects
     """
     # Import here to avoid circular dependency
-    from helion.language.ref_tile import RefTile
+    from .language.ref_tile import RefTile
 
     def _extract_slice(obj: object) -> object:
         return obj._slice if isinstance(obj, RefTile) else obj

helion/autotuner/base_search.py

@@ -47,8 +47,6 @@
 from .logger import format_triton_compile_failure
 from .progress_bar import iter_with_progress
 
-log = logging.getLogger(__name__)
-
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
@@ -58,6 +56,8 @@
     from ..runtime.settings import Settings
     from . import ConfigSpec
 
+log = logging.getLogger(__name__)
+
 
 class BaseAutotuner(abc.ABC):
     """

helion/language/memory_ops.py

@@ -9,7 +9,7 @@
 from .. import exc
 from .._compiler.indexing_strategy import SubscriptIndexing
 from . import _decorators
-from helion.language.stack_tensor import StackTensor
+from .stack_tensor import StackTensor
 
 if TYPE_CHECKING:
     from .._compiler.inductor_lowering import CodegenState

helion/language/signal_wait.py

@@ -9,7 +9,7 @@
 from .. import exc
 from .._compiler.indexing_strategy import SubscriptIndexing
 from . import _decorators
-from helion.language.stack_tensor import StackTensor
+from .stack_tensor import StackTensor
 
 if TYPE_CHECKING:
     import ast

helion/runtime/ref_mode.py

@@ -11,11 +11,11 @@
 import torch
 from torch.overrides import BaseTorchFunctionMode
 
-from helion._compiler.compile_environment import CompileEnvironment
-from helion._compiler.compile_environment import NoCurrentEnvironment
-from helion._compiler.compile_environment import tls as ce_tls
-from helion._utils import convert_size_arg
-from helion._utils import create_shape_matching_slices
+from .._compiler.compile_environment import CompileEnvironment
+from .._compiler.compile_environment import NoCurrentEnvironment
+from .._compiler.compile_environment import tls as ce_tls
+from .._utils import convert_size_arg
+from .._utils import create_shape_matching_slices
 
 if TYPE_CHECKING:
     from typing_extensions import Self

helion/runtime/settings.py

@@ -14,10 +14,10 @@
 import torch
 from torch._environment import is_fbcode
 
-from helion import exc
-from helion.autotuner.effort_profile import AutotuneEffort
-from helion.autotuner.effort_profile import get_effort_profile
-from helion.runtime.ref_mode import RefMode
+from .. import exc
+from ..autotuner.effort_profile import AutotuneEffort
+from ..autotuner.effort_profile import get_effort_profile
+from .ref_mode import RefMode
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager