feat: implement isinstance for objects and basic values

Author: govereau

KLR/Trace/Builtin.lean

@@ -33,9 +33,11 @@ def neuronxcc : Name := .str .anonymous "neuronxcc"
 def nki_ : Name := .str neuronxcc "nki"
 def nki_isa : Name := .str nki_ "isa"
 def nki_lang : Name := .str nki_ "language"
+def nki_typing : Name := .str nki_ "typing"
 
 def nl : String -> Name := .str nki_lang
 def nisa : String -> Name := .str nki_isa
+def nt : String -> Name := .str nki_typing
 
 -- conveience functions for creating environment entries
 

KLR/Trace/ISA.lean

@@ -88,6 +88,9 @@ def getTransposeOps(op: Option (List Int)) : Trace TransposeOps :=
 nki builtin.isa.get_nc_version := do
   lookup `arch
 
+nki builtin.typing.scalar (t : Term) := do
+  return .scalar t
+
 -- set_option linter.unusedVariables false
 
 nki builtin.isa.nc_matmul

KLR/Trace/NKI.lean

@@ -38,6 +38,7 @@ def NKIEnv : List (Name × Term) :=
   , module nki_
   , module nki_isa
   , module nki_lang
+  , module nki_typing
   , module `math
   , const_int (.str (nl "tile_size") "pmax") 128
   , const_int (.str (nl "tile_size") "gemm_stationary_fmax") 128

KLR/Trace/Python.lean

@@ -58,6 +58,34 @@ The builtin.python namespace is mapped to the top-level namespace.
 For example, builtin.python.f will appear as f.
 -/
 
+nki builtin.python.isinstance (t : Term) (ty : Term) := do
+  match t, ty with
+  | .object cls .., .source { name }
+  | .ref _ (.object cls), .source { name } => return .bool (cls == name)
+  | .none, .builtin `builtin.python.NoneType ..
+  | .bool .., .builtin `builtin.python.bool ..
+  | .int .., .builtin `builtin.python.int ..
+  | .float .., .builtin `builtin.python.float ..
+  | .string .., .builtin `builtin.python.str ..
+  | .tuple .., .builtin `builtin.python.tuple ..
+  | .list .., .builtin `builtin.python.list ..
+  | .ref _ .list, .builtin `builtin.python.list ..
+  | .dict .., .builtin `builtin.python.dict ..
+  | .ref _ .dict, .builtin `builtin.python.dict ..
+  | .scalar .., .builtin `builtin.typing.scalar ..
+  | .ellipsis, .builtin `builtin.python.ellipsis ..
+  | .slice .., .builtin `builtin.python.slice .. => return .bool true
+  | _, _ => return .bool false
+
+nki builtin.python.NoneType := do
+  return .none
+
+nki builtin.python.EllipsisType := do
+  return .ellipsis
+
+nki builtin.python.ellipsis := do
+  return .ellipsis
+
 nki builtin.python.slice (args : List Term) := do
   match args with
   | [e]     => return .slice (some 0) (<- fromNKI? e) (some 1)
@@ -91,6 +119,9 @@ nki builtin.python.abs (t : Term) := do
   | .float f => return .float f.abs
   | _ => throw "abs expects an integer or float number"
 
+nki builtin.python.tuple (l : List Term := []) := do
+  return .tuple l
+
 nki builtin.python.str (t : Term) := do
   return .string (<- t.toStr)
 

KLR/Trace/Term.lean

@@ -629,6 +629,7 @@ def builtinEnv : List (Name × Term) := Id.run do
     let names : List Name := match name with
       | .str `builtin.python n => [.str `builtins n, .str .anonymous n]
       | .str `builtin.isa n => [nisa n, name]
+      | .str `builtin.typing n => [nt n, name]
       | .str `builtin.lang n => [nl n, name]
       | _ => [name]
     names.map fun n => (n, fn)

KLR/Trace/Types.lean

@@ -112,6 +112,7 @@ inductive Term where
   | list     : Array Term -> Term
   | dict     : Array (String × Term) -> Term
   | tensor   : TensorLib.Tensor -> Term
+  | scalar   : Term -> Term
   -- indexing
   | ellipsis : Term
   | slice    : Option Int -> Option Int -> Option Int -> Term
@@ -138,6 +139,7 @@ def kindStr : Term → String
   | .list _ => "list"
   | .dict _ => "dict"
   | .tensor _ => "tensor"
+  | .scalar .. => "scalar"
   | .ellipsis => "ellipsis"
   | .slice _ _ _ => "slice"
   | .pointer _ => "pointer"
@@ -335,6 +337,7 @@ partial def toStr : Term -> Trace String
   | .slice a b c => return s!"slice({a},{b},{c})"
   | .pointer a => return s!"<Ptr({a.name})>"
   | .tensor .. => return "<Tensor>"
+  | .scalar .. => return "<scalar>"
 
 -- This is partial because the user could have created a heap graph
 partial def isTrue (t : Term) : Trace Bool := do
@@ -369,7 +372,8 @@ partial def isTrue (t : Term) : Trace Bool := do
      throw "The truth value of an array with more than one element is ambiguous"
     else
      throw "The truth value of an empty array is ambiguous"
-
+  | .scalar .. =>
+    throw "boolean conversion of scalar not supported"
 
 def isFalse (t : Term) : Trace Bool :=
   return not (<- t.isTrue)