Merge pull request #2190 from igchor/global_offset_fix

[L0] fix out-of-bounds access in setKernelGlobalOffset

Author: pbalcer

source/adapters/level_zero/command_buffer.cpp

@@ -627,32 +627,6 @@ urCommandBufferFinalizeExp(ur_exp_command_buffer_handle_t CommandBuffer) {
   return UR_RESULT_SUCCESS;
 }
 
-/**
- * Sets the global offset for a kernel command that will be appended to the
- * command buffer.
- * @param[in] CommandBuffer The CommandBuffer where the command will be
- * appended.
- * @param[in] Kernel The handle to the kernel that will be appended.
- * @param[in] GlobalWorkOffset The global offset value.
- * @return UR_RESULT_SUCCESS or an error code on failure
- */
-ur_result_t setKernelGlobalOffset(ur_exp_command_buffer_handle_t CommandBuffer,
-                                  ur_kernel_handle_t Kernel,
-                                  const size_t *GlobalWorkOffset) {
-
-  if (!CommandBuffer->Context->getPlatform()
-           ->ZeDriverGlobalOffsetExtensionFound) {
-    logger::debug("No global offset extension found on this driver");
-    return UR_RESULT_ERROR_INVALID_VALUE;
-  }
-
-  ZE2UR_CALL(zeKernelSetGlobalOffsetExp,
-             (Kernel->ZeKernel, GlobalWorkOffset[0], GlobalWorkOffset[1],
-              GlobalWorkOffset[2]));
-
-  return UR_RESULT_SUCCESS;
-}
-
 /**
  * Sets the kernel arguments for a kernel command that will be appended to the
  * command buffer.
@@ -754,7 +728,8 @@ ur_result_t urCommandBufferAppendKernelLaunchExp(
       Kernel->Mutex, Kernel->Program->Mutex, CommandBuffer->Mutex);
 
   if (GlobalWorkOffset != NULL) {
-    UR_CALL(setKernelGlobalOffset(CommandBuffer, Kernel, GlobalWorkOffset));
+    UR_CALL(setKernelGlobalOffset(CommandBuffer->Context, Kernel->ZeKernel,
+                                  WorkDim, GlobalWorkOffset));
   }
 
   // If there are any pending arguments set them now.

source/adapters/level_zero/helpers/kernel_helpers.cpp

@@ -67,16 +67,17 @@ ur_result_t getSuggestedLocalWorkSize(ur_device_handle_t hDevice,
 }
 
 ur_result_t setKernelGlobalOffset(ur_context_handle_t Context,
-                                  ze_kernel_handle_t Kernel,
+                                  ze_kernel_handle_t Kernel, uint32_t WorkDim,
                                   const size_t *GlobalWorkOffset) {
   if (!Context->getPlatform()->ZeDriverGlobalOffsetExtensionFound) {
     logger::debug("No global offset extension found on this driver");
     return UR_RESULT_ERROR_INVALID_VALUE;
   }
 
-  ZE2UR_CALL(
-      zeKernelSetGlobalOffsetExp,
-      (Kernel, GlobalWorkOffset[0], GlobalWorkOffset[1], GlobalWorkOffset[2]));
+  auto OffsetX = GlobalWorkOffset[0];
+  auto OffsetY = WorkDim > 1 ? GlobalWorkOffset[1] : 0;
+  auto OffsetZ = WorkDim > 2 ? GlobalWorkOffset[2] : 0;
+  ZE2UR_CALL(zeKernelSetGlobalOffsetExp, (Kernel, OffsetX, OffsetY, OffsetZ));
 
   return UR_RESULT_SUCCESS;
 }

source/adapters/level_zero/helpers/kernel_helpers.hpp

@@ -36,11 +36,12 @@ ur_result_t calculateKernelWorkDimensions(
  * command buffer.
  * @param[in] Context Context associated with the queue.
  * @param[in] Kernel The handle to the kernel that will be appended.
- * @param[in] GlobalWorkOffset The global offset value.
+ * @param[in] WorkDim The number of work dimensions.
+ * @param[in] GlobalWorkOffset Array of size WorkDim.
  * @return UR_RESULT_SUCCESS or an error code on failure
  */
 ur_result_t setKernelGlobalOffset(ur_context_handle_t Context,
-                                  ze_kernel_handle_t Kernel,
+                                  ze_kernel_handle_t Kernel, uint32_t WorkDim,
                                   const size_t *GlobalWorkOffset);
 
 /**

source/adapters/level_zero/kernel.cpp

@@ -96,14 +96,8 @@ ur_result_t urEnqueueKernelLaunch(
   std::scoped_lock<ur_shared_mutex, ur_shared_mutex, ur_shared_mutex> Lock(
       Queue->Mutex, Kernel->Mutex, Kernel->Program->Mutex);
   if (GlobalWorkOffset != NULL) {
-    if (!Queue->Device->Platform->ZeDriverGlobalOffsetExtensionFound) {
-      logger::error("No global offset extension found on this driver");
-      return UR_RESULT_ERROR_INVALID_VALUE;
-    }
-
-    ZE2UR_CALL(zeKernelSetGlobalOffsetExp,
-               (ZeKernel, GlobalWorkOffset[0], GlobalWorkOffset[1],
-                GlobalWorkOffset[2]));
+    UR_CALL(setKernelGlobalOffset(Queue->Context, ZeKernel, WorkDim,
+                                  GlobalWorkOffset));
   }
 
   // If there are any pending arguments set them now.
@@ -257,14 +251,8 @@ ur_result_t urEnqueueCooperativeKernelLaunchExp(
   std::scoped_lock<ur_shared_mutex, ur_shared_mutex, ur_shared_mutex> Lock(
       Queue->Mutex, Kernel->Mutex, Kernel->Program->Mutex);
   if (GlobalWorkOffset != NULL) {
-    if (!Queue->Device->Platform->ZeDriverGlobalOffsetExtensionFound) {
-      logger::error("No global offset extension found on this driver");
-      return UR_RESULT_ERROR_INVALID_VALUE;
-    }
-
-    ZE2UR_CALL(zeKernelSetGlobalOffsetExp,
-               (ZeKernel, GlobalWorkOffset[0], GlobalWorkOffset[1],
-                GlobalWorkOffset[2]));
+    UR_CALL(setKernelGlobalOffset(Queue->Context, ZeKernel, WorkDim,
+                                  GlobalWorkOffset));
   }
 
   // If there are any pending arguments set them now.

source/adapters/level_zero/v2/queue_immediate_in_order.cpp

@@ -252,7 +252,8 @@ ur_result_t ur_queue_immediate_in_order_t::enqueueKernelLaunch(
       hKernel->Mutex, hKernel->getProgramHandle()->Mutex, this->Mutex);
 
   if (pGlobalWorkOffset != NULL) {
-    UR_CALL(setKernelGlobalOffset(hContext, hZeKernel, pGlobalWorkOffset));
+    UR_CALL(
+        setKernelGlobalOffset(hContext, hZeKernel, workDim, pGlobalWorkOffset));
   }
 
   ze_group_count_t zeThreadGroupDimensions{1, 1, 1};