[OpenMP][libomptarget] Add mutex to safely read/modify HostPinnedAllocs map

[kevinsala]

Summary

The access to the HostPinnedAllocs map should be protected since different threads may be allocating and deallocating data concurrently. A thread allocating data may need to modify the map, so no other thread should access (even reading) the map at that point.

[jhuber6]

I have a review up that removes this map entirely, for this and other reasons D133053. The review seemed to have stalled as @tianshilei1992 had concerns about changing the interface and what that meant for backwards compatibility.

[tianshilei1992]

In D133655#3782567, @jhuber6 wrote:

I have a review up that removes this map entirely, for this and other reasons D133053. The review seemed to have stalled as @tianshilei1992 had concerns about changing the interface and what that meant for backwards compatibility.

Yeah, we still didn't make the consensus we need.

[jhuber6]

In D133655#3782570, @tianshilei1992 wrote:

In D133655#3782567, @jhuber6 wrote:

I have a review up that removes this map entirely, for this and other reasons D133053. The review seemed to have stalled as @tianshilei1992 had concerns about changing the interface and what that meant for backwards compatibility.

Yeah, we still didn't make the consensus we need.

There was a user-visible change where omp_target_free didn't work anymore for the special allocation types, but that was against the standard anyway. Then there was the plugin change. The issue was changing the plugin interface right? I could still just introduce a new one that is called by the old one, but if we remove backwards compatibility it would be much easier all around.

[jhuber6]

Is this solving an immediate error you're experiencing? If not we should wait for a consensus on removing this map entirely, otherwise if it's urgent we could probably land this temporarily.

[kevinsala]

In D133655#3785431, @jhuber6 wrote:

Is this solving an immediate error you're experiencing? If not we should wait for a consensus on removing this map entirely, otherwise if it's urgent we could probably land this temporarily.

No, actually I didn't experience any error related to it. I was debugging something else and saw this concurrent access not being protected. Thus, it's not a problem to wait for a consensus on removing the map.

[jhuber6]

In D133655#3785490, @kevinsala wrote:

In D133655#3785431, @jhuber6 wrote:

Is this solving an immediate error you're experiencing? If not we should wait for a consensus on removing this map entirely, otherwise if it's urgent we could probably land this temporarily.

No, actually I didn't experience any error related to it. I was debugging something else and saw this concurrent access not being protected. Thus, it's not a problem to wait for a consensus on removing the map.

I landed D133655 which removes this map entirely, thanks for the patch regardless.

[kevinsala]

Great! Then I'm closing this review. Thanks @jhuber6

openmp/libomptarget/plugins/cuda/src/rtl.cpp

class DeviceRTLTy {
   enum class PeerAccessState : uint8_t { Unkown, Yes, No };
   std::vector<std::vector<PeerAccessState>> PeerAccessMatrix;
   std::mutex PeerAccessMatrixLock;
 
   /// A class responsible for interacting with device native runtime library to
   /// allocate and free memory.
   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
     std::unordered_map<void *, TargetAllocTy> HostPinnedAllocs;
+    std::mutex HostPinnedAllocsMutex;
 
   public:
     void *allocate(size_t Size, void *, TargetAllocTy Kind) override {
       if (Size == 0)
         return nullptr;
 
       void *MemAlloc = nullptr;
       CUresult Err;
       switch (Kind) {
       case TARGET_ALLOC_DEFAULT:
       case TARGET_ALLOC_DEVICE:
         CUdeviceptr DevicePtr;
         Err = cuMemAlloc(&DevicePtr, Size);
         MemAlloc = (void *)DevicePtr;
         if (!checkResult(Err, "Error returned from cuMemAlloc\n"))
           return nullptr;
         break;
       case TARGET_ALLOC_HOST:
         void *HostPtr;
         Err = cuMemAllocHost(&HostPtr, Size);
         MemAlloc = HostPtr;
         if (!checkResult(Err, "Error returned from cuMemAllocHost\n"))
           return nullptr;
+        {
+          std::lock_guard<std::mutex> Lock(HostPinnedAllocsMutex);
           HostPinnedAllocs[MemAlloc] = Kind;
+        }
         break;
       case TARGET_ALLOC_SHARED:
         CUdeviceptr SharedPtr;
         Err = cuMemAllocManaged(&SharedPtr, Size, CU_MEM_ATTACH_GLOBAL);
         MemAlloc = (void *)SharedPtr;
         if (!checkResult(Err, "Error returned from cuMemAllocManaged\n"))
           return nullptr;
         break;
       }
 
       return MemAlloc;
     }
 
     int free(void *TgtPtr) override {
       CUresult Err;
+      TargetAllocTy Kind;
+      {
+        std::lock_guard<std::mutex> Lock(HostPinnedAllocsMutex);
         // Host pinned memory must be freed differently.
-      TargetAllocTy Kind =
-          (HostPinnedAllocs.find(TgtPtr) == HostPinnedAllocs.end())
-              ? TARGET_ALLOC_DEFAULT
-              : TARGET_ALLOC_HOST;
+        Kind = (HostPinnedAllocs.find(TgtPtr) == HostPinnedAllocs.end())
+            ? TARGET_ALLOC_DEFAULT : TARGET_ALLOC_HOST;
+      }
+
       switch (Kind) {
       case TARGET_ALLOC_DEFAULT:
       case TARGET_ALLOC_DEVICE:
       case TARGET_ALLOC_SHARED:
         Err = cuMemFree((CUdeviceptr)TgtPtr);
         if (!checkResult(Err, "Error returned from cuMemFree\n"))
           return OFFLOAD_FAIL;
         break;
if (const char *EnvStr = getenv("LIBOMPTARGET_SHARED_MEMORY_SIZE")) {
          DynamicMemorySize);
     }
     if (const char *EnvStr = getenv("LIBOMPTARGET_NUM_INITIAL_STREAMS")) {
       // LIBOMPTARGET_NUM_INITIAL_STREAMS has been set
       NumInitialStreams = std::stoi(EnvStr);
       DP("Parsed LIBOMPTARGET_NUM_INITIAL_STREAMS=%d\n", NumInitialStreams);
     }
 
-    for (int I = 0; I < NumberOfDevices; ++I)
-      DeviceAllocators.emplace_back();
+    // The CUDADeviceAllocatorTy has a std::mutex and these cannot be moved or
+    // copied. The swap operation guarantees no move, copy or swap on the
+    // individual elements.
+    {
+      std::vector<CUDADeviceAllocatorTy> DeviceAllocatorsTemp(NumberOfDevices);
+      DeviceAllocators.swap(DeviceAllocatorsTemp);
+    }
 
     // Get the size threshold from environment variable
     std::pair<size_t, bool> Res = MemoryManagerTy::getSizeThresholdFromEnv();
     UseMemoryManager = Res.second;
     size_t MemoryManagerThreshold = Res.first;
 
     if (UseMemoryManager)
       for (int I = 0; I < NumberOfDevices; ++I)