[SYCL] Implement register_host_memory extension API

llvm/include/llvm/SYCLLowerIR/DeviceConfigFile.td

@@ -105,6 +105,7 @@ def Aspectext_intel_xe_cores_per_cluster : Aspect<"ext_intel_xe_cores_per_cluste
 def Aspectext_intel_eus_per_xe_core : Aspect<"ext_intel_eus_per_xe_core">;
 def Aspectext_intel_max_lanes_per_hw_thread : Aspect<"ext_intel_max_lanes_per_hw_thread">;
 def Aspectext_oneapi_ipc_physical_memory : Aspect<"ext_oneapi_ipc_physical_memory">;
+def Aspectext_oneapi_register_host_memory : Aspect<"ext_oneapi_register_host_memory">;
 
 // Deprecated aspects
 def AspectInt64_base_atomics : Aspect<"int64_base_atomics">;
@@ -191,7 +192,8 @@ def : TargetInfo<"__TestAspectList",
     Aspectext_intel_xe_cores_per_cluster,
     Aspectext_intel_eus_per_xe_core,
     Aspectext_intel_max_lanes_per_hw_thread,
-    Aspectext_oneapi_ipc_physical_memory],
+    Aspectext_oneapi_ipc_physical_memory,
+    Aspectext_oneapi_register_host_memory],
     []>;
 // This definition serves the only purpose of testing whether the deprecated aspect list defined in here and in SYCL RT
 // match.

sycl/include/sycl/ext/oneapi/experimental/register_host_memory.hpp

