Investigate refactoring opportunities for batch management in Plugin and Compiler - review

DariaMityagina · DariaMityagina · commit 3d11a351edce · 2025-10-13T07:15:45.000Z
diff --git a/src/plugins/intel_npu/src/backend/include/zero_infer_request.hpp b/src/plugins/intel_npu/src/backend/include/zero_infer_request.hpp
@@ -23,6 +23,11 @@ class ZeroInferRequest final : public SyncInferRequest {
     explicit ZeroInferRequest(const std::shared_ptr<ZeroInitStructsHolder>& initStructs,
                               const std::shared_ptr<const ICompiledModel>& compiledModel,
                               const Config& config);
+    std::optional<size_t> determine_dynamic_batch_size(const IODescriptor& desc,
+                                                       const size_t index,
+                                                       const bool isInput,
+                                                       const std::shared_ptr<ov::ITensor>& tensor,
+                                                       const std::optional<size_t> batchSize);
 
     ov::SoPtr<ov::ITensor> get_tensor(const ov::Output<const ov::Node>& port) const override;
     void set_tensor(const ov::Output<const ov::Node>& port, const ov::SoPtr<ov::ITensor>& tensor) override;
diff --git a/src/plugins/intel_npu/src/backend/src/zero_infer_request.cpp b/src/plugins/intel_npu/src/backend/src/zero_infer_request.cpp
@@ -73,16 +73,21 @@ void check_level_zero_attributes_match(const IODescriptor& ioDescriptor, const A
     }
 }
 
