Decouple external weight processing from shared context and add support for stop context sharing

Author: javier-intel

onnxruntime/core/providers/openvino/backend_manager.cc

@@ -83,22 +83,23 @@ BackendManager::BackendManager(SessionContext& session_context,
   }
   std::string device_type = session_context_.device_type;
 
-  auto& sw = shared_context_.shared_weights;
-  if (session_context_.so_share_ep_contexts) {
-    std::filesystem::path weight_filename = session_context_.onnx_model_path_name.parent_path();
-    if (sw.external_weight_filename.empty() && !sw.metadata.empty()) {
-      // Reasonable assumption that all metadata entries have the same external file location
-      sw.external_weight_filename = sw.metadata.begin()->second.location;
-    }
-    weight_filename /= sw.external_weight_filename;
-    std::ifstream weight_file(weight_filename);
+  // Check if model is using external weights
+  if (auto filename = backend_utils::GetExternalWeightFilename(subgraph)) {
+    std::filesystem::path weights_filepath = session_context_.onnx_model_path_name.parent_path() / filename.value();
 
-    if (weight_file) {
-      if (!sw.mapped_weights) {
-        sw.mapped_weights = std::make_unique<SharedContext::SharedWeights::WeightsFile>(weight_filename);
-      }
-      backend_utils::CreateOVTensors(session_context_.device_type, sw.metadata, *sw.mapped_weights);
+    // Initialize external weights with fully qualified path
+    if (!std::filesystem::exists(weights_filepath)) {
+      ORT_THROW("Error: Failed to locate weight file at ", weights_filepath.string());
     }
+
+    external_weights_.emplace(weights_filepath);
+  }
+
+  if (session_context_.so_share_ep_contexts) {
+    ORT_ENFORCE(external_weights_.has_value(), "Expected external weight object to be valid");
+    backend_utils::CreateOVTensors(session_context_.device_type,
+                                   shared_context_.shared_weights.metadata,
+                                   external_weights_.value());
   }
 
   if (ModelHasSymbolicInputDims(subgraph)) {
@@ -324,7 +325,7 @@ static bool IsQDQGraph(const onnxruntime::GraphViewer& graph_viewer) {
 static void DumpOpenVINOEPModel([[maybe_unused]] const std::filesystem::path& onnx_model_path_name,
                                 [[maybe_unused]] ONNX_NAMESPACE::ModelProto* model_proto,
                                 [[maybe_unused]] const onnxruntime::Node& fused_node) {
-#ifndef  RELEASE
+#ifdef NOT_RELEASE
   if (openvino_ep::backend_utils::IsDebugEnabled()) {
     auto model_name = onnx_model_path_name.empty() ? "unknown.onnx" : onnx_model_path_name.filename();
 
@@ -384,7 +385,12 @@ BackendManager::GetModelProtoFromFusedNode(const onnxruntime::Node& fused_node,
   if (session_context_.device_type.find("NPU") != std::string::npos &&
       (enable_ovep_qdq_optimizer || session_context_.so_share_ep_contexts)) {
     std::unique_ptr<onnxruntime::Model> model;
-    Status status = CreateModelWithStrippedQDQNodes(subgraph, logger, session_context_.so_share_ep_contexts, enable_ovep_qdq_optimizer, model, shared_context_.shared_weights);
+    Status status = CreateModelWithStrippedQDQNodes(subgraph,
+                                                    logger,
+                                                    session_context_.so_share_ep_contexts,
+                                                    enable_ovep_qdq_optimizer,
+                                                    model,
+                                                    shared_context_.shared_weights.metadata);
     auto model_proto = model->ToProto();
     model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION);
     print_model_proto_duration();

onnxruntime/core/providers/openvino/backend_manager.h

@@ -54,6 +54,7 @@ class BackendManager {
   EPCtxHandler& ep_ctx_handle_;
   SessionContext& session_context_;
   SharedContext& shared_context_;
+  std::optional<fs::path> external_weights_;
 };
 
 }  // namespace openvino_ep

onnxruntime/core/providers/openvino/backend_utils.cc

@@ -4,6 +4,7 @@
 #include <sstream>
 #include <fstream>
 #include <utility>
+#include <string>
 
 #include <filesystem>
 #include <stdexcept>
@@ -20,22 +21,7 @@ using Exception = ov::Exception;
 namespace onnxruntime {
 namespace openvino_ep {
 
-SharedContext::SharedWeights::WeightsFile::WeightsFile(std::filesystem::path filename) : file_(filename, std::ios::in | std::ios::binary) {
-  try {
-    file_.exceptions(std::ifstream::failbit | std::ifstream::badbit);
-    weights_size_ = file_.seekg(0, std::ios::end).tellg();
-  } catch (std::ifstream::failure& e) {
-    ORT_THROW("Error: Failed to open weight file at ", filename.string(), " ", e.what());
-  }
-}
-
-void SharedContext::SharedWeights::WeightsFile::load_weights(size_t file_offset, void* data, size_t size) {
-  ORT_ENFORCE(file_offset < weights_size_ && size <= weights_size_ && (file_offset <= weights_size_ - size), "Error: File offset is out of bounds.");
-  file_.seekg(file_offset);
-  file_.read(reinterpret_cast<char*>(data), size);
-}
-
-std::ostream& operator<<(std::ostream& stream, const SharedContext::SharedWeights::Metadata::Map& metadata) {
+std::ostream& operator<<(std::ostream& stream, const Metadata::Map& metadata) {
   try {
     stream << metadata.size();
 
@@ -69,14 +55,14 @@ std::ostream& operator<<(std::ostream& stream, const SharedContext::SharedWeight
   return stream;
 }
 
-std::istream& operator>>(std::istream& stream, SharedContext::SharedWeights::Metadata::Map& metadata) {
+std::istream& operator>>(std::istream& stream, Metadata::Map& metadata) {
   size_t map_size{0};
   try {
     stream >> map_size;
 
     while (!stream.eof()) {
-      SharedContext::SharedWeights::Metadata::Key key;
-      SharedContext::SharedWeights::Metadata::Value value;
+      Metadata::Key key;
+      Metadata::Value value;
       stream >> key.name;
       stream >> value.location;
       stream >> value.data_offset;
@@ -399,8 +385,19 @@ ov::element::Type GetOpenVINOElementType(ONNX_NAMESPACE::TensorProto_DataType dt
 
 // Function to handle tensor creation from external data
 void CreateOVTensors(const std::string& device_name,
-                     SharedContext::SharedWeights::Metadata::Map& metadata_map,
-                     SharedContext::SharedWeights::WeightsFile& weights) {
+                     Metadata::Map& metadata_map,
+                     std::filesystem::path& weights_filepath) {
+  // File is guaranteed to exist at this point
+  std::ifstream file(weights_filepath, std::ios::in | std::ios::binary);
+  file.exceptions(std::ifstream::failbit | std::ifstream::badbit);
+  size_t weights_size = std::filesystem::file_size(weights_filepath);
+
+  const auto load_weights = [&file, weights_size](size_t file_offset, void* data, size_t size) {
+    ORT_ENFORCE(file_offset < weights_size && size <= weights_size && (file_offset <= weights_size - size), "Error: File offset is out of bounds.");
+    file.seekg(file_offset);
+    file.read(reinterpret_cast<char*>(data), size);
+  };
+
   for (auto& [key, value] : metadata_map) {
     if (value.tensor) continue;
 
@@ -416,18 +413,18 @@ void CreateOVTensors(const std::string& device_name,
       auto&& remote_tensor = npu_context.create_l0_host_tensor(ov_elementType, value.dimensions, ov::intel_npu::TensorType::INPUT);
 
       // Copy data to remote tensor
-      weights.load_weights(value.data_offset, remote_tensor.get(), value.size);
+      load_weights(value.data_offset, remote_tensor.get(), value.size);
       value.tensor = std::make_shared<ov::Tensor>(remote_tensor);
     } else {
       // Use vanilla tensors
       value.tensor = std::make_shared<ov::Tensor>(ov_elementType, value.dimensions);
-      weights.load_weights(value.data_offset, value.tensor->data(), value.size);
+      load_weights(value.data_offset, value.tensor->data(), value.size);
     }
     ORT_ENFORCE(value.tensor->get_byte_size() == value.size, "Unexpected tensor size mismatch");
   }
 }
 
-void DestroyOVTensors(SharedContext::SharedWeights::Metadata::Map& metadata_map) {
+void DestroyOVTensors(Metadata::Map& metadata_map) {
   for (auto& [key, value] : metadata_map) {
     if (value.tensor) {
       value.tensor.reset();

onnxruntime/core/providers/openvino/backend_utils.h

@@ -67,9 +67,9 @@ CreateOVModel(std::string&& model,
               std::map<std::string, std::shared_ptr<ov::Node>>& const_outputs_map);
 
 void CreateOVTensors(const std::string& device_name,
-                     SharedContext::SharedWeights::Metadata::Map& metadata_map,
-                     SharedContext::SharedWeights::WeightsFile& weights);
-void DestroyOVTensors(SharedContext::SharedWeights::Metadata::Map& metadata_map);
+                     Metadata::Map& metadata_map,
+                     std::filesystem::path& weights_filepath);
+void DestroyOVTensors(Metadata::Map& metadata_map);
 
 void printPerformanceCounts(const std::vector<OVProfilingInfo>& performanceMap,
                             std::ostream& stream, std::string deviceName);

onnxruntime/core/providers/openvino/backends/basic_backend.cc

@@ -125,10 +125,12 @@ BasicBackend::BasicBackend(std::unique_ptr<ONNX_NAMESPACE::ModelProto>& model_pr
   std::function<void(OVInferRequestPtr)> initializer = [](OVInferRequestPtr) {};
   auto metadata = shared_context_.shared_weights.metadata;
   if (session_context_.so_share_ep_contexts) {
+    // When shared ep contexts is set external weight references are transformed to model inputs. This
+    // creates an initializer to populate/bind input weight tensors to each inference request
     initializer = [&metadata](OVInferRequestPtr ir_ptr) {
       const auto input_count = ir_ptr->GetNumInputs();
       for (auto i = 0u; i < input_count; i++) {
-        using Key = SharedContext::SharedWeights::Metadata::Key;
+        using Key = Metadata::Key;
         const auto tensor_key = Key{ir_ptr->GetInputTensorName(i)};
         if (metadata.contains(tensor_key)) {
           auto& value = metadata.at(tensor_key);
@@ -137,6 +139,8 @@ BasicBackend::BasicBackend(std::unique_ptr<ONNX_NAMESPACE::ModelProto>& model_pr
       }
     };
   }
+
+  // Create inference request queue and initialize according to passed function
   inferRequestsQueue_ = std::unique_ptr<InferRequestsQueue>(new InferRequestsQueue(exe_network_, num_infer_req, std::move(initializer)));
 }
 

onnxruntime/core/providers/openvino/contexts.h

@@ -18,52 +18,42 @@ namespace openvino_ep {
 
 namespace fs = std::filesystem;
 
+struct Metadata {
+  struct Key {
+    std::string name;
+    bool operator==(const Key&) const = default;
+  };
+  struct Hash {
+    std::size_t operator()(const Key& key) const noexcept {
+      return std::hash<std::string>()(key.name);
+    }
+  };
+  struct Value {
+    std::string location;
+    unsigned int data_offset;
+    unsigned int size;
+    std::vector<size_t> dimensions;
+    std::int32_t element_type;
+    std::shared_ptr<ov::Tensor> tensor;
+  };
+  using Map = std::unordered_map<Key, Value, Hash>;
+  friend std::ostream& operator<<(std::ostream& right, const Metadata::Map& metadata);
+  friend std::istream& operator>>(std::istream& right, Metadata::Map& metadata);
+};
+
 class SharedContext : public WeakSingleton<SharedContext> {
   // Keep the core alive as long as the shared SharedContext are alive.
   std::shared_ptr<OVCore> OVCore_;
 
  public:
   SharedContext() : OVCore_(OVCore::Get()) {}
   struct SharedWeights {
-    struct Metadata {
-      struct Key {
-        std::string name;
-        bool operator==(const Key&) const = default;
-      };
-      struct Hash {
-        std::size_t operator()(const Key& key) const noexcept {
-          return std::hash<std::string>()(key.name);
-        }
-      };
-      struct Value {
-        std::string location;
-        unsigned int data_offset;
-        unsigned int size;
-        std::vector<size_t> dimensions;
-        std::int32_t element_type;
-        std::shared_ptr<ov::Tensor> tensor;
-      };
-      using Map = std::unordered_map<Key, Value, Hash>;
-      friend std::ostream& operator<<(std::ostream& right, const Metadata::Map& metadata);
-      friend std::istream& operator>>(std::istream& right, Metadata::Map& metadata);
-    };
-
-    struct WeightsFile {
-      ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WeightsFile);
-      WeightsFile() = delete;
-      explicit WeightsFile(std::filesystem::path filename);
-
-      void load_weights(size_t file_offset, void* data, size_t size);
-
-     private:
-      std::ifstream file_;
-      size_t weights_size_;
-    };
-
-    fs::path external_weight_filename;
-    std::unique_ptr<WeightsFile> mapped_weights;
     Metadata::Map metadata;
   } shared_weights;
+
+  void clear() {  // Deletes the data stored in the SharedContext
+    shared_weights.metadata.clear();
+  }
 };
 
 using config_t = std::map<std::string, ov::AnyMap>;
@@ -102,6 +92,7 @@ struct ProviderInfo {
   bool so_context_embed_mode{false};       // ORT session option
   bool so_share_ep_contexts{false};        // ORT session option
   fs::path so_context_file_path{};         // ORT session option
+  bool so_stop_share_ep_contexts{false};   // ORT session option
   const ConfigOptions* config_options{NULL};
   const std::unordered_set<std::string> valid_provider_keys = {"device_type", "device_id", "device_luid", "cache_dir", "precision",
                                                                "load_config", "context", "num_of_threads", "model_priority", "num_streams", "enable_opencl_throttling", "enable_qdq_optimizer",

onnxruntime/core/providers/openvino/openvino_execution_provider.cc

@@ -65,6 +65,7 @@ OpenVINOExecutionProvider::~OpenVINOExecutionProvider() {
     backend_manager.ShutdownBackendManager();
   }
   backend_managers_.clear();
+  shared_context_.reset();
 }
 
 std::vector<std::unique_ptr<ComputeCapability>>
@@ -106,7 +107,12 @@ common::Status OpenVINOExecutionProvider::Compile(
   auto& metadata = shared_context_->shared_weights.metadata;
   if (session_context_.so_share_ep_contexts && metadata.empty()) {
     // Metadata is always read from model location, this could be a source or epctx model
-    fs::path metadata_filename = session_context_.onnx_model_path_name.parent_path() / "metadata.bin";
+    fs::path metadata_filename;
+    if (session_context_.so_context_file_path.empty()) {
+      metadata_filename = session_context_.onnx_model_path_name.parent_path() / "metadata.bin";
+    } else {
+      metadata_filename = session_context_.so_context_file_path.parent_path() / "metadata.bin";
+    }
     std::ifstream file(metadata_filename, std::ios::binary);
     if (file) {
       file >> metadata;
@@ -191,6 +197,10 @@ common::Status OpenVINOExecutionProvider::Compile(
     }
   }
 
+  if (session_context_.so_stop_share_ep_contexts) {
+    shared_context_->clear();
+  }
+
   return status;
 }
 

onnxruntime/core/providers/openvino/openvino_provider_factory.cc

@@ -26,6 +26,7 @@ void ParseConfigOptions(ProviderInfo& pi) {
   pi.so_context_embed_mode = pi.config_options->GetConfigOrDefault(kOrtSessionOptionEpContextEmbedMode, "0") == "1";
   pi.so_share_ep_contexts = pi.config_options->GetConfigOrDefault(kOrtSessionOptionShareEpContexts, "0") == "1";
   pi.so_context_file_path = pi.config_options->GetConfigOrDefault(kOrtSessionOptionEpContextFilePath, "");
+  pi.so_stop_share_ep_contexts = pi.config_options->GetConfigOrDefault(kOrtSessionOptionStopShareEpContexts, "0") == "1";
 
   if (pi.so_share_ep_contexts) {
     ov::AnyMap map;

onnxruntime/core/providers/openvino/qdq_transformations/qdq_stripping.cc

@@ -11,6 +11,7 @@
 #include <algorithm>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 
 #include "core/providers/shared_library/provider_api.h"
 #include "core/providers/openvino/qdq_transformations/qdq_stripping.h"
@@ -683,10 +684,10 @@ static void AddInitializerAsInput(onnxruntime::Graph& dst_graph,
 // Creates a new model without the DQ/Q operators in the src graph.
 Status CreateModelWithStrippedQDQNodes(const GraphViewer& src_graph,
                                        const logging::Logger& logger,
-                                       bool enable_ovep_weight_sharing,
+                                       bool transform_weight_as_input,
                                        bool enable_ovep_qdq_optimizer,
                                        /*out*/ std::unique_ptr<onnxruntime::Model>& model,
-                                       /*out*/ sw& shared_weights) {
+                                       /*out*/ Metadata::Map& weight_metadata) {
   // NOTE: This function is a re-implementation of GraphViewerToProto() in core/graph/graph_proto_serializer.cc
   // with the following differences:
   //   - Uses onnxruntime::Graph APIs instead of onnx::GraphProto APIs.
@@ -777,7 +778,7 @@ Status CreateModelWithStrippedQDQNodes(const GraphViewer& src_graph,
       continue;  // Already handled this node unit
     }
 
-    bool IsWeightSharingWithoutOVEPQDQStripping = enable_ovep_weight_sharing && !enable_ovep_qdq_optimizer;
+    bool IsWeightSharingWithoutOVEPQDQStripping = transform_weight_as_input && !enable_ovep_qdq_optimizer;
 
     if (node_unit->UnitType() == NodeUnit::Type::SingleNode) {
       AddStandaloneNodeUnit(dst_graph, src_graph, *node_unit, initializers_to_keep, IsWeightSharingWithoutOVEPQDQStripping, logger);
@@ -802,11 +803,9 @@ Status CreateModelWithStrippedQDQNodes(const GraphViewer& src_graph,
   std::sort(const_inits.begin(), const_inits.end());
 
   // initialize map for creating metadata for initilizers with external weights
-  auto& metadata = shared_weights.metadata;
-
-  const auto& insert_metadata = [&metadata](const ONNX_NAMESPACE::TensorProto& proto) {
-    sw::Metadata::Map::key_type key{proto.name()};
-    sw::Metadata::Map::mapped_type value{};
+  const auto& insert_metadata = [&weight_metadata](const ONNX_NAMESPACE::TensorProto& proto) {
+    Metadata::Map::key_type key{proto.name()};
+    Metadata::Map::mapped_type value{};
 
     using mutable_proto_t = ONNX_NAMESPACE::TensorProto*;
     auto& mutable_proto = *const_cast<mutable_proto_t>(&proto);
@@ -829,7 +828,7 @@ Status CreateModelWithStrippedQDQNodes(const GraphViewer& src_graph,
       dim = proto.dims()[index++];
     }
 
-    metadata.emplace(key, std::move(value));
+    weight_metadata.emplace(key, std::move(value));
   };
 
   // Handle constant initializers
@@ -839,7 +838,7 @@ Status CreateModelWithStrippedQDQNodes(const GraphViewer& src_graph,
     // Check if the initializer has external data
     if (initializer_tensor.has_data_location() &&
         initializer_tensor.data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL &&
-        enable_ovep_weight_sharing) {
+        transform_weight_as_input) {
       insert_metadata(initializer_tensor);
 
       // Add initializer with external data as input
@@ -867,7 +866,7 @@ Status CreateModelWithStrippedQDQNodes(const GraphViewer& src_graph,
         // Check if the initializer has external data
         if (initializer_tensor.has_data_location() &&
             initializer_tensor.data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL &&
-            enable_ovep_weight_sharing) {
+            transform_weight_as_input) {
           insert_metadata(initializer_tensor);
 
           // Add initializer as input if it has external data