Arm backend: Add support for all padding modes (pytorch#18521)

The pad operator with different modes was decomposed in the to-edge-step
into a number of operation which were not handled well by the arm
backend.

Instead avoid the default decomposition and decompose replicate circular
and reflect modes using slice + concat operators.

Signed-off-by: Adrian Lundell <adrian.lundell@arm.com>

Author: AdrianLundell

backends/arm/_passes/rewrite_pad.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from collections.abc import Sequence
 from typing import Set, Type
 
 import torch
@@ -20,18 +21,10 @@ class RewritePadPass(ArmPass):
     _passes_required_after: Set[Type[ExportPass]] = set()
     targeted_ops = {
         exir_ops.edge.aten.constant_pad_nd.default,
+        exir_ops.edge.aten.pad.default,
     }
 
-    def call_operator(self, op, args, kwargs, meta, updated=False):
-        if op not in self.targeted_ops:
-            return super().call_operator(op, args, kwargs, meta)
-
-        if len(args) == 3:
-            input_tensor, pad, value = args
-        else:
-            input_tensor, pad = args
-            value = 0
-
+    def _rewrite_constant_pad(self, input_tensor, pad, value, meta):
         output_dtype = meta["val"].dtype
         if output_dtype in (torch.int8, torch.int16):
             input_qparams = meta.data.get("input_qparams", {})
@@ -65,3 +58,120 @@ def call_operator(self, op, args, kwargs, meta, updated=False):
             meta,
             True,
         )
+
+    def _slice_idx(self, x, dim: int, idx: int, meta):
+        return super().call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor,
+            (x, dim, idx, idx + 1),
+            {},
+            meta,
+            True,
+        )
+
+    def _pad_along_dim(
+        self,
+        x,
+        dim: int,
+        left: int,
+        right: int,
+        mode: str,
+        meta,
+    ):
+        if left == 0 and right == 0:
+            return x
+
+        size = x.data.shape[dim]
+        if isinstance(size, torch.SymInt):
+            raise ValueError(f"Pad mode '{mode}' does not support symbolic shape yet.")
+        if not isinstance(size, int):
+            raise ValueError(f"Expected integer dim size for pad rewrite, got {size}.")
+
+        left_tensors = []
+        right_tensors = []
+
+        if mode == "replicate":
+            left_tensors = [self._slice_idx(x, dim, 0, meta) for _ in range(left)]
+            right_tensors = [
+                self._slice_idx(x, dim, size - 1, meta) for _ in range(right)
+            ]
+        elif mode == "circular":
+            left_tensors = [
+                self._slice_idx(x, dim, size - left + i, meta) for i in range(left)
+            ]
+            right_tensors = [self._slice_idx(x, dim, i, meta) for i in range(right)]
+        elif mode == "reflect":
+            if left >= size or right >= size:
+                raise ValueError(
+                    f"Pad mode 'reflect' requires pad < input size, got left={left}, right={right}, size={size}."
+                )
+            left_tensors = [
+                self._slice_idx(x, dim, left - i, meta) for i in range(left)
+            ]
+            right_tensors = [
+                self._slice_idx(x, dim, size - 2 - i, meta) for i in range(right)
+            ]
+        else:
+            raise ValueError(f"Unsupported pad mode '{mode}'.")
+
+        return super().call_operator(
+            exir_ops.edge.aten.cat.default,
+            (left_tensors + [x] + right_tensors, dim),
+            {},
+            meta,
+            True,
+        )
+
+    def _rewrite_non_constant_pad(
+        self,
+        input_tensor,
+        pad: Sequence[int],
+        mode: str,
+        meta,
+    ):
+        if len(pad) % 2 != 0:
+            raise ValueError(f"Invalid pad spec length {len(pad)} for mode '{mode}'.")
+
+        output = input_tensor
+        pairs = [(pad[i], pad[i + 1]) for i in range(0, len(pad), 2)]
+        rank = len(input_tensor.data.shape)
+        for pair_idx, (left, right) in enumerate(pairs):
+            if not isinstance(left, int) or not isinstance(right, int):
+                raise ValueError(
+                    f"Pad mode '{mode}' expects integer pad values, got ({left}, {right})."
+                )
+            # F.pad pad tuples are ordered from the innermost dimension outward.
+            dim = rank - 1 - pair_idx
+            output = self._pad_along_dim(output, dim, left, right, mode, meta)
+        return output
+
+    def call_operator(self, op, args, kwargs, meta, updated=False):
+        if op not in self.targeted_ops:
+            return super().call_operator(op, args, kwargs, meta)
+
+        if op == exir_ops.edge.aten.constant_pad_nd.default:
+            if len(args) == 3:
+                input_tensor, pad, value = args
+            else:
+                input_tensor, pad = args
+                value = 0
+            return self._rewrite_constant_pad(input_tensor, pad, value, meta)
+
+        if len(args) < 2:
+            raise ValueError(
+                f"Expected at least 2 args for aten.pad.default, got {args}"
+            )
+
+        input_tensor, pad = args[:2]
+        mode = args[2] if len(args) > 2 else kwargs.get("mode", "constant")
+        value = args[3] if len(args) > 3 else kwargs.get("value", 0)
+
+        if not isinstance(mode, str):
+            raise ValueError(f"Expected string mode in aten.pad.default, got {mode}")
+
+        if mode == "constant":
+            return self._rewrite_constant_pad(input_tensor, pad, value, meta)
+
+        if mode in ("reflect", "replicate", "circular"):
+            return self._rewrite_non_constant_pad(input_tensor, pad, mode, meta)
+
+        raise ValueError(f"Unsupported pad mode '{mode}' in aten.pad.default.")

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -98,6 +98,7 @@
     exir_ops.edge.quantized_decomposed.quantize_per_channel.default,
     exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
     exir_ops.edge.quantized_decomposed.dequantize_per_channel.default,
