AlpakaInterface/interface/getDeviceCachingAllocator.h

0001 #ifndef HeterogeneousCore_AlpakaInterface_interface_getDeviceCachingAllocator_h
0002 #define HeterogeneousCore_AlpakaInterface_interface_getDeviceCachingAllocator_h
0003
0004 #include <cassert>
0005 #include <memory>
0006
0007 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0008 #include "HeterogeneousCore/AlpakaInterface/interface/AllocatorConfig.h"
0009 #include "HeterogeneousCore/AlpakaInterface/interface/CachingAllocator.h"
0010 #include "HeterogeneousCore/AlpakaInterface/interface/devices.h"
0011 #include "HeterogeneousCore/AlpakaInterface/interface/traits.h"
0012
0013 namespace cms::alpakatools {
0014
0015   namespace detail {
0016
0017     template <typename TDev,
0018               typename TQueue,
0019               typename = std::enable_if_t<cms::alpakatools::is_device_v<TDev> and cms::alpakatools::is_queue_v<TQueue>>>
0020     auto allocate_device_allocators() {
0021       using Allocator = CachingAllocator<TDev, TQueue>;
0022       auto const& devices = cms::alpakatools::devices<alpaka::Pltf<TDev>>();
0023       ssize_t const size = devices.size();
0024
0025       // allocate the storage for the objects
0026       auto ptr = std::allocator<Allocator>().allocate(size);
0027
0028       // construct the objects in the storage
0029       ptrdiff_t index = 0;
0030       try {
0031         for (; index < size; ++index) {
0032 #if __cplusplus >= 202002L
0033           std::construct_at(
0034 #else
0035           std::allocator<Allocator>().construct(
0036 #endif
0037               ptr + index,
0038               devices[index],
0039               config::binGrowth,
0040               config::minBin,
0041               config::maxBin,
0042               config::maxCachedBytes,
0043               config::maxCachedFraction,
0044               true,    // reuseSameQueueAllocations
0045               false);  // debug
0046         }
0047       } catch (...) {
0048         --index;
0049         // destroy any object that had been succesfully constructed
0050         while (index >= 0) {
0051           std::destroy_at(ptr + index);
0052           --index;
0053         }
0054         // deallocate the storage
0055         std::allocator<Allocator>().deallocate(ptr, size);
0056         // rethrow the exception
0057         throw;
0058       }
0059
0060       // use a custom deleter to destroy all objects and deallocate the memory
0061       auto deleter = [size](Allocator* ptr) {
0062         for (size_t i = size; i > 0; --i) {
0063           std::destroy_at(ptr + i - 1);
0064         }
0065         std::allocator<Allocator>().deallocate(ptr, size);
0066       };
0067
0068       return std::unique_ptr<Allocator[], decltype(deleter)>(ptr, deleter);
0069     }
0070
0071   }  // namespace detail
0072
0073   template <typename TDev,
0074             typename TQueue,
0075             typename = std::enable_if_t<cms::alpakatools::is_device_v<TDev> and cms::alpakatools::is_queue_v<TQueue>>>
0076   inline CachingAllocator<TDev, TQueue>& getDeviceCachingAllocator(TDev const& device) {
0077     // initialise all allocators, one per device
0078     CMS_THREAD_SAFE static auto allocators = detail::allocate_device_allocators<TDev, TQueue>();
0079
0080     size_t const index = alpaka::getNativeHandle(device);
0081     assert(index < cms::alpakatools::devices<alpaka::Pltf<TDev>>().size());
0082
0083     // the public interface is thread safe
0084     return allocators[index];
0085   }
0086
0087 }  // namespace cms::alpakatools
0088
0089 #endif  // HeterogeneousCore_AlpakaInterface_interface_getDeviceCachingAllocator_h