[OpenMP][CUDA] Use one event pool per device

[jdoerfert]

Summary

An event pool, similar to the stream pool, needs to be kept per device.
For one, events are associated with cuda contexts which means we cannot
destroy the former after the latter. Also, CUDA documentation states
streams and events need to be associated with the same context, which
we did not ensure at all.

[tianshilei1992]

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

181–182

I don't think it's good to put CUcontext in allocator. That template is only a reference for other specialization.

[jdoerfert]

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

181–182

It was but it is not anymore. It is a base class now. Can you elaborate why you think it's not good to put a CUcontext in here? After all, all the associated allocators do allocate things that are bound to a/this context.

[tianshilei1992]

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

181–182

Well, it's not about correctness or necessity. It's just different styles. Like I said, it's just a reference/interface. All other allocators can implement in their own way, including some resources that don't require context. But of course you could argue we have not encountered any resource that doesn't require context. That's true and that's why I think it's just about styles. I'd like to keep things more flexible, even currently I don't see any use.
This is not a blocker. I'm fine if you think it's the good way.

[jdoerfert]

Actually make this compile

[tianshilei1992]

LG

[jdoerfert]

format + namespace

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

if (Err != CUDA_SUCCESS) {
     DP("Error when syncing event = " DPxMOD "\n", DPxPTR(Event));
     CUDA_ERR_STRING(Err);
     return OFFLOAD_FAIL;
   }
 
   return OFFLOAD_SUCCESS;
 }
 
