POC: Create a Python wheel

This creates a wheel in the pypi test repository.
It is OSX only (or whatever architecture you run it on)
so it is not ready for prime time, but can be used for
demos and gives the guts of what a real solution will look
like.

.gitignore

@@ -3,3 +3,4 @@ __pycache__/
 *.so
 *.dSYM
 /.vscode/
+.wheel/

bin/make-wheel

@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e -u -o pipefail
+trap "kill 0" SIGINT SIGTERM
+
+ROOT=$(dirname $(dirname $(readlink -f $0)))
+WHEEL_DIR=$ROOT/.wheel/klr
+rm -rf $WHEEL_DIR
+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
+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
+
+# To upload the wheel, with the proper token in ~/.pypirc, run
+#
+#     python3 -m twine upload --verbose --repository testpypi dist/*
+#
+# from .wheel/klr

interop/MANIFEST.in

@@ -0,0 +1 @@
+include klr/bin/*

interop/README.md

@@ -0,0 +1,5 @@
+Python bindings for KLR
+
+# Usage
+
+TODO

interop/klr/parser.py

@@ -7,8 +7,8 @@
 import json
 import numpy as np
 import os
+import pathlib
 import subprocess
-import sys
 import tempfile
 import types
 
@@ -18,15 +18,14 @@
 
 
 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__))))
   bin = project_root + '/bin/klr'
-  if os.path.exists(bin):
-    subprocess.run(bin + ' parse-json ' + infile.name, stdout=outfile, shell=True, check=False)
-  else:
-    # This mess needs to be figured out once we have a proper wheel
-    # binary_path = files('klr') # using importlib.files
-    print("Can't find klr exe")
-    sys.exit(1)
+  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')
+  subprocess.run([bin, 'parse-json', infile.name], stdout=outfile, check=True)
 
 
 # This is a custom JSON encoder for use with AST nodes.
@@ -129,23 +128,29 @@ def apply_args(self, *args, **kwargs):
     self.kwargs = d
 
   def __call__(self, *args, **kwargs):
-    self.apply_args(*args, **kwargs)
-    json_kernel = self.json()
-    with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as temp_file:
-      temp_file.write(json_kernel)
-      temp_file.flush()
-      temp_filename = temp_file.name
-      klr_filename = temp_filename + ".klr"
-      with open(klr_filename, 'w') as klr_file:
-        dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-        run_klr(temp_file, klr_file)
-        with open(klr_filename, 'r') as file:
-          klr = file.read()
-          # This isn't great, but when the Lean exe discovers an error, e.g. an undefined
-          # variable, it just returns an empty string. We'll fix this.
-          if klr == '':
-            raise Exception("tracing failed")
-          return klr
+    klr_filename = None
+    try:
+      self.apply_args(*args, **kwargs)
+      json_kernel = self.json()
+      # temp_file will be deleted
+      with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as temp_file:
+        temp_file.write(json_kernel)
+        temp_file.flush()
+        temp_filename = temp_file.name
+        # klr_filename needs to be deleted after running klr
+        klr_filename = temp_filename + ".klr"
+        with open(klr_filename, 'w') as klr_file:
+          run_klr(temp_file, klr_file)
+          with open(klr_filename, 'r') as file:
+            klr = file.read()
+            # This isn't great, but when the Lean exe discovers an error, e.g. an undefined
+            # variable, it just returns an empty string. We'll fix this.
+            if klr == '':
+              raise Exception("tracing failed")
+            return klr
+    finally:
+      if os.path.exists(klr_filename):
+        os.remove(klr_filename)
 
   def ref_global(self, refname, val):
     return self.reference(self.globals, refname, val)

interop/pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "klr"
+version = "0.0.2"
+description = "Intermediate langauge for tensor compilers"
+authors = [
+  {name = "Paul Govereau", email = "[email protected]"},
+  {name = "Sean McLaughlin", email = "[email protected]"},
+]
+readme = "README.md"
+license = { file = "LICENSE" }
+keywords = ["trainium", "tpu", "pallas", "triton", "gpu"]
+dependencies = [
+  "numpy",
+]
+
+[project.urls]
+Repository = "https://github.com/leanprover/KLR"