diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index 8a769a8383e..cc278702446 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -22,7 +22,9 @@ 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, ) @@ -30,12 +32,98 @@ 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 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: 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) diff --git a/backends/xnnpack/test/ops/test_slice_copy.py b/backends/xnnpack/test/ops/test_slice_copy.py index 25a5051e587..a595ab68db9 100644 --- a/backends/xnnpack/test/ops/test_slice_copy.py +++ b/backends/xnnpack/test/ops/test_slice_copy.py @@ -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 @@ def forward(self, x): .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)