[L0 v2] add inital memory buffer support

Currently there is no support for buffer migration.
This will be done in next PRs.

Author: igchor

source/adapters/level_zero/CMakeLists.txt

@@ -118,6 +118,7 @@ if(UR_BUILD_ADAPTER_L0)
         ${CMAKE_CURRENT_SOURCE_DIR}/queue.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/sampler.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/kernel_helpers.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/helpers/memory_helpers.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/ur_level_zero.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/context.cpp
@@ -136,6 +137,7 @@ if(UR_BUILD_ADAPTER_L0)
         ${CMAKE_CURRENT_SOURCE_DIR}/sampler.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/image.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/kernel_helpers.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/helpers/memory_helpers.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/../../ur/ur.cpp
     )
 
@@ -199,13 +201,15 @@ if(UR_BUILD_ADAPTER_L0_V2)
         ${CMAKE_CURRENT_SOURCE_DIR}/platform.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/program.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/kernel_helpers.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/helpers/memory_helpers.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/adapter.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/ur_interface_loader.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/platform.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/program.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/helpers/kernel_helpers.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/helpers/memory_helpers.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/../../ur/ur.cpp
         # v2-only sources
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/command_list_cache.hpp
@@ -217,6 +221,7 @@ if(UR_BUILD_ADAPTER_L0_V2)
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event_provider.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/kernel.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/v2/memory.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/queue_api.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/queue_immediate_in_order.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/usm.hpp
@@ -229,6 +234,7 @@ if(UR_BUILD_ADAPTER_L0_V2)
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event_provider_normal.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/event.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/kernel.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/v2/memory.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/queue_api.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/queue_create.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/v2/queue_immediate_in_order.cpp

source/adapters/level_zero/helpers/memory_helpers.cpp

@@ -0,0 +1,33 @@
+//===--------- memory_helpers.cpp - Level Zero Adapter -------------------===//
+//
+// Copyright (C) 2024 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 "memory_helpers.hpp"
+#include "../common.hpp"
+
+ze_memory_type_t getMemoryType(ze_context_handle_t hContext, void *ptr) {
+  // TODO: use UMF once
+  // https://github.com/oneapi-src/unified-memory-framework/issues/687 is
+  // implemented
+  ZeStruct<ze_memory_allocation_properties_t> zeMemoryAllocationProperties;
+  ZE2UR_CALL_THROWS(zeMemGetAllocProperties,
+                    (hContext, ptr, &zeMemoryAllocationProperties, nullptr));
+  return zeMemoryAllocationProperties.type;
+}
+
+bool maybeImportUSM(ze_driver_handle_t hTranslatedDriver,
+                    ze_context_handle_t hContext, void *ptr, size_t size) {
+  if (ZeUSMImport.Enabled && ptr != nullptr &&
+      getMemoryType(hContext, ptr) == ZE_MEMORY_TYPE_UNKNOWN) {
+    // Promote the host ptr to USM host memory
+    ZeUSMImport.doZeUSMImport(hTranslatedDriver, ptr, size);
+    return true;
+  }
+  return false;
+}

source/adapters/level_zero/helpers/memory_helpers.hpp

@@ -0,0 +1,23 @@
+//===--------- memory_helpers.hpp - Level Zero Adapter -------------------===//
+//
+// Copyright (C) 2024 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
+//
+//===----------------------------------------------------------------------===//
+#pragma once
+
+#include <ur_api.h>
+#include <ze_api.h>
+
+// If USM Import feature is enabled and hostptr is supplied,
+// import the hostptr if not already imported into USM.
+// Data transfer rate is maximized when both source and destination
+// are USM pointers. Promotion of the host pointer to USM thus
+// optimizes data transfer performance.
+bool maybeImportUSM(ze_driver_handle_t hTranslatedDriver,
+                    ze_context_handle_t hContext, void *ptr, size_t size);
+
+ze_memory_type_t getMemoryType(ze_context_handle_t hContext, void *ptr);

source/adapters/level_zero/memory.cpp

@@ -15,6 +15,7 @@
 
 #include "context.hpp"
 #include "event.hpp"
+#include "helpers/memory_helpers.hpp"
 #include "image.hpp"
 #include "logger/ur_logger.hpp"
 #include "queue.hpp"
