[NPUW] Fix for weights bank UIDs (openvinotoolkit#28297)

Previous PR openvinotoolkit#28282
sgolebiewski-intel · Jan 7, 2025 · c0650a0 · c0650a0
1 parent fd29e26
commit c0650a0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 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
@@ -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
@@ -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