z

leanprover · Feb 8, 2025 · 6eff035 · 6eff035
1 parent 620775a
commit 6eff035
Show file tree

Hide file tree

Showing 49 changed files with 13 additions and 20 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -15,7 +15,7 @@ jobs:
         # macos-13 is x86
         # macos-latest is arm64
         # ubuntu-latest is x64
-        os: [macos-latest]
+        os: [ubuntu-latest, macos-latest]
 
     runs-on: ${{ matrix.os }}
 
@@ -25,23 +25,20 @@ 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
-        working-directory: ./Lean
         run: lake exe klr
       # Run pytest
       - uses: actions/setup-python@v5
         with:
           python-version: '3.x'
           cache: 'pip'
       - name: Install dependencies
-        working-directory: ./Python
+        working-directory: ./interop
         run: |
           python -m pip install --upgrade pip
           pip install -r requirements.txt
       - name: Run pytest
-        working-directory: ./Python
+        working-directory: ./interop
         run: |
           pytest
       - name: Make a wheel

diff --git a/Lean/Export.lean → Export.lean b/Lean/Export.lean → Export.lean
diff --git a/Lean/KLR.lean → KLR.lean b/Lean/KLR.lean → KLR.lean
diff --git a/Lean/KLR/Core.lean → KLR/Core.lean b/Lean/KLR/Core.lean → KLR/Core.lean
diff --git a/Lean/KLR/Core/Basic.lean → KLR/Core/Basic.lean b/Lean/KLR/Core/Basic.lean → KLR/Core/Basic.lean
diff --git a/Lean/KLR/Core/Encode.lean → KLR/Core/Encode.lean b/Lean/KLR/Core/Encode.lean → KLR/Core/Encode.lean
diff --git a/Lean/KLR/Core/Oper.lean → KLR/Core/Oper.lean b/Lean/KLR/Core/Oper.lean → KLR/Core/Oper.lean
diff --git a/Lean/KLR/Core/Pretty.lean → KLR/Core/Pretty.lean b/Lean/KLR/Core/Pretty.lean → KLR/Core/Pretty.lean
diff --git a/Lean/KLR/Python.lean → KLR/Python.lean b/Lean/KLR/Python.lean → KLR/Python.lean
diff --git a/Lean/KLR/Trace.lean → KLR/Trace.lean b/Lean/KLR/Trace.lean → KLR/Trace.lean
diff --git a/Lean/KLR/Trace/Basic.lean → KLR/Trace/Basic.lean b/Lean/KLR/Trace/Basic.lean → KLR/Trace/Basic.lean
diff --git a/Lean/KLR/Trace/Builtin.lean → KLR/Trace/Builtin.lean b/Lean/KLR/Trace/Builtin.lean → KLR/Trace/Builtin.lean
diff --git a/Lean/KLR/Trace/NKI.lean → KLR/Trace/NKI.lean b/Lean/KLR/Trace/NKI.lean → KLR/Trace/NKI.lean
diff --git a/Lean/KLR/Trace/Python.lean → KLR/Trace/Python.lean b/Lean/KLR/Trace/Python.lean → KLR/Trace/Python.lean
diff --git a/Lean/KLR/Trace/Tensor.lean → KLR/Trace/Tensor.lean b/Lean/KLR/Trace/Tensor.lean → KLR/Trace/Tensor.lean
diff --git a/Lean/KLR/Trace/Types.lean → KLR/Trace/Types.lean b/Lean/KLR/Trace/Types.lean → KLR/Trace/Types.lean
diff --git a/Lean/KLR/Util.lean → KLR/Util.lean b/Lean/KLR/Util.lean → KLR/Util.lean
diff --git a/Lean/Main.lean → Main.lean b/Lean/Main.lean → Main.lean
diff --git a/bin/klr b/bin/klr
@@ -2,4 +2,4 @@
 
 ROOT=$(dirname $(dirname $(readlink -f $0)))
 
-$ROOT/Lean/.lake/build/bin/klr $@
+$ROOT/.lake/build/bin/klr $@
diff --git a/bin/make-wheel b/bin/make-wheel
@@ -4,19 +4,13 @@ trap "kill 0" SIGINT SIGTERM
 
 ROOT=$(dirname $(dirname $(readlink -f $0)))
 WHEEL_DIR=$ROOT/.wheel/
-LEAN_DIR=$ROOT/Lean
-PYTHON_DIR=$ROOT/Python
+PYTHON_DIR=$ROOT/interop
 rm -rf $WHEEL_DIR
 mkdir -p $WHEEL_DIR/{klr,bin}
-
-# Build Lean code
-pushd .
-cd $LEAN_DIR
 lake build
-popd
 
 # Get klr binary
