reduce cudaDeviceSynchronize calls

Author: rongou

thrust/system/cuda/detail/par.h

@@ -69,31 +69,6 @@ struct execute_on_stream_base : execution_policy<Derived>
   {
     return exec.stream;
   }
-
-  friend __host__ __device__
-  cudaError_t
-  synchronize_stream(execute_on_stream_base &exec)
-  {
-    cudaError_t result;
-    if (THRUST_IS_HOST_CODE) {
-      #if THRUST_INCLUDE_HOST_CODE
-        cudaStreamSynchronize(exec.stream);
-        result = cudaGetLastError();
-      #endif
-    } else {
-      #if THRUST_INCLUDE_DEVICE_CODE
-        #if __THRUST_HAS_CUDART__
-          THRUST_UNUSED_VAR(exec);
-          cudaDeviceSynchronize();
-          result = cudaGetLastError();
-        #else
-          THRUST_UNUSED_VAR(exec);
-          result = cudaSuccess;
-        #endif
-      #endif
-    }
-    return result;
-  }
 };
 
 struct execute_on_stream : execute_on_stream_base<execute_on_stream>

thrust/system/cuda/detail/util.h

@@ -72,14 +72,27 @@ __thrust_exec_check_disable__
 template <class Derived>
 __host__ __device__
 cudaError_t
-synchronize_stream(execution_policy<Derived> &)
+synchronize_stream(execution_policy<Derived> &policy)
 {
-  #if __THRUST_HAS_CUDART__
-    cudaDeviceSynchronize();
-    return cudaGetLastError();
-  #else
-    return cudaSuccess;
-  #endif
+  cudaError_t result;
+  if (THRUST_IS_HOST_CODE) {
+    #if THRUST_INCLUDE_HOST_CODE
+      cudaStreamSynchronize(stream(policy));
+      result = cudaGetLastError();
+    #endif
+  } else {
+    #if THRUST_INCLUDE_DEVICE_CODE
+      #if __THRUST_HAS_CUDART__
+        THRUST_UNUSED_VAR(policy);
+        cudaDeviceSynchronize();
+        result = cudaGetLastError();
+      #else
+        THRUST_UNUSED_VAR(policy);
+        result = cudaSuccess;
+      #endif
+    #endif
+  }
+  return result;
 }
 
 // Entry point/interface.