From b44fe983f5f6e05106622b3f934bce5d9661d9ab Mon Sep 17 00:00:00 2001
From: Alexey Smirnov <alexey.smirnov@intel.com>
Date: Tue, 7 Jan 2025 16:05:01 +0000
Subject: [PATCH] [NPUW] Fix for weights bank UIDs (#28297)

Previous PR https://github.com/openvinotoolkit/openvino/pull/28282
---
 src/plugins/intel_npu/src/plugin/npuw/weights_bank.cpp | 6 ++++++
 src/plugins/intel_npu/src/plugin/npuw/weights_bank.hpp | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/plugins/intel_npu/src/plugin/npuw/weights_bank.cpp b/src/plugins/intel_npu/src/plugin/npuw/weights_bank.cpp
index 7fded56c28cb50..4cc804f7b7e399 100644
--- a/src/plugins/intel_npu/src/plugin/npuw/weights_bank.cpp
+++ b/src/plugins/intel_npu/src/plugin/npuw/weights_bank.cpp
@@ -41,6 +41,8 @@ int64_t Bank::registerLT(const LazyTensor& tensor, const std::string& device) {
     std::lock_guard<std::mutex> guard(m_mutex);
 
     auto& device_bank = m_device_banks[device_for_alloc];
+    std::unique_lock dev_guard(device_bank.mutex);
+
     auto iter_registered = device_bank.registered_tensors.find(tensor);
     if (iter_registered == device_bank.registered_tensors.end()) {
         auto uid = uid_count++;
@@ -79,10 +81,14 @@ void Bank::evaluate_and_allocate() {
         auto& device_bank = bank.second;
 
         std::vector<LazyTensor> vec;
+
+        std::unique_lock storage_guard(device_bank.mutex);
         vec.reserve(device_bank.storage.size());
         for (const auto& el : device_bank.storage) {
             vec.push_back(el.second.lt);
         }
+        storage_guard.unlock();
+
         ov::parallel_for(vec.size(), [&](std::size_t idx) {
             const auto& lt = vec[idx];
             std::unique_lock dev_guard(device_bank.mutex);
diff --git a/src/plugins/intel_npu/src/plugin/npuw/weights_bank.hpp b/src/plugins/intel_npu/src/plugin/npuw/weights_bank.hpp
index cc5489c6ae6628..0d1d84b490c5e2 100644
--- a/src/plugins/intel_npu/src/plugin/npuw/weights_bank.hpp
+++ b/src/plugins/intel_npu/src/plugin/npuw/weights_bank.hpp
@@ -29,7 +29,7 @@ class Bank {
     // Register LazyTensor in a bank if it's not there. Returns LazyTensor's unique id
     int64_t registerLT(const LazyTensor& tensor, const std::string& device);
 
-    // Allocate and evaluate a registered tensor on a specified device (if needed) and return it from the bank
+    // Get registered, allocated and evaluated tensor on a specified device
     ov::Tensor get(int64_t uid, const std::string& device);
 
     // Evaluate and allocate all LazyTensors in the bank