Implement interface for external allocators

include/alpaka/mem/buf/Traits.hpp

@@ -1,4 +1,4 @@
-/* Copyright 2025 Alexander Matthes, Benjamin Worpitz, Andrea Bocci, Bernhard Manfred Gruber, Jan Stephan,
+/* Copyright 2026 Alexander Matthes, Benjamin Worpitz, Andrea Bocci, Bernhard Manfred Gruber, Jan Stephan,
  *                Christian Kaever, Maria Michailidi, Simone Balducci
  * SPDX-License-Identifier: MPL-2.0
  */
@@ -10,11 +10,29 @@
 #include "alpaka/mem/view/Traits.hpp"
 #include "alpaka/platform/Traits.hpp"
 
+#include <concepts>
+#include <cstddef>
+#include <type_traits>
+
 namespace alpaka
 {
     //! The CPU device handle.
     class DevCpu;
 
+    namespace concepts
+    {
+
+        template<typename TAllocator>
+        concept CachingAllocator = std::copyable<TAllocator>
+                                   && requires(TAllocator alloc, void* ptr, std::size_t bytes, std::size_t align) {
+                                          {
+                                              alloc.allocate(bytes, align)
+                                          } -> std::same_as<void*>;
+                                          alloc.deallocate(ptr);
+                                      };
+
+    } // namespace concepts
+
     //! The buffer traits.
     namespace trait
     {
@@ -54,6 +72,26 @@ namespace alpaka
         template<typename TPlatform, typename TElem, typename TDim, typename TIdx>
         struct BufAllocManaged;
 
+        //! The caching-allocator-aware memory allocator trait.
+        template<
+            typename TElem,
+            typename TDim,
+            typename TIdx,
+            typename TDev,
+            concepts::CachingAllocator TAllocator,
+            typename TSfinae = void>
+        struct BufAllocCached;
+
+        //! The caching-allocator-aware stream-ordered memory allocator trait.
+        template<
+            typename TElem,
+            typename TDim,
+            typename TIdx,
+            typename TDev,
+            concepts::CachingAllocator TAllocator,
+            typename TSfinae = void>
+        struct AsyncBufAllocCached;
+
         //! The trait to transform a mutable buffer into a constant one.
         template<typename TBuf>
         struct MakeConstBuf;
@@ -71,6 +109,7 @@ namespace alpaka
     //! Allocates memory on the given device.
     //!
     //! \tparam TElem The element type of the returned buffer.
+    //! \tparam TIdx The linear index type of the buffer.
     //! \tparam TExtent The extent type of the buffer.
     //! \tparam TDev The type of device the buffer is allocated on.
     //! \param dev The device to allocate the buffer on.
@@ -80,10 +119,31 @@ namespace alpaka
     ALPAKA_FN_HOST auto allocBuf(TDev const& dev, TExtent const& extent = TExtent())
     {
         using Idx = std::conditional_t<std::is_void_v<TIdx>, Idx<TExtent>, TIdx>;
-
         return trait::BufAlloc<TElem, Dim<TExtent>, Idx, TDev>::allocBuf(dev, extent);
     }
 
+    //! Allocates memory on the given device using a caching allocator.
+    //!
+    //! \tparam TElem The element type of the returned buffer.
+    //! \tparam TIdx The linear index type of the buffer.
+    //! \tparam TExtent The extent type of the buffer.
+    //! \tparam TDev The type of device the buffer is allocated on.
+    //! \tparam TAllocator The type of the caching allocator wrapper.
+    //! \param dev The device to allocate the buffer on.
+    //! \param extent The extent of the buffer.
+    //! \param allocator The caching allocator wrapper to use.
+    //! \return The newly allocated buffer.
+    template<typename TElem, typename TIdx = void, typename TExtent = void, typename TDev = void, typename TAllocator>
+    requires concepts::CachingAllocator<std::remove_cvref_t<TAllocator>>
+    ALPAKA_FN_HOST auto allocBuf(TDev const& dev, TExtent const& extent, TAllocator&& allocator)
+    {
+        using Idx = std::conditional_t<std::is_void_v<TIdx>, Idx<TExtent>, TIdx>;
+        return trait::BufAllocCached<TElem, Dim<TExtent>, Idx, TDev, std::remove_cvref_t<TAllocator>>::allocBuf(
+            dev,
+            extent,
+            std::forward<TAllocator>(allocator));
+    }
+
     //! Allocates stream-ordered memory on the given device.
     //!
     //! \tparam TElem The element type of the returned buffer.
@@ -101,8 +161,24 @@ namespace alpaka
         return trait::AsyncBufAlloc<TElem, Dim<TExtent>, Idx, alpaka::Dev<TQueue>>::allocAsyncBuf(queue, extent);
     }
 
+    //! Allocates stream-ordered memory on the given device using a caching allocator.
+    template<
+        typename TElem,
+        typename TIdx = void,
+        typename TExtent = void,
+        typename TQueue = void,
+        typename TAllocator>
+    requires concepts::CachingAllocator<std::remove_cvref_t<TAllocator>>
+    ALPAKA_FN_HOST auto allocAsyncBuf(TQueue queue, TExtent const& extent, TAllocator&& allocator)
+    {
+        using Idx = std::conditional_t<std::is_void_v<TIdx>, Idx<TExtent>, TIdx>;
+        return trait::
+            AsyncBufAllocCached<TElem, Dim<TExtent>, Idx, alpaka::Dev<TQueue>, std::remove_cvref_t<TAllocator>>::
+                allocAsyncBuf(std::move(queue), extent, std::forward<TAllocator>(allocator));
+    }
+
     /* TODO: Remove this pragma block once support for clang versions <= 13 is removed. These versions are unable to
-       figure out that the template parameters are attached to a C++17 inline variable. */
+         figure out that the template parameters are attached to a C++17 inline variable. */
 #if ALPAKA_COMP_CLANG
 #    pragma clang diagnostic push
 #    pragma clang diagnostic ignored "-Wdocumentation"
@@ -147,6 +223,31 @@ namespace alpaka
         ALPAKA_UNREACHABLE(allocBuf<TElem, TIdx>(getDev(queue), extent));
     }
 
