Arm backend: Make sure DW-conv weights are reshaped once (pytorch#16071)

Make sure RewriteConv2DPass only reshapes shared weights once.

### Test plan
Functionality is tested in
backends/arm/test/misc/test_dw_convs_with_shared_weights.py.


Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/_passes/rewrite_conv2d_pass.py

@@ -90,7 +90,7 @@ def _is_depthwise_conv2d(self, node: torch.fx.Node) -> bool:
             return False
         groups = node.args[-1]
         in_channels = get_first_fake_tensor(node.all_input_nodes[0]).shape[1]
-        out_channels = get_first_fake_tensor(node.all_input_nodes[1]).shape[0]
+        out_channels = get_first_fake_tensor(node).shape[1]
         return (in_channels == groups) and (out_channels % in_channels) == 0
 
     def _reshape_weights(self, weight_node: torch.fx.Node, in_channels: int) -> None:
@@ -103,6 +103,7 @@ def _reshape_weights(self, weight_node: torch.fx.Node, in_channels: int) -> None
             raise RuntimeError(
                 f"Weight node {weight_node.name} is not a parameter or buffer"
             )
+
         reshaped_weight_tensor = (
             weight_tensor.permute(HWCM_ORDER)
             .reshape(
@@ -118,14 +119,19 @@ def _reshape_weights(self, weight_node: torch.fx.Node, in_channels: int) -> None
             param_name = self.exported_program.graph_signature.inputs_to_buffers[
                 weight_node.name
             ]
+            reshaped_weight_tensor = torch.nn.Buffer(reshaped_weight_tensor)
         elif is_param(self.exported_program, weight_node):
             param_name = self.exported_program.graph_signature.inputs_to_parameters[
                 weight_node.name
             ]
+            reshaped_weight_tensor = torch.nn.Parameter(
+                reshaped_weight_tensor, requires_grad=False
+            )
         else:
             raise RuntimeError(
                 f"Weight node {weight_node.name} is neither a parameter nor a buffer"
             )
+
         self.exported_program.state_dict[param_name] = reshaped_weight_tensor
         weight_node.meta["val"] = weight_node.meta["val"].reshape(
             weight_tensor.shape[2],
@@ -243,7 +249,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
             if self._is_depthwise_conv2d(node):
                 target_op = exir_ops.backend.tosa.DEPTHWISE_CONV2D.default
-                self._reshape_weights(weight, input_fake_tensor.shape[1])
+                # If there are any TOSA.DEPTHWISE_CONV2D nodes using the weights, we've already reshaped them.
+                if all(user.target != target_op for user in weight.users):
+                    self._reshape_weights(weight, input_fake_tensor.shape[1])
                 weight_fake_tensor = get_first_fake_tensor(weight)
             else:
                 target_op = exir_ops.backend.tosa.CONV2D.default

backends/arm/test/misc/test_dw_convs_with_shared_weights.py

@@ -0,0 +1,58 @@
+# Copyright 2025 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 Any, Tuple
+
+import torch
+from executorch.backends.arm._passes.rewrite_conv2d_pass import RewriteConv2dPass
+from executorch.backends.arm.test.tester.test_pipeline import (
+    PassPipeline,
+    TosaPipelineFP,
+    TosaPipelineINT,
+)
+
+input_t = Tuple[torch.Tensor]
+
+
+class DWConvsModule(torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        conv = torch.nn.Conv2d(6, 6, kernel_size=(2, 2), groups=6)
+        relu = torch.nn.ReLU()
+        self.sequential = torch.nn.ModuleList([conv, relu, conv])
+
+    def forward(self, x) -> torch.Tensor:
+        for m in self.sequential:
+            x = m(x)
+        return x
+
+    def get_inputs(self) -> input_t:
+        return (torch.randn(1, 6, 24, 24),)
+
+
+def test_convs_tosa_fp():
+    module = DWConvsModule()
+    pipeline = TosaPipelineFP[input_t](
+        module, module.get_inputs(), aten_op=[], exir_op=[]
+    )
+    pipeline.run()
+
+
+def test_convs_tosa_int():
+    module = DWConvsModule()
+    pipeline = TosaPipelineINT[input_t](
+        module, module.get_inputs(), aten_op=[], exir_op=[]
+    )
+    pipeline.run()
+
+
+def test_rewrite_conv_pass():
+    module = DWConvsModule()
+    pipeline = PassPipeline(
+        module, module.get_inputs(), passes_with_exported_program=[RewriteConv2dPass]
+    )
+    # We can't run TOSA backend dialect operators in eager mode
+    pipeline.pop_stage("run_method_and_compare_outputs")
+    pipeline.run()