From 034a43e7f73b9c7e9fc461dc5311f48caf5b0fe0 Mon Sep 17 00:00:00 2001
From: YashasSamaga <yashas_2010@yahoo.com>
Date: Tue, 17 Mar 2020 16:08:04 +0530
Subject: [PATCH] release and relock on wrapper resize

---
 modules/dnn/src/cuda4dnn/csl/memory.hpp | 14 +++++++++++---
 modules/dnn/src/op_cuda.hpp             | 11 +++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/modules/dnn/src/cuda4dnn/csl/memory.hpp b/modules/dnn/src/cuda4dnn/csl/memory.hpp
index 05c8446bee..40918cd4b3 100644
--- a/modules/dnn/src/cuda4dnn/csl/memory.hpp
+++ b/modules/dnn/src/cuda4dnn/csl/memory.hpp
@@ -276,14 +276,22 @@ namespace cv { namespace dnn { namespace cuda4dnn { namespace csl {
 
         MemoryLockGuard& operator=(const MemoryLockGuard&) = delete;
         MemoryLockGuard& operator=(MemoryLockGuard&& other) noexcept {
-            ptr = other.ptr;
-            other.ptr = nullptr;
+            if (&other != this) {
+                if(ptr != nullptr) {
+                    /* cudaHostUnregister does not throw for a valid ptr */
+                    CUDA4DNN_CHECK_CUDA(cudaHostUnregister(ptr));
+                }
+                ptr = other.ptr;
+                other.ptr = nullptr;
+            }
             return *this;
         }
 
         ~MemoryLockGuard() {
-            if(ptr != nullptr)
+            if(ptr != nullptr) {
+                /* cudaHostUnregister does not throw for a valid ptr */
                 CUDA4DNN_CHECK_CUDA(cudaHostUnregister(ptr));
+            }
         }
 
     private:
diff --git a/modules/dnn/src/op_cuda.hpp b/modules/dnn/src/op_cuda.hpp
index 880ed71634..27f071a382 100644
--- a/modules/dnn/src/op_cuda.hpp
+++ b/modules/dnn/src/op_cuda.hpp
@@ -308,7 +308,18 @@ namespace cv { namespace dnn {
 
             auto numel = total(shape_);
             if (numel > shared_block->device.size())
+            {
+                /* if the host memory was already page-locked, release it and register again with the new size */
+                shared_block->memGuard = cuda4dnn::csl::MemoryLockGuard();
+                try {
+                    CV_Assert(shared_block->host.type() == CV_32F);
+                    shared_block->memGuard = cuda4dnn::csl::MemoryLockGuard(shared_block->host.data, numel * sizeof(float));
+                } catch (...) {
+                    /* a common reason for failure is that the host system (for example, a Jetson device) does not support it */
+                    /* we ignore the failure as this is just an optimization and not a requirement */
+                }
                 shared_block->device.reset(numel);
+            }
         }
 
         static Ptr<BackendWrapper> create(Mat& m) {
-- 
GitLab