Merge pull request #29595 from abhishek-gola:handle_reinitialization

Replace global finalizeLayers flag with per-layer re-initialization

modules/dnn/include/opencv2/dnn/dnn.hpp

@@ -285,6 +285,9 @@ CV__DNN_INLINE_NS_BEGIN
 
         //! List of learned parameters must be stored here to allow read them by using Net::getParam().
         CV_PROP_RW std::vector<Mat> blobs;
+
+        //! Bumped when a blob is replaced; executors then re-run Layer::prepackWeights().
+        unsigned weightEpoch = 1;
         std::vector<Arg> inputs;
         std::vector<Arg> outputs;
         void* netimpl = nullptr;
@@ -505,8 +508,19 @@ CV__DNN_INLINE_NS_BEGIN
          */
         virtual void unsetAttached();
 
+        /** @brief One-time, shape-independent weight packing; runs once per
+         *  LayerInfo::weightEpoch, unlike finalize(). Default: no-op.
+         */
+        virtual void prepackWeights();
+
         CV_PROP int preferableTarget; //!< prefer target for layer forwarding
 
+        //! Executor-side bookkeeping for per-layer (re)initialization.
+        unsigned packedWeightEpoch = 0;        //!< LayerInfo::weightEpoch prepackWeights() last ran for
+        bool finalizedOnce = false;
+        std::vector<MatShape> lastInpShapes;   //!< input shapes finalize() last ran for
+        std::vector<int> lastInpTypes;         //!< input types finalize() last ran for
+
         Layer();
         explicit Layer(const LayerParams &params);      //!< Initializes only #name, #type and #blobs fields.
         virtual ~Layer();

modules/dnn/src/kv_cache_manager.cpp

@@ -119,7 +119,7 @@ void KVCacheManager::initPastTensors()
 
         Mat& past_t = netimpl->__tensors__.at(route.second);
         past_t = Mat(shape_vec, dtype, Scalar(0));
-        netimpl->finalizeLayers = true;
+        // The consumer's signature changed; the op loop re-finalizes it.
     }
 }
 
@@ -131,8 +131,7 @@ void KVCacheManager::applyRoutes()
         Mat& past_t = netimpl->__tensors__.at(route.second);
         if (present_t.empty())
             continue;
-        if (past_t.shape() != present_t.shape() || past_t.type() != present_t.type())
+        // Grown buffer changes the consumer's signature; the op loop re-finalizes it.
-            netimpl->finalizeLayers = true;
         present_t.copyTo(past_t);
     }
 }

modules/dnn/src/layer.cpp

@@ -190,6 +190,10 @@ std::vector<Mat> Layer::finalize(const std::vector<Mat>& inputs)
     return outputs;
 }
 
+void Layer::prepackWeights()
+{
+}
+
 void Layer::forward(std::vector<Mat*>& input, std::vector<Mat>& output, std::vector<Mat>& internals)
 {
     // We kept this method for compatibility. DNN calls it now only to support users' implementations.

modules/dnn/src/layers/attention_layer.cpp

@@ -153,8 +153,6 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
         output_ndims = params.get<int>("output_ndims", 3);
 
         do_rotary = params.get<bool>("do_rotary", false);
-
-        is_prepacked = false;
     }
 
     virtual bool supportBackend(int backendId) CV_OVERRIDE {
@@ -242,6 +240,27 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
         return flops;
     }
 
