[OpenMP] Wait for kernel prior to memory deallocation

[tianshilei1992]

Summary

In the function target, memory deallocation and target_data_end is called
immediately returning from launching kernel. This might cause a race condition
that the corresponding memory is still being used by the kernel and a potential
issue that when the kernel starts to execute, its required data have already
been deallocated, especially when multiple kernels running concurrently. Since
nevertheless, we will block the thread issuing the target offloading at the end
of the target, we just added a synchronization before memory deallocation
to make sure the correctness.

[jdoerfert]

LGTM. Typo in the commit message. Also shorten the commit title, e.g., wait for kernel prior to memory deallocation

[tianshilei1992]

Updated based on comments

[ye-luo]

Does it mean the D2H will always run synchronously after this change?
Does it also mean that target_data_end should be split into data transfer and data free parts?

[tianshilei1992]

In D84381#2168483, @ye-luo wrote:

Does it mean the D2H will always run synchronously after this change?
Does it also mean that target_data_end should be split into data transfer and data free parts?

That is good point. There is a critical issue in this patch. D2H is still async but the synchronization is lost.

[ye-luo]

Indeed, target_data_begin should be split as well. cudaMalloc blocks the whole device. Alternating cudaMalloc and transfer only makes the whole process further slower. Better to make all the allocation and then start queuing the transfer.

[tianshilei1992]

Fixed an issue that target may return before D2H is still in progress

[ye-luo]

OK. It is less broken now.
target_data_end still does Device.deallocTgtPtr and needs a sync before it.
To fully fix this issue, target_data_end must be spitted.

[tianshilei1992]

This comment was removed by tianshilei1992.

[tianshilei1992]

Will get it fix in another patch.

openmp/libomptarget/src/omptarget.cpp

rc = Device.run_region(TargetTable->EntriesBegin[TM->Index].addr,
                            &tgt_args[0], &tgt_offsets[0], tgt_args.size(),
                            &AsyncInfo);
   }
   if (rc != OFFLOAD_SUCCESS) {
     DP ("Executing target region abort target.\n");
     return OFFLOAD_FAIL;
   }
 
+  // This synchronization makes sure that before we deallocate any memory, the
+  // kernel has already finished. Otherwise, data will be freed when they're
+  // still being used.
+  // TODO: If we can separate the memory free and data transfer, we can avoid
+  // this synchronization.
+  if (Device.RTL->synchronize) {
+    rc = Device.RTL->synchronize(device_id, &AsyncInfo);
+    if (rc != OFFLOAD_SUCCESS) {
+      DP("Failed to synchronize.\n");
+      return OFFLOAD_FAIL;
+    }
+  }
+
   // Deallocate (first-)private arrays
   for (auto it : fpArrays) {
     int rt = Device.RTL->data_delete(Device.RTLDeviceID, it);
     if (rt != OFFLOAD_SUCCESS) {
       DP("Deallocation of (first-)private arrays failed.\n");
       return OFFLOAD_FAIL;
     }
   }
 
   // Move data from device.
   int rt = target_data_end(Device, arg_num, args_base, args, arg_sizes,
                            arg_types, arg_mappers, &AsyncInfo);
   if (rt != OFFLOAD_SUCCESS) {
     DP("Call to target_data_end failed, abort targe.\n");
     return OFFLOAD_FAIL;
   }
 
+  // This synchronization makes sure that all data transfer have been done
+  // before returning from this function.
   if (Device.RTL->synchronize)
     return Device.RTL->synchronize(device_id, &AsyncInfo);
 
   return OFFLOAD_SUCCESS;
 }