[Openmp-commits] [PATCH] D122014: [OpenMP][CUDA] Fix potential program crash caused by double free resources

[Shilei Tian via Phabricator via Openmp-commits]

tianshilei1992 updated this revision to Diff 416816.
tianshilei1992 added a comment.

rebase and fix comments


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D122014/new/

https://reviews.llvm.org/D122014

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
@@ -19,6 +19,7 @@
 #include <mutex>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "Debug.h"
@@ -224,22 +225,32 @@
   size_t Next = 0;
   /// Mutex to guard the pool.
   std::mutex Mutex;
-  /// Pool of resources.
+  /// Pool of resources. The difference between \p Resources and \p Pool is,
+  /// when a resource is acquired and released, it is all on \p Resources. When
+  /// a batch of new resources are needed, they are both added to \p Resources
+  /// and \p Pool. The reason for this setting is, \p Resources could contain
+  /// redundant elements because resources are not released, which can cause
+  /// double free. This setting makes sure that \p Pool always has every
+  /// resource allocated from the device.
   std::vector<ElementTy> Resources;
+  std::vector<ElementTy> Pool;
   /// A reference to the corresponding allocator.
   AllocTy Allocator;
 
   /// If `Resources` is used up, we will fill in more resources. It assumes that
   /// the new size `Size` should be always larger than the current size.
   bool resize(size_t Size) {
+    assert(Resources.size() == Pool.size() && "size mismatch");
     auto CurSize = Resources.size();
     assert(Size > CurSize && "Unexpected smaller size");
+    Pool.reserve(Size);
     Resources.reserve(Size);
     for (auto I = CurSize; I < Size; ++I) {
       ElementTy NewItem;
       int Ret = Allocator.create(NewItem);
       if (Ret != OFFLOAD_SUCCESS)
         return false;
+      Pool.push_back(NewItem);
       Resources.push_back(NewItem);
     }
     return true;
@@ -300,8 +311,9 @@
   /// Released all stored resources and clear the pool.
   /// Note: This function is not thread safe. Be sure to guard it if necessary.
   void clear() noexcept {
-    for (auto &R : Resources)
+    for (auto &R : Pool)
       (void)Allocator.destroy(R);
+    Pool.clear();
     Resources.clear();
   }
 };


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D122014.416816.patch
Type: text/x-patch
Size: 2142 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20220321/0c5bc367/attachment.bin>