[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

92

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

[jdoerfert]

openmp/libomptarget/include/device.h

92

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

92

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

92

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

[tianshilei1992]

openmp/libomptarget/include/device.h

92

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

348

Please delete copy constructor and copy assign operator.

376

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

419

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

373

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

348

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

373

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.

376

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.

419

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

348

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

376

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

419

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

 // 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;
+};
+
 /// 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.
struct StatesTy {
     /// 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;
+    /// A map of shadow pointers associated with this entry, the keys are host
+    /// pointer addresses to identify stale entries.
+    std::map<void **, ShadowPtrInfoTy> ShadowPtrInfos;
 
     /// 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;
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;
+  /// 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.emplace(ShadowPtrInfo.HstPtrAddr, ShadowPtrInfo);
+    if (Pair.second)
+      return true;
+    if (Pair.first->second.TgtPtrVal == ShadowPtrInfo.TgtPtrVal)
+      return false;
+    Pair.first->second = ShadowPtrInfo;
+    return true;
   }
-  bool getMayContainAttachedPointers() const {
-    return States->MayContainAttachedPointers;
+
+  /// 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.second) == OFFLOAD_FAIL)
+        return OFFLOAD_FAIL;
+    return OFFLOAD_SUCCESS;
   }
 
   /// Increment the delete counter indicating that this thread plans to delete
   /// the entry.
   void incDeleteCount() const { ++States->DeleteCount; }
 
   /// Decrement the delete counter prior to deletion. If the return value is 0
   /// deletion should take place, otherwise the entry is still needed.
   unsigned decDeleteCount() const { return --States->DeleteCount; }
 
   void lock() const { States->UpdateMtx.lock(); }
 
   void unlock() const { States->UpdateMtx.unlock(); }
 };
 
 /// 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;
 }
struct LookupResult {
 
   /// 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.
 class TargetPointerResultTy {
   struct FlagTy {
     /// Flag indicating that the map table entry was just created.
     unsigned IsNewEntry : 1;
     /// Flag indicating that the pointer is actually a host pointer, used when
     /// unified memory enabled.
     unsigned IsHostPointer : 1;
     /// 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};
 
   /// The corresponding map table entry which is stable.
   HostDataToTargetTy *Entry = nullptr;
 
   /// The corresponding target pointer
   void *TargetPointer = nullptr;
 
 public:
+  TargetPointerResultTy(const TargetPointerResultTy &) = delete;
+  TargetPointerResultTy &operator=(const TargetPointerResultTy &TPR) = delete;
+
   TargetPointerResultTy() {}
 
   TargetPointerResultTy(FlagTy Flags, HostDataToTargetTy *Entry,
                         void *TargetPointer)
-      : Flags(Flags), Entry(Entry), TargetPointer(TargetPointer) {}
+      : Flags(Flags), Entry(Entry), TargetPointer(TargetPointer) {
+    if (Entry)
+      Entry->lock();
+  }
+
+  TargetPointerResultTy(TargetPointerResultTy &&TPR)
+      : Flags(TPR.Flags), Entry(TPR.Entry), TargetPointer(TPR.TargetPointer) {
+    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 isHostPtr() const { return Flags.IsHostPointer; }
   void setIsHostPtr(bool IHP) { Flags.IsHostPointer = IHP; }
 
   bool isNew() const { return Flags.IsNewEntry; }
   void setIsNew(bool IN) { Flags.IsNewEntry = IN; }
 
   bool isLast() const { return Flags.IsLast; }
   void setIsLast(bool IL) { Flags.IsLast = IL; }
 
   void *getTargetPointer() const { return TargetPointer; }
   void setTargetPointer(void *TP) { TargetPointer = TP; }
 
   HostDataToTargetTy *getEntry() const { return Entry; }
-  void setEntry(HostDataToTargetTy *HDTTT) { Entry = HDTTT; }
+  void setEntry(HostDataToTargetTy *HDTTT, bool Lock = true) {
+    if (Entry)
+      Entry->unlock();
+    Entry = HDTTT;
+    if (Entry && Lock)
+      Entry->lock();
+  }
+
+  void reset() { *this = TargetPointerResultTy(); }
 };
 
