Add MLX integer support for aten.bitwise_and (pytorch#18979)

### Summary

Fixes pytorch#18925

`aten.bitwise_and` currently goes through the bool-only `LogicalAndNode`
path in the MLX delegate, which means integer tensors do not lower
correctly. This switches `aten.bitwise_and` to a dedicated
`BitwiseAndNode` while keeping `aten.logical_and` on the existing
logical path.

This change:
- adds `BitwiseAndNode` to the MLX schema
- adds `exec_bitwise_and()` to the MLX interpreter runtime
- registers `aten.bitwise_and.Tensor` and `aten.bitwise_and.Scalar` as
table-driven MLX binary ops
- keeps `aten.logical_and` on the existing bool-only logical path
- adds MLX op tests for bool, integer, and scalar `bitwise_and`

### Test plan

- `python3 -m py_compile backends/mlx/ops.py
backends/mlx/test/test_ops.py`
- `python3 backends/mlx/serialization/generate.py`
- `PYTHONPATH=src python3 -m executorch.backends.mlx.test.run_all_tests
--list | rg 'bitwise_and'`
  - in a local dev env, this registers:
    - `bitwise_and_bool`
    - `bitwise_and_int`
    - `bitwise_and_scalar`


cc @metascroy

Author: yadferhad

backends/mlx/ops.py

@@ -50,6 +50,7 @@
     AsStridedNode,
     AsTypeNode,
     Atan2Node,
+    BitwiseAndNode,
     BitwiseInvertNode,
     BroadcastToNode,
     CeilNode,
@@ -481,7 +482,14 @@ def _isnan_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     ([torch.ops.aten.minimum.default], MinimumNode, "aten.minimum", False),
     ([torch.ops.aten.atan2.default], Atan2Node, "aten.atan2", False),
     ([torch.ops.aten.logaddexp.default], LogAddExpNode, "aten.logaddexp", False),
+    ([torch.ops.aten.logical_and.default], LogicalAndNode, "aten.logical_and", False),
     ([torch.ops.aten.logical_or.default], LogicalOrNode, "aten.logical_or", False),
+    (
+        [torch.ops.aten.bitwise_and.Tensor, torch.ops.aten.bitwise_and.Scalar],
+        BitwiseAndNode,
+        "aten.bitwise_and",
+        True,
+    ),
     (
         [torch.ops.aten.lt.Tensor, torch.ops.aten.lt.Scalar],
         LessNode,
@@ -3143,34 +3151,6 @@ def _bitwise_not_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
-@REGISTRY.register(
-    target=[torch.ops.aten.logical_and.default, torch.ops.aten.bitwise_and.Tensor]
-)
-def _logical_and_handler(P: MLXProgramBuilder, n: Node) -> Slot:
-    """Handle aten.logical_and / aten.bitwise_and on bool tensors."""
-    args = P.args(n)
-    require_args(args, 2, 2, "aten.logical_and/bitwise_and")
-    require_kwargs(P.kwargs(n), set(), "aten.logical_and/bitwise_and")
-
-    # bitwise_and is only equivalent to logical_and for bool tensors.
-    if n.target == torch.ops.aten.bitwise_and.Tensor:
-        dtype = n.args[0].meta.get("val", None)
-        if dtype is not None and hasattr(dtype, "dtype") and dtype.dtype != torch.bool:
-            raise ValueError(
-                f"aten.bitwise_and on non-bool dtype {dtype.dtype} is not supported; "
-                "only bool tensors can be lowered via LogicalAndNode"
-            )
-    out = P.make_or_get_slot(n)
-    P.emit(
-        LogicalAndNode(
-            a=P.slot_to_tid(args[0]),
-            b=P.slot_to_tid(args[1]),
-            out=P.slot_to_tid(out),
-        )
-    )
-    return out
-
-
 @REGISTRY.register(target=[torch.ops.aten.scalar_tensor.default])
 def _scalar_tensor_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """This is equivalent to torch.full([], scalar, dtype=dtype)."""

backends/mlx/runtime/MLXInterpreter.h

@@ -1402,6 +1402,15 @@ exec_logical_or(const LogicalOrNode& n, ExecutionState& st, StreamOrDevice s) {
       n.out, logical_or(st.const_tensor_ref(n.a), st.const_tensor_ref(n.b), s));
 }
 
+inline void exec_bitwise_and(
+    const BitwiseAndNode& n,
+    ExecutionState& st,
+    StreamOrDevice s) {
+  st.set_tensor(
+      n.out,
+      bitwise_and(st.const_tensor_ref(n.a), st.const_tensor_ref(n.b), s));
+}
+
 inline void exec_tri(const TriNode& n, ExecutionState& st, StreamOrDevice s) {
   int rows = resolve_int(n.n, st);
   int cols = resolve_int(n.m, st);
@@ -2052,6 +2061,9 @@ class Interpreter {
       case OpCode::LOGICAL_OR:
         ops::exec_logical_or(std::get<LogicalOrNode>(instr.node), st, s);
         break;
+      case OpCode::BITWISE_AND:
+        ops::exec_bitwise_and(std::get<BitwiseAndNode>(instr.node), st, s);
+        break;
       case OpCode::TRI:
         ops::exec_tri(std::get<TriNode>(instr.node), st, s);
         break;

backends/mlx/serialization/schema.fbs

@@ -579,6 +579,12 @@ table LogicalOrNode {
     out: Tid (required);
 }
 
+table BitwiseAndNode {
+    a: Tid (required);
+    b: Tid (required);
+    out: Tid (required);
+}
+
 // Triangular matrix ops
 table TriNode {
     out: Tid (required);
@@ -1130,7 +1136,8 @@ union OpNode {
     ScanNode,
     MetalKernelNode,
     BitwiseInvertNode,
-    RollNode
+    RollNode,
+    BitwiseAndNode
     // BC: Add new op nodes here (append only)
 }
 

backends/mlx/test/test_ops.py

@@ -4274,6 +4274,8 @@ def create_model(self) -> nn.Module:
     {"op_name": "equal",         "op_fn": torch.eq, "shapes": [(2, 3, 4), (10,)], "dtypes": [torch.float32]},
     {"op_name": "not_equal",     "op_fn": torch.ne, "shapes": [(2, 3, 4), (10,)], "dtypes": [torch.float32]},
     # logical
+    {"op_name": "bitwise_and_bool", "op_fn": torch.bitwise_and, "shapes": _SHAPES_3, "dtypes": [torch.bool], "input_fn_a": _bool_input_fn(), "input_fn_b": _bool_input_fn()},
+    {"op_name": "bitwise_and_int",  "op_fn": torch.bitwise_and, "shapes": _SHAPES_3, "dtypes": [torch.int32, torch.int64], "input_fn_a": _int_input_fn(0, 256), "input_fn_b": _int_input_fn(0, 256)},
     {"op_name": "logical_and",   "op_fn": torch.logical_and, "shapes": [(2, 3, 4), (10,), (4, 8)], "dtypes": [torch.bool], "input_fn_a": _bool_input_fn(), "input_fn_b": _bool_input_fn()},
     {"op_name": "logical_or",    "op_fn": torch.logical_or,  "shapes": [(2, 3, 4), (10,), (4, 8)], "dtypes": [torch.bool], "input_fn_a": _bool_input_fn(), "input_fn_b": _bool_input_fn()},
 ]
@@ -4286,6 +4288,51 @@ def create_model(self) -> nn.Module:
     globals()[_cls.__name__] = _cls
 
 
+class BitwiseAndScalarModel(nn.Module):
+    def __init__(self, scalar):
+        super().__init__()
+        self.scalar = scalar
+
+    def forward(self, a: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_and(a, self.scalar)
+
+
+@register_test
+class BitwiseAndScalarTest(OpTestCase):
+    """Test case for aten.bitwise_and op (Tensor_Scalar variant)."""
+
+    name = "bitwise_and_scalar"
+
+    def __init__(
+        self,
+        shape: Tuple[int, ...],
+        dtype: torch.dtype,
+        scalar,
+    ):
+        self.shape = shape
+        self.dtype = dtype
+        self.scalar = scalar
+        shape_str = "x".join(str(s) for s in shape)
+        dtype_str = str(dtype).replace("torch.", "")
+        self.name = f"bitwise_and_scalar_{shape_str}_{dtype_str}"
+
+    @classmethod
+    def get_test_configs(cls) -> List["BitwiseAndScalarTest"]:
+        return [
+            cls(shape=(16,), dtype=torch.bool, scalar=True),
+            cls(shape=(4, 4), dtype=torch.int32, scalar=7),
+            cls(shape=(2, 3, 4), dtype=torch.int64, scalar=13),
+        ]
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        if self.dtype == torch.bool:
+            return _bool_input_fn()(self.shape, self.dtype)
+        return _int_input_fn(0, 256)(self.shape, self.dtype)
+
+    def create_model(self) -> nn.Module:
+        return BitwiseAndScalarModel(self.scalar)
+
+
 @register_test
 class PowerScalarTest(OpTestCase):
     """Test case for aten.pow op (Tensor_Scalar variant)."""