Properly layer the project directory structure.

.github/workflows/lint.yml

@@ -0,0 +1,38 @@
+name: Lint
+
+on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  black:
+    strategy:
+      matrix:
+        version: [3.11]
+        os: [ubuntu-latest]
+    runs-on: ${{matrix.os}}
+    steps:
+      - name: Checking out repository
+        uses: actions/checkout@8f4b7f84864484a7bf31766abe9204da3cbe65b3 # v3.5.0
+      - name: Setting up python
+        uses: actions/setup-python@d27e3f3d7c64b4bbf8e4abfb9b63b83e846e0435 # v4.5.0
+      - name: Fetching Base Branch
+        # We have to explicitly fetch the base branch as well
+        run: git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
+      - name: Install black
+        run: |
+          python3 -m pip install black==23.3
+      - name: Check if modified files are formatted
+        run: |
+          # The filter lowercase `d` means to exclude deleted files.
+          git diff "${GITHUB_BASE_REF?}" --name-only --diff-filter=d \
+            -- '*.py' \
+            | xargs --no-run-if-empty black --check --diff --verbose
+      - name: Instructions for fixing the above linting errors
+        if: failure()
+        run: |
+          printf "You can fix formatting by running 'black' on the modified python files:\n"
+          printf "  git diff ${GITHUB_BASE_REF?} --name-only -- '*.py' ':!third_party' | xargs black\n"

.github/workflows/test.yml

@@ -1,4 +1,4 @@
-name: Test
+name: Turbine Core Test
 
 on:
   workflow_dispatch:
@@ -39,40 +39,11 @@ jobs:
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --index-url https://download.pytorch.org/whl/cpu \
-            -r pytorch-cpu-requirements.txt \
-            -r torchvision-requirements.txt          
-          pip install --upgrade -r requirements.txt
-          pip install -e  .[testing]
+            -r core/pytorch-cpu-requirements.txt \
+            -r core/torchvision-requirements.txt          
+          pip install --upgrade -r core/requirements.txt
+          pip install -e  core[testing]
 
       - name: Run tests
         run: |
-          pytest -n 4 tests/
-
-  black:
-    strategy:
-      matrix:
-        version: [3.11]
-        os: [ubuntu-latest]
-    runs-on: ${{matrix.os}}
-    steps:
-      - name: Checking out repository
-        uses: actions/checkout@8f4b7f84864484a7bf31766abe9204da3cbe65b3 # v3.5.0
-      - name: Setting up python
-        uses: actions/setup-python@d27e3f3d7c64b4bbf8e4abfb9b63b83e846e0435 # v4.5.0
-      - name: Fetching Base Branch
-        # We have to explicitly fetch the base branch as well
-        run: git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
-      - name: Install black
-        run: |
-          python3 -m pip install black==23.3
-      - name: Check if modified files are formatted
-        run: |
-          # The filter lowercase `d` means to exclude deleted files.
-          git diff "${GITHUB_BASE_REF?}" --name-only --diff-filter=d \
-            -- '*.py' \
-            | xargs --no-run-if-empty black --check --diff --verbose
-      - name: Instructions for fixing the above linting errors
-        if: failure()
-        run: |
-          printf "You can fix formatting by running 'black' on the modified python files:\n"
-          printf "  git diff ${GITHUB_BASE_REF?} --name-only -- '*.py' ':!third_party' | xargs black\n"
+          pytest -n 4 core/tests/

.github/workflows/test_models.yml

@@ -39,20 +39,20 @@ jobs:
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --index-url https://download.pytorch.org/whl/cpu \
-            -r pytorch-cpu-requirements.txt \
-            -r torchvision-requirements.txt          
-          pip install --upgrade -r requirements.txt
-          pip install -e  .[testing]
-          pip install -e python/turbine_models
+            -r core/pytorch-cpu-requirements.txt \
+            -r core/torchvision-requirements.txt          
+          pip install --upgrade -r core/requirements.txt
+          pip install -e core[testing]
+          pip install -e models
 
       - name: Show current free memory
         run: |
           free -mh
 
       - name: Run stateless_llama tests
         run: |
-          pytest python/turbine_models/tests/stateless_llama_test.py
+          pytest models/turbine_models/tests/stateless_llama_test.py
 
       - name: Run sd tests
         run: |
-          pytest python/turbine_models/tests/sd_test.py
+          pytest models/turbine_models/tests/sd_test.py

README.md

@@ -39,15 +39,15 @@ please reach out to us on the `#turbine` channel of the
 
 ```
 pip install shark-turbine
-# Or editable: pip install -e .
+# Or editable: pip install -e core
 ```
 
 The above does install some unecessary cuda/cudnn packages for cpu use. To avoid this you
 can specify pytorch-cpu and install via:
 ```
 pip install --index-url https://download.pytorch.org/whl/cpu \
-    -r pytorch-cpu-requirements.txt \
-    -r torchvision-requirements.txt
+    -r core/pytorch-cpu-requirements.txt \
+    -r core torchvision-requirements.txt
 pip install shark-turbine
 ```
 
@@ -87,14 +87,14 @@ versions and has all flags. This can be installed prior to the package:
 Installing into a venv is highly recommended.
 
 ```
-pip install --upgrade -r requirements.txt
-pip install --upgrade -e .[torch-cpu-nightly,testing]
+pip install --upgrade -r core/requirements.txt
+pip install --upgrade -e core[torch-cpu-nightly,testing]
 ```
 
 Run tests:
 
 ```
-pytest
+pytest core/
 ```
 
 ### Using a development compiler

core/README.md