 /// Map for shadow pointers
 struct ShadowPtrValTy {
   void *HstPtrVal;
   void *TgtPtrAddr;
   void *TgtPtrVal;
 };
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;
 
   // NOTE: Once libomp gains full target-task support, this state should be
   // moved into the target task in libomp.
   std::map<int32_t, uint64_t> LoopTripCnt;
 
   DeviceTy(RTLInfoTy *RTL);
   // DeviceTy is not copyable
   DeviceTy(const DeviceTy &D) = delete;
struct DeviceTy {
   /// 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(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);
+  TargetPointerResultTy getTargetPointer(
+      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, 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);
 
   TargetPointerResultTy getTgtPtrBegin(void *HstPtrBegin, int64_t Size,
                                        bool UpdateRefCount,
struct DeviceTy {
   /// Deallocates memory which \p TgtPtrBegin points at and returns
   /// OFFLOAD_SUCCESS/OFFLOAD_FAIL when succeeds/fails.
   int32_t deleteData(void *TgtPtrBegin);
 
   // 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);
 
   int32_t runRegion(void *TgtEntryPtr, void **TgtVarsPtr, ptrdiff_t *TgtOffsets,
                     int32_t TgtVarsSize, AsyncInfoTy &AsyncInfo);
   int32_t runTeamRegion(void *TgtEntryPtr, void **TgtVarsPtr,
                         ptrdiff_t *TgtOffsets, int32_t TgtVarsSize,

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) {
 
   return OFFLOAD_SUCCESS;
 }
 
 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);
-      return OFFLOAD_SUCCESS;
-    } else {
-      REPORT("Trying to disassociate a pointer which was not mapped via "
-             "omp_target_associate_ptr\n");
-    }
-  } else {
-    REPORT("Association not found\n");
+  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()) {
   }
 
-  // Mapping not found
+  REPORT("Trying to disassociate a pointer which was not mapped via "
+         "omp_target_associate_ptr\n");
   return OFFLOAD_FAIL;
 }
 
 LookupResult DeviceTy::lookupMapping(HDTTMapAccessorTy &HDTTMap,
                                      void *HstPtrBegin, int64_t Size) {
 
   uintptr_t hp = (uintptr_t)HstPtrBegin;
   LookupResult lr;
LookupResult DeviceTy::lookupMapping(HDTTMapAccessorTy &HDTTMap,
 
   return lr;
 }
 
 TargetPointerResultTy DeviceTy::getTargetPointer(
     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) {
+    bool HasPresentModifier, bool HasHoldModifier, AsyncInfoTy &AsyncInfo,
+    HostDataToTargetTy *OwnedTPR, bool ReleaseHDTTMap) {
   TargetPointerResultTy TPR;
 
   LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
 
-  TPR.setEntry(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.
+  TPR.setEntry(LR.Entry, /* Lock */ LR.Entry != OwnedTPR);
 
   // Not that the entry can be null.
   HostDataToTargetTy *Entry = TPR.getEntry();
 
   // 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 ||
     TPR.setIsNew(true);
     uintptr_t Ptr = (uintptr_t)allocData(Size, HstPtrBegin);
     Entry = HDTTMap
                 ->emplace(new HostDataToTargetTy(
                     (uintptr_t)HstPtrBase, (uintptr_t)HstPtrBegin,
                     (uintptr_t)HstPtrBegin + Size, Ptr, HasHoldModifier,
                     HstPtrName))
                 .first->HDTT;
+    // Release the mapping table lock only after the entry is locked by
+    // attaching it to TPR.
     TPR.setEntry(Entry);
 
     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(),
          (HstPtrName) ? getNameFromMapping(HstPtrName).c_str() : "unknown");
     TPR.setTargetPointer((void *)Ptr);
   }
 
+  // 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 (TPR.getTargetPointer() && !TPR.isHostPtr() && HasFlagTo &&
       (TPR.isNew() || HasFlagAlways)) {
     DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n", Size,
        DPxPTR(HstPtrBegin), DPxPTR(TPR.getTargetPointer()));
 
