feat: simple pretty-printer for KLR

This change adds two related things: a pretty-printer for KLR terms,
and tensor names. Tensor names make the pretty printing nicer, but
have a second purpose. By naming all of the tensors, we can scan a KLR
kernel to collect up all of the input, output, and intermediate
tensors that will be needed to run the kernel. For argument tensors,
the generated tensor names are changed to the argument variable names;
this is just for readability.

NKL/FFI.lean

@@ -4,10 +4,12 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Paul Govereau
 -/
 import Lean
+import NKL.KLR.Pretty
 import NKL.Python
 import NKL.Trace
 
 namespace NKL
+open NKL.KLR
 
 local instance : MonadLift (Except String) IO where
   monadLift
@@ -19,4 +21,4 @@ def parse_json (s : String) : IO Unit := do
   let kernel <- Python.Parsing.parse s
   let stmts <- NKL.Trace.runNKIKernel kernel
   for s in stmts do
-    IO.println s!"{repr s}"
+    IO.println ("  " ++ Lean.format s) --s!"{s}\n{repr s}"

NKL/KLR/Basic.lean

@@ -14,16 +14,22 @@ portable format, a.k.a. Kernel Language Representation (KLR).
 
 namespace NKL.KLR
 
--- TODO switch to tensor lib
+-- TODO switch to TensorLib's version of these types
 --export TensorLib (Tensor Dtype Shape)
--- Mostly, NKL deals with empty tensors, so just check dtype and shape
--- TODO: talk to Sean about a more general BEq for Tensor
---instance : BEq Tensor where
---  beq t₁ t₂ := t₁.dtype == t₂.dtype && t₁.shape == t₂.shape
 
 abbrev Dtype := String
 abbrev Shape := List Int
-structure Tensor where
+
+/-
+A TensorName is essentially a typed variable, where the type
+must be a tensor type. When we flush out Typ below we may replace
+this with `Expr.var name (Typ.tensor dtype shape)`. For now, this
+only refers to dynamic tensors, or compile-time tensors, not
+trace-time tensors.
+-/
+
+structure TensorName where
+  name  : String
   dtype : Dtype
   shape : Shape
   deriving Repr, BEq
@@ -71,7 +77,7 @@ def toInt : Const -> Except String Int
 
 end Const
 
--- This correspondes to the "Quasi-Affine Expressions" in Neuron.
+-- This corresponds to the "Quasi-Affine Expressions" in Neuron.
 -- Note, `floor` is the usual integer division.
 inductive IndexExpr where
   | var (name : String)
@@ -94,7 +100,7 @@ inductive Index where
 inductive Expr where
   | var (x : String)
   | const (c : Const)
-  | tensor (t : Tensor)
+  | tensor (t : TensorName)
   | access (t : Expr) (ix : List Index)
   | call (f : Expr) (args : List Expr) (kwargs : List (String × Expr))
   deriving Repr, BEq

NKL/KLR/Encode.lean

@@ -265,7 +265,7 @@ private def chkIndex (i : Index) : Bool :=
 
 partial def encExpr : Expr -> ByteArray
   | .var s        => tag 0x30 [encString s]
-  | .tensor t     => tag 0x31 [encString t.dtype, encList encInt t.shape]
+  | .tensor t     => tag 0x31 [encString t.name, encString t.dtype, encList encInt t.shape]
   | .const c      => tag 0x32 [encConst c]
   | .access e ix  => tag 0x33 [encExpr e, encList encIndex ix]
   | .call f ax kw => tag 0x34 [encExpr f, encList encExpr ax, encList encKeyword kw]
@@ -276,7 +276,7 @@ where
 partial def decExpr : DecodeM Expr := do
   match (<- next) with
   | 0x30 => return .var (<- decString)
