[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.

[JonChesterfield]

Is it possible to split this into a refactor and a behaviour change patch? Hard to tell what the semantic change is here

[pdhaliwal]

Split this patch into two,

1. Current patch now only includes changes for simplifying the host pool lookup
2. D109875 includes plain refactoring.

[JonChesterfield]

I think this patch includes the refactor so I'm going to hold off on trying to review it until it can be rebased on main after the refactor lands. It's also possible to make a diff relative to another one though I'm not sure on the proper workflow for doing so.

[pdhaliwal]

This is not dependent on D109875, instead it is the other case. Apologies for not mentioning here. Diff shown here is messed up.

[pdhaliwal]

Rebase

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;
       {