Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_Vector.hpp
1//@HEADER
2// ************************************************************************
3//
4// Kokkos v. 4.0
5// Copyright (2022) National Technology & Engineering
6// Solutions of Sandia, LLC (NTESS).
7//
8// Under the terms of Contract DE-NA0003525 with NTESS,
9// the U.S. Government retains certain rights in this software.
10//
11// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12// See https://kokkos.org/LICENSE for license information.
13// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14//
15//@HEADER
16
17#ifndef KOKKOS_VECTOR_HPP
18#define KOKKOS_VECTOR_HPP
19#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
20#define KOKKOS_IMPL_PUBLIC_INCLUDE
21#define KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_VECTOR
22#endif
23
24#include <Kokkos_Core_fwd.hpp>
25#include <Kokkos_DualView.hpp>
26
27/* Drop in replacement for std::vector based on Kokkos::DualView
28 * Most functions only work on the host (it will not compile if called from
29 * device kernel)
30 *
31 */
32namespace Kokkos {
33
34template <class Scalar, class Arg1Type = void>
35class vector : public DualView<Scalar*, LayoutLeft, Arg1Type> {
36 public:
37 using value_type = Scalar;
38 using pointer = Scalar*;
39 using const_pointer = const Scalar*;
40 using reference = Scalar&;
41 using const_reference = const Scalar&;
42 using iterator = Scalar*;
43 using const_iterator = const Scalar*;
44 using size_type = size_t;
45
46 private:
47 size_t _size;
48 float _extra_storage;
49 using DV = DualView<Scalar*, LayoutLeft, Arg1Type>;
50
51 public:
52#ifdef KOKKOS_ENABLE_CUDA_UVM
53 KOKKOS_INLINE_FUNCTION reference operator()(int i) const {
54 return DV::h_view(i);
55 };
56 KOKKOS_INLINE_FUNCTION reference operator[](int i) const {
57 return DV::h_view(i);
58 };
59#else
60 inline reference operator()(int i) const { return DV::h_view(i); };
61 inline reference operator[](int i) const { return DV::h_view(i); };
62#endif
63
64 /* Member functions which behave like std::vector functions */
65
66 vector() : DV() {
67 _size = 0;
68 _extra_storage = 1.1;
69 }
70
71 vector(int n, Scalar val = Scalar())
72 : DualView<Scalar*, LayoutLeft, Arg1Type>("Vector", size_t(n * (1.1))) {
73 _size = n;
74 _extra_storage = 1.1;
75 DV::modified_flags(0) = 1;
76
77 assign(n, val);
78 }
79
80 void resize(size_t n) {
81 if (n >= span()) DV::resize(size_t(n * _extra_storage));
82 _size = n;
83 }
84
85 void resize(size_t n, const Scalar& val) { assign(n, val); }
86
87 void assign(size_t n, const Scalar& val) {
88 /* Resize if necessary (behavior of std:vector) */
89
90 if (n > span()) DV::resize(size_t(n * _extra_storage));
91 _size = n;
92
93 /* Assign value either on host or on device */
94
95 if (DV::template need_sync<typename DV::t_dev::device_type>()) {
96 set_functor_host f(DV::h_view, val);
97 parallel_for("Kokkos::vector::assign", n, f);
98 typename DV::t_host::execution_space().fence(
99 "Kokkos::vector::assign: fence after assigning values");
100 DV::template modify<typename DV::t_host::device_type>();
101 } else {
102 set_functor f(DV::d_view, val);
103 parallel_for("Kokkos::vector::assign", n, f);
104 typename DV::t_dev::execution_space().fence(
105 "Kokkos::vector::assign: fence after assigning values");
106 DV::template modify<typename DV::t_dev::device_type>();
107 }
108 }
109
110 void reserve(size_t n) { DV::resize(size_t(n * _extra_storage)); }
111
112 void push_back(Scalar val) {
113 DV::template sync<typename DV::t_host::device_type>();
114 DV::template modify<typename DV::t_host::device_type>();
115 if (_size == span()) {
116 size_t new_size = _size * _extra_storage;
117 if (new_size == _size) new_size++;
118 DV::resize(new_size);
119 }
120
121 DV::h_view(_size) = val;
122 _size++;
123 }
124
125 void pop_back() { _size--; }
126
127 void clear() { _size = 0; }
128
129 iterator insert(iterator it, const value_type& val) {
130 return insert(it, 1, val);
131 }
132
133 iterator insert(iterator it, size_type count, const value_type& val) {
134 if ((size() == 0) && (it == begin())) {
135 resize(count, val);
136 DV::sync_host();
137 return begin();
138 }
139 DV::sync_host();
140 DV::modify_host();
141 if (std::less<>()(it, begin()) || std::less<>()(end(), it))
142 Kokkos::abort("Kokkos::vector::insert : invalid insert iterator");
143 if (count == 0) return it;
144 ptrdiff_t start = std::distance(begin(), it);
145 auto org_size = size();
146 resize(size() + count);
147
148 std::copy_backward(begin() + start, begin() + org_size,
149 begin() + org_size + count);
150 std::fill_n(begin() + start, count, val);
151
152 return begin() + start;
153 }
154
155 private:
156 template <class T>
157 struct impl_is_input_iterator
158 : /* TODO replace this */ std::bool_constant<
159 !std::is_convertible<T, size_type>::value> {};
160
161 public:
162 // TODO: can use detection idiom to generate better error message here later
163 template <typename InputIterator>
164 std::enable_if_t<impl_is_input_iterator<InputIterator>::value, iterator>
165 insert(iterator it, InputIterator b, InputIterator e) {
166 ptrdiff_t count = std::distance(b, e);
167
168 DV::sync_host();
169 DV::modify_host();
170 if (std::less<>()(it, begin()) || std::less<>()(end(), it))
171 Kokkos::abort("Kokkos::vector::insert : invalid insert iterator");
172
173 ptrdiff_t start = std::distance(begin(), it);
174 auto org_size = size();
175
176 // Note: resize(...) invalidates it; use begin() + start instead
177 resize(size() + count);
178
179 std::copy_backward(begin() + start, begin() + org_size,
180 begin() + org_size + count);
181 std::copy(b, e, begin() + start);
182
183 return begin() + start;
184 }
185
186 KOKKOS_INLINE_FUNCTION constexpr bool is_allocated() const {
187 return DV::is_allocated();
188 }
189
190 size_type size() const { return _size; }
191 size_type max_size() const { return 2000000000; }
192 size_type span() const { return DV::span(); }
193 bool empty() const { return _size == 0; }
194
195 pointer data() const { return DV::h_view.data(); }
196
197 iterator begin() const { return DV::h_view.data(); }
198
199 iterator end() const {
200 return _size > 0 ? DV::h_view.data() + _size : DV::h_view.data();
201 }
202
203 reference front() { return DV::h_view(0); }
204
205 reference back() { return DV::h_view(_size - 1); }
206
207 const_reference front() const { return DV::h_view(0); }
208
209 const_reference back() const { return DV::h_view(_size - 1); }
210
211 /* std::algorithms which work originally with iterators, here they are
212 * implemented as member functions */
213
214 size_t lower_bound(const size_t& start, const size_t& theEnd,
215 const Scalar& comp_val) const {
216 int lower = start; // FIXME (mfh 24 Apr 2014) narrowing conversion
217 int upper =
218 _size > theEnd
219 ? theEnd
220 : _size - 1; // FIXME (mfh 24 Apr 2014) narrowing conversion
221 if (upper <= lower) {
222 return theEnd;
223 }
224
225 Scalar lower_val = DV::h_view(lower);
226 Scalar upper_val = DV::h_view(upper);
227 size_t idx = (upper + lower) / 2;
228 Scalar val = DV::h_view(idx);
229 if (val > upper_val) return upper;
230 if (val < lower_val) return start;
231
232 while (upper > lower) {
233 if (comp_val > val) {
234 lower = ++idx;
235 } else {
236 upper = idx;
237 }
238 idx = (upper + lower) / 2;
239 val = DV::h_view(idx);
240 }
241 return idx;
242 }
243
244 bool is_sorted() {
245 for (int i = 0; i < _size - 1; i++) {
246 if (DV::h_view(i) > DV::h_view(i + 1)) return false;
247 }
248 return true;
249 }
250
251 iterator find(Scalar val) const {
252 if (_size == 0) return end();
253
254 int upper, lower, current;
255 current = _size / 2;
256 upper = _size - 1;
257 lower = 0;
258
259 if ((val < DV::h_view(0)) || (val > DV::h_view(_size - 1))) return end();
260
261 while (upper > lower) {
262 if (val > DV::h_view(current))
263 lower = current + 1;
264 else
265 upper = current;
266 current = (upper + lower) / 2;
267 }
268
269 if (val == DV::h_view(current))
270 return &DV::h_view(current);
271 else
272 return end();
273 }
274
275 /* Additional functions for data management */
276
277 void device_to_host() { deep_copy(DV::h_view, DV::d_view); }
278 void host_to_device() const { deep_copy(DV::d_view, DV::h_view); }
279
280 void on_host() { DV::template modify<typename DV::t_host::device_type>(); }
281 void on_device() { DV::template modify<typename DV::t_dev::device_type>(); }
282
283 void set_overallocation(float extra) { _extra_storage = 1.0 + extra; }
284
285 public:
286 struct set_functor {
287 using execution_space = typename DV::t_dev::execution_space;
288 typename DV::t_dev _data;
289 Scalar _val;
290
291 set_functor(typename DV::t_dev data, Scalar val) : _data(data), _val(val) {}
292
293 KOKKOS_INLINE_FUNCTION
294 void operator()(const int& i) const { _data(i) = _val; }
295 };
296
297 struct set_functor_host {
298 using execution_space = typename DV::t_host::execution_space;
299 typename DV::t_host _data;
300 Scalar _val;
301
302 set_functor_host(typename DV::t_host data, Scalar val)
303 : _data(data), _val(val) {}
304
305 KOKKOS_INLINE_FUNCTION
306 void operator()(const int& i) const { _data(i) = _val; }
307 };
308};
309
310} // namespace Kokkos
311#ifdef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_VECTOR
312#undef KOKKOS_IMPL_PUBLIC_INCLUDE
313#undef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_VECTOR
314#endif
315#endif
Declaration and definition of Kokkos::DualView.