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

Author: againull

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

@@ -305,6 +305,14 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
   }
   pending_allocations.clear();
 
+  // Apply any pending raw pointer arguments (USM pointers) for this device.
+  for (auto &pending : pending_pointer_args) {
+    void *Ptr = const_cast<void *>(pending.ptrArgValue);
+    ZE2UR_CALL(zeKernelSetArgumentValue,
+               (hZeKernel, pending.argIndex, sizeof(void *), &Ptr));
+  }
+  pending_pointer_args.clear();
+
   return UR_RESULT_SUCCESS;
 }
 
@@ -319,6 +327,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,
@@ -405,14 +425,17 @@ ur_result_t urKernelSetArgPointer(
     ur_kernel_handle_t hKernel, ///< [in] handle of the kernel object
     uint32_t argIndex, ///< [in] argument index in range [0, num args - 1]
     const ur_kernel_arg_pointer_properties_t
-        *pProperties, ///< [in][optional] argument properties
+        *, ///< [in][optional] argument properties
     const void
         *pArgValue ///< [in] argument value represented as matching arg type.
     ) try {
   TRACK_SCOPE_LATENCY("urKernelSetArgPointer");
 
   std::scoped_lock<ur_shared_mutex> guard(hKernel->Mutex);
-  return hKernel->setArgPointer(argIndex, pProperties, pArgValue);
+  // Store the raw pointer value and defer setting the argument until
+  // we know the device where kernel is being submitted.
+  hKernel->addPendingPointerArgument(argIndex, pArgValue);
+  return UR_RESULT_SUCCESS;
 } catch (...) {
   return exceptionToResult(std::current_exception());
 }

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

@@ -82,6 +82,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,
@@ -115,6 +119,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