@@ -1599,30 +1600,11 @@ ur_result_t urMemBufferCreate(
     Host = Properties->pHost;
   }
 
-  // If USM Import feature is enabled and hostptr is supplied,
-  // import the hostptr if not already imported into USM.
-  // Data transfer rate is maximized when both source and destination
-  // are USM pointers. Promotion of the host pointer to USM thus
-  // optimizes data transfer performance.
   bool HostPtrImported = false;
-  if (ZeUSMImport.Enabled && Host != nullptr &&
-      (Flags & UR_MEM_FLAG_USE_HOST_POINTER) != 0) {
-    // Query memory type of the host pointer
-    ze_device_handle_t ZeDeviceHandle;
-    ZeStruct<ze_memory_allocation_properties_t> ZeMemoryAllocationProperties;
-    ZE2UR_CALL(zeMemGetAllocProperties,
-               (Context->ZeContext, Host, &ZeMemoryAllocationProperties,
-                &ZeDeviceHandle));
-
-    // If not shared of any type, we can import the ptr
-    if (ZeMemoryAllocationProperties.type == ZE_MEMORY_TYPE_UNKNOWN) {
-      // Promote the host ptr to USM host memory
-      ze_driver_handle_t driverHandle =
-          Context->getPlatform()->ZeDriverHandleExpTranslated;
-      ZeUSMImport.doZeUSMImport(driverHandle, Host, Size);
-      HostPtrImported = true;
-    }
-  }
+  if (Flags & UR_MEM_FLAG_USE_HOST_POINTER)
+    HostPtrImported =
+        maybeImportUSM(Context->getPlatform()->ZeDriverHandleExpTranslated,
+                       Context->ZeContext, Host, Size);
 
   _ur_buffer *Buffer = nullptr;
   auto HostPtrOrNull = (Flags & UR_MEM_FLAG_USE_HOST_POINTER)

source/adapters/level_zero/v2/api.cpp

