SWDEV-504084 - Make hipModuleGetFunction use the device the module is loaded on

If a module is loaded on one device, hipModuleGetFunction and other similar APIs should be able to run successfully from another device.

Change-Id: I96084cbd6c6dcf2a81019779a6ab1842ef2f35d1


[ROCm/clr commit: c46f843b99]

projects/clr/hipamd/src/hip_code_object.cpp

@@ -1068,6 +1068,8 @@ hipError_t DynCO::loadCodeObject(const char* fname, const void* image) {
   // No Lazy loading for DynCO
   IHIP_RETURN_ONFAIL(fb_info_->BuildProgram(ihipGetDevice()));
 
+  module_ = fb_info_->Module(device_id_);
+
   // Define Global variables
   IHIP_RETURN_ONFAIL(populateDynGlobalVars());
 
@@ -1114,23 +1116,19 @@ DynCO::~DynCO() {
 hipError_t DynCO::getDeviceVar(DeviceVar** dvar, std::string var_name) {
   amd::ScopedLock lock(dclock_);
 
-  CheckDeviceIdMatch();
-
   auto it = vars_.find(var_name);
   if (it == vars_.end()) {
     LogPrintfError("Cannot find the Var: %s ", var_name.c_str());
     return hipErrorNotFound;
   }
 
-  hipError_t err = it->second->getDeviceVar(dvar, device_id_, module());
+  hipError_t err = it->second->getDeviceVar(dvar, device_id_, module_);
   return err;
 }
 
 hipError_t DynCO::getDynFunc(hipFunction_t* hfunc, std::string func_name) {
   amd::ScopedLock lock(dclock_);
 
-  CheckDeviceIdMatch();
-
   if (hfunc == nullptr) {
     return hipErrorInvalidValue;
   }
@@ -1142,7 +1140,7 @@ hipError_t DynCO::getDynFunc(hipFunction_t* hfunc, std::string func_name) {
   }
 
   /* See if this could be solved */
-  return it->second->getDynFunc(hfunc, module());
+  return it->second->getDynFunc(hfunc, module_);
 }
 
 bool DynCO::isValidDynFunc(const void* hfunc) {
@@ -1214,7 +1212,7 @@ hipError_t DynCO::populateDynGlobalVars() {
                                      ->getDeviceProgram(*hip::getCurrentDevice()->devices()[0]);
 
   if (!dev_program->getGlobalVarFromCodeObj(&var_names)) {
-    LogPrintfError("Could not get Global vars from Code Obj for Module: 0x%x", module());
+    LogPrintfError("Could not get Global vars from Code Obj for Module: 0x%x", module_);
     return hipErrorSharedObjectSymbolNotFound;
   }
 
@@ -1242,7 +1240,7 @@ hipError_t DynCO::populateDynGlobalFuncs() {
 
   // Get all the global func names from COMGR
   if (!dev_program->getGlobalFuncFromCodeObj(&func_names)) {
-    LogPrintfError("Could not get Global Funcs from Code Obj for Module: 0x%x", module());
+    LogPrintfError("Could not get Global Funcs from Code Obj for Module: 0x%x", module_);
     return hipErrorSharedObjectSymbolNotFound;
   }
 
@@ -1316,7 +1314,7 @@ hipError_t StatCO::removeFatBinary(FatBinaryInfo** module) {
       hipError_t err;
       for (auto dev : g_devices) {
         DeviceVar* dvar = nullptr;
-        IHIP_RETURN_ONFAIL((*it)->getDeviceVarPtr(&dvar, dev->deviceId())); 
+        IHIP_RETURN_ONFAIL((*it)->getDeviceVarPtr(&dvar, dev->deviceId()));
         if (dvar != nullptr) {
           // free also deletes the device ptr
           err = ihipFree(dvar->device_ptr());

projects/clr/hipamd/src/hip_code_object.hpp

@@ -113,12 +113,12 @@ class DynCO : public CodeObject {
   amd::Monitor dclock_{true};
 
 public:
-  DynCO() : device_id_(ihipGetDevice()), fb_info_(nullptr) {}
+  DynCO() : device_id_(ihipGetDevice()), fb_info_(nullptr), module_(nullptr) {}
   virtual ~DynCO();
 
   //LoadsCodeObject and its data
   hipError_t loadCodeObject(const char* fname, const void* image=nullptr);
-  hipModule_t module() const { return fb_info_->Module(ihipGetDevice()); };
+  hipModule_t getModule() const { return module_; };
 
   //Gets GlobalVar/Functions from a dynamically loaded code object
   hipError_t getDynFunc(hipFunction_t* hfunc, std::string func_name);
@@ -137,15 +137,11 @@ public:
     }
     return hipSuccess;
   }
-  // Device ID Check to check if module is launched in the same device it was loaded.
-  inline void CheckDeviceIdMatch() const {
-    guarantee(device_id_ == ihipGetDevice(), "Device mismatch from where this module is loaded,"
-              "device_id: %d ihipGetDevice:%d", device_id_, ihipGetDevice());
-  }
 
 private:
   int device_id_;
   FatBinaryInfo* fb_info_;
+  hipModule_t module_;
 
   //Maps for vars/funcs, could be keyed in with std::string name
   std::unordered_map<std::string, Function*> functions_;

projects/clr/hipamd/src/hip_module.cpp

@@ -643,6 +643,22 @@ hipError_t ihipModuleLaunchCooperativeKernelMultiDevice(hipFunctionLaunchParams*
     }
   }
 
+  // Grid and Block dimensions should match across devices, as well as sharedMemBytes
+  for (uint32_t i = 1; i < numDevices; ++i) {
+    if (launchParamsList[i - 1].gridDimX != launchParamsList[i].gridDimX ||
+        launchParamsList[i - 1].gridDimY != launchParamsList[i].gridDimY ||
+        launchParamsList[i - 1].gridDimZ != launchParamsList[i].gridDimZ ||
+        launchParamsList[i - 1].blockDimX != launchParamsList[i].blockDimX ||
+        launchParamsList[i - 1].blockDimY != launchParamsList[i].blockDimY ||
+        launchParamsList[i - 1].blockDimZ != launchParamsList[i].blockDimZ) {
+      return hipErrorInvalidValue;
+    }
+
+    if (launchParamsList[i - 1].sharedMemBytes != launchParamsList[i].sharedMemBytes) {
+      return hipErrorInvalidValue;
+    }
+  }
+
   for (int i = 0; i < numDevices; ++i) {
     const hipFunctionLaunchParams& launch = launchParamsList[i];
     hip::Stream* hip_stream = reinterpret_cast<hip::Stream*>(launch.hStream);

projects/clr/hipamd/src/hip_platform.cpp

@@ -760,7 +760,7 @@ hipError_t PlatformState::loadModule(hipModule_t* module, const char* fname, con
     return hip_error;
   }
 
-  *module = dynCo->module();
+  *module = dynCo->getModule();
   assert(*module != nullptr);
 
   amd::ScopedLock lock(lock_);