Add mod prim operator (#4057)

Summary:

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

Reviewed By: meta-emilian

Differential Revision: D58973824

Author: tarun292

examples/selective_build/test_selective_build.sh

@@ -32,9 +32,9 @@ test_buck2_select_ops_in_list() {
     ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="add_mul"
 
     echo "Running selective build test"
-    # set max_kernel_num=19: 17 primops, add, mul
+    # set max_kernel_num=20: 17 primops, add, mul
     $BUCK run //examples/selective_build:selective_build_test \
-        --config=executorch.max_kernel_num=19 \
+        --config=executorch.max_kernel_num=20 \
         --config=executorch.select_ops=list \
         -- --model_path=./add_mul.pte
 
@@ -100,11 +100,11 @@ test_cmake_select_ops_in_list() {
 
     local example_dir=examples/selective_build
     local build_dir=cmake-out/${example_dir}
-    # set MAX_KERNEL_NUM=19: 17 primops, add, mul
+    # set MAX_KERNEL_NUM=20: 17 primops, add, mul
     rm -rf ${build_dir}
     retry cmake -DBUCK2="$BUCK" \
             -DCMAKE_BUILD_TYPE=Release \
-            -DMAX_KERNEL_NUM=19 \
+            -DMAX_KERNEL_NUM=20 \
             -DEXECUTORCH_SELECT_OPS_LIST="aten::convolution.out,\
 aten::_native_batch_norm_legit_no_training.out,aten::hardtanh.out,aten::add.out,\
 aten::mean.out,aten::view_copy.out,aten::permute_copy.out,aten::addmm.out,\

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);
 }