[OpenMP] Introduce asynchronous malloc and free

[jdoerfert]

Summary

If the plugin supports asynchronous malloc and free we can avoid early
synchronization and hide more runtime work as part of the ongoing kernel
execution.

[jdoerfert]

NOTE: Missing the cmake magic to determine CUDA_VERSION.

[jdoerfert]

Ensure synchronization for CUDA_VERSION < ...

[grokos]

Diff shows rtl.h and device.h have been deleted?

[jdoerfert]

Add files back, use the CUDART_VERSION macro to avoid need for cmake

[jdoerfert]

Go back to CUDA_VERSION

[grokos]

openmp/libomptarget/plugins/common/MemoryManager/MemoryManager.h

146–152

We aren't we passing an AsyncInfo object as an extra argument to deleteOnDevice and subsequently call DeviceAllocator.free() with that AsyncInfo instead of nullptr? I suspect it's because we need to make room for subsequent allocations before attempting to re-allocate, if that's the case can you add a comment to make it clearer?

296

If there is a reason why deleteOnDevice() cannot be called with an extra AsyncInfo argument, we are we adding it here? It's not used anywhere in the body of free().

[jdoerfert]

Address comments

[jdoerfert]

Took me a while to get back to this. Way behind on certain things. Apologies.

openmp/libomptarget/plugins/common/MemoryManager/MemoryManager.h

146–152

I'll add a comment to explain that this operation should not be delayed.

296

I added it to keep the interface the same but I'll remove it now

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

1517

why would we need a sync before calling data alloc?

[jdoerfert]

Address comment

[jdoerfert]

put sync for delete back

[jdoerfert]

Ensure atomicity (via events) of async allocations

[jdoerfert]

Fix issues, sync w/o pending work

[ye-luo]

openmp/libomptarget/include/device.h

305

What about async malloc?

openmp/libomptarget/plugins/cuda/src/rtl.cpp

321

pass StreamPoolTy& as an argument and there is no need of DeviceId

371

This class is one object per device. I would expect removing DeviceId replace vectors with a single element.

openmp/libomptarget/src/device.cpp

290

The current Event is intended for D2H, you expanded its usage. In case of one record after alloc and another record after the transfer. I think this will go wrong. Once the event is fulfilled after alloc. other thread may think the transfer has been completed.

[jdoerfert]

openmp/libomptarget/include/device.h

305

I'll rename.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

321

sure.

371

That is somewhat orthogonal, a cleanup for all this is following.

openmp/libomptarget/src/device.cpp

290

Right.

So checking if we also do D2H would partially help but not completely. The stickiness of the event is still a problem. It might also be for the pure D2H case with multiple transfers as only the first is properly guarded.
I think we need to:

1. only enter it here if we don't do D2H in the same directive (thus below).
2. remove events after we synchronized the AsyncInfo object it was recorded in.

[ye-luo]

openmp/libomptarget/src/device.cpp

290

Sorry I meant H2D not D2H. But seems you got my point. Both 1 and 2 are desired. I thought about 2 when the event was introduced but I didn't immediately figure out what is the best location to do 2.

[tianshilei1992]

cuDeviceGetAttribute has to be used to check if the device supports memory pool; otherwise runtime error will be emitted.

openmp/libomptarget/include/device.h

struct DeviceTy {
 
   HostDataToTargetListTy HostDataToTargetMap;
   PendingCtorsDtorsPerLibrary PendingCtorsDtors;
 
   ShadowPtrListTy ShadowPtrMap;
 
   std::mutex DataMapMtx, PendingGlobalsMtx, ShadowMtx;
 
+  /// Flag to indicate if the device supports asynchronous frees.
+  bool SupportsAsyncFree;
+
   // NOTE: Once libomp gains full target-task support, this state should be
   // moved into the target task in libomp.
   std::map<int32_t, uint64_t> LoopTripCnt;
 
   DeviceTy(RTLInfoTy *RTL);
   // DeviceTy is not copyable
   DeviceTy(const DeviceTy &D) = delete;
   DeviceTy &operator=(const DeviceTy &D) = delete;
TargetPointerResultTy getTgtPtrBegin(void *HstPtrBegin, int64_t Size,
                                        bool ForceDelete = false);
   /// For the map entry for \p HstPtrBegin, decrement the reference count
   /// specified by \p HasHoldModifier and, if the the total reference count is
   /// then zero, deallocate the corresponding device storage and remove the map
   /// entry.  Return \c OFFLOAD_SUCCESS if the map entry existed, and return
   /// \c OFFLOAD_FAIL if not.  It is the caller's responsibility to skip calling
   /// this function if the map entry is not expected to exist because
   /// \p HstPtrBegin uses shared memory.
-  int deallocTgtPtr(void *HstPtrBegin, int64_t Size, bool HasHoldModifier);
+  int deallocTgtPtr(void *HstPtrBegin, int64_t Size, bool HasHoldModifier,
+                    AsyncInfoTy &AsyncInfo);
   int associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size);
   int disassociatePtr(void *HstPtrBegin);
 
   // calls to RTL
   int32_t initOnce();
   __tgt_target_table *load_binary(void *Img);
 
   // device memory allocation/deallocation routines
   /// Allocates \p Size bytes on the device, host or shared memory space
   /// (depending on \p Kind) and returns the address/nullptr when
   /// succeeds/fails. \p HstPtr is an address of the host data which the
   /// allocated target data will be associated with. If it is unknown, the
   /// default value of \p HstPtr is nullptr. Note: this function doesn't do
   /// pointer association. Actually, all the __tgt_rtl_data_alloc
   /// implementations ignore \p HstPtr. \p Kind dictates what allocator should
   /// be used (host, shared, device).
-  void *allocData(int64_t Size, void *HstPtr = nullptr,
-                  int32_t Kind = TARGET_ALLOC_DEFAULT);
+  void *allocData(int64_t Size, void *HstPtr, int32_t Kind,
+                  __tgt_async_info *AsyncInfo = nullptr);
+
   /// Deallocates memory which \p TgtPtrBegin points at and returns
   /// OFFLOAD_SUCCESS/OFFLOAD_FAIL when succeeds/fails.
-  int32_t deleteData(void *TgtPtrBegin);
+  int32_t deleteData(void *TgtPtrBegin, __tgt_async_info *AsyncInfoPtr);
 
   // Data transfer. When AsyncInfo is nullptr, the transfer will be
   // synchronous.
   // Copy data from host to device
   int32_t submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
                      AsyncInfoTy &AsyncInfo);
   // Copy data from device back to host
   int32_t retrieveData(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size,

openmp/libomptarget/include/omptarget.h

public:
   /// Synchronize all pending actions.
   ///
   /// \returns OFFLOAD_FAIL or OFFLOAD_SUCCESS appropriately.
   int synchronize();
 
   /// Return a void* reference with a lifetime that is at least as long as this
   /// AsyncInfoTy object. The location can be used as intermediate buffer.
   void *&getVoidPtrLocation();
+
+  /// Return true if the associated __tgt_async_info object has a queue, thus is
+  /// initialized and has "pending" asynchronous work.
+  bool hasPotentiallyPendingAsyncWork() const { return AsyncInfo.Queue; }
 };
 
 /// This struct is a record of non-contiguous information
 struct __tgt_target_non_contig {
   uint64_t Offset;
   uint64_t Count;
   uint64_t Stride;
 };

openmp/libomptarget/include/omptargetplugin.h

 // allocated data on the target that will be used by libomptarget.so to
 // initialize the target data mapping structures. These addresses are
 // used to generate a table of target variables to pass to
 // __tgt_rtl_run_region(). The __tgt_rtl_data_alloc() returns NULL in
 // case an error occurred on the target device. Kind dictates what allocator
 // to use (e.g. shared, host, device).
 void *__tgt_rtl_data_alloc(int32_t ID, int64_t Size, void *HostPtr,
                            int32_t Kind);
+void *__tgt_rtl_data_alloc_async(int32_t ID, int64_t Size, void *HostPtr,
+                                 int32_t Kind, __tgt_async_info *AsyncInfo);
 
 // Pass the data content to the target device using the target address. In case
 // of success, return zero. Otherwise, return an error code.
 int32_t __tgt_rtl_data_submit(int32_t ID, void *TargetPtr, void *HostPtr,
                               int64_t Size);
 
 int32_t __tgt_rtl_data_submit_async(int32_t ID, void *TargetPtr, void *HostPtr,
                                     int64_t Size, __tgt_async_info *AsyncInfo);

 // Asynchronous version of __tgt_rtl_data_exchange
 int32_t __tgt_rtl_data_exchange_async(int32_t SrcID, void *SrcPtr,
                                       int32_t DesID, void *DstPtr, int64_t Size,
                                       __tgt_async_info *AsyncInfo);
 
 // De-allocate the data referenced by target ptr on the device. In case of
 // success, return zero. Otherwise, return an error code.
 int32_t __tgt_rtl_data_delete(int32_t ID, void *TargetPtr);
