[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

 //===--- amdgpu/src/rtl.cpp --------------------------------------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for hsa machine
 //
 //===----------------------------------------------------------------------===//
 
 #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"

 
 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;
-        });
+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;
+  std::tie(Err, IsValid) = isValidMemoryPool(MemoryPool);
+    DP("isValidMemoryPool failed: %s\n",
+       get_error_string(Err));
 
-    if (Err != HSA_STATUS_SUCCESS) {
-      DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
-         "Iterate all memory pools", get_error_string(Err));
-      return Err;
-    }
+  if (IsValid) {
+    Result->push_back(MemoryPool);
   }
 
   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]};
-}
-
 } // namespace
 } // namespace core
 
 struct EnvironmentVariables {
   int NumTeams;
   int TeamLimit;
   int TeamThreadLimit;
   int MaxTeamsDefault;
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;
+    uint32_t Size = 0;
     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 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("[%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 setupMemoryPools() {
-    using namespace std::placeholders;
-    hsa_status_t Err;
-    Err = core::collectMemoryPools(
-        CPUAgents, std::bind(&RTLDeviceInfoTy::addHostMemoryPool, this, _1, _2));
+  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("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));
+    }
+
+    // We need two fine-grained pools.
+    //  1. One with kernarg flag set for storing kernel arguments
+    //  2. Second for host allocations
+    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("HSA error in collecting memory pools for offload devices: %s\n",
-         get_error_string(Err));
+        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;
+        else
+          HostFineGrainedMemoryPool = MemoryPool;
+      }
+    }
+
     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("Setup for Device Memory Pools failed\n");
+      return;
+    }
+
+    err = setupHostMemoryPools(CPUAgents);
     if (err != HSA_STATUS_SUCCESS) {
-      DP("Error when setting up 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++) {