@@ -0,0 +1,10 @@
+# Turbine-Core Sub-Project
+
+This directory contains the core infrastructure for the project, sonsisting
+of the model export, runtime, and kernel development APIs. It is packaged
+and released as the 
+[`SHARK-Turbine` project on PyPi](https://pypi.org/project/shark-turbine/).
+
+It depends purely on PyTorch and the IREE compiler/runtime.
+
+See the repository-level README for further information.

examples/eager_mlp/mlp_eager_simple.py → core/examples/eager_mlp/mlp_eager_simple.py

@@ -13,7 +13,9 @@
 import torchvision.transforms as transforms
 import torchvision.datasets as datasets
 
-torch._dynamo.config.dynamic_shapes = False # TODO: https://github.com/nod-ai/SHARK-Turbine/issues/93
+torch._dynamo.config.dynamic_shapes = (
+    False  # TODO: https://github.com/nod-ai/SHARK-Turbine/issues/93
+)
 
 
 class MNISTDataLoader:
@@ -39,7 +41,10 @@ def get_train_loader(self):
 
     def get_test_loader(self):
         return DataLoader(
-            dataset=self.mnist_testset, batch_size=self.batch_size, shuffle=False, drop_last=True,
+            dataset=self.mnist_testset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            drop_last=True,
         )
 
 

mypy.ini → core/mypy.ini

@@ -1,4 +1,4 @@
 [mypy]
 
-mypy_path = $MYPY_CONFIG_FILE_DIR/python
+mypy_path = $MYPY_CONFIG_FILE_DIR
 packages = shark_turbine.aot,shark_turbine.dynamo,shark_turbine.support

setup.py → core/setup.py

@@ -12,12 +12,13 @@
 from setuptools import find_namespace_packages, setup
 
 THIS_DIR = os.path.realpath(os.path.dirname(__file__))
-VERSION_INFO_FILE = os.path.join(THIS_DIR, "version_info.json")
+REPO_DIR = os.path.dirname(THIS_DIR)
+VERSION_INFO_FILE = os.path.join(REPO_DIR, "version_info.json")
 
 
 with open(
     os.path.join(
-        THIS_DIR,
+        REPO_DIR,
         "README.md",
     ),
     "rt",
@@ -38,7 +39,6 @@ def load_version_info():
         "shark_turbine",
         "shark_turbine.*",
     ],
-    where="python",
 )
 
 print("Found packages:", packages)
@@ -90,9 +90,6 @@ def initialize_options(self):
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
     ],
-    package_dir={
-        "": "python",
-    },
     packages=packages,
     entry_points={
         "torch_dynamo_backends": [

python/shark_turbine/transforms/quantization/mm_group_quant.py → core/shark_turbine/transforms/quantization/mm_group_quant.py

@@ -63,9 +63,7 @@ def match(self, op: Operation):
             m=m,
             n=n,
             k=k,
-            element_type=self.builder.get_tensor_element_type(
-                op.operands[0].type
-            ),
+            element_type=self.builder.get_tensor_element_type(op.operands[0].type),
         )
 
 
@@ -134,9 +132,7 @@ def __init__(self, root_op: Operation, *, group_size: int = 128):
 
     def run(self):
         globals = self.globals
-        mms = match_children(
-            self.funcs, TransposedMMMatcher(globals, self.builder)
-        )
+        mms = match_children(self.funcs, TransposedMMMatcher(globals, self.builder))
 
         for mr in mms:
             if mr.k is None or mr.n is None:
@@ -165,12 +161,10 @@ def rewrite(self, mr: TransposedMMResult):
                 element_type=mr.element_type,
             )
 
-            inline_module = Operation.parse(
-                inline_module_asm, context=self.context
+            inline_module = Operation.parse(inline_module_asm, context=self.context)
+            actual_callee_name = self.merge_module(inline_module).translate_symbol(
+                "compute_mm_group_quant"
             )
-            actual_callee_name = self.merge_module(
-                inline_module
-            ).translate_symbol("compute_mm_group_quant")
             with InsertionPoint(mr.op), mr.op.location:
                 results = self.builder.call_native(
                     actual_callee_name, [mr.op.result.type], mr.op.operands[0]

python/turbine_models/README.md → models/README.md

@@ -16,11 +16,11 @@ YOUR_HF_TOKEN="insert token for headless"
 sudo apt install -y git
 git clone https://github.com/nod-ai/SHARK-Turbine.git
 cd SHARK-Turbine
-pip install -r requirements.txt
-pip install -r turbine-models-requirements.txt
+pip install -r core/requirements.txt
+pip install -r models/requirements.txt
 
 # do an editable install from the cloned SHARK-Turbine
-pip install --editable .
+pip install --editable core models
 
 # Log in with Hugging Face CLI if token setup is required
 if [[ $YOUR_HF_TOKEN == hf_* ]]; then

python/turbine_models/setup.py → models/setup.py

@@ -9,7 +9,7 @@
 
 
 TURBINE_MODELS_DIR = os.path.realpath(os.path.dirname(__file__))
-TURBINE_ROOT_DIR = Path(TURBINE_MODELS_DIR).parent.parent
+TURBINE_ROOT_DIR = Path(TURBINE_MODELS_DIR).parent
 print(TURBINE_ROOT_DIR)
 VERSION_INFO_FILE = os.path.join(TURBINE_ROOT_DIR, "version_info.json")
 
@@ -47,15 +47,11 @@ def load_version_info():
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
     ],
-    package_dir={
-        "": f"{str(TURBINE_ROOT_DIR)}/python",
-    },
     packages=find_namespace_packages(
         include=[
             "turbine_models",
             "turbine_models.*",
         ],
-        where=f"{str(TURBINE_ROOT_DIR)}/python",
     ),
     install_requires=[
         "Shark-Turbine",