diff --git a/Export.lean b/Export.lean new file mode 100644 index 0000000..c7bda25 --- /dev/null +++ b/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 + ] diff --git a/NKL.lean b/NKL.lean new file mode 100644 index 0000000..0c2d9e6 --- /dev/null +++ b/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 diff --git a/NKL/FFI.lean b/NKL/FFI.lean new file mode 100644 index 0000000..0cb4d08 --- /dev/null +++ b/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 () diff --git a/NKL/NKI.lean b/NKL/NKI.lean new file mode 100644 index 0000000..b6362b6 --- /dev/null +++ b/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 diff --git a/lakefile.lean b/lakefile.lean index f6fbe2b..6c5c9a8 100644 --- a/lakefile.lean +++ b/lakefile.lean @@ -3,6 +3,11 @@ open Lake DSL package "NKL" where +lean_lib "NKL" where + defaultFacets := #[LeanLib.staticFacet] + +lean_lib "Export" where + @[default_target] lean_exe "nkl" where root := `Main diff --git a/lean-toolchain b/lean-toolchain index 7f0ea50..eff86fd 100644 --- a/lean-toolchain +++ b/lean-toolchain @@ -1 +1 @@ -leanprover/lean4:v4.10.0 +leanprover/lean4:v4.13.0-rc3 diff --git a/python/lean.py b/python/lean.py new file mode 100644 index 0000000..83241d4 --- /dev/null +++ b/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 diff --git a/python/lean_rffi.c b/python/lean_rffi.c new file mode 100644 index 0000000..7d24e1a --- /dev/null +++ b/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 + +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 + +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); +} diff --git a/python/mk.sh b/python/mk.sh new file mode 100644 index 0000000..3a5c301 --- /dev/null +++ b/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} + diff --git a/python/test.py b/python/test.py new file mode 100644 index 0000000..51c47ea --- /dev/null +++ b/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)) diff --git a/python/test.sh b/python/test.sh new file mode 100644 index 0000000..f3b7df7 --- /dev/null +++ b/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 +