[QNN EP] Dump QNN json graph (microsoft#22843)

### Description
Adds QNN provider options to save the generated QNN graphs into JSON
files. These JSON files contain information about the QNN nodes,
tensors, data types used in the corresponding graph. These JSON files
can optionally be loaded by [QNN
Netron](https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-50/tools.html#qnn-netron-beta)
to visualize the QNN graphs generated by QNN EP.

New QNN provider options:
- `"dump_json_qnn_graph"`: Set to "1" to dump QNN graphs generated by
QNN EP as JSON files. Each graph partition assigned to QNN EP is dumped
to a separate file.
- `"json_qnn_graph_dir"`: Directory in which to dump QNN JSON graphs. If
not specified, QNN graphs are dumped in the program's current working
directory. Ignored if `"dump_json_qnn_graph"` is not set.


### Motivation and Context
Makes it easier to debug, visualize, and understand the QNN graphs
generated by QNN EP.

---------

Co-authored-by: Hector Li <[email protected]>

Author: adrianlizarraga

cmake/onnxruntime_providers_qnn.cmake

@@ -22,7 +22,8 @@
     onnxruntime_add_static_library(onnxruntime_providers_qnn ${onnxruntime_providers_qnn_srcs})
     onnxruntime_add_include_to_target(onnxruntime_providers_qnn onnxruntime_common onnxruntime_framework onnx
                                                                 onnx_proto protobuf::libprotobuf-lite
-                                                                flatbuffers::flatbuffers Boost::mp11)
+                                                                flatbuffers::flatbuffers Boost::mp11
+								nlohmann_json::nlohmann_json)
     add_dependencies(onnxruntime_providers_qnn onnx ${onnxruntime_EXTERNAL_DEPENDENCIES})
     set_target_properties(onnxruntime_providers_qnn PROPERTIES CXX_STANDARD_REQUIRED ON)
     set_target_properties(onnxruntime_providers_qnn PROPERTIES FOLDER "ONNXRuntime")
@@ -50,7 +51,8 @@
     source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_qnn_srcs})
     onnxruntime_add_shared_library_module(onnxruntime_providers_qnn ${onnxruntime_providers_qnn_srcs})
     onnxruntime_add_include_to_target(onnxruntime_providers_qnn ${ONNXRUNTIME_PROVIDERS_SHARED} ${GSL_TARGET} onnx
-	                                                        onnxruntime_common Boost::mp11 safeint_interface)
+	                                                        onnxruntime_common Boost::mp11 safeint_interface
+								nlohmann_json::nlohmann_json)
     target_link_libraries(onnxruntime_providers_qnn PRIVATE ${ONNXRUNTIME_PROVIDERS_SHARED} ${ABSEIL_LIBS} ${CMAKE_DL_LIBS})
     add_dependencies(onnxruntime_providers_qnn onnxruntime_providers_shared ${onnxruntime_EXTERNAL_DEPENDENCIES})
     target_include_directories(onnxruntime_providers_qnn PRIVATE ${ONNXRUNTIME_ROOT}

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -3674,6 +3674,10 @@ struct OrtApi {
    *   be available.
    *     - "0": Default. Disabled.
    *     - "1": Enabled.
+   *   "dump_json_qnn_graph": Set to "1" to dump QNN graphs generated by QNN EP as JSON files. Each graph partition
+   *      assigned to QNN EP is dumped to a separate file.
+   *   "json_qnn_graph_dir": Directory in which to dump QNN JSON graphs. If not specified, QNN graphs are dumped in the
+   *      program's current working directory. Ignored if "dump_json_qnn_graph" is not set.
    *
    * SNPE supported keys:
    *   "runtime": SNPE runtime engine, options: "CPU", "CPU_FLOAT32", "GPU", "GPU_FLOAT32_16_HYBRID", "GPU_FLOAT16",

onnxruntime/core/providers/qnn/builder/qnn_def.cc

@@ -3,6 +3,7 @@
 
 #include "core/providers/qnn/builder/qnn_def.h"
 #include "core/providers/qnn/builder/qnn_utils.h"
+#include <functional>
 #include <memory>
 #include <ostream>
 #include <cstring>
@@ -432,6 +433,15 @@ Status CompareQnnQuantParams(const Qnn_QuantizeParams_t& qparam0, const Qnn_Quan
   return Status::OK();
 }
 
+uint32_t CalcQnnTensorNumElems(const Qnn_Tensor_t& qnn_tensor) {
+  uint32_t* qnn_tensor_dims = GetQnnTensorDims(qnn_tensor);
+  uint32_t qnn_tensor_rank = GetQnnTensorRank(qnn_tensor);
+  return std::accumulate(qnn_tensor_dims,
+                         qnn_tensor_dims + qnn_tensor_rank,
+                         1,
+                         std::multiplies<uint32_t>());
+}
+
 bool CreateTensorInQnnGraph(const QNN_INTERFACE_VER_TYPE& qnn_interface,
                             const Qnn_GraphHandle_t& graph,
                             const std::string& node_name,
@@ -466,12 +476,7 @@ bool CreateTensorInQnnGraph(const QNN_INTERFACE_VER_TYPE& qnn_interface,
       return false;
     }
     // verify size expressed by the dims matches the raw tensor size
-    auto qnn_tensor_dims = GetQnnTensorDims(qnn_tensor);
-    auto qnn_tensor_rank = GetQnnTensorRank(qnn_tensor);
-    uint32_t qnn_tensor_size = std::accumulate(qnn_tensor_dims,
-                                               qnn_tensor_dims + qnn_tensor_rank,
-                                               static_cast<uint32_t>(data_size),
-                                               std::multiplies<uint32_t>());
+    uint32_t qnn_tensor_size = CalcQnnTensorNumElems(qnn_tensor) * gsl::narrow_cast<uint32_t>(data_size);
     auto qnn_tensor_buf_size = GetQnnTensorClientBuf(qnn_tensor).dataSize;
     if (qnn_tensor_size != qnn_tensor_buf_size) {
       ss << "Data length mismatch for static tensor. node_name: " << node_name

onnxruntime/core/providers/qnn/builder/qnn_def.h

@@ -122,6 +122,7 @@ Qnn_DataType_t GetQnnTensorDataType(const Qnn_Tensor_t& qnn_tensor);
 Qnn_TensorMemType_t GetQnnTensorMemType(const Qnn_Tensor_t& qnn_tensor);
 uint32_t GetQnnTensorRank(const Qnn_Tensor_t& qnn_tensor);
 uint32_t* GetQnnTensorDims(const Qnn_Tensor_t& qnn_tensor);
+uint32_t CalcQnnTensorNumElems(const Qnn_Tensor_t& qnn_tensor);
 const Qnn_ClientBuffer_t& GetQnnTensorClientBuf(const Qnn_Tensor_t& qnn_tensor);
 Qnn_MemHandle_t GetQnnTensorMemHandle(const Qnn_Tensor_t& qnn_tensor);
 const Qnn_QuantizeParams_t& GetQnnTensorQParams(const Qnn_Tensor_t& qnn_tensor);

onnxruntime/core/providers/qnn/builder/qnn_model.cc

@@ -4,6 +4,7 @@
 #include "qnn_model.h"
 
 #include <iostream>
+#include <fstream>
 #include <gsl/gsl>
 #include "QnnOpDef.h"
 
@@ -94,7 +95,8 @@ Status QnnModel::ComposeGraph(const GraphViewer& graph_viewer,
                               const onnxruntime::Node& fused_node,
                               const qnn::ModelSettings& model_settings,
                               const logging::Logger& logger,
-                              const QnnGraph_Config_t** graph_configs) {
+                              const QnnGraph_Config_t** graph_configs,
+                              const std::string& json_qnn_graph_path) {
   LOGS(logger, VERBOSE) << "ComposeGraph Graph name: " << graph_viewer.Name();
 
   // Holder for the NodeUnits in the graph, this will guarantee the NodeUnits is
@@ -137,7 +139,20 @@ Status QnnModel::ComposeGraph(const GraphViewer& graph_viewer,
     }
   }
 
-  ORT_RETURN_IF_NOT(qnn_model_wrapper.ComposeQnnGraph(), "Failed to compose Qnn graph.");
+  const bool build_json_graph = !json_qnn_graph_path.empty();
+  ORT_RETURN_IF_NOT(qnn_model_wrapper.ComposeQnnGraph(build_json_graph), "Failed to compose Qnn graph.");
+
+  if (build_json_graph) {
+    const nlohmann::json& json_graph = qnn_model_wrapper.GetQnnJSONGraph();
+    std::ofstream ofs(json_qnn_graph_path);
+
+    if (ofs.is_open()) {
+      ofs << json_graph.dump();
+      ofs.close();
+    } else {
+      LOGS(logger, WARNING) << "Could not open JSON graph file: " << json_qnn_graph_path;
+    }
+  }
 
   rt = GetGraphInfoFromModel(qnn_model_wrapper, logger);
   if (!rt) {

onnxruntime/core/providers/qnn/builder/qnn_model.h

@@ -35,7 +35,8 @@ class QnnModel {
                       const onnxruntime::Node& fused_node,
                       const qnn::ModelSettings& model_settings,
                       const logging::Logger& logger,
-                      const QnnGraph_Config_t** graph_configs = nullptr);
+                      const QnnGraph_Config_t** graph_configs = nullptr,
+                      const std::string& json_qnn_graph_path = "");
 
   Status FinalizeGraphs(const logging::Logger& logger);
 

onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.cc

@@ -267,7 +267,7 @@ bool QnnModelWrapper::CreateQnnNode(const std::string& qnn_node_name,
   }
 }
 
-bool QnnModelWrapper::ComposeQnnGraph() {
+bool QnnModelWrapper::ComposeQnnGraph(bool build_json_qnn_graph) {
   LOGS(logger_, VERBOSE) << "Compose Qnn Graph.";
   // ORT_RETURN_IF(qnn_op_property_list_.empty(), "Empty Qnn op list, no graph to compose.");
   if (qnn_op_property_list_.empty()) {
@@ -306,6 +306,10 @@ bool QnnModelWrapper::ComposeQnnGraph() {
       LOGS(logger_, ERROR) << error_msg;
       return false;
     }
+
+    if (build_json_qnn_graph) {
+      json_qnn_graph_.AddOp(op_config_wrapper);
+    }
   }
 
   return true;

onnxruntime/core/providers/qnn/builder/qnn_model_wrapper.h

@@ -8,10 +8,12 @@
 #include <vector>
 
 #include "QnnInterface.h"
-#include "qnn_def.h"
+#include "nlohmann/json.hpp"
 
 #include "core/providers/qnn/ort_api.h"
+#include "core/providers/qnn/builder/qnn_def.h"
 #include "core/providers/qnn/builder/qnn_quant_params_wrapper.h"
+#include "core/providers/qnn/builder/qnn_utils.h"
 
 namespace onnxruntime {
 namespace qnn {
@@ -91,7 +93,7 @@ class QnnModelWrapper {
                      std::vector<std::string>&& param_tensor_names,
                      bool do_op_validation = false);
 
-  bool ComposeQnnGraph();
+  bool ComposeQnnGraph(bool build_json_qnn_graph = false);
 
   Qnn_GraphHandle_t GetQnnGraph() const { return graph_; }
 
@@ -129,6 +131,10 @@ class QnnModelWrapper {
     return input_index_map_.find(tensor_name) != input_index_map_.end();
   }
 
+  const nlohmann::json& GetQnnJSONGraph() {
+    return json_qnn_graph_.Finalize();
+  }
+
   Qnn_TensorType_t GetTensorType(const std::string& tensor_name) const {
     if (IsInitializerInput(tensor_name)) {
       return QNN_TENSOR_TYPE_STATIC;
@@ -321,6 +327,7 @@ class QnnModelWrapper {
   const std::unordered_set<std::string>& initializer_lookup_;
   QnnBackendType qnn_backend_type_ = QnnBackendType::CPU;
   ModelSettings model_settings_ = {};
+  utils::QnnJSONGraph json_qnn_graph_;
 };  // QnnModelWrapper
 
 }  // namespace qnn