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

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.

Author: againull

sycl/test-e2e/MultiDevice/set_arg_pointer.cpp

@@ -0,0 +1,63 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+// UNSUPPORTED-TRACKER: CMPLRLLVM-67039
+// UNSUPPORTED: level_zero_v2_adapter
+
+// 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/kernel.cpp

@@ -125,13 +125,20 @@ 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));
@@ -733,9 +740,13 @@ ur_result_t urKernelSetArgPointer(
     /// value. If null then argument value is considered null.
     const void *ArgValue) {
 
-  // KernelSetArgValue is expecting a pointer to the argument
-  UR_CALL(ur::level_zero::urKernelSetArgValue(
-      Kernel, ArgIndex, sizeof(const void *), nullptr, &ArgValue));
+  // 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.
+  std::scoped_lock<ur_shared_mutex> Guard(Kernel->Mutex);
+  Kernel->PendingArguments.push_back(
+      {ArgIndex, sizeof(const void *), ArgValue, ur_mem_handle_t_::unknown});
   return UR_RESULT_SUCCESS;
 }
 
@@ -842,9 +853,8 @@ ur_result_t urKernelSetArgMemObj(
       return UR_RESULT_ERROR_INVALID_ARGUMENT;
     }
   }
-  auto Arg = UrMem ? UrMem : nullptr;
   Kernel->PendingArguments.push_back(
-      {ArgIndex, sizeof(void *), Arg, UrAccessMode});
+      {ArgIndex, sizeof(const void *), UrMem, UrAccessMode});
 
   return UR_RESULT_SUCCESS;
 }

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)