[SYCL][UR][L0 v2] get rid of std::function in memory.hpp

There migration logic is always the same, there's no
need to pass callback to every getDevicePtr/map/unmap
function.

Author: igchor

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

@@ -39,9 +39,9 @@ ur_result_t getMemPtr(ur_mem_handle_t memObj,
       urAccessMode =
           ur_mem_buffer_t::getDeviceAccessMode(properties->memoryAccess);
     }
-    ptr = ur_cast<char *>(
-        memBuffer->getDevicePtr(device, urAccessMode, 0, memBuffer->getSize(),
-                                [&](void *, void *, size_t) {}));
+    wait_list_view emptyWaitList(nullptr, 0);
+    ptr = ur_cast<char *>(memBuffer->getDevicePtr(
+        device, urAccessMode, 0, memBuffer->getSize(), nullptr, emptyWaitList));
   }
   assert(ptrStorage != nullptr);
   ptrStorage->push_back(std::make_unique<char *>(ptr));

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

@@ -14,6 +14,7 @@
 #include "../ur_interface_loader.hpp"
 #include "context.hpp"
 #include "kernel.hpp"
+#include "memory.hpp"
 
 ur_command_list_manager::ur_command_list_manager(
     ur_context_handle_t context, ur_device_handle_t device,
@@ -44,12 +45,7 @@ ur_result_t ur_command_list_manager::appendGenericFillUnlocked(
 
   auto pDst = ur_cast<char *>(dst->getDevicePtr(
       device, ur_mem_buffer_t::device_access_mode_t::read_only, offset, size,
-      [&](void *src, void *dst, size_t size) {
-        ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
-                          (zeCommandList.get(), dst, src, size, nullptr,
-                           waitListView.num, waitListView.handles));
-        waitListView.clear();
-      }));
+      zeCommandList.get(), waitListView));
 
   // PatternSize must be a power of two for zeCommandListAppendMemoryFill.
   // When it's not, the fill is emulated with zeCommandListAppendMemoryCopy.
