diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7dd32cb..b6d2bdd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,7 +2,7 @@ name: CI on: push: - branches: ["main"] + branches: ["**"] # For now, let's build all branches. Roll this back if it gets too slow or we exhaust our quota. pull_request: branches: ["**"] # * does not match '/' workflow_dispatch: # For manually triggering a build: https://docs.github.com/en/actions/writing-workflows/choosing-when-your-workflow-runs/events-that-trigger-workflows#workflow_dispatch @@ -12,9 +12,10 @@ jobs: strategy: matrix: # We mainly care about hardware rather than OS - # macos-13 is x86. All later ones are arm64 + # macos-13 is x86 + # macos-latest is arm64 # ubuntu-latest is x64 - os: [macos-13, macos-latest, ubuntu-latest] + os: [ubuntu-latest, macos-latest, macos-13] runs-on: ${{ matrix.os }} @@ -26,7 +27,6 @@ jobs: - uses: leanprover/lean-action@v1 - name: Run tests run: lake exe klr - # Run pytest - uses: actions/setup-python@v5 with: @@ -41,3 +41,18 @@ jobs: working-directory: ./interop 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_SKIP: pp* # The build doesn't work on pypy: auditwheel: error: cannot repair "/tmp/cibuildwheel/built_wheel/klr-0.0.3-cp38-cp38-linux_i686.whl" to "manylinux2014_i686" ABI because of the presence of too-recent versioned symbols. You'll need to compile the wheel on an older toolchain. + run: | + pip install cibuildwheel + bin/make-wheel + - uses: actions/upload-artifact@v4 + with: + name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }} + 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/bin/make-wheel b/bin/make-wheel index ead4f3a..b49f339 100755 --- a/bin/make-wheel +++ b/bin/make-wheel @@ -3,20 +3,31 @@ 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/ +PYTHON_DIR=$ROOT/interop rm -rf $WHEEL_DIR -mkdir -p $WHEEL_DIR/klr/bin +mkdir -p $WHEEL_DIR/{klr,bin} 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 + +# Get klr binary +cp .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/pyproject.toml $WHEEL_DIR -cp $ROOT/interop/MANIFEST.in $WHEEL_DIR cd $WHEEL_DIR -python -m build + +if [ -z "${GITHUB_RUN_ID+x}" ]; then + echo "Building wheel locally" + 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 ./wheelhouse +fi # To upload the wheel, with the proper token in ~/.pypirc, run # diff --git a/interop/MANIFEST.in b/interop/MANIFEST.in index 944d6af..04f0dcc 100644 --- a/interop/MANIFEST.in +++ b/interop/MANIFEST.in @@ -1 +1,2 @@ -include klr/bin/* +# All files must be below the python project root +include bin/klr diff --git a/interop/README.md b/interop/README.md index cdb2ee2..48f1a09 100644 --- a/interop/README.md +++ b/interop/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/parser.py b/interop/klr/parser.py index 56fa7d9..7fdc6c9 100644 --- a/interop/klr/parser.py +++ b/interop/klr/parser.py @@ -7,7 +7,6 @@ import json import numpy as np import os -import pathlib import subprocess import tempfile import types @@ -16,10 +15,16 @@ 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 diff --git a/interop/noop.c b/interop/noop.c new file mode 100644 index 0000000..d1adf6d --- /dev/null +++ b/interop/noop.c @@ -0,0 +1,8 @@ +// A dummy file to trick Python's `build` into makeing an arch-specific wheel. + +#include + +int main() { + printf("Hello, World!\n"); + return 0; +} diff --git a/interop/pyproject.toml b/interop/pyproject.toml index 65cec98..c16f65d 100644 --- a/interop/pyproject.toml +++ b/interop/pyproject.toml @@ -1,11 +1,17 @@ +# NB: All this setuptools stuff took about 1000 hours to figure out so never delete this. +[build-system] +#requires = ["setuptools>=61.0"] +requires = ["setuptools"] +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" license = { file = "LICENSE" } keywords = ["trainium", "tpu", "pallas", "triton", "gpu"] @@ -20,6 +26,10 @@ dependencies = [ "ml_dtypes", "numpy", ] +requires-python = ">= 3.8" # cibuildwheel can't seem to handle eariler versions + +[project.urls] +Repository = "https://github.com/leanprover/KLR" [tool.pytest.ini_options] pythonpath = "." # Needed for tests to pass @@ -27,5 +37,13 @@ testpaths = [ "test", ] -[project.urls] -Repository = "https://github.com/leanprover/KLR" +[tool.setuptools] +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"]