[AMDGPU][Libomptarget] Move Kernel/Symbol info tables to RTLDeviceInfoTy

[pdhaliwal]

Summary

Two globals KernelInfoTable & SymbolInfoTable are moved
into RTLDeviceInfoTy class.
This builds on the top of D102691.
[2/2]

[JonChesterfield]

The vector<> part is for N GPUs, and each function call is for a specific GPU. So we never need to pass the whole vector along, only the element corresponding to that GPU.

I'd prefer we populate a new map for each code object and then merge them, instead of mutating the existing map, but I guess that's not a significant difference.

PopulateInfoTablesContext looks like a hazard - it's a struct holding references, which is usually a bad sign for lifetimes. If that's only used to pass data across the void* boundary, let's use a lambda instance instead.

[pdhaliwal]

PopulateInfoTablesContext looks like a hazard - it's a struct holding references, which is usually a bad sign for lifetimes. If that's only used to pass data across the void* boundary, let's use a lambda instance instead.

The capturing lambda won't work here so that's why I chose to use a struct of references.

[pdhaliwal]

I could not find alternative to references in the PopulateInfoTablesContext. I could use pointers
but I don't see how that is different from having references.

Updated the patch to pass map directly instead of whole vector.

[pdhaliwal]

Remove RegisterModuleFromMemory declaration from Runtime

[JonChesterfield]

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

819–821

Minor point, output parameters at the end of the list is probably more common than at the start of the list for C++.

I'd rather pass mutable things by pointer, so the call site looks like function(42, &something_mutated), but that convention is from a different codebase to this one.

Better than either would be returning a pair<status, map> instead of mutating it in place. But given this change is moving from a global variable to a parameter, getting rid of the mutation at the same time seems a step too far for one patch.

993–994

Would lean towards dropping the void*, e.g. via D103030:

