From ca472b604f31c5e79661f77b7c906e396e2e1723 Mon Sep 17 00:00:00 2001 From: Rob Elliott Date: Wed, 17 Jun 2026 12:02:40 +0100 Subject: [PATCH] Arm backend: Decompose exact 1/16 bilinear downscale Signed-off-by: Rob Elliott Change-Id: Ib620350be2cf92ff8f4e5406a5ae934f2dc9f2e8 Signed-off-by: Rob Elliott --- backends/arm/_passes/__init__.py | 3 + backends/arm/_passes/arm_pass_manager.py | 2 + ...ompose_unsupported_bilinear_resize_pass.py | 159 +++++++++++ backends/arm/ethosu/partitioner.py | 6 +- .../tosa_supported_operators.py | 5 + .../misc/tosa_dialect/test_tosa_resize.py | 31 +++ .../arm/test/ops/test_upsample_bilinear2d.py | 21 +- ...ompose_unsupported_bilinear_resize_pass.py | 249 ++++++++++++++++++ backends/arm/tosa/partitioner.py | 49 +++- backends/arm/tosa/resize_utils.py | 65 ++++- backends/arm/vgf/partitioner.py | 6 +- 11 files changed, 585 insertions(+), 11 deletions(-) create mode 100644 backends/arm/_passes/decompose_unsupported_bilinear_resize_pass.py create mode 100644 backends/arm/test/passes/test_decompose_unsupported_bilinear_resize_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index f50991287dc..2d4ff8b2baf 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -101,6 +101,9 @@ ) from .decompose_tril_pass import DecomposeTrilPass # noqa from .decompose_unfold_to_gather_pass import DecomposeUnfoldToGatherPass # noqa +from .decompose_unsupported_bilinear_resize_pass import ( # noqa + DecomposeUnsupportedBilinearResizePass, +) from .decompose_var_pass import DecomposeVarPass # noqa from .decompose_where_scalar_other_pass import DecomposeWhereScalarOtherPass # noqa from .decorate_fp32_to_int32_casting_pass import DecorateFp32toInt32CastingPass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 41d0cc53e17..0066d78d313 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -99,6 +99,7 @@ DecomposeTOSAUnsupportedClampPass, DecomposeTrilPass, DecomposeUnfoldToGatherPass, + DecomposeUnsupportedBilinearResizePass, DecomposeVarPass, DecomposeWhereScalarOtherPass, DecorateFp32toInt32CastingPass, @@ -600,6 +601,7 @@ def _tosa_pipeline( DecomposeAsStridedCopyPass(), DecomposeMaxPool2dPass(), SizeAdjustInputPass(), + DecomposeUnsupportedBilinearResizePass(self.tosa_spec), RewriteAdaptiveAvgPool2dPass(), RewriteAvgPool2dPass(), ComputeConstantOpsAOTPass(exported_program), diff --git a/backends/arm/_passes/decompose_unsupported_bilinear_resize_pass.py b/backends/arm/_passes/decompose_unsupported_bilinear_resize_pass.py new file mode 100644 index 00000000000..85f1d0784ca --- /dev/null +++ b/backends/arm/_passes/decompose_unsupported_bilinear_resize_pass.py @@ -0,0 +1,159 @@ +# 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 Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import ( + create_node, + get_first_fake_tensor, +) +from executorch.backends.arm._passes.rewrite_avg_pool2d_pass import RewriteAvgPool2dPass +from executorch.backends.arm._passes.rewrite_upsample import RewriteUpsamplePass +from executorch.backends.arm.common.type import ensure_type +from executorch.backends.arm.tosa.resize_utils import ( + is_exact_tosa_resize_boundary_downscale_case, +) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx.passes.shape_prop import _extract_tensor_metadata + +_EXACT_BOUNDARY_DOWNSCALE = 16 +_EQUIVALENT_AVG_POOL_KERNEL = 2 +_EQUIVALENT_AVG_POOL_OFFSET = _EXACT_BOUNDARY_DOWNSCALE // 2 - 1 + + +def is_exact_tosa_boundary_bilinear_downscale( + node: torch.fx.Node, + tosa_spec: TosaSpecification, +) -> bool: + if ( + node.op != "call_function" + or node.target != exir_ops.edge.aten.upsample_bilinear2d.vec + ): + return False + + input_node = ensure_type(torch.fx.Node, node.args[0]) + align_corners = ensure_type(bool, node.args[2]) + if align_corners: + return False + input_size_yx = get_first_fake_tensor(input_node).shape[2:] + output_size_yx = get_first_fake_tensor(node).shape[2:] + + scale_y_n, scale_y_d, offset_y, border_y = ( + RewriteUpsamplePass.get_resize_parameters_1d( + input_size_yx[0], output_size_yx[0], align_corners + ) + ) + scale_x_n, scale_x_d, offset_x, border_x = ( + RewriteUpsamplePass.get_resize_parameters_1d( + input_size_yx[1], output_size_yx[1], align_corners + ) + ) + return is_exact_tosa_resize_boundary_downscale_case( + input_hw=input_size_yx, + output_hw=output_size_yx, + scale=[scale_y_n, scale_y_d, scale_x_n, scale_x_d], + offset=[offset_y, offset_x], + border=[border_y, border_x], + tosa_spec=tosa_spec, + ) + + +class DecomposeUnsupportedBilinearResizePass(ArmPass): + targeted_ops = (exir_ops.edge.aten.upsample_bilinear2d.vec,) + _passes_required_after: Set[Type[ExportPass]] = {RewriteAvgPool2dPass} + + def __init__(self, tosa_spec: TosaSpecification) -> None: + super().__init__() + self.tosa_spec = tosa_spec + + @staticmethod + def _set_fake_tensor_meta(node: torch.fx.Node, value: torch.Tensor) -> None: + node.meta["val"] = value + node.meta["tensor_meta"] = _extract_tensor_metadata(value) + + def call(self, graph_module): + graph = graph_module.graph + modified = False + for node in list(graph.nodes): + if not is_exact_tosa_boundary_bilinear_downscale(node, self.tosa_spec): + continue + + input_node = ensure_type(torch.fx.Node, node.args[0]) + input_val = input_node.meta["val"] + input_hw = [int(dim) for dim in get_first_fake_tensor(input_node).shape[2:]] + + with graph.inserting_before(node): + cropped_h = create_node( + graph, + op_target=exir_ops.edge.aten.slice_copy.Tensor, + args=( + input_node, + 2, + _EQUIVALENT_AVG_POOL_OFFSET, + input_hw[0] - _EQUIVALENT_AVG_POOL_OFFSET, + 1, + ), + from_node=node, + ) + cropped_h_val = exir_ops.edge.aten.slice_copy.Tensor( + input_val, + 2, + _EQUIVALENT_AVG_POOL_OFFSET, + input_hw[0] - _EQUIVALENT_AVG_POOL_OFFSET, + 1, + ) + self._set_fake_tensor_meta(cropped_h, cropped_h_val) + + cropped_hw = create_node( + graph, + op_target=exir_ops.edge.aten.slice_copy.Tensor, + args=( + cropped_h, + 3, + _EQUIVALENT_AVG_POOL_OFFSET, + input_hw[1] - _EQUIVALENT_AVG_POOL_OFFSET, + 1, + ), + from_node=node, + ) + cropped_hw_val = exir_ops.edge.aten.slice_copy.Tensor( + cropped_h_val, + 3, + _EQUIVALENT_AVG_POOL_OFFSET, + input_hw[1] - _EQUIVALENT_AVG_POOL_OFFSET, + 1, + ) + self._set_fake_tensor_meta(cropped_hw, cropped_hw_val) + + pooled = create_node( + graph, + op_target=exir_ops.edge.aten.avg_pool2d.default, + args=( + cropped_hw, + [_EQUIVALENT_AVG_POOL_KERNEL, _EQUIVALENT_AVG_POOL_KERNEL], + [_EXACT_BOUNDARY_DOWNSCALE, _EXACT_BOUNDARY_DOWNSCALE], + ), + from_node=node, + ) + pooled_val = torch.nn.functional.avg_pool2d( + cropped_hw_val, + kernel_size=_EQUIVALENT_AVG_POOL_KERNEL, + stride=_EXACT_BOUNDARY_DOWNSCALE, + ) + self._set_fake_tensor_meta(pooled, pooled_val) + + node.replace_all_uses_with(pooled) + graph.erase_node(node) + modified = True + + if modified: + graph.lint() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index 63bab44dc8c..a36584a975a 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/backends/arm/ethosu/partitioner.py @@ -6,7 +6,10 @@ from typing import final, Optional, Sequence from executorch.backends.arm.ethosu import EthosUBackend, EthosUCompileSpec -from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.partitioner import ( + DecomposableResizeSupported, + TOSAPartitioner, +) from executorch.exir.backend.partitioner import DelegationSpec from torch._ops import OpOverload from torch.fx.passes.operator_support import OperatorSupportBase @@ -33,5 +36,6 @@ def __init__( ) self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec + self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index e5d251b3fd9..3580a47d834 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -423,6 +423,7 @@ def tosa_support_factory( exported_program: ExportedProgram, reporter: WhyNoPartitionReporter, additional_checks: Optional[Sequence[OperatorSupportBase]] = None, + additional_positive_checks: Optional[Sequence[OperatorSupportBase]] = None, ) -> OperatorSupportBase: """Create an OperatorSupport composite for a TOSA spec. @@ -435,12 +436,16 @@ def tosa_support_factory( reporter (WhyNoPartitionReporter): Reporter for rejections. additional_checks (Optional[Sequence[OperatorSupportBase]]): Extra negative checks to apply. + additional_positive_checks (Optional[Sequence[OperatorSupportBase]]): + Extra positive checks to add to the support list. Returns: OperatorSupportBase: Composite checker for the given spec. """ positive_checks = _positive_checks(tosa_spec, exported_program, reporter) + if additional_positive_checks: + positive_checks.extend(additional_positive_checks) negative_checks = _negative_checks( tosa_spec, exported_program, reporter, additional_checks ) diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_resize.py b/backends/arm/test/misc/tosa_dialect/test_tosa_resize.py index eddb69a8caf..baccee6a58a 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_resize.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_resize.py @@ -9,6 +9,9 @@ import torch from executorch.backends.arm.tosa.dialect.lib import TosaValueError +from executorch.backends.arm.tosa.resize_utils import ( + is_exact_tosa_resize_boundary_downscale_case, +) from executorch.backends.arm.tosa.specification import ( TosaLoweringContext, TosaSpecification, @@ -53,6 +56,34 @@ def test_resize_rejects_exact_one_sixteenth_downscale(resize_mode: str): ) +def test_exact_boundary_resize_case_helper_only_matches_boundary_rule(): + spec = TosaSpecification.create_from_string("TOSA-1.0+INT") + assert is_exact_tosa_resize_boundary_downscale_case( + input_hw=(256, 448), + output_hw=(16, 28), + scale=[2, 32, 2, 32], + offset=[15, 15], + border=[-15, -15], + tosa_spec=spec, + ) + assert not is_exact_tosa_resize_boundary_downscale_case( + input_hw=(20000, 448), + output_hw=(16, 28), + scale=[2, 32, 2, 32], + offset=[15, 15], + border=[-15, -15], + tosa_spec=spec, + ) + assert not is_exact_tosa_resize_boundary_downscale_case( + input_hw=(256, 448), + output_hw=(15, 28), + scale=[2, 32, 2, 32], + offset=[15, 15], + border=[-15, -15], + tosa_spec=spec, + ) + + def test_resize_rejects_scale_numerator_over_tosa_limit(): with TosaLoweringContext( TosaSpecification.create_from_string("TOSA-1.0+INT") diff --git a/backends/arm/test/ops/test_upsample_bilinear2d.py b/backends/arm/test/ops/test_upsample_bilinear2d.py index 1ebb38c2368..46e4f4c556e 100644 --- a/backends/arm/test/ops/test_upsample_bilinear2d.py +++ b/backends/arm/test/ops/test_upsample_bilinear2d.py @@ -305,14 +305,27 @@ def test_upsample_bilinear2d_vec_tosa_FP_Interpolate( pipeline.run() -def test_upsample_bilinear2d_vec_tosa_does_not_delegate_exact_one_sixteenth_downscale(): - pipeline = OpNotSupportedPipeline[input_t1]( +def test_upsample_bilinear2d_vec_tosa_delegates_exact_one_sixteenth_downscale(): + pipeline = TosaPipelineFP[input_t1]( InterpolateAlignCornersFalse(size=None, scale_factor=1.0 / 16.0), (torch.randn(1, 3, 256, 448),), - {exir_op: 1}, - n_expected_delegates=0, + aten_op, + exir_op=[], ) + pipeline.pop_stage(-1) + pipeline.run() + +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vec_vgf_delegates_exact_one_sixteenth_downscale(): + pipeline = VgfPipeline[input_t1]( + InterpolateAlignCornersFalse(size=None, scale_factor=1.0 / 16.0), + (torch.randn(1, 3, 256, 448),), + aten_op, + exir_op, + quantize=False, + ) + pipeline.pop_stage(-1) pipeline.run() diff --git a/backends/arm/test/passes/test_decompose_unsupported_bilinear_resize_pass.py b/backends/arm/test/passes/test_decompose_unsupported_bilinear_resize_pass.py new file mode 100644 index 00000000000..daab0458379 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_unsupported_bilinear_resize_pass.py @@ -0,0 +1,249 @@ +# 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. + +import pytest +import torch +import torch.nn.functional as F + +from executorch.backends.arm._passes import DecomposeUnsupportedBilinearResizePass +from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export + + +class ExactBoundaryBilinearDownscale(torch.nn.Module): + def forward(self, x): + return F.interpolate( + x, + scale_factor=1.0 / 16.0, + mode="bilinear", + align_corners=False, + ) + + +class LegalBilinearDownscale(torch.nn.Module): + def forward(self, x): + return F.interpolate( + x, + scale_factor=1.0 / 8.0, + mode="bilinear", + align_corners=False, + ) + + +class ExactBoundaryBilinearDownscaleByOutputSize(torch.nn.Module): + def forward(self, x): + return F.interpolate( + x, + size=(16, 16), + mode="bilinear", + align_corners=False, + ) + + +def _export_exact_boundary_bilinear_downscale(): + model = ExactBoundaryBilinearDownscale() + example_inputs = (torch.randn(1, 15, 256, 256),) + return to_edge(export(model, example_inputs)) + + +def _get_upsample_nodes(edge_model): + return [ + node + for node in edge_model.exported_program().graph.nodes + if node.target == exir_ops.edge.aten.upsample_bilinear2d.vec + ] + + +def _get_avg_pool_nodes(edge_model): + return [ + node + for node in edge_model.exported_program().graph.nodes + if node.target == exir_ops.edge.aten.avg_pool2d.default + ] + + +def _get_slice_nodes(edge_model): + return [ + node + for node in edge_model.exported_program().graph.nodes + if node.target == exir_ops.edge.aten.slice_copy.Tensor + ] + + +def _transform_with_resize_decomposition(edge_model): + return edge_model.transform( + [ + DecomposeUnsupportedBilinearResizePass( + TosaSpecification.create_from_string("TOSA-1.0+FP") + ) + ] + ) + + +def test_tosa_partitioner_accepts_exact_boundary_bilinear_downscale(): + edge_model = _export_exact_boundary_bilinear_downscale() + + partition_result = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")).partition( + edge_model.exported_program() + ) + upsample_node = next( + node + for node in partition_result.tagged_exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.upsample_bilinear2d.vec + ) + + assert upsample_node.meta["delegation_tag"].startswith("tag") + + +def test_vgf_partitioner_accepts_exact_boundary_bilinear_downscale(): + edge_model = _export_exact_boundary_bilinear_downscale() + + partition_result = VgfPartitioner(VgfCompileSpec("TOSA-1.0+FP")).partition( + edge_model.exported_program() + ) + upsample_node = next( + node + for node in partition_result.tagged_exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.upsample_bilinear2d.vec + ) + + assert upsample_node.meta["delegation_tag"].startswith("tag") + + +def test_decompose_exact_boundary_bilinear_downscale_into_exact_avg_pool(): + edge_model = _export_exact_boundary_bilinear_downscale() + + transformed = _transform_with_resize_decomposition(edge_model) + avg_pool_nodes = _get_avg_pool_nodes(transformed) + slice_nodes = _get_slice_nodes(transformed) + + assert len(_get_upsample_nodes(transformed)) == 0 + assert len(avg_pool_nodes) == 1 + assert len(slice_nodes) == 2 + assert slice_nodes[0].args[1:] == (2, 7, 249, 1) + assert slice_nodes[1].args[1:] == (3, 7, 249, 1) + assert avg_pool_nodes[0].args[1] == [2, 2] + assert avg_pool_nodes[0].args[2] == [16, 16] + assert list(get_first_fake_tensor(avg_pool_nodes[0]).shape[2:]) == [16, 16] + + +def test_decompose_exact_boundary_bilinear_downscale_by_output_size(): + model = ExactBoundaryBilinearDownscaleByOutputSize() + edge_model = to_edge(export(model, (torch.randn(1, 15, 256, 256),))) + + transformed = _transform_with_resize_decomposition(edge_model) + avg_pool_nodes = _get_avg_pool_nodes(transformed) + slice_nodes = _get_slice_nodes(transformed) + + assert len(_get_upsample_nodes(transformed)) == 0 + assert len(avg_pool_nodes) == 1 + assert len(slice_nodes) == 2 + assert avg_pool_nodes[0].args[1] == [2, 2] + assert avg_pool_nodes[0].args[2] == [16, 16] + + +def test_decompose_resize_pass_is_noop_for_legal_downscale(): + model = LegalBilinearDownscale() + edge_model = to_edge(export(model, (torch.randn(1, 15, 256, 256),))) + + transformed = _transform_with_resize_decomposition(edge_model) + upsample_nodes = _get_upsample_nodes(transformed) + + assert len(upsample_nodes) == 1 + assert upsample_nodes[0].args[1] is None + + +def _decomposed_boundary_bilinear(x: torch.Tensor) -> torch.Tensor: + return F.avg_pool2d( + x[:, :, 7:-7, 7:-7], + kernel_size=2, + stride=16, + ) + + +def _make_numeric_case(kind: str) -> torch.Tensor: + shape = (1, 3, 256, 448) + match kind: + case "zeros": + return torch.zeros(shape) + case "ones": + return torch.ones(shape) + case "ramp_x": + return ( + torch.linspace(0, 1, shape[-1]).view(1, 1, 1, shape[-1]).expand(shape) + ) + case "ramp_y": + return ( + torch.linspace(0, 1, shape[-2]).view(1, 1, shape[-2], 1).expand(shape) + ) + case "checkerboard": + x = torch.arange(shape[-1]).view(1, 1, 1, shape[-1]) + y = torch.arange(shape[-2]).view(1, 1, shape[-2], 1) + return ((x + y) % 2).float().expand(shape) + case "horizontal_edge": + x = torch.zeros(shape) + x[:, :, shape[-2] // 2 :, :] = 1.0 + return x + case "vertical_edge": + x = torch.zeros(shape) + x[:, :, :, shape[-1] // 2 :] = 1.0 + return x + case "corner_impulse_tl": + x = torch.zeros(shape) + x[:, :, 0, 0] = 1.0 + return x + case "corner_impulse_tr": + x = torch.zeros(shape) + x[:, :, 0, shape[-1] - 1] = 1.0 + return x + case "corner_impulse_bl": + x = torch.zeros(shape) + x[:, :, shape[-2] - 1, 0] = 1.0 + return x + case "corner_impulse_br": + x = torch.zeros(shape) + x[:, :, shape[-2] - 1, shape[-1] - 1] = 1.0 + return x + case "center_impulse": + x = torch.zeros(shape) + x[:, :, shape[-2] // 2, shape[-1] // 2] = 1.0 + return x + case _: + raise AssertionError(f"Unhandled numeric case: {kind}") + + +@pytest.mark.parametrize( + "case_name", + [ + "zeros", + "ones", + "ramp_x", + "ramp_y", + "checkerboard", + "horizontal_edge", + "vertical_edge", + "corner_impulse_tl", + "corner_impulse_tr", + "corner_impulse_bl", + "corner_impulse_br", + "center_impulse", + ], +) +def test_exact_boundary_bilinear_downscale_numeric_edge_cases(case_name: str): + x = _make_numeric_case(case_name) + direct = F.interpolate( + x, + scale_factor=1.0 / 16.0, + mode="bilinear", + align_corners=False, + ) + decomposed = _decomposed_boundary_bilinear(x) + torch.testing.assert_close(decomposed, direct, atol=1e-6, rtol=1e-6) diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 4dd144936e9..09836a2121e 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -18,13 +18,16 @@ from collections import deque from itertools import count from pathlib import Path -from typing import Callable, cast, List, Optional, Sequence, Tuple +from typing import Callable, cast, List, Mapping, Optional, Sequence, Tuple import torch from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( calculate_multiples, ) +from executorch.backends.arm._passes.decompose_unsupported_bilinear_resize_pass import ( + is_exact_tosa_boundary_bilinear_downscale, +) from executorch.backends.arm.common.type import ensure_type from executorch.backends.arm.constants import DQ_OPS, Q_OPS @@ -33,6 +36,7 @@ ) from executorch.backends.arm.tosa.backend import TOSABackend from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.specification import TosaSpecification from executorch.exir.backend.partitioner import ( DelegationSpec, Partitioner, @@ -49,6 +53,31 @@ logger = logging.getLogger(__name__) +class DecomposableResizeSupported(OperatorSupportBase): + """Accept exact boundary bilinear downscales. + + These are decomposed later. + + """ + + def __init__(self, tosa_spec: TosaSpecification) -> None: + """Initialize the check with the active TOSA specification.""" + self.tosa_spec = tosa_spec + + def is_node_supported( + self, + submodules: Mapping[str, torch.nn.Module], + node: torch.fx.Node, + ) -> bool: + """Return True when the resize matches the decomposable boundary case. + + This keeps the node delegatable so a later pass can rewrite it. + + """ + del submodules + return is_exact_tosa_boundary_bilinear_downscale(node, self.tosa_spec) + + def _is_custom_partition_op( custom_ops: set[torch._ops.OpOverload], target: object ) -> bool: @@ -263,6 +292,7 @@ def __init__( ) self.tosa_spec = compile_spec.tosa_spec self.additional_checks = additional_checks + self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[torch._ops.OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() @@ -419,9 +449,7 @@ def _tag_module( # noqa "Got overlapping tags in two different modules, this shouldn't happen." ) tags = tags | submodule_tags - operator_support = tosa_support_factory( - self.tosa_spec, containing_program, reporter, self.additional_checks - ) + operator_support = self._create_operator_support(containing_program, reporter) if self._custom_partition_ops: custom_ops = set(self._custom_partition_ops) @@ -534,6 +562,19 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: tags.remove(active_tag) return tags + def _create_operator_support( + self, + containing_program: ExportedProgram, + reporter: WhyNoPartitionReporter, + ) -> OperatorSupportBase: + return tosa_support_factory( + self.tosa_spec, + containing_program, + reporter, + self.additional_checks, + additional_positive_checks=[self._decomposable_resize_support], + ) + def partition(self, exported_program: ExportedProgram) -> PartitionResult: """Partition the program and tag TOSA-compatible subgraphs. diff --git a/backends/arm/tosa/resize_utils.py b/backends/arm/tosa/resize_utils.py index 23be6ff42fc..f0d17b72cfe 100644 --- a/backends/arm/tosa/resize_utils.py +++ b/backends/arm/tosa/resize_utils.py @@ -39,6 +39,26 @@ def _max_scale(tosa_spec: TosaSpecification) -> int: return _MAX_SCALE_LEVEL_8K if getattr(tosa_spec, "level_8k", False) else _MAX_SCALE +def _is_downscale_at_or_below_tosa_minimum(scale_ints: Sequence[int]) -> bool: + scale_y_n, scale_y_d, scale_x_n, scale_x_d = scale_ints + return scale_y_d >= 16 * scale_y_n or scale_x_d >= 16 * scale_x_n + + +def is_exact_tosa_resize_boundary_downscale( + scale: Sequence[int | torch.SymInt], +) -> bool: + """Return True for the exact 1/16 downscale boundary TOSA rejects.""" + scale_ints = _as_concrete_ints(scale) + if scale_ints is None: + return False + + scale_y_n, scale_y_d, scale_x_n, scale_x_d = scale_ints + if min(scale_y_n, scale_y_d, scale_x_n, scale_x_d) <= 0: + return False + + return scale_y_d == 16 * scale_y_n and scale_x_d == 16 * scale_x_n + + def _validate_dimensions( input_hw: Sequence[int | torch.SymInt], output_hw: Sequence[int | torch.SymInt] | None, @@ -89,6 +109,8 @@ def get_tosa_resize_output_hw_validation_error( def _validate_scale( scale: Sequence[int | torch.SymInt], tosa_spec: TosaSpecification, + *, + allow_exact_boundary_downscale: bool = False, ) -> str | None: invalid_scale = _first_outside_range( _concrete_int_values(scale), _INT16_MIN, _INT16_MAX @@ -126,7 +148,11 @@ def _validate_scale( # The scale values are already in the doubled rational representation that # TOSA RESIZE lowering emits, so the lower-bound downscale rule can be # checked directly against them. - if scale_y_d >= 16 * scale_y_n or scale_x_d >= 16 * scale_x_n: + if _is_downscale_at_or_below_tosa_minimum(scale_ints): + if allow_exact_boundary_downscale and is_exact_tosa_resize_boundary_downscale( + scale_ints + ): + return None return ( "RESIZE downscale must be strictly greater than 1/16; " f"got y={scale_y_n}/{scale_y_d}, x={scale_x_n}/{scale_x_d}" @@ -134,6 +160,43 @@ def _validate_scale( return None +def is_exact_tosa_resize_boundary_downscale_case( + *, + input_hw: Sequence[int | torch.SymInt], + output_hw: Sequence[int | torch.SymInt] | None, + scale: Sequence[int | torch.SymInt], + offset: Sequence[int | torch.SymInt], + border: Sequence[int | torch.SymInt], + tosa_spec: TosaSpecification, +) -> bool: + if not is_exact_tosa_resize_boundary_downscale(scale): + return False + + scale_ints = _as_concrete_ints(scale) + + validation_error = _validate_dimensions(input_hw, output_hw) + if validation_error is not None: + return False + validation_error = _validate_scale( + scale, + tosa_spec, + allow_exact_boundary_downscale=True, + ) + if validation_error is not None: + return False + if scale_ints is None: + return False + + for validation_error in ( + _validate_offset(offset, scale_ints), + _validate_border(border, scale_ints), + _validate_output_shape(input_hw, output_hw, scale, offset, border), + ): + if validation_error is not None: + return False + return True + + def _validate_offset( offset: Sequence[int | torch.SymInt], scale_ints: list[int], diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index 04d6a23607c..8ed8d8941e3 100644 --- a/backends/arm/vgf/partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -5,7 +5,10 @@ from typing import final, Optional, Sequence -from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.partitioner import ( + DecomposableResizeSupported, + TOSAPartitioner, +) from executorch.backends.arm.vgf import VgfBackend, VgfCompileSpec from executorch.exir.backend.partitioner import DelegationSpec from executorch.exir.dialects._ops import ops as exir_ops @@ -34,6 +37,7 @@ def __init__( ) self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec + self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() # Preserve grid_sampler_2d for the VGF custom-lowering path only.