     int Ret = submitData(TPR.getTargetPointer(), HstPtrBegin, Size, AsyncInfo);
if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {
     return (void *)tp;
   }
 
   return NULL;
 }
 
 int DeviceTy::deallocTgtPtr(HDTTMapAccessorTy &HDTTMap, LookupResult LR,
                             int64_t Size) {
+  HostDataToTargetTy &HT = *LR.Entry;
+
   // Check if the pointer is contained in any sub-nodes.
   if (!(LR.Flags.IsContained || LR.Flags.ExtendsBefore ||
         LR.Flags.ExtendsAfter)) {
     REPORT("Section to delete (hst addr " DPxMOD ") does not exist in the"
            " allocated memory\n",
-           DPxPTR(LR.Entry->HstPtrBegin));
+           DPxPTR(HT.HstPtrBegin));
     return OFFLOAD_FAIL;
   }
 
-  auto &HT = *LR.Entry;
-  // Verify this thread is still in charge of deleting the entry.
+  DP("Deleting tgt data " DPxMOD " of size %" PRId64 "\n",
+     DPxPTR(HT.TgtPtrBegin), Size);
+
   assert(HT.getTotalRefCount() == 0 &&
          "Trying to delete entry that is in use or owned by another thread.");
 
-  DP("Deleting tgt data " DPxMOD " of size %" PRId64 "\n",
-     DPxPTR(HT.TgtPtrBegin), Size);
   deleteData((void *)HT.TgtPtrBegin);
+
   INFO(OMP_INFOTYPE_MAPPING_CHANGED, DeviceID,
        "Removing map entry with HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD
        ", Size=%" PRId64 ", Name=%s\n",
        DPxPTR(HT.HstPtrBegin), DPxPTR(HT.TgtPtrBegin), Size,
        (HT.HstPtrName) ? getNameFromMapping(HT.HstPtrName).c_str() : "unknown");
-  void *Event = LR.Entry->getEvent();
+  void *Event = HT.getEvent();
   HDTTMap->erase(LR.Entry);
-  delete LR.Entry;
+  delete &HT;
 
   int Ret = OFFLOAD_SUCCESS;
   if (Event && destroyEvent(Event) != OFFLOAD_SUCCESS) {
     REPORT("Failed to destroy event " DPxMOD "\n", DPxPTR(Event));
     Ret = OFFLOAD_FAIL;
   }
 
   return Ret;

 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) {
   return RTL->data_delete(RTLDeviceID, TgtPtrBegin);
 }
 
+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 %s to %s, %sPtr=" DPxMOD ", %sPtr=" DPxMOD
+       ", Size=%" PRId64 ", Name=%s\n",
+       H2D ? "host" : "device", H2D ? "device" : "host", H2D ? "Hst" : "Tgt",
+       DPxPTR(SrcPtrBegin), H2D ? "Tgt" : "Hst", DPxPTR(DstPtrBegin), Size,
+       (HT && HT->HstPtrName) ? getNameFromMapping(HT->HstPtrName).c_str()
+                              : "unknown");
+}
+
 // Submit data to device
 int32_t DeviceTy::submitData(void *TgtPtrBegin, void *HstPtrBegin, int64_t Size,
-                             AsyncInfoTy &AsyncInfo) {
+                             AsyncInfoTy &AsyncInfo,
+                             HostDataToTargetTy *Entry) {
   if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
+    if (!Entry) {
       HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
       LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
-    auto *HT = &*LR.Entry;
-
-    INFO(OMP_INFOTYPE_DATA_TRANSFER, DeviceID,
-         "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD
-         ", Size=%" PRId64 ", Name=%s\n",
-         DPxPTR(HstPtrBegin), DPxPTR(TgtPtrBegin), Size,
-         (HT && HT->HstPtrName) ? getNameFromMapping(HT->HstPtrName).c_str()
-                                : "unknown");
+      Entry = LR.Entry;
+    }
+    printCopyInfo(DeviceID, /* H2D */ true, TgtPtrBegin, HstPtrBegin, Size,
+                  Entry);
   }
 
   if (!AsyncInfo || !RTL->data_submit_async || !RTL->synchronize)
     return RTL->data_submit(RTLDeviceID, TgtPtrBegin, HstPtrBegin, Size);
