[Openmp-commits] [PATCH] D133655: [OpenMP][libomptarget] Add mutex to safely read/modify HostPinnedAllocs map

[Kevin Sala Penadés via Phabricator via Openmp-commits]

kevinsala created this revision.
kevinsala added reviewers: jdoerfert, tianshilei1992, jhuber6.
kevinsala added a project: OpenMP.
Herald added subscribers: guansong, yaxunl.
Herald added a project: All.
kevinsala requested review of this revision.
Herald added subscribers: openmp-commits, sstefan1.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D133655

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


Index: openmp/libomptarget/plugins/cuda/src/rtl.cpp
===================================================================
--- openmp/libomptarget/plugins/cuda/src/rtl.cpp
+++ openmp/libomptarget/plugins/cuda/src/rtl.cpp
@@ -367,6 +367,7 @@
   /// 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 {
@@ -390,7 +391,10 @@
         MemAlloc = HostPtr;
         if (!checkResult(Err, "Error returned from cuMemAllocHost\n"))
           return nullptr;
-        HostPinnedAllocs[MemAlloc] = Kind;
+        {
+          std::lock_guard<std::mutex> Lock(HostPinnedAllocsMutex);
+          HostPinnedAllocs[MemAlloc] = Kind;
+        }
         break;
       case TARGET_ALLOC_SHARED:
         CUdeviceptr SharedPtr;
@@ -406,11 +410,14 @@
 
     int free(void *TgtPtr) override {
       CUresult Err;
-      // Host pinned memory must be freed differently.
-      TargetAllocTy Kind =
-          (HostPinnedAllocs.find(TgtPtr) == HostPinnedAllocs.end())
-              ? TARGET_ALLOC_DEFAULT
-              : TARGET_ALLOC_HOST;
+      TargetAllocTy Kind;
+      {
+        std::lock_guard<std::mutex> Lock(HostPinnedAllocsMutex);
+        // Host pinned memory must be freed differently.
+        Kind = (HostPinnedAllocs.find(TgtPtr) == HostPinnedAllocs.end())
+            ? TARGET_ALLOC_DEFAULT : TARGET_ALLOC_HOST;
+      }
+
       switch (Kind) {
       case TARGET_ALLOC_DEFAULT:
       case TARGET_ALLOC_DEVICE:
@@ -564,8 +571,13 @@
       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();


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D133655.459321.patch
Type: text/x-patch
Size: 2297 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20220911/f2116eb4/attachment.bin>