[AsyncAlloc][SYCL][CUDA][Exp] Initial device side implementation for the sycl_ext_oneapi_async_memory_alloc extension

llvm/include/llvm/SYCLLowerIR/DeviceConfigFile.td

@@ -87,6 +87,7 @@ def AspectExt_intel_fpga_task_sequence : Aspect<"ext_intel_fpga_task_sequence">;
 def AspectExt_oneapi_atomic16 : Aspect<"ext_oneapi_atomic16">;
 def AspectExt_oneapi_virtual_functions : Aspect<"ext_oneapi_virtual_functions">;
 def AspectExt_intel_spill_memory_size : Aspect<"ext_intel_spill_memory_size">;
+def AspectExt_oneapi_async_memory_alloc : Aspect<"ext_oneapi_async_memory_alloc">;
 // Deprecated aspects
 def AspectInt64_base_atomics : Aspect<"int64_base_atomics">;
 def AspectInt64_extended_atomics : Aspect<"int64_extended_atomics">;
@@ -152,7 +153,8 @@ def : TargetInfo<"__TestAspectList",
     AspectExt_intel_fpga_task_sequence,
     AspectExt_oneapi_atomic16,
     AspectExt_oneapi_virtual_functions,
-    AspectExt_intel_spill_memory_size],
+    AspectExt_intel_spill_memory_size,
+    AspectExt_oneapi_async_memory_alloc],
     []>;
 // This definition serves the only purpose of testing whether the deprecated aspect list defined in here and in SYCL RT
 // match.

sycl/include/sycl/context.hpp

@@ -16,6 +16,7 @@
 #include <sycl/detail/owner_less_base.hpp>    // for OwnerLessBase
 #include <sycl/platform.hpp>                  // for platform
 #include <sycl/property_list.hpp>             // for property_list
+#include <sycl/usm/usm_enums.hpp>             // for usm::alloc
 #include <ur_api.h>                           // for ur_native_handle_t
 
 #ifdef __SYCL_INTERNAL_API
