Allow to specify cuda stream

psrdada_cpp/effelsberg/edd/SKRfiReplacementCuda.cuh

@@ -30,7 +30,7 @@ public:
      */
     void replace_rfi_data(const thrust::device_vector<int> &rfi_status,
                           thrust::device_vector<thrust::complex<float>> &data,
-                          std::size_t clean_windows = 5);
+                          std::size_t clean_windows = 5, cudaStream_t stream=0);
 
 private:
 

psrdada_cpp/effelsberg/edd/SpectralKurtosisCuda.cuh

@@ -38,7 +38,7 @@ public:
      * @param(out) stats         RFI statistics
      *
      */
-    void compute_sk_thrust(thrust::device_vector<thrust::complex<float>> const& data, RFIStatistics &stats);
+    void compute_sk_thrust(thrust::device_vector<thrust::complex<float>> const& data, RFIStatistics &stats, cudaStream_t _stream = 0);
 
     /**
      * @brief      computes spectral kurtosis (using optimized kernel function) for the given data and returns its rfi statistics.
@@ -47,7 +47,7 @@ public:
      * @param(out) stats         RFI statistics
      *
      */
-    void compute_sk(thrust::device_vector<thrust::complex<float>> &data, RFIStatistics &stats);
+    void compute_sk(thrust::device_vector<thrust::complex<float>> &data, RFIStatistics &stats, cudaStream_t _stream = 0);
 
 private:
     /**
@@ -61,7 +61,6 @@ private:
     std::size_t _sample_size; //size of input data
     float _sk_min, _sk_max;
     thrust::device_vector<float> _d_s1, _d_s2;
-    cudaStream_t _stream;
 };
 } //edd
 } //effelsberg

psrdada_cpp/effelsberg/edd/src/SKRfiReplacementCuda.cu

@@ -53,8 +53,9 @@ SKRfiReplacementCuda::~SKRfiReplacementCuda() {
 
 void SKRfiReplacementCuda::replace_rfi_data(const thrust::device_vector<int> &rfi_status,
                                             thrust::device_vector<thrust::complex<float>> &data,
-                                            std::size_t clean_windows) {
+                                            std::size_t clean_windows, cudaStream_t stream) {
     nvtxRangePushA("replace_rfi_data");
+    thrust::cuda::par.on(stream);
     thrust::device_vector<thrust::complex<float>> replacement_data;
     //initialize data members of the class
 

psrdada_cpp/effelsberg/edd/src/SpectralKurtosisCuda.cu

@@ -19,6 +19,7 @@ struct compute_power{
     }
 };
 
+
 struct power_square{
     __host__ __device__
     float operator()(thrust::complex<float> z)
@@ -29,6 +30,7 @@ struct power_square{
     }
 };
 
+
 struct check_rfi{
     const std::size_t M; //_window_size
     const float sk_min;
@@ -46,8 +48,10 @@ struct check_rfi{
    }
 };
 
+
 __device__ int rfi_count = 0;
 
+
 __device__ void warpReduce(volatile float *shmem_ptr1, volatile float *shmem_ptr2, int tid){
     shmem_ptr1[tid] += shmem_ptr1[tid + 32];
     shmem_ptr2[tid] += shmem_ptr2[tid + 32];
@@ -63,6 +67,7 @@ __device__ void warpReduce(volatile float *shmem_ptr1, volatile float *shmem_ptr
     shmem_ptr2[tid] += shmem_ptr2[tid + 1];
 }
 
+
 __global__ void compute_sk_kernel(const thrust::complex<float>* __restrict__ data, std::size_t sample_size, std::size_t window_size,
                                   float sk_max, float sk_min, int* __restrict__ rfi_status)
 {
@@ -97,6 +102,7 @@ __global__ void compute_sk_kernel(const thrust::complex<float>* __restrict__ dat
     }
 }
 
+
 SpectralKurtosisCuda::SpectralKurtosisCuda(std::size_t nchannels, std::size_t window_size, float sk_min, float sk_max)
     : _nchannels(nchannels),
       _window_size(window_size),
@@ -104,15 +110,15 @@ SpectralKurtosisCuda::SpectralKurtosisCuda(std::size_t nchannels, std::size_t wi
       _sk_max(sk_max)
 {
     BOOST_LOG_TRIVIAL(debug) << "Creating new SpectralKurtosisCuda instance... \n";
-    cudaStreamCreate(&_stream);
 }
 
+
 SpectralKurtosisCuda::~SpectralKurtosisCuda()
 {
     BOOST_LOG_TRIVIAL(debug) << "Destroying SpectralKurtosisCuda instance... \n";
-    cudaStreamDestroy(_stream);
 }
 
+
 void SpectralKurtosisCuda::init()
 {
     if((_sample_size % _window_size) != 0){
@@ -123,8 +129,10 @@ void SpectralKurtosisCuda::init()
     _nwindows = _sample_size /_window_size;
 }
 
-void SpectralKurtosisCuda::compute_sk_thrust(const thrust::device_vector<thrust::complex<float>> &data, RFIStatistics &stats){
+
+void SpectralKurtosisCuda::compute_sk_thrust(const thrust::device_vector<thrust::complex<float>> &data, RFIStatistics &stats, cudaStream_t _stream){
     nvtxRangePushA("compute_sk");
+    thrust::cuda::par.on(_stream);
     _sample_size = data.size();
     BOOST_LOG_TRIVIAL(debug) << "Computing SK (thrust version) for sample_size " << _sample_size
                              << " and window_size " << _window_size <<".\n";
@@ -158,7 +166,7 @@ void SpectralKurtosisCuda::compute_sk_thrust(const thrust::device_vector<thrust:
     nvtxRangePop();
 }
 
-void SpectralKurtosisCuda::compute_sk(thrust::device_vector<thrust::complex<float>> &data, RFIStatistics &stats){
+void SpectralKurtosisCuda::compute_sk(thrust::device_vector<thrust::complex<float>> &data, RFIStatistics &stats, cudaStream_t _stream){
     nvtxRangePushA("compute_sk_kernel");
     _sample_size = data.size();
     BOOST_LOG_TRIVIAL(debug) << "Computing SK (kernel version) for sample_size " << _sample_size