Repository: microsoft/onnxruntime
File: onnxruntime/test/providers/qnn/qnn_basic_test.cc (C++)

Code before:
  std::unique_ptr<ModelAndBuilder> model;
  std::vector<float> input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
  std::vector<int64_t> shape = {1, 3, 2};

  CreateModelInMemory(model,
                      QDQBuildAdd3Tensors<uint8_t>(TestInputDef<float>(shape, false, input_data),
                                                   TestInputDef<float>(shape, false, input_data),
                                                   TestInputDef<float>(shape, false, input_data)),
                      "add3.qdq");

  SessionOptions session_opts;
  session_opts.session_logid = "logger0";

  RunOptions run_opts;
  run_opts.run_tag = session_opts.session_logid;

  InferenceSession session_obj{session_opts, GetEnvironment()};
  onnxruntime::ProviderOptions options;

  options["backend_type"] = "htp";
  options["offload_graph_io_quantization"] = "0";
  options["htp_arch"] = "81";
  options["extended_udma"] = "1";

  auto qnn_ep = QnnExecutionProviderWithOptions(options, &session_opts);
  EXPECT_TRUE(session_obj.RegisterExecutionProvider(std::move(qnn_ep)).IsOK());

  auto status = session_obj.Load(model->model_data.data(), static_cast<int>(model->model_data.size()));
  ASSERT_TRUE(status.IsOK());
  status = session_obj.Initialize();
  ASSERT_TRUE(status.IsOK());
  std::vector<OrtValue> fetches;
  status = session_obj.Run(run_opts, model->builder.feeds_, model->builder.output_names_, &fetches);
  ASSERT_TRUE(status.IsOK());
}

// Test option for offloading quantization of graph inputs and dequantization of graph outputs to the CPU EP.
TEST_F(QnnHTPBackendTests, EPOffloadsGraphIOQuantDequant) {
  // Returns a function that checks that the Q/DQ ops at the graph IO boundary are offloaded to CPU
  // if the corresponding provider option is enabled.
  auto graph_checker_builder = [](bool offload_graph_io_quantization) -> std::function<void(const Graph&)> {
    return [offload_graph_io_quantization](const Graph& graph) {
      size_t num_q = 0;
      size_t num_dq = 0;
      size_t num_qnn_fused_node = 0;

      for (const Node& node : graph.Nodes()) {
        const std::string& ep_name = node.GetExecutionProviderType();
        const std::string& op_type = node.OpType();

        if (offload_graph_io_quantization && op_type == "QuantizeLinear") {

Review comment (suggestion):
nit: use `EXPECT_STATUS_OK` or `ASSERT_STATUS_OK`. that will also show the status error message. also, this check should probably be ASSERT instead of EXPECT.

Code after:
TEST_F(QnnHTPBackendTests, ExtendedUdmaModeTest) {
  // Create provider options with extended UDMA mode enabled
  ProviderOptions options;
  options["backend_type"] = "htp";
  options["offload_graph_io_quantization"] = "0";
  options["htp_arch"] = "81";
  options["extended_udma"] = "1";

  // Define a simple model with Add operation
  auto input_defs = {TestInputDef<float>({1, 3, 4, 4}, false, -10.0f, 10.0f),
                     TestInputDef<float>({1, 3, 4, 4}, false, -10.0f, 10.0f)};

  // Run the test - this should succeed because v81 supports extended UDMA
  RunQnnModelTest(BuildOpTestCase<float>("Add", input_defs, {}, {}, kOnnxDomain),
                  options,
                  13,
                  ExpectedEPNodeAssignment::All,
                  0.008f);
}

// Test option for offloading quantization of graph inputs and dequantization of graph outputs to the CPU EP.
TEST_F(QnnHTPBackendTests, EPOffloadsGraphIOQuantDequant) {
  // Returns a function that checks that the Q/DQ ops at the graph IO boundary are offloaded to CPU
  // if the corresponding provider option is enabled.
  auto graph_checker_builder = [](bool offload_graph_io_quantization) -> std::function<void(const Graph&)> {
    return [offload_graph_io_quantization](const Graph& graph) {
      size_t num_q = 0;
      size_t num_dq = 0;
      size_t num_qnn_fused_node = 0;

      for (const Node& node : graph.Nodes()) {
        const std::string& ep_name = node.GetExecutionProviderType();
        const std::string& op_type = node.OpType();

        if (offload_graph_io_quantization && op_type == "QuantizeLinear") {
          const bool consumes_graph_input = graph.IsInputsIncludingInitializers(node.InputDefs()[0]);
          EXPECT_EQ(ep_name, kCpuExecutionProvider);
          EXPECT_TRUE(consumes_graph_input);
          num_q += 1;
        } else if (offload_graph_io_quantization && op_type == "DequantizeLinear") {
          const bool produces_graph_output = graph.IsOutput(node.OutputDefs()[0]);
          EXPECT_EQ(ep_name, kCpuExecutionProvider);
          EXPECT_TRUE(produces_graph_output);
          num_dq += 1;
        } else {
          EXPECT_EQ(ep_name, kQnnExecutionProvider);
          num_qnn_fused_node += 1;
        }
      }

      EXPECT_EQ(num_q, static_cast<size_t>(offload_graph_io_quantization));