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,