Sym not prim op

exir/emit/test/test_emit.py

@@ -2405,6 +2405,53 @@ def forward(self, x):
                 self.assertTrue(expected.shape == et_result.shape)
                 self.assertTrue(torch.allclose(expected, et_result))
 
+    def test_emit_sym_not(self) -> None:
+        class SymNotModel(nn.Module):
+            def forward(self, x):
+                n = x.shape[0]
+                flag = n > 5
+                neg = torch.sym_not(flag)
+                val = torch.sym_float(neg)
+                return x + val
+
+        model = SymNotModel()
+        model.eval()
+        test_inputs = [
+            torch.randn(3, 4),  # n<=5: sym_not(False)=True, float(True)=1.0
+            torch.randn(8, 4),  # n>5: sym_not(True)=False, float(False)=0.0
+        ]
+        reference_outputs = []
+        with torch.no_grad():
+            for inp in test_inputs:
+                reference_outputs.append(model(inp))
+
+        batch_dim = Dim("batch", min=1, max=20)
+        dynamic_shapes = {"x": {0: batch_dim}}
+        exported_program = torch.export.export(
+            model, (test_inputs[0],), dynamic_shapes=dynamic_shapes
+        )
+        sym_not_nodes = [
+            n
+            for n in exported_program.graph.nodes
+            if n.op == "call_function" and n.target is torch.sym_not
+        ]
+        self.assertGreater(
+            len(sym_not_nodes), 0, "sym_not should appear in exported graph"
+        )
+
+        edge_program = to_edge(
+            exported_program,
+            compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
+        )
+        et_program = edge_program.to_executorch()
+        program_buffer = et_program.buffer
+        et_module = _load_for_executorch_from_buffer(program_buffer)
+        for inp, expected in zip(test_inputs, reference_outputs):
+            et_output = et_module.forward([inp])
+            et_result = et_output[0]
+            self.assertTrue(expected.shape == et_result.shape)
+            self.assertTrue(torch.allclose(expected, et_result))
+
     def test_emit_channels_last_constant(self) -> None:
         """Test that channels-last constant tensors are emitted correctly.
 

exir/passes/executorch_prim_ops_registry.py

@@ -124,6 +124,11 @@ def sym_min(a: _SymScalar, b: _SymScalar) -> bool:
     return min(a, b)  # pyre-ignore
 
 
+@bind_pattern_to_op(executorch_prims_lib, "sym_not.Scalar(Scalar a) -> bool")
+def sym_not(a: _SymScalar) -> bool:
+    return not a  # pyre-ignore
+
+
 _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS: Dict[Any, OpOverload] = {
     builtins.round: ops.backend.executorch_prim.round.Scalar,
     math.ceil: ops.backend.executorch_prim.ceil.Scalar,
@@ -143,6 +148,7 @@ def sym_min(a: _SymScalar, b: _SymScalar) -> bool:
     torch.sym_float: ops.backend.executorch_prim.sym_float.Scalar,
     torch.sym_max: ops.backend.executorch_prim.sym_max.Scalar,
     torch.sym_min: ops.backend.executorch_prim.sym_min.Scalar,
+    torch.sym_not: ops.backend.executorch_prim.sym_not.Scalar,
 }
 
 

kernels/prim_ops/register_prim_ops.cpp

@@ -391,6 +391,8 @@ static Kernel prim_ops[] = {
           } else if (a.isDouble()) {
             // TODO: This should be impossible
             out = EValue(a.toDouble());
+          } else if (a.isBool()) {
+            out = EValue(static_cast<double>(a.toBool()));
           } else {
             ET_KERNEL_CHECK_MSG(
                 context, false, InvalidType, /* void */, "%zu", (size_t)a.tag);
@@ -475,6 +477,36 @@ static Kernel prim_ops[] = {
         }),
 #endif
 
+#if !defined(EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD) || \
+    defined(INCLUDE_EXECUTORCH_PRIM_SYM_NOT_SCALAR)
+    // executorch_prim::sym_not.Scalar(bool a) -> bool
+    Kernel(
+        "executorch_prim::sym_not.Scalar",
+        [](KernelRuntimeContext& context, Span<EValue*> stack) {
+          ET_KERNEL_CHECK_MSG(
+              context,
+              stack.size() == 2,
+              InvalidProgram,
+              /* void */,
+              "Expected %zu args, got %zu",
+              (size_t)2,
+              stack.size());
+          EValue& a = *stack[0];
+          EValue& out = *stack[1];
+          if (a.isBool()) {
+            out = EValue(!a.toBool());
+          } else {
+            ET_KERNEL_CHECK_MSG(
+                context,
+                false,
+                InvalidType,
+                /* void */,
+                "sym_not only supports bool inputs, got %zu",
+                (size_t)a.tag);
+          }
+        }),
+#endif
+
 #if !defined(EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD) || \
     defined(INCLUDE_EXECUTORCH_PRIM_FLOORDIV_INT)
     // executorch_prim::floordiv.int(int, int) -> int

kernels/prim_ops/test/prim_ops_test.cpp

@@ -40,6 +40,7 @@ TEST_F(RegisterPrimOpsTest, OpRegistered) {
   EXPECT_TRUE(hasOpsFn("aten::sym_numel"));
   EXPECT_TRUE(hasOpsFn("executorch_prim::sym_max.Scalar"));
   EXPECT_TRUE(hasOpsFn("executorch_prim::sym_min.Scalar"));
+  EXPECT_TRUE(hasOpsFn("executorch_prim::sym_not.Scalar"));
 }
 
 TEST_F(RegisterPrimOpsTest, SymSizeReturnsCorrectValue) {
@@ -166,6 +167,43 @@ TEST_F(RegisterPrimOpsTest, SymMinReturnsCorrectValue) {
   EXPECT_EQ(stack[2]->toInt(), -5);
 }
 
+TEST_F(RegisterPrimOpsTest, SymNotReturnsCorrectValue) {
+  EValue values[2];
+  values[0] = EValue(true);
+  values[1] = EValue(false);
+
+  EValue* stack[2];
+  for (size_t i = 0; i < 2; i++) {
+    stack[i] = &values[i];
+  }
+
+  getOpsFn("executorch_prim::sym_not.Scalar")(context_, Span<EValue*>(stack));
+  EXPECT_EQ(stack[1]->toBool(), false);
+
+  // Test not(false) -> true
+  values[0] = EValue(false);
+  values[1] = EValue(false);
+  getOpsFn("executorch_prim::sym_not.Scalar")(context_, Span<EValue*>(stack));
+  EXPECT_EQ(stack[1]->toBool(), true);
+}
+
+TEST_F(RegisterPrimOpsTest, SymFloatHandlesBoolInput) {
+  EValue values[2];
+  values[0] = EValue(true);
+  values[1] = EValue(0.0);
+  EValue* stack[2];
+  for (size_t i = 0; i < 2; i++) {
+    stack[i] = &values[i];
+  }
+  getOpsFn("executorch_prim::sym_float.Scalar")(context_, Span<EValue*>(stack));
+  EXPECT_FLOAT_EQ(stack[1]->toDouble(), 1.0);
+
+  values[0] = EValue(false);
+  values[1] = EValue(0.0);
+  getOpsFn("executorch_prim::sym_float.Scalar")(context_, Span<EValue*>(stack));
+  EXPECT_FLOAT_EQ(stack[1]->toDouble(), 0.0);
+}
+
 TEST_F(RegisterPrimOpsTest, TestAlgebraOps) {
   EValue values[3];
   int64_t a = 3;