[OpenMP][FIX] Simplify the entry deletion scheme

[jdoerfert]

Summary

The scheme introduced with D121058 did not work properly. There are
three reasons, two are fixed here:

1. The thread ID scheme is not needed and can result in dangling pointers if the thread in charge is faster than another one interested in deletion. We now keep the entry until the last thread that wants to delete the entry makes it to the deletion point.
2. The PostProcessingPtrs set can contain an entry multiple times. The deletion can happen before we visit the entry again in a subsequent iteration. That resulted again in a dangling pointer. We track the deleted entries now to avoid this.

[tianshilei1992]

openmp/libomptarget/src/omptarget.cpp

910

In what cases the same entry is in PostProcessingPtrs more than once?

[jdoerfert]

openmp/libomptarget/src/omptarget.cpp

910

We have a test case where this happens. I forgot which one. The entry shows up 3 times. The first and last time we do not want to delete it, the middle one is a deletion.

[tianshilei1992]

LG

[ye-luo]

openmp/libomptarget/include/device.h

102–108

Each target task instead of thread.

224

Where does it get called?

[jdoerfert]

openmp/libomptarget/include/device.h

224

https://reviews.llvm.org/D123443#3452644

[ye-luo]

openmp/libomptarget/src/omptarget.cpp

869

In another thread which is working on increase the refcount may have a race with the current thread deleting the entry because the other thread released HDTTMapAccessorTy and only have the entry locked.

[jdoerfert]

Fix rebase issues

[jdoerfert]

openmp/libomptarget/src/omptarget.cpp

869

No. HDTTMap is still locked here. In the next patch we introduce the consistent locking for entries. So HDTTMap is locked, and then we lock entry, and then we do the lookup. No other thread can modify the entry at that point and the ref count and delete count will be valid.

[jdoerfert]

was fixed upstream in a better way

openmp/libomptarget/include/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),
-          MayContainAttachedPointers(false), DeleteThreadId(std::thread::id()) {
-    }
+        : DynRefCount(DRC), HoldRefCount(HRC) {}
     /// The dynamic reference count is the standard reference count as of OpenMP
     /// 4.5.  The hold reference count is an OpenMP extension for the sake of
     /// OpenACC support.
     ///
     /// The 'ompx_hold' map type modifier is permitted only on "omp target" and
     /// "omp target data", and "delete" is permitted only on "omp target exit
     /// data" and associated runtime library routines.  As a result, we really
     /// need to implement "reset" functionality only for the dynamic reference
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;
+    /// Counter used to coordinate the deletion of an entry. Each target task
+    /// that decides the entry is to be deleted will increment the counter. As
+    /// there is a gap between deciding on deletion and actual deletion, e.g.,
+    /// to wait for the copy back, the task has to decrement the counter before
+    /// deletion is issued. Only if it is 0, and the ref count is 0 too, we
+    /// delete the entry.
+    unsigned short DeleteCount = 0;
   };
   // 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;
public:
 
   void setMayContainAttachedPointers() const {
     States->MayContainAttachedPointers = true;
   }
   bool getMayContainAttachedPointers() const {
     return States->MayContainAttachedPointers;
   }
 
+  /// 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
class TargetPointerResultTy {
   /// The corresponding target pointer
   void *TargetPointer = nullptr;
 
 public:
   TargetPointerResultTy() {}
 
   TargetPointerResultTy(FlagTy Flags, HostDataToTargetTy *Entry,
                         void *TargetPointer)
-      : Flags(Flags), Entry(Entry), TargetPointer(TargetPointer) {
-  }
+      : Flags(Flags), Entry(Entry), TargetPointer(TargetPointer) {}
 
   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; }

openmp/libomptarget/src/device.cpp

if (ForceDelete) {
       assert(TPR.isLast() == HT.decShouldRemove(UseHoldRefCount) &&
              "expected correct IsLast prediction for reset");
     }
 
     const char *RefCountAction;
     if (!UpdateRefCount) {
       RefCountAction = " (update suppressed)";
     } else if (TPR.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.incDeleteCount();
       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 {
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.");
 
   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",

openmp/libomptarget/src/omptarget.cpp

struct PostProcessingInfo {
   void *HstPtrBegin;
 
   /// Size of the data
   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;
 
+  /// The entry we currently copy back and which we might delete.
+  HostDataToTargetTy *Entry;
+
   PostProcessingInfo(void *HstPtr, int64_t Size, int64_t ArgType, bool DelEntry,
-                     TargetPointerResultTy TPR)
-      : HstPtrBegin(HstPtr), DataSize(Size), ArgType(ArgType), TPR(TPR),
-        DelEntry(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,
-                                    const TargetPointerResultTy &TPR) {
+                                    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 (!TPR.getEntry() || !TPR.getEntry()->getMayContainAttachedPointers())
+  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();) {
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.
+      if (auto *Entry = TPR.getEntry())
         PostProcessingPtrs.emplace_back(HstPtrBegin, DataSize, ArgTypes[I],
-                                      DelEntry && !TPR.isHostPtr(), TPR);
+                                        DelEntry && !TPR.isHostPtr(), Entry);
     }
   }
 
   // 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
+  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;
     DeviceTy::HDTTMapAccessorTy HDTTMap =
         Device.HostDataToTargetMap.getExclusiveAccessor(!Info.DelEntry);
 
     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->getDeleteThreadId() != std::this_thread::get_id()) {
+          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
auto CB = [&](ShadowPtrListTy::iterator &Itr) {
         DP("Removing shadow pointer " DPxMOD "\n", DPxPTR((void **)Itr->first));
         Itr = Device.ShadowPtrMap.erase(Itr);
       } else {
         ++Itr;
       }
       return OFFLOAD_SUCCESS;
     };
     applyToShadowMapEntries(Device, CB, Info.HstPtrBegin, Info.DataSize,
-                            Info.TPR);
+                            Info.Entry);
 
     // 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 (!Info.DelEntry ||
+        (FromMapperBase && FromMapperBase == Info.HstPtrBegin)) {
+      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 (Ret != OFFLOAD_SUCCESS) {
-        REPORT("Deallocating data from device failed.\n");
-        break;
-      }
-    }
-  }
 
   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);
auto CB = [&](ShadowPtrListTy::iterator &Itr) {
         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));
       ++Itr;
       return OFFLOAD_SUCCESS;
     };
-    applyToShadowMapEntries(Device, CB, HstPtrBegin, ArgSize, TPR);
+    applyToShadowMapEntries(Device, CB, HstPtrBegin, ArgSize, TPR.getEntry());
   }
 
   if (ArgType & OMP_TGT_MAPTYPE_TO) {
     DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n",
        ArgSize, DPxPTR(HstPtrBegin), DPxPTR(TgtPtrBegin));
     int Ret = Device.submitData(TgtPtrBegin, HstPtrBegin, ArgSize, AsyncInfo);
     if (Ret != OFFLOAD_SUCCESS) {
       REPORT("Copying data to device failed.\n");
       return OFFLOAD_FAIL;
     }
 
     auto CB = [&](ShadowPtrListTy::iterator &Itr) {
       DP("Restoring original target pointer value " DPxMOD " for target "
          "pointer " DPxMOD "\n",
          DPxPTR(Itr->second.TgtPtrVal), DPxPTR(Itr->second.TgtPtrAddr));
       Ret = Device.submitData(Itr->second.TgtPtrAddr, &Itr->second.TgtPtrVal,
                               sizeof(void *), AsyncInfo);
       if (Ret != OFFLOAD_SUCCESS)
         REPORT("Copying data to device failed.\n");
       ++Itr;
       return Ret;
     };
-    applyToShadowMapEntries(Device, CB, HstPtrBegin, ArgSize, TPR);
+    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,