[Openmp-commits] [openmp] 357c803 - [OpenMP][Plugin] Minor adjustments to ResourcePool

[Shilei Tian via Openmp-commits]

Author: Shilei Tian
Date: 2021-12-28T16:11:03-05:00
New Revision: 357c8031ff29299ad4b823074ff7fdda23654cc8

URL: https://github.com/llvm/llvm-project/commit/357c8031ff29299ad4b823074ff7fdda23654cc8
DIFF: https://github.com/llvm/llvm-project/commit/357c8031ff29299ad4b823074ff7fdda23654cc8.diff

LOG: [OpenMP][Plugin] Minor adjustments to ResourcePool

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.

Reviewed By: jdoerfert

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/plugins/cuda/src/rtl.cpp b/openmp/libomptarget/plugins/cuda/src/rtl.cpp
index ed26f2f7731f6..1afee7ce3a02f 100644
--- a/openmp/libomptarget/plugins/cuda/src/rtl.cpp
+++ b/openmp/libomptarget/plugins/cuda/src/rtl.cpp
@@ -262,13 +262,11 @@ template <typename T> class ResourcePoolTy {
 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 @@ template <typename T> class ResourcePoolTy {
   ///       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 @@ template <typename T> class ResourcePoolTy {
     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 @@ class DeviceRTLTy {
   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 @@ class DeviceRTLTy {
         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 @@ class DeviceRTLTy {
     // 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;