AlpakaInterface/interface/SimpleVector.h

0001 #ifndef HeterogeneousCore_AlpakaInterface_interface_SimpleVector_h
0002 #define HeterogeneousCore_AlpakaInterface_interface_SimpleVector_h
0003
0004 //  author: Felice Pantaleo, CERN, 2018
0005 //  alpaka integration in cmssw: Adriano Di Florio, 2022
0006
0007 #include <type_traits>
0008 #include <utility>
0009 #include <alpaka/alpaka.hpp>
0010
0011 namespace cms::alpakatools {
0012
0013   template <class T>
0014   struct SimpleVector {
0015     constexpr SimpleVector() = default;
0016
0017     // ownership of m_data stays within the caller
0018     constexpr void construct(int capacity, T *data) {
0019       m_size = 0;
0020       m_capacity = capacity;
0021       m_data = data;
0022     }
0023
0024     constexpr int push_back_unsafe(const T &element) {
0025       auto previousSize = m_size;
0026       m_size++;
0027       if (previousSize < m_capacity) {
0028         m_data[previousSize] = element;
0029         return previousSize;
0030       } else {
0031         --m_size;
0032         return -1;
0033       }
0034     }
0035
0036     template <class... Ts>
0037     constexpr int emplace_back_unsafe(Ts &&...args) {
0038       auto previousSize = m_size;
0039       m_size++;
0040       if (previousSize < m_capacity) {
0041         (new (&m_data[previousSize]) T(std::forward<Ts>(args)...));
0042         return previousSize;
0043       } else {
0044         --m_size;
0045         return -1;
0046       }
0047     }
0048
0049     ALPAKA_FN_ACC inline T &back() { return m_data[m_size - 1]; }
0050
0051     ALPAKA_FN_ACC inline const T &back() const {
0052       if (m_size > 0) {
0053         return m_data[m_size - 1];
0054       } else
0055         return T();  //undefined behaviour
0056     }
0057
0058     // thread-safe version of the vector, when used in a CUDA kernel
0059     template <typename TAcc>
0060     ALPAKA_FN_ACC int push_back(const TAcc &acc, const T &element) {
0061       auto previousSize = alpaka::atomicAdd(acc, &m_size, 1, alpaka::hierarchy::Blocks{});
0062       if (previousSize < m_capacity) {
0063         m_data[previousSize] = element;
0064         return previousSize;
0065       } else {
0066         alpaka::atomicSub(acc, &m_size, 1, alpaka::hierarchy::Blocks{});
0067         return -1;
0068       }
0069     }
0070
0071     template <typename TAcc, class... Ts>
0072     ALPAKA_FN_ACC int emplace_back(const TAcc &acc, Ts &&...args) {
0073       auto previousSize = alpaka::atomicAdd(acc, &m_size, 1, alpaka::hierarchy::Blocks{});
0074       if (previousSize < m_capacity) {
0075         (new (&m_data[previousSize]) T(std::forward<Ts>(args)...));
0076         return previousSize;
0077       } else {
0078         alpaka::atomicSub(acc, &m_size, 1, alpaka::hierarchy::Blocks{});
0079         return -1;
0080       }
0081     }
0082
0083     // thread safe version of resize
0084     template <typename TAcc>
0085     ALPAKA_FN_ACC int extend(const TAcc &acc, int size = 1) {
0086       auto previousSize = alpaka::atomicAdd(acc, &m_size, size, alpaka::hierarchy::Blocks{});
0087       if (previousSize < m_capacity) {
0088         return previousSize;
0089       } else {
0090         alpaka::atomicSub(acc, &m_size, size, alpaka::hierarchy::Blocks{});
0091         return -1;
0092       }
0093     }
0094
0095     template <typename TAcc>
0096     ALPAKA_FN_ACC int shrink(const TAcc &acc, int size = 1) {
0097       auto previousSize = alpaka::atomicSub(acc, &m_size, size, alpaka::hierarchy::Blocks{});
0098       if (previousSize >= size) {
0099         return previousSize - size;
0100       } else {
0101         alpaka::atomicAdd(acc, &m_size, size, alpaka::hierarchy::Blocks{});
0102         return -1;
0103       }
0104     }
0105
0106     inline constexpr bool empty() const { return m_size <= 0; }
0107     inline constexpr bool full() const { return m_size >= m_capacity; }
0108     inline constexpr T &operator[](int i) { return m_data[i]; }
0109     inline constexpr const T &operator[](int i) const { return m_data[i]; }
0110     inline constexpr void reset() { m_size = 0; }
0111     inline constexpr int size() const { return m_size; }
0112     inline constexpr int capacity() const { return m_capacity; }
0113     inline constexpr T const *data() const { return m_data; }
0114     inline constexpr void resize(int size) { m_size = size; }
0115     inline constexpr void set_data(T *data) { m_data = data; }
0116
0117   private:
0118     int m_size;
0119     int m_capacity;
0120
0121     T *m_data;
0122   };
0123
0124   // ownership of m_data stays within the caller
0125   template <class T>
0126   SimpleVector<T> make_SimpleVector(int capacity, T *data) {
0127     SimpleVector<T> ret;
0128     ret.construct(capacity, data);
0129     return ret;
0130   }
0131
0132   // ownership of m_data stays within the caller
0133   template <class T>
0134   SimpleVector<T> *make_SimpleVector(SimpleVector<T> *mem, int capacity, T *data) {
0135     auto ret = new (mem) SimpleVector<T>();
0136     ret->construct(capacity, data);
0137     return ret;
0138   }
0139
0140 }  // namespace cms::alpakatools
0141
0142 #endif  // AlpakaCore_SimpleVector_h