celeritas/global/TrackExecutor.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2022-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/TrackExecutor.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include "corecel/Assert.hh"
0011 #include "corecel/Types.hh"
0012 #include "corecel/math/Algorithms.hh"
0013 #include "corecel/sys/ThreadId.hh"
0014 #include "celeritas/track/SimFunctors.hh"
0015
0016 #include "CoreTrackData.hh"
0017 #include "CoreTrackDataFwd.hh"
0018 #include "CoreTrackView.hh"
0019
0020 namespace celeritas
0021 {
0022 //---------------------------------------------------------------------------//
0023 /*!
0024  * Call a \c CoreTrackView executor for a given ThreadId.
0025  *
0026  * This class can be used to call a functor that applies to \c CoreTrackView
0027  * using a \c ThreadId, so that the tracks can be easily looped over as a
0028  * group on CPU or GPU. It applies a remapping from \em thread to \em slot if
0029  * the tracks are sorted. Otherwise, thread and track slot have the same
0030  * numerical value.
0031  *
0032  * This is primarily used by \c ActionLauncher .
0033  *
0034  * \code
0035 void foo_kernel(CoreParamsPtr const params,
0036                 CoreStatePtr const state)
0037 {
0038     TrackExecutor execute{params, state, MyTrackApplier{}};
0039
0040     for (auto tid : range(ThreadID{123}))
0041     {
0042         step(tid);
0043     }
0044 }
0045 \endcode
0046  *
0047  * \todo Rename to ThreadExecutor. The template parameter, which must operate
0048  * on a core track view, is a track executor.
0049  */
0050 template<class T>
0051 class TrackExecutor
0052 {
0053   public:
0054     //!@{
0055     //! \name Type aliases
0056     using ParamsPtr = CoreParamsPtr<MemSpace::native>;
0057     using StatePtr = CoreStatePtr<MemSpace::native>;
0058     using Applier = T;
0059     //!@}
0060
0061   public:
0062     //! Construct with core data and executor
0063     CELER_FUNCTION
0064     TrackExecutor(ParamsPtr params, StatePtr state, T&& execute_track)
0065         : params_{params}
0066         , state_{state}
0067         , execute_track_{celeritas::forward<T>(execute_track)}
0068     {
0069     }
0070
0071     //! Call the underlying function, using indirection array if needed
0072     CELER_FUNCTION void operator()(ThreadId thread)
0073     {
0074         CELER_EXPECT(thread < state_->size());
0075         CoreTrackView track(*params_, *state_, thread);
0076         return execute_track_(track);
0077     }
0078
0079   private:
0080     ParamsPtr const params_;
0081     StatePtr const state_;
0082     T execute_track_;
0083 };
0084
0085 //---------------------------------------------------------------------------//
0086 /*!
0087  * Launch the track only when a certain condition applies to the sim state.
0088  *
0089  * The condition \c C must have the signature \code
0090  * (SimTrackView const&) -> bool
0091   \endcode
0092  *
0093  * see \c make_active_track_executor for an example where this is used to apply
0094  * only to active (or killed) tracks.
0095  */
0096 template<class C, class T>
0097 class ConditionalTrackExecutor
0098 {
0099   public:
0100     //!@{
0101     //! \name Type aliases
0102     using ParamsPtr = CoreParamsPtr<MemSpace::native>;
0103     using StatePtr = CoreStatePtr<MemSpace::native>;
0104     using Applier = T;
0105     //!@}
0106
0107   public:
0108     //! Construct with condition and operator
0109     CELER_FUNCTION
0110     ConditionalTrackExecutor(ParamsPtr params,
0111                              StatePtr state,
0112                              C&& applies,
0113                              T&& execute_track)
0114         : params_{params}
0115         , state_{state}
0116         , applies_{celeritas::forward<C>(applies)}
0117         , execute_track_{celeritas::forward<T>(execute_track)}
0118     {
0119     }
0120
0121     //! Launch the given thread if the track meets the condition
0122     CELER_FUNCTION void operator()(ThreadId thread)
0123     {
0124         CELER_EXPECT(thread < state_->size());
0125         CoreTrackView track(*params_, *state_, thread);
0126         if (!applies_(track.make_sim_view()))
0127         {
0128             return;
0129         }
0130
0131         return execute_track_(track);
0132     }
0133
0134   private:
0135     ParamsPtr const params_;
0136     StatePtr const state_;
0137     C applies_;
0138     T execute_track_;
0139 };
0140
0141 //---------------------------------------------------------------------------//
0142 // DEDUCTION GUIDES
0143 //---------------------------------------------------------------------------//
0144 template<class T>
0145 CELER_FUNCTION TrackExecutor(CoreParamsPtr<MemSpace::native>,
0146                              CoreStatePtr<MemSpace::native>,
0147                              T&&) -> TrackExecutor<T>;
0148
0149 template<class C, class T>
0150 CELER_FUNCTION ConditionalTrackExecutor(CoreParamsPtr<MemSpace::native>,
0151                                         CoreStatePtr<MemSpace::native>,
0152                                         C&&,
0153                                         T&&) -> ConditionalTrackExecutor<C, T>;
0154
0155 //---------------------------------------------------------------------------//
0156 // FREE FUNCTIONS
0157 //---------------------------------------------------------------------------//
0158 /*!
0159  * Return a track executor that only applies to active, non-errored tracks.
0160  */
0161 template<class T>
0162 inline CELER_FUNCTION decltype(auto)
0163 make_active_track_executor(CoreParamsPtr<MemSpace::native> params,
0164                            CoreStatePtr<MemSpace::native> const& state,
0165                            T&& apply_track)
0166 {
0167     return ConditionalTrackExecutor{
0168         params, state, AppliesValid{}, celeritas::forward<T>(apply_track)};
0169 }
0170
0171 //---------------------------------------------------------------------------//
0172 /*!
0173  * Return a track executor that only applies if the action ID matches.
0174  *
0175  * \note This should generally only be used for post-step actions and other
0176  * cases where the IDs *explicitly* are set. Many explicit actions apply to all
0177  * threads, active or not.
0178  */
0179 template<class T>
0180 inline CELER_FUNCTION decltype(auto)
0181 make_action_track_executor(CoreParamsPtr<MemSpace::native> params,
0182                            CoreStatePtr<MemSpace::native> state,
0183                            ActionId action,
0184                            T&& apply_track)
0185 {
0186     CELER_EXPECT(action);
0187     return ConditionalTrackExecutor{params,
0188                                     state,
0189                                     IsStepActionEqual{action},
0190                                     celeritas::forward<T>(apply_track)};
0191 }
0192
0193 //---------------------------------------------------------------------------//
0194 /*!
0195  * Return a track executor that only applies for the given along-step action.
0196  */
0197 template<class T>
0198 inline CELER_FUNCTION decltype(auto)
0199 make_along_step_track_executor(CoreParamsPtr<MemSpace::native> params,
0200                                CoreStatePtr<MemSpace::native> state,
0201                                ActionId action,
0202                                T&& apply_track)
0203 {
0204     CELER_EXPECT(action);
0205     return ConditionalTrackExecutor{params,
0206                                     state,
0207                                     IsAlongStepActionEqual{action},
0208                                     celeritas::forward<T>(apply_track)};
0209 }
0210
0211 //---------------------------------------------------------------------------//
0212 }  // namespace celeritas