[UR][L0] Set pointer kernel arguments only for queue's associated device (#20177)

In multi-device context, ensure that pointer kernel arguments are set
only for the device associated with the queue being used for kernel
launch. Previously, arguments were set for all devices in the kernel's
device map, which was unnecessary and potentially incorrect when
launching on a specific device.

Same will be done for v2 adapter here:
#20179

Author: againull

sycl/test-e2e/MultiDevice/set_arg_pointer.cpp

@@ -0,0 +1,63 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+// UNSUPPORTED: level_zero_v2_adapter
+// UNSUPPORTED-TRACKER: CMPLRLLVM-67039
+
+// Test that usm device pointer can be used in a kernel compiled for a context
+// with multiple devices.
+
+#include <iostream>
+#include <sycl/detail/core.hpp>
+#include <sycl/kernel_bundle.hpp>
+#include <sycl/platform.hpp>
+#include <sycl/usm.hpp>
+#include <vector>
+
+using namespace sycl;
+
+class AddIdxKernel;
+
+int main() {
+  sycl::platform plt;
+  std::vector<sycl::device> devices = plt.get_devices();
+  if (devices.size() < 2) {
+    std::cout << "Need at least 2 GPU devices for this test.\n";
+    return 0;
+  }
+
+  std::vector<sycl::device> ctx_devices{devices[0], devices[1]};
+  sycl::context ctx(ctx_devices);
+
+  constexpr size_t N = 16;
+  std::vector<std::vector<int>> results(ctx_devices.size(),
+                                        std::vector<int>(N, 0));
+
+  // Create a kernel bundle compiled for both devices in the context
+  auto kb = sycl::get_kernel_bundle<sycl::bundle_state::executable>(ctx);
+
+  // For each device, create a queue and run a kernel using device USM
+  for (size_t i = 0; i < ctx_devices.size(); ++i) {
+    sycl::queue q(ctx, ctx_devices[i]);
+    int *data = sycl::malloc_device<int>(N, q);
+    q.fill(data, 1, N).wait();
+    q.submit([&](sycl::handler &h) {
+       h.use_kernel_bundle(kb);
+       h.parallel_for<AddIdxKernel>(
+           sycl::range<1>(N), [=](sycl::id<1> idx) { data[idx] += idx[0]; });
+     }).wait();
+    q.memcpy(results[i].data(), data, N * sizeof(int)).wait();
+    sycl::free(data, q);
+  }
+
+  for (size_t i = 0; i < ctx_devices.size(); ++i) {
+    std::cout << "Device " << i << " results: ";
+    for (size_t j = 0; j < N; ++j) {
+      if (results[i][j] != 1 + static_cast<int>(j)) {
+        return -1;
+      }
+      std::cout << results[i][j] << " ";
+    }
+  }
+  return 0;
+}

unified-runtime/source/adapters/level_zero/command_buffer.cpp

