Fix synchronization issues in RNG checkpointing utils (#5273)

This PR fixes two synchronization bugs. Both problems happen when you're restoring from checkpoint in random GPU operators and are due to lack of synchronization after H2D & D2D copies.

---------

Signed-off-by: Szymon Karpiński <skarpinski@nvidia.com>

dali/operators/random/rng_checkpointing_utils.h

@@ -63,6 +63,8 @@ class RngCheckpointUtils<GPUBackend, curand_states> {
   static void SaveState(OpCheckpoint &cpt, AccessOrder order, const curand_states &rng) {
     cpt.SetOrder(order);
     cpt.MutableCheckpointState() = rng.copy(order);
+    // The pipeline will perform host synchronization before serializing the checkpoints.
+    // TODO(skarpinski) Move synchronization out from pipeline's GetCheckpoint.
   }
 
   static void RestoreState(const OpCheckpoint &cpt, curand_states &rng) {
@@ -85,10 +87,11 @@ class RngCheckpointUtils<GPUBackend, curand_states> {
   static void DeserializeCheckpoint(OpCheckpoint &cpt, const std::string &data) {
     auto deserialized = SnapshotSerializer().Deserialize<std::vector<curandState>>(data);
     curand_states states(deserialized.size());
-    CUDA_CALL(cudaMemcpy(states.states(),
-                         deserialized.data(),
-                         sizeof(curandState) * deserialized.size(),
-                         cudaMemcpyHostToDevice));
+    CUDA_CALL(cudaMemcpyAsync(states.states(),
+                              deserialized.data(),
+                              sizeof(curandState) * deserialized.size(),
+                              cudaMemcpyHostToDevice, cudaStreamDefault));
+    CUDA_CALL(cudaStreamSynchronize(cudaStreamDefault));
     cpt.MutableCheckpointState() = states;
   }
 };

dali/operators/util/randomizer.cuh

@@ -59,8 +59,9 @@ struct curand_states {
   }
 
   DALI_HOST inline void set(const curand_states &other) {
-    CUDA_CALL(cudaMemcpy(states_, other.states_, sizeof(curandState) * len_,
-                         cudaMemcpyDeviceToDevice));
+    CUDA_CALL(cudaMemcpyAsync(states_, other.states_, sizeof(curandState) * len_,
+                              cudaMemcpyDeviceToDevice, cudaStreamDefault));
+    CUDA_CALL(cudaStreamSynchronize(cudaStreamDefault));
   }
 
  private: