[OpenMP][Offloading] Fix data race in data mapping by using two locks

[tianshilei1992]

Summary

This patch tries to partially fix one of the two data race issues reported in
[1] by introducing a per-entry mutex. Additional discussion can also be found in
D104418, which will also be refined to fix another data race problem.

Here is how it works. Like before, DataMapMtx is still being used for mapping
table lookup and update. In any case, we will get a table entry. If we need to
make a data transfer (update the data on the device), we need to lock the entry
right before releasing DataMapMtx, and the issue of data transfer should be
after releasing DataMapMtx, and the entry is unlocked afterwards. This can
guarantee that: 1) issue of data movement is not in critical region, which will
not affect performance too much, and also will not affect other threads that don't
touch the same entry; 2) if another thread accesses the same entry, the state of
data movement is consistent (which requires that a thread must first get the
update lock before getting data movement information).

For a target that doesn't support async data transfer, issue of data movement is
data transfer. This two-lock design can potentially improve concurrency compared
with the design that guards data movement with DataMapMtx as well. For a target
that supports async data movement, we could simply attach the event between the
issue of data movement and unlock the entry. For a thread that wants to get the
event, it must first get the lock. This can also get rid of the busy wait until
the event pointer is valid.

Reference:
[1] https://bugs.llvm.org/show_bug.cgi?id=49940

[tianshilei1992]

openmp/libomptarget/src/device.cpp

20

Hmm, not sure why it is here. Might be included by clangd automatically.

openmp/libomptarget/src/device.h

222

IsHostPtr will be removed later by changing the return type of this function to a data structure containing flags, iterator to the entry, etc., as shown in D104382.

[gValarini]

openmp/libomptarget/src/device.cpp

181

As far as I can understand, this lock is protecting both the queries/updates to the data map tables and the data transfer issue, which indeed fixes the problem described in the patch.

But I think this can reduce the possibility of concurrent data transfers to the same device when there is no overlapping between the data regions (as discussed in 49940), especially when the device only implements the synchronous API. What do you think about implementing two locking regions?

1. A DataMapMtx (the one already present) would be locked only when querying/updating the data map. This is mandatory since the mapping tables are global to each device and must be synchronized among the many runtime threads;
2. A DataMapEntry (one per mapped entry, stored at the HostDataToTargetTy) would be locked during the data transfer issue stage, avoiding the data corruption problem.

This could allow multiple threads to issue (or even execute, on the sync. case) multiple data transfer while still being correctly synchronized.

Note: the second locking region must start right before DataMapMtx is unlocked, so the same thread that first allocated the data is the one that issues the data transfer. This is important because other threads that use the same data won't issue such transfer (unless the always modifier is used).

[tianshilei1992]

openmp/libomptarget/src/device.cpp

181

Note: the second locking region must start right before DataMapMtx is unlocked, so the same thread that first allocated the data is the one that issues the data transfer. This is important because other threads that use the same data won't issue such transfer (unless the always modifier is used).

I'm not sure I understand your idea. If the second lock is inside the first lock, what's the point of it? For synchronous data movement, there is no issue stage. When the data movement is called, it just starts to copy and return when it's done.

[tianshilei1992]

rebase

[ye-luo]

openmp/libomptarget/src/device.cpp

181

I think @gValarini means the second lock overlaps with the first lock not inside.
T1 : lock table -> lock entry A -> unlock table > issue transfer A -> unlock entry A
T2: lock table (only need to wait for the unlock of table from T1 instead of the unlock of entry A) -> lock entry B -> unlock table > issue transfer B -> unlock entry B.
In this way, transfer A and B can happen concurrently.

[tianshilei1992]

openmp/libomptarget/src/device.cpp

181

Aha, interesting! Yeah, I get it. Pretty tricky. ;-)

[tianshilei1992]

introduced a per-entry lock to reduce resource contention

[tianshilei1992]

rebase

[tianshilei1992]

rebase and ping

[tianshilei1992]

rebase and ping

