Code Review 2

Author: winskuo-quic

backends/qualcomm/_passes/canonicalize_conv.py

@@ -34,6 +34,7 @@ def __init__(self, edge_program: torch.export.ExportedProgram):
         self.transpose_conv_set = {
             torch.ops.aten.conv_transpose1d.default,
             torch.ops.aten.conv_transpose2d.input,
+            torch.ops.aten.conv_transpose3d.input,
         }
 
     def dilate(self, tensor, dilation):

backends/qualcomm/builders/op_conv.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import warnings
 from typing import cast, Dict, List
 
 import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper
@@ -128,7 +127,8 @@ def define_node(
 
         filter_node = self.get_node(node.args[1])
         filter_tensor = get_parameter(filter_node, self.edge_program)
-        # weight of pytorch OIHW(conv2d) | IOHW(conv_transpose2d), yet QNN is HWIO
+        # weight of pytorch OIHW(conv2d) / IODHW(conv3d) or IOHW(conv_transpose2d) / IODHW(conv_transpose3d),
+        # yet QNN is HWIO or DHWIO
         is_transpose_conv = cast(bool, node.args[6])
         if is_conv2d:
             filter_axis_order = (2, 3, 0, 1) if is_transpose_conv else (2, 3, 1, 0)
@@ -193,23 +193,10 @@ def define_node(
         dilation_shape = [len(dilation)]
         output_padding_shape = [len(output_padding)]
 
-        # Some transpose_conv3d restrictions by QNN
-        if not is_conv2d and is_transpose_conv:
-            if dilation[0] != 1:
-
-                warnings.warn(
-                    "[QNN Delegate Op Builder]: As of QNN 2.37, TransposeConv3d only supports dilation of 1 in the depth dimension.",
-                    stacklevel=1,
-                )
-
-            # If using dilation, then stride height and width must = 1
-            if any(x != 1 for x in dilation) and (stride[1] != 1 or stride[2] != 1):
-                warnings.warn(
-                    "[QNN Delegate Op Builder]: As of QNN 2.37, TransposeConv3d, only supports dilation if stride in height and width = 1.",
-                    stacklevel=1,
-                )
-
         if is_transpose_conv:
+            assert all(
+                val == 1 for val in dilation
+            ), "CanonicalizeConv pass should perform dilate for transpose_conv."
             op_class = OpTransposeConv2d if is_conv2d else OpTransposeConv3d
         elif is_depthwise_conv:
             assert is_conv2d, "DepthWise only supports Conv2d"

backends/qualcomm/tests/models.py

@@ -765,8 +765,8 @@ def __init__(self, bias=True, dilation=1):
         self.conv_transpose = torch.nn.ConvTranspose3d(
             in_channels=1,
             out_channels=3,
-            kernel_size=1,
-            stride=(1, 1, 1),
+            kernel_size=3,
+            stride=2,
             padding=1,
             dilation=dilation,
             bias=bias,

backends/qualcomm/tests/test_qnn_delegate.py

@@ -317,7 +317,8 @@ def test_qnn_backend_conv_transpose3d(self):
         modules = [
             ConvTranspose3dSingle(),  # noqa: F405
             ConvTranspose3dSingle(bias=False),  # noqa: F405
-            ConvTranspose3dSingle(dilation=(1, 2, 3)),  # noqa: F405
+            ConvTranspose3dSingle(dilation=2),  # noqa: F405
+            ConvTranspose3dSingle(dilation=(3, 2, 3)),  # noqa: F405
         ]
         sample_input = (torch.randn([1, 1, 3, 3, 3]),)
         for i, module in enumerate(modules):
@@ -1897,7 +1898,8 @@ def test_qnn_backend_conv_transpose3d(self):
         modules = [
             ConvTranspose3dSingle(),  # noqa: F405
             ConvTranspose3dSingle(bias=False),  # noqa: F405
-            ConvTranspose3dSingle(dilation=(1, 2, 3)),  # noqa: F405
+            ConvTranspose3dSingle(dilation=2),  # noqa: F405
+            ConvTranspose3dSingle(dilation=(3, 2, 3)),  # noqa: F405
         ]
         sample_input = (torch.randn([1, 1, 3, 3, 3]),)
         for i, module in enumerate(modules):