[QNN-EP] Enable largest attr and remove sorted check for topk. (microsoft#26150)

quic-muchhsu · web-flow · commit b608f7987eab · 2025-09-25T11:02:18.000-07:00
### Description
Enable largest attr and remove sorted check for QNN topk op builder.


### Motivation and Context

Signed-off-by: Mu-Chein Hsu &lt;quic_muchhsu@quicinc.com&gt;
diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/topk.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/topk.cc
@@ -60,15 +60,6 @@ Status TopKOpBuilder::ExplictOpCheck(QnnModelWrapper& qnn_model_wrapper, const N
   if (!qnn_model_wrapper.IsConstantInput(input_1)) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "The number of top elements to retrieve must be specified as constant input.");
   }
-  NodeAttrHelper node_helper(node_unit);
-  auto largest = node_helper.Get("largest", 1);
-  auto sorted = node_helper.Get("sorted", 1);
-  if (0 == sorted) {
-    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "QNN TopK output is always sorted");
-  }
-  if (0 == largest) {
-    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "QNN TopK output is always largest values");
-  }
 
   return Status::OK();
 }
@@ -147,6 +138,16 @@ Status TopKOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wra
   qnn_model_wrapper.AddParamWrapper(std::move(k_param));
   std::vector<std::string> param_tensor_names{k_param_name};
 
+  // Add largest to TopK attr
+  uint8_t largest = static_cast<uint8_t>(NodeAttrHelper(node_unit).Get("largest", 1));
+  Qnn_Scalar_t qnn_largest_k = QNN_SCALAR_INIT;
+  qnn_largest_k.dataType = QNN_DATATYPE_BOOL_8;
+  qnn_largest_k.bool8Value = largest;
+  QnnParamWrapper k_largest(node_unit.Index(), node_unit.Name(), QNN_OP_TOP_K_PARAM_LARGEST, qnn_largest_k);
+  std::string k_largest_name = k_largest.GetParamTensorName();
+  qnn_model_wrapper.AddParamWrapper(std::move(k_largest));
+  param_tensor_names.push_back(k_largest_name);
+
   // 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));
diff --git a/onnxruntime/test/providers/qnn/topk_op_test.cc b/onnxruntime/test/providers/qnn/topk_op_test.cc
@@ -41,15 +41,19 @@ static void RunTopKTestOnCPU(const TestInputDef<DataType>& input_def,
                              const TestInputDef<int64_t>& k_def,
                              const std::vector<ONNX_NAMESPACE::AttributeProto>& attrs,
                              ExpectedEPNodeAssignment expected_ep_assignment,
-                             int opset = 19) {
+                             int opset = 19,
+                             bool verify_outputs = true) {
   ProviderOptions provider_options;
 
   provider_options["backend_type"] = "cpu";
 
   RunQnnModelTest(BuildTopKTestCase<DataType>(input_def, k_def, attrs),
                   provider_options,
                   opset,
-                  expected_ep_assignment);
+                  expected_ep_assignment,
+                  /*fp32_abs_err*/ 1e-5f,
+                  /*log_severity*/ logging::Severity::kERROR,
+                  /*verify_outputs*/ verify_outputs);
 }
 
 //
@@ -72,12 +76,13 @@ TEST_F(QnnCPUBackendTests, TopK_NonLastAxis) {
                           ExpectedEPNodeAssignment::All);
 }
 
-// Test that TopK that returns the top k minimum values is not supported by QNN EP.
-TEST_F(QnnCPUBackendTests, TopK_MinValues_Unsupported) {
+// Test that TopK with an axis attribute that is not the last dimension. Set largest to 0.
+TEST_F(QnnCPUBackendTests, TopK_NonLastAxis_Largest_0) {
   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("largest", static_cast<int64_t>(0))},
-                          ExpectedEPNodeAssignment::None);  // Should not be assigned to QNN EP.
+                          {utils::MakeAttribute("axis", static_cast<int64_t>(1)),
+                           utils::MakeAttribute("largest", static_cast<int64_t>(0))},
+                          ExpectedEPNodeAssignment::All);
 }
 
 // Test TopK on CPU backend: top 2 largest floats from last axis
@@ -88,6 +93,25 @@ TEST_F(QnnCPUBackendTests, TopK_LargestFloats_LastAxis) {
                           ExpectedEPNodeAssignment::All);
 }
 
