From 1af26c75335d6989ea74e59f3fbab135e7c4c916 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Mon, 6 Jul 2026 10:13:47 -0700
Subject: [PATCH 01/12] Insert write-back copy_ nodes at the earliest safe
 point

The insert_write_back_for_buffers pass placed every write-back copy_ at
the end of the graph, arbitrarily extending the lifetime of the value
being written back and wasting space in the memory plan.

Now each copy_(buffer, value) is inserted at the earliest point that
preserves the end-of-graph semantics: after the value is computed, after
every reader of the buffer or any alias of it (they must observe the old
contents), and after any mutation of the value or any alias of it (so we
snapshot the final value). Aliases are found with a forward walk using
schema alias_info, treating getitem, submodule calls, and schema-less
targets conservatively. If the value written back by one copy may alias
the buffer mutated by another, all copies fall back to the old
end-of-graph placement in their original order.

Fixes #7345
---
 .../insert_write_back_for_buffers_pass.py     | 156 +++++++++++++++++-
 exir/tests/test_passes.py                     |  73 +++++++-
 2 files changed, 223 insertions(+), 6 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 4dce40ae57c..39e6d2dbfab 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -4,7 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Dict, List, Optional, Tuple
+import operator
+from typing import Dict, List, Optional, Set, Tuple
 
 import torch
 from executorch.exir.operator.convert import is_inplace_variant
@@ -21,6 +22,112 @@
 from torchgen.model import SchemaKind
 
 
