[AMDGPU][Libomptarget] Simplify host memory pool lookup

[pdhaliwal]

Summary

Using lambda for host memory pool was causing stack smashing
error on CI. Using plain for loop fixed the issue.

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

 //===----------------------------------------------------------------------===//
 
 #include <algorithm>
 #include <assert.h>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <functional>
 #include <libelf.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <unordered_map>
 #include <vector>
 
 #include "interop_hsa.h"
FindKernargPool(const std::vector<hsa_agent_t> &HSAAgents) {
   if (KernArgPools.empty()) {
     DP("Unable to find any valid kernarg pool\n");
     return {HSA_STATUS_ERROR, hsa_amd_memory_pool_t{}};
   }
 
   return {HSA_STATUS_SUCCESS, KernArgPools[0]};
 }
 
+hsa_status_t addMemoryPool(hsa_amd_memory_pool_t MemoryPool, void *Data) {
+  std::vector<hsa_amd_memory_pool_t> *Result =
+      static_cast<std::vector<hsa_amd_memory_pool_t> *>(Data);
+  bool IsValid = false;
+  hsa_status_t Err;
+  std::tie(Err, IsValid) = isValidMemoryPool(MemoryPool);
+  if (Err != HSA_STATUS_SUCCESS) {
+    DP("isValidMemoryPool failed: %s\n", get_error_string(Err));
+    return Err;
+  }
+
+  if (IsValid) {
+    Result->push_back(MemoryPool);
+  }
+
+  return HSA_STATUS_SUCCESS;
+}
+
 } // namespace
 } // namespace core
 
 struct EnvironmentVariables {
   int NumTeams;
   int TeamLimit;
   int TeamThreadLimit;
   int MaxTeamsDefault;
if (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED) {
       DeviceFineGrainedMemoryPools[DeviceId] = MemoryPool;
     } else if (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED) {
       DeviceCoarseGrainedMemoryPools[DeviceId] = MemoryPool;
     }
 
     return HSA_STATUS_SUCCESS;
   }
 
-  hsa_status_t addHostMemoryPool(hsa_amd_memory_pool_t MemoryPool,
-                                 int DeviceId) {
+  hsa_status_t setupMemoryPools() {
+    using namespace std::placeholders;
+    hsa_status_t Err;
+    Err = core::collectMemoryPools(
+        HSAAgents, std::bind(&RTLDeviceInfoTy::addDeviceMemoryPool, this, _1, _2));
+    if (Err != HSA_STATUS_SUCCESS) {
+      DP("HSA error in collecting memory pools for offload devices: %s\n",
+         get_error_string(Err));
+      return Err;
+    }
+    return HSA_STATUS_SUCCESS;
+  }
+
+  hsa_status_t setupHostMemoryPools(std::vector<hsa_agent_t> &Agents) {
+    std::vector<hsa_amd_memory_pool_t> HostPools;
+
+    // collect all the "valid" pools for all the given agents.
+    for (const auto &Agent : Agents) {
+      hsa_status_t Err = hsa_amd_agent_iterate_memory_pools(
+          Agent, core::addMemoryPool, static_cast<void *>(&HostPools));
+      if (Err != HSA_STATUS_SUCCESS) {
+        DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
+           "Iterate all memory pools", get_error_string(Err));
+        return Err;
+      }
+    }
+
+    for (const auto &MemoryPool : HostPools) {
+      hsa_status_t Err = HSA_STATUS_SUCCESS;
       uint32_t GlobalFlags = 0;
-    hsa_status_t Err = hsa_amd_memory_pool_get_info(
+      Err = hsa_amd_memory_pool_get_info(
           MemoryPool, HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &GlobalFlags);
-
       if (Err != HSA_STATUS_SUCCESS) {
+        DP("Get memory pool info failed: %s\n", get_error_string(Err));
         return Err;
       }
 
-    uint32_t Size;
+      size_t Size = 0;
       Err = hsa_amd_memory_pool_get_info(MemoryPool,
                                          HSA_AMD_MEMORY_POOL_INFO_SIZE, &Size);
       if (Err != HSA_STATUS_SUCCESS) {
+        DP("Get memory pool size failed: %s\n", get_error_string(Err));
         return Err;
       }
 
-    if (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED &&
+      if ((GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED) &&
           Size > 0) {
         HostFineGrainedMemoryPool = MemoryPool;
       }
-
-    return HSA_STATUS_SUCCESS;
     }
 
-  hsa_status_t setupMemoryPools() {
-    using namespace std::placeholders;
-    hsa_status_t Err;
-    Err = core::collectMemoryPools(
-        CPUAgents, std::bind(&RTLDeviceInfoTy::addHostMemoryPool, this, _1, _2));
-    if (Err != HSA_STATUS_SUCCESS) {
-      DP("HSA error in collecting memory pools for CPU: %s\n",
-         get_error_string(Err));
-      return Err;
-    }
-    Err = core::collectMemoryPools(
-        HSAAgents, std::bind(&RTLDeviceInfoTy::addDeviceMemoryPool, this, _1, _2));
-    if (Err != HSA_STATUS_SUCCESS) {
-      DP("HSA error in collecting memory pools for offload devices: %s\n",
-         get_error_string(Err));
-      return Err;
-    }
     return HSA_STATUS_SUCCESS;
   }
 
   hsa_amd_memory_pool_t getDeviceMemoryPool(int DeviceId) {
     assert(DeviceId >= 0 && DeviceId < DeviceCoarseGrainedMemoryPools.size() &&
            "Invalid device Id");
     return DeviceCoarseGrainedMemoryPools[DeviceId];
   }
RTLDeviceInfoTy() {
     DeviceCoarseGrainedMemoryPools.resize(NumberOfDevices);
     DeviceFineGrainedMemoryPools.resize(NumberOfDevices);
 
     err = setupMemoryPools();
     if (err != HSA_STATUS_SUCCESS) {
       DP("Error when setting up memory pools");
       return;
     }
+    err = setupHostMemoryPools(CPUAgents);
+    if (err != HSA_STATUS_SUCCESS) {
+      DP("Error when setting up host memory pools");
+      return;
+    }
 
     for (int i = 0; i < NumberOfDevices; i++) {
       HSAQueues[i] = nullptr;
     }
 
     for (int i = 0; i < NumberOfDevices; i++) {
       uint32_t queue_size = 0;
       {