Replace core state counter synchronization with device kernels and functors

src/accel/CMakeLists.txt

@@ -51,6 +51,7 @@ list(APPEND SOURCES
 )
 
 celeritas_polysource(ExceptionConverter)
+celeritas_polysource(LocalOpticalGenOffload)
 
 if(Geant4_VERSION VERSION_GREATER_EQUAL 11.0)
   list(APPEND SOURCES

src/accel/LocalOpticalGenOffload.cc

@@ -14,10 +14,11 @@
 #include "corecel/sys/ScopedProfiling.hh"
 #include "geocel/GeantUtils.hh"
 #include "celeritas/global/CoreParams.hh"
-#include "celeritas/optical/CoreParams.hh"
-#include "celeritas/optical/CoreState.hh"
+#include "celeritas/optical/TrackExecutor.hh"
 #include "celeritas/optical/Transporter.hh"
+#include "celeritas/optical/action/ActionLauncher.hh"
 #include "celeritas/optical/gen/GeneratorAction.hh"
+#include "celeritas/optical/gen/detail/UpdatePendingExecutor.hh"
 #include "celeritas/phys/GeneratorRegistry.hh"
 
 #include "SetupOptions.hh"
@@ -105,17 +106,14 @@ void LocalOpticalGenOffload::InitializeEvent(int id)
     CELER_EXPECT(id >= 0);
 
     event_id_ = id_cast<UniqueEventId>(id);
-    if constexpr (CELERITAS_RESEED == CELERITAS_RESEED_TRACKSLOT)
+
+    if (!(G4Threading::IsMultithreadedApplication()
+          && G4MTRunManager::SeedOncePerCommunication()))
     {
-        if (!(G4Threading::IsMultithreadedApplication()
-              && G4MTRunManager::SeedOncePerCommunication()))
-        {
-            // Since Geant4 schedules events dynamically, reseed the Celeritas
-            // RNGs using the Geant4 event ID for reproducibility. This
-            // guarantees that an event can be reproduced given the event ID.
-            state_->reseed(transport_->params()->rng(),
-                           id_cast<UniqueEventId>(id));
-        }
+        // Since Geant4 schedules events dynamically, reseed the Celeritas RNGs
+        // using the Geant4 event ID for reproducibility. This guarantees that
+        // an event can be reproduced given the event ID.
+        state_->reseed(transport_->params()->rng(), id_cast<UniqueEventId>(id));
     }
 }
 
@@ -186,9 +184,19 @@ void LocalOpticalGenOffload::Flush()
     // Copy the buffered distributions to device
     generate_->insert(*state_, make_span(buffer_));
 
-    auto counters = state_->sync_get_counters();
-    counters.num_pending += num_photons_;
-    state_->sync_put_counters(counters);
+    // Update the number of primaries waiting to be generated based on the
+    // number of photons.
+    if (celeritas::device())
+    {
+        auto* s = dynamic_cast<optical::CoreState<MemSpace::device>*>(&*state_);
+        this->update_primaries(*s);
+    }
+    else
+    {
+        auto* s = dynamic_cast<optical::CoreState<MemSpace::host>*>(&*state_);
+        this->update_primaries(*s);
+    }
+
     num_photons_ = 0;
     buffer_.clear();
 
@@ -243,5 +251,32 @@ void LocalOpticalGenOffload::Finalize()
     CELER_ENSURE(!*this);
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Call the UpdatePending functor to update the number of primaries to be
+ * generated to include the buffered optical photons; use only one host thread.
+ */
+void LocalOpticalGenOffload::update_primaries(
+    optical::CoreState<MemSpace::host>& state) const
+{
+    auto const& optical_params = *transport_->params();
+    auto execute_thread = make_single_track_executor(
+        optical_params.ptr<MemSpace::native>(),
+        state.ptr(),
+        optical::detail::UpdatePendingExecutor{num_photons_});
+    launch_action(1, execute_thread);
+}
+
+//---------------------------------------------------------------------------//
+// DEVICE-DISABLED IMPLEMENTATION
+//---------------------------------------------------------------------------//
+#if !CELER_USE_DEVICE
+inline void LocalOpticalGenOffload::update_primaries(
+    optical::CoreState<MemSpace::device>&) const
+{
+    CELER_NOT_CONFIGURED("CUDA OR HIP");
+}
+#endif
+
 //---------------------------------------------------------------------------//
 }  // namespace celeritas

src/accel/LocalOpticalGenOffload.cu