+def _may_alias_input(node: torch.fx.Node) -> bool:
+    """
+    Whether the value produced by this node may alias one of its inputs. When
+    we cannot tell (no schema, getitem, submodule calls, etc.) we
+    conservatively answer True.
+    """
+    if node.op != "call_function":
+        return True
+    if node.target is operator.getitem:
+        return True
+    schema = getattr(node.target, "_schema", None)
+    if schema is None:
+        return True
+    return any(ret.alias_info is not None for ret in schema.returns)
+
+
+def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool:
+    """
+    Whether this node may mutate the value passed to it as input_node. When we
+    cannot tell we conservatively answer True.
+    """
+    if node.op == "output":
+        return False
+    if node.op != "call_function":
+        return True
+    schema = getattr(node.target, "_schema", None)
+    if schema is None:
+        return True
+    for i, arg in enumerate(node.args):
+        if arg is input_node and i < len(schema.arguments):
+            alias_info = schema.arguments[i].alias_info
+            if alias_info is not None and alias_info.is_write:
+                return True
+    schema_kwargs = {a.name: a for a in schema.arguments}
+    for name, arg in node.kwargs.items():
+        if arg is input_node and name in schema_kwargs:
+            alias_info = schema_kwargs[name].alias_info
+            if alias_info is not None and alias_info.is_write:
+                return True
+    return False
+
+
+def _collect_aliases(
+    seed: torch.fx.Node, node_order: Dict[torch.fx.Node, int]
+) -> Set[torch.fx.Node]:
+    """
+    The set of nodes whose values may alias the value of seed, found by
+    walking forward through the graph.
+    """
+    aliases = {seed}
+    for node in node_order:
+        if node in aliases:
+            continue
+        if any(arg in aliases for arg in node.all_input_nodes) and _may_alias_input(
+            node
+        ):
+            aliases.add(node)
+    return aliases
+
+
+def _insertion_point(
+    mutated_node: torch.fx.Node,
+    return_node: torch.fx.Node,
+    node_order: Dict[torch.fx.Node, int],
+    last_placeholder: torch.fx.Node,
+) -> torch.fx.Node:
+    """
+    The earliest node after which it is safe to insert
+    copy_(mutated_node, return_node), preserving the semantics of inserting it
+    at the end of the graph. The copy_ must come after:
+
+     * return_node itself, and any node that may mutate it (or an alias of
+       it), so that we write back the final value;
+     * every reader of mutated_node or an alias of it, since they must observe
+       the old value of the buffer (this also orders us after anything that
+       may mutate the buffer);
+     * all placeholders.
+    """
+    latest = last_placeholder
+    if node_order[return_node] > node_order[latest]:
+        latest = return_node
+
+    for alias in _collect_aliases(mutated_node, node_order):
+        for user in alias.users:
+            # Users not in node_order are copy_ nodes inserted by us for other
+            # buffers; ordering with respect to them is handled by the
+            # independence check in _insert_copy.
+            if (
+                user.op != "output"
+                and user in node_order
+                and node_order[user] > node_order[latest]
+            ):
+                latest = user
+
+    for alias in _collect_aliases(return_node, node_order):
+        for user in alias.users:
+            if (
+                user in node_order
+                and _mutates_input(user, alias)
+                and node_order[user] > node_order[latest]
+            ):
+                latest = user
+
+    return latest
+
+
 def _insert_copy(
     gm: torch.fx.GraphModule,
     mutated_outputs: List[Optional[str]],
@@ -28,15 +135,26 @@ def _insert_copy(
 ):
     """
     Find the all the buffers and inputs that were mutated and insert copy_
-    operators to reflect mutations.
+    operators to reflect mutations. Each copy_ is inserted at the earliest
+    point at which it is safe, rather than at the end of the graph, so that
+    the memory planner does not have to arbitrarily extend the lifetime of the
+    value written back.
     """
     output_node = gm.graph.output_node()
     assert output_node is not None
     outputs = pytree.tree_flatten(output_node.args)[0]
     assert len(outputs) == len(mutated_outputs)
 
+    node_order: Dict[torch.fx.Node, int] = {
+        node: i for i, node in enumerate(gm.graph.nodes)
+    }
+    last_placeholder = [node for node in gm.graph.nodes if node.op == "placeholder"][
+        -1
+    ]
+
+    # Pair up the returns with the nodes they mutate.
+    copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = []
     user_output_nodes = []
-    buffer_output_nodes = []
     for return_node, mutated_node_name in zip(outputs, mutated_outputs):
         # User output, leave alone
         if mutated_node_name is None:
@@ -50,9 +168,37 @@ def _insert_copy(
             raise RuntimeError(
                 f"Could not find {mutated_node_name} in either buffer or input nodes"
             )
+        copies.append((mutated_node, return_node))
+
+    # The copies themselves mutate the buffers. If the value written back by
+    # one copy may alias the buffer mutated by another, then the order of the
+    # copies (and their position relative to everything else) matters in ways
+    # the insertion points below do not track, so fall back to inserting all
+    # of them at the end of the graph, in their original order, as before.
+    independent = True
+    if len(copies) > 1:
+        mutated_aliases: Set[torch.fx.Node] = set()
+        return_aliases: Set[torch.fx.Node] = set()
+        for i, (mutated_node, return_node) in enumerate(copies):
+            mutated_alias = _collect_aliases(mutated_node, node_order)
+            return_alias = _collect_aliases(return_node, node_order)
+            if mutated_alias & return_aliases or return_alias & mutated_aliases:
+                independent = False
+                break
+            mutated_aliases |= mutated_alias
+            return_aliases |= return_alias
 
-        # insert copy
-        with gm.graph.inserting_before(output_node):
+    # insert the copies
+    buffer_output_nodes = []
+    for mutated_node, return_node in copies:
+        if independent:
+            insert_after = _insertion_point(
+                mutated_node, return_node, node_order, last_placeholder
+            )
+            insertion = gm.graph.inserting_after(insert_after)
+        else:
+            insertion = gm.graph.inserting_before(output_node)
+        with insertion:
             buffer_output = gm.graph.call_function(
                 torch.ops.aten.copy_.default, (mutated_node, return_node)
             )
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 20906fe92e9..a93c8c28636 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -1935,14 +1935,85 @@ def forward(self, x):
         #     %b_direct_copy_from_input : [num_users=1] = placeholder[target=b_direct_copy_from_input]
         #     %_lifted_tensor_constant2 : [num_users=1] = placeholder[target=_lifted_tensor_constant2]
         #     %x : [num_users=2] = placeholder[target=x]
+        #     %copy__default_1 : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%b_direct_copy_from_input, %x), kwargs = {})
         #     %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%x, %b_state), kwargs = {})
         #     %dim_order_ops__to_dim_order_copy_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.dim_order_ops._to_dim_order_copy.default](args = (%_lifted_tensor_constant2,), kwargs = {dtype: torch.float32, dim_order: []})
         #     %aten_add_tensor_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%b_state, %dim_order_ops__to_dim_order_copy_default), kwargs = {})
         #     %copy__default : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%b_state, %aten_add_tensor_1), kwargs = {})
-        #     %copy__default_1 : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%b_direct_copy_from_input, %x), kwargs = {})
         #     return (copy__default, copy__default_1, aten_add_tensor)
         self.assertEqual(count_copies(gm), 2)
 
+    def test_mutable_buffers_write_back_is_inserted_early(self) -> None:
+        class EarlyMutationModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(1))
+
+            def forward(self, x):
+                # The buffer's new value is computed at the very start, so its
+                # write-back can happen immediately, letting the memory
+                # planner reuse the space during the rest of the graph.
+                self.state.add_(1)
+                y = x + 1
+                y = y + 1
+                y = y + 1
+                return y
+
+        model = to_edge(
+            export(EarlyMutationModule(), (torch.zeros(1),), strict=True)
+        )
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
+        copies = [
+            node
+            for node in gm.graph.nodes
+            if node.target == torch.ops.aten.copy_.default
+        ]
+        self.assertEqual(len(copies), 1)
+        copy = copies[0]
+        # The copy_ comes right after the value it writes back, not at the end
+        # of the graph: every user computation (the adds on x) is after it.
+        self.assertEqual(node_order[copy], node_order[copy.args[1]] + 1)
+        output_node = gm.graph.output_node()
+        user_return = output_node.args[0][1]
+        self.assertLess(node_order[copy], node_order[user_return])
+
+    def test_mutable_buffers_write_back_after_old_value_reads(self) -> None:
+        class ReadOldValueModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(1))
+
+            def forward(self, x):
+                # The buffer's new value is computed before the old value is
+                # read, so the write-back must not simply follow the value it
+                # writes: it must wait for the read of the old value.
+                new_state = x * 2
+                old_plus = self.state + x
+                self.state.copy_(new_state)
+                return old_plus
+
+        model = to_edge(
+            export(ReadOldValueModule(), (torch.zeros(1),), strict=True)
+        )
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
+        copies = [
+            node
+            for node in gm.graph.nodes
+            if node.target == torch.ops.aten.copy_.default
+        ]
+        self.assertEqual(len(copies), 1)
+        copy = copies[0]
+        buffer_placeholder = copy.args[0]
+        self.assertEqual(buffer_placeholder.op, "placeholder")
+        # Every read of the buffer's old value stays before the write-back.
+        for user in buffer_placeholder.users:
+            if user is not copy and user.op != "output":
+                self.assertLess(node_order[user], node_order[copy])
+
     def test_remove_quantized_op_noop_pass(self) -> None:
         class TestAddSliceNoop(torch.nn.Module):
             def __init__(self):

From fe8b19e51e90949d13bb213bd375ea8835962be7 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Fri, 10 Jul 2026 11:51:54 -0700
Subject: [PATCH 02/12] Fix lint

