diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py index 239a424965b..1796b5bdbbe 100644 --- a/backends/xnnpack/partition/config/gemm_configs.py +++ b/backends/xnnpack/partition/config/gemm_configs.py @@ -25,7 +25,6 @@ is_per_tensor, is_qparam, is_quant, - tag_as_implicit_q_dq, ) from executorch.backends.xnnpack.utils.utils import ( get_input_node, @@ -139,11 +138,6 @@ def get_deps( valid_deps = valid_bias and valid_weight and valid_act and valid_output deps = list(chain(bias_deps, weight_deps, act_deps, output_deps)) - # Tag q/dq nodes as implicit q/dq nodes - for dep in deps: - if is_dequant(dep) or is_quant(dep): - tag_as_implicit_q_dq(dep) - return valid_deps, deps def _get_weight_deps( @@ -224,9 +218,12 @@ def _get_output_deps( node_users = list(node.users.keys()) if len(node_users) == 1: n_output = node_users[0] + # After a decomposed linear, the activation can have more dims than + # XNNPACK allows even when the 2-D matmul does not. Leave it out then. if ( n_output.op == "call_function" and format_target_name(n_output.target.__name__) in self.fused_acts + and self._check_tensor_ranks(n_output) ): gemm_deps.append(n_output) @@ -570,17 +567,23 @@ def find_partition_args(input_node): node.args = tuple(fake_args) node.users = src_partition.output_nodes[0].users valid_deps, deps = super().get_deps(node, ep) + _, precision = self._overwrite_precision(node) # Reset addmm node back to old args and users node.args = old_args node.users = old_users - # When using force_non_static_weights_for_f32_linear, we want to get_deps to overwrite the source partition nodes. - # Else we want to be greedy. + # A float linear with force_non_static_weights_for_f32_linear takes only the + # deps inside the source partition. Else we want to be greedy, but leave out + # views XNNPACK cannot hold. ret_deps = ( list(set(deps) & set(src_partition.nodes)) if self.force_non_static_weights_for_f32_linear - else list(set(deps) | set(src_partition.nodes)) + and precision == ConfigPrecisionType.FP32 + else list( + set(deps) + | {n for n in src_partition.nodes if self._check_tensor_ranks(n)} + ) ) return valid_deps, ret_deps diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index 3c87e47c656..a9a323d9453 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -16,11 +16,7 @@ ConfigPrecisionType, XNNPartitionerConfig, ) -from executorch.backends.xnnpack.utils.quant_utils import ( - is_dequant, - is_quant, - tag_as_implicit_q_dq, -) +from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant from executorch.backends.xnnpack.utils.utils import ( get_input_node, get_param_tensor, @@ -165,15 +161,6 @@ def get_node_and_deps( # Expected node --> fused_act (optional) --> dequant return deps - # Tag input nodes (dq nodes) as implicit q/dq nodes - for dq_input in node.all_input_nodes: - if is_dequant(dq_input): - tag_as_implicit_q_dq(dq_input) - - # Tag node_output (q node) as an implicit q/dq node - if is_quant(node_output): - tag_as_implicit_q_dq(node_output) - quantized_deps.append(node_output) return deps + quantized_deps diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index e0a575e000a..b5267d2405a 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/backends/xnnpack/partition/config/xnnpack_config.py @@ -14,7 +14,9 @@ is_dequant, is_qparam, is_quant, + tag_as_implicit_q_dq, ) +from executorch.backends.xnnpack.utils.xnnpack_constants import XNN_MAX_TENSOR_DIMS from executorch.exir.backend.canonical_partitioners.config_partitioner import ( format_target_name, PartitionerConfig, @@ -66,14 +68,31 @@ def get_partition( with the name of the XNNPartitionerConfig class which return this set of nodes. This enforces that all partitions returned by XNNPartitioner configs are labeled with the partitioner config which returned them + + Returns an empty list if any of the nodes has more dims than XNNPACK supports. + Otherwise marks every quantize and dequantize node in the partition, other + than the given node, as implicit, so that it is folded into the op it feeds + or follows. """ partitioned_nodes = self.get_node_and_deps(node, ep) + # A partition can take nodes that their own configs reject. + for n in partitioned_nodes: + if not self._check_tensor_ranks(n): + why(node, reason=f"its partition takes {n}, which has too many dims") + return [] + + # Fold only now, so a rejected node keeps its q/dq nodes as separate ops. A + # q/dq node partitioned on its own stays a separate op as well. + for n in partitioned_nodes: + if n is not node and (is_quant(n) or is_dequant(n)): + tag_as_implicit_q_dq(n) + # label partitioned nodes with the name of the partitioner config - for node in partitioned_nodes: - if "xnn_partitioner_config" in node.meta: - node.meta["xnn_partitioner_config"].append(self.__class__.__name__) + for n in partitioned_nodes: + if "xnn_partitioner_config" in n.meta: + n.meta["xnn_partitioner_config"].append(self.__class__.__name__) else: - node.meta["xnn_partitioner_config"] = [self.__class__.__name__] + n.meta["xnn_partitioner_config"] = [self.__class__.__name__] return partitioned_nodes @@ -94,7 +113,9 @@ def get_node_and_deps( ) -> List[torch.fx.Node]: """ Takes in a node and its exported program and returns a list of nodes - and its dependencies that need to be partitioned together + and its dependencies that need to be partitioned together. Quantize and + dequantize nodes in the list, other than the node itself, are folded into + the partition's op once get_partition takes it. Args: node: Node to be partitioned @@ -152,6 +173,20 @@ def check_common_constraints( why(node, reason="invalid dtype") return False + return self._check_tensor_ranks(node) + + def _check_tensor_ranks(self, node: torch.fx.Node) -> bool: + vals = [arg.meta.get("val") for arg in node.all_input_nodes] + node_val = node.meta.get("val") + vals.extend(node_val if isinstance(node_val, (list, tuple)) else [node_val]) + for val in vals: + if isinstance(val, torch.Tensor) and val.dim() > XNN_MAX_TENSOR_DIMS: + why( + node, + reason=f"XNNPACK supports at most {XNN_MAX_TENSOR_DIMS} dims, got {val.dim()}", + ) + return False + return True def _check_inputs_are_valid_dtypes(self, node, valid_dtypes): diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer.py b/backends/xnnpack/quantizer/xnnpack_quantizer.py index b84a15901fe..db4f4796942 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer.py @@ -257,8 +257,6 @@ class QuantPattern: torch.ops.aten.conv_transpose1d.default, torch.ops.aten.conv_transpose2d, torch.ops.aten.conv_transpose2d.input, - torch.ops.aten.conv_transpose3d, - torch.ops.aten.conv_transpose3d.input, } LINEAR_TARGETS = { diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index 8868d742b7d..1602c30561a 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -240,6 +240,11 @@ def _annotate_linear_relu( return annotated_partitions +def _is_1d_or_2d_conv(weight_shape: Optional[torch.Size]) -> bool: + # XNNPACK runs no other convs. + return weight_shape is None or len(weight_shape) <= 4 + + def _do_annotate_conv( gm: torch.fx.GraphModule, quantization_config: Optional[QuantizationConfig], @@ -274,9 +279,11 @@ def _do_annotate_conv( assert isinstance(weight, Node) weight_qspec = get_weight_qspec(quantization_config) num_groups = get_groups_from_conv(conv_node) + weight_shape = getattr(weight.meta.get("val", None), "shape", None) # skip if transposed conv has more than 1 group skip = skip or (is_conv_transpose and num_groups != 1) + skip = skip or not _is_1d_or_2d_conv(weight_shape) if is_conv_transpose: # transposed convs per output channel quantization @@ -291,8 +298,6 @@ def _do_annotate_conv( # Only annotate dynamically quantized conv if it's 2D and not depthwise if is_dynamic: - weight_val = weight.meta.get("val", None) - weight_shape = getattr(weight_val, "shape", None) # Skip if not a 4D weight tensor (i.e. not conv2d) skip = skip or (weight_shape is not None and len(weight_shape) != 4) # Skip if depthwise (default to groups=1 since it's not an arg) @@ -356,6 +361,7 @@ def _do_annotate_conv_relu( assert isinstance(weight, Node) weight_qspec = get_weight_qspec(quantization_config) groups = get_groups_from_conv(conv_node) + weight_shape = getattr(weight.meta.get("val", None), "shape", None) if is_conv_transpose: # transposed convs per output channel quantization weight_qspec = change_quantization_config(weight_qspec, ch_axis=1) @@ -371,7 +377,7 @@ def _do_annotate_conv_relu( if _is_annotated(partition): continue - if is_conv_transpose and groups != 1: + if (is_conv_transpose and groups != 1) or not _is_1d_or_2d_conv(weight_shape): continue if filter_fn and not filter_fn(conv_node): diff --git a/backends/xnnpack/test/BUCK b/backends/xnnpack/test/BUCK index 0eb2af1bf9b..3e28be745d0 100644 --- a/backends/xnnpack/test/BUCK +++ b/backends/xnnpack/test/BUCK @@ -96,6 +96,7 @@ fbcode_target(_kind = runtime.python_test, "test_xnnpack_utils.py", ], supports_static_listing = False, + preload_deps = ["//executorch/kernels/quantized:aot_lib"], deps = [ "fbsource//third-party/pypi/parameterized:parameterized", "//executorch/backends/test/harness:tester", diff --git a/backends/xnnpack/test/ops/test_conv2d.py b/backends/xnnpack/test/ops/test_conv2d.py index 0db84852422..2daacfd44cf 100644 --- a/backends/xnnpack/test/ops/test_conv2d.py +++ b/backends/xnnpack/test/ops/test_conv2d.py @@ -6,6 +6,7 @@ import itertools import unittest +from dataclasses import replace from typing import Optional import torch @@ -355,6 +356,34 @@ def test_fp32_conv3d_single_element_stride_doesnt_partition(self) -> None: .run_method_and_compare_outputs() ) + def test_qs8_quantizer_skips_conv3d(self) -> None: + # XNNPACK has no conv3d support, so its quantizer leaves 3d convs in float, + # with or without a fused relu. The same model in 2d gets its activations + # and weights quantized. The weights are quantized ahead of time, so only + # their dequantize nodes remain. + quant_config = get_symmetric_quantization_config(is_per_channel=True) + ops = "torch.ops.quantized_decomposed" + for dims, conv, conv_transpose, quantized_activations, quantized_weights in ( + (2, torch.nn.Conv2d, torch.nn.ConvTranspose2d, 4, 3), + (3, torch.nn.Conv3d, torch.nn.ConvTranspose3d, 0, 0), + ): + m = torch.nn.Sequential( + conv(2, 2, 2), torch.nn.ReLU(), conv(2, 2, 2), conv_transpose(2, 2, 2) + ) + ( + Tester(m.eval(), (torch.randn([1, 2] + [4] * dims),)) + .quantize(Quantize(quantization_config=quant_config)) + .export() + .check_count( + { + f"{ops}.quantize_per_tensor.default": quantized_activations, + f"{ops}.dequantize_per_tensor.default": quantized_activations, + f"{ops}.quantize_per_channel.default": 0, + f"{ops}.dequantize_per_channel.default": quantized_weights, + } + ) + ) + def test_fp32_conv2d_permute(self) -> None: for transpose in (True, False): for perm_order in list(itertools.permutations([0, 1, 2, 3])): @@ -375,6 +404,26 @@ def test_qs8_conv2d_per_channel(self) -> None: quant_config=get_symmetric_quantization_config(is_per_channel=True), ) + def test_qs8_conv2d_float_output(self) -> None: + # XNNPACK leaves out a quantized conv whose output is not quantized. The + # quantize and dequantize of its input then run in a delegate of their own. + quant_config = replace( + get_symmetric_quantization_config(is_per_channel=True), + output_activation=None, + ) + m = Conv2d() + ( + Tester(m.eval(), m.get_inputs()) + .quantize(Quantize(quantization_config=quant_config)) + .export() + .to_edge_transform_and_lower() + .check(["executorch_exir_dialects_edge__ops_aten_convolution_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_qs8_conv2d_even_kernel_same_padding(self) -> None: # An even-kernel 'same'-padding conv decomposes into # dequant -> constant_pad_nd -> convolution. The pad and conv are pulled @@ -832,14 +881,27 @@ def get_inputs(self): conv_count=1, ) else: - with self.assertRaises(RuntimeError): - self._test( - PerChannelConvTranspose2d( - 3 * groups, 5 * groups, groups, ch_axis - ), # ch_axis=0 - quant_config=None, - conv_count=1, + # XNNPACK leaves this conv out, and the quantize and dequantize + # nodes around it then run in delegates of their own. + m = PerChannelConvTranspose2d( + 3 * groups, 5 * groups, groups, ch_axis + ) + ( + Tester(m.eval(), m.get_inputs()) + .export() + .to_edge_transform_and_lower() + .check( + [ + "executorch_exir_dialects_edge__ops_aten_convolution_default" + ] ) + .check_count( + {"torch.ops.higher_order.executorch_call_delegate": 2} + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=1) + ) def test_fp32_tconv_output_padding(self): """Test transposed convolution with non-zero output padding.""" diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index 7e4cb4ab586..e4f7448e4be 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -6,6 +6,7 @@ # pyre-unsafe import unittest +from dataclasses import replace from itertools import product from typing import Callable, Dict, List, Optional, Tuple @@ -33,6 +34,13 @@ from torch.export.graph_signature import ExportGraphSignature, InputKind +# Buck preloads the quantized kernels instead of shipping this module. +try: + import executorch.extension.pybindings.portable_lib # noqa[F401] + import executorch.kernels.quantized # noqa[F401] +except ImportError: + pass + try: from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import ( @@ -840,6 +848,26 @@ def test_qs8_linear(self): quant_type="per_tensor", ) + def test_qs8_linear_float_output(self): + # XNNPACK leaves out a quantized linear whose output is not quantized. The + # quantize and dequantize of its input then run in a delegate of their own. + quant_config = replace( + get_symmetric_quantization_config(is_per_channel=True), + output_activation=None, + ) + module = BaseLinear(use_bias=True) + ( + Tester(module.eval(), module.get_inputs(rank=2)) + .quantize(Quantize(quantization_config=quant_config)) + .export() + .to_edge_transform_and_lower() + .check(["executorch_exir_dialects_edge__ops_aten_addmm_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + # Tests for q[dp]8-f16-qc8w def test_qd8_f16_per_channel_linear(self): self._test_qd8_per_channel_linear(dtype=torch.half) @@ -1152,6 +1180,32 @@ def test_linear_qd8_as_fp32(self): }, ) + def test_quantized_linear_with_force_non_static_weights_for_f32_linear(self): + # The flag is for float linears. A quantized linear that is decomposed into + # a matmul keeps its quantize and dequantize nodes and is delegated whole. + partitioner = XnnpackPartitioner(force_non_static_weights_for_f32_linear=True) + for is_dynamic in (False, True): + module = BaseLinear(use_bias=True) + quant_config = get_symmetric_quantization_config( + is_per_channel=True, is_dynamic=is_dynamic + ) + ( + Tester(module.eval(), module.get_inputs(rank=2)) + .quantize(Quantize(quantization_config=quant_config)) + .export() + .to_edge() + .partition(Partition(partitioner=partitioner)) + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_addmm_default", + "quantized_decomposed", + ] + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(atol=5e-2) + ) + def test_linear_with_force_non_static_weights_for_f32_linear(self): def check_signature( signature: ExportGraphSignature, diff --git a/backends/xnnpack/test/ops/test_permute.py b/backends/xnnpack/test/ops/test_permute.py index e0171f2178f..055108ab472 100644 --- a/backends/xnnpack/test/ops/test_permute.py +++ b/backends/xnnpack/test/ops/test_permute.py @@ -7,7 +7,22 @@ import unittest import torch -from executorch.backends.xnnpack.test.tester import Tester +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, +) +from executorch.backends.xnnpack.test.tester import ( + Quantize, + Tester, + ToEdgeTransformAndLower, +) + +# Buck preloads the quantized kernels instead of shipping this module. +try: + import executorch.extension.pybindings.portable_lib # noqa[F401] + import executorch.kernels.quantized # noqa[F401] +except ImportError: + pass class TestPermute(unittest.TestCase): @@ -83,6 +98,78 @@ def test_fp32_permute_copy(self): .run_method_and_compare_outputs() ) + class LinearPermute8d(torch.nn.Module): + def __init__(self, relu=False): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + self.relu = relu + + def forward(self, x): + y = self.linear(x.reshape(2, 2, 2, 2, 2, 2, 2, 4)) + if self.relu: + y = torch.relu(y) + return y.permute(0, 2, 5, 3, 6, 1, 4, 7).reshape(2, 64, -1) * 2 + + def test_fp32_permute_max_rank(self): + # XNNPACK allows at most 6 dims, so a 6-D permute is delegated and a 7-D + # one is not. + delegate = "torch.ops.higher_order.executorch_call_delegate" + permute = "executorch_exir_dialects_edge__ops_aten_permute_copy_default" + for rank, delegates in ((6, 1), (7, 0)): + dims = list(reversed(range(rank))) + ( + Tester(self.Permute(dims), (torch.randn([2] * rank),)) + .export() + .to_edge_transform_and_lower() + .check_count({delegate: delegates, permute: 1 - delegates}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def _test_permute_above_max_rank(self, quant_config=None, relu=False): + # A delegate that holds a tensor above 6 dims fails to load, so the 8-D + # views and permute stay outside and the mul is delegated. The linear's + # 2-D matmul is delegated too, unless its activation is quantized in 8-D. + for force_non_static_weights in (False, True): + tester = Tester(self.LinearPermute8d(relu).eval(), (torch.randn(2, 64, 4),)) + if quant_config is not None: + tester.quantize(Quantize(quantization_config=quant_config)) + partitioner = XnnpackPartitioner( + force_non_static_weights_for_f32_linear=force_non_static_weights + ) + tester.export().to_edge_transform_and_lower( + ToEdgeTransformAndLower([partitioner]) + ) + addmm = ["executorch_exir_dialects_edge__ops_aten_addmm_default"] + if quant_config is None: + tester.check_not(addmm).check_count( + {"torch.ops.higher_order.executorch_call_delegate": 2} + ) + else: + tester.check(addmm) + ( + tester.check_count( + {"executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1} + ) + .check_not(["executorch_exir_dialects_edge__ops_aten_mul_Tensor"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(qtol=1) + ) + + def test_fp32_permute_above_max_rank(self): + for relu in (False, True): + self._test_permute_above_max_rank(relu=relu) + + def test_qs8_permute_above_max_rank(self): + self._test_permute_above_max_rank(get_symmetric_quantization_config()) + + def test_qd8_permute_above_max_rank(self): + self._test_permute_above_max_rank( + get_symmetric_quantization_config(is_per_channel=True, is_dynamic=True) + ) + def test_qs8_permute(self): inputs = (torch.randn(1, 1, 4, 4),) (