@@ -49,46 +49,13 @@ ur_result_t urMemImageCreate(ur_context_handle_t hContext, ur_mem_flags_t flags,
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
-ur_result_t urMemBufferCreate(ur_context_handle_t hContext,
-                              ur_mem_flags_t flags, size_t size,
-                              const ur_buffer_properties_t *pProperties,
-                              ur_mem_handle_t *phBuffer) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
-ur_result_t urMemRetain(ur_mem_handle_t hMem) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
-ur_result_t urMemRelease(ur_mem_handle_t hMem) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
-ur_result_t urMemBufferPartition(ur_mem_handle_t hBuffer, ur_mem_flags_t flags,
-                                 ur_buffer_create_type_t bufferCreateType,
-                                 const ur_buffer_region_t *pRegion,
-                                 ur_mem_handle_t *phMem) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
 ur_result_t urMemGetNativeHandle(ur_mem_handle_t hMem,
                                  ur_device_handle_t hDevice,
                                  ur_native_handle_t *phNativeMem) {
   logger::error("{} function not implemented!", __FUNCTION__);
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
-ur_result_t urMemBufferCreateWithNativeHandle(
-    ur_native_handle_t hNativeMem, ur_context_handle_t hContext,
-    const ur_mem_native_properties_t *pProperties, ur_mem_handle_t *phMem) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
 ur_result_t urMemImageCreateWithNativeHandle(
     ur_native_handle_t hNativeMem, ur_context_handle_t hContext,
     const ur_image_format_t *pImageFormat, const ur_image_desc_t *pImageDesc,
@@ -258,14 +225,6 @@ urKernelSetArgSampler(ur_kernel_handle_t hKernel, uint32_t argIndex,
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
-ur_result_t
-urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
-                     const ur_kernel_arg_mem_obj_properties_t *pProperties,
-                     ur_mem_handle_t hArgValue) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
 ur_result_t urKernelSetSpecializationConstants(
     ur_kernel_handle_t hKernel, uint32_t count,
     const ur_specialization_constant_info_t *pSpecConstants) {

source/adapters/level_zero/v2/kernel.cpp

@@ -12,13 +12,14 @@
 
 #include "context.hpp"
 #include "kernel.hpp"
+#include "memory.hpp"
 
 #include "../device.hpp"
 #include "../platform.hpp"
 #include "../program.hpp"
 #include "../ur_interface_loader.hpp"
 
-ur_single_device_kernel_t::ur_single_device_kernel_t(ze_device_handle_t hDevice,
+ur_single_device_kernel_t::ur_single_device_kernel_t(ur_device_handle_t hDevice,
                                                      ze_kernel_handle_t hKernel,
                                                      bool ownZeHandle)
     : hDevice(hDevice), hKernel(hKernel, ownZeHandle) {
@@ -54,7 +55,7 @@ ur_kernel_handle_t_::ur_kernel_handle_t_(ur_program_handle_t hProgram,
     assert(urDevice != hProgram->Context->getDevices().end());
     auto deviceId = (*urDevice)->Id.value();
 
-    deviceKernels[deviceId].emplace(zeDevice, zeKernel, true);
+    deviceKernels[deviceId].emplace(*urDevice, zeKernel, true);
   }
   completeInitialization();
 }
@@ -118,7 +119,7 @@ ur_kernel_handle_t_::getZeHandle(ur_device_handle_t hDevice) {
     auto &kernel = deviceKernels[0].value();
 
     // hDevice is nullptr for native handle
-    if ((kernel.hDevice != nullptr && kernel.hDevice != hDevice->ZeDevice)) {
+    if ((kernel.hDevice != nullptr && kernel.hDevice != hDevice)) {
       throw UR_RESULT_ERROR_INVALID_DEVICE;
     }
 
@@ -239,6 +240,16 @@ ur_result_t ur_kernel_handle_t_::setExecInfo(ur_kernel_exec_info_t propName,
   return UR_RESULT_SUCCESS;
 }
 
+std::vector<ur_device_handle_t> ur_kernel_handle_t_::getDevices() const {
+  std::vector<ur_device_handle_t> devices;
+  for (size_t i = 0; i < deviceKernels.size(); ++i) {
+    if (deviceKernels[i].has_value()) {
+      devices.push_back(deviceKernels[i].value().hDevice);
+    }
+  }
+  return devices;
+}
+
 namespace ur::level_zero {
 ur_result_t urKernelCreate(ur_program_handle_t hProgram,
                            const char *pKernelName,
@@ -291,6 +302,28 @@ ur_result_t urKernelSetArgPointer(
   return hKernel->setArgPointer(argIndex, pProperties, pArgValue);
 }
 
+ur_result_t
+urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
+                     const ur_kernel_arg_mem_obj_properties_t *pProperties,
+                     ur_mem_handle_t hArgValue) {
+  TRACK_SCOPE_LATENCY("ur_kernel_handle_t_::setArgMemObj");
+
+  // TODO: support properties
+  std::ignore = pProperties;
+
+  auto kernelDevices = hKernel->getDevices();
+  if (kernelDevices.size() == 1) {
+    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;
+  }
+}
+
 ur_result_t urKernelSetExecInfo(
     ur_kernel_handle_t hKernel,     ///< [in] handle of the kernel object
     ur_kernel_exec_info_t propName, ///< [in] name of the execution attribute

source/adapters/level_zero/v2/kernel.hpp

@@ -15,11 +15,11 @@
 #include "common.hpp"
 
 struct ur_single_device_kernel_t {
-  ur_single_device_kernel_t(ze_device_handle_t hDevice,
+  ur_single_device_kernel_t(ur_device_handle_t hDevice,
                             ze_kernel_handle_t hKernel, bool ownZeHandle);
   ur_result_t release();
 
-  ze_device_handle_t hDevice;
+  ur_device_handle_t hDevice;
   v2::raii::ze_kernel_handle_t hKernel;
   mutable ZeCache<ZeStruct<ze_kernel_properties_t>> zeKernelProperties;
 };
@@ -40,6 +40,9 @@ struct ur_kernel_handle_t_ : _ur_object {
   // Get program handle of the kernel.
   ur_program_handle_t getProgramHandle() const;
 
+  // Get devices the kernel is built for.
+  std::vector<ur_device_handle_t> getDevices() const;
+
   // Get name of the kernel.
   const std::string &getName() const;