From ca472b604f31c5e79661f77b7c906e396e2e1723 Mon Sep 17 00:00:00 2001 From: Rob Elliott Date: Wed, 17 Jun 2026 12:02:40 +0100 Subject: [PATCH 1/4] 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. From da17e51a7b37350ae0740c70078a11c9a44da49b Mon Sep 17 00:00:00 2001 From: Rob Elliott Date: Fri, 19 Jun 2026 08:51:21 +0100 Subject: [PATCH 2/4] Arm backend: Fix grid sampler dispatch shape Signed-off-by: Rob Elliott Change-Id: I4728474acfc882a7cc13a2478e03c397b44ffbf2 Signed-off-by: Rob Elliott --- .../test/misc/test_custom_shader_payload.py | 60 ++++++++++++++++++- .../test/misc/test_custom_shader_payloads.py | 5 +- .../test/ops/test_custom_shader_lowering.py | 29 +++++++++ ...ewrite_grid_sampler_to_tosa_custom_pass.py | 53 +++++++++++++++- .../rewrite_grid_sampler_to_tosa_custom.py | 6 ++ backends/arm/vgf/shaders/grid_sampler.glsl | 26 ++++++-- backends/arm/vgf/shaders/grid_sampler.py | 29 ++++++++- .../arm/vgf/shaders/grid_sampler.spirv.b64 | 25 +------- examples/arm/custom_operators.md | 23 +++++++ 9 files changed, 218 insertions(+), 38 deletions(-) diff --git a/backends/arm/test/misc/test_custom_shader_payload.py b/backends/arm/test/misc/test_custom_shader_payload.py index bfafc143c57..2d1c3685bfd 100644 --- a/backends/arm/test/misc/test_custom_shader_payload.py +++ b/backends/arm/test/misc/test_custom_shader_payload.py @@ -27,7 +27,6 @@ GRID_SAMPLER_2D_SHADER_LANGUAGE, GRID_SAMPLER_2D_SHADER_SOURCE, GRID_SAMPLER_2D_VK_FORMAT, - GRID_SAMPLER_2D_WORKGROUP_SIZES, ) @@ -45,11 +44,12 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_round_trip(): interpolation_mode=0, padding_mode=2, align_corners=True, + output_shape=(1, 4, 8, 8), ) decoded = decode_payload(encode_payload(payload)) assert decoded["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT - assert decoded["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES + assert decoded["workgroup_sizes"] == [1, 1, 1] assert decoded["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE assert base64.b64decode(decoded["shader_code"])[:4] == b"\x03\x02\x23\x07" assert decoded["input_0_type"] == "Tensor" @@ -58,7 +58,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_round_trip(): assert decoded["input_0_binding"] == 0 assert decoded["input_1_type"] == "Tensor" assert decoded["input_1_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT - assert decoded["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER" + assert decoded["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM" assert decoded["input_1_binding"] == 1 assert decoded["output_0_type"] == "Tensor" assert decoded["output_0_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT @@ -72,10 +72,12 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_uses_sampler_for_c4(): padding_mode=0, align_corners=False, input_shape=(1, 4, 8, 8), + output_shape=(1, 4, 4, 4), input_dtype=torch.float32, ) assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE + assert payload["workgroup_sizes"] == [1, 1, 1] assert base64.b64decode(payload["shader_code"])[:4] == b"\x03\x02\x23\x07" assert payload["input_0_type"] == "Image" assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT @@ -104,11 +106,13 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_uses_int8_sampler_for_c padding_mode=0, align_corners=False, input_shape=(1, 4, 8, 8), + output_shape=(1, 4, 4, 4), input_dtype=torch.int8, output_dtype=torch.int8, ) assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE + assert payload["workgroup_sizes"] == [1, 1, 1] assert base64.b64decode(payload["shader_code"])[:4] == b"\x03\x02\x23\x07" assert payload["input_0_type"] == "Image" assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_INT8_VK_FORMAT @@ -130,23 +134,58 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_keeps_c3_on_buffer(): padding_mode=0, align_corners=False, input_shape=(1, 3, 8, 8), + output_shape=(1, 3, 4, 4), input_dtype=torch.float32, ) assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE + assert payload["workgroup_sizes"] == [1, 1, 1] assert payload["input_0_type"] == "Tensor" assert payload["input_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER" + assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM" assert payload["output_0_type"] == "Tensor" assert payload["output_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER" assert "input_0_sampler" not in payload +def test_grid_sampler_2d_custom_shader_payload_sampler_dispatch_rounds_up_output(): + payload = build_grid_sampler_2d_payload( + interpolation_mode=0, + padding_mode=0, + align_corners=False, + input_shape=(1, 4, 32, 32), + output_shape=(1, 4, 17, 9), + input_dtype=torch.float32, + ) + + assert payload["input_0_type"] == "Image" + assert payload["output_0_type"] == "Image" + assert payload["workgroup_sizes"] == [2, 3, 1] + + +def test_grid_sampler_2d_custom_shader_payload_buffer_dispatch_rounds_up_output(): + payload = build_grid_sampler_2d_payload( + interpolation_mode=2, + padding_mode=0, + align_corners=False, + input_shape=(1, 4, 32, 32), + output_shape=(1, 4, 17, 9), + input_dtype=torch.float32, + ) + + assert payload["input_0_type"] == "Tensor" + assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM" + assert payload["output_0_type"] == "Tensor" + assert payload["workgroup_sizes"] == [2, 3, 1] + + def test_grid_sampler_2d_custom_shader_payload_no_target_align_corners_sampler(): payload = build_grid_sampler_2d_payload( interpolation_mode=0, padding_mode=0, align_corners=True, input_shape=(1, 4, 8, 8), + output_shape=(1, 4, 8, 8), input_dtype=torch.float32, ) @@ -170,6 +209,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_int8_align_corners_samp padding_mode=0, align_corners=True, input_shape=(1, 4, 8, 8), + output_shape=(1, 4, 8, 8), input_dtype=torch.int8, output_dtype=torch.int8, ) @@ -194,6 +234,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_bicubic_buffer(): padding_mode=0, align_corners=False, input_shape=(1, 4, 8, 8), + output_shape=(1, 4, 8, 8), input_dtype=torch.float32, ) @@ -201,6 +242,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_bicubic_buffer(): assert payload["input_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER" assert payload["output_0_type"] == "Tensor" assert payload["output_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER" + assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM" assert "input_0_sampler" not in payload @@ -209,6 +251,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_uses_spirv(): interpolation_mode=0, padding_mode=0, align_corners=False, + output_shape=(1, 4, 8, 8), ) shader_binary = base64.b64decode(payload["shader_code"]) @@ -253,6 +296,7 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_rejects_bad_modes(): interpolation_mode=99, padding_mode=0, align_corners=False, + output_shape=(1, 4, 8, 8), ) with pytest.raises(ValueError, match="Unsupported padding_mode"): @@ -260,4 +304,14 @@ def test_grid_sampler_2d_custom_shader_payload_no_target_rejects_bad_modes(): interpolation_mode=0, padding_mode=99, align_corners=False, + output_shape=(1, 4, 8, 8), + ) + + +def test_grid_sampler_2d_custom_shader_payload_requires_output_shape(): + with pytest.raises(ValueError, match="requires output_shape for dispatch"): + build_grid_sampler_2d_payload( + interpolation_mode=0, + padding_mode=0, + align_corners=False, ) diff --git a/backends/arm/test/misc/test_custom_shader_payloads.py b/backends/arm/test/misc/test_custom_shader_payloads.py index 5c7120d14de..f55289e0882 100644 --- a/backends/arm/test/misc/test_custom_shader_payloads.py +++ b/backends/arm/test/misc/test_custom_shader_payloads.py @@ -33,7 +33,6 @@ GRID_SAMPLER_2D_SHADER_LANGUAGE, GRID_SAMPLER_2D_SHADER_SOURCE, GRID_SAMPLER_2D_VK_FORMAT, - GRID_SAMPLER_2D_WORKGROUP_SIZES, ) from torch.export import export @@ -107,12 +106,13 @@ def test_buffer_shader_payload_vgf_encodes_bindings_and_formats(): interpolation_mode=0, padding_mode=0, align_corners=False, + output_shape=(1, 4, 8, 8), ) ) ) assert payload["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT - assert payload["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES + assert payload["workgroup_sizes"] == [1, 1, 1] assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE assert payload["input_0_binding"] == 0 assert payload["input_1_binding"] == 1 @@ -150,6 +150,7 @@ def test_shader_payload_vgf_uses_expected_glsl_and_spirv_asset(): interpolation_mode=0, padding_mode=0, align_corners=False, + output_shape=(1, 4, 8, 8), ) assert GRID_SAMPLER_2D_SHADER_SOURCE == "grid_sampler.glsl" diff --git a/backends/arm/test/ops/test_custom_shader_lowering.py b/backends/arm/test/ops/test_custom_shader_lowering.py index fed9f9e2e8c..0fc2d7e0824 100644 --- a/backends/arm/test/ops/test_custom_shader_lowering.py +++ b/backends/arm/test/ops/test_custom_shader_lowering.py @@ -37,6 +37,7 @@ RewriteGridSamplerToTosaCustomPass, ) from executorch.backends.arm.vgf.shaders.grid_sampler import ( + _dispatch_shape_for_output_shape, decode_payload, grid_sampler_2d_operator_name, ) @@ -256,3 +257,31 @@ def test_shader_lowering_vgf_decodes_expected_implementation_attrs(): assert payload["input_0_binding"] == 0 assert payload["input_1_binding"] == 1 assert payload["output_0_binding"] == 2 + + +def test_grid_sampler_dispatch_shape_includes_batch_dimension(): + assert _dispatch_shape_for_output_shape((2, 3, 5, 9)) == [2, 1, 2] + + +def test_in_tree_grid_sampler_vgf_payload_uses_batched_dispatch_shape(): + edge_model = to_edge( + export( + _GridSampleModule(), + ( + torch.randn(2, 3, 11, 13), + torch.randn(2, 5, 9, 2), + ), + ) + ) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")): + transformed = edge_model.transform([RewriteGridSamplerToTosaCustomPass()]) + + custom_node = next( + node + for node in transformed.exported_program().graph.nodes + if node.target == exir_ops.backend.tosa.CUSTOM.default + ) + payload = decode_payload(custom_node.kwargs["implementation_attrs"]) + + assert payload["workgroup_sizes"] == [2, 1, 2] diff --git a/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py b/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py index eb4b5f23660..9b7e1b68919 100644 --- a/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py +++ b/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py @@ -28,7 +28,6 @@ GRID_SAMPLER_2D_SHADER_ENTRY_POINT, GRID_SAMPLER_2D_SHADER_LANGUAGE, GRID_SAMPLER_2D_VK_FORMAT, - GRID_SAMPLER_2D_WORKGROUP_SIZES, ) from executorch.exir import to_edge from executorch.exir.dialects._ops import ops as exir_ops @@ -88,7 +87,7 @@ def test_rewrite_grid_sampler_to_tosa_custom_vgf_no_target(): payload = decode_payload(custom_node.kwargs["implementation_attrs"]) assert payload["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT - assert payload["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES + assert payload["workgroup_sizes"] == [1, 1, 1] assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE assert payload["input_0_type"] == "Image" assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT @@ -111,6 +110,28 @@ def test_rewrite_grid_sampler_to_tosa_custom_vgf_no_target(): assert any(node.target == exir_ops.edge.aten.slice_copy.Tensor for node in nodes) +def test_rewrite_grid_sampler_to_tosa_custom_sampler_dispatch_rounds_up_output(): + model = GridSampler2d() + example_inputs = ( + torch.randn(1, 4, 32, 32), + torch.randn(1, 17, 9, 2), + ) + + edge_model = to_edge(export(model, example_inputs)) + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")): + edge_model = edge_model.transform([RewriteGridSamplerToTosaCustomPass()]) + nodes = list(edge_model.exported_program().graph.nodes) + + custom_node = next( + node for node in nodes if node.target == exir_ops.backend.tosa.CUSTOM.default + ) + payload = decode_payload(custom_node.kwargs["implementation_attrs"]) + + assert payload["input_0_type"] == "Image" + assert payload["output_0_type"] == "Image" + assert payload["workgroup_sizes"] == [2, 3, 1] + + def test_rewrite_grid_sampler_to_tosa_custom_no_target_uses_sampler_for_c4(): model = GridSampler2d() example_inputs = ( @@ -129,6 +150,7 @@ def test_rewrite_grid_sampler_to_tosa_custom_no_target_uses_sampler_for_c4(): payload = decode_payload(custom_node.kwargs["implementation_attrs"]) assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE + assert payload["workgroup_sizes"] == [1, 1, 1] assert ( payload["input_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER" @@ -202,6 +224,7 @@ def test_rewrite_grid_sampler_to_tosa_custom_c3_pad_for_align_corners(): ) payload = decode_payload(custom_node.kwargs["implementation_attrs"]) + assert payload["workgroup_sizes"] == [1, 1, 1] assert payload["input_0_type"] == "Image" assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT assert ( @@ -233,7 +256,33 @@ def test_rewrite_grid_sampler_to_tosa_custom_no_c3_pad_for_bicubic(): payload = decode_payload(custom_node.kwargs["implementation_attrs"]) assert payload["input_0_type"] == "Tensor" + assert payload["workgroup_sizes"] == [1, 1, 1] + assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM" assert not any(node.target == exir_ops.edge.aten.cat.default for node in nodes) assert not any( node.target == exir_ops.edge.aten.slice_copy.Tensor for node in nodes ) + + +def test_rewrite_grid_sampler_to_tosa_custom_buffer_dispatch_rounds_up_output(): + model = GridSampler2d() + model.interpolation_mode_ = 2 + example_inputs = ( + torch.randn(1, 4, 32, 32), + torch.randn(1, 17, 9, 2), + ) + + edge_model = to_edge(export(model, example_inputs)) + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")): + edge_model = edge_model.transform([RewriteGridSamplerToTosaCustomPass()]) + nodes = list(edge_model.exported_program().graph.nodes) + + custom_node = next( + node for node in nodes if node.target == exir_ops.backend.tosa.CUSTOM.default + ) + payload = decode_payload(custom_node.kwargs["implementation_attrs"]) + + assert payload["input_0_type"] == "Tensor" + assert payload["output_0_type"] == "Tensor" + assert payload["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_TENSOR_ARM" + assert payload["workgroup_sizes"] == [2, 3, 1] diff --git a/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py b/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py index efff1730914..3e74be05f45 100644 --- a/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py +++ b/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py @@ -158,16 +158,21 @@ def _encode_payload( padding_mode: int, align_corners: bool, input_tensor: torch.fx.Node, + output_tensor: torch.fx.Node, output_dtype: torch.dtype | None = None, ) -> list[int]: input_val = input_tensor.meta.get("val") if input_val is None: raise RuntimeError("grid_sampler_2d input is missing tensor metadata") + output_val = output_tensor.meta.get("val") + if output_val is None: + raise RuntimeError("grid_sampler_2d output is missing tensor metadata") payload = build_grid_sampler_2d_payload( interpolation_mode=interpolation_mode, padding_mode=padding_mode, align_corners=align_corners, input_shape=tuple(input_val.shape), + output_shape=tuple(output_val.shape), input_dtype=input_val.dtype, output_dtype=output_dtype, ) @@ -258,6 +263,7 @@ def call(self, graph_module): padding_mode=padding_mode, align_corners=align_corners, input_tensor=custom_input, + output_tensor=node, output_dtype=output_dtype, ) nhwc_input = create_node( diff --git a/backends/arm/vgf/shaders/grid_sampler.glsl b/backends/arm/vgf/shaders/grid_sampler.glsl index 30d22a98920..4a197cafdd2 100644 --- a/backends/arm/vgf/shaders/grid_sampler.glsl +++ b/backends/arm/vgf/shaders/grid_sampler.glsl @@ -4,6 +4,7 @@ // LICENSE file in the root directory of this source tree. #version 450 +#extension GL_ARM_tensors : require layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in; @@ -11,15 +12,30 @@ layout(set = 0, binding = 0) readonly buffer Input0 { float input0[]; }; -layout(set = 0, binding = 1) readonly buffer Input1 { - float input1[]; -}; +layout(set = 0, binding = 1) uniform tensorARM grid; layout(set = 0, binding = 2) writeonly buffer Output0 { float output0[]; }; void main() { - uint index = gl_GlobalInvocationID.x; - output0[index] = input0[index]; + uvec3 gid = gl_GlobalInvocationID.xyz; + uint output_batch = tensorSizeARM(grid, 0); + uint output_height = tensorSizeARM(grid, 1); + uint output_width = tensorSizeARM(grid, 2); + if (gid.x >= output_width || gid.y >= output_height || gid.z >= output_batch) { + return; + } + + uint output_spatial_size = output_width * output_height; + if (output_spatial_size == 0u) { + return; + } + + uint channels = output0.length() / (output_spatial_size * output_batch); + uint base = + ((gid.z * output_spatial_size) + (gid.y * output_width) + gid.x) * channels; + for (uint channel = 0u; channel < channels; ++channel) { + output0[base + channel] = 0.0; + } } diff --git a/backends/arm/vgf/shaders/grid_sampler.py b/backends/arm/vgf/shaders/grid_sampler.py index a82198060d0..abfb6190019 100644 --- a/backends/arm/vgf/shaders/grid_sampler.py +++ b/backends/arm/vgf/shaders/grid_sampler.py @@ -97,6 +97,7 @@ def build_grid_sampler_2d_payload( padding_mode: int, align_corners: bool, input_shape: tuple[int, ...] | None = None, + output_shape: tuple[int, ...] | None = None, input_dtype: Any | None = None, output_dtype: Any | None = None, ) -> dict[str, Any]: @@ -108,6 +109,8 @@ def build_grid_sampler_2d_payload( align_corners (bool): Whether grid_sample aligns tensor corners. input_shape (tuple[int, ...] | None): Input tensor shape, used to select sampler-backed shader metadata when supported. + output_shape (tuple[int, ...] | None): Output tensor shape, required + to derive dispatch counts for the shader launch. input_dtype (Any | None): Input tensor dtype, used to select sampler Vulkan formats when supported. output_dtype (Any | None): Output tensor dtype. Defaults to @@ -127,6 +130,8 @@ def build_grid_sampler_2d_payload( _PADDING_MODE_NAMES, "padding_mode", ) + if output_shape is None: + raise ValueError("grid_sampler payload requires output_shape for dispatch") if output_dtype is None: output_dtype = input_dtype @@ -150,7 +155,10 @@ def build_grid_sampler_2d_payload( payload = { "entry_point": GRID_SAMPLER_2D_SHADER_ENTRY_POINT, - "workgroup_sizes": GRID_SAMPLER_2D_WORKGROUP_SIZES, + # Current runtime consumes this field as dispatch counts, not local + # shader workgroup size. The current grid-sample shaders use a 2D + # output-space work model with an 8x8 work volume per workgroup. + "workgroup_sizes": _dispatch_shape_for_output_shape(output_shape), "shader_language": GRID_SAMPLER_2D_SHADER_LANGUAGE, "shader_code": shader_code, "input_0_binding": 0, @@ -186,7 +194,7 @@ def build_grid_sampler_2d_payload( "input_0_type": "Tensor", "input_0_vkformat": GRID_SAMPLER_2D_VK_FORMAT, "input_0_vkdescriptortype": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER", - "input_1_vkdescriptortype": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER", + "input_1_vkdescriptortype": "VK_DESCRIPTOR_TYPE_TENSOR_ARM", "output_0_type": "Tensor", "output_0_vkformat": GRID_SAMPLER_2D_VK_FORMAT, "output_0_vkdescriptortype": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER", @@ -195,6 +203,23 @@ def build_grid_sampler_2d_payload( return payload +def _dispatch_shape_for_output_shape(output_shape: tuple[int, ...]) -> list[int]: + if len(output_shape) != 4: + raise ValueError( + "grid_sampler output_shape must be rank 4 NCHW, " + f"got shape {output_shape}" + ) + output_batch = int(output_shape[0]) + output_height = int(output_shape[2]) + output_width = int(output_shape[3]) + group_x, group_y, group_z = GRID_SAMPLER_2D_WORKGROUP_SIZES + return [ + (output_width + group_x - 1) // group_x, + (output_height + group_y - 1) // group_y, + (output_batch + group_z - 1) // group_z, + ] + + def _sampler_vk_format(input_dtype: Any | None, output_dtype: Any | None) -> str | None: if str(input_dtype) != str(output_dtype): return None diff --git a/backends/arm/vgf/shaders/grid_sampler.spirv.b64 b/backends/arm/vgf/shaders/grid_sampler.spirv.b64 index 59750d3204b..942d74d1158 100644 --- a/backends/arm/vgf/shaders/grid_sampler.spirv.b64 +++ b/backends/arm/vgf/shaders/grid_sampler.spirv.b64 @@ -1,24 +1 @@ -AwIjBwAAAQALAA0AKAAAAAAAAAARAAIAAQAAAAsABgABAAAAR0xTTC5zdGQuNDUwAAAAAA4AAwAAAAAA -AQAAAA8ABgAFAAAABAAAAG1haW4AAAAACwAAABAABgAEAAAAEQAAAAgAAAAIAAAAAQAAAAMAAwACAAAA -wgEAAAQACgBHTF9HT09HTEVfY3BwX3N0eWxlX2xpbmVfZGlyZWN0aXZlAAAEAAgAR0xfR09PR0xFX2lu -Y2x1ZGVfZGlyZWN0aXZlAAUABAAEAAAAbWFpbgAAAAAFAAQACAAAAGluZGV4AAAABQAIAAsAAABnbF9H -bG9iYWxJbnZvY2F0aW9uSUQAAAAFAAQAEgAAAE91dHB1dDAABgAFABIAAAAAAAAAb3V0cHV0MAAFAAMA -FAAAAAAAAAAFAAQAGQAAAElucHV0MAAABgAFABkAAAAAAAAAaW5wdXQwAAAFAAMAGwAAAAAAAAAFAAQA -JQAAAElucHV0MQAABgAFACUAAAAAAAAAaW5wdXQxAAAFAAMAJwAAAAAAAABHAAQACwAAAAsAAAAcAAAA -RwAEABEAAAAGAAAABAAAAEcAAwASAAAAAwAAAEgABAASAAAAAAAAABkAAABIAAUAEgAAAAAAAAAjAAAA -AAAAAEcAAwAUAAAAGQAAAEcABAAUAAAAIQAAAAIAAABHAAQAFAAAACIAAAAAAAAARwAEABgAAAAGAAAA -BAAAAEcAAwAZAAAAAwAAAEgABAAZAAAAAAAAABgAAABIAAUAGQAAAAAAAAAjAAAAAAAAAEcAAwAbAAAA -GAAAAEcABAAbAAAAIQAAAAAAAABHAAQAGwAAACIAAAAAAAAARwAEACMAAAALAAAAGQAAAEcABAAkAAAA -BgAAAAQAAABHAAMAJQAAAAMAAABIAAQAJQAAAAAAAAAYAAAASAAFACUAAAAAAAAAIwAAAAAAAABHAAMA -JwAAABgAAABHAAQAJwAAACEAAAABAAAARwAEACcAAAAiAAAAAAAAABMAAgACAAAAIQADAAMAAAACAAAA -FQAEAAYAAAAgAAAAAAAAACAABAAHAAAABwAAAAYAAAAXAAQACQAAAAYAAAADAAAAIAAEAAoAAAABAAAA -CQAAADsABAAKAAAACwAAAAEAAAArAAQABgAAAAwAAAAAAAAAIAAEAA0AAAABAAAABgAAABYAAwAQAAAA -IAAAAB0AAwARAAAAEAAAAB4AAwASAAAAEQAAACAABAATAAAAAgAAABIAAAA7AAQAEwAAABQAAAACAAAA -FQAEABUAAAAgAAAAAQAAACsABAAVAAAAFgAAAAAAAAAdAAMAGAAAABAAAAAeAAMAGQAAABgAAAAgAAQA -GgAAAAIAAAAZAAAAOwAEABoAAAAbAAAAAgAAACAABAAdAAAAAgAAABAAAAArAAQABgAAACEAAAAIAAAA -KwAEAAYAAAAiAAAAAQAAACwABgAJAAAAIwAAACEAAAAhAAAAIgAAAB0AAwAkAAAAEAAAAB4AAwAlAAAA -JAAAACAABAAmAAAAAgAAACUAAAA7AAQAJgAAACcAAAACAAAANgAFAAIAAAAEAAAAAAAAAAMAAAD4AAIA -BQAAADsABAAHAAAACAAAAAcAAABBAAUADQAAAA4AAAALAAAADAAAAD0ABAAGAAAADwAAAA4AAAA+AAMA -CAAAAA8AAAA9AAQABgAAABcAAAAIAAAAPQAEAAYAAAAcAAAACAAAAEEABgAdAAAAHgAAABsAAAAWAAAA -HAAAAD0ABAAQAAAAHwAAAB4AAABBAAYAHQAAACAAAAAUAAAAFgAAABcAAAA+AAMAIAAAAB8AAAD9AAEA -OAABAA== +AwIjBwAAAQALAA0AdQAAAAAAAAARAAIAAQAAABEAAgBOEAAACgAFAFNQVl9BUk1fdGVuc29ycwAKAAsAU1BWX0tIUl9zdG9yYWdlX2J1ZmZlcl9zdG9yYWdlX2NsYXNzAAAAAAsABgABAAAAR0xTTC5zdGQuNDUwAAAAAA4AAwAAAAAAAQAAAA8ABgAFAAAABAAAAG1haW4AAAAACwAAABAABgAEAAAAEQAAAAgAAAAIAAAAAQAAAAMAAwACAAAAwgEAAAQABQBHTF9BUk1fdGVuc29ycwAABAAKAEdMX0dPT0dMRV9jcHBfc3R5bGVfbGluZV9kaXJlY3RpdmUAAAQACABHTF9HT09HTEVfaW5jbHVkZV9kaXJlY3RpdmUABQAEAAQAAABtYWluAAAAAAUAAwAJAAAAZ2lkAAUACAALAAAAZ2xfR2xvYmFsSW52b2NhdGlvbklEAAAABQAGAA4AAABvdXRwdXRfYmF0Y2gAAAAABQAEABMAAABncmlkAAAAAAUABgAXAAAAb3V0cHV0X2hlaWdodAAAAAUABgAbAAAAb3V0cHV0X3dpZHRoAAAAAAUABwA3AAAAb3V0cHV0X3NwYXRpYWxfc2l6ZQAFAAUAQAAAAGNoYW5uZWxzAAAAAAUABABCAAAAT3V0cHV0MAAGAAUAQgAAAAAAAABvdXRwdXQwAAUAAwBEAAAAAAAAAAUABABNAAAAYmFzZQAAAAAFAAQAXAAAAGNoYW5uZWwABQAEAHIAAABJbnB1dDAAAAYABQByAAAAAAAAAGlucHV0MAAABQADAHQAAAAAAAAARwAEAAsAAAALAAAAHAAAAEcABAATAAAAIQAAAAEAAABHAAQAEwAAACIAAAAAAAAARwAEAEEAAAAGAAAABAAAAEcAAwBCAAAAAgAAAEgABABCAAAAAAAAABkAAABIAAUAQgAAAAAAAAAjAAAAAAAAAEcAAwBEAAAAGQAAAEcABABEAAAAIQAAAAIAAABHAAQARAAAACIAAAAAAAAARwAEAHAAAAALAAAAGQAAAEcABABxAAAABgAAAAQAAABHAAMAcgAAAAIAAABIAAQAcgAAAAAAAAAYAAAASAAFAHIAAAAAAAAAIwAAAAAAAABHAAMAdAAAABgAAABHAAQAdAAAACEAAAAAAAAARwAEAHQAAAAiAAAAAAAAABMAAgACAAAAIQADAAMAAAACAAAAFQAEAAYAAAAgAAAAAAAAABcABAAHAAAABgAAAAMAAAAgAAQACAAAAAcAAAAHAAAAIAAEAAoAAAABAAAABwAAADsABAAKAAAACwAAAAEAAAAgAAQADQAAAAcAAAAGAAAAFgADAA8AAAAgAAAAKwAEAAYAAAAQAAAABAAAAEMQBAARAAAADwAAABAAAAAgAAQAEgAAAAAAAAARAAAAOwAEABIAAAATAAAAAAAAACsABAAGAAAAFQAAAAAAAAArAAQABgAAABkAAAABAAAAKwAEAAYAAAAdAAAAAgAAABQAAgAfAAAAHQADAEEAAAAPAAAAHgADAEIAAABBAAAAIAAEAEMAAAAMAAAAQgAAADsABABDAAAARAAAAAwAAAAVAAQARgAAACAAAAABAAAAKwAEAEYAAABlAAAAAAAAACsABAAPAAAAaQAAAAAAAAAgAAQAagAAAAwAAAAPAAAAKwAEAEYAAABtAAAAAQAAACsABAAGAAAAbwAAAAgAAAAsAAYABwAAAHAAAABvAAAAbwAAABkAAAAdAAMAcQAAAA8AAAAeAAMAcgAAAHEAAAAgAAQAcwAAAAwAAAByAAAAOwAEAHMAAAB0AAAADAAAADYABQACAAAABAAAAAAAAAADAAAA+AACAAUAAAA7AAQACAAAAAkAAAAHAAAAOwAEAA0AAAAOAAAABwAAADsABAANAAAAFwAAAAcAAAA7AAQADQAAABsAAAAHAAAAOwAEAA0AAAA3AAAABwAAADsABAANAAAAQAAAAAcAAAA7AAQADQAAAE0AAAAHAAAAOwAEAA0AAABcAAAABwAAAD0ABAAHAAAADAAAAAsAAAA+AAMACQAAAAwAAAA9AAQAEQAAABQAAAATAAAARhAFAAYAAAAWAAAAFAAAABUAAAA+AAMADgAAABYAAAA9AAQAEQAAABgAAAATAAAARhAFAAYAAAAaAAAAGAAAABkAAAA+AAMAFwAAABoAAAA9AAQAEQAAABwAAAATAAAARhAFAAYAAAAeAAAAHAAAAB0AAAA+AAMAGwAAAB4AAABBAAUADQAAACAAAAAJAAAAFQAAAD0ABAAGAAAAIQAAACAAAAA9AAQABgAAACIAAAAbAAAArgAFAB8AAAAjAAAAIQAAACIAAACoAAQAHwAAACQAAAAjAAAA9wADACYAAAAAAAAA+gAEACQAAAAlAAAAJgAAAPgAAgAlAAAAQQAFAA0AAAAnAAAACQAAABkAAAA9AAQABgAAACgAAAAnAAAAPQAEAAYAAAApAAAAFwAAAK4ABQAfAAAAKgAAACgAAAApAAAA+QACACYAAAD4AAIAJgAAAPUABwAfAAAAKwAAACMAAAAFAAAAKgAAACUAAACoAAQAHwAAACwAAAArAAAA9wADAC4AAAAAAAAA+gAEACwAAAAtAAAALgAAAPgAAgAtAAAAQQAFAA0AAAAvAAAACQAAAB0AAAA9AAQABgAAADAAAAAvAAAAPQAEAAYAAAAxAAAADgAAAK4ABQAfAAAAMgAAADAAAAAxAAAA+QACAC4AAAD4AAIALgAAAPUABwAfAAAAMwAAACsAAAAmAAAAMgAAAC0AAAD3AAMANQAAAAAAAAD6AAQAMwAAADQAAAA1AAAA+AACADQAAAD9AAEA+AACADUAAAA9AAQABgAAADgAAAAbAAAAPQAEAAYAAAA5AAAAFwAAAIQABQAGAAAAOgAAADgAAAA5AAAAPgADADcAAAA6AAAAPQAEAAYAAAA7AAAANwAAAKoABQAfAAAAPAAAADsAAAAVAAAA9wADAD4AAAAAAAAA+gAEADwAAAA9AAAAPgAAAPgAAgA9AAAA/QABAPgAAgA+AAAARAAFAAYAAABFAAAARAAAAAAAAAB8AAQARgAAAEcAAABFAAAAfAAEAAYAAABIAAAARwAAAD0ABAAGAAAASQAAADcAAAA9AAQABgAAAEoAAAAOAAAAhAAFAAYAAABLAAAASQAAAEoAAACGAAUABgAAAEwAAABIAAAASwAAAD4AAwBAAAAATAAAAEEABQANAAAATgAAAAkAAAAdAAAAPQAEAAYAAABPAAAATgAAAD0ABAAGAAAAUAAAADcAAACEAAUABgAAAFEAAABPAAAAUAAAAEEABQANAAAAUgAAAAkAAAAZAAAAPQAEAAYAAABTAAAAUgAAAD0ABAAGAAAAVAAAABsAAACEAAUABgAAAFUAAABTAAAAVAAAAIAABQAGAAAAVgAAAFEAAABVAAAAQQAFAA0AAABXAAAACQAAABUAAAA9AAQABgAAAFgAAABXAAAAgAAFAAYAAABZAAAAVgAAAFgAAAA9AAQABgAAAFoAAABAAAAAhAAFAAYAAABbAAAAWQAAAFoAAAA+AAMATQAAAFsAAAA+AAMAXAAAABUAAAD5AAIAXQAAAPgAAgBdAAAA9gAEAF8AAABgAAAAAAAAAPkAAgBhAAAA+AACAGEAAAA9AAQABgAAAGIAAABcAAAAPQAEAAYAAABjAAAAQAAAALAABQAfAAAAZAAAAGIAAABjAAAA+gAEAGQAAABeAAAAXwAAAPgAAgBeAAAAPQAEAAYAAABmAAAATQAAAD0ABAAGAAAAZwAAAFwAAACAAAUABgAAAGgAAABmAAAAZwAAAEEABgBqAAAAawAAAEQAAABlAAAAaAAAAD4AAwBrAAAAaQAAAPkAAgBgAAAA+AACAGAAAAA9AAQABgAAAGwAAABcAAAAgAAFAAYAAABuAAAAbAAAAG0AAAA+AAMAXAAAAG4AAAD5AAIAXQAAAPgAAgBfAAAA/QABADgAAQA= diff --git a/examples/arm/custom_operators.md b/examples/arm/custom_operators.md index 24375e5b937..12641db8fc8 100644 --- a/examples/arm/custom_operators.md +++ b/examples/arm/custom_operators.md @@ -90,3 +90,26 @@ For a minimal end-to-end example showing the required pieces in Python, see - If you need image semantics for 3-channel data, you must either: - pad to 4 channels explicitly before the custom node, or - stay on a tensor/buffer path + + +## Dispatch Shape + +### Current Limitation +- The payload field named `workgroup_sizes` is currently consumed as the dispatch shape for `vkCmdDispatch(...)`. +- The shader's actual local workgroup size still comes from the GLSL/SPIR-V declaration. +- So today this field is effectively misnamed: it carries workgroup counts, not local workgroup size. + +### What To Set +- Set `workgroup_sizes` to the dispatch grid you want the runtime to launch. +- Derive it from: + - the total work your shader must cover, and + - the shader-local workgroup size declared in the shader. +- For a 2D shader with local size `(lx, ly, lz)` covering an output of `(W, H)`, set: + - `dispatch_x = ceil(W / lx)` + - `dispatch_y = ceil(H / ly)` + - `dispatch_z = 1` unless the shader maps work across `z` + +### Practical Rule +- Treat shader local size and payload `workgroup_sizes` as two separate concepts. +- Local size is fixed by the shader. +- Payload `workgroup_sizes` must be computed from the work volume for that operator instance. From f2190749e7346e3a07f6860131dc87c1f7a74ff0 Mon Sep 17 00:00:00 2001 From: Rob Elliott Date: Fri, 19 Jun 2026 09:26:04 +0100 Subject: [PATCH 3/4] Arm backend: Dequantize quantized grid sample coordinates Signed-off-by: Rob Elliott Change-Id: Ibe61128651763aa2b83d9bf187b656fe11487aed Signed-off-by: Rob Elliott --- backends/arm/TARGETS | 1 + backends/arm/quantizer/arm_quantizer.py | 47 ++++++++ .../arm/quantizer/quantization_annotator.py | 15 ++- backends/arm/quantizer/quantization_config.py | 18 ++- ...ewrite_grid_sampler_to_tosa_custom_pass.py | 67 ++++++++++- .../test/quantizer/test_generic_annotater.py | 109 +++++++++++++++++- backends/arm/vgf/_passes/__init__.py | 3 + .../insert_grid_sampler_grid_dequant_pass.py | 95 +++++++++++++++ .../rewrite_grid_sampler_to_tosa_custom.py | 11 +- backends/arm/vgf/backend.py | 28 +++-- .../backends/arm-vgf/arm-vgf-quantization.md | 37 +----- 11 files changed, 370 insertions(+), 61 deletions(-) create mode 100644 backends/arm/vgf/_passes/insert_grid_sampler_grid_dequant_pass.py diff --git a/backends/arm/TARGETS b/backends/arm/TARGETS index 6254c73e188..bdb90c33880 100644 --- a/backends/arm/TARGETS +++ b/backends/arm/TARGETS @@ -88,6 +88,7 @@ runtime.python_library( srcs = [ "vgf/__init__.py", "vgf/_passes/__init__.py", + "vgf/_passes/insert_grid_sampler_grid_dequant_pass.py", "vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py", "vgf/backend.py", "vgf/check_env.py", diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py index bc10b142d32..04f96d562c0 100644 --- a/backends/arm/quantizer/arm_quantizer.py +++ b/backends/arm/quantizer/arm_quantizer.py @@ -23,6 +23,7 @@ from executorch.backends.arm.quantizer.quantization_config import ( QuantizationConfig, TOSAQuantizationConfig, + VGFQuantizationConfig, ) from executorch.backends.arm.quantizer.quantizer_support import ( TOSA_QUANTIZER_SUPPORT_DICT, @@ -107,6 +108,24 @@ logger = logging.getLogger(__name__) +def _wrap_vgf_quantization_config( + quantization_config: Optional[QuantizationConfig], +) -> Optional[QuantizationConfig]: + if ( + quantization_config is None + or isinstance(quantization_config, VGFQuantizationConfig) + or not isinstance(quantization_config, TOSAQuantizationConfig) + ): + return quantization_config + return VGFQuantizationConfig( + quantization_config.input_activation, + quantization_config.output_activation, + quantization_config.weight, + quantization_config.bias, + quantization_config.label, + ) + + def get_cond_while_submodules_ao( graph_module: GraphModule, apply_quantization: bool = False, @@ -1322,3 +1341,31 @@ def __init__( use_composable_quantizer: bool = False, ) -> None: super().__init__(compile_spec, use_composable_quantizer) + + def set_global( + self, quantization_config: Optional[QuantizationConfig] + ) -> TOSAQuantizer: + """Set the global quantization config for VGF lowering.""" + return super().set_global(_wrap_vgf_quantization_config(quantization_config)) + + def set_module_type( + self, module_type: Callable, quantization_config: Optional[QuantizationConfig] + ) -> TOSAQuantizer: + """Set the quantization config for a specific module type.""" + return super().set_module_type( + module_type, _wrap_vgf_quantization_config(quantization_config) + ) + + def set_module_name( + self, module_name: str, quantization_config: Optional[QuantizationConfig] + ) -> TOSAQuantizer: + """Set the quantization config for a specific module name.""" + return super().set_module_name( + module_name, _wrap_vgf_quantization_config(quantization_config) + ) + + def set_io( + self, quantization_config: Optional[QuantizationConfig] + ) -> TOSAQuantizer: + """Set the quantization config used for model inputs and outputs.""" + return super().set_io(_wrap_vgf_quantization_config(quantization_config)) diff --git a/backends/arm/quantizer/quantization_annotator.py b/backends/arm/quantizer/quantization_annotator.py index 58cd422bc6f..14d13236ee7 100644 --- a/backends/arm/quantizer/quantization_annotator.py +++ b/backends/arm/quantizer/quantization_annotator.py @@ -474,8 +474,11 @@ def _match_pattern( 8: _QParams((0.999 - (-0.999)) / (1 << 8), 0), 16: _QParams((0.99999 - (-0.99999)) / (1 << 16), 0), }, - # grid_sampler image input/output use SNORM-compatible qparams. The grid - # coordinate tensor is intentionally left unquantized. + # grid_sampler image input/output use SNORM-compatible qparams. The broader + # quantized graph currently quantizes the grid-producing path as well, so + # input 1 follows the standard activation qspec and lowering materializes a + # dequant boundary before the shader. This is a functional stopgap; we may + # want to preserve float grid coordinates or use a higher-precision path. torch.ops.aten.grid_sampler.default: { 8: _QParams(1.0 / 127.0, 0, -127, 127), }, @@ -824,13 +827,21 @@ def any_or_hardtanh_min_zero(n: Node): quant_properties.quant_output = _QuantProperty(0, output_act_qspec) elif node.target == torch.ops.aten.grid_sampler.default: image_node = ensure_type(Node, node.args[0]) + grid_node = ensure_type(Node, node.args[1]) grid_sampler_image_qspec = quantization_config.get_input_act_qspec( node, image_node ) + grid_sampler_grid_qspec = quantization_config.get_input_act_qspec( + node, grid_node + ) grid_sampler_output_qspec = quantization_config.get_output_act_qspec(node) if grid_sampler_image_qspec is None or grid_sampler_output_qspec is None: return None quant_properties.quant_inputs = [_QuantProperty(0, grid_sampler_image_qspec)] + if grid_sampler_grid_qspec is not None: + quant_properties.quant_inputs.append( + _QuantProperty(1, grid_sampler_grid_qspec) + ) quant_properties.quant_output = _QuantProperty(0, grid_sampler_output_qspec) elif node.target in (torch.ops.aten.where.self,): true_node = ensure_type(Node, node.args[1]) diff --git a/backends/arm/quantizer/quantization_config.py b/backends/arm/quantizer/quantization_config.py index b202ae28000..13507698dc5 100644 --- a/backends/arm/quantizer/quantization_config.py +++ b/backends/arm/quantizer/quantization_config.py @@ -259,6 +259,8 @@ def _derive_qparams_fn( class TOSAQuantizationConfig(QuantizationConfig): """Configures quantization, while enforcing TOSA specific constraints.""" + quantize_grid_sampler_grid = False + SHARED_OUTPUT_ACT_QSPEC_PATTERNS = { torch.ops.aten.adaptive_avg_pool2d.default, torch.ops.aten.upsample_bilinear2d.vec, @@ -307,12 +309,14 @@ def get_input_act_qspec(self, node=None, input_node=None): else: return SharedQuantizationSpec((node.args[0], node)) elif node.target == torch.ops.aten.grid_sampler.default: - if input_node != node.args[0]: + if input_node == node.args[0]: + input_act_qspec = super().get_input_act_qspec(node, input_node) + return _get_fixed_qparams_qspec( + node.target, _fixed_input_qspec_ops, input_act_qspec + ) + if not self.quantize_grid_sampler_grid: return None - input_act_qspec = super().get_input_act_qspec(node, input_node) - return _get_fixed_qparams_qspec( - node.target, _fixed_input_qspec_ops, input_act_qspec - ) + return super().get_input_act_qspec(node, input_node) elif node.target in _fixed_input_qspec_ops: input_act_qspec = super().get_input_act_qspec(node, input_node) return _get_fixed_qparams_qspec( @@ -321,6 +325,10 @@ def get_input_act_qspec(self, node=None, input_node=None): return super().get_input_act_qspec(node, input_node) + +class VGFQuantizationConfig(TOSAQuantizationConfig): + quantize_grid_sampler_grid = True + def get_weight_qspec( self, node: Optional[Node] = None ) -> Optional[QuantizationSpecBase]: diff --git a/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py b/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py index 9b7e1b68919..578f4427752 100644 --- a/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py +++ b/backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py @@ -8,15 +8,18 @@ import torch import torch.nn.functional as F from executorch.backends.arm._passes import FoldAndAnnotateQParamsPass +from executorch.backends.arm._passes.quant_args import QuantArgs from executorch.backends.arm.quantizer.arm_quantizer import ( get_symmetric_quantization_config, - TOSAQuantizer, + VgfQuantizer, ) from executorch.backends.arm.tosa.specification import ( TosaLoweringContext, TosaSpecification, ) -from executorch.backends.arm.vgf._passes.rewrite_grid_sampler_to_tosa_custom import ( +from executorch.backends.arm.vgf import VgfCompileSpec +from executorch.backends.arm.vgf._passes import ( + InsertGridSamplerGridDequantPass, RewriteGridSamplerToTosaCustomPass, ) from executorch.backends.arm.vgf.shaders.grid_sampler import ( @@ -171,8 +174,8 @@ def test_quantized_grid_sampler_uses_int8_sampler_payload( torch.randn(1, channels, 8, 8), torch.rand(1, 4, 4, 2), ) - quantizer = TOSAQuantizer( - TosaSpecification.create_from_string("TOSA-1.0+INT"), + quantizer = VgfQuantizer( + VgfCompileSpec("TOSA-1.0+INT"), use_composable_quantizer=use_composable_quantizer, ) quantizer.set_global(get_symmetric_quantization_config(is_per_channel=False)) @@ -185,7 +188,11 @@ def test_quantized_grid_sampler_uses_int8_sampler_payload( edge_model = to_edge(export(converted, example_inputs, strict=True)) with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP+INT")): edge_model = edge_model.transform( - [FoldAndAnnotateQParamsPass(), RewriteGridSamplerToTosaCustomPass()] + [ + FoldAndAnnotateQParamsPass(), + InsertGridSamplerGridDequantPass(), + RewriteGridSamplerToTosaCustomPass(), + ] ) nodes = list(edge_model.exported_program().graph.nodes) @@ -193,6 +200,7 @@ def test_quantized_grid_sampler_uses_int8_sampler_payload( node for node in nodes if node.target == exir_ops.backend.tosa.CUSTOM.default ) payload = decode_payload(custom_node.kwargs["implementation_attrs"]) + grid_input = custom_node.args[0][1] assert payload["input_0_type"] == "Image" assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_INT8_VK_FORMAT @@ -200,12 +208,61 @@ def test_quantized_grid_sampler_uses_int8_sampler_payload( assert payload["input_1_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT assert payload["output_0_type"] == "Image" assert payload["output_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_INT8_VK_FORMAT + assert grid_input.meta["val"].dtype == torch.float32 + assert grid_input.target in ( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + ) assert custom_node.meta["input_qparams"][0].qmin == -127 assert custom_node.meta["input_qparams"][0].qmax == 127 + assert 1 not in custom_node.meta["input_qparams"] assert next(iter(custom_node.meta["output_qparams"].values())).qmin == -127 assert next(iter(custom_node.meta["output_qparams"].values())).qmax == 127 +def test_quantized_grid_sampler_rejects_dequantized_grid_with_int8_image_payload(): + model = GridSampler2d().eval() + example_inputs = ( + torch.randn(1, 4, 8, 8), + torch.rand(1, 4, 4, 2), + ) + quantizer = VgfQuantizer(VgfCompileSpec("TOSA-1.0+INT")) + quantizer.set_global(get_symmetric_quantization_config(is_per_channel=False)) + + exported = export(model, example_inputs, strict=True) + prepared = prepare_pt2e(exported.module(), quantizer) + prepared(*example_inputs) + converted = convert_pt2e(prepared) + + edge_model = to_edge(export(converted, example_inputs, strict=True)) + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP+INT")): + edge_model = edge_model.transform([FoldAndAnnotateQParamsPass()]) + + exported_program = edge_model.exported_program() + grid_sampler_node = next( + node + for node in exported_program.graph.nodes + if node.target == exir_ops.edge.aten.grid_sampler_2d.default + ) + grid_node = grid_sampler_node.args[1] + grid_node.meta["val"] = torch.zeros_like(grid_node.meta["val"], dtype=torch.int8) + grid_sampler_node.meta["input_qparams"][1] = QuantArgs( + scale=[0.01, 0.01], + zp=[0, 0], + qmin=-128, + qmax=127, + dtype=torch.int8, + axis=3, + per_channel=True, + ) + + with pytest.raises( + RuntimeError, + match="grid_sampler grid dequant only supports per-tensor qparams", + ): + InsertGridSamplerGridDequantPass()(exported_program.graph_module) + + def test_rewrite_grid_sampler_to_tosa_custom_c3_pad_for_align_corners(): model = GridSampler2d() model.align_corners_ = True diff --git a/backends/arm/test/quantizer/test_generic_annotater.py b/backends/arm/test/quantizer/test_generic_annotater.py index b5cfd1efdc6..e3d32a1372a 100644 --- a/backends/arm/test/quantizer/test_generic_annotater.py +++ b/backends/arm/test/quantizer/test_generic_annotater.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025 Arm Limited and/or its affiliates. +# Copyright 2024-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. @@ -7,11 +7,20 @@ from typing import Any, Callable, Tuple import torch +import torch.nn.functional as F from executorch.backends.arm.quantizer import is_annotated +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_quantization_config, + VgfQuantizer, +) +from executorch.backends.arm.quantizer.quantization_annotator import annotate_graph from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineINT +from executorch.backends.arm.vgf import VgfCompileSpec from executorch.backends.test.harness.stages import StageType +from torch.export import export from torch.fx.passes.utils.source_matcher_utils import get_source_partitions +from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY input_t1 = Tuple[torch.Tensor] # Input x @@ -142,3 +151,101 @@ def test_concat_tosa_INT(): torch.concatenate, ((torch.randn(2, 3), torch.randn(2, 3)),), dim=0 ), ) + + +class GridSampleModule(torch.nn.Module): + def forward(self, x: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + return F.grid_sample( + x, + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=False, + ) + + +class GridFloatQuantizationConfig: + def __init__(self) -> None: + self.base = get_symmetric_quantization_config() + + def get_input_act_qspec(self, node=None, input_node=None): + if ( + node is not None + and input_node is not None + and node.target == torch.ops.aten.grid_sampler.default + and input_node == node.args[1] + ): + return None + return self.base.get_input_act_qspec(node, input_node) + + def get_output_act_qspec(self, node=None): + return self.base.get_output_act_qspec(node) + + def get_weight_qspec(self, node=None): + return self.base.get_weight_qspec(node) + + def get_bias_qspec(self, node=None): + return self.base.get_bias_qspec(node) + + +def test_grid_sampler_annotation_keeps_float_grid_when_grid_qspec_is_none(): + module = GridSampleModule().eval() + example_inputs = (torch.randn(1, 4, 8, 8), torch.randn(1, 4, 4, 2)) + gm = export(module, example_inputs).graph_module + + annotate_graph(gm, GridFloatQuantizationConfig()) + + grid_sampler_node = next( + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.aten.grid_sampler.default + ) + image_node = grid_sampler_node.args[0] + grid_node = grid_sampler_node.args[1] + annotation = grid_sampler_node.meta[Q_ANNOTATION_KEY] + + assert is_annotated(grid_sampler_node) + assert image_node in annotation.input_qspec_map + assert grid_node not in annotation.input_qspec_map + assert annotation.output_qspec is not None + + +def test_grid_sampler_annotation_keeps_default_tosa_grid_float(): + module = GridSampleModule().eval() + example_inputs = (torch.randn(1, 4, 8, 8), torch.randn(1, 4, 4, 2)) + gm = export(module, example_inputs).graph_module + + annotate_graph(gm, get_symmetric_quantization_config()) + + grid_sampler_node = next( + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.aten.grid_sampler.default + ) + grid_node = grid_sampler_node.args[1] + annotation = grid_sampler_node.meta[Q_ANNOTATION_KEY] + + assert grid_node not in annotation.input_qspec_map + + +def test_vgf_quantizer_quantizes_grid_sampler_grid_coords(): + module = GridSampleModule().eval() + example_inputs = (torch.randn(1, 4, 8, 8), torch.randn(1, 4, 4, 2)) + gm = export(module, example_inputs).graph_module + + quantizer = VgfQuantizer(VgfCompileSpec("TOSA-1.0+INT")) + quantizer.set_global(get_symmetric_quantization_config()) + quantizer.annotate(gm) + + grid_sampler_node = next( + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.aten.grid_sampler.default + ) + grid_node = grid_sampler_node.args[1] + annotation = grid_sampler_node.meta[Q_ANNOTATION_KEY] + + assert grid_node in annotation.input_qspec_map diff --git a/backends/arm/vgf/_passes/__init__.py b/backends/arm/vgf/_passes/__init__.py index 4733d218c47..444f3a08e6f 100644 --- a/backends/arm/vgf/_passes/__init__.py +++ b/backends/arm/vgf/_passes/__init__.py @@ -3,6 +3,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from .insert_grid_sampler_grid_dequant_pass import ( # noqa + InsertGridSamplerGridDequantPass, +) from .rewrite_grid_sampler_to_tosa_custom import ( # noqa RewriteGridSamplerToTosaCustomPass, ) diff --git a/backends/arm/vgf/_passes/insert_grid_sampler_grid_dequant_pass.py b/backends/arm/vgf/_passes/insert_grid_sampler_grid_dequant_pass.py new file mode 100644 index 00000000000..708f03a735f --- /dev/null +++ b/backends/arm/vgf/_passes/insert_grid_sampler_grid_dequant_pass.py @@ -0,0 +1,95 @@ +# 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 torch +from executorch.backends.arm._passes import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import create_node, set_node_arg +from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_input_qparams, +) +from executorch.backends.arm._passes.quant_args import QuantArgs +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 + + +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 _get_per_tensor_dequant_args(grid: torch.fx.Node, grid_qparams: QuantArgs) -> tuple: + return ( + grid, + grid_qparams.scale, + grid_qparams.zp, + grid_qparams.qmin, + grid_qparams.qmax, + grid_qparams.dtype, + ) + + +class InsertGridSamplerGridDequantPass(ArmPass): + """Insert an explicit float boundary for quantized grid_sample grid inputs. + + This runs before quant-node decomposition so the standard Arm quant passes + can legalize the inserted dequant op, and later VGF custom-op rewriting sees + the expected float grid contract. Per-channel grid qparams are rejected + because the follow-up decompose pass only legalizes per-tensor + quantized_decomposed ops. + + """ + + _passes_required_after: set[type[ExportPass]] = set() + + def call(self, graph_module: torch.fx.GraphModule): + modified = False + for node in list(graph_module.graph.nodes): + if ( + node.op != "call_function" + or node.target != exir_ops.edge.aten.grid_sampler_2d.default + ): + continue + + grid = node.args[1] + if not isinstance(grid, torch.fx.Node): + continue + if grid.meta["val"].dtype.is_floating_point: + continue + + grid_qparams = get_input_qparams(node).get(1) + if grid_qparams is None: + raise RuntimeError( + "Quantized grid_sampler grid input is missing input qparams" + ) + if grid_qparams.per_channel: + raise RuntimeError( + "grid_sampler grid dequant only supports per-tensor qparams" + ) + + with graph_module.graph.inserting_before(node): + dequant_grid = create_node( + graph_module.graph, + op_target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=_get_per_tensor_dequant_args(grid, grid_qparams), + from_node=node, + ) + _set_fake_tensor_meta( + dequant_grid, grid_qparams.dequantize_value(grid.meta["val"]) + ) + + set_node_arg(node, 1, dequant_grid) + if "input_qparams" in node.meta: + node.meta["input_qparams"] = { + idx: qargs + for idx, qargs in node.meta["input_qparams"].items() + if idx != 1 + } + modified = True + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py b/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py index 3e74be05f45..8d05ea3cafb 100644 --- a/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py +++ b/backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py @@ -245,6 +245,11 @@ def call(self, graph_module): interpolation_mode, align_corners, ) + if not grid.meta["val"].dtype.is_floating_point: + raise RuntimeError( + "grid_sampler rewrite expected float grid input; " + "InsertGridSamplerGridDequantPass should run before this pass" + ) operator_name = grid_sampler_2d_operator_name( interpolation_mode=interpolation_mode, @@ -278,11 +283,12 @@ def call(self, graph_module): custom_input.meta["val"], list(NHWC_ORDER) ), ) + custom_grid = grid custom_node = create_node( graph_module.graph, op_target=exir_ops.backend.tosa.CUSTOM.default, - args=([nhwc_input, grid],), + args=([nhwc_input, custom_grid],), kwargs={ "operator_name": operator_name, "domain_name": CUSTOM_SHADER_DOMAIN_NAME, @@ -291,7 +297,6 @@ def call(self, graph_module): from_node=node, inherit_qparams=True, ) - with graph_module.graph.inserting_after(custom_node): getitem_node = graph_module.graph.create_node( "call_function", @@ -300,7 +305,7 @@ def call(self, graph_module): kwargs={}, ) custom_output = _grid_sampler_2d_custom_fake_impl( - [nhwc_input.meta["val"], grid.meta["val"]], + [nhwc_input.meta["val"], custom_grid.meta["val"]], operator_name, CUSTOM_SHADER_DOMAIN_NAME, implementation_attrs, diff --git a/backends/arm/vgf/backend.py b/backends/arm/vgf/backend.py index 280427f0694..7a4bc59905f 100644 --- a/backends/arm/vgf/backend.py +++ b/backends/arm/vgf/backend.py @@ -20,7 +20,7 @@ from dataclasses import dataclass from typing import Any, final, List -from executorch.backends.arm._passes import RewriteConvPass +from executorch.backends.arm._passes import DecomposeQuantNodesPass, RewriteConvPass from executorch.backends.arm._passes.arm_pass_manager import ( _registered_pass_insertions, PassInsertions, @@ -30,7 +30,8 @@ arm_get_first_delegation_tag, TOSABackend, ) -from executorch.backends.arm.vgf._passes.rewrite_grid_sampler_to_tosa_custom import ( # type: ignore[import-not-found] +from executorch.backends.arm.vgf._passes import ( # type: ignore[import-not-found] + InsertGridSamplerGridDequantPass, RewriteGridSamplerToTosaCustomPass, ) @@ -48,6 +49,7 @@ from executorch.exir.backend.compile_spec_schema import ( # type: ignore[import-not-found] CompileSpec, ) +from executorch.exir.pass_base import ExportPass from torch.export.exported_program import ExportedProgram # debug functionality @@ -143,18 +145,20 @@ def check_vgf_runtime_backend_environment() -> VgfRuntimeEnvironmentCheck: ) -def _register_grid_sampler_rewrite_pass() -> None: - """Register VGF-only custom shader lowering passes.""" - existing_insertions = _registered_pass_insertions.get(RewriteConvPass) +def _register_pass_before(target_pass_type: type, pass_: ExportPass) -> None: + existing_insertions = _registered_pass_insertions.get(target_pass_type) if existing_insertions is not None and any( - isinstance(pass_, RewriteGridSamplerToTosaCustomPass) - for pass_ in existing_insertions.before_passes + isinstance(existing_pass, type(pass_)) + for existing_pass in existing_insertions.before_passes ): return - register_pass_insertions_before( - RewriteConvPass, - [RewriteGridSamplerToTosaCustomPass()], - ) + register_pass_insertions_before(target_pass_type, [pass_]) + + +def _register_vgf_rewrite_passes() -> None: + """Register VGF-only custom shader lowering passes.""" + _register_pass_before(DecomposeQuantNodesPass, InsertGridSamplerGridDequantPass()) + _register_pass_before(RewriteConvPass, RewriteGridSamplerToTosaCustomPass()) def _snapshot_registered_pass_insertions() -> dict[type, PassInsertions]: @@ -228,7 +232,7 @@ def preprocess( insertions_snapshot = _snapshot_registered_pass_insertions() try: - _register_grid_sampler_rewrite_pass() + _register_vgf_rewrite_passes() compile_spec = VgfCompileSpec._from_list(compile_specs) # deduce TOSA compile_spec from VGF compile spec. We get a new # compile spec list, containing only elements relevant for the diff --git a/docs/source/backends/arm-vgf/arm-vgf-quantization.md b/docs/source/backends/arm-vgf/arm-vgf-quantization.md index b10dca2f51e..679e45c53b5 100644 --- a/docs/source/backends/arm-vgf/arm-vgf-quantization.md +++ b/docs/source/backends/arm-vgf/arm-vgf-quantization.md @@ -56,51 +56,22 @@ Insert a quantizer with highest precedence. ```python def VgfQuantizer.set_global(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer': ``` -Set quantization_config for submodules not matched by other filters. - -Args: -- **quantization_config (Optional[QuantizationConfig])**: Configuration to - apply to modules that are not captured by name or type filters. - ``None`` indicates no quantization. +Set the global quantization config for VGF lowering. ```python def VgfQuantizer.set_io(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer': ``` -Set quantization_config for input and output nodes. - -Args: -- **quantization_config (Optional[QuantizationConfig])**: Configuration - describing activation quantization for model inputs and outputs. - ``None`` indicates no quantization. +Set the quantization config used for model inputs and outputs. ```python def VgfQuantizer.set_module_name(self, module_name: 'str', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer': ``` -Set quantization_config for submodules with a given module name. - -For example, calling set_module_name("blocks.sub") quantizes supported -patterns for that submodule with the provided quantization_config. - -Args: -- **module_name (str)**: Fully qualified module name to configure. -- **quantization_config (Optional[QuantizationConfig])**: Configuration - applied to the named submodule. ``None`` indicates no - quantization. +Set the quantization config for a specific module name. ```python def VgfQuantizer.set_module_type(self, module_type: 'Callable', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer': ``` -Set quantization_config for submodules with a given module type. - -For example, calling set_module_type(Softmax) quantizes supported -patterns in each Softmax instance with the provided quantization_config. - -Args: -- **module_type (Callable)**: Type whose submodules should use the - provided quantization configuration. -- **quantization_config (Optional[QuantizationConfig])**: Configuration to - apply to submodules of the given type. ``None`` indicates no - quantization. +Set the quantization config for a specific module type. ```python def VgfQuantizer.set_node_finder(self, quantization_config: 'Optional[QuantizationConfig]', node_finder: 'NodeFinder') -> 'TOSAQuantizer': From a1884e33fd87e8e9449800baee9172f1e8a677dd Mon Sep 17 00:00:00 2001 From: Rob Elliott Date: Thu, 16 Jul 2026 14:39:47 +0100 Subject: [PATCH 4/4] Retest CI