corecel/data/DeviceAllocation.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2020-2024 UT-Battelle, LLC, and other Celeritas developers.
0003 // See the top-level COPYRIGHT file for details.
0004 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0005 //---------------------------------------------------------------------------//
0006 //! \file corecel/data/DeviceAllocation.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <cstdlib>
0011 #include <memory>
0012 #include <utility>
0013
0014 #include "corecel/Macros.hh"
0015 #include "corecel/Types.hh"
0016 #include "corecel/cont/InitializedValue.hh"
0017 #include "corecel/cont/Span.hh"
0018 #include "corecel/sys/ThreadId.hh"
0019
0020 namespace celeritas
0021 {
0022 //---------------------------------------------------------------------------//
0023 /*!
0024  * Allocate raw uninitialized memory.
0025  *
0026  * This class is intended to be used by host-compiler \c .hh code as a bridge
0027  * to device memory. It allows Storage classes to allocate and manage device
0028  * memory without using \c thrust, which requires NVCC and propagates that
0029  * requirement into all downstream code.
0030  *
0031  * TODO: remove the stream constructor data members and rely on \c Copier or
0032  * \c thrust to do streamed async operations?
0033  */
0034 class DeviceAllocation
0035 {
0036   public:
0037     //!@{
0038     //! \name Type aliases
0039     using size_type = std::size_t;
0040     using SpanBytes = Span<std::byte>;
0041     using SpanConstBytes = Span<std::byte const>;
0042     //!@}
0043
0044   public:
0045     // Construct in unallocated state
0046     DeviceAllocation() = default;
0047
0048     // Construct in unallocated state
0049     explicit DeviceAllocation(StreamId stream);
0050
0051     // Construct and allocate a number of bytes
0052     explicit DeviceAllocation(size_type num_bytes);
0053
0054     // Construct and allocate a number of bytes
0055     DeviceAllocation(size_type num_bytes, StreamId stream);
0056
0057     // Swap with another allocation
0058     inline void swap(DeviceAllocation& other) noexcept;
0059
0060     //// ACCESSORS ////
0061
0062     //! Get the number of bytes allocated
0063     size_type size() const { return size_; }
0064
0065     //! Whether memory is allocated
0066     bool empty() const { return size_ == 0; }
0067
0068     //! Access the stream, set for asynchonous allocation/copy
0069     StreamId stream_id() const { return stream_; }
0070
0071     //// DEVICE ACCESSORS ////
0072
0073     // Get the device pointer
0074     inline SpanBytes device_ref();
0075
0076     // Get the device pointer
0077     inline SpanConstBytes device_ref() const;
0078
0079     // Copy data to device
0080     void copy_to_device(SpanConstBytes bytes);
0081
0082     // Copy data to host
0083     void copy_to_host(SpanBytes bytes) const;
0084
0085   private:
0086     struct DeviceFreeDeleter
0087     {
0088         StreamId stream_;
0089         void operator()(std::byte*) const noexcept(CELER_USE_DEVICE);
0090     };
0091     using DeviceUniquePtr = std::unique_ptr<std::byte[], DeviceFreeDeleter>;
0092
0093     //// DATA ////
0094
0095     InitializedValue<size_type> size_;
0096     StreamId stream_;
0097     DeviceUniquePtr data_;
0098 };
0099
0100 // Swap two allocations
0101 inline void swap(DeviceAllocation& a, DeviceAllocation& b) noexcept;
0102
0103 //---------------------------------------------------------------------------//
0104 // INLINE DEFINITIONS
0105 //---------------------------------------------------------------------------//
0106 /*!
0107  * Swap with another allocation.
0108  */
0109 void DeviceAllocation::swap(DeviceAllocation& other) noexcept
0110 {
0111     using std::swap;
0112     swap(this->data_, other.data_);
0113     swap(this->size_, other.size_);
0114     swap(this->stream_, other.stream_);
0115 }
0116
0117 //---------------------------------------------------------------------------//
0118 /*!
0119  * Get a view to the owned device memory.
0120  */
0121 auto DeviceAllocation::device_ref() -> SpanBytes
0122 {
0123     return {data_.get(), size_};
0124 }
0125
0126 //---------------------------------------------------------------------------//
0127 /*!
0128  * Get a view to the owned device memory.
0129  */
0130 auto DeviceAllocation::device_ref() const -> SpanConstBytes
0131 {
0132     return {data_.get(), size_};
0133 }
0134
0135 //---------------------------------------------------------------------------//
0136 /*!
0137  * Swap two allocations.
0138  */
0139 void swap(DeviceAllocation& a, DeviceAllocation& b) noexcept
0140 {
0141     return a.swap(b);
0142 }
0143
0144 //---------------------------------------------------------------------------//
0145 }  // namespace celeritas