[UR][L0 v2] Set pointer kernel arguments only for queue's associated device

unified-runtime/source/adapters/level_zero/helpers/kernel_helpers.cpp

@@ -156,3 +156,25 @@ ur_result_t calculateKernelWorkDimensions(
 
   return UR_RESULT_SUCCESS;
 }
+
+ur_result_t setArgValueOnZeKernel(ze_kernel_handle_t hZeKernel,
+                                  uint32_t argIndex, size_t argSize,
+                                  const void *pArgValue) {
+  // OpenCL: "the arg_value pointer can be NULL or point to a NULL value
+  // in which case a NULL value will be used as the value for the argument
+  // declared as a pointer to global or constant memory in the kernel"
+  //
+  // We don't know the type of the argument but it seems that the only time
+  // SYCL RT would send a pointer to NULL in 'arg_value' is when the argument
+  // is a NULL pointer. Treat a pointer to NULL in 'arg_value' as a NULL.
+  if (argSize == sizeof(void *) && pArgValue &&
+      *(void **)(const_cast<void *>(pArgValue)) == nullptr) {
+    pArgValue = nullptr;
+  }
+
+  ze_result_t ZeResult = ZE_CALL_NOCHECK(
+      zeKernelSetArgumentValue, (hZeKernel, argIndex, argSize, pArgValue));
+  if (ZeResult == ZE_RESULT_ERROR_INVALID_ARGUMENT)
+    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE;
+  return ze2urResult(ZeResult);
+}

unified-runtime/source/adapters/level_zero/helpers/kernel_helpers.hpp

@@ -71,3 +71,15 @@ inline void postSubmit(ze_kernel_handle_t hZeKernel,
     zeKernelSetGlobalOffsetExp(hZeKernel, 0, 0, 0);
   }
 }
+
+/**
+ * Helper to set kernel argument for ze_kernel_handle_t.
+ * @param[in] hZeKernel The handle to the Level-Zero kernel.
+ * @param[in] argIndex The index of the argument to set.
+ * @param[in] argSize The size of the argument to set.
+ * @param[in] pArgValue The pointer to the argument value.
+ * @return UR_RESULT_SUCCESS or an error code on failure
+ */
+ur_result_t setArgValueOnZeKernel(ze_kernel_handle_t hZeKernel,
+                                  uint32_t argIndex, size_t argSize,
+                                  const void *pArgValue);

unified-runtime/source/adapters/level_zero/v2/kernel.cpp

@@ -33,6 +33,17 @@ ur_single_device_kernel_t::ur_single_device_kernel_t(ur_device_handle_t hDevice,
       };
 }
 
+ur_result_t ur_single_device_kernel_t::setArgValue(uint32_t argIndex,
+                                                   size_t argSize,
+                                                   const void *pArgValue) {
+  return setArgValueOnZeKernel(hKernel.get(), argIndex, argSize, pArgValue);
+}
+
+ur_result_t ur_single_device_kernel_t::setArgPointer(uint32_t argIndex,
+                                                     const void *pArgValue) {
+  return setArgValue(argIndex, sizeof(void *), &pArgValue);
+}
+
 ur_result_t ur_single_device_kernel_t::release() {
   hKernel.reset();
   return UR_RESULT_SUCCESS;
@@ -184,19 +195,6 @@ ur_result_t ur_kernel_handle_t_::setArgValue(
     uint32_t argIndex, size_t argSize,
     const ur_kernel_arg_value_properties_t * /*pProperties*/,
     const void *pArgValue) {
-
-  // OpenCL: "the arg_value pointer can be NULL or point to a NULL value
-  // in which case a NULL value will be used as the value for the argument
-  // declared as a pointer to global or constant memory in the kernel"
-  //
-  // We don't know the type of the argument but it seems that the only time
-  // SYCL RT would send a pointer to NULL in 'arg_value' is when the argument
-  // is a NULL pointer. Treat a pointer to NULL in 'arg_value' as a NULL.
-  if (argSize == sizeof(void *) && pArgValue &&
-      *(void **)(const_cast<void *>(pArgValue)) == nullptr) {
-    pArgValue = nullptr;
-  }
-
   if (argIndex > zeCommonProperties->numKernelArgs - 1) {
     return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
   }
@@ -206,15 +204,8 @@ ur_result_t ur_kernel_handle_t_::setArgValue(
       continue;
     }
 
-    auto zeResult = ZE_CALL_NOCHECK(zeKernelSetArgumentValue,
-                                    (singleDeviceKernel.value().hKernel.get(),
-                                     argIndex, argSize, pArgValue));
-
-    if (zeResult == ZE_RESULT_ERROR_INVALID_ARGUMENT) {
-      return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE;
-    } else if (zeResult != ZE_RESULT_SUCCESS) {
-      return ze2urResult(zeResult);
-    }
+    UR_CALL(setArgValueOnZeKernel(singleDeviceKernel.value().hKernel.get(),
+                                  argIndex, argSize, pArgValue));
   }
   return UR_RESULT_SUCCESS;
 }
