conv2d padding

Differential Revision: D84071939

Author: cccclai

backends/qualcomm/builders/op_pad.py

@@ -18,7 +18,12 @@
 
 @register_node_visitor
 class Pad(NodeVisitor):
-    target = ["aten.constant_pad_nd.default"]
+    target = [
+        "aten.constant_pad_nd.default",
+        "aten.pad.default",  # handles reflect/replicate modes
+        "aten.reflection_pad2d.default",
+        "aten.replication_pad2d.default",
+    ]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)
@@ -28,6 +33,8 @@ def define_node(
         node: torch.fx.Node,
         nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper],
     ) -> PyQnnWrapper.PyQnnOpWrapper:
+
+        # ---- Input tensor ----
         input_node = self.get_node(node.args[0])
         input_tensor = self.get_tensor(input_node, node)
         pad_inp_tensor_wrapper = self.define_tensor(
@@ -39,6 +46,7 @@ def define_node(
         )
         pad_input_tensors = [pad_inp_tensor_wrapper]
 
+        # ---- Output tensor ----
         output_tensor = self.get_tensor(node, node)
         output_tensor_wrapper = self.define_tensor(
             node,
@@ -49,21 +57,43 @@ def define_node(
         )
         pad_output_tensors = [output_tensor_wrapper]
 
+        # ---- Pad amount handling ----
+        pad_list = cast(List[int], node.args[1])
         pad_amount_shape = [input_tensor.dim(), 2]
-        # pytorch padding start from the last index
-        pad_amount = np.reshape(cast(List[int], node.args[1]), (-1, 2))[::-1].astype(
-            np.uint32
-        )
-        # fulfill the pad amount for each idex of tensor
+
+        # PyTorch pad order: [last_dim, ..., first_dim]
+        pad_amount = np.reshape(pad_list, (-1, 2))[::-1].astype(np.uint32)
+
+        # Expand to full rank if needed
         if zero_amounts := pad_amount_shape[0] - pad_amount.shape[0]:
             pad_amount = np.concatenate(
                 (np.array([(0, 0)] * zero_amounts), pad_amount)
             ).astype(np.uint32)
 
+        # Apply axis reordering if necessary
         if QCOM_AXIS_ORDER in node.meta:
             pad_amount = pad_amount[list(node.meta[QCOM_AXIS_ORDER])]
-        pad_amount_val = node.args[2]
 
+        # ---- Determine mode ----
+        if len(node.args) >= 3 and isinstance(node.args[2], str):
+            mode = node.args[2]
+        elif "reflection" in node.target:
+            mode = "reflect"
+        elif "replication" in node.target:
+            mode = "replicate"
+        else:
+            mode = "constant"
+
+        scheme_map = {
+            "constant": OpPad.Scheme.CONSTANT,
+            "reflect": OpPad.Scheme.MIRROR_REFLECT,
+            "replicate": OpPad.Scheme.EDGE,
+        }
+
+        if mode not in scheme_map:
+            raise ValueError(f"[QNN][Pad] Unsupported pad mode: {mode}")
+
+        # ---- Create QNN op ----
         pad_op = PyQnnWrapper.PyQnnOpWrapper(
             node.name,
             QNN_OP_PACKAGE_NAME_QTI_AISW,
@@ -72,19 +102,29 @@ def define_node(
         pad_op.AddInputTensors(pad_input_tensors)
         pad_op.AddOutputTensors(pad_output_tensors)
 
-        # For now, we only support constant (0) padding due to torch implementation
+        # scheme param
         pad_op.AddScalarParam(
             OpPad.param_scheme,
             PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32,
-            {QCOM_DATA: np.uint32(OpPad.Scheme.CONSTANT)},
+            {QCOM_DATA: np.uint32(scheme_map[mode])},
         )
 
-        pad_op.AddScalarParam(
-            OpPad.param_pad_constant_value,
-            QNN_TENSOR_TYPE_MAP[type(pad_amount_val)],
-            {QCOM_DATA: pad_amount_val},
-        )
+        # pad_constant_value param (only for constant mode)
+        if mode == "constant":
+            # torch.constant_pad_nd takes optional pad value, default = 0.0
+            pad_value = node.kwargs.get("value", None)
+            if pad_value is None and len(node.args) > 2 and not isinstance(node.args[2], str):
+                pad_value = node.args[2]
+            if pad_value is None:
+                pad_value = 0.0
+
+            pad_op.AddScalarParam(
+                OpPad.param_pad_constant_value,
+                QNN_TENSOR_TYPE_MAP[type(pad_value)],
+                {QCOM_DATA: pad_value},
+            )
 
+        # pad_amount tensor param
         pad_op.AddTensorParam(
             OpPad.param_pad_amount,
             PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32,

backends/qualcomm/partition/utils.py

@@ -52,6 +52,7 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]:
         torch.ops.aten.leaky_relu.default,
         torch.ops.aten.linear.default,
         torch.ops.aten.matmul.default,
+        torch.ops.aten.pad.default,
         torch.ops.aten.pixel_shuffle.default,
         torch.ops.aten.pixel_unshuffle.default,
         torch.ops.aten.prelu.default,

backends/qualcomm/tests/models.py

@@ -490,6 +490,35 @@ def forward(self, x):
         return x
 
 
+class Conv2d(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        out_channels=6,
+        kernel_size: Union[int, Tuple[int, int]] = 3,
+        stride: Union[int, Tuple[int, int]] = 1,
+        padding: Union[int, Tuple[int, int]] = 0,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
 class Conv2dArgmin(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -318,6 +318,13 @@ def test_qnn_backend_conv2d(self):
             with self.subTest(i=i):
                 self.lower_module_and_test_output(module, sample_input)
 
+
+    def test_qnn_backend_conv2d_mode(self):
+        sample_input = (torch.randn(4, 3, 16, 16),)
+        for mode in ["zeros", "reflect", "replicate", "circular"]:
+            module = Conv2d(padding=1, padding_mode=mode)  # noqa: F405
+            self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_conv2d_channel_last(self):
         modules = [
             Conv2dSequential(channel_last=True),  # noqa: F405
@@ -1996,6 +2003,13 @@ def test_qnn_backend_conv2d(self):
                 module = self.get_qdq_module(module, sample_input)
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_conv2d_mode(self):
+        sample_input = (torch.randn(4, 3, 16, 16),)
+        for mode in ["zeros", "reflect", "replicate", "circular"]:
+            module = Conv2d(padding=1, padding_mode=mode)  # noqa: F405
+            module = self.get_qdq_module(module, sample_input)
+            self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_conv2d_block(self):
         o_ch, i_ch, kernel, padding = 32, 512, (1, 1), 0