From 620775a8ac76a34d04456061d74e9a7e69d50636 Mon Sep 17 00:00:00 2001 From: Sean McLaughlin Date: Fri, 7 Feb 2025 12:36:05 -0800 Subject: [PATCH] y --- .github/workflows/ci.yml | 16 ++++++---- .gitignore | 6 ++-- Export.lean => Lean/Export.lean | 0 KLR.lean => Lean/KLR.lean | 0 {KLR => Lean/KLR}/Core.lean | 0 {KLR => Lean/KLR}/Core/Basic.lean | 0 {KLR => Lean/KLR}/Core/Encode.lean | 0 {KLR => Lean/KLR}/Core/Oper.lean | 0 {KLR => Lean/KLR}/Core/Pretty.lean | 0 {KLR => Lean/KLR}/Python.lean | 0 {KLR => Lean/KLR}/Trace.lean | 0 {KLR => Lean/KLR}/Trace/Basic.lean | 0 {KLR => Lean/KLR}/Trace/Builtin.lean | 0 {KLR => Lean/KLR}/Trace/NKI.lean | 0 {KLR => Lean/KLR}/Trace/Python.lean | 0 {KLR => Lean/KLR}/Trace/Tensor.lean | 0 {KLR => Lean/KLR}/Trace/Types.lean | 0 {KLR => Lean/KLR}/Util.lean | 0 Main.lean => Lean/Main.lean | 0 lake-manifest.json => Lean/lake-manifest.json | 0 lakefile.lean => Lean/lakefile.lean | 0 lean-toolchain => Lean/lean-toolchain | 0 Python/MANIFEST.in | 2 ++ {interop => Python}/README.md | 9 ++++++ {interop => Python}/klr/__init__.py | 0 {interop => Python}/klr/parser.py | 12 +++++-- {interop => Python}/nki/__init__.py | 0 {interop => Python}/nki/isa/__init__.py | 0 {interop => Python}/nki/language/__init__.py | 0 {interop => Python}/noop.c | 2 ++ {interop => Python}/pyproject.toml | 31 +++++++------------ {interop => Python}/requirements.txt | 0 .../tests}/examples/__init__.py | 0 .../tests}/examples/average_pool.py | 0 .../tests}/examples/fused_mamba.py | 0 .../tests}/examples/getting_started.py | 0 .../test => Python/tests}/examples/index.py | 0 .../tests}/examples/layernorm.py | 0 .../test => Python/tests}/examples/layout.py | 0 .../test => Python/tests}/examples/matmul.py | 0 {interop/test => Python/tests}/examples/mm.py | 0 .../test => Python/tests}/examples/prof.py | 0 .../test => Python/tests}/examples/rmsnorm.py | 0 .../tests}/examples/sd_attention.py | 0 .../tests}/examples/tensor_addition.py | 0 .../tests}/examples/transpose2d.py | 0 {interop/test => Python/tests}/test_basic.py | 0 .../test => Python/tests}/test_examples.py | 0 bin/klr | 2 +- bin/make-wheel | 30 +++++++++++------- interop/MANIFEST.in | 1 - 51 files changed, 67 insertions(+), 44 deletions(-) rename Export.lean => Lean/Export.lean (100%) rename KLR.lean => Lean/KLR.lean (100%) rename {KLR => Lean/KLR}/Core.lean (100%) rename {KLR => Lean/KLR}/Core/Basic.lean (100%) rename {KLR => Lean/KLR}/Core/Encode.lean (100%) rename {KLR => Lean/KLR}/Core/Oper.lean (100%) rename {KLR => Lean/KLR}/Core/Pretty.lean (100%) rename {KLR => Lean/KLR}/Python.lean (100%) rename {KLR => Lean/KLR}/Trace.lean (100%) rename {KLR => Lean/KLR}/Trace/Basic.lean (100%) rename {KLR => Lean/KLR}/Trace/Builtin.lean (100%) rename {KLR => Lean/KLR}/Trace/NKI.lean (100%) rename {KLR => Lean/KLR}/Trace/Python.lean (100%) rename {KLR => Lean/KLR}/Trace/Tensor.lean (100%) rename {KLR => Lean/KLR}/Trace/Types.lean (100%) rename {KLR => Lean/KLR}/Util.lean (100%) rename Main.lean => Lean/Main.lean (100%) rename lake-manifest.json => Lean/lake-manifest.json (100%) rename lakefile.lean => Lean/lakefile.lean (100%) rename lean-toolchain => Lean/lean-toolchain (100%) create mode 100644 Python/MANIFEST.in rename {interop => Python}/README.md (85%) rename {interop => Python}/klr/__init__.py (100%) rename {interop => Python}/klr/parser.py (96%) rename {interop => Python}/nki/__init__.py (100%) rename {interop => Python}/nki/isa/__init__.py (100%) rename {interop => Python}/nki/language/__init__.py (100%) rename {interop => Python}/noop.c (50%) rename {interop => Python}/pyproject.toml (63%) rename {interop => Python}/requirements.txt (100%) rename {interop/test => Python/tests}/examples/__init__.py (100%) rename {interop/test => Python/tests}/examples/average_pool.py (100%) rename {interop/test => Python/tests}/examples/fused_mamba.py (100%) rename {interop/test => Python/tests}/examples/getting_started.py (100%) rename {interop/test => Python/tests}/examples/index.py (100%) rename {interop/test => Python/tests}/examples/layernorm.py (100%) rename {interop/test => Python/tests}/examples/layout.py (100%) rename {interop/test => Python/tests}/examples/matmul.py (100%) rename {interop/test => Python/tests}/examples/mm.py (100%) rename {interop/test => Python/tests}/examples/prof.py (100%) rename {interop/test => Python/tests}/examples/rmsnorm.py (100%) rename {interop/test => Python/tests}/examples/sd_attention.py (100%) rename {interop/test => Python/tests}/examples/tensor_addition.py (100%) rename {interop/test => Python/tests}/examples/transpose2d.py (100%) rename {interop/test => Python/tests}/test_basic.py (100%) rename {interop/test => Python/tests}/test_examples.py (100%) delete mode 100644 interop/MANIFEST.in diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ce70e8d..701536b 100644 --- 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: [ubuntu-latest] + os: [macos-latest] runs-on: ${{ matrix.os }} @@ -25,33 +25,37 @@ 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: ./interop + working-directory: ./Python run: | python -m pip install --upgrade pip pip install -r requirements.txt - name: Run pytest - working-directory: ./interop + working-directory: ./Python run: | pytest - name: Make a wheel # https://github.com/pypa/cibuildwheel # Hit this: https://github.com/pypa/cibuildwheel/discussions/1926 env: + # https://github.com/leanprover/lean4/pull/6631/files + MACOSX_DEPLOYMENT_TARGET: 99.0 CIBW_BUILD_VERBOSITY: 1 + # CIBW_BUILD: 'cp39-manylinux_x86_64 cp39-macosx_arm64' run: | pip install cibuildwheel bin/make-wheel - - uses: actions/upload-artifact@v4 with: name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }} - path: ./wheelhouse/*.whl + path: ./.wheel/wheelhouse/*.whl diff --git a/.gitignore b/.gitignore index 9acca94..2516042 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,6 @@ -.lake/** +\.lake/ __pycache__/ *.so *.dSYM -/.vscode/ -.wheel/ +.vscode/ +/.wheel/ diff --git a/Export.lean b/Lean/Export.lean similarity index 100% rename from Export.lean rename to Lean/Export.lean diff --git a/KLR.lean b/Lean/KLR.lean similarity index 100% rename from KLR.lean rename to Lean/KLR.lean diff --git a/KLR/Core.lean b/Lean/KLR/Core.lean similarity index 100% rename from KLR/Core.lean rename to Lean/KLR/Core.lean diff --git a/KLR/Core/Basic.lean b/Lean/KLR/Core/Basic.lean similarity index 100% rename from KLR/Core/Basic.lean rename to Lean/KLR/Core/Basic.lean diff --git a/KLR/Core/Encode.lean b/Lean/KLR/Core/Encode.lean similarity index 100% rename from KLR/Core/Encode.lean rename to Lean/KLR/Core/Encode.lean diff --git a/KLR/Core/Oper.lean b/Lean/KLR/Core/Oper.lean similarity index 100% rename from KLR/Core/Oper.lean rename to Lean/KLR/Core/Oper.lean diff --git a/KLR/Core/Pretty.lean b/Lean/KLR/Core/Pretty.lean similarity index 100% rename from KLR/Core/Pretty.lean rename to Lean/KLR/Core/Pretty.lean diff --git a/KLR/Python.lean b/Lean/KLR/Python.lean similarity index 100% rename from KLR/Python.lean rename to Lean/KLR/Python.lean diff --git a/KLR/Trace.lean b/Lean/KLR/Trace.lean similarity index 100% rename from KLR/Trace.lean rename to Lean/KLR/Trace.lean diff --git a/KLR/Trace/Basic.lean b/Lean/KLR/Trace/Basic.lean similarity index 100% rename from KLR/Trace/Basic.lean rename to Lean/KLR/Trace/Basic.lean diff --git a/KLR/Trace/Builtin.lean b/Lean/KLR/Trace/Builtin.lean similarity index 100% rename from KLR/Trace/Builtin.lean rename to Lean/KLR/Trace/Builtin.lean diff --git a/KLR/Trace/NKI.lean b/Lean/KLR/Trace/NKI.lean similarity index 100% rename from KLR/Trace/NKI.lean rename to Lean/KLR/Trace/NKI.lean diff --git a/KLR/Trace/Python.lean b/Lean/KLR/Trace/Python.lean similarity index 100% rename from KLR/Trace/Python.lean rename to Lean/KLR/Trace/Python.lean diff --git a/KLR/Trace/Tensor.lean b/Lean/KLR/Trace/Tensor.lean similarity index 100% rename from KLR/Trace/Tensor.lean rename to Lean/KLR/Trace/Tensor.lean diff --git a/KLR/Trace/Types.lean b/Lean/KLR/Trace/Types.lean similarity index 100% rename from KLR/Trace/Types.lean rename to Lean/KLR/Trace/Types.lean diff --git a/KLR/Util.lean b/Lean/KLR/Util.lean similarity index 100% rename from KLR/Util.lean rename to Lean/KLR/Util.lean diff --git a/Main.lean b/Lean/Main.lean similarity index 100% rename from Main.lean rename to Lean/Main.lean diff --git a/lake-manifest.json b/Lean/lake-manifest.json similarity index 100% rename from lake-manifest.json rename to Lean/lake-manifest.json diff --git a/lakefile.lean b/Lean/lakefile.lean similarity index 100% rename from lakefile.lean rename to Lean/lakefile.lean diff --git a/lean-toolchain b/Lean/lean-toolchain similarity index 100% rename from lean-toolchain rename to Lean/lean-toolchain diff --git a/Python/MANIFEST.in b/Python/MANIFEST.in new file mode 100644 index 0000000..04f0dcc --- /dev/null +++ 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 b/Python/README.md similarity index 85% rename from interop/README.md rename to Python/README.md index cdb2ee2..48f1a09 100644 --- a/interop/README.md +++ b/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 b/Python/klr/__init__.py similarity index 100% rename from interop/klr/__init__.py rename to Python/klr/__init__.py diff --git a/interop/klr/parser.py b/Python/klr/parser.py similarity index 96% rename from interop/klr/parser.py rename to Python/klr/parser.py index 56fa7d9..0d42202 100644 --- a/interop/klr/parser.py +++ b/Python/klr/parser.py @@ -7,7 +7,6 @@ import json import numpy as np import os -import pathlib import subprocess import tempfile import types @@ -16,15 +15,22 @@ from textwrap import dedent from importlib.resources import files +def up(f, n): + d = os.path.dirname(f) + 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 - project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + project_root = up(os.path.abspath(__file__), 2) bin = project_root + '/bin/klr' 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') + #bin = files('klr').joinpath('bin/klr') + raise Exception("no file: " + bin) subprocess.run([bin, 'parse-json', infile.name], stdout=outfile, check=True) diff --git a/interop/nki/__init__.py b/Python/nki/__init__.py similarity index 100% rename from interop/nki/__init__.py rename to Python/nki/__init__.py diff --git a/interop/nki/isa/__init__.py b/Python/nki/isa/__init__.py similarity index 100% rename from interop/nki/isa/__init__.py rename to Python/nki/isa/__init__.py diff --git a/interop/nki/language/__init__.py b/Python/nki/language/__init__.py similarity index 100% rename from interop/nki/language/__init__.py rename to Python/nki/language/__init__.py diff --git a/interop/noop.c b/Python/noop.c similarity index 50% rename from interop/noop.c rename to Python/noop.c index f26b97c..d1adf6d 100644 --- a/interop/noop.c +++ b/Python/noop.c @@ -1,3 +1,5 @@ +// A dummy file to trick Python's `build` into makeing an arch-specific wheel. + #include int main() { diff --git a/interop/pyproject.toml b/Python/pyproject.toml similarity index 63% rename from interop/pyproject.toml rename to Python/pyproject.toml index 84d556b..9fc85c4 100644 --- a/interop/pyproject.toml +++ b/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 = "govereau@amazon.com"}, {name = "Sean McLaughlin", email = "seanmcl@amazon.com"}, ] +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 b/Python/requirements.txt similarity index 100% rename from interop/requirements.txt rename to Python/requirements.txt diff --git a/interop/test/examples/__init__.py b/Python/tests/examples/__init__.py similarity index 100% rename from interop/test/examples/__init__.py rename to Python/tests/examples/__init__.py diff --git a/interop/test/examples/average_pool.py b/Python/tests/examples/average_pool.py similarity index 100% rename from interop/test/examples/average_pool.py rename to Python/tests/examples/average_pool.py diff --git a/interop/test/examples/fused_mamba.py b/Python/tests/examples/fused_mamba.py similarity index 100% rename from interop/test/examples/fused_mamba.py rename to Python/tests/examples/fused_mamba.py diff --git a/interop/test/examples/getting_started.py b/Python/tests/examples/getting_started.py similarity index 100% rename from interop/test/examples/getting_started.py rename to Python/tests/examples/getting_started.py diff --git a/interop/test/examples/index.py b/Python/tests/examples/index.py similarity index 100% rename from interop/test/examples/index.py rename to Python/tests/examples/index.py diff --git a/interop/test/examples/layernorm.py b/Python/tests/examples/layernorm.py similarity index 100% rename from interop/test/examples/layernorm.py rename to Python/tests/examples/layernorm.py diff --git a/interop/test/examples/layout.py b/Python/tests/examples/layout.py similarity index 100% rename from interop/test/examples/layout.py rename to Python/tests/examples/layout.py diff --git a/interop/test/examples/matmul.py b/Python/tests/examples/matmul.py similarity index 100% rename from interop/test/examples/matmul.py rename to Python/tests/examples/matmul.py diff --git a/interop/test/examples/mm.py b/Python/tests/examples/mm.py similarity index 100% rename from interop/test/examples/mm.py rename to Python/tests/examples/mm.py diff --git a/interop/test/examples/prof.py b/Python/tests/examples/prof.py similarity index 100% rename from interop/test/examples/prof.py rename to Python/tests/examples/prof.py diff --git a/interop/test/examples/rmsnorm.py b/Python/tests/examples/rmsnorm.py similarity index 100% rename from interop/test/examples/rmsnorm.py rename to Python/tests/examples/rmsnorm.py diff --git a/interop/test/examples/sd_attention.py b/Python/tests/examples/sd_attention.py similarity index 100% rename from interop/test/examples/sd_attention.py rename to Python/tests/examples/sd_attention.py diff --git a/interop/test/examples/tensor_addition.py b/Python/tests/examples/tensor_addition.py similarity index 100% rename from interop/test/examples/tensor_addition.py rename to Python/tests/examples/tensor_addition.py diff --git a/interop/test/examples/transpose2d.py b/Python/tests/examples/transpose2d.py similarity index 100% rename from interop/test/examples/transpose2d.py rename to Python/tests/examples/transpose2d.py diff --git a/interop/test/test_basic.py b/Python/tests/test_basic.py similarity index 100% rename from interop/test/test_basic.py rename to Python/tests/test_basic.py diff --git a/interop/test/test_examples.py b/Python/tests/test_examples.py similarity index 100% rename from interop/test/test_examples.py rename to Python/tests/test_examples.py diff --git a/bin/klr b/bin/klr index 9fa7d3d..0165271 100755 --- a/bin/klr +++ b/bin/klr @@ -2,4 +2,4 @@ ROOT=$(dirname $(dirname $(readlink -f $0))) -$ROOT/.lake/build/bin/klr $@ +$ROOT/Lean/.lake/build/bin/klr $@ diff --git a/bin/make-wheel b/bin/make-wheel index d9f5598..80eab68 100755 --- 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 deleted file mode 100644 index d85df01..0000000 --- a/interop/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -include bin/*