[QNN-EP] Support non-last axis TopK. (microsoft#24881)

minfhong-quic · javier-intel · commit 7f7f45baa67f · 2025-06-14T22:46:44.000-07:00
### Description
In TopK op builder, add Transpose around TopK to permute the axis to the last before and permute back after.
Additionally, since TopK's second output is indices which may have INT64 dtype, add Cast to cast transformed INT32 back to INT64 if is graph output.

### Motivation and Context
QNN only accepts TopK on the last axis but ONNX/ORT's TopK has axis attribute. Complement TopK op builder to avoid falling back to CPU for non-last axis TopK.
diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/softmax_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/softmax_op_builder.cc
@@ -56,22 +56,6 @@ std::vector<uint32_t> FlattenShapeFromAxis(const std::vector<uint32_t>& input_sh
   return output_shape;
 }
 
-std::vector<uint32_t> GetTransposePermToUseLastAxis(uint32_t input_rank, uint32_t axis) {
-  assert(axis < input_rank);
-  std::vector<uint32_t> transpose_perm;
-  transpose_perm.reserve(input_rank);
-
-  for (uint32_t dim = 0; dim < input_rank; dim++) {
-    transpose_perm.push_back(dim);
-  }
-
-  // Swap axis dim with last dim.
-  transpose_perm[axis] = input_rank - 1;
-  transpose_perm[input_rank - 1] = axis;
-
-  return transpose_perm;
-}
-
 Status SoftmaxOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
                                        const NodeUnit& node_unit,
                                        const logging::Logger& logger,
@@ -131,8 +115,10 @@ Status SoftmaxOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
     QNN EP is able to support arbitrary axis attribute by wrapping transposes around the operator.
     */
     std::string transpose_output_name = input_name + "_ort_qnn_ep_transpose";
-    std::vector<uint32_t> transpose_perm = GetTransposePermToUseLastAxis(static_cast<uint32_t>(input_rank),
-                                                                         static_cast<uint32_t>(axis));
+    std::vector<uint32_t> transpose_perm;
+    ORT_RETURN_IF_ERROR(utils::GetPermToLastAxis(static_cast<uint32_t>(axis),
+                                                 static_cast<uint32_t>(input_rank),
+                                                 transpose_perm));
 
     std::vector<uint32_t> transpose_output_shape = input_info.shape;
     transpose_output_shape[input_rank - 1] = input_info.shape[axis];
@@ -243,8 +229,10 @@ Status SoftmaxOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_
                       "Failed to add node.");
 
     const bool is_graph_output = qnn_model_wrapper.IsGraphOutput(orig_output_name);
-    std::vector<uint32_t> transpose_perm = GetTransposePermToUseLastAxis(static_cast<uint32_t>(output_rank),
-                                                                         static_cast<uint32_t>(axis));
+    std::vector<uint32_t> transpose_perm;
+    ORT_RETURN_IF_ERROR(utils::GetPermToLastAxis(static_cast<uint32_t>(axis),
+                                                 static_cast<uint32_t>(output_rank),
+                                                 transpose_perm));
 
     ORT_RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(),
                                                            transpose_input_name,
diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/topk.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/topk.cc
@@ -1,12 +1,21 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
 #include "core/providers/qnn/builder/opbuilder/base_op_builder.h"
 #include "core/providers/qnn/builder/op_builder_factory.h"
 #include "core/providers/qnn/builder/qnn_utils.h"