@@ -36,6 +37,10 @@ inline namespace _V1 {
 class device;
 class platform;
 
+namespace ext::oneapi::experimental {
+class memory_pool;
+} // namespace ext::oneapi::experimental
+
 namespace detail {
 class context_impl;
 }
@@ -245,6 +250,20 @@ class __SYCL_EXPORT context : public detail::OwnerLessBase<context> {
   /// \return a vector of valid SYCL device instances.
   std::vector<device> get_devices() const;
 
+  /// Gets default memory pool associated with a device and context.
+  ///
+  /// \return a memory pool for a particular device and context.
+  sycl::ext::oneapi::experimental::memory_pool
+  ext_oneapi_get_default_memory_pool(const device &dev,
+                                     const sycl::usm::alloc &kind) const;
+
+  /// Gets default memory pool associated with the context -- with the
+  /// allocation kind usm::alloc::host.
+  ///
+  /// \return a memory pool associated with this context.
+  sycl::ext::oneapi::experimental::memory_pool
+  ext_oneapi_get_default_memory_pool() const;
+
 private:
   /// Constructs a SYCL context object from a valid context_impl instance.
   context(std::shared_ptr<detail::context_impl> Impl);

sycl/include/sycl/detail/cg_types.hpp

@@ -65,6 +65,8 @@ enum class CGType : unsigned int {
   SemaphoreSignal = 25,
   ProfilingTag = 26,
   EnqueueNativeCommand = 27,
+  AsyncAlloc = 28,
+  AsyncFree = 29,
 };
 
 template <typename, typename T> struct check_fn_signature {

sycl/include/sycl/detail/property_helper.hpp

@@ -51,8 +51,10 @@ enum DataLessPropKind {
   GraphDependOnAllLeaves = 24,
   GraphUpdatable = 25,
   GraphEnableProfiling = 26,
+  MemPoolReadOnly = 27,
+  MemPoolZeroInit = 28,
   // Indicates the last known dataless property.
-  LastKnownDataLessPropKind = 26,
+  LastKnownDataLessPropKind = 28,
   // Exceeding 32 may cause ABI breaking change on some of OSes.
   DataLessPropKindSize = 32
 };
@@ -67,7 +69,9 @@ enum PropWithDataKind {
   AccPropBufferLocation = 5,
   QueueComputeIndex = 6,
   GraphNodeDependencies = 7,
-  PropWithDataKindSize = 8
+  MemPoolInitialThreshold = 8,
+  MemPoolMaximumSize = 9,
+  PropWithDataKindSize = 10
 };
 
 // Base class for dataless properties, needed to check that the type of an

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

@@ -0,0 +1,94 @@
+//==----------- async_alloc.hpp --- SYCL asynchronous allocation -----------==//
+//
+// 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/common.hpp> // for code_location
+#include <sycl/handler.hpp>       // for handler
+#include <sycl/queue.hpp>         // for queue
+#include <sycl/usm/usm_enums.hpp> // for usm::alloc
+
+namespace sycl {
+inline namespace _V1 {
+namespace ext::oneapi::experimental {
+
+// Forward declare memory_pool.
+class memory_pool;
+
+/**
+ * @brief  Asynchronousy allocate memory from a default pool.
+ *
+ * @param  q The queue with which to enqueue the asynchronous allocation.
+ * @param  kind The kind of memory pool allocation - device, host, shared, etc.
+ * @param  size The size in bytes to allocate.
+ *
+ * @return Generic pointer to allocated USM memory.
+ */
+__SYCL_EXPORT void *async_malloc(const sycl::queue &q, sycl::usm::alloc kind,
+                                 size_t size,
+                                 const sycl::detail::code_location &CodeLoc =
+                                     sycl::detail::code_location::current());
+
+/**
+ * @brief  Asynchronously allocate memory from a default pool.
+ *
+ * @param  h The handler with which to enqueue the asynchronous allocation.
+ * @param  kind The kind of memory pool allocation - device, host, shared, etc.
+ * @param  size The size in bytes to allocate.
+ *
+ * @return Generic pointer to allocated USM memory.
+ */
+__SYCL_EXPORT void *async_malloc(sycl::handler &h, sycl::usm::alloc kind,
+                                 size_t size);
+
+/**
+ * @brief  Asynchronously allocate memory from a specified pool.
+ *
+ * @param  q The queue with which to enqueue the asynchronous allocation.
+ * @param  size The size in bytes to allocate.
+ * @param  pool The pool with which to allocate from.
+ *
+ * @return Generic pointer to allocated USM memory.
+ */
+__SYCL_EXPORT void *
+async_malloc_from_pool(const sycl::queue &q, size_t size, memory_pool &pool,
+                       const sycl::detail::code_location &CodeLoc =
+                           sycl::detail::code_location::current());
+
+/**
+ * @brief  Asynchronously allocate memory from a specified pool.
+ *
+ * @param  h The handler with which to enqueue the asynchronous allocation.
+ * @param  size The size in bytes to allocate.
+ * @param  pool The pool with which to allocate from.
+ *
+ * @return Generic pointer to allocated USM memory.
+ */
+__SYCL_EXPORT void *async_malloc_from_pool(sycl::handler &h, size_t size,
+                                           memory_pool &pool);
+
+/**
+ * @brief  Asynchronously free memory.
+ *
+ * @param  q The queue with which to enqueue the asynchronous free.
+ * @param  ptr The generic pointer to be freed.
+ */
+__SYCL_EXPORT void async_free(const sycl::queue &q, void *ptr,
+                              const sycl::detail::code_location &CodeLoc =
+                                  sycl::detail::code_location::current());
+
+/**
+ * @brief  Asynchronously free memory.
+ *
+ * @param  h The handler with which to enqueue the asynchronous free.
+ * @param  ptr The generic pointer to be freed.
+ */
+__SYCL_EXPORT void async_free(sycl::handler &h, void *ptr);
+
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/ext/oneapi/experimental/async_alloc/memory_pool.hpp

@@ -0,0 +1,199 @@
+//==----------- memory_pool.hpp --- SYCL asynchronous allocation -----------==//
+//
+// 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/context.hpp>       // for context
+#include <sycl/device.hpp>        // for device
+#include <sycl/queue.hpp>         // for queue
+#include <sycl/usm/usm_enums.hpp> // for usm::alloc
+
+namespace sycl {
+inline namespace _V1 {
+namespace ext::oneapi::experimental {
+namespace property {
+
+// Property that determines the initial threshold of a memory pool.
+struct initial_threshold : public sycl::detail::PropertyWithData<
+                               sycl::detail::MemPoolInitialThreshold> {
+public:
+  initial_threshold(size_t initialThreshold)
+      : initialThreshold(initialThreshold) {};
+  size_t get_initial_threshold() { return initialThreshold; }
+
+private:
+  size_t initialThreshold;
+};
+
+// Property that determines the maximum size of a memory pool.
+struct maximum_size
+    : public sycl::detail::PropertyWithData<sycl::detail::MemPoolMaximumSize> {
+public:
+  maximum_size(size_t maxSize) : maxSize(maxSize) {};
+  size_t get_maximum_size() { return maxSize; }
+
+private:
+  size_t maxSize;
+};
+
+// Property that provides a performance hint that all allocations from this pool
+// will only be read from within SYCL kernel functions.
+struct read_only
+    : public sycl::detail::DataLessProperty<sycl::detail::MemPoolReadOnly> {
+public:
+  read_only() = default;
+};
+
+// Property that initial allocations to a pool (not subsequent allocations from
+// prior frees) are iniitialised to zero.
+struct zero_init
+    : public sycl::detail::DataLessProperty<sycl::detail::MemPoolZeroInit> {
+public:
+  zero_init() = default;
+};
+} // namespace property
+
+namespace detail {
+class memory_pool_impl {
+public:
+  memory_pool_impl(const sycl::context &ctx, const sycl::device &dev,
+                   const sycl::usm::alloc kind, const property_list &props);
+  memory_pool_impl(const sycl::context &ctx, const sycl::device &dev,
+                   const sycl::usm::alloc kind, ur_usm_pool_handle_t poolHandle,
+                   const bool isDefaultPool, const property_list &props);
+
+  ~memory_pool_impl();
+
+  memory_pool_impl(const memory_pool_impl &) = delete;
+  memory_pool_impl &operator=(const memory_pool_impl &) = delete;
+
+  ur_usm_pool_handle_t get_handle() const { return MPoolHandle; }
+  sycl::device get_device() const { return MDevice; }
+  sycl::context get_context() const {
+    return sycl::detail::createSyclObjFromImpl<sycl::context>(MContextImplPtr);
+  }
+  sycl::usm::alloc get_alloc_kind() const { return MKind; }
+  const property_list &getPropList() const { return MPropList; }
+
+  // Returns backend specific values.
+  size_t get_threshold() const;
+  size_t get_reserved_size_current() const;
+  size_t get_reserved_size_high() const;
+  size_t get_used_size_current() const;
+  size_t get_used_size_high() const;
+
+  void set_new_threshold(size_t newThreshold);
+  void reset_reserved_size_high();
+  void reset_used_size_high();
+  void trim_to(size_t minBytesToKeep);
+
+private:
+  std::shared_ptr<sycl::detail::context_impl> MContextImplPtr;
+  sycl::device MDevice;
+  sycl::usm::alloc MKind;
+  ur_usm_pool_handle_t MPoolHandle{0};
+  bool MIsDefaultPool = false;
+  property_list MPropList;
+};
+} // namespace detail
+
+/// Memory pool
+class __SYCL_EXPORT memory_pool {
+
+public:
+  memory_pool(const sycl::context &ctx, const property_list &props = {});
+
+  memory_pool(const sycl::context &ctx, const sycl::device &dev,
+              const sycl::usm::alloc kind, const property_list &props = {});
+
+  memory_pool(const sycl::queue &q, const sycl::usm::alloc kind,
+              const property_list &props = {});
+
+  memory_pool(const sycl::context &ctx, const void *ptr, size_t size,
+              const property_list &props = {});
+
+  ~memory_pool() = default;
+
+  // Copy constructible/assignable, move constructible/assignable.
+  memory_pool(const memory_pool &) = default;
+  memory_pool(memory_pool &&) = default;
+  memory_pool &operator=(const memory_pool &) = default;
+  memory_pool &operator=(memory_pool &&) = default;
+
+  // Equality comparison.
+  bool operator==(const memory_pool &rhs) const { return impl == rhs.impl; }
+  bool operator!=(const memory_pool &rhs) const { return !(*this == rhs); }
+
+  // Impl handles getters and setters.
+  sycl::context get_context() const { return impl->get_context(); }
+  sycl::device get_device() const { return impl->get_device(); }
+  sycl::usm::alloc get_alloc_kind() const { return impl->get_alloc_kind(); }
+
+  size_t get_threshold() const;
+  size_t get_reserved_size_current() const;
+  size_t get_reserved_size_high() const;
+  size_t get_used_size_current() const;
+  size_t get_used_size_high() const;
+
+  void set_new_threshold(size_t newThreshold);
+  void reset_reserved_size_high();
+  void reset_used_size_high();
+  void trim_to(size_t minBytesToKeep);
+
+  // Property getters.
+  template <typename propertyT> bool has_property() const noexcept {
+    return getPropList().template has_property<propertyT>();
+  }
+  template <typename propertyT> propertyT get_property() const {
+    return getPropList().template get_property<propertyT>();
+  }
+
+protected:
+  std::shared_ptr<detail::memory_pool_impl> impl;
+
+  template <class Obj>
+  friend const decltype(Obj::impl) &
+  sycl::detail::getSyclObjImpl(const Obj &SyclObject);
+
+  template <class T>
+  friend T sycl::detail::createSyclObjFromImpl(decltype(T::impl) ImplObj);
+
+  const property_list &getPropList() const;
+
+  memory_pool(std::shared_ptr<detail::memory_pool_impl> Impl) : impl(Impl) {}
+};
+
+} // namespace ext::oneapi::experimental
+
+template <>
+struct is_property<sycl::ext::oneapi::experimental::property::initial_threshold>
+    : std::true_type {};
+
+template <>
+struct is_property<sycl::ext::oneapi::experimental::property::maximum_size>
+    : std::true_type {};
+
+template <>
+struct is_property<sycl::ext::oneapi::experimental::property::read_only>
+    : std::true_type {};
+
+template <>
+struct is_property<sycl::ext::oneapi::experimental::property::zero_init>
+    : std::true_type {};
+} // namespace _V1
+} // namespace sycl
+
+namespace std {
+template <> struct hash<sycl::ext::oneapi::experimental::memory_pool> {
+  size_t operator()(
+      const sycl::ext::oneapi::experimental::memory_pool &mem_pool) const {
+    return hash<std::shared_ptr<
+        sycl::ext::oneapi::experimental::detail::memory_pool_impl>>()(
+        sycl::detail::getSyclObjImpl(mem_pool));
+  }
+};
+} // namespace std