[OpenMP][libomptarget] Implement memory lock/unlock API in NextGen plugins

[kevinsala]

Summary

This patch implements the memory lock/unlock API, introduced in patch https://reviews.llvm.org/D139208, in the NextGen plugins.

The patch also re-organizes the map of host pinned allocations in PluginInterface.

[tschuett]

https://github.com/openucx/ucx/blob/master/src/ucs/datastruct/pgtable.h super optimized page table for a similar use case.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

295

I would return std::optional<const EntryTy*> to avoid using nullptr to model failure.

328

I would return std::optional<void *> to avoid passing mutable pointers.

[jhuber6]

In D141227#4034338, @tschuett wrote:

https://github.com/openucx/ucx/blob/master/src/ucs/datastruct/pgtable.h super optimized page table for a similar use case.

That's interesting. We have another use inside of libomptarget where we need to determine if a pointer lies inside of an already mapped memory region. We could profile it and see if we could get some better performance with a more optimal implementation.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1832–1833

Nit, why not one line.

[jdoerfert]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

295

Why would we want that? It seems to me null is a fine, didn't find the entry response. We don't need to box things just to have boxed them.

328

This makes more sense. Though, again, why boxing. Returning the devptr content avoids mutable arguments just fine

[tschuett]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

295

Because these are C idoms in C++ code. LLVM is inconsistent whether returning true or false denotes success. The boxing overhead is negligible. The intent of your APIs is much clearer. tryFindIntersecting is a fallible API. òptional` is a great tool to for fallible APIs.

[jdoerfert]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

295

• Pointers/nullptr is arguably a C++ concept, we don't need to box things because C++ supports objects.
• I didn't argue about overheads, no need to start that discussion.
• There is no true or false return here so no confusion, it's EntryTy * that is nullptr or not.
• The API intent is to find an EntryTy, nullptr is arguably not an EntryTy object, hence none was found. How much clearer could it get if we would check an optional. Further, what would the set but nullptr case even mean?
• As far as I can see, our coding standards page does not include the word optional. If you want to have this as part of the canonical design, feel free to write an RFC. In the nextgen plugins we only use optional in the JIT, hence this is not an outlier.

[jdoerfert]

In D141227#4034338, @tschuett wrote:

https://github.com/openucx/ucx/blob/master/src/ucs/datastruct/pgtable.h super optimized page table for a similar use case.

The license doesn't look like something we could copy. If there is a paper or similar we could reimplement it I assume.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1829

Brief documentation seems to be used elsewhere, add it here and below.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

328

And remove the const cast at the use site. Cast it here if necessary.

[jdoerfert]

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1912

It's not const. Change isHostPinnedBuffer to getHostPinnedBuffer and assign it in the conditional.

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

498

It looks like we need https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#group__CUDART__MEMORY_1ge8d5c17670f16ac4fc8fcb4181cb490c
And https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#group__CUDART__MEMORY_1gc00502b44e5f1bdc0b424487ebb08db0

[kevinsala]

Fixing issues and format. Still missing registering the host memory in CUDA plugin.

[kevinsala]

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1912

Done, however, I renamed it to getDevicePtrFromPinnedBuffer() since it returns the device pointer.

[ye-luo]

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1912

DevicePtr usually refers to the memory on the device. I would call it DeviceAccessiblePtr

[jdoerfert]

LG, with 3 nits:

• Ye's comment.
• The one below.
• @carlo.bertolli needs to modify the plugin interface and if we land this first we need to remember to change it as the other patch makes it in. Otherwise we can wait with this one.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1907–1910

[carlo.bertolli]

I changed the old plugin interface for tgt_rtl_data_lock to return an error code. It now returns the lockedptr as function argument. Let me know if this is not what was called for.
Thanks for this extension!

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

584–585

Just a nit: for amdgpu's we don't need to keep a table of locked pointers. This is already done by ROCr. I would consider making this optional for amdgpu's.

[kevinsala]

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

584–585

You're right. But for the moment, we want to keep this information cached into a map at the plugin level. In this way, we can access the pointer info "faster" and keep track of the pinned memory buffers that are OpenMP-related.

[kevinsala]

Fixing pending comments, updating code documentation, and including API changes in https://reviews.llvm.org/D139208.

[kevinsala]

Still missing the pin/unpin calls in the CUDA plugin; it will be added in the following update of this patch.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

1912

Just renamed the function and related code documentation.

[JonChesterfield]

What is locked memory, and could we define it in a comment/documentation somewhere if it isn't already?

Is it the same as mmap LOCKED? That plus extra?

Does it imply GPUs can read/write it? If so, can they read/write within a kernel execution, or can we implement locked memory by copy at the start and end of the execution?

[carlo.bertolli]

In D141227#4041177, @JonChesterfield wrote:

What is locked memory, and could we define it in a comment/documentation somewhere if it isn't already?

Is it the same as mmap LOCKED? That plus extra?

Does it imply GPUs can read/write it? If so, can they read/write within a kernel execution, or can we implement locked memory by copy at the start and end of the execution?

Good question!
I always refer back to this article for a definition of locked/pinned:
https://lwn.net/Articles/600502/
but it might be stale, as my knowledge of OS is.
Aside from what I know or not know about OS, this is a kernel concept that we are reflecting in the plugin's.
In general, it is a concept used CUDA and HSA, so perhaps we can refer back to it?

We care about locking host memory because - roughly speaking - hsa_amd_memory_async_copy (asynchronous H2D and D2H memory copy) has a fast path for prelocked/pinned memory.
As locking memory is expensive, putting some "space" between lock operations, memory copies, and unlock operations is an effective optimization strategy. That's why we are giving users an API to prelock pointers when they know they can, and we avoid relocking them in the plugin's.
I hope this helps.

[carlo.bertolli]

In D141227#4041529, @carlo.bertolli wrote:

In D141227#4041177, @JonChesterfield wrote:

What is locked memory, and could we define it in a comment/documentation somewhere if it isn't already?

Is it the same as mmap LOCKED? That plus extra?

Does it imply GPUs can read/write it? If so, can they read/write within a kernel execution, or can we implement locked memory by copy at the start and end of the execution?

I forgot to answer this: a locked pointer is passed a set of agents that can access it. It is best described as a "agent accessible pointer". In the "first gen" plugin, we use the agent accessible pointer (via call to lock) to perform a memory_async_copy, passing in the GPU agent involved in the copy.

Good question!
I always refer back to this article for a definition of locked/pinned:
https://lwn.net/Articles/600502/
but it might be stale, as my knowledge of OS is.
Aside from what I know or not know about OS, this is a kernel concept that we are reflecting in the plugin's.
In general, it is a concept used CUDA and HSA, so perhaps we can refer back to it?

We care about locking host memory because - roughly speaking - hsa_amd_memory_async_copy (asynchronous H2D and D2H memory copy) has a fast path for prelocked/pinned memory.
As locking memory is expensive, putting some "space" between lock operations, memory copies, and unlock operations is an effective optimization strategy. That's why we are giving users an API to prelock pointers when they know they can, and we avoid relocking them in the plugin's.
I hope this helps.

[JonChesterfield]

For HSA we can either have coarse grain semantics on a locked pointer or fine grain depending on which hsa calls we make. That surfaces a user visible distinction - if their kernel writes to this memory, can the host see that write while the kernel executes? I don't know which answer is correct but would like us to choose one and document it, and probably have the same behaviour on nvptx.

[ye-luo]

Ping. What is left before merging?

[JonChesterfield]

I think it's good to go. If this is something you've tried in qmcpack so we have a real world example of it improving things then even better.

[ye-luo]

In D141227#4071198, @JonChesterfield wrote:

I think it's good to go. If this is something you've tried in qmcpack so we have a real world example of it improving things then even better.

QMCPACK doesn't allocate pinned memroy via OpenMP. So this patch doesn't immediately benefit QMCPACK. This patch lays the foundation to a further optimization and the 16 release branch will be created very soon. So I'm pushing for merging this in time.

[JonChesterfield]

If you don't mind keeping an eye on CI and reverting it on breakage feel free to land it, I'm away from my desk

[ye-luo]

This patch needs a rebase. Applying patch failed.

[kevinsala]

I'll update the patch today. I'm improving it to feature ref counting and certain overlapping of locked areas (partial overlapping with extension of an already locked area is forbidden).

[kevinsala]

Rebasing and adding support for ref counting on locked buffers and allow certain overlapping. Given an already locked buffer A, other buffers that are fully contained inside A can be locked, even if they are smaller than A. Extending an existing locked buffer is not allowed. The original region is unlocked once all its users have released the locked buffer and sub-buffers.

[ye-luo]

Ref-counting lock/unlock makes a lot sense.

[jhuber6]

If this fixes these failing with the nextgen plugins, please remove the line disabling the nextgen in the test.

Failed Tests (2):
  libomptarget :: amdgcn-amd-amdhsa :: mapping/prelock.cpp
  libomptarget :: amdgcn-amd-amdhsa-LTO :: mapping/prelock.cpp

[kevinsala]

In D141227#4078564, @jhuber6 wrote:

If this fixes these failing with the nextgen plugins, please remove the line disabling the nextgen in the test.

Failed Tests (2):
  libomptarget :: amdgcn-amd-amdhsa :: mapping/prelock.cpp
  libomptarget :: amdgcn-amd-amdhsa-LTO :: mapping/prelock.cpp

It doesn't fix these tests because they are apparently wrong. We cannot map a device accessible buffer, which is the agent_ptr returned when locking the buffer. (More info: https://reviews.llvm.org/D142399). We can fix it using is_device_ptr instead.

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

Error queryAsyncImpl(__tgt_async_info &AsyncInfo) override {
     // AsyncInfo. This is to make sure the synchronization only works for its
     // own tasks.
     AMDGPUStreamManager.returnResource(Stream);
     AsyncInfo.Queue = nullptr;
 
     return Plugin::success();
   }
 
+  /// Pin the host buffer and return the device pointer that should be used for
+  /// device transfers.
+  Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) override {
+    void *PinnedPtr = nullptr;
+
+    hsa_status_t Status =
+        hsa_amd_memory_lock(HstPtr, Size, nullptr, 0, &PinnedPtr);
+    if (auto Err = Plugin::check(Status, "Error in hsa_amd_memory_lock: %s\n"))
+      return Err;
+
+    return PinnedPtr;
+  }
+
+  /// Unpin the host buffer.
+  Error dataUnlockImpl(void *HstPtr) override {
+    hsa_status_t Status = hsa_amd_memory_unlock(HstPtr);
+    return Plugin::check(Status, "Error in hsa_amd_memory_unlock: %s\n");
+  }
+
   /// Submit data to the device (host to device transfer).
   Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
                        AsyncInfoWrapperTy &AsyncInfoWrapper) override {
-
     // Use one-step asynchronous operation when host memory is already pinned.
-    if (isHostPinnedMemoryBuffer(HstPtr)) {
+    if (void *PinnedPtr =
+            PinnedAllocs.getDeviceAccessiblePtrFromPinnedBuffer(HstPtr)) {
       AMDGPUStreamTy &Stream = getStream(AsyncInfoWrapper);
-      return Stream.pushPinnedMemoryCopyAsync(TgtPtr, HstPtr, Size);
+      return Stream.pushPinnedMemoryCopyAsync(TgtPtr, PinnedPtr, Size);
     }
 
     void *PinnedHstPtr = nullptr;
 
     // For large transfers use synchronous behavior.
     if (Size >= OMPX_MaxAsyncCopyBytes) {
       if (AsyncInfoWrapper.hasQueue())
         if (auto Err = synchronize(AsyncInfoWrapper))
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
     AMDGPUStreamTy &Stream = getStream(AsyncInfoWrapper);
     return Stream.pushMemoryCopyH2DAsync(TgtPtr, HstPtr, PinnedHstPtr, Size,
                                          PinnedMemoryManager);
   }
 
   /// Retrieve data from the device (device to host transfer).
   Error dataRetrieveImpl(void *HstPtr, const void *TgtPtr, int64_t Size,
                          AsyncInfoWrapperTy &AsyncInfoWrapper) override {
 
     // Use one-step asynchronous operation when host memory is already pinned.
-    if (isHostPinnedMemoryBuffer(HstPtr)) {
-      // Use one-step asynchronous operation when host memory is already pinned.
+    if (void *PinnedPtr =
+            PinnedAllocs.getDeviceAccessiblePtrFromPinnedBuffer(HstPtr)) {
       AMDGPUStreamTy &Stream = getStream(AsyncInfoWrapper);
-      return Stream.pushPinnedMemoryCopyAsync(HstPtr, TgtPtr, Size);
+      return Stream.pushPinnedMemoryCopyAsync(PinnedPtr, TgtPtr, Size);
     }
 
     void *PinnedHstPtr = nullptr;
 
     // For large transfers use synchronous behavior.
     if (Size >= OMPX_MaxAsyncCopyBytes) {
       if (AsyncInfoWrapper.hasQueue())
         if (auto Err = synchronize(AsyncInfoWrapper))

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

 protected:
   /// The preferred number of threads to run the kernel.
   uint32_t PreferredNumThreads;
 
   /// The maximum number of threads which the kernel could leverage.
   uint32_t MaxNumThreads;
 };
 
+/// Class representing a map of host pinned allocations. We track these pinned
+/// allocations, so memory tranfers invloving these buffers can be optimized.
+class PinnedAllocationMapTy {
+
+  /// Struct representing a map entry.
+  struct EntryTy {
+    /// The host pointer of the pinned allocation.
+    void *HstPtr;
+
+    /// The pointer that devices' driver should use to transfer data from/to the
+    /// pinned allocation. In most plugins, this pointer will be the same as the
+    /// host pointer above.
+    void *DevAccessiblePtr;
+
+    /// The size of the pinned allocation.
+    size_t Size;
+
+    /// The number of references to the pinned allocation. The allocation should
+    /// remain pinned and registered to the map until the number of references
+    /// becomes zero.
+    mutable size_t References;
+
+    /// Create an entry with the host and device acessible pointers, and the
+    /// buffer size.
+    EntryTy(void *HstPtr, void *DevAccessiblePtr, size_t Size)
+        : HstPtr(HstPtr), DevAccessiblePtr(DevAccessiblePtr), Size(Size),
+          References(1) {}
+
+    /// Utility constructor used for std::set searches.
+    EntryTy(void *HstPtr)
+        : HstPtr(HstPtr), DevAccessiblePtr(nullptr), Size(0), References(0) {}
+  };
+
+  /// Comparator of mep entries. Use the host pointer to enforce an order
+  /// between entries.
+  struct EntryCmpTy {
+    bool operator()(const EntryTy &Left, const EntryTy &Right) const {
+      return Left.HstPtr < Right.HstPtr;
+    }
+  };
+
+  typedef std::set<EntryTy, EntryCmpTy> PinnedAllocSetTy;
+
+  /// The map of host pinned allocations.
+  PinnedAllocSetTy Allocs;
+
+  /// The mutex to protect accesses to the map.
+  mutable std::shared_mutex Mutex;
+
+  /// Reference to the corresponding device.
+  GenericDeviceTy &Device;
+
+  /// Find an allocation that intersects with \p Buffer pointer. Assume
+  /// the map's mutex is acquired.
+  PinnedAllocSetTy::iterator findIntersecting(const void *Buffer) const {
+    if (Allocs.empty())
+      return Allocs.end();
+
+    // Search the first allocation with starting address that is not less than
+    // the buffer address.
+    auto It = Allocs.lower_bound({const_cast<void *>(Buffer)});
+
+    // Direct match of starting addresses.
+    if (It != Allocs.end() && It->HstPtr == Buffer)
+      return It;
+
+    // Not direct match but may be a previous pinned allocation in the map which
+    // contains the buffer. Return false if there is no such a previous
+    // allocation.
+    if (It == Allocs.begin())
+      return Allocs.end();
+
+    // Move to the previous pinned allocation.
+    --It;
+
+    // The buffer is not contained in the pinned allocation.
+    if (advanceVoidPtr(It->HstPtr, It->Size) > Buffer)
+      return It;
+
+    // None found.
+    return Allocs.end();
+  }
+
+public:
+  /// Create the map of pinned allocations corresponding to a specific device.
+  PinnedAllocationMapTy(GenericDeviceTy &Device) : Device(Device) {}
+
+  /// Register a host buffer that was recently locked. None of the already
+  /// registered pinned allocations should intersect with this new one. The
+  /// registration requires the host pointer in \p HstPtr, the pointer that the
+  /// devices should use when transferring data from/to the allocation in
+  /// \p DevAccessiblePtr, and the size of the allocation in \p Size. Notice
+  /// that some plugins may use the same pointer for the \p HstPtr and
+  /// \p DevAccessiblePtr. The allocation must be unregistered using the
+  /// unregisterHostBuffer function.
+  Error registerHostBuffer(void *HstPtr, void *DevAccessiblePtr, size_t Size);
+
+  /// Unregister a host pinned allocation passing the host pointer which was
+  /// previously registered using the registerHostBuffer function. When calling
+  /// this function, the pinned allocation cannot have any other user.
+  Error unregisterHostBuffer(void *HstPtr);
+
+  /// Lock the host buffer at \p HstPtr or register a new user if it intersects
+  /// with an already existing one. A partial overlapping with extension is not
+  /// allowed. The function returns the device accessible pointer of the pinned
+  /// buffer. The buffer must be unlocked using the unlockHostBuffer function.
+  Expected<void *> lockHostBuffer(void *HstPtr, size_t Size);
+
+  /// Unlock the host buffer at \p HstPtr or unregister a user if other users
+  /// are still using the pinned allocation. If this was the last user, the
+  /// pinned allocation is removed from the map and the memory is unlocked.
+  Error unlockHostBuffer(void *HstPtr);
+
+  /// Return the device accessible pointer associated to the host pinned
+  /// allocation which the \p HstPtr belongs, if any. Return null in case the
+  /// \p HstPtr does not belong to any host pinned allocation. The device
+  /// accessible pointer is the one that devices should use for data transfers
+  /// that involve a host pinned buffer.
+  void *getDeviceAccessiblePtrFromPinnedBuffer(const void *HstPtr) const {
+    std::shared_lock<std::shared_mutex> Lock(Mutex);
+
+    // Find the intersecting allocation if any.
+    auto It = findIntersecting(HstPtr);
+    if (It == Allocs.end())
+      return nullptr;
+
+    const EntryTy &Entry = *It;
+    return advanceVoidPtr(Entry.DevAccessiblePtr,
+                          getPtrDiff(HstPtr, Entry.HstPtr));
+  }
+
+  /// Check whether a buffer belongs to a registered host pinned allocation.
+  bool isHostPinnedBuffer(const void *HstPtr) const {
+    std::shared_lock<std::shared_mutex> Lock(Mutex);
+
+    // Return whether there is an intersecting allocation.
+    return (findIntersecting(const_cast<void *>(HstPtr)) != Allocs.end());
+  }
+};
+
 /// Class implementing common functionalities of offload devices. Each plugin
 /// should define the specific device class, derive from this generic one, and
 /// implement the necessary virtual function members.
 struct GenericDeviceTy : public DeviceAllocatorTy {
   /// Construct a device with its device id within the plugin, the number of
   /// devices in the plugin and the grid values for that kind of device.
   GenericDeviceTy(int32_t DeviceId, int32_t NumDevices,
                   const llvm::omp::GV &GridValues);
struct GenericDeviceTy : public DeviceAllocatorTy {
   virtual Error queryAsyncImpl(__tgt_async_info &AsyncInfo) = 0;
 
   /// Allocate data on the device or involving the device.
   Expected<void *> dataAlloc(int64_t Size, void *HostPtr, TargetAllocTy Kind);
 
   /// Deallocate data from the device or involving the device.
   Error dataDelete(void *TgtPtr, TargetAllocTy Kind);
 
+  /// Pin host memory to optimize transfers and return the device accessible
+  /// pointer that devices should use for memory transfers involving the host
+  /// pinned allocation.
+  Expected<void *> dataLock(void *HstPtr, int64_t Size) {
+    return PinnedAllocs.lockHostBuffer(HstPtr, Size);
+  }
+
+  virtual Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) = 0;
+
+  /// Unpin a host memory buffer that was previously pinned.
+  Error dataUnlock(void *HstPtr) {
+    return PinnedAllocs.unlockHostBuffer(HstPtr);
+  }
+
+  virtual Error dataUnlockImpl(void *HstPtr) = 0;
+
   /// Submit data to the device (host to device transfer).
   Error dataSubmit(void *TgtPtr, const void *HstPtr, int64_t Size,
                    __tgt_async_info *AsyncInfo);
   virtual Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
                                AsyncInfoWrapperTy &AsyncInfoWrapper) = 0;
 
   /// Retrieve data from the device (device to host transfer).
   Error dataRetrieve(void *HstPtr, const void *TgtPtr, int64_t Size,
private:
   virtual Error getDeviceHeapSize(uint64_t &V) = 0;
   virtual Error setDeviceHeapSize(uint64_t V) = 0;
 
   /// Indicate whether the device should setup the device environment. Notice
   /// that returning false in this function will change the behavior of the
   /// setupDeviceEnvironment() function.
   virtual bool shouldSetupDeviceEnvironment() const { return true; }
 
-  /// Register a host buffer as host pinned allocation.
-  Error registerHostPinnedMemoryBuffer(const void *Buffer, size_t Size);
-
-  /// Unregister a host pinned allocations.
-  Error unregisterHostPinnedMemoryBuffer(const void *Buffer);
-
   /// Pointer to the memory manager or nullptr if not available.
   MemoryManagerTy *MemoryManager;
 
   /// Environment variables defined by the OpenMP standard.
   Int32Envar OMP_TeamLimit;
   Int32Envar OMP_NumTeams;
   Int32Envar OMP_TeamsThreadLimit;
 
   /// Environment variables defined by the LLVM OpenMP implementation.
   Int32Envar OMPX_DebugKind;
   UInt32Envar OMPX_SharedMemorySize;
   UInt64Envar OMPX_TargetStackSize;
   UInt64Envar OMPX_TargetHeapSize;
 
-  /// Map of host pinned allocations. We track these pinned allocations so that
-  /// memory transfers involving these allocations can be optimized.
-  std::map<const void *, size_t> HostAllocations;
-  mutable std::shared_mutex HostAllocationsMutex;
-
 protected:
-  /// Check whether a buffer has been registered as host pinned memory.
-  bool isHostPinnedMemoryBuffer(const void *Buffer) const {
-    std::shared_lock<std::shared_mutex> Lock(HostAllocationsMutex);
-
-    if (HostAllocations.empty())
-      return false;
-
-    // Search the first allocation with starting address that is not less than
-    // the buffer address.
-    auto It = HostAllocations.lower_bound(Buffer);
-
-    // Direct match of starting addresses.
-    if (It != HostAllocations.end() && It->first == Buffer)
-      return true;
-
-    // Not direct match but may be a previous pinned allocation in the map which
-    // contains the buffer. Return false if there is no such a previous
-    // allocation.
-    if (It == HostAllocations.begin())
-      return false;
-
-    // Move to the previous pinned allocation.
-    --It;
-
-    // Evaluate whether the buffer is contained in the pinned allocation.
-    return (advanceVoidPtr(It->first, It->second) > (const char *)Buffer);
-  }
-
   /// Return the execution mode used for kernel \p Name.
   Expected<OMPTgtExecModeFlags> getExecutionModeForKernel(StringRef Name,
                                                           DeviceImageTy &Image);
 
   /// Environment variables defined by the LLVM OpenMP implementation
   /// regarding the initial number of streams and events.
   UInt32Envar OMPX_InitialNumStreams;
   UInt32Envar OMPX_InitialNumEvents;
protected:
   enum class PeerAccessState : uint8_t { AVAILABLE, UNAVAILABLE, PENDING };
 
   /// Array of peer access states with the rest of devices. This means that if
   /// the device I has a matrix PeerAccesses with PeerAccesses[J] == AVAILABLE,
   /// the device I can access device J's memory directly. However, notice this
   /// does not mean that device J can access device I's memory directly.
   llvm::SmallVector<PeerAccessState> PeerAccesses;
   std::mutex PeerAccessesLock;
+
+  /// Map of host pinned allocations used for optimize device transfers.
+  PinnedAllocationMapTy PinnedAllocs;
 };
 
 /// Class implementing common functionalities of offload plugins. Each plugin
 /// should define the specific plugin class, derive from this generic one, and
 /// implement the necessary virtual function members.
 struct GenericPluginTy {
 
   /// Construct a plugin instance.

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

: MemoryManager(nullptr), OMP_TeamLimit("OMP_TEAM_LIMIT"),
       // Do not initialize the following two envars since they depend on the
       // device initialization. These cannot be consulted until the device is
       // initialized correctly. We intialize them in GenericDeviceTy::init().
       OMPX_TargetStackSize(), OMPX_TargetHeapSize(),
       // By default, the initial number of streams and events are 32.
       OMPX_InitialNumStreams("LIBOMPTARGET_NUM_INITIAL_STREAMS", 32),
       OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 32),
       DeviceId(DeviceId), GridValues(OMPGridValues),
-      PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock() {
+      PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
+      PinnedAllocs(*this) {
   if (OMP_NumTeams > 0)
     GridValues.GV_Max_Teams =
         std::min(GridValues.GV_Max_Teams, uint32_t(OMP_NumTeams));
 
   if (OMP_TeamsThreadLimit > 0)
     GridValues.GV_Max_WG_Size =
         std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
 }
GenericDeviceTy::getExecutionModeForKernel(StringRef Name,
 
   // Check that the retrieved execution mode is valid.
   if (!GenericKernelTy::isValidExecutionMode(ExecModeGlobal.getValue()))
     return Plugin::error("Invalid execution mode %d for '%s'",
                          ExecModeGlobal.getValue(), Name.data());
 
   return ExecModeGlobal.getValue();
 }
 
-Error GenericDeviceTy::registerHostPinnedMemoryBuffer(const void *Buffer,
-                                                      size_t Size) {
-  std::lock_guard<std::shared_mutex> Lock(HostAllocationsMutex);
+Error PinnedAllocationMapTy::registerHostBuffer(void *HstPtr,
+                                                void *DevAccessiblePtr,
+  assert(HstPtr && "Invalid pointer");
+  assert(DevAccessiblePtr && "Invalid pointer");
 
-  auto Res = HostAllocations.insert({Buffer, Size});
+  std::lock_guard<std::shared_mutex> Lock(Mutex);
+
+  // No pinned allocation should intersect.
+  auto Res = Allocs.insert({HstPtr, DevAccessiblePtr, Size});
   if (!Res.second)
-    return Plugin::error("Registering an already registered pinned buffer");
+    return Plugin::error("Cannot register locked buffer");
 
   return Plugin::success();
 }
 
-Error GenericDeviceTy::unregisterHostPinnedMemoryBuffer(const void *Buffer) {
-  std::lock_guard<std::shared_mutex> Lock(HostAllocationsMutex);
+Error PinnedAllocationMapTy::unregisterHostBuffer(void *HstPtr) {
+  assert(HstPtr && "Invalid pointer");
+
+  std::lock_guard<std::shared_mutex> Lock(Mutex);
+
+  // Find the pinned allocation starting at the host pointer address.
+  auto It = Allocs.find({HstPtr});
+  if (It == Allocs.end())
+    return Plugin::error("Cannot find locked buffer");
+
+  const EntryTy &Entry = *It;
+
+  // There should be no other references to the pinned allocation.
+  if (Entry.References > 1)
+    return Plugin::error("The locked buffer is still being used");
+
+  // Remove the entry from the map.
+  Allocs.erase(It);
+
+  return Plugin::success();
+}
+
+Expected<void *> PinnedAllocationMapTy::lockHostBuffer(void *HstPtr,
+                                                       size_t Size) {
+  assert(HstPtr && "Invalid pointer");
+
+  std::lock_guard<std::shared_mutex> Lock(Mutex);
+
+  auto It = findIntersecting(HstPtr);
+
+  // No intersecting registered allocation found in the map. We must lock and
+  // register the memory buffer into the map.
+  if (It == Allocs.end()) {
+    // First, lock the host buffer and retrieve the device accessible pointer.
+    auto PinnedPtrOrErr = Device.dataLockImpl(HstPtr, Size);
+    if (!PinnedPtrOrErr)
+      return PinnedPtrOrErr.takeError();
+
+    // Then, insert the host buffer entry into the map.
+    auto Res = Allocs.insert({HstPtr, *PinnedPtrOrErr, Size});
+    if (!Res.second)
+      return Plugin::error("Cannot register locked buffer");
+
+    // Return the device accessible pointer.
+    return *PinnedPtrOrErr;
+  }
+
+  const EntryTy &Entry = *It;
+
+#ifdef OMPTARGET_DEBUG
+  // Do not allow partial overlapping among host pinned buffers.
+  if (advanceVoidPtr(HstPtr, Size) > advanceVoidPtr(Entry.HstPtr, Entry.Size))
+    return Plugin::error("Partial overlapping not allowed in locked memory");
+#endif
+
+  // Increase the number of references.
+  Entry.References++;
+
+  // Return the device accessible pointer after applying the correct offset.
+  return advanceVoidPtr(Entry.DevAccessiblePtr,
+                        getPtrDiff(HstPtr, Entry.HstPtr));
+}
+
+Error PinnedAllocationMapTy::unlockHostBuffer(void *HstPtr) {
+  assert(HstPtr && "Invalid pointer");
 
-  size_t Erased = HostAllocations.erase(Buffer);
+  std::lock_guard<std::shared_mutex> Lock(Mutex);
+
+  auto It = findIntersecting(HstPtr);
+  if (It == Allocs.end())
+    return Plugin::error("Cannot find locked buffer");
+
+  const EntryTy &Entry = *It;
+
+  // Decrease the number of references. No need to do anything if there are
+  // others using the allocation.
+  if (--Entry.References > 0)
+    return Plugin::success();
+
+  // This was the last user of the allocation. Unlock the original locked memory
+  // buffer, which is the host pointer stored in the entry.
+  if (auto Err = Device.dataUnlockImpl(Entry.HstPtr))
+    return Err;
+
+  // Remove the entry from the map.
+  size_t Erased = Allocs.erase(Entry);
   if (!Erased)
-    return Plugin::error("Cannot find a registered host pinned buffer");
+    return Plugin::error("Cannot find locked buffer");
 
   return Plugin::success();
 }
 
 Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo) {
   if (!AsyncInfo || !AsyncInfo->Queue)
     return Plugin::error("Invalid async info queue");
 
Expected<void *> GenericDeviceTy::dataAlloc(int64_t Size, void *HostPtr,
   // Report error if the memory manager or the device allocator did not return
   // any memory buffer.
   if (!Alloc)
     return Plugin::error("Invalid target data allocation kind or requested "
                          "allocator not implemented yet");
 
   // Register allocated buffer as pinned memory if the type is host memory.
   if (Kind == TARGET_ALLOC_HOST)
-    if (auto Err = registerHostPinnedMemoryBuffer(Alloc, Size))
+    if (auto Err = PinnedAllocs.registerHostBuffer(Alloc, Alloc, Size))
       return Err;
 
   return Alloc;
 }
 
 Error GenericDeviceTy::dataDelete(void *TgtPtr, TargetAllocTy Kind) {
   // Free is a noop when recording or replaying.
   if (RecordReplay.isRecordingOrReplaying())
     return Plugin::success();
 
   int Res;
   if (MemoryManager)
     Res = MemoryManager->free(TgtPtr);
   else
     Res = free(TgtPtr, Kind);
 
   if (Res)
     return Plugin::error("Failure to deallocate device pointer %p", TgtPtr);
 
   // Unregister deallocated pinned memory buffer if the type is host memory.
   if (Kind == TARGET_ALLOC_HOST)
-    if (auto Err = unregisterHostPinnedMemoryBuffer(TgtPtr))
+    if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
       return Err;
 
   return Plugin::success();
 }
 
 Error GenericDeviceTy::dataSubmit(void *TgtPtr, const void *HstPtr,
                                   int64_t Size, __tgt_async_info *AsyncInfo) {
   auto Err = Plugin::success();
if (Err) {
     REPORT("Failure to deallocate device pointer %p: %s\n", TgtPtr,
            toString(std::move(Err)).data());
     return OFFLOAD_FAIL;
   }
 
   return OFFLOAD_SUCCESS;
 }
 
+int32_t __tgt_rtl_data_lock(int32_t DeviceId, void *Ptr, int64_t Size,
+                            void **LockedPtr) {
+  auto LockedPtrOrErr = Plugin::get().getDevice(DeviceId).dataLock(Ptr, Size);
+  if (!LockedPtrOrErr) {
+    auto Err = LockedPtrOrErr.takeError();
+    REPORT("Failure to lock memory %p: %s\n", Ptr,
+           toString(std::move(Err)).data());
+    return OFFLOAD_FAIL;
+  }
+
+  if (!(*LockedPtrOrErr)) {
+    REPORT("Failure to lock memory %p: obtained a null locked pointer\n", Ptr);
+    return OFFLOAD_FAIL;
+  }
+  *LockedPtr = *LockedPtrOrErr;
+
+  return OFFLOAD_SUCCESS;
+}
+
+int32_t __tgt_rtl_data_unlock(int32_t DeviceId, void *Ptr) {
+  auto Err = Plugin::get().getDevice(DeviceId).dataUnlock(Ptr);
+  if (Err) {
+    REPORT("Failure to unlock memory %p: %s\n", Ptr,
+           toString(std::move(Err)).data());
+    return OFFLOAD_FAIL;
+  }
+
+  return OFFLOAD_SUCCESS;
+}
+
 int32_t __tgt_rtl_data_submit(int32_t DeviceId, void *TgtPtr, void *HstPtr,
                               int64_t Size) {
   return __tgt_rtl_data_submit_async(DeviceId, TgtPtr, HstPtr, Size,
                                      /* AsyncInfoPtr */ nullptr);
 }
 
 int32_t __tgt_rtl_data_submit_async(int32_t DeviceId, void *TgtPtr,
                                     void *HstPtr, int64_t Size,

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

Error queryAsyncImpl(__tgt_async_info &AsyncInfo) override {
     // occurs), return it to stream pool and reset AsyncInfo. This is to make
     // sure the synchronization only works for its own tasks.
     CUDAStreamManager.returnResource(Stream);
     AsyncInfo.Queue = nullptr;
 
     return Plugin::check(Res, "Error in cuStreamQuery: %s");
   }
 
+  Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) override {
+    // TODO: Register the buffer as CUDA host memory.
+    return HstPtr;
+  }
+
+  Error dataUnlockImpl(void *HstPtr) override { return Plugin::success(); }
+
   /// Submit data to the device (host to device transfer).
   Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
                        AsyncInfoWrapperTy &AsyncInfoWrapper) override {
     if (auto Err = setContext())
       return Err;
 
     CUstream Stream = getStream(AsyncInfoWrapper);
     if (!Stream)

openmp/libomptarget/plugins-nextgen/generic-elf-64bit/src/rtl.cpp

struct GenELF64DeviceTy : public GenericDeviceTy {
   }
 
   /// Free the memory. Use std::free in all cases.
   int free(void *TgtPtr, TargetAllocTy Kind) override {
     std::free(TgtPtr);
     return OFFLOAD_SUCCESS;
   }
 
+  /// This plugin does nothing to lock buffers. Do not return an error, just
+  /// return the same pointer as the device pointer.
+  Expected<void *> dataLockImpl(void *HstPtr, int64_t Size) override {
+    return HstPtr;
+  }
+
+  /// Nothing to do when unlocking the buffer.
+  Error dataUnlockImpl(void *HstPtr) override { return Plugin::success(); }
+
   /// Submit data to the device (host to device transfer).
   Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
                        AsyncInfoWrapperTy &AsyncInfoWrapper) override {
     std::memcpy(TgtPtr, HstPtr, Size);
     return Plugin::success();
   }
 
   /// Retrieve data from the device (device to host transfer).