Switch test/util/device_atomics.cu over to cudaMemcpyAsync

rapidsai · Jun 29, 2023 · 9d94386 · 9d94386
1 parent 7d928c4
commit 9d94386
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/cpp/test/util/device_atomics.cu b/cpp/test/util/device_atomics.cu
@@ -52,11 +52,11 @@ TEST(Raft, AtomicIncWarp)
   test_atomic_inc_warp_kernel<<<num_blocks, threads_per_block, 0, s>>>(counter.data(),
                                                                        out_device.data());
   // Copy data to host
-  RAFT_CUDA_TRY(cudaStreamSynchronize(s)); //ensure stream is done before copy
-  RAFT_CUDA_TRY(cudaMemcpy(out_host.data(),
-                           (const void*)out_device.data(),
-                           num_elts * sizeof(int),
-                           cudaMemcpyDeviceToHost));
+  RAFT_CUDA_TRY(cudaMemcpyAsync(out_host.data(),
+                                (const void*)out_device.data(),
+                                num_elts * sizeof(int),
+                                cudaMemcpyDeviceToHost,
+                                s));
 
   // Check that count is correct and that each thread index is contained in the
   // array exactly once.