y

leanprover · Feb 7, 2025 · 7b1244a · 7b1244a
1 parent d3828ca
commit 7b1244a
Show file tree

Hide file tree

Showing 50 changed files with 49 additions and 35 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -25,9 +25,10 @@ jobs:
       # Build (and test) Lean. Tests are all via #guard macros
       # now so you can't really build without testing.
       - uses: leanprover/lean-action@v1
+        with:
+          lake-package-directory: ./Lean
       - name: Run tests
         run: lake exe klr
-
       # Run pytest
       - uses: actions/setup-python@v5
         with:

diff --git a/.gitignore b/.gitignore
@@ -1,6 +1,6 @@
-.lake/**
+\.lake/
 __pycache__/
 *.so
 *.dSYM
-/.vscode/
-.wheel/
+.vscode/
+/.wheel/
diff --git a/Export.lean → Lean/Export.lean b/Export.lean → Lean/Export.lean
diff --git a/KLR.lean → Lean/KLR.lean b/KLR.lean → Lean/KLR.lean
diff --git a/KLR/Core.lean → Lean/KLR/Core.lean b/KLR/Core.lean → Lean/KLR/Core.lean
diff --git a/KLR/Core/Basic.lean → Lean/KLR/Core/Basic.lean b/KLR/Core/Basic.lean → Lean/KLR/Core/Basic.lean
diff --git a/KLR/Core/Encode.lean → Lean/KLR/Core/Encode.lean b/KLR/Core/Encode.lean → Lean/KLR/Core/Encode.lean
diff --git a/KLR/Core/Oper.lean → Lean/KLR/Core/Oper.lean b/KLR/Core/Oper.lean → Lean/KLR/Core/Oper.lean
diff --git a/KLR/Core/Pretty.lean → Lean/KLR/Core/Pretty.lean b/KLR/Core/Pretty.lean → Lean/KLR/Core/Pretty.lean
diff --git a/KLR/Python.lean → Lean/KLR/Python.lean b/KLR/Python.lean → Lean/KLR/Python.lean
diff --git a/KLR/Trace.lean → Lean/KLR/Trace.lean b/KLR/Trace.lean → Lean/KLR/Trace.lean
diff --git a/KLR/Trace/Basic.lean → Lean/KLR/Trace/Basic.lean b/KLR/Trace/Basic.lean → Lean/KLR/Trace/Basic.lean
diff --git a/KLR/Trace/Builtin.lean → Lean/KLR/Trace/Builtin.lean b/KLR/Trace/Builtin.lean → Lean/KLR/Trace/Builtin.lean
diff --git a/KLR/Trace/NKI.lean → Lean/KLR/Trace/NKI.lean b/KLR/Trace/NKI.lean → Lean/KLR/Trace/NKI.lean
diff --git a/KLR/Trace/Python.lean → Lean/KLR/Trace/Python.lean b/KLR/Trace/Python.lean → Lean/KLR/Trace/Python.lean
diff --git a/KLR/Trace/Tensor.lean → Lean/KLR/Trace/Tensor.lean b/KLR/Trace/Tensor.lean → Lean/KLR/Trace/Tensor.lean
diff --git a/KLR/Trace/Types.lean → Lean/KLR/Trace/Types.lean b/KLR/Trace/Types.lean → Lean/KLR/Trace/Types.lean
diff --git a/KLR/Util.lean → Lean/KLR/Util.lean b/KLR/Util.lean → Lean/KLR/Util.lean
diff --git a/Main.lean → Lean/Main.lean b/Main.lean → Lean/Main.lean
diff --git a/lake-manifest.json → Lean/lake-manifest.json b/lake-manifest.json → Lean/lake-manifest.json
diff --git a/lakefile.lean → Lean/lakefile.lean b/lakefile.lean → Lean/lakefile.lean
diff --git a/lean-toolchain → Lean/lean-toolchain b/lean-toolchain → Lean/lean-toolchain
diff --git a/Python/MANIFEST.in b/Python/MANIFEST.in
@@ -0,0 +1,2 @@
+# All files must be below the python project root
+include bin/klr
diff --git a/interop/README.md → Python/README.md b/interop/README.md → Python/README.md
@@ -1,3 +1,4 @@
+
 Python bindings for KLR
 
 # Usage
@@ -19,3 +20,11 @@ We do this to remove a dependency on the `neuronxcc` Python package that does no
    require it since it's large and unnecessary.
 
 2. Automatically sync `nki` examples with the github examples
+
+# Notes
+
+I keep getting __pycache__ directories created in the source tree, even with
+
+    export PYTHONDONTWRITEBYTECODE=1
+
+This is annoying.
diff --git a/interop/klr/__init__.py → Python/klr/__init__.py b/interop/klr/__init__.py → Python/klr/__init__.py
diff --git a/interop/klr/parser.py → Python/klr/parser.py b/interop/klr/parser.py → Python/klr/parser.py
diff --git a/interop/nki/__init__.py → Python/nki/__init__.py b/interop/nki/__init__.py → Python/nki/__init__.py
diff --git a/interop/nki/isa/__init__.py → Python/nki/isa/__init__.py b/interop/nki/isa/__init__.py → Python/nki/isa/__init__.py
diff --git a/interop/nki/language/__init__.py → Python/nki/language/__init__.py b/interop/nki/language/__init__.py → Python/nki/language/__init__.py
diff --git a/interop/noop.c → Python/noop.c b/interop/noop.c → Python/noop.c
@@ -1,3 +1,5 @@
+// A dummy file to trick Python's `build` into makeing an arch-specific wheel.
+
 #include <stdio.h>
 
 int main() {

diff --git a/interop/pyproject.toml → Python/pyproject.toml b/interop/pyproject.toml → Python/pyproject.toml
@@ -1,12 +1,18 @@
+# NB: All this setuptools stuff took about 1000 hours to figure out so never delete this.
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
 [project]
 name = "klr"
-version = "0.0.2"
-description = "Intermediate langauge for tensor compilers"
+version = "0.0.3"
 authors = [
   {name = "Paul Govereau", email = "[email protected]"},
   {name = "Sean McLaughlin", email = "[email protected]"},
 ]
+description = "Intermediate langauge for tensor compilers"
 readme = "README.md"
+requires-python = ">=3.8"
 license = { file = "LICENSE" }
 keywords = ["trainium", "tpu", "pallas", "triton", "gpu"]
 
@@ -21,31 +27,18 @@ dependencies = [
   "numpy",
 ]
 
+[project.urls]
+Repository = "https://github.com/leanprover/KLR"
+
 [tool.pytest.ini_options]
 pythonpath = "." # Needed for tests to pass
 testpaths = [
   "test",
 ]
 
-# NB: All this setuptools stuff took about 1000 hours to figure out so never delete this.
 [tool.setuptools]
-# This makes it ok for there to be both klr and wheelhouse (where we generates wheels on github actions) in the same directory
-#packages = ["klr"]
-# This tricks setuptools into knowing this is a binary wheel
-# The output says this is experimental and likely to change in future releases, but
-# we will keep our fingers crossed. Pretty sure we can make a setup.py file that does
-# this and it won't be that different.
-#ext-modules = [
-#  {name = "klr.noop", sources = ["noop.c"]}
-#]
+packages = ["bin", "klr", "nki", 'nki.isa', 'nki.language'] # Required to have 2 top level packages
 
 [[tool.setuptools.ext-modules]]
 name = "noop"
 sources = ["noop.c"]
-
-[tool.setuptools.package-data]
-# This part is essential to get the binary bin/klr included in the wheel
-klr = ["bin/klr"]
-
-[project.urls]
-Repository = "https://github.com/leanprover/KLR"
diff --git a/interop/requirements.txt → Python/requirements.txt b/interop/requirements.txt → Python/requirements.txt
diff --git a/interop/test/examples/__init__.py → Python/tests/examples/__init__.py b/interop/test/examples/__init__.py → Python/tests/examples/__init__.py
diff --git a/interop/test/examples/average_pool.py → Python/tests/examples/average_pool.py b/interop/test/examples/average_pool.py → Python/tests/examples/average_pool.py
diff --git a/interop/test/examples/fused_mamba.py → Python/tests/examples/fused_mamba.py b/interop/test/examples/fused_mamba.py → Python/tests/examples/fused_mamba.py
diff --git a/interop/test/examples/getting_started.py → Python/tests/examples/getting_started.py b/interop/test/examples/getting_started.py → Python/tests/examples/getting_started.py
diff --git a/interop/test/examples/index.py → Python/tests/examples/index.py b/interop/test/examples/index.py → Python/tests/examples/index.py
diff --git a/interop/test/examples/layernorm.py → Python/tests/examples/layernorm.py b/interop/test/examples/layernorm.py → Python/tests/examples/layernorm.py
diff --git a/interop/test/examples/layout.py → Python/tests/examples/layout.py b/interop/test/examples/layout.py → Python/tests/examples/layout.py
diff --git a/interop/test/examples/matmul.py → Python/tests/examples/matmul.py b/interop/test/examples/matmul.py → Python/tests/examples/matmul.py
diff --git a/interop/test/examples/mm.py → Python/tests/examples/mm.py b/interop/test/examples/mm.py → Python/tests/examples/mm.py
diff --git a/interop/test/examples/prof.py → Python/tests/examples/prof.py b/interop/test/examples/prof.py → Python/tests/examples/prof.py
diff --git a/interop/test/examples/rmsnorm.py → Python/tests/examples/rmsnorm.py b/interop/test/examples/rmsnorm.py → Python/tests/examples/rmsnorm.py
diff --git a/interop/test/examples/sd_attention.py → Python/tests/examples/sd_attention.py b/interop/test/examples/sd_attention.py → Python/tests/examples/sd_attention.py
diff --git a/interop/test/examples/tensor_addition.py → Python/tests/examples/tensor_addition.py b/interop/test/examples/tensor_addition.py → Python/tests/examples/tensor_addition.py
diff --git a/interop/test/examples/transpose2d.py → Python/tests/examples/transpose2d.py b/interop/test/examples/transpose2d.py → Python/tests/examples/transpose2d.py
diff --git a/interop/test/test_basic.py → Python/tests/test_basic.py b/interop/test/test_basic.py → Python/tests/test_basic.py
diff --git a/interop/test/test_examples.py → Python/tests/test_examples.py b/interop/test/test_examples.py → Python/tests/test_examples.py
diff --git a/bin/make-wheel b/bin/make-wheel
@@ -3,28 +3,36 @@ set -e -u -o pipefail
 trap "kill 0" SIGINT SIGTERM
 
 ROOT=$(dirname $(dirname $(readlink -f $0)))
-WHEEL_DIR=$ROOT/.wheel/klr
+WHEEL_DIR=$ROOT/.wheel/
+LEAN_DIR=$ROOT/Lean
+PYTHON_DIR=$ROOT/Python
 rm -rf $WHEEL_DIR
-mkdir -p $WHEEL_DIR/klr/bin
+mkdir -p $WHEEL_DIR/{klr,bin}
+
+# Build Lean code
+pushd .
+cd $LEAN_DIR
 lake build
-cp $ROOT/.lake/build/bin/klr $WHEEL_DIR/klr/bin/klr
-cp -R $ROOT/interop/nki $WHEEL_DIR/klr/nki
-cp -R $ROOT/interop/klr $WHEEL_DIR/klr/klr
+popd
+
+# Get klr binary
+cp $LEAN_DIR/.lake/build/bin/klr $WHEEL_DIR/bin/klr
+
+# Get Python packages
+cp -R $PYTHON_DIR/* $WHEEL_DIR
 cp $ROOT/LICENSE $WHEEL_DIR
-cp $ROOT/interop/README.md $WHEEL_DIR
-cp $ROOT/interop/noop.c $WHEEL_DIR
-cp $ROOT/interop/pyproject.toml $WHEEL_DIR
-cp $ROOT/interop/MANIFEST.in $WHEEL_DIR
 
 cd $WHEEL_DIR
 
 if [ -z "${GITHUB_RUN_ID+x}" ]; then
   echo "Building wheel locally"
-  python -m build -w
+  python -m build
 else
   # https://github.com/pypa/cibuildwheel
+  # NB: I could never get this to run on my mac locally. I got it to
+  # work on linux with no trouble.
   echo "Building wheel via GitHub Actions"
-  python -m cibuildwheel --output-dir $ROOT/wheelhouse
+  python -m cibuildwheel --output-dir ./wheelhouse
 fi
 
 # To upload the wheel, with the proper token in ~/.pypirc, run

diff --git a/interop/MANIFEST.in b/interop/MANIFEST.in
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		# All files must be below the python project root
		include bin/klr