+    // Geometry derives from the weight shape alone, so this is safe before finalize().
+    void packQKV(const Mat& weight) {
+        opt.init();
+        const auto weight_shape = shape(weight);
+        input_hidden_size = static_cast<size_t>(weight_shape[0]);
+        hidden_size = weight_shape[1];
+        qkv_hidden_sizes[2] = hidden_size - qkv_hidden_sizes[0] - qkv_hidden_sizes[1];
+        qkv_head_sizes[2] = static_cast<size_t>(qkv_hidden_sizes[2] / num_heads);
+
+        const auto *weight_data = weight.ptr<const float>();
+        packWeight(num_heads, qkv_head_sizes[0], input_hidden_size, weight_data,                                             hidden_size, packed_weight_q, opt);
+        packWeight(num_heads, qkv_head_sizes[1], input_hidden_size, weight_data + qkv_hidden_sizes[0],                       hidden_size, packed_weight_k, opt);
+        packWeight(num_heads, qkv_head_sizes[2], input_hidden_size, weight_data + qkv_hidden_sizes[0] + qkv_hidden_sizes[1], hidden_size, packed_weight_v, opt);
+    }
+
+    // Dynamic weights (blobs empty) aren't available yet; forward() packs those.
+    void prepackWeights() CV_OVERRIDE {
+        if (!blobs.empty())
+            packQKV(blobs.front());
+    }
+
     virtual void finalize(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr) CV_OVERRIDE {
         opt.init();
 
@@ -257,15 +276,6 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
         hidden_size = weight_shape[1];
         qkv_hidden_sizes[2] = hidden_size - qkv_hidden_sizes[0] - qkv_hidden_sizes[1];
         qkv_head_sizes[2] = static_cast<size_t>(qkv_hidden_sizes[2] / num_heads);
-
-        if (!blobs.empty()) {
-            const auto *weight_data = weight.ptr<const float>();
-            packWeight(num_heads, qkv_head_sizes[0], input_hidden_size, weight_data,                                             hidden_size, packed_weight_q, opt);
-            packWeight(num_heads, qkv_head_sizes[1], input_hidden_size, weight_data + qkv_hidden_sizes[0],                       hidden_size, packed_weight_k, opt);
-            packWeight(num_heads, qkv_head_sizes[2], input_hidden_size, weight_data + qkv_hidden_sizes[0] + qkv_hidden_sizes[1], hidden_size, packed_weight_v, opt);
-
-            is_prepacked = true;
-        }
     }
 
     void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE {
@@ -283,16 +293,12 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
         outputs_arr.getMatVector(outputs);
         internals_arr.getMatVector(internals);
 
-        // prepack weights
+        // Dynamic weight (blobs empty) may change each call, so repack it.
-        if (!is_prepacked) {
+        if (blobs.empty())
-            const auto &weight = blobs.empty() ? inputs[1] : blobs.front();
+            packQKV(inputs[1]);
-            const auto *weight_data = weight.ptr<const float>();
+        // Const weight: pack once as fallback when prepackWeights() wasn't called.
-            packWeight(num_heads, qkv_head_sizes[0], input_hidden_size, weight_data,                                             hidden_size, packed_weight_q, opt);
+        else if (packed_weight_q.empty())
-            packWeight(num_heads, qkv_head_sizes[1], input_hidden_size, weight_data + qkv_hidden_sizes[0],                       hidden_size, packed_weight_k, opt);
+            packQKV(blobs.front());
-            packWeight(num_heads, qkv_head_sizes[2], input_hidden_size, weight_data + qkv_hidden_sizes[0] + qkv_hidden_sizes[1], hidden_size, packed_weight_v, opt);
-
-            is_prepacked = true;
-        }
 
         float *packed_weights[3] = {packed_weight_q.data(), packed_weight_k.data(), packed_weight_v.data()};
         size_t packed_weights_size[3] = {packed_weight_q.size() / num_heads, packed_weight_k.size() / num_heads, packed_weight_v.size() / num_heads};
@@ -555,7 +561,6 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
     size_t hidden_size;
 
     bool do_rotary;
-    bool is_prepacked;
     std::vector<float> packed_weight_q;
     std::vector<float> packed_weight_k;
     std::vector<float> packed_weight_v;

modules/dnn/src/net_impl.cpp

@@ -153,7 +153,6 @@ void Net::Impl::clear()
     bufidxs.push_back(-1);
 
     prepared = false;
-    finalizeLayers = true;
     finalized = false;
     fusedSnapshotValid = false;
     fusedSnapshot.clear();
@@ -1653,6 +1652,12 @@ Mat Net::Impl::getParam(int layer, int numParam) const
     return layerBlobs[numParam];
 }
 
+// Bump only the epoch: the executor holding the packed weights may be another object.
+static void markLayerWeightsChanged(const Ptr<LayerInfo>& layer)
+{
+    layer->weightEpoch++;
+}
+
 void Net::Impl::setParam(int layer, int numParam, const Mat& blob)
 {
     // FIXIT we should not modify "execution" instance
@@ -1661,7 +1666,7 @@ void Net::Impl::setParam(int layer, int numParam, const Mat& blob)
     // we don't make strong checks, use this function carefully
     layerBlobs[numParam] = blob;
     if (mainGraph)
-        finalizeLayers = true;
+        markLayerWeightsChanged(getLayer(layer));
 }
 
 void Net::Impl::setParam(const std::string& outputTensorName, int numParam, const Mat& blob)
