[L0 v2] support multi-device memory buffers

Author: igchor

source/adapters/level_zero/v2/kernel.cpp

@@ -316,11 +316,16 @@ urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
     auto zePtr = hArgValue->getPtr(kernelDevices.front());
     return hKernel->setArgPointer(argIndex, nullptr, zePtr);
   } else {
-    // TODO: Implement this for multi-device kernels.
-    // Do this the same way as in legacy (keep a pending Args vector and
-    // do actual allocation on kernel submission) or allocate the memory
-    // immediately (only for small allocations?)
-    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    // TODO: if devices do not have p2p capabilities, we need to have allocation
+    // on each device. Do this the same way as in legacy (keep a pending Args
+    // vector and do actual allocation on kernel submission) or allocate the
+    // memory immediately (only for small allocations?).
+
+    // Get memory that is accessible by the first device.
+    // If kernel is submitted to a different device the memory
+    // will be accessed trough the link or migrated in enqueueKernelLaunch.
+    auto zePtr = hArgValue->getPtr(kernelDevices.front());
+    return hKernel->setArgPointer(argIndex, nullptr, zePtr);
   }
 }
 

source/adapters/level_zero/v2/memory.cpp

@@ -28,10 +28,8 @@ ur_host_mem_handle_t::ur_host_mem_handle_t(ur_context_handle_t hContext,
   }
 
   if (!hostPtrImported) {
-    // TODO: use UMF
-    ZeStruct<ze_host_mem_alloc_desc_t> hostDesc;
-    ZE2UR_CALL_THROWS(zeMemAllocHost, (hContext->getZeHandle(), &hostDesc, size,
-                                       0, &this->ptr));
+    UR_CALL_THROWS(hContext->getDefaultUSMPool()->allocate(
+        hContext, nullptr, nullptr, UR_USM_TYPE_HOST, size, &this->ptr));
 
     if (hostPtr) {
       std::memcpy(this->ptr, hostPtr, size);
@@ -40,9 +38,11 @@ ur_host_mem_handle_t::ur_host_mem_handle_t(ur_context_handle_t hContext,
 }
 
 ur_host_mem_handle_t::~ur_host_mem_handle_t() {
-  // TODO: use UMF API here
   if (ptr) {
-    ZE_CALL_NOCHECK(zeMemFree, (hContext->getZeHandle(), ptr));
+    auto ret = hContext->getDefaultUSMPool()->free(ptr);
+    if (ret != UR_RESULT_SUCCESS) {
+      logger::error("Failed to free host memory: {}", ret);
+    }
   }
 }
 
@@ -51,55 +51,80 @@ void *ur_host_mem_handle_t::getPtr(ur_device_handle_t hDevice) {
   return ptr;
 }
 
+ur_result_t ur_device_mem_handle_t::migrateBufferTo(ur_device_handle_t hDevice,
+                                                    void *src, size_t size) {
+  auto Id = hDevice->Id.value();
+
+  if (!deviceAllocations[Id]) {
+    UR_CALL(hContext->getDefaultUSMPool()->allocate(hContext, hDevice, nullptr,
+                                                    UR_USM_TYPE_DEVICE, size,
+                                                    &deviceAllocations[Id]));
+  }
+
+  auto commandList = hContext->commandListCache.getImmediateCommandList(
+      hDevice->ZeDevice, true,
+      hDevice
+          ->QueueGroup[ur_device_handle_t_::queue_group_info_t::type::Compute]
+          .ZeOrdinal,
+      ZE_COMMAND_QUEUE_MODE_SYNCHRONOUS, ZE_COMMAND_QUEUE_PRIORITY_NORMAL,
+      std::nullopt);
+
+  ZE2UR_CALL(zeCommandListAppendMemoryCopy,
+             (commandList.get(), deviceAllocations[Id], src, size, nullptr, 0,
+              nullptr));
+
+  activeAllocationDevice = hDevice;
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_device_mem_handle_t::ur_device_mem_handle_t(ur_context_handle_t hContext,
                                                void *hostPtr, size_t size)
     : ur_mem_handle_t_(hContext, size),
-      deviceAllocations(hContext->getPlatform()->getNumDevices()) {
-  // Legacy adapter allocated the memory directly on a device (first on the
-  // contxt) and if the buffer is used on another device, memory is migrated
-  // (depending on an env var setting).
-  //
-  // TODO: port this behavior or figure out if it makes sense to keep the memory
-  // in a host buffer (e.g. for smaller sizes).
+      deviceAllocations(hContext->getPlatform()->getNumDevices()),
+      activeAllocationDevice(nullptr) {
   if (hostPtr) {
-    buffer.assign(reinterpret_cast<char *>(hostPtr),
-                  reinterpret_cast<char *>(hostPtr) + size);
+    auto initialDevice = hContext->getDevices()[0];
+    UR_CALL_THROWS(migrateBufferTo(initialDevice, hostPtr, size));
   }
 }
 
 ur_device_mem_handle_t::~ur_device_mem_handle_t() {
-  // TODO: use UMF API here
   for (auto &ptr : deviceAllocations) {
     if (ptr) {
-      ZE_CALL_NOCHECK(zeMemFree, (hContext->getZeHandle(), ptr));
+      auto ret = hContext->getDefaultUSMPool()->free(ptr);
+      if (ret != UR_RESULT_SUCCESS) {
+        logger::error("Failed to free device memory: {}", ret);
+      }
     }
   }
 }
 
 void *ur_device_mem_handle_t::getPtr(ur_device_handle_t hDevice) {
   std::lock_guard lock(this->Mutex);
 
-  auto &ptr = deviceAllocations[hDevice->Id.value()];
-  if (!ptr) {
-    ZeStruct<ze_device_mem_alloc_desc_t> deviceDesc;
-    ZE2UR_CALL_THROWS(zeMemAllocDevice, (hContext->getZeHandle(), &deviceDesc,
-                                         size, 0, hDevice->ZeDevice, &ptr));
-
-    if (!buffer.empty()) {
-      auto commandList = hContext->commandListCache.getImmediateCommandList(
-          hDevice->ZeDevice, true,
-          hDevice
-              ->QueueGroup
-                  [ur_device_handle_t_::queue_group_info_t::type::Compute]
-              .ZeOrdinal,
-          ZE_COMMAND_QUEUE_MODE_SYNCHRONOUS, ZE_COMMAND_QUEUE_PRIORITY_NORMAL,
-          std::nullopt);
-      ZE2UR_CALL_THROWS(
-          zeCommandListAppendMemoryCopy,
-          (commandList.get(), ptr, buffer.data(), size, nullptr, 0, nullptr));
-    }
+  if (!activeAllocationDevice) {
+    UR_CALL_THROWS(hContext->getDefaultUSMPool()->allocate(
+        hContext, hDevice, nullptr, UR_USM_TYPE_DEVICE, getSize(),
+        &deviceAllocations[hDevice->Id.value()]));
+    activeAllocationDevice = hDevice;
   }
-  return ptr;
+
+  if (activeAllocationDevice == hDevice) {
+    return deviceAllocations[hDevice->Id.value()];
+  }
+
+  auto &p2pDevices = hContext->getP2PDevices(hDevice);
+  auto p2pAccessible = std::find(p2pDevices.begin(), p2pDevices.end(),
+                                 activeAllocationDevice) != p2pDevices.end();
+
+  if (!p2pAccessible) {
+    // TODO: migrate buffer through the host
+    throw UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
+
+  // TODO: see if it's better to migrate the memory to the specified device
+  return deviceAllocations[activeAllocationDevice->Id.value()];
 }
 
 namespace ur::level_zero {

source/adapters/level_zero/v2/memory.hpp

@@ -12,6 +12,7 @@
 
 #include <ur_api.h>
 
+#include "../device.hpp"
 #include "common.hpp"
 
 struct ur_mem_handle_t_ : _ur_object {
@@ -48,8 +49,13 @@ struct ur_device_mem_handle_t : public ur_mem_handle_t_ {
   void *getPtr(ur_device_handle_t) override;
 
 private:
-  std::vector<char> buffer;
-
   // Vector of per-device allocations indexed by device->Id
   std::vector<void *> deviceAllocations;
+
+  // Specifies device on which the latest allocation resides.
+  // If null, there is no allocation.
+  ur_device_handle_t activeAllocationDevice;
+
+  ur_result_t migrateBufferTo(ur_device_handle_t hDevice, void *src,
+                              size_t size);
 };

source/adapters/level_zero/v2/queue_immediate_in_order.cpp

@@ -239,6 +239,8 @@ ur_result_t ur_queue_immediate_in_order_t::enqueueKernelLaunch(
   auto [pWaitEvents, numWaitEvents] =
       getWaitListView(phEventWaitList, numEventsInWaitList, handler);
 
+  // TODO: consider migrating memory to the device if memory buffers are used
+
   TRACK_SCOPE_LATENCY(
       "ur_queue_immediate_in_order_t::zeCommandListAppendLaunchKernel");
   ZE2UR_CALL(zeCommandListAppendLaunchKernel,

source/adapters/level_zero/v2/usm.cpp

@@ -184,6 +184,10 @@ ur_result_t ur_usm_pool_handle_t_::allocate(
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t ur_usm_pool_handle_t_::free(void *ptr) {
+  return umf::umf2urResult(umfFree(ptr));
+}
+
 namespace ur::level_zero {
 ur_result_t urUSMPoolCreate(
     ur_context_handle_t hContext, ///< [in] handle of the context object

source/adapters/level_zero/v2/usm.hpp

@@ -24,6 +24,7 @@ struct ur_usm_pool_handle_t_ : _ur_object {
   ur_result_t allocate(ur_context_handle_t hContext, ur_device_handle_t hDevice,
                        const ur_usm_desc_t *pUSMDesc, ur_usm_type_t type,
                        size_t size, void **ppRetMem);
+  ur_result_t free(void *ptr);
 
 private:
   ur_context_handle_t hContext;