+// Test TopK on CPU backend: top 2 largest floats from last axis. Largest to 0.
+TEST_F(QnnCPUBackendTests, TopK_LargestFloats_LastAxis_Largest_0) {
+  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("largest", static_cast<int64_t>(0))},  // Attributes
+                          ExpectedEPNodeAssignment::All);
+}
+
+// Test TopK on CPU backend: top 2 largest floats from last axis. Set largest to 0. Set sorted to false.
+TEST_F(QnnCPUBackendTests, TopK_LargestFloats_LastAxis_Largest_0_Sorted) {
+  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("largest", static_cast<int64_t>(0)),
+                           utils::MakeAttribute("sorted", static_cast<int64_t>(0))},  // Attributes
+                          ExpectedEPNodeAssignment::All,
+                          /*opset*/ 19,
+                          /*verify_outputs*/ false);  // Disable verify output. Unsorted doesn't guarantee output sequence
+}
+
 // Test TopK on CPU backend: top 2 largest int32s from last axis
 TEST_F(QnnCPUBackendTests, TopK_LargestInt32s_LastAxis) {
   std::vector<int32_t> input_data = {-6, -5, -4, -3, -2, 0, 1, 2, 3, 4, 5, 6};
@@ -97,6 +121,27 @@ TEST_F(QnnCPUBackendTests, TopK_LargestInt32s_LastAxis) {
                             ExpectedEPNodeAssignment::All);
 }
 
+// Test TopK on CPU backend: top 2 largest int32s from last axis. Set largest to 0.
+TEST_F(QnnCPUBackendTests, TopK_LargestInt32s_LastAxis_Largest_0) {
+  std::vector<int32_t> input_data = {-6, -5, -4, -3, -2, 0, 1, 2, 3, 4, 5, 6};
+  RunTopKTestOnCPU<int32_t>(TestInputDef<int32_t>({1, 2, 2, 3}, false, input_data),
+                            TestInputDef<int64_t>({1}, true /* is_initializer */, {2}),
+                            {utils::MakeAttribute("largest", static_cast<int64_t>(0))},  // Attributes
+                            ExpectedEPNodeAssignment::All);
+}
+
+// Test TopK on CPU backend: top 2 largest int32s from last axis. Set largest to 0. Set sorted to false.
+TEST_F(QnnCPUBackendTests, TopK_LargestInt32s_LastAxis_Largest_0_Sorted) {
+  std::vector<int32_t> input_data = {-6, -5, -4, -3, -2, 0, 1, 2, 3, 4, 5, 6};
+  RunTopKTestOnCPU<int32_t>(TestInputDef<int32_t>({1, 2, 2, 3}, false, input_data),
+                            TestInputDef<int64_t>({1}, true /* is_initializer */, {2}),
+                            {utils::MakeAttribute("largest", static_cast<int64_t>(0)),
+                             utils::MakeAttribute("sorted", static_cast<int64_t>(0))},  // Attributes
+                            ExpectedEPNodeAssignment::All,
+                            /*opset*/ 19,
+                            /*verify_outputs*/ false);  // Disable verify output. Unsorted doesn't guarantee output sequence
+}
+
 #if defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__)
 //
 // HTP tests:
@@ -166,6 +211,14 @@ TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U8_LastAxis) {
                                ExpectedEPNodeAssignment::All);
 }
 
+// Test 8-bit QDQ TopK on HTP backend: top 2 largest floats from last axis. Set largest to 0.
+TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U8_LastAxis_Largest_0) {
+  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("largest", static_cast<int64_t>(0))},  // Attributes
+                               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)),
@@ -174,6 +227,15 @@ TEST_F(QnnHTPBackendTests, TopK_U8_NonLastAxis) {
                                ExpectedEPNodeAssignment::All);
 }
 
+// Test 8-bit QDQ TopK on HTP backend: non-last axis. Set largest to 0.
+TEST_F(QnnHTPBackendTests, TopK_U8_NonLastAxis_Largest_0) {
+  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)),
+                                utils::MakeAttribute("largest", static_cast<int64_t>(0))},  // 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)),
@@ -183,6 +245,15 @@ TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U16_LastAxis) {
                                 21);  // opset
 }
 
+// Test 16-bit QDQ TopK on HTP backend: top 2 largest floats from last axis. Set largest to 0.
+TEST_F(QnnHTPBackendTests, TopK_LargestFloats_U16_LastAxis_Largest_0) {
+  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("largest", static_cast<int64_t>(0))},  // Attributes
+                                ExpectedEPNodeAssignment::All,
+                                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)),
@@ -192,6 +263,15 @@ TEST_F(QnnHTPBackendTests, TopK_U16_NonLastAxis) {
                                 21);  // opset
 }
 
+// Test 16-bit QDQ TopK on HTP backend: non-last axis. Set largest to 0.
+TEST_F(QnnHTPBackendTests, TopK_U16_NonLastAxis_Largest_0) {
+  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)),
+                                 utils::MakeAttribute("largest", static_cast<int64_t>(0))},  // Attributes
+                                ExpectedEPNodeAssignment::All,
+                                21);  // opset
+}
 #endif  // defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__)
 }  // namespace test
 }  // namespace onnxruntime
