[OpenMP] Avoid costly shadow map traversals whenever possible

[jhuber6]

Summary

In the OpenMC app we saw omp target update spending an awful lot of
time in the shadow map traversal without ever doing any update there.
There are two cases that allow us to avoid the traversal completely.
The simplest thing is that small updates cannot (reasonably) contain
an attached pointer part. The other case requires to track in the
mapping table if an entry might contain an attached pointer as part.
Given that we have a single location shadow map entries are created,
the latter is actually fairly easy as well.

[jdoerfert]

Update api.cpp

[jdoerfert]

@grokos wdyt?

[grokos]

LGTM with two minor nits.

openmp/libomptarget/src/device.cpp

293

Can you add a comment explaining that IsNew is always going to be false because getTgtPtrBegin is called from targetDataEnd, targetDataContiguous and processDataBefore, i.e. for data which is expected to be mapped already. And maybe it's more elegant to hardcode false in the return statement directly, thus avoiding defining a variable which is never used in the body of this function.

openmp/libomptarget/src/omptarget.cpp

619

an be --> can be

[grokos]

Ping. I'm preparing some more ShadowPtrMap optimizations which I'd like to base on top of this patch. Is there any planned change or can we land the patch?

[jhuber6]

Commandeering to land for @jdoerfert.

[jhuber6]

Addressing comments and fixing errors in debug build.

[JonChesterfield]

This seems to have broken some x64 openmp tests, e.g. https://lab.llvm.org/buildbot/#/builders/193/builds/2940

[ronlieb]

i see the same failures in my local build, and when i revert this commit, and rebuild/test, failures go away

[jhuber6]

In D113124#3186283, @JonChesterfield wrote:

This seems to have broken some x64 openmp tests, e.g. https://lab.llvm.org/buildbot/#/builders/193/builds/2940

In D113124#3186301, @ronlieb wrote:

i see the same failures in my local build, and when i revert this commit, and rebuild/test, failures go away

Yeah, I can see those failures now. I can revert until it's figured out.

[jdoerfert]

Fix errors, uninitialized iterator and wrong TPR (obj was used, ptr is needed).

openmp/libomptarget/src/api.cpp

if (DevicesSize <= (size_t)device_num) {
     DP("Call to omp_target_is_present with invalid device ID, returning "
        "false\n");
     return false;
   }
 
   DeviceTy &Device = *PM->Devices[device_num];
   bool IsLast; // not used
   bool IsHostPtr;
-  void *TgtPtr = Device.getTgtPtrBegin(const_cast<void *>(ptr), 0, IsLast,
-                                       /*UpdateRefCount=*/false,
-                                       /*UseHoldRefCount=*/false, IsHostPtr);
-  int rc = (TgtPtr != NULL);
+  TargetPointerResultTy TPR =
+      Device.getTgtPtrBegin(const_cast<void *>(ptr), 0, IsLast,
+  int rc = (TPR.TargetPointer != NULL);
   // Under unified memory the host pointer can be returned by the
   // getTgtPtrBegin() function which means that there is no device
   // corresponding point for ptr. This function should return false
   // in that situation.
   if (PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY)
     rc = !IsHostPtr;
   DP("Call to omp_target_is_present returns %d\n", rc);
   return rc;

openmp/libomptarget/src/device.h

 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) {}
+        : DynRefCount(DRC), HoldRefCount(HRC),
+          MayContainAttachedPointers(false) {}
     /// 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;
   };
bool decShouldRemove(bool UseHoldRefCount, bool AfterReset = false) const {
     if (AfterReset)
       return ThisRefCount != INFRefCount;
     return ThisRefCount == 1;
   }
 
   void lock() const { States->UpdateMtx.lock(); }
 
   void unlock() const { States->UpdateMtx.unlock(); }
