Arm backend: Split shared get_attr nodes before dedupe (#20420)

DeduplicateGetAttrPass gives duplicate get_attr nodes different backing
attrs before PT2E folding. It missed the case where one get_attr node
has multiple users, for example:

    w = get_attr("weight")
    conv1 = conv2d(x, w)
    conv2 = conv2d(y, w)

Some passes preceding DeduplicateGetAttrPass hid this issue because they
rebuilt the graph through ExportPass. That rebuild could split the
shared get_attr node into one node per use, even when the pass did not
change any operators.

Handle that case in DeduplicateGetAttrPass. Keep the original get_attr
node for the first user, then create new get_attr nodes for the other
users with the same target and copied metadata.

The existing dedupe logic can then give those get_attr nodes different
backing attrs. This keeps shared parameter handling correct without
relying on graph rebuild side effects.

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>

Author: Sebastian-Larsson

backends/arm/_passes/deduplicate_get_attr_pass.py

@@ -9,6 +9,7 @@
 from executorch.backends.arm._passes import ArmPass
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule, Node
+from torch.fx.node import map_arg
 from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix
 
 
@@ -24,6 +25,13 @@ class DeduplicateGetAttrPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
+    def _replace_input_node(self, node: Node, old_node: Node, new_node: Node) -> None:
+        def maybe_replace_node(arg: Any) -> Any:
+            return new_node if arg is old_node else arg
+
+        node.args = map_arg(node.args, maybe_replace_node)
+        node.kwargs = map_arg(node.kwargs, maybe_replace_node)
+
     def _get_attr(self, graph_module: GraphModule, target: str) -> Any:
         attr: Any = graph_module
         for target_atom in target.split("."):
@@ -51,9 +59,26 @@ def _copy_attr(self, graph_module: GraphModule, node: Node) -> str:
 
         return attr_name
 
+    def _split_shared_get_attrs(self, graph_module: GraphModule) -> bool:
+        modified = False
+
+        for node in list(graph_module.graph.find_nodes(op="get_attr")):
+            users = list(node.users)
+            if len(users) <= 1:
+                continue
+
+            for user in users[1:]:
+                with graph_module.graph.inserting_before(user):
+                    new_node = graph_module.graph.get_attr(node.target)
+                    new_node.meta.update(node.meta)
+                self._replace_input_node(user, node, new_node)
+                modified = True
+
+        return modified
+
     def call(self, graph_module: GraphModule) -> PassResult:
         seen_targets: set[str] = set()
-        modified = False
+        modified = self._split_shared_get_attrs(graph_module)
 
         for node in graph_module.graph.find_nodes(op="get_attr"):
 

backends/arm/test/passes/test_deduplicate_get_attr_pass.py

@@ -0,0 +1,37 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from executorch.backends.arm._passes.deduplicate_get_attr_pass import (
+    DeduplicateGetAttrPass,
+)
+from torch.fx import Graph, GraphModule
+
+
+def test_deduplicate_get_attr_splits_shared_node_users() -> None:
+    root = torch.nn.Module()
+    shared = torch.ones(2, 2)
+    root.register_buffer("shared", shared)
+
+    graph = Graph()
+    x = graph.placeholder("x")
+    attr = graph.get_attr("shared")
+    first = graph.call_function(torch.ops.aten.add.Tensor, (x, attr))
+    second = graph.call_function(torch.ops.aten.sub.Tensor, (first, attr))
+    graph.output(second)
+    graph_module = GraphModule(root, graph)
+
+    result = DeduplicateGetAttrPass()(graph_module)
+
+    assert result is not None
+    assert result.modified
+
+    get_attrs = list(graph_module.graph.find_nodes(op="get_attr"))
+    assert len(get_attrs) == 2
+    assert len({node.target for node in get_attrs}) == 2
+    assert first.args[1] is get_attrs[0]
+    assert second.args[1] is get_attrs[1]
+    assert getattr(graph_module, get_attrs[0].target) is shared
+    assert getattr(graph_module, get_attrs[1].target) is shared