+    //! If supported, allocates stream-ordered memory using a caching allocator; otherwise falls back to cached
+    //! allocBuf.
+    template<
+        typename TElem,
+        typename TIdx = void,
+        typename TExtent = void,
+        typename TQueue = void,
+        typename TAllocator>
+    requires concepts::CachingAllocator<std::remove_cvref_t<TAllocator>>
+    ALPAKA_FN_HOST auto allocAsyncBufIfSupported(TQueue queue, TExtent const& extent, TAllocator&& allocator)
+    {
+        using Idx = std::conditional_t<std::is_void_v<TIdx>, Idx<TExtent>, TIdx>;
+
+        if constexpr(hasAsyncBufSupport<alpaka::Dev<TQueue>, Dim<TExtent>>)
+        {
+            return allocAsyncBuf<TElem, Idx>(std::move(queue), extent, std::forward<TAllocator>(allocator));
+        }
+        else
+        {
+            return allocBuf<TElem, Idx>(getDev(queue), extent, std::forward<TAllocator>(allocator));
+        }
+
+        ALPAKA_UNREACHABLE(allocBuf<TElem, TIdx>(getDev(queue), extent));
+    }
+
     //! Allocates pinned/mapped host memory, accessible by all devices in the given platform.
     //!
     //! \tparam TElem The element type of the returned buffer.
@@ -188,7 +289,7 @@ namespace alpaka
     }
 
     /* TODO: Remove this pragma block once support for clang versions <= 13 is removed. These versions are unable to
-       figure out that the template parameters are attached to a C++17 inline variable. */
+         figure out that the template parameters are attached to a C++17 inline variable. */
 #if ALPAKA_COMP_CLANG
 #    pragma clang diagnostic push
 #    pragma clang diagnostic ignored "-Wdocumentation"

include/alpaka/mem/buf/cpu/traits/BufCpuTraits.hpp

@@ -165,12 +165,54 @@ namespace alpaka::trait
         }
     };
 
+    //! The CPU caching memory allocation trait specialization.
+    template<typename TElem, typename TDim, typename TIdx, typename TAllocator>
+    struct BufAllocCached<TElem, TDim, TIdx, DevCpu, TAllocator>
+    {
+        template<typename TExtent>
+        ALPAKA_FN_HOST static auto allocBuf(DevCpu const& dev, TExtent const& extent, TAllocator allocator)
+            -> BufCpu<TElem, TDim, TIdx>
+        {
+            ALPAKA_DEBUG_MINIMAL_LOG_SCOPE;
+
+            std::size_t const bytes = static_cast<std::size_t>(getExtentProduct(extent)) * sizeof(TElem);
+            void* const memPtr = allocator.allocate(bytes, alignof(TElem));
+            auto deleter = [alloc = std::move(allocator)](TElem* ptr) mutable { alloc.deallocate(ptr); };
+            return BufCpu<TElem, TDim, TIdx>(dev, static_cast<TElem*>(memPtr), std::move(deleter), extent);
+        }
+    };
+
     //! The BufCpu stream-ordered memory allocation capability trait specialization.
     template<typename TDim>
     struct HasAsyncBufSupport<TDim, DevCpu> : public std::true_type
     {
     };
 