+    exir_ops.edge.aten.pad.default,
     exir_ops.edge.aten.constant_pad_nd.default,
     exir_ops.edge.aten.amax.default,
     exir_ops.edge.aten.amin.default,
@@ -219,6 +220,7 @@
     exir_ops.edge.aten.pow.Tensor_Scalar,
     exir_ops.edge.aten.pow.Tensor_Tensor,
     operator.getitem,
+    exir_ops.edge.aten.pad.default,
     exir_ops.edge.aten.constant_pad_nd.default,
     exir_ops.edge.aten.amax.default,
     exir_ops.edge.aten.amin.default,

backends/arm/test/ops/test_constant_pad_nd.py

@@ -22,47 +22,113 @@
 input_t1 = Tuple[torch.Tensor]  # Input x
 
 test_data_suite = {
-    "4dim_last1dim": lambda: (torch.rand(1, 1, 16, 16), (1, 1, 0, 0, 0, 0, 0, 0), 1),
-    "4dim_last2dim": lambda: (torch.rand(1, 1, 16, 16), (1, 0, 1, 0, 0, 0, 0, 0), 2),
-    "4dim_last3dim": lambda: (torch.rand(1, 1, 16, 16), (1, 1, 0, 2, 0, 2, 0, 0), 3),
-    "4dim_last4dim": lambda: (torch.rand(1, 1, 16, 16), (1, 0, 1, 1, 0, 2, 0, 2), 4),
-    "3dim_last1dim": lambda: (torch.rand(1, 1, 16), (1, 1, 0, 0, 0, 0), 1),
-    "3dim_last2dim": lambda: (torch.rand(1, 1, 16), (1, 0, 1, 1, 0, 0), 2),
-    "3dim_last3dim": lambda: (torch.rand(1, 1, 16), (1, 0, 1, 0, 1, 1), 3),
-    "2dim_last1dim": lambda: (torch.rand(1, 1, 16), (1, 1, 0, 0), 1),
-    "2dim_last2dim": lambda: (torch.rand(1, 1, 16), (1, 0, 1, 1), 2),
+    "4dim_last1dim": lambda: (
+        torch.rand(1, 1, 16, 16),
+        (1, 1, 0, 0, 0, 0, 0, 0),
+        1,
+        "constant",
+    ),
+    "4dim_last2dim": lambda: (
+        torch.rand(1, 1, 16, 16),
+        (1, 0, 1, 0, 0, 0, 0, 0),
+        2,
+        "constant",
+    ),
+    "4dim_last3dim": lambda: (
+        torch.rand(1, 1, 16, 16),
+        (1, 1, 0, 2, 0, 2, 0, 0),
+        3,
+        "constant",
+    ),
+    "4dim_last4dim": lambda: (
+        torch.rand(1, 1, 16, 16),
+        (1, 0, 1, 1, 0, 2, 0, 2),
+        4,
+        "constant",
+    ),
+    "3dim_last1dim": lambda: (torch.rand(1, 1, 16), (1, 1, 0, 0, 0, 0), 1, "constant"),
+    "3dim_last2dim": lambda: (torch.rand(1, 1, 16), (1, 0, 1, 1, 0, 0), 2, "constant"),
+    "3dim_last3dim": lambda: (torch.rand(1, 1, 16), (1, 0, 1, 0, 1, 1), 3, "constant"),
+    "2dim_last1dim": lambda: (torch.rand(1, 1, 16), (1, 1, 0, 0), 1, "constant"),
+    "2dim_last2dim": lambda: (torch.rand(1, 1, 16), (1, 0, 1, 1), 2, "constant"),
+    "4dim_reflect": lambda: (
+        torch.rand(6, 6, 6, 6),
+        (3, 3, 3, 3, 3, 3),
+        None,
+        "reflect",
+    ),
+    "4dim_replicate": lambda: (
+        torch.rand(3, 3, 3, 3),
+        (3, 3, 3, 3, 3, 3),
+        None,
+        "replicate",
+    ),
+    "4dim_circular": lambda: (
+        torch.rand(3, 3, 3, 3),
+        (3, 3, 3, 3, 3, 3),
+        None,
+        "circular",
+    ),
+    "2dim_reflect": lambda: (
+        torch.rand(6, 6),
+        (3, 3),
+        None,
+        "reflect",
+    ),
+    "2dim_replicate": lambda: (
+        torch.rand(3, 3),
+        (3, 3),
+        None,
+        "replicate",
+    ),
+    "2dim_circular": lambda: (
+        torch.rand(3, 3),
+        (3, 3),
+        None,
+        "circular",
+    ),
 }
 
 test_data_suite_bf16 = {
     "4dim_last1dim_bf16": lambda: (
         torch.rand(1, 1, 8, 8, dtype=torch.bfloat16),
         (1, 1, 0, 0, 0, 0, 0, 0),
         1.0,
+        "constant",
     ),
     "3dim_last1dim_bf16": lambda: (
         torch.rand(1, 1, 8, dtype=torch.bfloat16),
         (1, 0, 1, 0, 0, 0),
         -0.5,
+        "constant",
     ),
 }
 test_data_suite_fp16 = {
     "4dim_last1dim_fp16": lambda: (
         torch.rand(1, 1, 8, 8, dtype=torch.float16),
         (1, 1, 0, 0, 0, 0, 0, 0),
         1.0,
+        "constant",
     ),
     "3dim_last1dim_fp16": lambda: (
         torch.rand(1, 1, 8, dtype=torch.float16),
         (1, 0, 1, 0, 0, 0),
         -0.5,
+        "constant",
     ),
 }
 
 
 class ConstantPadND(torch.nn.Module):
