[Wave] Migrate iree.turbine.support -> wave_lang.support (#48)

harsh-nod · web-flow · commit cd0b5d47fe71 · 2025-07-20T13:17:34.000-07:00
Signed-off-by: Harsh Menon &lt;harsh@nod-labs.com&gt;
diff --git a/docs/core/support.rst b/docs/core/support.rst
@@ -1,15 +1,15 @@
-`iree.turbine.support`
+`wave_lang.support`
 ======================
 
-.. automodule:: iree.turbine.support
+.. automodule:: wave_lang.support
   :imported-members:
   :members:
   :undoc-members:
 
 conversions
 --------------
 
-.. py:module:: iree.turbine.support.conversions
+.. py:module:: wave_lang.support.conversions
 
 .. autofunction:: dtype_to_element_type
 .. autofunction:: torch_dtype_to_numpy
@@ -18,7 +18,7 @@ conversions
 debugging
 --------------
 
-.. py:module:: iree.turbine.support.debugging
+.. py:module:: wave_lang.support.debugging
 
 .. autoclass:: DebugFlags
   :members:
@@ -27,15 +27,15 @@ debugging
 logging
 --------------
 
-.. py:module:: iree.turbine.support.logging
+.. py:module:: wave_lang.support.logging
 
 .. autoclass:: DefaultFormatter
 .. autofunction:: get_logger
 
 tools
 --------------
 
-.. py:module:: iree.turbine.support.tools
+.. py:module:: wave_lang.support.tools
 
 .. autofunction:: iree_tool_format_cli_input_arg
 .. autofunction:: iree_tool_prepare_input_args
diff --git a/iree/turbine/kernel/wave/analysis/index_sequence_analysis.py b/iree/turbine/kernel/wave/analysis/index_sequence_analysis.py
@@ -57,7 +57,7 @@
     print_trace,
     try_apply_pass,
 )
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 
 from copy import deepcopy, copy
 from enum import Enum
diff --git a/iree/turbine/kernel/wave/analysis/partition_strided_operators.py b/iree/turbine/kernel/wave/analysis/partition_strided_operators.py
@@ -19,7 +19,7 @@
 )
 from ..._support.indexing import IndexSequence, IndexSymbol
 from ..._support.tracing import CapturedTrace
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from ..utils.general_utils import (
     all_equal,
 )
diff --git a/iree/turbine/kernel/wave/expansion/expansion.py b/iree/turbine/kernel/wave/expansion/expansion.py
@@ -49,7 +49,7 @@
     get_users,
     get_inputs,
 )
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from copy import deepcopy
 import math
 
diff --git a/iree/turbine/kernel/wave/hoisting.py b/iree/turbine/kernel/wave/hoisting.py
@@ -6,7 +6,7 @@
 
 from .constraints import Constraint
 from .utils.general_utils import get_induction_variable
-from ...support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from iree.turbine.kernel._support.tracing import CapturedTrace
 import torch.fx as fx
 from ..ops.wave_ops import *
diff --git a/iree/turbine/kernel/wave/iree_utils.py b/iree/turbine/kernel/wave/iree_utils.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import torch
-from ...support.conversions import TORCH_DTYPE_TO_IREE_TYPE_ASM
+from wave_lang.support.conversions import TORCH_DTYPE_TO_IREE_TYPE_ASM
 from iree.turbine.runtime.launch import Launchable
 
 
diff --git a/iree/turbine/kernel/wave/memory_analysis/solver.py b/iree/turbine/kernel/wave/memory_analysis/solver.py
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 
 logger = get_logger("wave.memory_analysis.solver")
 
diff --git a/iree/turbine/kernel/wave/profiling.py b/iree/turbine/kernel/wave/profiling.py
@@ -15,7 +15,7 @@
 
 import torch
 import iree.runtime
-from ...support.conversions import TORCH_DTYPE_TO_SIGNED_MLIR_TYPE_ASM
+from wave_lang.support.conversions import TORCH_DTYPE_TO_SIGNED_MLIR_TYPE_ASM
 from .compile_options import WaveCompileOptions
 
 BenchmarkResult = namedtuple(
diff --git a/iree/turbine/kernel/wave/promotion.py b/iree/turbine/kernel/wave/promotion.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._support.tracing import CapturedTrace
-from ...support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from .._support.indexing import IndexingContext
 from ..ops.wave_ops import *
 from ..lang.global_symbols import *
diff --git a/iree/turbine/kernel/wave/scheduling/loop_reconstruction.py b/iree/turbine/kernel/wave/scheduling/loop_reconstruction.py
@@ -18,7 +18,7 @@
 import torch.fx as fx
 from collections import deque, defaultdict
 from ..visualization import visualize_mapped_graphs, visualize_graph
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from .loop_reconstruction_utils import (
     ArgumentContext,
     create_fill_stage_schedule,
diff --git a/iree/turbine/kernel/wave/scheduling/loop_reconstruction_utils.py b/iree/turbine/kernel/wave/scheduling/loop_reconstruction_utils.py
@@ -16,7 +16,7 @@
 import math
 from collections import defaultdict, deque, ChainMap
 from ..visualization import visualize_mapped_graphs
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from ...lang.global_symbols import SHARED_ADDRESS_SPACE
 import random
 from typing import Optional
diff --git a/iree/turbine/kernel/wave/scheduling/modulo_scheduling.py b/iree/turbine/kernel/wave/scheduling/modulo_scheduling.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import torch.fx as fx
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 from .graph_utils import (
     Edge,
     EdgeWeight,
diff --git a/iree/turbine/kernel/wave/scheduling/schedule.py b/iree/turbine/kernel/wave/scheduling/schedule.py
@@ -41,7 +41,7 @@
     SHUFFLE_UNITS,
 )
 import torch.fx as fx
-from ....support.logging import get_logger
+from wave_lang.support.logging import get_logger
 import json
 import os
 
diff --git a/iree/turbine/kernel/wave/type_inference.py b/iree/turbine/kernel/wave/type_inference.py
@@ -7,7 +7,7 @@
 from ..ops.wave_ops import *
 from .._support.tracing import CapturedTrace
 import torch.fx as fx
-from ...support.logging import get_logger
+from wave_lang.support.logging import get_logger
 
 logger = get_logger("turbine.wave.type_inference")
 
diff --git a/iree/turbine/runtime/device.py b/iree/turbine/runtime/device.py
@@ -31,7 +31,7 @@
     get_driver,
 )
 
-from ..support.conversions import (
+from wave_lang.support.conversions import (
     dtype_to_element_type,
     torch_dtype_to_numpy,
 )
@@ -58,7 +58,7 @@ def __init__(self, torch_device):
         )
 
 
-from ..support.logging import runtime_logger as logger
+from wave_lang.support.logging import runtime_logger as logger
 
 __all__ = [
     "get_vm_instance",
diff --git a/iree/turbine/runtime/launch.py b/iree/turbine/runtime/launch.py
@@ -29,7 +29,7 @@
     VmVariantList,
 )
 
-from ..support.logging import runtime_logger as logger
+from wave_lang.support.logging import runtime_logger as logger
 
 from .device import (
     get_device_from_torch,
diff --git a/iree/turbine/runtime/op_reg/base.py b/iree/turbine/runtime/op_reg/base.py
@@ -20,7 +20,7 @@
 import torch
 from torch import Tensor
 
-from ...support.ir_imports import (
+from wave_lang.support.ir_imports import (
     Block,
     Context,
     FunctionType,
@@ -37,9 +37,9 @@
     func_d,
 )
 
-from ...support.logging import runtime_logger as logger
+from wave_lang.support.logging import runtime_logger as logger
 
-from ...support.conversions import (
+from wave_lang.support.conversions import (
     TORCH_DTYPE_TO_IREE_TYPE_ASM,
 )
 
diff --git a/iree/turbine/runtime/op_reg/impl_helper.py b/iree/turbine/runtime/op_reg/impl_helper.py
@@ -24,9 +24,9 @@ def generate(kb: KernelBuilder):
 import logging
 import textwrap
 
-from ...support.logging import runtime_logger as logger
+from wave_lang.support.logging import runtime_logger as logger
 
-from ...support.ir_imports import (
+from wave_lang.support.ir_imports import (
     FlatSymbolRefAttr,
     FunctionType,
     MLIRError,
diff --git a/lit_tests/lit.cfg.py b/lit_tests/lit.cfg.py
@@ -6,7 +6,7 @@
 import lit.llvm
 import lit.util
 
-from iree.turbine.support.logging import get_logger
+from wave_lang.support.logging import get_logger
 
 logger = get_logger("turbine.lit_tests")
 
diff --git a/wave_lang/aot/builtins/jittable.py b/wave_lang/aot/builtins/jittable.py
@@ -27,7 +27,7 @@
     GraphNodeImporter,
     InputInfo,
 )
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     FlatSymbolRefAttr,
     FunctionType,
     Operation,
@@ -38,7 +38,7 @@
     func_d,
     util_d,
 )
-from iree.turbine.support.logging import aot_logger as logger
+from wave_lang.support.logging import aot_logger as logger
 
 from ..decompositions import current_aot_decompositions
 from ..passes import (
diff --git a/wave_lang/aot/compiled_module.py b/wave_lang/aot/compiled_module.py
@@ -17,7 +17,7 @@
 
 from torch.export import ExportedProgram
 
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     Context,
     Location,
     MLIRError,
@@ -26,7 +26,7 @@
     PassManager,
     StringAttr,
 )
-from iree.turbine.support.logging import aot_logger as logger
+from wave_lang.support.logging import aot_logger as logger
 from wave_lang.transforms.general.custom_op_expansion import ExpandCustomOpsPass
 
 from . import builtins
diff --git a/wave_lang/aot/exporter.py b/wave_lang/aot/exporter.py
@@ -18,7 +18,7 @@
     Session,
     Source,
 )
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     Context,
     Operation,
 )
diff --git a/wave_lang/aot/support/ir_utils.py b/wave_lang/aot/support/ir_utils.py
@@ -20,10 +20,13 @@
     ContextCache,
     RefTracker,
 )
-from iree.turbine.support.conversions import (
+from wave_lang.dynamo.type_conversion import (
+    NativeTypeConverter,
+)
+from wave_lang.support.conversions import (
     TORCH_DTYPE_TO_IREE_TYPE,
 )
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     ArrayAttr,
     Attribute,
     BF16Type,
@@ -61,10 +64,7 @@
     func_d,
     tensor_d,
 )
-from iree.turbine.support.logging import aot_logger as logger
-from wave_lang.dynamo.type_conversion import (
-    NativeTypeConverter,
-)
+from wave_lang.support.logging import aot_logger as logger
 
 from ..tensor_traits import (
     DeviceAffinity,
diff --git a/wave_lang/aot/support/procedural/base.py b/wave_lang/aot/support/procedural/base.py
@@ -18,7 +18,7 @@
 import torch
 from torch.utils._pytree import tree_map
 
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     F32Type,
     F64Type,
     IndexType,
diff --git a/wave_lang/aot/support/procedural/exported_program.py b/wave_lang/aot/support/procedural/exported_program.py
@@ -28,7 +28,7 @@
     GraphNodeImporter,
     InputInfo,
 )
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     FlatSymbolRefAttr,
     FunctionType,
     IrType,
@@ -39,7 +39,7 @@
     func_d,
     util_d,
 )
