Enable optical photon track offloading from Geant4 to Celeritas

src/accel/CMakeLists.txt

@@ -31,6 +31,7 @@ list(APPEND SOURCES
   GeantStepDiagnostic.cc
   IntegrationBase.cc
   LocalOpticalGenOffload.cc
+  LocalOpticalTrackOffload.cc
   LocalTransporter.cc
   Logger.cc
   PGPrimaryGeneratorAction.cc

src/accel/LocalOpticalTrackOffload.cc

@@ -0,0 +1,209 @@
+//------------------------------- -*- C++ -*- -------------------------------//
+// Copyright Celeritas contributors: see top-level COPYRIGHT file for details
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file accel/LocalOpticalTrackOffload.cc
+//---------------------------------------------------------------------------//
+#include "LocalOpticalTrackOffload.hh"
+
+#include <CLHEP/Units/SystemOfUnits.h>
+#include <G4EventManager.hh>
+#include <G4MTRunManager.hh>
+
+#include "corecel/Assert.hh"
+#include "corecel/sys/ScopedProfiling.hh"
+#include "geocel/GeantUtils.hh"
+#include "geocel/g4/Convert.hh"
+#include "celeritas/ext/GeantUnits.hh"
+#include "celeritas/global/CoreParams.hh"
+#include "celeritas/optical/CoreParams.hh"
+#include "celeritas/optical/Transporter.hh"
+
+#include "SetupOptions.hh"
+#include "SharedParams.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Offload Geant4 optical photon tracks to Celeritas
+ */
+LocalOpticalTrackOffload::LocalOpticalTrackOffload(SetupOptions const& options,
+                                                   SharedParams& params)
+{
+    CELER_VALIDATE(params.mode() == SharedParams::Mode::enabled,
+                   << "cannot create local optical track offload when "
+                      "Celeritas "
+                      "offloading is disabled");
+
+    // Save a pointer to the optical transporter
+    transport_ = params.optical_problem_loaded().transporter;
+
+    // Save a pointer to the direct generator action to insert tracks
+    direct_gen_
+        = std::dynamic_pointer_cast<optical::DirectGeneratorAction const>(
+            params.optical_problem_loaded().generator);
+
+    CELER_ASSERT(transport_);
+    CELER_ASSERT(transport_->params());
+
+    auto const& optical_params = *transport_->params();
+
+    // Check the thread ID and MT model
+    validate_geant_threading(optical_params.max_streams());
+
+    CELER_EXPECT(options.optical);
+    auto const& capacity = options.optical->capacity;
+    auto_flush_ = capacity.tracks;
+
+    auto stream_id = id_cast<StreamId>(get_geant_thread_id());
+
+    // Allocate thread-local state data
+    auto memspace = celeritas::device() ? MemSpace::device : MemSpace::host;
+    if (memspace == MemSpace::device)
+    {
+        state_ = std::make_shared<optical::CoreState<MemSpace::device>>(
+            optical_params, stream_id, capacity.tracks);
+    }
+    else
+    {
+        state_ = std::make_shared<optical::CoreState<MemSpace::host>>(
+            optical_params, stream_id, capacity.tracks);
+    }
+
+    // Allocate auxiliary data
+    if (optical_params.aux_reg())
+    {
+        state_->aux() = std::make_shared<AuxStateVec>(
+            *optical_params.aux_reg(), memspace, stream_id, capacity.tracks);
+    }
+
+    CELER_ENSURE(*this);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Initialize with options and shared data.
+ */
+void LocalOpticalTrackOffload::Initialize(SetupOptions const& options,
+                                          SharedParams& params)
+{
+    *this = LocalOpticalTrackOffload(options, params);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Set the event ID and reseed the Celeritas RNG at the start of an event.
+ */
+void LocalOpticalTrackOffload::InitializeEvent(int id)
+{
+    CELER_EXPECT(*this);
+    CELER_EXPECT(id >= 0);
+
+    event_id_ = id_cast<UniqueEventId>(id);
+    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));
+    }
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Buffer optical tracks.
+ */
+void LocalOpticalTrackOffload::Push(G4Track& g4track)
+{
+    CELER_EXPECT(*this);
+
+    ++num_pushed_;
+
+    CELER_EXPECT(g4track.GetDefinition());
+    CELER_EXPECT(g4track.GetDefinition()->GetParticleName() == "opticalphoton");
+
+    // Convert Geant4 track to optical::TrackInitializer
+    TrackData init;
+
+    init.energy = units::MevEnergy(
+        convert_from_geant(g4track.GetKineticEnergy(), CLHEP::MeV));
+
+    init.position = convert_from_geant(g4track.GetPosition(), CLHEP::cm);
+
+    init.direction = convert_from_geant(g4track.GetMomentumDirection(), 1);
+
+    init.time = convert_from_geant(g4track.GetGlobalTime(), CLHEP::second);
+    init.polarization = convert_from_geant(g4track.GetPolarization(), 1);
+
+    ScopedProfiling profile_this{"push"};
+
+    buffer_.push_back(init);
+
+    if (buffer_.size() >= auto_flush_)
+    {
+        this->Flush();
+    }
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Flush buffered optical photon tracks.
+ */
+void LocalOpticalTrackOffload::Flush()
+{
+    CELER_EXPECT(*this);
+
+    if (buffer_.empty())
+    {
+        return;
+    }
+
+    ScopedProfiling profile_this("flush");
+
+    // Insert tracks
+    if (direct_gen_)
+    {
+        direct_gen_->insert(*state_, make_span(buffer_));
+    }
+
+    // Transport tracks
+    (*transport_)(*state_);
+
+    ++num_flushed_;
+    buffer_.clear();
+}
+//---------------------------------------------------------------------------//
+auto LocalOpticalTrackOffload::GetActionTime() const -> MapStrDbl
+{
+    CELER_EXPECT(*this);
+    // TODO Add Per-track optical transport action timing once
+    // optical track insertion and transport are implemented.
+    return transport_->get_action_times(*state_->aux());
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Finalize the local optical track offload state
+ */
+void LocalOpticalTrackOffload::Finalize()
+{
+    CELER_EXPECT(*this);
+
+    CELER_VALIDATE(buffer_.empty(),
+                   << buffer_.size() << " optical tracks were not flushed");
+
+    CELER_LOG(info) << "Finalizing Celeritas after " << num_pushed_
+                    << " optical tracks pushed (over " << num_flushed_
+                    << " ) flushes";
+
+    *this = {};
+
+    CELER_ENSURE(!*this);
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace celeritas

src/accel/LocalOpticalTrackOffload.hh

@@ -0,0 +1,102 @@
+//------------------------------- -*- C++ -*- -------------------------------//
+// Copyright Celeritas contributors: see top-level COPYRIGHT file for details
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file accel/LocalOpticalTrackOffload.hh
+//---------------------------------------------------------------------------//
+#pragma once
+
+#include "corecel/Types.hh"
+#include "celeritas/Types.hh"
+#include "celeritas/optical/TrackInitializer.hh"
+#include "celeritas/optical/gen/DirectGeneratorAction.hh"
+
+#include "TrackOffloadInterface.hh"
+
+namespace celeritas
+{
+namespace optical
+{
+class CoreStateBase;
+class Transporter;
+}  // namespace optical
+
+struct SetupOptions;
+class SharedParams;
+
+//---------------------------------------------------------------------------//
+/*!
+ * Offload Geant4 optical photon tracks to Celeritas.
+ */
+class LocalOpticalTrackOffload final : public TrackOffloadInterface
+{
+  public:
+    //!@{
+    //! \name Type aliases
+    using TrackData = optical::TrackInitializer;
+    //!@}
+
+  public:
+    // Construct in an invalid state
+    LocalOpticalTrackOffload() = default;
+
+    // Construct with shared (across threads) params
+    LocalOpticalTrackOffload(SetupOptions const& options, SharedParams& params);
+
+    //!@{
+    //! \name TrackOffloadInterface
+
+    // Initialize with options and shared data
+    void Initialize(SetupOptions const&, SharedParams&) final;
+
+    // Set the event ID and reseed the Celeritas RNG at the start of an event
+    void InitializeEvent(int) final;
+
+    // Transport all buffered tracks to completion
+    void Flush() final;
+
+    // Clear local data and return to an invalid state
+    void Finalize() final;
+
+    // Whether the class instance is initialized
+    bool Initialized() const final { return static_cast<bool>(transport_); }
+
+    // Number of buffered tracks
+    size_type GetBufferSize() const final { return buffer_.size(); }
+
+    // Get accumulated action times
+    MapStrDbl GetActionTime() const final;
+    //!@}
+
+    // Offload optical distribution track to Celeritas
+    void Push(G4Track&) final;
+    // Number of optical tracks pushed to offload
+    size_type num_pushed() const { return num_pushed_; }
+
+  private:
+    // Transport pending optical tracks
+    std::shared_ptr<optical::Transporter> transport_;
+
+    // Thread-local state data
+    std::shared_ptr<optical::CoreStateBase> state_;
+
+    std::shared_ptr<optical::DirectGeneratorAction const> direct_gen_;
+
+    // Buffered tracks for offloading
+    std::vector<TrackData> buffer_;
+
+    // Number of photons tracks to buffer before offloading
+    size_type auto_flush_{};
+
+    // Accumulated number of optical photon tracks
+    size_type num_pushed_{};
+
+    // Accumulated number of tracks pushed over flushes
+    size_type num_flushed_{};
+
+    //  Current event ID for obtaining it
+    UniqueEventId event_id_;
+};
+
+//---------------------------------------------------------------------------//
+}  // namespace celeritas

src/accel/SetupOptions.hh

@@ -128,6 +128,8 @@ struct OpticalSetupOptions
     inp::OpticalGenerator generator;
     //! Limits for the optical stepping loop
     inp::OpticalTrackingLimits limits;
+
+    bool offload_optical_tracks = false;
 };
 
 //---------------------------------------------------------------------------//

src/accel/UserActionIntegration.cc

@@ -13,10 +13,10 @@
 #include "corecel/sys/Stopwatch.hh"
 
 #include "ExceptionConverter.hh"
+#include "G4OpticalPhoton.hh"
 #include "TimeOutput.hh"  // IWYU pragma: keep
 
 #include "detail/IntegrationSingleton.hh"
-
 namespace celeritas
 {
 //---------------------------------------------------------------------------//
@@ -75,6 +75,21 @@ void UserActionIntegration::PreUserTrackingAction(G4Track* track)
 
         // Either "pushed" or we're in kill_offload mode
         track->SetTrackStatus(fStopAndKill);
+        return;
+    }
+    if (track->GetDefinition() == G4OpticalPhoton::Definition())
+    {
+        if (mode == SharedParams::Mode::enabled)
+        {
+            // Celeritas is transporting this optical photon
+            CELER_TRY_HANDLE(detail::IntegrationSingleton::instance()
+                                 .local_track_offload()
+                                 .Push(*track),
+                             ExceptionConverter("celer.track.push",
+                                                +&singleton.shared_params()));
+        }
+        // Kill the optical photon (offloaded or not)
+        track->SetTrackStatus(fStopAndKill);
     }
 }
 

src/accel/detail/IntegrationSingleton.cc

@@ -6,6 +6,7 @@
 //---------------------------------------------------------------------------//
 #include "IntegrationSingleton.hh"
 
+#include <memory>
 #include <G4RunManager.hh>
 #include <G4Threading.hh>
 
@@ -14,6 +15,7 @@
 #include "corecel/io/Logger.hh"
 #include "corecel/sys/ScopedMpiInit.hh"
 #include "geocel/GeantUtils.hh"
+#include "accel/LocalOpticalTrackOffload.hh"
 
 #include "LoggerImpl.hh"
 #include "../ExceptionConverter.hh"
@@ -104,8 +106,17 @@ LocalOffloadInterface& IntegrationSingleton::local_offload()
             && std::holds_alternative<inp::OpticalOffloadGenerator>(
                 options_.optical->generator))
         {
+            CELER_LOG(info) << "optical gen offloading enabled";
             offload = std::make_unique<LocalOpticalGenOffload>();
         }
+        else if (options_.optical
+                 && std::holds_alternative<inp::OpticalTrackOffload>(
+                     options_.optical->generator))
+
+        {
+            CELER_LOG(info) << "optical track offloading enabled";
+            offload = std::make_unique<LocalOpticalTrackOffload>();
+        }
         else
         {
             // TODO: if offloading direct optical tracks, return optical