Arm backend: Dequantize quantized grid sample coordinates (#20962)

Arm backend: Dequantize quantized grid sample coordinates

Functional progress for quantized grid sample to enable the runtime
codegen to handle de-q into samplers in shaders.

Signed-off-by: Rob Elliott <Robert.Elliott@arm.com>
Change-Id: Ibe61128651763aa2b83d9bf187b656fe11487aed

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @rascani

---------

Signed-off-by: Rob Elliott <Robert.Elliott@arm.com>

Author: robell

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/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,
@@ -1404,3 +1423,31 @@ def __init__(
         use_composable_quantizer: bool = True,
     ) -> 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(
@@ -414,3 +418,7 @@ def get_output_act_qspec(
         if len(node.args) == 0:
             return super().get_output_act_qspec()
         return SharedQuantizationSpec((cast(Node, node.args[0]), node))
+
+
+class VGFQuantizationConfig(TOSAQuantizationConfig):
+    quantize_grid_sampler_grid = True

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,27 +188,81 @@ 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)
 
     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"])
+    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
     assert payload["input_1_type"] == "Tensor"
     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

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

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