-from iree.turbine.support.logging import aot_logger as logger
+from wave_lang.support.logging import aot_logger as logger
 
 from ...tensor_traits import (
     DeviceAffinity,
diff --git a/wave_lang/aot/support/procedural/globals.py b/wave_lang/aot/support/procedural/globals.py
@@ -19,13 +19,13 @@
     tree_unflatten,
 )
 
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     IrType,
     Operation,
     Value,
     util_d,
 )
-from iree.turbine.support.logging import aot_logger as logger
+from wave_lang.support.logging import aot_logger as logger
 
 from ..ir_utils import (
     GlobalAttributes,
diff --git a/wave_lang/aot/support/procedural/iree_emitter.py b/wave_lang/aot/support/procedural/iree_emitter.py
@@ -12,10 +12,10 @@
 
 import torch
 
-from iree.turbine.support.conversions import (
+from wave_lang.support.conversions import (
     TORCH_DTYPE_TO_IREE_TYPE,
 )
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     IndexType,
     IntegerType,
     IrType,
diff --git a/wave_lang/aot/support/procedural/primitives.py b/wave_lang/aot/support/procedural/primitives.py
@@ -17,7 +17,7 @@
 
 import torch
 
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     F32Type,
     IrType,
     RankedTensorType,
diff --git a/wave_lang/aot/support/procedural/tracer.py b/wave_lang/aot/support/procedural/tracer.py
@@ -20,14 +20,14 @@
     treespec_dumps,
 )
 
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     DictAttr,
     Location,
     StringAttr,
     Value,
     func_d,
 )
