Arm backend: Fix transposed conv lowering (pytorch#19391)

Avoid accidentally slicing transpose convolutions.

Input slicing is used for regular convolutions to trim unused edge
regions for TOSA lowering, but that is not valid for transpose
convolutions. Shrinking the input to a transpose convolution changes the
expanded output shape and can break output_padding cases.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell

Signed-off-by: Christoffer J.L <christoffer.johanssonlundqvist@arm.com>
Co-authored-by: Sebastian Larsson <38941629+Sebastian-Larsson@users.noreply.github.com>

Author: Christoffer-JL

backends/arm/_passes/size_adjust_input_pass.py

@@ -63,9 +63,13 @@ def _greater_than(input: SymIntLike, other: int) -> bool | torch.SymBool:
 
 
 def get_slices_convolution(conv_node: torch.fx.Node) -> Slices:
-    slices = []
+    slices: Slices = []
 
-    input_node, weight, _, stride_hw, pad_hw, dilation_hw, _, _, _ = conv_node.args
+    input_node, weight, _, stride_hw, pad_hw, dilation_hw, transposed, _, _ = (
+        conv_node.args
+    )
+    if transposed:
+        return slices
     weight_shape = cast(torch.fx.Node, weight).meta["val"].shape
     input_shape = cast(torch.fx.Node, input_node).meta["val"].shape
     spatial_rank = len(input_shape) - 2

backends/arm/test/passes/test_size_adjust_input_pass.py

@@ -42,6 +42,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.conv(x)
 
 
+class TransposeConvModule(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.conv = torch.nn.ConvTranspose2d(
+            in_channels=3,
+            out_channels=6,
+            kernel_size=3,
+            stride=2,
+            output_padding=1,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
 def _needs_truncation(input_length, kernel_size, stride, padding):
     return _greater_than((input_length + 2 * padding - kernel_size) % stride, padding)
 
@@ -115,6 +130,30 @@ def test_size_adjust_input_static_conv_no_adjustment_needed():
     ), "No slice nodes should be inserted when no adjustment is needed"
 
 
+def test_size_adjust_input_skips_transpose_conv2d() -> None:
+    model = TransposeConvModule()
+    example_inputs = (torch.randn(1, 3, 16, 16),)
+    edge_model = to_edge(export(model, example_inputs))
+    edge_model = edge_model.transform([SizeAdjustInputPass()])
+    gm = edge_model.exported_program().graph_module
+
+    conv_node = next(
+        n
+        for n in gm.graph.nodes
+        if n.op == "call_function"
+        and n.target == exir_ops.edge.aten.convolution.default
+    )
+    input_node = conv_node.args[0]
+    assert input_node.meta["val"].shape == example_inputs[0].shape
+
+    slice_nodes = [
+        n
+        for n in gm.graph.nodes
+        if n.op == "call_function" and n.target == exir_ops.edge.aten.slice_copy.Tensor
+    ]
+    assert len(slice_nodes) == 0
+
+
 def test_size_adjust_input_dynamic_conv2d():
     kernel_size, stride, padding = 3, 3, 1
     model = ConvModule(kernel_size=kernel_size, stride=stride, padding=padding)