Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-09-16 08:52:40

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
0066     explicit inline CollectionStateStore(size_type size);
0067 
0068     // Construct from values by capture
0069     explicit inline CollectionStateStore(S<Ownership::value, M>&& other);
0070 
0071     // Copy construction from state data (convenience for unit tests)
0072     template<Ownership W2, MemSpace M2>
0073     explicit inline CollectionStateStore(S<W2, M2> const& other);
0074 
0075     // Copy assignment from state data (convenience for unit tests)
0076     template<Ownership W2, MemSpace M2>
0077     inline CollectionStateStore& operator=(S<W2, M2> const& other);
0078 
0079     //! Default move, delete copy (since ref "points to" val)
0080     CELER_DEFAULT_MOVE_DELETE_COPY(CollectionStateStore);
0081 
0082     //! Whether any data is being stored
0083     explicit operator bool() const { return static_cast<bool>(val_); }
0084 
0085     //! Number of elements
0086     size_type size() const { return val_.size(); }
0087 
0088     // Get a reference to the mutable state data
0089     inline Ref& ref();
0090 
0091     // Get a reference to the mutable state data
0092     inline Ref const& ref() const;
0093 
0094   private:
0095     Value val_;
0096     Ref ref_;
0097 
0098     template<template<Ownership, MemSpace> class S2, MemSpace M2>
0099     friend class CollectionStateStore;
0100 };
0101 
0102 //---------------------------------------------------------------------------//
0103 /*!
0104  * Construct from parameter data.
0105  *
0106  * Most states are constructed with a \c resize function that takes host
0107  * parameter data and the number of states.
0108  */
0109 template<template<Ownership, MemSpace> class S, MemSpace M>
0110 template<template<Ownership, MemSpace> class P>
0111 CollectionStateStore<S, M>::CollectionStateStore(HostCRef<P> const& p,
0112                                                  StreamId sid,
0113                                                  size_type size)
0114 {
0115     CELER_EXPECT(sid);
0116     CELER_EXPECT(size > 0);
0117     resize(&val_, p, sid, size);
0118     CELER_ASSERT(val_);
0119 
0120     // Save reference
0121     ref_ = val_;
0122 }
0123 
0124 //---------------------------------------------------------------------------//
0125 /*!
0126  * Construct from parameter data.
0127  *
0128  * Most states are constructed with a \c resize function that takes host
0129  * parameter data and the number of states.
0130  */
0131 template<template<Ownership, MemSpace> class S, MemSpace M>
0132 template<template<Ownership, MemSpace> class P>
0133 CollectionStateStore<S, M>::CollectionStateStore(HostCRef<P> const& p,
0134                                                  size_type size)
0135 {
0136     CELER_EXPECT(size > 0);
0137     resize(&val_, p, size);
0138     CELER_ASSERT(val_);
0139 
0140     // Save reference
0141     ref_ = val_;
0142 }
0143 
0144 //---------------------------------------------------------------------------//
0145 /*!
0146  * Construct without parameters.
0147  *
0148  * A few states are constructed with a \c resize function that doesn't depend
0149  * on any parameter data.
0150  */
0151 template<template<Ownership, MemSpace> class S, MemSpace M>
0152 CollectionStateStore<S, M>::CollectionStateStore(size_type size)
0153 {
0154     CELER_EXPECT(size > 0);
0155     resize(&val_, size);
0156     CELER_ASSERT(val_);
0157 
0158     // Save reference
0159     ref_ = val_;
0160 }
0161 
0162 //---------------------------------------------------------------------------//
0163 /*!
0164  * Construct from values by capture.
0165  */
0166 template<template<Ownership, MemSpace> class S, MemSpace M>
0167 CollectionStateStore<S, M>::CollectionStateStore(S<Ownership::value, M>&& other)
0168     : val_(std::move(other))
0169 {
0170     CELER_EXPECT(val_);
0171     // Save reference
0172     ref_ = val_;
0173 }
0174 
0175 //---------------------------------------------------------------------------//
0176 /*!
0177  * Construct from a state.
0178  */
0179 template<template<Ownership, MemSpace> class S, MemSpace M>
0180 template<Ownership W2, MemSpace M2>
0181 CollectionStateStore<S, M>::CollectionStateStore(S<W2, M2> const& other)
0182 {
0183     CELER_EXPECT(other);
0184     // Assign using const-cast because state copy operators have to be mutable
0185     // even when they're just copying...
0186     val_ = const_cast<S<W2, M2>&>(other);
0187     CELER_ASSERT(val_);
0188     // Save reference
0189     ref_ = val_;
0190 }
0191 
0192 //---------------------------------------------------------------------------//
0193 /*!
0194  * Copy assign from a state.
0195  */
0196 template<template<Ownership, MemSpace> class S, MemSpace M>
0197 template<Ownership W2, MemSpace M2>
0198 auto CollectionStateStore<S, M>::operator=(S<W2, M2> const& other)
0199     -> CollectionStateStore<S, M>&
0200 {
0201     CELER_EXPECT(other);
0202     // Assign
0203     val_ = const_cast<S<W2, M2>&>(other);
0204     CELER_ASSERT(val_);
0205     // Save reference
0206     ref_ = val_;
0207     return *this;
0208 }
0209 
0210 //---------------------------------------------------------------------------//
0211 /*!
0212  * Get a reference to the mutable state data.
0213  */
0214 template<template<Ownership, MemSpace> class S, MemSpace M>
0215 auto CollectionStateStore<S, M>::ref() -> Ref&
0216 {
0217     CELER_EXPECT(*this);
0218     return ref_;
0219 }
0220 
0221 //---------------------------------------------------------------------------//
0222 /*!
0223  * Get a reference to the mutable state data.
0224  */
0225 template<template<Ownership, MemSpace> class S, MemSpace M>
0226 auto CollectionStateStore<S, M>::ref() const -> Ref const&
0227 {
0228     CELER_EXPECT(*this);
0229     return ref_;
0230 }
0231 
0232 //---------------------------------------------------------------------------//
0233 }  // namespace celeritas