feat: implement in and is operators

KLR/NKI/Annotations.lean

@@ -41,6 +41,7 @@ abbrev Ann := Pass Unit
 
 private def isValidName' : Name -> Bool
   | .str `neuronxcc.nki._pre_prod_kernels _
+  | .str `neuronxcc.nki._private_nkl _
   | .str `neuronxcc.nki._pre_prod_nkl _ => true
   | .str _ "neuronxcc" => false
   | .str n _ => isValidName' n

KLR/NKI/Simplify.lean

@@ -202,6 +202,14 @@ private def expr' (e' : Python.Expr') : Simplify Expr' :=
   | .boolOp op l => return (<- booleanOp op (<- exprs l)).expr
   | .binOp op l r => return .binOp (<- binOp op) (<- expr l) (<- expr r)
   | .unaryOp op e => return (<- unaryOp op) (<- expr e)
+  | .compare a [.is] [⟨.const .none, pos⟩] =>
+      return .binOp .eq (<- expr a) ⟨ .value .none, pos ⟩
+  | .compare a [.isNot] [⟨.const .none, pos⟩] =>
+      return .binOp .ne (<- expr a) ⟨ .value .none, pos ⟩
+  | .compare a [.isIn] [b] =>
+      return .call ⟨ .var `builtin.op.in, a.pos ⟩ [<- expr a, <- expr b] []
+  | .compare a [.notIn] [b] =>
+      return .call ⟨ .var `builtin.op.notin, a.pos ⟩ [<- expr a, <- expr b] []
   | .compare a ops l => do
       let a <- expr a
       let ops <- ops.mapM cmpOp
@@ -396,10 +404,10 @@ private def params (args : Python.Args) : Simplify (List Param) := do
     throw "varargs are not supported in NKI"
   if args.kwarg.isSome then
     warn "variable keyword arguments are not supported in NKI"
-  if args.posonlyargs.length > 0 then
-    warn "position-only arguments are not supported in NKI"
-  if args.kwonlyargs.length > 0 then
-    warn "keyword-only arguments are not supported in NKI"
+  --if args.posonlyargs.length > 0 then
+  --  warn "position-only arguments are not supported in NKI"
+  --if args.kwonlyargs.length > 0 then
+  --  warn "keyword-only arguments are not supported in NKI"
   let defaults := args.all_defaults
   let mut params := []
   for name in args.names do

KLR/Trace/Python.lean

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 -/
 
+import Batteries.Data.String
 import KLR.Core
 import KLR.Trace.ISA
 import KLR.Trace.Types
@@ -24,6 +25,7 @@ Python related builtins
 -/
 
 namespace KLR.Trace
+open Substring (containsSubstr)
 open Core
 
 /-
@@ -53,6 +55,22 @@ nki builtin.op.invert (t : Term) := do
   let i : Int <- fromNKI? t
   return .int i.toInt32.complement.toInt
 
+private def isin (t : Term) (l : Term) : Trace Bool := do
+  let l <- match l with
+    | .ref name _ => lookup name
+    | _ => pure l
+  match t, l with
+  | _, .tuple l => return l.contains t
+  | _, .list a => return a.contains t
+  | .string t, .string s => return containsSubstr s t
+  | _ , _ => throw "in operator not support on types {kindStr t} and {kindStr l}"
+
+nki builtin.op.in (t : Term) (l : Term) := do
+  return .bool (<- isin t l)
+
+nki builtin.op.notin (t : Term) (l : Term) := do
+  return .bool (<- isin t l).not
+
 /-
 The builtin.python namespace is mapped to the top-level namespace.
 For example, builtin.python.f will appear as f.