+    //! The CPU caching stream-ordered memory allocation trait specialization.
+    template<typename TElem, typename TDim, typename TIdx, typename TAllocator>
+    struct AsyncBufAllocCached<TElem, TDim, TIdx, DevCpu, TAllocator>
+    {
+        template<typename TQueue, typename TExtent>
+        ALPAKA_FN_HOST static auto allocAsyncBuf(TQueue queue, TExtent const& extent, TAllocator allocator)
+            -> BufCpu<TElem, TDim, TIdx>
+        {
+            ALPAKA_DEBUG_MINIMAL_LOG_SCOPE;
+
+            static_assert(
+                std::is_same_v<Dev<TQueue>, DevCpu>,
+                "The BufCpu buffer can only be used with a queue on a DevCpu device!");
+            auto const dev = getDev(queue);
+            std::size_t const bytes = static_cast<std::size_t>(getExtentProduct(extent)) * sizeof(TElem);
+            void* const memPtr = allocator.allocate(bytes, alignof(TElem));
+            auto deleter = [l_queue = std::move(queue), alloc = std::move(allocator)](TElem* ptr) mutable
+            {
+                alpaka::wait(l_queue);
+                alloc.deallocate(ptr);
+            };
+            return BufCpu<TElem, TDim, TIdx>(dev, static_cast<TElem*>(memPtr), std::move(deleter), extent);
+        }
+    };
+
     //! The BufCpu stream-ordered memory allocation trait specialization.
     template<typename TElem, typename TDim, typename TIdx>
     struct AsyncBufAlloc<TElem, TDim, TIdx, DevCpu>

include/alpaka/mem/buf/sycl/traits/BufGenericSyclTraits.hpp

@@ -194,12 +194,63 @@ namespace alpaka::trait
         }
     };
 
+    //! The SYCL caching memory allocation trait specialization.
+    template<typename TElem, typename TDim, typename TIdx, concepts::Tag TTag, typename TAllocator>
+    struct BufAllocCached<TElem, TDim, TIdx, DevGenericSycl<TTag>, TAllocator>
+    {
+        template<typename TExtent>
+        ALPAKA_FN_HOST static auto allocBuf(
+            DevGenericSycl<TTag> const& dev,
+            TExtent const& extent,
+            TAllocator allocator) -> BufGenericSycl<TElem, TDim, TIdx, TTag>
+        {
+            ALPAKA_DEBUG_MINIMAL_LOG_SCOPE;
+
+            std::size_t const bytes = static_cast<std::size_t>(getExtentProduct(extent)) * sizeof(TElem);
+            void* const memPtr = allocator.allocate(bytes, alignof(TElem));
+            auto deleter = [alloc = std::move(allocator)](TElem* ptr) mutable { alloc.deallocate(ptr); };
+            return BufGenericSycl<TElem, TDim, TIdx, TTag>(
+                dev,
+                static_cast<TElem*>(memPtr),
+                std::move(deleter),
+                extent);
+        }
+    };
+
     //! The BufGenericSycl stream-ordered memory allocation capability trait specialization.
     template<typename TDim, concepts::Tag TTag>
     struct HasAsyncBufSupport<TDim, DevGenericSycl<TTag>> : std::true_type
     {
     };
 
+    //! The SYCL caching stream-ordered memory allocation trait specialization.
+    template<typename TElem, typename TDim, typename TIdx, concepts::Tag TTag, typename TAllocator>
+    struct AsyncBufAllocCached<TElem, TDim, TIdx, DevGenericSycl<TTag>, TAllocator>
+    {
+        template<bool TBlocking, typename TExtent>
+        ALPAKA_FN_HOST static auto allocAsyncBuf(
+            alpaka::detail::QueueGenericSyclBase<TTag, TBlocking> queue,
+            TExtent const& extent,
+            TAllocator allocator) -> BufGenericSycl<TElem, TDim, TIdx, TTag>
+        {
+            ALPAKA_DEBUG_MINIMAL_LOG_SCOPE;
+
+            auto const dev = getDev(queue);
+            std::size_t const bytes = static_cast<std::size_t>(getExtentProduct(extent)) * sizeof(TElem);
+            void* const memPtr = allocator.allocate(bytes, alignof(TElem));
+            auto deleter = [l_queue = std::move(queue), alloc = std::move(allocator)](TElem* ptr) mutable
+            {
+                alpaka::wait(l_queue);
+                alloc.deallocate(ptr);
+            };
+            return BufGenericSycl<TElem, TDim, TIdx, TTag>(
+                dev,
+                static_cast<TElem*>(memPtr),
+                std::move(deleter),
+                extent);
+        }
+    };
+
     //! The BufGenericSycl stream-ordered memory allocation trait specialization.
     template<typename TElem, typename TDim, typename TIdx, concepts::Tag TTag>
     struct AsyncBufAlloc<TElem, TDim, TIdx, DevGenericSycl<TTag>>

