diff --git a/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp b/openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp index 905bc51bcd214..403083f81a42a 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 46f1e7863981a..ff6d86e7725be 100644 --- a/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h +++ b/openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h @@ -1182,21 +1182,31 @@ template 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 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(); }