@@ -0,0 +1,36 @@
+//------------------------------ -*- cuda -*- -------------------------------//
+// Copyright Celeritas contributors: see top-level COPYRIGHT file for details
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file accel/LocalOpticalGenOffload.cu
+//---------------------------------------------------------------------------//
+#include "LocalOpticalGenOffload.hh"
+
+#include "celeritas/global/CoreParams.hh"
+#include "celeritas/optical/TrackExecutor.hh"
+#include "celeritas/optical/Transporter.hh"
+#include "celeritas/optical/action/ActionLauncher.device.hh"
+#include "celeritas/optical/gen/detail/UpdatePendingExecutor.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Call the UpdatePending functor to update number of primaries to be generated
+ * to include the buffered optical photons; use only one device thread.
+ */
+void LocalOpticalGenOffload::update_primaries(
+    optical::CoreState<MemSpace::device>& state) const
+{
+    auto const& optical_params = *transport_->params();
+    auto execute_thread = make_single_track_executor(
+        optical_params.ptr<MemSpace::native>(),
+        state.ptr(),
+        optical::detail::UpdatePendingExecutor{num_photons_});
+    static KernelLauncher<decltype(execute_thread)> const launch_kernel(
+        "update-pending");
+    launch_kernel(1, state.stream_id(), execute_thread);
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace celeritas

src/accel/LocalOpticalGenOffload.hh

@@ -11,6 +11,8 @@
 #include "corecel/Types.hh"
 #include "celeritas/Types.hh"
 #include "celeritas/inp/Control.hh"
+#include "celeritas/optical/CoreParams.hh"
+#include "celeritas/optical/CoreState.hh"
 #include "celeritas/optical/gen/GeneratorData.hh"
 
 #include "LocalOffloadInterface.hh"
@@ -80,6 +82,11 @@ class LocalOpticalGenOffload final : public LocalOffloadInterface
     explicit operator bool() const { return this->Initialized(); }
 
   private:
+    // Update the number of primaries waiting to be generated on host/device
+    // Called by Flush()
+    void update_primaries(optical::CoreState<MemSpace::host>&) const;
+    void update_primaries(optical::CoreState<MemSpace::device>&) const;
+
     // Transport pending optical tracks
     std::shared_ptr<optical::Transporter> transport_;
 

src/celeritas/CMakeLists.txt

@@ -397,24 +397,28 @@ celeritas_polysource(em/model/RelativisticBremModel)
 celeritas_polysource(em/model/SeltzerBergerModel)
 celeritas_polysource(em/model/CoulombScatteringModel)
 celeritas_polysource(geo/detail/BoundaryAction)
+celeritas_polysource(global/Stepper)
 celeritas_polysource(global/detail/KillActive)
 celeritas_polysource(global/detail/TrackSlotUtils)
 celeritas_polysource(mucf/model/DTMixMucfModel)
 celeritas_polysource(neutron/model/ChipsNeutronElasticModel)
 celeritas_polysource(neutron/model/NeutronCaptureModel)
 celeritas_polysource(neutron/model/NeutronInelasticModel)
 celeritas_polysource(optical/model/AbsorptionModel)
+celeritas_polysource(optical/Runner)
 celeritas_polysource(optical/model/MieModel)
 celeritas_polysource(optical/model/RayleighModel)
 celeritas_polysource(optical/model/WavelengthShiftModel)
 celeritas_polysource(optical/action/AlongStepAction)
 celeritas_polysource(optical/action/DetectorAction)
 celeritas_polysource(optical/action/DiscreteSelectAction)
+celeritas_polysource(optical/action/LocateVacanciesAction)
 celeritas_polysource(optical/action/PreStepAction)
 celeritas_polysource(optical/action/StepDiagnostic)
 celeritas_polysource(optical/action/TrackingCutAction)
 celeritas_polysource(optical/action/detail/TrackInitAlgorithms)
 celeritas_polysource(optical/gen/GeneratorAction)
+celeritas_polysource(optical/gen/GeneratorBase)
 celeritas_polysource(optical/gen/OffloadAction)
 celeritas_polysource(optical/gen/OffloadGatherAction)
 celeritas_polysource(optical/gen/PrimaryGeneratorAction)

src/celeritas/global/CoreTrackView.hh

@@ -120,8 +120,8 @@ class CoreTrackView
     inline CELER_FUNCTION CoreStateCounters const& counters() const;
 
   private:
-    StateRef const& states_;
     ParamsRef const& params_;
+    StateRef const& states_;
     ThreadId const thread_id_;
     TrackSlotId track_slot_id_;
 };
@@ -136,7 +136,7 @@ CELER_FUNCTION
 CoreTrackView::CoreTrackView(ParamsRef const& params,
                              StateRef const& states,
                              ThreadId thread)
