[AMDGPU][Libomptarget] Refactor device/host memory pools setup

[pdhaliwal]

Summary

• Simplifies logic for device memory pool lookup
• Merges kernarg host pool lookup into setupHostMemoryPools
• Removes redundant code.

[pdhaliwal]

run clang-format

[JonChesterfield]

We've lost some error checking in the refactor but otherwise it looks like an improvement

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

651

This is going to return success if it finds a pool for either kernarg or fine grain. I think we should only be claiming success if we've found (at least?) one for each. We could handle the control flow for that by initialising both pools to {0} and checking neither is {0} before returning, or by keeping track of whether both have been initialized

[pdhaliwal]

Check if both values are set properly.

[pdhaliwal]

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

651

Thanks for pointing this out. It actually found an issue when only one memory pool is available and it is both kernarg and fine grained.

[JonChesterfield]

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

651

That's interesting. If we have only one pool for both then we must use it for both.

What if we have one pool for both and one for only fine grain? Presumably we should then use different ones as they may be backed by different allocators.

Iteration order is probably difficult to predict. Maybe first kernarg pool found is used, and first non-kernarg pool is used for fine grain, but before erroring out we also use the kernarg pool for fine grain iff it hasn't been assigned yet?

[pdhaliwal]

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

651

That makes sense. Using "if-else" instead of "only if" can make sure that in each iteration only one of the kernarg and fine grained pool is assigned. After the for loop, kernarg should be always assigned (I guess?). If FineGrainedMemoryPool is not set then kernarg pool can be reused. Issues/Suggestions for this approach?

[JonChesterfield]

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

651

Sounds good to me

[pdhaliwal]

update logic for fine-grained/kernarg pools

[JonChesterfield]

LG, thanks!

[JonChesterfield]

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

18

functional may now be unused (think it was only here for placeholders)

[pdhaliwal]

Abandoning this in favor of D110902 due to dependency issues.

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>
 

 
 uint16_t create_header() {
   uint16_t header = HSA_PACKET_TYPE_KERNEL_DISPATCH << HSA_PACKET_HEADER_TYPE;
   header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_ACQUIRE_FENCE_SCOPE;
   header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_RELEASE_FENCE_SCOPE;
   return header;
 }
 
-hsa_status_t addKernArgPool(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 AllocAllowed = false;
-  hsa_status_t err = hsa_amd_memory_pool_get_info(
-      MemoryPool, HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED,
-      &AllocAllowed);
-  if (err != HSA_STATUS_SUCCESS) {
-    DP("Alloc allowed in memory pool check failed: %s\n",
-       get_error_string(err));
-    return err;
-  }
-
-  if (!AllocAllowed) {
-    // nothing needs to be done here.
-    return HSA_STATUS_SUCCESS;
-  }
-
-  uint32_t GlobalFlags = 0;
-  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;
-  }
-
-  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) &&
-      (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_KERNARG_INIT) &&
-      size > 0) {
-    Result->push_back(MemoryPool);
-  }
-
-  return HSA_STATUS_SUCCESS;
-}
-
 std::pair<hsa_status_t, bool>
 isValidMemoryPool(hsa_amd_memory_pool_t MemoryPool) {
   bool AllocAllowed = false;
   hsa_status_t Err = hsa_amd_memory_pool_get_info(
       MemoryPool, HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED,
       &AllocAllowed);
   if (Err != HSA_STATUS_SUCCESS) {
     DP("Alloc allowed in memory pool check failed: %s\n",
        get_error_string(Err));
     return {Err, false};
   }
 
   return {HSA_STATUS_SUCCESS, AllocAllowed};
 }
 
