Merge pull request #1810 from nrspruit/fix_l0_kernel_cts

[L0] Fix kernel error handling and enumeration checking

Author: Omar Ahmed

source/adapters/level_zero/kernel.cpp

@@ -710,7 +710,14 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelCreate(
     ZeKernelDesc.pKernelName = KernelName;
 
     ze_kernel_handle_t ZeKernel;
-    ZE2UR_CALL(zeKernelCreate, (ZeModule, &ZeKernelDesc, &ZeKernel));
+    auto ZeResult =
+        ZE_CALL_NOCHECK(zeKernelCreate, (ZeModule, &ZeKernelDesc, &ZeKernel));
+    // Gracefully handle the case that kernel create fails.
+    if (ZeResult != ZE_RESULT_SUCCESS) {
+      delete *RetKernel;
+      *RetKernel = nullptr;
+      return ze2urResult(ZeResult);
+    }
 
     auto ZeDevice = It.first;
 
@@ -764,20 +771,29 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgValue(
     PArgValue = nullptr;
   }
 
+  if (ArgIndex > Kernel->ZeKernelProperties->numKernelArgs - 1) {
+    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+  }
+
   std::scoped_lock<ur_shared_mutex> Guard(Kernel->Mutex);
+  ze_result_t ZeResult = ZE_RESULT_SUCCESS;
   if (Kernel->ZeKernelMap.empty()) {
     auto ZeKernel = Kernel->ZeKernel;
-    ZE2UR_CALL(zeKernelSetArgumentValue,
-               (ZeKernel, ArgIndex, ArgSize, PArgValue));
+    ZeResult = ZE_CALL_NOCHECK(zeKernelSetArgumentValue,
+                               (ZeKernel, ArgIndex, ArgSize, PArgValue));
   } else {
     for (auto It : Kernel->ZeKernelMap) {
       auto ZeKernel = It.second;
-      ZE2UR_CALL(zeKernelSetArgumentValue,
-                 (ZeKernel, ArgIndex, ArgSize, PArgValue));
+      ZeResult = ZE_CALL_NOCHECK(zeKernelSetArgumentValue,
+                                 (ZeKernel, ArgIndex, ArgSize, PArgValue));
     }
   }
 
-  return UR_RESULT_SUCCESS;
+  if (ZeResult == ZE_RESULT_ERROR_INVALID_ARGUMENT) {
+    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE;
+  }
+
+  return ze2urResult(ZeResult);
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgLocal(
@@ -826,6 +842,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(
     } catch (...) {
       return UR_RESULT_ERROR_UNKNOWN;
     }
+  case UR_KERNEL_INFO_NUM_REGS:
   case UR_KERNEL_INFO_NUM_ARGS:
     return ReturnValue(uint32_t{Kernel->ZeKernelProperties->numKernelArgs});
   case UR_KERNEL_INFO_REFERENCE_COUNT:
@@ -1076,6 +1093,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgSampler(
 ) {
   std::ignore = Properties;
   std::scoped_lock<ur_shared_mutex> Guard(Kernel->Mutex);
+  if (ArgIndex > Kernel->ZeKernelProperties->numKernelArgs - 1) {
+    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+  }
   ZE2UR_CALL(zeKernelSetArgumentValue, (Kernel->ZeKernel, ArgIndex,
                                         sizeof(void *), &ArgValue->ZeSampler));
 
@@ -1095,6 +1115,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgMemObj(
   // The ArgValue may be a NULL pointer in which case a NULL value is used for
   // the kernel argument declared as a pointer to global or constant memory.
 
+  if (ArgIndex > Kernel->ZeKernelProperties->numKernelArgs - 1) {
+    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+  }
+
   ur_mem_handle_t_ *UrMem = ur_cast<ur_mem_handle_t_ *>(ArgValue);
 
   ur_mem_handle_t_::access_mode_t UrAccessMode = ur_mem_handle_t_::read_write;

test/conformance/kernel/kernel_adapter_level_zero-v2.match

@@ -1,17 +1,10 @@
-urKernelGetInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_KERNEL_INFO_NUM_REGS
-urKernelSetArgLocalTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgMemObjTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetArgPointerTest.SuccessHost/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetArgPointerTest.SuccessDevice/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetArgPointerTest.SuccessShared/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgPointerNegativeTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgSamplerTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgValueTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgValueTest.InvalidKernelArgumentSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoTest.SuccessIndirectAccess/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoUSMPointersTest.SuccessHost/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoUSMPointersTest.SuccessDevice/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoUSMPointersTest.SuccessShared/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelGetSuggestedLocalWorkSizeTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelGetSuggestedLocalWorkSizeTest.Success2D/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelGetSuggestedLocalWorkSizeTest.Success3D/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
+urKernelGetSuggestedLocalWorkSizeTest.Success3D/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_

test/conformance/kernel/kernel_adapter_level_zero.match

@@ -1,10 +1,3 @@
-urKernelGetInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_KERNEL_INFO_NUM_REGS
-urKernelSetArgLocalTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgMemObjTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgPointerNegativeTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgSamplerTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgValueTest.InvalidKernelArgumentIndex/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-urKernelSetArgValueTest.InvalidKernelArgumentSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoTest.SuccessIndirectAccess/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoUSMPointersTest.SuccessHost/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urKernelSetExecInfoUSMPointersTest.SuccessDevice/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_