Arm backend: Fix grid_sample when Q/DQ have been fused by exir

pytorch#22905 removes
some Q/DQ nodes before running the Arm passes. In the
FuseGridSamplerFlowOffsetPass pass, we expected to have a
Q/DQ pair before the grid_sample, but that pair is now removed
at the time of running the exir passes.
Change the guard in FuseGridSamplerFlowOffsetPass to check for
concat before the grid_sample to account for the fact we now
have a cat.

Signed-off-by: George Gekov <george.gekov@arm.com>
Change-Id: Icd968d84c6ec499626ddf5839a5728afbb00a731

Author: gggekov

backends/arm/test/passes/test_fuse_grid_sampler_flow_offset_pass.py

@@ -334,6 +334,49 @@ def test_match_flow_offset_grid_accepts_identity_base_grid_requantization():
     assert _match_flow_offset_grid(grid_sampler) == (flow, 0)
 
 
+def test_match_flow_offset_grid_accepts_folded_identity_base_grid():
+    grid_sampler, flow = _build_flow_offset_grid_sampler()
+    grid = grid_sampler.args[1]
+    assert isinstance(grid, torch.fx.Node)
+    add = grid.args[0]
+    assert isinstance(add, torch.fx.Node)
+    quantize = add.args[0]
+    assert isinstance(quantize, torch.fx.Node)
+    dequantize = quantize.args[0]
+    assert isinstance(dequantize, torch.fx.Node)
+    base_grid = dequantize.args[0]
+    assert isinstance(base_grid, torch.fx.Node)
+
+    qparams = QuantArgs(0.007870171219110489, 0, -127, 127, torch.int8)
+    base_grid.meta["input_qparams"] = {0: qparams}
+    base_grid.meta["output_qparams"] = {0: qparams}
+    add.replace_input_with(quantize, base_grid)
+
+    assert _match_flow_offset_grid(grid_sampler) == (flow, 0)
+
+
+def test_match_flow_offset_grid_rejects_folded_requantized_base_grid():
+    grid_sampler, _ = _build_flow_offset_grid_sampler()
+    grid = grid_sampler.args[1]
+    assert isinstance(grid, torch.fx.Node)
+    add = grid.args[0]
+    assert isinstance(add, torch.fx.Node)
+    quantize = add.args[0]
+    assert isinstance(quantize, torch.fx.Node)
+    dequantize = quantize.args[0]
+    assert isinstance(dequantize, torch.fx.Node)
+    base_grid = dequantize.args[0]
+    assert isinstance(base_grid, torch.fx.Node)
+
+    base_grid.meta["input_qparams"] = {
+        0: QuantArgs(0.007870171219110489, 0, -127, 127, torch.int8)
+    }
+    base_grid.meta["output_qparams"] = {0: QuantArgs(0.01, 0, -127, 127, torch.int8)}
+    add.replace_input_with(quantize, base_grid)
+
+    assert _match_flow_offset_grid(grid_sampler) is None
+
+
 @pytest.mark.parametrize(
     "kwargs",
     [

backends/arm/vgf/_passes/fuse_grid_sampler_flow_offset.py

@@ -10,11 +10,12 @@
     grid = permute(base_grid + flow[:, i:i + 2], [0, 2, 3, 1])
     output = grid_sample(image, grid, bilinear, border, align_corners=True)
 
-The base grid may pass through an identity dequantize/quantize pair. Images
-must have shape ``[1, 3|4, H, W]``, flow must have shape ``[1, 4, H, W]``,
-and ``i`` must be zero or two. Image and output quantization ranges must be
-``[-127, 127]`` because the custom shader stores them in SNORM images. The
-flow tensor may use the full int8 range because the shader reads it directly.
+The base grid may be direct or pass through an identity dequantize/quantize
+pair. Images must have shape ``[1, 3|4, H, W]``, flow must have shape
+``[1, 4, H, W]``, and ``i`` must be zero or two. Image and output
+quantization ranges must be ``[-127, 127]`` because the custom shader stores
+them in SNORM images. The flow tensor may use the full int8 range because the
+shader reads it directly.
 
 """
 
@@ -133,17 +134,43 @@ def _is_identity_requantization(
     return True
 
 
+def _has_identity_qparams(node: Node) -> bool:
+    """Return whether a node preserves one quantization representation.
+
+    Args:
+        node (Node): Candidate node with folded quantization metadata.
+
+    Returns:
+        bool: Whether every input and output qparam is identical.
+
+    """
+    try:
+        input_qparams = get_input_qparams(node)
+        output_qparams = get_output_qparams(node)
+    except ValueError:
+        return False
+    return all(
+        input_qparam == output_qparam
+        for input_qparam in input_qparams.values()
+        for output_qparam in output_qparams.values()
+    )
+
+
 def _unwrap_quantized_base_grid(node: Node) -> Node | None:
-    """Unwrap an identity dequantize-quantize pair around a base grid.
+    """Return a direct base grid or unwrap an identity Q/DQ pair.
 
     Args:
-        node (Node): Candidate quantize node.
+        node (Node): Candidate base-grid or quantize node.
 
     Returns:
-        Node | None: Unwrapped base-grid node, or ``None`` when the wrapper
-        changes integer values or has an unsupported form.
+        Node | None: Base-grid node, or ``None`` when the wrapper changes
+        integer values or has an unsupported form.
 
     """
+    # The exir FoldRedundantDequantizeQuantizePass has removed a DQ/Q pair after
+    # the cat. Ensure the cat has the same input & output qparams before returning.
+    if _target_is(node, "aten.cat.default") and _has_identity_qparams(node):
+        return node
     if not _target_is(node, "quantized_decomposed.quantize_per_tensor.default"):
         return None
     dequantize = node.args[0]