Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-02-22 10:31:23

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2023-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 celeritas/global/CoreState.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009 
0010 #include <vector>
0011 
0012 #include "corecel/Types.hh"
0013 #include "corecel/data/AuxInterface.hh"
0014 #include "corecel/data/AuxStateData.hh"
0015 #include "corecel/data/AuxStateVec.hh"
0016 #include "corecel/data/Collection.hh"
0017 #include "corecel/data/CollectionStateStore.hh"
0018 #include "corecel/data/DeviceVector.hh"
0019 #include "corecel/data/Ref.hh"
0020 #include "corecel/sys/ThreadId.hh"
0021 #include "celeritas/track/CoreStateCounters.hh"
0022 
0023 #include "CoreTrackData.hh"
0024 
0025 #include "detail/CoreStateThreadOffsets.hh"
0026 
0027 namespace celeritas
0028 {
0029 class CoreParams;
0030 //---------------------------------------------------------------------------//
0031 /*!
0032  * Abstract base class for CoreState.
0033  */
0034 class CoreStateInterface
0035 {
0036   public:
0037     //!@{
0038     //! \name Type aliases
0039     using size_type = TrackSlotId::size_type;
0040     //!@}
0041 
0042   public:
0043     // Support polymorphic deletion
0044     virtual ~CoreStateInterface();
0045 
0046     //! Thread/stream ID
0047     virtual StreamId stream_id() const = 0;
0048 
0049     //! Number of track slots
0050     virtual size_type size() const = 0;
0051 
0052     //! Access track initialization counters
0053     virtual CoreStateCounters const& counters() const = 0;
0054 
0055     //! Access auxiliary state data
0056     virtual AuxStateVec const& aux() const = 0;
0057 
0058   protected:
0059     CoreStateInterface() = default;
0060     CELER_DEFAULT_COPY_MOVE(CoreStateInterface);
0061 };
0062 
0063 //---------------------------------------------------------------------------//
0064 /*!
0065  * Store all state data for a single thread.
0066  *
0067  * When the state lives on the device, we maintain a separate copy of the
0068  * device "ref" in device memory: otherwise we'd have to copy the entire state
0069  * in launch arguments and access it through constant memory.
0070  *
0071  * \todo Encapsulate all the action management accessors in a helper class.
0072  */
0073 template<MemSpace M>
0074 class CoreState final : public CoreStateInterface
0075 {
0076   public:
0077     //!@{
0078     //! \name Type aliases
0079     template<template<Ownership, MemSpace> class S>
0080     using StateRef = S<Ownership::reference, M>;
0081 
0082     using Ref = StateRef<CoreStateData>;
0083     using Ptr = ObserverPtr<Ref, M>;
0084     //!@}
0085 
0086   public:
0087     // Construct from CoreParams
0088     CoreState(CoreParams const& params,
0089               StreamId stream_id,
0090               size_type num_track_slots);
0091 
0092     // Default destructor
0093     ~CoreState();
0094 
0095     // Prevent move/copy
0096     CELER_DELETE_COPY_MOVE(CoreState);
0097 
0098     //! Thread/stream ID
0099     StreamId stream_id() const final { return this->ref().stream_id; }
0100 
0101     //! Number of track slots
0102     size_type size() const final { return states_.size(); }
0103 
0104     // Set a warmup flag
0105     void warming_up(bool);
0106 
0107     //! Whether the state is being transported with no active particles
0108     bool warming_up() const { return warming_up_; }
0109 
0110     //// CORE DATA ////
0111 
0112     //! Get a reference to the mutable state data
0113     Ref& ref() { return states_.ref(); }
0114 
0115     //! Get a reference to the mutable state data
0116     Ref const& ref() const { return states_.ref(); }
0117 
0118     //! Get a native-memspace pointer to the mutable state data
0119     Ptr ptr() { return ptr_; }
0120 
0121     //! Reset the state data
0122     void reset();
0123 
0124     //// COUNTERS ////
0125 
0126     //! Track initialization counters
0127     CoreStateCounters& counters() { return counters_; }
0128 
0129     //! Track initialization counters
0130     CoreStateCounters const& counters() const final { return counters_; }
0131 
0132     //// USER DATA ////
0133 
0134     //! Access auxiliary state data
0135     AuxStateVec const& aux() const final { return aux_state_; }
0136 
0137     //! Access auxiliary state data (mutable)
0138     AuxStateVec& aux() { return aux_state_; }
0139 
0140     // Convenience function to access auxiliary "collection group" data
0141     template<template<Ownership, MemSpace> class S>
0142     inline StateRef<S>& aux_data(AuxId auxid);
0143 
0144     //// TRACK SORTING ////
0145 
0146     //! Return whether tracks can be sorted by action
0147     bool has_action_range() const { return !offsets_.empty(); }
0148 
0149     // Get a range of sorted track slots about to undergo a given action
0150     Range<ThreadId> get_action_range(ActionId action_id) const;
0151 
0152     // Access the range of actions to apply for all track IDs
0153     inline auto& action_thread_offsets();
0154 
0155     // Access the range of actions to apply for all track IDs
0156     inline auto const& action_thread_offsets() const;
0157 
0158     // Access action offsets for computation (native memory space)
0159     inline auto& native_action_thread_offsets();
0160 
0161   private:
0162     // State data
0163     CollectionStateStore<CoreStateData, M> states_;
0164 
0165     // Copy of state ref in device memory, if M == MemSpace::device
0166     DeviceVector<Ref> device_ref_vec_;
0167 
0168     // Native pointer to ref data
0169     Ptr ptr_;
0170 
0171     // Counters for track initialization and activity
0172     CoreStateCounters counters_;
0173 
0174     // User-added data associated with params
0175     AuxStateVec aux_state_;
0176 
0177     // Indices of first thread assigned to a given action
0178     detail::CoreStateThreadOffsets<M> offsets_;
0179 
0180     // Whether no primaries should be generated
0181     bool warming_up_{false};
0182 };
0183 
0184 //---------------------------------------------------------------------------//
0185 /*!
0186  * Convenience function to access auxiliary "collection group" data.
0187  */
0188 template<MemSpace M>
0189 template<template<Ownership, MemSpace> class S>
0190 auto CoreState<M>::aux_data(AuxId auxid) -> StateRef<S>&
0191 {
0192     CELER_EXPECT(auxid < aux_state_.size());
0193 
0194     // TODO: use "checked static cast" for better runtime performance
0195     auto* state = dynamic_cast<AuxStateData<S, M>*>(&aux_state_.at(auxid));
0196     CELER_ASSERT(state);
0197 
0198     CELER_ENSURE(*state);
0199     return state->ref();
0200 }
0201 
0202 //---------------------------------------------------------------------------//
0203 /*!
0204  * Access the range of actions to apply for all track IDs.
0205  */
0206 template<MemSpace M>
0207 auto& CoreState<M>::action_thread_offsets()
0208 {
0209     return offsets_.host_action_thread_offsets();
0210 }
0211 
0212 //---------------------------------------------------------------------------//
0213 /*!
0214  * Access the range of actions to apply for all track IDs.
0215  */
0216 template<MemSpace M>
0217 auto const& CoreState<M>::action_thread_offsets() const
0218 {
0219     return offsets_.host_action_thread_offsets();
0220 }
0221 
0222 //---------------------------------------------------------------------------//
0223 /*!
0224  * Access action offsets for computation (native memory space).
0225  */
0226 template<MemSpace M>
0227 auto& CoreState<M>::native_action_thread_offsets()
0228 {
0229     return offsets_.native_action_thread_offsets();
0230 }
0231 
0232 //---------------------------------------------------------------------------//
0233 // EXPLICIT INSTANTIATION
0234 //---------------------------------------------------------------------------//
0235 
0236 extern template class CoreState<MemSpace::host>;
0237 extern template class CoreState<MemSpace::device>;
0238 
0239 //---------------------------------------------------------------------------//
0240 }  // namespace celeritas