Arm backend: Delegate reflection padding on U55

Enable legal static ReflectionPad1d, ReflectionPad2d, and
ReflectionPad3d cases proven through Vela and Corstone-300.

Reflection padding mirrors edge values when extending a tensor:

    Original:          [A B C D]
    Pad 2 each side:   [C B | A B C D | C B]

Cover ranks, padding shapes, batching, and quantization modes, and
leave dynamic shapes undelegated when padding validity cannot be
proven at compile time.

Authored with Codex.

Change-Id: Ib0b833c2c866b13e504920671303ee22ceac01b6
Signed-off-by: Per Held <per.held@arm.com>

Author: perheld

backends/arm/operator_support/ethos_u55_support.py

@@ -208,9 +208,6 @@ class EthosU55NotSupported(OperatorSupportBase):
         exir_ops.edge.aten.scatter_reduce.two,
         exir_ops.edge.aten.scatter_add.default,
         exir_ops.edge.aten.upsample_bilinear2d.vec,  # RESIZE
-        exir_ops.edge.aten.reflection_pad1d.default,  # REVERSE
-        exir_ops.edge.aten.reflection_pad2d.default,  # REVERSE
-        exir_ops.edge.aten.reflection_pad3d.default,  # REVERSE
         exir_ops.edge.aten.where.self,  # SELECT
     ]
 
@@ -355,6 +352,35 @@ def is_node_supported(
             )
             return False
 
