Test External Info

onnxruntime/core/framework/tensor_external_data_info.cc

@@ -116,7 +116,7 @@ void ExternalDataInfo::SetExternalLocationToProto(const std::filesystem::path& e
   length->set_value(std::to_string(tensor_bytes_size));
 }
 
-std::ostream& ExternalDataInfo::AddPrepackedEntriesToProto(
+std::ostream& ExternalDataInfo::WritePrepackedToFileAndAddToProto(
     const PrepackedForSerialization::BlobsInderect& prepacked_for_write, bool align, int64_t allocation_granularity,
     std::ostream& os, int64_t& external_offset, ::ONNX_NAMESPACE::TensorProto& proto) {
   for (const auto& iter : prepacked_for_write) {

onnxruntime/core/framework/tensor_external_data_info.h

@@ -62,11 +62,11 @@ class ExternalDataInfo {
     return stream;
   }
 
-  static std::ostream& AddPrepackedEntriesToProto(const PrepackedForSerialization::BlobsInderect& prepacked_for_write,
-                                                  bool align, int64_t allocation_granularity,
-                                                  std::ostream& os,
-                                                  int64_t& external_offset,
-                                                  ::ONNX_NAMESPACE::TensorProto& proto);
+  static std::ostream& WritePrepackedToFileAndAddToProto(const PrepackedForSerialization::BlobsInderect& prepacked_for_write,
+                                                         bool align, int64_t allocation_granularity,
+                                                         std::ostream& os,
+                                                         int64_t& external_offset,
+                                                         ::ONNX_NAMESPACE::TensorProto& proto);
 
   using PrepackedInfo = std::tuple<OFFSET_TYPE, size_t, std::string>;
   using PrepackedInfos = std::unordered_map<std::string, std::vector<PrepackedInfo>>;

onnxruntime/core/graph/graph.cc

@@ -4182,7 +4182,7 @@ Status Graph::ToGraphProtoWithExternalInitiallizersImpl(
       if (prepacked_parent_graph != nullptr) {
         const auto* iters_to_blobs = prepacked_parent_graph->GetBlobsForWeight(initializer.name());
         if (iters_to_blobs != nullptr && !iters_to_blobs->empty()) {
-          ORT_RETURN_IF_NOT(ExternalDataInfo::AddPrepackedEntriesToProto(
+          ORT_RETURN_IF_NOT(ExternalDataInfo::WritePrepackedToFileAndAddToProto(
               *iters_to_blobs, model_saving_options.align_offset,
               model_saving_options.allocation_granularity,
               external_stream, external_offset, *output_proto));

onnxruntime/test/framework/session_state_test.cc

@@ -378,7 +378,7 @@ class PrePackingTestOpKernel : public OpKernel {
     ORT_UNUSED_PARAMETER(tensor);
     ORT_UNUSED_PARAMETER(input_idx);
 
-    size_t weight_packed_len = 8;
+    constexpr const size_t weight_packed_len = sizeof(float) * 2;
     weight_packed_ = IAllocator::MakeUniquePtr<void>(alloc, weight_packed_len, true);
     float* data_weights_packed = reinterpret_cast<float*>(weight_packed_.get());
     data_weights_packed[0] = 1.2345f;

onnxruntime/test/framework/tensorutils_test.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "core/framework/prepacked_weights.h"
+#include "core/framework/prepacked_weights_container.h"
 #include "core/framework/tensorprotoutils.h"
 #include "core/graph/onnx_protobuf.h"
 #include "test/util/include/asserts.h"
@@ -19,6 +21,72 @@ using namespace ONNX_NAMESPACE;
 namespace onnxruntime {
 namespace test {
 
+// Test ExternalData functionality
+TEST(TensorProtoUtilsTest, SetExternalDataInformation) {
+  ONNX_NAMESPACE::TensorProto tensor_proto;
+  const std::filesystem::path kExternalDataPath("test.bin");
+  const int64_t init_offset = 100;
+  const size_t init_length = 200;
+
+  ExternalDataInfo::SetExternalLocationToProto(kExternalDataPath, init_offset, init_length, tensor_proto);
+
+  ASSERT_EQ(tensor_proto.data_location(), ONNX_NAMESPACE::TensorProto_DataLocation::TensorProto_DataLocation_EXTERNAL);
+  ASSERT_EQ(tensor_proto.external_data_size(), 3);
+  ASSERT_EQ(tensor_proto.external_data(0).key(), "location");
+  ASSERT_EQ(tensor_proto.external_data(0).value(), ToUTF8String(kExternalDataPath.native()));
+  ASSERT_EQ(tensor_proto.external_data(1).key(), "offset");
+  ASSERT_EQ(tensor_proto.external_data(1).value(), std::to_string(init_offset));
+  ASSERT_EQ(tensor_proto.external_data(2).key(), "length");
+  ASSERT_EQ(tensor_proto.external_data(2).value(), std::to_string(init_length));
+
+  PrepackedForSerialization prepacked_for_serialization;
+  PrePackedWeights prepacked_weights;
+  constexpr size_t buffer_size = 100;
+  const std::string init_name = "test_initializer";
+  const std::string blob_key = "test_key";
+
+  auto delete_array = [](void* p) { delete[] reinterpret_cast<char*>(p); };
+
+  prepacked_weights.buffers_.push_back(BufferUniquePtr(new char[buffer_size],
+                                                       delete_array));
+  prepacked_weights.buffer_sizes_.push_back(buffer_size);
+  // Write a second entry like this
+  prepacked_weights.buffers_.push_back(BufferUniquePtr(new char[buffer_size],
+                                                       delete_array));
+  prepacked_weights.buffer_sizes_.push_back(buffer_size);
+
+  prepacked_for_serialization.MainGraph().WritePackedForSaving(init_name, blob_key, std::move(prepacked_weights));
+
+  const int64_t starting_offset = 300;
+  int64_t external_offset = starting_offset;
+  std::stringstream ss;
+  const auto* blobs_for_weight = prepacked_for_serialization.MainGraph().GetBlobsForWeight(init_name);
+  ASSERT_TRUE(blobs_for_weight != nullptr);
+  ASSERT_TRUE(ExternalDataInfo::WritePrepackedToFileAndAddToProto(*blobs_for_weight,
+                                                                  true, 0, ss, external_offset, tensor_proto));
+
+  auto external_data_info = std::make_unique<ExternalDataInfo>();
+  ASSERT_STATUS_OK(ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info));
+
+  // This should have prepacked_data entry with two blobs for a single key.
+  ASSERT_TRUE(external_data_info->HasPrepackedInfo());
+  auto prepacked_infos = external_data_info->TakePrepackedInfos();
+  ASSERT_EQ(prepacked_infos.size(), 1U);
+  ASSERT_TRUE(prepacked_infos.count(blob_key) > 0);
+
+  int64_t final_offset = starting_offset;
+  for (const auto& blob_info : prepacked_infos[blob_key]) {
+    int64_t offset = std::get<0>(blob_info);
+    ASSERT_GT(offset, final_offset);
+    size_t length = std::get<1>(blob_info);
+    std::string checksum = std::get<2>(blob_info);  // currently "0"
+    final_offset = offset + length;
+    ASSERT_EQ(length, buffer_size);
+    ASSERT_EQ(checksum, "0");
+  }
+  ASSERT_EQ(final_offset, external_offset);
+}
+
 // T must be float for double, and it must match with the 'type' argument
 template <typename T>
 void TestUnpackFloatTensor(TensorProto_DataType type, const std::filesystem::path& model_path) {