CI: Make a wheel

This took a lot of time to figure out, but we can generate
linux and osx wheels now on GitHub Actions that seem functional
on both systems. Will probably hit problems, e.g. with different
Mac OSs, but I suspect that the right tweaks cibuildwheel arguments will
be able to fix those.

.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

.gitignore

@@ -1,6 +1,6 @@
-.lake/**
+\.lake/
 __pycache__/
 *.so
 *.dSYM
-/.vscode/
-.wheel/
+.vscode/
+/.wheel/

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
 #

interop/MANIFEST.in

@@ -1 +1,2 @@
-include klr/bin/*
+# All files must be below the python project root
+include bin/klr

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.

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

interop/noop.c

@@ -0,0 +1,8 @@
+// A dummy file to trick Python's `build` into makeing an arch-specific wheel.
+
+#include <stdio.h>
+
+int main() {
+    printf("Hello, World!\n");
+    return 0;
+}

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 = "[email protected]"},
   {name = "Sean McLaughlin", email = "[email protected]"},
 ]
+description = "Intermediate langauge for tensor compilers"
 readme = "README.md"
 license = { file = "LICENSE" }
 keywords = ["trainium", "tpu", "pallas", "triton", "gpu"]
@@ -20,12 +26,24 @@ 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
 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"]