From 8d3850ac028337979d4f5682430228ee2cb16080 Mon Sep 17 00:00:00 2001
From: Vladislav Vinogradov <vlad.vinogradov@itseez.com>
Date: Thu, 12 Nov 2015 13:07:30 +0300
Subject: [PATCH] add cv::gpu::StreamAccessor::wrapStream method

it allows to import existed CUDA stream to OpenCV
---
 .../include/opencv2/gpu/stream_accessor.hpp   |  1 +
 modules/gpu/src/cudastream.cpp                | 18 +++++++++++++++-
 modules/gpu/test/test_stream.cpp              | 21 +++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/modules/gpu/include/opencv2/gpu/stream_accessor.hpp b/modules/gpu/include/opencv2/gpu/stream_accessor.hpp
index 9f4c6ee1e7..bcd58ba3a1 100644
--- a/modules/gpu/include/opencv2/gpu/stream_accessor.hpp
+++ b/modules/gpu/include/opencv2/gpu/stream_accessor.hpp
@@ -57,6 +57,7 @@ namespace cv
         struct StreamAccessor
         {
             CV_EXPORTS static cudaStream_t getStream(const Stream& stream);
+            CV_EXPORTS static Stream wrapStream(cudaStream_t stream);
         };
     }
 }
diff --git a/modules/gpu/src/cudastream.cpp b/modules/gpu/src/cudastream.cpp
index efc5956ad6..66ee7f9755 100644
--- a/modules/gpu/src/cudastream.cpp
+++ b/modules/gpu/src/cudastream.cpp
@@ -89,6 +89,7 @@ struct Stream::Impl
     }
 
     cudaStream_t stream;
+    bool own_stream;
     int ref_counter;
 };
 
@@ -335,6 +336,7 @@ void cv::gpu::Stream::create()
     impl = (Stream::Impl*) fastMalloc(sizeof(Stream::Impl));
 
     impl->stream = stream;
+    impl->own_stream = true;
     impl->ref_counter = 1;
 }
 
@@ -342,9 +344,23 @@ void cv::gpu::Stream::release()
 {
     if (impl && CV_XADD(&impl->ref_counter, -1) == 1)
     {
-        cudaSafeCall( cudaStreamDestroy(impl->stream) );
+        if (impl->own_stream)
+        {
+            cudaSafeCall( cudaStreamDestroy(impl->stream) );
+        }
         cv::fastFree(impl);
     }
 }
 
+Stream StreamAccessor::wrapStream(cudaStream_t stream)
+{
+    Stream::Impl* impl = (Stream::Impl*) fastMalloc(sizeof(Stream::Impl));
+
+    impl->stream = stream;
+    impl->own_stream = false;
+    impl->ref_counter = 1;
+
+    return Stream(impl);
+}
+
 #endif /* !defined (HAVE_CUDA) */
diff --git a/modules/gpu/test/test_stream.cpp b/modules/gpu/test/test_stream.cpp
index a34a4e8409..364139dc68 100644
--- a/modules/gpu/test/test_stream.cpp
+++ b/modules/gpu/test/test_stream.cpp
@@ -44,6 +44,8 @@
 
 #ifdef HAVE_CUDA
 
+#include "opencv2/gpu/stream_accessor.hpp"
+
 using namespace cvtest;
 
 #if CUDART_VERSION >= 5000
@@ -125,6 +127,25 @@ GPU_TEST_P(Async, Convert)
     stream.waitForCompletion();
 }
 
+GPU_TEST_P(Async, WrapStream)
+{
+    cudaStream_t cuda_stream = NULL;
+    ASSERT_EQ(cudaSuccess, cudaStreamCreate(&cuda_stream));
+
+    cv::gpu::Stream stream = cv::gpu::StreamAccessor::wrapStream(cuda_stream);
+
+    stream.enqueueUpload(src, d_src);
+    stream.enqueueConvert(d_src, d_dst, CV_32S);
+    stream.enqueueDownload(d_dst, dst);
+
+    Async* test = this;
+    stream.enqueueHostCallback(checkConvert, test);
+
+    stream.waitForCompletion();
+
+    ASSERT_EQ(cudaSuccess, cudaStreamDestroy(cuda_stream));
+}
+
 INSTANTIATE_TEST_CASE_P(GPU_Stream, Async, ALL_DEVICES);
 
 #endif