+namespace {
+
 // Structure contains per-device data
 struct DeviceDataTy {
   /// List that contains all the kernels.
   std::list<KernelTy> KernelsList;
 
   std::list<FuncOrGblEntryTy> FuncGblEntries;
 
   CUcontext Context = nullptr;
   // Device properties
   int ThreadsPerBlock = 0;
   int BlocksPerGrid = 0;
   int WarpSize = 0;
   // OpenMP properties
   int NumTeams = 0;
   int NumThreads = 0;
 };
 
 /// Resource allocator where \p T is the resource type.
 /// Functions \p create and \p destroy return OFFLOAD_SUCCESS and OFFLOAD_FAIL
 /// accordingly. The implementation should not raise any exception.
-template <typename T> class AllocatorTy {
-public:
+template <typename T> struct AllocatorTy {
+  AllocatorTy(CUcontext C) noexcept : Context(C) {}
+  using ElementTy = T;
+
+  virtual ~AllocatorTy() {}
+
   /// Create a resource and assign to R.
-  int create(T &R) noexcept;
+  virtual int create(T &R) noexcept = 0;
   /// Destroy the resource.
-  int destroy(T) noexcept;
-};
+  virtual int destroy(T) noexcept = 0;
 
-/// Allocator for CUstream.
-template <> class AllocatorTy<CUstream> {
+protected:
   CUcontext Context;
+};
 
-public:
-  AllocatorTy(CUcontext C) noexcept : Context(C) {}
+/// Allocator for CUstream.
+struct StreamAllocatorTy final : public AllocatorTy<CUstream> {
+  StreamAllocatorTy(CUcontext C) noexcept : AllocatorTy<CUstream>(C) {}
 
   /// See AllocatorTy<T>::create.
-  int create(CUstream &Stream) noexcept {
+  int create(CUstream &Stream) noexcept override {
     if (!checkResult(cuCtxSetCurrent(Context),
                      "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
     if (!checkResult(cuStreamCreate(&Stream, CU_STREAM_NON_BLOCKING),
                      "Error returned from cuStreamCreate\n"))
       return OFFLOAD_FAIL;
 
     return OFFLOAD_SUCCESS;
   }
 
   /// See AllocatorTy<T>::destroy.
-  int destroy(CUstream Stream) noexcept {
+  int destroy(CUstream Stream) noexcept override {
     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;
   }
 };
 
 /// Allocator for CUevent.
-template <> class AllocatorTy<CUevent> {
-public:
+struct EventAllocatorTy final : public AllocatorTy<CUevent> {
+  EventAllocatorTy(CUcontext C) noexcept : AllocatorTy<CUevent>(C) {}
+
   /// See AllocatorTy<T>::create.
-  int create(CUevent &Event) noexcept {
+  int create(CUevent &Event) noexcept override {
     if (!checkResult(cuEventCreate(&Event, CU_EVENT_DEFAULT),
                      "Error returned from cuEventCreate\n"))
       return OFFLOAD_FAIL;
 
     return OFFLOAD_SUCCESS;
   }
 
   /// See AllocatorTy<T>::destroy.
-  int destroy(CUevent Event) noexcept {
+  int destroy(CUevent Event) noexcept override {
     if (!checkResult(cuEventDestroy(Event),
                      "Error returned from cuEventDestroy\n"))
       return OFFLOAD_FAIL;
 
     return OFFLOAD_SUCCESS;
   }
 };
 
 /// A generic pool of resources where \p T is the resource type.
 /// \p T should be copyable as the object is stored in \p std::vector .
-template <typename T> class ResourcePoolTy {
+template <typename AllocTy> class ResourcePoolTy {
+  using ElementTy = typename AllocTy::ElementTy;
   /// Index of the next available resource.
   size_t Next = 0;
   /// Mutex to guard the pool.
   std::mutex Mutex;
   /// Pool of resources.
-  std::vector<T> Resources;
+  std::vector<ElementTy> Resources;
   /// A reference to the corresponding allocator.
-  AllocatorTy<T> 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) {
     auto CurSize = Resources.size();
     assert(Size > CurSize && "Unexpected smaller size");
     Resources.reserve(Size);
     for (auto I = CurSize; I < Size; ++I) {
-      T NewItem;
+      ElementTy NewItem;
       int Ret = Allocator.create(NewItem);
       if (Ret != OFFLOAD_SUCCESS)
         return false;
       Resources.push_back(NewItem);
     }
     return true;
   }
 
 public:
-  ResourcePoolTy(AllocatorTy<T> &&A, size_t Size = 0) noexcept
+  ResourcePoolTy(AllocTy &&A, size_t Size = 0) noexcept
       : Allocator(std::move(A)) {
     if (Size)
       (void)resize(Size);
   }
 
   ~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
   /// 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(T &R) noexcept {
+  int acquire(ElementTy &R) noexcept {
     std::lock_guard<std::mutex> LG(Mutex);
     if (Next == Resources.size()) {
       auto NewSize = Resources.size() ? Resources.size() * 2 : 1;
       if (!resize(NewSize))
         return OFFLOAD_FAIL;
     }
 
     assert(Next < Resources.size());
public:
   /// 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(T R) noexcept {
+  void release(ElementTy R) noexcept {
     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();
   }
 };
 
+} // namespace
+
 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.
+
+  /// Number of initial streams for each device.
   int NumInitialStreams = 32;
 
+  /// Number of initial events for each device.
+  int NumInitialEvents = 8;
+
   static constexpr const int32_t HardThreadLimit = 1024;
   static constexpr const int32_t DefaultNumTeams = 128;
   static constexpr const int32_t DefaultNumThreads = 128;
 
-  using StreamPoolTy = ResourcePoolTy<CUstream>;
+  using StreamPoolTy = ResourcePoolTy<StreamAllocatorTy>;
   std::vector<std::unique_ptr<StreamPoolTy>> StreamPool;
 
-  ResourcePoolTy<CUevent> EventPool;
+  using EventPoolTy = ResourcePoolTy<EventAllocatorTy>;
+  std::vector<std::unique_ptr<EventPoolTy>> EventPool;
 
   std::vector<DeviceDataTy> DeviceData;
   std::vector<CUmodule> Modules;
 
   /// Vector of flags indicating the initalization status of all associated
   /// devices.
   std::vector<bool> InitializedFlags;
 
public:
 
   // This class should not be copied
   DeviceRTLTy(const DeviceRTLTy &) = delete;
   DeviceRTLTy(DeviceRTLTy &&) = delete;
 
   DeviceRTLTy()
       : NumberOfDevices(0), EnvNumTeams(-1), EnvTeamLimit(-1),
         EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED),
-        DynamicMemorySize(0), EventPool(AllocatorTy<CUevent>()) {
+        DynamicMemorySize(0) {
 
     DP("Start initializing CUDA\n");
 
     CUresult Err = cuInit(0);
     if (Err == CUDA_ERROR_INVALID_HANDLE) {
       // Can't call cuGetErrorString if dlsym failed
       DP("Failed to load CUDA shared library\n");
       return;
     }
     if (!checkResult(Err, "Error returned from cuInit\n")) {
       return;
     }
 
     Err = cuDeviceGetCount(&NumberOfDevices);
     if (!checkResult(Err, "Error returned from cuDeviceGetCount\n"))
       return;
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     DeviceData.resize(NumberOfDevices);
     StreamPool.resize(NumberOfDevices);
+    EventPool.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")) {
int initDevice(const int DeviceId) {
     Err = cuDevicePrimaryCtxRetain(&DeviceData[DeviceId].Context, Device);
     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
+    // Initialize the stream pool.
     if (!StreamPool[DeviceId])
       StreamPool[DeviceId] = std::make_unique<StreamPoolTy>(
-          AllocatorTy<CUstream>(DeviceData[DeviceId].Context),
-          NumInitialStreams);
+          StreamAllocatorTy(DeviceData[DeviceId].Context), NumInitialStreams);
+
+    // Initialize the event pool.
+    if (!EventPool[DeviceId])
+      EventPool[DeviceId] = std::make_unique<EventPoolTy>(
+          EventAllocatorTy(DeviceData[DeviceId].Context), NumInitialEvents);
 
     // 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);
int deinitDevice(const int DeviceId) {
     if (UseMemoryManager)
       MemoryManagers[DeviceId].release();
 
     // Close module
     if (CUmodule &M = Modules[DeviceId])
       checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
 
     StreamPool[DeviceId].reset();
-
-    // The event pool is shared, we initialize it once all devices have been
-    // deinitialized.
-    if (std::none_of(InitializedFlags.begin(), InitializedFlags.end(),
-                     [](bool IsInitialized) { return IsInitialized; }))
-      EventPool.clear();
+    EventPool[DeviceId].reset();
 
     // Destroy context
     DeviceDataTy &D = DeviceData[DeviceId];
     if (D.Context) {
       if (checkResult(cuCtxSetCurrent(D.Context),
                       "Error returned from cuCtxSetCurrent\n")) {
         CUdevice Device;
         if (checkResult(cuCtxGetDevice(&Device),
checkResult(
         "Error returned from cuDeviceGetAttribute\n");
     checkResult(
         cuDeviceGetAttribute(
             &TmpInt2, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, Device),
         "Error returned from cuDeviceGetAttribute\n");
     printf("    Compute Capabilities: \t\t%d%d \n", TmpInt, TmpInt2);
   }
 
-  int createEvent(void **P) {
+  int createEvent(int DeviceId, void **P) {
     CUevent Event = nullptr;
-    if (EventPool.acquire(Event) != OFFLOAD_SUCCESS)
+    if (EventPool[DeviceId]->acquire(Event) != OFFLOAD_SUCCESS)
       return OFFLOAD_FAIL;
     *P = Event;
     return OFFLOAD_SUCCESS;
   }
 
-  int destroyEvent(void *EventPtr) {
-    EventPool.release(reinterpret_cast<CUevent>(EventPtr));
+  int destroyEvent(int DeviceId, void *EventPtr) {
+    EventPool[DeviceId]->release(reinterpret_cast<CUevent>(EventPtr));
     return OFFLOAD_SUCCESS;
   }
 
   int waitEvent(const int DeviceId, __tgt_async_info *AsyncInfo,
                 void *EventPtr) const {
     CUstream Stream = getStream(DeviceId, AsyncInfo);
     CUevent Event = reinterpret_cast<CUevent>(EventPtr);
 

 
 void __tgt_rtl_print_device_info(int32_t device_id) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
   DeviceRTL.printDeviceInfo(device_id);
 }
 
 int32_t __tgt_rtl_create_event(int32_t device_id, void **event) {
   assert(event && "event is nullptr");
-  return DeviceRTL.createEvent(event);
+  return DeviceRTL.createEvent(device_id, event);
 }
 
 int32_t __tgt_rtl_record_event(int32_t device_id, void *event_ptr,
                                __tgt_async_info *async_info_ptr) {
   assert(async_info_ptr && "async_info_ptr is nullptr");
   assert(async_info_ptr->Queue && "async_info_ptr->Queue is nullptr");
   assert(event_ptr && "event_ptr is nullptr");
 
int32_t __tgt_rtl_sync_event(int32_t device_id, void *event_ptr) {
   assert(event_ptr && "event is nullptr");
 
   return syncEvent(event_ptr);
 }
 
 int32_t __tgt_rtl_destroy_event(int32_t device_id, void *event_ptr) {
   assert(event_ptr && "event is nullptr");
 
-  return DeviceRTL.destroyEvent(event_ptr);
+  return DeviceRTL.destroyEvent(device_id, event_ptr);
 }
 
 int32_t __tgt_rtl_release_async_info(int32_t device_id,
                                      __tgt_async_info *async_info) {
   assert(DeviceRTL.isValidDeviceId(device_id) && "device_id is invalid");
   assert(async_info && "async_info is nullptr");
 
   return DeviceRTL.releaseAsyncInfo(device_id, async_info);