Tracing for python source functions

NKL/FFI.lean

@@ -5,6 +5,7 @@ Authors: Paul Govereau
 -/
 import Lean
 import NKL.Python
+import NKL.Trace
 
 namespace NKL
 
@@ -16,12 +17,6 @@ local instance : MonadLift (Except String) IO where
 @[export parse_json]
 def parse_json (s : String) : IO Unit := do
   let kernel <- Python.Parsing.parse s
-  let names := kernel.funcs.map fun x => x.fst
-  let names := String.intercalate "," names
-  IO.println s!"Found functions: {names}"
-  for x in kernel.args do
-    IO.println s!"arg: {repr x}"
-  for x in kernel.kwargs do
-    IO.println s!"arg: {repr x}"
-  for x in kernel.globals do
-    IO.println s!"global: {repr x}"
+  let stmts <- NKL.Trace.runNKIKernel kernel
+  for s in stmts do
+    IO.println s!"{repr s}"

NKL/Python.lean

@@ -106,11 +106,11 @@ then the structure will be populated with:
   defaults = [1, 2]
   vararg = "args"
   kwonlyargs = [d, e]
-  kw_defaults = [None, 3]
+  kw_defaults = [("e", 3)]
   kwarg = "kwargs"
 
-Note, defaults and kw_defaults are inconsistent in how they treat
-missing arguments, but this is just how it works in the python AST.
+Note, this is slightly different from the official Python AST, which
+encodes the kw_defaults as a list with None for missing defaults.
 -/
 structure Args where
   posonlyargs : List String
@@ -122,17 +122,13 @@ structure Args where
   kwarg : Option String
   deriving Repr
 
-def Args.names (ax : Args) : List String :=
-  let xs := ax.posonlyargs.append ax.args
-  let xs := match ax.vararg with | none => xs | some x => xs.append [x]
-  let xs := xs.append ax.kwonlyargs
-  let xs := match ax.kwarg  with | none => xs | some x => xs.append [x]
-  xs
-
-def Args.all_defaults (ax : Args) : List (String × Expr') :=
-  let args := ax.posonlyargs ++ ax.args
-  let dflt := args.reverse.zip ax.defaults.reverse
-  dflt ++ ax.kw_defaults
+def Args.names (args : Args) : List String :=
+  args.posonlyargs ++ args.args ++ args.kwonlyargs
+
+def Args.all_defaults (args : Args) : List (String × Expr') :=
+  let pargs := args.posonlyargs ++ args.args
+  let dflt  := pargs.reverse.zip args.defaults.reverse
+  dflt ++ args.kw_defaults
 
 structure Fun where
   source : String

NKL/Trace.lean

@@ -3,7 +3,18 @@ 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.KLR
+import NKL.Python
 import NKL.Trace.Types
 import NKL.Trace.Basic
 import NKL.Trace.Builtin
---import NKL.Trace.Python
+import NKL.Trace.Python
+import NKL.Trace.NKI
+
+namespace NKL.Trace
+
+def runNKIKernel (k : NKL.Python.Kernel) : Except String (List NKL.KLR.Stmt) :=
+  tracer ⟨ .ofList NKIEnv, #[] ⟩ do
+    traceKernel k
+    let g <- get
+    return g.body.toList

NKL/Trace/Basic.lean

@@ -118,11 +118,8 @@ def unOp : String -> Term -> TraceM Term
   | op, _ => throw s!"unimp {op}"
 
 -- Comparison operators
--- TODO: need to think about comparison of tensors, in NKI this is object equality,
--- but I suspect this doesn't make sense and may be a source of bugs.
 
 def cmpOp : String -> Term -> Term -> TraceM Bool
-  | "Eq", .expr l _, .expr r _ => return l == r
   | s, l, r => throw s!"unsupported comparison operator {s} {repr l} {repr r}"
 
 def compare : Term -> List String -> List Term -> TraceM Term

NKL/Trace/NKI.lean

@@ -0,0 +1,49 @@
+/-
+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.KLR
+import NKL.Trace.Types
+import NKL.Trace.Builtin
+
+/-
+# NKI built-ins
+
+This module defines the builtin constants used by tracing for NKI kernels.
+-/
+namespace NKL.Trace
+open NKL.KLR
+
+private def module (s : String) : Name × Item :=
+  let name := s.toName
+  (name, .module name)
+
+private def const_var (s : String) : Name × Item :=
+  let name := s.toName
+  (name, .term (.expr (.var s) (.any name)))
+
+/-
+Note: this object contains a bunch of architecture parameters that
+need to be set according to which HW we are compiling for.
+TODO: figure out the mechanism for this.
+-/
+def tile_size : Global :=
+  let name := "nki.langauge.tile_size".toName
+  { name := name
+  , attr := attrs
+  , call := uncallable name
+  }
+where
+  attrs : GlobalAttr
+  | "pmax" => return .expr (.const $ .int 128) .int
+  | a => throw s!"unsupported attribute {a}"
+
+def NKIEnv : List (Name × Item) :=
+  [ module "nki"
+  , module "nki.language"
+  , const_var "nki.language.add"
+  , const_var "nki.language.load"
+  , const_var "nki.language.store"
+  , ("nki.language.tile_size".toName, .global tile_size)
+  ]