+        reflection_pad_constraints = {
+            exir_ops.edge.aten.reflection_pad1d.default: ((2, 3), (2,)),
+            exir_ops.edge.aten.reflection_pad2d.default: ((3, 4), (2, 4)),
+            exir_ops.edge.aten.reflection_pad3d.default: ((4, 5), (6,)),
+        }
+        if node.target in reflection_pad_constraints:
+            input_shape = get_first_fake_tensor(node.all_input_nodes[0]).shape
+            padding = typing.cast(typing.Sequence[int], node.args[1])
+            supported_ranks, supported_padding_lengths = reflection_pad_constraints[
+                node.target
+            ]
+            if (
+                len(input_shape) in supported_ranks
+                and len(padding) in supported_padding_lengths
+            ):
+                spatial_sizes = tuple(reversed(input_shape[-(len(padding) // 2) :]))
+                pad_pairs = tuple(zip(padding[::2], padding[1::2]))
+                if all(
+                    isinstance(size, int) and 0 <= before < size and 0 <= after < size
+                    for (before, after), size in zip(pad_pairs, spatial_sizes)
+                ):
+                    return True
+            self.reporter.report_reject(
+                node,
+                "U55 reflection padding requires a supported static input rank "
+                "and nonnegative padding smaller than its spatial dimension.",
+            )
+            return False
+
         return True
 
 

backends/arm/operator_support/tosa_supported_operators.py

@@ -762,6 +762,15 @@ def is_node_supported(
             input_node = node.all_input_nodes[0]
             input_quantized = FuseQuantizedActivationPass._is_fuseable_input(input_node)
 
+        if any(
+            isinstance(input_node.meta["val"], torch.SymInt)
+            for input_node in node.all_input_nodes
+        ):
+            self.reporter.report_reject(
+                node, "Symbolic scalar inputs cannot be delegated."
+            )
+            return False
+
         input_quantized = input_quantized or all(
             (input_node.target in DQ_OPS)
             or _is_integer_dtype(get_first_fake_tensor(input_node).dtype)

backends/arm/test/ops/test_reflection_pad1d.py

@@ -0,0 +1,75 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT
+
+
+input_t1 = Tuple[torch.Tensor]
+
+test_data_suite_u55 = {
+    "rank2_symmetric": lambda: (torch.rand(2, 5), (1, 1)),
+    "rank3_symmetric": lambda: (torch.rand(1, 2, 5), (1, 1)),
+    "asymmetric": lambda: (torch.rand(1, 2, 5), (1, 3)),
+    "maximum_legal": lambda: (torch.rand(1, 2, 5), (4, 4)),
+    "batched": lambda: (torch.rand(2, 2, 5), (1, 1)),
+}
+
+
+class ReflectionPad1d(torch.nn.Module):
+    def __init__(self, padding):
+        super().__init__()
+        self.padding = padding
+
+    def forward(self, x):
+        return torch.nn.functional.pad(x, self.padding, mode="reflect")
+
+
+@common.parametrize("test_data", test_data_suite_u55)
+@common.XfailIfNoCorstone300
+def test_reflection_pad1d_u55_INT(test_data):
+    data, padding = test_data()
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad1d(padding),
+        (data,),
+        aten_ops=[],
+        exir_ops=[],
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+def test_reflection_pad1d_u55_INT_a16w8():
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad1d((1, 1)),
+        (torch.rand(1, 2, 5),),
+        aten_ops=[],
+        exir_ops=[],
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+def test_reflection_pad1d_u55_INT_symbolic_width_not_delegated():
+    width = torch.export.Dim("width", min=3, max=8)
+    tester = ArmTester(
+        ReflectionPad1d((1, 1)),
+        (torch.rand(1, 2, 5),),
+        common.get_u55_compile_spec(),
+        dynamic_shapes={"x": {2: width}},
+    )
+    tester.quantize().export().to_edge().partition()
+
+    targets = {
+        node.target
+        for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
+    }
+    assert torch.ops.aten.scalar_tensor.default in targets
+    assert torch.ops.higher_order.executorch_call_delegate not in targets

backends/arm/test/ops/test_reflection_pad2d.py

@@ -0,0 +1,87 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT
+
+
+input_t1 = Tuple[torch.Tensor]
+
+test_data_suite_u55 = {
+    "both_axes": lambda: (torch.rand(1, 2, 4, 3), (1, 1, 1, 1)),
+    "width_only": lambda: (torch.rand(1, 2, 4, 3), (1, 1, 0, 0)),
+    "height_only": lambda: (torch.rand(1, 2, 4, 3), (0, 0, 1, 1)),
+    "asymmetric": lambda: (torch.rand(1, 2, 5, 4), (1, 2, 3, 1)),
+    "maximum_legal": lambda: (torch.rand(1, 2, 4, 3), (2, 2, 3, 3)),
+    "batched": lambda: (torch.rand(2, 2, 4, 3), (1, 1, 1, 1)),
+}
+
+
+class ReflectionPad2d(torch.nn.Module):
+    def __init__(self, padding):
+        super().__init__()
+        self.padding = padding
+
+    def forward(self, x):
+        return torch.nn.functional.pad(x, self.padding, mode="reflect")
+
+
+@common.parametrize("test_data", test_data_suite_u55)
+@common.XfailIfNoCorstone300
+def test_reflection_pad2d_u55_INT(test_data):
+    data, padding = test_data()
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad2d(padding),
+        (data,),
+        aten_ops=[],
+        exir_ops=[],
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+def test_reflection_pad2d_u55_INT_a16w8():
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad2d((1, 1, 1, 1)),
+        (torch.rand(1, 2, 4, 3),),
+        aten_ops=[],
+        exir_ops=[],
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+def test_reflection_pad2d_u55_INT_rank3():
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad2d((1, 1, 1, 1)),
+        (torch.rand(2, 4, 3),),
+        aten_ops=[],
+        exir_ops=[],
+    )
+    pipeline.run()
+
+
+def test_reflection_pad2d_u55_INT_symbolic_width_not_delegated():
+    width = torch.export.Dim("width", min=3, max=8)
+    tester = ArmTester(
+        ReflectionPad2d((1, 1, 1, 1)),
+        (torch.rand(1, 2, 4, 5),),
+        common.get_u55_compile_spec(),
+        dynamic_shapes={"x": {3: width}},
+    )
+    tester.quantize().export().to_edge().partition()
+
+    targets = {
+        node.target
+        for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
+    }
+    assert torch.ops.aten.scalar_tensor.default in targets
+    assert torch.ops.higher_order.executorch_call_delegate not in targets

backends/arm/test/ops/test_reflection_pad3d.py

@@ -0,0 +1,75 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT
+
+
+input_t1 = Tuple[torch.Tensor]
+
+test_data_suite_u55 = {
+    "rank4_symmetric": lambda: (torch.rand(2, 4, 4, 4), (1, 1, 1, 1, 1, 1)),
+    "rank5_symmetric": lambda: (torch.rand(1, 2, 4, 4, 4), (1, 1, 1, 1, 1, 1)),
+    "asymmetric": lambda: (torch.rand(1, 2, 5, 5, 5), (1, 2, 2, 1, 3, 1)),
+    "maximum_legal": lambda: (torch.rand(1, 2, 4, 4, 4), (3, 3, 3, 3, 3, 3)),
+    "batched": lambda: (torch.rand(2, 2, 4, 4, 4), (1, 1, 1, 1, 1, 1)),
+}
+
+
+class ReflectionPad3d(torch.nn.Module):
+    def __init__(self, padding):
+        super().__init__()
+        self.padding = padding
+
+    def forward(self, x):
+        return torch.nn.functional.pad(x, self.padding, mode="reflect")
+
+
+@common.parametrize("test_data", test_data_suite_u55)
+@common.XfailIfNoCorstone300
+def test_reflection_pad3d_u55_INT(test_data):
+    data, padding = test_data()
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad3d(padding),
+        (data,),
+        aten_ops=[],
+        exir_ops=[],
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+def test_reflection_pad3d_u55_INT_a16w8():
+    pipeline = EthosU55PipelineINT[input_t1](
+        ReflectionPad3d((1, 1, 1, 1, 1, 1)),
+        (torch.rand(1, 2, 4, 4, 4),),
+        aten_ops=[],
+        exir_ops=[],
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+def test_reflection_pad3d_u55_INT_symbolic_width_not_delegated():
+    width = torch.export.Dim("width", min=3, max=8)
+    tester = ArmTester(
+        ReflectionPad3d((1, 1, 1, 1, 1, 1)),
+        (torch.rand(1, 2, 4, 4, 5),),
+        common.get_u55_compile_spec(),
+        dynamic_shapes={"x": {4: width}},
+    )
+    tester.quantize().export().to_edge().partition()
+
+    targets = {
+        node.target
+        for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
+    }
+    assert torch.ops.aten.scalar_tensor.default in targets
+    assert torch.ops.higher_order.executorch_call_delegate not in targets