[OpenMP][Plugin] Introduce generic resource pool

[tianshilei1992]

Summary

Currently CUDA streams are managed by StreamManagerTy. It works very well. Now
we have the need that some resources, such as CUDA stream and event, will be
hold by libomptarget. It is always good to buffer those resources. What's more
important, given the way that libomptarget and plugins are connected, we cannot
make sure whether plugins are still alive when libomptarget is destroyed. That
leads to an issue that those resouces hold by libomptarget might not be
released correctly. As a result, we need an unified management of all the resources
that can be shared between libomptarget and plugins.

ResourcePoolTy is designed to manage the type of resource for one device.
It has to work with an allocator which is supposed to provide create and
destroy. In this way, when the plugin is destroyed, we can make sure that
all resources allocated from native runtime library will be released correctly,
no matter whether libomptarget starts its destroy.

[ye-luo]

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

196

Too many managers. How about ResourcePoolTy

277

allocate and deallocate are really meant for pointers. How about create/destroy for single object.

321–324

I think we can avoid StreamManagerTy abstraction and have
std::vector<std::unique_ptr<RMTy>> StreamPools;
directly here.

DeviceData, Modules, DeviceAllocators, MemoryManagers don't need additional layer.

[ye-luo]

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

264

You may consider acquire/release instead of put/get. I'm open to suggestions.

[ye-luo]

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

226

Whe CUDA call fails. Allocator.allocate() returns nullptr and the code continues running. This doesn't make sense.

277

Error handling and returning the object needs to be separate.

[tianshilei1992]

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

226

