[AMDGPU][Libomptarget] Refactor device and host pools setup into separate methods

[pdhaliwal]

Summary

Mostly NFC, this patch simplifies memory pools setup.

[pdhaliwal]

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

633

This check was already done in isValidMemoryPool

[JonChesterfield]

Is the only difference between the two blocks calling addDeviceMemoryPool vs addHostMemoryPool? If so passing that in as a parameter seems better than copy&paste, which seems to already exist as collectMemoryPools

Calling once on the device agents and once on the host agents seems an improvement, but doing so by copy/paste/specialise collectMemoryPools doesn't seem to be necessary for that.

err = setupDevicePools(HSAAgents);

would become
err = collectMemoryPools(addDeviceMemoryPool, HSAAgents)
or similar

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

624

This doesn't look right. I think iteration will stop when the passed function returns success. There's a 'continue' in the enum that we could return if you want to keep looking after isValid fails, which I think is necessary to skip over zero size pools

644

same iteration here

[JonChesterfield]

Is this patch still current? Looking at current main I think it's been implemented in 05ba9ff6a6d243a

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

hsa_status_t addKernArgPool(hsa_amd_memory_pool_t MemoryPool, void *Data) {
   if ((GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED) &&
       (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_KERNARG_INIT)) {
     Result->push_back(MemoryPool);
   }
 
   return HSA_STATUS_SUCCESS;
 }
 
-template <typename AccumulatorFunc>
-hsa_status_t collectMemoryPools(const std::vector<hsa_agent_t> &Agents,
-                                AccumulatorFunc Func) {
-  for (int DeviceId = 0; DeviceId < Agents.size(); DeviceId++) {
-    hsa_status_t Err = hsa::amd_agent_iterate_memory_pools(
-        Agents[DeviceId], [&](hsa_amd_memory_pool_t MemoryPool) {
-          hsa_status_t Err;
-          if ((Err = isValidMemoryPool(MemoryPool)) != HSA_STATUS_SUCCESS) {
-            DP("Skipping memory pool: %s\n", get_error_string(Err));
-          } else
-            Func(MemoryPool, DeviceId);
-          return HSA_STATUS_SUCCESS;
-        });
-
-    if (Err != HSA_STATUS_SUCCESS) {
-      DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
-         "Iterate all memory pools", get_error_string(Err));
-      return Err;
-    }
-  }
-
-  return HSA_STATUS_SUCCESS;
-}
-
 std::pair<hsa_status_t, hsa_amd_memory_pool_t>
 FindKernargPool(const std::vector<hsa_agent_t> &HSAAgents) {
   std::vector<hsa_amd_memory_pool_t> KernArgPools;
   for (const auto &Agent : HSAAgents) {
     hsa_status_t err = HSA_STATUS_SUCCESS;
     err = hsa_amd_agent_iterate_memory_pools(
         Agent, addKernArgPool, static_cast<void *>(&KernArgPools));
     if (err != HSA_STATUS_SUCCESS) {
hsa_status_t addHostMemoryPool(hsa_amd_memory_pool_t MemoryPool,
     uint32_t GlobalFlags = 0;
     hsa_status_t Err = hsa_amd_memory_pool_get_info(
         MemoryPool, HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &GlobalFlags);
 
     if (Err != HSA_STATUS_SUCCESS) {
       return Err;
     }
 
-    uint32_t Size;
-    Err = hsa_amd_memory_pool_get_info(MemoryPool,
-                                       HSA_AMD_MEMORY_POOL_INFO_SIZE, &Size);
-    if (Err != HSA_STATUS_SUCCESS) {
-      return Err;
+    if (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED) {
+      HostFineGrainedMemoryPool = MemoryPool;
     }
 
-    if (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED &&
-        Size > 0) {
-      HostFineGrainedMemoryPool = MemoryPool;
-    }
+    return HSA_STATUS_SUCCESS;
 
+  hsa_status_t setupDevicePools(const std::vector<hsa_agent_t> &Agents) {
+    for (int DeviceId = 0; DeviceId < Agents.size(); DeviceId++) {
+      hsa_status_t Err = hsa::amd_agent_iterate_memory_pools(
+          Agents[DeviceId], [&](hsa_amd_memory_pool_t MemoryPool) {
+            hsa_status_t ValidStatus = core::isValidMemoryPool(MemoryPool);
+            if (ValidStatus != HSA_STATUS_SUCCESS) {
+              DP("Alloc allowed in memory pool check failed: %s\n",
+                 get_error_string(ValidStatus));
               return HSA_STATUS_SUCCESS;
             }
+            return addDeviceMemoryPool(MemoryPool, DeviceId);
+          });
 
-  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));
+        DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
+           "Iterate all memory pools", get_error_string(Err));
         return Err;
       }
-    Err = core::collectMemoryPools(
-        HSAAgents, std::bind(&RTLDeviceInfoTy::addDeviceMemoryPool, this, _1, _2));
+    }
+    return HSA_STATUS_SUCCESS;
+  }
+
+  hsa_status_t setupHostPools(const std::vector<hsa_agent_t> &Agents) {
+    for (int DeviceId = 0; DeviceId < Agents.size(); DeviceId++) {
+      hsa_status_t Err = hsa::amd_agent_iterate_memory_pools(Agents[DeviceId], [&](hsa_amd_memory_pool_t MemoryPool) {
+        hsa_status_t Err = core::isValidMemoryPool(MemoryPool);
+        if (Err != HSA_STATUS_SUCCESS) {
+          DP("isValidMemoryPool returned: %s\n", get_error_string(Err));
+          return HSA_STATUS_SUCCESS;
+        }
+        return addHostMemoryPool(MemoryPool, DeviceId);
+      });
       if (Err != HSA_STATUS_SUCCESS) {
-      DP("HSA error in collecting memory pools for offload devices: %s\n",
-         get_error_string(Err));
+        DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
+           "Iterate all memory pools", 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() {
     NumTeams.resize(NumberOfDevices);
     NumThreads.resize(NumberOfDevices);
     deviceStateStore.resize(NumberOfDevices);
     KernelInfoTable.resize(NumberOfDevices);
     SymbolInfoTable.resize(NumberOfDevices);
     DeviceCoarseGrainedMemoryPools.resize(NumberOfDevices);
     DeviceFineGrainedMemoryPools.resize(NumberOfDevices);
 
-    err = setupMemoryPools();
+    err = setupHostPools(CPUAgents);
+    if (err != HSA_STATUS_SUCCESS) {
+      DP("Error when setting up host memory pools");
+      return;
+    }
+
+    err = setupDevicePools(HSAAgents);
     if (err != HSA_STATUS_SUCCESS) {
-      DP("Error when setting up memory pools");
+      DP("Error when setting up device memory pools");
       return;
     }
 
     for (int i = 0; i < NumberOfDevices; i++) {
       uint32_t queue_size = 0;
       {
         hsa_status_t err = hsa_agent_get_info(
             HSAAgents[i], HSA_AGENT_INFO_QUEUE_MAX_SIZE, &queue_size);