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

Author: DariaMityagina

src/plugins/intel_npu/src/backend/src/zero_infer_request.cpp

@@ -84,7 +84,9 @@ std::optional<size_t> determine_dynamic_batch_size(const IODescriptor& desc,
     }
 
     auto batchFromModel = ioShape[intel_npu::utils::BATCH_AXIS];
-    if (!batchFromModel.is_dynamic()) {
+    auto batchModelFromIR =
+        desc.shapeFromIRModel.has_value() && desc.shapeFromIRModel.value()[intel_npu::utils::BATCH_AXIS].is_dynamic();
+    if (!batchFromModel.is_dynamic() && !batchModelFromIR) {
         return std::nullopt;
     }
 

src/plugins/intel_npu/src/compiler_adapter/include/graph.hpp

@@ -64,6 +64,7 @@ class Graph : public IGraph {
 
 protected:
     bool release_blob(const Config& config);
+    std::optional<size_t> determine_batch_size();
 
     std::shared_ptr<ZeGraphExtWrappers> _zeGraphExt;
 

src/plugins/intel_npu/src/compiler_adapter/src/graph.cpp

@@ -232,6 +232,10 @@ void Graph::initialize(const Config& config) {
     //  releasing it here to avoid unnecessary memory usage.
     _blobIsReleased = release_blob(config);
 
+    if (!_batchSize.has_value()) {
+        _batchSize = determine_batch_size();
+    }
+
     if (_zeroInitStruct->getCommandQueueDdiTable().version() < ZE_MAKE_VERSION(1, 1) &&
         config.get<RUN_INFERENCES_SEQUENTIALLY>()) {
         auto numberOfCommandLists = _batchSize.has_value() ? *_batchSize : 1;
@@ -288,6 +292,58 @@ uint32_t Graph::get_last_submitted_id() const {
     return _lastSubmittedId;
 }
 
+std::optional<size_t> Graph::determine_batch_size() {
+    if (!_metadata.outputs.at(0).shapeFromIRModel.has_value()) {
+        _logger.debug("Batching on the plugin is not used, batching is handled by the compiler");
+        return std::nullopt;
+    }
+
+    const ov::PartialShape& firstShape = *_metadata.outputs.at(0).shapeFromIRModel;
+    if (firstShape.is_dynamic() || firstShape.rank().get_length() == 0) {
+        return std::nullopt;
+    }
+
+    const size_t candidateBatchSize = firstShape[utils::BATCH_AXIS].get_max_length();
+    if (candidateBatchSize == 0 || candidateBatchSize == utils::DEFAULT_BATCH_SIZE) {
+        _logger.debug("Batching on the plugin is not used, batching is handled by the compiler");
+        return std::nullopt;
+    }
+
+    auto checkDescriptorsUseCandidateBatchSize = [candidateBatchSize](const std::vector<IODescriptor>& descriptors) {
+        for (const IODescriptor& descriptor : descriptors) {
+            OPENVINO_ASSERT(descriptor.shapeFromIRModel.has_value(),
+                            "Missing value for the \"shapeFromIRModel\" attribute, I/O descriptor");
+
+            const ov::PartialShape& shapeFromCompiler = descriptor.shapeFromCompiler;
+            const ov::PartialShape& shapeFromIRModel = *descriptor.shapeFromIRModel;
+
+            if (shapeFromCompiler.is_dynamic() || shapeFromCompiler.rank().get_length() == 0 ||
+                *shapeFromCompiler.begin() != utils::DEFAULT_BATCH_SIZE) {
+                return false;
+            }
+
+            if (!descriptor.isStateInput && !descriptor.isStateOutput && !descriptor.isShapeTensor) {
+                if (shapeFromIRModel.is_dynamic() || shapeFromIRModel.rank().get_length() == 0 ||
+                    *shapeFromIRModel.begin() != candidateBatchSize) {
+                    return false;
+                }
+            }
+        }
+
+        return true;
+    };
+
+    if (!checkDescriptorsUseCandidateBatchSize(_metadata.inputs) ||
+        !checkDescriptorsUseCandidateBatchSize(_metadata.outputs)) {
+        _logger.debug("Batching on the plugin is not used, batching is handled by the compiler");
+        return std::nullopt;
+    }
+
+    _logger.debug("Batching is handled by the plugin");
+
+    return candidateBatchSize;
+}
+
 const std::optional<std::size_t> Graph::get_batch_size() const {
     return _batchSize;
 }

src/plugins/intel_npu/src/plugin/src/plugin.cpp

@@ -708,7 +708,7 @@ std::shared_ptr<ov::ICompiledModel> Plugin::compile_model(const std::shared_ptr<
     }
 
     std::optional<int64_t> batch = std::nullopt;
-    if (originalBatch.has_value()) {
+    if (originalBatch.has_value() && successfullyDebatched) {
         batch = originalBatch.value().is_static() ? originalBatch.value().get_length() : -1;
         if (batch > 0) {
             // Initial batch setup for static cases

src/plugins/intel_npu/src/plugin/src/transformations.cpp

@@ -235,13 +235,15 @@ std::tuple<std::shared_ptr<ov::Model>, bool> handlePluginBatching(
             logger.info("The model has been debatched successfully");
             successfullyDebatched = true;
         }
+        if (batchModeIsAvailable) {
+            // If we have successfully debatched the model on the PLUGIN side, we should
+            // avoid repeating the same in the compiler by resetting the batch mode
+            updateBatchMode(ov::intel_npu::BatchMode::COMPILER);
+        }
     } catch (const std::exception& ex) {
         logger.info("Couldn't validate and reshape the model. Batching will be handled by compiler. Error: %s",
                     ex.what());
     }
-    if (batchModeIsAvailable) {
-        updateBatchMode(ov::intel_npu::BatchMode::COMPILER);
-    }
     return {reshapedModel, successfullyDebatched};
 }