[OpenMP][libomptarget] Retrieve multiple resources from resource mana…

…gers

This patch extends the plugin resource managers to return more than one resource
per call. The return function is not extended since we do not return more than
one resource anywhere.

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

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -1146,13 +1146,11 @@ struct AMDGPUStreamTy {
   Error pushMemoryCopyD2HAsync(void *Dst, const void *Src, void *Inter,
                                uint64_t CopySize,
                                AMDGPUMemoryManagerTy &MemoryManager) {
-    // TODO: Managers should define a function to retrieve multiple resources
-    // in a single call.
     // Retrieve available signals for the operation's outputs.
     AMDGPUSignalTy *OutputSignals[2] = {};
-    for (auto &Signal : OutputSignals) {
-      if (auto Err = SignalManager.getResource(Signal))
-        return Err;
+    if (auto Err = SignalManager.getResources(/*Num=*/2, OutputSignals))
+      return Err;
+    for (auto Signal : OutputSignals) {
       Signal->reset();
       Signal->increaseUseCount();
     }
@@ -1218,9 +1216,9 @@ struct AMDGPUStreamTy {
                                AMDGPUMemoryManagerTy &MemoryManager) {
     // Retrieve available signals for the operation's outputs.
     AMDGPUSignalTy *OutputSignals[2] = {};
-    for (auto &Signal : OutputSignals) {
-      if (auto Err = SignalManager.getResource(Signal))
-        return Err;
+    if (auto Err = SignalManager.getResources(/*Num=*/2, OutputSignals))
+      return Err;
+    for (auto Signal : OutputSignals) {
       Signal->reset();
       Signal->increaseUseCount();
     }

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

@@ -1182,21 +1182,31 @@ template <typename ResourceRef> class GenericDeviceResourceManagerTy {
     return Plugin::success();
   }
 
-  /// Get resource from the pool or create new resources. If the function
-  /// succeeeds, the handle to the resource is saved in \p Handle.
+  /// Get a resource from the pool or create new ones. If the function succeeds,
+  /// the handle to the resource is saved in \p Handle.
   Error getResource(ResourceHandleTy &Handle) {
+    return getResources(1, &Handle);
+  }
+
+  /// Get multiple resources from the pool or create new ones. If the function
+  /// succeeeds, the handles to the resources are saved in \p Handles.
+  Error getResources(uint32_t Num, ResourceHandleTy *Handles) {
     const std::lock_guard<std::mutex> Lock(Mutex);
 
     assert(NextAvailable <= ResourcePool.size() &&
            "Resource pool is corrupted");
 
-    if (NextAvailable == ResourcePool.size())
-      // By default we double the resource pool every time.
-      if (auto Err = ResourcePoolTy::resizeResourcePool(NextAvailable * 2))
+    if (NextAvailable + Num > ResourcePool.size())
+      // Double the resource pool or resize it to provide the requested ones.
+      if (auto Err = ResourcePoolTy::resizeResourcePool(
+              std::max(NextAvailable * 2, NextAvailable + Num)))
         return Err;
 
-    // Save the handle in the output parameter.
-    Handle = ResourcePool[NextAvailable++];
+    // Save the handles in the output array parameter.
+    for (uint32_t r = 0; r < Num; ++r)
+      Handles[r] = ResourcePool[NextAvailable + r];
+
+    NextAvailable += Num;
 
     return Plugin::success();
   }