[NPU] Refactor Plugin::import model in order to allow integrating the upcoming blob format in an OOP friendly manner

src/plugins/intel_npu/src/common/include/intel_npu/common/iparser.hpp

@@ -4,9 +4,17 @@
 
 #pragma once
 
+#include <variant>
+
 #include "intel_npu/common/filtered_config.hpp"
 #include "intel_npu/common/igraph.hpp"
 
+namespace ov {
+
+class ICore;
+
+}
+
 namespace intel_npu {
 
 class IParser {
@@ -19,16 +27,21 @@ class IParser {
      * @param config Used to influence the downstream flow of the implementation based on preferences.
      * @param initBlobs Optional. If provided, the "weights separation" flow is enabled and the binary objects
      * corresponding to the init schedules will be parsed as well.
-     * @param model Optional, but required if "initBlobs" is provided. The "ov::Model" object is leveraged in the
-     * "weights separation" implementation in order to extract the buffers of the weights.
+     * @param weightsSource The source of weights that should be used if the weights separation feature is enabled.
+     * "monostate" indicates the absence of weights. The "string_view" & "core" pair corresponds to the "weights path"
+     * config option.
+     * @param compatibilityDescriptor A string that describes the compatibility requirements of the blob from compiler's
+     * pov.
      * @return A wrapper over the corresponding L0 graph handles (multiple only if "initBlobs" has been provided). This
      * wrapper further details the compiled model and brings it in a state closer to execution.
      */
     virtual std::shared_ptr<IGraph> parse(
         const ov::Tensor& mainBlob,
         const FilteredConfig& config,
+        std::variant<std::monostate,
+                     std::shared_ptr<const ov::Model>,
+                     std::pair<std::string_view, std::shared_ptr<ov::ICore>>>&& weightsSource,
         const std::optional<std::vector<ov::Tensor>>& initBlobs = std::nullopt,
-        std::optional<std::shared_ptr<const ov::Model>>&& model = std::nullopt,
         const std::optional<std::string>& compatibilityDescriptor = std::nullopt) const = 0;
 
     virtual ~IParser() = default;

src/plugins/intel_npu/src/compiler_adapter/include/parser.hpp

@@ -21,8 +21,10 @@ class Parser final : public IParser {
     std::shared_ptr<IGraph> parse(
         const ov::Tensor& mainBlob,
         const FilteredConfig& config,
+        std::variant<std::monostate,
+                     std::shared_ptr<const ov::Model>,
+                     std::pair<std::string_view, std::shared_ptr<ov::ICore>>>&& weightsSource,
         const std::optional<std::vector<ov::Tensor>>& initBlobs = std::nullopt,
-        std::optional<std::shared_ptr<const ov::Model>>&& model = std::nullopt,
         const std::optional<std::string>& compatibilityDescriptor = std::nullopt) const override;
 
 private:

src/plugins/intel_npu/src/compiler_adapter/include/weightless_graph.hpp

@@ -13,6 +13,12 @@
 #include "openvino/op/constant.hpp"
 #include "openvino/runtime/iremote_context.hpp"
 
+namespace ov {
+
+class ICore;
+
+}
+
 namespace intel_npu {
 
 /**
@@ -31,7 +37,9 @@ class WeightlessGraph final : public Graph {
                     const std::vector<GraphDescriptor>& initGraphDesc,
                     std::vector<NetworkMetadata> initMetadata,
                     std::optional<std::vector<ov::Tensor>> initBlobs,
-                    std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>>&& constants,
+                    std::variant<std::monostate,
+                                 std::shared_ptr<const ov::Model>,
+                                 std::pair<std::string_view, std::shared_ptr<ov::ICore>>>&& weightsSource,
                     const FilteredConfig& config,
                     const bool blobIsPersistent = false);
 

src/plugins/intel_npu/src/compiler_adapter/include/weightless_utils.hpp

@@ -4,24 +4,10 @@
 
 #pragma once
 
-#include <memory>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
 #include "intel_npu/common/network_metadata.hpp"
-#include "openvino/core/model.hpp"
-#include "openvino/op/constant.hpp"
 
 namespace intel_npu {
 
 bool isInitMetadata(const NetworkMetadata& networkMetadata);
 
-std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> get_all_constants_in_topological_order(
-    const std::shared_ptr<const ov::Model>& model);
-
-std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> get_all_constants_memory_mapped(
-    const std::string& weightsPath,
-    const std::vector<NetworkMetadata>& initNetworkMetadata);
-
 }  // namespace intel_npu

src/plugins/intel_npu/src/compiler_adapter/src/driver_compiler_adapter.cpp

@@ -204,11 +204,6 @@ std::shared_ptr<IGraph> DriverCompilerAdapter::compileWS(std::shared_ptr<ov::Mod
         _logger.info("Compilation memory usage: Peak %lld KB", compile_model_mem_end - compile_model_mem_start);
     }
 
-    auto constants = get_all_constants_in_topological_order(model);
-    // Note: Delete model prematurely, constants are still valid due to
-    // shared_ptr semantics.
-    model = nullptr;
-
     return std::make_shared<WeightlessGraph>(_zeGraphExt,
                                              _zeroInitStruct,
                                              mainGraphHandle,
@@ -217,7 +212,7 @@ std::shared_ptr<IGraph> DriverCompilerAdapter::compileWS(std::shared_ptr<ov::Mod
                                              initGraphDescriptors,
                                              std::move(initNetworkMetadata),
                                              /* initBlobs = */ std::nullopt,
-                                             std::move(constants),
+                                             std::move(model),
                                              updatedConfig);
 }
 

src/plugins/intel_npu/src/compiler_adapter/src/parser.cpp

@@ -25,11 +25,14 @@ Parser::Parser(const std::shared_ptr<ZeroInitStructsHolder>& zeroInitStruct)
                     "Failed to create ZeGraphExtWrappers in Parser. Please check if the driver is properly installed.");
 }
 
-std::shared_ptr<IGraph> Parser::parse(const ov::Tensor& mainBlob,
-                                      const FilteredConfig& config,
-                                      const std::optional<std::vector<ov::Tensor>>& initBlobs,
-                                      std::optional<std::shared_ptr<const ov::Model>>&& model,
-                                      const std::optional<std::string>& compatibilityDescriptor) const {
+std::shared_ptr<IGraph> Parser::parse(
+    const ov::Tensor& mainBlob,
+    const FilteredConfig& config,
+    std::variant<std::monostate,
+                 std::shared_ptr<const ov::Model>,
+                 std::pair<std::string_view, std::shared_ptr<ov::ICore>>>&& weightsSource,
+    const std::optional<std::vector<ov::Tensor>>& initBlobs,
+    const std::optional<std::string>& compatibilityDescriptor) const {
     OV_ITT_TASK_CHAIN(PARSE_BLOB, itt::domains::NPUPlugin, "Parser", "parse");
 
     // Detect blob format
@@ -57,11 +60,6 @@ std::shared_ptr<IGraph> Parser::parse(const ov::Tensor& mainBlob,
     OV_ITT_TASK_NEXT(PARSE_BLOB, "getNetworkMetaMainGraph");
     mainNetworkMetadata = _zeGraphExt->getNetworkMeta(mainGraphDesc);
     _logger.debug("main schedule parse end");
-    if (model.has_value()) {
-        mainNetworkMetadata.name = model.value()->get_friendly_name();
-    } else {
-        _logger.debug("networkMeta name is empty in parse!");
-    }
 
     // exporting the blob when we get it from cache or ov::hint::compiled_blob property
     // shall be available
@@ -97,13 +95,6 @@ std::shared_ptr<IGraph> Parser::parse(const ov::Tensor& mainBlob,
     }
     _logger.debug("inits schedule parse end");
 
-    auto constants = model.has_value()
-                         ? get_all_constants_in_topological_order(model.value())
-                         : get_all_constants_memory_mapped(config.get<WEIGHTS_PATH>(), initNetworkMetadata);
-    // Note: Delete model prematurely, constants are still valid due to
-    // shared_ptr semantics.
-    model = std::nullopt;
-
     return std::make_shared<WeightlessGraph>(_zeGraphExt,
                                              _zeroInitStruct,
                                              mainGraphDesc,
@@ -112,7 +103,7 @@ std::shared_ptr<IGraph> Parser::parse(const ov::Tensor& mainBlob,
                                              initGraphDescriptors,
                                              std::move(initNetworkMetadata),
                                              initBlobs,
-                                             std::move(constants),
+                                             std::move(weightsSource),
                                              config,
                                              blobIsPersistent);
 }

src/plugins/intel_npu/src/compiler_adapter/src/plugin_compiler_adapter.cpp

@@ -233,11 +233,6 @@ std::shared_ptr<IGraph> PluginCompilerAdapter::compileWS(std::shared_ptr<ov::Mod
 
     _logger.debug("compile end");
 
-    auto constants = get_all_constants_in_topological_order(model);
-    // Note: Delete model prematurely, constants are still valid due to
-    // shared_ptr semantics.
-    model = nullptr;
-
     return std::make_shared<WeightlessGraph>(
         _zeGraphExt,
         _zeroInitStruct,
@@ -247,7 +242,7 @@ std::shared_ptr<IGraph> PluginCompilerAdapter::compileWS(std::shared_ptr<ov::Mod
         initGraphDescriptors,
         std::move(initNetworkMetadata),
         tensorsInits,
-        std::move(constants),
+        std::move(model),
         localConfig,
         /* persistentBlob = */ true);  // exporting the blob shall be available in such a scenario
 }

src/plugins/intel_npu/src/compiler_adapter/src/weightless_graph.cpp

@@ -16,9 +16,13 @@
 #include "intel_npu/utils/zero/zero_cmd_queue_pool.hpp"
 #include "intel_npu/utils/zero/zero_utils.hpp"
 #include "openvino/core/memory_util.hpp"
+#include "openvino/core/rt_info/weightless_caching_attributes.hpp"
 #include "openvino/core/weight_sharing_util.hpp"
+#include "openvino/runtime/icore.hpp"
 #include "openvino/runtime/make_tensor.hpp"
+#include "openvino/runtime/shared_buffer.hpp"
 #include "openvino/util/common_util.hpp"
+#include "openvino/util/mmap_object.hpp"
 
 #define USE_SINGLE_THREADED_RUN_INIT 1
 
@@ -27,6 +31,119 @@ namespace intel_npu {
 namespace {
 
 constexpr uint8_t MAIN_SCHEDULE_INDEX = 0;
+constexpr std::string_view WEIGHTS_IR_EXTENSION = ".bin";
+constexpr std::string_view ONNX_EXTENSION = ".onnx";
+
+std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> get_all_constants_in_topological_order(
+    const std::shared_ptr<const ov::Model>& model) {
+    std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> constants;
+
+    OPENVINO_ASSERT(model != nullptr, "Model is required to extract constants in topological order.");
+    bool at_least_one_wca_found = false;
+
+    // Match the inputs of the "init" model with the Constant nodes of the original model
+    for (auto&& node : model->get_ops()) {
+        if (!ov::is_type<ov::op::v0::Constant>(node)) {
+            continue;
+        }
+
+        auto constantNode = std::static_pointer_cast<ov::op::v0::Constant>(node);
+        ov::RTMap& runtimeInfoMap = constantNode->get_rt_info();
+        const auto& weightlessCacheAttrIt = runtimeInfoMap.find(ov::WeightlessCacheAttribute::get_type_info_static());
+        if (weightlessCacheAttrIt != runtimeInfoMap.end()) {
+            at_least_one_wca_found = true;
+            auto& weightlessCacheAttr = weightlessCacheAttrIt->second.as<ov::WeightlessCacheAttribute>();
+
+            auto& constant = constants[weightlessCacheAttr.bin_offset];
+            if (constant != nullptr) {
+                // if multiple constants point to the same buffer, ensure that
+                // their binary sizes are the same
+                OPENVINO_ASSERT(constant->get_byte_size() == constantNode->get_byte_size(),
+                                "Found ov::Constant that points to the common buffer but has mismatching byte size. "
+                                "This may indicate a bug in OV model compression.");
+                continue;
+            }
+            constant = std::move(constantNode);
+        }
+    }
+
+    OPENVINO_ASSERT(at_least_one_wca_found,
+                    "No \"WeightlessCacheAttribute\" has been found in any of the model's Constant nodes. This "
+                    "attribute is required for running the \"weights separation\" flow.");
+
+    return constants;
+}
+
+std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> get_all_constants_memory_mapped(
+    std::string_view weightsPath,
+    const std::vector<NetworkMetadata>& initNetworkMetadata) {
+    std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> constants;
+
+    auto mapped_memory = ov::load_mmap_object(weightsPath);
+    for (const auto& initMetadata : initNetworkMetadata) {
+        for (const IODescriptor& descriptor : initMetadata.inputs) {
+            const auto& opt = ov::util::view_to_number<size_t>(descriptor.nameFromCompiler);
+            OPENVINO_ASSERT(opt.has_value(), "Failed to parse id for constant: ", descriptor.nameFromCompiler);
+
+            const size_t id = opt.value();
+            const size_t byte_size =
+                ov::util::get_memory_size(descriptor.precision, shape_size(descriptor.shapeFromCompiler.to_shape()));
+            OPENVINO_ASSERT(id <= mapped_memory->size() && byte_size <= mapped_memory->size() - id,
+                            "Constant offset/size is out of bounds for mapped weights file: offset=",
+                            id,
+                            ", size=",
+                            byte_size,
+                            ", file_size=",
+                            mapped_memory->size(),
+                            ", name=",
+                            descriptor.nameFromCompiler);
+            auto weight_buffer =
+                std::make_shared<ov::SharedBuffer<std::shared_ptr<ov::MappedMemory>>>(mapped_memory->data() + id,
+                                                                                      byte_size,
+                                                                                      mapped_memory);
+
+            auto [it, inserted] =
+                constants.emplace(id,
+                                  std::make_shared<ov::op::v0::Constant>(descriptor.precision,
+                                                                         descriptor.shapeFromCompiler.to_shape(),
+                                                                         weight_buffer));
+            if (!inserted) {
+                OPENVINO_ASSERT(it->second->get_byte_size() == byte_size,
+                                "Duplicate constant offset found with mismatching byte size: offset=",
+                                id);
+            }
+        }
+    }
+
+    return constants;
+}
+
+std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>> extract_constants_map(
+    std::variant<std::monostate,
+                 std::shared_ptr<const ov::Model>,
+                 std::pair<std::string_view, std::shared_ptr<ov::ICore>>>&& weightsSource,
+    const std::vector<NetworkMetadata>& initNetworkMetadata) {
+    if (const std::shared_ptr<const ov::Model>* model = std::get_if<std::shared_ptr<const ov::Model>>(&weightsSource)) {
+        return get_all_constants_in_topological_order(*model);
+    } else if (const std::pair<std::string_view, std::shared_ptr<ov::ICore>>* weightsPathAndCore =
+                   std::get_if<std::pair<std::string_view, std::shared_ptr<ov::ICore>>>(&weightsSource)) {
+        auto [weightsPath, core] = *weightsPathAndCore;
+
+        auto ext = ov::util::path_to_string(ov::util::make_path(weightsPath).extension());
+        if (ext == ONNX_EXTENSION) {
+            const auto model = core->read_model(weightsPath, weightsPath, {});
+            return get_all_constants_in_topological_order(model);
+        } else if (ext == WEIGHTS_IR_EXTENSION) {
+            return get_all_constants_memory_mapped(weightsPath, initNetworkMetadata);
+        } else {
+            OPENVINO_THROW("Invalid path to the weights: ",
+                           weightsPath,
+                           ". A \".bin\" or \".onnx\" extension was expected.");
+        }
+    }
+
+    OPENVINO_THROW("No source of weights has been provided to the \"WeightlessGraph\" ctor");
+}
 
 struct QueueData {
     int64_t initIndex = -1;
@@ -137,7 +254,9 @@ WeightlessGraph::WeightlessGraph(const std::shared_ptr<ZeGraphExtWrappers>& zeGr
                                  const std::vector<GraphDescriptor>& initGraphDesc,
                                  std::vector<NetworkMetadata> initMetadata,
                                  std::optional<std::vector<ov::Tensor>> initBlobs,
-                                 std::unordered_map<size_t, std::shared_ptr<ov::op::v0::Constant>>&& constants,
+                                 std::variant<std::monostate,
+                                              std::shared_ptr<const ov::Model>,
+                                              std::pair<std::string_view, std::shared_ptr<ov::ICore>>>&& weightsSource,
                                  const FilteredConfig& config,
                                  const bool blobIsPersistent)
     : Graph(zeGraphExt,
@@ -152,7 +271,7 @@ WeightlessGraph::WeightlessGraph(const std::shared_ptr<ZeGraphExtWrappers>& zeGr
       _initsGraphDesc(initGraphDesc),
       _initBlobs(std::move(initBlobs)),
       _initsMetadata(std::move(initMetadata)),
-      _constants(std::move(constants)),
+      _constants(extract_constants_map(std::move(weightsSource), _initsMetadata)),
       _wgLogger("WeightlessGraph", config.get<LOG_LEVEL>()) {
     if (!config.get<CREATE_EXECUTOR>() || config.get<DEFER_WEIGHTS_LOAD>()) {
         _wgLogger.info("Graph initialize is deferred from the \"WeightlessGraph\" constructor");