[OpenMP][Fix] Track all threads that may delete an entry

[gValarini]

Summary

The entries inside a "target data end" is processed in three steps:

1. Query internal data maps for the entries and dispatch any necessary device-side operations (i.e., data retrieval);
2. Synchronize the such operations;
3. Update the host-side pointers and remove any entry which reference counter reached zero.

Such steps may be executed by multiple threads which may even operate on
the same entries. The current implementation (D121058) tries to
synchronize these threads by tracking the "owner" for the deletion of
each entry using their thread ID. Unfortunately it may failed to do so
because of the following reasons:

1. The owner is always assigned at the first step only if the reference count is 0 when the map is queried. This does not work when such owner thread is faster than a previous one that is also processing the same entry on another "target data end", leading to user-after-free problems.
2. The entry is only added for post-processing (step 3) if its reference count was 0 at query time (step 1). This does not allow for threads to exchange responsibility for the deletion, leading again to user-after-free problems.
3. An entry may appear multiple times in the arguments array of a "target data end", which may lead to deleting the entry prematurely, leading, again, to user-after-free problems.

This patch addresses these problems by tracking all the threads that are
using an entry at "target data end" region through a counter, ensuring
only the last one deletes it when needed. It also ensures that all
entries that are successfully found inside the data maps in step 1 are
also processed in step 3, regardless if their reference count was zeroed
or not at query time. This ensures the deletion ownership may be passed
to any thread that is using such entry.

[jdoerfert]

openmp/libomptarget/include/device.h

76

Make a separate header, this is reusable.

openmp/libomptarget/src/omptarget.cpp

687

I think, this can leave an entry intact on the device:

T1: arrives here, TPR is still life so DeleterThreadCount is 3, won't destory the entry, HDTTMap is released.
T2: arrives here, TPR is still life so DeleterThreadCount is 3, won't destory the entry. HDTTMap is released.
Both exit this function, TPRs are destroyed, DeleterThreadCount is decremented twice and reaches 1, but nobody executing the code that would delete the entry.

[ye-luo]

openmp/libomptarget/src/omptarget.cpp

687

Can this be protected by not calling HDTTMap.destroy()?

[ye-luo]

openmp/libomptarget/src/omptarget.cpp

687

Second thought on this. If the mutex protection from HDTTMap is anyway needed, std::atomic seems likely not needed.
Probably need to revisit D123443, D123444 and D123446.

[gValarini]

Update the reference counting scheme and which entries are included in the post-processing.

• The counters are now explicitly managed. No more scope problems!
• All entries are added for post-processing to allow any thread that references an entry at a target data end to become the deleter of such entry.

[jdoerfert]

I have minor questions but after reading it once it seems to make sense.

openmp/libomptarget/src/omptarget.cpp

737

IsNotLastUser, right?

751

What's this block about?

[gValarini]

Fixed a variable name

[gValarini]

openmp/libomptarget/src/omptarget.cpp

687

@ye-luo we have a new version of this solution, I think your previous comments are now fixed.

687

I believe this does not happen anymore.

737

Ops, IsNotLastUser is the correct one. Thanks!

751

To be completely honest, I cannot recall. Most of the checks that can affect DelEntry were transferred from targetDataEnd. Although I understand the other ones related to struct and ptr/obj mapping, I don´t get why the first element of a mapper-based entry should not have its reference count updated or should not be removed.

I know this was implemented in https://reviews.llvm.org/D100673, do you remember the purpose of this check? We could add a comment here explaining it better.

[ye-luo]

openmp/libomptarget/src/device.cpp

364

getTgtPtrBegin arguments are very correlated. Better to be scrutinized in the the future.

openmp/libomptarget/src/omptarget.cpp

803

More readable with
/*FromDataEnd=*/true

[gValarini]

Update documentation and variable names

[gValarini]

openmp/libomptarget/src/device.cpp

364

I agree. Unfortunately, the comment below is a little misleading, since UpdateRefCount is not always set on targetDataEnd.