-  else
-    return RTL->data_submit_async(RTLDeviceID, TgtPtrBegin, HstPtrBegin, Size,
-                                  AsyncInfo);
 }
 
 // Retrieve data from device
 int32_t DeviceTy::retrieveData(void *HstPtrBegin, void *TgtPtrBegin,
-                               int64_t Size, AsyncInfoTy &AsyncInfo) {
+                               int64_t Size, AsyncInfoTy &AsyncInfo,
+                               HostDataToTargetTy *Entry) {
   if (getInfoLevel() & OMP_INFOTYPE_DATA_TRANSFER) {
+    if (!Entry) {
       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");
+      Entry = LR.Entry;
+    }
+    printCopyInfo(DeviceID, /* H2D */ false, TgtPtrBegin, HstPtrBegin, Size,
+                  Entry);
   }
-
   if (!RTL->data_retrieve_async || !RTL->synchronize)
     return RTL->data_retrieve(RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size);
-  else
-    return RTL->data_retrieve_async(RTLDeviceID, HstPtrBegin, TgtPtrBegin, Size,
-                                    AsyncInfo);
 }
 
 // Copy data from current device to destination device directly
 int32_t DeviceTy::dataExchange(void *SrcPtr, DeviceTy &DstDev, void *DstPtr,
                                int64_t Size, AsyncInfoTy &AsyncInfo) {
   if (!AsyncInfo || !RTL->data_exchange_async || !RTL->synchronize) {
     assert(RTL->data_exchange && "RTL->data_exchange is nullptr");
     return RTL->data_exchange(RTLDeviceID, SrcPtr, DstDev.RTLDeviceID, DstPtr,
                               Size);
-  } else
-    return RTL->data_exchange_async(RTLDeviceID, SrcPtr, DstDev.RTLDeviceID,
-                                    DstPtr, Size, AsyncInfo);
+  }
 }
 
 // Run region on device
 int32_t DeviceTy::runRegion(void *TgtEntryPtr, void **TgtVarsPtr,
                             ptrdiff_t *TgtOffsets, int32_t TgtVarsSize,
                             AsyncInfoTy &AsyncInfo) {
   if (!RTL->run_region || !RTL->synchronize)
     return RTL->run_region(RTLDeviceID, TgtEntryPtr, TgtVarsPtr, TgtOffsets,

openmp/libomptarget/src/omptarget.cpp

if (arg_types[i] & OMP_TGT_MAPTYPE_PTR_AND_OBJ) {
       // 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.
       Pointer_TPR = Device.getTargetPointer(
           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 = (void **)Pointer_TPR.getTargetPointer();
       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),
          (Pointer_TPR.isNew() ? "" : " not"));
       PointerHstPtrBegin = (void **)HstPtrBase;
       // modify current entry.
       HstPtrBase = *PointerHstPtrBegin;
       // 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 = arg_types[i] & OMP_TGT_MAPTYPE_TO;
     const bool HasFlagAlways = arg_types[i] & OMP_TGT_MAPTYPE_ALWAYS;
+    // Note that HDTTMap will be released in getTargetPointer.
     auto TPR = Device.getTargetPointer(
         HDTTMap, HstPtrBegin, HstPtrBase, data_size, HstPtrName, HasFlagTo,
         HasFlagAlways, IsImplicit, UpdateRef, HasCloseModifier,
-        HasPresentModifier, HasHoldModifier, AsyncInfo);
+        HasPresentModifier, HasHoldModifier, AsyncInfo, Pointer_TPR.getEntry());
 
     void *TgtPtrBegin = TPR.getTargetPointer();
     // 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 && (data_size || 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",
        data_size, DPxPTR(TgtPtrBegin), (TPR.isNew() ? "" : " not"));
 
     if (arg_types[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));
       args_base[i] = TgtPtrBase;
     }
 
     if (arg_types[i] & OMP_TGT_MAPTYPE_PTR_AND_OBJ && !TPR.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 PointerHstPtrBegin
-      // 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};
-        Pointer_TPR.getEntry()->setMayContainAttachedPointers();
-        UpdateDevPtr = true;
-      }
-
-      if (UpdateDevPtr) {
-        std::lock_guard<decltype(*Pointer_TPR.getEntry())> LG(
-            *Pointer_TPR.getEntry());
-        Device.ShadowMtx.unlock();
-
+      if (Pointer_TPR.getEntry()->addShadowPointer(
+              ShadowPtrInfoTy{PointerHstPtrBegin, HstPtrBase,
+                              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);
+                                    sizeof(void *), AsyncInfo, TPR.getEntry());
         if (Ret != OFFLOAD_SUCCESS) {
           REPORT("Copying data to device failed.\n");
           return OFFLOAD_FAIL;
         }
         if (Pointer_TPR.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 {
   HostDataToTargetTy *Entry;
 
   PostProcessingInfo(void *HstPtr, int64_t Size, int64_t ArgType, bool DelEntry,
                      HostDataToTargetTy *Entry)
       : HstPtrBegin(HstPtr), DataSize(Size), ArgType(ArgType),
         DelEntry(DelEntry), Entry(Entry) {}
 };
 
-/// 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, HostDataToTargetTy *Entry) {
-  // 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 (!Entry || !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
 
 /// 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;
if ((ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) || DelEntry) {
           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);
+                                    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
int targetDataEnd(ident_t *loc, DeviceTy &Device, int32_t ArgNum,
   std::set<HostDataToTargetTy *> DeletedEntries;
   for (PostProcessingInfo &Info : PostProcessingPtrs) {
     if (DeletedEntries.count(Info.Entry))
       continue;
 
     // 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 Info.DelEntry is not set.
-    LookupResult LR;
+    // thread reusing the entry now.
     DeviceTy::HDTTMapAccessorTy HDTTMap =
-        Device.HostDataToTargetMap.getExclusiveAccessor(!Info.DelEntry);
+        Device.HostDataToTargetMap.getExclusiveAccessor();
+    LookupResult LR;
+
+    // We cannot use a lock guard because we may end up delete the mutex.
+    Info.Entry->lock();
 
     if (Info.DelEntry) {
       LR = Device.lookupMapping(HDTTMap, Info.HstPtrBegin, Info.DataSize);
       assert(LR.Entry == Info.Entry && "Entry disappeared unexpectedly");
       if (LR.Entry->getTotalRefCount() != 0 ||
           LR.Entry->decDeleteCount() != 0) {
         // The thread is not in charge of deletion anymore. Give up access to
         // the HDTT map and unset the deletion flag.
         HDTTMap.destroy();
         Info.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.
-    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
-      // If we copied the struct to the host, we need to restore the pointer.
-      if (Info.ArgType & OMP_TGT_MAPTYPE_FROM) {
-        void **ShadowHstPtrAddr = (void **)Itr->first;
-        *ShadowHstPtrAddr = Itr->second.HstPtrVal;
+      // If we copied the struct to the host, we need to restore the pointer.
+      Info.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 (Info.DelEntry) {
-        DP("Removing shadow pointer " DPxMOD "\n", DPxPTR((void **)Itr->first));
-        Itr = Device.ShadowPtrMap.erase(Itr);
-      } else {
-        ++Itr;
-      }
+               DPxPTR(ShadowPtr.HstPtrVal), DPxPTR(ShadowPtr.HstPtrAddr));
             return OFFLOAD_SUCCESS;
-    };
-    applyToShadowMapEntries(Device, CB, Info.HstPtrBegin, Info.DataSize,
-                            Info.Entry);
+          });
+    }
 
     // If we are deleting the entry the DataMapMtx is locked and we own the
     // entry.
     if (!Info.DelEntry ||
         (FromMapperBase && FromMapperBase == Info.HstPtrBegin)) {
+      Info.Entry->unlock();
       continue;
     }
 
     // Remember we deleted the entry in case it shows up in the
     // PostProcessingPtrs again.
     DeletedEntries.insert(Info.Entry);
     auto Ret = Device.deallocTgtPtr(HDTTMap, LR, Info.DataSize);
 
if (TPR.isHostPtr()) {
     DP("hst data:" DPxMOD " unified and shared, becomes a noop\n",
        DPxPTR(HstPtrBegin));
     return OFFLOAD_SUCCESS;
   }
 
   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;
     }
 
-    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
-      void **ShadowHstPtrAddr = (void **)Itr->first;
+    if (TPR.getEntry()) {
+      int Ret = TPR.getEntry()->foreachShadowPointerInfo(
+          [&](const ShadowPtrInfoTy &ShadowPtr) {
             // Wait for device-to-host memcopies for whole struct to complete,
             // before restoring the correct host pointer.
             if (AsyncInfo.synchronize() != OFFLOAD_SUCCESS)
               return OFFLOAD_FAIL;
-      *ShadowHstPtrAddr = Itr->second.HstPtrVal;
+            *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.getEntry());
+          });
+      if (Ret != OFFLOAD_SUCCESS) {
+        DP("Updating shadow map failed\n");
+        return Ret;
+      }
+    }
   }
 
   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);
