Tracing for python source functions #20
Conversation
Default function arguments were not being handled the same way as function arguments. It was possible that an expression that could be passed as an argument could not have been used as a default value for the same argument. This patch fixes this issue, and simplifies some of the parsing code for arguments and function signatures.
Remove some unnecessary parts of KLR, and reorganize the source files.
Basic tracing of Terms in the inner TraceM monad.
This patch adds basic tracing for user python functions. The main code is in Python.lean, and depends on definitions in Basic.lean and NKI.lean, which are incomplete. As more primitives are implemented, more user kernels will be supported.
|
|
||
| def runNKIKernel (k : NKL.Python.Kernel) : Except String (List NKL.KLR.Stmt) := | ||
| tracer ⟨ .ofList nki_env, #[] ⟩ do | ||
| traceKernel k |
There was a problem hiding this comment.
This is a kind of annoying aspect of Lean imo. Using the namespace mechanism, I now have no idea where to look for traceKernel without UI support. I wonder if we just get used to it, or have some conventions around files vs namespace organization.
| let name := s.toName | ||
| (name, .term (.expr (.var s) (.any name))) | ||
|
|
||
| def nki_env : List (Name × Item) := |
|
|
||
| def indexExpr' : Expr' -> Tracer KLR.IndexExpr | ||
| | .const (.int i) => return .int i | ||
| | .const c => throw s!"invalid constant {repr c} in index expression" |
There was a problem hiding this comment.
Wouldn't .none be np.newaxis? That's used a lot in kernels.
|
|
||
| def list_access (l : List Term) : List KLR.Index -> TraceM Term | ||
| | [.coord (some (.int i))] => | ||
| let n := i.toNat |
There was a problem hiding this comment.
This ignores Python negative integer interpretations. Does NKI not allow things like
In [146]: [1,2,3][-1]
Out[146]: 3
| for (n, f) in k.funcs do | ||
| let n := n.toName | ||
| if not (s.env.contains n) then | ||
| extend_global n (.source f) |
There was a problem hiding this comment.
I can't find this definition. This updates the TraceM inner monad? Maybe I'll take a break and catch up when you merge some of the earlier PRs.
| let names := f.args.names | ||
| if args.length + kwargs.length > names.length then | ||
| throw "too many arguments supplied (varargs not supported)" | ||
| let argmap <- f.args.names.enum.mapM fun (i,x) => do |
There was a problem hiding this comment.
Nice to know how python gets argument values now!
| | .assign xs e => assign xs e | ||
| | .augAssign x op e => do | ||
| stmt' (.assign [x] (.exprPos (.binOp op x e) (<- getPos))) | ||
| | .annAssign _ _ .none => return () |
| -- When we perform an assignment, we will either add to the environment | ||
| -- the term found on the RHS, or the variable itself. The latter case | ||
| -- allows us to lookup and find the variable without substituting | ||
| -- it's definition. |
| | .expr e ty => xs.forM fun x => do | ||
| extend x.toName (.expr (.var xs[0]!) ty) | ||
| add_stmt (KLR.Stmt.assign x e) | ||
| | t => xs.forM fun x => do extend x.toName t |
There was a problem hiding this comment.
Think you can remove the do here.
| | some e => indexExpr e | ||
|
|
||
| def index : Expr -> Tracer KLR.Index | ||
| | .exprPos (.const .ellipsis) _ => return .ellipsis |
There was a problem hiding this comment.
Maybe we should make more of an effort to get rid of ... if possible.
This patch adds basic tracing for user python functions. The main code
is in Python.lean, and depends on definitions in Basic.lean and
NKI.lean, which are incomplete. As more primitives are implemented,
more user kernels will be supported.
Note to reviewers: it may be helpful to check the tests in
interop/tests/test_basic.py.