[SYCL][UR][L0 v2] Fix urMemBufferCreateWithNativeHandle (#17698)

The implementation incorrectly assumed that hDevice is always set for
memory type other than host which resulted in nullptr dereference in
ur_discrete_buffer_handle_t ctor.

The assumption is not true for shared allocations.

Implement a separate handle type to handle shared allocations: the
implementation will just use the allocation directly.

Author: igchor

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

@@ -233,6 +233,7 @@ ur_discrete_buffer_handle_t::ur_discrete_buffer_handle_t(
     hDevice = hDevice ? hDevice : hContext->getDevices()[0];
     devicePtr = allocateOnDevice(hDevice, size);
   } else {
+    assert(hDevice);
     this->IsInteropNativeHandle = interopNativeHandle;
     deviceAllocations[hDevice->Id.value()] = usm_unique_ptr_t(
         devicePtr, [this, hContext = this->hContext, ownZePtr](void *ptr) {
@@ -363,6 +364,34 @@ void ur_discrete_buffer_handle_t::unmapHostPtr(
   hostAllocations.erase(hostAlloc);
 }
 
+ur_shared_buffer_handle_t::ur_shared_buffer_handle_t(
+    ur_context_handle_t hContext, void *sharedPtr, size_t size,
+    device_access_mode_t accesMode, bool ownDevicePtr)
+    : ur_mem_buffer_t(hContext, size, accesMode),
+      ptr(sharedPtr, [hContext, ownDevicePtr](void *ptr) {
+        if (!ownDevicePtr || !checkL0LoaderTeardown()) {
+          return;
+        }
+        ZE_CALL_NOCHECK(zeMemFree, (hContext->getZeHandle(), ptr));
+      }) {}
+
+void *ur_shared_buffer_handle_t::getDevicePtr(
+    ur_device_handle_t, device_access_mode_t, size_t offset, size_t,
+    std::function<void(void *src, void *dst, size_t)>) {
+  return reinterpret_cast<char *>(ptr.get()) + offset;
+}
+
+void *ur_shared_buffer_handle_t::mapHostPtr(
+    ur_map_flags_t, size_t offset, size_t,
+    std::function<void(void *src, void *dst, size_t)>) {
+  return reinterpret_cast<char *>(ptr.get()) + offset;
+}
+
+void ur_shared_buffer_handle_t::unmapHostPtr(
+    void *, std::function<void(void *src, void *dst, size_t)>) {
+  // nop
+}
+
 static bool useHostBuffer(ur_context_handle_t hContext) {
   // We treat integrated devices (physical memory shared with the CPU)
   // differently from discrete devices (those with distinct memories).
@@ -587,6 +616,10 @@ ur_result_t urMemBufferCreateWithNativeHandle(
         hContext, ptr, size, accessMode, ownNativeHandle, true);
     // if useHostBuffer(hContext) is true but the allocation is on device, we'll
     // treat it as discrete memory
+  } else if (memoryAttrs.type == ZE_MEMORY_TYPE_SHARED) {
+    // For shared allocation, we can use it directly
+    *phMem = ur_mem_handle_t_::create<ur_shared_buffer_handle_t>(
+        hContext, ptr, size, accessMode, ownNativeHandle);
   } else {
     if (memoryAttrs.type == ZE_MEMORY_TYPE_HOST) {
       // For host allocation, we need to copy the data to a device buffer
@@ -595,7 +628,8 @@ ur_result_t urMemBufferCreateWithNativeHandle(
           hContext, hDevice, nullptr, size, accessMode, ptr, ownNativeHandle,
           true);
     } else {
-      // For device/shared allocation, we can use it directly
+      // For device allocation, we can use it directly
+      assert(hDevice);
       *phMem = ur_mem_handle_t_::create<ur_discrete_buffer_handle_t>(
           hContext, hDevice, ptr, size, accessMode, nullptr, ownNativeHandle,
           true);

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

@@ -161,6 +161,24 @@ struct ur_discrete_buffer_handle_t : ur_mem_buffer_t {
                               size_t size);
 };
 
+struct ur_shared_buffer_handle_t : ur_mem_buffer_t {
+  ur_shared_buffer_handle_t(ur_context_handle_t hContext, void *devicePtr,
+                            size_t size, device_access_mode_t accesMode,
+                            bool ownDevicePtr);
+
+  void *
+  getDevicePtr(ur_device_handle_t, device_access_mode_t, size_t offset,
+               size_t size,
+               std::function<void(void *src, void *dst, size_t)>) override;
+  void *mapHostPtr(ur_map_flags_t, size_t offset, size_t size,
+                   std::function<void(void *src, void *dst, size_t)>) override;
+  void unmapHostPtr(void *pMappedPtr,
+                    std::function<void(void *src, void *dst, size_t)>) override;
+
+private:
+  usm_unique_ptr_t ptr;
+};
+
 struct ur_mem_sub_buffer_t : ur_mem_buffer_t {
   ur_mem_sub_buffer_t(ur_mem_handle_t hParent, size_t offset, size_t size,
                       device_access_mode_t accesMode);
@@ -263,6 +281,7 @@ struct ur_mem_handle_t_ {
       : mem(std::in_place_type<T>, std::forward<Args>(args)...) {}
 
   std::variant<ur_usm_handle_t, ur_integrated_buffer_handle_t,
-               ur_discrete_buffer_handle_t, ur_mem_sub_buffer_t, ur_mem_image_t>
+               ur_discrete_buffer_handle_t, ur_shared_buffer_handle_t,
+               ur_mem_sub_buffer_t, ur_mem_image_t>
       mem;
 };

unified-runtime/test/adapters/level_zero/CMakeLists.txt

@@ -94,6 +94,15 @@ function(add_adapter_tests adapter)
         ENVIRONMENT
             "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_${adapter}>\""
         )
+
+    add_adapter_test(${adapter}_mem_buffer_create_with_native_handle
+        FIXTURE DEVICES
+        SOURCES
+            urMemBufferCreateWithNativeHandleShared.cpp
+        ENVIRONMENT
+            "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_${adapter}>\""
+        )
+    target_link_libraries(test-adapter-${adapter}_mem_buffer_create_with_native_handle PRIVATE LevelZeroLoader LevelZeroLoader-Headers)
 endfunction()
 
 if(UR_BUILD_ADAPTER_L0)

unified-runtime/test/adapters/level_zero/urMemBufferCreateWithNativeHandleShared.cpp

@@ -0,0 +1,61 @@
+// Copyright (C) 2025 Intel Corporation
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+// Exceptions. See LICENSE.TXT
+//
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "ur_api.h"
+#include "uur/checks.h"
+#include "uur/raii.h"
+#include "ze_api.h"
+#include <uur/fixtures.h>
+
+using urMemBufferCreateWithNativeHandleTest = uur::urQueueTest;
+UUR_INSTANTIATE_DEVICE_TEST_SUITE(urMemBufferCreateWithNativeHandleTest);
+
+TEST_P(urMemBufferCreateWithNativeHandleTest, SharedBufferIsUsedDirectly) {
+  UUR_KNOWN_FAILURE_ON(uur::LevelZero{});
+
+  // Initialize Level Zero driver is required if this test is linked statically
+  // with Level Zero loader, the driver will not be init otherwise.
+  zeInit(ZE_INIT_FLAG_GPU_ONLY);
+
+  ur_native_handle_t nativeContext;
+  ASSERT_SUCCESS(urContextGetNativeHandle(context, &nativeContext));
+
+  ur_native_handle_t nativeDevice;
+  ASSERT_SUCCESS(urDeviceGetNativeHandle(device, &nativeDevice));
+
+  ze_device_mem_alloc_desc_t DeviceDesc = {};
+  DeviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC;
+  DeviceDesc.ordinal = 0;
+  DeviceDesc.flags = 0;
+  DeviceDesc.pNext = nullptr;
+
+  ze_host_mem_alloc_desc_t HostDesc = {};
+  HostDesc.stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC;
+  HostDesc.pNext = nullptr;
+  HostDesc.flags = 0;
+
+  void *SharedBuffer = nullptr;
+  ASSERT_EQ(
+      zeMemAllocShared(reinterpret_cast<ze_context_handle_t>(nativeContext),
+                       &DeviceDesc, &HostDesc, 12 * sizeof(int), 1, nullptr,
+                       &SharedBuffer),
+      ZE_RESULT_SUCCESS);
+
+  uur::raii::Mem buffer;
+  ASSERT_SUCCESS(urMemBufferCreateWithNativeHandle(
+      reinterpret_cast<ur_native_handle_t>(SharedBuffer), context, nullptr,
+      buffer.ptr()));
+
+  void *mappedPtr;
+  ASSERT_SUCCESS(urEnqueueMemBufferMap(queue, buffer.get(), true, 0, 0,
+                                       12 * sizeof(int), 0, nullptr, nullptr,
+                                       &mappedPtr));
+
+  ASSERT_EQ(mappedPtr, SharedBuffer);
+  ASSERT_EQ(zeMemFree(reinterpret_cast<ze_context_handle_t>(nativeContext),
+                      SharedBuffer),
+            ZE_RESULT_SUCCESS);
+}