SWDEV-216213 - Lookup module functions from PlatformState::functions_.

Change-Id: I91dfe327f2ebdcf4c9b39ddd14d60aa0ce2fa9f4

vdi/hip_internal.hpp

@@ -194,6 +194,7 @@ public:
     std::string deviceName;
     std::vector< std::pair< hipModule_t, bool > >* modules;
     std::vector<hipFunction_t> functions;
+    bool dyn_mod;
   };
   struct DeviceVar {
     void* shadowVptr;
@@ -216,6 +217,7 @@ public:
     return *platform_;
   }
 
+  bool unregisterFunc(hipModule_t hmod);
   std::vector< std::pair<hipModule_t, bool> >* unregisterVar(hipModule_t hmod);
 
 
@@ -223,6 +225,8 @@ public:
   void registerVar(const void* hostvar, const DeviceVar& var);
   void registerFunction(const void* hostFunction, const DeviceFunction& func);
 
+  bool findModFunc(hipFunction_t* hfunc, hipModule_t hmod, const char* name);
+  bool createFunc(hipFunction_t* hfunc, hipModule_t hmod, const char* name);
   hipFunction_t getFunc(const void* hostFunction, int deviceId);
   bool getFuncAttr(const void* hostFunction, hipFuncAttributes* func_attr);
   bool getGlobalVar(const void* hostVar, int deviceId, hipModule_t hmod,

vdi/hip_module.cpp

@@ -92,6 +92,10 @@ hipError_t hipModuleUnload(hipModule_t hmod)
 
   amd::Program* program = as_amd(reinterpret_cast<cl_program>(hmod));
 
+  if(!PlatformState::instance().unregisterFunc(hmod)) {
+    HIP_RETURN(hipErrorInvalidSymbol);
+  }
+
   if(!ihipModuleUnregisterGlobal(hmod)) {
     HIP_RETURN(hipErrorInvalidSymbol);
   }
@@ -122,7 +126,7 @@ extern hipError_t __hipExtractCodeObjectFromFatBinary(const void* data,
                                                 const std::vector<const char*>& devices,
                                                 std::vector<std::pair<const void*, size_t>>& code_objs);
 
-bool ihipModuleRegisterUndefined(amd::Program* program, hipModule_t* module) {
+inline bool ihipModuleRegisterUndefined(amd::Program* program, hipModule_t* module) {
 
   std::vector<std::string> undef_vars;
   device::Program* dev_program
@@ -150,7 +154,34 @@ bool ihipModuleRegisterUndefined(amd::Program* program, hipModule_t* module) {
   return true;
 }
 
-bool ihipModuleRegisterGlobal(amd::Program* program, hipModule_t* module) {
+inline bool ihipModuleRegisterFunc(amd::Program* program, hipModule_t* module) {
+
+  std::vector<std::string> func_names;
+  device::Program* dev_program
+    = program->getDeviceProgram(*hip::getCurrentDevice()->devices()[0]);
+
+
+  // Get all the global func names from COMGR
+  if (!dev_program->getGlobalFuncFromCodeObj(&func_names)) {
+    return false;
+  }
+
+  for (auto it = func_names.begin(); it != func_names.end(); ++it) {
+    auto modules = new std::vector<std::pair<hipModule_t, bool> >(g_devices.size());
+    for (size_t dev = 0; dev < g_devices.size(); ++dev) {
+      modules->at(dev) = std::make_pair(*module, true);
+    }
+
+    // Create a new pointer, since the hostFunction* wont be available
+    // if device code not in the same file as host code.
+    PlatformState::DeviceFunction dfunc{*it, modules, std::vector<hipFunction_t>{ g_devices.size() }, true};
+    PlatformState::instance().registerFunction(new std::string(it->c_str()), dfunc);
+  }
+
+  return true;
+}
+
+inline bool ihipModuleRegisterGlobal(amd::Program* program, hipModule_t* module) {
 
   size_t var_size = 0;
   hipDeviceptr_t device_ptr = nullptr;
@@ -212,6 +243,10 @@ hipError_t ihipModuleLoadData(hipModule_t *module, const void *image)
     return hipErrorSharedObjectInitFailed;
   }
 
+  if (!ihipModuleRegisterFunc(program, module)) {
+    return hipErrorSharedObjectSymbolNotFound;
+  }
+
   return hipSuccess;
 }
 
@@ -219,21 +254,9 @@ hipError_t hipModuleGetFunction(hipFunction_t *hfunc, hipModule_t hmod, const ch
 {
   HIP_INIT_API(hipModuleGetFunction, hfunc, hmod, name);
 
-  amd::Program* program = as_amd(reinterpret_cast<cl_program>(hmod));
-
-  const amd::Symbol* symbol = program->findSymbol(name);
-  if (!symbol) {
+  if (!PlatformState::instance().findModFunc(hfunc, hmod, name)) {
     HIP_RETURN(hipErrorNotFound);
   }
-
-  amd::Kernel* kernel = new amd::Kernel(*program, *symbol, name);
-  if (!kernel) {
-    HIP_RETURN(hipErrorOutOfMemory);
-  }
-
-  hip::Function* f = new hip::Function(kernel);
-  *hfunc = f->asHipFunction();
-
   HIP_RETURN(hipSuccess);
 }
 

vdi/hip_platform.cpp

@@ -188,6 +188,30 @@ void PlatformState::init()
   }
 }
 
+bool PlatformState::unregisterFunc(hipModule_t hmod) {
+  amd::ScopedLock lock(lock_);
+  auto it = functions_.begin();
+  while (it  != functions_.end()) {
+    DeviceFunction& dfunc = it->second;
+    if ((*dfunc.modules)[0].first == hmod) {
+      if (dfunc.dyn_mod) {
+        std::string *s = reinterpret_cast<std::string*>(const_cast<void*>(it->first));
+        delete s;
+      }
+      for (size_t dev = 0; dev < g_devices.size(); ++dev) {
+        if (dfunc.functions[dev] != 0) {
+          hip::Function* f = reinterpret_cast<hip::Function*>(dfunc.functions[dev]);
+          delete f;
+        }
+      }
+      functions_.erase(it++);
+    } else {
+      ++it;
+    }
+  }
+  return true;
+}
+
 std::vector< std::pair<hipModule_t, bool> >* PlatformState::unregisterVar(hipModule_t hmod) {
   amd::ScopedLock lock(lock_);
   std::vector< std::pair<hipModule_t, bool> >* rmodules = nullptr;
@@ -292,6 +316,43 @@ bool CL_CALLBACK getSvarInfo(cl_program program, std::string var_name, void** va
                                                     var_addr, var_size);
 }
 
+bool PlatformState::findModFunc(hipFunction_t* hfunc, hipModule_t hmod, const char* name) {
+  amd::ScopedLock lock(lock_);
+  for (auto it = functions_.begin(); it != functions_.end(); ++it) {
+    PlatformState::DeviceFunction& devFunc = it->second;
+    if ((devFunc.deviceName == name) && (hmod == (*devFunc.modules)[ihipGetDevice()].first)) {
+      if (devFunc.functions[ihipGetDevice()] == 0) {
+        if(!createFunc(&devFunc.functions[ihipGetDevice()], hmod, name)) {
+          return false;
+        }
+      }
+      *hfunc = devFunc.functions[ihipGetDevice()];
+      return true;
+    }
+  }
+  return false;
+}
+
+bool PlatformState::createFunc(hipFunction_t* hfunc, hipModule_t hmod, const char* name) {
+  amd::Program* program = as_amd(reinterpret_cast<cl_program>(hmod));
+
+  const amd::Symbol* symbol = program->findSymbol(name);
+  if (!symbol) {
+    return false;
+  }
+
+  amd::Kernel* kernel = new amd::Kernel(*program, *symbol, name);
+  if (!kernel) {
+    return false;
+  }
+
+  hip::Function* f = new hip::Function(kernel);
+  *hfunc = f->asHipFunction();
+
+  return true;
+}
+
+
 hipFunction_t PlatformState::getFunc(const void* hostFunction, int deviceId) {
   amd::ScopedLock lock(lock_);
   const auto it = functions_.find(hostFunction);
@@ -308,7 +369,7 @@ hipFunction_t PlatformState::getFunc(const void* hostFunction, int deviceId) {
         (*devFunc.modules)[deviceId].second = true;
       }
       hipFunction_t function = nullptr;
-      if (hipSuccess == hipModuleGetFunction(&function, module, devFunc.deviceName.c_str()) &&
+      if (createFunc(&function, module, devFunc.deviceName.c_str()) &&
           function != nullptr) {
         devFunc.functions[deviceId] = function;
       }
@@ -435,7 +496,7 @@ extern "C" void __hipRegisterFunction(
   dim3*        gridDim,
   int*         wSize)
 {
-  PlatformState::DeviceFunction func{ std::string{deviceName}, modules, std::vector<hipFunction_t>{g_devices.size()}};
+  PlatformState::DeviceFunction func{ std::string{deviceName}, modules, std::vector<hipFunction_t>{g_devices.size()}, false};
   PlatformState::instance().registerFunction(hostFunction, func);
 //  for (size_t i = 0; i < g_devices.size(); ++i) {
 //    PlatformState::instance().getFunc(hostFunction, i);