vg/detail/VecgeomNavCollection.hh

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