[OpenMP][FIX] Ensure exclusive access to the HDTT map

This patch solves two problems with the `HostDataToTargetMap` (HDTT
map) which caused races and crashes before:

1) Any access to the HDTT map needs to be exclusive access. This was not
   the case for the "dump table" traversals that could collide with
   updates by other threads. The new `Accessor` and `ProtectedObject`
   wrappers will ensure we have a hard time introducing similar races in
   the future. Note that we could allow multiple concurrent
   read-accesses but that feature can be added to the `Accessor` API
   later.
2) The elements of the HDTT map were `HostDataToTargetTy` objects which
   meant that they could be copied/moved/deleted as the map was changed.
   However, we sometimes kept pointers to these elements around after we
   gave up the map lock which caused potential races again. The new
   indirection through `HostDataToTargetMapKeyTy` will allows us to
   modify the map while keeping the (interesting part of the) entries
   valid. To offset potential cost we duplicate the ordering key of the
   entry which avoids an additional indirect lookup.

We should replace more objects with "protected objects" as we go.

Differential Revision: https://reviews.llvm.org/D121057

openmp/libomptarget/include/ExclusiveAccess.h

@@ -0,0 +1,100 @@
+//===---- ExclusiveAccess.h - Helper for exclusive access data structures -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OMPTARGET_EXCLUSIVE_ACCESS
+#define OMPTARGET_EXCLUSIVE_ACCESS
+
+#include <cstddef>
+#include <cstdint>
+#include <mutex>
+
+/// Forward declaration.
+template <typename Ty> struct Accessor;
+
+/// A protected object is a simple wrapper to allocate an object of type \p Ty
+/// together with a mutex that guards accesses to the object. The only way to
+/// access the object is through the "exclusive accessor" which will lock the
+/// mutex accordingly.
+template <typename Ty> struct ProtectedObj {
+  using AccessorTy = Accessor<Ty>;
+
+  /// Get an exclusive access Accessor object. \p DoNotGetAccess allows to
+  /// create an accessor that is not owning anything based on a boolean
+  /// condition.
+  AccessorTy &&getExclusiveAccessor(bool DoNotGetAccess = false);
+
+private:
+  Ty Obj;
+  std::mutex Mtx;
+  friend struct Accessor<Ty>;
+};
+
+/// Helper to provide transparent exclusive access to protected objects.
+template <typename Ty> struct Accessor {
+  /// Default constructor does not own anything and cannot access anything.
+  Accessor() : Ptr(nullptr) {}
+
+  /// 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; }
+
+  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() {
+    unlock();
+    Ptr = nullptr;
+  }
+
+  /// Provide transparent access to the underlying object.
+  Ty &operator*() {
+    assert(Ptr && "Trying to access an object through a non-owning (or "
+                  "destroyed) accessor!");
+    return Ptr->Obj;
+  }
+  Ty *operator->() {
+    assert(Ptr && "Trying to access an object through a non-owning (or "
+                  "destroyed) accessor!");
+    return &Ptr->Obj;
+  }
+
+private:
+  /// Lock the underlying object if there is one.
+  void lock() {
+    if (Ptr)
+      Ptr->Mtx.lock();
+  }
+
+  /// Unlock the underlying object if there is one.
+  void unlock() {
+    if (Ptr)
+      Ptr->Mtx.unlock();
+  }
+
+  /// Pointer to the underlying object or null if the accessor lost access,
+  /// e.g., after a destroy call.
+  ProtectedObj<Ty> *Ptr;
+};
+
+template <typename Ty>
+Accessor<Ty> &&ProtectedObj<Ty>::getExclusiveAccessor(bool DoNotGetAccess) {
+  if (DoNotGetAccess)
+    return std::move(Accessor<Ty>());
+  return std::move(Accessor<Ty>(*this));
+}
+
+#endif

openmp/libomptarget/include/device.h

@@ -15,13 +15,15 @@
 
 #include <cassert>
 #include <cstddef>
+#include <cstdint>
 #include <list>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <set>
 #include <vector>
 
+#include "ExclusiveAccess.h"
 #include "omptarget.h"
 #include "rtl.h"
 
@@ -214,28 +216,41 @@ struct HostDataToTargetTy {
   void unlock() const { States->UpdateMtx.unlock(); }
 };
 
