[Wave] Migrate iree.turbine.transforms -> wave.transforms (#40)

Signed-off-by: Harsh Menon <[email protected]>

Author: harsh-nod

iree/turbine/aot/compiled_module.py

@@ -29,7 +29,7 @@
     StringAttr,
 )
 from ..support.logging import aot_logger as logger
-from ..transforms.general.custom_op_expansion import ExpandCustomOpsPass
+from wave.transforms.general.custom_op_expansion import ExpandCustomOpsPass
 
 from .support.procedural import (
     GlobalsDef,

iree/turbine/runtime/op_reg/impl_helper.py

@@ -36,7 +36,7 @@ def generate(kb: KernelBuilder):
     Value,
 )
 
-from ...transforms.merger import Merger
+from wave.transforms.merger import Merger
 
 from .base import (
     KernelBuilder,

lit_tests/kernel/wave/sharktank_integration.py

@@ -2,6 +2,7 @@
 
 import textwrap
 from typing import Optional
+from wave.transforms.merger import Merger
 
 import torch
 from jinja2 import BaseLoader, Environment
@@ -43,7 +44,6 @@
 from iree.turbine.runtime.op_reg.impl_helper import (
     call_function,
 )
-from iree.turbine.transforms.merger import Merger
 
 _JINJA2_ENVIRONMENT: Optional[Environment] = None
 

iree/turbine/transforms/general/custom_op_expansion.py renamed to wave/transforms/general/custom_op_expansion.py

@@ -5,40 +5,37 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from typing import Callable, Optional
+from wave.dynamo.type_conversion import (
+    NativeTypeConverter,
+)
 
 import torch
 from torch import Tensor
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
-from wave.dynamo.type_conversion import (
-    NativeTypeConverter,
-)
-
-from ...runtime.op_reg.base import (
+from iree.turbine.runtime.op_reg.base import (
     ALL_CUSTOM_OP_REGS,
     AttrArg,
+    CustomOp,
     EmptyOptionalTensorArg,
     IntArg,
-    CustomOp,
     KernelBuilder,
     KernelSelection,
     TensorArg,
     TensorListArg,
 )
-
-from ...support.conversions import (
+from iree.turbine.support.conversions import (
     MLIR_TYPE_ASM_TO_TORCH_DTYPE,
 )
-
-from ...support.ir_imports import (
+from iree.turbine.support.ir_imports import (
     Block,
     FloatAttr,
-    IrType,
     InsertionPoint,
     IntegerAttr,
-    OpResult,
+    IrType,
     Operation,
+    OpResult,
     RankedTensorType,
     StringAttr,
     SymbolTable,
@@ -70,7 +67,7 @@ def funcs(self):
 
     def erase_unused_op(self, op: Operation):
         """Recursively erases any unused torch ops, starting with op."""
-        from ...support.ir_imports import OpResult
+        from iree.turbine.support.ir_imports import OpResult
 
         worklist = set()
         worklist.add(op)

iree/turbine/transforms/merger.py renamed to wave/transforms/merger.py

@@ -14,7 +14,6 @@
     SymbolTable,
 )
 
-
 __all__ = [
     "Merger",
 ]