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