fix predict issue in PR36682

Author: DanLiu2Intel

src/plugins/intel_npu/src/backend/include/zero_dynamic_pipeline.hpp

@@ -18,10 +18,8 @@ struct DynamicArguments {
     std::vector<npu_vm_runtime_mem_ref_handle_t> _inputMemRefHandles;
     std::vector<npu_vm_runtime_mem_ref_handle_t> _outputMemRefHandles;
 
-    // Share runtime_execution_context during VM execution and forecasting
-    npu_vm_runtime_execution_context_handle_t _executionContext = nullptr;
-    // Set by the caller after the first successful @c npuVMRuntimeExecute.
-    bool _executedOnce = false;
+    // Share runtime_execution_context in param during VM execution and forecasting
+    npu_vm_runtime_execute_params_t _executeParams = {};
 
     DynamicArguments() = default;
     DynamicArguments(const DynamicArguments&) = delete;

src/plugins/intel_npu/src/backend/src/zero_dynamic_pipeline.cpp

@@ -153,17 +153,18 @@ void DynamicArguments::setArgumentProperties(uint32_t argi,
 }
 
 DynamicArguments::~DynamicArguments() {
-    if (_executionContext != nullptr) {
-        npuVMRuntimeDestroyExecutionContext(_executionContext);
-        _executionContext = nullptr;
+    if (_executeParams.executionContext != nullptr) {
+        npuVMRuntimeDestroyExecutionContext(_executeParams.executionContext);
+        _executeParams.executionContext = nullptr;
     }
 }
 
 void DynamicArguments::ensureExecutionContext(npu_vm_runtime_handle_t vmRuntime) {
-    if (_executionContext != nullptr) {
+    if (_executeParams.executionContext != nullptr) {
         return;
     }
-    if (npuVMRuntimeCreateExecutionContext(vmRuntime, &_executionContext) != NPU_VM_RUNTIME_RESULT_SUCCESS) {
+    if (npuVMRuntimeCreateExecutionContext(vmRuntime, &_executeParams.executionContext) !=
+        NPU_VM_RUNTIME_RESULT_SUCCESS) {
         OPENVINO_THROW("Failed to create a VM execution context");
     }
 }
@@ -339,20 +340,24 @@ void DynamicPipeline::execute_vm_runtime(npu_vm_runtime_handle_t vmRuntime,
                                          ze_fence_handle_t fence,
                                          ze_event_handle_t event) {
     _logger.debug("Start to execute graph with runtime engine");
+
+    npu_vm_runtime_execute_params_t* params = &args._executeParams;
+    bool firstInference = params->graphDdiTableExt == nullptr;
     bool noTensorChange = true;
-    // _executedOnce is true only after a successful npuVMRuntimeExecute below
-    const bool firstExecution = !args._executedOnce;
 
     auto processMemRefs = [&](auto& memRefs, auto& targetMemRefHandles) {
         targetMemRefHandles.clear();
         targetMemRefHandles.reserve(memRefs.size());
+
         for (auto& memref : memRefs) {
-            auto impl = std::static_pointer_cast<MemRefTypeImpl>(memref._impl);
+            std::shared_ptr<MemRefTypeImpl> impl = std::static_pointer_cast<MemRefTypeImpl>(memref._impl);
             if (impl == nullptr) {
                 impl = std::make_shared<MemRefTypeImpl>();
                 memref._impl = impl;
             }
             impl->UpdateMemRefHandleStatus(memref);
+
+            // VM runtime execute path always needs current memref handles.
             targetMemRefHandles.push_back(impl->_memRef);
 
             if (impl->_ptrUpdated || impl->_shapeUpdated || impl->_strideUpdated) {
@@ -364,7 +369,7 @@ void DynamicPipeline::execute_vm_runtime(npu_vm_runtime_handle_t vmRuntime,
     processMemRefs(args._inputsMemRef, args._inputMemRefHandles);
     processMemRefs(args._outputsMemRef, args._outputMemRefHandles);
 
-    if (!firstExecution && noTensorChange) {
+    if (!firstInference && noTensorChange) {
         _logger.debug("Reuse command list without update since no tensor change detected");
         auto result = zeCommandQueueExecuteCommandLists(commandQueue,
                                                         static_cast<uint32_t>(commandLists.size()),
@@ -388,31 +393,28 @@ void DynamicPipeline::execute_vm_runtime(npu_vm_runtime_handle_t vmRuntime,
 
     // Create the VM execution context (owned by args._impl, destroyed with it).
     args.ensureExecutionContext(vmRuntime);
-
-    npu_vm_runtime_execute_params_t params{};
-    params.executionContext = args._executionContext;
-    params.pInputs = args._inputMemRefHandles.data();
-    params.numOfInputs = static_cast<uint32_t>(args._inputMemRefHandles.size());
-    params.pOutputs = args._outputMemRefHandles.data();
-    params.numOfOutputs = static_cast<uint32_t>(args._outputMemRefHandles.size());
-    params.ctx = _init_structs->getContext();
-    params.device = _init_structs->getDevice();
-    params.graphDdiTableExt = _init_structs->getGraphDdiTable().getImpl();
-    params.commandLists = commandLists.data();
-    params.numCommandLists = static_cast<uint64_t>(commandLists.size());
-    params.commandQueue = commandQueue;
-    params.inferenceFence = fence;
-    params.event = event;
+    params->executionContext = args._executeParams.executionContext;
+
+    params->pInputs = args._inputMemRefHandles.data();
+    params->numOfInputs = static_cast<uint32_t>(args._inputMemRefHandles.size());
+    params->pOutputs = args._outputMemRefHandles.data();
+    params->numOfOutputs = static_cast<uint32_t>(args._outputMemRefHandles.size());
+    params->ctx = _init_structs->getContext();
+    params->device = _init_structs->getDevice();
+    params->graphDdiTableExt = _init_structs->getGraphDdiTable().getImpl();
+    params->commandLists = commandLists.data();
+    params->numCommandLists = static_cast<uint64_t>(commandLists.size());
+    params->commandQueue = commandQueue;
+    params->inferenceFence = fence;
+    params->event = event;
 
     _logger.debug("Execute graph with runtime engine");
-    if (npuVMRuntimeExecute(vmRuntime, &params) != NPU_VM_RUNTIME_RESULT_SUCCESS) {
+    if (npuVMRuntimeExecute(vmRuntime, params) != NPU_VM_RUNTIME_RESULT_SUCCESS) {
         OPENVINO_THROW("Failed to execute VM runtime engine");
     } else {
         _logger.debug("Execution runtime engine is created successfully.");
     }
 
-    args._executedOnce = true;
-
     _logger.debug("Completed to execute graph with runtime engine");
 }
 
@@ -426,9 +428,13 @@ void DynamicPipeline::predict_output_shape(const IGraph& graph,
     const npu_vm_runtime_handle_t vmRuntime = static_cast<npu_vm_runtime_handle_t>(graph.get_handle());
     OPENVINO_ASSERT(vmRuntime != nullptr, "predict_output_shape requires a valid VM runtime engine");
 
-    auto processMemRefs = [&](auto& memRefs, auto& targetMemRefHandles) {
-        targetMemRefHandles.clear();
-        targetMemRefHandles.reserve(memRefs.size());
+    std::vector<std::shared_ptr<MemRefTypeImpl>> inputMemRefImpls, outputMemRefImpls;
+
+    auto processMemRefs = [&](auto& memRefs, auto& destMemRefHandles, auto& destMemRefImpls) {
+        destMemRefHandles.clear();
+        destMemRefHandles.reserve(memRefs.size());
+        destMemRefImpls.clear();
+        destMemRefImpls.reserve(memRefs.size());
 
         for (auto& memref : memRefs) {
             auto impl = std::static_pointer_cast<MemRefTypeImpl>(memref._impl);
@@ -437,12 +443,13 @@ void DynamicPipeline::predict_output_shape(const IGraph& graph,
                 memref._impl = impl;
             }
             impl->UpdateMemRefHandleStatus(memref);
-            targetMemRefHandles.push_back(impl->_memRef);
+            destMemRefHandles.push_back(impl->_memRef);
+            destMemRefImpls.push_back(impl);
         }
     };
 
-    processMemRefs(inputsMemRefs, args._inputMemRefHandles);
-    processMemRefs(outputsMemRefs, args._outputMemRefHandles);
+    processMemRefs(inputsMemRefs, args._inputMemRefHandles, inputMemRefImpls);
+    processMemRefs(outputsMemRefs, args._outputMemRefHandles, outputMemRefImpls);
 
     npu_vm_runtime_result_t result = NPU_VM_RUNTIME_RESULT_SUCCESS;
     npu_vm_runtime_version_t version{};
@@ -467,16 +474,18 @@ void DynamicPipeline::predict_output_shape(const IGraph& graph,
         params.numOfInputs = static_cast<uint32_t>(args._inputMemRefHandles.size());
         params.pOutputs = args._outputMemRefHandles.data();
         params.numOfOutputs = static_cast<uint32_t>(args._outputMemRefHandles.size());
-        params.executionContext = args._executionContext;
+        params.executionContext = args._executeParams.executionContext;
 
         result = npuVMRuntimePredictOutputShape2(vmRuntime, &params);
     }
 
     if (result != NPU_VM_RUNTIME_RESULT_SUCCESS) {
         OPENVINO_THROW("Failed to predict output shape with VM runtime engine, error code: ", result);
     } else {
-        for (auto& out : outputsMemRefs) {
-            std::shared_ptr<MemRefTypeImpl> outImpl = std::static_pointer_cast<MemRefTypeImpl>(out._impl);
+        for (size_t i = 0; i < outputsMemRefs.size(); ++i) {
+            auto& out = outputsMemRefs[i];
+            std::shared_ptr<MemRefTypeImpl> outImpl = outputMemRefImpls[i];
+
             if (outImpl == nullptr) {
                 OPENVINO_THROW("MemRefType implementation is broken, unknown error happens in shape prediction.");
             }

src/plugins/intel_npu/src/utils/include/intel_npu/utils/vm/npu_vm_runtime.hpp

@@ -10,6 +10,7 @@
 #endif
 
 #include <level_zero/ze_api.h>
+
 #include "ze_graph_ext.h"
 
 #if defined(__cplusplus)