Add transposed convolution

Summary: Continued support of cadence python references

Differential Revision: D83602808

Author: Andrew Grebenisan

backends/cadence/aot/ref_implementations.py

@@ -932,7 +932,6 @@ def quantized_conv1d_nlc_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ...
 @quantized_conv_variant("nhwc", torch.uint8, torch.uint8, is_1d=True)
 def quantized_conv1d_nlc_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ...
 
-
 @impl(m, "convolution")
 def convolution(
     input_tensor: torch.Tensor,
@@ -960,6 +959,7 @@ def convolution(
     _stride: tuple[int, int] | int = stride
     _padding: tuple[int, int] | int = padding
     _dilation: tuple[int, int] | int = dilation
+
     if conv_is_1d:
         conv = torch.nn.functional.conv1d
         _stride = stride[0]
@@ -977,6 +977,52 @@ def convolution(
 
     return conv_out
 
+@impl(m, "transposed_convolution")
+def transposed_convolution(
+    input_tensor: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    stride: tuple[int, int],
+    padding: tuple[int, int],
+    dilation: tuple[int, int],
+    output_padding: tuple[int, int],
+    groups: int,
+    channel_last: bool = False,
+) -> torch.Tensor:
+
+    conv_is_1d = len(input_tensor.shape) == 3
+    if channel_last:
+        if conv_is_1d:
+            input_tensor = input_tensor.movedim(-1, 1).contiguous()
+            if len(weight.shape) != 3:
+                raise ValueError("Weight tensor must be 3D if input is 3D")
+            weight = weight.movedim(-1, 1).contiguous()
+        else:
+            input_tensor = input_tensor.movedim(-1, -3)
+            if len(weight.shape) != 4:
+                raise ValueError("Weight tensor must be 4D if input is nd > 3")
+            weight = torch.permute(weight, (0, -1, 1, 2)).contiguous()
+
+    _stride: tuple[int, int] | int = stride
+    _padding: tuple[int, int] | int = padding
+    _dilation: tuple[int, int] | int = dilation
+
+    if conv_is_1d:
+        conv = torch.nn.functional.conv_transpose1d
+        _stride = stride[0]
+        _padding = padding[0]
+        _dilation = dilation[0]
+    else:
+        conv = torch.nn.functional.conv_transpose2d
+
+    conv_out = conv(input_tensor, weight, bias, _stride, _padding, output_padding, groups, _dilation)
+    if channel_last:
+        if conv_is_1d:
+            conv_out = conv_out.movedim(1, -1).contiguous()
+        else:
+            conv_out = conv_out.movedim(-3, -1).contiguous()
+
+    return conv_out
 
 @impl(m, "avg_pool2d")
 def avg_pool2d(

backends/cadence/aot/tests/test_ref_implementations.py

@@ -1259,7 +1259,7 @@ def test_rope(
 
     @expand(
         [
-            # Test case 1: Basic 2D convolution (NCHW format)
+            # Basic 2D convolution (NCHW format)
             (
                 "basic_2d_nchw",
                 torch.tensor(
@@ -1274,11 +1274,12 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor(
                     [[[[5.0]]]], dtype=torch.float32
                 ),  # expected: 1*1 + 4*1 = 5
             ),
-            # Test case 2: Basic 2D convolution (NHWC format)
+            # Basic 2D convolution (NHWC format)
             (
                 "basic_2d_nhwc",
                 torch.tensor(
@@ -1293,11 +1294,12 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 True,  # channel_last
+
                 torch.tensor(
                     [[[[5.0]]]], dtype=torch.float32
                 ),  # expected: 1*1 + 4*1 = 5
             ),
-            # Test case 3: 2D convolution with stride=2
+            # 2D convolution with stride=2
             (
                 "conv2d_stride2",
                 torch.tensor(
@@ -1322,9 +1324,10 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor([[[[14.0, 22.0], [46.0, 54.0]]]], dtype=torch.float32),
             ),
-            # Test case 4: 2D convolution with padding=1
+            # 2D convolution with padding=1
             (
                 "conv2d_padding1",
                 torch.tensor(
@@ -1339,12 +1342,13 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor(
                     [[[[1.0, 2.0, 0.0], [3.0, 5.0, 2.0], [0.0, 3.0, 4.0]]]],
                     dtype=torch.float32,
                 ),  # expected with padding
             ),
-            # Test case 5: 2D convolution with dilation=2
+            # 2D convolution with dilation=2
             (
                 "conv2d_dilation2",
                 torch.tensor(
@@ -1369,9 +1373,10 @@ def test_rope(
                 (2, 2),  # dilation=2
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor([[[[24.0, 28.0], [40.0, 44.0]]]], dtype=torch.float32),
             ),
-            # Test case 6: 2D grouped convolution (groups=2)
+            # 2D grouped convolution (groups=2)
             (
                 "conv2d_groups2",
                 torch.tensor(
@@ -1396,9 +1401,10 @@ def test_rope(
                 (1, 1),  # dilation
                 2,  # groups=2
                 False,  # channel_last
+
                 torch.tensor([[[[10.0]], [[14.0]]]], dtype=torch.float32),
             ),
-            # Test case 7: 1D convolution (NCL format)
+            # 1D convolution (NCL format)
             (
                 "conv1d_ncl",
                 torch.tensor(
@@ -1411,11 +1417,12 @@ def test_rope(
                 (1, 1),  # dilation (only dilation[1] is used for 1D)
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor(
                     [[[3.0, 5.0, 7.0]]], dtype=torch.float32
                 ),  # expected: [1+2, 2+3, 3+4]
             ),
-            # Test case 8: 1D convolution (NLC format)
+            # 1D convolution (NLC format)
             (
                 "conv1d_nlc",
                 torch.tensor(
@@ -1430,9 +1437,10 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 True,  # channel_last
+
                 torch.tensor([[[3.0], [5.0], [7.0]]], dtype=torch.float32),
             ),
-            # Test case 9: Multi-channel input and output
+            # Multi-channel input and output
             (
                 "multi_channel",
                 torch.tensor(
@@ -1469,9 +1477,10 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor([[[[10.0]], [[11.0]]]], dtype=torch.float32),
             ),
-            # Test case 10: Convolution with non-zero bias
+            # Convolution with non-zero bias
             (
                 "conv2d_with_bias",
                 torch.tensor(
@@ -1486,6 +1495,7 @@ def test_rope(
                 (1, 1),  # dilation
                 1,  # groups
                 False,  # channel_last
+
                 torch.tensor(
                     [[[[15.0]]]], dtype=torch.float32
                 ),  # expected: 5 + 10 = 15
@@ -1534,6 +1544,75 @@ def test_convolution(
             f"Output values don't match expected in {name}. Got {output}, expected {expected_output}",
         )
 
+    @expand(
+        [
+            (
+                "conv2d_transposed_stride2",
+                torch.tensor(
+                    [[[[1.0, 2.0], [3.0, 4.0]]]], dtype=torch.float32
+                ),  # input: 1x1x2x2
+                torch.tensor(
+                    [[[[1.0, 1.0], [1.0, 1.0]]]], dtype=torch.float32
+                ),  # weight: 1x1x2x2
+                torch.tensor([0.0], dtype=torch.float32),  # bias
+                (1, 1),  # stride=2
+                (0, 0),  # padding
+                (1, 1),  # dilation
+                1,  # groups
+                (0, 0), # output_padding
+                False,  # channel_last
+                torch.tensor(
+                    [[[[1.0,  3.0, 2.0],
+                       [4.0, 10.0, 6.0],
+                       [3.0,  7.0, 4.0]]]], dtype=torch.float32
+                ),
+            ),
+        ]
+    )
+    def test_transposed_convolution(
+        self,
+        name: str,
+        input_tensor: torch.Tensor,
+        weight: torch.Tensor,
+        bias: torch.Tensor,
+        stride: tuple[int, int],
+        padding: tuple[int, int],
+        dilation: tuple[int, int],
+        groups: int,
+        output_padding: tuple[int, int],
+        channel_last: bool,
+        expected_output: torch.Tensor,
+    ) -> None:
+        output = torch.ops.cadence.transposed_convolution(
+            input_tensor,
+            weight,
+            bias,
+            stride,
+            padding,
+            dilation,
+            output_padding,
+            groups,
+            channel_last,
+        )
+
+        # Verify output properties
+        self.assertEqual(
+            output.dtype,
+            input_tensor.dtype,
+            f"Output dtype should match input dtype in {name}",
+        )
+        self.assertEqual(
+            output.shape,
+            expected_output.shape,
+            f"Output shape should match expected shape in {name}",
+        )
+
+        # Verify output matches expected values
+        self.assertTrue(
+            torch.equal(output, expected_output),
+            f"Output values don't match expected in {name}. Got {output}, expected {expected_output}",
+        )
+
     @expand(
         [
             # Basic non-quantized average pooling