include/alpaka/mem/buf/uniformCudaHip/traits/BufUniformCudaHipRtTraits.hpp

@@ -255,12 +255,69 @@ namespace alpaka::trait
         }
     };
 
+    //! The CUDA/HIP caching memory allocation trait specialization.
+    template<typename TApi, typename TElem, typename Dim, typename TIdx, typename TAllocator>
+    struct BufAllocCached<TElem, Dim, TIdx, DevUniformCudaHipRt<TApi>, TAllocator>
+    {
+        template<typename TExtent>
+        ALPAKA_FN_HOST static auto allocBuf(
+            DevUniformCudaHipRt<TApi> const& dev,
+            TExtent const& extent,
+            TAllocator allocator) -> BufUniformCudaHipRt<TApi, TElem, Dim, TIdx>
+        {
+            ALPAKA_DEBUG_MINIMAL_LOG_SCOPE;
+
+            ALPAKA_UNIFORM_CUDA_HIP_RT_CHECK(TApi::setDevice(dev.getNativeHandle()));
+
+            std::size_t const pitchBytes = static_cast<std::size_t>(getWidth(extent)) * sizeof(TElem);
+            std::size_t const bytes = static_cast<std::size_t>(getExtentProduct(extent)) * sizeof(TElem);
+            void* const memPtr = allocator.allocate(bytes, alignof(TElem));
+            auto deleter = [alloc = std::move(allocator)](TElem* ptr) mutable { alloc.deallocate(ptr); };
+            return BufUniformCudaHipRt<TApi, TElem, Dim, TIdx>(
+                dev,
+                static_cast<TElem*>(memPtr),
+                std::move(deleter),
+                extent,
+                pitchBytes);
+        }
+    };
+
     //! The CUDA/HIP stream-ordered memory allocation capability trait specialization.
     template<typename TApi, typename TDim>
     struct HasAsyncBufSupport<TDim, DevUniformCudaHipRt<TApi>> : std::true_type
     {
     };
 
+    //! The CUDA/HIP caching stream-ordered memory allocation trait specialization.
+    template<typename TApi, typename TElem, typename TDim, typename TIdx, typename TAllocator>
+    struct AsyncBufAllocCached<TElem, TDim, TIdx, DevUniformCudaHipRt<TApi>, TAllocator>
+    {
+        template<typename TQueue, typename TExtent>
+        ALPAKA_FN_HOST static auto allocAsyncBuf(TQueue queue, TExtent const& extent, TAllocator allocator)
+            -> BufUniformCudaHipRt<TApi, TElem, TDim, TIdx>
+        {
+            ALPAKA_DEBUG_MINIMAL_LOG_SCOPE;
+
+            auto const dev = getDev(queue);
+            ALPAKA_UNIFORM_CUDA_HIP_RT_CHECK(TApi::setDevice(dev.getNativeHandle()));
+
+            std::size_t const pitchBytes = static_cast<std::size_t>(getWidth(extent)) * sizeof(TElem);
+            std::size_t const bytes = static_cast<std::size_t>(getExtentProduct(extent)) * sizeof(TElem);
+            void* const memPtr = allocator.allocate(bytes, alignof(TElem));
+            auto deleter = [l_queue = std::move(queue), alloc = std::move(allocator)](TElem* ptr) mutable
+            {
+                alpaka::wait(l_queue);
+                alloc.deallocate(ptr);
+            };
+            return BufUniformCudaHipRt<TApi, TElem, TDim, TIdx>(
+                dev,
+                static_cast<TElem*>(memPtr),
+                std::move(deleter),
+                extent,
+                pitchBytes);
+        }
+    };
+
     //! The CUDA/HIP stream-ordered memory allocation trait specialization.
     template<typename TApi, typename TElem, typename TDim, typename TIdx>
     struct AsyncBufAlloc<TElem, TDim, TIdx, DevUniformCudaHipRt<TApi>>