celeritas/global/ActionLauncher.device.hh

0001 //---------------------------------*-CUDA-*----------------------------------//
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/ActionLauncher.device.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <type_traits>
0011
0012 #include "corecel/DeviceRuntimeApi.hh"
0013
0014 #include "corecel/Assert.hh"
0015 #include "corecel/Macros.hh"
0016 #include "corecel/Types.hh"
0017 #include "corecel/cont/Range.hh"
0018 #include "corecel/sys/KernelLauncher.device.hh"
0019 #include "corecel/sys/ThreadId.hh"
0020 #include "celeritas/track/TrackInitParams.hh"
0021
0022 #include "ActionInterface.hh"
0023 #include "CoreParams.hh"
0024 #include "CoreState.hh"
0025
0026 namespace celeritas
0027 {
0028 //---------------------------------------------------------------------------//
0029 /*!
0030  * Profile and launch core stepping loop kernels.
0031  *
0032  * This is an extension to \c KernelLauncher which uses an action's label and
0033  * takes core params/state to determine the launch size and/or action range.
0034  *
0035  * Example:
0036  * \code
0037  void FooAction::step(CoreParams const& params,
0038                       CoreStateDevice& state) const
0039  {
0040    auto execute_thread = make_blah_executor(blah);
0041    static ActionLauncher<decltype(execute_thread)> const launch_kernel(*this);
0042    launch_kernel(state, execute_thread);
0043  }
0044  * \endcode
0045  */
0046 template<class F>
0047 class ActionLauncher : public KernelLauncher<F>
0048 {
0049     static_assert((std::is_trivially_copyable_v<F> || CELERITAS_USE_HIP
0050                    || CELER_COMPILER == CELER_COMPILER_CLANG)
0051                       && !std::is_pointer_v<F> && !std::is_reference_v<F>,
0052                   "Launched action must be a trivially copyable function "
0053                   "object");
0054     using StepActionT = CoreStepActionInterface;
0055
0056   public:
0057     // Create a launcher from a string
0058     using KernelLauncher<F>::KernelLauncher;
0059
0060     // Create a launcher from an action
0061     explicit ActionLauncher(StepActionT const& action);
0062
0063     // Create a launcher with a string extension
0064     ActionLauncher(StepActionT const& action, std::string_view ext);
0065
0066     // Launch a kernel for a thread range or number of threads
0067     using KernelLauncher<F>::operator();
0068
0069     // Launch a kernel for the wrapped executor
0070     void operator()(CoreState<MemSpace::device> const& state,
0071                     F const& call_thread) const;
0072
0073     // Launch with reduced grid size for when tracks are sorted
0074     void operator()(StepActionT const& action,
0075                     CoreParams const& params,
0076                     CoreState<MemSpace::device> const& state,
0077                     F const& call_thread) const;
0078 };
0079
0080 //---------------------------------------------------------------------------//
0081 /*!
0082  * Create a launcher from an action.
0083  */
0084 template<class F>
0085 ActionLauncher<F>::ActionLauncher(StepActionT const& action)
0086     : ActionLauncher{action.label()}
0087 {
0088 }
0089
0090 //---------------------------------------------------------------------------//
0091 /*!
0092  * Create a launcher with a string extension.
0093  */
0094 template<class F>
0095 ActionLauncher<F>::ActionLauncher(StepActionT const& action,
0096                                   std::string_view ext)
0097     : ActionLauncher{std::string(action.label()) + "-" + std::string(ext)}
0098 {
0099 }
0100
0101 //---------------------------------------------------------------------------//
0102 /*!
0103  * Launch a kernel for the wrapped executor.
0104  */
0105 template<class F>
0106 void ActionLauncher<F>::operator()(CoreState<MemSpace::device> const& state,
0107                                    F const& call_thread) const
0108 {
0109     return (*this)(
0110         range(ThreadId{state.size()}), state.stream_id(), call_thread);
0111 }
0112
0113 //---------------------------------------------------------------------------//
0114 /*!
0115  * Launch with reduced grid size for when tracks are sorted.
0116  *
0117  * These argument should be consistent with those in \c ActionLauncher.hh .
0118  */
0119 template<class F>
0120 void ActionLauncher<F>::operator()(StepActionT const& action,
0121                                    CoreParams const& params,
0122                                    CoreState<MemSpace::device> const& state,
0123                                    F const& call_thread) const
0124 {
0125     if (state.has_action_range()
0126         && is_action_sorted(action.order(), params.init()->track_order()))
0127     {
0128         // Launch on a subset of threads
0129         return (*this)(state.get_action_range(action.action_id()),
0130                        state.stream_id(),
0131                        call_thread);
0132     }
0133     else
0134     {
0135         // Not partitioned by action: launch on all threads
0136         return (*this)(state, call_thread);
0137     }
0138 }
0139
0140 //---------------------------------------------------------------------------//
0141 }  // namespace celeritas