traccc/utils/array_wrapper.hpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2021 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 #pragma once
0009
0010 #include <type_traits>
0011 #include <utility>
0012
0013 #include <traccc/definitions/qualifiers.hpp>
0014 #include <traccc/utils/functor.hpp>
0015 #include <vecmem/memory/memory_resource.hpp>
0016 #include <vecmem/memory/unique_ptr.hpp>
0017
0018 namespace traccc {
0019 /**
0020  * @brief Tuple of plain-old-data types.
0021  *
0022  * This is just a super-simple re-implementation of `std::tuple` that ensures
0023  * that the type is trivially constructible.
0024  */
0025 template <typename... Ts>
0026 struct pod {};
0027
0028 template <>
0029 struct pod<> {};
0030
0031 template <typename T, typename... Ts>
0032 struct pod<T, Ts...> {
0033   static_assert(std::is_standard_layout_v<T> && std::is_trivial_v<T>,
0034                 "Types in POD-tuple must each be POD.");
0035
0036   T v;
0037   pod<Ts...> r;
0038 };
0039
0040 /**
0041  * @brief Get a given value from a POD tuple.
0042  *
0043  * @tparam I The index in the POD type to retrieve.
0044  * @tparam Ts The types of the POD tuple.
0045  *
0046  * @param p A valid of the POD tuple type described by `Ts...`.
0047  *
0048  * @return A value of type `Ts[I]` as contained in `p`.
0049  */
0050 template <std::size_t I, typename... Ts>
0051 constexpr auto& pod_get(pod<Ts...>& p) {
0052   if constexpr (I == 0) {
0053     return p.v;
0054   } else {
0055     return pod_get<I - 1>(p.r);
0056   }
0057 }
0058
0059 /**
0060  * @brief An array wrapper that can swap between SoA and AoS layouts almost
0061  * transparently.
0062  *
0063  * @tparam F The higher-order layout type, either `soa` or `aos`.
0064  * @tparam T The struct type to embed in memory.
0065  */
0066 template <template <typename...> typename F,
0067           template <template <typename> typename> typename T>
0068 struct array_wrapper {
0069   struct owner {
0070     owner(vecmem::memory_resource& mr, std::size_t n) : data(mr, n) {}
0071
0072     typename details::functor::reapply<
0073         F, typename T<details::functor::identity>::tuple_t>::type::owner data;
0074   };
0075
0076   struct handle {
0077     using handle_t = typename details::functor::reapply<
0078         F, typename T<details::functor::identity>::tuple_t>::type::handle;
0079
0080     handle(const owner& o) : data(o.data) {}
0081
0082     TRACCC_HOST_DEVICE std::size_t size() const { return data.size(); }
0083
0084     template <std::size_t I>
0085     constexpr TRACCC_HOST_DEVICE auto& get(std::size_t i) {
0086       return data.get<I>(i);
0087     }
0088
0089     template <std::size_t I>
0090     constexpr TRACCC_HOST_DEVICE auto get(std::size_t i) const {
0091       return data.get<I>(i);
0092     }
0093
0094     template <std::size_t... Ns>
0095     constexpr TRACCC_HOST_DEVICE T<details::functor::identity>
0096     _construct_helper_identity(std::index_sequence<Ns...>,
0097                                std::size_t i) const {
0098       return T<details::functor::identity>{get<Ns>(i)...};
0099     }
0100
0101     template <std::size_t... Ns>
0102     constexpr TRACCC_HOST_DEVICE T<details::functor::reference>
0103     _construct_helper_reference(std::index_sequence<Ns...>, std::size_t i) {
0104       return T<details::functor::reference>{get<Ns>(i)...};
0105     }
0106
0107     constexpr TRACCC_HOST_DEVICE T<details::functor::identity> operator[](
0108         std::size_t i) const {
0109       return _construct_helper_identity(
0110           std::make_index_sequence<std::tuple_size_v<
0111               typename T<details::functor::identity>::tuple_t>>(),
0112           i);
0113     }
0114
0115     constexpr TRACCC_HOST_DEVICE T<details::functor::reference> operator[](
0116         std::size_t i) {
0117       return _construct_helper_reference(
0118           std::make_index_sequence<std::tuple_size_v<
0119               typename T<details::functor::identity>::tuple_t>>(),
0120           i);
0121     }
0122
0123     handle_t data;
0124   };
0125 };
0126
0127 /**
0128  * @brief Helper function to convert a list of unique pointers to raw pointers.
0129  */
0130 template <std::size_t... Ns, typename... Ts>
0131 std::tuple<Ts*...> _get_ptrs(
0132     std::index_sequence<Ns...>,
0133     const std::tuple<vecmem::unique_alloc_ptr<Ts[]>...>& o) {
0134   return {std::get<Ns>(o).get()...};
0135 }
0136
0137 /**
0138  * @brief A struct-of-arrays (SoA) layout scheme.
0139  *
0140  * Struct-of-arrays is a common approach in vectorized programming both on CPUs
0141  * as well as CPUs. If only parts of the struct are accessed at the same time,
0142  * an SoA layout may increase locality and allow for coalescing of accesses.
0143  *
0144  * @tparam Ts The types contained in the POD struct.
0145  */
0146 template <typename... Ts>
0147 struct soa {
0148   struct owner {
0149     owner(vecmem::memory_resource& mr, std::size_t n)
0150         : _size(n), _ptrs{vecmem::make_unique_alloc<Ts[]>(mr, n)...} {}
0151
0152     constexpr std::size_t size() const { return _size; }
0153
0154     const std::tuple<vecmem::unique_alloc_ptr<Ts[]>...>& pointers() const {
0155       return _ptrs;
0156     }
0157
0158    private:
0159     std::size_t _size;
0160     std::tuple<vecmem::unique_alloc_ptr<Ts[]>...> _ptrs;
0161   };
0162
0163   struct handle {
0164    private:
0165     using tuple_t = std::tuple<Ts*...>;
0166
0167    public:
0168     handle(const owner& o)
0169         : _size(o.size()),
0170           _ptrs(_get_ptrs(std::make_index_sequence<sizeof...(Ts)>(),
0171                           o.pointers())) {}
0172
0173     constexpr TRACCC_HOST_DEVICE std::size_t size() const { return _size; }
0174
0175     template <std::size_t I>
0176     constexpr TRACCC_HOST_DEVICE auto& get(std::size_t i) {
0177       return std::get<I>(_ptrs)[i];
0178     }
0179
0180     template <std::size_t I>
0181     constexpr TRACCC_HOST_DEVICE const auto& get(std::size_t i) const {
0182       return std::get<I>(_ptrs)[i];
0183     }
0184
0185    private:
0186     std::size_t _size;
0187     std::tuple<Ts*...> _ptrs;
0188   };
0189 };
0190
0191 /**
0192  * @brief An array-of-structs (Aos) layout scheme.
0193  *
0194  * Array-of-structs is the default layout scheme in C++. It is useful in cases
0195  * where all elements of a struct need to be accessed simultaneously.
0196  *
0197  * @tparam Ts The types contained in the POD struct.
0198  */
0199 template <typename... Ts>
0200 struct aos {
0201   struct owner {
0202     constexpr owner(vecmem::memory_resource& mr, std::size_t n)
0203         : _size(n), _ptr{vecmem::make_unique_alloc<pod<Ts...>[]>(mr, n)} {}
0204
0205     constexpr std::size_t size() const { return _size; }
0206
0207     constexpr const vecmem::unique_alloc_ptr<pod<Ts...>[]>& pointer() const {
0208       return _ptr;
0209     }
0210
0211    private:
0212     std::size_t _size;
0213     vecmem::unique_alloc_ptr<pod<Ts...>[]> _ptr;
0214   };
0215
0216   struct handle {
0217    private:
0218     using tuple_t = pod<Ts...>;
0219
0220    public:
0221     constexpr handle(const owner& o)
0222         : _size(o.size()), _ptr(o.pointer().get()) {}
0223
0224     constexpr TRACCC_HOST_DEVICE std::size_t size() const { return _size; }
0225
0226     template <std::size_t I>
0227     constexpr TRACCC_HOST_DEVICE auto& get(std::size_t i) {
0228       return pod_get<I>(_ptr[i]);
0229     }
0230
0231     template <std::size_t I>
0232     constexpr TRACCC_HOST_DEVICE const auto& get(std::size_t i) const {
0233       return pod_get<I>(_ptr[i]);
0234     }
0235
0236    private:
0237     std::size_t _size;
0238     tuple_t* _ptr;
0239   };
0240 };
0241 }  // namespace traccc