Address review comments

Author: smirnov-alexey

src/plugins/intel_npu/src/plugin/npuw/base_sync_infer_request.cpp

@@ -188,15 +188,20 @@ ov::SoPtr<ov::ITensor> ov::npuw::IBaseInferRequest::get_tensor(const ov::Output<
 
 void ov::npuw::IBaseInferRequest::set_tensor(const ov::Output<const ov::Node>& port,
                                              const ov::SoPtr<ov::ITensor>& tensor) {
-    if (is_not_stored_io(port)) {
-        m_port_to_tensor[port] = TensorStorage{tensor, true, false, true};
+    if (!is_stored(port)) {
+        // TODO: might be useful to check if the tensor is allocated on the device
+        m_port_to_tensor[port] = TensorStorage{tensor, false, false, true};
     } else {
         m_port_to_tensor.at(port).tensor = tensor;
+        m_port_to_tensor.at(port).set_from_outside = true;
+    }
+
+    if (is_io(port)) {
+        m_port_to_tensor.at(port).persistent = true;
     }
 
     // Check if setting input tensor
     if (m_port_to_tensor.at(port).persistent) {
-        m_port_to_tensor.at(port).set_from_outside = true;
         handle_set_remote_input(port, tensor);
     }
 }
@@ -217,25 +222,20 @@ void ov::npuw::IBaseInferRequest::check_tensors() const {
     return;
 }
 
-bool ov::npuw::IBaseInferRequest::is_not_stored_io(const ov::Output<const ov::Node>& port) const {
-    // Due to lazy I/O allocation we need to create stored object here
-    if (m_port_to_tensor.find(port) == m_port_to_tensor.end()) {
-        // Only I/O set_tensor() is allowed for this class - check it
-        bool is_io = false;
-        for (std::size_t i = 0; i < m_npuw_model->inputs().size(); ++i) {
-            if (m_npuw_model->inputs()[i] == port) {
-                is_io = true;
-                break;
-            }
+bool ov::npuw::IBaseInferRequest::is_stored(const ov::Output<const ov::Node>& port) const {
+    return m_port_to_tensor.find(port) != m_port_to_tensor.end();
+}
+
+bool ov::npuw::IBaseInferRequest::is_io(const ov::Output<const ov::Node>& port) const {
+    for (std::size_t i = 0; i < m_npuw_model->inputs().size(); ++i) {
+        if (m_npuw_model->inputs()[i] == port) {
+            return true;
         }
-        for (std::size_t i = 0; i < m_npuw_model->outputs().size(); ++i) {
-            if (m_npuw_model->outputs()[i] == port) {
-                is_io = true;
-                break;
-            }
+    }
+    for (std::size_t i = 0; i < m_npuw_model->outputs().size(); ++i) {
+        if (m_npuw_model->outputs()[i] == port) {
+            return true;
         }
-        NPUW_ASSERT(is_io && "Only I/O tensors might be left unset at this point. Internal error!");
-        return true;
     }
     return false;
 }
@@ -551,7 +551,7 @@ void ov::npuw::IBaseInferRequest::bind_global_params(std::size_t idx, RqPtr requ
         LOG_DEBUG("Processing " << param_idx << " -> " << sub_in_idx << std::endl);
 
         const auto& g_port = m_npuw_model->inputs()[param_idx];
-        const auto& g_tnsr = is_not_stored_io(g_port) ? get_tensor(g_port) : m_port_to_tensor.at(g_port).tensor;
+        const auto& g_tnsr = is_stored(g_port) ? m_port_to_tensor.at(g_port).tensor : get_tensor(g_port);
         const auto& s_port = request->get_inputs()[sub_in_idx];
         LOG_DEBUG("Processing " << g_port << " -> " << s_port << "...");
         LOG_BLOCK();
@@ -770,7 +770,7 @@ void ov::npuw::IBaseInferRequest::bind_global_results(std::size_t idx, RqPtr req
         std::tie(result_idx, sub_out_idx) = it;
         const auto& g_port = m_npuw_model->outputs()[result_idx];
         const auto& s_port = request->get_outputs()[sub_out_idx];
-        request->set_tensor(s_port, is_not_stored_io(g_port) ? get_tensor(g_port) : m_port_to_tensor.at(g_port).tensor);
+        request->set_tensor(s_port, is_stored(g_port) ? m_port_to_tensor.at(g_port).tensor : get_tensor(g_port));
     }
 
     LOG_DEBUG("Done");

src/plugins/intel_npu/src/plugin/npuw/base_sync_infer_request.hpp

@@ -104,8 +104,10 @@ class IBaseInferRequest : public ov::ISyncInferRequest {
     mutable std::map<ov::Output<const ov::Node>, TensorStorage>
         m_port_to_tensor;  // mutable due to lazy I/O allocation in get_tensor()
 
-    // Check to verify that m_port_to_tensor doesn't have anything stored at the port and it's I/O
-    bool is_not_stored_io(const ov::Output<const ov::Node>& port) const;
+    // Check that m_port_to_tensor does have a tensor stored at the port
+    bool is_stored(const ov::Output<const ov::Node>& port) const;
+    // Check the port is I/O
+    bool is_io(const ov::Output<const ov::Node>& port) const;
 
     struct QuantGatherTensors {
         ov::Tensor w, z, s;

src/plugins/intel_npu/src/plugin/npuw/just_sync_infer_request.cpp

@@ -389,9 +389,8 @@ ov::npuw::JustInferRequest::JustInferRequest(const std::shared_ptr<ov::npuw::Com
 
 void ov::npuw::JustInferRequest::set_tensor(const ov::Output<const ov::Node>& port,
                                             const ov::SoPtr<ov::ITensor>& tensor) {
-    if (is_not_stored_io(port)) {
-        m_port_to_tensor[port] = TensorStorage{tensor, true, false, true};
-    }
+    NPUW_ASSERT(is_io(port));
+    m_port_to_tensor[port] = TensorStorage{tensor, true, false, true};
 
     // Check if setting output tensor
     for (std::size_t i = 0; i < m_npuw_model->outputs().size(); ++i) {