+
 namespace onnxruntime {
 namespace qnn {
+
 const int TOPK_MIN_INPUT = 2;
 const int TOPK_MAX_INPUT = 2;
+
 class TopKOpBuilder : public BaseOpBuilder {
  public:
   TopKOpBuilder() : BaseOpBuilder("TopKOpBuilder") {}
@@ -41,8 +50,11 @@ class TopKOpBuilder : public BaseOpBuilder {
 
 Status TopKOpBuilder::ExplictOpCheck(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit) const {
   size_t input_count = node_unit.Inputs().size();
+  size_t output_count = node_unit.Outputs().size();
   ORT_RETURN_IF_NOT(input_count >= TOPK_MIN_INPUT && input_count <= TOPK_MAX_INPUT,
                     "For ONNX TopK operation the expected number of inputs is 2.");
+  ORT_RETURN_IF_NOT(output_count == 2, "QNN TopK expects exactly 2 outputs.");
+
   // Skip the first input. The second input needs to be an initializer.
   const auto& input_1 = node_unit.Inputs()[1].node_arg.Name();
   if (!qnn_model_wrapper.IsConstantInput(input_1)) {
@@ -57,14 +69,6 @@ Status TopKOpBuilder::ExplictOpCheck(QnnModelWrapper& qnn_model_wrapper, const N
   if (0 == largest) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "QNN TopK output is always largest values");
   }
-  auto& input_0 = node_unit.Inputs()[0];
-  std::vector<uint32_t> input_shape;
-  ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(input_0.node_arg, input_shape), "Cannot get shape");
-  auto rank = input_shape.size();
-  auto axis = node_helper.Get("axis", -1);
-
-  ORT_RETURN_IF_NOT(axis == -1 || axis == static_cast<int32_t>(rank - 1),
-                    "QNN TopK's axis is always the last dimension");
 
   return Status::OK();
 }
@@ -81,6 +85,40 @@ Status TopKOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
   const auto& inputs = node_unit.Inputs();
   ORT_RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, inputs[0], logger, input_names));
 
+  // HTP only supports TopK at the last axis, and thus check whether extra Transpose is required.
+  TensorInfo input_info = {};
+  ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(node_unit.Inputs()[0], input_info));
+
+  size_t input_rank = input_info.shape.size();
+  int32_t axis = NodeAttrHelper(node_unit).Get("axis", -1);
+  if (axis == -1 || axis == static_cast<int32_t>(input_rank - 1)) {
+    return Status::OK();
+  }
+
+  // Add Transpose to permute axis to the last.
+  std::string transpose_output_name = input_names[0] + "_ort_qnn_ep_transpose";
+  std::vector<uint32_t> transpose_perm;
+  ORT_RETURN_IF_ERROR(utils::GetPermToLastAxis(static_cast<uint32_t>(axis),
+                                               static_cast<uint32_t>(input_rank),
+                                               transpose_perm));
+
+  std::vector<uint32_t> transpose_output_shape = input_info.shape;
+  transpose_output_shape[input_rank - 1] = input_info.shape[axis];
+  transpose_output_shape[axis] = input_info.shape[input_rank - 1];
+
+  ORT_RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(),
+                                                         input_names[0],
+                                                         transpose_output_name,
+                                                         input_info.shape,
+                                                         transpose_perm,
+                                                         transpose_output_shape,
+                                                         input_info.qnn_data_type,
+                                                         input_info.quant_param,
+                                                         do_op_validation,
+                                                         false,
+                                                         false));
+  input_names[0] = transpose_output_name;
+
   return Status::OK();
 }
 
@@ -108,9 +146,125 @@ Status TopKOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wra
   std::string k_param_name = k_param.GetParamTensorName();
   qnn_model_wrapper.AddParamWrapper(std::move(k_param));
   std::vector<std::string> param_tensor_names{k_param_name};