@@ -87,21 +83,11 @@ ur_result_t ur_command_list_manager::appendGenericCopyUnlocked(
 
   auto pSrc = ur_cast<char *>(src->getDevicePtr(
       device, ur_mem_buffer_t::device_access_mode_t::read_only, srcOffset, size,
-      [&](void *src, void *dst, size_t size) {
-        ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
-                          (zeCommandList.get(), dst, src, size, nullptr,
-                           waitListView.num, waitListView.handles));
-        waitListView.clear();
-      }));
+      zeCommandList.get(), waitListView));
 
   auto pDst = ur_cast<char *>(dst->getDevicePtr(
       device, ur_mem_buffer_t::device_access_mode_t::write_only, dstOffset,
-      size, [&](void *src, void *dst, size_t size) {
-        ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
-                          (zeCommandList.get(), dst, src, size, nullptr,
-                           waitListView.num, waitListView.handles));
-        waitListView.clear();
-      }));
+      size, zeCommandList.get(), waitListView));
 
   ZE2UR_CALL(zeCommandListAppendMemoryCopy,
              (zeCommandList.get(), pDst, pSrc, size, zeSignalEvent,
@@ -130,20 +116,10 @@ ur_result_t ur_command_list_manager::appendRegionCopyUnlocked(
 
   auto pSrc = ur_cast<char *>(src->getDevicePtr(
       device, ur_mem_buffer_t::device_access_mode_t::read_only, 0,
-      src->getSize(), [&](void *src, void *dst, size_t size) {
-        ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
-                          (zeCommandList.get(), dst, src, size, nullptr,
-                           waitListView.num, waitListView.handles));
-        waitListView.clear();
-      }));
+      src->getSize(), zeCommandList.get(), waitListView));
   auto pDst = ur_cast<char *>(dst->getDevicePtr(
       device, ur_mem_buffer_t::device_access_mode_t::write_only, 0,
-      dst->getSize(), [&](void *src, void *dst, size_t size) {
-        ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
-                          (zeCommandList.get(), dst, src, size, nullptr,
-                           waitListView.num, waitListView.handles));
-        waitListView.clear();
-      }));
+      dst->getSize(), zeCommandList.get(), waitListView));
 
   ZE2UR_CALL(zeCommandListAppendMemoryCopyRegion,
              (zeCommandList.get(), pDst, &zeParams.dstRegion, zeParams.dstPitch,
@@ -213,13 +189,6 @@ ur_result_t ur_command_list_manager::appendKernelLaunch(
 
   auto waitListView = getWaitListView(phEventWaitList, numEventsInWaitList);
 
-  auto memoryMigrate = [&](void *src, void *dst, size_t size) {
-    ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
-                      (zeCommandList.get(), dst, src, size, nullptr,
-                       waitListView.num, waitListView.handles));
-    waitListView.clear();
-  };
-
   // If the offset is {0, 0, 0}, pass NULL instead.
   // This allows us to skip setting the offset.
   bool hasOffset = false;
@@ -232,7 +201,7 @@ ur_result_t ur_command_list_manager::appendKernelLaunch(
 
   UR_CALL(hKernel->prepareForSubmission(context, device, pGlobalWorkOffset,
                                         workDim, WG[0], WG[1], WG[2],
-                                        memoryMigrate));
+                                        zeCommandList.get(), waitListView));
 
   TRACK_SCOPE_LATENCY(
       "ur_command_list_manager::zeCommandListAppendLaunchKernel");

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

@@ -13,10 +13,11 @@
 #include "common.hpp"
 #include "context.hpp"
 #include "event_pool_cache.hpp"
-#include "memory.hpp"
 #include "queue_api.hpp"
 #include <ze_api.h>
 
+struct ur_mem_buffer_t;
+
 struct wait_list_view {
   ze_event_handle_t *handles;
   uint32_t num;

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

@@ -277,7 +277,7 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
     ur_context_handle_t hContext, ur_device_handle_t hDevice,
     const size_t *pGlobalWorkOffset, uint32_t workDim, uint32_t groupSizeX,
     uint32_t groupSizeY, uint32_t groupSizeZ,
-    std::function<void(void *, void *, size_t)> migrate) {
+    ze_command_list_handle_t commandList, wait_list_view &waitListView) {
   auto hZeKernel = getZeHandle(hDevice);
 
   if (pGlobalWorkOffset != NULL) {
@@ -293,8 +293,9 @@ ur_result_t ur_kernel_handle_t_::prepareForSubmission(
     if (pending.hMem) {
       if (!pending.hMem->isImage()) {
         auto hBuffer = pending.hMem->getBuffer();
-        zePtr = hBuffer->getDevicePtr(hDevice, pending.mode, 0,
-                                      hBuffer->getSize(), migrate);
+        zePtr =
+            hBuffer->getDevicePtr(hDevice, pending.mode, 0, hBuffer->getSize(),
+                                  commandList, waitListView);
       } else {
         auto hImage = static_cast<ur_mem_image_t *>(pending.hMem->getImage());
         zePtr = reinterpret_cast<void *>(hImage->getZeImage());

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

@@ -83,12 +83,13 @@ struct ur_kernel_handle_t_ : ur_object {
 
   // Set all required values for the kernel before submission (including pending
   // memory allocations).
-  ur_result_t
-  prepareForSubmission(ur_context_handle_t hContext, ur_device_handle_t hDevice,
-                       const size_t *pGlobalWorkOffset, uint32_t workDim,
-                       uint32_t groupSizeX, uint32_t groupSizeY,
-                       uint32_t groupSizeZ,
-                       std::function<void(void *, void *, size_t)> migrate);
+  ur_result_t prepareForSubmission(ur_context_handle_t hContext,
+                                   ur_device_handle_t hDevice,
+                                   const size_t *pGlobalWorkOffset,
+                                   uint32_t workDim, uint32_t groupSizeX,
+                                   uint32_t groupSizeY, uint32_t groupSizeZ,
+                                   ze_command_list_handle_t cmdList,
+                                   wait_list_view &waitListView);
 
 private:
   // Keep the program of the kernel.

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

@@ -32,22 +32,24 @@ ur_usm_handle_t::ur_usm_handle_t(ur_context_handle_t hContext, size_t size,
     : ur_mem_buffer_t(hContext, size, device_access_mode_t::read_write),
       ptr(const_cast<void *>(ptr)) {}
 
-void *ur_usm_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*/) {
+void *ur_usm_handle_t::getDevicePtr(ur_device_handle_t /*hDevice*/,
+                                    device_access_mode_t /*access*/,
+                                    size_t offset, size_t /*size*/,
+                                    ze_command_list_handle_t /*cmdList*/,
+                                    wait_list_view & /*waitListView*/) {
   return ur_cast<char *>(ptr) + offset;
 }
 
-void *
-ur_usm_handle_t::mapHostPtr(ur_map_flags_t /*flags*/, size_t offset,
-                            size_t /*size*/,
-                            std::function<void(void *src, void *dst, size_t)>) {
+void *ur_usm_handle_t::mapHostPtr(ur_map_flags_t /*flags*/, size_t offset,
+                                  size_t /*size*/,
+                                  ze_command_list_handle_t /*cmdList*/,
+                                  wait_list_view & /*waitListView*/) {
   return ur_cast<char *>(ptr) + offset;
 }
 
-void ur_usm_handle_t::unmapHostPtr(
-    void * /*pMappedPtr*/, std::function<void(void *src, void *dst, size_t)>) {
+void ur_usm_handle_t::unmapHostPtr(void * /*pMappedPtr*/,
+                                   ze_command_list_handle_t cmdList,
+                                   wait_list_view & /*waitListView*/) {
   /* nop */
 }
 
@@ -106,14 +108,14 @@ ur_integrated_buffer_handle_t::~ur_integrated_buffer_handle_t() {
 
 void *ur_integrated_buffer_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*/) {
+    size_t offset, size_t /*size*/, ze_command_list_handle_t /*cmdList*/,
+    wait_list_view & /*waitListView*/) {
   return ur_cast<char *>(ptr.get()) + offset;
 }
 
 void *ur_integrated_buffer_handle_t::mapHostPtr(
     ur_map_flags_t /*flags*/, size_t offset, size_t /*size*/,
-    std::function<void(void *src, void *dst, size_t)> /*migrate*/) {
+    ze_command_list_handle_t /*cmdList*/, wait_list_view & /*waitListView*/) {
   // TODO: if writeBackPtr is set, we should map to that pointer
   // because that's what SYCL expects, SYCL will attempt to call free
   // on the resulting pointer leading to double free with the current
@@ -122,7 +124,8 @@ void *ur_integrated_buffer_handle_t::mapHostPtr(
 }
 
 void ur_integrated_buffer_handle_t::unmapHostPtr(
-    void * /*pMappedPtr*/, std::function<void(void *src, void *dst, size_t)>) {
+    void * /*pMappedPtr*/, ze_command_list_handle_t /*cmdList*/,
+    wait_list_view & /*waitListView*/) {
   // TODO: if writeBackPtr is set, we should copy the data back
   /* nop */
 }
@@ -250,8 +253,8 @@ void *ur_discrete_buffer_handle_t::getActiveDeviceAlloc(size_t offset) {
 
 void *ur_discrete_buffer_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*/) {
+    size_t /*size*/, ze_command_list_handle_t /*cmdList*/,
+    wait_list_view & /*waitListView*/) {
   TRACK_SCOPE_LATENCY("ur_discrete_buffer_handle_t::getDevicePtr");
 
   if (!activeAllocationDevice) {
@@ -283,9 +286,22 @@ void *ur_discrete_buffer_handle_t::getDevicePtr(
   return getActiveDeviceAlloc(offset);
 }
 
-void *ur_discrete_buffer_handle_t::mapHostPtr(
-    ur_map_flags_t flags, size_t offset, size_t size,
-    std::function<void(void *src, void *dst, size_t)> migrate) {
+static void migrateMemory(ze_command_list_handle_t cmdList, void *src,
+                          void *dst, size_t size,
+                          wait_list_view &waitListView) {
+  if (!cmdList) {
+    throw UR_RESULT_ERROR_INVALID_NULL_HANDLE;
+  }
+  ZE2UR_CALL_THROWS(zeCommandListAppendMemoryCopy,
+                    (cmdList, dst, src, size, nullptr, waitListView.num,
+                     waitListView.handles));
+  waitListView.clear();
+}
+
+void *ur_discrete_buffer_handle_t::mapHostPtr(ur_map_flags_t flags,
+                                              size_t offset, size_t size,
+                                              ze_command_list_handle_t cmdList,
+                                              wait_list_view &waitListView) {
   TRACK_SCOPE_LATENCY("ur_discrete_buffer_handle_t::mapHostPtr");
   // TODO: use async alloc?
 
@@ -309,15 +325,16 @@ void *ur_discrete_buffer_handle_t::mapHostPtr(
 
   if (activeAllocationDevice && (flags & UR_MAP_FLAG_READ)) {
     auto srcPtr = getActiveDeviceAlloc(offset);
-    migrate(srcPtr, hostAllocations.back().ptr.get(), size);
+    migrateMemory(cmdList, srcPtr, hostAllocations.back().ptr.get(), size,
+                  waitListView);
   }
 
   return hostAllocations.back().ptr.get();
 }
 
-void ur_discrete_buffer_handle_t::unmapHostPtr(
-    void *pMappedPtr,
-    std::function<void(void *src, void *dst, size_t)> migrate) {
+void ur_discrete_buffer_handle_t::unmapHostPtr(void *pMappedPtr,
+                                               ze_command_list_handle_t cmdList,
+                                               wait_list_view &waitListView) {
   TRACK_SCOPE_LATENCY("ur_discrete_buffer_handle_t::unmapHostPtr");
 
   auto hostAlloc =
@@ -341,8 +358,9 @@ void ur_discrete_buffer_handle_t::unmapHostPtr(
   // 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);
+    migrateMemory(cmdList, hostAlloc->ptr.get(),
+                  getActiveDeviceAlloc(hostAlloc->offset), hostAlloc->size,
+                  waitListView);
   }
 
   hostAllocations.erase(hostAlloc);
@@ -361,18 +379,20 @@ ur_shared_buffer_handle_t::ur_shared_buffer_handle_t(
 
 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)>) {
+    ze_command_list_handle_t /*cmdList*/, wait_list_view & /*waitListView*/) {
   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)>) {
+void *
+ur_shared_buffer_handle_t::mapHostPtr(ur_map_flags_t, size_t offset, size_t,
+                                      ze_command_list_handle_t /*cmdList*/,
+                                      wait_list_view & /*waitListView*/) {
   return reinterpret_cast<char *>(ptr.get()) + offset;
 }
 
 void ur_shared_buffer_handle_t::unmapHostPtr(
-    void *, std::function<void(void *src, void *dst, size_t)>) {
+    void *, ze_command_list_handle_t /*cmdList*/,
+    wait_list_view & /*waitListView*/) {
   // nop
 }
 
@@ -403,24 +423,27 @@ ur_mem_sub_buffer_t::~ur_mem_sub_buffer_t() {
   ur::level_zero::urMemRelease(hParent);
 }
 
-void *ur_mem_sub_buffer_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) {
+void *ur_mem_sub_buffer_t::getDevicePtr(ur_device_handle_t hDevice,
+                                        device_access_mode_t access,
+                                        size_t offset, size_t size,
+                                        ze_command_list_handle_t cmdList,
+                                        wait_list_view &waitListView) {
   return hParent->getBuffer()->getDevicePtr(
-      hDevice, access, offset + this->offset, size, std::move(migrate));
+      hDevice, access, offset + this->offset, size, cmdList, waitListView);
 }
 
-void *ur_mem_sub_buffer_t::mapHostPtr(
-    ur_map_flags_t flags, size_t offset, size_t size,
-    std::function<void(void *src, void *dst, size_t)> migrate) {
+void *ur_mem_sub_buffer_t::mapHostPtr(ur_map_flags_t flags, size_t offset,
+                                      size_t size,
+                                      ze_command_list_handle_t cmdList,
+                                      wait_list_view &waitListView) {
   return hParent->getBuffer()->mapHostPtr(flags, offset + this->offset, size,
-                                          std::move(migrate));
+                                          cmdList, waitListView);
 }
 
-void ur_mem_sub_buffer_t::unmapHostPtr(
-    void *pMappedPtr,
-    std::function<void(void *src, void *dst, size_t)> migrate) {
-  return hParent->getBuffer()->unmapHostPtr(pMappedPtr, std::move(migrate));
+void ur_mem_sub_buffer_t::unmapHostPtr(void *pMappedPtr,
+                                       ze_command_list_handle_t cmdList,
+                                       wait_list_view &waitListView) {
+  return hParent->getBuffer()->unmapHostPtr(pMappedPtr, cmdList, waitListView);
 }
 
 ur_shared_mutex &ur_mem_sub_buffer_t::getMutex() {
@@ -690,9 +713,10 @@ ur_result_t urMemGetNativeHandle(ur_mem_handle_t hMem,
 
   std::scoped_lock<ur_shared_mutex> lock(hBuffer->getMutex());
 
+  wait_list_view emptyWaitListView(nullptr, 0);
   auto ptr = hBuffer->getDevicePtr(
       hDevice, ur_mem_buffer_t::device_access_mode_t::read_write, 0,
-      hBuffer->getSize(), nullptr);
+      hBuffer->getSize(), nullptr, emptyWaitListView);
   *phNativeMem = reinterpret_cast<ur_native_handle_t>(ptr);
   return UR_RESULT_SUCCESS;
 } catch (...) {