exir and MLX: fix folded-output clones, shared constants, and constant boolean masks

backends/mlx/builder/program_builder.py

@@ -462,7 +462,7 @@ def _make_io_slots(self):  # noqa: C901
         )
 
         output_kind_targets = defaultdict(set)
-        constant_tensors = {}
+        constant_tensors = set()
         user_inputs = []
         user_outputs = []
         mutable_buffers = []
@@ -491,17 +491,17 @@ def _make_io_slots(self):  # noqa: C901
             if isinstance(arg, TensorArgument):
                 if kind == InputKind.PARAMETER:
                     # Parameters are treated as constants (not mutated)
-                    constant_tensors[name] = target
+                    constant_tensors.add(name)
                 elif kind == InputKind.BUFFER:
                     if target in output_kind_targets[OutputKind.BUFFER_MUTATION]:
                         mutable_buffers.append(name)
                     else:
                         # Non-mutated buffers (like lifted tensor constants) are constants
-                        constant_tensors[name] = target
+                        constant_tensors.add(name)
                 elif kind == InputKind.USER_INPUT:
                     user_inputs.append(name)
                 elif kind == InputKind.CONSTANT_TENSOR:
-                    constant_tensors[name] = target
+                    constant_tensors.add(name)
                 else:
                     raise NotImplementedError(
                         f"Support for input {arg} is not implemented"
@@ -520,21 +520,12 @@ def _make_io_slots(self):  # noqa: C901
             else:
                 raise NotImplementedError(f"Support for input {arg} is not implemented")
 
-        target_to_constant_slot = {}
         for node in self.ep.graph.nodes:
             if node.op == "placeholder":
                 if node.users == {}:
                     continue
                 if node.name in constant_tensors:
-                    # A constant shared with other partitions is duplicated once per
-                    # user, so several placeholders here can name the same target.
-                    target = constant_tensors[node.name]
-                    if target in target_to_constant_slot:
-                        self.set_slot(node, target_to_constant_slot[target])
-                    else:
-                        target_to_constant_slot[target] = self.make_or_get_slot(
-                            node, id_space=IdSpace.Constant
-                        )
+                    self.make_or_get_slot(node, id_space=IdSpace.Constant)
                 elif node.name in user_inputs:
                     val = node.meta.get("val", None)
                     if isinstance(val, torch.Tensor) and not val.is_contiguous():

backends/mlx/ops.py

@@ -2406,6 +2406,51 @@ def _roll_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
+def _constant_mask_as_positions(
+    P: MLXProgramBuilder, n: Node, non_none: List[Tuple[int, Slot]]
+) -> List[Tuple[int, Slot]]:
+    """Replace a constant boolean index of aten.index.Tensor with its positions.
+
+    The MLX runtime gathers only with integer indices. A boolean mask that is
+    a constant of the program, used as the only index, becomes one integer
+    index per mask dimension. Any other boolean index is refused, so the CPU
+    runs it.
+    """
+    axis, mask = next(
+        (
+            (axis, idx)
+            for axis, idx in non_none
+            if getattr(n.args[1][axis].meta.get("val"), "dtype", None) == torch.bool
+        ),
+        (None, None),
+    )
+    if mask is None:
+        return non_none
+    mask_node = n.args[1][axis]
+    # The partitioner sees parameters and persistent buffers as fake tensors,
+    # so only a mask in ep.constants has values it can read.
+    if (
+        len(non_none) > 1
+        or mask.id_space != IdSpace.Constant
+        or P.get_placeholder_target(mask_node) not in P.ep.constants
+    ):
+        raise NotImplementedError(
+            "aten.index.Tensor: the MLX runtime gathers only with integer "
+            "indices, so only a constant boolean mask used as the only index "
+            "is lowered"
+        )
+    target, value = P.get_placeholder_target_and_tensor(mask_node)
+    if not value.any():
+        raise NotImplementedError(
+            "aten.index.Tensor: an all-false mask selects nothing, and the MLX "
+            "runtime cannot write an empty output"
+        )
+    return [
+        (axis + i, P.make_or_get_constant(f"{target}_positions_{i}", positions))
+        for i, positions in enumerate(value.nonzero(as_tuple=True))
+    ]
+
+
 @REGISTRY.register(target=[torch.ops.aten.index.Tensor])
 def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     args = P.args(n)
@@ -2427,6 +2472,8 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     if len(non_none) == 0:
         raise ValueError("aten.index.Tensor: all indices are None")
 
+    non_none = _constant_mask_as_positions(P, n, non_none)
+
     if len(non_none) == 1:
         axis, idx = non_none[0]
         idx_meta = n.args[1][axis].meta.get("val")

backends/mlx/test/test_partitioner.py

@@ -20,6 +20,7 @@
 from executorch.backends.mlx.passes import get_default_passes
 from executorch.backends.mlx.test.test_utils import get_mlx_node_counts
 from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower
+from executorch.exir.program import get_fake_program
 from executorch.runtime import Runtime
 from torch.export import export
 
@@ -79,7 +80,9 @@ def _run(model, inputs, transform_passes=None):
         counts = get_mlx_node_counts(path)
         method = Runtime.get().load_program(path).load_method("forward")
         out = method.execute(list(inputs))[0]
-    return counts, delegates, (out - ref).abs().max().item()
+    assert out.shape == ref.shape, (out.shape, ref.shape)
+    diff = (out - ref).abs()
+    return counts, delegates, diff.max().item() if diff.numel() else 0.0
 
 
 class TestMLXPartitionerConstantInputs(unittest.TestCase):
@@ -443,6 +446,126 @@ def forward(self, x):
         self.assertLess(err, 1e-6)
 
 
+class TestMLXPartitionerBoolIndex(unittest.TestCase):
+    """The MLX runtime gathers only with integer indices."""
+
+    class ConstantMask(nn.Module):
+        def __init__(self, mask, persistent=None, axis=0, scale=True):
+            super().__init__()
+            if persistent is None:
+                self.mask = mask
+            else:
+                self.register_buffer("mask", mask, persistent=persistent)
+            self.axis = axis
+            self.scale = scale
+
+        def forward(self, x):
+            y = (x + 1)[(slice(None),) * self.axis + (self.mask,)]
+            if self.axis:
+                # The leading dim's stride depends on how many elements the
+                # mask selects, which exir spec propagation rejects, so keep it
+                # in the delegate.
+                y = y.sum(self.axis)
+            return y * 2 if self.scale else y
+
+    def test_constant_mask_becomes_positions(self):
+        mask = torch.tensor([1, 0, 1, 1, 0, 1, 0, 1]) == 1
+        cases = {
+            "1d": (self.ConstantMask(mask), (8, 16)),
+            "2d": (self.ConstantMask(mask.reshape(1, 8)), (1, 8, 16)),
+            "3d": (self.ConstantMask(mask.reshape(2, 2, 2)), (2, 2, 2, 5)),
+            "buffer": (self.ConstantMask(mask, persistent=False), (8, 16)),
+            "axis_1": (self.ConstantMask(mask.reshape(2, 4), axis=1), (3, 2, 4, 5)),
+        }
+        for name, (model, shape) in cases.items():
+            with self.subTest(name):
+                _, delegates, err = _run(model.eval(), (torch.randn(shape),))
+                self.assertEqual(delegates, 1)
+                self.assertLess(err, 1e-5)
+
+    def test_other_bool_indices_run_on_the_cpu(self):
+        class ComputedMask(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.ids = torch.tensor([[1, 7, 7, 2, 7, 3, 7, 4]])
+
+            def forward(self, x):
+                return (x + 1)[self.ids == 7] * 2
+
+        class MaskInput(nn.Module):
+            def forward(self, x, mask):
+                return (x + 1)[mask] * 2
+
+        class MaskWithIndex(nn.Module):
+            def __init__(self, mask):
+                super().__init__()
+                self.mask = mask
+
+            def forward(self, x, i):
+                return (x + 1)[self.mask, i] * 2
+
+        class MaskThenUpdate(nn.Module):
+            def __init__(self, mask):
+                super().__init__()
+                self.register_buffer("mask", mask, persistent=False)
+
+            def forward(self, x):
+                y = (x + 1)[self.mask] * 2
+                self.mask.copy_(x[:, :, 0] > 0)
+                return y
+
+        mask = torch.tensor([[1, 0, 1, 1, 0, 1, 0, 1]]) == 1
+        x = torch.randn(1, 8, 16)
+        # An all-false mask gives an empty result, so it is the graph output
+        # here: an MLX op after it would have to write an empty output too.
+        all_false = self.ConstantMask(torch.zeros(1, 8, dtype=torch.bool), scale=False)
+        cases = {
+            "computed_mask": (ComputedMask(), (x,), 2),
+            "mask_input": (MaskInput(), (x, mask), 2),
+            "mask_with_index": (MaskWithIndex(mask), (x, torch.tensor([3])), 2),
+            "persistent_buffer": (self.ConstantMask(mask, persistent=True), (x,), 2),
+            "all_false": (all_false, (x,), 1),
+        }
+        for name, (model, inputs, expected) in cases.items():
+            with self.subTest(name):
+                _, delegates, err = _run(model.eval(), inputs)
+                self.assertEqual(delegates, expected)
+                self.assertLess(err, 1e-6)
+
+        with self.subTest("mutated_buffer"):
+            # Not run: the eager reference call would overwrite the mask first.
+            program = _lower(MaskThenUpdate(mask).eval(), (x,))
+            self.assertEqual(_delegate_count(program), 2)
+
+    def test_handler_refuses_masks_it_cannot_lower(self):
+        mask = torch.tensor([1, 0, 1, 1, 0, 1, 0, 1]) == 1
+        cases = {
+            # A persistent buffer has no values in the fake copy, so the handler
+            # must refuse it rather than fail reading it.
+            "persistent_buffer": (
+                self.ConstantMask(mask, persistent=True),
+                "only a constant boolean mask",
+            ),
+            "all_false": (
+                self.ConstantMask(torch.zeros(8, dtype=torch.bool)),
+                "all-false",
+            ),
+        }
+        for name, (model, reason) in cases.items():
+            with self.subTest(name):
+                # The same fake copy of the program that to_backend gives the
+                # partitioner.
+                ep = get_fake_program(export(model.eval(), (torch.randn(8, 16),)))
+                index = next(
+                    n for n in ep.graph.nodes if n.target == torch.ops.aten.index.Tensor
+                )
+                builder = MLXProgramBuilder(ep)
+                builder.check_support_only()
+                info = builder.node_info[index]
+                self.assertFalse(info.supported)
+                self.assertIn(reason, info.unsupported_reason)
+
+
 class TestMLXPartitionerMixedSupport(unittest.TestCase):
     """An operator is preserved from decomposition per operator, not per call."""
 

backends/transforms/test/test_create_delete_constant_placeholder.py

@@ -211,6 +211,10 @@ def forward(self, x):
         export(SharedWeight(), (torch.zeros(1),), strict=True)
     ).exported_program()
     (weight_name,) = exported_program.graph_signature.inputs_to_parameters
+    weight = next(n for n in exported_program.graph.nodes if n.name == weight_name)
+    # A copy is made for a user in another delegate.
+    for i, user in enumerate(weight.users):
+        user.meta["delegation_tag"] = f"tag{i}"
     duplicate_constant_node(exported_program, weight_name)
 
     weight_nodes = [

exir/backend/backend_api.py

@@ -423,8 +423,9 @@ def to_backend(
 
     update_to_real_program(tagged_exported_program, edge_program)
 
-    for tag, _ in partitioner_result.partition_tags.items():
-        _maybe_duplicate_constant_nodes(tagged_exported_program, tag)
+    _maybe_duplicate_constant_nodes(
+        tagged_exported_program, partitioner_result.partition_tags.keys()
+    )
 
     tagged_graph_module = _partition_and_lower(
         tagged_exported_program.graph_module,
@@ -750,8 +751,9 @@ def to_backend(
 
         update_to_real_program(tagged_exported_program, edge_program)
 
-        for tag, _ in partitioner_result.partition_tags.items():
-            _maybe_duplicate_constant_nodes(tagged_exported_program, tag)
+        _maybe_duplicate_constant_nodes(
+            tagged_exported_program, partitioner_result.partition_tags.keys()
+        )
 
         backend_id_to_call_submodule_nodes = _create_partitions(
             tagged_exported_program.graph_module,