+
+  void setMayContainAttachedPointers() const {
+    States->MayContainAttachedPointers = true;
+  }
+  bool getMayContainAttachedPointers() const {
+    return States->MayContainAttachedPointers;
+  }
 };
 
 typedef uintptr_t HstPtrBeginTy;
 inline bool operator<(const HostDataToTargetTy &lhs, const HstPtrBeginTy &rhs) {
   return lhs.HstPtrBegin < rhs;
 }
 inline bool operator<(const HstPtrBeginTy &lhs, const HostDataToTargetTy &rhs) {
   return lhs < rhs.HstPtrBegin;
struct DeviceTy {
   /// - 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);
   void *getTgtPtrBegin(void *HstPtrBegin, int64_t Size);
-  void *getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
-                       bool UpdateRefCount, bool UseHoldRefCount,
-                       bool &IsHostPtr, bool MustContain = false,
+  TargetPointerResultTy getTgtPtrBegin(void *HstPtrBegin, int64_t Size,
+                                       bool &IsLast, bool UpdateRefCount,
+                                       bool UseHoldRefCount, bool &IsHostPtr,
+                                       bool MustContain = false,
                                        bool ForceDelete = false);
   /// For the map entry for \p HstPtrBegin, decrement the reference count
   /// specified by \p HasHoldModifier and, if the the total reference count is
   /// then zero, deallocate the corresponding device storage and remove the map
   /// entry.  Return \c OFFLOAD_SUCCESS if the map entry existed, and return
   /// \c OFFLOAD_FAIL if not.  It is the caller's responsibility to skip calling
   /// this function if the map entry is not expected to exist because
   /// \p HstPtrBegin uses shared memory.
   int deallocTgtPtr(void *HstPtrBegin, int64_t Size, bool HasHoldModifier);

openmp/libomptarget/src/device.cpp

if (TargetPointer && !IsHostPtr && HasFlagTo && (IsNew || HasFlagAlways)) {
     DataMapMtx.unlock();
   }
 
   return {{IsNew, IsHostPtr}, Entry, TargetPointer};
 }
 
 // Used by targetDataBegin, targetDataEnd, targetDataUpdate and target.
 // Return the target pointer begin (where the data will be moved).
-// Decrement the reference counter if called from targetDataEnd.
-void *DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
-                               bool UpdateRefCount, bool UseHoldRefCount,
-                               bool &IsHostPtr, bool MustContain,
-                               bool ForceDelete) {
-  void *rc = NULL;
+// Decrement the reference counter if called from targetDataEnd. The data is
+// excepted to be mapped already, so the result will never be new.
+TargetPointerResultTy
+DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
+                         bool &IsHostPtr, bool MustContain, bool ForceDelete) {
+  void *TargetPointer = NULL;
   IsHostPtr = false;
   IsLast = false;
   DataMapMtx.lock();
   LookupResult lr = lookupMapping(HstPtrBegin, Size);
 
   if (lr.Flags.IsContained ||
       (!MustContain && (lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter))) {
     auto &HT = *lr.Entry;
if (lr.Flags.IsContained ||
     const char *DynRefCountAction = UseHoldRefCount ? "" : RefCountAction;
     const char *HoldRefCountAction = UseHoldRefCount ? RefCountAction : "";
     uintptr_t tp = HT.TgtPtrBegin + ((uintptr_t)HstPtrBegin - HT.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);
-    rc = (void *)tp;
+    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);
     IsHostPtr = true;
-    rc = HstPtrBegin;
+    TargetPointer = HstPtrBegin;
   }
 
   DataMapMtx.unlock();
-  return rc;
+  return {{false, IsHostPtr}, lr.Entry, TargetPointer};
 }
 
 // Return the target pointer begin (where the data will be moved).
 // Lock-free version called when loading global symbols from the fat binary.
 void *DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size) {
   uintptr_t hp = (uintptr_t)HstPtrBegin;
   LookupResult lr = lookupMapping(HstPtrBegin, Size);
   if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {

openmp/libomptarget/src/omptarget.cpp

 //===----------------------------------------------------------------------===//
 
 #include "omptarget.h"
 #include "device.h"
 #include "private.h"
 #include "rtl.h"
 
 #include <cassert>
+#include <cstdint>
 #include <vector>
 
 int AsyncInfoTy::synchronize() {
   int Result = OFFLOAD_SUCCESS;
   if (AsyncInfo.Queue) {
     // If we have a queue we need to synchronize it now.
     Result = Device.synchronize(*this);
     assert(AsyncInfo.Queue == nullptr &&
if (arg_types[i] & OMP_TGT_MAPTYPE_PTR_AND_OBJ && !IsHostPtr) {
       // 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[Pointer_HstPtrBegin] = {
             HstPtrBase, PointerTgtPtrBegin, ExpectedTgtPtrBase};
+        TPR.MapTableEntry->setMayContainAttachedPointers();
         UpdateDevPtr = true;
       }
 
       if (UpdateDevPtr) {
         Pointer_TPR.MapTableEntry->lock();
         Device.ShadowMtx.unlock();
 
         DP("Update pointer (" DPxMOD ") -> [" DPxMOD "]\n",
struct DeallocTgtPtrInfo {
   /// Size of the data
   int64_t DataSize;
   /// Whether it has \p ompx_hold modifier
   bool HasHoldModifier;
 
   DeallocTgtPtrInfo(void *HstPtr, int64_t Size, bool HasHoldModifier)
       : HstPtrBegin(HstPtr), DataSize(Size), HasHoldModifier(HasHoldModifier) {}
 };
+
+/// 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.MapTableEntry->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.
+  Device.ShadowMtx.lock();
+  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;
+  }
+  Device.ShadowMtx.unlock();
+}
+
 } // 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;
for (int32_t I = ArgNum - 1; I >= 0; --I) {
     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
-    void *TgtPtrBegin = Device.getTgtPtrBegin(
+    TargetPointerResultTy TPR = Device.getTgtPtrBegin(
         HstPtrBegin, DataSize, IsLast, UpdateRef, HasHoldModifier, IsHostPtr,
         !IsImplicit, ForceDelete);
+    void *TgtPtrBegin = TPR.TargetPointer;
     if (!TgtPtrBegin && (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
         // or target exit data construct with a present map-type-modifier then
         // on entry to the region if the corresponding list item does not appear
if ((ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) || DelEntry) {
         DelEntry = false;
         FromMapperBase = HstPtrBegin;
       }
 
       // 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.
-      uintptr_t LB = (uintptr_t)HstPtrBegin;
-      uintptr_t UB = (uintptr_t)HstPtrBegin + DataSize;
-      Device.ShadowMtx.lock();
-      for (ShadowPtrListTy::iterator Itr = Device.ShadowPtrMap.begin();
-           Itr != Device.ShadowPtrMap.end();) {
-        void **ShadowHstPtrAddr = (void **)Itr->first;
-
-        // An STL map is sorted on its keys; use this property
-        // to quickly determine when to break out of the loop.
-        if ((uintptr_t)ShadowHstPtrAddr < LB) {
-          ++Itr;
-          continue;
-        }
-        if ((uintptr_t)ShadowHstPtrAddr >= UB)
-          break;
-
+      auto CB = [&](ShadowPtrListTy::iterator &Itr) {
         // If we copied the struct to the host, we need to restore the pointer.
         if (ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) {
+          void **ShadowHstPtrAddr = (void **)Itr->first;
           DP("Restoring original host pointer value " DPxMOD " for host "
              "pointer " DPxMOD "\n",
              DPxPTR(Itr->second.HstPtrVal), DPxPTR(ShadowHstPtrAddr));
           *ShadowHstPtrAddr = Itr->second.HstPtrVal;
         }
         // If the struct is to be deallocated, remove the shadow entry.
         if (DelEntry) {
-          DP("Removing shadow pointer " DPxMOD "\n", DPxPTR(ShadowHstPtrAddr));
+          DP("Removing shadow pointer " DPxMOD "\n",
+             DPxPTR(Itr->second.HstPtrVal));
           Itr = Device.ShadowPtrMap.erase(Itr);
         } else {
           ++Itr;
         }
-      }
-      Device.ShadowMtx.unlock();
+        return OFFLOAD_SUCCESS;
+      };
+      applyToShadowMapEntries(Device, CB, HstPtrBegin, DataSize, TPR);
 
       // Add pointer to the buffer for later deallocation
       if (DelEntry && !IsHostPtr)
         DeallocTgtPtrs.emplace_back(HstPtrBegin, DataSize, HasHoldModifier);
     }
   }
 
   // TODO: We should not synchronize here but pass the AsyncInfo object to the
int targetDataEnd(ident_t *loc, DeviceTy &Device, int32_t ArgNum,
   return OFFLOAD_SUCCESS;
 }
 
 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;
-  void *TgtPtrBegin = Device.getTgtPtrBegin(
+  TargetPointerResultTy TPR = Device.getTgtPtrBegin(
       HstPtrBegin, ArgSize, IsLast, /*UpdateRefCount=*/false,
       /*UseHoldRefCount=*/false, IsHostPtr, /*MustContain=*/true);
+  void *TgtPtrBegin = TPR.TargetPointer;
   if (!TgtPtrBegin) {
     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;
     }
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);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Copying data from device failed.\n");
       return OFFLOAD_FAIL;
     }
 
-    uintptr_t LB = (uintptr_t)HstPtrBegin;
-    uintptr_t UB = (uintptr_t)HstPtrBegin + ArgSize;
-    Device.ShadowMtx.lock();
-    for (ShadowPtrListTy::iterator IT = Device.ShadowPtrMap.begin();
-         IT != Device.ShadowPtrMap.end(); ++IT) {
-      void **ShadowHstPtrAddr = (void **)IT->first;
-      if ((uintptr_t)ShadowHstPtrAddr < LB)
-        continue;
-      if ((uintptr_t)ShadowHstPtrAddr >= UB)
-        break;
+    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
+      void **ShadowHstPtrAddr = (void **)Itr->first;
       DP("Restoring original host pointer value " DPxMOD
          " for host pointer " DPxMOD "\n",
-         DPxPTR(IT->second.HstPtrVal), DPxPTR(ShadowHstPtrAddr));
-      *ShadowHstPtrAddr = IT->second.HstPtrVal;
-    }
-    Device.ShadowMtx.unlock();
+         DPxPTR(Itr->second.HstPtrVal), DPxPTR(ShadowHstPtrAddr));
+      *ShadowHstPtrAddr = Itr->second.HstPtrVal;
+      return OFFLOAD_SUCCESS;
+    };
+    applyToShadowMapEntries(Device, CB, HstPtrBegin, ArgSize, TPR);
   }
 
   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;
     }
 
