Arm backend: Fix transposed conv lowering

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.

Signed-off-by: Christoffer J.L <christoffer.johanssonlundqvist@arm.com>
Change-Id: Ib3e349310bf1672ea1cce366154101aed1415e6e

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)