+int32_t __tgt_rtl_data_delete_async(int32_t ID, void *TargetPtr,
+                                    __tgt_async_info *AsyncInfo);
 
 // Transfer control to the offloaded entry Entry on the target device.
 // Args and Offsets are arrays of NumArgs size of target addresses and
 // offsets. An offset should be added to the target address before passing it
 // to the outlined function on device side. If AsyncInfo is nullptr, it is
 // synchronous; otherwise it is asynchronous. However, AsyncInfo may be
 // ignored on some platforms, like x86_64. In that case, it is synchronous. In
 // case of success, return zero. Otherwise, return an error code.

openmp/libomptarget/include/rtl.h

 //===------------ rtl.h - Target independent OpenMP target RTL ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Declarations for handling RTL plugins.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_RTL_H
 #define _OMPTARGET_RTL_H
 
 #include "omptarget.h"
+#include <cstdint>
 #include <list>
 #include <map>
 #include <mutex>
 #include <string>
 #include <vector>
 
 // Forward declarations.
 struct DeviceTy;
 struct __tgt_bin_desc;
 
 struct RTLInfoTy {
   typedef int32_t(is_valid_binary_ty)(void *);
   typedef int32_t(is_data_exchangable_ty)(int32_t, int32_t);
   typedef int32_t(number_of_devices_ty)();
   typedef int32_t(init_device_ty)(int32_t);
   typedef __tgt_target_table *(load_binary_ty)(int32_t, void *);
   typedef void *(data_alloc_ty)(int32_t, int64_t, void *, int32_t);
+  typedef void *(data_alloc_async_ty)(int32_t, int64_t, void *, int32_t,
+                                      __tgt_async_info *);
   typedef int32_t(data_submit_ty)(int32_t, void *, void *, int64_t);
   typedef int32_t(data_submit_async_ty)(int32_t, void *, void *, int64_t,
                                         __tgt_async_info *);
   typedef int32_t(data_retrieve_ty)(int32_t, void *, void *, int64_t);
   typedef int32_t(data_retrieve_async_ty)(int32_t, void *, void *, int64_t,
                                           __tgt_async_info *);
   typedef int32_t(data_exchange_ty)(int32_t, void *, int32_t, void *, int64_t);
   typedef int32_t(data_exchange_async_ty)(int32_t, void *, int32_t, void *,
                                           int64_t, __tgt_async_info *);
   typedef int32_t(data_delete_ty)(int32_t, void *);
+  typedef int32_t(data_delete_async_ty)(int32_t, void *, __tgt_async_info *);
   typedef int32_t(run_region_ty)(int32_t, void *, void **, ptrdiff_t *,
                                  int32_t);
   typedef int32_t(run_region_async_ty)(int32_t, void *, void **, ptrdiff_t *,
                                        int32_t, __tgt_async_info *);
   typedef int32_t(run_team_region_ty)(int32_t, void *, void **, ptrdiff_t *,
                                       int32_t, int32_t, int32_t, uint64_t);
   typedef int32_t(run_team_region_async_ty)(int32_t, void *, void **,
                                             ptrdiff_t *, int32_t, int32_t,
#endif
 
   // Functions implemented in the RTL.
   is_valid_binary_ty *is_valid_binary = nullptr;
   is_data_exchangable_ty *is_data_exchangable = nullptr;
   number_of_devices_ty *number_of_devices = nullptr;
   init_device_ty *init_device = nullptr;
   load_binary_ty *load_binary = nullptr;
   data_alloc_ty *data_alloc = nullptr;
+  data_alloc_async_ty *data_alloc_async = nullptr;
   data_submit_ty *data_submit = nullptr;
   data_submit_async_ty *data_submit_async = nullptr;
   data_retrieve_ty *data_retrieve = nullptr;
   data_retrieve_async_ty *data_retrieve_async = nullptr;
   data_exchange_ty *data_exchange = nullptr;
   data_exchange_async_ty *data_exchange_async = nullptr;
   data_delete_ty *data_delete = nullptr;
+  data_delete_async_ty *data_delete_async = nullptr;
   run_region_ty *run_region = nullptr;
   run_region_async_ty *run_region_async = nullptr;
   run_team_region_ty *run_team_region = nullptr;
   run_team_region_async_ty *run_team_region_async = nullptr;
   init_requires_ty *init_requires = nullptr;
   synchronize_ty *synchronize = nullptr;
   register_lib_ty register_lib = nullptr;
   register_lib_ty unregister_lib = nullptr;

openmp/libomptarget/plugins/common/MemoryManager/MemoryManager.h

 #include <functional>
 #include <list>
 #include <mutex>
 #include <set>
 #include <unordered_map>
 #include <vector>
 
 #include "Debug.h"
+#include "omptarget.h"
 #include "omptargetplugin.h"
 
 /// Base class of per-device allocator.
 class DeviceAllocatorTy {
 public:
   virtual ~DeviceAllocatorTy() = default;
 
   /// Allocate a memory of size \p Size . \p HstPtr is used to assist the
-  /// allocation.
-  virtual void *allocate(size_t Size, void *HstPtr, TargetAllocTy Kind) = 0;
-
-  /// Delete the pointer \p TgtPtr on the device
-  virtual int free(void *TgtPtr) = 0;
+  /// allocation. \p AsyncInfo is used to make it asynchronous, if the device
+  /// (driver) supports it.
+  virtual void *allocate(size_t Size, void *HstPtr, TargetAllocTy Kind,
+                         __tgt_async_info *AsyncInfo) = 0;
+
+  /// Delete the pointer \p TgtPtr on the device. \p AsyncInfo is used to make
+  /// it asynchronous, if the device (driver) supports it.
+  virtual int free(void *TgtPtr, __tgt_async_info *AsyncInfo) = 0;
 };
 
 /// Class of memory manager. The memory manager is per-device by using
 /// per-device allocator. Therefore, each plugin using memory manager should
 /// have an allocator for each device.
 class MemoryManagerTy {
   static constexpr const size_t BucketSize[] = {
       0,       1U << 2, 1U << 3,  1U << 4,  1U << 5,  1U << 6, 1U << 7,
#endif
   DeviceAllocatorTy &DeviceAllocator;
 
   /// The threshold to manage memory using memory manager. If the request size
   /// is larger than \p SizeThreshold, the allocation will not be managed by the
   /// memory manager.
   size_t SizeThreshold = 1U << 13;
 
   /// Request memory from target device
-  void *allocateOnDevice(size_t Size, void *HstPtr) const {
-    return DeviceAllocator.allocate(Size, HstPtr, TARGET_ALLOC_DEVICE);
+  void *allocateOnDevice(size_t Size, void *HstPtr,
+                         __tgt_async_info *AsyncInfo) const {
+    return DeviceAllocator.allocate(Size, HstPtr, TARGET_ALLOC_DEVICE,
+                                    AsyncInfo);
   }
 
   /// Deallocate data on device
-  int deleteOnDevice(void *Ptr) const { return DeviceAllocator.free(Ptr); }
+  int deleteOnDevice(void *Ptr) const {
+    // If we need to make space on the device we don't want it to be
+    // asynchronous as it only happens if we run out of memory. Freeing
+    // it is priority now.
+    return DeviceAllocator.free(Ptr, /* AsyncInfo */ nullptr);
+  }
 
   /// This function is called when it tries to allocate memory on device but the
   /// device returns out of memory. It will first free all memory in the
   /// FreeList and try to allocate again.
-  void *freeAndAllocate(size_t Size, void *HstPtr) {
+  void *freeAndAllocate(size_t Size, void *HstPtr,
+                        __tgt_async_info *AsyncInfo) {
     std::vector<void *> RemoveList;
 
     // Deallocate all memory in FreeList
     for (int I = 0; I < NumBuckets; ++I) {
       FreeListTy &List = FreeLists[I];
       std::lock_guard<std::mutex> Lock(FreeListLocks[I]);
       if (List.empty())
         continue;
       for (const NodeTy &N : List) {
         deleteOnDevice(N.Ptr);
         RemoveList.push_back(N.Ptr);
       }
       FreeLists[I].clear();
     }
 
     // Remove all nodes in the map table which have been released
     if (!RemoveList.empty()) {
       std::lock_guard<std::mutex> LG(MapTableLock);
       for (void *P : RemoveList)
         PtrToNodeTable.erase(P);
     }
 
     // Try allocate memory again
-    return allocateOnDevice(Size, HstPtr);
+    return allocateOnDevice(Size, HstPtr, AsyncInfo);
   }
 
   /// The goal is to allocate memory on the device. It first tries to
   /// allocate directly on the device. If a \p nullptr is returned, it might
   /// be because the device is OOM. In that case, it will free all unused
   /// memory and then try again.
-  void *allocateOrFreeAndAllocateOnDevice(size_t Size, void *HstPtr) {
-    void *TgtPtr = allocateOnDevice(Size, HstPtr);
+  void *allocateOrFreeAndAllocateOnDevice(size_t Size, void *HstPtr,
+                                          __tgt_async_info *AsyncInfo) {
+    void *TgtPtr = allocateOnDevice(Size, HstPtr, AsyncInfo);
     // We cannot get memory from the device. It might be due to OOM. Let's
     // free all memory in FreeLists and try again.
     if (TgtPtr == nullptr) {
       DP("Failed to get memory on device. Free all memory in FreeLists and "
          "try again.\n");
-      TgtPtr = freeAndAllocate(Size, HstPtr);
+      TgtPtr = freeAndAllocate(Size, HstPtr, AsyncInfo);
     }
 
     if (TgtPtr == nullptr)
       DP("Still cannot get memory on device probably because the device is "
          "OOM.\n");
 
     return TgtPtr;
   }
for (auto Itr = PtrToNodeTable.begin(); Itr != PtrToNodeTable.end();
          ++Itr) {
       assert(Itr->second.Ptr && "nullptr in map table");
       deleteOnDevice(Itr->second.Ptr);
     }
   }
 
   /// Allocate memory of size \p Size from target device. \p HstPtr is used to
   /// assist the allocation.
-  void *allocate(size_t Size, void *HstPtr) {
+  void *allocate(size_t Size, void *HstPtr, __tgt_async_info *AsyncInfo) {
     // If the size is zero, we will not bother the target device. Just return
     // nullptr directly.
     if (Size == 0)
       return nullptr;
 
     DP("MemoryManagerTy::allocate: size %zu with host pointer " DPxMOD ".\n",
        Size, DPxPTR(HstPtr));
 
     // If the size is greater than the threshold, allocate it directly from
     // device.
     if (Size > SizeThreshold) {
       DP("%zu is greater than the threshold %zu. Allocate it directly from "
          "device\n",
          Size, SizeThreshold);
-      void *TgtPtr = allocateOrFreeAndAllocateOnDevice(Size, HstPtr);
+      void *TgtPtr = allocateOrFreeAndAllocateOnDevice(Size, HstPtr, AsyncInfo);
 
       DP("Got target pointer " DPxMOD ". Return directly.\n", DPxPTR(TgtPtr));
 
       return TgtPtr;
     }
 
     NodeTy *NodePtr = nullptr;
 
void *allocate(size_t Size, void *HstPtr, __tgt_async_info *AsyncInfo) {
     if (NodePtr != nullptr)
       DP("Find one node " DPxMOD " in the bucket.\n", DPxPTR(NodePtr));
 
     // We cannot find a valid node in FreeLists. Let's allocate on device and
     // create a node for it.
     if (NodePtr == nullptr) {
       DP("Cannot find a node in the FreeLists. Allocate on device.\n");
       // Allocate one on device
-      void *TgtPtr = allocateOrFreeAndAllocateOnDevice(Size, HstPtr);
+      void *TgtPtr = allocateOrFreeAndAllocateOnDevice(Size, HstPtr, AsyncInfo);
 
       if (TgtPtr == nullptr)
         return nullptr;
 
       // Create a new node and add it into the map table
       {
         std::lock_guard<std::mutex> Guard(MapTableLock);
         auto Itr = PtrToNodeTable.emplace(TgtPtr, NodeTy(Size, TgtPtr));
         NodePtr = &Itr.first->second;
       }
 
       DP("Node address " DPxMOD ", target pointer " DPxMOD ", size %zu\n",
          DPxPTR(NodePtr), DPxPTR(TgtPtr), Size);
     }
 
     assert(NodePtr && "NodePtr should not be nullptr at this point");
 
     return NodePtr->Ptr;
   }
 
   /// Deallocate memory pointed by \p TgtPtr
   int free(void *TgtPtr) {
     DP("MemoryManagerTy::free: target memory " DPxMOD ".\n", DPxPTR(TgtPtr));
 
     NodeTy *P = nullptr;
 
     // Look it up into the table
     {
       std::lock_guard<std::mutex> G(MapTableLock);
       auto Itr = PtrToNodeTable.find(TgtPtr);

openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.h

 
 CUresult cuGetErrorString(CUresult, const char **);
 CUresult cuInit(unsigned);
 CUresult cuLaunchKernel(CUfunction, unsigned, unsigned, unsigned, unsigned,
                         unsigned, unsigned, unsigned, CUstream, void **,
                         void **);
 
 CUresult cuMemAlloc(CUdeviceptr *, size_t);
+CUresult cuMemAllocAsync(CUdeviceptr *, size_t, CUstream);
 CUresult cuMemAllocHost(void **, size_t);
 CUresult cuMemAllocManaged(CUdeviceptr *, size_t, unsigned int);
 
 CUresult cuMemcpyDtoDAsync(CUdeviceptr, CUdeviceptr, size_t, CUstream);
 CUresult cuMemcpyDtoH(void *, CUdeviceptr, size_t);
 CUresult cuMemcpyDtoHAsync(void *, CUdeviceptr, size_t, CUstream);
 CUresult cuMemcpyHtoD(CUdeviceptr, const void *, size_t);
 CUresult cuMemcpyHtoDAsync(CUdeviceptr, const void *, size_t, CUstream);
 
 CUresult cuMemFree(CUdeviceptr);
+CUresult cuMemFreeAsync(CUdeviceptr, CUstream);
 CUresult cuMemFreeHost(void *);
 
 CUresult cuModuleGetFunction(CUfunction *, CUmodule, const char *);
 CUresult cuModuleGetGlobal(CUdeviceptr *, size_t *, CUmodule, const char *);
 
 CUresult cuModuleUnload(CUmodule);
 CUresult cuStreamCreate(CUstream *, unsigned);
 CUresult cuStreamDestroy(CUstream);

openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp

 DLWRAP(cuDeviceGetName, 3);
 DLWRAP(cuDeviceTotalMem, 2);
 DLWRAP(cuDriverGetVersion, 1);
 
 DLWRAP(cuGetErrorString, 2);
 DLWRAP(cuLaunchKernel, 11);
 
 DLWRAP(cuMemAlloc, 2);
+DLWRAP(cuMemAllocAsync, 3);
 DLWRAP(cuMemAllocHost, 2);
 DLWRAP(cuMemAllocManaged, 3);
 
 DLWRAP(cuMemcpyDtoDAsync, 4);
 DLWRAP(cuMemcpyDtoH, 3);
 DLWRAP(cuMemcpyDtoHAsync, 4);
 DLWRAP(cuMemcpyHtoD, 3);
 DLWRAP(cuMemcpyHtoDAsync, 4);
 
 DLWRAP(cuMemFree, 1);
+DLWRAP(cuMemFreeAsync, 2);
 DLWRAP(cuMemFreeHost, 1);
 DLWRAP(cuModuleGetFunction, 3);
 DLWRAP(cuModuleGetGlobal, 4);
 
 DLWRAP(cuModuleUnload, 1);
 DLWRAP(cuStreamCreate, 2);
 DLWRAP(cuStreamDestroy, 1);
 DLWRAP(cuStreamSynchronize, 1);

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 #include <memory>
 #include <mutex>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
 #include "Debug.h"
 #include "DeviceEnvironment.h"
+#include "omptarget.h"
 #include "omptargetplugin.h"
 
 #define TARGET_NAME CUDA
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
 
 #include "MemoryManager.h"
 
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
public:
   /// Note: This function is not thread safe. Be sure to guard it if necessary.
   void clear() noexcept {
     for (auto &R : Resources)
       (void)Allocator.destroy(R);
     Resources.clear();
   }
 };
 
+using StreamPoolTy = ResourcePoolTy<CUstream>;
+using StreamPoolVectorTy = std::vector<std::unique_ptr<StreamPoolTy>>;
+
+static CUstream getStream(const int DeviceId, __tgt_async_info *AsyncInfo,
+                          const StreamPoolVectorTy &StreamPoolVector) {
+  assert(AsyncInfo && "AsyncInfo is nullptr");
+
+  if (!AsyncInfo->Queue) {
+    CUstream S;
+    if (StreamPoolVector[DeviceId]->acquire(S) != OFFLOAD_SUCCESS)
+      return nullptr;
+
+    AsyncInfo->Queue = S;
+  }
+
+  return reinterpret_cast<CUstream>(AsyncInfo->Queue);
+}
+
 class DeviceRTLTy {
   int NumberOfDevices;
   // OpenMP environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
   int EnvTeamThreadLimit;
   // OpenMP requires flags
   int64_t RequiresFlags;
   // Amount of dynamic shared memory to use at launch.
   uint64_t DynamicMemorySize;
   // Number of initial streams for each device.
   int NumInitialStreams = 32;
 
   static constexpr const int HardTeamLimit = 1U << 16U; // 64k
   static constexpr const int HardThreadLimit = 1024;
   static constexpr const int DefaultNumTeams = 128;
   static constexpr const int DefaultNumThreads = 128;
 
-  using StreamPoolTy = ResourcePoolTy<CUstream>;
-  std::vector<std::unique_ptr<StreamPoolTy>> StreamPool;
+  StreamPoolVectorTy StreamPoolVector;
 
   ResourcePoolTy<CUevent> EventPool;
 
   std::vector<DeviceDataTy> DeviceData;
   std::vector<CUmodule> Modules;
 
   /// A class responsible for interacting with device native runtime library to
   /// allocate and free memory.
   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
     const int DeviceId;
     const std::vector<DeviceDataTy> &DeviceData;
     std::unordered_map<void *, TargetAllocTy> HostPinnedAllocs;
 
+    StreamPoolVectorTy &StreamPoolVector;
+
   public:
-    CUDADeviceAllocatorTy(int DeviceId, std::vector<DeviceDataTy> &DeviceData)
-        : DeviceId(DeviceId), DeviceData(DeviceData) {}
+    CUDADeviceAllocatorTy(int DeviceId, std::vector<DeviceDataTy> &DeviceData,
+                          StreamPoolVectorTy &StreamPoolVector)
+        : DeviceId(DeviceId), DeviceData(DeviceData),
+          StreamPoolVector(StreamPoolVector) {}
 
-    void *allocate(size_t Size, void *, TargetAllocTy Kind) override {
+    void *allocate(size_t Size, void *, TargetAllocTy Kind,
+                   __tgt_async_info *AsyncInfo) override {
       if (Size == 0)
         return nullptr;
 
       CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
       if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
         return nullptr;
 
       void *MemAlloc = nullptr;
       switch (Kind) {
       case TARGET_ALLOC_DEFAULT:
       case TARGET_ALLOC_DEVICE:
         CUdeviceptr DevicePtr;
+#if CUDA_VERSION >= 11020
+        if (AsyncInfo) {
+          CUstream Stream = getStream(DeviceId, AsyncInfo, StreamPoolVector);
+          Err = cuMemAllocAsync(&DevicePtr, Size, Stream);
+        } else {
           Err = cuMemAlloc(&DevicePtr, Size);
+        }
+#else
+        Err = cuMemAlloc(&DevicePtr, Size);
+#endif
         MemAlloc = (void *)DevicePtr;
         if (!checkResult(Err, "Error returned from cuMemAlloc\n"))
           return nullptr;
         break;
       case TARGET_ALLOC_HOST:
         void *HostPtr;
         Err = cuMemAllocHost(&HostPtr, Size);
         MemAlloc = HostPtr;
         if (!checkResult(Err, "Error returned from cuMemAllocHost\n"))
           return nullptr;
         HostPinnedAllocs[MemAlloc] = Kind;
         break;
       case TARGET_ALLOC_SHARED:
         CUdeviceptr SharedPtr;
         Err = cuMemAllocManaged(&SharedPtr, Size, CU_MEM_ATTACH_GLOBAL);
         MemAlloc = (void *)SharedPtr;
         if (!checkResult(Err, "Error returned from cuMemAllocManaged\n"))
           return nullptr;
         break;
       }
 
       return MemAlloc;
     }
 
-    int free(void *TgtPtr) override {
+    int free(void *TgtPtr, __tgt_async_info *AsyncInfo) override {
       CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
       if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
         return OFFLOAD_FAIL;
 
       // Host pinned memory must be freed differently.
       TargetAllocTy Kind =
           (HostPinnedAllocs.find(TgtPtr) == HostPinnedAllocs.end())
               ? TARGET_ALLOC_DEFAULT
               : TARGET_ALLOC_HOST;
       switch (Kind) {
       case TARGET_ALLOC_DEFAULT:
       case TARGET_ALLOC_DEVICE:
       case TARGET_ALLOC_SHARED:
+#if CUDA_VERSION >= 11020
+        if (AsyncInfo) {
+          CUstream Stream = getStream(DeviceId, AsyncInfo, StreamPoolVector);
+          Err = cuMemFreeAsync((CUdeviceptr)TgtPtr, Stream);
+        } else {
           Err = cuMemFree((CUdeviceptr)TgtPtr);
+        }
+#else
+        Err = cuMemFree((CUdeviceptr)TgtPtr);
+#endif
         if (!checkResult(Err, "Error returned from cuMemFree\n"))
           return OFFLOAD_FAIL;
         break;
       case TARGET_ALLOC_HOST:
         Err = cuMemFreeHost(TgtPtr);
         if (!checkResult(Err, "Error returned from cuMemFreeHost\n"))
           return OFFLOAD_FAIL;
         break;
#endif
   // Clear entries table for a device
   void clearOffloadEntriesTable(const int DeviceId) {
     DeviceData[DeviceId].FuncGblEntries.emplace_back();
     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = nullptr;
   }
 
-  CUstream getStream(const int DeviceId, __tgt_async_info *AsyncInfo) const {
-    assert(AsyncInfo && "AsyncInfo is nullptr");
-
-    if (!AsyncInfo->Queue) {
-      CUstream S;
-      if (StreamPool[DeviceId]->acquire(S) != OFFLOAD_SUCCESS)
-        return nullptr;
-
-      AsyncInfo->Queue = S;
-    }
-
-    return reinterpret_cast<CUstream>(AsyncInfo->Queue);
-  }
-
 public:
   // This class should not be copied
   DeviceRTLTy(const DeviceRTLTy &) = delete;
   DeviceRTLTy(DeviceRTLTy &&) = delete;
 
   DeviceRTLTy()
       : NumberOfDevices(0), EnvNumTeams(-1), EnvTeamLimit(-1),
         EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED),
if (!checkResult(Err, "Error returned from cuDeviceGetCount\n"))
       return;
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     DeviceData.resize(NumberOfDevices);
-    StreamPool.resize(NumberOfDevices);
+    StreamPoolVector.resize(NumberOfDevices);
 
     // Get environment variables regarding teams
     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     }
     if (const char *EnvStr = getenv("OMP_TEAMS_THREAD_LIMIT")) {
DeviceRTLTy()
     }
     if (const char *EnvStr = getenv("LIBOMPTARGET_NUM_INITIAL_STREAMS")) {
       // LIBOMPTARGET_NUM_INITIAL_STREAMS has been set
       NumInitialStreams = std::stoi(EnvStr);
       DP("Parsed LIBOMPTARGET_NUM_INITIAL_STREAMS=%d\n", NumInitialStreams);
     }
 
     for (int I = 0; I < NumberOfDevices; ++I)
-      DeviceAllocators.emplace_back(I, DeviceData);
+      DeviceAllocators.emplace_back(I, DeviceData, StreamPoolVector);
 
     // Get the size threshold from environment variable
     std::pair<size_t, bool> Res = MemoryManagerTy::getSizeThresholdFromEnv();
     UseMemoryManager = Res.second;
     size_t MemoryManagerThreshold = Res.first;
 
     if (UseMemoryManager)
       for (int I = 0; I < NumberOfDevices; ++I)
         MemoryManagers.emplace_back(std::make_unique<MemoryManagerTy>(
             DeviceAllocators[I], MemoryManagerThreshold));
   }
 
   ~DeviceRTLTy() {
     // We first destruct memory managers in case that its dependent data are
     // destroyed before it.
     for (auto &M : MemoryManagers)
       M.release();
 
     for (CUmodule &M : Modules)
       // Close module
       if (M)
         checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
 
-    for (auto &S : StreamPool)
+    for (auto &S : StreamPoolVector)
       S.reset();
 
     EventPool.clear();
 
     for (DeviceDataTy &D : DeviceData) {
       // Destroy context
       if (D.Context) {
         checkResult(cuCtxSetCurrent(D.Context),
int initDevice(const int DeviceId) {
     if (!checkResult(Err, "Error returned from cuDevicePrimaryCtxRetain\n"))
       return OFFLOAD_FAIL;
 
     Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
     // Initialize stream pool
-    if (!StreamPool[DeviceId])
-      StreamPool[DeviceId] = std::make_unique<StreamPoolTy>(
+    if (!StreamPoolVector[DeviceId])
+      StreamPoolVector[DeviceId] = std::make_unique<StreamPoolTy>(
           AllocatorTy<CUstream>(DeviceData[DeviceId].Context),
           NumInitialStreams);
 
     // Query attributes to determine number of threads/block and blocks/grid.
     int MaxGridDimX;
     Err = cuDeviceGetAttribute(&MaxGridDimX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
                                Device);
     if (Err != CUDA_SUCCESS) {
// send device environment data to the device
            "environment setting.\n");
       }
     }
 
     return getOffloadEntriesTable(DeviceId);
   }
 
   void *dataAlloc(const int DeviceId, const int64_t Size,
-                  const TargetAllocTy Kind) {
+                  const TargetAllocTy Kind, __tgt_async_info *AsyncInfo) {
     switch (Kind) {
     case TARGET_ALLOC_DEFAULT:
     case TARGET_ALLOC_DEVICE:
       if (UseMemoryManager)
-        return MemoryManagers[DeviceId]->allocate(Size, nullptr);
+        return MemoryManagers[DeviceId]->allocate(Size, nullptr, AsyncInfo);
       else
-        return DeviceAllocators[DeviceId].allocate(Size, nullptr, Kind);
+        return DeviceAllocators[DeviceId].allocate(Size, nullptr, Kind,
+                                                   AsyncInfo);
     case TARGET_ALLOC_HOST:
     case TARGET_ALLOC_SHARED:
-      return DeviceAllocators[DeviceId].allocate(Size, nullptr, Kind);
+      return DeviceAllocators[DeviceId].allocate(Size, nullptr, Kind,
+                                                 AsyncInfo);
     }
 
     REPORT("Invalid target data allocation kind or requested allocator not "
            "implemented yet\n");
 
     return nullptr;
   }
 
   int dataSubmit(const int DeviceId, const void *TgtPtr, const void *HstPtr,
                  const int64_t Size, __tgt_async_info *AsyncInfo) const {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
-    CUstream Stream = getStream(DeviceId, AsyncInfo);
+    CUstream Stream = getStream(DeviceId, AsyncInfo, StreamPoolVector);
 
     Err = cuMemcpyHtoDAsync((CUdeviceptr)TgtPtr, HstPtr, Size, Stream);
     if (Err != CUDA_SUCCESS) {
       DP("Error when copying data from host to device. Pointers: host "
          "= " DPxMOD ", device = " DPxMOD ", size = %" PRId64 "\n",
          DPxPTR(HstPtr), DPxPTR(TgtPtr), Size);
       CUDA_ERR_STRING(Err);
       return OFFLOAD_FAIL;
     }
 
     return OFFLOAD_SUCCESS;
   }
 
   int dataRetrieve(const int DeviceId, void *HstPtr, const void *TgtPtr,
                    const int64_t Size, __tgt_async_info *AsyncInfo) const {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
-    CUstream Stream = getStream(DeviceId, AsyncInfo);
+    CUstream Stream = getStream(DeviceId, AsyncInfo, StreamPoolVector);
 
     Err = cuMemcpyDtoHAsync(HstPtr, (CUdeviceptr)TgtPtr, Size, Stream);
     if (Err != CUDA_SUCCESS) {
       DP("Error when copying data from device to host. Pointers: host "
          "= " DPxMOD ", device = " DPxMOD ", size = %" PRId64 "\n",
          DPxPTR(HstPtr), DPxPTR(TgtPtr), Size);
       CUDA_ERR_STRING(Err);
       return OFFLOAD_FAIL;
     }
 
     return OFFLOAD_SUCCESS;
   }
 
   int dataExchange(int SrcDevId, const void *SrcPtr, int DstDevId, void *DstPtr,
                    int64_t Size, __tgt_async_info *AsyncInfo) const {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
     CUresult Err = cuCtxSetCurrent(DeviceData[SrcDevId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
-    CUstream Stream = getStream(SrcDevId, AsyncInfo);
+    CUstream Stream = getStream(SrcDevId, AsyncInfo, StreamPoolVector);
 
     // If they are two devices, we try peer to peer copy first
     if (SrcDevId != DstDevId) {
       int CanAccessPeer = 0;
       Err = cuDeviceCanAccessPeer(&CanAccessPeer, SrcDevId, DstDevId);
       if (Err != CUDA_SUCCESS) {
         REPORT("Error returned from cuDeviceCanAccessPeer. src = %" PRId32
                ", dst = %" PRId32 "\n",
if (SrcDevId != DstDevId) {
          ", src_id =%" PRId32 ", dst_ptr = " DPxMOD ", dst_id =%" PRId32 "\n",
          DPxPTR(SrcPtr), SrcDevId, DPxPTR(DstPtr), DstDevId);
       CUDA_ERR_STRING(Err);
     }
 
     return memcpyDtoD(SrcPtr, DstPtr, Size, Stream);
   }
 
-  int dataDelete(const int DeviceId, void *TgtPtr) {
+  int dataDelete(const int DeviceId, void *TgtPtr,
+                 __tgt_async_info *AsyncInfo) {
     if (UseMemoryManager)
       return MemoryManagers[DeviceId]->free(TgtPtr);
 
-    return DeviceAllocators[DeviceId].free(TgtPtr);
+    return DeviceAllocators[DeviceId].free(TgtPtr, AsyncInfo);
   }
 
   int runTargetTeamRegion(const int DeviceId, void *TgtEntryPtr, void **TgtArgs,
                           ptrdiff_t *TgtOffsets, const int ArgNum,
                           const int TeamNum, const int ThreadLimit,
                           const unsigned int LoopTripCount,
                           __tgt_async_info *AsyncInfo) const {
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
int runTargetTeamRegion(const int DeviceId, void *TgtEntryPtr, void **TgtArgs,
     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
          "Launching kernel %s with %d blocks and %d threads in %s mode\n",
          (getOffloadEntry(DeviceId, TgtEntryPtr))
              ? getOffloadEntry(DeviceId, TgtEntryPtr)->name
              : "(null)",
          CudaBlocksPerGrid, CudaThreadsPerBlock,
          (!IsSPMDMode ? (IsGenericMode ? "Generic" : "SPMD-Generic") : "SPMD"));
 
-    CUstream Stream = getStream(DeviceId, AsyncInfo);
+    CUstream Stream = getStream(DeviceId, AsyncInfo, StreamPoolVector);
     Err = cuLaunchKernel(KernelInfo->Func, CudaBlocksPerGrid, /* gridDimY */ 1,
                          /* gridDimZ */ 1, CudaThreadsPerBlock,
                          /* blockDimY */ 1, /* blockDimZ */ 1,
                          DynamicMemorySize, Stream, &Args[0], nullptr);
     if (!checkResult(Err, "Error returned from cuLaunchKernel\n"))
       return OFFLOAD_FAIL;
 
     DP("Launch of entry point at " DPxMOD " successful!\n",
        DPxPTR(TgtEntryPtr));
 
     return OFFLOAD_SUCCESS;
   }
 
   int synchronize(const int DeviceId, __tgt_async_info *AsyncInfo) const {
     CUstream Stream = reinterpret_cast<CUstream>(AsyncInfo->Queue);
     CUresult Err = cuStreamSynchronize(Stream);
 
     // Once the stream is synchronized, return it to stream pool and reset
     // AsyncInfo. This is to make sure the synchronization only works for its
     // own tasks.
-    StreamPool[DeviceId]->release(reinterpret_cast<CUstream>(AsyncInfo->Queue));
+    StreamPoolVector[DeviceId]->release(
+        reinterpret_cast<CUstream>(AsyncInfo->Queue));
     AsyncInfo->Queue = nullptr;
 
     if (Err != CUDA_SUCCESS) {
       DP("Error when synchronizing stream. stream = " DPxMOD
          ", async info ptr = " DPxMOD "\n",
          DPxPTR(Stream), DPxPTR(AsyncInfo));
       CUDA_ERR_STRING(Err);
     }
public:
 
   int destroyEvent(void *EventPtr) {
     EventPool.release(reinterpret_cast<CUevent>(EventPtr));
     return OFFLOAD_SUCCESS;
   }
 
   int waitEvent(const int DeviceId, __tgt_async_info *AsyncInfo,
                 void *EventPtr) const {
-    CUstream Stream = getStream(DeviceId, AsyncInfo);
+    CUstream Stream = getStream(DeviceId, AsyncInfo, StreamPoolVector);
     CUevent Event = reinterpret_cast<CUevent>(EventPtr);
 
     // We don't use CU_EVENT_WAIT_DEFAULT here as it is only available from
     // specific CUDA version, and defined as 0x0. In previous version, per CUDA
     // API document, that argument has to be 0x0.
     CUresult Err = cuStreamWaitEvent(Stream, Event, 0);
     if (Err != CUDA_SUCCESS) {
       DP("Error when waiting event. stream = " DPxMOD ", event = " DPxMOD "\n",

 
 __tgt_target_table *__tgt_rtl_load_binary(int32_t device_id,
                                           __tgt_device_image *image) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   return DeviceRTL.loadBinary(device_id, image);
 }
 
-void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *,
+void *__tgt_rtl_data_alloc(int32_t device_id, int64_t size, void *hst_ptr,
                            int32_t kind) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
+  __tgt_async_info AsyncInfo;
+  void *ptr =
+      __tgt_rtl_data_alloc_async(device_id, size, hst_ptr, kind, &AsyncInfo);
+  if (AsyncInfo.Queue)
+    __tgt_rtl_synchronize(device_id, &AsyncInfo);
+  return ptr;
+}
 
-  return DeviceRTL.dataAlloc(device_id, size, (TargetAllocTy)kind);
+void *__tgt_rtl_data_alloc_async(int32_t device_id, int64_t size, void *hst_ptr,
+                                 int32_t kind,
+                                 __tgt_async_info *async_info_ptr) {
+  assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
+  return DeviceRTL.dataAlloc(device_id, size, (TargetAllocTy)kind,
+                             async_info_ptr);
 }
 
 int32_t __tgt_rtl_data_submit(int32_t device_id, void *tgt_ptr, void *hst_ptr,
                               int64_t size) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
   __tgt_async_info AsyncInfo;
   const int32_t rc = __tgt_rtl_data_submit_async(device_id, tgt_ptr, hst_ptr,
if (rc != OFFLOAD_SUCCESS)
     return OFFLOAD_FAIL;
 
   return __tgt_rtl_synchronize(src_dev_id, &AsyncInfo);
 }
 
 int32_t __tgt_rtl_data_delete(int32_t device_id, void *tgt_ptr) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
 
-  return DeviceRTL.dataDelete(device_id, tgt_ptr);
+  __tgt_async_info AsyncInfo;
+  const int32_t rc =
+      __tgt_rtl_data_delete_async(device_id, tgt_ptr, &AsyncInfo);
+  if (rc != OFFLOAD_SUCCESS)
+    return OFFLOAD_FAIL;
+
+  if (AsyncInfo.Queue)
+    return __tgt_rtl_synchronize(device_id, &AsyncInfo);
+  return OFFLOAD_SUCCESS;
+}
+
+int32_t __tgt_rtl_data_delete_async(int32_t device_id, void *tgt_ptr,
+                                    __tgt_async_info *async_info_ptr) {
+  assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
+#if CUDA_VERSION < 11020
+  if (async_info_ptr && async_info_ptr->Queue)
+    __tgt_rtl_synchronize(device_id, async_info_ptr);
+#endif
+  return DeviceRTL.dataDelete(device_id, tgt_ptr, async_info_ptr);
 }
 
 int32_t __tgt_rtl_run_target_team_region(int32_t device_id, void *tgt_entry_ptr,
                                          void **tgt_args,
                                          ptrdiff_t *tgt_offsets,
                                          int32_t arg_num, int32_t team_num,
                                          int32_t thread_limit,
                                          uint64_t loop_tripcount) {

openmp/libomptarget/plugins/exports

 VERS1.0 {
   global:
     __tgt_rtl_is_valid_binary;
     __tgt_rtl_is_data_exchangable;
     __tgt_rtl_number_of_devices;
     __tgt_rtl_init_requires;
     __tgt_rtl_init_device;
     __tgt_rtl_load_binary;
     __tgt_rtl_data_alloc;
+    __tgt_rtl_data_alloc_async;
     __tgt_rtl_data_submit;
     __tgt_rtl_data_submit_async;
     __tgt_rtl_data_retrieve;
     __tgt_rtl_data_retrieve_async;
     __tgt_rtl_data_exchange;
     __tgt_rtl_data_exchange_async;
     __tgt_rtl_data_delete;
+    __tgt_rtl_data_delete_async;
     __tgt_rtl_run_target_team_region;
     __tgt_rtl_run_target_team_region_async;
     __tgt_rtl_run_target_region;
     __tgt_rtl_run_target_region_async;
     __tgt_rtl_synchronize;
     __tgt_rtl_register_lib;
     __tgt_rtl_unregister_lib;
     __tgt_rtl_supports_empty_images;

openmp/libomptarget/src/api.cpp

if (device_num == omp_get_initial_device()) {
     return;
   }
 
   if (!device_is_ready(device_num)) {
     DP("omp_target_free returns, nothing to do\n");
     return;
   }
 
-  PM->Devices[device_num]->deleteData(device_ptr);
+  PM->Devices[device_num]->deleteData(device_ptr, /* AsyncInfoPtr */ nullptr);
   DP("omp_target_free deallocated device ptr\n");
 }
 
 EXTERN int omp_target_is_present(const void *ptr, int device_num) {
   TIMESCOPE();
   DP("Call to omp_target_is_present for device %d and address " DPxMOD "\n",
      device_num, DPxPTR(ptr));
 

openmp/libomptarget/src/device.cpp

 //===--------- device.cpp - Target independent OpenMP target RTL ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Functionality for managing devices that are handled by RTL plugins.
 //
 //===----------------------------------------------------------------------===//
 
 #include "device.h"
+#include "omptarget.h"
 #include "private.h"
 #include "rtl.h"
 
 #include <cassert>
 #include <climits>
 #include <cstdio>
 #include <string>
 
-int HostDataToTargetTy::addEventIfNecessary(
-    DeviceTy &Device, AsyncInfoTy &AsyncInfo) const {
+int HostDataToTargetTy::addEventIfNecessary(DeviceTy &Device,
+                                            AsyncInfoTy &AsyncInfo) const {
   // First, check if the user disabled atomic map transfer/malloc/dealloc.
-  if (!PM->UseEventsForAtomicTransfers)
+  if (!AsyncInfo.hasPotentiallyPendingAsyncWork() || !PM->UseEventsForAtomicTransfers)
     return OFFLOAD_SUCCESS;
 
   void *Event = getEvent();
   bool NeedNewEvent = Event == nullptr;
   if (NeedNewEvent && Device.createEvent(&Event) != OFFLOAD_SUCCESS) {
     REPORT("Failed to create event\n");
     return OFFLOAD_FAIL;
   }
if (NeedNewEvent)
     setEvent(Event);
 
   return OFFLOAD_SUCCESS;
 }
 
 DeviceTy::DeviceTy(RTLInfoTy *RTL)
     : DeviceID(-1), RTL(RTL), RTLDeviceID(-1), IsInit(false), InitFlag(),
       HasPendingGlobals(false), HostDataToTargetMap(), PendingCtorsDtors(),
-      ShadowPtrMap(), DataMapMtx(), PendingGlobalsMtx(), ShadowMtx() {}
+      ShadowPtrMap(), DataMapMtx(), PendingGlobalsMtx(), ShadowMtx() {
+  SupportsAsyncFree = RTL->data_delete_async && RTL->synchronize;
+}
 
 DeviceTy::~DeviceTy() {
   if (DeviceID == -1 || !(getInfoLevel() & OMP_INFOTYPE_DUMP_TABLE))
     return;
 
   ident_t loc = {0, 0, 0, 0, ";libomptarget;libomptarget;0;0;;"};
   dumpTargetPointerMappings(&loc, *this);
 }
LookupResult DeviceTy::lookupMapping(void *HstPtrBegin, int64_t Size) {
   if (lr.Flags.ExtendsAfter) {
     DP("WARNING: Pointer is already mapped but section extends beyond mapped "
        "region\n");
   }
 
   return lr;
 }
 
-TargetPointerResultTy
-DeviceTy::getTargetPointer(void *HstPtrBegin, void *HstPtrBase, int64_t Size,
-                           map_var_info_t HstPtrName, bool HasFlagTo,
-                           bool HasFlagAlways, bool IsImplicit,
-                           bool UpdateRefCount, bool HasCloseModifier,
-                           bool HasPresentModifier, bool HasHoldModifier,
-                           AsyncInfoTy &AsyncInfo) {
+TargetPointerResultTy DeviceTy::getTargetPointer(
+    void *HstPtrBegin, void *HstPtrBase, int64_t Size,
+    map_var_info_t HstPtrName, bool HasFlagTo, bool HasFlagAlways,
+    bool IsImplicit, bool UpdateRefCount, bool HasCloseModifier,
+    bool HasPresentModifier, bool HasHoldModifier, AsyncInfoTy &AsyncInfo) {
   void *TargetPointer = nullptr;
   bool IsHostPtr = false;
   bool IsNew = false;
 
   DataMapMtx.lock();
 
   LookupResult LR = lookupMapping(HstPtrBegin, Size);
   auto Entry = LR.Entry;
DP("Mapping required by 'present' map type modifier does not exist for "
        "HstPtrBegin=" DPxMOD ", Size=%" PRId64 "\n",
        DPxPTR(HstPtrBegin), Size);
     MESSAGE("device mapping required by 'present' map type modifier does not "
             "exist for host address " DPxMOD " (%" PRId64 " bytes)",
             DPxPTR(HstPtrBegin), Size);
   } else if (Size) {
     // If it is not contained and Size > 0, we should create a new entry for it.
     IsNew = true;
-    uintptr_t Ptr = (uintptr_t)allocData(Size, HstPtrBegin);
+    uintptr_t Ptr = (uintptr_t)allocData(Size, HstPtrBegin,
+                                         TARGET_ALLOC_DEFAULT, AsyncInfo);
     Entry = HostDataToTargetMap
                 .emplace((uintptr_t)HstPtrBase, (uintptr_t)HstPtrBegin,
                          (uintptr_t)HstPtrBegin + Size, Ptr, HasHoldModifier,
                          HstPtrName)
                 .first;
     INFO(OMP_INFOTYPE_MAPPING_CHANGED, DeviceID,
          "Creating new map entry with "
          "HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD ", Size=%ld, "
          "DynRefCount=%s, HoldRefCount=%s, Name=%s\n",
          DPxPTR(HstPtrBegin), DPxPTR(Ptr), Size,
          Entry->dynRefCountToStr().c_str(), Entry->holdRefCountToStr().c_str(),
          (HstPtrName) ? getNameFromMapping(HstPtrName).c_str() : "unknown");
     TargetPointer = (void *)Ptr;
+
+    // Ensure the atomicity of the allocation, if required. Allocation have to
+    // look atomic even if they are async and not accompanied with a data
+    // transfer.
+    if (SupportsAsyncFree) {
+      HostDataToTargetTy::LockGuard LG(*Entry);
+      if (Entry->addEventIfNecessary(*this, AsyncInfo) != OFFLOAD_SUCCESS)
+        return {{false /* IsNewEntry */, false /* IsHostPointer */},
+                {} /* MapTableEntry */,
+                nullptr /* TargetPointer */};
+    }
   }
 
   // If the target pointer is valid, and we need to transfer data, issue the
   // data transfer.
   if (TargetPointer && !IsHostPtr && HasFlagTo && (IsNew || HasFlagAlways)) {
     // Lock the entry before releasing the mapping table lock such that another
     // thread that could issue data movement will get the right result.
     HostDataToTargetTy::LockGuard LG(*Entry);
     // Release the mapping table lock right after the entry is locked.
     DataMapMtx.unlock();
 
     DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n", Size,
        DPxPTR(HstPtrBegin), DPxPTR(TargetPointer));
 
     int Ret = submitData(TargetPointer, HstPtrBegin, Size, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Copying data to device failed.\n");
       // We will also return nullptr if the data movement fails because that
       // pointer points to a corrupted memory region so it doesn't make any
       // sense to continue to use it.
       TargetPointer = nullptr;
-    } else if (Entry->addEventIfNecessary(*this, AsyncInfo) !=
-        OFFLOAD_SUCCESS)
+    } else if (Entry->addEventIfNecessary(*this, AsyncInfo) != OFFLOAD_SUCCESS)
       return {{false /* IsNewEntry */, false /* IsHostPointer */},
               {} /* MapTableEntry */,
               nullptr /* TargetPointer */};
   } else {
     // Release the mapping table lock directly.
     DataMapMtx.unlock();
     // If not a host pointer and no present modifier, we need to wait for the
     // event if it exists.
if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {
     uintptr_t tp = HT.TgtPtrBegin + (hp - HT.HstPtrBegin);
     return (void *)tp;
   }
 
   return NULL;
 }
 
 int DeviceTy::deallocTgtPtr(void *HstPtrBegin, int64_t Size,
-                            bool HasHoldModifier) {
+                            bool HasHoldModifier, AsyncInfoTy &AsyncInfo) {
   // Check if the pointer is contained in any sub-nodes.
   int Ret = OFFLOAD_SUCCESS;
   DataMapMtx.lock();
   LookupResult lr = lookupMapping(HstPtrBegin, Size);
   if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {
     auto &HT = *lr.Entry;
     if (HT.decRefCount(HasHoldModifier) == 0) {
       DP("Deleting tgt data " DPxMOD " of size %" PRId64 "\n",
          DPxPTR(HT.TgtPtrBegin), Size);
-      deleteData((void *)HT.TgtPtrBegin);
+      deleteData((void *)HT.TgtPtrBegin, AsyncInfo);
       INFO(OMP_INFOTYPE_MAPPING_CHANGED, DeviceID,
            "Removing map entry with HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD
            ", Size=%" PRId64 ", Name=%s\n",
            DPxPTR(HT.HstPtrBegin), DPxPTR(HT.TgtPtrBegin), Size,
            (HT.HstPtrName) ? getNameFromMapping(HT.HstPtrName).c_str()
                            : "unknown");
       void *Event = lr.Entry->getEvent();
       HostDataToTargetMap.erase(lr.Entry);

 // Load binary to device.
 __tgt_target_table *DeviceTy::load_binary(void *Img) {
   RTL->Mtx.lock();
   __tgt_target_table *rc = RTL->load_binary(RTLDeviceID, Img);
   RTL->Mtx.unlock();
   return rc;
 }
 
-void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {
+void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind,
+                          __tgt_async_info *AsyncInfo) {
+  if (!AsyncInfo || !SupportsAsyncFree)
     return RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);
+  return RTL->data_alloc_async(RTLDeviceID, Size, HstPtr, Kind, AsyncInfo);
 }
 
-int32_t DeviceTy::deleteData(void *TgtPtrBegin) {
+int32_t DeviceTy::deleteData(void *TgtPtrBegin, __tgt_async_info *AsyncInfo) {
+  if (!AsyncInfo || !SupportsAsyncFree)
     return RTL->data_delete(RTLDeviceID, TgtPtrBegin);
+  return RTL->data_delete_async(RTLDeviceID, TgtPtrBegin, AsyncInfo);
 }
 
 // Submit data to device
 int32_t DeviceTy::submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
                              AsyncInfoTy &AsyncInfo) {
   if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
     LookupResult LR = lookupMapping(HstPtrBegin, Size);
     auto *HT = &*LR.Entry;
 
     INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceID,
          "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD
          ", Size=%" PRId64 ", Name=%s\n",
          DPxPTR(HstPtrBegin), DPxPTR(TgtPtrBegin), Size,
          (HT && HT->HstPtrName) ? getNameFromMapping(HT->HstPtrName).c_str()
                                 : "unknown");
   }
 
   if (!AsyncInfo || !RTL->data_submit_async || !RTL->synchronize)
     return RTL->data_submit(RTLDeviceID, TgtPtrBegin, HstPtrBegin, Size);
-  else
-    return RTL->data_submit_async(RTLDeviceID, TgtPtrBegin, HstPtrBegin, Size,
-                                  AsyncInfo);
 }
 
 // Retrieve data from device
 int32_t DeviceTy::retrieveData(void *HstPtrBegin, void *TgtPtrBegin,
                                int64_t Size, AsyncInfoTy &AsyncInfo) {
   if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
     LookupResult LR = lookupMapping(HstPtrBegin, Size);
     auto *HT = &*LR.Entry;
     INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceID,
          "Copying data from device to host, TgtPtr=" DPxMOD ", HstPtr=" DPxMOD
          ", Size=%" PRId64 ", Name=%s\n",
          DPxPTR(TgtPtrBegin), DPxPTR(HstPtrBegin), Size,
          (HT && HT->HstPtrName) ? getNameFromMapping(HT->HstPtrName).c_str()
                                 : "unknown");
   }
 
   if (!RTL->data_retrieve_async || !RTL->synchronize)
     return RTL->data_retrieve(RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size);
-  else
-    return RTL->data_retrieve_async(RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size,
-                                    AsyncInfo);
 }
 
 // Copy data from current device to destination device directly
 int32_t DeviceTy::dataExchange(void *SrcPtr, DeviceTy &DstDev, void *DstPtr,
                                int64_t Size, AsyncInfoTy &AsyncInfo) {
   if (!AsyncInfo || !RTL->data_exchange_async || !RTL->synchronize) {
     assert(RTL->data_exchange && "RTL->data_exchange is nullptr");
     return RTL->data_exchange(RTLDeviceID, SrcPtr, DstDev.RTLDeviceID, DstPtr,
                               Size);
-  } else
-    return RTL->data_exchange_async(RTLDeviceID, SrcPtr, DstDev.RTLDeviceID,
-                                    DstPtr, Size, AsyncInfo);
+  }
 }
 
 // Run region on device
 int32_t DeviceTy::runRegion(void *TgtEntryPtr, void **TgtVarsPtr,
                             ptrdiff_t *TgtOffsets, int32_t TgtVarsSize,
                             AsyncInfoTy &AsyncInfo) {
   if (!RTL->run_region || !RTL->synchronize)
     return RTL->run_region(RTLDeviceID, TgtEntryPtr, TgtVarsPtr, TgtOffsets,
                            TgtVarsSize);
-  else
-    return RTL->run_region_async(RTLDeviceID, TgtEntryPtr, TgtVarsPtr,
-                                 TgtOffsets, TgtVarsSize, AsyncInfo);
+  return RTL->run_region_async(RTLDeviceID, TgtEntryPtr, TgtVarsPtr, TgtOffsets,
+                               TgtVarsSize, AsyncInfo);
 }
 
 // Run region on device
 bool DeviceTy::printDeviceInfo(int32_t RTLDevId) {
   if (!RTL->print_device_info)
     return false;
   RTL->print_device_info(RTLDevId);
   return true;
 }
 
 // Run team region on device.
 int32_t DeviceTy::runTeamRegion(void *TgtEntryPtr, void **TgtVarsPtr,
                                 ptrdiff_t *TgtOffsets, int32_t TgtVarsSize,
                                 int32_t NumTeams, int32_t ThreadLimit,
                                 uint64_t LoopTripCount,
                                 AsyncInfoTy &AsyncInfo) {
   if (!RTL->run_team_region_async || !RTL->synchronize)
     return RTL->run_team_region(RTLDeviceID, TgtEntryPtr, TgtVarsPtr,
                                 TgtOffsets, TgtVarsSize, NumTeams, ThreadLimit,
                                 LoopTripCount);
-  else
-    return RTL->run_team_region_async(RTLDeviceID, TgtEntryPtr, TgtVarsPtr,
-                                      TgtOffsets, TgtVarsSize, NumTeams,
-                                      ThreadLimit, LoopTripCount, AsyncInfo);
 }
 
 // Whether data can be copied to DstDevice directly
 bool DeviceTy::isDataExchangable(const DeviceTy &DstDevice) {
   if (RTL != DstDevice.RTL || !RTL->is_data_exchangable)
     return false;
 
   if (RTL->is_data_exchangable(RTLDeviceID, DstDevice.RTLDeviceID))

openmp/libomptarget/src/omptarget.cpp

if ((ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) || DelEntry) {
       applyToShadowMapEntries(Device, CB, HstPtrBegin, DataSize, TPR);
 
       // Add pointer to the buffer for later deallocation
       if (DelEntry && !IsHostPtr)
         DeallocTgtPtrs.emplace_back(HstPtrBegin, DataSize, HasHoldModifier);
     }
   }
 
-  // TODO: We should not synchronize here but pass the AsyncInfo object to the
-  //       allocate/deallocate device APIs.
-  //
-  // We need to synchronize before deallocating data.
+  // We need to synchronize before deallocating data if the device does not
+  // support asynchronous frees.
+  if (!Device.SupportsAsyncFree) {
     Ret = AsyncInfo.synchronize();
     if (Ret != OFFLOAD_SUCCESS)
       return OFFLOAD_FAIL;
+  }
 
   // Deallocate target pointer
   for (DeallocTgtPtrInfo &Info : DeallocTgtPtrs) {
     if (FromMapperBase && FromMapperBase == Info.HstPtrBegin)
       continue;
     Ret = Device.deallocTgtPtr(Info.HstPtrBegin, Info.DataSize,
-                               Info.HasHoldModifier);
+                               Info.HasHoldModifier, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Deallocating data from device failed.\n");
       return OFFLOAD_FAIL;
     }
   }
 
   return OFFLOAD_SUCCESS;
 }
int addArg(void *HstPtr, int64_t ArgSize, int64_t ArgOffset,
              bool IsFirstPrivate, void *&TgtPtr, int TgtArgsIndex,
              const map_var_info_t HstPtrName = nullptr,
              const bool AllocImmediately = false) {
     // If the argument is not first-private, or its size is greater than a
     // predefined threshold, we will allocate memory and issue the transfer
     // immediately.
     if (ArgSize > FirstPrivateArgSizeThreshold || !IsFirstPrivate ||
         AllocImmediately) {
-      TgtPtr = Device.allocData(ArgSize, HstPtr);
+      TgtPtr =
+          Device.allocData(ArgSize, HstPtr, TARGET_ALLOC_DEFAULT, AsyncInfo);
       if (!TgtPtr) {
         DP("Data allocation for %sprivate array " DPxMOD " failed.\n",
            (IsFirstPrivate ? "first-" : ""), DPxPTR(HstPtr));
         return OFFLOAD_FAIL;
       }
 #ifdef OMPTARGET_DEBUG
       void *TgtPtrBase = (void *)((intptr_t)TgtPtr + ArgOffset);
       DP("Allocated %" PRId64 " bytes of target memory at " DPxMOD
if (!FirstPrivateArgInfo.empty()) {
       auto Itr = FirstPrivateArgBuffer.begin();
       // Copy all host data to this buffer
       for (FirstPrivateArgInfoTy &Info : FirstPrivateArgInfo) {
         std::copy(Info.HstPtrBegin, Info.HstPtrEnd, Itr);
         Itr = std::next(Itr, Info.AlignedSize);
       }
       // Allocate target memory
       void *TgtPtr =
-          Device.allocData(FirstPrivateArgSize, FirstPrivateArgBuffer.data());
+          Device.allocData(FirstPrivateArgSize, FirstPrivateArgBuffer.data(),
+                           TARGET_ALLOC_DEFAULT, AsyncInfo);
       if (TgtPtr == nullptr) {
         DP("Failed to allocate target memory for private arguments.\n");
         return OFFLOAD_FAIL;
       }
       TgtPtrs.push_back(TgtPtr);
       DP("Allocated %" PRId64 " bytes of target memory at " DPxMOD "\n",
          FirstPrivateArgSize, DPxPTR(TgtPtr));
       // Transfer data to target device
int packAndTransfer(std::vector<void *> &TgtArgs) {
     }
 
     return OFFLOAD_SUCCESS;
   }
 
   /// Free all target memory allocated for private arguments
   int free() {
     for (void *P : TgtPtrs) {
-      int Ret = Device.deleteData(P);
+      int Ret = Device.deleteData(P, AsyncInfo);
       if (Ret != OFFLOAD_SUCCESS) {
         DP("Deallocation of (first-)private arrays failed.\n");
         return OFFLOAD_FAIL;
       }
     }
 
     TgtPtrs.clear();
 
Ret = processDataAfter(loc, DeviceId, HostPtr, ArgNum, ArgBases, Args,
                            ArgSizes, ArgTypes, ArgNames, ArgMappers,
                            PrivateArgumentManager, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Failed to process data after launching the kernel.\n");
       return OFFLOAD_FAIL;
     }
   }
 
-  return OFFLOAD_SUCCESS;
+  // Ensure we synchronize at the very end of a target directive. If we
+  // synchronized before this might be a no-op as the "queue" is a nullptr now.
+  return AsyncInfo.synchronize();
 }

openmp/libomptarget/src/rtl.cpp

 #endif
 
     DP("Registering RTL %s supporting %d devices!\n", R.RTLName.c_str(),
        R.NumberOfDevices);
 
     // Optional functions
     *((void **)&R.init_requires) =
         dlsym(dynlib_handle, "__tgt_rtl_init_requires");
+    *((void **)&R.data_alloc_async) =
+        dlsym(dynlib_handle, "__tgt_rtl_data_alloc_async");
+    *((void **)&R.data_delete_async) =
+        dlsym(dynlib_handle, "__tgt_rtl_data_delete_async");
     *((void **)&R.data_submit_async) =
         dlsym(dynlib_handle, "__tgt_rtl_data_submit_async");
     *((void **)&R.data_retrieve_async) =
         dlsym(dynlib_handle, "__tgt_rtl_data_retrieve_async");
     *((void **)&R.run_region_async) =
         dlsym(dynlib_handle, "__tgt_rtl_run_target_region_async");
     *((void **)&R.run_team_region_async) =
         dlsym(dynlib_handle, "__tgt_rtl_run_target_team_region_async");