Arm backend: Dequantize quantized grid sample coordinates

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",

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

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),

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)

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()

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
     )

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))

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])

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]: