Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-02-14 04:11:48

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