-from iree.turbine.support.logging import aot_logger as logger
+from wave_lang.support.logging import aot_logger as logger
 
 from ...tensor_traits import DeviceAffinity
 from ..ir_utils import ModuleBuilder, attributes_from_argument_device_affinities
diff --git a/wave_lang/dynamo/type_conversion.py b/wave_lang/dynamo/type_conversion.py
@@ -13,7 +13,7 @@
 import re
 from typing import List, Optional
 
-from iree.turbine.support.ir_imports import (
+from wave_lang.support.ir_imports import (
     Context,
     F64Type,
     IntegerType,
diff --git a/wave_lang/support/__init__.py b/wave_lang/support/__init__.py
diff --git a/wave_lang/support/conversions.py b/wave_lang/support/conversions.py
diff --git a/wave_lang/support/debugging.py b/wave_lang/support/debugging.py
diff --git a/wave_lang/support/ir_imports.py b/wave_lang/support/ir_imports.py
diff --git a/wave_lang/support/logging.py b/wave_lang/support/logging.py
diff --git a/wave_lang/tools/interpreter.py b/wave_lang/tools/interpreter.py
diff --git a/wave_lang/transforms/general/custom_op_expansion.py b/wave_lang/transforms/general/custom_op_expansion.py

Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,7 @@`
`57`	`57`	`print_trace,`
`58`	`58`	`try_apply_pass,`
`59`	`59`	`)`
`60`		`-from ....support.logging import get_logger`
	`60`	`+from wave_lang.support.logging import get_logger`
`61`	`61`
`62`	`62`	`from copy import deepcopy, copy`
`63`	`63`	`from enum import Enum`
Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@`
`19`	`19`	`)`
`20`	`20`	`from ..._support.indexing import IndexSequence, IndexSymbol`
`21`	`21`	`from ..._support.tracing import CapturedTrace`
`22`		`-from ....support.logging import get_logger`
	`22`	`+from wave_lang.support.logging import get_logger`
`23`	`23`	`from ..utils.general_utils import (`
`24`	`24`	`all_equal,`
`25`	`25`	`)`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,7 @@`
`49`	`49`	`get_users,`
`50`	`50`	`get_inputs,`
`51`	`51`	`)`
`52`		`-from ....support.logging import get_logger`
	`52`	`+from wave_lang.support.logging import get_logger`
`53`	`53`	`from copy import deepcopy`
`54`	`54`	`import math`
`55`	`55`