Add flag to insert host codegen call

.github/workflows/ci-tk.yaml

@@ -0,0 +1,74 @@
+name: "TK CI"
+
+on:
+  pull_request:
+  push:
+    branches:
+    - main
+
+concurrency:
+  # A PR number if a pull request and otherwise the commit hash. This cancels
+  # queued and in-progress runs for the same PR (presubmit) or commit
+  # (postsubmit). The workflow name is prepended to avoid conflicts between
+  # different workflows.
+  group: ${{ github.workflow }}-${{ github.event.number || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    name: "Unit Tests and Type Checking"
+    strategy:
+      fail-fast: false
+      matrix:
+        version: [3.11]
+        os: [ubuntu-latest, nodai-amdgpu-mi300-x86-64]
+    runs-on: ${{matrix.os}}
+    env:
+      PIP_CACHE_DIR: "${{ github.workspace }}/.pip-cache"
+    steps:
+    - name: "Setting up Python"
+      id: setup_python
+      uses: actions/setup-python@v3
+      with:
+        python-version: ${{matrix.version}}
+
+    - name: "Checkout Code"
+      uses: actions/checkout@v3
+
+    - name: Cache Pip Packages
+      uses: actions/cache@v4
+      id: cache-pip
+      with:
+        path: ${{ env.PIP_CACHE_DIR }}
+        key: pip-${{ steps.setup_python.outputs.python-version }}-${{ hashFiles('*requirements.txt') }}
+
+    - name: Install pip deps
+      run: |
+        python -m pip install --no-compile --upgrade pip
+        # Note: We install in three steps in order to satisfy requirements
+        # from non default locations first. Installing the PyTorch CPU
+        # wheels saves multiple minutes and a lot of bandwidth on runner setup.
+        pip install --no-compile -r pytorch-cpu-requirements.txt
+        pip install --no-cache-dir -r iree-requirements-ci.txt
+        pip install -r requirements.txt -e .
+
+    - name: Run unit tests
+      if: ${{ !cancelled() }}
+      run: |
+        pytest -n 4 --capture=tee-sys -vv ./tests/kernel/wave/
+
+    - name: Run e2e tests on MI300
+      if: "contains(matrix.os, 'mi300') && !cancelled()"
+      run: |
+        export WAVE_RUN_E2E_TESTS=1
+        pytest -n 4 --capture=tee-sys -vv ./tests/kernel/wave/
+
+    - name: Run LIT tests
+      if: ${{ !cancelled() }}
+      run: |
+        lit lit_tests/ -v
+
+    - name: MyPy Type Checking
+      if: ${{ !cancelled() }}
+      run: |
+        mypy

.github/workflows/ci.yaml

@@ -18,9 +18,10 @@ jobs:
   test:
     name: "Unit Tests and Type Checking"
     strategy:
+      fail-fast: false
       matrix:
         version: [3.11]
-        os: [ubuntu-latest, nodai-amdgpu-mi300-x86-64]
+        os: [ubuntu-latest]
     runs-on: ${{matrix.os}}
     env:
       PIP_CACHE_DIR: "${{ github.workspace }}/.pip-cache"
@@ -54,14 +55,7 @@ jobs:
     - name: Run unit tests
       if: ${{ !cancelled() }}
       run: |
-        pytest -n 4 .
-
-    - name: Run e2e tests on MI300
-      if: "contains(matrix.os, 'mi300') && !cancelled()"
-      run: |
-        export WAVE_RUN_E2E_TESTS=1
-        export TEST_PARAMS_PATH=./tests/kernel/wave/test_param.json
-        pytest -n 4 ./tests/kernel/wave/
+        pytest -n 4 --capture=tee-sys -vv .
 
     - name: Run LIT tests
       if: ${{ !cancelled() }}

.github/workflows/perf.yaml

@@ -21,9 +21,10 @@ jobs:
   test:
     name: "Unit Tests and Type Checking"
     strategy:
+      fail-fast: false
       matrix:
         version: [3.11]
-        os: [ubuntu-latest, nodai-amdgpu-mi300-x86-64]
+        os: [nodai-amdgpu-mi300-x86-64]
     runs-on: ${{matrix.os}}
     env:
       PIP_CACHE_DIR: "${{ github.workspace }}/.pip-cache"
@@ -53,21 +54,10 @@ jobs:
         pip install --no-compile -r pytorch-cpu-requirements.txt
         pip install --no-cache-dir -r iree-requirements-ci.txt
         pip install -r requirements.txt -e .
-    - name: Run unit tests
-      if: ${{ !cancelled() }}
-      run: |
-        pytest -n 4 .
+
     - name: Run e2e tests on MI300
       if: "contains(matrix.os, 'mi300') && !cancelled()"
       run: |
         export WAVE_RUN_E2E_TESTS=1
         export TEST_PARAMS_PATH="tests/kernel/wave/test_param.json"
-        pytest -n 1 ./tests/kernel/wave/
-    - name: Run LIT tests
-      if: ${{ !cancelled() }}
-      run: |
-        lit lit_tests/ -v
-    - name: MyPy Type Checking
-      if: ${{ !cancelled() }}
-      run: |
-        mypy
+        pytest -n 1 --capture=tee-sys -vv ./tests/kernel/wave/

.github/workflows/test_build_release.yml

@@ -19,6 +19,7 @@ jobs:
   test:
     name: "Test Build Release Process"
     strategy:
+      fail-fast: false
       matrix:
         version: [3.11]
         os: [ubuntu-latest]

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

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.

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__":

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):

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):