[grokos]

What about other instances where we have access to the HostDataToTargetMap followed be data transfers, e.g. when removing a mapping and copy data back to the host? Shouldn't we apply the same logic there?

openmp/libomptarget/src/device.cpp

232–233

Redundant, if we use USM and don't have the close modifier, then MoveData will be false already.

openmp/libomptarget/src/omptarget.cpp

519–531

This part must be updated to match the changes introduced in D105121. Also, why do we set MoveData = true only when the ALWAYS flag is set? It should also be true if the mapping is new, but the problem here is that we don't know whether or not we have a new mapping because getTargetPointer (ex getOrAllocTgtPtr) hasn't been called yet. So I think MoveData should be a 3-state integer (maybe use an enum for it) which will pass instructions from targetDataBegin to getTargetPointer about what to do:

• Case 0: false, the TO flag hasn't been set (or we are using USM without the close modifier etc.), so don't move data no matter what
• Case 1: true, the TO and ALWAYS flags have been set, so move data
• Case 2: undecided, the TO flag has been set but targetDataBegin doesn't know whether the mapping is new, so move data if RefCount==1

So this part of the logic should look something like this:

enum MoveState MoveData = MoveState::false;
if (arg_types[i] & OMP_TGT_MAPTYPE_TO) {
  if (!(PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) ||
      HasCloseModifier) {
    if (arg_types[i] & OMP_TGT_MAPTYPE_ALWAYS)
      MoveData = MoveState::true;
    else
      MoveData = MoveState::undecided;
  }
}

551–552

Here we also need to take care of this data transfer. Copying the pointee object onto the device and updating the device pointer should behave as a single atomic operation.

[tianshilei1992]

openmp/libomptarget/src/omptarget.cpp

519–531

My intention is to cover new entry in getTargetPointer, but I did forget to set MoveData to true in the function if there is a new entry.

[tianshilei1992]

rebase

[tianshilei1992]

In D104555#2874750, @grokos wrote:

What about other instances where we have access to the HostDataToTargetMap followed be data transfers, e.g. when removing a mapping and copy data back to the host? Shouldn't we apply the same logic there?

I didn't come up with a case that can cause data race with D2H because we have a mandatory synchronization at the end of target region.

[grokos]

openmp/libomptarget/src/device.cpp

244

This is not correct either. getTargetPointer() doesn't know whether the TO flag is set, yet it will perform a data transfer even if we just map(alloc : ...) as long as we are mapping a new entry. That's why I proposed that you use a 3-state integer instead of a bool for MoveData.

[tianshilei1992]

openmp/libomptarget/src/device.cpp

244

Oh, you're right. I'll update it.

[tianshilei1992]

update the logic to determine if data movement is needed

[tianshilei1992]

openmp/libomptarget/src/omptarget.cpp

519–531

I refined the logic but seems like it almost broke all declare mapper tests…Need to investigate further.

[tianshilei1992]

fixed issues

[tianshilei1992]

fixed shadow mapping

[tianshilei1992]

Ping. Hope the bug fix could go into 13.

[grokos]

This patch looks good now! Let's commit it and then move on with D104418.

openmp/libomptarget/src/device.h

struct HostDataToTargetTy {
   map_var_info_t HstPtrName; // Optional source name of mapped variable.
 
   uintptr_t TgtPtrBegin; // target info.
 
 private:
   /// use mutable to allow modification via std::set iterator which is const.
   mutable uint64_t RefCount;
   static const uint64_t INFRefCount = ~(uint64_t)0;
+  /// 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::shared_ptr<std::mutex> UpdateMtx;
 
 public:
   HostDataToTargetTy(uintptr_t BP, uintptr_t B, uintptr_t E, uintptr_t TB,
                      map_var_info_t Name = nullptr, bool IsINF = false)
       : HstPtrBase(BP), HstPtrBegin(B), HstPtrEnd(E), HstPtrName(Name),
-        TgtPtrBegin(TB), RefCount(IsINF ? INFRefCount : 1) {}
+        TgtPtrBegin(TB), RefCount(IsINF ? INFRefCount : 1),
+        UpdateMtx(std::make_shared<std::mutex>()) {}
 