@@ -278,7 +269,11 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
     const size_t *pGlobalWorkOffset, uint32_t workDim, uint32_t groupSizeX,
     uint32_t groupSizeY, uint32_t groupSizeZ,
     ze_command_list_handle_t commandList, wait_list_view &waitListView) {
-  auto hZeKernel = getZeHandle(hDevice);
+  auto &deviceKernelOpt = deviceKernels[deviceIndex(hDevice)];
+  if (!deviceKernelOpt.has_value())
+    return UR_RESULT_ERROR_INVALID_KERNEL;
+  auto &deviceKernel = deviceKernelOpt.value();
+  auto hZeKernel = deviceKernel.hKernel.get();
 
   if (pGlobalWorkOffset != NULL) {
     UR_CALL(
@@ -301,10 +296,17 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
         zePtr = reinterpret_cast<void *>(hImage->getZeImage());
       }
     }
-    UR_CALL(setArgPointer(pending.argIndex, nullptr, zePtr));
+    // Set the argument only on this device's kernel.
+    UR_CALL(deviceKernel.setArgPointer(pending.argIndex, zePtr));
   }
   pending_allocations.clear();
 
+  // Apply any pending raw pointer arguments (USM pointers) for this device.
+  for (auto &pending : pending_pointer_args) {
+    UR_CALL(deviceKernel.setArgPointer(pending.argIndex, pending.ptrArgValue));
+  }
+  pending_pointer_args.clear();
+
   return UR_RESULT_SUCCESS;
 }
 
@@ -319,6 +321,18 @@ ur_result_t ur_kernel_handle_t_::addPendingMemoryAllocation(
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t
+ur_kernel_handle_t_::addPendingPointerArgument(uint32_t argIndex,
+                                               const void *pArgValue) {
+  if (argIndex > zeCommonProperties->numKernelArgs - 1) {
+    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+  }
+
+  pending_pointer_args.push_back({argIndex, pArgValue});
+
+  return UR_RESULT_SUCCESS;
+}
+
 std::vector<char> ur_kernel_handle_t_::getSourceAttributes() const {
   uint32_t size;
   ZE2UR_CALL_THROWS(zeKernelGetSourceAttributes,
@@ -412,6 +426,13 @@ ur_result_t urKernelSetArgPointer(
   TRACK_SCOPE_LATENCY("urKernelSetArgPointer");
 
   std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
+  // In multi-device context store the raw pointer value and defer setting the
+  // argument until we know the device where kernel is being submitted.
+  if (hKernel->getContext()->getDevices().size() > 1) {
+    hKernel->addPendingPointerArgument(argIndex, pArgValue);
+    return UR_RESULT_SUCCESS;
+  }
+
   return hKernel->setArgPointer(argIndex, pProperties, pArgValue);
 } catch (...) {
   return exceptionToResult(std::current_exception());

unified-runtime/source/adapters/level_zero/v2/kernel.hpp

@@ -21,6 +21,14 @@ struct ur_single_device_kernel_t {
                             ze_kernel_handle_t hKernel, bool ownZeHandle);
   ur_result_t release();
 
+  // Set argument value on this device's ze kernel only.
+  ur_result_t setArgValue(uint32_t argIndex, size_t argSize,
+                          const void *pArgValue);
+
+  // Convenience for pointer args: sets a pointer-sized argument on this
+  // device's ze kernel.
+  ur_result_t setArgPointer(uint32_t argIndex, const void *pArgValue);
+
   ur_device_handle_t hDevice;
   v2::raii::ze_kernel_handle_t hKernel;
   mutable ZeCache<ZeStruct<ze_kernel_properties_t>> zeKernelProperties;
@@ -82,6 +90,10 @@ struct ur_kernel_handle_t_ : ur_object {
   ur_result_t
   addPendingMemoryAllocation(pending_memory_allocation_t allocation);
 
+  // Add a pending pointer argument for which device is not yet known.
+  ur_result_t addPendingPointerArgument(uint32_t argIndex,
+                                        const void *pArgValue);
+
   // Set all required values for the kernel before submission (including pending
   // memory allocations).
   ur_result_t prepareForSubmission(ur_context_handle_t hContext,
@@ -92,6 +104,9 @@ struct ur_kernel_handle_t_ : ur_object {
                                    ze_command_list_handle_t cmdList,
                                    wait_list_view &waitListView);
 
+  // Get context of the kernel.
+  ur_context_handle_t getContext() const { return hProgram->Context; }
+
   ur::RefCount RefCount;
 
 private:
@@ -115,6 +130,14 @@ struct ur_kernel_handle_t_ : ur_object {
 
   std::vector<pending_memory_allocation_t> pending_allocations;
 
+  struct pending_pointer_arg_t {
+    uint32_t argIndex;
+    const void *ptrArgValue;
+  };
+
+  // Pointer arguments that need to be applied per-device at submission time.
+  std::vector<pending_pointer_arg_t> pending_pointer_args;
+
   void completeInitialization();
 
   // pointer to any non-null kernel in deviceKernels