Setup Lean/Python interop

.github/workflows/ci.yml

@@ -15,4 +15,4 @@ jobs:
       - uses: actions/checkout@v4
       - uses: leanprover/lean-action@v1
       - name: Run tests
-        run: lake exe anpu
+        run: lake exe nkl

.gitignore

@@ -0,0 +1 @@
+.lake/**

Export.lean

@@ -0,0 +1,90 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Govereau
+-/
+import Lean
+import NKL
+
+/-
+Generate python files form lean definitions.
+Note: this library is only used at compile-time.
+-/
+
+open Lean Meta
+
+abbrev Handle := IO.FS.Handle
+
+-- Place double-quotes around a string
+private def dq (str : String) := s!"\"{str}\""
+
+-- Extract the rightmost string from a Name: A.B.c ==> c
+-- We shouldn't find any anonymous or numerical names
+-- (such names will generate JSON that mismatches with Python)
+private def cname : Name -> MetaM String
+  | .str _ s => return s
+  | n => throwError s!"Invalid Constructor Name {n}"
+
+-- Print python namedtuple representing a single constructor
+private def printTuple
+  (h : Handle) (isStruct : Bool)
+  (name : String) (fields : List String) : MetaM Unit :=
+  do
+    let fields := List.map dq fields
+    let fn := if isStruct then "struct" else "cons"
+    h.putStrLn s!"{name.capitalize} = {fn}({dq name}, {fields})"
+
+-- Generate namedtuple for a structure
+private def genStructure (h : Handle) (si : StructureInfo) : MetaM Unit := do
+  let name <- cname si.structName
+  let ns <- List.mapM cname si.fieldNames.toList
+  printTuple h True name ns
+
+-- Generate namedtuple's for an inductive type
+-- Note, we assume the inductive type does not have the
+-- same name as any of its constructors.
+private def genInductive (h : Handle) (tc : Name) : MetaM Unit := do
+  let mut names : Array String := #[]
+  let tci <- getConstInfoInduct tc
+  for c in tci.ctors do
+    let ci <- getConstInfoCtor c
+    let name <- cname ci.name
+    names := names.push name
+    forallTelescopeReducing ci.type fun xs _ => do
+      let mut ns := []
+      for i in [:ci.numFields] do
+        let ld <- xs[ci.numParams + i]!.fvarId!.getDecl
+        ns := .cons ld.userName.toString ns
+      printTuple h False name ns.reverse
+  let rhs := String.intercalate " | " (.map .capitalize names.toList)
+  h.putStrLn s!"{<- cname tci.name} = {rhs}"
+
+private def genPython (h : Handle) (name : Name) : MetaM Unit := do
+  h.putStrLn ""
+  match getStructureInfo? (<- getEnv) name with
+  | some si => genStructure h si
+  | none => genInductive h name
+
+private def header :=
+"# This file is automatically generated, do no edit
+from functools import namedtuple
+
+def cons(name, args):
+  return namedtuple(name, args + [\"struct\"], defaults=[False])
+
+def struct(name, args):
+  return namedtuple(name, args + [\"struct\"], defaults=[True])
+"
+
+run_meta
+  let h <- IO.FS.Handle.mk "python/lean_types.py" IO.FS.Mode.write
+  h.putStr header
+  flip List.forM (genPython h)
+    [ `NKL.Const
+    , `NKL.BinOp
+    , `NKL.Expr
+    , `NKL.Index
+    , `NKL.Stmt
+    , `NKL.Arg
+    , `NKL.Fun
+    ]

Main.lean

@@ -1,4 +1,2 @@
-import ANPU
-
 def main : IO Unit :=
-  IO.println s!"Hello, {ANPU.hello}!"
+  IO.println s!"Hello, NKL!"

NKL.lean

@@ -0,0 +1,7 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Govereau
+-/
+import NKL.NKI
+import NKL.FFI

NKL/FFI.lean

@@ -0,0 +1,22 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Govereau
+-/
+import Lean
+import NKL.NKI
+
+namespace NKL
+
+-- temporary for testing
+
+@[export parse_json]
+def parse_json (json : String) : IO Unit := do
+  match Lean.Json.parse json with
+  | .error str => throw $ .userError str
+  | .ok jsn => do
+    match Lean.fromJson? jsn with
+    | .error str => throw $ .userError str
+    | .ok (_:Fun) => do
+      IO.println "parse successsful"
+      return ()

NKL/NKI.lean

@@ -0,0 +1,64 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Govereau
+-/
+import Lean
+
+/-!
+# Concrete Syntax of NKI kernels
+
+Representation of the "concrete" syntax of NKI kernels
+generated by the python frontend.
+-/
+
+namespace NKL
+
+inductive Const where
+  | nil
+  | bool (value: Bool)
+  | int (value: Int)
+  | float (value: Float)
+  | string (value: String)
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson
+
+inductive BinOp where
+  | And | Or
+  | Eq | NotEq | Lt | LtE | Gt | GtE
+  | Add | Sub | Mul | Div
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson
+
+mutual
+inductive Expr where
+  | value (c: Const)
+  | bvar (name: String)
+  | var (name: String)
+  | subscript (tensor: String) (ix: Array Index)
+  | binop (op: BinOp) (left right: Expr)
+  | call (f: String) (args: Array Expr)
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson
+
+inductive Index where
+  | coord (i : Expr)
+  | slice (l u step: Expr)
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson
+end
+
+inductive Stmt where
+  | ret(e: Expr)
+  | assign (x: String) (e: Expr)
+  | forloop (x: String) (iter: Expr) (body: List Stmt)
+  | gridcall (f: String) (ix: Array Index) (args: Array Expr)
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson
+
+structure Arg where
+  name  : String
+  type  : Option String := .none
+  value : Option Const := .none
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson
+
+structure Fun where
+  name : String
+  args : Array Arg
+  body : List Stmt
+  deriving Repr, BEq, Lean.ToJson, Lean.FromJson

README.md

@@ -1 +1 @@
-# ANPU
+# Neuron Kernel Language (NKL)

lakefile.lean

@@ -1,12 +1,13 @@
 import Lake
 open Lake DSL
 
-package "ANPU" where
-  -- add package configuration options here
+package "NKL" where
 
-lean_lib «ANPU» where
-  -- add library configuration options here
+lean_lib "NKL" where
+  defaultFacets := #[LeanLib.staticFacet]
+
+lean_lib "Export" where
 
 @[default_target]
-lean_exe "anpu" where
+lean_exe "nkl" where
   root := `Main

lean-toolchain

@@ -1 +1 @@
-leanprover/lean4:v4.10.0
+leanprover/lean4:v4.13.0-rc3

python/lean.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Released under Apache 2.0 license as described in the file LICENSE.
+# Authors: Paul Govereau
+
+from lean_types import *
+from lean_rffi import *
+
+def to_json_dict(obj):
+  if isinstance(obj, tuple) and hasattr(obj, '_fields'):
+    d = {k:to_json_dict(v) for k,v in obj._asdict().items()}
+    if not d.pop('struct'):
+      d = {obj.__class__.__name__: d}
+    return d
+  return obj

python/lean_rffi.c

@@ -0,0 +1,74 @@
+/*
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Paul Govereau
+*/
+
+// -----------------------------------------------------------------------------
+// Lean part
+
+#include <lean/lean.h>
+
+extern void lean_initialize_runtime_module();
+extern lean_object* initialize_NKL(uint8_t builtin, lean_object*);
+extern lean_object* parse_json(lean_object*, lean_object*);
+
+static lean_object *world = NULL;
+
+int lean_init() {
+  lean_initialize_runtime_module();
+  world = lean_io_mk_world();
+  lean_object *res = initialize_NKL(1, world);
+  if (!lean_io_result_is_ok(res)) {
+      lean_io_result_show_error(res);
+      lean_dec(res);
+      return 1;
+  }
+  lean_dec_ref(res);
+  lean_io_mark_end_initialization();
+  return 0;
+}
+
+int parse(const char *json) {
+  // lean_mk_string will copy the string
+  lean_object *s = lean_mk_string(json);
+  lean_object *res = parse_json(s, world);
+  if (!lean_io_result_is_ok(res)) {
+      // TODO: raise python exception rather than printing
+      lean_io_result_show_error(res);
+      lean_dec(res);
+      return 1;
+  }
+  lean_dec_ref(res);
+  return 0;
+}
+
+// -----------------------------------------------------------------------------
+// Python part
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+static PyObject* py_to_lean(PyObject *self, PyObject *args) {
+    const char *json;
+    if (!PyArg_ParseTuple(args, "s", &json))
+        return NULL;
+    // TODO: raise python exception on error
+    int res = parse(json);
+    return PyLong_FromLong(res);
+}
+
+static PyMethodDef methods[] = {
+    {"py_to_lean",  py_to_lean, METH_VARARGS, "Test python to lean"},
+    {NULL, NULL, 0, NULL}
+};
+
+static struct PyModuleDef module = {
+    PyModuleDef_HEAD_INIT, "lean_rffi", NULL, -1, methods
+};
+
+PyMODINIT_FUNC PyInit_lean_rffi(void) {
+    if (lean_init())
+        return NULL;
+    return PyModule_Create(&module);
+}

