[Openmp-commits] [openmp] 523ac0f - [OpenMP][libomptarget] Retrieve multiple resources from resource managers

[Kevin Sala via Openmp-commits]

Author: Kevin Sala
Date: 2023-07-28T00:37:08+02:00
New Revision: 523ac0fcdf7f37a25d646894fbed03d8611f66c2

URL: https://github.com/llvm/llvm-project/commit/523ac0fcdf7f37a25d646894fbed03d8611f66c2
DIFF: https://github.com/llvm/llvm-project/commit/523ac0fcdf7f37a25d646894fbed03d8611f66c2.diff

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

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

Added: 
    

Modified: 
    openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
    openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

Removed: 
    


################################################################################
diff  --git a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp b/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
index 905bc51bcd2142..403083f81a42a9 100644
--- a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/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();
     }

diff  --git a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
index 46f1e7863981ab..ff6d86e7725be6 100644
--- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h
+++ b/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();
   }