Refactor track initialization code (#240)

* Use collections for track initializer interface
* Separate track initialization data storage from functionality
* Template TrackInterface on Ownership and MemSpace and remove State/ParamStore
* Address reviewer feedback

app/demo-loop/LDemoInterface.hh

@@ -152,6 +152,7 @@ resize(StateData<Ownership::value, M>*                               data,
     CELER_EXPECT(data);
     CELER_EXPECT(params);
     CELER_EXPECT(size > 0);
+    resize(&data->geometry, params.geometry, size);
     resize(&data->materials, params.materials, size);
     resize(&data->particles, params.particles, size);
     resize(&data->physics, params.physics, size);

src/CMakeLists.txt

@@ -134,9 +134,9 @@ if(CELERITAS_USE_VecGeom)
     geometry/GeoMaterialParams.cc
     geometry/GeoParams.cc
     geometry/detail/VGNavCollection.cc
-    sim/ParamStore.cc
-    sim/StateStore.cc
-    sim/TrackInitializerStore.cc
+    sim/TrackInitInterface.cc
+    sim/TrackInitParams.cc
+    sim/TrackInitUtils.cc
   )
   list(APPEND PRIVATE_DEPS VecGeom::vgdml VecGeom::vecgeom)
   if(CELERITAS_USE_CUDA)

src/random/RngParams.hh

@@ -44,7 +44,7 @@ class RngParams
 };
 
 //---------------------------------------------------------------------------//
-// INLINE DEFINITOINS
+// INLINE DEFINITIONS
 //---------------------------------------------------------------------------//
 /*!
  * Construct with seed.

src/sim/SimInterface.hh

@@ -8,14 +8,12 @@
 #pragma once
 
 #include "base/Collection.hh"
+#include "base/CollectionBuilder.hh"
 #include "base/Macros.hh"
 #include "base/Types.hh"
 #include "detail/SimStateInit.hh"
 #include "Types.hh"
 
-#ifndef __CUDA_ARCH__
-#    include "base/CollectionBuilder.hh"
-#endif
 
 namespace celeritas
 {
@@ -67,7 +65,6 @@ struct SimStateData
     }
 };
 
-#ifndef __CUDA_ARCH__
 //---------------------------------------------------------------------------//
 /*!
  * Resize simulation states and set \c alive to be false.
@@ -79,7 +76,6 @@ void resize(SimStateData<Ownership::value, M>* data, size_type size)
     make_builder(&data->state).resize(size);
     detail::sim_state_init(make_ref(*data));
 }
-#endif
 
 //---------------------------------------------------------------------------//
 } // namespace celeritas

src/sim/TrackInitInterface.cc

@@ -0,0 +1,77 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2021 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file TrackInitInterface.cc
+//---------------------------------------------------------------------------//
+#include "TrackInitInterface.hh"
+
+#include "base/Assert.hh"
+#include "base/CollectionBuilder.hh"
+#include "comm/Device.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Resize and initialize track initializer data on device.
+ */
+void resize(
+    TrackInitStateData<Ownership::value, MemSpace::device>* data,
+    const TrackInitParamsData<Ownership::const_reference, MemSpace::host>& params,
+    size_type size)
+{
+    CELER_EXPECT(params);
+    CELER_EXPECT(size > 0);
+    CELER_EXPECT(celeritas::device());
+
+    // Allocate device data
+    auto capacity = params.storage_factor * size;
+    make_builder(&data->initializers.storage).resize(capacity);
+    make_builder(&data->parents.storage).resize(capacity);
+    make_builder(&data->secondary_counts).resize(size);
+
+    // Start with an empty vector of track initializers and parent thread IDs
+    data->initializers.resize(0);
+    data->parents.resize(0);
+
+    // Initialize vacancies to mark all track slots as empty
+    StateCollection<size_type, Ownership::value, MemSpace::host> vacancies;
+    make_builder(&vacancies).resize(size);
+    for (auto i : range(ThreadId{size}))
+        vacancies[i] = i.get();
+    data->vacancies.storage = vacancies;
+    data->vacancies.resize(size);
+
+    // Initialize the track counter for each event as the number of primary
+    // particles in that event
+    Collection<TrackId::size_type, Ownership::value, MemSpace::host, EventId>
+        track_counters;
+    for (const auto& p : params.primaries[AllItems<Primary, MemSpace::host>{}])
+    {
+        const auto event_id = p.event_id;
+        if (!(event_id < track_counters.size()))
+        {
+            make_builder(&track_counters).resize(event_id.get() + 1);
+        }
+        ++track_counters[event_id];
+    }
+    data->track_counters = track_counters;
+    data->num_primaries  = params.primaries.size();
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Resize and initialize track initializer data on host.
+ */
+void resize(
+    TrackInitStateData<Ownership::value, MemSpace::host>*,
+    const TrackInitParamsData<Ownership::const_reference, MemSpace::host>&,
+    size_type)
+{
+    CELER_NOT_IMPLEMENTED("Host track initializer state");
+}
+
+//---------------------------------------------------------------------------//
+} // namespace celeritas