[OpenMP] Remove shadow pointer map and introduce consistent locking

[jdoerfert]

Summary

The shadow pointer map was problematic as we scanned an entire list if
an entry had shadow pointers. The new scheme stores the shadow pointers
inside the entries. This allows easy access without any search. It also
helps us, but also makes it necessary, to define a consistent locking
scheme. The implicit locking of entries via TargetPointerResultTy makes
this pretty effortless, however one has to:

• Lock HDTTMap before locking an entry.
• Do not lock HDTTMap while holding an entry lock.
• Hold the entry lock to read or modify an entry.

The changes to submitData and retrieveData have been made to ensure 2
when the LIBOMPTARGET_INFO flag is used. Most everything else happens by
itself as TargetPointerResultTy acts as a lock_guard for the entry. It
is a little complicated when we deal with multiple entries, especially
as they can be equal. However, one can still follow the rules with
reasonable effort.

LookupResult are now finally also looking the entry before it is
inspected. This is good even if we haven't run into a problem yet.

[jdoerfert]

Improve comment

[tianshilei1992]

openmp/libomptarget/include/device.h

101

We don't use the order here, right? If so, unordered_map can usually get better performance.

[jdoerfert]

openmp/libomptarget/include/device.h

101

We do not use the order, correct. I can change it though I doubt it makes much of a difference.

[tianshilei1992]

openmp/libomptarget/include/device.h

101

Would be better to do it. https://stackoverflow.com/questions/36392583/is-an-unordered-map-really-faster-than-a-map-in-practice

[jdoerfert]

openmp/libomptarget/include/device.h

101

That comparison is not helpful. Our expected size is *probably* < 10.

[tianshilei1992]

openmp/libomptarget/include/device.h

101

It also says:

TL;DR

Results: the unordered_map shows its superiority as soon as there is data in the map. The only time it exhibits worse performance than the ordered map is when the maps are empty.

The answer also shows the search speed is different, not just insertion speed.

Actually, based on all the uses of ShadowPtrInfos, it doesn't have to be a map at all because we never need the key-value relation. It can be simply an unordered_set<ShadowPtrInfoTy> hashed by ShadowPtrInfoTy::HstPtrAddr. The only place we need the "hash" feature is when adding an element.

[ye-luo]

The new scheme stores the shadow pointers inside the entries.

Makes a lot more sense now.

openmp/libomptarget/include/device.h

378

I don't get the intention of this function. Need documentation. It also seems buggy. Don't you need to unlock Entry before assign?

388

Please delete copy constructor and copy assign operator.

415

Do you know what this is? Add some documentation?

[ye-luo]

I'm afraid you are abusing existing mutex. Please be specific about what part of the Entry you are trying to protect with TPR lock/unlock.

openmp/libomptarget/include/device.h

375

void lock() const { States->UpdateMtx.lock(); }
void unlock() const { States->UpdateMtx.unlock(); }

the mutex is only intended to protect data transfer submission.

[jdoerfert]

openmp/libomptarget/include/device.h

375

The mutex is now used to protect the entire entry. I will update the comment.
While the data transfer is running the entry should not be looked at, similarly, while anything might modify the entry, the entry should not be looked at. Hence, one can only get the entry locked.

378

It's not buggy, IMHO. Assignment will swap the entry, which may or may not be locked. The new object goes out of scope immediately which will trigger the destructor. If there is an entry it will be unlocked, if there was none, all is already good.

This function is needed to give up ownership of an entry early, that is while the TPR is still live but effectively not needed anymore.

388

Aren't they deleted anyway? I can make it explicit, sure.

415

We do know what this is. And we should add some documentation. We should do so in a separate patch though.

[ye-luo]

openmp/libomptarget/include/device.h

378

You are right. The new object will be responsible to call the unlock via the destructor. So this is correct.

388

I'm not sure if they are deleted by default. So helpful to make them deleted.

415

Then do it separately. Probably also for PendingCtorsDtorsPerLibrary PendingCtorsDtors;

[jdoerfert]

Fix rebase issues, address comments

[jdoerfert]

Use set instead of map

[tianshilei1992]

LGTM. Like @ye-luo suggested, we need to update the documents of lock and unlock as they are being used in a different way now.

[jdoerfert]

Rebase, addressed comments, incorporated most parts of https://reviews.llvm.org/D123443.
The mapping stress tests slows down a little but running all AMD GPU tests is as fast as before.
Consistent locking of the entries should be worth it, and the stress test is unlikely anyway.
The set vs unordered set situation is resolved with a SmallSet of size 2, unlikely we attach more pointers into an object.

[ye-luo]

Got compiler warning.

/scratch3/opt/llvm-clang/llvm-project/openmp/libomptarget/src/omptarget.cpp: In lambda function:
/scratch3/opt/llvm-clang/llvm-project/openmp/libomptarget/src/omptarget.cpp:1058:56: warning: cast from type 'void* const*' to type 'void*' casts away qualifiers [-Wcast-qual]
 1058 |                                     (void *)&ShadowPtr.TgtPtrVal,
      |

openmp/libomptarget/include/device.h

 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_DEVICE_H
 #define _OMPTARGET_DEVICE_H
 
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <list>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <set>
 #include <thread>
 
 #include "ExclusiveAccess.h"
 #include "omptarget.h"
 #include "rtl.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 // Forward declarations.
 struct RTLInfoTy;
 struct __tgt_bin_desc;
 struct __tgt_target_table;
 
 using map_var_info_t = void *;
 
 // enum for OMP_TARGET_OFFLOAD; keep in sync with kmp.h definition
 enum kmp_target_offload_kind {
   tgt_disabled = 0,
   tgt_default = 1,
   tgt_mandatory = 2
 };
 typedef enum kmp_target_offload_kind kmp_target_offload_kind_t;
 
