ensure qparams match for slices

backends/xnnpack/partition/config/generic_node_configs.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 # Copyright 2026 Arm Limited and/or its affiliates.
@@ -22,20 +22,108 @@
     tag_as_implicit_q_dq,
 )
 from executorch.backends.xnnpack.utils.utils import (
+    get_param_tensor,
     get_input_node,
+    is_param_node,
     normalize_mean_dims,
     normalize_pool2d_args,
 )
 from executorch.exir.backend.canonical_partitioners.config_partitioner import (
     format_target_name,
 )
 from executorch.exir.backend.utils import is_shape_dynamic, WhyNoPartition
+from torch._subclasses.fake_tensor import FakeTensor
 from torch.export import ExportedProgram
 
 logger = logging.getLogger(__name__)
 why = WhyNoPartition(logger=logger)
 
 
+def _get_q_dq_quantization_params(
+    node: torch.fx.Node,
+) -> tuple[tuple[str, object], ...]:
+    params = []
+    for index, argument in enumerate(node.target._schema.arguments):
+        # These are dequant-only floating output types. The quantized storage
+        # dtype remains part of the parameters and must match.
+        if index == 0 or argument.name in {"out_dtype", "output_dtype"}:
+            continue
+        value = (
+            node.args[index]
+            if index < len(node.args)
+            else node.kwargs.get(argument.name, argument.default_value)
+        )
+        params.append((argument.name, value))
+    return tuple(params)
+
+
+def _q_dq_quantization_param_values_match(
+    dequant_value: object,
+    quant_value: object,
+    ep: ExportedProgram,
+) -> bool:
+    if isinstance(dequant_value, torch.fx.Node) and is_param_node(
+        ep, dequant_value
+    ):
+        dequant_value = get_param_tensor(ep, dequant_value)
+    if isinstance(quant_value, torch.fx.Node) and is_param_node(ep, quant_value):
+        quant_value = get_param_tensor(ep, quant_value)
+
+    # to_backend replaces state tensors with data-less FakeTensors before
+    # partitioning, so only object identity is safely comparable in that path.
+    if isinstance(dequant_value, FakeTensor) or isinstance(quant_value, FakeTensor):
+        return dequant_value is quant_value
+
+    if isinstance(dequant_value, torch.Tensor) or isinstance(
+        quant_value, torch.Tensor
+    ):
+        return (
+            isinstance(dequant_value, torch.Tensor)
+            and isinstance(quant_value, torch.Tensor)
+            and torch.equal(dequant_value, quant_value)
+        )
+    if isinstance(dequant_value, (list, tuple)) or isinstance(
+        quant_value, (list, tuple)
+    ):
+        return (
+            isinstance(dequant_value, (list, tuple))
+            and isinstance(quant_value, (list, tuple))
+            and len(dequant_value) == len(quant_value)
+            and all(
+                _q_dq_quantization_param_values_match(
+                    dequant_element, quant_element, ep
+                )
+                for dequant_element, quant_element in zip(
+                    dequant_value, quant_value
+                )
+            )
+        )
+    return dequant_value == quant_value
+
+
+def _q_dq_quantization_params_match(
+    dequant_node: torch.fx.Node,
+    quant_node: torch.fx.Node,
+    ep: ExportedProgram,
+) -> bool:
+    dequant_params = _get_q_dq_quantization_params(dequant_node)
+    quant_params = _get_q_dq_quantization_params(quant_node)
+    if len(dequant_params) != len(quant_params):
+        return False
+
+    for (dequant_name, dequant_value), (quant_name, quant_value) in zip(
+        dequant_params, quant_params
+    ):
+        if dequant_name != quant_name:
+            return False
+        if not _q_dq_quantization_param_values_match(
+            dequant_value, quant_value, ep
+        ):
+            return False
+
+    return True
+
+
 class GenericNodePartitionerConfig(XNNPartitionerConfig):
     def __init__(self, fused_act: Optional[List[str]] = None, **kwargs):
         """
@@ -639,6 +727,19 @@
         input_node = get_input_node(node, 0)
         output_node = node
 
+        # Only the single-user dq -> slice -> q pattern is serialized as a
+        # quantized slice. Multi-user slices are serialized as FP32.
+        if is_dequant(input_node) and len(node.users) == 1:
+            quant_node = next(iter(node.users))
+            if is_quant(quant_node) and not _q_dq_quantization_params_match(
+                input_node, quant_node, ep
+            ):
+                why(
+                    node,
+                    reason="XNNPACK static slice requires matching input and output quantization parameters",
+                )
+                return False
+
         input_shape = list(input_node.meta["val"].shape)
         output_shape = list(output_node.meta["val"].shape)
 

backends/xnnpack/test/ops/test_slice_copy.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -7,7 +7,13 @@
 import unittest
 
 import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
+    XnnpackPartitioner,
+)
 from executorch.backends.xnnpack.test.tester import Tester
+from executorch.backends.xnnpack.utils.utils import get_param_tensor
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.program._fake_program import get_fake_program
 
 
 class TestSliceCopy(unittest.TestCase):
@@ -168,3 +174,167 @@
             .serialize()
             .run_method_and_compare_outputs()
         )
+
+    def test_qs8_slice_copy_mismatched_qparams_falls_back(self):
+        class SliceCopy(torch.nn.Module):
+            def forward(self, x):
+                quantized = torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                    x, 0.25, 0, -128, 127, torch.int8
+                )
+                dequantized = (
+                    torch.ops.quantized_decomposed.dequantize_per_tensor.default(
+                        quantized, 0.25, 0, -128, 127, torch.int8
+                    )
+                )
+                sliced = torch.ops.aten.slice.Tensor(dequantized, 1, 0, 2)
+                requantized = (
+                    torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                        sliced, 0.5, 0, -128, 127, torch.int8
+                    )
+                )
+                return torch.ops.quantized_decomposed.dequantize_per_tensor.default(
+                    requantized, 0.5, 0, -128, 127, torch.int8
+                )
+
+        inputs = (torch.randn(1, 4, 3),)
+        (
+            Tester(SliceCopy(), inputs)
+            .export()
+            .to_edge()
+            .partition()
+            .check(["executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"])
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_qs8_per_channel_slice_copy_mismatched_qparams_falls_back(self):
+        class SliceCopy(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer(
+                    "input_scales", torch.tensor([0.25, 0.5, 0.75, 1.0])
+                )
+                self.register_buffer(
+                    "output_scales", torch.tensor([0.5, 0.75, 1.0, 1.25])
+                )
+                self.register_buffer("zero_points", torch.zeros(4, dtype=torch.int64))
+
+            def forward(self, x):
+                quantized = torch.ops.quantized_decomposed.quantize_per_channel.default(
+                    x, self.input_scales, self.zero_points, 1, -128, 127, torch.int8
+                )
+                dequantized = (
+                    torch.ops.quantized_decomposed.dequantize_per_channel.default(
+                        quantized,
+                        self.input_scales,
+                        self.zero_points,
+                        1,
+                        -128,
+                        127,
+                        torch.int8,
+                    )
+                )
+                sliced = torch.ops.aten.slice.Tensor(dequantized, 2, 0, 2)
+                requantized = (
+                    torch.ops.quantized_decomposed.quantize_per_channel.default(
+                        sliced,
+                        self.output_scales,
+                        self.zero_points,
+                        1,
+                        -128,
+                        127,
+                        torch.int8,
+                    )
+                )
+                return torch.ops.quantized_decomposed.dequantize_per_channel.default(
+                    requantized,
+                    self.output_scales,
+                    self.zero_points,
+                    1,
+                    -128,
+                    127,
+                    torch.int8,
+                )
+
+        inputs = (torch.randn(1, 4, 3),)
+        (
+            Tester(SliceCopy(), inputs)
+            .export()
+            .to_edge()
+            .partition()
+            .check(["executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"])
+        )
+
+    def test_qs8_per_channel_slice_copy_matching_distinct_constant_qparams(self):
+        class SliceCopy(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.input_scales = torch.tensor([0.25, 0.5, 0.75, 1.0])
+                self.output_scales = torch.tensor([0.25, 0.5, 0.75, 1.0])
+                self.zero_points = torch.zeros(4, dtype=torch.int64)
+
+            def forward(self, x):
+                quantized = torch.ops.quantized_decomposed.quantize_per_channel.default(
+                    x, self.input_scales, self.zero_points, 1, -128, 127, torch.int8
+                )
+                dequantized = (
+                    torch.ops.quantized_decomposed.dequantize_per_channel.default(
+                        quantized,
+                        self.input_scales,
+                        self.zero_points,
+                        1,
+                        -128,
+                        127,
+                        torch.int8,
+                    )
+                )
+                sliced = torch.ops.aten.slice.Tensor(dequantized, 2, 0, 2)
+                requantized = (
+                    torch.ops.quantized_decomposed.quantize_per_channel.default(
+                        sliced,
+                        self.output_scales,
+                        self.zero_points,
+                        1,
+                        -128,
+                        127,
+                        torch.int8,
+                    )
+                )
+                return torch.ops.quantized_decomposed.dequantize_per_channel.default(
+                    requantized,
+                    self.output_scales,
+                    self.zero_points,
+                    1,
+                    -128,
+                    127,
+                    torch.int8,
+                )
+
+        real_edge_program = (
+            Tester(SliceCopy(), (torch.randn(1, 4, 3),))
+            .export()
+            .to_edge()
+            .get_artifact()
+            .exported_program()
+        )
+        edge_program = get_fake_program(real_edge_program)
+        slice_node = next(
+            node
+            for node in edge_program.graph.nodes
+            if node.target == exir_ops.edge.aten.slice_copy.Tensor
+        )
+        dequant_node = slice_node.args[0]
+        quant_node = next(iter(slice_node.users))
+        input_scales = get_param_tensor(edge_program, dequant_node.args[1])
+        output_scales = get_param_tensor(edge_program, quant_node.args[1])
+
+        self.assertIsNot(dequant_node.args[1], quant_node.args[1])
+        if input_scales is None or output_scales is None:
+            self.fail("Expected lifted scale constants")
+        self.assertTrue(torch.equal(input_scales, output_scales))
+
+        partition_result = XnnpackPartitioner().partition(edge_program)
+        delegation_tag = slice_node.meta.get("delegation_tag")
+        self.assertIsNotNone(delegation_tag)
+        self.assertIn(delegation_tag, partition_result.partition_tags)