-  | 0x31 => return .tensor $ .mk (<- decString) (<- decList decInt)
+  | 0x31 => return .tensor $ .mk (<- decString) (<- decString) (<- decList decInt)
   | 0x32 => return .const (<- decConst)
   | 0x33 => return .access (<- decExpr) (<- decList decIndex)
   | 0x34 => return .call (<- decExpr) (<- decList decExpr) (<- decList decKeyword)
@@ -293,7 +293,7 @@ private def ixz := Index.coord (IndexExpr.int 0)
 
 #guard chkExpr nil
 #guard chkExpr (.var "var")
-#guard chkExpr (.tensor $ .mk "float32" [1,2,3])
+#guard chkExpr (.tensor $ .mk "t" "float32" [1,2,3])
 #guard chkExpr (.const (.int 1))
 #guard chkExpr (.access nil [ixz, ixz, ixz])
 #guard chkExpr (.call nil [nil, nil, nil] [("a", nil), ("b", nil)])

NKL/KLR/Pretty.lean

@@ -0,0 +1,80 @@
+/-
+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.Basic
+
+namespace NKL.KLR
+open Std
+
+/-
+This is a simple pretty printer for KLR terms. At some point, we may want to
+make this output valid python syntax that would parse and elaborate to the same
+KLR kernel. At the moment, there are too many unknowns to spend time on this.
+The format here is just for ease of debugging, feel free to modify as you wish.
+-/
+
+private def abracket (f : Format) : Format :=
+  Format.bracket "<" f ">"
+
+private def ppArgs [ToFormat a] (l : List a) : Format :=
+  Format.joinSep l ","
+
+def ppTensor (t : TensorName) : Format :=
+  "%" ++ t.name ++ abracket (t.dtype ++ ":" ++ ppArgs t.shape)
+
+def ppConst : Const -> Format
+  | .none       => "None"
+  | .bool true  => "True"
+  | .bool false => "False"
+  | .int i      => format i
+  | .float f    => format f
+  | .string s   => "\"" ++ s.push '"'
+
+private def addParens : Nat -> Format -> Format
+  | 0, f => f
+  | _, f => f.paren
+
+def ppIndexExpr (n : Nat) : IndexExpr -> Format
+  | .var x      => x
+  | .int i      => format i
+  | .neg e      => "-" ++ ppIndexExpr (n+1) e
+  | .add l r    => addParens n $ ppIndexExpr 1 l ++ "+" ++ ppIndexExpr 1 r
+  | .mul i e    => addParens n $ format i ++ "*" ++ ppIndexExpr 1 e
+  | .floor e i  => addParens n $ ppIndexExpr 1 e ++ "/" ++ format i
+  | .ceil e i   => "ceil" ++ Format.paren (ppIndexExpr 0 e ++","++ format i)
+  | .mod e i    => addParens n $ ppIndexExpr 1 e ++ "%" ++ format i
+
+def ppIndexExpr? : Option IndexExpr -> Format
+  | none   => "None"
+  | some e => ppIndexExpr 0 e
+
+def ppIndex : Index -> Format
+  | .ellipsis    => "..."
+  | .coord e     => ppIndexExpr? e
+  | .slice l u s => .joinSep ([l,u,s].map ppIndexExpr?) ":"
+
+partial def ppExpr : Expr -> Format
+  | .var x         => x
+  | .const c       => ppConst c
+  | .tensor t      => ppTensor t
+  | .access t ix   => .fill (ppExpr t ++ .sbracket (.joinSep (ix.map ppIndex) ","))
+  | .call f args kwargs =>
+      let args := args.map ppExpr
+      let kwargs := kwargs.map fun (x,e) => x ++ "=" ++ ppExpr e
+      .fill (ppExpr f ++ .paren (ppArgs (args ++ kwargs)))
+
+def ppStmt : Stmt -> Format
+  | .pass       => "pass"
+  | .expr e     => ppExpr e
+  | .ret e      => "ret" ++ ppExpr e
+  | .assign x e => x ++ " = " ++ ppExpr e
+  | .loop _ _ _ _ _ => "<loop>"
+
+instance : ToFormat TensorName where format := ppTensor
+instance : ToFormat Const      where format := ppConst
+instance : ToFormat IndexExpr  where format := ppIndexExpr 0
+instance : ToFormat Index      where format := ppIndex
+instance : ToFormat Expr       where format := ppExpr
+instance : ToFormat Stmt       where format := ppStmt

