From ea85c09eaf553caa3b9cd1445e557353597dbd34 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 18:18:07 -0800
Subject: [PATCH 1/4] Clone a folded output with an ATen op when the program is
 not Edge yet

When constant_prop_pass folds a program output into a constant, it puts
a clone in front of the output. That clone was always an Edge dialect op
(exir_ops.edge.aten.clone or _clone_dim_order), even when the pass ran
on a program that was still ATen dialect, before to_edge.

Running the pass that early is sometimes needed, for example to fold a
constant op that is not Core ATen before to_edge_transform_and_lower
checks the graph. The Edge op in an ATen program then breaks any step
that verifies the program before to_edge. The MLX partitioner does that
in transform_for_pre_decomposition, so lowering fails with:

    SpecViolationError: Operator '<EdgeOpOverload: aten.clone.default>:
    schema = ...' is not an allowed operator type

The pass now checks the program's dialect. An Edge program gets the same
clone as before. Any other program, ATen or the training dialect that
torch.export.export returns, gets torch.ops.aten.clone.default, which
to_edge later turns into the right Edge op.

The new test folds a constant output in an ATen, a training and an Edge
program, and checks the clone op. For the ATen and training programs it
also checks that they verify and that to_edge turns the clone into the
Edge op.
---
 exir/passes/constant_prop_pass.py | 13 ++++++++-----
 exir/tests/test_passes.py         | 29 +++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py