-  ORT_RETURN_IF_ERROR(ProcessOutputs(qnn_model_wrapper, node_unit, std::move(input_names),
-                                     std::move(param_tensor_names), logger, do_op_validation,
-                                     GetQnnOpType(node_unit.OpType())));
+
+  // HTP only supports TopK at the last axis, and thus check whether extra Transpose is required.
+  TensorInfo input_info = {};
+  ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(node_unit.Inputs()[0], input_info));
+
+  size_t input_rank = input_info.shape.size();
+  int32_t axis = NodeAttrHelper(node_unit).Get("axis", -1);
+  if (axis == -1 || axis == static_cast<int32_t>(input_rank - 1)) {
+    ORT_RETURN_IF_ERROR(ProcessOutputs(qnn_model_wrapper,
+                                       node_unit,
+                                       std::move(input_names),
+                                       std::move(param_tensor_names),
+                                       logger,
+                                       do_op_validation,
+                                       GetQnnOpType(node_unit.OpType())));
+    return Status::OK();
+  }
+
+  const auto& outputs = node_unit.Outputs();
+  std::vector<std::string> transpose_input_names;
+  std::vector<std::vector<std::uint32_t>> transpose_input_shapes;
+
+  // Add TopK outputs.
+  for (size_t output_idx = 0; output_idx < 2; ++output_idx) {
+    const auto& output = outputs[output_idx];
+
+    // Since user may not be aware of the additional Transpose, the original output name of TopK node must be used by
+    // the additional Transpose node which has the same output as original TopK node.
+    const std::string& output_name = output.node_arg.Name();
+    std::string transpose_input_name = output_name + "_ort_qnn_ep_transpose";
+    transpose_input_names.push_back(std::move(transpose_input_name));
+
+    // Since the input of TopK node is permuted, its output shape must be manually calculated.
+    TensorInfo output_info = {};
+    ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(output, output_info));
+    size_t output_rank = output_info.shape.size();
+
+    std::vector<uint32_t> transpose_input_shape = output_info.shape;
+    transpose_input_shape[output_rank - 1] = output_info.shape[axis];
+    transpose_input_shape[axis] = output_info.shape[output_rank - 1];
+    transpose_input_shapes.push_back(std::move(transpose_input_shape));
+
+    QnnTensorWrapper output_tensorwrapper(transpose_input_names[output_idx],
+                                          QNN_TENSOR_TYPE_NATIVE,
+                                          output_info.qnn_data_type,
+                                          output_info.quant_param.Copy(),
+                                          std::vector<uint32_t>(transpose_input_shapes[output_idx]));
+    ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(output_tensorwrapper)), "Failed to add tensor.");
+  }
+
+  // Add TopK node.
+  ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(utils::GetNodeName(node_unit),
+                                                    QNN_OP_PACKAGE_NAME_QTI_AISW,
+                                                    GetQnnOpType(node_unit.OpType()),
+                                                    std::move(input_names),
+                                                    std::vector<std::string>(transpose_input_names),
+                                                    std::move(param_tensor_names)),
+                    "Failed to add node.");
+
+  // Add Transpose nodes for each output to permute back.
+  for (size_t output_idx = 0; output_idx < 2; ++output_idx) {
+    const auto& output = outputs[output_idx];
+    const std::string& output_name = output.node_arg.Name();
+
+    TensorInfo output_info = {};
+    ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(output, output_info));
+    size_t output_rank = output_info.shape.size();
+
+    std::vector<uint32_t> transpose_perm;
+    ORT_RETURN_IF_ERROR(utils::GetPermToLastAxis(static_cast<uint32_t>(axis),
+                                                 static_cast<uint32_t>(output_rank),
+                                                 transpose_perm));
+
+    std::string transpose_output_name = output_name;
+    bool is_graph_output = qnn_model_wrapper.IsGraphOutput(output_name);
+
+    // TopK's second output is indices which could be INT64 dtype, and QnnTensorWrapper directly changes the dtype to
+    // INT32 during the wrapper construction. Nevertheless, if this output happens to be graph output, an additional
+    // Cast must be added to cast dtype from INT32 back to INT64.
+    bool is_cast_required = output_idx == 1 && output_info.qnn_data_type == QNN_DATATYPE_INT_64 && is_graph_output;
+    std::string cast_input_name = "";
+    if (is_cast_required) {
+      cast_input_name = transpose_output_name + "_ort_qnn_ep_cast";
+      // For the same reason described above, the original output name is now used by this Cast.
+      transpose_output_name = cast_input_name;
+      // Since additional Cast is added, below Transpose is no longer graph output.
+      is_graph_output = false;
+    }
+
+    ORT_RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(),
+                                                           transpose_input_names[output_idx],
+                                                           transpose_output_name,
+                                                           transpose_input_shapes[output_idx],
+                                                           transpose_perm,
+                                                           output_info.shape,
+                                                           output_info.qnn_data_type,
+                                                           output_info.quant_param,
+                                                           do_op_validation,
+                                                           false,
+                                                           is_graph_output));
+
+    if (is_cast_required) {
+      QnnTensorWrapper cast_output_tensorwrapper(output_name,
+                                                 QNN_TENSOR_TYPE_APP_READ,
+                                                 output_info.qnn_data_type,
+                                                 output_info.quant_param.Copy(),
+                                                 std::vector<uint32_t>(output_info.shape));
+      ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_output_tensorwrapper)),
+                        "Failed to add tensor.");
+      ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(cast_input_name,
+                                                        QNN_OP_PACKAGE_NAME_QTI_AISW,
+                                                        "Cast",
+                                                        {cast_input_name},
+                                                        {output_name},
+                                                        {}),
+                        "Failed to add node");
+    }
+  }
+
   return Status::OK();
 }
 