-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;
-          bool Valid = false;
-          std::tie(Err, Valid) = isValidMemoryPool(MemoryPool);
-          if (Err != HSA_STATUS_SUCCESS) {
-            return Err;
-          }
-          if (Valid)
-            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) {
-      DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
-         "Iterate all memory pools", get_error_string(err));
-      return {err, hsa_amd_memory_pool_t{}};
-    }
-  }
-
-  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));
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 setupMemoryPools() {
-    using namespace std::placeholders;
-    hsa_status_t Err;
-    Err = core::collectMemoryPools(
-        HSAAgents, std::bind(&RTLDeviceInfoTy::addDeviceMemoryPool, this, _1, _2));
+  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) {
+            bool AllocAllowed = false;
+            hsa_status_t ErrGetInfo;
+            std::tie(ErrGetInfo, AllocAllowed) =
+                core::isValidMemoryPool(MemoryPool);
+            if (ErrGetInfo != HSA_STATUS_SUCCESS) {
+              DP("Alloc allowed in memory pool check failed: %s\n",
+                 get_error_string(ErrGetInfo));
+              return ErrGetInfo;
+            }
+            if (AllocAllowed) {
+              return addDeviceMemoryPool(MemoryPool, DeviceId);
+            }
+
+            return HSA_STATUS_SUCCESS;
+          });
+
       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_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;
       }
     }
 
+    // We need two fine-grained pools.
+    //  1. One with kernarg flag set for storing kernel arguments
+    //  2. Second for host allocations
+    bool FineGrainedMemoryPoolSet = false;
+    bool KernArgPoolSet = false;
     for (const auto &MemoryPool : HostPools) {
       hsa_status_t Err = HSA_STATUS_SUCCESS;
       uint32_t GlobalFlags = 0;
       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;
       }
 
       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) &&
           Size > 0) {
+        if (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_KERNARG_INIT) {
+          KernArgPool = MemoryPool;
+          KernArgPoolSet = true;
+        } else {
           HostFineGrainedMemoryPool = MemoryPool;
+          FineGrainedMemoryPoolSet = true;
         }
       }
+    }
+
+    if (!KernArgPoolSet) {
+      DP("No fine-grained kernarg pool found\n");
+      return HSA_STATUS_ERROR;
+    }
+
+    if (!FineGrainedMemoryPoolSet) {
+      // only other option is to share the kernarg pool.
+      FineGrainedMemoryPool = KernArgPool;
+    }
 
     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() {
     NumberOfDevices = (int)HSAAgents.size();
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting HSA.\n");
       return;
     } else {
       DP("There are %d devices supporting HSA.\n", NumberOfDevices);
     }
-    std::tie(err, KernArgPool) = core::FindKernargPool(CPUAgents);
-    if (err != HSA_STATUS_SUCCESS) {
-      DP("Error when reading memory pools\n");
-      return;
-    }
 
     // Init the device info
     HSAQueues.resize(NumberOfDevices);
     FuncGblEntries.resize(NumberOfDevices);
     ThreadsPerGroup.resize(NumberOfDevices);
     ComputeUnits.resize(NumberOfDevices);
     GPUName.resize(NumberOfDevices);
     GroupsPerDevice.resize(NumberOfDevices);
     WarpSize.resize(NumberOfDevices);
     NumTeams.resize(NumberOfDevices);
     NumThreads.resize(NumberOfDevices);
     deviceStateStore.resize(NumberOfDevices);
     KernelInfoTable.resize(NumberOfDevices);
     SymbolInfoTable.resize(NumberOfDevices);
     DeviceCoarseGrainedMemoryPools.resize(NumberOfDevices);
     DeviceFineGrainedMemoryPools.resize(NumberOfDevices);
 
-    err = setupMemoryPools();
+    err = setupDevicePools(HSAAgents);
     if (err != HSA_STATUS_SUCCESS) {
-      DP("Error when setting up memory pools");
+      DP("Setup for Device Memory Pools failed\n");
       return;
     }
+
     err = setupHostMemoryPools(CPUAgents);
     if (err != HSA_STATUS_SUCCESS) {
-      DP("Error when setting up host memory pools");
+      DP("Setup for Host Memory Pools failed\n");
       return;
     }
 
     for (int i = 0; i < NumberOfDevices; i++) {
       HSAQueues[i] = nullptr;
     }
 
     for (int i = 0; i < NumberOfDevices; i++) {