corecel/data/Copier.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2021-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/Copier.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <cstddef>
0011 #include <type_traits>
0012
0013 #include "corecel/Assert.hh"
0014 #include "corecel/Types.hh"
0015 #include "corecel/cont/Span.hh"
0016 #include "corecel/sys/ThreadId.hh"
0017
0018 namespace celeritas
0019 {
0020 //---------------------------------------------------------------------------//
0021 /*!
0022  * Copy spans of data.
0023  *
0024  * The destination (which can be a reusable buffer) is the constructor
0025  * argument, and the source of the data to copy is the function argument.
0026  *
0027  * Example of copying data from device to host:
0028  * \code
0029     Copier<int, MemSpace::host> copy_to_host{host_ints};
0030     copy_to_host(MemSpace::device, device_ints);
0031  * \endcode
0032  */
0033 template<class T, MemSpace M>
0034 class Copier
0035 {
0036     static_assert(std::is_trivially_copyable<T>::value,
0037                   "Data is not trivially copyable");
0038
0039   public:
0040     //! Construct with the destination and the class's memspace
0041     explicit Copier(Span<T> dst) : dst_{dst} {};
0042
0043     //! Also construct with a stream ID to use for async copy
0044     Copier(Span<T> dst, StreamId stream) : dst_{dst}, stream_{stream} {};
0045
0046     inline void operator()(MemSpace srcmem, Span<T const> src) const;
0047
0048   private:
0049     Span<T> dst_;
0050     StreamId stream_;
0051     static constexpr auto dstmem = M;
0052 };
0053
0054 //---------------------------------------------------------------------------//
0055 /*!
0056  * Copy a value from device to host.
0057  *
0058  * The source of the data to copy is the function argument.
0059  */
0060 template<class T>
0061 class ItemCopier
0062 {
0063     static_assert(std::is_trivially_copyable<T>::value,
0064                   "Data is not trivially copyable");
0065
0066   public:
0067     //! Default constructor
0068     ItemCopier() = default;
0069
0070     //! Also construct with a stream ID to use for async copy
0071     explicit ItemCopier(StreamId stream) : stream_{stream} {};
0072
0073     inline T operator()(T const* src) const;
0074
0075   private:
0076     StreamId stream_;
0077 };
0078
0079 //---------------------------------------------------------------------------//
0080 // Copy bytes between two memory spaces
0081 void copy_bytes(MemSpace dstmem,
0082                 void* dst,
0083                 MemSpace srcmem,
0084                 void const* src,
0085                 std::size_t count);
0086
0087 // Asynchronously copy bytes between two memory spaces
0088 void copy_bytes(MemSpace dstmem,
0089                 void* dst,
0090                 MemSpace srcmem,
0091                 void const* src,
0092                 std::size_t count,
0093                 StreamId stream);
0094
0095 //---------------------------------------------------------------------------//
0096 /*!
0097  * Copy data from the given source and memory space.
0098  */
0099 template<class T, MemSpace M>
0100 void Copier<T, M>::operator()(MemSpace srcmem, Span<T const> src) const
0101 {
0102     CELER_EXPECT(src.size() == dst_.size());
0103     if (stream_)
0104     {
0105         copy_bytes(dstmem,
0106                    dst_.data(),
0107                    srcmem,
0108                    src.data(),
0109                    src.size() * sizeof(T),
0110                    stream_);
0111     }
0112     else
0113     {
0114         copy_bytes(
0115             dstmem, dst_.data(), srcmem, src.data(), src.size() * sizeof(T));
0116     }
0117 }
0118
0119 //---------------------------------------------------------------------------//
0120 /*!
0121  * Copy a value from device to host.
0122  */
0123 template<class T>
0124 T ItemCopier<T>::operator()(T const* src) const
0125 {
0126     T dst;
0127     if (stream_)
0128     {
0129         copy_bytes(
0130             MemSpace::host, &dst, MemSpace::device, src, sizeof(T), stream_);
0131     }
0132     else
0133     {
0134         copy_bytes(MemSpace::host, &dst, MemSpace::device, src, sizeof(T));
0135     }
0136     return dst;
0137 }
0138
0139 //---------------------------------------------------------------------------//
0140 }  // namespace celeritas