+/// Information about shadow pointers.
+struct ShadowPtrInfoTy {
+  void **HstPtrAddr = nullptr;
+  void *HstPtrVal = nullptr;
+  void **TgtPtrAddr = nullptr;
+  void *TgtPtrVal = nullptr;
+
+  bool operator==(const ShadowPtrInfoTy &Other) const {
+    return HstPtrAddr == Other.HstPtrAddr;
+  }
+};
+
+inline bool operator<(const ShadowPtrInfoTy &lhs, const ShadowPtrInfoTy &rhs) {
+  return lhs.HstPtrAddr < rhs.HstPtrAddr;
+}
+
 /// Map between host data and target data.
 struct HostDataToTargetTy {
   const uintptr_t HstPtrBase; // host info.
   const uintptr_t HstPtrBegin;
   const uintptr_t HstPtrEnd;       // non-inclusive.
   const map_var_info_t HstPtrName; // Optional source name of mapped variable.
 
   const uintptr_t TgtPtrBegin; // target info.
 
 private:
   static const uint64_t INFRefCount = ~(uint64_t)0;
   static std::string refCountToStr(uint64_t RefCount) {
     return RefCount == INFRefCount ? "INF" : std::to_string(RefCount);
   }
 
   struct StatesTy {
     StatesTy(uint64_t DRC, uint64_t HRC)
-        : DynRefCount(DRC), HoldRefCount(HRC),
-          MayContainAttachedPointers(false) {}
+        : DynRefCount(DRC), HoldRefCount(HRC) {}
     /// The dynamic reference count is the standard reference count as of OpenMP
     /// 4.5.  The hold reference count is an OpenMP extension for the sake of
     /// OpenACC support.
     ///
     /// The 'ompx_hold' map type modifier is permitted only on "omp target" and
     /// "omp target data", and "delete" is permitted only on "omp target exit
     /// data" and associated runtime library routines.  As a result, we really
     /// need to implement "reset" functionality only for the dynamic reference
     /// counter.  Likewise, only the dynamic reference count can be infinite
     /// because, for example, omp_target_associate_ptr and "omp declare target
     /// link" operate only on it.  Nevertheless, it's actually easier to follow
     /// the code (and requires less assertions for special cases) when we just
     /// implement these features generally across both reference counters here.
     /// Thus, it's the users of this class that impose those restrictions.
     ///
     uint64_t DynRefCount;
     uint64_t HoldRefCount;
 
-    /// Boolean flag to remember if any subpart of the mapped region might be
-    /// an attached pointer.
-    bool MayContainAttachedPointers;
-
-    /// This mutex will be locked when data movement is issued. For targets that
-    /// doesn't support async data movement, this mutex can guarantee that after
-    /// it is released, memory region on the target is update to date. For
-    /// targets that support async data movement, this can guarantee that data
-    /// movement has been issued. This mutex *must* be locked right before
-    /// releasing the mapping table lock.
-    std::mutex UpdateMtx;
+    /// A map of shadow pointers associated with this entry, the keys are host
+    /// pointer addresses to identify stale entries.
+    llvm::SmallSet<ShadowPtrInfoTy, 2> ShadowPtrInfos;
+
     /// Pointer to the event corresponding to the data update of this map.
     /// Note: At present this event is created when the first data transfer from
     /// host to device is issued, and only being used for H2D. It is not used
     /// for data transfer in another direction (device to host). It is still
     /// unclear whether we need it for D2H. If in the future we need similar
     /// mechanism for D2H, and if the event cannot be shared between them, Event
     /// should be written as <tt>void *Event[2]</tt>.
     void *Event = nullptr;
uint64_t OtherRefCount =
         UseHoldRefCount ? States->DynRefCount : States->HoldRefCount;
     if (OtherRefCount > 0)
       return false;
     if (AfterReset)
       return ThisRefCount != INFRefCount;
     return ThisRefCount == 1;
   }
 
-  void setMayContainAttachedPointers() const {
-    States->MayContainAttachedPointers = true;
-  }
-  bool getMayContainAttachedPointers() const {
-    return States->MayContainAttachedPointers;
+  /// Add the shadow pointer info \p ShadowPtrInfo to this entry but only if the
+  /// the target ptr value was not already present in the existing set of shadow
+  /// pointers. Return true if something was added.
+  bool addShadowPointer(const ShadowPtrInfoTy &ShadowPtrInfo) const {
+    auto Pair = States->ShadowPtrInfos.insert(ShadowPtrInfo);
+    if (Pair.second)
+      return true;
+    // Check for a stale entry, if found, replace the old one.
+    if ((*Pair.first).TgtPtrVal == ShadowPtrInfo.TgtPtrVal)
+      return false;
+    States->ShadowPtrInfos.erase(ShadowPtrInfo);
+    return addShadowPointer(ShadowPtrInfo);
   }
 
-  void lock() const { States->UpdateMtx.lock(); }
+  /// Apply \p CB to all shadow pointers of this entry. Returns OFFLOAD_FAIL if
+  /// \p CB returned OFFLOAD_FAIL for any of them, otherwise this returns
+  /// OFFLOAD_SUCCESS. The entry is locked for this operation.
+  template <typename CBTy> int foreachShadowPointerInfo(CBTy CB) const {
+    for (auto &It : States->ShadowPtrInfos)
+      if (CB(It) == OFFLOAD_FAIL)
+        return OFFLOAD_FAIL;
+    return OFFLOAD_SUCCESS;
+  }
+
+  /// Lock this entry for exclusive access. Ensure to get exclusive access to
+  /// HDTTMap first!
+  void lock() const { Mtx.lock(); }
 
-  void unlock() const { States->UpdateMtx.unlock(); }
+  /// Unlock this entry to allow other threads inspecting it.
+  void unlock() const { Mtx.unlock(); }
+
+private:
+  // Mutex that needs to be held before the entry is inspected or modified. The
+  // HDTTMap mutex needs to be held before trying to lock any HDTT Entry.
+  mutable std::mutex Mtx;
 };
 
 /// Wrapper around the HostDataToTargetTy to be used in the HDTT map. In
 /// addition to the HDTT pointer we store the key value explicitly. This
 /// allows the set to inspect (sort/search/...) this entry without an additional
 /// load of HDTT. HDTT is a pointer to allow the modification of the set without
 /// invalidating HDTT entries which can now be inspected at the same time.
 struct HostDataToTargetMapKeyTy {
   uintptr_t KeyValue;
 
   HostDataToTargetMapKeyTy(void *Key) : KeyValue(uintptr_t(Key)) {}
+  HostDataToTargetMapKeyTy(uintptr_t Key) : KeyValue(Key) {}
   HostDataToTargetMapKeyTy(HostDataToTargetTy *HDTT)
       : KeyValue(HDTT->HstPtrBegin), HDTT(HDTT) {}
   HostDataToTargetTy *HDTT;
 };
 inline bool operator<(const HostDataToTargetMapKeyTy &LHS,
                       const uintptr_t &RHS) {
   return LHS.KeyValue < RHS;
 }
 inline bool operator<(const uintptr_t &LHS,
                       const HostDataToTargetMapKeyTy &RHS) {
   return LHS < RHS.KeyValue;
 }
 inline bool operator<(const HostDataToTargetMapKeyTy &LHS,
                       const HostDataToTargetMapKeyTy &RHS) {
   return LHS.KeyValue < RHS.KeyValue;
 }
 
-struct LookupResult {
-  struct {
-    unsigned IsContained : 1;
-    unsigned ExtendsBefore : 1;
-    unsigned ExtendsAfter : 1;
-  } Flags;
-
-  /// The corresponding map table entry which is stable.
-  HostDataToTargetTy *Entry = nullptr;
-
-  LookupResult() : Flags({0, 0, 0}), Entry() {}
-};
-
 /// This struct will be returned by \p DeviceTy::getTargetPointer which provides
-/// more data than just a target pointer.
+/// more data than just a target pointer. A TargetPointerResultTy that has a non
+/// null Entry owns the entry. As long as the TargetPointerResultTy (TPR) exists
+/// the entry is locked. To give up ownership without destroying the TPR use the
+/// reset() function.
 struct TargetPointerResultTy {
-  struct {
+  struct FlagTy {
     /// If the map table entry is just created
     unsigned IsNewEntry : 1;
     /// If the pointer is actually a host pointer (when unified memory enabled)
     unsigned IsHostPointer : 1;
     /// If the pointer is present in the mapping table.
     unsigned IsPresent : 1;
-  } Flags = {0, 0, 0};
+    /// Flag indicating that this was the last user of the entry and the ref
+    /// count is now 0.
+    unsigned IsLast : 1;
+  } Flags = {0, 0, 0, 0};
+
+  TargetPointerResultTy(const TargetPointerResultTy &) = delete;
+  TargetPointerResultTy &operator=(const TargetPointerResultTy &TPR) = delete;
+  TargetPointerResultTy() {}
+
+  TargetPointerResultTy(FlagTy Flags, HostDataToTargetTy *Entry,
+                        void *TargetPointer)
+      : Flags(Flags), TargetPointer(TargetPointer), Entry(Entry) {
+    if (Entry)
+      Entry->lock();
+  }
+
+  TargetPointerResultTy(TargetPointerResultTy &&TPR)
+      : Flags(TPR.Flags), TargetPointer(TPR.TargetPointer), Entry(TPR.Entry) {
+    TPR.Entry = nullptr;
+  }
+
+  TargetPointerResultTy &operator=(TargetPointerResultTy &&TPR) {
+    if (&TPR != this) {
+      std::swap(Flags, TPR.Flags);
+      std::swap(Entry, TPR.Entry);
+      std::swap(TargetPointer, TPR.TargetPointer);
+    }
+    return *this;
+  }
+
+  ~TargetPointerResultTy() {
+    if (Entry)
+      Entry->unlock();
+  }
 
   bool isPresent() const { return Flags.IsPresent; }
 
   bool isHostPointer() const { return Flags.IsHostPointer; }
 
-  /// The corresponding map table entry which is stable.
-  HostDataToTargetTy *Entry = nullptr;
-
   /// The corresponding target pointer
   void *TargetPointer = nullptr;
+
+  HostDataToTargetTy *getEntry() const { return Entry; }
+  void setEntry(HostDataToTargetTy *HDTTT,
+                HostDataToTargetTy *OwnedTPR = nullptr) {
+    if (Entry)
+      Entry->unlock();
+    Entry = HDTTT;
+    if (Entry && Entry != OwnedTPR)
+      Entry->lock();
+  }
+
+  void reset() { *this = TargetPointerResultTy(); }
+
+private:
+  /// The corresponding map table entry which is stable.
+  HostDataToTargetTy *Entry = nullptr;
 };
 
-/// Map for shadow pointers
-struct ShadowPtrValTy {
-  void *HstPtrVal;
-  void *TgtPtrAddr;
-  void *TgtPtrVal;
+struct LookupResult {
+  struct {
+    unsigned IsContained : 1;
+    unsigned ExtendsBefore : 1;
+    unsigned ExtendsAfter : 1;
+  } Flags;
+
+  LookupResult() : Flags({0, 0, 0}), TPR() {}
+
+  TargetPointerResultTy TPR;
 };
-typedef std::map<void *, ShadowPtrValTy> ShadowPtrListTy;
 
 ///
 struct PendingCtorDtorListsTy {
   std::list<void *> PendingCtors;
   std::list<void *> PendingDtors;
 };
 typedef std::map<__tgt_bin_desc *, PendingCtorDtorListsTy>
     PendingCtorsDtorsPerLibrary;
 
 struct DeviceTy {
   int32_t DeviceID;
   RTLInfoTy *RTL;
   int32_t RTLDeviceID;
 
   bool IsInit;
struct DeviceTy {
   /// at a time.
   ProtectedObj<HostDataToTargetListTy> HostDataToTargetMap;
 
   /// The type used to access the HDTT map.
   using HDTTMapAccessorTy = decltype(HostDataToTargetMap)::AccessorTy;
 
   PendingCtorsDtorsPerLibrary PendingCtorsDtors;
 
-  ShadowPtrListTy ShadowPtrMap;
-
-  std::mutex PendingGlobalsMtx, ShadowMtx;
+  std::mutex PendingGlobalsMtx;
 
   DeviceTy(RTLInfoTy *RTL);
   // DeviceTy is not copyable
   DeviceTy(const DeviceTy &D) = delete;
   DeviceTy &operator=(const DeviceTy &D) = delete;
 
   ~DeviceTy();
 
   // Return true if data can be copied to DstDevice directly
   bool isDataExchangable(const DeviceTy &DstDevice);
 
   /// Lookup the mapping of \p HstPtrBegin in \p HDTTMap. The accessor ensures
   /// exclusive access to the HDTT map.
   LookupResult lookupMapping(HDTTMapAccessorTy &HDTTMap, void *HstPtrBegin,
-                             int64_t Size);
+                             int64_t Size,
+                             HostDataToTargetTy *OwnedTPR = nullptr);
 
   /// Get the target pointer based on host pointer begin and base. If the
   /// mapping already exists, the target pointer will be returned directly. In
   /// addition, if required, the memory region pointed by \p HstPtrBegin of size
   /// \p Size will also be transferred to the device. If the mapping doesn't
   /// exist, and if unified shared memory is not enabled, a new mapping will be
   /// created and the data will also be transferred accordingly. nullptr will be
   /// returned because of any of following reasons:
   /// - Data allocation failed;
   /// - The user tried to do an illegal mapping;
   /// - Data transfer issue fails.
-  TargetPointerResultTy
-  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 getTargetPointer(
+      HDTTMapAccessorTy &HDTTMap, void *HstPtrBegin, void *HstPtrBase,
+      int64_t Size, map_var_info_t HstPtrName, bool HasFlagTo,
+      bool HasCloseModifier, bool HasPresentModifier, bool HasHoldModifier,
+      AsyncInfoTy &AsyncInfo, HostDataToTargetTy *OwnedTPR = nullptr,
+      bool ReleaseHDTTMap = true);
 
   /// Return the target pointer for \p HstPtrBegin in \p HDTTMap. The accessor
   /// ensures exclusive access to the HDTT map.
   void *getTgtPtrBegin(HDTTMapAccessorTy &HDTTMap, void *HstPtrBegin,
                        int64_t Size);
 
   /// Return the target pointer begin (where the data will be moved).
   /// Used by targetDataBegin, targetDataEnd, targetDataUpdate and target.
   /// - \p UpdateRefCount and \p UseHoldRefCount controls which and if the entry
   /// reference counters will be decremented.
   /// - \p MustContain enforces that the query must not extend beyond an already
   /// mapped entry to be valid.
   /// - \p ForceDelete deletes the entry regardless of its reference counting
   /// (unless it is infinite).
   /// - \p FromDataEnd tracks the number of threads referencing the entry at
   /// targetDataEnd for delayed deletion purpose.
   [[nodiscard]] TargetPointerResultTy
-  getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
-                 bool UpdateRefCount, bool UseHoldRefCount, bool &IsHostPtr,
-                 bool MustContain = false, bool ForceDelete = false,
-                 bool FromDataEnd = false);
+  getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool UpdateRefCount,
+                 bool UseHoldRefCount, bool MustContain = false,
+                 bool ForceDelete = false, bool FromDataEnd = false);
 
   /// Remove the \p Entry from the data map. Expect the entry's total reference
   /// count to be zero and the caller thread to be the last one using it. \p
   /// HDTTMap ensure the caller holds exclusive access and can modify the map.
   /// 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.
struct DeviceTy {
   /// OFFLOAD_SUCCESS/OFFLOAD_FAIL when succeeds/fails. p Kind dictates what
   /// allocator should be used (host, shared, device).
   int32_t deleteData(void *TgtPtrBegin, int32_t Kind = TARGET_ALLOC_DEFAULT);
 
   // 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);
+                     AsyncInfoTy &AsyncInfo,
+                     HostDataToTargetTy *Entry = nullptr);
   // Copy data from device back to host
   int32_t retrieveData(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size,
-                       AsyncInfoTy &AsyncInfo);
+                       AsyncInfoTy &AsyncInfo,
+                       HostDataToTargetTy *Entry = nullptr);
   // Copy data from current device to destination device directly
   int32_t dataExchange(void *SrcPtr, DeviceTy &DstDev, void *DstPtr,
                        int64_t Size, AsyncInfoTy &AsyncInfo);
 
   /// Notify the plugin about a new mapping starting at the host address
   /// \p HstPtr and \p Size bytes.
   int32_t notifyDataMapped(void *HstPtr, int64_t Size);
 

