Add actions and executors for generating optical distribution data (c…

…eleritas-project#1184)

- Add pre-step gather action to cache state data needed for pre-generators
- Add post-step pre-generator action to create and buffer optical distribution data

src/celeritas/CMakeLists.txt

@@ -255,6 +255,8 @@ celeritas_polysource(global/alongstep/AlongStepNeutralAction)
 celeritas_polysource(global/alongstep/AlongStepUniformMscAction)
 celeritas_polysource(global/alongstep/AlongStepRZMapFieldMscAction)
 celeritas_polysource(neutron/model/ChipsNeutronElasticModel)
+celeritas_polysource(optical/detail/PreGenAction)
+celeritas_polysource(optical/detail/PreGenGatherAction)
 celeritas_polysource(phys/detail/DiscreteSelectAction)
 celeritas_polysource(phys/detail/PreStepAction)
 celeritas_polysource(random/RngReseed)

src/celeritas/optical/detail/OpticalGenStorage.hh

@@ -0,0 +1,31 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/optical/detail/OpticalGenStorage.hh
+//---------------------------------------------------------------------------//
+#pragma once
+
+#include <vector>
+
+#include "corecel/data/StreamStore.hh"
+
+#include "../OpticalGenData.hh"
+
+namespace celeritas
+{
+namespace detail
+{
+//---------------------------------------------------------------------------//
+struct OpticalGenStorage
+{
+    using StoreT = StreamStore<OpticalGenParamsData, OpticalGenStateData>;
+
+    StoreT obj;
+    std::vector<OpticalBufferSize> size;
+};
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas

src/celeritas/optical/detail/PreGenAction.cc

@@ -0,0 +1,170 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/optical/detail/PreGenAction.cc
+//---------------------------------------------------------------------------//
+#include "PreGenAction.hh"
+
+#include <algorithm>
+
+#include "corecel/Assert.hh"
+#include "celeritas/global/ActionLauncher.hh"
+#include "celeritas/global/CoreParams.hh"
+#include "celeritas/global/CoreState.hh"
+#include "celeritas/global/CoreTrackData.hh"
+#include "celeritas/global/TrackExecutor.hh"
+#include "celeritas/optical/CerenkovParams.hh"
+#include "celeritas/optical/OpticalPropertyParams.hh"
+#include "celeritas/optical/ScintillationParams.hh"
+
+#include "OpticalGenStorage.hh"
+#include "PreGenExecutor.hh"
+
+namespace celeritas
+{
+namespace detail
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Construct with action ID, optical properties, and storage.
+ */
+PreGenAction::PreGenAction(ActionId id,
+                           SPConstProperties properties,
+                           SPConstCerenkov cerenkov,
+                           SPConstScintillation scintillation,
+                           SPGenStorage storage)
+    : id_(id)
+    , properties_(std::move(properties))
+    , cerenkov_(std::move(cerenkov))
+    , scintillation_(std::move(scintillation))
+    , storage_(std::move(storage))
+{
+    CELER_EXPECT(id_);
+    CELER_EXPECT(scintillation_ || (cerenkov_ && properties_));
+    CELER_EXPECT(storage_);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Descriptive name of the action.
+ */
+std::string PreGenAction::description() const
+{
+    return "generate Cerenkov and scintillation optical distribution data";
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Execute the action with host data.
+ */
+void PreGenAction::execute(CoreParams const& params, CoreStateHost& state) const
+{
+    this->execute_impl(params, state);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Execute the action with device data.
+ */
+void PreGenAction::execute(CoreParams const& params,
+                           CoreStateDevice& state) const
+{
+    this->execute_impl(params, state);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Generate optical distribution data post-step.
+ *
+ * The distributions are stored in separate Cerenkov and scintillation buffers
+ * indexed by the current buffer size plus the track slot ID. The data is
+ * compacted at the end of each step by removing all invalid distributions. The
+ * order of the distributions in the buffers is guaranteed to be reproducible.
+ */
+template<MemSpace M>
+void PreGenAction::execute_impl(CoreParams const& core_params,
+                                CoreState<M>& core_state) const
+{
+    size_type state_size = core_state.size();
+    StreamId stream = core_state.stream_id();
+    auto& buffer_size = storage_->size[stream.get()];
+    auto const& state = storage_->obj.state<M>(stream, state_size);
+
+    CELER_VALIDATE(buffer_size.cerenkov + state_size <= state.cerenkov.size(),
+                   << "insufficient capacity (" << state.cerenkov.size()
+                   << ") for buffered Cerenkov distribution data (total "
+                      "capacity requirement of "
+                   << buffer_size.cerenkov + state_size << ")");
+    CELER_VALIDATE(
+        buffer_size.scintillation + state_size <= state.scintillation.size(),
+        << "insufficient capacity (" << state.scintillation.size()
+        << ") for buffered scintillation distribution data (total "
+           "capacity requirement of "
+        << buffer_size.scintillation + state_size << ")");
+
+    // Generate the optical distribution data
+    this->pre_generate(core_params, core_state);
+
+    // Compact the buffers
+    buffer_size.cerenkov = this->remove_if_invalid(
+        state.cerenkov, buffer_size.cerenkov, state_size, stream);
+    buffer_size.scintillation = this->remove_if_invalid(
+        state.scintillation, buffer_size.scintillation, state_size, stream);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Launch a (host) kernel to generate optical distribution data post-step.
+ */
+void PreGenAction::pre_generate(CoreParams const& core_params,
+                                CoreStateHost& core_state) const
+{
+    TrackExecutor execute{
+        core_params.ptr<MemSpace::native>(),
+        core_state.ptr(),
+        detail::PreGenExecutor{properties_->host_ref(),
+                               cerenkov_->host_ref(),
+                               scintillation_->host_ref(),
+                               storage_->obj.state<MemSpace::native>(
+                                   core_state.stream_id(), core_state.size()),
+                               storage_->size[core_state.stream_id().get()]}};
+    launch_action(*this, core_params, core_state, execute);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Remove all invalid distributions from the buffer.
+ */
+size_type
+PreGenAction::remove_if_invalid(ItemsRef<MemSpace::host> const& buffer,
+                                size_type offset,
+                                size_type size,
+                                StreamId) const
+{
+    auto* start = static_cast<OpticalDistributionData*>(buffer.data());
+    auto* stop
+        = std::remove_if(start + offset, start + offset + size, IsInvalid{});
+    return stop - start;
+}
+
+//---------------------------------------------------------------------------//
+#if !CELER_USE_DEVICE
+void PreGenAction::pre_generate(CoreParams const&, CoreStateDevice&) const
+{
+    CELER_NOT_CONFIGURED("CUDA OR HIP");
+}
+
+size_type PreGenAction::remove_if_invalid(ItemsRef<MemSpace::device> const&,
+                                          size_type,
+                                          size_type,
+                                          StreamId) const
+{
+    CELER_NOT_CONFIGURED("CUDA OR HIP");
+}
+#endif
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas

src/celeritas/optical/detail/PreGenAction.cu

@@ -0,0 +1,74 @@
+//---------------------------------*-CUDA-*----------------------------------//
+// Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/optical/detail/PreGenAction.cu
+//---------------------------------------------------------------------------//
+#include "PreGenAction.hh"
+
+#include <thrust/device_ptr.h>
+#include <thrust/execution_policy.h>
+#include <thrust/remove.h>
+
+#include "corecel/Assert.hh"
+#include "corecel/sys/ScopedProfiling.hh"
+#include "corecel/sys/Thrust.device.hh"
+#include "celeritas/global/ActionLauncher.device.hh"
+#include "celeritas/global/CoreParams.hh"
+#include "celeritas/global/CoreState.hh"
+#include "celeritas/global/TrackExecutor.hh"
+#include "celeritas/optical/CerenkovParams.hh"
+#include "celeritas/optical/OpticalPropertyParams.hh"
+#include "celeritas/optical/ScintillationParams.hh"
+
+#include "OpticalGenStorage.hh"
+#include "PreGenExecutor.hh"
+
+namespace celeritas
+{
+namespace detail
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Launch a kernel to generate optical distribution data post-step.
+ */
+void PreGenAction::pre_generate(CoreParams const& core_params,
+                                CoreStateDevice& core_state) const
+{
+    TrackExecutor execute{
+        core_params.ptr<MemSpace::native>(),
+        core_state.ptr(),
+        detail::PreGenExecutor{properties_->device_ref(),
+                               cerenkov_->device_ref(),
+                               scintillation_->device_ref(),
+                               storage_->obj.state<MemSpace::native>(
+                                   core_state.stream_id(), core_state.size()),
+                               storage_->size[core_state.stream_id().get()]}};
+    static ActionLauncher<decltype(execute)> const launch_kernel(*this);
+    launch_kernel(core_state, execute);
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Remove all invalid distributions from the buffer.
+ */
+size_type
+PreGenAction::remove_if_invalid(ItemsRef<MemSpace::device> const& buffer,
+                                size_type offset,
+                                size_type size,
+                                StreamId stream) const
+{
+    ScopedProfiling profile_this{"remove-if-invalid"};
+    auto start = thrust::device_pointer_cast(buffer.data().get());
+    auto stop = thrust::remove_if(thrust_execute_on(stream),
+                                  start + offset,
+                                  start + offset + size,
+                                  IsInvalid{});
+    CELER_DEVICE_CHECK_ERROR();
+    return stop - start;
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas

src/celeritas/optical/detail/PreGenAction.hh

@@ -0,0 +1,112 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/optical/detail/PreGenAction.hh
+//---------------------------------------------------------------------------//
+#pragma once
+
+#include <memory>
+
+#include "corecel/Macros.hh"
+#include "corecel/data/Collection.hh"
+#include "celeritas/global/ActionInterface.hh"
+#include "celeritas/optical/OpticalDistributionData.hh"
+
+namespace celeritas
+{
+class CerenkovParams;
+class OpticalPropertyParams;
+class ScintillationParams;
+
+namespace detail
+{
+struct OpticalGenStorage;
+//---------------------------------------------------------------------------//
+/*!
+ * Generate optical distribution data.
+ */
+class PreGenAction final : public ExplicitCoreActionInterface
+{
+  public:
+    //!@{
+    //! \name Type aliases
+    using SPConstCerenkov = std::shared_ptr<CerenkovParams const>;
+    using SPConstProperties = std::shared_ptr<OpticalPropertyParams const>;
+    using SPConstScintillation = std::shared_ptr<ScintillationParams const>;
+    using SPGenStorage = std::shared_ptr<detail::OpticalGenStorage>;
+    //!@}
+
+    //! Check if the distribution data is valid
+    struct IsInvalid
+    {
+        CELER_FUNCTION bool
+        operator()(OpticalDistributionData const& data) const
+        {
+            return !data;
+        }
+    };
+
+  public:
+    // Construct with action ID, optical properties, and storage
+    PreGenAction(ActionId id,
+                 SPConstProperties properties,
+                 SPConstCerenkov cerenkov,
+                 SPConstScintillation scintillation,
+                 SPGenStorage storage);
+
+    // Launch kernel with host data
+    void execute(CoreParams const&, CoreStateHost&) const final;
+
+    // Launch kernel with device data
+    void execute(CoreParams const&, CoreStateDevice&) const final;
+
+    //! ID of the model
+    ActionId action_id() const final { return id_; }
+
+    //! Short name for the action
+    std::string label() const final { return "optical-pre-generator-post"; }
+
+    // Name of the action (for user output)
+    std::string description() const final;
+
+    //! Dependency ordering of the action
+    ActionOrder order() const final { return ActionOrder::post_post; }
+
+  private:
+    //// TYPES ////
+
+    template<MemSpace M>
+    using ItemsRef
+        = Collection<OpticalDistributionData, Ownership::reference, M>;
+
+    //// DATA ////
+
+    ActionId id_;
+    SPConstProperties properties_;
+    SPConstCerenkov cerenkov_;
+    SPConstScintillation scintillation_;
+    SPGenStorage storage_;
+
+    //// HELPER FUNCTIONS ////
+
+    template<MemSpace M>
+    void execute_impl(CoreParams const&, CoreState<M>&) const;
+
+    void pre_generate(CoreParams const&, CoreStateHost&) const;
+    void pre_generate(CoreParams const&, CoreStateDevice&) const;
+
+    size_type remove_if_invalid(ItemsRef<MemSpace::host> const&,
+                                size_type,
+                                size_type,
+                                StreamId) const;
+    size_type remove_if_invalid(ItemsRef<MemSpace::device> const&,
+                                size_type,
+                                size_type,
+                                StreamId) const;
+};
+
+//---------------------------------------------------------------------------//
+}  // namespace detail
+}  // namespace celeritas