Remove unnecessary calls to cudaStreamSynchronize()

[hercules2020/kcf.git] / src / fft_cufft.cpp

diff --git a/src/fft_cufft.cpp b/src/fft_cufft.cpp
index 61bf1c9550bdc829f7eedeeb54540838bb2adb79..5c0da667c6916ac78582b5dd9a4578b5cfc85a6d 100644 (file)
--- a/src/fft_cufft.cpp
+++ b/src/fft_cufft.cpp
@@ -4,6 +4,7 @@ cuFFT::cuFFT()
 {
     CudaSafeCall(cudaSetDeviceFlags(cudaDeviceMapHost));
     cudaErrorCheck(cublasCreate(&cublas));
+    cudaErrorCheck(cublasSetStream(cublas, cudaStreamPerThread));
 }
 
 cufftHandle cuFFT::create_plan_fwd(uint howmany) const
@@ -99,17 +100,20 @@ void cuFFT::inverse(ComplexMat &complex_input, MatScales &real_result)
     Fft::inverse(complex_input, real_result);
 
     uint n_channels = complex_input.n_channels;
-    cufftComplex *in = reinterpret_cast<cufftComplex *>(complex_input.get_p_data());
+    cufftComplex *in = reinterpret_cast<cufftComplex *>(complex_input.get_dev_data());
     cufftReal *out = real_result.deviceMem();
     float alpha = 1.0 / (m_width * m_height);
 
     if (n_channels == 1)
         cudaErrorCheck(cufftExecC2R(plan_i_1ch, in, out));
 #ifdef BIG_BATCH
+    else
         cudaErrorCheck(cufftExecC2R(plan_i_all_scales, in, out));
 #endif
-    // TODO: Investigate whether this scalling is needed or not
     cudaErrorCheck(cublasSscal(cublas, real_result.total(), &alpha, out, 1));
+    // The result is a cv::Mat, which will be accesses by CPU, so we
+    // must synchronize with the GPU here
+    CudaSafeCall(cudaStreamSynchronize(cudaStreamPerThread));
 }
 
 cuFFT::~cuFFT()