diff --git a/onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.cc b/onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.cc
@@ -601,13 +601,14 @@ Status QnnModelWrapper::AddTransposeNode(NodeIndex node_index,
   ORT_RETURN_IF_NOT(AddTensorWrapper(std::move(output_tensorwrapper)), "Failed to add tensor.");
   const static std::string qnn_node_type = "Transpose";
 
-  CreateQnnNode(output_name,
-                QNN_OP_PACKAGE_NAME_QTI_AISW,
-                qnn_node_type,
-                {input_name},
-                {output_name},
-                {param_tensor_name},
-                do_op_validation);
+  ORT_RETURN_IF_NOT(CreateQnnNode(output_name,
+                                  QNN_OP_PACKAGE_NAME_QTI_AISW,
+                                  qnn_node_type,
+                                  {input_name},
+                                  {output_name},
+                                  {param_tensor_name},
+                                  do_op_validation),
+                    "QNN EP: Failed to create manually inserted Qnn Transpose node.");
 
   return Status::OK();
 }
diff --git a/onnxruntime/core/providers/qnn/builder/qnn_utils.cc b/onnxruntime/core/providers/qnn/builder/qnn_utils.cc
@@ -1299,6 +1299,21 @@ Status InsertConvertOp(QnnModelWrapper& qnn_model_wrapper,
   return Status::OK();
 }
 
+Status GetPermToLastAxis(uint32_t axis, uint32_t rank, std::vector<uint32_t>& perm) {
+  ORT_RETURN_IF_NOT(axis < rank, "Expected axis must be smaller than rank: ", axis, " >= ", rank);
+
+  perm.reserve(rank);
+  for (uint32_t dim = 0; dim < rank; ++dim) {
+    perm.push_back(dim);
+  }
+
+  // Swap axis with the last one.
+  perm[axis] = rank - 1;
+  perm[rank - 1] = axis;
+
+  return Status::OK();
+}
+
 }  // namespace utils
 }  // namespace qnn
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/qnn/builder/qnn_utils.h b/onnxruntime/core/providers/qnn/builder/qnn_utils.h
@@ -385,6 +385,16 @@ Status InsertConvertOp(QnnModelWrapper& qnn_model_wrapper,
                        bool output_symmetric,
                        bool do_op_validation);
 
+/**
+ * Get permutation to transpose given axis to the last one.
+ *
+ * @param[in] axis the current axis to be transposed
+ * @param[in] rank the expected rank for permutation
+ * @param[out] perm the permutation for transpose
+ * @return execution status of this function
+ */
+Status GetPermToLastAxis(uint32_t axis, uint32_t rank, std::vector<uint32_t>& perm);
+
 }  // namespace utils
 }  // namespace qnn
 }  // namespace onnxruntime
