[Openmp-commits] [PATCH] D116340: [OpenMP][Plugin] Minor adjustments to ResourcePool

[Shilei Tian via Phabricator via Openmp-commits]

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

This patch makes some minor adjustments to `ResourcePool`:

- Don't initialize the resources if `Size` is 0 which can avoid assertion.
- Add a new interface function `clear` to release all hold resources.
- If initial size is 0, resize to 1 when the first request is encountered.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D116340

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
@@ -262,13 +262,11 @@
 public:
   ResourcePoolTy(AllocatorTy<T> &&A, size_t Size = 0) noexcept
       : Allocator(std::move(A)) {
-    (void)resize(Size);
+    if (Size)
+      (void)resize(Size);
   }
 
-  ~ResourcePoolTy() noexcept {
-    for (auto &R : Resources)
-      (void)Allocator.destroy(R);
-  }
+  ~ResourcePoolTy() noexcept { clear(); }
 
   /// Get a resource from pool. `Next` always points to the next available
   /// resource. That means, `[0, next-1]` have been assigned, and `[id,]` are
@@ -283,8 +281,13 @@
   ///       Next
   int acquire(T &R) noexcept {
     std::lock_guard<std::mutex> LG(Mutex);
-    if (Next == Resources.size() && !resize(Resources.size() * 2))
-      return OFFLOAD_FAIL;
+    if (Next == Resources.size()) {
+      auto NewSize = Resources.size() ? Resources.size() * 2 : 1;
+      if (!resize(NewSize))
+        return OFFLOAD_FAIL;
+    }
+
+    assert(Next < Resources.size());
 
     R = Resources[Next++];
 
@@ -307,6 +310,14 @@
     std::lock_guard<std::mutex> LG(Mutex);
     Resources[--Next] = R;
   }
+
+  /// 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)
+      (void)Allocator.destroy(R);
+    Resources.clear();
+  }
 };
 
 class DeviceRTLTy {
@@ -328,7 +339,6 @@
   static constexpr const int DefaultNumThreads = 128;
 
   using StreamPoolTy = ResourcePoolTy<CUstream>;
-  using StreamAllocatorTy = AllocatorTy<CUstream>;
   std::vector<std::unique_ptr<StreamPoolTy>> StreamPool;
 
   std::vector<DeviceDataTy> DeviceData;
@@ -563,7 +573,7 @@
         checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
 
     for (auto &S : StreamPool)
-      S = nullptr;
+      S.reset();
 
     for (DeviceDataTy &D : DeviceData) {
       // Destroy context
@@ -631,7 +641,8 @@
     // Initialize stream pool
     if (!StreamPool[DeviceId])
       StreamPool[DeviceId] = std::make_unique<StreamPoolTy>(
-          StreamAllocatorTy(DeviceData[DeviceId].Context), NumInitialStreams);
+          AllocatorTy<CUstream>(DeviceData[DeviceId].Context),
+          NumInitialStreams);
 
     // Query attributes to determine number of threads/block and blocks/grid.
     int MaxGridDimX;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D116340.396420.patch
Type: text/x-patch
Size: 2511 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20211228/ba6c2f08/attachment.bin>