Merge pull request #19286 from diablodale:add-cuda-stream-constructor

alalek · alalek · commit 6ca46afa637d · 2021-02-03T09:55:02.000Z
diff --git a/modules/core/include/opencv2/core/cuda.hpp b/modules/core/include/opencv2/core/cuda.hpp
@@ -656,6 +656,18 @@ class CV_EXPORTS_W Stream
     //! 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
@@ -41,6 +41,7 @@
 //M*/
 
 #include "precomp.hpp"
+#include <cstdint>
 
 using namespace cv;
 using namespace cv::cuda;
@@ -293,6 +294,7 @@ class cv::cuda::Stream::Impl
 
     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
@@ -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 https://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
@@ -33,6 +33,8 @@ def test_cuda_interop(self):
         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()