Decrement the reference counter if called from targetDataEnd.

I updated the doc comment for now and I'll add a small patch removing some of the return arguments that should be placed inside TargetPointerResultTy later.

[jdoerfert]

Ready to merge?

[gValarini]

In D132676#4052127, @jdoerfert wrote:

Ready to merge?

I believe it is.

There is just that one block of code you commented about that tries to manage the deallocation of nested custom mappers that I still cannot wrap my head around. But I believe we would need to revise how we manage the reference counting of those since they are probably being left allocated by the end of the program (https://github.com/llvm/llvm-project/issues/60106).

Do you think we can merge the patch and take care of that later?

[ye-luo]

Quite close to the release. Let's merge this and deal with the remaining issue in a later patch.

[gValarini]

In D132676#4065499, @ye-luo wrote:

Quite close to the release. Let's merge this and deal with the remaining issue in a later patch.

Ok, I will land it. Thanks!

[gValarini]

Update summary.

[gValarini]

Rebase

openmp/libomptarget/include/ExclusiveAccess.h

template <typename Ty> struct Accessor {
 
   /// Constructor to get exclusive access by locking the mutex protecting the
   /// underlying object.
   Accessor(ProtectedObj<Ty> &PO) : Ptr(&PO) { lock(); }
 
   /// Constructor to get exclusive access by taking it from \p Other.
   Accessor(Accessor<Ty> &&Other) : Ptr(Other.Ptr) { Other.Ptr = nullptr; }
 
+  /// Move-assign to get exclusive access by taking it from \p Other.
+  Accessor &operator=(Accessor<Ty> &&Other) {
+    Ptr = Other.Ptr;
+    Other.Ptr = nullptr;
+  }
+
   Accessor(Accessor &Other) = delete;
 
   /// If the object is still owned when the lifetime ends we give up access.
   ~Accessor() { unlock(); }
 
   /// Give up access to the underlying object, virtually "destroying" the
   /// accessor even if the object is still life.
   void destroy() {

openmp/libomptarget/include/device.h

 //===----------- device.h - Target independent OpenMP target RTL ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Declarations for managing devices that are handled by RTL plugins.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_DEVICE_H
 #define _OMPTARGET_DEVICE_H
 
+#include <atomic>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <list>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <set>

 // 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;
 
+struct ScopedAtomicCounter {
+private:
+  std::atomic<int> *Counter = nullptr;
+
+public:
+  ScopedAtomicCounter() = default;
+
+  ScopedAtomicCounter(std::atomic<int> &C) : Counter(&C) {
+    if (Counter) {
+      (*Counter)++;
+      assert(*Counter > 0);
+    }
+  }
+
+  ~ScopedAtomicCounter() {
+    if (Counter) {
+      (*Counter)--;
+      assert(*Counter >= 0);
+    }
+  }
+
+  ScopedAtomicCounter(ScopedAtomicCounter &&Other) : Counter(Other.Counter) {
+    Other.Counter = nullptr;
+  }
+  ScopedAtomicCounter &operator=(ScopedAtomicCounter &&Other) {
+    Counter = Other.Counter;
+    Other.Counter = nullptr;
+    return *this;
+  }
+};
+
 /// 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), DeleteThreadId(std::thread::id()) {
+          MayContainAttachedPointers(false), DeleterThreadCount(0) {
     }
     /// 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
struct StatesTy {
     /// 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;
 
-    /// The id of the thread responsible for deleting this entry. This thread
-    /// set the reference count to zero *last*. Other threads might reuse the
-    /// entry while it is marked for deletion but not yet deleted (e.g., the
-    /// data is still being moved back). If another thread reuses the entry we
-    /// will have a non-zero reference count *or* the thread will have changed
-    /// this id, effectively taking over deletion responsibility.
-    std::thread::id DeleteThreadId;
+    /// TODO: Describe it!
+    std::atomic<int> DeleterThreadCount;
   };
   // When HostDataToTargetTy is used by std::set, std::set::iterator is const
   // use unique_ptr to make States mutable.
   const std::unique_ptr<StatesTy> States;
 
 public:
   HostDataToTargetTy(uintptr_t BP, uintptr_t B, uintptr_t E, uintptr_t TB,
                      bool UseHoldRefCount, map_var_info_t Name = nullptr,
public:
 
   /// Get the event bound to this data map.
   void *getEvent() const { return States->Event; }
 
   /// Add a new event, if necessary.
   /// Returns OFFLOAD_FAIL if something went wrong, OFFLOAD_SUCCESS otherwise.
   int addEventIfNecessary(DeviceTy &Device, AsyncInfoTy &AsyncInfo) const;
 
-  /// Indicate that the current thread expected to delete this entry.
-  void setDeleteThreadId() const {
-    States->DeleteThreadId = std::this_thread::get_id();
+  /// TODO: Describe it!
+  int getDeleterThreadCount() {
+    return States->DeleterThreadCount;
   }
 
-  /// Return the thread id of the thread expected to delete this entry.
-  std::thread::id getDeleteThreadId() const { return States->DeleteThreadId; }
+  /// TODO: Describe it!
+  ScopedAtomicCounter getScopedDeleterThreadCounter() {
+    return ScopedAtomicCounter(States->DeleterThreadCount);
+  }
 
   /// Set the event bound to this data map.
   void setEvent(void *Event) const { States->Event = Event; }
 
   /// Reset the specified reference count unless it's infinity.  Reset to 1
   /// (even if currently 0) so it can be followed by a decrement.
   void resetRefCount(bool UseHoldRefCount) const {
     uint64_t &ThisRefCount =
struct {
     unsigned IsHostPointer : 1;
   } Flags = {0, 0};
 
   /// The corresponding map table entry which is stable.
   HostDataToTargetTy *Entry = nullptr;
 
   /// The corresponding target pointer
   void *TargetPointer = nullptr;
+
+  /// TODO: Describe it!
+  ScopedAtomicCounter ThreadDeleteCounter = {};
 };
 
 /// Map for shadow pointers
 struct ShadowPtrValTy {
   void *HstPtrVal;
   void *TgtPtrAddr;
   void *TgtPtrVal;
 };
struct DeviceTy {
   /// ensures exclusive access to the HDTT map.
   void *getTgtPtrBegin(HDTTMapAccessorTy &HDTTMap, void *HstPtrBegin,
                        int64_t Size);
 
   TargetPointerResultTy getTgtPtrBegin(void *HstPtrBegin, int64_t Size,
                                        bool &IsLast, bool UpdateRefCount,
                                        bool UseHoldRefCount, bool &IsHostPtr,
                                        bool MustContain = false,
-                                       bool ForceDelete = false);
+                                       bool ForceDelete = false,
+                                       bool FromDataEnd = false);
 
+  /// TODO: Update it!
   /// Deallocate \p LR and remove the entry. Assume the total reference count is
   /// zero and the calling thread is the deleting thread for \p LR. \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.
-  int deallocTgtPtr(HDTTMapAccessorTy &HDTTMap, LookupResult LR, int64_t Size);
+  void eraseMapEntry(HDTTMapAccessorTy &HDTTMap, HostDataToTargetTy *Entry, int64_t Size);
+
+  /// TODO: Describe it!
+  int deallocTgtPtrAndEntry(HostDataToTargetTy *Entry, int64_t Size);
 
   int associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size);
   int disassociatePtr(void *HstPtrBegin);
 
   // calls to RTL
   int32_t initOnce();
   __tgt_target_table *loadBinary(void *Img);
 

openmp/libomptarget/src/device.cpp

 }
 
 // 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.
 TargetPointerResultTy
 DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
                          bool UpdateRefCount, bool UseHoldRefCount,
-                         bool &IsHostPtr, bool MustContain, bool ForceDelete) {
+                         bool &IsHostPtr, bool MustContain, bool ForceDelete,
+                         bool FromDataEnd) {
   HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
 
+  ScopedAtomicCounter DeleterScopedCounter;
   void *TargetPointer = NULL;
   bool IsNew = false;
   IsHostPtr = false;
   IsLast = false;
   LookupResult LR = lookupMapping(HDTTMap, HstPtrBegin, Size);
 
   if (LR.Flags.IsContained ||
       (!MustContain && (LR.Flags.ExtendsBefore || LR.Flags.ExtendsAfter))) {
     auto &HT = *LR.Entry;
     IsLast = HT.decShouldRemove(UseHoldRefCount, ForceDelete);
 
     if (ForceDelete) {
       HT.resetRefCount(UseHoldRefCount);
       assert(IsLast == HT.decShouldRemove(UseHoldRefCount) &&
              "expected correct IsLast prediction for reset");
     }
 
+    if (FromDataEnd) {
+      DeleterScopedCounter = LR.Entry->getScopedDeleterThreadCounter();
+    }
+
     const char *RefCountAction;
     if (!UpdateRefCount) {
       RefCountAction = " (update suppressed)";
     } else if (IsLast) {
-      // Mark the entry as to be deleted by this thread. Another thread might
-      // reuse the entry and take "ownership" for the deletion while this thread
-      // is waiting for data transfers. That is fine and the current thread will
-      // simply skip the deletion step then.
-      HT.setDeleteThreadId();
       HT.decRefCount(UseHoldRefCount);
       assert(HT.getTotalRefCount() == 0 &&
              "Expected zero reference count when deletion is scheduled");
       if (ForceDelete)
         RefCountAction = " (reset, delayed deletion)";
       else
         RefCountAction = " (decremented, delayed deletion)";
     } else {
if (LR.Flags.IsContained ||
     // use directly from the host.
     DP("Get HstPtrBegin " DPxMOD " Size=%" PRId64 " for unified shared "
        "memory\n",
        DPxPTR((uintptr_t)HstPtrBegin), Size);
     IsHostPtr = true;
     TargetPointer = HstPtrBegin;
   }
 
-  return {{IsNew, IsHostPtr}, LR.Entry, TargetPointer};
+  return {{IsNew, IsHostPtr},
+          LR.Entry,
+          TargetPointer,
+          std::move(DeleterScopedCounter)};
 }
 
 // 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);
     return (void *)TP;
   }
 
   return NULL;
 }
 
-int DeviceTy::deallocTgtPtr(HDTTMapAccessorTy &HDTTMap, LookupResult LR,
-                            int64_t Size) {
-  // 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));
-    return OFFLOAD_FAIL;
-  }
+void DeviceTy::eraseMapEntry(HDTTMapAccessorTy &HDTTMap, HostDataToTargetTy *Entry, int64_t Size) {
+  assert(Entry && "Trying to delete a null entry from the HDTT map.");
 
-  auto &HT = *LR.Entry;
-  // Verify this thread is still in charge of deleting the entry.
-  assert(HT.getTotalRefCount() == 0 &&
-         HT.getDeleteThreadId() == std::this_thread::get_id() &&
+  // Verify this thread is the only deleter holding a reference to the entry.
+  assert(Entry->getTotalRefCount() == 0 && Entry->getDeleterThreadCount() == 1 &&
          "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();
-  HDTTMap->erase(LR.Entry);
-  delete LR.Entry;
+       DPxPTR(Entry->HstPtrBegin), DPxPTR(Entry->TgtPtrBegin), Size,
+       (Entry->HstPtrName) ? getNameFromMapping(Entry->HstPtrName).c_str() : "unknown");
 
-  int Ret = OFFLOAD_SUCCESS;
+  HDTTMap->erase(Entry);
+}
+
+int DeviceTy::deallocTgtPtrAndEntry(HostDataToTargetTy *Entry, int64_t Size) {
+  assert(Entry && "Trying to deallocate a null entry.");
+
+  DP("Deleting tgt data " DPxMOD " of size %" PRId64 "\n",
+     DPxPTR(Entry->TgtPtrBegin), Size);
+  
+  void *Event = Entry->getEvent();
   if (Event && destroyEvent(Event) != OFFLOAD_SUCCESS) {
     REPORT("Failed to destroy event " DPxMOD "\n", DPxPTR(Event));
-    Ret = OFFLOAD_FAIL;
+    return OFFLOAD_FAIL;
   }
   
+  int Ret = deleteData((void *)Entry->TgtPtrBegin);
+  delete Entry;
+
   return Ret;
 }
 
 /// Init device, should not be called directly.
 void DeviceTy::init() {
   // Make call to init_requires if it exists for this plugin.
   if (RTL->init_requires)
     RTL->init_requires(PM->RTLs.RequiresFlags);

openmp/libomptarget/src/omptarget.cpp

struct PostProcessingInfo {
   int64_t DataSize;
 
   /// The mapping type (bitfield).
   int64_t ArgType;
 
   /// The target pointer information.
   TargetPointerResultTy TPR;
 
-  /// Are we expecting to delete this entry or not. Even if set, we might not
-  /// delete the entry if another thread reused the entry in the meantime.
-  bool DelEntry;
-
   PostProcessingInfo(void *HstPtr, int64_t Size, int64_t ArgType, bool DelEntry,
-                     TargetPointerResultTy TPR)
-      : HstPtrBegin(HstPtr), DataSize(Size), ArgType(ArgType), TPR(TPR),
-        DelEntry(DelEntry) {}
+                     TargetPointerResultTy &&TPR)
+      : HstPtrBegin(HstPtr), DataSize(Size), ArgType(ArgType),
+        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,
for (ShadowPtrListTy::iterator Itr = Device.ShadowPtrMap.begin();
     if (ShadowHstPtrAddr >= UB)
       break;
 
     if (CB(Itr) == OFFLOAD_FAIL)
       break;
   }
 }
 
+/// TODO: Describe it!
+static int targetDataEndPostProcessing(DeviceTy &Device, void *FromMapperBase,
+                                       void *HstPtrBegin, int64_t DataSize,
+                                       int64_t ArgType,
+                                       TargetPointerResultTy TPR) {
+  int Ret = OFFLOAD_SUCCESS;
+
+  // Only delete the entry when the total ref count is zero and the current
+  // thread is the sole deleter.
+  auto HDTTMap = Device.HostDataToTargetMap.getExclusiveAccessor();
+  bool DelEntry = true;
+  if (TPR.Flags.IsHostPointer || TPR.Entry->getTotalRefCount() != 0 ||
+      TPR.Entry->getDeleterThreadCount() != 1) {
+    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.
+  auto CB = [&](ShadowPtrListTy::iterator &Itr) {
+    // If we copied the struct to the host, we need to restore the pointer.
+    if (ArgType & OMP_TGT_MAPTYPE_FROM) {
+      void **ShadowHstPtrAddr = (void **)Itr->first;
+      *ShadowHstPtrAddr = Itr->second.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;
+    }
+    return OFFLOAD_SUCCESS;
+  };
+  applyToShadowMapEntries(Device, CB, HstPtrBegin, DataSize, TPR);
+
+  // If we are deleting the entry the DataMapMtx is locked and we own the
+  // entry.
+  if (DelEntry) {
+    if (!FromMapperBase || FromMapperBase != HstPtrBegin) {
+      Device.eraseMapEntry(HDTTMap, TPR.Entry, DataSize);
+      // Since all necessary operations on the HDTTMap are completed, there is
+      // no need to hold its lock any more.
+      HDTTMap.destroy();
+      Ret = Device.deallocTgtPtrAndEntry(TPR.Entry, DataSize);
+    }
+
+    if (Ret != OFFLOAD_SUCCESS) {
+      REPORT("Deallocating data from device failed.\n");
+    }
+  }
+
+  return Ret;
+}
 } // 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;
   SmallVector<PostProcessingInfo> PostProcessingPtrs;
   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 (ArgMappers && ArgMappers[I]) {
       // Instead of executing the regular path of targetDataEnd, call the
       // targetDataMapper variant which will call targetDataEnd again
       // with new arguments.
       DP("Calling targetDataMapper for the %dth argument\n", I);
 
       map_var_info_t ArgName = (!ArgNames) ? nullptr : ArgNames[I];
       Ret = targetDataMapper(Loc, Device, ArgBases[I], Args[I], ArgSizes[I],
                              ArgTypes[I], ArgName, ArgMappers[I], AsyncInfo,
for (int32_t I = ArgNum - 1; I >= 0; --I) {
     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;
+    const bool FromDataEnd = true;
 
     // If PTR_AND_OBJ, HstPtrBegin is address of pointee
     TargetPointerResultTy TPR = Device.getTgtPtrBegin(
         HstPtrBegin, DataSize, IsLast, UpdateRef, HasHoldModifier, IsHostPtr,
-        !IsImplicit, ForceDelete);
+        !IsImplicit, ForceDelete, FromDataEnd);
     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
if ((ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) || DelEntry) {
       }
       if (DelEntry && FromMapper && I == 0) {
         DelEntry = false;
         FromMapperBase = HstPtrBegin;
       }
 
       // Add pointer to the buffer for post-synchronize processing.
       PostProcessingPtrs.emplace_back(HstPtrBegin, DataSize, ArgTypes[I],
-                                      DelEntry && !IsHostPtr, TPR);
+                                      DelEntry && !IsHostPtr, std::move(TPR));
     }
   }
 
   // TODO: We should not synchronize here but pass the AsyncInfo object to the
   //       allocate/deallocate device APIs.
   //
   // We need to synchronize before deallocating data.
   Ret = AsyncInfo.synchronize();
   if (Ret != OFFLOAD_SUCCESS)
     return OFFLOAD_FAIL;
 
   // Deallocate target pointer
   for (PostProcessingInfo &Info : PostProcessingPtrs) {
-    // 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;
-    DeviceTy::HDTTMapAccessorTy HDTTMap =
-        Device.HostDataToTargetMap.getExclusiveAccessor(!Info.DelEntry);
-
-    if (Info.DelEntry) {
-      LR = Device.lookupMapping(HDTTMap, Info.HstPtrBegin, Info.DataSize);
-      if (LR.Entry->getTotalRefCount() != 0 ||
-          LR.Entry->getDeleteThreadId() != std::this_thread::get_id()) {
-        // 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;
-        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));
-        auto OldItr = Itr;
-        Itr++;
-        Device.ShadowPtrMap.erase(OldItr);
-      } else {
-        ++Itr;
-      }
-      return OFFLOAD_SUCCESS;
-    };
-    applyToShadowMapEntries(Device, CB, Info.HstPtrBegin, Info.DataSize,
-                            Info.TPR);
-
-    // If we are deleting the entry the DataMapMtx is locked and we own the
-    // entry.
-    if (Info.DelEntry) {
-      if (!FromMapperBase || FromMapperBase != Info.HstPtrBegin)
-        Ret = Device.deallocTgtPtr(HDTTMap, LR, Info.DataSize);
-
-      if (Ret != OFFLOAD_SUCCESS) {
-        REPORT("Deallocating data from device failed.\n");
-        break;
-      }
-    }
-  }
+    // Move TPR to the inner scope. This way, scoped managed variables (e.g.,
+    // Info.TPR.ThreadDeleteCounter) can be correctly deconstructed.
+    Ret = targetDataEndPostProcessing(Device, FromMapperBase, Info.HstPtrBegin,
+                                      Info.DataSize, Info.ArgType,
+                                      std::move(Info.TPR));
+    if (Ret != OFFLOAD_SUCCESS)
 
   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);