@@ -1695,21 +1700,21 @@ void Net::Impl::setParam(const std::string& outputTensorName, int numParam, cons
 
             if (numParam < (int)layer->blobs.size()) {
                 layer->blobs[numParam] = blob;
-                finalizeLayers = true;
+                markLayerWeightsChanged(layer);
                 return;
             }
 
             Conv2Layer* conv = dynamic_cast<Conv2Layer*>(layer.get());
             if (conv && numParam == 0) {
                 conv->setWeights(blob, Mat(), defaultC0, accuracy);
-                finalizeLayers = true;
+                markLayerWeightsChanged(layer);
                 return;
             }
 
             ConvTranspose2Layer* deconv = dynamic_cast<ConvTranspose2Layer*>(layer.get());
             if (deconv && numParam == 0) {
                 deconv->setWeights(blob, Mat(), defaultC0, accuracy);
-                finalizeLayers = true;
+                markLayerWeightsChanged(layer);
                 return;
             }
 

modules/dnn/src/net_impl.hpp

@@ -142,7 +142,6 @@ struct Net::Impl : public detail::NetImplBase
     int defaultC0;
     bool enableFP16, haveFP16;
     bool prepared; // need to rerun graph transformations/optimizations
-    bool finalizeLayers; // need to initialize each layer
     bool finalized = false; // executors have been selected for the current backend/target
 
     // Post-fusion (pre block-layout) snapshot so finalize() can re-run from a clean

modules/dnn/src/net_impl2.cpp

@@ -558,7 +558,6 @@ void Net::Impl::prepareForInference()
     if (this->ort_session)
     {
         prepared = true;
-        finalizeLayers = false;
         return;
     }
 #endif
@@ -577,7 +576,6 @@ void Net::Impl::prepareForInference()
         fuseBasic();
         totalLayers = updateGraphOfs(mainGraph, 0, true);
         prepared = true;
-        finalizeLayers = true;
     }
 }
 
@@ -643,6 +641,11 @@ void Net::Impl::finalizeGraph(const Ptr<Graph>& graph, bool useCUDA)
             backend = DNN_BACKEND_OPENCV;
         }
         CV_Assert(exec);
+        // Re-finalize can hand back the same object, so reset state for the new backend.
+        exec->packedWeightEpoch = 0;
+        exec->finalizedOnce = false;
+        exec->lastInpShapes.clear();
+        exec->lastInpTypes.clear();
         g->exec_[i] = exec;
         g->execBackend_[i] = backend;
         CV_LOG_INFO(NULL, cv::format("DNN/NewEngine: finalize op #%zu '%s' (%s) -> %s",
@@ -857,11 +860,6 @@ void Net::Impl::forwardMainGraph(InputArrayOfArrays inputs, OutputArrayOfArrays
 
     forwardGraph(mainGraph, inputs, outputs, true);
 
-    // reset finalizeLayer so that layers are only initialized once.
-    // [TODO] if a target or backend change or there are some other important
-    // global changes in configuration, finalizeLayers should be set to 'true' again
-    finalizeLayers = false;
-
     // Feed present.* outputs back as past_key_values.* inputs for the next step (causal-lm-with-past).
     if (useKVCache && kvCacheManager.hasRoutes)
         kvCacheManager.applyRoutes();
@@ -1281,8 +1279,7 @@ void Net::Impl::setGraphInput(Ptr<Graph>& graph, size_t idx, const Mat& m)
                                          typeToString(adata.type).c_str()));
         }
         Mat& inp_t = argTensor(inp);
-        if (inp_t.shape() != mshape || inp_t.type() != adata_type)
+        // The op loop detects signature changes per layer; no global flag needed.
-            finalizeLayers = true;
         inp_t.fit(mshape, adata_type);
 
         if (adata.type == CV_16BF && mtype == CV_16U)
@@ -1579,17 +1576,27 @@ void Net::Impl::forwardGraph(Ptr<Graph>& graph, InputArrayOfArrays inputs_,
 
         std::vector<Ptr<Graph> >* subgraphs = op->subgraphs();
         if (!subgraphs) {
+            // Blobs live on 'op', packed buffers on the executor 'layer'.
+            if (layer->packedWeightEpoch != op->weightEpoch) {
+                layer->prepackWeights();
+                layer->packedWeightEpoch = op->weightEpoch;
+                // New weights: re-finalize too, for layers that pack inside finalize().
+                layer->finalizedOnce = false;
+            }
+            // Re-finalize only when this layer's own input signature changed.
+            if (!layer->finalizedOnce || layer->lastInpShapes != inpShapes || layer->lastInpTypes != inpTypes) {
+                layer->finalize((InputArrayOfArrays)inpMats, (OutputArrayOfArrays)outMats);
+                layer->lastInpShapes = inpShapes;
+                layer->lastInpTypes = inpTypes;
+                layer->finalizedOnce = true;
+            }
 #ifdef HAVE_CUDA
             if (opBackend == DNN_BACKEND_CUDA) {
-                if (finalizeLayers)
-                    layer->finalize(inpMats, outMats);
                 forwardOpCUDA(this, gimpl, opidx, inputs, outputs, inpMats, outMats);
             } else
 #endif
             {
                 // Device-resident inputs were already synced to host in the capture loop above.
-                if (finalizeLayers)
-                    layer->finalize(inpMats, outMats);
                 layer->forward(inpMats, outMats, tempMats);
 #ifdef HAVE_CUDA
                 // CPU produced fresh host data; invalidate any stale device copy of its outputs.