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