Add BlobContainer class and derivates for each `std::vector<uint8_t…

…>` and `std::shared_ptr<ov::AlignedBuffer>` blob types
openvinotoolkit · Nov 29, 2024 · ed220a4 · ed220a4
1 parent bc4ca08
commit ed220a4
Show file tree

Hide file tree

Showing 12 changed files with 99 additions and 40 deletions.
diff --git a/src/plugins/intel_npu/src/common/include/intel_npu/common/igraph.hpp b/src/plugins/intel_npu/src/common/include/intel_npu/common/igraph.hpp
@@ -16,9 +16,66 @@
 
 namespace intel_npu {
 
+class BlobContainer {
+public:
+    virtual void* get_ptr() {
+        OPENVINO_THROW("const BlobContainer::get_ptr() method is not implemented!");
+    }
+
+    virtual size_t size() const {
+        OPENVINO_THROW("BlobContainer::size() method is not implemented!");
+    }
+
+    virtual bool release_from_memory() {
+        OPENVINO_THROW("BlobContainer::release_from_memory() method is not implemented!");
+    }
+};
+
+class BlobContainerVector : public BlobContainer {
+public:
+    BlobContainerVector(std::vector<uint8_t> blob) : _ownershipBlob(std::move(blob)) {}
+
+    void* get_ptr() override {
+        return reinterpret_cast<void*>(_ownershipBlob.data());
+    }
+
+    size_t size() const override {
+        return _ownershipBlob.size();
+    }
+
+    bool release_from_memory() override {
+        _ownershipBlob.clear();
+        _ownershipBlob.shrink_to_fit();
+        return true;
+    }
+
+private:
+    std::vector<uint8_t> _ownershipBlob;
+};
+
+class BlobContainerAlignedBuffer : public BlobContainer {
+public:
+    BlobContainerAlignedBuffer(const std::shared_ptr<ov::AlignedBuffer>& blobSO) : _ownershipBlob(blobSO) {}
+
+    void* get_ptr() override {
+        return _ownershipBlob->get_ptr();
+    }
+
+    size_t size() const override {
+        return _ownershipBlob->size();
+    }
+
+    bool release_from_memory() override {
+        return false;
+    }
+
+private:
+    std::shared_ptr<ov::AlignedBuffer> _ownershipBlob;
+};
+
 class IGraph : public std::enable_shared_from_this<IGraph> {
 public:
-    IGraph(ze_graph_handle_t handle, NetworkMetadata metadata, std::optional<std::shared_ptr<ov::AlignedBuffer>> blob)
+    IGraph(ze_graph_handle_t handle, NetworkMetadata metadata, std::optional<std::unique_ptr<BlobContainer>> blob)
         : _handle(handle),
           _metadata(std::move(metadata)) {
         if (blob.has_value()) {
@@ -98,7 +155,7 @@ class IGraph : public std::enable_shared_from_this<IGraph> {
     // first inference starts running
     std::mutex _mutex;
 
-    std::shared_ptr<ov::AlignedBuffer> _blob;
+    std::unique_ptr<BlobContainer> _blob;
 };
 
 }  // namespace intel_npu
diff --git a/src/plugins/intel_npu/src/common/include/intel_npu/common/npu.hpp b/src/plugins/intel_npu/src/common/include/intel_npu/common/npu.hpp
@@ -58,7 +58,7 @@ class ICompilerAdapter {
 public:
     virtual std::shared_ptr<IGraph> compile(const std::shared_ptr<const ov::Model>& model,
                                             const Config& config) const = 0;
-    virtual std::shared_ptr<IGraph> parse(std::shared_ptr<ov::AlignedBuffer> networkSOPtr, const Config& config) const = 0;
+    virtual std::shared_ptr<IGraph> parse(std::unique_ptr<BlobContainer> blobPtr, const Config& config) const = 0;
     virtual ov::SupportedOpsMap query(const std::shared_ptr<const ov::Model>& model, const Config& config) const = 0;
 
     virtual ~ICompilerAdapter() = default;

diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/driver_compiler_adapter.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/driver_compiler_adapter.hpp
@@ -26,7 +26,7 @@ class DriverCompilerAdapter final : public ICompilerAdapter {
 
     std::shared_ptr<IGraph> compile(const std::shared_ptr<const ov::Model>& model, const Config& config) const override;
 
-    std::shared_ptr<IGraph> parse(std::shared_ptr<ov::AlignedBuffer> networkSOPtr, const Config& config) const override;
+    std::shared_ptr<IGraph> parse(std::unique_ptr<BlobContainer> blobPtr, const Config& config) const override;
 
     ov::SupportedOpsMap query(const std::shared_ptr<const ov::Model>& model, const Config& config) const override;
 

diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/driver_graph.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/driver_graph.hpp
@@ -21,7 +21,7 @@ class DriverGraph final : public IGraph {
                 ze_graph_handle_t graphHandle,
                 NetworkMetadata metadata,
                 const Config& config,
-                std::optional<std::shared_ptr<ov::AlignedBuffer>> blob);
+                std::optional<std::unique_ptr<BlobContainer>> blob);
 
     void export_blob(std::ostream& stream) const override;
 

diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/plugin_compiler_adapter.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/plugin_compiler_adapter.hpp
@@ -21,7 +21,7 @@ class PluginCompilerAdapter final : public ICompilerAdapter {
 
     std::shared_ptr<IGraph> compile(const std::shared_ptr<const ov::Model>& model, const Config& config) const override;
 
-    std::shared_ptr<IGraph> parse(std::shared_ptr<ov::AlignedBuffer> networkSOPtr, const Config& config) const override;
+    std::shared_ptr<IGraph> parse(std::unique_ptr<BlobContainer> blobPtr, const Config& config) const override;
 
     ov::SupportedOpsMap query(const std::shared_ptr<const ov::Model>& model, const Config& config) const override;
 

diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/plugin_graph.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/plugin_graph.hpp
@@ -23,7 +23,7 @@ class PluginGraph final : public IGraph {
                 const std::shared_ptr<ZeroInitStructsHolder>& zeroInitStruct,
                 ze_graph_handle_t graphHandle,
                 NetworkMetadata metadata,
-                std::shared_ptr<ov::AlignedBuffer> blobSOPtr,
+                std::unique_ptr<BlobContainer> blobPtr,
                 const Config& config);
 
     void export_blob(std::ostream& stream) const override;

diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/driver_compiler_adapter.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/driver_compiler_adapter.cpp
@@ -206,11 +206,11 @@ std::shared_ptr<IGraph> DriverCompilerAdapter::compile(const std::shared_ptr<con
                                          std::nullopt);
 }
 
-std::shared_ptr<IGraph> DriverCompilerAdapter::parse(std::shared_ptr<ov::AlignedBuffer> networkSOPtr, const Config& config) const {
+std::shared_ptr<IGraph> DriverCompilerAdapter::parse(std::unique_ptr<BlobContainer> blobPtr, const Config& config) const {
     OV_ITT_TASK_CHAIN(PARSE_BLOB, itt::domains::NPUPlugin, "DriverCompilerAdapter", "parse");
 
     _logger.debug("parse start");
-    ze_graph_handle_t graphHandle = _zeGraphExt->getGraphHandle(reinterpret_cast<const uint8_t*>(networkSOPtr->get_ptr()), networkSOPtr->size());
+    ze_graph_handle_t graphHandle = _zeGraphExt->getGraphHandle(reinterpret_cast<const uint8_t*>(blobPtr->get_ptr()), blobPtr->size());
     _logger.debug("parse end");
 
     OV_ITT_TASK_NEXT(PARSE_BLOB, "getNetworkMeta");
@@ -221,7 +221,7 @@ std::shared_ptr<IGraph> DriverCompilerAdapter::parse(std::shared_ptr<ov::Aligned
                                          graphHandle,
                                          std::move(networkMeta),
                                          config,
-                                         std::optional<std::shared_ptr<ov::AlignedBuffer>>(std::move(networkSOPtr)));
+                                         std::optional<std::unique_ptr<BlobContainer>>(std::move(blobPtr)));
 }
 
 ov::SupportedOpsMap DriverCompilerAdapter::query(const std::shared_ptr<const ov::Model>& model,

diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/driver_graph.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/driver_graph.cpp
@@ -15,7 +15,7 @@ DriverGraph::DriverGraph(const std::shared_ptr<ZeGraphExtWrappers>& zeGraphExt,
                          ze_graph_handle_t graphHandle,
                          NetworkMetadata metadata,
                          const Config& config,
-                         std::optional<std::shared_ptr<ov::AlignedBuffer>> blob)
+                         std::optional<std::unique_ptr<BlobContainer>> blob)
     : IGraph(graphHandle, std::move(metadata), std::move(blob)),
       _zeGraphExt(zeGraphExt),
       _zeroInitStruct(zeroInitStruct),
@@ -142,13 +142,10 @@ bool DriverGraph::release_blob(const Config& config) {
         return false;
     }
 
-    if (_blob.use_count() > 1) {
-        // blob is not allocated by plugin, no need for memory optimization
-        return false;        
+    if(!_blob->release_from_memory()) {
+        return false;
     }
 
-    _blob.reset();
-
     _logger.debug("Blob is released");
 
     return true;

diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/plugin_compiler_adapter.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/plugin_compiler_adapter.cpp
@@ -83,36 +83,35 @@ std::shared_ptr<IGraph> PluginCompilerAdapter::compile(const std::shared_ptr<con
 
     _logger.debug("compile start");
     auto networkDesc = _compiler->compile(model, config);
-    auto networkSO = std::make_shared<std::vector<uint8_t>>(std::move(networkDesc.compiledNetwork));
+    auto blobPtr = std::make_unique<BlobContainerVector>(std::move(networkDesc.compiledNetwork));
     _logger.debug("compile end");
 
     ze_graph_handle_t graphHandle = nullptr;
 
     if (_zeGraphExt) {
         // Depending on the config, we may get an error when trying to get the graph handle from the compiled network
         try {
-            graphHandle = _zeGraphExt->getGraphHandle(networkSO->data(), networkSO->size());
+            graphHandle = _zeGraphExt->getGraphHandle(reinterpret_cast<const uint8_t*>(blobPtr->get_ptr()), blobPtr->size());
         } catch (...) {
             _logger.info("Failed to obtain the level zero graph handle. Inference requests for this model are not "
                          "allowed. Only exports are available");
         }
     }
-    auto networkSOPtr = std::make_shared<ov::SharedBuffer<std::shared_ptr<std::vector<uint8_t>>>>(reinterpret_cast<char*>(networkSO->data()), networkSO->size(), networkSO);
     return std::make_shared<PluginGraph>(_zeGraphExt,
                                          _compiler,
                                          _zeroInitStruct,
                                          graphHandle,
                                          std::move(networkDesc.metadata),
-                                         networkSOPtr,
+                                         std::move(blobPtr),
                                          config);
 }
 
-std::shared_ptr<IGraph> PluginCompilerAdapter::parse(std::shared_ptr<ov::AlignedBuffer> networkSOPtr, const Config& config) const {
+std::shared_ptr<IGraph> PluginCompilerAdapter::parse(std::unique_ptr<BlobContainer> blobPtr, const Config& config) const {
     OV_ITT_TASK_CHAIN(PARSE_BLOB, itt::domains::NPUPlugin, "PluginCompilerAdapter", "parse");
 
     _logger.debug("parse start");
-    std::vector<uint8_t> network(networkSOPtr->size());
-    network.assign(reinterpret_cast<uint8_t*>(networkSOPtr->get_ptr()), reinterpret_cast<uint8_t*>(networkSOPtr->get_ptr()) + networkSOPtr->size());
+    std::vector<uint8_t> network(blobPtr->size());
+    network.assign(reinterpret_cast<const uint8_t*>(blobPtr->get_ptr()), reinterpret_cast<const uint8_t*>(blobPtr->get_ptr()) + blobPtr->size());
     auto networkMeta = _compiler->parse(network, config);
     network.clear();
     network.shrink_to_fit();
@@ -121,15 +120,15 @@ std::shared_ptr<IGraph> PluginCompilerAdapter::parse(std::shared_ptr<ov::Aligned
     ze_graph_handle_t graphHandle = nullptr;
 
     if (_zeGraphExt) {
-        graphHandle = _zeGraphExt->getGraphHandle(reinterpret_cast<const uint8_t*>(networkSOPtr->get_ptr()), networkSOPtr->size());
+        graphHandle = _zeGraphExt->getGraphHandle(reinterpret_cast<const uint8_t*>(blobPtr->get_ptr()), blobPtr->size());
     }
 
     return std::make_shared<PluginGraph>(_zeGraphExt,
                                          _compiler,
                                          _zeroInitStruct,
                                          graphHandle,
                                          std::move(networkMeta),
-                                         std::move(networkSOPtr),
+                                         std::move(blobPtr),
                                          config);
 }
 

diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/plugin_graph.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/plugin_graph.cpp
@@ -15,9 +15,9 @@ PluginGraph::PluginGraph(const std::shared_ptr<ZeGraphExtWrappers>& zeGraphExt,
                          const std::shared_ptr<ZeroInitStructsHolder>& zeroInitStruct,
                          ze_graph_handle_t graphHandle,
                          NetworkMetadata metadata,
-                         std::shared_ptr<ov::AlignedBuffer> blobSOPtr,
+                         std::unique_ptr<BlobContainer> blobPtr,
                          const Config& config)
-    : IGraph(graphHandle, std::move(metadata), std::optional<std::shared_ptr<ov::AlignedBuffer>>(std::move(blobSOPtr))),
+    : IGraph(graphHandle, std::move(metadata), std::optional<std::unique_ptr<BlobContainer>>(std::move(blobPtr))),
       _zeGraphExt(zeGraphExt),
       _zeroInitStruct(zeroInitStruct),
       _compiler(compiler),
@@ -58,7 +58,7 @@ std::vector<ov::ProfilingInfo> PluginGraph::process_profiling_output(const std::
     // Only if we work with std::vector<uint8_t> blobs, but then IGraph needs to have 2 declarations for the same blob
     // Maybe if we templatize blob in IGraph to be either std::vector<uint8_t> or std::shared_ptr<ov::AlignedBuffer>?
     std::vector<uint8_t> blob(_blob->size());
-    blob.assign(reinterpret_cast<uint8_t*>(_blob->get_ptr()), reinterpret_cast<uint8_t*>(_blob->get_ptr()) + _blob->size());
+    blob.assign(reinterpret_cast<const uint8_t*>(_blob->get_ptr()), reinterpret_cast<const uint8_t*>(_blob->get_ptr()) + _blob->size());
     return _compiler->process_profiling_output(profData, blob, config);
 }
 

diff --git a/src/plugins/intel_npu/src/plugin/include/plugin.hpp b/src/plugins/intel_npu/src/plugin/include/plugin.hpp
@@ -44,14 +44,16 @@ class Plugin : public ov::IPlugin {
 
     std::shared_ptr<ov::ICompiledModel> import_model(std::istream& stream, const ov::AnyMap& properties) const override;
 
-    std::shared_ptr<ov::ICompiledModel> import_model(std::shared_ptr<ov::AlignedBuffer> model_buffer,
-                                                    const ov::AnyMap& properties) const override;
+    std::shared_ptr<ov::ICompiledModel> import_model(std::istream& /* unusedStream */,
+                                                     std::shared_ptr<ov::AlignedBuffer> model_buffer,
+                                                     const ov::AnyMap& properties) const override;
 
     std::shared_ptr<ov::ICompiledModel> import_model(std::istream& stream,
                                                      const ov::SoPtr<ov::IRemoteContext>& context,
                                                      const ov::AnyMap& properties) const override;
 
-    std::shared_ptr<ov::ICompiledModel> import_model(std::shared_ptr<ov::AlignedBuffer> model_buffer,
+    std::shared_ptr<ov::ICompiledModel> import_model(std::istream& stream,
+                                                     std::shared_ptr<ov::AlignedBuffer> model_buffer,
                                                      const ov::SoPtr<ov::IRemoteContext>& context,
                                                      const ov::AnyMap& properties) const override;
 

diff --git a/src/plugins/intel_npu/src/plugin/src/plugin.cpp b/src/plugins/intel_npu/src/plugin/src/plugin.cpp
@@ -769,15 +769,16 @@ std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& stream, c
         auto compiler = getCompiler(localConfig);
 
         auto graphSize = getFileSize(stream);
-        auto blobSO = std::make_shared<std::vector<uint8_t>>(graphSize);
-        stream.read(reinterpret_cast<char*>(blobSO->data()), graphSize);
+
+        std::vector<uint8_t> blob(graphSize);
+        stream.read(reinterpret_cast<char*>(blob.data()), graphSize);
         if (!stream) {
             OPENVINO_THROW("Failed to read data from stream!");
         }
         _logger.debug("Successfully read %zu bytes into blob.", graphSize);
 
-        auto blobSOPtr = std::make_shared<ov::SharedBuffer<std::shared_ptr<std::vector<uint8_t>>>>(reinterpret_cast<char*>(blobSO->data()), graphSize, blobSO);
-        auto graph = compiler->parse(std::move(blobSOPtr), localConfig);
+        auto blobContainerPtr = std::make_unique<BlobContainerVector>(std::move(blob));
+        auto graph = compiler->parse(std::move(blobContainerPtr), localConfig);
         graph->update_network_name("net" + std::to_string(_compiledModelLoadCounter++));
 
         const std::shared_ptr<ov::Model> modelDummy =
@@ -795,7 +796,9 @@ std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& stream, c
     return compiledModel;
 }
 
-std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::shared_ptr<ov::AlignedBuffer> model_buffer, const ov::AnyMap& properties) const { 
+std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& /* unusedStream */,
+                                                         std::shared_ptr<ov::AlignedBuffer> model_buffer, 
+                                                         const ov::AnyMap& properties) const { 
     OV_ITT_SCOPED_TASK(itt::domains::NPUPlugin, "Plugin::import_model");
     OV_ITT_TASK_CHAIN(PLUGIN_IMPORT_MODEL, itt::domains::NPUPlugin, "Plugin::import_model", "merge_configs");
 
@@ -820,8 +823,8 @@ std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::shared_ptr<ov::Ali
 
     try {
         auto compiler = getCompiler(localConfig);
-
-        auto graph = compiler->parse(model_buffer, localConfig);
+        auto blobContainerPtr = std::make_unique<BlobContainerAlignedBuffer>(model_buffer);
+        auto graph = compiler->parse(std::move(blobContainerPtr), localConfig);
         graph->update_network_name("net" + std::to_string(_compiledModelLoadCounter++));
 
         const std::shared_ptr<ov::Model> modelDummy =
@@ -850,15 +853,16 @@ std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& stream,
     return import_model(stream, context, properties);
 }
 
-std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::shared_ptr<ov::AlignedBuffer> model_buffer,
+std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& stream,
+                                                         std::shared_ptr<ov::AlignedBuffer> model_buffer,
                                                          const ov::SoPtr<ov::IRemoteContext>& context,
                                                          const ov::AnyMap& properties) const {
     auto casted = std::dynamic_pointer_cast<RemoteContextImpl>(context._ptr);
     if (casted == nullptr) {
         OPENVINO_THROW("Invalid remote context type. Can't cast to ov::intel_npu::RemoteContext type");
     }
 
-    return import_model(model_buffer, properties);
+    return import_model(stream, model_buffer, properties);
 }
 
 ov::SupportedOpsMap Plugin::query_model(const std::shared_ptr<const ov::Model>& model,