CI/repo changes for shark-turbine to iree-turbine rename (#260)

This commit makes the necessary changes to adapt to the new iree-turbine namespace. Related PR: iree-org/iree-turbine#197 Signed-off-by: saienduri <[email protected]>
nod-ai · Oct 7, 2024 · 9ca19a5 · 9ca19a5
1 parent daa3f19
commit 9ca19a5
Show file tree

Hide file tree

Showing 33 changed files with 39 additions and 39 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -54,7 +54,7 @@ jobs:
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=shark-turbine"
+            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
           pip install --no-compile -r requirements.txt -e sharktank/
 
       - name: Run sharktank tests

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
@@ -53,7 +53,7 @@ jobs:
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=shark-turbine"
+            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
           pip install --no-compile -r requirements.txt -e sharktank/ shortfin/
 
           # Try with the latest nightly releases, not what iree-turbine pins.
@@ -85,7 +85,7 @@ jobs:
           python -m pip install --no-compile --upgrade pip
           pip install --no-compile -r pytorch-rocm-requirements.txt
           pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=shark-turbine"
+            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
           pip install --no-compile -r requirements.txt -e sharktank/ shortfin/
 
       - name: Run punet tests

diff --git a/README.md b/README.md
@@ -45,7 +45,7 @@ pip install -r pytorch-rocm-requirements.txt
 ```
 # Clone and install editable iree-turbine dep in deps/
 pip install -f https://iree.dev/pip-release-links.html --src deps \
-  -e "git+https://github.com/iree-org/iree-turbine.git#egg=shark-turbine"
+  -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
 
 # Install editable local projects.
 pip install -r requirements.txt -e sharktank/ shortfin/

diff --git a/docs/model_cookbook.md b/docs/model_cookbook.md
@@ -176,7 +176,7 @@ source .venv/bin/activate
 # Install requirements.
 pip install -r pytorch-cpu-requirements.txt
 pip install -f https://iree.dev/pip-release-links.html --src deps \
-  -e "git+https://github.com/iree-org/iree-turbine.git#egg=shark-turbine"
+  -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
 
 # Install local projects.
 pip install -r requirements.txt -e sharktank/ shortfin/

diff --git a/docs/quantization.md b/docs/quantization.md
@@ -277,9 +277,9 @@ is everything). We're just starting to exploit some of this as the PyTorch
 level. Some examples:
 
 * Something as simple as a humble runtime
-[tensor trace/print](https://github.com/iree-org/iree-turbine/blob/main/shark_turbine/ops/iree.py#L52)
-* [Simple linalg based template expansion](https://github.com/iree-org/iree-turbine/blob/main/shark_turbine/ops/_jinja_test_ops.py#L28)
-  (see backing example [jinja template](https://github.com/iree-org/iree-turbine/blob/main/shark_turbine/ops/templates/test_add_jinja.mlir)).
+[tensor trace/print](https://github.com/iree-org/iree-turbine/blob/main/iree.turbine/ops/iree.py#L52)
+* [Simple linalg based template expansion](https://github.com/iree-org/iree-turbine/blob/main/iree.turbine/ops/_jinja_test_ops.py#L28)
+  (see backing example [jinja template](https://github.com/iree-org/iree-turbine/blob/main/iree.turbine/ops/templates/test_add_jinja.mlir)).
 * Optimal linalg-based [8-bit block scaled mmt for weight compression](https://github.com/nod-ai/sharktank/blob/main/sharktank/sharktank/kernels/mmt_block_scaled_q8.py)
   (see backing [jinja template](https://github.com/nod-ai/sharktank/blob/main/sharktank/sharktank/kernels/templates/mmt_block_scaled_q8_3d.mlir)).
 * DSL based [like this fused attention kernel](https://github.com/iree-org/iree-turbine/blob/main/tests/kernel/fused_attention_test.py#L20)

diff --git a/sharktank/setup.py b/sharktank/setup.py
@@ -95,7 +95,7 @@ def initialize_options(self):
         "sharktank": ["py.typed", "kernels/templates/*.mlir"],
     },
     install_requires=[
-        "shark-turbine",
+        "iree-turbine",
     ],
     extras_require={
         "testing": [

diff --git a/sharktank/sharktank/examples/export_paged_llm_v1.py b/sharktank/sharktank/examples/export_paged_llm_v1.py
@@ -9,7 +9,7 @@
 import json
 import torch
 
-from shark_turbine.aot import *
+from iree.turbine.aot import *
 
 from sharktank.layers import *
 from sharktank.types import *

diff --git a/sharktank/sharktank/examples/sharding/export_ffn_net.py b/sharktank/sharktank/examples/sharding/export_ffn_net.py
@@ -89,7 +89,7 @@ def main(raw_args=None):
     ds = Dataset.load(args.output_irpa_file)
 
     mdl = ShardedFFN(ds.root_theta)
-    from shark_turbine import aot
+    from iree.turbine import aot
 
     example_arg = torch.empty(bs, sl, primary_dim, dtype=torch.float16)
     ep = torch.export.export(mdl, (example_arg,))

diff --git a/sharktank/sharktank/examples/sharding/export_gemm.py b/sharktank/sharktank/examples/sharding/export_gemm.py
@@ -4,7 +4,7 @@
 import torch
 from torch import Tensor
 from sharktank import ops
-from shark_turbine import aot
+from iree.turbine import aot
 
 
 def export_gemm(

diff --git a/sharktank/sharktank/export_layer/export_moe.py b/sharktank/sharktank/export_layer/export_moe.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import torch
-from shark_turbine.aot import *
+from iree.turbine.aot import *
 from sharktank.models.llama.testing import make_moe_block_theta, make_rand_torch
 from sharktank.layers.mixture_of_experts_block import PreGatherMoeBlock
 from ..utils import cli

diff --git a/sharktank/sharktank/export_layer/export_paged_attention.py b/sharktank/sharktank/export_layer/export_paged_attention.py
@@ -13,7 +13,7 @@
 
 import torch.nn.functional as F
 
-from shark_turbine.aot import *
+from iree.turbine.aot import *
 
 from sharktank.layers import *
 from sharktank.types import *

diff --git a/sharktank/sharktank/kernels/base.py b/sharktank/sharktank/kernels/base.py
@@ -12,7 +12,7 @@
 
 from jinja2 import Environment, PackageLoader, select_autoescape
 
-from shark_turbine.support.ir_imports import (
+from iree.turbine.support.ir_imports import (
     FlatSymbolRefAttr,
     FunctionType,
     IrType,
@@ -24,15 +24,15 @@
     Value,
 )
 
-from shark_turbine.runtime.op_reg import (
+from iree.turbine.runtime.op_reg import (
     def_library,
     CustomOp,
     KernelBuilder,
     KernelSelection,
     TensorArg,
 )
 
-from shark_turbine.transforms.merger import Merger
+from iree.turbine.transforms.merger import Merger
 
 from ..utils.logging import get_logger
 

diff --git a/sharktank/sharktank/models/punet/tools/run_punet.py b/sharktank/sharktank/models/punet/tools/run_punet.py
@@ -9,7 +9,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 
 from ..model import Unet2DConditionModel, ClassifierFreeGuidanceUnetModel
 from ....utils.patching import SaveModuleResultTensorsPatch

diff --git a/sharktank/sharktank/ops/default_impls.py b/sharktank/sharktank/ops/default_impls.py
@@ -24,7 +24,7 @@
 from ..types.tensors import unbox_tensor, AnyTensor
 from ._registry import AllOfType, AllOfExprs, AllOfExprsVariadic, IsOfType
 from .signatures import *
-import shark_turbine.ops.iree
+import iree.turbine.ops.iree
 
 
 @cat.override(AllOfType(Tensor, PrimitiveTensor))
@@ -393,7 +393,7 @@ def to_default(tensor: Tensor, *args, **kwargs):
 
 @transfer_to_logical_device.override(Tensor)
 def transfer_to_logical_device_default(tensor: Tensor, ordinal: int):
-    return shark_turbine.ops.iree.transfer_to_logical_device(
+    return iree.turbine.ops.iree.transfer_to_logical_device(
         f"{ordinal}", unbox_tensor(tensor)
     )
 

diff --git a/sharktank/sharktank/types/gguf_interop/base.py b/sharktank/sharktank/types/gguf_interop/base.py
@@ -13,7 +13,7 @@
 
 from gguf import GGUFReader, GGUFValueType
 
-from shark_turbine.aot import (
+from iree.turbine.aot import (
     ExternalTensorTrait,
 )
 

diff --git a/sharktank/sharktank/types/tensors.py b/sharktank/sharktank/types/tensors.py
@@ -27,7 +27,7 @@
 from torch import Tensor
 from torch.utils._pytree import register_pytree_node, SequenceKey
 from ..utils.math import ceildiv
-from shark_turbine.aot import (
+from iree.turbine.aot import (
     ExternalTensorTrait,
 )
 from ..utils import tree as tree_utils

diff --git a/sharktank/sharktank/types/theta.py b/sharktank/sharktank/types/theta.py
@@ -15,7 +15,7 @@
 import torch
 import torch.nn.functional as F
 
-from shark_turbine.aot import (
+from iree.turbine.aot import (
     ExternalTensorTrait,
     ParameterArchive,
     ParameterArchiveEntry,

diff --git a/sharktank/sharktank/utils/io.py b/sharktank/sharktank/utils/io.py
@@ -6,7 +6,7 @@
 
 from pathlib import Path
 
-from shark_turbine.aot import (
+from iree.turbine.aot import (
     ParameterArchiveBuilder,
 )
 

diff --git a/sharktank/sharktank/utils/logging.py b/sharktank/sharktank/utils/logging.py
@@ -6,7 +6,7 @@
 
 import logging
 
-from shark_turbine.support.logging import get_logger
+from iree.turbine.support.logging import get_logger
 
 
 transform_logger: logging.Logger = get_logger("sharktank.transforms")
diff --git a/sharktank/tests/kernels/batch_matmul_transpose_b_test.py b/sharktank/tests/kernels/batch_matmul_transpose_b_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 
 

diff --git a/sharktank/tests/kernels/conv_2d_nchw_fchw_test.py b/sharktank/tests/kernels/conv_2d_nchw_fchw_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 from sharktank.ops.qconv_impls import _pad_last_2d
 

diff --git a/sharktank/tests/kernels/einsum_q4_test.py b/sharktank/tests/kernels/einsum_q4_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 from sharktank.types import layout_utils
 

diff --git a/sharktank/tests/kernels/mmt_block_scaled_offset_q4_test.py b/sharktank/tests/kernels/mmt_block_scaled_offset_q4_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 from sharktank.types import layout_utils
 

diff --git a/sharktank/tests/kernels/mmt_block_scaled_q8_test.py b/sharktank/tests/kernels/mmt_block_scaled_q8_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 
 

diff --git a/sharktank/tests/kernels/mmt_super_block_scaled_offset_q4_test.py b/sharktank/tests/kernels/mmt_super_block_scaled_offset_q4_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 from sharktank.types import layout_utils
 

diff --git a/sharktank/tests/kernels/mmtfp_test.py b/sharktank/tests/kernels/mmtfp_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 
 

diff --git a/sharktank/tests/kernels/pooling_nchw_sum_test.py b/sharktank/tests/kernels/pooling_nchw_sum_test.py
@@ -13,7 +13,7 @@
 
 import torch
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank import kernels
 from sharktank.ops.qconv_impls import _pad_last_2d
 

diff --git a/sharktank/tests/layers/sharded_conv2d_with_iree_test.py b/sharktank/tests/layers/sharded_conv2d_with_iree_test.py
@@ -9,7 +9,7 @@
 from pathlib import Path
 import tempfile
 import torch
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank.models.punet.layers import Conv2DLayer
 from sharktank import ops
 from sharktank.types import (

diff --git a/sharktank/tests/models/llama/moe_block_test.py b/sharktank/tests/models/llama/moe_block_test.py
@@ -8,7 +8,7 @@
 from typing import List
 
 import torch
-from shark_turbine.aot import *
+from iree.turbine.aot import *
 from sharktank.models.llama.testing import make_moe_block_theta, make_rand_torch
 from sharktank.layers.mixture_of_experts_block import PreGatherMoeBlock
 from sharktank import ops

diff --git a/sharktank/tests/models/llama/sharded_llama_test.py b/sharktank/tests/models/llama/sharded_llama_test.py
@@ -24,7 +24,7 @@
 import tempfile
 import torch
 from copy import deepcopy
-from shark_turbine.aot import FxProgramsBuilder, export
+from iree.turbine.aot import FxProgramsBuilder, export
 import iree.runtime
 from pathlib import Path
 

diff --git a/sharktank/tests/models/punet/sharded_resnet_block_with_iree_test.py b/sharktank/tests/models/punet/sharded_resnet_block_with_iree_test.py
@@ -10,7 +10,7 @@
 import torch
 
 
-from shark_turbine import aot
+from iree.turbine import aot
 from sharktank.models.punet.testing import make_resnet_block_2d_theta
 from sharktank.models.punet.layers import ResnetBlock2D
 from sharktank.models.punet.sharding import ResnetBlock2DSplitOutputChannelsSharding

diff --git a/sharktank/tests/types/dataset_test.py b/sharktank/tests/types/dataset_test.py
@@ -11,7 +11,7 @@
 
 import torch
 
-from shark_turbine.aot import ExternalTensorTrait
+from iree.turbine.aot import ExternalTensorTrait
 from sharktank.types import *
 
 

diff --git a/turbine-requirements.txt b/turbine-requirements.txt
@@ -1 +1 @@
--e "git+https://github.com/iree-org/iree-turbine.git#egg=shark-turbine"
+-e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"