-typedef uintptr_t HstPtrBeginTy;
-inline bool operator<(const HostDataToTargetTy &lhs, const HstPtrBeginTy &rhs) {
-  return lhs.HstPtrBegin < rhs;
+/// 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
+/// load of HDTT. HDTT is a pointer to allow the modification of the set without
+/// invalidating HDTT entries which can now be inspected at the same time.
+struct HostDataToTargetMapKeyTy {
+  uintptr_t KeyValue;
+
+  HostDataToTargetMapKeyTy(void *Key) : KeyValue(uintptr_t(Key)) {}
+  HostDataToTargetMapKeyTy(HostDataToTargetTy *HDTT)
+      : KeyValue(HDTT->HstPtrBegin), HDTT(HDTT) {}
+  HostDataToTargetTy *HDTT;
+};
+inline bool operator<(const HostDataToTargetMapKeyTy &lhs,
+                      const uintptr_t &rhs) {
+  return lhs.KeyValue < rhs;
 }
-inline bool operator<(const HstPtrBeginTy &lhs, const HostDataToTargetTy &rhs) {
-  return lhs < rhs.HstPtrBegin;
+inline bool operator<(const uintptr_t &lhs,
+                      const HostDataToTargetMapKeyTy &rhs) {
+  return lhs < rhs.KeyValue;
 }
-inline bool operator<(const HostDataToTargetTy &lhs,
-                      const HostDataToTargetTy &rhs) {
-  return lhs.HstPtrBegin < rhs.HstPtrBegin;
+inline bool operator<(const HostDataToTargetMapKeyTy &lhs,
+                      const HostDataToTargetMapKeyTy &rhs) {
+  return lhs.KeyValue < rhs.KeyValue;
 }
 
-typedef std::set<HostDataToTargetTy, std::less<>> HostDataToTargetListTy;
-
 struct LookupResult {
   struct {
     unsigned IsContained : 1;
     unsigned ExtendsBefore : 1;
     unsigned ExtendsAfter : 1;
   } Flags;
 
-  HostDataToTargetListTy::iterator Entry;
+  /// The corresponding map table entry which is stable.
+  HostDataToTargetTy *Entry = nullptr;
 
   LookupResult() : Flags({0, 0, 0}), Entry() {}
 };
@@ -250,8 +265,8 @@ struct TargetPointerResultTy {
     unsigned IsHostPointer : 1;
   } Flags = {0, 0};
 
-  /// The iterator to the corresponding map table entry
-  HostDataToTargetListTy::iterator MapTableEntry{};
+  /// The corresponding map table entry which is stable.
+  HostDataToTargetTy *Entry = nullptr;
 
   /// The corresponding target pointer
   void *TargetPointer = nullptr;
@@ -282,12 +297,24 @@ struct DeviceTy {
   std::once_flag InitFlag;
   bool HasPendingGlobals;
 
-  HostDataToTargetListTy HostDataToTargetMap;
+  /// Host data to device map type with a wrapper key indirection that allows
+  /// concurrent modification of the entries without invalidating the underlying
+  /// entries.
+  using HostDataToTargetListTy =
+      std::set<HostDataToTargetMapKeyTy, std::less<>>;
+
+  /// The HDTTMap is a protected object that can only be accessed by one thread
+  /// at a time.
+  ProtectedObj<HostDataToTargetListTy> HostDataToTargetMap;
+
+  /// The type used to access the HDTT map.
+  using HDTTMapAccessorTy = decltype(HostDataToTargetMap)::AccessorTy;
+
   PendingCtorsDtorsPerLibrary PendingCtorsDtors;
 
   ShadowPtrListTy ShadowPtrMap;
 
-  std::mutex DataMapMtx, PendingGlobalsMtx, ShadowMtx;
+  std::mutex PendingGlobalsMtx, ShadowMtx;
 
   // NOTE: Once libomp gains full target-task support, this state should be
   // moved into the target task in libomp.
@@ -303,7 +330,11 @@ struct DeviceTy {
   // Return true if data can be copied to DstDevice directly
   bool isDataExchangable(const DeviceTy &DstDevice);
 
-  LookupResult lookupMapping(void *HstPtrBegin, int64_t Size);
+  /// Lookup the mapping of \p HstPtrBegin in \p HDTTMap. The accessor ensures
+  /// exclusive access to the HDTT map.
+  LookupResult lookupMapping(HDTTMapAccessorTy &HDTTMap, void *HstPtrBegin,
+                             int64_t Size);
+
   /// 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 required, the memory region pointed by \p HstPtrBegin of size
@@ -320,7 +351,12 @@ struct DeviceTy {
                    bool HasFlagAlways, bool IsImplicit, bool UpdateRefCount,
                    bool HasCloseModifier, bool HasPresentModifier,
                    bool HasHoldModifier, AsyncInfoTy &AsyncInfo);
-  void *getTgtPtrBegin(void *HstPtrBegin, int64_t Size);
+
+  /// Return the target pointer for \p HstPtrBegin in \p HDTTMap. The accessor
+  /// 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,