@@ -1004,12 +1004,16 @@ ur_result_t setKernelPendingArguments(
     ze_kernel_handle_t ZeKernel) {
   // If there are any pending arguments set them now.
   for (auto &Arg : PendingArguments) {
-    // 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.
     char **ZeHandlePtr = nullptr;
-    if (Arg.Value) {
-      UR_CALL(Arg.Value->getZeHandlePtr(ZeHandlePtr, Arg.AccessMode, Device,
-                                        nullptr, 0u));
+    if (auto MemObjPtr = std::get_if<ur_mem_handle_t>(&Arg.Value)) {
+      ur_mem_handle_t MemObj = *MemObjPtr;
+      if (MemObj) {
+        UR_CALL(MemObj->getZeHandlePtr(ZeHandlePtr, Arg.AccessMode, Device,
+                                       nullptr, 0u));
+      }
+    } else {
+      auto Ptr = const_cast<void **>(&std::get<const void *>(Arg.Value));
+      ZeHandlePtr = reinterpret_cast<char **>(Ptr);
     }
     ZE2UR_CALL(zeKernelSetArgumentValue,
                (ZeKernel, Arg.Index, Arg.Size, ZeHandlePtr));

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/kernel.cpp

@@ -125,16 +125,22 @@ ur_result_t urEnqueueKernelLaunch(
 
   // If there are any pending arguments set them now.
   for (auto &Arg : Kernel->PendingArguments) {
-    // 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.
+    // The Arg.Value can be either a ur_mem_handle_t or a raw pointer
+    // (const void*). Resolve per-device: for mem handles obtain the device
+    // specific handle, otherwise pass the raw pointer value.
     char **ZeHandlePtr = nullptr;
-    if (Arg.Value) {
-      UR_CALL(Arg.Value->getZeHandlePtr(ZeHandlePtr, Arg.AccessMode,
-                                        Queue->Device, EventWaitList,
-                                        NumEventsInWaitList));
+    if (auto MemObjPtr = std::get_if<ur_mem_handle_t>(&Arg.Value)) {
+      ur_mem_handle_t MemObj = *MemObjPtr;
+      if (MemObj) {
+        UR_CALL(MemObj->getZeHandlePtr(ZeHandlePtr, Arg.AccessMode,
+                                       Queue->Device, EventWaitList,
+                                       NumEventsInWaitList));
+      }
+    } else {
+      auto Ptr = const_cast<void **>(&std::get<const void *>(Arg.Value));
+      ZeHandlePtr = reinterpret_cast<char **>(Ptr);
     }
-    ZE2UR_CALL(zeKernelSetArgumentValue,
-               (ZeKernel, Arg.Index, Arg.Size, ZeHandlePtr));
+    UR_CALL(setArgValueOnZeKernel(ZeKernel, Arg.Index, Arg.Size, ZeHandlePtr));
   }
   Kernel->PendingArguments.clear();
 
@@ -422,41 +428,21 @@ ur_result_t urKernelSetArgValue(
 
   UR_ASSERT(Kernel, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
 
-  // 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 > 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;
-    ZeResult = ZE_CALL_NOCHECK(zeKernelSetArgumentValue,
-                               (ZeKernel, ArgIndex, ArgSize, PArgValue));
+    UR_CALL(setArgValueOnZeKernel(ZeKernel, ArgIndex, ArgSize, PArgValue))
   } else {
     for (auto It : Kernel->ZeKernelMap) {
       auto ZeKernel = It.second;
-      ZeResult = ZE_CALL_NOCHECK(zeKernelSetArgumentValue,
-                                 (ZeKernel, ArgIndex, ArgSize, PArgValue));
+      UR_CALL(setArgValueOnZeKernel(ZeKernel, ArgIndex, ArgSize, PArgValue))
     }
   }
-
-  if (ZeResult == ZE_RESULT_ERROR_INVALID_ARGUMENT) {
-    return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE;
-  }
-
-  return ze2urResult(ZeResult);
+  return UR_RESULT_SUCCESS;
 }
 
 ur_result_t urKernelSetArgLocal(
@@ -732,6 +718,23 @@ ur_result_t urKernelSetArgPointer(
     /// [in][optional] SVM pointer to memory location holding the argument
     /// value. If null then argument value is considered null.
     const void *ArgValue) {
+  UR_ASSERT(Kernel, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
+  {
+    std::scoped_lock<ur_shared_mutex> Guard(Kernel->Mutex);
+    // In multi-device context instead of setting pointer arguments immediately
+    // across all device kernels, store them as pending so they can be resolved
+    // per-device at enqueue time. This ensures the correct handle is used for
+    // the device of the queue.
+    if (Kernel->Program->Context->getDevices().size() > 1) {
+      if (ArgIndex > Kernel->ZeKernelProperties->numKernelArgs - 1) {
+        return UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_INDEX;
+      }
+      Kernel->PendingArguments.push_back({ArgIndex, sizeof(const void *),
+                                          ArgValue, ur_mem_handle_t_::unknown});
+
+      return UR_RESULT_SUCCESS;
+    }
+  }
 
   // KernelSetArgValue is expecting a pointer to the argument
   UR_CALL(ur::level_zero::urKernelSetArgValue(

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

@@ -10,6 +10,7 @@
 #pragma once
 
 #include <unordered_set>
+#include <variant>
 
 #include "common.hpp"
 #include "common/ur_ref_count.hpp"
@@ -97,8 +98,10 @@ struct ur_kernel_handle_t_ : ur_object {
   struct ArgumentInfo {
     uint32_t Index;
     size_t Size;
-    // const ur_mem_handle_t_ *Value;
-    ur_mem_handle_t_ *Value;
+    // Value may be either a memory object or a raw pointer value (for pointer
+    // arguments). Resolve at enqueue time per-device to ensure correct handle
+    // is used for that device.
+    std::variant<ur_mem_handle_t, const void *> Value;
     ur_mem_handle_t_::access_mode_t AccessMode{ur_mem_handle_t_::unknown};
   };
   // Arguments that still need to be set (with zeKernelSetArgumentValue)