-cp $LEAN_DIR/.lake/build/bin/klr $WHEEL_DIR/bin/klr
+cp .lake/build/bin/klr $WHEEL_DIR/bin/klr
 
 # Get Python packages
 cp -R $PYTHON_DIR/* $WHEEL_DIR

diff --git a/Python/MANIFEST.in → interop/MANIFEST.in b/Python/MANIFEST.in → interop/MANIFEST.in
diff --git a/Python/README.md → interop/README.md b/Python/README.md → interop/README.md
diff --git a/Python/klr/__init__.py → interop/klr/__init__.py b/Python/klr/__init__.py → interop/klr/__init__.py
diff --git a/Python/klr/parser.py → interop/klr/parser.py b/Python/klr/parser.py → interop/klr/parser.py
@@ -20,7 +20,7 @@ def up(f, n):
   for _ in range(n):
     d = os.path.dirname(d)
   return d
-  
+
 
 def run_klr(infile, outfile):
   # For development, pick up the klr binary from the project dir
@@ -29,8 +29,7 @@ def run_klr(infile, outfile):
   if not os.path.isfile(bin):
     # For regular pip users, pick up the klr from the wheel. While the type of `bin` here is
     # PosixPath rather than string, they both work as the first argument to subprocess.run
-    #bin = files('klr').joinpath('bin/klr')
-   raise Exception("no file: " + bin)
+    bin = files('klr').joinpath('bin/klr')
   subprocess.run([bin, 'parse-json', infile.name], stdout=outfile, check=True)
 
 

diff --git a/Python/nki/__init__.py → interop/nki/__init__.py b/Python/nki/__init__.py → interop/nki/__init__.py
diff --git a/Python/nki/isa/__init__.py → interop/nki/isa/__init__.py b/Python/nki/isa/__init__.py → interop/nki/isa/__init__.py
diff --git a/Python/nki/language/__init__.py → interop/nki/language/__init__.py b/Python/nki/language/__init__.py → interop/nki/language/__init__.py
diff --git a/Python/noop.c → interop/noop.c b/Python/noop.c → interop/noop.c
diff --git a/Python/pyproject.toml → interop/pyproject.toml b/Python/pyproject.toml → interop/pyproject.toml
@@ -12,7 +12,6 @@ authors = [
 ]
 description = "Intermediate langauge for tensor compilers"
 readme = "README.md"
-requires-python = ">=3.8"
 license = { file = "LICENSE" }
 keywords = ["trainium", "tpu", "pallas", "triton", "gpu"]
 
@@ -37,8 +36,12 @@ testpaths = [
 ]
 
 [tool.setuptools]
-packages = ["bin", "klr", "nki", 'nki.isa', 'nki.language'] # Required to have 2 top level packages
+packages = ["bin", "klr", "nki", 'nki.isa', 'nki.language'] # This field is tequired to have more than 1 top level packages (nki/klr). It's not beautiful but it works.
 
+# Build a fake C file to trick the `build` command to make an arch-specific wheel
+# We need this because we hide our `klr` Lean binary in the wheel but Python tools
+# don't know about it. I couldn't find another way to stop Python from generating
+# a generic `any` wheel.
 [[tool.setuptools.ext-modules]]
 name = "noop"
 sources = ["noop.c"]
diff --git a/Python/requirements.txt → interop/requirements.txt b/Python/requirements.txt → interop/requirements.txt
diff --git a/Python/tests/examples/__init__.py → interop/test/examples/__init__.py b/Python/tests/examples/__init__.py → interop/test/examples/__init__.py
diff --git a/Python/tests/examples/average_pool.py → interop/test/examples/average_pool.py b/Python/tests/examples/average_pool.py → interop/test/examples/average_pool.py
diff --git a/Python/tests/examples/fused_mamba.py → interop/test/examples/fused_mamba.py b/Python/tests/examples/fused_mamba.py → interop/test/examples/fused_mamba.py
diff --git a/Python/tests/examples/getting_started.py → interop/test/examples/getting_started.py b/Python/tests/examples/getting_started.py → interop/test/examples/getting_started.py
diff --git a/Python/tests/examples/index.py → interop/test/examples/index.py b/Python/tests/examples/index.py → interop/test/examples/index.py
diff --git a/Python/tests/examples/layernorm.py → interop/test/examples/layernorm.py b/Python/tests/examples/layernorm.py → interop/test/examples/layernorm.py
diff --git a/Python/tests/examples/layout.py → interop/test/examples/layout.py b/Python/tests/examples/layout.py → interop/test/examples/layout.py
diff --git a/Python/tests/examples/matmul.py → interop/test/examples/matmul.py b/Python/tests/examples/matmul.py → interop/test/examples/matmul.py
diff --git a/Python/tests/examples/mm.py → interop/test/examples/mm.py b/Python/tests/examples/mm.py → interop/test/examples/mm.py
diff --git a/Python/tests/examples/prof.py → interop/test/examples/prof.py b/Python/tests/examples/prof.py → interop/test/examples/prof.py
diff --git a/Python/tests/examples/rmsnorm.py → interop/test/examples/rmsnorm.py b/Python/tests/examples/rmsnorm.py → interop/test/examples/rmsnorm.py
diff --git a/Python/tests/examples/sd_attention.py → interop/test/examples/sd_attention.py b/Python/tests/examples/sd_attention.py → interop/test/examples/sd_attention.py
diff --git a/Python/tests/examples/tensor_addition.py → interop/test/examples/tensor_addition.py b/Python/tests/examples/tensor_addition.py → interop/test/examples/tensor_addition.py
diff --git a/Python/tests/examples/transpose2d.py → interop/test/examples/transpose2d.py b/Python/tests/examples/transpose2d.py → interop/test/examples/transpose2d.py
diff --git a/Python/tests/test_basic.py → interop/test/test_basic.py b/Python/tests/test_basic.py → interop/test/test_basic.py
diff --git a/Python/tests/test_examples.py → interop/test/test_examples.py b/Python/tests/test_examples.py → interop/test/test_examples.py
diff --git a/Lean/lake-manifest.json → lake-manifest.json b/Lean/lake-manifest.json → lake-manifest.json
diff --git a/Lean/lakefile.lean → lakefile.lean b/Lean/lakefile.lean → lakefile.lean
diff --git a/Lean/lean-toolchain → lean-toolchain b/Lean/lean-toolchain → lean-toolchain
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2,4 +2,4 @@

		ROOT=$(dirname $(dirname $(readlink -f $0)))

		$ROOT/Lean/.lake/build/bin/klr $@
		$ROOT/.lake/build/bin/klr $@