Add environment variable to kill Celeritas-supported tracks in Geant4 (…

…#1060)
celeritas-project · Dec 13, 2023 · 8670a6f · 8670a6f
1 parent cb176ff
commit 8670a6f
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 21 deletions.
diff --git a/app/celer-g4/TrackingAction.cc b/app/celer-g4/TrackingAction.cc
@@ -50,7 +50,7 @@ void TrackingAction::PreUserTrackingAction(G4Track const* track)
                  == !SharedParams::CeleritasDisabled());
     CELER_EXPECT(static_cast<bool>(*params_) == static_cast<bool>(*transport_));
 
-    if (SharedParams::CeleritasDisabled())
+    if (SharedParams::CeleritasDisabled() && !SharedParams::KillOffloadTracks())
         return;
 
     auto const& allowed_particles = params_->OffloadParticles();
@@ -59,9 +59,12 @@ void TrackingAction::PreUserTrackingAction(G4Track const* track)
                   track->GetDefinition())
         != std::end(allowed_particles))
     {
-        // Celeritas is transporting this track
-        ExceptionConverter call_g4exception{"celer0003", params_.get()};
-        CELER_TRY_HANDLE(transport_->Push(*track), call_g4exception);
+        if (!SharedParams::CeleritasDisabled())
+        {
+            // Celeritas is transporting this track
+            ExceptionConverter call_g4exception{"celer0003", params_.get()};
+            CELER_TRY_HANDLE(transport_->Push(*track), call_g4exception);
+        }
         const_cast<G4Track*>(track)->SetTrackStatus(fStopAndKill);
     }
 }

diff --git a/app/celer-g4/celer-g4.cc b/app/celer-g4/celer-g4.cc
@@ -88,6 +88,7 @@ void print_usage(std::string_view exec_name)
                  "  CELER_DISABLE: nonempty disables offloading\n"
                  "  CELER_DISABLE_DEVICE: nonempty disables CUDA\n"
                  "  CELER_DISABLE_ROOT: nonempty disables ROOT I/O\n"
+                 "  CELER_KILL_OFFLOAD: nonempty kills offload tracks\n"
                  "  CELER_LOG: global logging level\n"
                  "  CELER_LOG_LOCAL: thread-local logging level\n"
               << std::endl;

diff --git a/doc/main/usage.rst b/doc/main/usage.rst
@@ -280,6 +280,7 @@ tell what variables are in use or may be useful.
  HEPMC3_VERBOSE          celeritas HepMC3 debug verbosity
  VECGEOM_VERBOSE         celeritas VecGeom CUDA verbosity
  CELER_DISABLE           accel     Disable Celeritas offloading entirely
+ CELER_KILL_OFFLOAD      accel     Kill Celeritas-supported tracks in Geant4
  CELER_STRIP_SOURCEDIR   accel     Strip directories from exception output
  ======================= ========= ==========================================
 

diff --git a/src/accel/SharedParams.cc b/src/accel/SharedParams.cc
@@ -14,8 +14,11 @@
 #include <utility>
 #include <vector>
 #include <CLHEP/Random/Random.h>
+#include <G4Electron.hh>
+#include <G4Gamma.hh>
 #include <G4ParticleDefinition.hh>
 #include <G4ParticleTable.hh>
+#include <G4Positron.hh>
 #include <G4RunManager.hh>
 #include <G4Threading.hh>
 
@@ -170,6 +173,23 @@ bool SharedParams::CeleritasDisabled()
     return result;
 }
 
+bool SharedParams::KillOffloadTracks()
+{
+    static bool const result = [] {
+        if (celeritas::getenv("CELER_KILL_OFFLOAD").empty())
+            return false;
+
+        if (CeleritasDisabled())
+        {
+            CELER_LOG(info) << "Killing Geant4 tracks supported by Celeritas "
+                               "offloading since the 'CELER_KILL_OFFLOAD' "
+                               "environment variable is present and non-empty";
+        }
+        return true;
+    }();
+    return result;
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Set up Celeritas using Geant4 data and existing output registery.
@@ -335,6 +355,28 @@ void SharedParams::Finalize()
     CELER_ENSURE(!*this);
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Get a vector of particles supported by Celeritas offloading.
+ */
+auto SharedParams::OffloadParticles() const -> VecG4ParticleDef const&
+{
+    if (!CeleritasDisabled())
+    {
+        // Get the supported particles from Celeritas
+        CELER_ASSERT(*this);
+        return particles_;
+    }
+
+    // In a Geant4-only simulation, use a hardcoded list of supported particles
+    static VecG4ParticleDef const particles = {
+        G4Gamma::Gamma(),
+        G4Electron::Electron(),
+        G4Positron::Positron(),
+    };
+    return particles;
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Lazily obtained number of streams.

diff --git a/src/accel/SharedParams.hh b/src/accel/SharedParams.hh
@@ -65,6 +65,9 @@ class SharedParams
     // True if Celeritas is globally disabled using the CELER_DISABLE env
     static bool CeleritasDisabled();
 
+    // Whether to kill tracks that would have been offloaded
+    static bool KillOffloadTracks();
+
     // Construct in an uninitialized state
     SharedParams() = default;
 
@@ -88,7 +91,7 @@ class SharedParams
     inline SPConstParams Params() const;
 
     // Get a vector of particles supported by Celeritas offloading
-    inline VecG4ParticleDef const& OffloadParticles() const;
+    VecG4ParticleDef const& OffloadParticles() const;
 
     //! Whether Celeritas core params have been created
     explicit operator bool() const { return static_cast<bool>(params_); }
@@ -172,18 +175,6 @@ auto SharedParams::Params() const -> SPConstParams
     return params_;
 }
 
-//---------------------------------------------------------------------------//
-/*!
- * Get a vector of particles supported by Celeritas offloading.
- *
- * This can only be called after \c Initialize.
- */
-auto SharedParams::OffloadParticles() const -> VecG4ParticleDef const&
-{
-    CELER_EXPECT(*this);
-    return particles_;
-}
-
 //---------------------------------------------------------------------------//
 /*!
  * Hit manager, to be used only by LocalTransporter.

diff --git a/src/accel/SimpleOffload.cc b/src/accel/SimpleOffload.cc
@@ -105,17 +105,20 @@ void SimpleOffload::BeginOfEventAction(G4Event const* event)
  */
 void SimpleOffload::PreUserTrackingAction(G4Track* track)
 {
-    if (!*this)
+    if (!*this && !SharedParams::KillOffloadTracks())
         return;
 
     if (std::find(params_->OffloadParticles().begin(),
                   params_->OffloadParticles().end(),
                   track->GetDefinition())
         != params_->OffloadParticles().end())
     {
-        // Celeritas is transporting this track
-        ExceptionConverter call_g4exception{"celer0003", params_};
-        CELER_TRY_HANDLE(local_->Push(*track), call_g4exception);
+        if (!SharedParams::CeleritasDisabled())
+        {
+            // Celeritas is transporting this track
+            ExceptionConverter call_g4exception{"celer0003", params_};
+            CELER_TRY_HANDLE(local_->Push(*track), call_g4exception);
+        }
         track->SetTrackStatus(fStopAndKill);
     }
 }