vg/detail/VecgeomNavCollection.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2021-2024 UT-Battelle, LLC, and other Celeritas developers.
0003 // See the top-level COPYRIGHT file for details.
0004 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0005 //---------------------------------------------------------------------------//
0006 //! \file geocel/vg/detail/VecgeomNavCollection.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <memory>
0011 #include <vector>
0012 #include <VecGeom/base/Config.h>
0013 #include <VecGeom/base/Cuda.h>
0014 #include <VecGeom/navigation/NavStateFwd.h>
0015 #include <VecGeom/navigation/NavStatePool.h>
0016 #include <VecGeom/navigation/NavigationState.h>
0017
0018 #include "corecel/Assert.hh"
0019 #include "corecel/OpaqueId.hh"
0020 #include "corecel/Types.hh"
0021 #include "corecel/cont/Span.hh"
0022 #include "corecel/sys/ThreadId.hh"
0023
0024 namespace celeritas
0025 {
0026 namespace detail
0027 {
0028 //---------------------------------------------------------------------------//
0029 /*!
0030  * Collection-like container for managing VecGeom navigation states.
0031  *
0032  * Since reference and value all behave differently for host and device, we
0033  * only *declare* the class, and provide specializations for each type. The
0034  * specializations are also explicitly declared before their definitions, since
0035  * many of the definitions may not be available depending on which compiler or
0036  * which compilation phase is active.
0037  */
0038 template<Ownership W, MemSpace M>
0039 struct VecgeomNavCollection;
0040
0041 //---------------------------------------------------------------------------//
0042 // HOST MEMSPACE
0043 //---------------------------------------------------------------------------//
0044 /*!
0045  * Manage navigation states in host memory.
0046  */
0047 template<>
0048 struct VecgeomNavCollection<Ownership::value, MemSpace::host>
0049 {
0050     using NavState = vecgeom::cxx::NavigationState;
0051     using UPNavState = std::unique_ptr<NavState>;
0052
0053     std::vector<UPNavState> nav_state;
0054
0055     // Resize with a number of states
0056     void resize(int max_depth, size_type size);
0057     // Whether the collection is assigned
0058     explicit operator bool() const { return !nav_state.empty(); }
0059 };
0060
0061 //---------------------------------------------------------------------------//
0062 /*!
0063  * Reference a host-owned navigation state.
0064  */
0065 template<>
0066 struct VecgeomNavCollection<Ownership::reference, MemSpace::host>
0067 {
0068     using NavState = vecgeom::cxx::NavigationState;
0069     using UPNavState = std::unique_ptr<NavState>;
0070
0071     Span<UPNavState> nav_state;
0072
0073     // Default construction and copy construction
0074     VecgeomNavCollection() = default;
0075     VecgeomNavCollection(VecgeomNavCollection const&) = default;
0076
0077     // Obtain reference from host memory
0078     VecgeomNavCollection&
0079     operator=(VecgeomNavCollection<Ownership::value, MemSpace::host>& other);
0080     // Default copy assignment
0081     VecgeomNavCollection& operator=(VecgeomNavCollection const&) = default;
0082
0083     // Get the navigation state for a given track slot
0084     NavState& at(int, TrackSlotId tid) const;
0085     //! True if the collection is assigned/valiid
0086     explicit operator bool() const { return !nav_state.empty(); }
0087 };
0088
0089 //---------------------------------------------------------------------------//
0090 // DEVICE MEMSPACE
0091 //---------------------------------------------------------------------------//
0092 /*!
0093  * Delete a VecGeom pool.
0094  *
0095  * Due to VecGeom macros, the definition of this function can only be compiled
0096  * from a .cc file.
0097  */
0098 struct NavStatePoolDeleter
0099 {
0100     using arg_type = vecgeom::cxx::NavStatePool*;
0101     void operator()(arg_type) const;
0102 };
0103
0104 //---------------------------------------------------------------------------//
0105 /*!
0106  * Manage a pool of device-side geometry states.
0107  *
0108  * Construction and destruction of the NavStatePool has to be in a host
0109  * compilation unit due to VecGeom macro magic. We hide this class to keep
0110  * NavStatePool and smart pointer usage from the NVCC device compiler.
0111  */
0112 template<>
0113 struct VecgeomNavCollection<Ownership::value, MemSpace::device>
0114 {
0115     using UPNavStatePool
0116         = std::unique_ptr<vecgeom::cxx::NavStatePool, NavStatePoolDeleter>;
0117
0118     UPNavStatePool pool;
0119     void* ptr = nullptr;
0120     int max_depth = 0;
0121     size_type size = 0;
0122
0123     // Resize based on geometry params and state size
0124     void resize(int max_depth, size_type size);
0125     //! True if the collection is assigned/valid
0126     explicit CELER_FUNCTION operator bool() const { return ptr; }
0127 };
0128
0129 //---------------------------------------------------------------------------//
0130 /*!
0131  * Reference on-device memory owned by VecgeomNavCollection<value, device>.
0132  *
0133  * The NavStatePool underpinning the storage returns a void pointer that must
0134  * be manually manipulated to get a single state pointer. The max_depth
0135  * argument must be the same as the GeoParams.
0136  */
0137 template<>
0138 struct VecgeomNavCollection<Ownership::reference, MemSpace::device>
0139 {
0140     using NavState = vecgeom::NavigationState;
0141
0142     vecgeom::NavStatePoolView pool_view = {nullptr, 0, 0};
0143
0144     // Default construct and copy construct
0145     VecgeomNavCollection() = default;
0146     VecgeomNavCollection(VecgeomNavCollection const& other) = default;
0147
0148     // Assign from device value
0149     VecgeomNavCollection&
0150     operator=(VecgeomNavCollection<Ownership::value, MemSpace::device>& other);
0151     // Assign from device reference
0152     VecgeomNavCollection& operator=(VecgeomNavCollection const& other)
0153         = default;
0154
0155     // Get the navigation state for the given track slot
0156     inline CELER_FUNCTION NavState& at(int max_depth, TrackSlotId tid) const;
0157
0158     //! True if the collection is assigned/valid
0159     explicit CELER_FUNCTION operator bool() const
0160     {
0161         return pool_view.IsValid();
0162     }
0163 };
0164
0165 //---------------------------------------------------------------------------//
0166 /*!
0167  * Get the navigation state at the given track slot.
0168  *
0169  * The max_depth_param is used for error checking against the allocated
0170  * max_depth.
0171  */
0172 CELER_FUNCTION auto
0173 VecgeomNavCollection<Ownership::reference, MemSpace::device>::at(
0174     int max_depth_param, TrackSlotId tid) const -> NavState&
0175 {
0176     CELER_EXPECT(this->pool_view.IsValid());
0177     CELER_EXPECT(tid < this->pool_view.Capacity());
0178     CELER_EXPECT(max_depth_param == this->pool_view.Depth());
0179
0180     return *const_cast<NavState*>((this->pool_view)[tid.get()]);
0181 }
0182
0183 //---------------------------------------------------------------------------//
0184 }  // namespace detail
0185 }  // namespace celeritas