Merge pull request #2106 from winstonzhang-intel/URLZA-391

[L0] Fix urEnqueueDeviceGlobalVariableRead to use device specific modules
oneapi-src · Sep 20, 2024 · cbb1935 · cbb1935
1 parent 672ac18
commit cbb1935
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/source/adapters/level_zero/kernel.cpp b/source/adapters/level_zero/kernel.cpp
@@ -661,11 +661,19 @@ ur_result_t ur_queue_handle_legacy_t_::enqueueDeviceGlobalVariableRead(
 
   std::scoped_lock<ur_shared_mutex> lock(Queue->Mutex);
 
+  ze_module_handle_t ZeModule{};
+  auto It = Program->ZeModuleMap.find(Queue->Device->ZeDevice);
+  if (It != Program->ZeModuleMap.end()) {
+    ZeModule = It->second;
+  } else {
+    ZeModule = Program->ZeModule;
+  }
+
   // Find global variable pointer
   size_t GlobalVarSize = 0;
   void *GlobalVarPtr = nullptr;
   ZE2UR_CALL(zeModuleGetGlobalPointer,
-             (Program->ZeModule, Name, &GlobalVarSize, &GlobalVarPtr));
+             (ZeModule, Name, &GlobalVarSize, &GlobalVarPtr));
   if (GlobalVarSize < Offset + Count) {
     setErrorMessage("Read from device global variable is out of range.",
                     UR_RESULT_ERROR_INVALID_VALUE,