Arm backend: Add bounded U55 unfold_copy support

backends/arm/_passes/arm_pass_manager.py

@@ -584,7 +584,7 @@ def _tosa_pipeline(
                 DecomposeGeluPass(),
                 DecomposeAddSubAlphaPass(),
                 DecomposeGroupedConvPass(),
-                DecomposeUnfoldToGatherPass(),
+                DecomposeUnfoldToGatherPass(use_slice=self.tosa_spec.is_U55_subset),
                 DecomposeEmbeddingPass(),
                 DecomposeIndexSelectToGatherPass(),
                 CastInt64BuffersToInt32Pass(exported_program),

backends/arm/_passes/decompose_unfold_to_gather_pass.py

@@ -97,6 +97,10 @@ class DecomposeUnfoldToGatherPass(ArmOpTargetedPass):
         exir_ops.edge.aten.unfold_copy.default,
     }
 
+    def __init__(self, use_slice: bool = False, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.use_slice = use_slice
+
     _UnfoldCopyInfo = tuple[
         torch.Tensor,  # x_val (FakeTensor)
         int,  # C
@@ -166,6 +170,47 @@ def call_operator(self, op, args, kwargs, meta):
             needs_bool_cast,
         ) = self._compute_unfold_copy_params(x, dim, size, step)
 
+        if self.use_slice:
+            rank = len(x_val.shape)
+            dim_norm = dim % rank
+            perm = list(range(dim_norm)) + list(range(dim_norm + 1, rank)) + [dim_norm]
+            windows = []
+            for window in range(U):
+                start = window * S
+                sliced = super().call_operator(
+                    exir_ops.edge.aten.slice_copy.Tensor,
+                    (x, dim_norm, start, start + C),
+                    {},
+                    meta,
+                    updated=True,
+                )
+                transposed = super().call_operator(
+                    exir_ops.edge.aten.permute_copy.default,
+                    (sliced, perm),
+                    {},
+                    meta,
+                    updated=True,
+                )
+                windows.append(
+                    super().call_operator(
+                        exir_ops.edge.aten.unsqueeze_copy.default,
+                        (transposed, dim_norm),
+                        {},
+                        meta,
+                        updated=True,
+                    )
+                )
+
+            if len(windows) == 1:
+                return windows[0]
+            return super().call_operator(
+                exir_ops.edge.aten.cat.default,
+                (windows, dim_norm),
+                {},
+                meta,
+                updated=True,
+            )
+
         (
             to_copy_op,
             view_op,

backends/arm/operator_support/ethos_u55_support.py

@@ -207,7 +207,6 @@ class EthosU55NotSupported(OperatorSupportBase):
         exir_ops.edge.aten.select_scatter.default,
         exir_ops.edge.aten.scatter_reduce.two,
         exir_ops.edge.aten.scatter_add.default,
-        exir_ops.edge.aten.unfold_copy.default,  # GATHER
         exir_ops.edge.aten.upsample_bilinear2d.vec,  # RESIZE
         exir_ops.edge.aten.reflection_pad1d.default,  # REVERSE
         exir_ops.edge.aten.reflection_pad2d.default,  # REVERSE
@@ -359,6 +358,50 @@ def is_node_supported(
         return True
 
 
+class EthosU55UnfoldCopyCheck(OperatorSupportBase):
+    """Accept bounded static unfold_copy cases that lower to slices."""
+
+    max_windows = 16
+
+    def __init__(self, reporter: WhyNoPartitionReporter):
+        self.reporter = reporter
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        del submodules
+        if node.target != exir_ops.edge.aten.unfold_copy.default:
+            return True
+
+        input_arg, dim, size, step = node.args
+        input_node = typing.cast(fx.Node, input_arg)
+        input_tensor = get_first_fake_tensor(input_node)
+        input_shape = input_tensor.shape
+        if (
+            input_tensor.dtype == torch.bool
+            or not all(isinstance(arg, int) for arg in (dim, size, step))
+            or any(not isinstance(value, int) for value in input_shape)
+        ):
+            self.reporter.report_reject(
+                node, "U55 unfold_copy requires static non-BOOL input."
+            )
+            return False
+
+        rank = len(input_shape)
+        dim = typing.cast(int, dim) % rank
+        size = typing.cast(int, size)
+        step = typing.cast(int, step)
+        windows = (input_shape[dim] - size) // step + 1
+        if windows > self.max_windows:
+            self.reporter.report_reject(
+                node,
+                f"U55 unfold_copy supports at most {self.max_windows} windows.",
+            )
+            return False
+
+        return True
+
+
 class EthosU55CastCheck(OperatorSupportBase):
     """Reject unsupported casts on U55.
 

backends/arm/operator_support/tosa_supported_operators.py

@@ -42,6 +42,7 @@
     EthosU55NotSupported,
     EthosU55ResizeCheck,
     EthosU55ReverseCheck,
+    EthosU55UnfoldCopyCheck,
 )
 from executorch.backends.arm.operator_support.tosa_profile_supported_op_lists import (
     TOSA_PRO_FP_SupportList,
@@ -411,6 +412,7 @@ def _negative_checks(
         checks.append(EthosU55NotSupported(reporter))
         checks.append(EthosU55ResizeCheck(reporter))
         checks.append(EthosU55ReverseCheck(reporter))
+        checks.append(EthosU55UnfoldCopyCheck(reporter))
         checks.append(EthosU55DtypeSupport(reporter))
         checks.append(EthosU55CastCheck(reporter))
 

backends/arm/test/ops/test_unfold_copy.py

@@ -8,13 +8,16 @@
 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,
     EthosU85PipelineINT,
     OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
 )
+from executorch.exir.dialects._ops import ops as exir_ops
 
 
 class UnfoldCopy(torch.nn.Module):
@@ -206,10 +209,47 @@ def test_unfold_copy_tosa_INT(test_data: input_params):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_int | test_data_fp)
+test_data_u55 = {
+    "rank1_dim0": test_data_fp["test_fp32_1d_dim0"],
+    "rank2_dim1": test_data_fp["test_fp32_2d_dim1"],
+    "rank2_max_windows": (torch.rand(1, 17), 1, 2, 1),
+    "rank3_dim1": test_data_int["test_int8_3d_dim1"],
+    "rank3_dim_neg1": test_data_fp["test_fp32_3d_dim_neg1"],
+}
+
+
+@common.parametrize("test_data", test_data_u55)
 @common.XfailIfNoCorstone300
 def test_unfold_copy_u55_INT(test_data: input_params):
-    # Gather op is not supported on U55
+    pipeline = EthosU55PipelineINT[input_params](
+        UnfoldCopy(),
+        test_data,
+        aten_ops=[],
+        exir_ops=[],
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+def test_unfold_copy_u55_INT_a16w8():
+    pipeline = EthosU55PipelineINT[input_params](
+        UnfoldCopy(),
+        test_data_fp["test_fp32_2d_dim1"],
+        aten_ops=[],
+        exir_ops=[],
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize(
+    "test_data",
+    {
+        "bool": test_data_int["test_bool_1d_dim_neg1"],
+        "too_many_windows": (torch.rand(1, 20), 1, 2, 1),
+    },
+)
+def test_unfold_copy_u55_INT_not_delegated(test_data: input_params):
     pipeline = OpNotSupportedPipeline[input_params](
         UnfoldCopy(),
         test_data,
@@ -221,6 +261,29 @@ def test_unfold_copy_u55_INT(test_data: input_params):
     pipeline.run()
 
 
+def test_unfold_copy_u55_INT_symbolic_dim_not_delegated():
+    length = torch.export.Dim("length", min=4, max=12)
+    tester = ArmTester(
+        UnfoldCopy(),
+        (torch.rand(2, 6), 1, 3, 1),
+        common.get_u55_compile_spec(),
+        dynamic_shapes={
+            "input_": {1: length},
+            "dim_": None,
+            "size_": None,
+            "step_": None,
+        },
+    )
+    tester.quantize().export().to_edge().partition()
+
+    targets = {
+        node.target
+        for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
+    }
+    assert exir_ops.edge.aten.unfold_copy.default in targets
+    assert torch.ops.higher_order.executorch_call_delegate not in targets
+
+
 @common.parametrize(
     "test_data",
     test_data_int | test_data_fp,