corecel/data/AuxStateData.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 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/AuxStateData.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <memory>
0011
0012 #include "corecel/Types.hh"
0013
0014 #include "AuxInterface.hh"
0015 #include "CollectionStateStore.hh"
0016 #include "ParamsDataInterface.hh"
0017
0018 namespace celeritas
0019 {
0020 //---------------------------------------------------------------------------//
0021 /*!
0022  * Helper class for retrieving templated state data on a single stream.
0023  *
0024  * This class is most easily used with \c make_aux_state to create a
0025  * "collection group"-style state (see \ref collections) associated with a
0026  * \c AuxParamsInterface subclass.
0027  *
0028  * The state class \c S must have a \c resize method that's constructable with
0029  * a templated params data class \c P, a stream ID, and a state size.
0030  */
0031 template<template<Ownership, MemSpace> class S, MemSpace M>
0032 class AuxStateData final : public AuxStateInterface
0033 {
0034   public:
0035     //!@{
0036     //! \name Type aliases
0037     using Ref = S<Ownership::reference, M>;
0038     //!@}
0039
0040   public:
0041     // Construct by resizing and passing host params
0042     template<template<Ownership, MemSpace> class P>
0043     inline AuxStateData(HostCRef<P> const& p,
0044                         StreamId stream_id,
0045                         size_type size);
0046
0047     //! Whether any data is being stored
0048     explicit operator bool() const { return static_cast<bool>(store_); }
0049
0050     //! Number of elements in the state
0051     size_type size() const { return store_.size(); }
0052
0053     //! Get a reference to the mutable state data
0054     Ref& ref() { return store_.ref(); }
0055
0056     //! Get a reference to immutable state data
0057     Ref const& ref() const { return store_.ref(); }
0058
0059   private:
0060     CollectionStateStore<S, M> store_;
0061 };
0062
0063 //---------------------------------------------------------------------------//
0064 /*!
0065  * Create a auxiliary state given a runtime memory space.
0066  *
0067  * Example:
0068  * \code
0069     return make_aux_state<ParticleTallyStateData>(
0070         *this, memspace, stream, size);
0071  * \endcode
0072  */
0073 template<template<Ownership, MemSpace> class S, template<Ownership, MemSpace> class P>
0074 std::unique_ptr<AuxStateInterface>
0075 make_aux_state(ParamsDataInterface<P> const& params,
0076                MemSpace m,
0077                StreamId stream_id,
0078                size_type size)
0079 {
0080     if (m == MemSpace::host)
0081     {
0082         return std::make_unique<AuxStateData<S, MemSpace::host>>(
0083             params.host_ref(), stream_id, size);
0084     }
0085     else if (m == MemSpace::device)
0086     {
0087         return std::make_unique<AuxStateData<S, MemSpace::device>>(
0088             params.host_ref(), stream_id, size);
0089     }
0090     CELER_ASSERT_UNREACHABLE();
0091 }
0092
0093 //---------------------------------------------------------------------------//
0094 // INLINE DEFINITIONS
0095 //---------------------------------------------------------------------------//
0096 /*!
0097  * Construct by resizing and passing host params.
0098  */
0099 template<template<Ownership, MemSpace> class S, MemSpace M>
0100 template<template<Ownership, MemSpace> class P>
0101 AuxStateData<S, M>::AuxStateData(HostCRef<P> const& p,
0102                                  StreamId stream_id,
0103                                  size_type size)
0104     : store_{p, stream_id, size}
0105 {
0106 }
0107
0108 //---------------------------------------------------------------------------//
0109 }  // namespace celeritas