Update handler function signature to take in the ExternalDataInfo for…

… original initializer

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -551,36 +551,37 @@ typedef OrtStatus*(ORT_API_CALL* OrtWriteBufferFunc)(_In_ void* state,
  * written to an external file or stored within the model. ORT calls this function for every initializer when
  * generating a model.
  *
- * If the function sets the `is_external` output parameter to false, ORT stores initializer data within the model.
+ * If the function implementation sets the `new_external_info` output parameter to NULL, ORT stores the initializer data
+ * within the generated model.
  *
- * Otherwise, if `is_external` is set to false, ORT assumes that this function stores the initializer data to a file.
- * In this case, ORT configures the model's initializer to point to the location and offset returned from this function
- * via the `location` and `offset` output parameters.
+ * Otherwise, if the function implementation sets `new_external_info` to a valid OrtExternalInitializerInfo instance,
+ * ORT assumes that this function stores the initializer data in a file. In this case, ORT configures the model's
+ * initializer to point to the location specified by the `new_external_info` output parameter.
  *
  * \param[in] state Opaque pointer holding the user's state.
  * \param[in] initializer_name The initializer's name as a null-terminated string.
- * \param[in] initializer_data Pointer to the initializer's raw data (contiguous).
- * \param[in] initializer_num_bytes The size in bytes of the initializer's data.
- * \param[in] initializer_type The type and shape information for the initializer.
- * \param[out] is_external Output parameter set to true if the initializer data is to be stored externally.
- *                         The function implemented is responsible for writing the initializer data to file.
- *                         If set to false, ORT stores the initializers within the model.
- * \param[out] location Output parameter set to the location (i.e., file path) into which the initializer data is stored
- *                      by the function implementer. Ignored if `is_external` is set to false.
- * \param[out] offset Output parameter set to the location offset into which the initializer data is stored
- *                    by the function implementer. Ignored if `is_external` is set to false.
+ * \param[in] initializer_value OrtValue containing the initializer's data, type, and shape.
+ * \param[in] external_info If the initializer is originally stored in an external file, `external_info` contains
+ *                          the file path, file offset, and the data's byte size within the file. Otherwise,
+ *                          `external_info` is NULL if the initializer is not originally stored in a file.
+ * \param[out] new_external_info Output parameter set to a new OrtExternalInitializerInfo instance indicating the
+ *                               location where the function implementation stored the initializer data.
+ *                               The function implementation must use `OrtApi::CreateExternalInitializerInfo()` to
+ *                               create the instance.
+ *                               If the function implementation sets `new_external_info` to NULL,
+ *                               ORT stores the initializers within the model.
+ *
+ * \note ORT takes ownership of the `new_external_info` output parameter.
  *
  * \return OrtStatus* Write status. Return nullptr on success.
  *                    Use CreateStatus to provide error info. Use ORT_FAIL as the error code.
  *                    ORT will release the OrtStatus* if not null.
  */
 typedef OrtStatus*(ORT_API_CALL* OrtHandleInitializerDataFunc)(_In_ void* state,
                                                                _In_ const char* initializer_name,
-                                                               _In_ const void* initializer_data,
-                                                               _In_ size_t initializer_num_bytes,
-                                                               _In_ const OrtTypeInfo* initializer_type,
-                                                               _Out_ bool* is_external,
-                                                               _Out_ const ORTCHAR_T** location, _Out_ int64_t* offset);
+                                                               _In_ const OrtValue* initializer_value,
+                                                               _In_opt_ const OrtExternalInitializerInfo* external_info,
+                                                               _Outptr_result_maybenull_ OrtExternalInitializerInfo** new_external_info);
 
 /** \brief Algorithm to use for cuDNN Convolution Op
  */