diff --git a/onnxruntime/test/providers/qnn/topk_op_test.cc b/onnxruntime/test/providers/qnn/topk_op_test.cc
@@ -4,12 +4,13 @@
 #if !defined(ORT_MINIMAL_BUILD)
 
 #include <string>
+#include <vector>
 
-#include "test/providers/qnn/qnn_test_utils.h"
-#include "core/graph/node_attr_utils.h"
+#include "gtest/gtest.h"
 
+#include "core/graph/node_attr_utils.h"
 #include "core/graph/onnx_protobuf.h"
-#include "gtest/gtest.h"
+#include "test/providers/qnn/qnn_test_utils.h"
 
 namespace onnxruntime {
 namespace test {
@@ -63,12 +64,12 @@ TEST_F(QnnCPUBackendTests, TopK_DynamicK_Unsupported) {
                           ExpectedEPNodeAssignment::None);  // Should not be assigned to QNN EP.
 }
 
-// Test that TopK with an axis attribute that is not the last dimension is not supported by QNN EP.
-TEST_F(QnnCPUBackendTests, TopK_NonLastAxis_Unsupported) {
+// Test that TopK with an axis attribute that is not the last dimension.
+TEST_F(QnnCPUBackendTests, TopK_NonLastAxis) {
   RunTopKTestOnCPU<float>(TestInputDef<float>({1, 3, 4, 4}, false, GetFloatDataInRange(-10.0f, 10.0f, 48)),
                           TestInputDef<int64_t>({1}, true /* is_initializer */, {2}),
                           {utils::MakeAttribute("axis", static_cast<int64_t>(1))},
-                          ExpectedEPNodeAssignment::None);  // Should not be assigned to QNN EP.
+                          ExpectedEPNodeAssignment::All);
 }
 
 // Test that TopK that returns the top k minimum values is not supported by QNN EP.
@@ -165,6 +166,14 @@ TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U8_LastAxis) {
                                ExpectedEPNodeAssignment::All);
 }
 
+// Test 8-bit QDQ TopK on HTP backend: non-last axis
+TEST_F(QnnHTPBackendTests, TopK_U8_NonLastAxis) {
+  RunQDQTopKTestOnHTP<uint8_t>(TestInputDef<float>({1, 3, 4, 4}, false, GetFloatDataInRange(-10.0f, 10.0f, 48)),
+                               TestInputDef<int64_t>({1}, true /* is_initializer */, {2}),
+                               {utils::MakeAttribute("axis", static_cast<int64_t>(1))},  // Attributes
+                               ExpectedEPNodeAssignment::All);
+}
+
 // Test 16-bit QDQ TopK on HTP backend: top 2 largest floats from last axis
 TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U16_LastAxis) {
   RunQDQTopKTestOnHTP<uint16_t>(TestInputDef<float>({1, 3, 4, 4}, false, GetFloatDataInRange(-20.0f, 20.0f, 48)),
@@ -174,6 +183,15 @@ TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U16_LastAxis) {
                                 21);  // opset
 }
 
+// Test 16-bit QDQ TopK on HTP backend: non-last axis
+TEST_F(QnnHTPBackendTests, TopK_U16_NonLastAxis) {
+  RunQDQTopKTestOnHTP<uint16_t>(TestInputDef<float>({1, 3, 4, 4}, false, GetFloatDataInRange(-20.0f, 20.0f, 48)),
+                                TestInputDef<int64_t>({1}, true /* is_initializer */, {2}),
+                                {utils::MakeAttribute("axis", static_cast<int64_t>(1))},  // Attributes
+                                ExpectedEPNodeAssignment::All,
+                                21);  // opset
+}
+
 #endif  // defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__)
 }  // namespace test
 }  // namespace onnxruntime
