Qualcomm AI Engine Direct - Adding QNN backend support for the sort core ATen op

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,

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,

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

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,

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,

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__()

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):

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__()