Add mod prim operator

Summary: Adding the missing mod operator to our prim library which meta-emilian ran into.

Differential Revision: D58973824

Author: tarun292

exir/passes/executorch_prim_ops_registry.py

@@ -81,6 +81,11 @@ def eq(a: _SymScalar, b: _SymScalar) -> bool:
     return a == b
 
 
+@bind_pattern_to_op(executorch_prims_lib, "mod.Scalar(SymInt a, SymInt b) -> SymInt")
+def mod(a: SymInt, b: SymInt) -> SymInt:
+    return SymInt(int(a) % int(b))
+
+
 _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS: Dict[OpOverload, OpOverload] = {
     operator.sub: ops.backend.executorch_prim.sub.Scalar,
     operator.mul: ops.backend.executorch_prim.mul.Scalar,
@@ -92,6 +97,7 @@ def eq(a: _SymScalar, b: _SymScalar) -> bool:
     operator.lt: ops.backend.executorch_prim.lt.Scalar,
     operator.ge: ops.backend.executorch_prim.ge.Scalar,
     operator.le: ops.backend.executorch_prim.le.Scalar,
+    operator.mod: ops.backend.executorch_prim.mod.Scalar,
     torch.sym_float: ops.backend.executorch_prim.sym_float.Scalar,
 }
 

kernels/prim_ops/register_prim_ops.cpp

@@ -260,6 +260,17 @@ static Kernel prim_ops[] = {
           out = EValue(a.toInt() / b.toInt());
         }),
 
+    // executorch_prim::mod.int(int, int) -> int
+    Kernel(
+        "executorch_prim::mod.int",
+        [](RuntimeContext& context, EValue** stack) {
+          (void)context;
+          EValue& a = *stack[0];
+          EValue& b = *stack[1];
+          EValue& out = *stack[2];
+          out = EValue(a.toInt() % b.toInt());
+        }),
+
     // executorch_prim::et_copy_index.tensor(tensor, tensor) -> tensor
     Kernel("executorch_prim::et_copy_index.tensor", &et_copy_index),
     // executorch_prim::et_view.default(Tensor, int[]) -> Tensor

kernels/prim_ops/test/prim_ops_test.cpp

@@ -112,6 +112,9 @@ TEST_F(RegisterPrimOpsTest, TestAlgebraOps) {
   getOpsFn("executorch_prim::truediv.Scalar")(context, stack);
   EXPECT_FLOAT_EQ(stack[2]->toDouble(), 0.75);
 
+  getOpsFn("executorch_prim::mod.int")(context, stack);
+  EXPECT_EQ(stack[2]->toInt(), 3);
+
   getOpsFn("executorch_prim::sym_float.Scalar")(context, stack);
   EXPECT_FLOAT_EQ(stack[1]->toDouble(), 3.0);
 }