From 34c3f0f4953b351d3da423245a946f92c0e81af4 Mon Sep 17 00:00:00 2001
From: Dale Phurrough <dale@hidale.com>
Date: Thu, 28 Jan 2021 16:14:01 +0100
Subject: [PATCH] add cuda::Stream constructor with cuda flags

---
 modules/core/include/opencv2/core/cuda.hpp | 12 ++++++++++++
 modules/core/src/cuda_stream.cpp           | 19 +++++++++++++++++++
 modules/core/test/test_cuda.cpp            | 21 +++++++++++++++++++++
 modules/python/test/test_cuda.py           |  2 ++
 4 files changed, 54 insertions(+)
 create mode 100755 modules/core/test/test_cuda.cpp

diff --git a/modules/core/include/opencv2/core/cuda.hpp b/modules/core/include/opencv2/core/cuda.hpp
index 5fa09682e3..82558846d6 100644
--- a/modules/core/include/opencv2/core/cuda.hpp
+++ b/modules/core/include/opencv2/core/cuda.hpp
@@ -656,6 +656,18 @@ public:
     //! creates a new asynchronous stream with custom allocator
     CV_WRAP Stream(const Ptr<GpuMat::Allocator>& allocator);
 
+    /** @brief creates a new Stream using the cudaFlags argument to determine the behaviors of the stream
+
+    @note The cudaFlags parameter is passed to the underlying api cudaStreamCreateWithFlags() and
+    supports the same parameter values.
+    @code
+        // creates an OpenCV cuda::Stream that manages an asynchronous, non-blocking,
+        // non-default CUDA stream
+        cv::cuda::Stream cvStream(cudaStreamNonBlocking);
+    @endcode
+     */
+    CV_WRAP Stream(const size_t cudaFlags);
+
     /** @brief Returns true if the current stream queue is finished. Otherwise, it returns false.
     */
     CV_WRAP bool queryIfComplete() const;
diff --git a/modules/core/src/cuda_stream.cpp b/modules/core/src/cuda_stream.cpp
index 5fb873a369..3680e0720a 100644
--- a/modules/core/src/cuda_stream.cpp
+++ b/modules/core/src/cuda_stream.cpp
@@ -41,6 +41,7 @@
 //M*/
 
 #include "precomp.hpp"
+#include <cstdint>
 
 using namespace cv;
 using namespace cv::cuda;
@@ -293,6 +294,7 @@ public:
 
     Impl();
     Impl(const Ptr<GpuMat::Allocator>& allocator);
+    Impl(const unsigned int cudaFlags);
     explicit Impl(cudaStream_t stream);
 
     ~Impl();
@@ -312,6 +314,13 @@ cv::cuda::Stream::Impl::Impl(const Ptr<GpuMat::Allocator>& allocator) : stream(0
     ownStream = true;
 }
 
+cv::cuda::Stream::Impl::Impl(const unsigned int cudaFlags) : stream(0), ownStream(false)
+{
+    cudaSafeCall(cudaStreamCreateWithFlags(&stream, cudaFlags));
+    ownStream = true;
+    allocator = makePtr<StackAllocator>(stream);
+}
+
 cv::cuda::Stream::Impl::Impl(cudaStream_t stream_) : stream(stream_), ownStream(false)
 {
     allocator = makePtr<StackAllocator>(stream);
@@ -450,6 +459,16 @@ cv::cuda::Stream::Stream(const Ptr<GpuMat::Allocator>& allocator)
 #endif
 }
 
+cv::cuda::Stream::Stream(const size_t cudaFlags)
+{
+#ifndef HAVE_CUDA
+    CV_UNUSED(cudaFlags);
+    throw_no_cuda();
+#else
+    impl_ = makePtr<Impl>(cudaFlags & UINT_MAX);
+#endif
+}
+
 bool cv::cuda::Stream::queryIfComplete() const
 {
 #ifndef HAVE_CUDA
diff --git a/modules/core/test/test_cuda.cpp b/modules/core/test/test_cuda.cpp
new file mode 100755
index 0000000000..a3e0a9034b
--- /dev/null
+++ b/modules/core/test/test_cuda.cpp
@@ -0,0 +1,21 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+#if defined(HAVE_CUDA)
+
+#include "test_precomp.hpp"
+#include <cuda_runtime.h>
+#include "opencv2/core/cuda.hpp"
+
+namespace opencv_test { namespace {
+
+TEST(CUDA_Stream, construct_cudaFlags)
+{
+    cv::cuda::Stream stream(cudaStreamNonBlocking);
+    EXPECT_NE(stream.cudaPtr(), nullptr);
+}
+
+}} // namespace
+
+#endif
diff --git a/modules/python/test/test_cuda.py b/modules/python/test/test_cuda.py
index 8eba9076a2..4bcd4108f1 100644
--- a/modules/python/test/test_cuda.py
+++ b/modules/python/test/test_cuda.py
@@ -33,6 +33,8 @@ class cuda_test(NewOpenCVTests):
         self.assertTrue(cuMat.cudaPtr() != 0)
         stream = cv.cuda_Stream()
         self.assertTrue(stream.cudaPtr() != 0)
+        asyncstream = cv.cuda_Stream(1)  # cudaStreamNonBlocking
+        self.assertTrue(asyncstream.cudaPtr() != 0)
 
 if __name__ == '__main__':
     NewOpenCVTests.bootstrap()