[OpenMP][libomptarget] Retrieve multiple resources from resource managers

[kevinsala]

Summary

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.

Depends on D155621

[kevinsala]

Simplifying patch.

[jdoerfert]

LG

[tianshilei1992]

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

1149–1150

[jhuber6]

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

1149–1150

Should be /*Num=*/ since clang-format understands it.

[kevinsala]

Fixing comments.

[kevinsala]

Rebasing

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

public:
   /// memory buffer. The operation consists of a two-step copy from the
   /// device buffer to an intermediate pinned host buffer, and then, to a
   /// unpinned host buffer. Both operations are asynchronous and dependant.
   /// The intermediate pinned buffer will be released to the specified memory
   /// manager once the operation completes.
   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))
+    for (auto Signal : OutputSignals) {
       Signal->reset();
       Signal->increaseUseCount();
     }
 
     std::lock_guard<std::mutex> Lock(Mutex);
 
     // Consume stream slot and compute dependencies.
     auto [Curr, InputSignal] = consume(OutputSignals[0]);
public:
   /// the pinned host buffer. Both operations are asynchronous and dependant.
   /// The intermediate pinned buffer will be released to the specified memory
   /// manager once the operation completes.
   Error pushMemoryCopyH2DAsync(void *Dst, const void *Src, void *Inter,
                                uint64_t CopySize,
                                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))
+    for (auto Signal : OutputSignals) {
       Signal->reset();
       Signal->increaseUseCount();
     }
 
     AMDGPUSignalTy *OutputSignal = OutputSignals[0];
 
     std::lock_guard<std::mutex> Lock(Mutex);
 

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

Error deinit() {
     if (auto Err = ResourcePoolTy::resizeResourcePool(NextAvailable))
       return Err;
 
     ResourcePool.clear();
 
     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();
   }
 
   /// Return resource to the pool.
   Error returnResource(ResourceHandleTy Handle) {
     const std::lock_guard<std::mutex> Lock(Mutex);