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

backends/mlx/builder/program_builder.py

@@ -520,21 +520,21 @@ def _make_io_slots(self):  # noqa: C901
             else:
                 raise NotImplementedError(f"Support for input {arg} is not implemented")
 
-        target_to_constant_slot = {}
+        constant_targets = set()
         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.
+                    # Constant data is named by target, so a second slot would get none.
                     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
+                    if target in constant_targets:
+                        raise ValueError(
+                            f"constant '{target}' has more than one placeholder, "
+                            f"'{node.name}' is a second one"
                         )
+                    constant_targets.add(target)
+                    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,53 @@ 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, Any]]
+) -> List[Tuple[int, Slot, Any]]:
+    """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.
+    """
+    masks = [
+        (axis, idx)
+        for axis, idx, val in non_none
+        if isinstance(val, torch.Tensor) and val.dtype == torch.bool
+    ]
+    if not masks:
+        return non_none
+    axis, mask = masks[0]
+    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 (target := 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"
+        )
+    value = P.ep.constants[target]
+    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),
+            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)
@@ -2422,14 +2469,19 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     x_ndim = len(x_meta.shape) if x_meta is not None else None
 
     # Filter out None indices and track which axes they correspond to
-    non_none = [(i, idx) for i, idx in enumerate(idx_list) if idx is not None]
+    non_none = [
+        (i, idx, n.args[1][i].meta.get("val"))
+        for i, idx in enumerate(idx_list)
+        if idx is not None
+    ]
 
     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")
+        axis, idx, idx_meta = non_none[0]
         ndim_match = (
             x_meta is not None
             and idx_meta is not None
@@ -2464,8 +2516,8 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
             "aten.index.Tensor with multiple indices requires input shape metadata"
         )
 
-    indices = [P.slot_to_tid(idx) for _, idx in non_none]
-    axes = [i for i, _ in non_none]
+    indices = [P.slot_to_tid(idx) for _, idx, _ in non_none]
+    axes = [i for i, _, _ in non_none]
 
     # slice_sizes: 1 for indexed axes, full static size for non-indexed axes.
     indexed_axes = set(axes)

backends/mlx/test/test_partitioner.py

@@ -20,8 +20,10 @@
 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
+from torch.export.graph_signature import InputKind, InputSpec, TensorArgument
 
 
 class TestMLXPartitionerRejectsToEdge(unittest.TestCase):
@@ -79,7 +81,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):
@@ -353,6 +357,35 @@ def forward(self, x):
         self.assertGreater(delegates, 1)
         self.assertLess(err, 1e-5)
 
+    def test_second_placeholder_for_one_constant_raises(self):
+        class TwoReads(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("c", torch.randn(4))
+
+            def forward(self, x):
+                return x + self.c, x * self.c
+
+        ep = export(TwoReads(), (torch.randn(4),), strict=False)
+        original = next(n for n in ep.graph.nodes if n.name == "b_c")
+        with ep.graph.inserting_after(original):
+            second = ep.graph.placeholder("b_c_second")
+        second.meta = dict(original.meta)
+        mul = next(n for n in ep.graph.nodes if n.target == torch.ops.aten.mul.Tensor)
+        mul.replace_input_with(original, second)
+        ep.graph_signature.input_specs.insert(
+            1, InputSpec(InputKind.BUFFER, TensorArgument("b_c_second"), "c", True)
+        )
+
+        for check_only in (True, False):
+            with self.subTest(check_only=check_only):
+                builder = MLXProgramBuilder(ep)
+                with self.assertRaisesRegex(ValueError, "more than one placeholder"):
+                    if check_only:
+                        builder.check_support_only()
+                    else:
+                        builder.build()
+
 
 class TestMLXPartitionerGroupedKeys(unittest.TestCase):
     """Repeated KV heads are fused only when the caller enables the shared passes."""
@@ -443,6 +476,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,