NKL/Trace/Python.lean

@@ -116,7 +116,8 @@ partial def expr' : Expr' -> Tracer Item
   | .const c => return .term (<- const c)
   | .tensor s dty => do
       let shape <- s.mapM integer
-      return .term (.expr (.tensor ⟨ dty, shape ⟩) (.tensor dty shape))
+      let name <- genName "t".toName
+      return .term (.expr (.tensor ⟨ name.toString, dty, shape ⟩) (.tensor dty shape))
   | .name id _ => lookup_item id.toName
   | .attr (.exprPos e p) id _ => do withPos p ((<- expr' e).attr id)
   | .tuple l _ => return .term (.tuple (<- l.mapM term))
@@ -190,6 +191,7 @@ partial def stmt' : Stmt' -> Tracer Unit
 partial def bind_args (f : Fun)
                       (args : List Term)
                       (kwargs : List (String × Term))
+                      (rename : Bool := false)
                       : Tracer (List (String × Term)) := do
   if f.args.vararg != none || f.args.kwarg != none then
     throw "var args not supported"
@@ -208,16 +210,21 @@ partial def bind_args (f : Fun)
       return (x, <- term' e)
     else
       throw s!"argument {x} not supplied"
+  -- rename tensors if asked to
+  let argmap := if rename then argmap.map renameTensors else argmap
   return argmap
+where
+  renameTensors : String × Term -> String × Term
+  | (s, .expr (.tensor t) ty) => (s, .expr (.tensor {t with name := s}) ty)
+  | other => other
 
 -- For a function call, first evaluate the argument in the current environment.
 -- Then enter a new environment and evaluate the function statements.
 partial def function_call (f : Fun)
                           (args : List Term)
                           (kwargs : List (String × Term))
                           : Tracer Unit := do
-  let args <- bind_args f args kwargs
-  --let args <- args.mapM fun (x,e) => return (x, e)
+  let args <- bind_args f args kwargs (rename:=true)
   withSrc f.source $ enterFun $ do
     args.forM fun (x,e) => do extend x.toName e
     f.body.forM stmt

NKL/Trace/Tensor.lean

@@ -23,21 +23,21 @@ private def tensor_call (op : String) (args : List Expr) : Term :=
 
 -- Unary operations on tensors
 
-def tensor_op (op : String) (t : Tensor) : TraceM Term :=
+def tensor_op (op : String) (t : TensorName) : TraceM Term :=
   return tensor_call op [.tensor t]
 
 -- Binary operations on tensors / scalars
 
-def tensor_tensor (op : String) (l r : Tensor) : TraceM Term :=
+def tensor_tensor (op : String) (l r : TensorName) : TraceM Term :=
   return tensor_call op [.tensor l, .tensor r]
 
-private def broadcast (t : Tensor) (c : Const) : Expr :=
+private def broadcast (t : TensorName) (c : Const) : Expr :=
   let args := t.shape.map fun i => Expr.const (.int i)
   let args := .const c :: args
   .call (.var "broadcast") args []
 
-def tensor_scalar (op : String) (t : Tensor) (c : Const) : TraceM Term :=
+def tensor_scalar (op : String) (t : TensorName) (c : Const) : TraceM Term :=
   return tensor_call op [ .tensor t, broadcast t c]
 
-def scalar_tensor (op : String) (c : Const) (t : Tensor) : TraceM Term :=
+def scalar_tensor (op : String) (c : Const) (t : TensorName) : TraceM Term :=
   return tensor_call op [ .tensor t, broadcast t c]