Skip to content

Commit 8b48cc6

Browse files
committed
Arm backend: Delegate nearest upsample on U55
Allow batch-size-1 nearest-neighbor upsampling with equal 2x, 4x, or 8x spatial scales to lower through TOSA RESIZE on Ethos-U55. Validate provided scale factors directly so fractional factors whose rounded output shapes resemble supported ratios continue to fall back to CPU. Exercise supported scales through Upsample, interpolate, and UpsamplingNearest2d using the Ethos-U55 pipeline. Add fallback coverage for batch sizes greater than one and for scale factor 2.25, whose rounded 2x2-to-4x4 output previously appeared to be a supported 2x resize. Authored with Codex. Signed-off-by: Per Held <per.held@arm.com> Change-Id: I37a8680d9b5bfe03e8d858f6a34b0afde81b9376
1 parent ddc96d0 commit 8b48cc6

3 files changed

Lines changed: 190 additions & 21 deletions

File tree

‎backends/arm/operator_support/ethos_u55_support.py‎

Lines changed: 91 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,6 @@ class EthosU55NotSupported(OperatorSupportBase):
209209
exir_ops.edge.aten.scatter_reduce.two,
210210
exir_ops.edge.aten.scatter_add.default,
211211
exir_ops.edge.aten.unfold_copy.default, # GATHER
212-
exir_ops.edge.aten.upsample_nearest2d.vec, # RESIZE
213212
exir_ops.edge.aten.upsample_bilinear2d.vec, # RESIZE
214213
exir_ops.edge.aten.reflection_pad1d.default, # REVERSE
215214
exir_ops.edge.aten.reflection_pad2d.default, # REVERSE
@@ -246,6 +245,97 @@ def is_node_supported(
246245
return True
247246

248247

248+
class EthosU55ResizeCheck(OperatorSupportBase):
249+
"""Accept nearest-neighbor upscales supported by Ethos-U55.
250+
251+
Ethos-U55 supports nearest-neighbor TOSA RESIZE when both spatial dimensions
252+
have batch size 1 and use the same power-of-two scale in the range 2x
253+
through 8x.
254+
255+
"""
256+
257+
def __init__(self, reporter: WhyNoPartitionReporter):
258+
"""Initialize the check with a reporter.
259+
260+
Args:
261+
reporter (WhyNoPartitionReporter): Reporter for rejection reasons.
262+
263+
"""
264+
self.reporter = reporter
265+
266+
def is_node_supported(
267+
self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
268+
) -> bool:
269+
"""Return True when a resize satisfies the U55 constraints.
270+
271+
Args:
272+
submodules (typing.Mapping[str, torch.nn.Module]): Exported modules.
273+
node (fx.Node): FX node to check.
274+
275+
Returns:
276+
bool: True if supported; otherwise, False.
277+
278+
"""
279+
if node.target != exir_ops.edge.aten.upsample_nearest2d.vec:
280+
return True
281+
282+
input_shape = get_first_fake_tensor(node.all_input_nodes[0]).shape
283+
output_shape = get_first_fake_tensor(node).shape
284+
input_batch = input_shape[0]
285+
output_batch = output_shape[0]
286+
if isinstance(input_batch, torch.SymInt) or isinstance(
287+
output_batch, torch.SymInt
288+
):
289+
self.reporter.report_reject(
290+
node,
291+
"U55 nearest-neighbor resize requires a static batch size of 1.",
292+
)
293+
return False
294+
if input_batch != 1 or output_batch != 1:
295+
self.reporter.report_reject(
296+
node, "U55 nearest-neighbor resize requires batch size 1."
297+
)
298+
return False
299+
300+
scale_factors_arg = node.args[2]
301+
if scale_factors_arg is not None:
302+
scale_factors = typing.cast(typing.Sequence[float], scale_factors_arg)
303+
if len(scale_factors) != 2 or not (
304+
scale_factors[0] == scale_factors[1] and scale_factors[0] in (2, 4, 8)
305+
):
306+
self.reporter.report_reject(
307+
node,
308+
"U55 nearest-neighbor resize requires equal 2x, 4x, or 8x "
309+
"scale factors.",
310+
)
311+
return False
312+
return True
313+
314+
input_height, input_width = input_shape[-2:]
315+
output_height, output_width = output_shape[-2:]
316+
if any(
317+
isinstance(dim, torch.SymInt)
318+
for dim in (input_height, input_width, output_height, output_width)
319+
):
320+
self.reporter.report_reject(
321+
node,
322+
"U55 nearest-neighbor resize with an explicit size requires "
323+
"static spatial dimensions.",
324+
)
325+
return False
326+
if any(
327+
output_height == input_height * scale
328+
and output_width == input_width * scale
329+
for scale in (2, 4, 8)
330+
):
331+
return True
332+
333+
self.reporter.report_reject(
334+
node, "U55 nearest-neighbor resize requires a 2x, 4x, or 8x upscale."
335+
)
336+
return False
337+
338+
249339
class EthosU55CastCheck(OperatorSupportBase):
250340
"""Reject unsupported casts on U55.
251341

‎backends/arm/operator_support/tosa_supported_operators.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
EthosU55CastCheck,
4141
EthosU55DtypeSupport,
4242
EthosU55NotSupported,
43+
EthosU55ResizeCheck,
4344
)
4445
from executorch.backends.arm.operator_support.tosa_profile_supported_op_lists import (
4546
TOSA_PRO_FP_SupportList,
@@ -363,6 +364,7 @@ def _negative_checks(
363364

364365
if tosa_spec.is_U55_subset:
365366
checks.append(EthosU55NotSupported(reporter))
367+
checks.append(EthosU55ResizeCheck(reporter))
366368
checks.append(EthosU55DtypeSupport(reporter))
367369
checks.append(EthosU55CastCheck(reporter))
368370

‎backends/arm/test/ops/test_upsample_nearest2d.py‎

Lines changed: 97 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,24 @@
66
from typing import Optional, Tuple
77

88
import torch
9+
from executorch.backends.arm.operator_support.ethos_u55_support import (
10+
EthosU55ResizeCheck,
11+
)
912
from executorch.backends.arm.quantizer.arm_quantizer import (
1013
get_symmetric_a16w8_quantization_config,
1114
)
1215
from executorch.backends.arm.test import common
16+
from executorch.backends.arm.test.tester.arm_tester import ArmTester
1317

1418
from executorch.backends.arm.test.tester.test_pipeline import (
19+
EthosU55PipelineINT,
1520
OpNotSupportedPipeline,
1621
TosaPipelineFP,
1722
TosaPipelineINT,
1823
VgfPipeline,
1924
)
25+
from executorch.exir.backend.utils import WhyNoPartitionReporter
26+
from executorch.exir.dialects._ops import ops as exir_ops
2027

2128
aten_op = "torch.ops.aten.upsample_nearest2d.vec"
2229
exir_op = "executorch_exir_dialects_edge__ops_aten_upsample_nearest2d_vec"
@@ -75,7 +82,15 @@
7582
}
7683

7784
test_data_u55 = {
78-
"rand_double_size": lambda: (torch.rand(2, 4, 8, 3), (16, 6), None, True),
85+
"rand_double_scale": lambda: (torch.rand(1, 4, 8, 3), None, 2.0, True),
86+
"rand_double_scale_different_shape": lambda: (
87+
torch.rand(1, 2, 4, 3),
88+
None,
89+
2.0,
90+
True,
91+
),
92+
"rand_quadruple_size": lambda: (torch.rand(1, 4, 8, 3), (32, 12), None, True),
93+
"rand_octuple_size": lambda: (torch.rand(1, 4, 8, 3), (64, 24), None, True),
7994
}
8095

8196
test_data_suite_dynamic = {
@@ -447,58 +462,120 @@ def test_upsample_nearest2d_vec_vgf_quant_a16w8_interpolate(
447462

448463
@common.parametrize("test_data", test_data_u55)
449464
@common.XfailIfNoCorstone300
450-
def test_upsample_nearest2d_vec_u55_INT_Upsample_not_delegated(
451-
test_data: torch.Tensor,
452-
):
465+
def test_upsample_nearest2d_vec_u55_INT_Upsample(test_data: torch.Tensor):
453466
test_data, size, scale_factor, compare_outputs = test_data()
454-
pipeline = OpNotSupportedPipeline[input_t1](
467+
pipeline = EthosU55PipelineINT[input_t1](
455468
Upsample(size, scale_factor),
456469
(test_data,),
457-
{exir_op: 1},
458-
n_expected_delegates=0,
459-
quantize=True,
460-
u55_subset=True,
470+
aten_op,
471+
exir_op,
461472
)
473+
pipeline.run()
462474

475+
476+
@common.parametrize("test_data", test_data_u55)
477+
@common.XfailIfNoCorstone300
478+
def test_upsample_nearest2d_vec_u55_INT_Interpolate(test_data: torch.Tensor):
479+
test_data, size, scale_factor, compare_outputs = test_data()
480+
pipeline = EthosU55PipelineINT[input_t1](
481+
Interpolate(size, scale_factor),
482+
(test_data,),
483+
aten_op,
484+
exir_op,
485+
)
463486
pipeline.run()
464487

465488

466489
@common.parametrize("test_data", test_data_u55)
467490
@common.XfailIfNoCorstone300
468-
def test_upsample_nearest2d_vec_u55_INT_Interpolate_not_delegated(
491+
def test_upsample_nearest2d_vec_u55_INT_UpsamplingNearest2d(
469492
test_data: torch.Tensor,
470493
):
471494
test_data, size, scale_factor, compare_outputs = test_data()
495+
pipeline = EthosU55PipelineINT[input_t1](
496+
UpsamplingNearest2d(size, scale_factor),
497+
(test_data,),
498+
aten_op,
499+
exir_op,
500+
)
501+
pipeline.run()
502+
503+
504+
def test_upsample_nearest2d_vec_u55_INT_unsupported_scale_not_delegated():
505+
# 2.25 rounds a 2x2 input to 4x4, which looks like a supported 2x resize from shapes alone.
506+
test_data = torch.rand(1, 4, 2, 2)
472507
pipeline = OpNotSupportedPipeline[input_t1](
473-
Interpolate(size, scale_factor),
508+
Interpolate(None, 2.25),
474509
(test_data,),
475510
{exir_op: 1},
476511
n_expected_delegates=0,
477512
quantize=True,
478513
u55_subset=True,
479514
)
480-
481515
pipeline.run()
482516

483517

484-
@common.parametrize("test_data", test_data_u55)
485-
@common.XfailIfNoCorstone300
486-
def test_upsample_nearest2d_vec_u55_INT_UpsamplingBilinear2d_not_delegated(
487-
test_data: torch.Tensor,
488-
):
489-
test_data, size, scale_factor, compare_outputs = test_data()
518+
def test_upsample_nearest2d_vec_u55_INT_unsupported_batch_not_delegated():
519+
test_data = torch.rand(2, 4, 8, 3)
490520
pipeline = OpNotSupportedPipeline[input_t1](
491-
UpsamplingNearest2d(size, scale_factor),
521+
Interpolate(None, 2.0),
492522
(test_data,),
493523
{exir_op: 1},
494524
n_expected_delegates=0,
495525
quantize=True,
496526
u55_subset=True,
497527
)
498-
499528
pipeline.run()
500529

501530

531+
def test_upsample_nearest2d_vec_u55_INT_dynamic_batch_not_delegated():
532+
test_data = torch.rand(2, 4, 8, 3)
533+
batch = torch.export.Dim("batch", min=1, max=4)
534+
tester = ArmTester(
535+
Interpolate(None, 2.0),
536+
(test_data,),
537+
common.get_u55_compile_spec(),
538+
dynamic_shapes={"x": {0: batch}},
539+
)
540+
tester.quantize().export().to_edge().partition()
541+
542+
targets = {
543+
node.target
544+
for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
545+
}
546+
assert exir_ops.edge.aten.upsample_nearest2d.vec in targets
547+
assert torch.ops.higher_order.executorch_call_delegate not in targets
548+
549+
550+
def test_upsample_nearest2d_vec_u55_INT_dynamic_spatial_size_not_delegated():
551+
test_data = torch.rand(1, 4, 8, 3)
552+
input_height = torch.export.Dim("input_height", min=4, max=12)
553+
input_width = torch.export.Dim("input_width", min=2, max=6)
554+
tester = ArmTester(
555+
Interpolate((16, 6), None),
556+
(test_data,),
557+
common.get_u55_compile_spec(),
558+
dynamic_shapes={"x": {2: input_height, 3: input_width}},
559+
)
560+
tester.quantize().export().to_edge()
561+
resize_node = next(
562+
node
563+
for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
564+
if node.target == exir_ops.edge.aten.upsample_nearest2d.vec
565+
)
566+
assert not EthosU55ResizeCheck(WhyNoPartitionReporter()).is_node_supported(
567+
{}, resize_node
568+
)
569+
570+
tester.partition()
571+
targets = {
572+
node.target
573+
for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes
574+
}
575+
assert exir_ops.edge.aten.upsample_nearest2d.vec in targets
576+
assert torch.ops.higher_order.executorch_call_delegate not in targets
577+
578+
502579
@common.parametrize("test_data", test_data_suite_dynamic)
503580
def test_upsample_nearest2d_vec_tosa_FP_dynamic_nearest(test_data: torch.Tensor):
504581
test_data, size, scale_factor, compare_outputs = test_data()

0 commit comments

Comments
 (0)