CUDAUtilities/interface/SimpleVector.h

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