From 96d4e589f1e01854610c278039c0102c815e2ff3 Mon Sep 17 00:00:00 2001
From: "Seth R. Johnson" <johnsonsr@ornl.gov>
Date: Fri, 10 Jan 2025 10:53:35 -0500
Subject: [PATCH] Add support for Geant4 navigator through offload interface
 (#1558)

---
 .github/workflows/build-spack.yml |  1 -
 app/celer-g4/CMakeLists.txt       |  2 +-
 src/accel/LocalTransporter.cc     | 26 +++++++++++++++-----------
 src/accel/detail/HitManager.cc    |  6 ++++--
 4 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/build-spack.yml b/.github/workflows/build-spack.yml
index e4acf2c9fd..e9abe6a1ec 100644
--- a/.github/workflows/build-spack.yml
+++ b/.github/workflows/build-spack.yml
@@ -260,7 +260,6 @@ jobs:
             ${{
               (   (matrix.special == 'minimal')
                || (matrix.special == 'float')
-               || (matrix.geometry == 'geant4')
                || !matrix.geant
               ) && '1' || ''
             }}
diff --git a/app/celer-g4/CMakeLists.txt b/app/celer-g4/CMakeLists.txt
index eda68dd9cf..b603e6f44d 100644
--- a/app/celer-g4/CMakeLists.txt
+++ b/app/celer-g4/CMakeLists.txt
@@ -88,7 +88,7 @@ if(NOT CELERITAS_USE_HepMC3 OR NOT CELERITAS_USE_Geant4)
   set(_props DISABLED true)
 endif()
 if(CELERITAS_CORE_GEO STREQUAL "Geant4")
-  # FIXME: geant4 navigator doesn't work at scale, nor through accel
+  # FIXME: geant4 navigator doesn't work with field navigator
   set(_props DISABLED true)
 endif()
 if(CELERITAS_REAL_TYPE STREQUAL "float")
diff --git a/src/accel/LocalTransporter.cc b/src/accel/LocalTransporter.cc
index a6020ccd55..eba2b82747 100644
--- a/src/accel/LocalTransporter.cc
+++ b/src/accel/LocalTransporter.cc
@@ -131,17 +131,6 @@ LocalTransporter::LocalTransporter(SetupOptions const& options,
         }
     }
 
-    if (CELERITAS_CORE_GEO == CELERITAS_CORE_GEO_GEANT4)
-    {
-        /*!
-         * \todo Add support for Geant4 navigation wrapper, which requires
-         * calling \c state.ref().geometry.reset() on the local transporter
-         * thread due to thread-allocated navigator data.
-         */
-        CELER_NOT_IMPLEMENTED(
-            "offloading when using Celeritas Geant4 navigation wrapper");
-    }
-
     // Create hit processor on the local thread so that it's deallocated when
     // this object is destroyed
     StreamId stream_id{static_cast<size_type>(thread_id)};
@@ -168,6 +157,8 @@ LocalTransporter::LocalTransporter(SetupOptions const& options,
 
     // Save state for reductions at the end
     params.set_state(stream_id.get(), step_->sp_state());
+
+    CELER_ENSURE(*this);
 }
 
 //---------------------------------------------------------------------------//
@@ -311,6 +302,19 @@ void LocalTransporter::Finalize()
                    << "offloaded tracks (" << buffer_.size()
                    << " in buffer) were not flushed");
 
+    if constexpr (CELERITAS_CORE_GEO == CELERITAS_CORE_GEO_GEANT4)
+    {
+        // Geant4 navigation states *MUST* be deallocated on the thread in
+        // which they're allocated
+        auto state = std::dynamic_pointer_cast<CoreState<MemSpace::host>>(
+            step_->sp_state());
+        CELER_ASSERT(state);
+#if CELERITAS_CORE_GEO == CELERITAS_CORE_GEO_GEANT4
+        CELER_LOG_LOCAL(debug) << "Deallocating navigation states";
+        state->ref().geometry.reset();
+#endif
+    }
+
     // Reset all data
     CELER_LOG_LOCAL(debug) << "Resetting local transporter";
     *this = {};
diff --git a/src/accel/detail/HitManager.cc b/src/accel/detail/HitManager.cc
index 4d43a05cd8..ca56534b09 100644
--- a/src/accel/detail/HitManager.cc
+++ b/src/accel/detail/HitManager.cc
@@ -83,8 +83,10 @@ HitManager::HitManager(SPConstGeo geo,
         }
     }
 
-    // Hit processors *must* be allocated on the thread they're used because of
-    // geant4 thread-local SDs. There must be one per thread.
+    // Hit processors MUST be allocated on the thread they're used because of
+    // geant4 thread-local SDs. They MUST also be DEallocated on the same
+    // thread they're created due to Geant4 thread-local allocators.
+    // There must be one hit processor per thread.
     processor_weakptrs_.resize(num_streams);
     processors_.resize(num_streams);