Reduce IR Constant Memory Retention

src/core/tests/xml_util/custom_ir.cpp

@@ -257,7 +257,13 @@ class XmlDeserializer : public ov::util::XmlDeserializer {
                              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)
-        : ov::util::XmlDeserializer(node, origin_weights, opsets, extensions, variables, version),
+        : ov::util::XmlDeserializer(
+              node,
+              origin_weights ? std::make_shared<ov::util::BufferWeightsProvider>(origin_weights) : nullptr,
+              opsets,
+              extensions,
+              variables,
+              version),
           m_origin_weights{origin_weights},
           m_weights_map{std::ref(weights_map)} {}
 
@@ -349,13 +355,12 @@ class XmlDeserializer : public ov::util::XmlDeserializer {
 private:
     std::unique_ptr<ov::util::XmlDeserializer> make_visitor(
         const pugi::xml_node& node,
-        const std::shared_ptr<ov::AlignedBuffer>& origin_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) const override {
         return std::make_unique<XmlDeserializer>(node,
-                                                 origin_weights,
+                                                 m_origin_weights,
                                                  m_weights_map,
                                                  opsets,
                                                  extensions,

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>
@@ -45,10 +47,46 @@ 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 BufferWeightsProvider : public WeightsProvider {
+public:
+    explicit BufferWeightsProvider(std::shared_ptr<ov::AlignedBuffer> weights);
+
+    std::shared_ptr<ov::AlignedBuffer> load_region(size_t offset, size_t size) const override;
+    size_t size() const override;
+
+private:
+    std::shared_ptr<ov::AlignedBuffer> m_weights;
+};
+
+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:
+    using WeightsRegionKey = std::pair<size_t, size_t>;
+
+    std::filesystem::path m_weights_path;
+    size_t m_weights_size = 0;
+    size_t m_weights_source_id = 0;
+    std::shared_ptr<ov::AlignedBuffer> m_weights_source_handle;
+    mutable std::map<WeightsRegionKey, std::shared_ptr<ov::AlignedBuffer>> m_loaded_weights_regions;
+};
+
 class XmlDeserializer : public ov::AttributeVisitor {
 public:
     explicit XmlDeserializer(const pugi::xml_node& node,
-                             const std::shared_ptr<ov::AlignedBuffer>& weights,
+                             std::shared_ptr<WeightsProvider> weights_provider,
                              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,
@@ -77,8 +115,8 @@ class XmlDeserializer : public ov::AttributeVisitor {
     virtual void set_constant_num_buffer(ov::AttributeAdapter<std::shared_ptr<ov::AlignedBuffer>>& adapter);
 
     const pugi::xml_node& get_node() const;
-    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:
@@ -106,10 +144,8 @@ class XmlDeserializer : public ov::AttributeVisitor {
 
     /// \brief Traverses xml node representation in order to create ov function for it.
     /// \param node xml node representation
-    /// \param weights weights attached to current node
     /// \return shared pointer to function representing input node
-    std::shared_ptr<ov::Model> parse_function(const pugi::xml_node& root,
-                                              const std::shared_ptr<ov::AlignedBuffer>& weights);
+    std::shared_ptr<ov::Model> parse_function(const pugi::xml_node& root);
     /// \brief Traverses xml node representation in order to get the purpose attribute of
     /// inputs/outputs in the body of Loop op. \param node xml node representation \return struct
     /// with value of purpuse attribute
@@ -119,7 +155,6 @@ class XmlDeserializer : public ov::AttributeVisitor {
 
     std::shared_ptr<ov::Node> create_node(const ov::OutputVector& inputs,
                                           const pugi::xml_node& node,
-                                          const std::shared_ptr<ov::AlignedBuffer>& weights,
                                           const GenericLayerParams& params);
 
     void read_meta_data(const std::shared_ptr<ov::Model>& model, const pugi::xml_node& meta_section);
@@ -130,17 +165,19 @@ class XmlDeserializer : public ov::AttributeVisitor {
 
     virtual std::unique_ptr<XmlDeserializer> make_visitor(
         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) const {
-        return std::make_unique<XmlDeserializer>(node, weights, opsets, extensions, variables, version);
+        return std::make_unique<XmlDeserializer>(node, get_weights_provider(), opsets, extensions, variables, version);
     }
 
+    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;
@@ -150,7 +187,6 @@ class XmlDeserializer : public ov::AttributeVisitor {
     /// it will be used during Inputs/Outputs Description creation in SubGraph processing
     ///
     IoMap io_map;
-
     int64_t m_version;
 };
 

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,86 @@
 
 namespace ov::util {
 
+namespace {
+
+class FileRegionBuffer : public ov::AlignedBuffer {
+public:
+    FileRegionBuffer(size_t size, size_t source_id, size_t offset, std::shared_ptr<ov::AlignedBuffer> source_handle)
+        : ov::AlignedBuffer(size),
+          m_source_handle(std::move(source_handle)),
+          m_descriptor(ov::create_base_descriptor(source_id, offset, m_source_handle)) {}
+
+    std::shared_ptr<ov::IBufferDescriptor> get_descriptor() const override {
+        return m_descriptor;
+    }
+
+private:
+    std::shared_ptr<ov::AlignedBuffer> m_source_handle;
+    std::shared_ptr<ov::IBufferDescriptor> m_descriptor;
+};
+
+}  // namespace
+
+BufferWeightsProvider::BufferWeightsProvider(std::shared_ptr<ov::AlignedBuffer> weights)
+    : m_weights(std::move(weights)) {}
+
+std::shared_ptr<ov::AlignedBuffer> BufferWeightsProvider::load_region(size_t offset, size_t size) const {
+    OPENVINO_ASSERT(m_weights != nullptr, "Empty weights data in bin file or bin file cannot be found!");
+    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);
+}
+
+size_t BufferWeightsProvider::size() const {
+    return m_weights ? m_weights->size() : 0;
+}
+
+FileWeightsProvider::FileWeightsProvider(std::filesystem::path weights_path)
+    : m_weights_path(std::move(weights_path)),
+      m_weights_source_id(std::hash<std::decay_t<decltype(m_weights_path.native())>>{}(m_weights_path.native())),
+      m_weights_source_handle(std::make_shared<ov::AlignedBuffer>()) {
+    std::ifstream weights_stream(m_weights_path, std::ios::binary | std::ios::ate);
+    OPENVINO_ASSERT(weights_stream.is_open(), "Weights file ", m_weights_path, " cannot be opened!");
+
+    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!");
+
+    const FileWeightsProvider::WeightsRegionKey key{offset, size};
+    if (const auto found = m_loaded_weights_regions.find(key); found != m_loaded_weights_regions.end()) {
+        return found->second;
+    }
+
+    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<FileRegionBuffer>(size, m_weights_source_id, offset, m_weights_source_handle);
+
+    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);
+
+    m_loaded_weights_regions.emplace(key, buffer);
+
+    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);
@@ -407,18 +488,30 @@ class RTInfoDeserializer : public ov::AttributeVisitor {
 };
 
 XmlDeserializer::XmlDeserializer(const pugi::xml_node& node,
-                                 const std::shared_ptr<ov::AlignedBuffer>& weights,
+                                 std::shared_ptr<WeightsProvider> weights_provider,
                                  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)
     : 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 {
+    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_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")) {
@@ -823,13 +916,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)) {
@@ -868,17 +958,16 @@ void XmlDeserializer::on_adapter(const std::string& name, ov::ValueAccessor<std:
         if (body_node.empty()) {
             OPENVINO_THROW("TensorIterator has no body.");
         }
-        model = parse_function(m_node.child(name.c_str()), m_weights);
+        model = parse_function(m_node.child(name.c_str()));
     } else if (!name.compare("net")) {
-        model = parse_function(m_node, m_weights);
+        model = parse_function(m_node);
     } else {
         OPENVINO_THROW("Error: not recognized adapter name: ", name, ".");
     }
     adapter.set(model);
 }
 
 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");
@@ -892,13 +981,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)) {
@@ -913,13 +1004,12 @@ 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);
     }
 }
 
-std::shared_ptr<ov::Model> XmlDeserializer::parse_function(const pugi::xml_node& root,
-                                                           const std::shared_ptr<ov::AlignedBuffer>& weights) {
+std::shared_ptr<ov::Model> XmlDeserializer::parse_function(const pugi::xml_node& root) {
     // OV_ITT_SCOPE_CHAIN(FIRST_INFERENCE, taskChain, itt::domains::V10Reader_RT, "V10Parser", "Parse");
 
     struct FunctionNodes {
@@ -1025,7 +1115,7 @@ std::shared_ptr<ov::Model> XmlDeserializer::parse_function(const pugi::xml_node&
             inputs[realInputPortId] = input_node->output(p_output.get_real_output_port_id(e.fromPortId));
         }
 
-        auto node = create_node(inputs, p.xml, weights, p.params);
+        auto node = create_node(inputs, p.xml, p.params);
         id_to_node[layer_id] = node;
 
         if (const auto& parameter_node = ov::as_type_ptr<ov::op::v0::Parameter>(node)) {
@@ -1213,7 +1303,6 @@ static const std::string& translate_type_name(const std::string& name) {
 
 std::shared_ptr<ov::Node> XmlDeserializer::create_node(const std::vector<ov::Output<ov::Node>>& inputs,
                                                        const pugi::xml_node& node,
-                                                       const std::shared_ptr<ov::AlignedBuffer>& weights,
                                                        const GenericLayerParams& params) {
     // Check that inputs are correctly defined
     for (size_t i = 0; i < inputs.size(); i++) {
@@ -1234,7 +1323,7 @@ std::shared_ptr<ov::Node> XmlDeserializer::create_node(const std::vector<ov::Out
     const auto extensionIt = m_extensions.find(type);
 
     if (extensionIt != m_extensions.end()) {
-        auto visitor = make_visitor(node, weights, m_opsets, m_extensions, m_variables, m_version);
+        auto visitor = make_visitor(node, m_opsets, m_extensions, m_variables, m_version);
         ovNode = (*extensionIt->second).create(inputs, *visitor).at(0).get_node_shared_ptr();
     }
 
@@ -1285,7 +1374,7 @@ std::shared_ptr<ov::Node> XmlDeserializer::create_node(const std::vector<ov::Out
             constant->alloc_buffer_on_visit_attributes(false);
         }
         ovNode->set_arguments(inputs);
-        auto visitor = make_visitor(node, weights, m_opsets, m_extensions, m_variables, m_version);
+        auto visitor = make_visitor(node, m_opsets, m_extensions, m_variables, m_version);
         if (ovNode->visit_attributes(*visitor)) {
             ovNode->constructor_validate_and_infer_types();
         }
@@ -1297,7 +1386,7 @@ std::shared_ptr<ov::Node> XmlDeserializer::create_node(const std::vector<ov::Out
         ovNode = std::make_shared<ov::op::util::FrameworkNode>(inputs);
         // XmlDeserializer visitor(node, weights, m_opsets, m_extensions, m_variables, m_version);
         // ovNode->visit_attributes(visitor);
-        auto visitor = make_visitor(node, weights, m_opsets, m_extensions, m_variables, m_version);
+        auto visitor = make_visitor(node, m_opsets, m_extensions, m_variables, m_version);
         ovNode->visit_attributes(*visitor);
 
         size_t index{0};