edm/details/device_container.hpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2021-2022 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 #pragma once
0009
0010 // Project include(s).
0011 #include "traccc/definitions/qualifiers.hpp"
0012 #include "traccc/edm/details/container_base.hpp"
0013
0014 // VecMem include(s).
0015 #include <vecmem/containers/device_vector.hpp>
0016 #include <vecmem/containers/jagged_device_vector.hpp>
0017
0018 // System include(s).
0019 #include <cassert>
0020
0021 namespace traccc {
0022
0023 /// Host container describing objects in a given event
0024 template <typename header_t, typename item_t>
0025 class device_container
0026     : public container_base<header_t, item_t, vecmem::device_vector,
0027                             vecmem::jagged_device_vector> {
0028  public:
0029   /// Base class type
0030   using base_type = container_base<header_t, item_t, vecmem::device_vector,
0031                                    vecmem::jagged_device_vector>;
0032
0033   /// Inherit all of the base class's constructors
0034   using base_type::base_type;
0035
0036   /// Constructor from a "view object"
0037   ///
0038   /// Note that it may also be a "data" or "buffer" object. It just needs to
0039   /// look like a "view object".
0040   ///
0041   template <template <typename, typename> class view_t>
0042   TRACCC_HOST_DEVICE device_container(const view_t<header_t, item_t>& view)
0043       : base_type(view.headers, view.items) {}
0044
0045   /**
0046    * @brief Bounds-checking mutable element accessor.
0047    */
0048   TRACCC_HOST_DEVICE
0049   typename base_type::element_view at(typename base_type::size_type i) {
0050     return {this->m_headers.at(i), this->m_items.at(i)};
0051   }
0052
0053   /**
0054    * @brief Bounds-checking immutable element accessor.
0055    */
0056   TRACCC_HOST_DEVICE
0057   typename base_type::const_element_view at(
0058       typename base_type::size_type i) const {
0059     return {this->m_headers.at(i), this->m_items.at(i)};
0060   }
0061
0062   /**
0063    * @brief Bounds-checking mutable item vector element accessor.
0064    */
0065   TRACCC_HOST_DEVICE
0066   typename base_type::item_type& at(const typename base_type::link_type& link) {
0067     return this->m_items.at(link.first).at(link.second);
0068   }
0069
0070   /**
0071    * @brief Bounds-checking immutable item vector element accessor.
0072    */
0073   TRACCC_HOST_DEVICE
0074   const typename base_type::item_type& at(
0075       const typename base_type::link_type& link) const {
0076     return this->m_items.at(link.first).at(link.second);
0077   }
0078
0079   /**
0080    * @brief Mutable element accessor.
0081    */
0082   TRACCC_HOST_DEVICE
0083   typename base_type::element_view operator[](typename base_type::size_type i) {
0084     return {this->m_headers[i], this->m_items[i]};
0085   }
0086
0087   /**
0088    * @brief Immutable element accessor.
0089    */
0090   TRACCC_HOST_DEVICE
0091   typename base_type::const_element_view operator[](
0092       typename base_type::size_type i) const {
0093     return {this->m_headers[i], this->m_items[i]};
0094   }
0095
0096   /**
0097    * @brief Mutable item vector element accessor.
0098    */
0099   TRACCC_HOST_DEVICE
0100   typename base_type::item_type& operator[](
0101       const typename base_type::link_type& link) {
0102     return this->m_items[link.first][link.second];
0103   }
0104
0105   /**
0106    * @brief Immutable item vector element accessor.
0107    */
0108   TRACCC_HOST_DEVICE
0109   const typename base_type::item_type& operator[](
0110       const typename base_type::link_type& link) const {
0111     return this->m_items[link.first][link.second];
0112   }
0113
0114   /**
0115    * @brief Return the size of the container.
0116    *
0117    * In principle, the size of the two internal vectors should always be
0118    * equal, but we can assert this at runtime for debug builds.
0119    */
0120   TRACCC_HOST_DEVICE
0121   typename base_type::size_type size(void) const {
0122     assert(this->m_headers.size() == this->m_items.size());
0123     return this->m_headers.size();
0124   }
0125
0126   /**
0127    * @breif Get number of items of jagged vector
0128    */
0129   TRACCC_HOST_DEVICE
0130   typename base_type::size_type total_size() const {
0131     typename base_type::size_type ret = 0;
0132     for (const typename base_type::item_vector::value_type& item :
0133          this->m_items) {
0134       ret += item.size();
0135     }
0136     return ret;
0137   }
0138
0139 };  // class device_container
0140
0141 }  // namespace traccc