@@ -6236,6 +6237,21 @@ struct OrtApi {
    */
   ORT_CLASS_RELEASE(ExternalInitializerInfo);
 
+  /** \brief Creates an OrtExternalInitializerInfo instance.
+   *
+   * \param[in] filepath The relative path to the file that stores the initializer's data. ORT copies this path string.
+   * \param[in] file_offset The byte offset where the initializer's data is stored within the file.
+   * \param[in] byte_size The size in bytes of the initializer's data within the file.
+   * \param[out] out Output parameter set to the new OrtExternalInitializerInfo instance.
+   *                 Must be released by calling ReleaseExternalInitializerInfo().
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.23.
+   */
+  ORT_API2_STATUS(CreateExternalInitializerInfo, _In_ const ORTCHAR_T* filepath, _In_ int64_t file_offset,
+                  _In_ size_t byte_size, _Outptr_ OrtExternalInitializerInfo** out);
+
   /** \brief Get the relative path to the file that stores the initializer's data.
    *
    * \note The path is relative to the filesystem directory where the ONNX model was stored.

onnxruntime/core/framework/tensor_external_data_info.h

@@ -25,6 +25,10 @@ class ExternalDataInfo {
   using OFFSET_TYPE = off_t;
 #endif
 
+  ExternalDataInfo() = default;
+  ExternalDataInfo(const PathString& rel_path, OFFSET_TYPE offset, size_t length)
+      : rel_path_(rel_path), offset_(offset), length_(length) {}
+
   const PathString& GetRelPath() const { return rel_path_; }
 
   OFFSET_TYPE GetOffset() const { return offset_; }

onnxruntime/core/graph/graph.cc

@@ -4650,14 +4650,17 @@ Status Graph::ToGraphProtoWithInitializerHandlerImpl(OrtHandleInitializerDataFun
   // This loop processes subgraphs bottom up.
   for (const auto& node : Nodes()) {
     if (node.ContainsSubgraph()) {
-      // Let find this node in the output_graph_proto
+      // Let's find this node in the output_graph_proto
+      // The node name is optional, so we may need to check by the output value name
+      // given that they can only assigned once.
       auto hit = std::find_if(output_graph_proto.mutable_node()->begin(),
                               output_graph_proto.mutable_node()->end(),
                               [&node](const ONNX_NAMESPACE::NodeProto& proto) {
-                                // ONNX allows empty node names. However, Graph::Resolve() ensures node names are
-                                // unique. Therefore, we can match on node name since there will be at most one node
-                                // with an empty name.
-                                return proto.name() == node.Name();
+                                const auto& node_name = node.Name();
+                                if (!node_name.empty())
+                                  return proto.name() == node_name;
+                                return (proto.output_size() > 0 &&
+                                        proto.output(0) == node.OutputDefs()[0]->Name());
                               });
       ORT_RETURN_IF_NOT(hit != output_graph_proto.mutable_node()->end(), "Node ", node.Name(),
                         " not found in output_graph_proto");
@@ -4714,39 +4717,41 @@ Status Graph::ToGraphProtoWithInitializerHandlerImpl(OrtHandleInitializerDataFun
       }
       output_proto->set_doc_string(initializer->doc_string());
 
-      // Get an OrtValue with the initializer data. If this Graph has already loaded the initializer into an
-      // OrtValue, we use that. Otherwise, if the initializer points to external data, we load it
-      // (potentially via memory mapping). Lastly, we resort to copying the initializer into an OrtValue.
       OrtValue ort_value;
-      bool has_ort_value = GetOrtValueInitializer(initializer->name(), ort_value, /*check_outer_scope*/ false);
+      std::unique_ptr<ExternalDataInfo> original_ext_data_info = nullptr;
 