---
 exir/passes/insert_write_back_for_buffers_pass.py | 4 +---
 exir/tests/test_passes.py                         | 8 ++------
 2 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 39e6d2dbfab..d53dce2091f 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -148,9 +148,7 @@ def _insert_copy(
     node_order: Dict[torch.fx.Node, int] = {
         node: i for i, node in enumerate(gm.graph.nodes)
     }
-    last_placeholder = [node for node in gm.graph.nodes if node.op == "placeholder"][
-        -1
-    ]
+    last_placeholder = [node for node in gm.graph.nodes if node.op == "placeholder"][-1]
 
     # Pair up the returns with the nodes they mutate.
     copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = []
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index a93c8c28636..2c0586b244b 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -1959,9 +1959,7 @@ def forward(self, x):
                 y = y + 1
                 return y
 
-        model = to_edge(
-            export(EarlyMutationModule(), (torch.zeros(1),), strict=True)
-        )
+        model = to_edge(export(EarlyMutationModule(), (torch.zeros(1),), strict=True))
         gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
 
         node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
@@ -1994,9 +1992,7 @@ def forward(self, x):
                 self.state.copy_(new_state)
                 return old_plus
 
-        model = to_edge(
-            export(ReadOldValueModule(), (torch.zeros(1),), strict=True)
-        )
+        model = to_edge(export(ReadOldValueModule(), (torch.zeros(1),), strict=True))
         gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
 
         node_order = {node: i for i, node in enumerate(gm.graph.nodes)}

From ff2892d46f3f712a587e65a51e193fedb67a19da Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Sat, 11 Jul 2026 16:58:03 -0700
Subject: [PATCH 03/12] Fix lint (flake8 B007)

---
 exir/passes/insert_write_back_for_buffers_pass.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index d53dce2091f..296226bd5f0 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -177,7 +177,7 @@ def _insert_copy(
     if len(copies) > 1:
         mutated_aliases: Set[torch.fx.Node] = set()
         return_aliases: Set[torch.fx.Node] = set()
-        for i, (mutated_node, return_node) in enumerate(copies):
+        for mutated_node, return_node in copies:
             mutated_alias = _collect_aliases(mutated_node, node_order)
             return_alias = _collect_aliases(return_node, node_order)
             if mutated_alias & return_aliases or return_alias & mutated_aliases:

From b6457f4818a9b77b12afc980fe13736dd6579bba Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Wed, 16 Sep 2026 15:26:43 -0700
Subject: [PATCH 04/12] Handle graphs with no placeholders; address
 alias-analysis review comments

---
 .../insert_write_back_for_buffers_pass.py     | 33 +++++++++++++++----
 exir/tests/test_passes.py                     | 12 +++++++
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 296226bd5f0..d398398327a 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -28,6 +28,9 @@ def _may_alias_input(node: torch.fx.Node) -> bool:
     we cannot tell (no schema, getitem, submodule calls, etc.) we
     conservatively answer True.
     """
+    if node.op == "output":
+        # The output node produces no value of its own, so it cannot alias.
+        return False
     if node.op != "call_function":
         return True
     if node.target is operator.getitem:
@@ -38,6 +41,18 @@ def _may_alias_input(node: torch.fx.Node) -> bool:
     return any(ret.alias_info is not None for ret in schema.returns)
 
 
+def _contains_node(value: object, input_node: torch.fx.Node) -> bool:
+    """
+    Whether input_node appears in value, looking through lists and tuples so
+    that container arguments (e.g. foreach-style ops) are handled.
+    """
+    if value is input_node:
+        return True
+    if isinstance(value, (list, tuple)):
+        return any(_contains_node(v, input_node) for v in value)
+    return False
+
+
 def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool:
     """
     Whether this node may mutate the value passed to it as input_node. When we
@@ -51,13 +66,13 @@ def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool:
     if schema is None:
         return True
     for i, arg in enumerate(node.args):
-        if arg is input_node and i < len(schema.arguments):
+        if _contains_node(arg, input_node) and i < len(schema.arguments):
             alias_info = schema.arguments[i].alias_info
             if alias_info is not None and alias_info.is_write:
                 return True
     schema_kwargs = {a.name: a for a in schema.arguments}
     for name, arg in node.kwargs.items():
-        if arg is input_node and name in schema_kwargs:
+        if _contains_node(arg, input_node) and name in schema_kwargs:
             alias_info = schema_kwargs[name].alias_info
             if alias_info is not None and alias_info.is_write:
                 return True
@@ -86,7 +101,7 @@ def _insertion_point(
     mutated_node: torch.fx.Node,
     return_node: torch.fx.Node,
     node_order: Dict[torch.fx.Node, int],
-    last_placeholder: torch.fx.Node,
+    last_placeholder: Optional[torch.fx.Node],
 ) -> torch.fx.Node:
     """
     The earliest node after which it is safe to insert
@@ -100,9 +115,12 @@ def _insertion_point(
        may mutate the buffer);
      * all placeholders.
     """
-    latest = last_placeholder
-    if node_order[return_node] > node_order[latest]:
-        latest = return_node
+    latest = return_node
+    if (
+        last_placeholder is not None
+        and node_order[last_placeholder] > node_order[latest]
+    ):
+        latest = last_placeholder
 
     for alias in _collect_aliases(mutated_node, node_order):
         for user in alias.users:
@@ -148,7 +166,8 @@ def _insert_copy(
     node_order: Dict[torch.fx.Node, int] = {
         node: i for i, node in enumerate(gm.graph.nodes)
     }
-    last_placeholder = [node for node in gm.graph.nodes if node.op == "placeholder"][-1]
+    placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
+    last_placeholder = placeholders[-1] if placeholders else None
 
     # Pair up the returns with the nodes they mutate.
     copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = []
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 2c0586b244b..4109f33473b 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -2010,6 +2010,18 @@ def forward(self, x):
             if user is not copy and user.op != "output":
                 self.assertLess(node_order[user], node_order[copy])
 
+    def test_mutable_buffers_write_back_no_inputs(self) -> None:
+        class NoInputModule(torch.nn.Module):
+            def forward(self):
+                return torch.ones(3) * 2
+
+        model = to_edge(export(NoInputModule(), (), strict=True))
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        # A graph with no placeholders has nothing to write back; the pass
+        # must complete cleanly rather than assume an input exists.
+        self.assertEqual(sum(node.op == "placeholder" for node in gm.graph.nodes), 0)
+
     def test_remove_quantized_op_noop_pass(self) -> None:
         class TestAddSliceNoop(torch.nn.Module):
             def __init__(self):

From a7d562ca4707f060dd3f68385176c9aefda79bda Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Wed, 16 Sep 2026 15:41:55 -0700
Subject: [PATCH 05/12] Treat view_copy as aliasing:
 ReplaceViewCopyWithViewPass later makes it a true alias

---
 .../insert_write_back_for_buffers_pass.py     |  6 ++++
 exir/tests/test_passes.py                     | 35 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index d398398327a..1832395f7e1 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -9,6 +9,7 @@
 
 import torch
 from executorch.exir.operator.convert import is_inplace_variant
+from executorch.exir.passes.replace_view_copy_with_view_pass import _is_view_copy
 
 from torch.export.exported_program import (
     ExportedProgram,
@@ -35,6 +36,11 @@ def _may_alias_input(node: torch.fx.Node) -> bool:
         return True
     if node.target is operator.getitem:
         return True
+    if _is_view_copy(node):
+        # view_copy produces a fresh tensor here, but ReplaceViewCopyWithViewPass
+        # later rewrites non-output view_copy nodes into true aliases, so readers
+        # through a view must stay ordered before any write-back into its base.
+        return True
     schema = getattr(node.target, "_schema", None)
     if schema is None:
         return True
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 4109f33473b..1226c4afa31 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -2010,6 +2010,41 @@ def forward(self, x):
             if user is not copy and user.op != "output":
                 self.assertLess(node_order[user], node_order[copy])
 
+    def test_mutable_buffers_write_back_after_view_reads(self) -> None:
+        class ViewReadModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(4))
+
+            def forward(self, x):
+                # The buffer's old value is read through a view. view_copy is
+                # a copy here, but ReplaceViewCopyWithViewPass later turns it
+                # into a true alias, so the write-back must still come after
+                # the read.
+                new_state = x * 2
+                old_through_view = self.state.view(2, 2).sum()
+                self.state.copy_(new_state)
+                return old_through_view
+
+        model = to_edge(export(ViewReadModule(), (torch.zeros(4),), strict=True))
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
+        copies = [
+            node
+            for node in gm.graph.nodes
+            if node.target == torch.ops.aten.copy_.default
+        ]
+        self.assertEqual(len(copies), 1)
+        copy = copies[0]
+        views = [node for node in gm.graph.nodes if "view_copy" in str(node.target)]
+        self.assertTrue(len(views) >= 1)
+        # Every reader through the view stays before the write-back.
+        for view in views:
+            for user in view.users:
+                if user is not copy and user.op != "output":
+                    self.assertLess(node_order[user], node_order[copy])
+
     def test_mutable_buffers_write_back_no_inputs(self) -> None:
         class NoInputModule(torch.nn.Module):
             def forward(self):

From a4d947a0c572bef8c596fac9ce4d48b3a4a80db2 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Wed, 16 Sep 2026 15:58:32 -0700
Subject: [PATCH 06/12] Add replace_view_copy_with_view_pass dep to Buck target
 for _is_view_copy import

---
 exir/passes/BUCK | 1 +
 1 file changed, 1 insertion(+)

diff --git a/exir/passes/BUCK b/exir/passes/BUCK
index a63ce43dbf6..d0a0ec4427d 100644
--- a/exir/passes/BUCK
+++ b/exir/passes/BUCK
@@ -90,6 +90,7 @@ fbcode_target(_kind = runtime.python_library,
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/exir/passes:replace_view_copy_with_view_pass",
     ],
 )
 

From f73e16fb62300929bf6e36a95b7cabf0b751e4e6 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Thu, 17 Sep 2026 15:16:44 -0700
Subject: [PATCH 07/12] Make alias closure bidirectional; cover the
 non-independent fallback with a test

---
 .../insert_write_back_for_buffers_pass.py     | 32 ++++++++++----
 exir/tests/test_passes.py                     | 42 +++++++++++++++++++
 2 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 1832395f7e1..2913b217351 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -89,17 +89,31 @@ def _collect_aliases(
     seed: torch.fx.Node, node_order: Dict[torch.fx.Node, int]
 ) -> Set[torch.fx.Node]:
     """
-    The set of nodes whose values may alias the value of seed, found by
-    walking forward through the graph.
+    The set of nodes whose values may alias the value of seed. The closure is
+    taken in both directions: a node that may alias its inputs pulls its
+    result into the set of its inputs (forward), and pulls its inputs into
+    the set of its result (backward). The backward direction matters for
+    views: the base of a view must count as an alias of the view's value,
+    since a mutation of the base is a mutation of the view once
+    ReplaceViewCopyWithViewPass has run.
     """
     aliases = {seed}
-    for node in node_order:
-        if node in aliases:
-            continue
-        if any(arg in aliases for arg in node.all_input_nodes) and _may_alias_input(
-            node
-        ):
-            aliases.add(node)
+    changed = True
+    while changed:
+        changed = False
+        for node in node_order:
+            if node in aliases:
+                if _may_alias_input(node):
+                    for arg in node.all_input_nodes:
+                        if arg not in aliases:
+                            aliases.add(arg)
+                            changed = True
+                continue
+            if any(arg in aliases for arg in node.all_input_nodes) and _may_alias_input(
+                node
+            ):
+                aliases.add(node)
+                changed = True
     return aliases
 
 
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 6d7a7d4e292..6ffc5f86ed0 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -2123,6 +2123,48 @@ def forward(self, x):
                 if user is not copy and user.op != "output":
                     self.assertLess(node_order[user], node_order[copy])
 
+    def test_mutable_buffers_write_back_aliased_fallback(self) -> None:
+        class AliasedWriteBacksModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("a", torch.zeros(4))
+                self.register_buffer("b", torch.zeros(2, 2))
+
+            def forward(self, x):
+                # b's new value is a view of a, and a is mutated as well, so
+                # the two write-backs are not independent: after the view
+                # rewrite the value written into b aliases the buffer the
+                # other copy_ mutates. The pass must fall back to inserting
+                # both copies at the end of the graph in their original
+                # order.
+                self.b.copy_(self.a.view(2, 2))
+                self.a.add_(x)
+                return x + 1
+
+        model = to_edge(
+            export(AliasedWriteBacksModule(), (torch.zeros(4),), strict=True)
+        )
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
+        copies = [
+            node
+            for node in gm.graph.nodes
+            if node.target == torch.ops.aten.copy_.default
+        ]
+        self.assertEqual(len(copies), 2)
+        # Both copies sit at the end of the graph, after all other compute.
+        last_compute = max(
+            node_order[node]
+            for node in gm.graph.nodes
+            if node.op == "call_function" and node not in copies
+        )
+        for copy in copies:
+            self.assertGreater(node_order[copy], last_compute)
+        # The copies keep their original (output-spec) order.
+        output_args = gm.graph.output_node().args[0]
+        self.assertEqual([output_args[0], output_args[1]], copies)
+
     def test_mutable_buffers_write_back_no_inputs(self) -> None:
         class NoInputModule(torch.nn.Module):
             def forward(self):

From 130379643dcc53745cb4e2c5efc1f4462b7f4dff Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Thu, 17 Sep 2026 15:28:54 -0700
Subject: [PATCH 08/12] Alias analysis: schema-precise alias edges, adjacency
 built once, cached BFS closures

---
 .../insert_write_back_for_buffers_pass.py     | 136 ++++++++++++------
 1 file changed, 92 insertions(+), 44 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 2913b217351..36d3bc889ac 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -23,28 +23,71 @@
 from torchgen.model import SchemaKind
 
 
-def _may_alias_input(node: torch.fx.Node) -> bool:
+def _fx_nodes_in(value: object) -> List[torch.fx.Node]:
+    """The FX nodes contained in value, looking through lists and tuples."""
+    if isinstance(value, torch.fx.Node):
+        return [value]
+    if isinstance(value, (list, tuple)):
+        return [n for v in value for n in _fx_nodes_in(v)]
+    return []
+
+
+def _aliasing_inputs(node: torch.fx.Node) -> List[torch.fx.Node]:
     """
-    Whether the value produced by this node may alias one of its inputs. When
-    we cannot tell (no schema, getitem, submodule calls, etc.) we
-    conservatively answer True.
+    The subset of node's FX inputs that the value produced by node may alias.
+    When we cannot tell (no schema, getitem, submodule calls, etc.) we
+    conservatively answer all inputs; for schema-annotated ops we answer only
+    the inputs whose alias set is shared with a return, so that e.g. the
+    shape-supplying argument of expand_as does not count as an alias.
     """
     if node.op == "output":
         # The output node produces no value of its own, so it cannot alias.
-        return False
+        return []
     if node.op != "call_function":
-        return True
+        return list(node.all_input_nodes)
     if node.target is operator.getitem:
-        return True
+        return list(node.all_input_nodes)
     if _is_view_copy(node):
         # view_copy produces a fresh tensor here, but ReplaceViewCopyWithViewPass
-        # later rewrites non-output view_copy nodes into true aliases, so readers
-        # through a view must stay ordered before any write-back into its base.
-        return True
+        # later rewrites non-output view_copy nodes into true aliases of their
+        # base, the first argument.
+        return _fx_nodes_in(node.args[0] if node.args else None)
     schema = getattr(node.target, "_schema", None)
     if schema is None:
-        return True
-    return any(ret.alias_info is not None for ret in schema.returns)
+        return list(node.all_input_nodes)
+    ret_sets: Set[str] = set()
+    for ret in schema.returns:
+        if ret.alias_info is not None:
+            ret_sets |= set(ret.alias_info.before_set)
+            ret_sets |= set(ret.alias_info.after_set)
+    if not ret_sets:
+        return []
+    if "*" in ret_sets:
+        # A wildcard return may alias any input.
+        return list(node.all_input_nodes)
+    aliasing: List[torch.fx.Node] = []
+    for i, arg in enumerate(node.args):
+        if i < len(schema.arguments):
+            alias_info = schema.arguments[i].alias_info
+            arg_sets = (
+                set(alias_info.before_set) | set(alias_info.after_set)
+                if alias_info is not None
+                else set()
+            )
+            if arg_sets & ret_sets or "*" in arg_sets:
+                aliasing.extend(_fx_nodes_in(arg))
+    schema_kwargs = {a.name: a for a in schema.arguments}
+    for name, arg in node.kwargs.items():
+        if name in schema_kwargs:
+            alias_info = schema_kwargs[name].alias_info
+            arg_sets = (
+                set(alias_info.before_set) | set(alias_info.after_set)
+                if alias_info is not None
+                else set()
+            )
+            if arg_sets & ret_sets or "*" in arg_sets:
+                aliasing.extend(_fx_nodes_in(arg))
+    return aliasing
 
 
 def _contains_node(value: object, input_node: torch.fx.Node) -> bool:
@@ -85,36 +128,39 @@ def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool:
     return False
 
 
-def _collect_aliases(
-    seed: torch.fx.Node, node_order: Dict[torch.fx.Node, int]
-) -> Set[torch.fx.Node]:
+class _AliasIndex:
     """
-    The set of nodes whose values may alias the value of seed. The closure is
-    taken in both directions: a node that may alias its inputs pulls its
-    result into the set of its inputs (forward), and pulls its inputs into
-    the set of its result (backward). The backward direction matters for
-    views: the base of a view must count as an alias of the view's value,
-    since a mutation of the base is a mutation of the view once
+    Alias closures over the graph. The undirected adjacency (each node joined
+    to the inputs its value may alias) is built once, and each closure is a
+    breadth-first walk cached per seed. The closure is symmetric on purpose:
+    the base of a view must count as an alias of the view's value, since a
+    mutation of the base is a mutation of the view once
     ReplaceViewCopyWithViewPass has run.
     """
-    aliases = {seed}
-    changed = True
-    while changed:
-        changed = False
-        for node in node_order:
-            if node in aliases:
-                if _may_alias_input(node):
-                    for arg in node.all_input_nodes:
-                        if arg not in aliases:
-                            aliases.add(arg)
-                            changed = True
-                continue
-            if any(arg in aliases for arg in node.all_input_nodes) and _may_alias_input(
-                node
-            ):
-                aliases.add(node)
-                changed = True
-    return aliases
+
+    def __init__(self, nodes: List[torch.fx.Node]) -> None:
+        self._adjacency: Dict[torch.fx.Node, List[torch.fx.Node]] = {}
+        for node in nodes:
+            for arg in _aliasing_inputs(node):
+                self._adjacency.setdefault(node, []).append(arg)
+                self._adjacency.setdefault(arg, []).append(node)
+        self._cache: Dict[torch.fx.Node, Set[torch.fx.Node]] = {}
+
+    def aliases(self, seed: torch.fx.Node) -> Set[torch.fx.Node]:
+        cached = self._cache.get(seed)
+        if cached is not None:
+            return cached
+        seen = {seed}
+        frontier = [seed]
+        while frontier:
+            node = frontier.pop()
+            for other in self._adjacency.get(node, ()):
+                if other not in seen:
+                    seen.add(other)
+                    frontier.append(other)
+        for node in seen:
+            self._cache[node] = seen
+        return seen
 
 
 def _insertion_point(
@@ -122,6 +168,7 @@ def _insertion_point(
     return_node: torch.fx.Node,
     node_order: Dict[torch.fx.Node, int],
     last_placeholder: Optional[torch.fx.Node],
+    alias_index: _AliasIndex,
 ) -> torch.fx.Node:
     """
     The earliest node after which it is safe to insert
@@ -142,7 +189,7 @@ def _insertion_point(
     ):
         latest = last_placeholder
 
-    for alias in _collect_aliases(mutated_node, node_order):
+    for alias in alias_index.aliases(mutated_node):
         for user in alias.users:
             # Users not in node_order are copy_ nodes inserted by us for other
             # buffers; ordering with respect to them is handled by the
@@ -154,7 +201,7 @@ def _insertion_point(
             ):
                 latest = user
 
-    for alias in _collect_aliases(return_node, node_order):
+    for alias in alias_index.aliases(return_node):
         for user in alias.users:
             if (
                 user in node_order
@@ -188,6 +235,7 @@ def _insert_copy(
     }
     placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
     last_placeholder = placeholders[-1] if placeholders else None
+    alias_index = _AliasIndex(list(gm.graph.nodes))
 
     # Pair up the returns with the nodes they mutate.
     copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = []
@@ -217,8 +265,8 @@ def _insert_copy(
         mutated_aliases: Set[torch.fx.Node] = set()
         return_aliases: Set[torch.fx.Node] = set()
         for mutated_node, return_node in copies:
-            mutated_alias = _collect_aliases(mutated_node, node_order)
-            return_alias = _collect_aliases(return_node, node_order)
+            mutated_alias = alias_index.aliases(mutated_node)
+            return_alias = alias_index.aliases(return_node)
             if mutated_alias & return_aliases or return_alias & mutated_aliases:
                 independent = False
                 break
@@ -230,7 +278,7 @@ def _insert_copy(
     for mutated_node, return_node in copies:
         if independent:
             insert_after = _insertion_point(
-                mutated_node, return_node, node_order, last_placeholder
+                mutated_node, return_node, node_order, last_placeholder, alias_index
             )
             insertion = gm.graph.inserting_after(insert_after)
         else:

From 63639060c83fd1a75ce2c515d45c7c8732256099 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Thu, 17 Sep 2026 16:18:55 -0700
Subject: [PATCH 09/12] Lint: split schema-less dispatch and pairing helpers
 out of _aliasing_inputs (C901); rename shadowing loop var (F402)

---
 .../insert_write_back_for_buffers_pass.py     | 79 +++++++++++--------
 exir/tests/test_passes.py                     |  4 +-
 2 files changed, 49 insertions(+), 34 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 36d3bc889ac..75d3e8be868 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -32,13 +32,33 @@ def _fx_nodes_in(value: object) -> List[torch.fx.Node]:
     return []
 
 
-def _aliasing_inputs(node: torch.fx.Node) -> List[torch.fx.Node]:
+def _alias_sets(alias_info: object) -> Set[str]:
+    """The alias-set annotations of a schema argument or return."""
+    if alias_info is None:
+        return set()
+    return set(alias_info.before_set) | set(alias_info.after_set)  # pyre-ignore[16]
+
+
+def _schema_paired_args(node: torch.fx.Node, schema) -> List[Tuple[object, object]]:
+    """Each FX arg/kwarg of node paired with its schema argument."""
+    schema_kwargs = {a.name: a for a in schema.arguments}
+    return [
+        (arg, schema.arguments[i])
+        for i, arg in enumerate(node.args)
+        if i < len(schema.arguments)
+    ] + [
+        (arg, schema_kwargs[name])
+        for name, arg in node.kwargs.items()
+        if name in schema_kwargs
+    ]
+
+
+def _schemaless_aliasing_inputs(
+    node: torch.fx.Node,
+) -> Optional[List[torch.fx.Node]]:
     """
-    The subset of node's FX inputs that the value produced by node may alias.
-    When we cannot tell (no schema, getitem, submodule calls, etc.) we
-    conservatively answer all inputs; for schema-annotated ops we answer only
-    the inputs whose alias set is shared with a return, so that e.g. the
-    shape-supplying argument of expand_as does not count as an alias.
+    Aliasing inputs for the nodes that cannot be answered from a schema, or
+    None when the node has a schema to consult.
     """
     if node.op == "output":
         # The output node produces no value of its own, so it cannot alias.
@@ -52,41 +72,36 @@ def _aliasing_inputs(node: torch.fx.Node) -> List[torch.fx.Node]:
         # later rewrites non-output view_copy nodes into true aliases of their
         # base, the first argument.
         return _fx_nodes_in(node.args[0] if node.args else None)
-    schema = getattr(node.target, "_schema", None)
-    if schema is None:
+    if getattr(node.target, "_schema", None) is None:
         return list(node.all_input_nodes)
+    return None
+
+
+def _aliasing_inputs(node: torch.fx.Node) -> List[torch.fx.Node]:
+    """
+    The subset of node's FX inputs that the value produced by node may alias.
+    When we cannot tell (no schema, getitem, submodule calls, etc.) we
+    conservatively answer all inputs; for schema-annotated ops we answer only
+    the inputs whose alias set is shared with a return, so that e.g. the
+    shape-supplying argument of expand_as does not count as an alias.
+    """
+    special = _schemaless_aliasing_inputs(node)
+    if special is not None:
+        return special
+    schema = node.target._schema  # pyre-ignore[16]
     ret_sets: Set[str] = set()
     for ret in schema.returns:
-        if ret.alias_info is not None:
-            ret_sets |= set(ret.alias_info.before_set)
-            ret_sets |= set(ret.alias_info.after_set)
+        ret_sets |= _alias_sets(ret.alias_info)
     if not ret_sets:
         return []
     if "*" in ret_sets:
         # A wildcard return may alias any input.
         return list(node.all_input_nodes)
     aliasing: List[torch.fx.Node] = []
-    for i, arg in enumerate(node.args):
-        if i < len(schema.arguments):
-            alias_info = schema.arguments[i].alias_info
-            arg_sets = (
-                set(alias_info.before_set) | set(alias_info.after_set)
-                if alias_info is not None
-                else set()
-            )
-            if arg_sets & ret_sets or "*" in arg_sets:
-                aliasing.extend(_fx_nodes_in(arg))
-    schema_kwargs = {a.name: a for a in schema.arguments}
-    for name, arg in node.kwargs.items():
-        if name in schema_kwargs:
-            alias_info = schema_kwargs[name].alias_info
-            arg_sets = (
-                set(alias_info.before_set) | set(alias_info.after_set)
-                if alias_info is not None
-                else set()
-            )
-            if arg_sets & ret_sets or "*" in arg_sets:
-                aliasing.extend(_fx_nodes_in(arg))
+    for arg, schema_arg in _schema_paired_args(node, schema):
+        arg_sets = _alias_sets(schema_arg.alias_info)
+        if arg_sets & ret_sets or "*" in arg_sets:
+            aliasing.extend(_fx_nodes_in(arg))
     return aliasing
 
 
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
index 6ffc5f86ed0..9cdd83b18f2 100644
--- a/exir/tests/test_passes.py
+++ b/exir/tests/test_passes.py
@@ -2159,8 +2159,8 @@ def forward(self, x):
             for node in gm.graph.nodes
             if node.op == "call_function" and node not in copies
         )
-        for copy in copies:
-            self.assertGreater(node_order[copy], last_compute)
+        for copy_node in copies:
+            self.assertGreater(node_order[copy_node], last_compute)
         # The copies keep their original (output-spec) order.
         output_args = gm.graph.output_node().args[0]
         self.assertEqual([output_args[0], output_args[1]], copies)

From 82d0a1fd139b55e196be8ecdc6ccb68513089877 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Thu, 17 Sep 2026 16:28:21 -0700
Subject: [PATCH 10/12] Type alias-info helpers for mypy; skip alias analysis
 when there are no write-backs

---
 .../insert_write_back_for_buffers_pass.py     | 93 ++++++++++---------
 1 file changed, 50 insertions(+), 43 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 75d3e8be868..e92358f2096 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -32,14 +32,16 @@ def _fx_nodes_in(value: object) -> List[torch.fx.Node]:
     return []
 
 
-def _alias_sets(alias_info: object) -> Set[str]:
+def _alias_sets(alias_info: Optional[torch._C._AliasInfo]) -> Set[str]:
     """The alias-set annotations of a schema argument or return."""
     if alias_info is None:
         return set()
-    return set(alias_info.before_set) | set(alias_info.after_set)  # pyre-ignore[16]
+    return set(alias_info.before_set) | set(alias_info.after_set)
 
 
-def _schema_paired_args(node: torch.fx.Node, schema) -> List[Tuple[object, object]]:
+def _schema_paired_args(
+    node: torch.fx.Node, schema: torch.FunctionSchema
+) -> List[Tuple[object, torch.Argument]]:
     """Each FX arg/kwarg of node paired with its schema argument."""
     schema_kwargs = {a.name: a for a in schema.arguments}
     return [
@@ -245,13 +247,6 @@ def _insert_copy(
     outputs = pytree.tree_flatten(output_node.args)[0]
     assert len(outputs) == len(mutated_outputs)
 
-    node_order: Dict[torch.fx.Node, int] = {
-        node: i for i, node in enumerate(gm.graph.nodes)
-    }
-    placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
-    last_placeholder = placeholders[-1] if placeholders else None
-    alias_index = _AliasIndex(list(gm.graph.nodes))
-
     # Pair up the returns with the nodes they mutate.
     copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = []
     user_output_nodes = []
@@ -270,40 +265,52 @@ def _insert_copy(
             )
         copies.append((mutated_node, return_node))
 
-    # The copies themselves mutate the buffers. If the value written back by
-    # one copy may alias the buffer mutated by another, then the order of the
-    # copies (and their position relative to everything else) matters in ways
-    # the insertion points below do not track, so fall back to inserting all
-    # of them at the end of the graph, in their original order, as before.
-    independent = True
-    if len(copies) > 1:
-        mutated_aliases: Set[torch.fx.Node] = set()
-        return_aliases: Set[torch.fx.Node] = set()
-        for mutated_node, return_node in copies:
-            mutated_alias = alias_index.aliases(mutated_node)
-            return_alias = alias_index.aliases(return_node)
-            if mutated_alias & return_aliases or return_alias & mutated_aliases:
-                independent = False
-                break
-            mutated_aliases |= mutated_alias
-            return_aliases |= return_alias
-
     # insert the copies
-    buffer_output_nodes = []
-    for mutated_node, return_node in copies:
-        if independent:
-            insert_after = _insertion_point(
-                mutated_node, return_node, node_order, last_placeholder, alias_index
-            )
-            insertion = gm.graph.inserting_after(insert_after)
-        else:
-            insertion = gm.graph.inserting_before(output_node)
-        with insertion:
-            buffer_output = gm.graph.call_function(
-                torch.ops.aten.copy_.default, (mutated_node, return_node)
-            )
-            # add output of copy to graph outputs
-            buffer_output_nodes.append(buffer_output)
+    buffer_output_nodes: List[torch.fx.Node] = []
+    # The alias analysis is only needed to place copies, so graphs with no
+    # write-backs (the common case for models without mutable state) skip its
+    # cost entirely.
+    if copies:
+        node_order: Dict[torch.fx.Node, int] = {
+            node: i for i, node in enumerate(gm.graph.nodes)
+        }
+        placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"]
+        last_placeholder = placeholders[-1] if placeholders else None
+        alias_index = _AliasIndex(list(gm.graph.nodes))
+
+        # The copies themselves mutate the buffers. If the value written back
+        # by one copy may alias the buffer mutated by another, then the order
+        # of the copies (and their position relative to everything else)
+        # matters in ways the insertion points below do not track, so fall
+        # back to inserting all of them at the end of the graph, in their
+        # original order, as before.
+        independent = True
+        if len(copies) > 1:
+            mutated_aliases: Set[torch.fx.Node] = set()
+            return_aliases: Set[torch.fx.Node] = set()
+            for mutated_node, return_node in copies:
+                mutated_alias = alias_index.aliases(mutated_node)
+                return_alias = alias_index.aliases(return_node)
+                if mutated_alias & return_aliases or return_alias & mutated_aliases:
+                    independent = False
+                    break
+                mutated_aliases |= mutated_alias
+                return_aliases |= return_alias
+
+        for mutated_node, return_node in copies:
+            if independent:
+                insert_after = _insertion_point(
+                    mutated_node, return_node, node_order, last_placeholder, alias_index
+                )
+                insertion = gm.graph.inserting_after(insert_after)
+            else:
+                insertion = gm.graph.inserting_before(output_node)
+            with insertion:
+                buffer_output = gm.graph.call_function(
+                    torch.ops.aten.copy_.default, (mutated_node, return_node)
+                )
+                # add output of copy to graph outputs
+                buffer_output_nodes.append(buffer_output)
 
     with gm.graph.inserting_before(output_node):
         buffer_output_nodes.extend(user_output_nodes)

From f4057aeeb69112c3a948e84ccdb58960aa8ca876 Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Thu, 17 Sep 2026 16:40:49 -0700
Subject: [PATCH 11/12] Trust only aten:: schemas for alias/mutation
 introspection, matching cse_pass policy

---
 exir/passes/insert_write_back_for_buffers_pass.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index e92358f2096..0988187a178 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -32,6 +32,16 @@ def _fx_nodes_in(value: object) -> List[torch.fx.Node]:
     return []
 
 
+def _schema_is_trusted(schema: torch.FunctionSchema) -> bool:
+    """
+    Only aten:: schemas are trusted for alias/mutation introspection. Custom
+    op schemas (mlx::, torchao::, etc.) may not accurately annotate mutation
+    or aliasing (the same policy cse_pass.py applies), so they are treated as
+    unknown.
+    """
+    return schema.name.startswith("aten::")
+
+
 def _alias_sets(alias_info: Optional[torch._C._AliasInfo]) -> Set[str]:
     """The alias-set annotations of a schema argument or return."""
     if alias_info is None:
@@ -74,7 +84,8 @@ def _schemaless_aliasing_inputs(
         # later rewrites non-output view_copy nodes into true aliases of their
         # base, the first argument.
         return _fx_nodes_in(node.args[0] if node.args else None)
-    if getattr(node.target, "_schema", None) is None:
+    schema = getattr(node.target, "_schema", None)
+    if schema is None or not _schema_is_trusted(schema):
         return list(node.all_input_nodes)
     return None
 
@@ -129,7 +140,7 @@ def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool:
     if node.op != "call_function":
         return True
     schema = getattr(node.target, "_schema", None)
-    if schema is None:
+    if schema is None or not _schema_is_trusted(schema):
         return True
     for i, arg in enumerate(node.args):
         if _contains_node(arg, input_node) and i < len(schema.arguments):

From 136a0638101b2f28b125504ff89f727ff5561a4e Mon Sep 17 00:00:00 2001
From: Jacky Li <86073892+JPL11@users.noreply.github.com>
Date: Thu, 17 Sep 2026 16:50:23 -0700
Subject: [PATCH 12/12] Independence check: destinations that may alias each
 other also force the end-of-graph fallback

---
 exir/passes/insert_write_back_for_buffers_pass.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py
index 0988187a178..304ef95a600 100644
--- a/exir/passes/insert_write_back_for_buffers_pass.py
+++ b/exir/passes/insert_write_back_for_buffers_pass.py
@@ -302,7 +302,13 @@ def _insert_copy(
             for mutated_node, return_node in copies:
                 mutated_alias = alias_index.aliases(mutated_node)
                 return_alias = alias_index.aliases(return_node)
-                if mutated_alias & return_aliases or return_alias & mutated_aliases:
+                if (
+                    mutated_alias & return_aliases
+                    or return_alias & mutated_aliases
+                    # Two destinations that may share storage must also keep
+                    # their original write order.
+                    or mutated_alias & mutated_aliases
+                ):
                     independent = False
                     break
                 mutated_aliases |= mutated_alias