Arm backend: Support dynamic negative pad values (#24084)

For both cases, negative values on both sides mean that the pad is just
turned into a slice.

When just one side is negative:

- Reflective and replicate mode:
The original border is reflected/replicated, then the negative side is
sliced. This means values that are sliced away can appear in the padded
border.

- Circluar mode:
The negative side is sliced, then the resulting slice is rotated.

Add regression cases for cropping on either side, mixed signs, shifts,
and multiple dimensions. Cover reflection and replication where cropping
removes the original dimension but leaves padded elements.

This commit was authored with Codex.


cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/arm/_passes/rewrite_pad.py

@@ -141,7 +141,24 @@ def _rewrite_non_constant_pad(
                 )
             # 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)
+            # Reflect and replicate use the original borders; circular wraps
+            # the cropped input.
+            if mode != "circular":
+                output = self._pad_along_dim(
+                    output, dim, max(left, 0), max(right, 0), mode, meta
+                )
+            if left < 0 or right < 0:
+                output = super().call_operator(
+                    exir_ops.edge.aten.slice_copy.Tensor,
+                    (output, dim, max(-left, 0), right if right < 0 else None),
+                    {},
+                    meta,
+                    True,
+                )
+            if mode == "circular":
+                output = self._pad_along_dim(
+                    output, dim, max(left, 0), max(right, 0), mode, meta
+                )
         return output
 
     def call_operator(self, op, args, kwargs, meta, updated=False):

backends/arm/test/ops/test_constant_pad_nd.py

@@ -100,6 +100,38 @@
         None,
         "circular",
     ),
+    **{
+        f"{mode}_{case}": lambda mode=mode, shape=shape, pad=pad: (
+            torch.rand(shape),
+            pad,
+            None,
+            mode,
+        )
+        for mode in ("reflect", "replicate", "circular")
+        for case, (shape, pad) in {
+            "crop_end": ((1, 3, 6), (0, -2)),
+            "crop_start": ((1, 3, 6), (-2, 0)),
+            "crop_both": ((1, 3, 6), (-1, -2)),
+            "crop_and_pad": ((1, 3, 6), (-2, 3)),
+            "pad_and_crop": ((1, 3, 6), (3, -2)),
+            "crop_shift": ((1, 3, 6), (-1, 1)),
+            "crop_multidim": ((1, 3, 6, 6), (-1, 2, 3, -2)),
+            "crop_penultimate_dim": ((1, 3, 6, 6), (0, 0, -2, 1)),
+        }.items()
+    },
+    **{
+        f"{mode}_{case}": lambda mode=mode, pad=pad: (
+            torch.rand(1, 3, 6),
+            pad,
+            None,
+            mode,
+        )
+        for mode in ("reflect", "replicate")
+        for case, pad in {
+            "crop_all_left": (-6, 2),
+            "crop_all_right": (2, -6),
+        }.items()
+    },
 }
 
 test_data_suite_bf16 = {
@@ -159,10 +191,8 @@ def __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:
-                nonzero_idx = i
-                break
+        while nonzero_idx >= 2 and pad[nonzero_idx - 2 : nonzero_idx] == (0, 0):
+            nonzero_idx -= 2
         self.pad = pad[:nonzero_idx]
 
     def forward(self, x: torch.Tensor):