XNNPACK: Remove no-op expand_copy before partitioning (#19978)

Remove aten.expand_copy nodes when input and output metadata have the
same dtype and shape. Static export can leave these shape-preserving
expands as portable copy kernels even though they are identities for the
lowered graph.

Run the cleanup in the normal XNNPACK transform pass path so it can
remove inter-delegate expand_copy nodes before partitioning.

For EdgeTAM mask decoder, expand_copy ops drop from 32 to 0,
non-delegate kernel calls drop from 114 to 82, and delegate calls drop
by 1, resulting in a ~9% speedup on a measured SVE2 and SME2 Android
devices.



cc @GregoryComer @digantdesai @cbilgin @freddan80 @per @zingo
@oscarandersson8218 @Sebastian-Larsson @robell @rascani

Signed-off-by: Måns Nilsson <mans.nilsson@arm.com>

Author: mansnils

backends/xnnpack/_passes/remove_noop_expand_copy_pass.py

@@ -0,0 +1,59 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# 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.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+class RemoveNoopExpandCopyPass(ExportPass):
+    """
+    Remove ``expand_copy`` nodes that do not change tensor shape or dtype.
+
+    In static XNNPACK export flows, shape-specialization can turn an expand into
+    a materialized copy whose input and output metadata are identical. Such a
+    node is an identity for the lowered graph and can be bypassed. The pass
+    leaves nodes in place whenever the output shape differs from the input
+    shape.
+    """
+
+    def _is_noop_expand_copy(self, node: torch.fx.Node) -> bool:
+        # TODO: Investigate moving this to a shared backend transform. Other
+        # backends already carry equivalent no-op expand handling.
+        if node.target != exir_ops.edge.aten.expand_copy.default:
+            return False
+
+        input_node = node.args[0]
+        if not isinstance(input_node, torch.fx.Node):
+            return False
+
+        input_value = input_node.meta.get("val")
+        output_value = node.meta.get("val")
+        if input_value is None or output_value is None:
+            return False
+
+        return (
+            input_value.dtype == output_value.dtype
+            and input_value.shape == output_value.shape
+        )
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        graph = graph_module.graph
+
+        for node in list(graph.nodes):
+            if not self._is_noop_expand_copy(node):
+                continue
+
+            node.replace_all_uses_with(node.args[0])
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, True)

backends/xnnpack/test/passes/test_remove_noop_expand_copy_pass.py

@@ -0,0 +1,75 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# 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 unittest
+
+import torch
+from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
+    RemoveNoopExpandCopyPass,
+)
+from executorch.backends.xnnpack.test.tester import RunPasses, Tester
+from executorch.backends.xnnpack.utils.configs import (
+    get_transform_passes,
+    get_xnnpack_edge_compile_config,
+)
+from executorch.exir import to_edge_transform_and_lower
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+class TestRemoveNoopExpandCopyPass(unittest.TestCase):
+    PassStage = RunPasses([RemoveNoopExpandCopyPass])
+    expand_copy_name = "executorch_exir_dialects_edge__ops_aten_expand_copy_default"
+
+    def setUp(self):
+        torch._dynamo.reset()
+
+    class NoopExpand(torch.nn.Module):
+        def forward(self, x):
+            y = x.expand(x.shape)
+            return y + 1
+
+    class BroadcastExpand(torch.nn.Module):
+        def forward(self, x):
+            y = x.expand(2, 3)
+            return y + 1
+
+    def test_removes_same_shape_expand_copy(self):
+        (
+            Tester(self.NoopExpand(), (torch.randn(2, 3),))
+            .export()
+            .to_edge()
+            .check_count({self.expand_copy_name: 1})
+            .run_passes(self.PassStage)
+            .check_count({self.expand_copy_name: 0})
+            .run_method_and_compare_outputs()
+        )
+
+    def test_keeps_broadcasting_expand_copy(self):
+        (
+            Tester(self.BroadcastExpand(), (torch.randn(1, 3),))
+            .export()
+            .to_edge()
+            .check_count({self.expand_copy_name: 1})
+            .run_passes(self.PassStage)
+            .check_count({self.expand_copy_name: 1})
+            .run_method_and_compare_outputs()
+        )
+
+    def test_transform_passes_remove_same_shape_expand_copy(self):
+        edge_program = to_edge_transform_and_lower(
+            torch.export.export(self.NoopExpand(), (torch.randn(2, 3),), strict=True),
+            transform_passes=get_transform_passes(),
+            compile_config=get_xnnpack_edge_compile_config(),
+        )
+        graph = edge_program.exported_program().graph_module.graph
+
+        self.assertFalse(
+            any(
+                node.target == exir_ops.edge.aten.expand_copy.default
+                for node in graph.nodes
+            )
+        )

backends/xnnpack/utils/configs.py

@@ -1,12 +1,17 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# 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.
 
 from typing import List
 
 import executorch.exir as exir
+
+from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
+    RemoveNoopExpandCopyPass,
+)
 from executorch.exir.pass_manager import PassType
 
 
@@ -20,7 +25,9 @@ def get_xnnpack_edge_compile_config(
 
 
 def get_transform_passes(additional_passes=None) -> List[PassType]:
-    passes = additional_passes if additional_passes else []
+    passes = [RemoveNoopExpandCopyPass()]
+    if additional_passes:
+        passes.extend(additional_passes)
     return passes