static hsa_status_t populate_InfoTables(hsa_executable_symbol_t symbol, 
  std::map<std::string, atl_kernel_info_t> &KernelInfoTable,
  std::map<std::string, atl_symbol_info_t> &SymbolInfoTable) {

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

89

This seems orthogonal to moving the global.

Also, looks like we now have declarations in rtl.cpp and definitions in atmi_interop_hsa.cpp, which seems bad as we lose the compiler checking that the call site matches the definition

[pdhaliwal]

Remove struct. Following D103030 approach.

[pdhaliwal]

Revert delete atmi_hsa_interop.h

[pdhaliwal]

Rebased

[JonChesterfield]

except for the lost error check above, I think this is now a non-functional change. LG with that error path restored

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

1127

lost an err = assignment and check here

[JonChesterfield]

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

501

This is dead - resize the vector will populate it with default constructed map<string>, and clear on a default constructed map is a no-op

574

The KernelInfoTable and SymbolInfoTable calls here are also dead. They're free store allocated memory, don't need to tear them down before closing hsa (unlike kernel arg pool, which contains pointers from hsa allocated memory)

[JonChesterfield]

passing by const& is much better than passing by &, I think we should do that in this patch

openmp/libomptarget/plugins/amdgpu/impl/atmi_interop_hsa.h

80–81

This is only called from one place, with info == HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE, so can be simplified to a map lookup. Orthogonal to this patch.

This and get_symbol_info only do lookups in the map so we should pass the map to them by const&.

[pdhaliwal]

Removed dead code and pass tables as const &

[JonChesterfield]

LG, thanks!

[pdhaliwal]

Rebased.

openmp/libomptarget/plugins/amdgpu/impl/atmi_interop_hsa.h

 /*===--------------------------------------------------------------------------
  *              ATMI (Asynchronous Task and Memory Interface)
  *
  * This file is distributed under the MIT License. See LICENSE.txt for details.
  *===------------------------------------------------------------------------*/
 #ifndef INCLUDE_ATMI_INTEROP_HSA_H_
 #define INCLUDE_ATMI_INTEROP_HSA_H_
 
 #include "atmi_runtime.h"
 #include "hsa.h"
 #include "hsa_ext_amd.h"
+#include "internal.h"
+
+#include <map>
+#include <string>
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 /** \defgroup interop_hsa_functions ATMI-HSA Interop
  *  @{
  */
 

  * @retval ::ATMI_STATUS_SUCCESS The function has executed successfully.
  *
  * @retval ::ATMI_STATUS_ERROR If @p symbol, @p var_addr or @p var_size are
  * invalid
  * location in the current node, or if ATMI is not initialized.
  *
  * @retval ::ATMI_STATUS_UNKNOWN The function encountered errors.
  */
-atmi_status_t atmi_interop_hsa_get_symbol_info(atmi_mem_place_t place,
-                                               const char *symbol,
-                                               void **var_addr,
-                                               unsigned int *var_size);
-
+atmi_status_t atmi_interop_hsa_get_symbol_info(
+    const std::map<std::string, atl_symbol_info_t> &SymbolInfoTable,
+    atmi_mem_place_t place, const char *symbol, void **var_addr,
 /**
  * @brief Get the HSA-specific kernel info from a kernel name
  *
  * @detail Use this function to query the HSA-specific kernel info from the
  * kernel name.
  * This function is meaningful only after calling one
  * of the @p atmi_module_register functions.
  *

  * @retval ::ATMI_STATUS_SUCCESS The function has executed successfully.
  *
  * @retval ::ATMI_STATUS_ERROR If @p symbol, @p var_addr or @p var_size are
  * invalid
  * location in the current node, or if ATMI is not initialized.
  *
  * @retval ::ATMI_STATUS_UNKNOWN The function encountered errors.
  */
 atmi_status_t atmi_interop_hsa_get_kernel_info(
+    const std::map<std::string, atl_kernel_info_t> &KernelInfoTable,
     atmi_mem_place_t place, const char *kernel_name,
     hsa_executable_symbol_info_t info, uint32_t *value);
+
 /** @} */
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // INCLUDE_ATMI_INTEROP_HSA_H_

openmp/libomptarget/plugins/amdgpu/impl/atmi_interop_hsa.cpp

 /*===--------------------------------------------------------------------------
  *              ATMI (Asynchronous Task and Memory Interface)
  *
  * This file is distributed under the MIT License. See LICENSE.txt for details.
  *===------------------------------------------------------------------------*/
 #include "atmi_interop_hsa.h"
 #include "internal.h"
 
 using core::atl_is_atmi_initialized;
 
-atmi_status_t atmi_interop_hsa_get_symbol_info(atmi_mem_place_t place,
-                                               const char *symbol,
-                                               void **var_addr,
-                                               unsigned int *var_size) {
+atmi_status_t atmi_interop_hsa_get_symbol_info(
+    const std::map<std::string, atl_symbol_info_t> &SymbolInfoTable,
+    atmi_mem_place_t place, const char *symbol, void **var_addr,
   /*
      // Typical usage:
      void *var_addr;
      size_t var_size;
      atmi_interop_hsa_get_symbol_addr(gpu_place, "symbol_name", &var_addr,
      &var_size);
      atmi_memcpy(signal, host_add, var_addr, var_size);
   */
 
   if (!atl_is_atmi_initialized())
     return ATMI_STATUS_ERROR;
   atmi_machine_t *machine = atmi_machine_get_info();
   if (!symbol || !var_addr || !var_size || !machine)
     return ATMI_STATUS_ERROR;
   if (place.dev_id < 0 ||
       place.dev_id >= machine->device_count_by_type[place.dev_type])
     return ATMI_STATUS_ERROR;
 
   // get the symbol info
   std::string symbolStr = std::string(symbol);
-  if (SymbolInfoTable[place.dev_id].find(symbolStr) !=
-      SymbolInfoTable[place.dev_id].end()) {
-    atl_symbol_info_t info = SymbolInfoTable[place.dev_id][symbolStr];
+  auto It = SymbolInfoTable.find(symbolStr);
+  if (It != SymbolInfoTable.end()) {
+    atl_symbol_info_t info = It->second;
     *var_addr = reinterpret_cast<void *>(info.addr);
     *var_size = info.size;
     return ATMI_STATUS_SUCCESS;
   } else {
     *var_addr = NULL;
     *var_size = 0;
     return ATMI_STATUS_ERROR;
   }
 }
 
 atmi_status_t atmi_interop_hsa_get_kernel_info(
+    const std::map<std::string, atl_kernel_info_t> &KernelInfoTable,
     atmi_mem_place_t place, const char *kernel_name,
     hsa_executable_symbol_info_t kernel_info, uint32_t *value) {
   /*
      // Typical usage:
      uint32_t value;
      atmi_interop_hsa_get_kernel_addr(gpu_place, "kernel_name",
                                   HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE,
                                   &val);
   */
 
   if (!atl_is_atmi_initialized())
     return ATMI_STATUS_ERROR;
   atmi_machine_t *machine = atmi_machine_get_info();
   if (!kernel_name || !value || !machine)
     return ATMI_STATUS_ERROR;
   if (place.dev_id < 0 ||
       place.dev_id >= machine->device_count_by_type[place.dev_type])
     return ATMI_STATUS_ERROR;
 
   atmi_status_t status = ATMI_STATUS_SUCCESS;
   // get the kernel info
   std::string kernelStr = std::string(kernel_name);
-  if (KernelInfoTable[place.dev_id].find(kernelStr) !=
-      KernelInfoTable[place.dev_id].end()) {
-    atl_kernel_info_t info = KernelInfoTable[place.dev_id][kernelStr];
+  auto It = KernelInfoTable.find(kernelStr);
+  if (It != KernelInfoTable.end()) {
+    atl_kernel_info_t info = It->second;
     switch (kernel_info) {
     case HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_GROUP_SEGMENT_SIZE:
       *value = info.group_segment_size;
       break;
     case HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_PRIVATE_SEGMENT_SIZE:
       *value = info.private_segment_size;
       break;
     case HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE:

openmp/libomptarget/plugins/amdgpu/impl/internal.h

typedef struct atl_kernel_info_s {
   std::vector<uint64_t> arg_sizes;
 } atl_kernel_info_t;
 
 typedef struct atl_symbol_info_s {
   uint64_t addr;
   uint32_t size;
 } atl_symbol_info_t;
 
-extern std::vector<std::map<std::string, atl_kernel_info_t>> KernelInfoTable;
-extern std::vector<std::map<std::string, atl_symbol_info_t>> SymbolInfoTable;
-
 // ---------------------- Kernel End -------------
 
 namespace core {
 class TaskgroupImpl;
 class TaskImpl;
 class Kernel;
 class KernelImpl;
 } // namespace core

openmp/libomptarget/plugins/amdgpu/impl/system.cpp

 };
 
 // global variables. TODO: Get rid of these
 atmi_machine_t g_atmi_machine;
 ATLMachine g_atl_machine;
 
 std::vector<hsa_amd_memory_pool_t> atl_gpu_kernarg_pools;
 
-std::vector<std::map<std::string, atl_kernel_info_t>> KernelInfoTable;
-std::vector<std::map<std::string, atl_symbol_info_t>> SymbolInfoTable;
-
 bool g_atmi_initialized = false;
 
 /*
    atlc is all internal global values.
    The structure atl_context_t is defined in atl_internal.h
    Most references will use the global structure prefix atlc.
 */
 atl_context_t atlc = {.struct_initialized = false};
atmi_status_t Runtime::Initialize() {
   }
 
   atl_set_atmi_initialized();
   return ATMI_STATUS_SUCCESS;
 }
 
 atmi_status_t Runtime::Finalize() {
   atmi_status_t rc = ATMI_STATUS_SUCCESS;
-  for (uint32_t i = 0; i < SymbolInfoTable.size(); i++) {
-    SymbolInfoTable[i].clear();
-  }
-  SymbolInfoTable.clear();
-  for (uint32_t i = 0; i < KernelInfoTable.size(); i++) {
-    KernelInfoTable[i].clear();
-  }
-  KernelInfoTable.clear();
-
   atl_reset_atmi_initialized();
   hsa_status_t err = hsa_shut_down();
   if (err != HSA_STATUS_SUCCESS) {
     printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__, "Shutting down HSA",
            get_error_string(err));
     rc = ATMI_STATUS_ERROR;
   }
 
if (atlc.g_hsa_initialized == false) {
     if (err != HSA_STATUS_SUCCESS)
       return err;
     if (err != HSA_STATUS_SUCCESS) {
       printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
              "After initializing compute and memory", get_error_string(err));
       return err;
     }
 
-    int gpu_count = g_atl_machine.processorCount<ATLGPUProcessor>();
-    KernelInfoTable.resize(gpu_count);
-    SymbolInfoTable.resize(gpu_count);
-    for (uint32_t i = 0; i < SymbolInfoTable.size(); i++)
-      SymbolInfoTable[i].clear();
-    for (uint32_t i = 0; i < KernelInfoTable.size(); i++)
-      KernelInfoTable[i].clear();
     atlc.g_hsa_initialized = true;
     DEBUG_PRINT("done\n");
   }
   return HSA_STATUS_SUCCESS;
 }
 
 void init_tasks() {
   if (atlc.g_tasks_initialized != false)
if (message_is_string(key, ".name")) {
         }
       });
     }
   });
   return error;
 }
 } // namespace
 
