From b07b56459bec61949d1861e6cc5ef0d4d86e2497 Mon Sep 17 00:00:00 2001 From: Arik Horodniceanu Date: Fri, 10 Jul 2026 14:23:24 -0700 Subject: [PATCH] Qualcomm AI Engine Direct - Adding QNN backend support for sort core ATen op --- backends/qualcomm/_passes/layout_transform.py | 1 + backends/qualcomm/builders/__init__.py | 2 + backends/qualcomm/builders/op_sort.py | 124 ++++++++++++++++++ .../quantizer/annotators/htp_rules.py | 31 ++++- .../quantizer/annotators/lpai_rules.py | 31 ++++- backends/qualcomm/tests/models.py | 39 ++++++ .../qualcomm/tests/rework/htp/op/v68/test.py | 12 ++ backends/qualcomm/tests/rework/src/op.py | 32 +++++ backends/qualcomm/tests/test_qnn_delegate.py | 48 +++++++ 9 files changed, 316 insertions(+), 4 deletions(-) create mode 100644 backends/qualcomm/builders/op_sort.py diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 5b9c13e6ef4..def28a93779 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -124,6 +124,7 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.sigmoid.default, exir_ops.edge.aten.sign.default, exir_ops.edge.aten.slice_copy.Tensor, + exir_ops.edge.aten.sort.default, exir_ops.edge.aten.split_with_sizes.default, exir_ops.edge.aten.split_with_sizes_copy.default, exir_ops.edge.aten.sqrt.default, diff --git a/backends/qualcomm/builders/__init__.py b/backends/qualcomm/builders/__init__.py index dd69515b1c1..e8c38f28bf2 100644 --- a/backends/qualcomm/builders/__init__.py +++ b/backends/qualcomm/builders/__init__.py @@ -99,6 +99,7 @@ op_slice_copy, op_slice_scatter, op_softmax, + op_sort, op_space_to_depth, op_split_with_sizes, op_sqrt, @@ -214,6 +215,7 @@ op_slice_copy, op_slice_scatter, op_softmax, + op_sort, op_space_to_depth, op_split_with_sizes, op_squeeze, diff --git a/backends/qualcomm/builders/op_sort.py b/backends/qualcomm/builders/op_sort.py new file mode 100644 index 00000000000..0333cd71849 --- /dev/null +++ b/backends/qualcomm/builders/op_sort.py @@ -0,0 +1,124 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +import warnings +from typing import cast, Dict + +import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import ( + QCOM_AXIS_ORDER, + QCOM_DATA, + QCOM_QUANT_ATTRS, +) + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpTopK, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Sort(NodeVisitor): + """Implement aten.sort via QNN TopK with k = dim_size. + + Limitations: + - Only supports sorting along the last dimension (channel dim in NHWC). + Sorting on other dimensions will return None (op not delegated). + - Sort indices are output as int32 by QNN. In quantized models, consuming + sort indices with ops like gather (which expect int64) may cause dtype mismatches. + For quantized use cases, prefer consuming only the sorted values output. + """ + + target = ["aten.sort.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], + ) -> PyQnnManager.PyQnnOpWrapper: + + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + nodes_to_wrappers, + ) + + # aten.sort signature: sort(input, dim=-1, descending=False) + dim = cast(int, node.args[1]) if len(node.args) > 1 else -1 + descending = cast(bool, node.args[2]) if len(node.args) > 2 else False + + if dim < 0: + dim = dim % len(input_tensor.shape) + if QCOM_AXIS_ORDER in node.meta: + dim = node.meta[QCOM_AXIS_ORDER].index(dim) + + # Sort is TopK with k = full dimension size + k = input_tensor.shape[dim] + + if dim != len(input_tensor.shape) - 1: + warnings.warn( + "[QNN Delegate Op Builder]: QNN currently only supports channel as dimension for TopK/Sort.", + stacklevel=1, + ) + return + + sort_input_tensors = [input_tensor_wrapper] + + output_val_tensor = self.get_tensor(node, node, 0) + output_idx_tensor = self.get_tensor(node, node, 1).to(torch.int32) + + # QNN constraint, topk/sort output_0 requires having the same quant config as input + node.meta[QCOM_QUANT_ATTRS] = input_node.meta.get(QCOM_QUANT_ATTRS) + output_val_tensor_wrapper = self.define_tensor( + node, + node, + output_val_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + # sort output_1 is index, do not quantize it. + node.meta.pop(QCOM_QUANT_ATTRS, None) + output_index_tensor_wrapper = self.define_tensor( + node, + node, + output_idx_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + wrapper_idx=1, + ) + sort_output_tensors = [output_val_tensor_wrapper, output_index_tensor_wrapper] + + topk_op = PyQnnManager.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpTopK.op_name, + ) + topk_op.AddInputTensors(sort_input_tensors) + topk_op.AddOutputTensors(sort_output_tensors) + + topk_op.AddScalarParam( + OpTopK.param_k, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(k)}, + ) + + # descending=True in sort corresponds to largest=True in TopK + topk_op.AddScalarParam( + OpTopK.param_largest, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_BOOL_8, + {QCOM_DATA: descending}, + ) + + return topk_op diff --git a/backends/qualcomm/quantizer/annotators/htp_rules.py b/backends/qualcomm/quantizer/annotators/htp_rules.py index ca8abb246bf..79ff4c59d56 100644 --- a/backends/qualcomm/quantizer/annotators/htp_rules.py +++ b/backends/qualcomm/quantizer/annotators/htp_rules.py @@ -648,8 +648,11 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: return out_act_quantization_spec = quantization_config.output_activation - # QNN constraint, topk output_0 requires having the same quant config as input - if node.args[0].target == torch.ops.aten.topk.default: + # QNN constraint, topk/sort output_0 requires having the same quant config as input + if node.args[0].target in ( + torch.ops.aten.topk.default, + torch.ops.aten.sort.default, + ): out_act_quantization_spec = SharedQuantizationSpec(node.args[0]) node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( output_qspec=out_act_quantization_spec, @@ -1464,6 +1467,30 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: ) +@register_annotator([torch.ops.aten.sort.default], QnnConstants.OpTopK.op_name) +class Sort(GeneralOpDef): + @staticmethod + def annotate(node: Node, quantization_config: QuantizationConfig) -> None: + if _is_annotated([node]): + return + + input_qspec_map = {} + input_act_qspec = quantization_config.input_activation + out_act_quantization_spec = None + if input_act_qspec is not None: + if _is_float_tensor(node.args[0]): + input_act = node.args[0] + assert isinstance(input_act, Node) + input_qspec_map[input_act] = input_act_qspec + out_act_quantization_spec = SharedQuantizationSpec((input_act, node)) + + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=out_act_quantization_spec, + _annotated=True, + ) + + @register_annotator( [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default], QnnConstants.OpSigmoid.op_name, diff --git a/backends/qualcomm/quantizer/annotators/lpai_rules.py b/backends/qualcomm/quantizer/annotators/lpai_rules.py index 6e5b343c5c7..83ffaa96094 100644 --- a/backends/qualcomm/quantizer/annotators/lpai_rules.py +++ b/backends/qualcomm/quantizer/annotators/lpai_rules.py @@ -414,8 +414,11 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: return out_act_quantization_spec = quantization_config.output_activation - # QNN constraint, topk output_0 requires having the same quant config as input - if node.args[0].target == torch.ops.aten.topk.default: + # QNN constraint, topk/sort output_0 requires having the same quant config as input + if node.args[0].target in ( + torch.ops.aten.topk.default, + torch.ops.aten.sort.default, + ): out_act_quantization_spec = SharedQuantizationSpec(node.args[0]) node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( output_qspec=out_act_quantization_spec, @@ -911,6 +914,30 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: ) +@register_annotator([torch.ops.aten.sort.default], QnnConstants.OpTopK.op_name) +class Sort(GeneralOpDef): + @staticmethod + def annotate(node: Node, quantization_config: QuantizationConfig) -> None: + if _is_annotated([node]): + return + + input_qspec_map = {} + input_act_qspec = quantization_config.input_activation + out_act_quantization_spec = None + if input_act_qspec is not None: + if _is_float_tensor(node.args[0]): + input_act = node.args[0] + assert isinstance(input_act, Node) + input_qspec_map[input_act] = input_act_qspec + out_act_quantization_spec = SharedQuantizationSpec((input_act, node)) + + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=out_act_quantization_spec, + _annotated=True, + ) + + @register_annotator( [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default], QnnConstants.OpSigmoid.op_name, diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 6decea516c0..9d40827ebc6 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -885,6 +885,18 @@ def forward(self, x): return topk_values +class Conv2dSort(torch.nn.Module): + def __init__(self, descending=True): + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3) + self.descending = descending + + def forward(self, x): + x = self.conv(x) + sort_values, sort_indices = torch.sort(x, dim=1, descending=self.descending) + return sort_values + + class Conv2dUnbind(torch.nn.Module): def __init__(self): super().__init__() @@ -2604,6 +2616,33 @@ def forward(self, x): return torch.nn.functional.softmax(x, dim=self.dim) +class Sort(torch.nn.Module): + def __init__(self, dim=-1, descending=True): + super().__init__() + self.dim = dim + self.descending = descending + + def forward(self, x): + values, indices = torch.sort(x, dim=self.dim, descending=self.descending) + return values + + +class SortAndIndex(torch.nn.Module): + def __init__(self, shape, dim=-1, descending=True): + super().__init__() + self.idx_source = torch.rand(*shape) + self.dim = dim + self.descending = descending + self.rank = len(shape) + + def forward(self, x): + normalized_dim = self.dim if self.dim >= 0 else self.rank + self.dim + + values, indices = torch.sort(x, dim=normalized_dim, descending=self.descending) + gathered = torch.gather(self.idx_source, normalized_dim, indices) + return values + gathered + + class Sqrt(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/rework/htp/op/v68/test.py b/backends/qualcomm/tests/rework/htp/op/v68/test.py index 19a7843de65..3ed8d946754 100644 --- a/backends/qualcomm/tests/rework/htp/op/v68/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v68/test.py @@ -1219,6 +1219,18 @@ def test_softmax(request, kwargs): Softmax.test(request, kwargs) # noqa: F405 +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + @enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_split(request, kwargs): diff --git a/backends/qualcomm/tests/rework/src/op.py b/backends/qualcomm/tests/rework/src/op.py index 14d853ed919..5c2d608715d 100644 --- a/backends/qualcomm/tests/rework/src/op.py +++ b/backends/qualcomm/tests/rework/src/op.py @@ -4323,6 +4323,38 @@ def test(subtests, qnn_config, quantizer, compile_spec, expected): ) +class Sort(torch.nn.Module): + def __init__(self, dim, descending): + super().__init__() + self.dim = dim + self.descending = descending + + def forward(self, x): + sorted_values, sorted_indices = torch.sort( + x, dim=self.dim, descending=self.descending + ) + return sorted_values + + @staticmethod + @unpack_fixtures + def test(subtests, qnn_config, quantizer, compile_spec, expected): + inputs = (torch.randn(1, 2, 3, 4),) + # QNN only supports sort along the last dimension + dims = [-1] + descendings = [True, False] + for dim, descending in itertools.product(dims, descendings): + with subtests.test(msg=f"dim:{dim}, descending:{descending}"): + with expected as metrics: + export_and_verify( + module=__class__(dim=dim, descending=descending), + inputs=inputs, + qnn_config=qnn_config, + quantizer=quantizer, + compile_specs=compile_spec, + metrics=metrics, + ) + + class Split(torch.nn.Module): def __init__(self, split_size_or_sections, dim): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index cbb7c82d111..53f8d75c2ea 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -2358,6 +2358,43 @@ def test_qnn_backend_softmax(self): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_sort(self): + modules = [ + Conv2dSort(descending=True), # noqa: F405 + Conv2dSort(descending=False), # noqa: F405 + ] + sample_input = (torch.randn(1, 3, 32, 32),) + for i, module in enumerate(modules): + with self.subTest(i=i): + self.lower_module_and_test_output(module, sample_input) + + def test_qnn_backend_sort_and_index(self): + test_comb = [ + { + QCOM_MODULE: SortAndIndex( # noqa: F405 + shape=(3, 10), dim=-1, descending=True + ), + QCOM_SAMPLE_INPUTS: (torch.randn(3, 10),), + }, + { + QCOM_MODULE: SortAndIndex( # noqa: F405 + shape=(2, 4, 8), dim=-1, descending=True + ), + QCOM_SAMPLE_INPUTS: (torch.randn(2, 4, 8),), + }, + { + QCOM_MODULE: SortAndIndex( # noqa: F405 + shape=(1, 4, 8, 10), dim=-1, descending=True + ), + QCOM_SAMPLE_INPUTS: (torch.randn(1, 4, 8, 10),), + }, + ] + for i, test in enumerate(test_comb): + with self.subTest(i=i): + self.lower_module_and_test_output( + test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS] + ) + def test_qnn_backend_squared_relu(self): module = SquaredReLU() # noqa: F405 sample_input = (torch.randn([2, 5, 1, 3]),) @@ -5497,6 +5534,17 @@ def test_qnn_backend_softmax(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_sort(self): + modules = [ + Conv2dSort(descending=True), # noqa: F405 + Conv2dSort(descending=False), # noqa: F405 + ] + sample_input = (torch.randn(1, 3, 32, 32),) + for i, module in enumerate(modules): + with self.subTest(i=i): + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(qdq_module, sample_input) + def test_qnn_backend_squared_relu(self): module = SquaredReLU() # noqa: F405 sample_input = (torch.randn([2, 5, 1, 3]),)