corecel/data/CollectionStateStore.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/CollectionStateStore.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  * Helper class for storing Collection classes on host or device.
0019  *
0020  * This can be used for unit tests (MemSpace is host) as well as production
0021  * code. States generally shouldn't be copied between host and device, so the
0022  * only "production use case" construction argument is the size. Other
0023  * constructors are implemented for convenience in unit tests.
0024  *
0025  * The State class must be templated on ownership and memory space, and
0026  * additionally must have an operator bool(), a templated operator=, and a
0027  * size() accessor. It must also define a free function "resize" that takes:
0028  * - \b REQUIRED: a pointer to the state with \c Ownership::value semantics
0029  * - \b OPTIONAL: a \c Ownership::const_reference instance of \c MemSpace::host
0030  *   params data
0031  * - \b OPTIONAL: a \c StreamId for setting up thread/task-local data
0032  * - \b REQUIRED: a \c size_type for specifying the size of the new state.
0033  *
0034  * \code
0035     CollectionStateStore<ParticleStateData, MemSpace::device> pstates(
0036         *particle_params, num_tracks);
0037     state_data.particle = pstates.ref();
0038    \endcode
0039  */
0040 template<template<Ownership, MemSpace> class S, MemSpace M>
0041 class CollectionStateStore
0042 {
0043   public:
0044     //!@{
0045     //! \name Type aliases
0046     using Value = S<Ownership::value, M>;
0047     using Ref = S<Ownership::reference, M>;
0048     using size_type = TrackSlotId::size_type;
0049     //!@}
0050
0051   public:
0052     CollectionStateStore() = default;
0053     ~CollectionStateStore() = default;
0054
0055     // Construct from parameters and stream ID
0056     template<template<Ownership, MemSpace> class P>
0057     inline CollectionStateStore(HostCRef<P> const& p,
0058                                 StreamId stream_id,
0059                                 size_type size);
0060
0061     // Construct from just parameters
0062     template<template<Ownership, MemSpace> class P>
0063     inline CollectionStateStore(HostCRef<P> const& p, size_type size);
0064
0065     // Construct without parameters and with stream ID
0066     explicit inline CollectionStateStore(StreamId stream_id, size_type size);
0067
0068     // Construct without parameters
0069     explicit inline CollectionStateStore(size_type size);
0070
0071     // Construct from values by capture
0072     explicit inline CollectionStateStore(S<Ownership::value, M>&& other);
0073
0074     // Copy construction from state data (convenience for unit tests)
0075     template<Ownership W2, MemSpace M2>
0076     explicit inline CollectionStateStore(S<W2, M2> const& other);
0077
0078     // Copy assignment from state data (convenience for unit tests)
0079     template<Ownership W2, MemSpace M2>
0080     inline CollectionStateStore& operator=(S<W2, M2> const& other);
0081
0082     //! Default move, delete copy (since ref "points to" val)
0083     CELER_DEFAULT_MOVE_DELETE_COPY(CollectionStateStore);
0084
0085     //! Whether any data is being stored
0086     explicit operator bool() const { return static_cast<bool>(val_); }
0087
0088     //! Number of elements
0089     size_type size() const { return val_.size(); }
0090
0091     // Get a reference to the mutable state data
0092     inline Ref& ref();
0093
0094     // Get a reference to the mutable state data
0095     inline Ref const& ref() const;
0096
0097   private:
0098     Value val_;
0099     Ref ref_;
0100
0101     template<template<Ownership, MemSpace> class S2, MemSpace M2>
0102     friend class CollectionStateStore;
0103 };
0104
0105 //---------------------------------------------------------------------------//
0106 /*!
0107  * Construct from parameter data.
0108  *
0109  * Most states are constructed with a \c resize function that takes host
0110  * parameter data and the number of states.
0111  */
0112 template<template<Ownership, MemSpace> class S, MemSpace M>
0113 template<template<Ownership, MemSpace> class P>
0114 CollectionStateStore<S, M>::CollectionStateStore(HostCRef<P> const& p,
0115                                                  StreamId sid,
0116                                                  size_type size)
0117 {
0118     CELER_EXPECT(sid);
0119     CELER_EXPECT(size > 0);
0120     resize(&val_, p, sid, size);
0121     CELER_ASSERT(val_);
0122
0123     // Save reference
0124     ref_ = val_;
0125 }
0126
0127 //---------------------------------------------------------------------------//
0128 /*!
0129  * Construct from parameter data.
0130  *
0131  * Most states are constructed with a \c resize function that takes host
0132  * parameter data and the number of states.
0133  */
0134 template<template<Ownership, MemSpace> class S, MemSpace M>
0135 template<template<Ownership, MemSpace> class P>
0136 CollectionStateStore<S, M>::CollectionStateStore(HostCRef<P> const& p,
0137                                                  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 CollectionStateStore<S, M>::CollectionStateStore(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 CollectionStateStore<S, M>::CollectionStateStore(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 CollectionStateStore<S, M>::CollectionStateStore(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 CollectionStateStore<S, M>::CollectionStateStore(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 CollectionStateStore<S, M>::operator=(S<W2, M2> const& other)
0220     -> CollectionStateStore<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 reference to the mutable state data.
0234  */
0235 template<template<Ownership, MemSpace> class S, MemSpace M>
0236 auto CollectionStateStore<S, M>::ref() -> Ref&
0237 {
0238     CELER_EXPECT(*this);
0239     return ref_;
0240 }
0241
0242 //---------------------------------------------------------------------------//
0243 /*!
0244  * Get a reference to the mutable state data.
0245  */
0246 template<template<Ownership, MemSpace> class S, MemSpace M>
0247 auto CollectionStateStore<S, M>::ref() const -> Ref const&
0248 {
0249     CELER_EXPECT(*this);
0250     return ref_;
0251 }
0252
0253 //---------------------------------------------------------------------------//
0254 }  // namespace celeritas