-      if (!has_ort_value) {
-        if (utils::HasExternalData(*initializer)) {
-          ORT_RETURN_IF_ERROR(utils::GetExtDataFromTensorProto(Env::Default(), ModelPath(), *initializer, ort_value));
-        } else {
+      if (utils::HasExternalDataInFile(*initializer)) {
+        // Initializer has data in an external file. Load it into OrtValue (potentially via memory mapping).
+        ORT_RETURN_IF_ERROR(ExternalDataInfo::Create(initializer->external_data(), original_ext_data_info));
+        ORT_RETURN_IF_ERROR(utils::GetExtDataFromTensorProto(Env::Default(), ModelPath(), *initializer, ort_value));
+      } else {
+        // Initializer is either stored inline within the TensorProto or it is "external data in memory".
+        // Get an OrtValue (if already loaded by Graph) or copy into an OrtValue otherwise.
+        bool graph_has_ort_value = GetOrtValueInitializer(initializer->name(), ort_value, /*check_outer_scope*/ false);
+        if (!graph_has_ort_value) {
+          assert(!utils::HasExternalData(*initializer));
           ORT_RETURN_IF_ERROR(utils::TensorProtoToOrtValue(Env::Default(), ModelPath(), *initializer,
                                                            CPUAllocator::DefaultInstance(), ort_value));
         }
       }
 
-      const Tensor& tensor = ort_value.Get<Tensor>();
-      auto type_proto = utils::TypeProtoFromTensorProto(*initializer);
-      std::unique_ptr<OrtTypeInfo> type_info = OrtTypeInfo::FromTypeProto(type_proto);
-
-      // The following variables are set by the call to handle_initializer_func.
-      bool is_external = false;
-      const ORTCHAR_T* ext_location = nullptr;
-      int64_t ext_offset = 0;
-
-      ORT_RETURN_IF_ERROR(ToStatusAndRelease(handle_initializer_func(state, initializer->name().c_str(),
-                                                                     tensor.DataRaw(), tensor.SizeInBytes(),
-                                                                     type_info.get(), &is_external,
-                                                                     &ext_location, &ext_offset)));
-
-      if (is_external) {
-        ExternalDataInfo::SetExternalLocationToProto(ext_location, ext_offset,
-                                                     tensor.SizeInBytes(), *output_proto);
+      // Call the user's initializer handler function. If the user wants to store the initializer externally,
+      // the handler function will use OrtApi::CreateExternalInitializerInfo() to create a new
+      // OrtExternalInitializerInfo instance that indicates the location of the data.
+      ExternalDataInfo* new_external_info = nullptr;
+      Status status = ToStatusAndRelease(handle_initializer_func(state, initializer->name().c_str(),
+                                                                 &ort_value,
+                                                                 static_cast<OrtExternalInitializerInfo*>(original_ext_data_info.get()),
+                                                                 reinterpret_cast<OrtExternalInitializerInfo**>(&new_external_info)));
+
+      std::unique_ptr<ExternalDataInfo> new_external_info_holder(new_external_info);  // Take ownership
+      ORT_RETURN_IF_ERROR(status);
+
+      if (new_external_info != nullptr) {
+        ExternalDataInfo::SetExternalLocationToProto(new_external_info->GetRelPath(), new_external_info->GetOffset(),
+                                                     new_external_info->GetLength(), *output_proto);
       } else {
+        const Tensor& tensor = ort_value.Get<Tensor>();
         output_proto->clear_data_location();
         output_proto->set_raw_data(tensor.DataRaw(), tensor.SizeInBytes());
       }

onnxruntime/core/graph/model.cc

@@ -419,6 +419,15 @@ common::Status Model::ToGraphProtoWithInitializerHandler(OrtHandleInitializerDat
                                                          void* state,
                                                          /*out*/ ONNX_NAMESPACE::ModelProto& model_proto) const {
   model_proto = model_proto_;
+
+  // Sync current model_metadata_ back to protobuf metadata_props
+  model_proto.clear_metadata_props();
+  for (const auto& metadata : model_metadata_) {
+    const gsl::not_null<StringStringEntryProto*> prop{model_proto.add_metadata_props()};
+    prop->set_key(metadata.first);
+    prop->set_value(metadata.second);
+  }
+
   const auto& graph = *graph_;
   ORT_RETURN_IF_ERROR(graph.ToGraphProtoWithInitializerHandler(handle_initializer_func,
                                                                state, *model_proto.mutable_graph()));

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -2538,6 +2538,15 @@ ORT_API(void, OrtApis::ReleaseExternalInitializerInfo, _Frees_ptr_opt_ OrtExtern
   delete static_cast<onnxruntime::ExternalDataInfo*>(info);
 }
 
+ORT_API_STATUS_IMPL(OrtApis::CreateExternalInitializerInfo, _In_ const ORTCHAR_T* filepath,
+                    _In_ int64_t file_offset, _In_ size_t byte_size, _Outptr_ OrtExternalInitializerInfo** out) {
+  API_IMPL_BEGIN
+  auto ext_data_info = std::make_unique<onnxruntime::ExternalDataInfo>(filepath, file_offset, byte_size);
+  *out = static_cast<OrtExternalInitializerInfo*>(ext_data_info.release());
+  return nullptr;
+  API_IMPL_END
+}
+
 ORT_API(const ORTCHAR_T*, OrtApis::ExternalInitializerInfo_GetFilePath, _In_ const OrtExternalInitializerInfo* info) {
   return info->GetRelPath().c_str();
 }
@@ -4142,6 +4151,7 @@ static constexpr OrtApi ort_api_1_to_23 = {
     &OrtApis::Node_GetGraph,
     &OrtApis::Node_GetEpName,
     &OrtApis::ReleaseExternalInitializerInfo,
+    &OrtApis::CreateExternalInitializerInfo,
     &OrtApis::ExternalInitializerInfo_GetFilePath,
     &OrtApis::ExternalInitializerInfo_GetFileOffset,
     &OrtApis::ExternalInitializerInfo_GetByteSize,

onnxruntime/core/session/ort_apis.h

@@ -700,6 +700,8 @@ ORT_API_STATUS_IMPL(Node_GetEpName, _In_ const OrtNode* node, _Outptr_result_may
 
 // OrtExternalInitializerInfo
 ORT_API(void, ReleaseExternalInitializerInfo, _Frees_ptr_opt_ OrtExternalInitializerInfo* info);
+ORT_API_STATUS_IMPL(CreateExternalInitializerInfo, _In_ const ORTCHAR_T* filepath, _In_ int64_t file_offset,
+                    _In_ size_t byte_size, _Outptr_ OrtExternalInitializerInfo** out);
 ORT_API(const ORTCHAR_T*, ExternalInitializerInfo_GetFilePath, _In_ const OrtExternalInitializerInfo* info);
 ORT_API(int64_t, ExternalInitializerInfo_GetFileOffset, _In_ const OrtExternalInitializerInfo* info);
 ORT_API(size_t, ExternalInitializerInfo_GetByteSize, _In_ const OrtExternalInitializerInfo* info);

onnxruntime/test/providers/qnn/qnn_ep_context_test.cc

@@ -687,29 +687,62 @@ struct CustomInitializerHandlerState {
 
 static OrtStatus* ORT_API_CALL TestHandleInitializerDataFunc(void* state,
                                                              const char* initializer_name,
-                                                             const void* initializer_data,
-                                                             size_t initializer_num_bytes,
-                                                             const OrtTypeInfo* initializer_type,
-                                                             bool* is_external,
-                                                             const ORTCHAR_T** location, int64_t* offset) {
+                                                             const OrtValue* initializer_value,
+                                                             const OrtExternalInitializerInfo* external_info,
+                                                             OrtExternalInitializerInfo** new_external_info) {
+  const OrtApi& ort_api = Ort::GetApi();
   CustomInitializerHandlerState* custom_state = reinterpret_cast<CustomInitializerHandlerState*>(state);
-  (void)initializer_type;
 
   if (std::string("constant") == initializer_name) {
     // Keep a specific initializer in the model just to test both scenarios.
     // A real implementation may check the byte size and keep small initializers in the model.
-    *is_external = false;
+    *new_external_info = nullptr;
     return nullptr;
   }
 
-  // Other initializers are stored in an external file.
-  *is_external = true;
-  *offset = custom_state->outfile->tellp();
-  *location = custom_state->external_file_path;
+  //
+  // Store other initializers in an external file.
+  //
 
-  custom_state->outfile->write(static_cast<const char*>(initializer_data), initializer_num_bytes);
+  // If the original initializer was stored in an external file, keep it there (just for testing).
+  if (external_info != nullptr) {
+    const ORTCHAR_T* location = ort_api.ExternalInitializerInfo_GetFilePath(external_info);
+    int64_t offset = ort_api.ExternalInitializerInfo_GetFileOffset(external_info);
+    size_t byte_size = ort_api.ExternalInitializerInfo_GetByteSize(external_info);
+
+    if (OrtStatus* status = ort_api.CreateExternalInitializerInfo(location, offset, byte_size, new_external_info);
+        status != nullptr) {
+      return status;
+    }
+
+    return nullptr;
+  }
+
+  // Get initializer's byte size
+  size_t byte_size = 0;
+  if (OrtStatus* status = ort_api.GetTensorSizeInBytes(initializer_value, &byte_size); status != nullptr) {
+    return status;
+  }
+
+  // Get initializer's data.
+  const void* initializer_data = nullptr;
+  if (OrtStatus* status = ort_api.GetTensorData(initializer_value, &initializer_data); status != nullptr) {
+    return status;
+  }
+
+  // Write initializer data to some file.
+  int64_t offset = custom_state->outfile->tellp();
+  const ORTCHAR_T* location = custom_state->external_file_path;
+
+  custom_state->outfile->write(static_cast<const char*>(initializer_data), byte_size);
   custom_state->outfile->flush();
 
+  // Provide caller (ORT) with the new external info.
+  if (OrtStatus* status = ort_api.CreateExternalInitializerInfo(location, offset, byte_size, new_external_info);
+      status != nullptr) {
+    return status;
+  }
+
   return nullptr;
 }