[AMDGPU][Libomptarget] Collect allocatable memory pools using HSA

[pdhaliwal]

Summary

The logic is almost similar to that of system.cpp with one change that
instead of adding all the memory pools to a device struct it only
keeps a single pool. The existing approach also always allocated memory on
the first HSA pool found for a GPU.

This depends on D104691. The goal of this series of patches is to remove
_atl_machine global. The next patch will drop g_atl_machine entirely.

[JonChesterfield]

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

371

This looks like it can return a hsa_status_t directly

hsa_status_t Err = get_info();
if (Err != HSA_SUCCESS) return Err;
if (!AllocAllowed) return HSA_ERROR; // or a more specific enum
return HSA_SUCCESS;

382

control flow here is simpler if the bool Valid is dropped by isValidMemoryPool returning a hsa_status_t

477

Is this per-cpu-agent, or per-cpu-agent-with-nonzero-memory-pool? I wonder if we'd be better off with a scalar hsa_amd_memory_pool_t HostFineGrainedMemoryPool; for now

721

Ah, it includes the agents with zero available memory. Not sure distinguishing between CPU agents is helpful at present.

Also, doesn't this plus the (size > 0) check mean that a threadripper would have a vector of length 4 of the host pools, but only index 0 and 1 would be populated, and not correspond to the iteration order of CPU agents?

[pdhaliwal]

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

371

I was thinking that it would be a good idea to stop early if there any kind of error in hsa API.

477

It is per-cpu-agent. I haven't added any zero size check following the existing logic. We always use only one memory pool so using single scalar should not cause any issues.

721

There is no size >0 check so it would populate all the slots.

[pdhaliwal]

Fix clang-tidy errors and use scalar for host memory pool.

[JonChesterfield]

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

371

that's probably fair, OK

594

Can probably return the error without printing here, the caller will know it was addDeviceMemoryPool that failed anyway. Medium term it's probably worth consolidating all the misc print statements into rtl.cpp and routing them through the DP() macro

637

This looks like it should be returning a hsa_status_t, the caller may want to disable offloading if the memory pool query fails

[pdhaliwal]

Addressed review comments.

[JonChesterfield]

Thanks!

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

 
 #include <algorithm>
 #include <assert.h>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <elf.h>
 #include <fstream>
+#include <functional>
 #include <iostream>
 #include <libelf.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <thread>
 #include <unordered_map>

 template <typename C> hsa_status_t iterate_agents(C cb) {
   auto L = [](hsa_agent_t agent, void *data) -> hsa_status_t {
     C *unwrapped = static_cast<C *>(data);
     return (*unwrapped)(agent);
   };
   return hsa_iterate_agents(L, static_cast<void *>(&cb));
 }
 
+template <typename C>
+hsa_status_t amd_agent_iterate_memory_pools(hsa_agent_t Agent, C cb) {
+  auto L = [](hsa_amd_memory_pool_t MemoryPool, void *data) -> hsa_status_t {
+    C *unwrapped = static_cast<C *>(data);
+    return (*unwrapped)(MemoryPool);
+  };
+
+  return hsa_amd_agent_iterate_memory_pools(Agent, L, static_cast<void *>(&cb));
+}
+
 } // namespace hsa
 
 /// Keep entries table per device
 struct FuncOrGblEntryTy {
   __tgt_target_table Table;
   std::vector<__tgt_offload_entry> Entries;
 };
 
hsa_status_t addKernArgPool(hsa_amd_memory_pool_t MemoryPool, void *Data) {
   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) {
     fprintf(stderr, "Get memory pool info 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_t size = 0;
-    err = hsa_amd_memory_pool_get_info(MemoryPool,
-                                       HSA_AMD_MEMORY_POOL_INFO_SIZE, &size);
-    if (err != HSA_STATUS_SUCCESS) {
-      fprintf(stderr, "Get memory pool size failed: %s\n",
-              get_error_string(err));
-      return err;
-    }
-    if (size > 0)
-      Result->push_back(MemoryPool);
+  err = hsa_amd_memory_pool_get_info(MemoryPool, HSA_AMD_MEMORY_POOL_INFO_SIZE,
+                                     &size);
+    fprintf(stderr, "Get memory pool size failed: %s\n", get_error_string(err));
+
+  if ((GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED) &&
+      (GlobalFlags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_KERNARG_INIT) &&
+      size > 0) {
   }
 
   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) {
+    fprintf(stderr, "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) {
+      printf("[%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) {
public:
 
   std::vector<hsa_executable_t> HSAExecutables;
 
   std::vector<std::map<std::string, atl_kernel_info_t>> KernelInfoTable;
   std::vector<std::map<std::string, atl_symbol_info_t>> SymbolInfoTable;
 
   hsa_amd_memory_pool_t KernArgPool;
 
+  // fine grained memory pool for host allocations
+  hsa_amd_memory_pool_t HostFineGrainedMemoryPool;
+
+  // fine and coarse-grained memory pools per offloading device
+  std::vector<hsa_amd_memory_pool_t> DeviceFineGrainedMemoryPools;
+  std::vector<hsa_amd_memory_pool_t> DeviceCoarseGrainedMemoryPools;
+
   struct atmiFreePtrDeletor {
     void operator()(void *p) {
       core::Runtime::Memfree(p); // ignore failure to free
     }
   };
 
   // device_State shared across loaded binaries, error if inconsistent size
   std::vector<std::pair<std::unique_ptr<void, atmiFreePtrDeletor>, uint64_t>>
assert(device_id < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncGblEntries[device_id].emplace_back();
     FuncOrGblEntryTy &E = FuncGblEntries[device_id].back();
     // KernelArgPoolMap.clear();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = 0;
   }
 
+  hsa_status_t addDeviceMemoryPool(hsa_amd_memory_pool_t MemoryPool,
+                                   int DeviceId) {
+    assert(DeviceId < DeviceFineGrainedMemoryPools.size() && "Error here.");
+    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;
+    }
+
+    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) {
+    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 &&
+        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) {
+      fprintf(stderr, "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) {
+      fprintf(stderr,
+              "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];
+  }
+
+  hsa_amd_memory_pool_t getHostMemoryPool() {
+    return HostFineGrainedMemoryPool;
+  }
+
   RTLDeviceInfoTy() {
     // LIBOMPTARGET_KERNEL_TRACE provides a kernel launch trace to stderr
     // anytime. You do not need a debug library build.
     //  0 => no tracing
     //  1 => tracing dispatch only
     // >1 => verbosity increase
     if (char *envStr = getenv("LIBOMPTARGET_KERNEL_TRACE"))
       print_kernel_trace = atoi(envStr);
RTLDeviceInfoTy() {
     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();
+    if (err != HSA_STATUS_SUCCESS) {
+      DP("Error when setting up 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;
       {