openmp/libomptarget/src/api.cpp

EXTERN int omp_target_is_present(const void *Ptr, int DeviceNum) {
   PM->RTLsMtx.unlock();
   if (DevicesSize <= (size_t)DeviceNum) {
     DP("Call to omp_target_is_present with invalid device ID, returning "
        "false\n");
     return false;
   }
 
   DeviceTy &Device = *PM->Devices[DeviceNum];
-  bool IsLast; // not used
-  bool IsHostPtr;
   // omp_target_is_present tests whether a host pointer refers to storage that
   // is mapped to a given device. However, due to the lack of the storage size,
   // only check 1 byte. Cannot set size 0 which checks whether the pointer (zero
   // lengh array) is mapped instead of the referred storage.
-  TargetPointerResultTy TPR =
-      Device.getTgtPtrBegin(const_cast<void *>(Ptr), 1, IsLast,
+  TargetPointerResultTy TPR = Device.getTgtPtrBegin(const_cast<void *>(Ptr), 1,
                                                     /*UpdateRefCount=*/false,
-                            /*UseHoldRefCount=*/false, IsHostPtr);
+                                                    /*UseHoldRefCount=*/false);
   int Rc = TPR.isPresent();
   DP("Call to omp_target_is_present returns %d\n", Rc);
   return Rc;
 }
 
 EXTERN int omp_target_memcpy(void *Dst, const void *Src, size_t Length,
                              size_t DstOffset, size_t SrcOffset, int DstDevice,
                              int SrcDevice) {
if (DevicesSize <= DeviceNum) {
     return nullptr;
   }
 
   if (!deviceIsReady(DeviceNum)) {
     REPORT("Device %d is not ready, returning nullptr.\n", DeviceNum);
     return nullptr;
   }
 
-  bool IsLast = false;
-  bool IsHostPtr = false;
   auto &Device = *PM->Devices[DeviceNum];
-  TargetPointerResultTy TPR =
-      Device.getTgtPtrBegin(const_cast<void *>(Ptr), 1, IsLast,
+  TargetPointerResultTy TPR = Device.getTgtPtrBegin(const_cast<void *>(Ptr), 1,
                                                     /*UpdateRefCount=*/false,
-                            /*UseHoldRefCount=*/false, IsHostPtr);
+                                                    /*UseHoldRefCount=*/false);
   if (!TPR.isPresent()) {
     DP("Ptr " DPxMOD "is not present on device %d, returning nullptr.\n",
        DPxPTR(Ptr), DeviceNum);
     return nullptr;
   }
 
   DP("omp_get_mapped_ptr returns " DPxMOD ".\n", DPxPTR(TPR.TargetPointer));
 
   return TPR.TargetPointer;
 }

openmp/libomptarget/src/device.cpp

 #include "omptarget.h"
 #include "private.h"
 #include "rtl.h"
 
 #include <cassert>
 #include <climits>
 #include <cstdint>
 #include <cstdio>
+#include <mutex>
 #include <string>
 #include <thread>
 
 int HostDataToTargetTy::addEventIfNecessary(DeviceTy &Device,
                                             AsyncInfoTy &AsyncInfo) const {
   // First, check if the user disabled atomic map transfer/malloc/dealloc.
   if (!PM->UseEventsForAtomicTransfers)
     return OFFLOAD_SUCCESS;
int HostDataToTargetTy::addEventIfNecessary(DeviceTy &Device,
   if (NeedNewEvent)
     setEvent(Event);
 
   return OFFLOAD_SUCCESS;
 }
 
 DeviceTy::DeviceTy(RTLInfoTy *RTL)
     : DeviceID(-1), RTL(RTL), RTLDeviceID(-1), IsInit(false), InitFlag(),
-      HasPendingGlobals(false), PendingCtorsDtors(), ShadowPtrMap(),
-      PendingGlobalsMtx(), ShadowMtx() {}
+      HasPendingGlobals(false), PendingCtorsDtors(), PendingGlobalsMtx() {}
 
 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);
 }
 
 int DeviceTy::associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size) {
   HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
 
   // Check if entry exists
   auto It = HDTTMap->find(HstPtrBegin);
   if (It != HDTTMap->end()) {
     HostDataToTargetTy &HDTT = *It->HDTT;
+    std::lock_guard<HostDataToTargetTy> LG(HDTT);
     // Mapping already exists
     bool IsValid = HDTT.HstPtrEnd == (uintptr_t)HstPtrBegin + Size &&
                    HDTT.TgtPtrBegin == (uintptr_t)TgtPtrBegin;
     if (IsValid) {
       DP("Attempt to re-associate the same device ptr+offset with the same "
          "host ptr, nothing to do\n");
       return OFFLOAD_SUCCESS;
     }
int DeviceTy::associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size) {
   // Notify the plugin about the new mapping.
   return notifyDataMapped(HstPtrBegin, Size);
 }
 
 int DeviceTy::disassociatePtr(void *HstPtrBegin) {
   HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
 
   auto It = HDTTMap->find(HstPtrBegin);
-  if (It != HDTTMap->end()) {
-    HostDataToTargetTy &HDTT = *It->HDTT;
-    // Mapping exists
-    if (HDTT.getHoldRefCount()) {
-      // This is based on OpenACC 3.1, sec 3.2.33 "acc_unmap_data", L3656-3657:
-      // "It is an error to call acc_unmap_data if the structured reference
-      // count for the pointer is not zero."
-      REPORT("Trying to disassociate a pointer with a non-zero hold reference "
-             "count\n");
-    } else if (HDTT.isDynRefCountInf()) {
-      DP("Association found, removing it\n");
-      void *Event = HDTT.getEvent();
-      delete &HDTT;
-      if (Event)
-        destroyEvent(Event);
-      HDTTMap->erase(It);
+  if (It == HDTTMap->end()) {
+    REPORT("Association not found\n");
+    return OFFLOAD_FAIL;
+  }
+  HostDataToTargetTy &HDTT = *It->HDTT;
+  std::lock_guard<HostDataToTargetTy> LG(HDTT);
+
+    return OFFLOAD_FAIL;
+  }
+
+  if (HDTT.isDynRefCountInf()) {
+    return OFFLOAD_SUCCESS;
+  }
 
-      // Notify the plugin about the unmapped memory.
-      return notifyDataUnmapped(HstPtrBegin);
-    } else {
-      REPORT("Trying to disassociate a pointer which was not mapped via "
-             "omp_target_associate_ptr\n");
-    }
-  } else {
-    REPORT("Association not found\n");
-  }
-
-  // Mapping not found
   return OFFLOAD_FAIL;
 }
 
 LookupResult DeviceTy::lookupMapping(HDTTMapAccessorTy &HDTTMap,
-                                     void *HstPtrBegin, int64_t Size) {
+                                     void *HstPtrBegin, int64_t Size,
+                                     HostDataToTargetTy *OwnedTPR) {
 
   uintptr_t HP = (uintptr_t)HstPtrBegin;
   LookupResult LR;
 
   DP("Looking up mapping(HstPtrBegin=" DPxMOD ", Size=%" PRId64 ")...\n",
      DPxPTR(HP), Size);
 
   if (HDTTMap->empty())
     return LR;
 
   auto Upper = HDTTMap->upper_bound(HP);
 
   if (Size == 0) {
     // specification v5.1 Pointer Initialization for Device Data Environments
     // upper_bound satisfies
     //   std::prev(upper)->HDTT.HstPtrBegin <= hp < upper->HDTT.HstPtrBegin
     if (Upper != HDTTMap->begin()) {
-      LR.Entry = std::prev(Upper)->HDTT;
-      auto &HT = *LR.Entry;
+      LR.TPR.setEntry(std::prev(Upper)->HDTT, OwnedTPR);
       // the left side of extended address range is satisified.
-      // hp >= HT.HstPtrBegin || hp >= HT.HstPtrBase
-      LR.Flags.IsContained = HP < HT.HstPtrEnd || HP < HT.HstPtrBase;
+      // hp >= LR.TPR.getEntry()->HstPtrBegin || hp >=
+      // LR.TPR.getEntry()->HstPtrBase
+      LR.Flags.IsContained = HP < LR.TPR.getEntry()->HstPtrEnd ||
+                             HP < LR.TPR.getEntry()->HstPtrBase;
     }
 
     if (!LR.Flags.IsContained && Upper != HDTTMap->end()) {
-      LR.Entry = Upper->HDTT;
-      auto &HT = *LR.Entry;
+      LR.TPR.setEntry(Upper->HDTT, OwnedTPR);
       // the right side of extended address range is satisified.
-      // hp < HT.HstPtrEnd || hp < HT.HstPtrBase
-      LR.Flags.IsContained = HP >= HT.HstPtrBase;
+      // hp < LR.TPR.getEntry()->HstPtrEnd || hp < LR.TPR.getEntry()->HstPtrBase
+      LR.Flags.IsContained = HP >= LR.TPR.getEntry()->HstPtrBase;
     }
   } else {
     // check the left bin
     if (Upper != HDTTMap->begin()) {
-      LR.Entry = std::prev(Upper)->HDTT;
-      auto &HT = *LR.Entry;
+      LR.TPR.setEntry(std::prev(Upper)->HDTT, OwnedTPR);
       // Is it contained?
-      LR.Flags.IsContained = HP >= HT.HstPtrBegin && HP < HT.HstPtrEnd &&
-                             (HP + Size) <= HT.HstPtrEnd;
+      LR.Flags.IsContained = HP >= LR.TPR.getEntry()->HstPtrBegin &&
+                             HP < LR.TPR.getEntry()->HstPtrEnd &&
+                             (HP + Size) <= LR.TPR.getEntry()->HstPtrEnd;
       // Does it extend beyond the mapped region?
-      LR.Flags.ExtendsAfter = HP < HT.HstPtrEnd && (HP + Size) > HT.HstPtrEnd;
+      LR.Flags.ExtendsAfter = HP < LR.TPR.getEntry()->HstPtrEnd &&
+                              (HP + Size) > LR.TPR.getEntry()->HstPtrEnd;
     }
 
     // check the right bin
     if (!(LR.Flags.IsContained || LR.Flags.ExtendsAfter) &&
         Upper != HDTTMap->end()) {
-      LR.Entry = Upper->HDTT;
-      auto &HT = *LR.Entry;
+      LR.TPR.setEntry(Upper->HDTT, OwnedTPR);
       // Does it extend into an already mapped region?
-      LR.Flags.ExtendsBefore =
-          HP < HT.HstPtrBegin && (HP + Size) > HT.HstPtrBegin;
+      LR.Flags.ExtendsBefore = HP < LR.TPR.getEntry()->HstPtrBegin &&
+                               (HP + Size) > LR.TPR.getEntry()->HstPtrBegin;
       // Does it extend beyond the mapped region?
-      LR.Flags.ExtendsAfter = HP < HT.HstPtrEnd && (HP + Size) > HT.HstPtrEnd;
+      LR.Flags.ExtendsAfter = HP < LR.TPR.getEntry()->HstPtrEnd &&
+                              (HP + Size) > LR.TPR.getEntry()->HstPtrEnd;
     }
 
     if (LR.Flags.ExtendsBefore) {
       DP("WARNING: Pointer is not mapped but section extends into already "
          "mapped data\n");
     }
     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,
+    HDTTMapAccessorTy &HDTTMap, 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) {
-  HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
+    bool HasPresentModifier, bool HasHoldModifier, AsyncInfoTy &AsyncInfo,
+    HostDataToTargetTy *OwnedTPR, bool ReleaseHDTTMap) {
 
-  void *TargetPointer = nullptr;
-  bool IsHostPtr = false;
-  bool IsPresent = true;
-  bool IsNew = false;
+  LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size, OwnedTPR);
+  LR.TPR.Flags.IsPresent = true;
 
-  LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
-  auto *Entry = LR.Entry;
+  // Release the mapping table lock only after the entry is locked by
+  // attaching it to TPR. Once TPR is destroyed it will release the lock
+  // on entry. If it is returned the lock will move to the returned object.
+  // If LR.Entry is already owned/locked we avoid trying to lock it again.
 
   // Check if the pointer is contained.
   // If a variable is mapped to the device manually by the user - which would
   // lead to the IsContained flag to be true - then we must ensure that the
   // device address is returned even under unified memory conditions.
   if (LR.Flags.IsContained ||
       ((LR.Flags.ExtendsBefore || LR.Flags.ExtendsAfter) && IsImplicit)) {
-    auto &HT = *LR.Entry;
     const char *RefCountAction;
     if (UpdateRefCount) {
       // After this, reference count >= 1. If the reference count was 0 but the
       // entry was still there we can reuse the data on the device and avoid a
       // new submission.
-      HT.incRefCount(HasHoldModifier);
+      LR.TPR.getEntry()->incRefCount(HasHoldModifier);
       RefCountAction = " (incremented)";
     } else {
       // It might have been allocated with the parent, but it's still new.
-      IsNew = HT.getTotalRefCount() == 1;
+      LR.TPR.Flags.IsNewEntry = LR.TPR.getEntry()->getTotalRefCount() == 1;
       RefCountAction = " (update suppressed)";
     }
     const char *DynRefCountAction = HasHoldModifier ? "" : RefCountAction;
     const char *HoldRefCountAction = HasHoldModifier ? RefCountAction : "";
-    uintptr_t Ptr = HT.TgtPtrBegin + ((uintptr_t)HstPtrBegin - HT.HstPtrBegin);
+    uintptr_t Ptr = LR.TPR.getEntry()->TgtPtrBegin +
+                    ((uintptr_t)HstPtrBegin - LR.TPR.getEntry()->HstPtrBegin);
     INFO(OMP_INFOTYPE_MAPPING_EXISTS, DeviceID,
          "Mapping exists%s with HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD
          ", Size=%" PRId64 ", DynRefCount=%s%s, HoldRefCount=%s%s, Name=%s\n",
          (IsImplicit ? " (implicit)" : ""), DPxPTR(HstPtrBegin), DPxPTR(Ptr),
-         Size, HT.dynRefCountToStr().c_str(), DynRefCountAction,
-         HT.holdRefCountToStr().c_str(), HoldRefCountAction,
+         Size, LR.TPR.getEntry()->dynRefCountToStr().c_str(), DynRefCountAction,
+         LR.TPR.getEntry()->holdRefCountToStr().c_str(), HoldRefCountAction,
          (HstPtrName) ? getNameFromMapping(HstPtrName).c_str() : "unknown");
-    TargetPointer = (void *)Ptr;
+    LR.TPR.TargetPointer = (void *)Ptr;
   } else if ((LR.Flags.ExtendsBefore || LR.Flags.ExtendsAfter) && !IsImplicit) {
     // Explicit extension of mapped data - not allowed.
     MESSAGE("explicit extension not allowed: host address specified is " DPxMOD
             " (%" PRId64
             " bytes), but device allocation maps to host at " DPxMOD
             " (%" PRId64 " bytes)",
-            DPxPTR(HstPtrBegin), Size, DPxPTR(Entry->HstPtrBegin),
-            Entry->HstPtrEnd - Entry->HstPtrBegin);
+            DPxPTR(HstPtrBegin), Size, DPxPTR(LR.TPR.getEntry()->HstPtrBegin),
+            LR.TPR.getEntry()->HstPtrEnd - LR.TPR.getEntry()->HstPtrBegin);
     if (HasPresentModifier)
       MESSAGE("device mapping required by 'present' map type modifier does not "
               "exist for host address " DPxMOD " (%" PRId64 " bytes)",
               DPxPTR(HstPtrBegin), Size);
   } else if (PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY &&
              !HasCloseModifier) {
     // If unified shared memory is active, implicitly mapped variables that are
     // not privatized use host address. Any explicitly mapped variables also use
     // host address where correctness is not impeded. In all other cases maps
     // are respected.
     // In addition to the mapping rules above, the close map modifier forces the
     // mapping of the variable to the device.
     if (Size) {
       DP("Return HstPtrBegin " DPxMOD " Size=%" PRId64 " for unified shared "
          "memory\n",
          DPxPTR((uintptr_t)HstPtrBegin), Size);
-      IsPresent = false;
-      IsHostPtr = true;
-      TargetPointer = HstPtrBegin;
+      LR.TPR.Flags.IsPresent = false;
+      LR.TPR.Flags.IsHostPointer = true;
+      LR.TPR.TargetPointer = HstPtrBegin;
     }
   } else if (HasPresentModifier) {
     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;
+    LR.TPR.Flags.IsNewEntry = true;
     uintptr_t Ptr = (uintptr_t)allocData(Size, HstPtrBegin);
-    Entry = HDTTMap
+    // Release the mapping table lock only after the entry is locked by
+    // attaching it to TPR.
+    LR.TPR.setEntry(HDTTMap
                         ->emplace(new HostDataToTargetTy(
                             (uintptr_t)HstPtrBase, (uintptr_t)HstPtrBegin,
                             (uintptr_t)HstPtrBegin + Size, Ptr, HasHoldModifier,
                             HstPtrName))
-                .first->HDTT;
+                        .first->HDTT);
     INFO(OMP_INFOTYPE_MAPPING_CHANGED, DeviceID,
          "Creating new map entry with HstPtrBase=" DPxMOD
          ", HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD ", Size=%ld, "
          "DynRefCount=%s, HoldRefCount=%s, Name=%s\n",
          DPxPTR(HstPtrBase), DPxPTR(HstPtrBegin), DPxPTR(Ptr), Size,
-         Entry->dynRefCountToStr().c_str(), Entry->holdRefCountToStr().c_str(),
+         LR.TPR.getEntry()->dynRefCountToStr().c_str(),
+         LR.TPR.getEntry()->holdRefCountToStr().c_str(),
          (HstPtrName) ? getNameFromMapping(HstPtrName).c_str() : "unknown");
-    TargetPointer = (void *)Ptr;
+    LR.TPR.TargetPointer = (void *)Ptr;
 
     // Notify the plugin about the new mapping.
     if (notifyDataMapped(HstPtrBegin, Size))
       return {{false /* IsNewEntry */, false /* IsHostPointer */},
               nullptr /* Entry */,
               nullptr /* TargetPointer */};
   } else {
     // This entry is not present and we did not create a new entry for it.
-    IsPresent = false;
+    LR.TPR.Flags.IsPresent = false;
   }
 