-    def __init__(self, pad: Tuple, value: float | None = None):
+    def __init__(
+        self,
+        pad: Tuple,
+        value: float | None = None,
+        mode: str = "constant",
+    ):
         super().__init__()
         self.value = value
+        self.mode = mode
         nonzero_idx = len(pad)
         for i in range(0, len(pad), 2):
             if pad[i] + pad[i + 1] == 0:
@@ -71,18 +137,17 @@ def __init__(self, pad: Tuple, value: float | None = None):
         self.pad = pad[:nonzero_idx]
 
     def forward(self, x: torch.Tensor):
-        x = F.pad(x, pad=self.pad, mode="constant", value=self.value)
-        return x
+        return F.pad(x, pad=self.pad, mode=self.mode, value=self.value)
 
 
 @common.parametrize(
     "test_data",
     test_data_suite | test_data_suite_bf16 | test_data_suite_fp16,
 )
 def test_constant_pad_nd_tosa_FP(test_data: Tuple):
-    test_data, padding, value = test_data()
+    test_data, padding, value, mode = test_data()
     pipeline = TosaPipelineFP[input_t1](
-        ConstantPadND(padding, value),
+        ConstantPadND(padding, value, mode),
         (test_data,),
         aten_op,
         exir_op,
@@ -93,9 +158,9 @@ def test_constant_pad_nd_tosa_FP(test_data: Tuple):
 
 @common.parametrize("test_data", test_data_suite)
 def test_constant_pad_nd_tosa_INT(test_data: Tuple):
-    test_data, padding, value = test_data()
+    test_data, padding, value, mode = test_data()
     pipeline = TosaPipelineINT[input_t1](
-        ConstantPadND(padding, value),
+        ConstantPadND(padding, value, mode),
         (test_data,),
         aten_op,
         exir_op,
@@ -106,9 +171,9 @@ def test_constant_pad_nd_tosa_INT(test_data: Tuple):
 @common.parametrize("test_data", test_data_suite)
 def test_constant_pad_nd_tosa_INT_a16w8(test_data: Tuple):
     """Test constant_pad_nd op with int16 I/O quantization for TOSA INT."""
-    test_data, padding, value = test_data()
+    test_data, padding, value, mode = test_data()
     pipeline = TosaPipelineINT[input_t1](
-        ConstantPadND(padding, value),
+        ConstantPadND(padding, value, mode),
         (test_data,),
         aten_op,
         exir_op,
@@ -120,9 +185,9 @@ def test_constant_pad_nd_tosa_INT_a16w8(test_data: Tuple):
 @common.parametrize("test_data", test_data_suite | test_data_suite_fp16)
 @common.SkipIfNoModelConverter
 def test_constant_pad_nd_vgf_no_quant(test_data: Tuple):
-    inp, padding, value = test_data()
+    inp, padding, value, mode = test_data()
     pipeline = VgfPipeline[input_t1](
-        ConstantPadND(padding, value),
+        ConstantPadND(padding, value, mode),
         (inp,),
         aten_op,
         exir_op,
@@ -134,9 +199,9 @@ def test_constant_pad_nd_vgf_no_quant(test_data: Tuple):
 @common.parametrize("test_data", test_data_suite)
 @common.SkipIfNoModelConverter
 def test_constant_pad_nd_vgf_quant(test_data: Tuple):
-    inp, padding, value = test_data()
+    inp, padding, value, mode = test_data()
     pipeline = VgfPipeline[input_t1](
-        ConstantPadND(padding, value),
+        ConstantPadND(padding, value, mode),
         (inp,),
         aten_op,
         exir_op,