Rename shark-turbine -> iree.turbine (#197)

* Move files from files from `shark-turbine` to `iree/turbine`. * Update imports * Update `setup.py` * Make backward redirect `shark-turbine` -> `iree.turbine` (do we need this?) Progress on #28 --------- Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
iree-org · Oct 5, 2024 · 40016ad · 40016ad
1 parent 3312f73
commit 40016ad
Show file tree

Hide file tree

Showing 192 changed files with 250 additions and 807 deletions.
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -2,4 +2,4 @@ include README.md
 include requirements.txt
 include pytorch-cpu-requirements.txt
 include version_info.json
-include shark_turbine/ops/templates/*.mlir
+include iree/turbine/ops/templates/*.mlir
diff --git a/README.md b/README.md
@@ -8,7 +8,7 @@ Turbine provides a collection of tools:
 
 * *AOT Export*: For compiling one or more `nn.Module`s to compiled, deployment
   ready artifacts. This operates via both a simple one-shot export API (Already upstreamed to [torch-mlir](https://github.com/llvm/torch-mlir/blob/main/python/torch_mlir/extras/fx_importer.py))
-  for simple models and an underlying [advanced API](shark_turbine/aot/compiled_module.py) for complicated models
+  for simple models and an underlying [advanced API](iree/turbine/aot/compiled_module.py) for complicated models
   and accessing the full features of the runtime.
 * *Eager Execution*: A `torch.compile` backend is provided and a Turbine Tensor/Device
   is available for more native, interactive use within a PyTorch session.

diff --git a/build_tools/build_release.py b/build_tools/build_release.py
@@ -159,10 +159,8 @@ def main():
         print("Downloading remaining requirements")
         download_requirements(REPO_ROOT / "requirements.txt")
 
-    print("Building shark-turbine")
-    build_wheel(REPO_ROOT)
     print("Building iree-turbine")
-    build_wheel(REPO_ROOT, env={"TURBINE_PACKAGE_NAME": "iree-turbine"})
+    build_wheel(REPO_ROOT)
 
 
 if __name__ == "__main__":

diff --git a/examples/aot_mlp/mlp_export_dynamic.py b/examples/aot_mlp/mlp_export_dynamic.py
@@ -12,7 +12,7 @@
 import torch
 import torch.nn as nn
 
-import shark_turbine.aot as aot
+import iree.turbine.aot as aot
 
 
 class MLP(nn.Module):

diff --git a/examples/aot_mlp/mlp_export_simple.py b/examples/aot_mlp/mlp_export_simple.py
@@ -9,7 +9,7 @@
 import torch
 import torch.nn as nn
 
-import shark_turbine.aot as aot
+import iree.turbine.aot as aot
 
 
 class MLP(nn.Module):

diff --git a/examples/llama2_inference/README.md b/examples/llama2_inference/README.md
diff --git a/examples/llama2_inference/llama2.ipynb b/examples/llama2_inference/llama2.ipynb
diff --git a/examples/llama2_inference/llama2_state_schema.json b/examples/llama2_inference/llama2_state_schema.json
diff --git a/examples/llama2_inference/requirements.txt b/examples/llama2_inference/requirements.txt
diff --git a/examples/resnet-18/requirements.txt b/examples/resnet-18/requirements.txt
@@ -1,2 +1,2 @@
 transformers
-shark_turbine==0.9.2
+iree_turbine==0.9.2
diff --git a/examples/resnet-18/resnet-18.py b/examples/resnet-18/resnet-18.py
@@ -1,6 +1,6 @@
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 import torch
-from shark_turbine.aot import *
+from iree.turbine.aot import *
 import iree.runtime as rt
 
 # Loading feature extractor and pretrained model from huggingface

diff --git a/examples/runtime_torture/launchable_torture.py b/examples/runtime_torture/launchable_torture.py
@@ -12,9 +12,9 @@
 import torch
 import torch.nn as nn
 
-import shark_turbine.aot as aot
+import iree.turbine.aot as aot
 
-from shark_turbine.runtime import (
+from iree.turbine.runtime import (
     Launchable,
 )
 

diff --git a/iree/turbine/__init__.py b/iree/turbine/__init__.py
@@ -8,15 +8,3 @@
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-# TODO: This redirection layer exists while we are migrating from the
-# shark_turbine top-level package name to iree.turbine. It exports the
-# public API but not the internal details. In a future switch, all code
-# will be directly located here and the redirect will be done in the
-# shark_turbine namespace.
-
-from shark_turbine import aot
-from shark_turbine import dynamo
-from shark_turbine import kernel
-from shark_turbine import ops
-from shark_turbine import runtime
diff --git a/shark_turbine/aot/__init__.py → iree/turbine/aot/__init__.py b/shark_turbine/aot/__init__.py → iree/turbine/aot/__init__.py
diff --git a/shark_turbine/aot/builtins/__init__.py → iree/turbine/aot/builtins/__init__.py b/shark_turbine/aot/builtins/__init__.py → iree/turbine/aot/builtins/__init__.py
diff --git a/shark_turbine/aot/builtins/globals.py → iree/turbine/aot/builtins/globals.py b/shark_turbine/aot/builtins/globals.py → iree/turbine/aot/builtins/globals.py
diff --git a/shark_turbine/aot/builtins/jittable.py → iree/turbine/aot/builtins/jittable.py b/shark_turbine/aot/builtins/jittable.py → iree/turbine/aot/builtins/jittable.py
diff --git a/shark_turbine/aot/compiled_module.py → iree/turbine/aot/compiled_module.py b/shark_turbine/aot/compiled_module.py → iree/turbine/aot/compiled_module.py
diff --git a/shark_turbine/aot/decompositions.py → iree/turbine/aot/decompositions.py b/shark_turbine/aot/decompositions.py → iree/turbine/aot/decompositions.py
diff --git a/shark_turbine/aot/exporter.py → iree/turbine/aot/exporter.py b/shark_turbine/aot/exporter.py → iree/turbine/aot/exporter.py
diff --git a/shark_turbine/aot/fx_programs.py → iree/turbine/aot/fx_programs.py b/shark_turbine/aot/fx_programs.py → iree/turbine/aot/fx_programs.py
diff --git a/shark_turbine/aot/params.py → iree/turbine/aot/params.py b/shark_turbine/aot/params.py → iree/turbine/aot/params.py
diff --git a/shark_turbine/aot/passes/__init__.py → iree/turbine/aot/passes/__init__.py b/shark_turbine/aot/passes/__init__.py → iree/turbine/aot/passes/__init__.py
diff --git a/shark_turbine/aot/passes/functorch.py → iree/turbine/aot/passes/functorch.py b/shark_turbine/aot/passes/functorch.py → iree/turbine/aot/passes/functorch.py
diff --git a/shark_turbine/aot/support/ir_utils.py → iree/turbine/aot/support/ir_utils.py b/shark_turbine/aot/support/ir_utils.py → iree/turbine/aot/support/ir_utils.py
diff --git a/...urbine/aot/support/procedural/__init__.py → ...urbine/aot/support/procedural/__init__.py b/...urbine/aot/support/procedural/__init__.py → ...urbine/aot/support/procedural/__init__.py
diff --git a/shark_turbine/aot/support/procedural/base.py → iree/turbine/aot/support/procedural/base.py b/shark_turbine/aot/support/procedural/base.py → iree/turbine/aot/support/procedural/base.py
diff --git a/...ot/support/procedural/exported_program.py → ...ot/support/procedural/exported_program.py b/...ot/support/procedural/exported_program.py → ...ot/support/procedural/exported_program.py
diff --git a/...turbine/aot/support/procedural/globals.py → ...turbine/aot/support/procedural/globals.py b/...turbine/aot/support/procedural/globals.py → ...turbine/aot/support/procedural/globals.py
diff --git a/...ne/aot/support/procedural/iree_emitter.py → ...ne/aot/support/procedural/iree_emitter.py b/...ne/aot/support/procedural/iree_emitter.py → ...ne/aot/support/procedural/iree_emitter.py
diff --git a/...bine/aot/support/procedural/primitives.py → ...bine/aot/support/procedural/primitives.py b/...bine/aot/support/procedural/primitives.py → ...bine/aot/support/procedural/primitives.py
diff --git a/..._turbine/aot/support/procedural/tracer.py → .../turbine/aot/support/procedural/tracer.py b/..._turbine/aot/support/procedural/tracer.py → .../turbine/aot/support/procedural/tracer.py
diff --git a/shark_turbine/aot/tensor_traits.py → iree/turbine/aot/tensor_traits.py b/shark_turbine/aot/tensor_traits.py → iree/turbine/aot/tensor_traits.py
diff --git a/shark_turbine/dynamo/__init__.py → iree/turbine/dynamo/__init__.py b/shark_turbine/dynamo/__init__.py → iree/turbine/dynamo/__init__.py
diff --git a/shark_turbine/dynamo/backends/cpu.py → iree/turbine/dynamo/backends/cpu.py b/shark_turbine/dynamo/backends/cpu.py → iree/turbine/dynamo/backends/cpu.py
diff --git a/shark_turbine/dynamo/decompositions.py → iree/turbine/dynamo/decompositions.py b/shark_turbine/dynamo/decompositions.py → iree/turbine/dynamo/decompositions.py
diff --git a/shark_turbine/dynamo/executor.py → iree/turbine/dynamo/executor.py b/shark_turbine/dynamo/executor.py → iree/turbine/dynamo/executor.py
diff --git a/shark_turbine/dynamo/passes.py → iree/turbine/dynamo/passes.py b/shark_turbine/dynamo/passes.py → iree/turbine/dynamo/passes.py
diff --git a/shark_turbine/dynamo/tensor.py → iree/turbine/dynamo/tensor.py b/shark_turbine/dynamo/tensor.py → iree/turbine/dynamo/tensor.py
@@ -474,8 +474,8 @@ def _get_device_state() -> DeviceState:
     return DeviceState(driver="local-task")
 
 
-# Inspiration from https://github.com/nod-ai/SHARK-Turbine/blob/8293de5414889c72ff5cd10bf33c43fb0a3ea3ee/python/shark_turbine/aot/builtins/jittable.py#L212-L237
-# and https://github.com/nod-ai/SHARK-Turbine/blob/main/python/shark_turbine/dynamo/backends/cpu.py
+# Inspiration from https://github.com/nod-ai/SHARK-Turbine/blob/8293de5414889c72ff5cd10bf33c43fb0a3ea3ee/python/iree/turbine/aot/builtins/jittable.py#L212-L237
+# and https://github.com/nod-ai/SHARK-Turbine/blob/main/python/iree/turbine/dynamo/backends/cpu.py
 # TODO: Try to generalize for other devices.
 def compute_method(super_fn, *args, **kwargs):
     # Compute factory fns reserve the last arg as src_op

diff --git a/shark_turbine/dynamo/type_conversion.py → iree/turbine/dynamo/type_conversion.py b/shark_turbine/dynamo/type_conversion.py → iree/turbine/dynamo/type_conversion.py
diff --git a/shark_turbine/importers/README.md → iree/turbine/importers/README.md b/shark_turbine/importers/README.md → iree/turbine/importers/README.md
diff --git a/shark_turbine/importers/ir.py → iree/turbine/importers/ir.py b/shark_turbine/importers/ir.py → iree/turbine/importers/ir.py
diff --git a/shark_turbine/importers/utils.py → iree/turbine/importers/utils.py b/shark_turbine/importers/utils.py → iree/turbine/importers/utils.py
diff --git a/shark_turbine/kernel/__init__.py → iree/turbine/kernel/__init__.py b/shark_turbine/kernel/__init__.py → iree/turbine/kernel/__init__.py
diff --git a/shark_turbine/kernel/_support/context.py → iree/turbine/kernel/_support/context.py b/shark_turbine/kernel/_support/context.py → iree/turbine/kernel/_support/context.py
diff --git a/shark_turbine/kernel/_support/dtype.py → iree/turbine/kernel/_support/dtype.py b/shark_turbine/kernel/_support/dtype.py → iree/turbine/kernel/_support/dtype.py
diff --git a/shark_turbine/kernel/_support/indexing.py → iree/turbine/kernel/_support/indexing.py b/shark_turbine/kernel/_support/indexing.py → iree/turbine/kernel/_support/indexing.py
diff --git a/shark_turbine/kernel/_support/regions.py → iree/turbine/kernel/_support/regions.py b/shark_turbine/kernel/_support/regions.py → iree/turbine/kernel/_support/regions.py
diff --git a/shark_turbine/kernel/_support/shaped_type.py → iree/turbine/kernel/_support/shaped_type.py b/shark_turbine/kernel/_support/shaped_type.py → iree/turbine/kernel/_support/shaped_type.py
diff --git a/shark_turbine/kernel/_support/tracing.py → iree/turbine/kernel/_support/tracing.py b/shark_turbine/kernel/_support/tracing.py → iree/turbine/kernel/_support/tracing.py
diff --git a/shark_turbine/kernel/compiler/base.py → iree/turbine/kernel/compiler/base.py b/shark_turbine/kernel/compiler/base.py → iree/turbine/kernel/compiler/base.py
diff --git a/shark_turbine/kernel/compiler/builder.py → iree/turbine/kernel/compiler/builder.py b/shark_turbine/kernel/compiler/builder.py → iree/turbine/kernel/compiler/builder.py
diff --git a/...rbine/kernel/compiler/dispatch_codegen.py → ...rbine/kernel/compiler/dispatch_codegen.py b/...rbine/kernel/compiler/dispatch_codegen.py → ...rbine/kernel/compiler/dispatch_codegen.py
diff --git a/...k_turbine/kernel/compiler/host_codegen.py → iree/turbine/kernel/compiler/host_codegen.py b/...k_turbine/kernel/compiler/host_codegen.py → iree/turbine/kernel/compiler/host_codegen.py
diff --git a/shark_turbine/kernel/compiler/ir.py → iree/turbine/kernel/compiler/ir.py b/shark_turbine/kernel/compiler/ir.py → iree/turbine/kernel/compiler/ir.py
diff --git a/...turbine/kernel/compiler/kernel_codegen.py → ...turbine/kernel/compiler/kernel_codegen.py b/...turbine/kernel/compiler/kernel_codegen.py → ...turbine/kernel/compiler/kernel_codegen.py
diff --git a/shark_turbine/kernel/compiler/op_matchers.py → iree/turbine/kernel/compiler/op_matchers.py b/shark_turbine/kernel/compiler/op_matchers.py → iree/turbine/kernel/compiler/op_matchers.py
diff --git a/shark_turbine/kernel/compiler/utils.py → iree/turbine/kernel/compiler/utils.py b/shark_turbine/kernel/compiler/utils.py → iree/turbine/kernel/compiler/utils.py
diff --git a/...turbine/kernel/compiler/vector_codegen.py → ...turbine/kernel/compiler/vector_codegen.py b/...turbine/kernel/compiler/vector_codegen.py → ...turbine/kernel/compiler/vector_codegen.py
diff --git a/shark_turbine/kernel/gen/__init__.py → iree/turbine/kernel/gen/__init__.py b/shark_turbine/kernel/gen/__init__.py → iree/turbine/kernel/gen/__init__.py
diff --git a/shark_turbine/kernel/gen/kernel.py → iree/turbine/kernel/gen/kernel.py b/shark_turbine/kernel/gen/kernel.py → iree/turbine/kernel/gen/kernel.py
diff --git a/shark_turbine/kernel/gen/thread.py → iree/turbine/kernel/gen/thread.py b/shark_turbine/kernel/gen/thread.py → iree/turbine/kernel/gen/thread.py
diff --git a/shark_turbine/kernel/lang/__init__.py → iree/turbine/kernel/lang/__init__.py b/shark_turbine/kernel/lang/__init__.py → iree/turbine/kernel/lang/__init__.py
diff --git a/shark_turbine/kernel/lang/global_symbols.py → iree/turbine/kernel/lang/global_symbols.py b/shark_turbine/kernel/lang/global_symbols.py → iree/turbine/kernel/lang/global_symbols.py
diff --git a/shark_turbine/kernel/lang/grid.py → iree/turbine/kernel/lang/grid.py b/shark_turbine/kernel/lang/grid.py → iree/turbine/kernel/lang/grid.py
diff --git a/shark_turbine/kernel/lang/kernel_buffer.py → iree/turbine/kernel/lang/kernel_buffer.py b/shark_turbine/kernel/lang/kernel_buffer.py → iree/turbine/kernel/lang/kernel_buffer.py
diff --git a/shark_turbine/kernel/lang/prims.py → iree/turbine/kernel/lang/prims.py b/shark_turbine/kernel/lang/prims.py → iree/turbine/kernel/lang/prims.py
diff --git a/shark_turbine/kernel/lang/types.py → iree/turbine/kernel/lang/types.py b/shark_turbine/kernel/lang/types.py → iree/turbine/kernel/lang/types.py
diff --git a/shark_turbine/kernel/lang/wave_types.py → iree/turbine/kernel/lang/wave_types.py b/shark_turbine/kernel/lang/wave_types.py → iree/turbine/kernel/lang/wave_types.py
diff --git a/shark_turbine/kernel/ops/__init__.py → iree/turbine/kernel/ops/__init__.py b/shark_turbine/kernel/ops/__init__.py → iree/turbine/kernel/ops/__init__.py
diff --git a/shark_turbine/kernel/ops/base.py → iree/turbine/kernel/ops/base.py b/shark_turbine/kernel/ops/base.py → iree/turbine/kernel/ops/base.py
diff --git a/shark_turbine/kernel/ops/control_flow.py → iree/turbine/kernel/ops/control_flow.py b/shark_turbine/kernel/ops/control_flow.py → iree/turbine/kernel/ops/control_flow.py
diff --git a/shark_turbine/kernel/ops/core.py → iree/turbine/kernel/ops/core.py b/shark_turbine/kernel/ops/core.py → iree/turbine/kernel/ops/core.py
diff --git a/shark_turbine/kernel/ops/math.py → iree/turbine/kernel/ops/math.py b/shark_turbine/kernel/ops/math.py → iree/turbine/kernel/ops/math.py
diff --git a/shark_turbine/kernel/ops/memory.py → iree/turbine/kernel/ops/memory.py b/shark_turbine/kernel/ops/memory.py → iree/turbine/kernel/ops/memory.py
diff --git a/shark_turbine/kernel/ops/reduction.py → iree/turbine/kernel/ops/reduction.py b/shark_turbine/kernel/ops/reduction.py → iree/turbine/kernel/ops/reduction.py
diff --git a/..._turbine/kernel/ops/shape_manipulation.py → .../turbine/kernel/ops/shape_manipulation.py b/..._turbine/kernel/ops/shape_manipulation.py → .../turbine/kernel/ops/shape_manipulation.py
diff --git a/shark_turbine/kernel/ops/wave_ops.py → iree/turbine/kernel/ops/wave_ops.py b/shark_turbine/kernel/ops/wave_ops.py → iree/turbine/kernel/ops/wave_ops.py
diff --git a/shark_turbine/kernel/wave/README.md → iree/turbine/kernel/wave/README.md b/shark_turbine/kernel/wave/README.md → iree/turbine/kernel/wave/README.md
diff --git a/shark_turbine/kernel/wave/__init__.py → iree/turbine/kernel/wave/__init__.py b/shark_turbine/kernel/wave/__init__.py → iree/turbine/kernel/wave/__init__.py
diff --git a/shark_turbine/kernel/wave/barriers.py → iree/turbine/kernel/wave/barriers.py b/shark_turbine/kernel/wave/barriers.py → iree/turbine/kernel/wave/barriers.py
diff --git a/shark_turbine/kernel/wave/codegen.py → iree/turbine/kernel/wave/codegen.py b/shark_turbine/kernel/wave/codegen.py → iree/turbine/kernel/wave/codegen.py
@@ -40,10 +40,10 @@
     scf_d,
     vector_d,
 )
-from shark_turbine.aot.support.ir_utils import _is_float_type, _is_integer_like_type
+from iree.turbine.aot.support.ir_utils import _is_float_type, _is_integer_like_type
 
 # TK infrastructure imports.
-from shark_turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel.lang.global_symbols import *
 from ..ops.wave_ops import (
     write,
     broadcast,

diff --git a/shark_turbine/kernel/wave/constraints.py → iree/turbine/kernel/wave/constraints.py b/shark_turbine/kernel/wave/constraints.py → iree/turbine/kernel/wave/constraints.py
diff --git a/...rbine/kernel/wave/decompose_reduce_ops.py → ...rbine/kernel/wave/decompose_reduce_ops.py b/...rbine/kernel/wave/decompose_reduce_ops.py → ...rbine/kernel/wave/decompose_reduce_ops.py
diff --git a/..._turbine/kernel/wave/docs/gemm_example.md → .../turbine/kernel/wave/docs/gemm_example.md b/..._turbine/kernel/wave/docs/gemm_example.md → .../turbine/kernel/wave/docs/gemm_example.md
diff --git a/shark_turbine/kernel/wave/expansion.py → iree/turbine/kernel/wave/expansion.py b/shark_turbine/kernel/wave/expansion.py → iree/turbine/kernel/wave/expansion.py
diff --git a/shark_turbine/kernel/wave/hoisting.py → iree/turbine/kernel/wave/hoisting.py b/shark_turbine/kernel/wave/hoisting.py → iree/turbine/kernel/wave/hoisting.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ...support.logging import get_logger
-from shark_turbine.kernel._support.tracing import CapturedTrace
+from iree.turbine.kernel._support.tracing import CapturedTrace
 import torch.fx as fx
 from ..ops.wave_ops import *
 from ..lang.global_symbols import *

diff --git a/...ne/kernel/wave/index_sequence_analysis.py → ...ne/kernel/wave/index_sequence_analysis.py b/...ne/kernel/wave/index_sequence_analysis.py → ...ne/kernel/wave/index_sequence_analysis.py
diff --git a/shark_turbine/kernel/wave/iree_utils.py → iree/turbine/kernel/wave/iree_utils.py b/shark_turbine/kernel/wave/iree_utils.py → iree/turbine/kernel/wave/iree_utils.py
diff --git a/...bine/kernel/wave/minimize_global_loads.py → ...bine/kernel/wave/minimize_global_loads.py b/...bine/kernel/wave/minimize_global_loads.py → ...bine/kernel/wave/minimize_global_loads.py
diff --git a/shark_turbine/kernel/wave/promotion.py → iree/turbine/kernel/wave/promotion.py b/shark_turbine/kernel/wave/promotion.py → iree/turbine/kernel/wave/promotion.py
diff --git a/...urbine/kernel/wave/scheduling/__init__.py → ...urbine/kernel/wave/scheduling/__init__.py b/...urbine/kernel/wave/scheduling/__init__.py → ...urbine/kernel/wave/scheduling/__init__.py
diff --git a/...ine/kernel/wave/scheduling/graph_utils.py → ...ine/kernel/wave/scheduling/graph_utils.py b/...ine/kernel/wave/scheduling/graph_utils.py → ...ine/kernel/wave/scheduling/graph_utils.py
diff --git a/...el/wave/scheduling/loop_reconstruction.py → ...el/wave/scheduling/loop_reconstruction.py b/...el/wave/scheduling/loop_reconstruction.py → ...el/wave/scheduling/loop_reconstruction.py
diff --git a/...e/scheduling/loop_reconstruction_utils.py → ...e/scheduling/loop_reconstruction_utils.py b/...e/scheduling/loop_reconstruction_utils.py → ...e/scheduling/loop_reconstruction_utils.py
diff --git a/...rnel/wave/scheduling/modulo_scheduling.py → ...rnel/wave/scheduling/modulo_scheduling.py b/...rnel/wave/scheduling/modulo_scheduling.py → ...rnel/wave/scheduling/modulo_scheduling.py
diff --git a/...rbine/kernel/wave/scheduling/resources.py → ...rbine/kernel/wave/scheduling/resources.py b/...rbine/kernel/wave/scheduling/resources.py → ...rbine/kernel/wave/scheduling/resources.py
diff --git a/...urbine/kernel/wave/scheduling/schedule.py → ...urbine/kernel/wave/scheduling/schedule.py b/...urbine/kernel/wave/scheduling/schedule.py → ...urbine/kernel/wave/scheduling/schedule.py
diff --git a/...ine/kernel/wave/shared_memory_indexing.py → ...ine/kernel/wave/shared_memory_indexing.py b/...ine/kernel/wave/shared_memory_indexing.py → ...ine/kernel/wave/shared_memory_indexing.py
diff --git a/...bine/kernel/wave/thread_shape_analysis.py → ...bine/kernel/wave/thread_shape_analysis.py b/...bine/kernel/wave/thread_shape_analysis.py → ...bine/kernel/wave/thread_shape_analysis.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ...support.logging import get_logger
-from shark_turbine.kernel._support.tracing import CapturedTrace
+from iree.turbine.kernel._support.tracing import CapturedTrace
 import torch.fx as fx
 from ..ops.wave_ops import *
 from ..lang.global_symbols import *

diff --git a/shark_turbine/kernel/wave/utils.py → iree/turbine/kernel/wave/utils.py b/shark_turbine/kernel/wave/utils.py → iree/turbine/kernel/wave/utils.py
@@ -33,7 +33,7 @@
     TilingConstraint,
 )
 import torch.fx as fx
-import shark_turbine.kernel.lang as tkl
+import iree.turbine.kernel.lang as tkl
 
 
 import tempfile

diff --git a/shark_turbine/kernel/wave/visualization.py → iree/turbine/kernel/wave/visualization.py b/shark_turbine/kernel/wave/visualization.py → iree/turbine/kernel/wave/visualization.py
diff --git a/shark_turbine/kernel/wave/wave.py → iree/turbine/kernel/wave/wave.py b/shark_turbine/kernel/wave/wave.py → iree/turbine/kernel/wave/wave.py
@@ -42,7 +42,7 @@
 from .thread_shape_analysis import determine_thread_shapes
 from .scheduling.schedule import schedule_graph
 from .._support.indexing import IndexingContext, IndexExpr
-import shark_turbine.kernel.lang as tkl
+import iree.turbine.kernel.lang as tkl
 from .._support.tracing import (
     CapturedTrace,
     CompiledContext,

diff --git a/shark_turbine/kernel/wave/wave_sim.py → iree/turbine/kernel/wave/wave_sim.py b/shark_turbine/kernel/wave/wave_sim.py → iree/turbine/kernel/wave/wave_sim.py
diff --git a/shark_turbine/ops/__init__.py → iree/turbine/ops/__init__.py b/shark_turbine/ops/__init__.py → iree/turbine/ops/__init__.py
diff --git a/shark_turbine/ops/_jinja_test_ops.py → iree/turbine/ops/_jinja_test_ops.py b/shark_turbine/ops/_jinja_test_ops.py → iree/turbine/ops/_jinja_test_ops.py
diff --git a/shark_turbine/ops/_str_format_test_ops.py → iree/turbine/ops/_str_format_test_ops.py b/shark_turbine/ops/_str_format_test_ops.py → iree/turbine/ops/_str_format_test_ops.py
diff --git a/shark_turbine/ops/iree.py → iree/turbine/ops/iree.py b/shark_turbine/ops/iree.py → iree/turbine/ops/iree.py
diff --git a/...turbine/ops/templates/test_add_jinja.mlir → ...turbine/ops/templates/test_add_jinja.mlir b/...turbine/ops/templates/test_add_jinja.mlir → ...turbine/ops/templates/test_add_jinja.mlir
diff --git a/...ine/ops/templates/test_add_strformat.mlir → ...ine/ops/templates/test_add_strformat.mlir b/...ine/ops/templates/test_add_strformat.mlir → ...ine/ops/templates/test_add_strformat.mlir
diff --git a/...bine/ops/templates/test_syntax_error.mlir → ...bine/ops/templates/test_syntax_error.mlir b/...bine/ops/templates/test_syntax_error.mlir → ...bine/ops/templates/test_syntax_error.mlir
diff --git a/shark_turbine/runtime/__init__.py → iree/turbine/runtime/__init__.py b/shark_turbine/runtime/__init__.py → iree/turbine/runtime/__init__.py
diff --git a/shark_turbine/runtime/device.py → iree/turbine/runtime/device.py b/shark_turbine/runtime/device.py → iree/turbine/runtime/device.py
diff --git a/shark_turbine/runtime/launch.py → iree/turbine/runtime/launch.py b/shark_turbine/runtime/launch.py → iree/turbine/runtime/launch.py
diff --git a/shark_turbine/runtime/op_reg/__init__.py → iree/turbine/runtime/op_reg/__init__.py b/shark_turbine/runtime/op_reg/__init__.py → iree/turbine/runtime/op_reg/__init__.py
diff --git a/shark_turbine/runtime/op_reg/base.py → iree/turbine/runtime/op_reg/base.py b/shark_turbine/runtime/op_reg/base.py → iree/turbine/runtime/op_reg/base.py
diff --git a/shark_turbine/runtime/op_reg/compiler.py → iree/turbine/runtime/op_reg/compiler.py b/shark_turbine/runtime/op_reg/compiler.py → iree/turbine/runtime/op_reg/compiler.py
diff --git a/shark_turbine/runtime/op_reg/eager.py → iree/turbine/runtime/op_reg/eager.py b/shark_turbine/runtime/op_reg/eager.py → iree/turbine/runtime/op_reg/eager.py
diff --git a/shark_turbine/runtime/op_reg/impl_helper.py → iree/turbine/runtime/op_reg/impl_helper.py b/shark_turbine/runtime/op_reg/impl_helper.py → iree/turbine/runtime/op_reg/impl_helper.py
diff --git a/shark_turbine/runtime/tracing.py → iree/turbine/runtime/tracing.py b/shark_turbine/runtime/tracing.py → iree/turbine/runtime/tracing.py
diff --git a/shark_turbine/support/__init__.py → iree/turbine/support/__init__.py b/shark_turbine/support/__init__.py → iree/turbine/support/__init__.py
diff --git a/shark_turbine/support/conversions.py → iree/turbine/support/conversions.py b/shark_turbine/support/conversions.py → iree/turbine/support/conversions.py
diff --git a/shark_turbine/support/debugging.py → iree/turbine/support/debugging.py b/shark_turbine/support/debugging.py → iree/turbine/support/debugging.py
diff --git a/shark_turbine/support/exceptions.py → iree/turbine/support/exceptions.py b/shark_turbine/support/exceptions.py → iree/turbine/support/exceptions.py
diff --git a/shark_turbine/support/ir_imports.py → iree/turbine/support/ir_imports.py b/shark_turbine/support/ir_imports.py → iree/turbine/support/ir_imports.py
diff --git a/shark_turbine/support/logging.py → iree/turbine/support/logging.py b/shark_turbine/support/logging.py → iree/turbine/support/logging.py
diff --git a/shark_turbine/tools/__init__.py → iree/turbine/tools/__init__.py b/shark_turbine/tools/__init__.py → iree/turbine/tools/__init__.py
diff --git a/shark_turbine/tools/interpreter.py → iree/turbine/tools/interpreter.py b/shark_turbine/tools/interpreter.py → iree/turbine/tools/interpreter.py
diff --git a/shark_turbine/transforms/builder.py → iree/turbine/transforms/builder.py b/shark_turbine/transforms/builder.py → iree/turbine/transforms/builder.py
diff --git a/...urbine/transforms/general/add_metadata.py → ...urbine/transforms/general/add_metadata.py b/...urbine/transforms/general/add_metadata.py → ...urbine/transforms/general/add_metadata.py
@@ -12,7 +12,7 @@
 
 import re
 
-from shark_turbine.support.ir_imports import *
+from iree.turbine.support.ir_imports import *
 
 from ..rewriter import *
 from iree.compiler.ir import Context, DictAttr

diff --git a/...transforms/general/custom_op_expansion.py → ...transforms/general/custom_op_expansion.py b/...transforms/general/custom_op_expansion.py → ...transforms/general/custom_op_expansion.py
diff --git a/...e/transforms/general/rename_parameters.py → ...e/transforms/general/rename_parameters.py b/...e/transforms/general/rename_parameters.py → ...e/transforms/general/rename_parameters.py
diff --git a/shark_turbine/transforms/merger.py → iree/turbine/transforms/merger.py b/shark_turbine/transforms/merger.py → iree/turbine/transforms/merger.py
diff --git a/...transforms/quantization/mm_group_quant.py → ...transforms/quantization/mm_group_quant.py b/...transforms/quantization/mm_group_quant.py → ...transforms/quantization/mm_group_quant.py
diff --git a/shark_turbine/transforms/rewriter.py → iree/turbine/transforms/rewriter.py b/shark_turbine/transforms/rewriter.py → iree/turbine/transforms/rewriter.py
diff --git a/lit_tests/kernel/wave/barriers.py b/lit_tests/kernel/wave/barriers.py
@@ -3,18 +3,18 @@
 import logging
 from typing import Callable
 import unittest
-import shark_turbine.kernel as tk
-import shark_turbine.kernel.lang as tkl
-import shark_turbine.kernel.wave as tkw
-from shark_turbine.kernel.wave.promotion import promote_node, promote_placeholders
-from shark_turbine.kernel.wave.barriers import add_shared_memory_barriers
-from shark_turbine.kernel.wave.hoisting import hoist_allocs
-from shark_turbine.kernel.wave.expansion import expand_graph
-from shark_turbine.kernel.lang.global_symbols import *
-from shark_turbine.kernel._support.tracing import CapturedTrace
-from shark_turbine.kernel._support.indexing import IndexingContext
-from shark_turbine.kernel.ops.wave_ops import *
-from shark_turbine.kernel.wave.utils import run_test, print_trace
+import iree.turbine.kernel as tk
+import iree.turbine.kernel.lang as tkl
+import iree.turbine.kernel.wave as tkw
+from iree.turbine.kernel.wave.promotion import promote_node, promote_placeholders
+from iree.turbine.kernel.wave.barriers import add_shared_memory_barriers
+from iree.turbine.kernel.wave.hoisting import hoist_allocs
+from iree.turbine.kernel.wave.expansion import expand_graph
+from iree.turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel._support.tracing import CapturedTrace
+from iree.turbine.kernel._support.indexing import IndexingContext
+from iree.turbine.kernel.ops.wave_ops import *
+from iree.turbine.kernel.wave.utils import run_test, print_trace
 
 
 def get_read_nodes(graph: fx.Graph) -> list[CustomOp]:

diff --git a/lit_tests/kernel/wave/codegen.py b/lit_tests/kernel/wave/codegen.py
@@ -2,11 +2,11 @@
 
 import pytest
 from typing import Callable
-import shark_turbine.kernel as tk
-import shark_turbine.kernel.lang as tkl
-import shark_turbine.kernel.wave as tkw
-from shark_turbine.kernel.lang.global_symbols import *
-from shark_turbine.kernel.wave.utils import run_test
+import iree.turbine.kernel as tk
+import iree.turbine.kernel.lang as tkl
+import iree.turbine.kernel.wave as tkw
+from iree.turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel.wave.utils import run_test
 import torch
 
 M = tkl.sym.M

diff --git a/lit_tests/kernel/wave/expansion.py b/lit_tests/kernel/wave/expansion.py
@@ -2,13 +2,13 @@
 
 import logging
 import unittest
-import shark_turbine.kernel as tk
-import shark_turbine.kernel.lang as tkl
-import shark_turbine.kernel.wave as tkw
-from shark_turbine.kernel.wave.expansion import expand_graph
-from shark_turbine.kernel._support.indexing import IndexingContext
-from shark_turbine.kernel.lang.global_symbols import *
-from shark_turbine.kernel.wave.utils import run_test, print_trace
+import iree.turbine.kernel as tk
+import iree.turbine.kernel.lang as tkl
+import iree.turbine.kernel.wave as tkw
+from iree.turbine.kernel.wave.expansion import expand_graph
+from iree.turbine.kernel._support.indexing import IndexingContext
+from iree.turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel.wave.utils import run_test, print_trace
 import sympy
 
 # Input sizes

diff --git a/lit_tests/kernel/wave/index_sequence_analysis.py b/lit_tests/kernel/wave/index_sequence_analysis.py
@@ -2,22 +2,22 @@
 
 import logging
 import unittest
-import shark_turbine.kernel as tk
-import shark_turbine.kernel.lang as tkl
-import shark_turbine.kernel.wave as tkw
-from shark_turbine.kernel.wave.promotion import promote_placeholders
-from shark_turbine.kernel.wave.hoisting import hoist_allocs
-from shark_turbine.kernel.wave.expansion import expand_graph
-from shark_turbine.kernel.lang.global_symbols import *
-from shark_turbine.kernel._support.tracing import CapturedTrace
-from shark_turbine.kernel._support.indexing import IndexingContext
-from shark_turbine.kernel.ops.wave_ops import *
-from shark_turbine.kernel.wave.utils import run_test, print_trace
-from shark_turbine.kernel.wave.minimize_global_loads import minimize_global_loads
-from shark_turbine.kernel.wave.shared_memory_indexing import (
+import iree.turbine.kernel as tk
+import iree.turbine.kernel.lang as tkl
+import iree.turbine.kernel.wave as tkw
+from iree.turbine.kernel.wave.promotion import promote_placeholders
+from iree.turbine.kernel.wave.hoisting import hoist_allocs
+from iree.turbine.kernel.wave.expansion import expand_graph
+from iree.turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel._support.tracing import CapturedTrace
+from iree.turbine.kernel._support.indexing import IndexingContext
+from iree.turbine.kernel.ops.wave_ops import *
+from iree.turbine.kernel.wave.utils import run_test, print_trace
+from iree.turbine.kernel.wave.minimize_global_loads import minimize_global_loads
+from iree.turbine.kernel.wave.shared_memory_indexing import (
     apply_shared_memory_indexing_corrections,
 )
-from shark_turbine.kernel.wave.index_sequence_analysis import (
+from iree.turbine.kernel.wave.index_sequence_analysis import (
     partition_strided_operators,
 )
 

diff --git a/lit_tests/kernel/wave/minimize_global_loads.py b/lit_tests/kernel/wave/minimize_global_loads.py
@@ -3,21 +3,21 @@
 import logging
 from typing import Callable
 import unittest
-import shark_turbine.kernel as tk
-import shark_turbine.kernel.lang as tkl
-import shark_turbine.kernel.wave as tkw
-from shark_turbine.kernel.wave.promotion import promote_placeholders
-from shark_turbine.kernel.wave.hoisting import hoist_allocs
-from shark_turbine.kernel.wave.barriers import add_shared_memory_barriers
-from shark_turbine.kernel.wave.expansion import expand_graph
-from shark_turbine.kernel.lang.global_symbols import *
-from shark_turbine.kernel._support.tracing import CapturedTrace
-from shark_turbine.kernel._support.indexing import IndexingContext
-from shark_turbine.kernel.ops.wave_ops import *
-from shark_turbine.kernel.wave.utils import run_test, print_trace
-from shark_turbine.kernel.wave.minimize_global_loads import minimize_global_loads
-from shark_turbine.kernel.wave.visualization import visualize_graph
-from shark_turbine.kernel.wave.shared_memory_indexing import (
+import iree.turbine.kernel as tk
+import iree.turbine.kernel.lang as tkl
+import iree.turbine.kernel.wave as tkw
+from iree.turbine.kernel.wave.promotion import promote_placeholders
+from iree.turbine.kernel.wave.hoisting import hoist_allocs
+from iree.turbine.kernel.wave.barriers import add_shared_memory_barriers
+from iree.turbine.kernel.wave.expansion import expand_graph
+from iree.turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel._support.tracing import CapturedTrace
+from iree.turbine.kernel._support.indexing import IndexingContext
+from iree.turbine.kernel.ops.wave_ops import *
+from iree.turbine.kernel.wave.utils import run_test, print_trace
+from iree.turbine.kernel.wave.minimize_global_loads import minimize_global_loads
+from iree.turbine.kernel.wave.visualization import visualize_graph
+from iree.turbine.kernel.wave.shared_memory_indexing import (
     apply_shared_memory_indexing_corrections,
 )