Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-05-18 08:22:42

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