+    int Ret = Device.submitData(TgtPtrBegin, HstPtrBegin, ArgSize, AsyncInfo,
+                                TPR.getEntry());
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Copying data to device failed.\n");
       return OFFLOAD_FAIL;
     }
 
-    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
+    if (TPR.getEntry()) {
+      int Ret = TPR.getEntry()->foreachShadowPointerInfo(
+          [&](const ShadowPtrInfoTy &ShadowPtr) {
             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)
+               DPxPTR(ShadowPtr.TgtPtrVal), DPxPTR(ShadowPtr.TgtPtrAddr));
+            Ret = Device.submitData(ShadowPtr.TgtPtrAddr,
+                                    (void *)&ShadowPtr.TgtPtrVal,
+                                    sizeof(void *), AsyncInfo, TPR.getEntry());
+            if (Ret != OFFLOAD_SUCCESS) {
               REPORT("Copying data to device failed.\n");
-      ++Itr;
+              return OFFLOAD_FAIL;
+            }
+            return OFFLOAD_SUCCESS;
+          });
+      if (Ret != OFFLOAD_SUCCESS) {
+        DP("Updating shadow map failed\n");
         return Ret;
-    };
-    applyToShadowMapEntries(Device, CB, HstPtrBegin, ArgSize, TPR.getEntry());
+      }
+    }
   }
   return OFFLOAD_SUCCESS;
 }
 
 static int targetDataNonContiguous(ident_t *loc, DeviceTy &Device,
                                    void *ArgsBase,
                                    __tgt_target_non_contig *NonContig,
                                    uint64_t Size, int64_t ArgType,
if (!(ArgTypes[I] & OMP_TGT_MAPTYPE_TARGET_PARAM)) {
           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];