python/mk.sh

@@ -0,0 +1,28 @@
+#!/bin/sh
+
+set -x
+
+# TODO: for now using a simple script to build the python library
+# Need to decide which of lake or setuptools is better to use,
+# and how we will distribute everything
+
+# make sure libNKL.a and lean_types.py are generated
+(cd ..; lake build NKL Export)
+
+LEAN_CFLAGS="-I$(lean --print-prefix)/include"
+LEAN_LDFLAGS="-L$(lean --print-libdir) -L../.lake/build/lib"
+LEAN_LIBS="-lNKL -lInit_shared -lleanshared"
+
+# we can use the following to statically link the lean code
+#LEAN_LIBS="-lNKL -lLean -lStd -lInit -lleanrt -lleancpp -luv -lgmp -lc++"
+
+PY_EXT=$(python-config --extension-suffix)
+PY_CFLAGS=$(python-config --cflags)
+PY_LDFLAGS=$(python-config --ldflags)
+PY_LIBS="-lpython3.10"
+
+clang lean_rffi.c -dynamiclib -o lean_rffi${PY_EXT} \
+  ${LEAN_CFLAGS} ${PY_CFLAGS} \
+  ${LEAN_LDFLAGS} ${LEAN_LIBS} \
+  ${PY_LDFLAGS} ${PY_LIBS}
+

python/test.py

@@ -0,0 +1,15 @@
+import json
+from lean import *
+
+# A few quick tests TODO: setup pytest
+
+# this should fail
+py_to_lean("hello")
+
+# this should succeed
+py_to_lean('{"name":"name", "body":[], "args":[]}')
+
+# this should succeed
+f = Fun("name", [], [])
+j = to_json_dict(f)
+py_to_lean(json.dumps(j))

python/test.sh

@@ -0,0 +1,8 @@
+#!bin/sh
+
+DIR=$(lean --print-libdir)
+
+export LD_LIBRARY_PATH=$DIR   # for linux
+export DYLD_LIBRARY_PATH=$DIR # for OS/X
+python test.py
+