[OpenMP] Wait for kernel prior to memory deallocation

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 move the synchronization ahead a little bit to make sure
the correctness.

Reviewers: jdoerfert

Reviewed By: jdoerfert

Subscribers: yaxunl, guansong, sstefan1, openmp-commits

Tags: #openmp

Differential Revision: https://reviews.llvm.org/D84381

diff --git a/openmp/libomptarget/src/omptarget.cpp b/openmp/libomptarget/src/omptarget.cpp
index 47971b9..a613f2e 100644 (file)
--- a/openmp/libomptarget/src/omptarget.cpp
+++ b/openmp/libomptarget/src/omptarget.cpp
@@ -927,6 +927,14 @@ int target(int64_t device_id, void *host_ptr, int32_t arg_num,
     return OFFLOAD_FAIL;
   }
 
+  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);
@@ -944,8 +952,5 @@ int target(int64_t device_id, void *host_ptr, int32_t arg_num,
     return OFFLOAD_FAIL;
   }
 
-  if (Device.RTL->synchronize)
-    return Device.RTL->synchronize(device_id, &AsyncInfo);
-
   return OFFLOAD_SUCCESS;
 }