index 44b0121268c..2f4aa73d90b 100644
--- a/exir/passes/constant_prop_pass.py
+++ b/exir/passes/constant_prop_pass.py
@@ -463,11 +463,12 @@ def _update_output_node_and_specs(
     output_specs = exported_program.graph_signature.output_specs
     assert len(output_nodes) == len(output_specs)
 
-    clone_op = (
-        exir_ops.edge.aten.clone.default
-        if _skip_dim_order
-        else exir_ops.edge.dim_order_ops._clone_dim_order.default
-    )
+    if exported_program.dialect != "EDGE":
+        clone_op = torch.ops.aten.clone.default
+    elif _skip_dim_order:
+        clone_op = exir_ops.edge.aten.clone.default
+    else:
+        clone_op = exir_ops.edge.dim_order_ops._clone_dim_order.default
 
     for i in range(len(output_specs)):
         out_node = output_nodes[i]
@@ -494,6 +495,8 @@ def constant_prop_pass(
     """
     This pass is for constant propagation for Exported Program with lifted parameters,
     as the parameters will not be shown up as `get_attr` but as `placeholder` to the graph.
+    A graph output that is a constant gets a clone in front of it: an Edge clone when
+    exported_program.dialect is "EDGE", and torch.ops.aten.clone.default otherwise.
 
     Args:
         exported_program: The ExportedProgram to perform constant propagation on.
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 81c5298aab1..ad626ec35e6 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -1637,6 +1637,35 @@ def forward(self) -> torch.Tensor:
         edge.exported_program()._validate()
         edge.to_executorch()
 
+    def test_constant_prop_for_output_clones_in_program_dialect(self) -> None:
+        class M(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                return x * 2, torch.arange(4, dtype=torch.float32) + 1
+
+        exported = export(M(), (torch.ones(4),), strict=True)
+        for ep in (exported.run_decompositions({}), exported):
+            with self.subTest(dialect=ep.dialect):
+                ep = constant_prop_pass(ep)
+                [_, folded] = ep.graph.output_node().args[0]
+                self.assertEqual(folded.target, torch.ops.aten.clone.default)
+                ep._validate()
+                edge = to_edge(
+                    ep, compile_config=EdgeCompileConfig(_skip_dim_order=False)
+                )
+                [_, folded] = edge.exported_program().graph.output_node().args[0]
+                self.assertEqual(
+                    folded.target, exir_ops.edge.dim_order_ops._clone_dim_order.default
+                )
+
+        with self.subTest(dialect="EDGE"):
+            edge = to_edge(
+                export(M(), (torch.ones(4),), strict=True),
+                compile_config=EdgeCompileConfig(_skip_dim_order=True),
+            )
+            ep = constant_prop_pass(edge.exported_program())
+            [_, folded] = ep.graph.output_node().args[0]
+            self.assertEqual(folded.target, exir_ops.edge.aten.clone.default)
+
     def test_constant_prop_preserves_memory_alloc(self) -> None:
         class Add(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor:

From e07c9f2589ee7e5cfda33c65942a96bebe673978 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 19:26:12 -0800
Subject: [PATCH 2/4] Copy a shared constant once per program, not once per
 user

When ops in different partitions, or in a partition and on the CPU, read
the same constant, duplicate_constant_node gives every user but one its
own copy of the placeholder, all with the original target. Two users in
one partition then bring two placeholders of it into that delegate, and
two users left on the CPU keep two in the owning program.

A program that lists the same lifted tensor constant twice fails to
unlift: ExportedProgram.module() raises "`<name>` is not a buffer". For a
delegate, to_executorch raises it too, because it runs the delegate's
program when it retraces the delegate call. For example, a constant read
by an add and a mul in one delegate and by a sin on the CPU fails this
way. So does a second partitioner that takes the users left on the CPU
into one delegate.

Users with the same delegation tag, or with none, now share one
placeholder, and the first group keeps the original. So copies are made
only for users in different programs. On a graph without tags none are
made, and the tests that call duplicate_constant_node directly now tag
their users.

The MLX builder gave placeholders that name the same target one constant
slot, to cope with the copy per user. Each delegate now gets one
placeholder per constant, so that code is removed.

The new test lowers a constant (a lifted tensor constant, a buffer, and
a parameter) that two delegates each read twice and that ops left in the
owning program also read, through two to_backend calls, once with the
ops left in the owning program reading it first and once last. It also
lowers both delegates through one to_backend call. A direct call on
untagged users makes no copy. It checks that each program takes the
constant once and that to_executorch passes.
---
 backends/mlx/builder/program_builder.py       | 19 +---
 ...test_create_delete_constant_placeholder.py |  4 +
 .../duplicate_constant_node_pass.py           | 59 ++++++-----
 exir/backend/test/BUCK                        |  1 +
 exir/backend/test/test_passes.py              | 17 +++-
 exir/backend/test/test_utils.py               | 99 ++++++++++++++++++-
 6 files changed, 160 insertions(+), 39 deletions(-)

diff --git a/backends/mlx/builder/program_builder.py b/backends/mlx/builder/program_builder.py
index d5ab4124579..74cbaf044d7 100644
--- a/backends/mlx/builder/program_builder.py
+++ b/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():
diff --git a/backends/transforms/test/test_create_delete_constant_placeholder.py b/backends/transforms/test/test_create_delete_constant_placeholder.py
index b85a14b036f..d805b4d23f5 100644
--- a/backends/transforms/test/test_create_delete_constant_placeholder.py
+++ b/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 = [
diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
index 07c3eb5e57a..81998649fbd 100644
--- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
+++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
@@ -6,7 +6,7 @@
 
 import copy
 import logging
-from typing import Optional, Set
+from typing import Dict, List, Optional, Set
 
 import torch
 from torch._export.utils import get_buffer, get_lifted_tensor_constant, get_param
@@ -33,6 +33,20 @@ def _get_attribute_or_constants(
     return constant_or_attribute
 
 
+def _users_reading_a_copy(node: torch.fx.Node) -> List[List[torch.fx.Node]]:
+    """
+    The users of node that read a copy of it, in groups that share one. The
+    users with the same delegation tag, or with none, form a group, so each
+    program takes a constant once: for example, a program that lists the
+    same lifted constant twice fails to unlift. The first group keeps the
+    original.
+    """
+    groups: Dict[Optional[str], List[torch.fx.Node]] = {}
+    for user in node.users:
+        groups.setdefault(user.meta.get("delegation_tag"), []).append(user)
+    return list(groups.values())[1:]
+
+
 # TODO: add other passes to duplicate call_function nodes
 def duplicate_constant_node(
     exported_program: ExportedProgram, candidate_node: str
@@ -40,6 +54,7 @@ def duplicate_constant_node(
     """
     A pass to duplicate the attributes/constants node (the candidate_node) in the graph. Mostly used for duplicating light-weight data.
     If the data is too large, try tag it with "no_copy" to prevent high memory usage and make it as part of the output.
+    Users with the same delegation tag, or with none, share one placeholder.
 
     Args:
         exported_program: the exported program to be modified. If constants nodes are copied, they will be added as new
@@ -76,9 +91,8 @@ def duplicate_constant_node(
                 raise RuntimeError(
                     f"placeholder node for non-params, non-buffer, and non-tensor constants should not be tagged: {node} "
                 )
-            users = list(node.users.keys())
-            for ith in range(1, len(node.users)):
-                copy_constant_or_attribute_fqn = node.name + f"_copy_{ith - 1}"
+            for ith, group in enumerate(_users_reading_a_copy(node)):
+                copy_constant_or_attribute_fqn = node.name + f"_copy_{ith}"
                 with exported_program.graph.inserting_before(
                     constant_or_attribute_node
                 ):
@@ -96,26 +110,27 @@ def duplicate_constant_node(
                     copied_constant_or_attribute_node.meta["val"] = (
                         constant_or_attribute_node.meta["val"]
                     )
-                    new_args = tuple(
-                        [
-                            (
-                                arg
-                                if arg != constant_or_attribute_node
+                    for user in group:
+                        new_args = tuple(
+                            [
+                                (
+                                    arg
+                                    if arg != constant_or_attribute_node
+                                    else copied_constant_or_attribute_node
+                                )
+                                for arg in user.args
+                            ]
+                        )
+                        new_kwargs = {
+                            key: (
+                                value
+                                if value != constant_or_attribute_node
                                 else copied_constant_or_attribute_node
                             )
-                            for arg in users[ith].args
-                        ]
-                    )
-                    new_kwargs = {
-                        key: (
-                            value
-                            if value != constant_or_attribute_node
-                            else copied_constant_or_attribute_node
-                        )
-                        for key, value in users[ith].kwargs.items()
-                    }
-                    users[ith].args = new_args
-                    users[ith].kwargs = new_kwargs
+                            for key, value in user.kwargs.items()
+                        }
+                        user.args = new_args
+                        user.kwargs = new_kwargs
                     # The copy keeps the original target, so it reads the
                     # original tensor and needs no data of its own.
                     new_input_specs.append(
diff --git a/exir/backend/test/BUCK b/exir/backend/test/BUCK
index 99913f65989..0068adf13e7 100644
--- a/exir/backend/test/BUCK
+++ b/exir/backend/test/BUCK
@@ -378,6 +378,7 @@ fbcode_target(_kind = runtime.python_test,
         "//executorch/exir/backend:backend_api",
         "//executorch/exir/backend:partitioner",
         "//executorch/exir/backend:utils",
+        "//executorch/exir/backend/test/demos/rpc:executor_backend_preprocess",
         "//executorch/exir/dialects:lib",
         "//executorch/extension/pybindings:portable_lib",  # @manual
         "//executorch/runtime/executor/test:test_backend_compiler_lib",
diff --git a/exir/backend/test/test_passes.py b/exir/backend/test/test_passes.py
index 815a5b5bbed..33414cc615c 100644
--- a/exir/backend/test/test_passes.py
+++ b/exir/backend/test/test_passes.py
@@ -34,12 +34,22 @@ def forward(self, x):
         )
 
         const_nodes = [
-            node.name
+            node
             for node in edge.exported_program().graph.nodes
             if node.op == "placeholder" and is_buffer(edge.exported_program(), node)
         ]
+        # Users in one program share the original.
+        self.assertEqual(
+            duplicate_constant_node(edge.exported_program(), const_nodes[0].name),
+            set(),
+        )
+        # A copy is made for a user in another delegate.
+        for i, user in enumerate(const_nodes[0].users):
+            user.meta["delegation_tag"] = f"tag{i}"
 
-        copied_nodes = duplicate_constant_node(edge.exported_program(), const_nodes[0])
+        copied_nodes = duplicate_constant_node(
+            edge.exported_program(), const_nodes[0].name
+        )
         self.assertEqual(len(copied_nodes), 1)
 
         # Check that the new constant node is in the graph
@@ -76,6 +86,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         user = users[1]
         user.args = tuple(a for a in user.args if a is not buffer_node)
         user.kwargs = {"other": buffer_node}
+        # A copy is made for a user in another delegate.
+        for i, u in enumerate(users):
+            u.meta["delegation_tag"] = f"tag{i}"
 
         # Patch validation since we modified the graph
         edge.exported_program()._validate = lambda: None
diff --git a/exir/backend/test/test_utils.py b/exir/backend/test/test_utils.py
index ff11444930a..c31fb7e8b2b 100644
--- a/exir/backend/test/test_utils.py
+++ b/exir/backend/test/test_utils.py
@@ -4,13 +4,21 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import itertools
 import unittest
 
 import torch
 from executorch import exir
 from executorch.exir import to_edge
 from executorch.exir.backend.backend_api import to_backend
-from executorch.exir.backend.partitioner import Partitioner, PartitionResult
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.backend.test.demos.rpc.executor_backend_preprocess import (
+    ExecutorBackend,
+)
 from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
 from executorch.exir.backend.utils import (
     format_delegated_graph,
@@ -18,9 +26,11 @@
     get_delegates,
     get_non_lowered_nodes,
     is_identical_graph,
+    tag_constant_data,
 )
 
 from executorch.exir.dialects._ops import bind_pattern_to_op, ops as exir_ops
+from executorch.exir.lowered_backend_module import get_lowered_backend_modules
 from torch.export import export, ExportedProgram
 from torch.fx import symbolic_trace
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
@@ -390,3 +400,90 @@ def forward(self, a, x, b):
             payload_names,
             f"Delegate names mismatch: {delegate_names} vs {payload_names}",
         )
+
+    def test_to_backend_gives_each_program_one_copy_of_a_shared_constant(self):
+        class UseConst(torch.nn.Module):
+            def __init__(self, cpu_first, kind="constant"):
+                super().__init__()
+                const = torch.ones(2)
+                if kind == "buffer":
+                    self.register_buffer("const", const)
+                else:
+                    # A plain tensor attribute is lifted as a constant tensor.
+                    self.const = (
+                        torch.nn.Parameter(const) if kind == "parameter" else const
+                    )
+                self.cpu_first = cpu_first
+
+            def forward(self, x):
+                if self.cpu_first:
+                    s, c = torch.sin(self.const), torch.cos(self.const)
+                y = (x + self.const) * self.const
+                z = (x - self.const) / self.const
+                if not self.cpu_first:
+                    s, c = torch.sin(self.const), torch.cos(self.const)
+                return y, z, s, c
+
+        class TagNodes(Partitioner):
+            def __init__(self, tags):
+                super().__init__()
+                self.tags = tags
+
+            def partition(self, program: ExportedProgram) -> PartitionResult:
+                for node in program.graph.nodes:
+                    if node.name in self.tags:
+                        node.meta["delegation_tag"] = self.tags[node.name]
+                tag_constant_data(program)
+                return PartitionResult(
+                    tagged_exported_program=program,
+                    partition_tags=dict.fromkeys(
+                        self.tags.values(),
+                        DelegationSpec(ExecutorBackend.__name__, []),
+                    ),
+                )
+
+        def counts(program):
+            programs = [program] + [
+                module.original_module
+                for module in get_lowered_backend_modules(program.graph_module)
+            ]
+            return [
+                [spec.target for spec in p.graph_signature.input_specs].count("const")
+                for p in programs
+            ]
+
+        # Each delegate reads the constant twice, and sin and cos read it
+        # outside of them. Each program takes it once, whichever reads first.
+        for kind, cpu_first in itertools.product(
+            ("constant", "buffer", "parameter"), (False, True)
+        ):
+            with self.subTest(kind=kind, cpu_first=cpu_first):
+                lowered = to_edge(
+                    export(UseConst(cpu_first, kind), (torch.ones(2),), strict=True)
+                )
+                lowered = lowered.to_backend(
+                    TagNodes({"aten_add_tensor": "tag0", "aten_mul_tensor": "tag0"})
+                )
+                self.assertEqual(counts(lowered.exported_program()), [1, 1])
+                lowered.exported_program().module()
+                lowered = lowered.to_backend(
+                    TagNodes({"aten_sub_tensor": "tag1", "aten_div_tensor": "tag1"})
+                )
+                self.assertEqual(counts(lowered.exported_program()), [1, 1, 1])
+                lowered.to_executorch()
+
+        # One partitioner takes both delegates, so one call copies it twice.
+        with self.subTest(one_partitioner=True):
+            lowered = to_edge(export(UseConst(False), (torch.ones(2),), strict=True))
+            lowered = lowered.to_backend(
+                TagNodes(
+                    {
+                        "aten_add_tensor": "tag0",
+                        "aten_mul_tensor": "tag0",
+                        "aten_sub_tensor": "tag1",
+                        "aten_div_tensor": "tag1",
+                    }
+                )
+            )
+            self.assertEqual(counts(lowered.exported_program()), [1, 1, 1])
+            lowered.to_executorch()

From 91720c30bca3e60bf983725e16855804fe4332cd Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 3 Nov 2026 03:01:46 -0800
Subject: [PATCH 3/4] MLX: gather with the positions of a constant boolean mask

The MLX partitioner takes aten.index.Tensor whatever the dtype of its
indices. The MLX runtime gathers only with integer indices, so x[mask]
with a boolean mask lowers, then fails when it runs:

    MLX execute failed: [Gather] Boolean indices not supported

The mask is often a constant of the program, for example one that
constant_prop_pass folded before lowering. Its values are known at
lowering time, so the index handler now replaces it with its integer
positions, one index per mask dimension, and the gather stays in the
delegate.

Any other boolean index is refused, so the partitioner leaves it on the
CPU:

- a mask computed in the graph, such as x[x > 0]
- a mask that is a method input, or a buffer the program mutates
- a mask that is a parameter or a persistent buffer, because the
  partitioner sees those as fake tensors and cannot read the mask
- a constant mask used together with other indices
- an all-false mask, because the result is empty and the MLX runtime
  cannot write an empty output

The new tests run constant masks (1-D, 2-D, 3-D, a non-persistent
buffer, and a mask at axis 1), which must stay in one delegate. They
also run a mask computed in the graph, a mask passed as an input, a
constant mask used with another index, a persistent buffer, a buffer
the program mutates, and an all-false mask, none of which may lower the
index. A direct check on the fake program the partitioner sees confirms
that the handler itself refuses the persistent buffer and the all-false
mask, rather than failing on them.
---
 backends/mlx/ops.py                   |  47 ++++++++++
 backends/mlx/test/test_partitioner.py | 125 +++++++++++++++++++++++++-
 2 files changed, 171 insertions(+), 1 deletion(-)

diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index 5db12feef82..e356d4e7499 100644
--- a/backends/mlx/ops.py
+++ b/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")
diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py
index 5897fe58bee..a2c5a096383 100644
--- a/backends/mlx/test/test_partitioner.py
+++ b/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."""
 

From 2450a5a9ea73d938f828f32e1576af474edb45e0 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Tue, 3 Nov 2026 05:57:19 -0800
Subject: [PATCH 4/4] Copy the shared constants of all delegates in one pass

To give each delegate its own placeholder for a constant it shares with
another delegate, to_backend called duplicate_constant_node once per shared
constant, for each partition tag. Every call recompiled the graph module and
validated the whole program, so the cost grew with the number of shared
constants times the size of the graph. A model lowered to many delegates,
with thousands of folded constants, spent about an hour here.

Now to_backend collects the shared constants of all tags first and
copies them in one pass with the new duplicate_constant_nodes, which
recompiles and validates once. _maybe_duplicate_constant_nodes takes all
the tags, and raises a TypeError for a single tag string, which would
otherwise be read as one tag per character. The pass reads the kind of
each constant from its input spec, instead of through get_param,
get_buffer and get_lifted_tensor_constant, which rebuild the signature
lookups on every call. So a buffer whose data is not in the state dict,
which the verifier allows, no longer raises a KeyError here; its copy
needs no data. Each shared constant still gets one copy per other
program. The group of its first user in graph order keeps the original.
Before, when one op read two shared constants, which delegate kept the
original depended on the tag order, on the iteration order of a set of
names, and on the order in which earlier copies had rewired that op, so
it could change from run to run.
---
 exir/backend/backend_api.py                   | 10 ++-
 .../duplicate_constant_node_pass.py           | 68 +++++++--------
 exir/backend/test/test_utils.py               | 82 +++++++++++++++++++
 exir/backend/utils.py                         | 20 +++--
 4 files changed, 134 insertions(+), 46 deletions(-)

diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py
index 8104595bfe1..7c37f94bdfe 100644
--- a/exir/backend/backend_api.py
+++ b/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,
diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
index 81998649fbd..a6cfda16742 100644
--- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
+++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
@@ -9,40 +9,24 @@
 from typing import Dict, List, Optional, Set
 
 import torch
-from torch._export.utils import get_buffer, get_lifted_tensor_constant, get_param
 
 from torch.export import ExportedProgram
-from torch.export.exported_program import InputSpec, TensorArgument
+from torch.export.exported_program import InputKind, InputSpec, TensorArgument
 
 
-def _get_attribute_or_constants(
-    exported_program: ExportedProgram, node: torch.fx.Node
-) -> Optional[torch.Tensor]:
-    # get either attribute node or constant constant
-    maybe_param = get_param(exported_program, node)
-    maybe_buffer = get_buffer(exported_program, node)
-    maybe_lifted_tensor = get_lifted_tensor_constant(exported_program, node)
-
-    constant_or_attribute = None
-    if maybe_param is not None:
-        constant_or_attribute = maybe_param
-    elif maybe_buffer is not None:
-        constant_or_attribute = maybe_buffer
-    elif maybe_lifted_tensor is not None:
-        constant_or_attribute = maybe_lifted_tensor
-    return constant_or_attribute
-
-
-def _users_reading_a_copy(node: torch.fx.Node) -> List[List[torch.fx.Node]]:
+def _users_reading_a_copy(
+    node: torch.fx.Node, order: Dict[torch.fx.Node, int]
+) -> List[List[torch.fx.Node]]:
     """
     The users of node that read a copy of it, in groups that share one. The
     users with the same delegation tag, or with none, form a group, so each
     program takes a constant once: for example, a program that lists the
-    same lifted constant twice fails to unlift. The first group keeps the
-    original.
+    same lifted constant twice fails to unlift. The group of the first user
+    in graph order keeps the original.
     """
     groups: Dict[Optional[str], List[torch.fx.Node]] = {}
-    for user in node.users:
+    # Rewiring a user for another constant moves it to the end of node.users.
+    for user in sorted(node.users, key=order.__getitem__):
         groups.setdefault(user.meta.get("delegation_tag"), []).append(user)
     return list(groups.values())[1:]
 
@@ -64,17 +48,30 @@ def duplicate_constant_node(
     Returns:
         The set of the names of the new constant nodes
     """
-    to_be_copied = [
-        node
+    return duplicate_constant_nodes(exported_program, {candidate_node})
+
+
+def duplicate_constant_nodes(
+    exported_program: ExportedProgram, candidate_nodes: Set[str]
+) -> Set[str]:
+    """
+    duplicate_constant_node for several constant nodes at once. It recompiles
+    and validates the program once, rather than once per node, which matters
+    on a large graph with many shared constants.
+
+    Returns:
+        The set of the names of the new constant nodes
+    """
+    if not any(
+        node.op == "placeholder" and node.name in candidate_nodes
         for node in exported_program.graph.nodes
-        if node.name == candidate_node and node.op == "placeholder"
-    ]
-    if len(to_be_copied) == 0:
+    ):
         logging.info("no constant node to be copied")
         return set()
     new_input_specs = []
     old_signature = exported_program.graph_signature
     copied_nodes = set()
+    order = {node: i for i, node in enumerate(exported_program.graph.nodes)}
 
     placeholder_idx = -1
     for node in exported_program.graph.nodes:
@@ -84,14 +81,19 @@ def duplicate_constant_node(
         placeholder_idx += 1
         old_input_spec = old_signature.input_specs[placeholder_idx]
         old_input_spec_copy = copy.deepcopy(old_input_spec)
-        if node == to_be_copied[0]:
+        if node.name in candidate_nodes:
             constant_or_attribute_node = node
-            constant_or_attribute = _get_attribute_or_constants(exported_program, node)
-            if constant_or_attribute is None:
+            # The input spec kind answers this without get_param, get_buffer and
+            # get_lifted_tensor_constant, which rebuild a dict on every call.
+            if old_input_spec.kind not in (
+                InputKind.PARAMETER,
+                InputKind.BUFFER,
+                InputKind.CONSTANT_TENSOR,
+            ):
                 raise RuntimeError(
                     f"placeholder node for non-params, non-buffer, and non-tensor constants should not be tagged: {node} "
                 )
-            for ith, group in enumerate(_users_reading_a_copy(node)):
+            for ith, group in enumerate(_users_reading_a_copy(node, order)):
                 copy_constant_or_attribute_fqn = node.name + f"_copy_{ith}"
                 with exported_program.graph.inserting_before(
                     constant_or_attribute_node
diff --git a/exir/backend/test/test_utils.py b/exir/backend/test/test_utils.py
index c31fb7e8b2b..725edbc0364 100644
--- a/exir/backend/test/test_utils.py
+++ b/exir/backend/test/test_utils.py
@@ -6,6 +6,7 @@
 
 import itertools
 import unittest
+from unittest.mock import patch
 
 import torch
 from executorch import exir
@@ -21,6 +22,7 @@
 )
 from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
 from executorch.exir.backend.utils import (
+    _maybe_duplicate_constant_nodes,
     format_delegated_graph,
     get_delegated_payload,
     get_delegates,
@@ -487,3 +489,83 @@ def counts(program):
             )
             self.assertEqual(counts(lowered.exported_program()), [1, 1, 1])
             lowered.to_executorch()
+
+    def test_maybe_duplicate_constant_nodes_refuses_one_tag_string(self):
+        program = to_edge(
+            export(torch.nn.Linear(2, 2), (torch.ones(2),), strict=True)
+        ).exported_program()
+        # A str is an iterable of tags too, one per character.
+        with self.assertRaises(TypeError):
+            _maybe_duplicate_constant_nodes(program, "tag0")
+
+    def test_constants_shared_by_many_delegates_are_copied_in_one_pass(self):
+        class ShareConsts(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("a", torch.ones(2))
+                self.register_buffer("b", torch.ones(2))
+
+            def forward(self, x):
+                return x + self.a, x * self.a, self.a - self.b, x / self.b
+
+        program = to_edge(
+            export(ShareConsts(), (torch.ones(2),), strict=True)
+        ).exported_program()
+        ops = [node for node in program.graph.nodes if node.op == "call_function"]
+        for i, node in enumerate(ops):
+            node.meta["delegation_tag"] = f"tag{i}"
+        tag_constant_data(program)
+
+        # Recompiling and validating once per shared constant is quadratic in
+        # the graph size, which a large model with many delegates cannot afford.
+        with patch.object(
+            ExportedProgram,
+            "_validate",
+            autospec=True,
+            side_effect=ExportedProgram._validate,
+        ) as validate, patch.object(
+            torch.fx.GraphModule,
+            "recompile",
+            autospec=True,
+            side_effect=torch.fx.GraphModule.recompile,
+        ) as recompile:
+            _maybe_duplicate_constant_nodes(
+                program, [node.meta["delegation_tag"] for node in ops]
+            )
+        self.assertEqual(validate.call_count, 1)
+        self.assertEqual(recompile.call_count, 1)
+
+        # Copying b_a rewires sub, which must not change who keeps b_b: the
+        # first reader in graph order does.
+        self.assertEqual(
+            {
+                node.meta["delegation_tag"]: [arg.name for arg in node.args]
+                for node in ops
+            },
+            {
+                "tag0": ["x", "b_a"],
+                "tag1": ["x", "b_a_copy_0"],
+                "tag2": ["b_a_copy_1", "b_b"],
+                "tag3": ["x", "b_b_copy_0"],
+            },
+        )
+        for node in program.graph.nodes:
+            if node.op == "placeholder" and node.name != "x":
+                self.assertEqual(
+                    {user.meta["delegation_tag"] for user in node.users},
+                    {node.meta["delegation_tag"]},
+                )
+        self.assertEqual(
+            {
+                spec.arg.name: spec.target
+                for spec in program.graph_signature.input_specs
+            },
+            {
+                "b_a_copy_0": "a",
+                "b_a_copy_1": "a",
+                "b_a": "a",
+                "b_b_copy_0": "b",
+                "b_b": "b",
+                "x": None,
+            },
+        )
diff --git a/exir/backend/utils.py b/exir/backend/utils.py
index fc8d0ceeb70..45e0b64efff 100644
--- a/exir/backend/utils.py
+++ b/exir/backend/utils.py
@@ -16,7 +16,7 @@
 import torch
 from executorch.exir.backend.backend_details import ExportedProgram
 from executorch.exir.backend.canonical_partitioners.duplicate_constant_node_pass import (
-    duplicate_constant_node,
+    duplicate_constant_nodes,
 )
 from executorch.exir.common import setting_python_recursive_limit
 from executorch.exir.delegate import executorch_call_delegate
@@ -208,7 +208,7 @@ def _assign_new_tag(
 
 def _maybe_duplicate_constant_nodes(
     tagged_exported_program: ExportedProgram,
-    tag: str,
+    tags: Iterable[str],
 ) -> None:
     """
     If the constants node is shared by different tagged nodes, like
@@ -221,9 +221,13 @@ def _maybe_duplicate_constant_nodes(
 
     backend can estimate how much they want to duplicate the constant node, either error out or default to duplicate
     """
+    if isinstance(tags, str):
+        raise TypeError("tags is an iterable of tags, not one tag")
+    tags = set(tags)
     candidate_nodes = set()
     for node in tagged_exported_program.graph.nodes:
-        if node.meta.get("delegation_tag", "") == tag:
+        tag = node.meta.get("delegation_tag", None)
+        if tag in tags:
             if node.op == "placeholder":
                 for user in node.users:
                     users_tag = user.meta.get("delegation_tag", None)
@@ -235,12 +239,10 @@ def _maybe_duplicate_constant_nodes(
                             )
                         else:
                             candidate_nodes.add(node.name)
-    copied_nodes = set()
-    for candidate_node in candidate_nodes:
-        # Both tagged exported program and the owning program need to go through the same duplication pass
-        copied_nodes = copied_nodes.union(
-            duplicate_constant_node(tagged_exported_program, candidate_node)
-        )
+    if not candidate_nodes:
+        return
+    # Both tagged exported program and the owning program need to go through the same duplication pass
+    copied_nodes = duplicate_constant_nodes(tagged_exported_program, candidate_nodes)
     candidate_node_with_copies = candidate_nodes.union(copied_nodes)
     _assign_new_tag(tagged_exported_program, candidate_node_with_copies)