-static hsa_status_t get_code_object_custom_metadata(void *binary,
-                                                    size_t binSize, int gpu) {
+static hsa_status_t get_code_object_custom_metadata(
+    void *binary, size_t binSize, int gpu,
+    std::map<std::string, atl_kernel_info_t> &KernelInfoTable) {
   // parse code object with different keys from v2
   // also, the kernel name is not the same as the symbol name -- so a
   // symbol->name map is needed
 
   std::pair<unsigned char *, unsigned char *> metadata =
       find_metadata(binary, binSize);
   if (!metadata.first) {
     return HSA_STATUS_ERROR_INVALID_CODE_OBJECT;
for (size_t i = 0; i < kernelsSize; i++) {
     // implicit args by discounting the compiler set implicit args
     info.kernel_segment_size =
         (hasHiddenArgs ? kernel_explicit_args_size : kernel_segment_size) +
         sizeof(atmi_implicit_args_t);
     DEBUG_PRINT("[%s: kernarg seg size] (%lu --> %u)\n", kernelName.c_str(),
                 kernel_segment_size, info.kernel_segment_size);
 
     // kernel received, now add it to the kernel info table
-    KernelInfoTable[gpu][kernelName] = info;
+    KernelInfoTable[kernelName] = info;
   }
 
   return HSA_STATUS_SUCCESS;
 }
 
-static hsa_status_t populate_InfoTables(hsa_executable_symbol_t symbol,
-                                        int gpu) {
+static hsa_status_t
+populate_InfoTables(hsa_executable_symbol_t symbol, int gpu,
+                    std::map<std::string, atl_kernel_info_t> &KernelInfoTable,
+                    std::map<std::string, atl_symbol_info_t> &SymbolInfoTable) {
   hsa_symbol_kind_t type;
 
   uint32_t name_length;
   hsa_status_t err;
   err = hsa_executable_symbol_get_info(symbol, HSA_EXECUTABLE_SYMBOL_INFO_TYPE,
                                        &type);
   if (err != HSA_STATUS_SUCCESS) {
     printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
if (type == HSA_SYMBOL_KIND_KERNEL) {
     // remove the suffix .kd from symbol name.
     name[name_length - 3] = 0;
 
     atl_kernel_info_t info;
     std::string kernelName(name);
     // by now, the kernel info table should already have an entry
     // because the non-ROCr custom code object parsing is called before
     // iterating over the code object symbols using ROCr
-    if (KernelInfoTable[gpu].find(kernelName) == KernelInfoTable[gpu].end()) {
-      return HSA_STATUS_ERROR;
+    if (KernelInfoTable.find(kernelName) == KernelInfoTable.end()) {
+      if (HSA_STATUS_ERROR_INVALID_CODE_OBJECT != HSA_STATUS_SUCCESS) {
+        printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
+               "Finding the entry kernel info table",
+               get_error_string(HSA_STATUS_ERROR_INVALID_CODE_OBJECT));
+        exit(1);
+      }
     }
     // found, so assign and update
-    info = KernelInfoTable[gpu][kernelName];
+    info = KernelInfoTable[kernelName];
 
     /* Extract dispatch information from the symbol */
     err = hsa_executable_symbol_get_info(
         symbol, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT,
         &(info.kernel_object));
     if (err != HSA_STATUS_SUCCESS) {
       printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
              "Extracting the symbol from the executable",
if (type == HSA_SYMBOL_KIND_KERNEL) {
 
     DEBUG_PRINT(
         "Kernel %s --> %lx symbol %u group segsize %u pvt segsize %u bytes "
         "kernarg\n",
         kernelName.c_str(), info.kernel_object, info.group_segment_size,
         info.private_segment_size, info.kernel_segment_size);
 
     // assign it back to the kernel info table
-    KernelInfoTable[gpu][kernelName] = info;
+    KernelInfoTable[kernelName] = info;
     free(name);
   } else if (type == HSA_SYMBOL_KIND_VARIABLE) {
     err = hsa_executable_symbol_get_info(
         symbol, HSA_EXECUTABLE_SYMBOL_INFO_NAME_LENGTH, &name_length);
     if (err != HSA_STATUS_SUCCESS) {
       printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
              "Symbol info extraction", get_error_string(err));
       return err;
if (type == HSA_SYMBOL_KIND_KERNEL) {
     atmi_mem_place_t place = ATMI_MEM_PLACE(ATMI_DEVTYPE_GPU, gpu, 0);
     DEBUG_PRINT("Symbol %s = %p (%u bytes)\n", name, (void *)info.addr,
                 info.size);
     err = register_allocation(reinterpret_cast<void *>(info.addr),
                               (size_t)info.size, place);
     if (err != HSA_STATUS_SUCCESS) {
       return err;
     }
-    SymbolInfoTable[gpu][std::string(name)] = info;
+    SymbolInfoTable[std::string(name)] = info;
     free(name);
   } else {
     DEBUG_PRINT("Symbol is an indirect function\n");
   }
   return HSA_STATUS_SUCCESS;
 }
 
-atmi_status_t Runtime::RegisterModuleFromMemory(
+atmi_status_t RegisterModuleFromMemory(
+    std::map<std::string, atl_kernel_info_t> &KernelInfoTable,
+    std::map<std::string, atl_symbol_info_t> &SymbolInfoTable,
     void *module_bytes, size_t module_size, atmi_place_t place,
     atmi_status_t (*on_deserialized_data)(void *data, size_t size,
                                           void *cb_state),
     void *cb_state, std::vector<hsa_executable_t> &HSAExecutables) {
   hsa_status_t err;
   int gpu = place.device_id;
   assert(gpu >= 0);
 
atmi_status_t RegisterModuleFromMemory(
 
   bool module_load_success = false;
   do // Existing control flow used continue, preserve that for this patch
   {
     {
       // Some metadata info is not available through ROCr API, so use custom
       // code object metadata parsing to collect such metadata info
 
-      err = get_code_object_custom_metadata(module_bytes, module_size, gpu);
+      err = get_code_object_custom_metadata(module_bytes, module_size, gpu,
+                                            KernelInfoTable);
       if (err != HSA_STATUS_SUCCESS) {
         DEBUG_PRINT("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
                     "Getting custom code object metadata",
                     get_error_string(err));
         continue;
       }
 
       // Deserialize code object.
if (err != HSA_STATUS_SUCCESS) {
       printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
              "Freeze the executable", get_error_string(err));
       return ATMI_STATUS_ERROR;
     }
 
     err = hsa::executable_iterate_symbols(
         executable,
         [&](hsa_executable_t, hsa_executable_symbol_t symbol) -> hsa_status_t {
-          return populate_InfoTables(symbol, gpu);
+          return populate_InfoTables(symbol, gpu, KernelInfoTable,
+                                     SymbolInfoTable);
         });
-
     if (err != HSA_STATUS_SUCCESS) {
       printf("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
              "Iterating over symbols for execuatable", get_error_string(err));
       return ATMI_STATUS_ERROR;
     }
 
     // save the executable and destroy during finalize
     HSAExecutables.push_back(executable);
     return ATMI_STATUS_SUCCESS;
   } else {
     return ATMI_STATUS_ERROR;
   }
 }
 
 } // namespace core

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

#define check(msg, status)                                                     \
   }
 #else
 #define check(msg, status)                                                     \
   {}
 #endif
 
 #include "elf_common.h"
 
+namespace core {
+atmi_status_t RegisterModuleFromMemory(
+    std::map<std::string, atl_kernel_info_t> &KernelInfo,
+    std::map<std::string, atl_symbol_info_t> &SymbolInfoTable, void *, size_t,
+    atmi_place_t,
+    atmi_status_t (*on_deserialized_data)(void *data, size_t size,
+                                          void *cb_state),
+    void *cb_state, std::vector<hsa_executable_t> &HSAExecutables);
+}
+
 /// Keep entries table per device
 struct FuncOrGblEntryTy {
   __tgt_target_table Table;
   std::vector<__tgt_offload_entry> Entries;
 };
 
 enum ExecutionModeType {
   SPMD,    // constructors, destructors,
public:
 
   // Resource pools
   SignalPoolT FreeSignalPool;
 
   bool hostcall_required = false;
 
   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;
+
   struct atmiFreePtrDeletor {
     void operator()(void *p) {
       atmi_free(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>>
RTLDeviceInfoTy() {
     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);
 
     for (int i = 0; i < NumberOfDevices; i++) {
       HSAQueues[i] = nullptr;
     }
 
     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);
public:
   }
 
   ~RTLDeviceInfoTy() {
     DP("Finalizing the HSA-ATMI DeviceInfo.\n");
     // Run destructors on types that use HSA before
     // atmi_finalize removes access to it
     deviceStateStore.clear();
     KernelArgPoolMap.clear();
     // Terminate hostrpc before finalizing ATMI
     hostrpc_terminate();
 
     for (uint32_t I = 0; I < HSAExecutables.size(); I++) {
       hsa_status_t Err = hsa_executable_destroy(HSAExecutables[I]);
       if (Err != HSA_STATUS_SUCCESS) {
         DP("[%s:%d] %s failed: %s\n", __FILE__, __LINE__,
            "Destroying executable", get_error_string(Err));
       }
if (rc == 0) {
     return ATMI_STATUS_SUCCESS;
   } else {
     return ATMI_STATUS_ERROR;
   }
 }
 
 template <typename C>
 atmi_status_t module_register_from_memory_to_place(
+    std::map<std::string, atl_kernel_info_t> &KernelInfoTable,
+    std::map<std::string, atl_symbol_info_t> &SymbolInfoTable,
     void *module_bytes, size_t module_size, atmi_place_t place, C cb,
     std::vector<hsa_executable_t> &HSAExecutables) {
   auto L = [](void *data, size_t size, void *cb_state) -> atmi_status_t {
     C *unwrapped = static_cast<C *>(cb_state);
     return (*unwrapped)(data, size);
   };
-  return core::Runtime::RegisterModuleFromMemory(
-      module_bytes, module_size, place, L, static_cast<void *>(&cb),
-      HSAExecutables);
+  return core::RegisterModuleFromMemory(
+      KernelInfoTable, SymbolInfoTable, module_bytes, module_size, place, L,
+      static_cast<void *>(&cb), HSAExecutables);
 }
 } // namespace
 
 static uint64_t get_device_State_bytes(char *ImageStart, size_t img_size) {
   uint64_t device_State_bytes = 0;
   {
     // If this is the deviceRTL, get the state variable size
     symbol_info size_si;
#endif
   }
 
   atmi_status_t after_loading() {
     if (valid) {
       if (!in_image()) {
         DP("Setting global device environment after load (%u bytes)\n",
            si.size);
         int device_id = host_device_env.device_num;
-
+        auto &SymbolInfo = DeviceInfo.SymbolInfoTable[device_id];
         void *state_ptr;
         uint32_t state_ptr_size;
         atmi_status_t err = atmi_interop_hsa_get_symbol_info(
-            get_gpu_mem_place(device_id), sym(), &state_ptr, &state_ptr_size);
+            SymbolInfo, get_gpu_mem_place(device_id), sym(), &state_ptr,
+            &state_ptr_size);
         if (err != ATMI_STATUS_SUCCESS) {
           DP("failed to find %s in loaded image\n", sym());
           return err;
         }
 
         if (state_ptr_size != si.size) {
           DP("Symbol had size %u before loading, %u after\n", state_ptr_size,
              si.size);
__tgt_target_table *__tgt_rtl_load_binary_locked(int32_t device_id,
   if (!elf_machine_id_is_amdgcn(image)) {
     return NULL;
   }
 
   {
     auto env = device_environment(device_id, DeviceInfo.NumberOfDevices, image,
                                   img_size);
 
+    auto &KernelInfo = DeviceInfo.KernelInfoTable[device_id];
+    auto &SymbolInfo = DeviceInfo.SymbolInfoTable[device_id];
     atmi_status_t err = module_register_from_memory_to_place(
-        (void *)image->ImageStart, img_size, get_gpu_place(device_id),
+        KernelInfo, SymbolInfo, (void *)image->ImageStart, img_size,
+        get_gpu_place(device_id),
         [&](void *data, size_t size) {
           if (image_contains_symbol(data, size, "needs_hostcall_buffer")) {
             __atomic_store_n(&DeviceInfo.hostcall_required, true,
                              __ATOMIC_RELEASE);
           }
           return env.before_loading(data, size);
         },
         DeviceInfo.HSAExecutables);
__tgt_target_table *__tgt_rtl_load_binary_locked(int32_t device_id,
 
   {
     // the device_State array is either large value in bss or a void* that
     // needs to be assigned to a pointer to an array of size device_state_bytes
     // If absent, it has been deadstripped and needs no setup.
 
     void *state_ptr;
     uint32_t state_ptr_size;
+    auto &SymbolInfoMap = DeviceInfo.SymbolInfoTable[device_id];
     atmi_status_t err = atmi_interop_hsa_get_symbol_info(
-        get_gpu_mem_place(device_id), "omptarget_nvptx_device_State",
-        &state_ptr, &state_ptr_size);
+        SymbolInfoMap, get_gpu_mem_place(device_id),
+        "omptarget_nvptx_device_State", &state_ptr, &state_ptr_size);
 
     if (err != ATMI_STATUS_SUCCESS) {
       DP("No device_state symbol found, skipping initialization\n");
     } else {
       if (state_ptr_size < sizeof(void *)) {
         DP("unexpected size of state_ptr %u != %zu\n", state_ptr_size,
            sizeof(void *));
         return NULL;
for (__tgt_offload_entry *e = HostBegin; e != HostEnd; ++e) {
     }
 
     if (e->size) {
       __tgt_offload_entry entry = *e;
 
       void *varptr;
       uint32_t varsize;
 
+      auto &SymbolInfoMap = DeviceInfo.SymbolInfoTable[device_id];
       atmi_status_t err = atmi_interop_hsa_get_symbol_info(
-          get_gpu_mem_place(device_id), e->name, &varptr, &varsize);
+          SymbolInfoMap, get_gpu_mem_place(device_id), e->name, &varptr,
+          &varsize);
 
       if (err != ATMI_STATUS_SUCCESS) {
         // Inform the user what symbol prevented offloading
         DP("Loading global '%s' (Failed)\n", e->name);
         return NULL;
       }
 
       if (varsize != e->size) {
if (e->size) {
 
       continue;
     }
 
     DP("to find the kernel name: %s size: %lu\n", e->name, strlen(e->name));
 
     atmi_mem_place_t place = get_gpu_mem_place(device_id);
     uint32_t kernarg_segment_size;
+    auto &KernelInfoMap = DeviceInfo.KernelInfoTable[device_id];
     atmi_status_t err = atmi_interop_hsa_get_kernel_info(
-        place, e->name, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE,
+        KernelInfoMap, place, e->name,
+        HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE,
         &kernarg_segment_size);
 
     // each arg is a void * in this openmp implementation
     uint32_t arg_num = kernarg_segment_size / sizeof(void *);
     std::vector<size_t> arg_sizes(arg_num);
     for (std::vector<size_t>::iterator it = arg_sizes.begin();
          it != arg_sizes.end(); it++) {
       *it = sizeof(void *);
for (int32_t i = 0; i < arg_num; ++i) {
     ptrs[i] = (void *)((intptr_t)tgt_args[i] + tgt_offsets[i]);
     args[i] = &ptrs[i];
     DP("Offseted base: arg[%d]:" DPxMOD "\n", i, DPxPTR(ptrs[i]));
   }
 
   KernelTy *KernelInfo = (KernelTy *)tgt_entry_ptr;
 
   std::string kernel_name = std::string(KernelInfo->Name);
+  auto &KernelInfoTable = DeviceInfo.KernelInfoTable;
   if (KernelInfoTable[device_id].find(kernel_name) ==
       KernelInfoTable[device_id].end()) {
     DP("Kernel %s not found\n", kernel_name.c_str());
     return OFFLOAD_FAIL;
   }
 
   uint32_t group_segment_size;
   uint32_t sgpr_count, vgpr_count, sgpr_spill_count, vgpr_spill_count;