[OpenMP][FIX] Avoid races in the handling of to be deleted mapping entries

[jdoerfert]

Summary

If we decided to delete a mapping entry we did not act on it right away
but first issued and waited for memory copies. In the meantime some
other thread might reuse the entry. While there was some logic to avoid
colliding on the actual "deletion" part, there were two races happening:

1. The data transfer back of the thread deleting the entry and the data transfer back of the thread taking over the entry raced.
2. The update to the shadow map happened regardless if the entry was actually reused by another thread which left the shadow map in a inconsistent state.

To fix both issues we will now update the shadow map and delete the
entry only if we are sure the thread is responsible for deletion, hence
no other thread took over the entry and reused it. We also wait for a
potential former data transfer from the device to finish before we issue
another one that would race with it.

[tianshilei1992]

openmp/libomptarget/src/device.cpp

370

Unfortunately, __kmpc_global_thread_num cannot give you a unique id. I did a simple test:

#include <cstdint>
#include <cstdio>
#include <thread>
#include <vector>

extern "C" int32_t __kmpc_global_thread_num(void *);

void foo() {
  std::chrono::seconds timespan(1);
  std::this_thread::sleep_for(timespan);
  int gtid = __kmpc_global_thread_num(nullptr);
  printf("gtid = %d.\n", gtid);
}

int main(int argc, char *argv[]) {
  std::vector<std::thread> pool(8);
  for (int i = 0; i < 8; ++i) {
    pool[i] = std::thread(foo);
  }

  for (std::thread &t : pool) {
    if (t.joinable()) {
      t.join();
    }
  }

  return 0;
}

It prints:

gtid = 0.
gtid = 9.
gtid = 10.
gtid = 9.
gtid = 10.
gtid = 9.
gtid = 10.
gtid = 9.

Since libomptarget is not forced to use OpenMP thread, we cannot use functions in libomp to tell thread id. std::thread::id should be used here I think.

[jdoerfert]

openmp/libomptarget/src/device.cpp

370

Interesting. I was expecting only openmp threads but you are right. Let's do std::trhead::id then.

[jdoerfert]

Use std::thread::id

[tianshilei1992]

LG. That also means we have to hold the lock longer than before, but for the correctness, we don't have choice.

[ronlieb]

openmp/libomptarget/test/mapping/map_back_race.cpp

32

seeing sporadic failures on this test, in the amdgpu openmp buildbot. return code -11
i think we need a return 0; here

[jdoerfert]

openmp/libomptarget/test/mapping/map_back_race.cpp

32

main should return 0 by default, IIRC. return code -11 sounds more like we still have a bug -.-

[jdoerfert]

openmp/libomptarget/test/mapping/map_back_race.cpp

32

@ronlieb It seems we only fail for host (x86) offloading. One needs to check but I'd expect the error code -11 to be 139 (128 + 11 = 139), which would be a common segfault. given it's on the host we can probably use gdb to pinpoint it. Feel free to disable it for the host for now, that should make the bot green at least.

[ronlieb]

thanks for suggestion: https://reviews.llvm.org/D122658

openmp/libomptarget/include/device.h

 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <list>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <set>
+#include <thread>
 #include <vector>
 
 #include "ExclusiveAccess.h"
 #include "omptarget.h"
 #include "rtl.h"
 
 // Forward declarations.
 struct RTLInfoTy;
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) {}
+          MayContainAttachedPointers(false), DeleteThreadId(std::thread::id()) {
+    }
     /// 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
struct StatesTy {
     /// Pointer to the event corresponding to the data update of this map.
     /// Note: At present this event is created when the first data transfer from
     /// host to device is issued, and only being used for H2D. It is not used
     /// for data transfer in another direction (device to host). It is still
     /// unclear whether we need it for D2H. If in the future we need similar
     /// mechanism for D2H, and if the event cannot be shared between them, Event
     /// should be written as <tt>void *Event[2]</tt>.
     void *Event = nullptr;
+
+    /// 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;
   };
   // 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();
+  }
+
+  /// Return the thread id of the thread expected to delete this entry.
+  std::thread::id getDeleteThreadId() const { return States->DeleteThreadId; }
+
   /// 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 =
         UseHoldRefCount ? States->HoldRefCount : States->DynRefCount;
uint64_t &ThisRefCount =
         UseHoldRefCount ? States->HoldRefCount : States->DynRefCount;
     uint64_t OtherRefCount =
         UseHoldRefCount ? States->DynRefCount : States->HoldRefCount;
     (void)OtherRefCount;
     if (ThisRefCount != INFRefCount) {
       if (ThisRefCount > 0)
         --ThisRefCount;
       else
-        assert(OtherRefCount > 0 && "total refcount underflow");
+        assert(OtherRefCount >= 0 && "total refcount underflow");
     }
     return getTotalRefCount();
   }
 
   /// Is the dynamic (and thus the total) reference count infinite?
   bool isDynRefCountInf() const { return States->DynRefCount == INFRefCount; }
 
   /// Convert the dynamic reference count to a debug string.