-  // If the target pointer is valid, and we need to transfer data, issue the
-  // data transfer.
-  if (TargetPointer && !IsHostPtr && HasFlagTo && (IsNew || HasFlagAlways) &&
-      Size != 0) {
-    // Lock the entry before releasing the mapping table lock such that another
-    // thread that could issue data movement will get the right result.
-    std::lock_guard<decltype(*Entry)> LG(*Entry);
-    // Release the mapping table lock right after the entry is locked.
+  // All mapping table modifications have been made. If the user requested it we
+  // give up the lock.
+  if (ReleaseHDTTMap)
     HDTTMap.destroy();
 
+  // If the target pointer is valid, and we need to transfer data, issue the
+  // data transfer.
+  if (LR.TPR.TargetPointer && !LR.TPR.Flags.IsHostPointer && HasFlagTo &&
+      (LR.TPR.Flags.IsNewEntry || HasFlagAlways) && Size != 0) {
     DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n", Size,
-       DPxPTR(HstPtrBegin), DPxPTR(TargetPointer));
+       DPxPTR(HstPtrBegin), DPxPTR(LR.TPR.TargetPointer));
 
-    int Ret = submitData(TargetPointer, HstPtrBegin, Size, AsyncInfo);
+    int Ret = submitData(LR.TPR.TargetPointer, HstPtrBegin, Size, AsyncInfo,
+                         LR.TPR.getEntry());
     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)
+      LR.TPR.TargetPointer = nullptr;
+    } else if (LR.TPR.getEntry()->addEventIfNecessary(*this, AsyncInfo) !=
+               OFFLOAD_SUCCESS)
       return {{false /* IsNewEntry */, false /* IsHostPointer */},
               nullptr /* Entry */,
               nullptr /* TargetPointer */};
   } else {
-    // Release the mapping table lock directly.
-    HDTTMap.destroy();
     // If not a host pointer and no present modifier, we need to wait for the
     // event if it exists.
     // Note: Entry might be nullptr because of zero length array section.
-    if (Entry && !IsHostPtr && !HasPresentModifier) {
-      std::lock_guard<decltype(*Entry)> LG(*Entry);
-      void *Event = Entry->getEvent();
+    if (LR.TPR.getEntry() && !LR.TPR.Flags.IsHostPointer &&
+        !HasPresentModifier) {
+      void *Event = LR.TPR.getEntry()->getEvent();
       if (Event) {
         int Ret = waitEvent(Event, AsyncInfo);
         if (Ret != OFFLOAD_SUCCESS) {
           // If it fails to wait for the event, we need to return nullptr in
           // case of any data race.
           REPORT("Failed to wait for event " DPxMOD ".\n", DPxPTR(Event));
           return {{false /* IsNewEntry */, false /* IsHostPointer */},
                   nullptr /* Entry */,
                   nullptr /* TargetPointer */};
         }
       }
     }
   }
 
-  return {{IsNew, IsHostPtr, IsPresent}, Entry, TargetPointer};
+  return std::move(LR.TPR);
 }
 
 TargetPointerResultTy
-DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
-                         bool UpdateRefCount, bool UseHoldRefCount,
-                         bool &IsHostPtr, bool MustContain, bool ForceDelete,
-                         bool FromDataEnd) {
+DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool UpdateRefCount,
+                         bool UseHoldRefCount, bool MustContain,
+                         bool ForceDelete, bool FromDataEnd) {
   HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
 
-  void *TargetPointer = NULL;
-  bool IsNew = false;
-  bool IsPresent = true;
-  IsHostPtr = false;
-  IsLast = false;
   LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
 
+  LR.TPR.Flags.IsPresent = true;
+
   if (LR.Flags.IsContained ||
       (!MustContain && (LR.Flags.ExtendsBefore || LR.Flags.ExtendsAfter))) {
-    auto &HT = *LR.Entry;
-    IsLast = HT.decShouldRemove(UseHoldRefCount, ForceDelete);
+    LR.TPR.Flags.IsLast =
+        LR.TPR.getEntry()->decShouldRemove(UseHoldRefCount, ForceDelete);
 
     if (ForceDelete) {
-      HT.resetRefCount(UseHoldRefCount);
-      assert(IsLast == HT.decShouldRemove(UseHoldRefCount) &&
+      LR.TPR.getEntry()->resetRefCount(UseHoldRefCount);
+      assert(LR.TPR.Flags.IsLast ==
+                 LR.TPR.getEntry()->decShouldRemove(UseHoldRefCount) &&
              "expected correct IsLast prediction for reset");
     }
 
     // Increment the number of threads that is using the entry on a
     // targetDataEnd, tracking the number of possible "deleters". A thread may
     // come to own the entry deletion even if it was not the last one querying
     // for it. Thus, we must track every query on targetDataEnds to ensure only
     // the last thread that holds a reference to an entry actually deletes it.
     if (FromDataEnd)
-      HT.incDataEndThreadCount();
+      LR.TPR.getEntry()->incDataEndThreadCount();
 
     const char *RefCountAction;
     if (!UpdateRefCount) {
       RefCountAction = " (update suppressed)";
-    } else if (IsLast) {
-      HT.decRefCount(UseHoldRefCount);
-      assert(HT.getTotalRefCount() == 0 &&
+    } else if (LR.TPR.Flags.IsLast) {
+      LR.TPR.getEntry()->decRefCount(UseHoldRefCount);
+      assert(LR.TPR.getEntry()->getTotalRefCount() == 0 &&
              "Expected zero reference count when deletion is scheduled");
       if (ForceDelete)
         RefCountAction = " (reset, delayed deletion)";
       else
         RefCountAction = " (decremented, delayed deletion)";
     } else {
-      HT.decRefCount(UseHoldRefCount);
+      LR.TPR.getEntry()->decRefCount(UseHoldRefCount);
       RefCountAction = " (decremented)";
     }
     const char *DynRefCountAction = UseHoldRefCount ? "" : RefCountAction;
     const char *HoldRefCountAction = UseHoldRefCount ? RefCountAction : "";
-    uintptr_t TP = HT.TgtPtrBegin + ((uintptr_t)HstPtrBegin - HT.HstPtrBegin);
+    uintptr_t TP = LR.TPR.getEntry()->TgtPtrBegin +
+                   ((uintptr_t)HstPtrBegin - LR.TPR.getEntry()->HstPtrBegin);
     INFO(OMP_INFOTYPE_MAPPING_EXISTS, DeviceID,
          "Mapping exists with HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD ", "
          "Size=%" PRId64 ", DynRefCount=%s%s, HoldRefCount=%s%s\n",
-         DPxPTR(HstPtrBegin), DPxPTR(TP), Size, HT.dynRefCountToStr().c_str(),
-         DynRefCountAction, HT.holdRefCountToStr().c_str(), HoldRefCountAction);
-    TargetPointer = (void *)TP;
+         DPxPTR(HstPtrBegin), DPxPTR(TP), Size,
+         LR.TPR.getEntry()->dynRefCountToStr().c_str(), DynRefCountAction,
+         LR.TPR.getEntry()->holdRefCountToStr().c_str(), HoldRefCountAction);
+    LR.TPR.TargetPointer = (void *)TP;
   } else if (PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) {
     // If the value isn't found in the mapping and unified shared memory
     // is on then it means we have stumbled upon a value which we need to
     // use directly from the host.
     DP("Get HstPtrBegin " DPxMOD " Size=%" PRId64 " for unified shared "
        "memory\n",
        DPxPTR((uintptr_t)HstPtrBegin), Size);
-    IsPresent = false;
-    IsHostPtr = true;
-    TargetPointer = HstPtrBegin;
+    LR.TPR.Flags.IsPresent = false;
+    LR.TPR.Flags.IsHostPointer = true;
+    LR.TPR.TargetPointer = HstPtrBegin;
   } else {
     // OpenMP Specification v5.2: if a matching list item is not found, the
     // pointer retains its original value as per firstprivate semantics.
-    IsPresent = false;
-    IsHostPtr = false;
-    TargetPointer = HstPtrBegin;
+    LR.TPR.Flags.IsPresent = false;
+    LR.TPR.Flags.IsHostPointer = false;
+    LR.TPR.TargetPointer = HstPtrBegin;
   }
 
-  return {{IsNew, IsHostPtr, IsPresent}, LR.Entry, TargetPointer};
+  return std::move(LR.TPR);
 }
 
 // Return the target pointer begin (where the data will be moved).
 void *DeviceTy::getTgtPtrBegin(HDTTMapAccessorTy &HDTTMap, void *HstPtrBegin,
                                int64_t Size) {
   uintptr_t HP = (uintptr_t)HstPtrBegin;
   LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
   if (LR.Flags.IsContained || LR.Flags.ExtendsBefore || LR.Flags.ExtendsAfter) {
-    auto &HT = *LR.Entry;
-    uintptr_t TP = HT.TgtPtrBegin + (HP - HT.HstPtrBegin);
+    uintptr_t TP =
+        LR.TPR.getEntry()->TgtPtrBegin + (HP - LR.TPR.getEntry()->HstPtrBegin);
     return (void *)TP;
   }
 
   return NULL;
 }
 
 int DeviceTy::eraseMapEntry(HDTTMapAccessorTy &HDTTMap,
                             HostDataToTargetTy *Entry, int64_t Size) {

 void *DeviceTy::allocData(int64_t Size, void *HstPtr, int32_t Kind) {
   return RTL->data_alloc(RTLDeviceID, Size, HstPtr, Kind);
 }
 
 int32_t DeviceTy::deleteData(void *TgtPtrBegin, int32_t Kind) {
   return RTL->data_delete(RTLDeviceID, TgtPtrBegin, Kind);
 }
 
-// Submit data to device
-int32_t DeviceTy::submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
-                             AsyncInfoTy &AsyncInfo) {
-  if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
-    HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
-    LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
-    auto *HT = &*LR.Entry;
+static void printCopyInfo(int DeviceId, bool H2D, void *SrcPtrBegin,
+                          void *DstPtrBegin, int64_t Size,
+                          HostDataToTargetTy *HT) {
 
-    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");
-  }
+  INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceId,
+       "Copying data from %s to %s, %sPtr=" DPxMOD ", %sPtr=" DPxMOD
+       H2D ? "host" : "device", H2D ? "device" : "host", H2D ? "Hst" : "Tgt",
+       DPxPTR(SrcPtrBegin), H2D ? "Tgt" : "Hst", DPxPTR(DstPtrBegin), Size,
 
+// Submit data to device
+int32_t DeviceTy::submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
+                             AsyncInfoTy &AsyncInfo,
+                             HostDataToTargetTy *Entry) {
+  if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
+    HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor(Entry);
+    LookupResult LR;
+    if (!Entry) {
+      LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
+      Entry = LR.TPR.getEntry();
+    }
+    printCopyInfo(DeviceID, /* H2D */ true, HstPtrBegin, TgtPtrBegin, Size,
+                  Entry);
+  }
+
   if (!AsyncInfo || !RTL->data_submit_async || !RTL->synchronize)
     return RTL->data_submit(RTLDeviceID, TgtPtrBegin, HstPtrBegin, Size);
   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) {
+                               int64_t Size, AsyncInfoTy &AsyncInfo,
+                               HostDataToTargetTy *Entry) {
   if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
-    HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
-    LookupResult LR = lookupMapping(HDTTMap, 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");
+    HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor(Entry);
+    LookupResult LR;
+    if (!Entry) {
+      LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
+      Entry = LR.TPR.getEntry();
+    }
+    printCopyInfo(DeviceID, /* H2D */ false, TgtPtrBegin, HstPtrBegin, Size,
+                  Entry);
   }
 
   if (!RTL->data_retrieve_async || !RTL->synchronize)
     return RTL->data_retrieve(RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size);
   return RTL->data_retrieve_async(RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size,
                                   AsyncInfo);
 }
 

openmp/libomptarget/src/omptarget.cpp

for (int32_t I = 0; I < ArgNum; ++I) {
     bool HasHoldModifier = ArgTypes[I] & OMP_TGT_MAPTYPE_OMPX_HOLD;
     // UpdateRef is based on MEMBER_OF instead of TARGET_PARAM because if we
     // have reached this point via __tgt_target_data_begin and not __tgt_target
     // then no argument is marked as TARGET_PARAM ("omp target data map" is not
     // associated with a target region, so there are no target parameters). This
     // may be considered a hack, we could revise the scheme in the future.
     bool UpdateRef =
         !(ArgTypes[I] & OMP_TGT_MAPTYPE_MEMBER_OF) && !(FromMapper && I == 0);
+
+    DeviceTy::HDTTMapAccessorTy HDTTMap =
+        Device.HostDataToTargetMap.getExclusiveAccessor();
     if (ArgTypes[I] & OMP_TGT_MAPTYPE_PTR_AND_OBJ) {
       DP("Has a pointer entry: \n");
       // Base is address of pointer.
       //
       // Usually, the pointer is already allocated by this time.  For example:
       //
       //   #pragma omp target map(s.p[0:N])
       //
       // The map entry for s comes first, and the PTR_AND_OBJ entry comes
       // afterward, so the pointer is already allocated by the time the
       // PTR_AND_OBJ entry is handled below, and PointerTgtPtrBegin is thus
       // non-null.  However, "declare target link" can produce a PTR_AND_OBJ
       // entry for a global that might not already be allocated by the time the
       // PTR_AND_OBJ entry is handled below, and so the allocation might fail
       // when HasPresentModifier.
       PointerTpr = Device.getTargetPointer(
-          HstPtrBase, HstPtrBase, sizeof(void *), /*HstPtrName=*/nullptr,
+          HDTTMap, HstPtrBase, HstPtrBase, sizeof(void *),
+          /*HstPtrName=*/nullptr,
           /*HasFlagTo=*/false, /*HasFlagAlways=*/false, IsImplicit, UpdateRef,
-          HasCloseModifier, HasPresentModifier, HasHoldModifier, AsyncInfo);
+          HasCloseModifier, HasPresentModifier, HasHoldModifier, AsyncInfo,
+          /* OwnedTPR */ nullptr, /* ReleaseHDTTMap */ false);
       PointerTgtPtrBegin = PointerTpr.TargetPointer;
       IsHostPtr = PointerTpr.Flags.IsHostPointer;
       if (!PointerTgtPtrBegin) {
         REPORT("Call to getTargetPointer returned null pointer (%s).\n",
                HasPresentModifier ? "'present' map type modifier"
                                   : "device failure or illegal mapping");
         return OFFLOAD_FAIL;
       }
       DP("There are %zu bytes allocated at target address " DPxMOD " - is%s new"
          "\n",
          sizeof(void *), DPxPTR(PointerTgtPtrBegin),
          (PointerTpr.Flags.IsNewEntry ? "" : " not"));
       PointerHstPtrBegin = HstPtrBase;
       // modify current entry.
       HstPtrBase = *(void **)HstPtrBase;
       // No need to update pointee ref count for the first element of the
       // subelement that comes from mapper.
       UpdateRef =
           (!FromMapper || I != 0); // subsequently update ref count of pointee
     }
 
     const bool HasFlagTo = ArgTypes[I] & OMP_TGT_MAPTYPE_TO;
     const bool HasFlagAlways = ArgTypes[I] & OMP_TGT_MAPTYPE_ALWAYS;
