diff --git a/backends/xnnpack/_passes/insert_pad_qdq.py b/backends/xnnpack/_passes/insert_pad_qdq.py index 744a642105e..f5db7d9b620 100644 --- a/backends/xnnpack/_passes/insert_pad_qdq.py +++ b/backends/xnnpack/_passes/insert_pad_qdq.py @@ -4,81 +4,78 @@ # 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 cast, List + import torch from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass -from executorch.backends.xnnpack.utils.quant_utils import is_quant, tag_as_implicit_q_dq +from executorch.backends.xnnpack.utils.quant_utils import ( + is_dequant, + is_quant, + tag_as_implicit_q_dq, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import PassResult class InsertPadQDQPass(XNNPACKPass): """ - Completes the quantization of a constant_pad_nd that sits inside a quantized - region but was left with an fp32 output. - - An even-kernel 'same'-padding conv decomposes (after quantization) into - dequant -> constant_pad_nd -> convolution. Because the pad is introduced by - to_edge decomposition -- after the quantizer has run -- it is never annotated, - so no quantize follows it and its output would serialize as fp32. The - downstream conv would then reject its (now fp32) activation. - - A zero-valued pad preserves quantization, so we insert an implicit - quantize -> dequantize pair after the pad, reusing the feeding dequant's - params. The pad then delegates as a normal quantized XNNStaticConstantPad and - the conv sees a proper dequantized activation. The pad node itself is left in - place, so all graph shapes stay consistent through later retracing passes. + Inserts implicit quantize/dequantize pairs after constant_pad_nd nodes + that sit in a quantized context (input is a dequantize node), so the pad + can be serialized as a quantized static pad op. + + Skips pads whose output is already quantized (idempotent). + + Without this pass, a zero-valued constant_pad_nd between a dequantize and + a convolution would serialize as fp32 while the conv expects quantized + activation, causing a mismatch. """ - def _insert_qdq_after(self, graph, pad, q_params): - with graph.inserting_after(pad): - q = graph.create_node( - "call_function", - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - args=(), - ) - q.meta = pad.meta.copy() - tag_as_implicit_q_dq(q) - with graph.inserting_after(q): - dq = graph.create_node( - "call_function", - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - args=(q,) + q_params, - ) - dq.meta = q.meta.copy() - tag_as_implicit_q_dq(dq) - pad.replace_all_uses_with(dq) - # Set last so replace_all_uses_with above does not rewrite the quantize's - # own input. - q.args = (pad,) + q_params - - def call(self, graph_module: torch.fx.GraphModule): + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: graph = graph_module.graph - for pad in list(graph.nodes): + for node in list(graph.nodes): if ( - pad.op != "call_function" - or pad.target != exir_ops.edge.aten.constant_pad_nd.default + node.op != "call_function" + or node.target != exir_ops.edge.aten.constant_pad_nd.default ): continue - # Only per-tensor static activations are handled: _insert_qdq_after - # builds quantize_per_tensor.default, so the feeding dequant must have - # the matching per-tensor signature (a per-channel/per-token/affine - # dequant would supply mismatched args). - dq = pad.args[0] - if ( - not isinstance(dq, torch.fx.Node) - or dq.target - != exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ): + pad_input = node.args[0] + if not (isinstance(pad_input, torch.fx.Node) and is_dequant(pad_input)): + continue + + pad_value = cast(float, node.args[2]) if len(node.args) > 2 else 0.0 + if pad_value != 0.0: + continue + + pad_amounts = cast(List[int], node.args[1]) + if any(p < 0 for p in pad_amounts): continue - # Skip if the pad's output is already quantized. Requiring *no* user to - # be a quantize (rather than merely "not all") avoids double-quantizing - # pre-existing quant consumers when the pad has mixed users. - if not pad.users or any(is_quant(user) for user in pad.users): + if any(is_quant(user) for user in node.users): continue - self._insert_qdq_after(graph, pad, tuple(dq.args[1:])) + q_params = pad_input.args[1:] + + with graph.inserting_after(node): + q = graph.create_node( + "call_function", + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(node,) + q_params, + ) + q.meta = node.meta.copy() + tag_as_implicit_q_dq(q) + + with graph.inserting_after(q): + dq = graph.create_node( + "call_function", + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + args=(q,) + q_params, + ) + dq.meta = q.meta.copy() + tag_as_implicit_q_dq(dq) + + node.replace_all_uses_with(dq) + q.args = (node,) + q_params graph_module.recompile() return PassResult(graph_module, True) diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py index d0e9528809d..f093b60cf15 100644 --- a/backends/xnnpack/partition/config/gemm_configs.py +++ b/backends/xnnpack/partition/config/gemm_configs.py @@ -397,44 +397,73 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: return False return True - def supported_precision_types(self): - return [ - ConfigPrecisionType.FP32, - ConfigPrecisionType.STATIC_QUANT, - ConfigPrecisionType.DYNAMIC_QUANT, - ] - def _get_act_deps( self, node: torch.fx.Node, ep: ExportedProgram, precision: ConfigPrecisionType ) -> Tuple[bool, List[torch.fx.Node]]: - # An even-kernel 'same'-padding conv decomposes into - # dequant -> constant_pad_nd -> convolution. Pull the zero-valued spatial - # pad into this conv's partition so it delegates as a quantized - # XNNStaticConstantPad alongside the conv (InsertPadQDQPass completes the - # pad's quantization); otherwise the pad is orphaned and the conv is left - # with an unquantized activation. - # Only supporting 2D, non-transposed convs - is_transpose = node.args[6] - is_2d_conv = len(cast(List[int], node.args[4])) == 2 act_input = get_input_node(node, self.act_idx) + is_transpose = node.args[6] if ( precision != ConfigPrecisionType.FP32 and not is_transpose - and is_2d_conv - and act_input.target == exir_ops.edge.aten.constant_pad_nd.default - and len(act_input.users) == 1 - and is_dequant(get_input_node(act_input, 0)) + and is_node(act_input) ): - pad_value = act_input.args[2] if len(act_input.args) > 2 else 0 - pad_amounts = cast(List[int], act_input.args[1]) - spatial_only = len(pad_amounts) <= 4 or all(a == 0 for a in pad_amounts[4:]) - if pad_value == 0 and spatial_only: - valid, deps = super()._get_act_deps(act_input, ep, precision) - if valid: - return (True, [act_input, *deps]) + # Find a constant_pad_nd in the activation chain. It may be the direct + # input to the conv, or it may be behind the QDQ pair inserted by + # InsertPadQDQPass (dequant -> pad -> q -> dq -> conv). + pad_node = None + if act_input.target == exir_ops.edge.aten.constant_pad_nd.default: + pad_node = act_input + elif is_dequant(act_input) and len(act_input.users) == 1: + q_node = get_input_node(act_input, 0) + if isinstance(q_node, torch.fx.Node) and is_quant(q_node): + maybe_pad = get_input_node(q_node, 0) + if ( + isinstance(maybe_pad, torch.fx.Node) + and maybe_pad.target + == exir_ops.edge.aten.constant_pad_nd.default + ): + pad_node = maybe_pad + + if pad_node is not None and len(pad_node.users) == 1: + conv_padding = cast(List[int], node.args[4]) + is_1d = len(conv_padding) == 1 + is_2d = len(conv_padding) == 2 + + pad_value = pad_node.args[2] if len(pad_node.args) > 2 else 0 + if pad_value != 0.0: + return super()._get_act_deps(node, ep, precision) + + pad_amounts = cast(List[int], pad_node.args[1]) + + pad_ok = False + if is_1d: + temporal_only = len(pad_amounts) <= 2 or all( + a == 0 for a in pad_amounts[2:] + ) + pad_ok = ( + len(pad_amounts) % 2 == 0 + and all(a >= 0 for a in pad_amounts) + and temporal_only + ) + elif is_2d: + pad_ok = len(pad_amounts) <= 4 or all( + a == 0 for a in pad_amounts[4:] + ) + + if pad_ok: + valid, deps = super()._get_act_deps(pad_node, ep, precision) + if valid: + return (True, [pad_node, *deps]) return super()._get_act_deps(node, ep, precision) + def supported_precision_types(self): + return [ + ConfigPrecisionType.FP32, + ConfigPrecisionType.STATIC_QUANT, + ConfigPrecisionType.DYNAMIC_QUANT, + ] + class AddmmConfig(GEMMConfig): """ diff --git a/backends/xnnpack/test/ops/test_conv1d.py b/backends/xnnpack/test/ops/test_conv1d.py index 35d9bced512..20dc864bf87 100644 --- a/backends/xnnpack/test/ops/test_conv1d.py +++ b/backends/xnnpack/test/ops/test_conv1d.py @@ -90,6 +90,24 @@ def forward(self, x): z = torch.add(y, z) return z + class Conv1dSamePadding(torch.nn.Module): + def __init__(self, kernel_size: int, dilation: int = 1): + super().__init__() + self.conv1d = torch.nn.Conv1d( + in_channels=2, + out_channels=4, + kernel_size=kernel_size, + dilation=dilation, + padding="same", + bias=True, + ) + + def forward(self, x): + return self.conv1d(x) + + def _get_calibration_samples(self, inputs): + return [tuple(torch.randn_like(inputs[i]) for i in range(len(inputs)))] + def _test_conv1d( self, module, @@ -102,9 +120,7 @@ def _test_conv1d( skip_to_executorch=False, ): calibration_samples = ( - [tuple(torch.randn_like(inputs[i]) for i in range(len(inputs)))] - if quantized - else None + self._get_calibration_samples(inputs) if quantized else None ) tester = ( @@ -160,6 +176,43 @@ def test_qs8_conv1d(self): self.Conv1d(), inputs, 1, quantized=True, dynamic_shape=dynamic_shapes ) + def test_qs8_conv1d_even_kernel_same_padding(self): + inputs = (torch.randn(1, 2, 16),) + configs = [ + (2, 1), + (3, 1), + (4, 1), + (4, 2), + ] + for kernel_size, dilation in configs: + with self.subTest(kernel_size=kernel_size, dilation=dilation): + ( + Tester( + self.Conv1dSamePadding( + kernel_size=kernel_size, dilation=dilation + ), + inputs, + ) + .quantize( + Quantize( + calibration_samples=self._get_calibration_samples(inputs) + ) + ) + .export() + .check_count({"torch.ops.aten.conv1d.padding": 1}) + .to_edge_transform_and_lower() + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_convolution_default", + "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default", + ] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(num_runs=10, atol=0.04, rtol=0.02) + ) + def test_qs8_conv1d_batchnorm_seq(self): inputs = (torch.randn(2, 2, 4),) dynamic_shapes = ({0: torch.export.Dim("batch", min=2, max=10)},)