-std::optional<size_t> determine_dynamic_batch_size(const IODescriptor& desc,
-                                                   const std::shared_ptr<ov::ITensor>& tensor,
-                                                   const std::optional<size_t> batchSize) {
+}  // namespace
+
+std::optional<size_t> ZeroInferRequest::determine_dynamic_batch_size(const IODescriptor& desc,
+                                                                     const size_t index,
+                                                                     const bool isInput,
+                                                                     const std::shared_ptr<ov::ITensor>& tensor,
+                                                                     const std::optional<size_t> batchSize) {
     if (tensor == nullptr && !batchSize.has_value()) {
         return std::nullopt;
     }
 
-    auto dynamicBatchFromIR = desc.shapeFromIRModel.has_value() && (*desc.shapeFromIRModel).size() &&
-                              (*desc.shapeFromIRModel)[intel_npu::utils::BATCH_AXIS].is_dynamic();
-    if (!dynamicBatchFromIR) {
+    auto batchFromModel = isInput ? _compiledModel->inputs()[index].get_partial_shape()[intel_npu::utils::BATCH_AXIS]
+                                  : _compiledModel->outputs()[index].get_partial_shape()[intel_npu::utils::BATCH_AXIS];
+
+    if (!batchFromModel.is_dynamic()) {
         return std::nullopt;
     }
 
@@ -97,8 +102,6 @@ std::optional<size_t> determine_dynamic_batch_size(const IODescriptor& desc,
     return tensor->get_shape()[intel_npu::utils::BATCH_AXIS];
 }
 
-}  // namespace
-
 //------------------------------------------------------------------------------
 ZeroInferRequest::ZeroInferRequest(const std::shared_ptr<ZeroInitStructsHolder>& initStructs,
                                    const std::shared_ptr<const ICompiledModel>& compiledModel,
@@ -310,8 +313,11 @@ void ZeroInferRequest::set_tensor(const ov::Output<const ov::Node>& port, const
             return;
         }
 
-        auto batchSizeCandidate =
-            determine_dynamic_batch_size(_metadata.inputs.at(foundPort.idx), tensor._ptr, std::nullopt);
+        auto batchSizeCandidate = determine_dynamic_batch_size(_metadata.inputs.at(foundPort.idx),
+                                                               foundPort.idx,
+                                                               true,
+                                                               tensor._ptr,
+                                                               std::nullopt);
 
         if (batchSizeCandidate.has_value()) {
             if (!_dynamicBatchValueChanged) {
@@ -351,8 +357,11 @@ void ZeroInferRequest::set_tensor(const ov::Output<const ov::Node>& port, const
             return;
         }
 
-        auto batchSizeCandidate =
-            determine_dynamic_batch_size(_metadata.outputs.at(foundPort.idx), tensor._ptr, std::nullopt);
+        auto batchSizeCandidate = determine_dynamic_batch_size(_metadata.outputs.at(foundPort.idx),
+                                                               foundPort.idx,
+                                                               false,
+                                                               tensor._ptr,
+                                                               std::nullopt);
 
         if (batchSizeCandidate.has_value()) {
             if (!_dynamicBatchValueChanged) {
@@ -439,7 +448,8 @@ void ZeroInferRequest::set_tensors(const ov::Output<const ov::Node>& port,
 
     _logger.debug("ZeroInferRequest::set_tensors: %zu", tensors.size());
 
-    auto batchSizeCandidate = determine_dynamic_batch_size(_metadata.inputs.at(foundPort.idx), nullptr, tensors.size());
+    auto batchSizeCandidate =
+        determine_dynamic_batch_size(_metadata.inputs.at(foundPort.idx), foundPort.idx, true, nullptr, tensors.size());
 
     // Check if batch has been changed
     if (batchSizeCandidate.has_value()) {
diff --git a/src/plugins/intel_npu/src/common/include/intel_npu/common/icompiler_adapter.hpp b/src/plugins/intel_npu/src/common/include/intel_npu/common/icompiler_adapter.hpp
@@ -42,11 +42,11 @@ class ICompilerAdapter {
      * @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(ov::Tensor mainBlob,
-                                          const Config& config,
-                                          std::optional<std::vector<ov::Tensor>> initBlobs = std::nullopt,
-                                          const std::optional<std::shared_ptr<const ov::Model>>& model = std::nullopt,
-                                          std::optional<int64_t> batchSize = std::nullopt) const = 0;
+    virtual std::shared_ptr<IGraph> parse(
+        ov::Tensor mainBlob,
+        const Config& config,
+        std::optional<std::vector<ov::Tensor>> initBlobs = std::nullopt,
+        const std::optional<std::shared_ptr<const ov::Model>>& model = std::nullopt) const = 0;
 
     virtual ov::SupportedOpsMap query(const std::shared_ptr<const ov::Model>& model, const Config& config) const = 0;
     virtual uint32_t get_version() const = 0;
diff --git a/src/plugins/intel_npu/src/common/include/intel_npu/common/igraph.hpp b/src/plugins/intel_npu/src/common/include/intel_npu/common/igraph.hpp
@@ -36,8 +36,6 @@ class IGraph : public std::enable_shared_from_this<IGraph> {
 
     virtual void set_argument_value(uint32_t argi, const void* argv) const = 0;
 
-    virtual void set_metadata(NetworkMetadata metadata) = 0;
-
     virtual void initialize(const Config& config) = 0;
 
     virtual ~IGraph() = default;
diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/driver_compiler_adapter.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/driver_compiler_adapter.hpp
@@ -22,11 +22,11 @@ class DriverCompilerAdapter final : public ICompilerAdapter {
 
     std::shared_ptr<IGraph> compileWS(const std::shared_ptr<ov::Model>& model, const Config& config) const override;
 
-    std::shared_ptr<IGraph> parse(ov::Tensor mainBlob,
-                                  const Config& config,
-                                  std::optional<std::vector<ov::Tensor>> initBlobs = std::nullopt,
-                                  const std::optional<std::shared_ptr<const ov::Model>>& model = std::nullopt,
-                                  std::optional<int64_t> batchSize = std::nullopt) const override;
+    std::shared_ptr<IGraph> parse(
+        ov::Tensor mainBlob,
+        const Config& config,
+        std::optional<std::vector<ov::Tensor>> initBlobs = std::nullopt,
+        const std::optional<std::shared_ptr<const ov::Model>>& model = std::nullopt) const override;
 
     ov::SupportedOpsMap query(const std::shared_ptr<const ov::Model>& model, const Config& config) const override;
 
diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/graph.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/graph.hpp
@@ -35,8 +35,6 @@ class Graph : public IGraph {
 
     void set_argument_value(uint32_t argi, const void* argv) const override;
 
-    void set_metadata(NetworkMetadata metadata) override;
-
     void initialize(const Config& config) override;
 
     const NetworkMetadata& get_metadata() const override;
diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/plugin_compiler_adapter.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/plugin_compiler_adapter.hpp
@@ -23,11 +23,11 @@ class PluginCompilerAdapter final : public ICompilerAdapter {
 
     std::shared_ptr<IGraph> compileWS(const std::shared_ptr<ov::Model>& model, const Config& config) const override;
 
-    std::shared_ptr<IGraph> parse(ov::Tensor mainBlob,
-                                  const Config& config,
-                                  std::optional<std::vector<ov::Tensor>> initBlobs = std::nullopt,
-                                  const std::optional<std::shared_ptr<const ov::Model>>& model = std::nullopt,
-                                  std::optional<int64_t> batchSize = std::nullopt) const override;
+    std::shared_ptr<IGraph> parse(
+        ov::Tensor mainBlob,
+        const Config& config,
+        std::optional<std::vector<ov::Tensor>> initBlobs = std::nullopt,
+        const std::optional<std::shared_ptr<const ov::Model>>& model = std::nullopt) const override;
 
     ov::SupportedOpsMap query(const std::shared_ptr<const ov::Model>& model, const Config& config) const override;
 
diff --git a/src/plugins/intel_npu/src/compiler_adapter/include/ze_graph_ext_wrappers.hpp b/src/plugins/intel_npu/src/compiler_adapter/include/ze_graph_ext_wrappers.hpp
@@ -39,8 +39,7 @@ class ZeGraphExtWrappers {
 
     GraphDescriptor getGraphDescriptor(void* data, size_t size) const;
 
-    NetworkMetadata getNetworkMeta(GraphDescriptor& graphDescriptor,
-                                   std::optional<int64_t> batchSize = std::nullopt) const;
+    NetworkMetadata getNetworkMeta(GraphDescriptor& graphDescriptor) const;
 
     void destroyGraph(GraphDescriptor& graphDescriptor);
 
@@ -62,8 +61,7 @@ class ZeGraphExtWrappers {
     void getMetadata(ze_graph_handle_t graphHandle,
                      uint32_t index,
                      std::vector<IODescriptor>& inputs,
-                     std::vector<IODescriptor>& outputs,
-                     std::optional<int64_t> batchSize) const;
+                     std::vector<IODescriptor>& outputs) const;
 
     void initializeGraphThroughCommandList(ze_graph_handle_t graphHandle, uint32_t commandQueueGroupOrdinal) const;
 
diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/driver_compiler_adapter.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/driver_compiler_adapter.cpp
@@ -231,19 +231,19 @@ std::shared_ptr<IGraph> DriverCompilerAdapter::compileWS(const std::shared_ptr<o
                                              config);
 }
 
-std::shared_ptr<IGraph> DriverCompilerAdapter::parse(ov::Tensor mainBlob,
-                                                     const Config& config,
-                                                     std::optional<std::vector<ov::Tensor>> initBlobs,
-                                                     const std::optional<std::shared_ptr<const ov::Model>>& model,
-                                                     std::optional<int64_t> batchSize) const {
+std::shared_ptr<IGraph> DriverCompilerAdapter::parse(
+    ov::Tensor mainBlob,
+    const Config& config,
+    std::optional<std::vector<ov::Tensor>> initBlobs,
+    const std::optional<std::shared_ptr<const ov::Model>>& model) const {
     OV_ITT_TASK_CHAIN(PARSE_BLOB, itt::domains::NPUPlugin, "DriverCompilerAdapter", "parse");
 
     _logger.debug("parse start");
     auto mainGraphDesc = _zeGraphExt->getGraphDescriptor(mainBlob.data(), mainBlob.get_byte_size());
     _logger.debug("parse end");
 
     OV_ITT_TASK_NEXT(PARSE_BLOB, "getNetworkMeta");
-    auto networkMeta = _zeGraphExt->getNetworkMeta(mainGraphDesc, batchSize);
+    auto networkMeta = _zeGraphExt->getNetworkMeta(mainGraphDesc);
 
     // exporting the blob when we get it from cache or ov::hint::compiled_blob property
     // shall be available
diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/graph.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/graph.cpp
@@ -42,10 +42,6 @@ Graph::Graph(const std::shared_ptr<ZeGraphExtWrappers>& zeGraphExt,
     }
 }
 
-void Graph::set_metadata(NetworkMetadata metadata) {
-    _metadata = metadata;
-}
-
 const NetworkMetadata& Graph::get_metadata() const {
     return _metadata;
 }
diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/plugin_compiler_adapter.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/plugin_compiler_adapter.cpp
@@ -255,11 +255,11 @@ std::shared_ptr<IGraph> PluginCompilerAdapter::compileWS(const std::shared_ptr<o
         _compiler);
 }
 
-std::shared_ptr<IGraph> PluginCompilerAdapter::parse(ov::Tensor mainBlob,
-                                                     const Config& config,
-                                                     std::optional<std::vector<ov::Tensor>> initBlobs,
-                                                     const std::optional<std::shared_ptr<const ov::Model>>& model,
-                                                     std::optional<int64_t> batchSize) const {
+std::shared_ptr<IGraph> PluginCompilerAdapter::parse(
+    ov::Tensor mainBlob,
+    const Config& config,
+    std::optional<std::vector<ov::Tensor>> initBlobs,
+    const std::optional<std::shared_ptr<const ov::Model>>& model) const {
     OV_ITT_TASK_CHAIN(PARSE_BLOB, itt::domains::NPUPlugin, "PluginCompilerAdapter", "parse");
 
     _logger.debug("parse start");
@@ -308,19 +308,6 @@ std::shared_ptr<IGraph> PluginCompilerAdapter::parse(ov::Tensor mainBlob,
         network.clear();
         network.shrink_to_fit();
 
-        if (batchSize.has_value()) {
-            for (auto& in : networkMeta.inputs) {
-                if (in.shapeFromIRModel.has_value()) {
-                    in.shapeFromIRModel.value()[intel_npu::utils::BATCH_AXIS] = ov::Dimension(batchSize.value());
-                }
-            }
-            for (auto& out : networkMeta.outputs) {
-                if (out.shapeFromIRModel.has_value()) {
-                    out.shapeFromIRModel.value()[intel_npu::utils::BATCH_AXIS] = ov::Dimension(batchSize.value());
-                }
-            }
-        }
-
         if (_zeGraphExt) {
             auto initGraphDesc = _zeGraphExt->getGraphDescriptor(initBlob.data(), initBlob.get_byte_size());
 
diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/ze_graph_ext_wrappers.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/ze_graph_ext_wrappers.cpp
@@ -313,8 +313,7 @@ GraphDescriptor ZeGraphExtWrappers::getGraphDescriptor(void* blobData, size_t bl
  * @returns A descriptor object containing the metadata converted in OpenVINO specific structures.
  */
 static IODescriptor getIODescriptor(const ze_graph_argument_properties_3_t& arg,
-                                    const std::optional<ze_graph_argument_metadata_t>& metadata,
-                                    std::optional<int64_t> batchSize) {
+                                    const std::optional<ze_graph_argument_metadata_t>& metadata) {
     auto logger = Logger::global().clone("getIODescriptor");
     ov::element::Type_t precision = zeroUtils::toOVElementType(arg.devicePrecision);
     ov::Shape shapeFromCompiler;
@@ -331,9 +330,7 @@ static IODescriptor getIODescriptor(const ze_graph_argument_properties_3_t& arg,
         const auto dynamicDim = std::numeric_limits<uint64_t>::max();
         shapeFromIRModel.reserve(metadata->shape_size);
         for (uint32_t id = 0; id < metadata->shape_size; id++) {
-            if (batchSize.has_value() && id == utils::BATCH_AXIS) {
-                shapeFromIRModel.push_back(ov::Dimension(batchSize.value()));
-            } else if (metadata->shape[id] != dynamicDim) {
+            if (metadata->shape[id] != dynamicDim) {
                 shapeFromIRModel.push_back(metadata->shape[id]);
             } else {
                 // lower bound is ignored, so we set it to 1 just to satisfy the Dimension constructor,
@@ -398,8 +395,7 @@ static IODescriptor getIODescriptor(const ze_graph_argument_properties_3_t& arg,
 void ZeGraphExtWrappers::getMetadata(ze_graph_handle_t graphHandle,
                                      uint32_t index,
                                      std::vector<IODescriptor>& inputs,
-                                     std::vector<IODescriptor>& outputs,
-                                     std::optional<int64_t> batchSize) const {
+                                     std::vector<IODescriptor>& outputs) const {
     if (NotSupportArgumentMetadata(_graphExtVersion)) {
         ze_graph_argument_properties_3_t arg = {};
         _logger.debug("getMetadata - perform pfnGetArgumentProperties3");
@@ -408,10 +404,10 @@ void ZeGraphExtWrappers::getMetadata(ze_graph_handle_t graphHandle,
 
         switch (arg.type) {
         case ZE_GRAPH_ARGUMENT_TYPE_INPUT: {
-            inputs.push_back(getIODescriptor(arg, std::nullopt, batchSize));
+            inputs.push_back(getIODescriptor(arg, std::nullopt));
         } break;
         case ZE_GRAPH_ARGUMENT_TYPE_OUTPUT: {
-            outputs.push_back(getIODescriptor(arg, std::nullopt, batchSize));
+            outputs.push_back(getIODescriptor(arg, std::nullopt));
         } break;
         default: {
             OPENVINO_THROW("Invalid ze_graph_argument_type_t found in ze_graph_argument_properties_3_t object: ",
@@ -439,10 +435,10 @@ void ZeGraphExtWrappers::getMetadata(ze_graph_handle_t graphHandle,
 
         switch (arg.type) {
         case ZE_GRAPH_ARGUMENT_TYPE_INPUT: {
-            inputs.push_back(getIODescriptor(arg, optionalMetadata, batchSize));
+            inputs.push_back(getIODescriptor(arg, optionalMetadata));
         } break;
         case ZE_GRAPH_ARGUMENT_TYPE_OUTPUT: {
-            outputs.push_back(getIODescriptor(arg, optionalMetadata, batchSize));
+            outputs.push_back(getIODescriptor(arg, optionalMetadata));
         } break;
         default: {
             OPENVINO_THROW("Invalid ze_graph_argument_type_t found in ze_graph_argument_properties_3_t object: ",
@@ -452,8 +448,7 @@ void ZeGraphExtWrappers::getMetadata(ze_graph_handle_t graphHandle,
     }
 }
 
-NetworkMetadata ZeGraphExtWrappers::getNetworkMeta(GraphDescriptor& graphDescriptor,
-                                                   std::optional<int64_t> batchSize) const {
+NetworkMetadata ZeGraphExtWrappers::getNetworkMeta(GraphDescriptor& graphDescriptor) const {
     ze_graph_properties_t graphProperties = {};
     graphProperties.stype = ZE_STRUCTURE_TYPE_GRAPH_PROPERTIES;
 
@@ -462,7 +457,7 @@ NetworkMetadata ZeGraphExtWrappers::getNetworkMeta(GraphDescriptor& graphDescrip
     THROW_ON_FAIL_FOR_LEVELZERO_EXT("pfnGetProperties", result, _zeroInitStruct->getGraphDdiTable());
     NetworkMetadata meta;
     for (uint32_t index = 0; index < graphProperties.numGraphArgs; ++index) {
-        getMetadata(graphDescriptor._handle, index, meta.inputs, meta.outputs, batchSize);
+        getMetadata(graphDescriptor._handle, index, meta.inputs, meta.outputs);
     }
     // TODO: support this information in CiD [track: E#33479]
     meta.numStreams = 1;
diff --git a/src/plugins/intel_npu/src/plugin/include/compiled_model.hpp b/src/plugins/intel_npu/src/plugin/include/compiled_model.hpp
@@ -33,7 +33,8 @@ class CompiledModel final : public ICompiledModel {
                   const std::shared_ptr<const ov::IPlugin>& plugin,
                   const std::shared_ptr<IDevice>& device,
                   const std::shared_ptr<IGraph>& graph,
-                  const FilteredConfig& config);
+                  const FilteredConfig& config,
+                  std::optional<int64_t> batchSize);
 
     CompiledModel(const CompiledModel&) = delete;
 
@@ -57,6 +58,8 @@ class CompiledModel final : public ICompiledModel {
 
     const FilteredConfig& get_config() const override;
 
+    const std::optional<int64_t>& get_batch_size() const;
+
 private:
     void configure_stream_executors();
 
@@ -68,6 +71,8 @@ class CompiledModel final : public ICompiledModel {
     std::unique_ptr<Properties> _properties;
 
     std::shared_ptr<IGraph> _graph;
+
+    std::optional<int64_t> _batchSize;
 };
 
 }  //  namespace intel_npu
diff --git a/src/plugins/intel_npu/src/plugin/include/plugin.hpp b/src/plugins/intel_npu/src/plugin/include/plugin.hpp
@@ -63,10 +63,10 @@ class Plugin : public ov::IPlugin {
     ov::SupportedOpsMap query_model(const std::shared_ptr<const ov::Model>& model,
                                     const ov::AnyMap& properties) const override;
 
-    void handleDynamicBatching(std::shared_ptr<ov::Model>& modelForCompilation,
-                               Config& localConfig,
-                               const std::function<void(ov::intel_npu::BatchMode)>& updateBatchMode,
-                               std::optional<ov::Dimension>& originalBatch) const;
+    void handlePluginBatching(std::shared_ptr<ov::Model>& modelForCompilation,
+                              Config& localConfig,
+                              const std::function<void(ov::intel_npu::BatchMode)>& updateBatchMode,
+                              std::optional<ov::Dimension>& originalBatch) const;
 
 private:
     void init_options();
diff --git a/src/plugins/intel_npu/src/plugin/src/compiled_model.cpp b/src/plugins/intel_npu/src/plugin/src/compiled_model.cpp
diff --git a/src/plugins/intel_npu/src/plugin/src/metadata.cpp b/src/plugins/intel_npu/src/plugin/src/metadata.cpp
diff --git a/src/plugins/intel_npu/src/plugin/src/plugin.cpp b/src/plugins/intel_npu/src/plugin/src/plugin.cpp
diff --git a/src/plugins/intel_npu/src/utils/include/intel_npu/utils/utils.hpp b/src/plugins/intel_npu/src/utils/include/intel_npu/utils/utils.hpp

Original file line number	Diff line number	Diff line change
`@@ -42,10 +42,6 @@ Graph::Graph(const std::shared_ptr<ZeGraphExtWrappers>& zeGraphExt,`
`42`	`42`	`}`
`43`	`43`	`}`
`44`	`44`
`45`		`-void Graph::set_metadata(NetworkMetadata metadata) {`
`46`		`- _metadata = metadata;`
`47`		`-}`
`48`		`-`
`49`	`45`	`const NetworkMetadata& Graph::get_metadata() const {`
`50`	`46`	`return _metadata;`
`51`	`47`	`}`