+    // Note that HDTTMap will be released in getTargetPointer.
     auto TPR = Device.getTargetPointer(
-        HstPtrBegin, HstPtrBase, DataSize, HstPtrName, HasFlagTo, HasFlagAlways,
-        IsImplicit, UpdateRef, HasCloseModifier, HasPresentModifier,
-        HasHoldModifier, AsyncInfo);
+        HDTTMap, HstPtrBegin, HstPtrBase, DataSize, HstPtrName, HasFlagTo,
+        HasFlagAlways, IsImplicit, UpdateRef, HasCloseModifier,
+        HasPresentModifier, HasHoldModifier, AsyncInfo, PointerTpr.getEntry());
     void *TgtPtrBegin = TPR.TargetPointer;
     IsHostPtr = TPR.Flags.IsHostPointer;
     // If data_size==0, then the argument could be a zero-length pointer to
     // NULL, so getOrAlloc() returning NULL is not an error.
     if (!TgtPtrBegin && (DataSize || HasPresentModifier)) {
       REPORT("Call to getTargetPointer returned null pointer (%s).\n",
              HasPresentModifier ? "'present' map type modifier"
                                 : "device failure or illegal mapping");
       return OFFLOAD_FAIL;
     }
     DP("There are %" PRId64 " bytes allocated at target address " DPxMOD
        " - is%s new\n",
        DataSize, DPxPTR(TgtPtrBegin), (TPR.Flags.IsNewEntry ? "" : " not"));
 
     if (ArgTypes[I] & OMP_TGT_MAPTYPE_RETURN_PARAM) {
       uintptr_t Delta = (uintptr_t)HstPtrBegin - (uintptr_t)HstPtrBase;
       void *TgtPtrBase = (void *)((uintptr_t)TgtPtrBegin - Delta);
       DP("Returning device pointer " DPxMOD "\n", DPxPTR(TgtPtrBase));
       ArgsBase[I] = TgtPtrBase;
     }
 
     if (ArgTypes[I] & OMP_TGT_MAPTYPE_PTR_AND_OBJ && !IsHostPtr) {
-      // Check whether we need to update the pointer on the device
-      bool UpdateDevPtr = false;
 
       uint64_t Delta = (uint64_t)HstPtrBegin - (uint64_t)HstPtrBase;
       void *ExpectedTgtPtrBase = (void *)((uint64_t)TgtPtrBegin - Delta);
 
-      Device.ShadowMtx.lock();
-      auto Entry = Device.ShadowPtrMap.find(PointerHstPtrBegin);
-      // If this pointer is not in the map we need to insert it. If the map
-      // contains a stale entry, we need to update it (e.g. if the pointee was
-      // deallocated and later on is reallocated at another device address). The
-      // latter scenario is the subject of LIT test env/base_ptr_ref_count.c. An
-      // entry is removed from ShadowPtrMap only when the PTR of a PTR_AND_OBJ
-      // pair is deallocated, not when the OBJ is deallocated. In
-      // env/base_ptr_ref_count.c the PTR is a global "declare target" pointer,
-      // so it stays in the map for the lifetime of the application. When the
-      // OBJ is deallocated and later on allocated again (at a different device
-      // address), ShadowPtrMap still contains an entry for Pointer_HstPtrBegin
-      // which is stale, pointing to the old ExpectedTgtPtrBase of the OBJ.
-      if (Entry == Device.ShadowPtrMap.end() ||
-          Entry->second.TgtPtrVal != ExpectedTgtPtrBase) {
-        // create or update shadow pointers for this entry
-        Device.ShadowPtrMap[PointerHstPtrBegin] = {
-            HstPtrBase, PointerTgtPtrBegin, ExpectedTgtPtrBase};
-        PointerTpr.Entry->setMayContainAttachedPointers();
-        UpdateDevPtr = true;
-      }
-
-      if (UpdateDevPtr) {
-        std::lock_guard<decltype(*PointerTpr.Entry)> LG(*PointerTpr.Entry);
-        Device.ShadowMtx.unlock();
-
+      if (PointerTpr.getEntry()->addShadowPointer(ShadowPtrInfoTy{
+              (void **)PointerHstPtrBegin, HstPtrBase,
+              (void **)PointerTgtPtrBegin, ExpectedTgtPtrBase})) {
         DP("Update pointer (" DPxMOD ") -> [" DPxMOD "]\n",
            DPxPTR(PointerTgtPtrBegin), DPxPTR(TgtPtrBegin));
 
         void *&TgtPtrBase = AsyncInfo.getVoidPtrLocation();
         TgtPtrBase = ExpectedTgtPtrBase;
 
-        int Ret = Device.submitData(PointerTgtPtrBegin, &TgtPtrBase,
-                                    sizeof(void *), AsyncInfo);
+        int Ret =
+            Device.submitData(PointerTgtPtrBegin, &TgtPtrBase, sizeof(void *),
+                              AsyncInfo, PointerTpr.getEntry());
         if (Ret != OFFLOAD_SUCCESS) {
           REPORT("Copying data to device failed.\n");
           return OFFLOAD_FAIL;
         }
-        if (PointerTpr.Entry->addEventIfNecessary(Device, AsyncInfo) !=
+        if (PointerTpr.getEntry()->addEventIfNecessary(Device, AsyncInfo) !=
             OFFLOAD_SUCCESS)
           return OFFLOAD_FAIL;
-      } else
-        Device.ShadowMtx.unlock();
+      }
     }
   }
 
   return OFFLOAD_SUCCESS;
 }
 
 namespace {
 /// This structure contains information to deallocate a target pointer, aka.
struct PostProcessingInfo {
 
   /// Index of the argument in the data mapping scheme.
   int32_t ArgIndex;
 
   /// The target pointer information.
   TargetPointerResultTy TPR;
 
   PostProcessingInfo(void *HstPtr, int64_t Size, int64_t ArgType,
-                     int32_t ArgIndex, TargetPointerResultTy TPR)
+                     int32_t ArgIndex, TargetPointerResultTy &&TPR)
       : HstPtrBegin(HstPtr), DataSize(Size), ArgType(ArgType),
-        ArgIndex(ArgIndex), TPR(TPR) {}
+        ArgIndex(ArgIndex), TPR(std::move(TPR)) {}
 };
 
-/// Apply \p CB to the shadow map pointer entries in the range \p Begin, to
-/// \p Begin + \p Size. \p CB is called with a locked shadow pointer map and the
-/// passed iterator can be updated. If the callback returns OFFLOAD_FAIL the
-/// rest of the map is not checked anymore.
-template <typename CBTy>
-static void applyToShadowMapEntries(DeviceTy &Device, CBTy CB, void *Begin,
-                                    uintptr_t Size,
-                                    const TargetPointerResultTy &TPR) {
-  // If we have an object that is too small to hold a pointer subobject, no need
-  // to do any checking.
-  if (Size < sizeof(void *))
-    return;
-
-  // If the map entry for the object was never marked as containing attached
-  // pointers, no need to do any checking.
-  if (!TPR.Entry || !TPR.Entry->getMayContainAttachedPointers())
-    return;
-
-  uintptr_t LB = (uintptr_t)Begin;
-  uintptr_t UB = LB + Size;
-  // Now we are looking into the shadow map so we need to lock it.
-  std::lock_guard<decltype(Device.ShadowMtx)> LG(Device.ShadowMtx);
-  for (ShadowPtrListTy::iterator Itr = Device.ShadowPtrMap.begin();
-       Itr != Device.ShadowPtrMap.end();) {
-    uintptr_t ShadowHstPtrAddr = (uintptr_t)Itr->first;
-
-    // An STL map is sorted on its keys; use this property
-    // to quickly determine when to break out of the loop.
-    if (ShadowHstPtrAddr < LB) {
-      ++Itr;
-      continue;
-    }
-    if (ShadowHstPtrAddr >= UB)
-      break;
-
-    if (CB(Itr) == OFFLOAD_FAIL)
-      break;
-  }
-}
-
 } // namespace
 
 /// Applies the necessary post-processing procedures to entries listed in \p
 /// EntriesInfo after the execution of all device side operations from a target
 /// data end. This includes the update of pointers at the host and removal of
 /// device buffer when needed. It returns OFFLOAD_FAIL or OFFLOAD_SUCCESS
 /// according to the successfulness of the operations.
 [[nodiscard]] static int
 postProcessingTargetDataEnd(DeviceTy *Device,
-                            SmallVector<PostProcessingInfo> EntriesInfo,
+                            SmallVector<PostProcessingInfo> &EntriesInfo,
                             bool FromMapper) {
   int Ret = OFFLOAD_SUCCESS;
   void *FromMapperBase = nullptr;
 
   for (auto &[HstPtrBegin, DataSize, ArgType, ArgIndex, TPR] : EntriesInfo) {
     bool DelEntry = !TPR.isHostPointer();
 
     // If the last element from the mapper (for end transfer args comes in
for (auto &[HstPtrBegin, DataSize, ArgType, ArgIndex, TPR] : EntriesInfo) {
 
     // If we marked the entry to be deleted we need to verify no other
     // thread reused it by now. If deletion is still supposed to happen by
     // this thread LR will be set and exclusive access to the HDTT map
     // will avoid another thread reusing the entry now. Note that we do
     // not request (exclusive) access to the HDTT map if DelEntry is
     // not set.
     DeviceTy::HDTTMapAccessorTy HDTTMap =
-        Device->HostDataToTargetMap.getExclusiveAccessor(!DelEntry);
+        Device->HostDataToTargetMap.getExclusiveAccessor();
+
+    // We cannot use a lock guard because we may end up delete the mutex.
+    // We also explicitly unlocked the entry after it was put in the EntriesInfo
+    // so it can be reused.
+    TPR.getEntry()->lock();
+    auto *Entry = TPR.getEntry();
 
-    const bool IsNotLastUser = TPR.Entry->decDataEndThreadCount() != 0;
-    if (DelEntry && (TPR.Entry->getTotalRefCount() != 0 || IsNotLastUser)) {
+    const bool IsNotLastUser = Entry->decDataEndThreadCount() != 0;
+    if (DelEntry && (Entry->getTotalRefCount() != 0 || IsNotLastUser)) {
       // The thread is not in charge of deletion anymore. Give up access
       // to the HDTT map and unset the deletion flag.
       HDTTMap.destroy();
       DelEntry = false;
     }
 
     // If we copied back to the host a struct/array containing pointers,
     // we need to restore the original host pointer values from their
     // shadow copies. If the struct is going to be deallocated, remove any
     // remaining shadow pointer entries for this struct.
     const bool HasFrom = ArgType & OMP_TGT_MAPTYPE_FROM;
-    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
-      // If we copied the struct to the host, we need to restore the
-      // pointer.
-      if (HasFrom) {
-        void **ShadowHstPtrAddr = (void **)Itr->first;
-        *ShadowHstPtrAddr = Itr->second.HstPtrVal;
+      Entry->foreachShadowPointerInfo(
+          [&](const ShadowPtrInfoTy &ShadowPtr) {
+            *ShadowPtr.HstPtrAddr = ShadowPtr.HstPtrVal;
             DP("Restoring original host pointer value " DPxMOD " for host "
                "pointer " DPxMOD "\n",
-           DPxPTR(Itr->second.HstPtrVal), DPxPTR(ShadowHstPtrAddr));
-      }
-      // If the struct is to be deallocated, remove the shadow entry.
-      if (DelEntry) {
-        DP("Removing shadow pointer " DPxMOD "\n", DPxPTR((void **)Itr->first));
-        auto OldItr = Itr;
-        Itr++;
-        Device->ShadowPtrMap.erase(OldItr);
-      } else {
-        ++Itr;
-      }
+               DPxPTR(ShadowPtr.HstPtrVal), DPxPTR(ShadowPtr.HstPtrAddr));
             return OFFLOAD_SUCCESS;
-    };
-    applyToShadowMapEntries(*Device, CB, HstPtrBegin, DataSize, TPR);
+          });
+    }
+
+    // Give up the lock as we either don't need it anymore (e.g., done with
+    // TPR), or erase TPR.
+    TPR.setEntry(nullptr);
 
     if (!DelEntry || (FromMapperBase && FromMapperBase == HstPtrBegin))
       continue;
 
-    // If we are deleting the entry the DataMapMtx is locked and we own
-    // the entry.
-    Ret = Device->eraseMapEntry(HDTTMap, TPR.Entry, DataSize);
+    Ret = Device->eraseMapEntry(HDTTMap, Entry, DataSize);
     // Entry is already remove from the map, we can unlock it now.
     HDTTMap.destroy();
-    Ret |= Device->deallocTgtPtrAndEntry(TPR.Entry, DataSize);
+    Ret |= Device->deallocTgtPtrAndEntry(Entry, DataSize);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Deallocating data from device failed.\n");
       break;
     }
   }
 
