sys/detail/AsyncMemoryResource.hh

0001 //------------------------------- -*- C++ -*- -------------------------------//
0002 // Copyright Celeritas contributors: see top-level COPYRIGHT file for details
0003 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0004 //---------------------------------------------------------------------------//
0005 //! \file corecel/sys/detail/AsyncMemoryResource.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include "corecel/Macros.hh"
0010
0011 #if CELER_USE_DEVICE
0012 #    include <thrust/mr/memory_resource.h>
0013 #endif
0014 #include <cstdlib>
0015
0016 #include "corecel/DeviceRuntimeApi.hh"
0017
0018 namespace celeritas
0019 {
0020 namespace detail
0021 {
0022 //---------------------------------------------------------------------------//
0023 //! CUDA/HIP opaque stream handle
0024 #if CELER_USE_DEVICE
0025 using DeviceStream_t = CELER_DEVICE_API_SYMBOL(Stream_t);
0026 #else
0027 using DeviceStream_t = std::nullptr_t;
0028 #endif
0029
0030 //! Allocate memory asynchronously if supported
0031 void* malloc_async(std::size_t bytes, DeviceStream_t s);
0032 //! Free memory asynchronously if supported
0033 void free_async(void* ptr, DeviceStream_t s);
0034
0035 #if CELER_USE_DEVICE || defined(__DOXYGEN__)
0036 //---------------------------------------------------------------------------//
0037 /*!
0038  * Thrust async memory resource associated with a CUDA/HIP stream.
0039  */
0040 class AsyncMemoryResource final : public thrust::mr::memory_resource<void*>
0041 {
0042   public:
0043     //!@{
0044     //! \name Type aliases
0045     using pointer = void*;
0046     using StreamT = DeviceStream_t;
0047     //!@}
0048
0049     // Construct memory resource for the stream
0050     explicit AsyncMemoryResource(StreamT stream) : stream_{stream} {}
0051
0052     // Construct with default Stream
0053     AsyncMemoryResource() = default;
0054
0055     // Allocate device memory
0056     pointer do_allocate(std::size_t bytes, std::size_t) final;
0057
0058     // Deallocate device memory
0059     void do_deallocate(pointer p, std::size_t, std::size_t) final;
0060
0061   private:
0062     StreamT stream_{nullptr};
0063 };
0064 #else
0065 class AsyncMemoryResource;
0066 inline void* malloc_async(std::size_t, DeviceStream_t)
0067 {
0068     CELER_ASSERT_UNREACHABLE();
0069 }
0070 inline void free_async(void*, DeviceStream_t)
0071 {
0072     CELER_ASSERT_UNREACHABLE();
0073 }
0074
0075 #endif
0076
0077 //---------------------------------------------------------------------------//
0078 }  // namespace detail
0079 }  // namespace celeritas