-    : states_(states), params_(params), thread_id_(thread)
+    : params_(params), states_(states), thread_id_(thread)
 {
     CELER_EXPECT(states_.track_slots.empty()
                  || thread_id_ < states_.track_slots.size());
@@ -157,7 +157,7 @@ CELER_FUNCTION
 CoreTrackView::CoreTrackView(ParamsRef const& params,
                              StateRef const& states,
                              TrackSlotId track_slot)
-    : states_(states), params_(params), track_slot_id_(track_slot)
+    : params_(params), states_(states), track_slot_id_(track_slot)
 {
     CELER_EXPECT(track_slot_id_ < states_.size());
 }

src/celeritas/global/Stepper.cc

@@ -14,13 +14,16 @@
 #include "corecel/sys/ScopedProfiling.hh"
 #include "orange/OrangeData.hh"
 #include "celeritas/Types.hh"
+#include "celeritas/global/TrackExecutor.hh"
 #include "celeritas/random/RngReseed.hh"
 #include "celeritas/track/ExtendFromPrimariesAction.hh"
 #include "celeritas/track/TrackInitParams.hh"
 
+#include "ActionLauncher.hh"
 #include "CoreParams.hh"
 
 #include "detail/KillActive.hh"
+#include "detail/SetGeneratedExecutor.hh"
 
 namespace celeritas
 {
@@ -127,13 +130,10 @@ template<MemSpace M>
 auto Stepper<M>::operator()() -> result_type
 {
     ScopedProfiling profile_this{"step"};
-    auto counters = state_->sync_get_counters();
-    counters.num_generated = 0;
-    counters.num_cut = 0;
-    counters.num_errored = 0;
-    state_->sync_put_counters(counters);
+    // Initialize the num_generated counter to zero
+    this->set_generated();
     actions_->step(*params_, *state_);
-    counters = state_->sync_get_counters();
+    auto counters = state_->sync_get_counters();
 
     // Get the number of track initializers and active tracks
     result_type result;
@@ -169,9 +169,6 @@ auto Stepper<M>::operator()(SpanConstPrimary primaries) -> result_type
                    << "event number " << max_id->event_id.unchecked_get()
                    << " exceeds max_events=" << params_->init()->max_events());
 
-    auto counters = state_->sync_get_counters();
-    counters.num_pending = primaries.size();
-    state_->sync_put_counters(counters);
     primaries_action_->insert(*params_, *state_, primaries);
 
     return (*this)();
@@ -213,6 +210,31 @@ void Stepper<M>::reseed(UniqueEventId event_id)
     params_->init()->reset_track_ids(state_->stream_id(), &state_->ref().init);
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Set the num_pending counter to the number of generated primaries.
+ */
+template<>
+void Stepper<MemSpace::host>::set_generated()
+{
+    auto execute_thread
+        = make_single_track_executor(params_->ptr<MemSpace::native>(),
+                                     state_->ptr(),
+                                     detail::SetGeneratedExecutor{});
+    launch_core(1, "set-generated", *params_, *state_, execute_thread);
+}
+
+//---------------------------------------------------------------------------//
+// DEVICE-DISABLED IMPLEMENTATION
+//---------------------------------------------------------------------------//
+#if !CELER_USE_DEVICE
+template<>
+void Stepper<MemSpace::device>::set_generated()
+{
+    CELER_NOT_CONFIGURED("CUDA OR HIP");
+}
+#endif
+
 //---------------------------------------------------------------------------//
 // EXPLICIT INSTANTIATION
 //---------------------------------------------------------------------------//

src/celeritas/global/Stepper.cu

@@ -0,0 +1,38 @@
+//------------------------------ -*- cuda -*- -------------------------------//
+// Copyright Celeritas contributors: see top-level COPYRIGHT file for details
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/global/Stepper.cu
+//---------------------------------------------------------------------------//
+#include "Stepper.hh"
+
+#include "corecel/Assert.hh"
+#include "corecel/Types.hh"
+#include "corecel/sys/KernelLauncher.device.hh"
+#include "celeritas/global/TrackExecutor.hh"
+
+#include "CoreParams.hh"
+#include "CoreState.hh"
+
+#include "detail/SetGeneratedExecutor.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Set the num_pending counter to the number of generated primaries.
+ */
+template<>
+void Stepper<MemSpace::device>::set_generated()
+{
+    auto execute_thread
+        = make_single_track_executor(params_->ptr<MemSpace::native>(),
+                                     state_->ptr(),
+                                     detail::SetGeneratedExecutor{});
+    static KernelLauncher<decltype(execute_thread)> const launch_kernel(
+        "set-generated");
+    launch_kernel(1, state_->stream_id(), execute_thread);
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace celeritas

src/celeritas/global/Stepper.hh

@@ -148,6 +148,8 @@ class Stepper final : public StepperInterface
     //!@{
     //! \name Type aliases
     using StateRef = CoreStateData<Ownership::reference, M>;
+    using CoreStateHost = CoreState<MemSpace::host>;
+    using CoreStateDevice = CoreState<MemSpace::device>;
     //!@}
 
   public:
@@ -184,6 +186,9 @@ class Stepper final : public StepperInterface
     //! Reset the core state counters and data so it can be reused
     void reset_state() { state_->reset(); }
 
+    //! Reset the num_generated state counter to zero
+    void set_generated();
+
     //! Get a shared pointer to the state (TEMPORARY, DO NOT USE)
     SPState sp_state() final { return state_; }
 
@@ -199,11 +204,9 @@ class Stepper final : public StepperInterface
 };
 
 //---------------------------------------------------------------------------//
-// EXPLICIT INSTANTIATION
+// EXPLICIT INSTANTIATION removed but retained in Stepper.cc so that the
+// set_generated() member function can be specialized based on MemSpace
 //---------------------------------------------------------------------------//
 
-extern template class Stepper<MemSpace::host>;
-extern template class Stepper<MemSpace::device>;
-
 //---------------------------------------------------------------------------//
 }  // namespace celeritas