corecel/data/StateDataStore.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/StateDataStore.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include "corecel/Assert.hh"
0010 #include "corecel/OpaqueId.hh"
0011 #include "corecel/Types.hh"
0012 #include "corecel/sys/ThreadId.hh"
0013
0014 namespace celeritas
0015 {
0016 //---------------------------------------------------------------------------//
0017 /*!
0018  * Store and reference stateful collection groups on host and device.
0019  * \tparam S State data collection group
0020  *
0021  * This can be used for unit tests (MemSpace is host) as well as production
0022  * code. States generally shouldn't be copied between host and device, so the
0023  * only "production use case" construction argument is the size. Other
0024  * constructors are implemented for convenience in unit tests.
0025  *
0026  * The State class must be templated on ownership and memory space, and
0027  * additionally must have an operator bool(), a templated operator=, and a
0028  * size() accessor. It must also define a free function "resize" that takes:
0029  * - \b REQUIRED: a pointer to the state with \c Ownership::value semantics
0030  * - \b OPTIONAL: a \c Ownership::const_reference instance of \c MemSpace::host
0031  *   params data
0032  * - \b OPTIONAL: a \c StreamId for setting up thread/task-local data
0033  * - \b REQUIRED: a \c size_type for specifying the size of the new state.
0034  *
0035  * \par Example:
0036  * \code
0037     StateDataStore<ParticleStateData, MemSpace::device> pstates(
0038         *particle_params, num_tracks);
0039     state_data.particle = pstates.ref();
0040    \endcode
0041  */
0042 template<template<Ownership, MemSpace> class S, MemSpace M>
0043 class StateDataStore
0044 {
0045   public:
0046     //!@{
0047     //! \name Type aliases
0048     using Value = S<Ownership::value, M>;
0049     using Ref = S<Ownership::reference, M>;
0050     using size_type = TrackSlotId::size_type;
0051     //!@}
0052
0053   public:
0054     //! The default constructor initializes as empty
0055     StateDataStore() = default;
0056     ~StateDataStore() = default;
0057
0058     // Construct from parameters and stream ID
0059     template<template<Ownership, MemSpace> class P>
0060     inline StateDataStore(
0061         HostCRef<P> const& p, StreamId stream_id, size_type size);
0062
0063     // Construct from just parameters
0064     template<template<Ownership, MemSpace> class P>
0065     inline StateDataStore(HostCRef<P> const& p, size_type size);
0066
0067     // Construct without parameters and with stream ID
0068     explicit inline StateDataStore(StreamId stream_id, size_type size);
0069
0070     // Construct without parameters
0071     explicit inline StateDataStore(size_type size);
0072
0073     // Construct from values by capture
0074     explicit inline StateDataStore(S<Ownership::value, M>&& other);
0075
0076     // Copy construction from state data (convenience for unit tests)
0077     template<Ownership W2, MemSpace M2>
0078     explicit inline StateDataStore(S<W2, M2> const& other);
0079
0080     // Copy assignment from state data (convenience for unit tests)
0081     template<Ownership W2, MemSpace M2>
0082     inline StateDataStore& operator=(S<W2, M2> const& other);
0083
0084     //! Default move, delete copy (since ref "points to" val)
0085     CELER_DEFAULT_MOVE_DELETE_COPY(StateDataStore);
0086
0087     //! Whether any data is being stored
0088     explicit operator bool() const { return static_cast<bool>(val_); }
0089
0090     //! Number of elements
0091     size_type size() const { return val_.size(); }
0092
0093     // Get a reference to the mutable state data
0094     inline Ref& ref();
0095
0096     // Get a reference to the mutable state data
0097     inline Ref const& ref() const;
0098
0099   private:
0100     Value val_;
0101     Ref ref_;
0102
0103     template<template<Ownership, MemSpace> class S2, MemSpace M2>
0104     friend class StateDataStore;
0105 };
0106
0107 //---------------------------------------------------------------------------//
0108 /*!
0109  * Construct from parameter data.
0110  *
0111  * Most states are constructed with a \c resize function that takes host
0112  * parameter data and the number of states.
0113  */
0114 template<template<Ownership, MemSpace> class S, MemSpace M>
0115 template<template<Ownership, MemSpace> class P>
0116 StateDataStore<S, M>::StateDataStore(
0117     HostCRef<P> const& p, StreamId sid, size_type size)
0118 {
0119     CELER_EXPECT(sid);
0120     CELER_EXPECT(size > 0);
0121     resize(&val_, p, sid, size);
0122     CELER_ASSERT(val_);
0123
0124     // Save reference
0125     ref_ = val_;
0126 }
0127
0128 //---------------------------------------------------------------------------//
0129 /*!
0130  * Construct from parameter data.
0131  *
0132  * Most states are constructed with a \c resize function that takes host
0133  * parameter data and the number of states.
0134  */
0135 template<template<Ownership, MemSpace> class S, MemSpace M>
0136 template<template<Ownership, MemSpace> class P>
0137 StateDataStore<S, M>::StateDataStore(HostCRef<P> const& p, size_type size)
0138 {
0139     CELER_EXPECT(size > 0);
0140     resize(&val_, p, size);
0141     CELER_ASSERT(val_);
0142
0143     // Save reference
0144     ref_ = val_;
0145 }
0146
0147 //---------------------------------------------------------------------------//
0148 /*!
0149  * Construct with stream ID and without parameters.
0150  *
0151  * A few states are constructed with a \c resize function that doesn't depend
0152  * on any parameter data.
0153  */
0154 template<template<Ownership, MemSpace> class S, MemSpace M>
0155 StateDataStore<S, M>::StateDataStore(StreamId sid, size_type size)
0156 {
0157     CELER_EXPECT(size > 0);
0158     resize(&val_, sid, size);
0159     CELER_ASSERT(val_);
0160
0161     // Save reference
0162     ref_ = val_;
0163 }
0164
0165 //---------------------------------------------------------------------------//
0166 /*!
0167  * Construct without parameters.
0168  *
0169  * A few states are constructed with a \c resize function that doesn't depend
0170  * on any parameter data.
0171  */
0172 template<template<Ownership, MemSpace> class S, MemSpace M>
0173 StateDataStore<S, M>::StateDataStore(size_type size)
0174 {
0175     CELER_EXPECT(size > 0);
0176     resize(&val_, size);
0177     CELER_ASSERT(val_);
0178
0179     // Save reference
0180     ref_ = val_;
0181 }
0182
0183 //---------------------------------------------------------------------------//
0184 /*!
0185  * Construct from values by capture.
0186  */
0187 template<template<Ownership, MemSpace> class S, MemSpace M>
0188 StateDataStore<S, M>::StateDataStore(S<Ownership::value, M>&& other)
0189     : val_(std::move(other))
0190 {
0191     CELER_EXPECT(val_);
0192     // Save reference
0193     ref_ = val_;
0194 }
0195
0196 //---------------------------------------------------------------------------//
0197 /*!
0198  * Construct from a state.
0199  */
0200 template<template<Ownership, MemSpace> class S, MemSpace M>
0201 template<Ownership W2, MemSpace M2>
0202 StateDataStore<S, M>::StateDataStore(S<W2, M2> const& other)
0203 {
0204     CELER_EXPECT(other);
0205     // Assign using const-cast because state copy operators have to be mutable
0206     // even when they're just copying...
0207     val_ = const_cast<S<W2, M2>&>(other);
0208     CELER_ASSERT(val_);
0209     // Save reference
0210     ref_ = val_;
0211 }
0212
0213 //---------------------------------------------------------------------------//
0214 /*!
0215  * Copy assign from a state.
0216  */
0217 template<template<Ownership, MemSpace> class S, MemSpace M>
0218 template<Ownership W2, MemSpace M2>
0219 auto StateDataStore<S, M>::operator=(S<W2, M2> const& other)
0220     -> StateDataStore<S, M>&
0221 {
0222     CELER_EXPECT(other);
0223     // Assign
0224     val_ = const_cast<S<W2, M2>&>(other);
0225     CELER_ASSERT(val_);
0226     // Save reference
0227     ref_ = val_;
0228     return *this;
0229 }
0230
0231 //---------------------------------------------------------------------------//
0232 /*!
0233  * Get a mutable reference to the mutable state data.
0234  */
0235 template<template<Ownership, MemSpace> class S, MemSpace M>
0236 auto StateDataStore<S, M>::ref() -> Ref&
0237 {
0238     return ref_;
0239 }
0240
0241 //---------------------------------------------------------------------------//
0242 /*!
0243  * Get a const reference to the state data.
0244  */
0245 template<template<Ownership, MemSpace> class S, MemSpace M>
0246 auto StateDataStore<S, M>::ref() const -> Ref const&
0247 {
0248     return ref_;
0249 }
0250
0251 //---------------------------------------------------------------------------//
0252 }  // namespace celeritas