struct DeviceTy {
   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);
-  /// 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);
+
+  /// 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);
+
   int associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size);
   int disassociatePtr(void *HstPtrBegin);
 
   // calls to RTL
   int32_t initOnce();
   __tgt_target_table *load_binary(void *Img);
 
   // device memory allocation/deallocation routines

openmp/libomptarget/src/device.cpp

 #include "private.h"
 #include "rtl.h"
 
 #include <cassert>
 #include <climits>
 #include <cstdint>
 #include <cstdio>
 #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;
 
   void *Event = getEvent();
TargetPointerResultTy DeviceTy::getTargetPointer(
   // Check if the pointer is contained.
   // If a variable is mapped to the device manually by the user - which would
   // lead to the IsContained flag to be true - then we must ensure that the
   // device address is returned even under unified memory conditions.
   if (LR.Flags.IsContained ||
       ((LR.Flags.ExtendsBefore || LR.Flags.ExtendsAfter) && IsImplicit)) {
     auto &HT = *LR.Entry;
     const char *RefCountAction;
-    assert(HT.getTotalRefCount() > 0 && "expected existing RefCount > 0");
     if (UpdateRefCount) {
-      // After this, RefCount > 1.
+      // After this, reference count >= 1. If the reference count was 0 but the
+      // entry was still there we can reuse the data on the device and avoid a
+      // new submission.
       HT.incRefCount(HasHoldModifier);
       RefCountAction = " (incremented)";
     } else {
       // It might have been allocated with the parent, but it's still new.
       IsNew = HT.getTotalRefCount() == 1;
       RefCountAction = " (update suppressed)";
     }
     const char *DynRefCountAction = HasHoldModifier ? "" : RefCountAction;
DeviceTy::getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
   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;
-    // We do not zero the total reference count here.  deallocTgtPtr does that
-    // atomically with removing the mapping.  Otherwise, before this thread
-    // removed the mapping in deallocTgtPtr, another thread could retrieve the
-    // mapping, increment and decrement back to zero, and then both threads
-    // would try to remove the mapping, resulting in a double free.
     IsLast = HT.decShouldRemove(UseHoldRefCount, ForceDelete);
-    const char *RefCountAction;
-    if (!UpdateRefCount) {
-      RefCountAction = " (update suppressed)";
-    } else if (ForceDelete) {
+
+    if (ForceDelete) {
       HT.resetRefCount(UseHoldRefCount);
       assert(IsLast == HT.decShouldRemove(UseHoldRefCount) &&
              "expected correct IsLast prediction for reset");
-      if (IsLast)
-        RefCountAction = " (reset, deferred final decrement)";
-      else {
-        HT.decRefCount(UseHoldRefCount);
-        RefCountAction = " (reset)";
-      }
+
+    const char *RefCountAction;
+    if (!UpdateRefCount) {
+      RefCountAction = " (update suppressed)";
     } else if (IsLast) {
-      RefCountAction = " (deferred final decrement)";
+      // 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 {
       HT.decRefCount(UseHoldRefCount);
       RefCountAction = " (decremented)";
     }
     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,
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(void *HstPtrBegin, int64_t Size,
-                            bool HasHoldModifier) {
-  HDTTMapAccessorTy HDTTMap = HostDataToTargetMap.getExclusiveAccessor();
-
+int DeviceTy::deallocTgtPtr(HDTTMapAccessorTy &HDTTMap, LookupResult LR,
+                            int64_t Size) {
   // Check if the pointer is contained in any sub-nodes.
-  int Ret = OFFLOAD_SUCCESS;
-  LookupResult lr = lookupMapping(HDTTMap, HstPtrBegin, Size);
-  if (lr.Flags.IsContained || lr.Flags.ExtendsBefore || lr.Flags.ExtendsAfter) {
-    auto &HT = *lr.Entry;
-    if (HT.decRefCount(HasHoldModifier) == 0) {
-      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;
-      if (Event && destroyEvent(Event) != OFFLOAD_SUCCESS) {
-        REPORT("Failed to destroy event " DPxMOD "\n", DPxPTR(Event));
-        Ret = OFFLOAD_FAIL;
-      }
-    }
-  } else {
-    REPORT("Section to delete (hst addr " DPxMOD ") does not exist in the"
-           " allocated memory\n",
-           DPxPTR(HstPtrBegin));
-    Ret = OFFLOAD_FAIL;
-  }
+  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;
+  }
+
+  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() &&
+         "Trying to delete entry that is in use or owned by another thread.");
+
+       (HT.HstPtrName) ? getNameFromMapping(HT.HstPtrName).c_str() : "unknown");
+  void *Event = LR.Entry->getEvent();
+  HDTTMap->erase(LR.Entry);
+  delete LR.Entry;
+
+  int Ret = OFFLOAD_SUCCESS;
 
   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)

openmp/libomptarget/src/omptarget.cpp

for (int32_t i = 0; i < arg_num; ++i) {
     }
   }
 
   return OFFLOAD_SUCCESS;
 }
 
 namespace {
 /// This structure contains information to deallocate a target pointer, aka.
-/// used to call the function \p DeviceTy::deallocTgtPtr.
-struct DeallocTgtPtrInfo {
+/// used to fix up the shadow map and potentially delete the entry from the
+/// mapping table via \p DeviceTy::deallocTgtPtr.
+struct PostProcessingInfo {
   /// Host pointer used to look up into the map table
   void *HstPtrBegin;
+
   /// 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) {}
+  /// 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) {}
 };
 
 /// 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,

 } // 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;
-  std::vector<DeallocTgtPtrInfo> DeallocTgtPtrs;
+  std::vector<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;
for (int32_t I = ArgNum - 1; I >= 0; --I) {
 
     if ((ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) || DelEntry) {
       // Move data back to the host
       if (ArgTypes[I] & OMP_TGT_MAPTYPE_FROM) {
         bool Always = ArgTypes[I] & OMP_TGT_MAPTYPE_ALWAYS;
         if ((Always || IsLast) && !IsHostPtr) {
           DP("Moving %" PRId64 " bytes (tgt:" DPxMOD ") -> (hst:" DPxMOD ")\n",
              DataSize, DPxPTR(TgtPtrBegin), DPxPTR(HstPtrBegin));
+
+          std::lock_guard<decltype(*TPR.Entry)> LG(*TPR.Entry);
+          // Wait for any previous transfer if an event is present.
+          if (void *Event = TPR.Entry->getEvent()) {
+            if (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);
           if (Ret != OFFLOAD_SUCCESS) {
             REPORT("Copying data from device failed.\n");
             return OFFLOAD_FAIL;
           }
+
+          // As we are expecting to delete the entry the d2h copy might race
+          // with another one that also tries to delete the entry. This happens
+          // as the entry can be reused and the reuse might happen after the
+          // copy-back was issued but before it completed. Since the reuse might
+          // also copy-back a value we would race.
+          if (IsLast) {
+            if (TPR.Entry->addEventIfNecessary(Device, AsyncInfo) !=
+                OFFLOAD_SUCCESS)
+              return OFFLOAD_FAIL;
+          }
         }
       }
       if (DelEntry && FromMapper && I == 0) {
         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.
-      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;
-          // 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;
-          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));
-          Itr = Device.ShadowPtrMap.erase(Itr);
-        } else {
-          ++Itr;
-        }
-        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);
-    }
-  }
+      // Add pointer to the buffer for post-synchronize processing.
+      PostProcessingPtrs.emplace_back(HstPtrBegin, DataSize, ArgTypes[I],
+                                      DelEntry && !IsHostPtr, 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 (Info.ArgType & OMP_TGT_MAPTYPE_FROM) {
+      if (Info.DelEntry) {
+        DP("Removing shadow pointer " DPxMOD "\n", DPxPTR((void **)Itr->first));
+    applyToShadowMapEntries(Device, CB, Info.HstPtrBegin, Info.DataSize,
+                            Info.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;
+    // 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);
 
-  // Deallocate target pointer
-  for (DeallocTgtPtrInfo &Info : DeallocTgtPtrs) {
-    if (FromMapperBase && FromMapperBase == Info.HstPtrBegin)
-      continue;
-    Ret = Device.deallocTgtPtr(Info.HstPtrBegin, Info.DataSize,
-                               Info.HasHoldModifier);
       if (Ret != OFFLOAD_SUCCESS) {
         REPORT("Deallocating data from device failed.\n");
-      return OFFLOAD_FAIL;
+        break;
+      }
     }
   }
 
-  return 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);
   bool IsLast, IsHostPtr;
   TargetPointerResultTy TPR = Device.getTgtPtrBegin(

openmp/libomptarget/test/mapping/map_back_race.cpp

+// RUN: %libomptarget-compile-run-and-check-generic
+
+// Taken from https://github.com/llvm/llvm-project/issues/54216
+
+#include <algorithm>
+#include <cstdlib>
+#include <iostream>
+
+bool almost_equal(float x, float gold, float rel_tol = 1e-09,
+                  float abs_tol = 0.0) {
+  return std::abs(x - gold) <=
+         std::max(rel_tol * std::max(std::abs(x), std::abs(gold)), abs_tol);
+}
+void test_parallel_for__target() {
+  const int N0{32768};
+  const float expected_value{N0};
+  float counter_N0{};
+#pragma omp parallel for
+  for (int i0 = 0; i0 < N0; i0++) {
+#pragma omp target map(tofrom : counter_N0)
+    {
+#pragma omp atomic update
+      counter_N0 = counter_N0 + 1.;
+    }
+  }
+  if (!almost_equal(counter_N0, expected_value, 0.01)) {
+    std::cerr << "Expected: " << expected_value << " Got: " << counter_N0
+              << std::endl;
+    std::exit(112);
+  }
+}
+int main() { test_parallel_for__target(); }