@@ -0,0 +1,82 @@
+//==--- register_host_memory.hpp - SYCL host memory registration extension -==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <sycl/detail/export.hpp> // for __SYCL_EXPORT
+#include <sycl/ext/oneapi/properties.hpp>
+
+#include <cstddef> // for size_t
+#include <cstdint> // for uint32_t
+#include <type_traits>
+
+namespace sycl {
+inline namespace _V1 {
+class context;
+
+namespace ext::oneapi::experimental {
+
+// Indicates that device code will only read from the registered range. Device
+// writes to a range registered with this property are undefined behavior.
+struct read_only_key : detail::compile_time_property_key<
+                           detail::PropKind::RegisterHostMemoryReadOnly> {
+  using value_t = property_value<read_only_key>;
+};
+
+inline constexpr read_only_key::value_t read_only;
+
+namespace detail {
+// Implementation-internal flags describing a host memory registration. They
+// are translated to UR flags in the runtime library.
+enum register_host_memory_flags : uint32_t {
+  register_host_memory_flag_read_only = 1u << 0,
+};
+
+// Non-templated implementation entry points, defined in the SYCL runtime
+// library. Flags is a bitwise OR of register_host_memory_flags values.
+__SYCL_EXPORT void register_host_memory(void *Ptr, size_t NumBytes,
+                                        const context &Ctxt, uint32_t Flags);
+__SYCL_EXPORT void unregister_host_memory(void *Ptr, const context &Ctxt);
+
+// Lowers a compile-time property list to the runtime flags word.
+template <typename Properties> uint32_t getRegisterHostMemoryFlags() {
+  uint32_t Flags = 0;
+  if constexpr (std::decay_t<Properties>::template has_property<
+                    read_only_key>())
+    Flags |= register_host_memory_flag_read_only;
+  return Flags;
+}
+} // namespace detail
+
+/// Registers the existing host memory range \p ptr of \p numBytes bytes with
+/// \p ctxt so that it behaves like a USM host allocation. See
+/// sycl_ext_oneapi_register_host_memory for the full semantics.
+///
+/// \p ptr and \p numBytes must both be aligned to the host page size, \p ptr
+/// must not be null, \p numBytes must not be zero, and every device in \p ctxt
+/// must have aspect::ext_oneapi_register_host_memory.
+template <typename Properties = empty_properties_t>
+std::enable_if_t<is_property_list_v<std::decay_t<Properties>>>
+register_host_memory(void *ptr, size_t numBytes, const context &ctxt,
+                     Properties props = {}) {
+  (void)props;
+  detail::register_host_memory(
+      ptr, numBytes, ctxt, detail::getRegisterHostMemoryFlags<Properties>());
+}
+
+/// Unregisters a host memory range previously registered with
+/// register_host_memory. \p ptr must be the exact base pointer that was passed
+/// to register_host_memory with the same \p ctxt, and the registration must
+/// still be in effect. This does not free or unmap the underlying host memory.
+inline void unregister_host_memory(void *ptr, const context &ctxt) {
+  detail::unregister_host_memory(ptr, ctxt);
+}
+
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/ext/oneapi/properties/property.hpp

@@ -193,8 +193,9 @@ enum PropKind : uint32_t {
   ZeroInit = 48,
   FastLink = 49,
   PhysicalMemoryEnableIPC = 50,
+  RegisterHostMemoryReadOnly = 51,
   // PropKindSize must always be the last value.
-  PropKindSize = 51,
+  PropKindSize = 52,
 };
 
 template <typename PropertyT> struct PropertyToKind {

sycl/include/sycl/info/aspects.def

@@ -91,3 +91,4 @@ __SYCL_ASPECT(ext_intel_xe_cores_per_cluster, 93)
 __SYCL_ASPECT(ext_intel_eus_per_xe_core, 94)
 __SYCL_ASPECT(ext_intel_max_lanes_per_hw_thread, 95)
 __SYCL_ASPECT(ext_oneapi_ipc_physical_memory, 96)
+__SYCL_ASPECT(ext_oneapi_register_host_memory, 97)

sycl/include/sycl/sycl.hpp

@@ -127,6 +127,7 @@ can be disabled by setting SYCL_DISABLE_FSYCL_SYCLHPP_WARNING macro.")
 #include <sycl/ext/oneapi/experimental/profiling_tag.hpp>
 #include <sycl/ext/oneapi/experimental/raw_kernel_arg.hpp>
 #include <sycl/ext/oneapi/experimental/reduction_properties.hpp>
+#include <sycl/ext/oneapi/experimental/register_host_memory.hpp>
 #include <sycl/ext/oneapi/experimental/root_group.hpp>
 #include <sycl/ext/oneapi/experimental/syclbin_kernel_bundle.hpp>
 #include <sycl/ext/oneapi/experimental/syclbin_properties.hpp>

sycl/source/detail/device_impl.hpp

@@ -1391,6 +1391,11 @@ class device_impl {
                  UR_DEVICE_INFO_EXTERNAL_MEMORY_IMPORT_SUPPORT_EXP>()
           .value_or(0);
     }
+    CASE(ext_oneapi_register_host_memory) {
+      return get_info_impl_nocheck<
+                 UR_DEVICE_INFO_USM_HOST_ALLOC_REGISTER_SUPPORT_EXP>()
+          .value_or(0);
+    }
     CASE(ext_oneapi_external_semaphore_import) {
       return get_info_impl_nocheck<
                  UR_DEVICE_INFO_EXTERNAL_SEMAPHORE_IMPORT_SUPPORT_EXP>()

sycl/source/detail/ur_device_info_ret_types.inc

@@ -209,4 +209,5 @@ MAP(UR_DEVICE_INFO_CLOCK_DEVICE_SUPPORT_EXP, ur_bool_t)
 MAP(UR_DEVICE_INFO_IS_INTEGRATED_GPU, ur_bool_t)
 MAP(UR_DEVICE_INFO_DEVICE_WAIT_SUPPORT_EXP, ur_bool_t)
 MAP(UR_DEVICE_INFO_IPC_PHYSICAL_MEMORY_SUPPORT_EXP, ur_bool_t)
+MAP(UR_DEVICE_INFO_USM_HOST_ALLOC_REGISTER_SUPPORT_EXP, ur_bool_t)
 // clang-format on

sycl/source/detail/usm/usm_impl.cpp

@@ -14,14 +14,23 @@
 #include <sycl/detail/ur.hpp>
 #include <sycl/device.hpp>
 #include <sycl/ext/intel/experimental/usm_properties.hpp>
+#include <sycl/ext/oneapi/experimental/register_host_memory.hpp>
 #include <sycl/ext/oneapi/memcpy2d.hpp>
 #include <sycl/usm.hpp>
 
 #include <array>
 #include <cassert>
+#include <cstdint>
 #include <cstdlib>
+#include <limits>
 #include <memory>
 
+#ifdef _WIN32
+#include <windows.h>
+#else
+#include <unistd.h>
+#endif
+
 #ifdef XPTI_ENABLE_INSTRUMENTATION
 // Include the headers necessary for emitting
 // traces using the trace framework
@@ -609,6 +618,105 @@ void release_from_device_copy(const void *Ptr, const queue &Queue) {
   release_from_usm_device_copy(Ptr, Queue.get_context());
 }
 
+// Host memory registration APIs, see sycl_ext_oneapi_register_host_memory.
+
+namespace detail {
+
+// Throws errc::feature_not_supported unless every device in the context
+// reports aspect::ext_oneapi_register_host_memory.
+static void checkRegisterHostMemorySupport(const context &Ctxt) {
+  detail::context_impl &CtxtImpl = *detail::getSyclObjImpl(Ctxt);
+  for (detail::device_impl &Dev : CtxtImpl.getDevices()) {
+    if (!Dev.has(aspect::ext_oneapi_register_host_memory))
+      throw sycl::exception(
+          make_error_code(errc::feature_not_supported),
+          "At least one device in the context does not support registering "
+          "host memory (aspect::ext_oneapi_register_host_memory).");
+  }
+}
+
+// Maps a failed UR result from the host memory registration APIs to a
+// sycl::exception with the error code mandated by the extension specification.
+// Invalid argument conditions map to errc::invalid; anything else is a backend
+// error.
+static void throwRegisterHostMemoryError(ur_result_t Err, const char *What) {
+  errc Code;
+  switch (Err) {
+  case UR_RESULT_ERROR_INVALID_NULL_POINTER:
+  case UR_RESULT_ERROR_INVALID_VALUE:
+  case UR_RESULT_ERROR_INVALID_ARGUMENT:
+    Code = errc::invalid;
+    break;
+  default:
+    Code = errc::runtime;
+    break;
+  }
+  throw detail::set_ur_error(sycl::exception(make_error_code(Code), What), Err);
+}
+
+void register_host_memory(void *Ptr, size_t NumBytes, const context &Ctxt,
+                          uint32_t Flags) {
+  if (Ptr == nullptr)
+    throw sycl::exception(make_error_code(errc::invalid),
+                          "register_host_memory: pointer must not be null.");
+  if (NumBytes == 0)
+    throw sycl::exception(make_error_code(errc::invalid),
+                          "register_host_memory: size must not be zero.");
+  static const size_t PageSize = []() {
+#ifdef _WIN32
+    SYSTEM_INFO Info;
+    GetSystemInfo(&Info);
+    return static_cast<size_t>(Info.dwPageSize);
+#else
+    return static_cast<size_t>(sysconf(_SC_PAGESIZE));
+#endif
+  }();
+  if (reinterpret_cast<uintptr_t>(Ptr) % PageSize != 0)
+    throw sycl::exception(
+        make_error_code(errc::invalid),
+        "register_host_memory: pointer must be aligned to the host page size.");
+  if (NumBytes % PageSize != 0)
+    throw sycl::exception(
+        make_error_code(errc::invalid),
+        "register_host_memory: size must be a multiple of the host page size.");
+  if (NumBytes >
+      std::numeric_limits<uintptr_t>::max() - reinterpret_cast<uintptr_t>(Ptr))
+    throw sycl::exception(make_error_code(errc::invalid),
+                          "register_host_memory: range is not representable in "
+                          "the host address space.");
+  checkRegisterHostMemorySupport(Ctxt);
+
+  ur_exp_usm_host_alloc_register_properties_t Props = {
+      UR_STRUCTURE_TYPE_EXP_USM_HOST_ALLOC_REGISTER_PROPERTIES,
+      /*pNext=*/nullptr,
+      /*flags=*/0};
+  if (Flags & register_host_memory_flag_read_only)
+    Props.flags |= UR_EXP_USM_HOST_ALLOC_REGISTER_FLAG_READ_ONLY;
+
+  auto [urCtx, Adapter] = get_ur_handles(Ctxt);
+  ur_result_t Err =
+      Adapter->call_nocheck<detail::UrApiKind::urUSMHostAllocRegisterExp>(
+          urCtx, Ptr, NumBytes, &Props);
+  if (Err != UR_RESULT_SUCCESS)
+    throwRegisterHostMemoryError(Err, "register_host_memory failed.");
+}
+
+void unregister_host_memory(void *Ptr, const context &Ctxt) {
+  if (Ptr == nullptr)
+    throw sycl::exception(make_error_code(errc::invalid),
+                          "unregister_host_memory: pointer must not be null.");
+  checkRegisterHostMemorySupport(Ctxt);
+
+  auto [urCtx, Adapter] = get_ur_handles(Ctxt);
+  ur_result_t Err =
+      Adapter->call_nocheck<detail::UrApiKind::urUSMHostAllocUnregisterExp>(
+          urCtx, Ptr);
+  if (Err != UR_RESULT_SUCCESS)
+    throwRegisterHostMemoryError(Err, "unregister_host_memory failed.");
+}
+
+} // namespace detail
+
 void *malloc_device(size_t numBytes, const device &syclDevice,
                     const property_list &propList) {
   sycl::context ctxt = syclDevice.get_platform().khr_get_default_context();

sycl/source/feature_test.hpp.in

@@ -77,6 +77,7 @@ inline namespace _V1 {
 #define SYCL_EXT_ONEAPI_KERNEL_ARG_PROPERTIES 1
 #define SYCL_EXT_ONEAPI_ANNOTATED_PTR 1
 #define SYCL_EXT_ONEAPI_COPY_OPTIMIZE 1
+#define SYCL_EXT_ONEAPI_REGISTER_HOST_MEMORY 1
 #define SYCL_EXT_ONEAPI_VIRTUAL_MEM 1
 #define SYCL_EXT_ONEAPI_USM_MALLOC_PROPERTIES 1
 #define SYCL_EXT_INTEL_CACHE_CONFIG 1