-    uintptr_t LB = (uintptr_t)HstPtrBegin;
-    uintptr_t UB = (uintptr_t)HstPtrBegin + ArgSize;
-    Device.ShadowMtx.lock();
-    for (ShadowPtrListTy::iterator IT = Device.ShadowPtrMap.begin();
-         IT != Device.ShadowPtrMap.end(); ++IT) {
-      void **ShadowHstPtrAddr = (void **)IT->first;
-      if ((uintptr_t)ShadowHstPtrAddr < LB)
-        continue;
-      if ((uintptr_t)ShadowHstPtrAddr >= UB)
-        break;
+    auto CB = [&](ShadowPtrListTy::iterator &Itr) {
       DP("Restoring original target pointer value " DPxMOD " for target "
          "pointer " DPxMOD "\n",
-         DPxPTR(IT->second.TgtPtrVal), DPxPTR(IT->second.TgtPtrAddr));
-      Ret = Device.submitData(IT->second.TgtPtrAddr, &IT->second.TgtPtrVal,
+         DPxPTR(Itr->second.TgtPtrVal), DPxPTR(Itr->second.TgtPtrAddr));
+      Ret = Device.submitData(Itr->second.TgtPtrAddr, &Itr->second.TgtPtrVal,
                               sizeof(void *), AsyncInfo);
-      if (Ret != OFFLOAD_SUCCESS) {
+      if (Ret != OFFLOAD_SUCCESS)
         REPORT("Copying data to device failed.\n");
-        Device.ShadowMtx.unlock();
-        return OFFLOAD_FAIL;
-      }
-    }
-    Device.ShadowMtx.unlock();
+      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,
if (!(ArgTypes[I] & OMP_TGT_MAPTYPE_TARGET_PARAM)) {
         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();
-        PointerTgtPtrBegin = Device.getTgtPtrBegin(
+        TargetPointerResultTy TPR = Device.getTgtPtrBegin(
             HstPtrVal, ArgSizes[I], IsLast, /*UpdateRefCount=*/false,
             /*UseHoldRefCount=*/false, IsHostPtr);
+        PointerTgtPtrBegin = TPR.TargetPointer;
         if (!PointerTgtPtrBegin) {
           DP("No lambda captured variable mapped (" DPxMOD ") - ignored\n",
              DPxPTR(HstPtrVal));
           continue;
         }
         if (IsHostPtr) {
           DP("Unified memory is active, no need to map lambda captured"
              "variable (" DPxMOD ")\n",
if (!(ArgTypes[I] & OMP_TGT_MAPTYPE_TARGET_PARAM)) {
       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;
       const bool IsFirstPrivate = (ArgTypes[I] & OMP_TGT_MAPTYPE_TO);
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);
-      TgtPtrBegin = Device.getTgtPtrBegin(HstPtrBegin, ArgSizes[I], IsLast,
-                                          /*UpdateRefCount=*/false,
-                                          /*UseHoldRefCount=*/false, IsHostPtr);
+      TPR = Device.getTgtPtrBegin(HstPtrBegin, ArgSizes[I], IsLast,
+      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
     }
     TgtArgsPositions[I] = TgtArgs.size();