Adjust implementation of functions to new ur_program structure

Author: againull

source/adapters/level_zero/kernel.cpp

@@ -493,18 +493,11 @@ ur_result_t urEnqueueDeviceGlobalVariableWrite(
                ///< this particular kernel execution instance.
 ) {
   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;
+  ze_module_handle_t ZeModule =
+      Program->getZeModuleHandle(Queue->Device->ZeDevice);
   ZE2UR_CALL(zeModuleGetGlobalPointer,
              (ZeModule, Name, &GlobalVarSize, &GlobalVarPtr));
   if (GlobalVarSize < Offset + Count) {
@@ -555,15 +548,8 @@ ur_result_t urEnqueueDeviceGlobalVariableRead(
                ///< this particular kernel execution instance.
 ) {
   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;
-  }
-
+  ze_module_handle_t ZeModule =
+      Program->getZeModuleHandle(Queue->Device->ZeDevice);
   // Find global variable pointer
   size_t GlobalVarSize = 0;
   void *GlobalVarPtr = nullptr;
@@ -601,10 +587,6 @@ ur_result_t urKernelCreate(
         *RetKernel ///< [out] pointer to handle of kernel object created.
 ) {
   std::shared_lock<ur_shared_mutex> Guard(Program->Mutex);
-  if (Program->State != ur_program_handle_t_::state::Exe) {
-    return UR_RESULT_ERROR_INVALID_PROGRAM_EXECUTABLE;
-  }
-
   try {
     ur_kernel_handle_t_ *UrKernel = new ur_kernel_handle_t_(true, Program);
     *RetKernel = reinterpret_cast<ur_kernel_handle_t>(UrKernel);
@@ -614,8 +596,13 @@ ur_result_t urKernelCreate(
     return UR_RESULT_ERROR_UNKNOWN;
   }
 
-  for (auto It : Program->ZeModuleMap) {
-    auto ZeModule = It.second;
+  for (auto &Dev : Program->AssociatedDevices) {
+    auto ZeDevice = Dev->ZeDevice;
+    // Program may be associated with all devices from the context but built only for subset of devices.
+    if (Program->getState(ZeDevice) != ur_program_handle_t_::state::Exe)
+      continue;
+
+    auto ZeModule = Program->getZeModuleHandle(ZeDevice);
     ZeStruct<ze_kernel_desc_t> ZeKernelDesc;
     ZeKernelDesc.flags = 0;
     ZeKernelDesc.pKernelName = KernelName;
@@ -630,8 +617,6 @@ ur_result_t urKernelCreate(
       return ze2urResult(ZeResult);
     }
 
-    auto ZeDevice = It.first;
-
     // Store the kernel in the ZeKernelMap so the correct
     // kernel can be retrieved later for a specific device
     // where a queue is being submitted.
@@ -648,7 +633,10 @@ ur_result_t urKernelCreate(
     for (auto ZeSubDevice : ZeSubDevices) {
       (*RetKernel)->ZeKernelMap[ZeSubDevice] = ZeKernel;
     }
-  }
+  } 
+  // There is no any successfully built executable for program.
+  if ((*RetKernel)->ZeKernelMap.empty())
+    return UR_RESULT_ERROR_INVALID_PROGRAM_EXECUTABLE;
 
   (*RetKernel)->ZeKernel = (*RetKernel)->ZeKernelMap.begin()->second;