   uint64_t getRefCount() const { return RefCount; }
 
   uint64_t resetRefCount() const {
     if (RefCount != INFRefCount)
       RefCount = 1;
 
     return RefCount;
public:
 
   /// Should one decrement of the reference count (after resetting it if
   /// \c AfterReset) remove this mapping?
   bool decShouldRemove(bool AfterReset = false) const {
     if (AfterReset)
       return !isRefCountInf();
     return getRefCount() == 1;
   }
+
+  void lock() const { UpdateMtx->lock(); }
+
+  void unlock() const { UpdateMtx->unlock(); }
 };
 
 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 PendingCtorDtorListsTy {
   std::list<void *> PendingCtors;
   std::list<void *> PendingDtors;
 };
 typedef std::map<__tgt_bin_desc *, PendingCtorDtorListsTy>
     PendingCtorsDtorsPerLibrary;
 
+enum class MoveDataStateTy : uint32_t { REQUIRED, NONE, UNKNOWN };
+
 struct DeviceTy {
   int32_t DeviceID;
   RTLInfoTy *RTL;
   int32_t RTLDeviceID;
 
   bool IsInit;
   std::once_flag InitFlag;
   bool HasPendingGlobals;
struct DeviceTy {
   DeviceTy &operator=(const DeviceTy &D);
 
   ~DeviceTy();
 
   // Return true if data can be copied to DstDevice directly
   bool isDataExchangable(const DeviceTy &DstDevice);
 
   LookupResult lookupMapping(void *HstPtrBegin, int64_t Size);
-  TargetPointerResultTy getOrAllocTgtPtr(void *HstPtrBegin, void *HstPtrBase,
-                                         int64_t Size,
-                                         map_var_info_t HstPtrName,
-                                         bool IsImplicit, bool UpdateRefCount,
-                                         bool HasCloseModifier,
-                                         bool HasPresentModifier);
+  /// Get the target pointer based on host pointer begin and base. If the
+  /// mapping already exists, the target pointer will be returned directly. In
+  /// addition, if \p MoveData is true, the memory region pointed by \p
+  /// HstPtrBegin of size \p Size will also be transferred to the device. If the
+  /// mapping doesn't exist, and if unified memory is not enabled, a new mapping
+  /// will be created and the data will also be transferred accordingly. nullptr
+  /// will be returned because of any of following reasons:
+  /// - Data allocation failed;
+  /// - The user tried to do an illegal mapping;
+  /// - Data transfer issue fails.
+  TargetPointerResultTy
+  getTargetPointer(void *HstPtrBegin, void *HstPtrBase, int64_t Size,
+                   map_var_info_t HstPtrName, MoveDataStateTy MoveData,
+                   bool IsImplicit, bool UpdateRefCount, bool HasCloseModifier,
+                   bool HasPresentModifier, AsyncInfoTy &AsyncInfo);
   void *getTgtPtrBegin(void *HstPtrBegin, int64_t Size);
   void *getTgtPtrBegin(void *HstPtrBegin, int64_t Size, bool &IsLast,
                        bool UpdateRefCount, bool &IsHostPtr,
                        bool MustContain = false, bool ForceDelete = false);
   int deallocTgtPtr(void *TgtPtrBegin, int64_t Size,
                     bool HasCloseModifier = false);
   int associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size);
   int disassociatePtr(void *HstPtrBegin);

openmp/libomptarget/src/device.cpp

 
 #include "device.h"
 #include "private.h"
 #include "rtl.h"
 
 #include <cassert>
 #include <climits>
 #include <cstdio>
 #include <string>
 
