From f53f91c38ccff62934e2ad7907be2d19aa364a2f Mon Sep 17 00:00:00 2001 From: Per Held Date: Mon, 17 Aug 2026 12:44:27 +0200 Subject: [PATCH] Arm backend: Add bounded U55 unfold_copy support Lower static unfold_copy operations with at most 16 windows to slices, transposes, and concatenation. Keep BOOL, larger, and symbolic cases on portable fallback. Authored with Codex. Signed-off-by: Per Held Change-Id: Ie17547fe3c0c09207ae2442d601ea41c2c870516 --- backends/arm/_passes/arm_pass_manager.py | 2 +- .../decompose_unfold_to_gather_pass.py | 45 +++++++++++++ .../arm/operator_support/ethos_u55_support.py | 45 ++++++++++++- .../tosa_supported_operators.py | 2 + backends/arm/test/ops/test_unfold_copy.py | 67 ++++++++++++++++++- 5 files changed, 157 insertions(+), 4 deletions(-) diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index a0540d4a639..0e206b0b528 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/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), diff --git a/backends/arm/_passes/decompose_unfold_to_gather_pass.py b/backends/arm/_passes/decompose_unfold_to_gather_pass.py index 950290b3b83..d0fc8db2f77 100644 --- a/backends/arm/_passes/decompose_unfold_to_gather_pass.py +++ b/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, diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py index da1c5d2036c..35ce0e51df9 100644 --- a/backends/arm/operator_support/ethos_u55_support.py +++ b/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. diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 22bc948e7fa..c1fa7015623 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/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)) diff --git a/backends/arm/test/ops/test_unfold_copy.py b/backends/arm/test/ops/test_unfold_copy.py index 4b235ffa720..5e6b6aa612e 100644 --- a/backends/arm/test/ops/test_unfold_copy.py +++ b/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,