Merge pull request #997 from yxsamliu/mgpu

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

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;
     }

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);

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{};