+  delete &EntriesInfo;
   return Ret;
 }
 
 /// Internal function to undo the mapping and retrieve the data from the device.
 int targetDataEnd(ident_t *Loc, DeviceTy &Device, int32_t ArgNum,
                   void **ArgBases, void **Args, int64_t *ArgSizes,
                   int64_t *ArgTypes, map_var_info_t *ArgNames,
                   void **ArgMappers, AsyncInfoTy &AsyncInfo, bool FromMapper) {
   int Ret = OFFLOAD_SUCCESS;
-  SmallVector<PostProcessingInfo> PostProcessingPtrs;
+  auto *PostProcessingPtrs = new SmallVector<PostProcessingInfo>();
   void *FromMapperBase = nullptr;
   // process each input.
   for (int32_t I = ArgNum - 1; I >= 0; --I) {
     // Ignore private variables and arrays - there is no mapping for them.
     // Also, ignore the use_device_ptr directive, it has no effect here.
     if ((ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) ||
         (ArgTypes[I] & OMP_TGT_MAPTYPE_PRIVATE))
       continue;
if (getParentIndex(ArgTypes[I]) < 0 && NextI < ArgNum &&
         DP("Using a Padding of %" PRId64 " bytes for begin address " DPxMOD
            "\n",
            Padding, DPxPTR(HstPtrBegin));
         HstPtrBegin = (char *)HstPtrBegin - Padding;
         DataSize += Padding;
       }
     }
 
-    bool IsLast, IsHostPtr;
     bool IsImplicit = ArgTypes[I] & OMP_TGT_MAPTYPE_IMPLICIT;
     bool UpdateRef = (!(ArgTypes[I] & OMP_TGT_MAPTYPE_MEMBER_OF) ||
                       (ArgTypes[I] & OMP_TGT_MAPTYPE_PTR_AND_OBJ)) &&
                      !(FromMapper && I == 0);
     bool ForceDelete = ArgTypes[I] & OMP_TGT_MAPTYPE_DELETE;
     bool HasPresentModifier = ArgTypes[I] & OMP_TGT_MAPTYPE_PRESENT;
     bool HasHoldModifier = ArgTypes[I] & OMP_TGT_MAPTYPE_OMPX_HOLD;
 
     // If PTR_AND_OBJ, HstPtrBegin is address of pointee
-    TargetPointerResultTy TPR = Device.getTgtPtrBegin(
-        HstPtrBegin, DataSize, IsLast, UpdateRef, HasHoldModifier, IsHostPtr,
+    TargetPointerResultTy TPR =
+        Device.getTgtPtrBegin(HstPtrBegin, DataSize, UpdateRef, HasHoldModifier,
                               !IsImplicit, ForceDelete, /*FromDataEnd=*/true);
     void *TgtPtrBegin = TPR.TargetPointer;
     if (!TPR.isPresent() && !TPR.isHostPointer() &&
         (DataSize || HasPresentModifier)) {
       DP("Mapping does not exist (%s)\n",
          (HasPresentModifier ? "'present' map type modifier" : "ignored"));
       if (HasPresentModifier) {
         // OpenMP 5.1, sec. 2.21.7.1 "map Clause", p. 350 L10-13:
         // "If a map clause appears on a target, target data, target enter data
if (!TPR.isPresent() && !TPR.isHostPointer() &&
         MESSAGE("device mapping required by 'present' map type modifier does "
                 "not exist for host address " DPxMOD " (%" PRId64 " bytes)",
                 DPxPTR(HstPtrBegin), DataSize);
         return OFFLOAD_FAIL;
       }
     } else {
       DP("There are %" PRId64 " bytes allocated at target address " DPxMOD
          " - is%s last\n",
-         DataSize, DPxPTR(TgtPtrBegin), (IsLast ? "" : " not"));
+         DataSize, DPxPTR(TgtPtrBegin), (TPR.Flags.IsLast ? "" : " not"));
     }
 
     // OpenMP 5.1, sec. 2.21.7.1 "map Clause", p. 351 L14-16:
     // "If the map clause appears on a target, target data, or target exit data
     // construct and a corresponding list item of the original list item is not
     // present in the device data environment on exit from the region then the
     // list item is ignored."
     if (!TPR.isPresent())
       continue;
 
     // Move data back to the host
     const bool HasAlways = ArgTypes[I] & OMP_TGT_MAPTYPE_ALWAYS;
     const bool HasFrom = ArgTypes[I] & OMP_TGT_MAPTYPE_FROM;
-    if (HasFrom && (HasAlways || IsLast) && !IsHostPtr && DataSize != 0) {
+    if (HasFrom && (HasAlways || TPR.Flags.IsLast) &&
+        !TPR.Flags.IsHostPointer && DataSize != 0) {
       DP("Moving %" PRId64 " bytes (tgt:" DPxMOD ") -> (hst:" DPxMOD ")\n",
          DataSize, DPxPTR(TgtPtrBegin), DPxPTR(HstPtrBegin));
 
-      std::lock_guard<decltype(*TPR.Entry)> LG(*TPR.Entry);
       // Wait for any previous transfer if an event is present.
-      if (void *Event = TPR.Entry->getEvent()) {
+      if (void *Event = TPR.getEntry()->getEvent()) {
         if (Device.waitEvent(Event, AsyncInfo) != OFFLOAD_SUCCESS) {
           REPORT("Failed to wait for event " DPxMOD ".\n", DPxPTR(Event));
           return OFFLOAD_FAIL;
         }
       }
 
-      Ret = Device.retrieveData(HstPtrBegin, TgtPtrBegin, DataSize, AsyncInfo);
+      Ret = Device.retrieveData(HstPtrBegin, TgtPtrBegin, DataSize, AsyncInfo,
+                                TPR.getEntry());
       if (Ret != OFFLOAD_SUCCESS) {
         REPORT("Copying data from device failed.\n");
         return OFFLOAD_FAIL;
       }
 
       // As we are expecting to delete the entry the d2h copy might race
       // with another one that also tries to delete the entry. This happens
       // as the entry can be reused and the reuse might happen after the
       // copy-back was issued but before it completed. Since the reuse might
       // also copy-back a value we would race.
-      if (IsLast) {
-        if (TPR.Entry->addEventIfNecessary(Device, AsyncInfo) !=
+      if (TPR.Flags.IsLast) {
+        if (TPR.getEntry()->addEventIfNecessary(Device, AsyncInfo) !=
             OFFLOAD_SUCCESS)
           return OFFLOAD_FAIL;
       }
     }
 
     // Add pointer to the buffer for post-synchronize processing.
-    PostProcessingPtrs.emplace_back(HstPtrBegin, DataSize, ArgTypes[I], I, TPR);
+    PostProcessingPtrs->emplace_back(HstPtrBegin, DataSize, ArgTypes[I], I,
+                                     std::move(TPR));
+    PostProcessingPtrs->back().TPR.getEntry()->unlock();
   }
 
   // Add post-processing functions
   // TODO: We might want to remove `mutable` in the future by not changing the
   // captured variables somehow.
-  AsyncInfo.addPostProcessingFunction(
-      [=, Device = &Device,
-       PostProcessingPtrs = std::move(PostProcessingPtrs)]() mutable -> int {
-        return postProcessingTargetDataEnd(Device, PostProcessingPtrs,
-                                           FromMapperBase);
-      });
+  AsyncInfo.addPostProcessingFunction([=, Device = &Device]() mutable -> int {
+    return postProcessingTargetDataEnd(Device, *PostProcessingPtrs,
 
   return Ret;
 }
 
 static int targetDataContiguous(ident_t *Loc, DeviceTy &Device, void *ArgsBase,
                                 void *HstPtrBegin, int64_t ArgSize,
                                 int64_t ArgType, AsyncInfoTy &AsyncInfo) {
   TIMESCOPE_WITH_IDENT(Loc);
-  bool IsLast, IsHostPtr;
-  TargetPointerResultTy TPR = Device.getTgtPtrBegin(
-      HstPtrBegin, ArgSize, IsLast, /*UpdateRefCount=*/false,
-      /*UseHoldRefCount=*/false, IsHostPtr, /*MustContain=*/true);
+  TargetPointerResultTy TPR =
+      Device.getTgtPtrBegin(HstPtrBegin, ArgSize, /*UpdateRefCount=*/false,
+                            /*UseHoldRefCount=*/false, /*MustContain=*/true);
   void *TgtPtrBegin = TPR.TargetPointer;
   if (!TPR.isPresent()) {
     DP("hst data:" DPxMOD " not found, becomes a noop\n", DPxPTR(HstPtrBegin));
     if (ArgType & OMP_TGT_MAPTYPE_PRESENT) {
       MESSAGE("device mapping required by 'present' motion modifier does not "
               "exist for host address " DPxMOD " (%" PRId64 " bytes)",
               DPxPTR(HstPtrBegin), ArgSize);
       return OFFLOAD_FAIL;
     }
     return OFFLOAD_SUCCESS;
   }
 
-  if (IsHostPtr) {
+  if (TPR.Flags.IsHostPointer) {
     DP("hst data:" DPxMOD " unified and shared, becomes a noop\n",
        DPxPTR(HstPtrBegin));
     return OFFLOAD_SUCCESS;
   }
 
+  if (ArgType & OMP_TGT_MAPTYPE_TO) {
+    DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n",
+       ArgSize, DPxPTR(HstPtrBegin), DPxPTR(TgtPtrBegin));
+    int Ret = Device.submitData(TgtPtrBegin, HstPtrBegin, ArgSize, AsyncInfo,
+                                TPR.getEntry());
+    if (Ret != OFFLOAD_SUCCESS) {
+      REPORT("Copying data to device failed.\n");
+      return OFFLOAD_FAIL;
+    }
+    if (TPR.getEntry()) {
+      int Ret = TPR.getEntry()->foreachShadowPointerInfo(
+          [&](const ShadowPtrInfoTy &ShadowPtr) {
+            DP("Restoring original target pointer value " DPxMOD " for target "
+               "pointer " DPxMOD "\n",
+               DPxPTR(ShadowPtr.TgtPtrVal), DPxPTR(ShadowPtr.TgtPtrAddr));
+            Ret = Device.submitData(ShadowPtr.TgtPtrAddr,
+                                    (void *)&ShadowPtr.TgtPtrVal,
+                                    sizeof(void *), AsyncInfo);
+            if (Ret != OFFLOAD_SUCCESS) {
+              REPORT("Copying data to device failed.\n");
+              return OFFLOAD_FAIL;
+            }
+            return OFFLOAD_SUCCESS;
+          });
+      if (Ret != OFFLOAD_SUCCESS) {
+        DP("Updating shadow map failed\n");
+        return Ret;
+      }
+    }
+  }
+
   if (ArgType & OMP_TGT_MAPTYPE_FROM) {
     DP("Moving %" PRId64 " bytes (tgt:" DPxMOD ") -> (hst:" DPxMOD ")\n",
        ArgSize, DPxPTR(TgtPtrBegin), DPxPTR(HstPtrBegin));
-    int Ret = Device.retrieveData(HstPtrBegin, TgtPtrBegin, ArgSize, AsyncInfo);
+    int Ret = Device.retrieveData(HstPtrBegin, TgtPtrBegin, ArgSize, AsyncInfo,
+                                  TPR.getEntry());
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Copying data from device failed.\n");
       return OFFLOAD_FAIL;
     }
 
     // Wait for device-to-host memcopies for whole struct to complete,
     // before restoring the correct host pointer.
-    AsyncInfo.addPostProcessingFunction([=, Device = &Device]() -> int {
-      auto CB = [&](ShadowPtrListTy::iterator &Itr) {
-        void **ShadowHstPtrAddr = (void **)Itr->first;
-        *ShadowHstPtrAddr = Itr->second.HstPtrVal;
+    if (auto *Entry = TPR.getEntry()) {
+      AsyncInfo.addPostProcessingFunction([=]() -> int {
+        int Ret = Entry->foreachShadowPointerInfo(
+            [&](const ShadowPtrInfoTy &ShadowPtr) {
+              *ShadowPtr.HstPtrAddr = ShadowPtr.HstPtrVal;
               DP("Restoring original host pointer value " DPxMOD
                  " for host pointer " DPxMOD "\n",
-           DPxPTR(Itr->second.HstPtrVal), DPxPTR(ShadowHstPtrAddr));
-        ++Itr;
+                 DPxPTR(ShadowPtr.HstPtrVal), DPxPTR(ShadowPtr.HstPtrAddr));
               return OFFLOAD_SUCCESS;
-      };
-      applyToShadowMapEntries(*Device, CB, HstPtrBegin, ArgSize, TPR);
-
+            });
+        Entry->unlock();
+        if (Ret != OFFLOAD_SUCCESS) {
+          DP("Updating shadow map failed\n");
+          return Ret;
+        }
         return OFFLOAD_SUCCESS;
       });
     }
-
-  if (ArgType & OMP_TGT_MAPTYPE_TO) {
-    DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n",
-       ArgSize, DPxPTR(HstPtrBegin), DPxPTR(TgtPtrBegin));
-    int Ret = Device.submitData(TgtPtrBegin, HstPtrBegin, ArgSize, AsyncInfo);
-    if (Ret != OFFLOAD_SUCCESS) {
-      REPORT("Copying data to device failed.\n");
-      return OFFLOAD_FAIL;
-    }
 
-    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
-      DP("Restoring original target pointer value " DPxMOD " for target "
-         "pointer " DPxMOD "\n",
-         DPxPTR(Itr->second.TgtPtrVal), DPxPTR(Itr->second.TgtPtrAddr));
-      Ret = Device.submitData(Itr->second.TgtPtrAddr, &Itr->second.TgtPtrVal,
-                              sizeof(void *), AsyncInfo);
-      if (Ret != OFFLOAD_SUCCESS)
-        REPORT("Copying data to device failed.\n");
-      ++Itr;
-      return Ret;
-    };
-    applyToShadowMapEntries(Device, CB, HstPtrBegin, ArgSize, TPR);
-  }
   return OFFLOAD_SUCCESS;
 }
 
 static int targetDataNonContiguous(ident_t *Loc, DeviceTy &Device,
                                    void *ArgsBase,
                                    __tgt_target_non_contig *NonContig,
                                    uint64_t Size, int64_t ArgType,
                                    int CurrentDim, int DimSize, uint64_t Offset,
if (!(ArgTypes[I] & OMP_TGT_MAPTYPE_TARGET_PARAM)) {
         unsigned Idx = getParentIndex(ArgTypes[I]);
         int TgtIdx = TgtArgsPositions[Idx];
         assert(TgtIdx != -1 && "Base address must be translated already.");
         // The parent lambda must be processed already and it must be the last
         // in TgtArgs and TgtOffsets arrays.
         void *HstPtrVal = Args[I];
         void *HstPtrBegin = ArgBases[I];
         void *HstPtrBase = Args[Idx];
-        bool IsLast, IsHostPtr; // IsLast is unused.
         void *TgtPtrBase =
             (void *)((intptr_t)TgtArgs[TgtIdx] + TgtOffsets[TgtIdx]);
         DP("Parent lambda base " DPxMOD "\n", DPxPTR(TgtPtrBase));
         uint64_t Delta = (uint64_t)HstPtrBegin - (uint64_t)HstPtrBase;
         void *TgtPtrBegin = (void *)((uintptr_t)TgtPtrBase + Delta);
         void *&PointerTgtPtrBegin = AsyncInfo.getVoidPtrLocation();
         TargetPointerResultTy TPR = Device.getTgtPtrBegin(
-            HstPtrVal, ArgSizes[I], IsLast, /*UpdateRefCount=*/false,
-            /*UseHoldRefCount=*/false, IsHostPtr);
+            HstPtrVal, ArgSizes[I], /*UpdateRefCount=*/false,
+            /*UseHoldRefCount=*/false);
         PointerTgtPtrBegin = TPR.TargetPointer;
         if (!TPR.isPresent()) {
           DP("No lambda captured variable mapped (" DPxMOD ") - ignored\n",
              DPxPTR(HstPtrVal));
           continue;
         }
-        if (IsHostPtr) {
+        if (TPR.Flags.IsHostPointer) {
           DP("Unified memory is active, no need to map lambda captured"
              "variable (" DPxMOD ")\n",
              DPxPTR(HstPtrVal));
           continue;
         }
         DP("Update lambda reference (" DPxMOD ") -> [" DPxMOD "]\n",
            DPxPTR(PointerTgtPtrBegin), DPxPTR(TgtPtrBegin));
         Ret = Device.submitData(TgtPtrBegin, &PointerTgtPtrBegin,
-                                sizeof(void *), AsyncInfo);
+                                sizeof(void *), AsyncInfo, TPR.getEntry());
         if (Ret != OFFLOAD_SUCCESS) {
           REPORT("Copying data to device failed.\n");
           return OFFLOAD_FAIL;
         }
       }
       continue;
     }
     void *HstPtrBegin = Args[I];
     void *HstPtrBase = ArgBases[I];
     void *TgtPtrBegin;
     map_var_info_t HstPtrName = (!ArgNames) ? nullptr : ArgNames[I];
     ptrdiff_t TgtBaseOffset;
-    bool IsLast, IsHostPtr; // unused.
     TargetPointerResultTy TPR;
     if (ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) {
       DP("Forwarding first-private value " DPxMOD " to the target construct\n",
          DPxPTR(HstPtrBase));
       TgtPtrBegin = HstPtrBase;
       TgtBaseOffset = 0;
     } else if (ArgTypes[I] & OMP_TGT_MAPTYPE_PRIVATE) {
       TgtBaseOffset = (intptr_t)HstPtrBase - (intptr_t)HstPtrBegin;
if (ArgTypes[I] & OMP_TGT_MAPTYPE_LITERAL) {
       if (Ret != OFFLOAD_SUCCESS) {
         REPORT("Failed to process %sprivate argument " DPxMOD "\n",
                (IsFirstPrivate ? "first-" : ""), DPxPTR(HstPtrBegin));
         return OFFLOAD_FAIL;
       }
     } else {
       if (ArgTypes[I] & OMP_TGT_MAPTYPE_PTR_AND_OBJ)
         HstPtrBase = *reinterpret_cast<void **>(HstPtrBase);
-      TPR = Device.getTgtPtrBegin(HstPtrBegin, ArgSizes[I], IsLast,
+      TPR = Device.getTgtPtrBegin(HstPtrBegin, ArgSizes[I],
                                   /*UpdateRefCount=*/false,
-                                  /*UseHoldRefCount=*/false, IsHostPtr);
+                                  /*UseHoldRefCount=*/false);
       TgtPtrBegin = TPR.TargetPointer;
       TgtBaseOffset = (intptr_t)HstPtrBase - (intptr_t)HstPtrBegin;
 #ifdef OMPTARGET_DEBUG
       void *TgtPtrBase = (void *)((intptr_t)TgtPtrBegin + TgtBaseOffset);
       DP("Obtained target argument " DPxMOD " from host pointer " DPxMOD "\n",
          DPxPTR(TgtPtrBase), DPxPTR(HstPtrBegin));
 #endif
     }