 DeviceTy::DeviceTy(const DeviceTy &D)
     : DeviceID(D.DeviceID), RTL(D.RTL), RTLDeviceID(D.RTLDeviceID),
       IsInit(D.IsInit), InitFlag(), HasPendingGlobals(D.HasPendingGlobals),
       HostDataToTargetMap(D.HostDataToTargetMap),
       PendingCtorsDtors(D.PendingCtorsDtors), ShadowPtrMap(D.ShadowPtrMap),
       DataMapMtx(), PendingGlobalsMtx(), ShadowMtx(),
       LoopTripCnt(D.LoopTripCnt) {}
LookupResult DeviceTy::lookupMapping(void *HstPtrBegin, int64_t Size) {
   if (lr.Flags.ExtendsAfter) {
     DP("WARNING: Pointer is already mapped but section extends beyond mapped "
        "region\n");
   }
 
   return lr;
 }
 
-// Used by targetDataBegin
-// Return a struct containing target pointer begin (where the data will be
-// moved).
-// Allocate memory if this is the first occurrence of this mapping.
-// Increment the reference counter.
-// If the target pointer is NULL, then either data allocation failed or the user
-// tried to do an illegal mapping.
-// The returned struct also returns an iterator to the map table entry
-// corresponding to the host pointer (if exists), and two flags indicating
-// whether the entry is just created, and if the target pointer included is
-// actually a host pointer (when unified memory enabled).
 TargetPointerResultTy
-DeviceTy::getOrAllocTgtPtr(void *HstPtrBegin, void *HstPtrBase, int64_t Size,
-                           map_var_info_t HstPtrName, bool IsImplicit,
-                           bool UpdateRefCount, bool HasCloseModifier,
-                           bool HasPresentModifier) {
-  void *TargetPointer = NULL;
-  bool IsNew = false;
+DeviceTy::getTargetPointer(void *HstPtrBegin, void *HstPtrBase, int64_t Size,
+                           map_var_info_t HstPtrName, MoveDataStateTy MoveData,
+                           bool IsImplicit, bool UpdateRefCount,
+                           bool HasCloseModifier, bool HasPresentModifier,
+                           AsyncInfoTy &AsyncInfo) {
+  void *TargetPointer = nullptr;
   bool IsHostPtr = false;
+  bool IsNew = false;
+
   DataMapMtx.lock();
+
   LookupResult LR = lookupMapping(HstPtrBegin, Size);
   auto Entry = LR.Entry;
 
   // 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;
if (LR.Flags.IsContained ||
     // are respected.
     // In addition to the mapping rules above, the close map modifier forces the
     // mapping of the variable to the device.
     if (Size) {
       DP("Return HstPtrBegin " DPxMOD " Size=%" PRId64 " for unified shared "
          "memory\n",
          DPxPTR((uintptr_t)HstPtrBegin), Size);
       IsHostPtr = true;
       TargetPointer = HstPtrBegin;
     }
   } else if (HasPresentModifier) {
     DP("Mapping required by 'present' map type modifier does not exist for "
        "HstPtrBegin=" DPxMOD ", Size=%" PRId64 "\n",
        DPxPTR(HstPtrBegin), Size);
     MESSAGE("device mapping required by 'present' map type modifier does not "
             "exist for host address " DPxMOD " (%" PRId64 " bytes)",
             DPxPTR(HstPtrBegin), Size);
   } else if (Size) {
     // If it is not contained and Size > 0, we should create a new entry for it.
     IsNew = true;
+    if (MoveData == MoveDataStateTy::UNKNOWN)
+      MoveData = MoveDataStateTy::REQUIRED;
     uintptr_t Ptr = (uintptr_t)allocData(Size, HstPtrBegin);
     Entry = HostDataToTargetMap
                 .emplace((uintptr_t)HstPtrBase, (uintptr_t)HstPtrBegin,
                          (uintptr_t)HstPtrBegin + Size, Ptr, HstPtrName)
                 .first;
     INFO(OMP_INFOTYPE_MAPPING_CHANGED, DeviceID,
          "Creating new map entry with "
          "HstPtrBegin=" DPxMOD ", TgtPtrBegin=" DPxMOD ", Size=%ld, "
          "RefCount=%s, Name=%s\n",
          DPxPTR(HstPtrBegin), DPxPTR(Ptr), Size, Entry->refCountToStr().c_str(),
          (HstPtrName) ? getNameFromMapping(HstPtrName).c_str() : "unknown");
     TargetPointer = (void *)Ptr;
   }
 
+  // If the target pointer is valid, and we need to transfer data, issue the
+  // data transfer.
+  if (TargetPointer && (MoveData == MoveDataStateTy::REQUIRED)) {
+    // Lock the entry before releasing the mapping table lock such that another
+    // thread that could issue data movement will get the right result.
+    Entry->lock();
+    // Release the mapping table lock right after the entry is locked.
+    DataMapMtx.unlock();
+
+    DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n", Size,
+       DPxPTR(HstPtrBegin), DPxPTR(TargetPointer));
+
+    int Ret = submitData(TargetPointer, HstPtrBegin, Size, AsyncInfo);
+
+    // Unlock the entry immediately after the data movement is issued.
+    Entry->unlock();
+
+    if (Ret != OFFLOAD_SUCCESS) {
+      REPORT("Copying data to device failed.\n");
+      // We will also return nullptr if the data movement fails because that
+      // pointer points to a corrupted memory region so it doesn't make any
+      // sense to continue to use it.
+      TargetPointer = nullptr;
+    }
+  } else {
+    // Release the mapping table lock directly.
     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 &IsHostPtr,

openmp/libomptarget/src/omptarget.cpp

if (arg_types[i] & OMP_TGT_MAPTYPE_PTR_AND_OBJ) {
       //
       // The map entry for s comes first, and the PTR_AND_OBJ entry comes
       // afterward, so the pointer is already allocated by the time the
       // PTR_AND_OBJ entry is handled below, and PointerTgtPtrBegin is thus
       // non-null.  However, "declare target link" can produce a PTR_AND_OBJ
       // entry for a global that might not already be allocated by the time the
       // PTR_AND_OBJ entry is handled below, and so the allocation might fail
       // when HasPresentModifier.
-      Pointer_TPR = Device.getOrAllocTgtPtr(
-          HstPtrBase, HstPtrBase, sizeof(void *), nullptr, IsImplicit,
-          UpdateRef, HasCloseModifier, HasPresentModifier);
+      Pointer_TPR = Device.getTargetPointer(
+          HstPtrBase, HstPtrBase, sizeof(void *), nullptr,
+          MoveDataStateTy::NONE, IsImplicit, UpdateRef, HasCloseModifier,
+          HasPresentModifier, AsyncInfo);
       PointerTgtPtrBegin = Pointer_TPR.TargetPointer;
       IsHostPtr = Pointer_TPR.Flags.IsHostPointer;
       if (!PointerTgtPtrBegin) {
         REPORT("Call to getOrAllocTgtPtr returned null pointer (%s).\n",
                HasPresentModifier ? "'present' map type modifier"
                                   : "device failure or illegal mapping");
         return OFFLOAD_FAIL;
       }
       DP("There are %zu bytes allocated at target address " DPxMOD " - is%s new"
          "\n",
          sizeof(void *), DPxPTR(PointerTgtPtrBegin),
          (Pointer_TPR.Flags.IsNewEntry ? "" : " not"));
       Pointer_HstPtrBegin = HstPtrBase;
       // modify current entry.
       HstPtrBase = *(void **)HstPtrBase;
       // No need to update pointee ref count for the first element of the
       // subelement that comes from mapper.
       UpdateRef =
           (!FromMapper || i != 0); // subsequently update ref count of pointee
     }
 
-    auto TPR = Device.getOrAllocTgtPtr(HstPtrBegin, HstPtrBase, data_size,
-                                       HstPtrName, IsImplicit, UpdateRef,
-                                       HasCloseModifier, HasPresentModifier);
+    MoveDataStateTy MoveData = MoveDataStateTy::NONE;
+    const bool UseUSM = PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY;
+    const bool HasFlagTo = arg_types[i] & OMP_TGT_MAPTYPE_TO;
+    const bool HasFlagAlways = arg_types[i] & OMP_TGT_MAPTYPE_ALWAYS;
+    if (HasFlagTo && (!UseUSM || HasCloseModifier))
+      MoveData = HasFlagAlways ? MoveDataStateTy::REQUIRED
+                               : MoveData = MoveDataStateTy::UNKNOWN;
+
+    auto TPR = Device.getTargetPointer(
+        HstPtrBegin, HstPtrBase, data_size, HstPtrName, MoveData, IsImplicit,
+        UpdateRef, HasCloseModifier, HasPresentModifier, AsyncInfo);
     void *TgtPtrBegin = TPR.TargetPointer;
     IsHostPtr = TPR.Flags.IsHostPointer;
     // If data_size==0, then the argument could be a zero-length pointer to
     // NULL, so getOrAlloc() returning NULL is not an error.
     if (!TgtPtrBegin && (data_size || HasPresentModifier)) {
       REPORT("Call to getOrAllocTgtPtr returned null pointer (%s).\n",
              HasPresentModifier ? "'present' map type modifier"
                                 : "device failure or illegal mapping");
       return OFFLOAD_FAIL;
     }
     DP("There are %" PRId64 " bytes allocated at target address " DPxMOD
        " - is%s new\n",
        data_size, DPxPTR(TgtPtrBegin), (TPR.Flags.IsNewEntry ? "" : " not"));
 
     if (arg_types[i] & OMP_TGT_MAPTYPE_RETURN_PARAM) {
       uintptr_t Delta = (uintptr_t)HstPtrBegin - (uintptr_t)HstPtrBase;
       void *TgtPtrBase = (void *)((uintptr_t)TgtPtrBegin - Delta);
       DP("Returning device pointer " DPxMOD "\n", DPxPTR(TgtPtrBase));
       args_base[i] = TgtPtrBase;
     }
 
-    if (arg_types[i] & OMP_TGT_MAPTYPE_TO) {
-      bool copy = false;
-      if (!(PM->RTLs.RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY) ||
-          HasCloseModifier) {
-        if (TPR.Flags.IsNewEntry || (arg_types[i] & OMP_TGT_MAPTYPE_ALWAYS))
-          copy = true;
-      }
-
-      if (copy && !IsHostPtr) {
-        DP("Moving %" PRId64 " bytes (hst:" DPxMOD ") -> (tgt:" DPxMOD ")\n",
-           data_size, DPxPTR(HstPtrBegin), DPxPTR(TgtPtrBegin));
-        int rt =
-            Device.submitData(TgtPtrBegin, HstPtrBegin, data_size, AsyncInfo);
-        if (rt != OFFLOAD_SUCCESS) {
-          REPORT("Copying data to device failed.\n");
-          return OFFLOAD_FAIL;
-        }
-      }
-    }
-
     if (arg_types[i] & OMP_TGT_MAPTYPE_PTR_AND_OBJ && !IsHostPtr) {
       // Check whether we need to update the pointer on the device
       bool UpdateDevPtr = false;
 
       uint64_t Delta = (uint64_t)HstPtrBegin - (uint64_t)HstPtrBase;
       void *ExpectedTgtPtrBase = (void *)((uint64_t)TgtPtrBegin - Delta);
 
       Device.ShadowMtx.lock();
       auto Entry = Device.ShadowPtrMap.find(Pointer_HstPtrBegin);
       // If this pointer is not in the map we need to insert it. If the map
       // contains a stale entry, we need to update it (e.g. if the pointee was
       // deallocated and later on is reallocated at another device address). The
       // latter scenario is the subject of LIT test env/base_ptr_ref_count.c. An
       // entry is removed from ShadowPtrMap only when the PTR of a PTR_AND_OBJ