ResourceManagerTy doesn't know what is a valid return value. Ty can be pointer, object, etc. Of course, for CUDA, everything is pointer, but it doesn't sound right to have if (Ty == nullptr)`.

277

create and destroy sound good.

Error handling and returning the object needs to be separate.

I didn't follow.

[ye-luo]

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

226

Allocator.allocate(Ty&) should return OFFLOAD_SUCCESS, OFFLOAD_FAIL.
You have AllocatorTy which knows what Ty is. Indeed, Ty should be removed, it can be derived from AllocatorTy

277

I mean int32_t Allocator.create(Ty& res)
So the return is OFFLOAD_SUCCESS or OFFLOAD_FAIL.
The created resource is returned via res.

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

use create/destroy

[tianshilei1992]

reduce indirection

[tianshilei1992]

Free stream pool and allocator explicitly to avoid CUcontext is destroyed before them

[ye-luo]

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

192

What does it mean a resource is copyable? Did you mean a resource handle is copyable?

226

If Ty can be accessed from AllocatorTy, we'd better avoid passing Ty as a template argument.

[tianshilei1992]

fix comments

[tianshilei1992]

fix comments

[tianshilei1992]

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

226

I did it in another way: to avoid passing AllocatorTy as a template argument.

[ye-luo]

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

246

How about making this a member rather than a reference and remove StreamAllocator vector?
ResourcePoolTy(AllocatorTy<T> &&A, size_t Size = 0)

[tianshilei1992]

fix comments

[tianshilei1992]

fixed error caused by the order of destructions

[ye-luo]

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

557–558

Probabaly S.reset() is more C++ way.

[tianshilei1992]

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

557–558

Will make the change in another patch.

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

struct DeviceDataTy {
   int ThreadsPerBlock = 0;
   int BlocksPerGrid = 0;
   int WarpSize = 0;
   // OpenMP properties
   int NumTeams = 0;
   int NumThreads = 0;
 };
 
-class StreamManagerTy {
-  int NumberOfDevices;
-  // The initial size of stream pool
-  int EnvNumInitialStreams;
-  // Per-device stream mutex
-  std::vector<std::unique_ptr<std::mutex>> StreamMtx;
-  // Per-device stream Id indicates the next available stream in the pool
-  std::vector<int> NextStreamId;
-  // Per-device stream pool
-  std::vector<std::vector<CUstream>> StreamPool;
-  // Reference to per-device data
-  std::vector<DeviceDataTy> &DeviceData;
-
-  // If there is no CUstream left in the pool, we will resize the pool to
-  // allocate more CUstream. This function should be called with device mutex,
-  // and we do not resize to smaller one.
-  void resizeStreamPool(const int DeviceId, const size_t NewSize) {
-    std::vector<CUstream> &Pool = StreamPool[DeviceId];
-    const size_t CurrentSize = Pool.size();
-    assert(NewSize > CurrentSize && "new size is not larger than current size");
-
-    CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
-    if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n")) {
-      // We will return if cannot switch to the right context in case of
-      // creating bunch of streams that are not corresponding to the right
-      // device. The offloading will fail later because selected CUstream is
-      // nullptr.
-      return;
-    }
-
-    Pool.resize(NewSize, nullptr);
-
-    for (size_t I = CurrentSize; I < NewSize; ++I) {
-      checkResult(cuStreamCreate(&Pool[I], CU_STREAM_NON_BLOCKING),
-                  "Error returned from cuStreamCreate\n");
+/// A generic pool of resources.
+///
+/// `Ty` is the type of the resource. It should be copyable.
+/// `AllocatorTy` should provide two functions:
+/// - `int AllocatorTy::create(Ty &)`: create one resource.
+/// - `int Allocator::destroy(Ty)`: destroy the resource.
+/// Both functions return OFFLOAD_SUCCESS or OFFLOAD_FAIL accordingly.
+template <typename Ty, typename AllocatorTy> class ResourcePoolTy {
+  /// Index of the next available resource.
+  size_t Next = 0;
+  /// Mutex to guard the pool.
+  std::mutex Mutex;
+  /// Pool of resources.
+  std::vector<Ty> Resources;
+  /// A reference to the corresponding allocator.
+  AllocatorTy &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) {
+    auto CurSize = Resources.size();
+    assert(Size > CurSize && "Unexpected smaller size");
+    Resources.reserve(Size);
+    for (auto I = CurSize; I < Size; ++I) {
+      Ty NewItem;
+      int Ret = Allocator.create(NewItem);
+      if (Ret != OFFLOAD_SUCCESS)
+        return false;
+      Resources.push_back(NewItem);
     }
+    return true;
   }
 
 public:
-  StreamManagerTy(const int NumberOfDevices,
-                  std::vector<DeviceDataTy> &DeviceData)
-      : NumberOfDevices(NumberOfDevices), EnvNumInitialStreams(32),
-        DeviceData(DeviceData) {
-    StreamPool.resize(NumberOfDevices);
-    NextStreamId.resize(NumberOfDevices);
-    StreamMtx.resize(NumberOfDevices);
+  ResourcePoolTy(AllocatorTy &A, size_t Size = 0) noexcept : Allocator(A) {
+    (void)resize(Size);
+  }
 
-    if (const char *EnvStr = getenv("LIBOMPTARGET_NUM_INITIAL_STREAMS"))
-      EnvNumInitialStreams = std::stoi(EnvStr);
+  ~ResourcePoolTy() noexcept {
+    for (auto &R : Resources)
+      (void)Allocator.destroy(R);
+  }
+
+  /// Get a resource from pool. `Next` always points to the next available
+  /// resource. That means, `[0, next-1]` have been assigned, and `[id,]` are
+  /// still available. If there is no resource left, we will ask for more. Each
+  /// time a resource is assigned, the id will increase one.
+  /// xxxxxs+++++++++
+  ///      ^
+  ///      Next
+  /// After assignment, the pool becomes the following and s is assigned.
+  /// xxxxxs+++++++++
+  ///       ^
+  ///       Next
+  int acquire(Ty &R) noexcept {
+    std::lock_guard<std::mutex> LG(Mutex);
+    if (Next == Resources.size() && !resize(Resources.size() * 2))
+      return OFFLOAD_FAIL;
 
-    // Initialize the next stream id
-    std::fill(NextStreamId.begin(), NextStreamId.end(), 0);
+    R = Resources[Next++];
 
-    // Initialize stream mutex
-    for (std::unique_ptr<std::mutex> &Ptr : StreamMtx)
-      Ptr = std::make_unique<std::mutex>();
+    return OFFLOAD_SUCCESS;
   }
 
-  ~StreamManagerTy() {
-    // Destroy streams
-    for (int I = 0; I < NumberOfDevices; ++I) {
-      checkResult(cuCtxSetCurrent(DeviceData[I].Context),
-                  "Error returned from cuCtxSetCurrent\n");
-
-      for (CUstream &S : StreamPool[I]) {
-        if (S)
-          checkResult(cuStreamDestroy(S),
-                      "Error returned from cuStreamDestroy\n");
-      }
-    }
+  /// Return the resource back to the pool. When we return a resource, we need
+  /// to first decrease `Next`, and then copy the resource back. It is worth
+  /// noting that, the order of resources return might be different from that
+  /// they're assigned, that saying, at some point, there might be two identical
+  /// resources.
+  /// xxax+a+++++
+  ///     ^
+  ///     Next
+  /// However, it doesn't matter, because they're always on the two sides of
+  /// `Next`. The left one will in the end be overwritten by another resource.
+  /// Therefore, after several execution, the order of pool might be different
+  /// from its initial state.
+  void release(Ty R) noexcept {
+    std::lock_guard<std::mutex> LG(Mutex);
+    Resources[--Next] = R;
   }
+};
 
-  // Get a CUstream from pool. Per-device next stream id always points to the
-  // next available CUstream. That means, CUstreams [0, id-1] have been
-  // assigned, and [id,] are still available. If there is no CUstream left, we
-  // will ask more CUstreams from CUDA RT. Each time a CUstream is assigned,
-  // the id will increase one.
-  // xxxxxs+++++++++
-  //      ^
-  //      id
-  // After assignment, the pool becomes the following and s is assigned.
-  // xxxxxs+++++++++
-  //       ^
-  //       id
-  CUstream getStream(const int DeviceId) {
-    const std::lock_guard<std::mutex> Lock(*StreamMtx[DeviceId]);
-    int &Id = NextStreamId[DeviceId];
-    // No CUstream left in the pool, we need to request from CUDA RT
-    if (Id == static_cast<int>(StreamPool[DeviceId].size())) {
-      // By default we double the stream pool every time
-      resizeStreamPool(DeviceId, Id * 2);
-    }
-    return StreamPool[DeviceId][Id++];
-  }
-
-  // Return a CUstream back to pool. As mentioned above, per-device next
-  // stream is always points to the next available CUstream, so when we return
-  // a CUstream, we need to first decrease the id, and then copy the CUstream
-  // back.
-  // It is worth noting that, the order of streams return might be different
-  // from that they're assigned, that saying, at some point, there might be
-  // two identical CUstreams.
-  // xxax+a+++++
-  //     ^
-  //     id
-  // However, it doesn't matter, because they're always on the two sides of
-  // id. The left one will in the end be overwritten by another CUstream.
-  // Therefore, after several execution, the order of pool might be different
-  // from its initial state.
-  void returnStream(const int DeviceId, CUstream Stream) {
-    const std::lock_guard<std::mutex> Lock(*StreamMtx[DeviceId]);
-    int &Id = NextStreamId[DeviceId];
-    assert(Id > 0 && "Wrong stream ID");
-    StreamPool[DeviceId][--Id] = Stream;
-  }
+/// Allocator for CUstream.
+class StreamAllocatorTy {
+  CUcontext Context;
 
-  bool initializeDeviceStreamPool(const int DeviceId) {
-    assert(StreamPool[DeviceId].empty() && "stream pool has been initialized");
+public:
+  StreamAllocatorTy(CUcontext C) noexcept : Context(C) {}
 
-    resizeStreamPool(DeviceId, EnvNumInitialStreams);
+  int create(CUstream &Stream) noexcept {
+    if (!checkResult(cuCtxSetCurrent(Context),
+                     "Error returned from cuCtxSetCurrent\n"))
+      return OFFLOAD_FAIL;
 
-    // Check the size of stream pool
-    if (static_cast<int>(StreamPool[DeviceId].size()) != EnvNumInitialStreams)
-      return false;
+    if (!checkResult(cuStreamCreate(&Stream, CU_STREAM_NON_BLOCKING),
+                     "Error returned from cuStreamCreate\n"))
+      return OFFLOAD_FAIL;
 
-    // Check whether each stream is valid
-    for (CUstream &S : StreamPool[DeviceId])
-      if (!S)
-        return false;
+    return OFFLOAD_SUCCESS;
+  }
 
-    return true;
+  int destroy(CUstream Stream) noexcept {
+    if (!checkResult(cuCtxSetCurrent(Context),
+                     "Error returned from cuCtxSetCurrent\n"))
+      return OFFLOAD_FAIL;
+    if (!checkResult(cuStreamDestroy(Stream),
+                     "Error returned from cuStreamDestroy\n"))
+      return OFFLOAD_FAIL;
+
+    return OFFLOAD_SUCCESS;
   }
 };
 
 class DeviceRTLTy {
   int NumberOfDevices;
   // OpenMP environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
   int EnvTeamThreadLimit;
   // OpenMP requires flags
   int64_t RequiresFlags;
   // Amount of dynamic shared memory to use at launch.
   uint64_t DynamicMemorySize;
+  // Number of initial streams for each device.
+  int NumInitialStreams = 32;
 
   static constexpr const int HardTeamLimit = 1U << 16U; // 64k
   static constexpr const int HardThreadLimit = 1024;
   static constexpr const int DefaultNumTeams = 128;
   static constexpr const int DefaultNumThreads = 128;
 
-  std::unique_ptr<StreamManagerTy> StreamManager;
+  using StreamPoolTy = ResourcePoolTy<CUstream, StreamAllocatorTy>;
+  std::vector<std::unique_ptr<StreamPoolTy>> StreamPool;
+  std::vector<std::unique_ptr<StreamAllocatorTy>> StreamAllocator;
+
   std::vector<DeviceDataTy> DeviceData;
   std::vector<CUmodule> Modules;
 
   /// A class responsible for interacting with device native runtime library to
   /// allocate and free memory.
   class CUDADeviceAllocatorTy : public DeviceAllocatorTy {
     const int DeviceId;
     const std::vector<DeviceDataTy> &DeviceData;
void clearOffloadEntriesTable(const int DeviceId) {
     FuncOrGblEntryTy &E = DeviceData[DeviceId].FuncGblEntries.back();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = nullptr;
   }
 
   CUstream getStream(const int DeviceId, __tgt_async_info *AsyncInfo) const {
     assert(AsyncInfo && "AsyncInfo is nullptr");
 
-    if (!AsyncInfo->Queue)
-      AsyncInfo->Queue = StreamManager->getStream(DeviceId);
+    if (!AsyncInfo->Queue) {
+      CUstream S;
+      if (StreamPool[DeviceId]->acquire(S) != OFFLOAD_SUCCESS)
+        return nullptr;
+
+      AsyncInfo->Queue = S;
+    }
 
     return reinterpret_cast<CUstream>(AsyncInfo->Queue);
   }
 
 public:
   // This class should not be copied
   DeviceRTLTy(const DeviceRTLTy &) = delete;
   DeviceRTLTy(DeviceRTLTy &&) = delete;
if (!checkResult(Err, "Error returned from cuDeviceGetCount\n"))
       return;
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     DeviceData.resize(NumberOfDevices);
+    StreamAllocator.resize(NumberOfDevices);
+    StreamPool.resize(NumberOfDevices);
 
     // Get environment variables regarding teams
     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     }
     if (const char *EnvStr = getenv("OMP_TEAMS_THREAD_LIMIT")) {
       // OMP_TEAMS_THREAD_LIMIT has been set
       EnvTeamThreadLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAMS_THREAD_LIMIT=%d\n", EnvTeamThreadLimit);
     }
     if (const char *EnvStr = getenv("OMP_NUM_TEAMS")) {
       // OMP_NUM_TEAMS has been set
       EnvNumTeams = std::stoi(EnvStr);
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     }
     if (const char *EnvStr = getenv("LIBOMPTARGET_SHARED_MEMORY_SIZE")) {
       // LIBOMPTARGET_SHARED_MEMORY_SIZE has been set
       DynamicMemorySize = std::stoi(EnvStr);
       DP("Parsed LIBOMPTARGET_SHARED_MEMORY_SIZE = %" PRIu64 "\n",
          DynamicMemorySize);
     }
-
-    StreamManager =
-        std::make_unique<StreamManagerTy>(NumberOfDevices, DeviceData);
+    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(I, DeviceData);
 
     // 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)
         MemoryManagers.emplace_back(std::make_unique<MemoryManagerTy>(
             DeviceAllocators[I], MemoryManagerThreshold));
   }
 
   ~DeviceRTLTy() {
     // We first destruct memory managers in case that its dependent data are
     // destroyed before it.
     for (auto &M : MemoryManagers)
       M.release();
 
-    StreamManager = nullptr;
-
     for (CUmodule &M : Modules)
       // Close module
       if (M)
         checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
 
     for (DeviceDataTy &D : DeviceData) {
       // Destroy context
       if (D.Context) {
         checkResult(cuCtxSetCurrent(D.Context),
                     "Error returned from cuCtxSetCurrent\n");
         CUdevice Device;
         checkResult(cuCtxGetDevice(&Device),
                     "Error returned from cuCtxGetDevice\n");
         checkResult(cuDevicePrimaryCtxRelease(Device),
                     "Error returned from cuDevicePrimaryCtxRelease\n");
int initDevice(const int DeviceId) {
     if (!checkResult(Err, "Error returned from cuDevicePrimaryCtxRetain\n"))
       return OFFLOAD_FAIL;
 
     Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
     // Initialize stream pool
-    if (!StreamManager->initializeDeviceStreamPool(DeviceId))
-      return OFFLOAD_FAIL;
+    if (!StreamPool[DeviceId]) {
+      StreamAllocator[DeviceId] =
+          std::make_unique<StreamAllocatorTy>(DeviceData[DeviceId].Context);
+      StreamPool[DeviceId] = std::make_unique<StreamPoolTy>(
+          *StreamAllocator[DeviceId], NumInitialStreams);
+    }
 
     // Query attributes to determine number of threads/block and blocks/grid.
     int MaxGridDimX;
     Err = cuDeviceGetAttribute(&MaxGridDimX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
                                Device);
     if (Err != CUDA_SUCCESS) {
       DP("Error getting max grid dimension, use default value %d\n",
          DeviceRTLTy::DefaultNumTeams);
public:
 
   int synchronize(const int DeviceId, __tgt_async_info *AsyncInfo) const {
     CUstream Stream = reinterpret_cast<CUstream>(AsyncInfo->Queue);
     CUresult Err = cuStreamSynchronize(Stream);
 
     // Once the stream is synchronized, return it to stream pool and reset
     // AsyncInfo. This is to make sure the synchronization only works for its
     // own tasks.
-    StreamManager->returnStream(DeviceId,
-                                reinterpret_cast<CUstream>(AsyncInfo->Queue));
+    StreamPool[DeviceId]->release(reinterpret_cast<CUstream>(AsyncInfo->Queue));
     AsyncInfo->Queue = nullptr;
 
     if (Err != CUDA_SUCCESS) {
       DP("Error when synchronizing stream. stream = " DPxMOD
          ", async info ptr = " DPxMOD "\n",
          DPxPTR(Stream), DPxPTR(AsyncInfo));
       CUDA_ERR_STRING(Err);
     }