hip-clang: fix kernel not found on multi-gpu

__hipRegisterFunction is called during by .init functions during program initialization.
It calls hipModuleGetFunction to locate kernel symbol in code objects. hipModuleGetFunction
assumes current device when locating kernel symbols. This works for HCC but not for hip-clang,
since hip-clang needs to locate kernel symbols for different devices without switching
between devices.

This patch introduces a new hsa agent parameter to ihipModuleGetFunction, which allows
__hipRegisterFunction to choose the correct hsa agent when locating kernel symbols. By
default it uses this_agent(), therefore this patch has no impact on HCC.


[ROCm/clr commit: 8f5c812a68]

projects/clr/hipamd/src/hip_clang.cpp

@@ -132,11 +132,13 @@ extern "C" void __hipRegisterFunction(
   assert(modules && modules->size() >= g_deviceCnt);
   for (int deviceId = 0; deviceId < g_deviceCnt; ++deviceId) {
     hipFunction_t function;
-    if ((hipSuccess == hipModuleGetFunction(&function, modules->at(deviceId), deviceName) ||
+    hsa_agent_t agent = g_allAgents[deviceId + 1];
+    if ((hipSuccess == hipModuleGetFunctionEx(&function, modules->at(deviceId), deviceName, &agent) ||
         // With code-object-v3, we need to match the kernel descriptor symbol name
-        (hipSuccess == hipModuleGetFunction(
+        (hipSuccess == hipModuleGetFunctionEx(
                            &function, modules->at(deviceId),
-                           (std::string(deviceName) + std::string(".kd")).c_str()
+                           (std::string(deviceName) + std::string(".kd")).c_str(),
+                           &agent
                        ))) && function != nullptr) {
       functions[deviceId] = function;
     }

projects/clr/hipamd/src/hip_hcc_internal.h

@@ -943,6 +943,8 @@ extern hipError_t ihipDeviceSetState();
 
 extern ihipDevice_t* ihipGetDevice(int);
 ihipCtx_t* ihipGetPrimaryCtx(unsigned deviceIndex);
+hipError_t hipModuleGetFunctionEx(hipFunction_t* hfunc, hipModule_t hmod,
+                                  const char* name, hsa_agent_t *agent);
 
 
 hipStream_t ihipSyncAndResolveStream(hipStream_t);

projects/clr/hipamd/src/hip_module.cpp

@@ -360,13 +360,14 @@ inline hsa_status_t copy_agent_global_variables(hsa_executable_t, hsa_agent_t ag
     return HSA_STATUS_SUCCESS;
 }
 
-hsa_executable_symbol_t find_kernel_by_name(hsa_executable_t executable, const char* kname) {
+hsa_executable_symbol_t find_kernel_by_name(hsa_executable_t executable, const char* kname,
+                                            hsa_agent_t agent) {
     using namespace hip_impl;
 
     pair<const char*, hsa_executable_symbol_t> r{kname, {}};
 
     hsa_executable_iterate_agent_symbols(
-        executable, this_agent(),
+        executable, agent,
         [](hsa_executable_t, hsa_agent_t, hsa_executable_symbol_t x, void* s) {
             auto p = static_cast<pair<const char*, hsa_executable_symbol_t>*>(s);
 
@@ -384,6 +385,11 @@ hsa_executable_symbol_t find_kernel_by_name(hsa_executable_t executable, const c
     return r.second;
 }
 
+hsa_executable_symbol_t find_kernel_by_name(hsa_executable_t executable, const char* kname) {
+    return find_kernel_by_name(executable, kname, hip_impl::this_agent());
+}
+
+
 string read_elf_file_as_string(const void* file) {
     // Precondition: file points to an ELF image that was BITWISE loaded
     //               into process accessible memory, and not one loaded by
@@ -437,7 +443,8 @@ namespace hip_impl {
     }
 } // Namespace hip_impl.
 
-hipError_t ihipModuleGetFunction(hipFunction_t* func, hipModule_t hmod, const char* name) {
+hipError_t ihipModuleGetFunction(hipFunction_t* func, hipModule_t hmod, const char* name,
+                                 hsa_agent_t *agent = nullptr) {
     using namespace hip_impl;
 
     if (!func || !name) return hipErrorInvalidValue;
@@ -450,7 +457,10 @@ hipError_t ihipModuleGetFunction(hipFunction_t* func, hipModule_t hmod, const ch
 
     if (!*func) return hipErrorInvalidValue;
 
-    auto kernel = find_kernel_by_name(hmod->executable, name);
+    if (!agent)
+      *agent = this_agent();
+
+    auto kernel = find_kernel_by_name(hmod->executable, name, *agent);
 
     if (kernel.handle == 0u) return hipErrorNotFound;
 
@@ -462,11 +472,19 @@ hipError_t ihipModuleGetFunction(hipFunction_t* func, hipModule_t hmod, const ch
     return hipSuccess;
 }
 
+// Get kernel for the current hsa agent.
 hipError_t hipModuleGetFunction(hipFunction_t* hfunc, hipModule_t hmod, const char* name) {
     HIP_INIT_API(hipModuleGetFunction, hfunc, hmod, name);
     return ihipLogStatus(ihipModuleGetFunction(hfunc, hmod, name));
 }
 
+// Get kernel for the given hsa agent. Internal use only.
+hipError_t hipModuleGetFunctionEx(hipFunction_t* hfunc, hipModule_t hmod,
+                                  const char* name, hsa_agent_t *agent) {
+    HIP_INIT_API(hipModuleGetFunctionEx, hfunc, hmod, name);
+    return ihipLogStatus(ihipModuleGetFunction(hfunc, hmod, name, agent));
+}
+
 namespace {
 hipFuncAttributes make_function_attributes(const amd_kernel_code_t& header) {
     hipFuncAttributes r{};