Reduce IR Constant Memory Retention

src/core/xml_util/include/openvino/xml_util/xml_deserialize_util.hpp

@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <cstddef>
+#include <filesystem>
 #include <functional>
 #include <map>
 #include <pugixml.hpp>
@@ -38,14 +40,35 @@ void str_to_container(const std::string& value, T& res) {
 template <>
 void str_to_container<std::vector<std::string>>(const std::string& value, std::vector<std::string>& res);
 
+class WeightsProvider {
+public:
+    virtual ~WeightsProvider() = default;
+
+    virtual std::shared_ptr<ov::AlignedBuffer> load_region(size_t offset, size_t size) const = 0;
+    virtual size_t size() const = 0;
+};
+
+class FileWeightsProvider : public WeightsProvider {
+public:
+    explicit FileWeightsProvider(std::filesystem::path weights_path);
+
+    std::shared_ptr<ov::AlignedBuffer> load_region(size_t offset, size_t size) const override;
+    size_t size() const override;
+
+private:
+    std::filesystem::path m_weights_path;
+    size_t m_weights_size = 0;
+};
+
 class XmlDeserializer : public ov::AttributeVisitor {
 public:
     explicit XmlDeserializer(const pugi::xml_node& node,
                              const std::shared_ptr<ov::AlignedBuffer>& weights,
                              const std::unordered_map<std::string, ov::OpSet>& opsets,
                              const std::unordered_map<ov::DiscreteTypeInfo, ov::BaseOpExtension::Ptr>& extensions,
                              std::unordered_map<std::string, std::shared_ptr<ov::op::util::Variable>>& variables,
-                             size_t version);
+                             size_t version,
+                             std::shared_ptr<WeightsProvider> weights_provider = nullptr);
 
     void on_adapter(const std::string& name, ov::ValueAccessor<std::string>& value) override;
 
@@ -73,6 +96,9 @@ class XmlDeserializer : public ov::AttributeVisitor {
     const std::shared_ptr<ov::AlignedBuffer>& get_weights() const {
         return m_weights;
     }
+    const std::shared_ptr<WeightsProvider>& get_weights_provider() const {
+        return m_weights_provider;
+    }
 
 private:
     struct IoMap {
@@ -128,12 +154,22 @@ class XmlDeserializer : public ov::AttributeVisitor {
         const std::unordered_map<ov::DiscreteTypeInfo, ov::BaseOpExtension::Ptr>& extensions,
         std::unordered_map<std::string, std::shared_ptr<ov::op::util::Variable>>& variables,
         size_t version) const {
-        return std::make_unique<XmlDeserializer>(node, weights, opsets, extensions, variables, version);
+        return std::make_unique<XmlDeserializer>(node,
+                                                 weights,
+                                                 opsets,
+                                                 extensions,
+                                                 variables,
+                                                 version,
+                                                 get_weights_provider());
     }
 
+    std::shared_ptr<ov::AlignedBuffer> load_weights_region(size_t offset, size_t size) const;
+    size_t get_available_weights_size() const;
+
     // -- DATA --
     const pugi::xml_node m_node;
     const std::shared_ptr<ov::AlignedBuffer>& m_weights;
+    const std::shared_ptr<WeightsProvider> m_weights_provider;
     const std::unordered_map<std::string, ov::OpSet>& m_opsets;
     const std::unordered_map<ov::DiscreteTypeInfo, ov::BaseOpExtension::Ptr>& m_extensions;
     std::unordered_map<std::string, std::shared_ptr<ov::op::util::Variable>>& m_variables;

src/core/xml_util/src/xml_deserialize_util.cpp

@@ -4,6 +4,7 @@
 
 #include "openvino/xml_util/xml_deserialize_util.hpp"
 
+#include <fstream>
 #include <regex>
 #include <stack>
 #include <string_view>
@@ -30,6 +31,36 @@
 
 namespace ov::util {
 
+FileWeightsProvider::FileWeightsProvider(std::filesystem::path weights_path) : m_weights_path(std::move(weights_path)) {
+    m_weights_size = static_cast<size_t>(std::filesystem::file_size(m_weights_path));
+}
+
+std::shared_ptr<ov::AlignedBuffer> FileWeightsProvider::load_region(size_t offset, size_t size) const {
+    OPENVINO_ASSERT(offset <= m_weights_size && size <= m_weights_size - offset, "Incorrect weights in bin file!");
+
+    std::ifstream weights_stream(m_weights_path, std::ios::binary);
+    OPENVINO_ASSERT(weights_stream.is_open(), "Weights file ", m_weights_path, " cannot be opened!");
+
+    weights_stream.seekg(static_cast<std::streamoff>(offset), std::ios::beg);
+    OPENVINO_ASSERT(weights_stream.good(), "Failed to seek weights file ", m_weights_path, " to offset ", offset);
+
+    auto buffer = std::make_shared<ov::AlignedBuffer>(size);
+    weights_stream.read(buffer->get_ptr<char>(), static_cast<std::streamsize>(size));
+    OPENVINO_ASSERT(static_cast<size_t>(weights_stream.gcount()) == size,
+                    "Failed to read ",
+                    size,
+                    " bytes from weights file ",
+                    m_weights_path,
+                    " at offset ",
+                    offset);
+
+    return buffer;
+}
+
+size_t FileWeightsProvider::size() const {
+    return m_weights_size;
+}
+
 template <>
 void str_to_container<std::vector<std::string>>(const std::string& value, std::vector<std::string>& res) {
     std::stringstream ss(value);
@@ -422,14 +453,38 @@ XmlDeserializer::XmlDeserializer(const pugi::xml_node& node,
                                  const std::unordered_map<std::string, ov::OpSet>& opsets,
                                  const std::unordered_map<ov::DiscreteTypeInfo, ov::BaseOpExtension::Ptr>& extensions,
                                  std::unordered_map<std::string, std::shared_ptr<ov::op::util::Variable>>& variables,
-                                 size_t version)
+                                 size_t version,
+                                 std::shared_ptr<WeightsProvider> weights_provider)
     : m_node(node),
       m_weights(weights),
+      m_weights_provider(std::move(weights_provider)),
       m_opsets(opsets),
       m_extensions(extensions),
       m_variables(variables),
       m_version(version) {}
 
+std::shared_ptr<ov::AlignedBuffer> XmlDeserializer::load_weights_region(size_t offset, size_t size) const {
+    if (m_weights) {
+        OPENVINO_ASSERT(offset <= m_weights->size() && size <= m_weights->size() - offset,
+                        "Incorrect weights in bin file!");
+        auto* data = m_weights->get_ptr<char>() + offset;
+        return std::make_shared<ov::SharedBuffer<std::shared_ptr<ov::AlignedBuffer>>>(data, size, m_weights);
+    }
+
+    OPENVINO_ASSERT(m_weights_provider, "Empty weights data in bin file or bin file cannot be found!");
+    return m_weights_provider->load_region(offset, size);
+}
+
+size_t XmlDeserializer::get_available_weights_size() const {
+    if (m_weights) {
+        return m_weights->size();
+    }
+    if (m_weights_provider) {
+        return m_weights_provider->size();
+    }
+    return 0;
+}
+
 ov::Any XmlDeserializer::parse_weightless_cache_attribute(const pugi::xml_node& node) const {
     ov::Any wl_attr;
     if (const auto data_node = node.child("data")) {
@@ -835,13 +890,10 @@ void XmlDeserializer::on_adapter(const std::string& name, ov::ValueAccessor<void
             if (!getParameters<int64_t>(dn, "shape", shape))
                 return;
 
-            if (!m_weights)
-                OPENVINO_THROW("Empty weights data in bin file or bin file cannot be found!");
-            if (m_weights->size() < offset + size)
-                OPENVINO_THROW("Incorrect weights in bin file!");
-            char* data = m_weights->get_ptr<char>() + offset;
-            auto buffer =
-                ov::AttributeAdapter<std::shared_ptr<ov::StringAlignedBuffer>>::unpack_string_tensor(data, size);
+            auto raw_buffer = load_weights_region(offset, size);
+            auto buffer = ov::AttributeAdapter<std::shared_ptr<ov::StringAlignedBuffer>>::unpack_string_tensor(
+                raw_buffer->get_ptr<char>(),
+                size);
             a->set(buffer);
         }
     } else if (auto a = ov::as_type<ov::AttributeAdapter<ov::op::util::FrameworkNodeAttrs>>(&adapter)) {
@@ -890,7 +942,6 @@ void XmlDeserializer::on_adapter(const std::string& name, ov::ValueAccessor<std:
 }
 
 void XmlDeserializer::set_constant_num_buffer(ov::AttributeAdapter<std::shared_ptr<ov::AlignedBuffer>>& adapter) {
-    OPENVINO_ASSERT(m_weights, "Empty weights data in bin file or bin file cannot be found!");
     std::vector<int64_t> shape;
     std::string el_type_str;
     const auto& dn = m_node.child("data");
@@ -904,13 +955,15 @@ void XmlDeserializer::set_constant_num_buffer(ov::AttributeAdapter<std::shared_p
 
     const auto size = static_cast<size_t>(pugixml::get_uint64_attr(dn, "size"));
     const auto offset = static_cast<size_t>(pugixml::get_uint64_attr(dn, "offset"));
-    OPENVINO_ASSERT(m_weights->size() >= offset + size, "Incorrect weights in bin file!");
-
-    char* data = m_weights->get_ptr<char>() + offset;
+    OPENVINO_ASSERT(offset <= get_available_weights_size() && size <= get_available_weights_size() - offset,
+                    "Incorrect weights in bin file!");
 
     const auto el_type = ov::element::Type(el_type_str);
     if (el_type == element::string) {
-        auto buffer = ov::AttributeAdapter<std::shared_ptr<ov::StringAlignedBuffer>>::unpack_string_tensor(data, size);
+        auto raw_buffer = load_weights_region(offset, size);
+        auto buffer = ov::AttributeAdapter<std::shared_ptr<ov::StringAlignedBuffer>>::unpack_string_tensor(
+            raw_buffer->get_ptr<char>(),
+            size);
         adapter.set(buffer);
     } else {
         if (size < ((ov::shape_size(shape) * el_type.bitwidth() + 7) >> 3)) {
@@ -925,7 +978,7 @@ void XmlDeserializer::set_constant_num_buffer(ov::AttributeAdapter<std::shared_p
                            ov::util::get_memory_size(el_type, ov::shape_size(shape)));
         }
 
-        auto buffer = std::make_shared<ov::SharedBuffer<std::shared_ptr<ov::AlignedBuffer>>>(data, size, m_weights);
+        auto buffer = load_weights_region(offset, size);
         adapter.set(buffer);
     }
 }

src/frontends/ir/src/frontend.cpp

@@ -5,6 +5,7 @@
 #include "openvino/frontend/ir/frontend.hpp"
 
 #include <array>
+#include <cstdlib>
 #include <optional>
 #include <pugixml.hpp>
 #include <vector>
@@ -18,6 +19,7 @@
 #include "openvino/util/file_util.hpp"
 #include "openvino/util/mmap_object.hpp"
 #include "openvino/util/xml_parse_utils.hpp"
+#include "openvino/xml_util/xml_deserialize_util.hpp"
 #include "transformations/resolve_names_collisions.hpp"
 #include "utils.hpp"
 
@@ -41,6 +43,19 @@ size_t get_ir_version(const pugi::xml_document& doc) {
 
 constexpr size_t HEADER_SIZE_LIM = 512lu;
 
+bool use_separate_const_weights_loading() {
+    const auto* flag = std::getenv("OV_IR_SEPARATE_CONST_WEIGHTS");
+    if (!flag) {
+        return false;
+    }
+
+    std::string value(flag);
+    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char character) {
+        return std::tolower(character);
+    });
+    return value == "1" || value == "true" || value == "yes" || value == "on";
+}
+
 /**
  * @brief Extracts IR version from model stream
  * @param model Model's stream
@@ -151,8 +166,9 @@ void FrontEnd::add_extension(const ov::Extension::Ptr& ext) {
 InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& variants) const {
     std::ifstream local_model_stream;
     std::istream* provided_model_stream = nullptr;
-    std::shared_ptr<ov::AlignedBuffer> model_buf;
-    std::shared_ptr<ov::AlignedBuffer> weights;
+    std::shared_ptr<ov::AlignedBuffer> model_buf;  //
+    std::shared_ptr<ov::AlignedBuffer> weights;    //
+    std::shared_ptr<ov::util::WeightsProvider> weights_provider;
 
     auto create_extensions_map = [&]() -> std::unordered_map<ov::DiscreteTypeInfo, ov::BaseOpExtension::Ptr> {
         std::unordered_map<ov::DiscreteTypeInfo, ov::BaseOpExtension::Ptr> exts;
@@ -168,16 +184,22 @@ InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& variants) const
             return std::make_shared<InputModel>(*provided_model_stream,
                                                 weights,
                                                 create_extensions_map(),
-                                                std::move(weights_path));
+                                                std::move(weights_path),
+                                                weights_provider);
         } else if (local_model_stream.is_open()) {
             auto input_model = std::make_shared<InputModel>(local_model_stream,
                                                             weights,
                                                             create_extensions_map(),
-                                                            std::move(weights_path));
+                                                            std::move(weights_path),
+                                                            weights_provider);
             local_model_stream.close();
             return input_model;
         } else if (model_buf) {
-            return std::make_shared<InputModel>(model_buf, weights, create_extensions_map(), std::move(weights_path));
+            return std::make_shared<InputModel>(model_buf,
+                                                weights,
+                                                create_extensions_map(),
+                                                std::move(weights_path),
+                                                weights_provider);
         }
         return nullptr;
     };
@@ -215,31 +237,36 @@ InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& variants) const
     }
 
     if (!weights_path.empty()) {
-        const auto enable_mmap = variants.back().is<bool>() ? variants.back().as<bool>() : false;
-        if (enable_mmap) {
-            auto mapped_memory = ov::load_mmap_object(weights_path);
-            weights = std::make_shared<ov::SharedBuffer<std::shared_ptr<MappedMemory>>>(mapped_memory->data(),
-                                                                                        mapped_memory->size(),
-                                                                                        mapped_memory);
-        } else if (std::ifstream bin_stream(weights_path, std::ios::binary); bin_stream.is_open()) {
-            bin_stream.seekg(0, std::ios::end);
-            size_t file_size = bin_stream.tellg();
-            bin_stream.seekg(0, std::ios::beg);
-
-            auto aligned_weights_buffer = std::make_shared<ov::AlignedBuffer>(file_size);
-            bin_stream.read(aligned_weights_buffer->get_ptr<char>(), aligned_weights_buffer->size());
-
-            weights = std::make_shared<ov::SharedBuffer<std::shared_ptr<ov::AlignedBuffer>>>(
-                aligned_weights_buffer->get_ptr<char>(),
-                aligned_weights_buffer->size(),
-                aligned_weights_buffer,
-                ov::create_base_descriptor(
-                    std::hash<std::decay_t<decltype(weights_path.native())>>{}(weights_path.native()),
-                    0,
-                    aligned_weights_buffer));
-
+        const auto separate_const_weights_loading = use_separate_const_weights_loading() && !weights;
+        if (separate_const_weights_loading) {
+            weights_provider = std::make_shared<ov::util::FileWeightsProvider>(weights_path);
         } else {
-            OPENVINO_THROW("Weights file ", weights_path, " cannot be opened!");
+            const auto enable_mmap = variants.back().is<bool>() ? variants.back().as<bool>() : false;
+            if (enable_mmap) {
+                auto mapped_memory = ov::load_mmap_object(weights_path);
+                weights = std::make_shared<ov::SharedBuffer<std::shared_ptr<MappedMemory>>>(mapped_memory->data(),
+                                                                                            mapped_memory->size(),
+                                                                                            mapped_memory);
+            } else if (std::ifstream bin_stream(weights_path, std::ios::binary); bin_stream.is_open()) {
+                bin_stream.seekg(0, std::ios::end);
+                size_t file_size = bin_stream.tellg();
+                bin_stream.seekg(0, std::ios::beg);
+
+                auto aligned_weights_buffer = std::make_shared<ov::AlignedBuffer>(file_size);
+                bin_stream.read(aligned_weights_buffer->get_ptr<char>(), aligned_weights_buffer->size());
+
+                weights = std::make_shared<ov::SharedBuffer<std::shared_ptr<ov::AlignedBuffer>>>(
+                    aligned_weights_buffer->get_ptr<char>(),
+                    aligned_weights_buffer->size(),
+                    aligned_weights_buffer,
+                    ov::create_base_descriptor(
+                        std::hash<std::decay_t<decltype(weights_path.native())>>{}(weights_path.native()),
+                        0,
+                        aligned_weights_buffer));
+
+            } else {
+                OPENVINO_THROW("Weights file ", weights_path, " cannot be opened!");
+            }
         }
     }