[CTS][L0 v2] adjust urEnqueueMemBufferMap

igchor · igchor · commit 8448f78be849 · 2025-01-17T18:20:59.000Z
to match SYCL expectations. When pHost is set in buffer
properties, urEnqueueMemBufferMap should map memory to
that pointer (instead of creating a new allocation).
diff --git a/source/adapters/level_zero/v2/memory.cpp b/source/adapters/level_zero/v2/memory.cpp
@@ -209,7 +209,7 @@ ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(
     device_access_mode_t accessMode)
     : ur_mem_handle_t_(hContext, size, accessMode),
       deviceAllocations(hContext->getPlatform()->getNumDevices()),
-      activeAllocationDevice(nullptr), hostAllocations() {
+      activeAllocationDevice(nullptr), mapToPtr(hostPtr), hostAllocations() {
   if (hostPtr) {
     auto initialDevice = hContext->getDevices()[0];
     UR_CALL_THROWS(migrateBufferTo(initialDevice, hostPtr, size));
@@ -246,12 +246,18 @@ ur_discrete_mem_handle_t::~ur_discrete_mem_handle_t() {
   if (!activeAllocationDevice || !writeBackPtr)
     return;
 
-  auto srcPtr = ur_cast<char *>(
-      deviceAllocations[activeAllocationDevice->Id.value()].get());
+  auto srcPtr = getActiveDeviceAlloc();
   synchronousZeCopy(hContext, activeAllocationDevice, writeBackPtr, srcPtr,
                     getSize());
 }
 
+void *ur_discrete_mem_handle_t::getActiveDeviceAlloc(size_t offset) {
+  assert(activeAllocationDevice);
+  return ur_cast<char *>(
+             deviceAllocations[activeAllocationDevice->Id.value()].get()) +
+         offset;
+}
+
 void *ur_discrete_mem_handle_t::getDevicePtr(
     ur_device_handle_t hDevice, device_access_mode_t access, size_t offset,
     size_t size, std::function<void(void *src, void *dst, size_t)> migrate) {
@@ -272,10 +278,8 @@ void *ur_discrete_mem_handle_t::getDevicePtr(
     hDevice = activeAllocationDevice;
   }
 
-  char *ptr;
   if (activeAllocationDevice == hDevice) {
-    ptr = ur_cast<char *>(deviceAllocations[hDevice->Id.value()].get());
-    return ptr + offset;
+    return getActiveDeviceAlloc(offset);
   }
 
   auto &p2pDevices = hContext->getP2PDevices(hDevice);
@@ -288,9 +292,7 @@ void *ur_discrete_mem_handle_t::getDevicePtr(
   }
 
   // TODO: see if it's better to migrate the memory to the specified device
-  return ur_cast<char *>(
-             deviceAllocations[activeAllocationDevice->Id.value()].get()) +
-         offset;
+  return getActiveDeviceAlloc(offset);
 }
 
 void *ur_discrete_mem_handle_t::mapHostPtr(
@@ -299,55 +301,60 @@ void *ur_discrete_mem_handle_t::mapHostPtr(
   TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::mapHostPtr");
   // TODO: use async alloc?
 
-  void *ptr;
-  UR_CALL_THROWS(hContext->getDefaultUSMPool()->allocate(
-      hContext, nullptr, nullptr, UR_USM_TYPE_HOST, size, &ptr));
+  void *ptr = mapToPtr;
+  if (!ptr) {
+    UR_CALL_THROWS(hContext->getDefaultUSMPool()->allocate(
+        hContext, nullptr, nullptr, UR_USM_TYPE_HOST, size, &ptr));
+  }
 
-  hostAllocations.emplace_back(ptr, size, offset, flags);
+  usm_unique_ptr_t mappedPtr =
+      usm_unique_ptr_t(ptr, [ownsAlloc = bool(mapToPtr), this](void *p) {
+        if (ownsAlloc) {
+          UR_CALL_THROWS(hContext->getDefaultUSMPool()->free(p));
+        }
+      });
+
+  hostAllocations.emplace_back(std::move(mappedPtr), size, offset, flags);
 
   if (activeAllocationDevice && (flags & UR_MAP_FLAG_READ)) {
-    auto srcPtr =
-        ur_cast<char *>(
-            deviceAllocations[activeAllocationDevice->Id.value()].get()) +
-        offset;
-    migrate(srcPtr, hostAllocations.back().ptr, size);
+    auto srcPtr = getActiveDeviceAlloc(offset);
+    migrate(srcPtr, hostAllocations.back().ptr.get(), size);
   }
 
-  return hostAllocations.back().ptr;
+  return hostAllocations.back().ptr.get();
 }
 
 void ur_discrete_mem_handle_t::unmapHostPtr(
     void *pMappedPtr,
     std::function<void(void *src, void *dst, size_t)> migrate) {
   TRACK_SCOPE_LATENCY("ur_discrete_mem_handle_t::unmapHostPtr");
 
-  for (auto &hostAllocation : hostAllocations) {
-    if (hostAllocation.ptr == pMappedPtr) {
-      void *devicePtr = nullptr;
-      if (activeAllocationDevice) {
-        devicePtr =
-            ur_cast<char *>(
-                deviceAllocations[activeAllocationDevice->Id.value()].get()) +
-            hostAllocation.offset;
-      } else if (!(hostAllocation.flags &
-                   UR_MAP_FLAG_WRITE_INVALIDATE_REGION)) {
-        devicePtr = ur_cast<char *>(getDevicePtr(
-            hContext->getDevices()[0], device_access_mode_t::read_only,
-            hostAllocation.offset, hostAllocation.size, migrate));
-      }
+  auto hostAlloc =
+      std::find_if(hostAllocations.begin(), hostAllocations.end(),
+                   [pMappedPtr](const host_allocation_desc_t &desc) {
+                     return desc.ptr.get() == pMappedPtr;
+                   });
 
-      if (devicePtr) {
-        migrate(hostAllocation.ptr, devicePtr, hostAllocation.size);
-      }
+  if (hostAlloc == hostAllocations.end()) {
+    throw UR_RESULT_ERROR_INVALID_ARGUMENT;
+  }
 
-      // TODO: use async free here?
-      UR_CALL_THROWS(hContext->getDefaultUSMPool()->free(hostAllocation.ptr));
-      return;
-    }
+  bool shouldMigrateToDevice =
+      !(hostAlloc->flags & UR_MAP_FLAG_WRITE_INVALIDATE_REGION);
+
+  if (!activeAllocationDevice && shouldMigrateToDevice) {
+    allocateOnDevice(hContext->getDevices()[0], getSize());
+  }
+
+  // TODO: tests require that memory is migrated even for
+  // UR_MAP_FLAG_WRITE_INVALIDATE_REGION when there is an active device
+  // allocation. is this correct?
+  if (activeAllocationDevice) {
+    migrate(hostAlloc->ptr.get(), getActiveDeviceAlloc(hostAlloc->offset),
+            hostAlloc->size);
   }
 
-  // No mapping found
-  throw UR_RESULT_ERROR_INVALID_ARGUMENT;
+  hostAllocations.erase(hostAlloc);
 }
 
 static bool useHostBuffer(ur_context_handle_t hContext) {
diff --git a/source/adapters/level_zero/v2/memory.hpp b/source/adapters/level_zero/v2/memory.hpp
@@ -98,11 +98,11 @@ struct ur_integrated_mem_handle_t : public ur_mem_handle_t_ {
 };
 
 struct host_allocation_desc_t {
-  host_allocation_desc_t(void *ptr, size_t size, size_t offset,
+  host_allocation_desc_t(usm_unique_ptr_t ptr, size_t size, size_t offset,
                          ur_map_flags_t flags)
-      : ptr(ptr), size(size), offset(offset), flags(flags) {}
+      : ptr(std::move(ptr)), size(size), offset(offset), flags(flags) {}
 
-  void *ptr;
+  usm_unique_ptr_t ptr;
   size_t size;
   size_t offset;
   ur_map_flags_t flags;
@@ -146,10 +146,13 @@ struct ur_discrete_mem_handle_t : public ur_mem_handle_t_ {
   // If not null, copy the buffer content back to this memory on release.
   void *writeBackPtr = nullptr;
 
+  // If not null, mapHostPtr should map memory to this ptr
+  void *mapToPtr = nullptr;
+
   std::vector<host_allocation_desc_t> hostAllocations;
 
+  void *getActiveDeviceAlloc(size_t offset = 0);
   void *allocateOnDevice(ur_device_handle_t hDevice, size_t size);
-
   ur_result_t migrateBufferTo(ur_device_handle_t hDevice, void *src,
                               size_t size);
 };
diff --git a/test/adapters/level_zero/CMakeLists.txt b/test/adapters/level_zero/CMakeLists.txt
@@ -86,6 +86,14 @@ function(add_adapter_tests adapter)
             ur_umf
         )
     endif()
+
+    add_adapter_test(${adapter}_mem_buffer_map
+        FIXTURE DEVICES
+        SOURCES
+            urEnqueueMemBufferMapHostPtr.cpp
+        ENVIRONMENT
+            "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_${adapter}>\""
+        )
 endfunction()
 
 if(UR_BUILD_ADAPTER_L0)
diff --git a/test/adapters/level_zero/urEnqueueMemBufferMapHostPtr.cpp b/test/adapters/level_zero/urEnqueueMemBufferMapHostPtr.cpp
@@ -0,0 +1,38 @@
+// Copyright (C) 2023 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 "../../enqueue/helpers.h"
+#include <uur/fixtures.h>
+#include <uur/known_failure.h>
+#include <uur/raii.h>
+
+using urEnqueueMemBufferMapTestWithParamL0 =
+    uur::urMemBufferQueueTestWithParam<uur::mem_buffer_test_parameters_t>;
+
+UUR_DEVICE_TEST_SUITE_P(
+    urEnqueueMemBufferMapTestWithParamL0,
+    ::testing::ValuesIn(uur::mem_buffer_test_parameters),
+    uur::printMemBufferTestString<urEnqueueMemBufferMapTestWithParamL0>);
+
+TEST_P(urEnqueueMemBufferMapTestWithParamL0, MapWithHostPtr) {
+  uur::raii::Mem buffer;
+
+  char *ptr = new char[4096];
+
+  ur_buffer_properties_t props;
+  props.pHost = ptr;
+
+  ASSERT_SUCCESS(urMemBufferCreate(context, UR_MEM_FLAG_USE_HOST_POINTER, 4096,
+                                   &props, buffer.ptr()));
+
+  void *mappedPtr = nullptr;
+  ASSERT_SUCCESS(urEnqueueMemBufferMap(
+      queue, buffer.get(), true, UR_MAP_FLAG_READ | UR_MAP_FLAG_WRITE, 0, 4096,
+      0, nullptr, nullptr, reinterpret_cast<void **>(&mappedPtr)));
+
+  ASSERT_EQ(ptr, mappedPtr);
+
+  delete[] ptr;
+}
