Line data Source code
1 : //* This file is part of the MOOSE framework
2 : //* https://www.mooseframework.org
3 : //*
4 : //* All rights reserved, see COPYRIGHT for full restrictions
5 : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
6 : //*
7 : //* Licensed under LGPL 2.1, please see LICENSE for details
8 : //* https://www.gnu.org/licenses/lgpl-2.1.html
9 :
10 : #pragma once
11 :
12 : #include "KokkosArray.h"
13 :
14 : namespace Moose::Kokkos
15 : {
16 :
17 : /**
18 : * A simple object holding the dimension information of an inner array.
19 : * @tparam inner The inner array dimension size
20 : */
21 : template <unsigned int inner>
22 : struct JaggedArrayInnerDim
23 : {
24 : /**
25 : * Size of each dimension
26 : */
27 : unsigned int dim[inner] = {0};
28 : /**
29 : * Stride of each dimension
30 : */
31 : unsigned int stride[inner + 1] = {0};
32 :
33 : #ifdef MOOSE_KOKKOS_SCOPE
34 : /**
35 : * Get the size of a dimension
36 : * @param i The dimension index
37 : * @returns The size of the dimension
38 : */
39 5529 : KOKKOS_FUNCTION unsigned int operator[](unsigned int i) const { return dim[i]; }
40 : #endif
41 : };
42 :
43 : #ifdef MOOSE_KOKKOS_SCOPE
44 : /**
45 : * The inner array wrapper class. This object provides access to an inner array with
46 : * local multi-dimensional indexing and is created as a temporary object by the jagged array class.
47 : * @tparam T The data type
48 : * @tparam inner The inner array dimension size
49 : * @tparam layout The memory layout type
50 : */
51 : template <typename T, unsigned int inner, LayoutType layout>
52 : class JaggedArrayInnerData
53 : {
54 : public:
55 : /**
56 : * Constructor
57 : * @param data The pointer to the inner array data
58 : * @param dim The inner array dimension information
59 : */
60 3799 : KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<inner> dim)
61 2075 : : _data(data), _dim(dim)
62 : {
63 3799 : }
64 :
65 : /**
66 : * Get the total inner array size
67 : * @returns The total inner array size
68 : */
69 4 : KOKKOS_FUNCTION unsigned int size() const { return _dim.stride[inner]; }
70 : /**
71 : * Get the size of a dimension of the inner array
72 : * @param dim The dimension index
73 : * @returns The size of the dimension of the inner array
74 : */
75 3963 : KOKKOS_FUNCTION unsigned int n(unsigned int dim) const { return _dim[dim]; }
76 : /**
77 : * Get an array entry
78 : * @param i The dimensionless inner array index
79 : */
80 3432 : KOKKOS_FUNCTION T & operator[](unsigned int i) const
81 : {
82 : KOKKOS_ASSERT(i < _dim.stride[inner]);
83 :
84 3432 : return _data[i];
85 : }
86 : /**
87 : * Get an array entry
88 : * @param i The inner array index of each dimension
89 : */
90 : template <typename... indices>
91 : KOKKOS_FUNCTION T & operator()(indices... i) const;
92 :
93 : protected:
94 : /**
95 : * Pointer to the inner array data
96 : */
97 : T * _data;
98 : /**
99 : * Inner array dimension information
100 : */
101 : JaggedArrayInnerDim<inner> _dim;
102 : };
103 :
104 : template <typename T, unsigned int inner, LayoutType layout>
105 : template <typename... indices>
106 : KOKKOS_FUNCTION T &
107 1716 : JaggedArrayInnerData<T, inner, layout>::operator()(indices... i) const
108 : {
109 : static_assert((std::is_convertible<indices, unsigned int>::value && ...),
110 : "All arguments must be convertible to unsigned int");
111 : static_assert(sizeof...(i) == inner, "Number of arguments should match array dimension");
112 :
113 : #ifndef NDEBUG
114 : {
115 : unsigned int idx[inner] = {static_cast<unsigned int>(i)...};
116 :
117 : for (unsigned int d = 0; d < sizeof...(i); ++d)
118 : KOKKOS_ASSERT(idx[d] < _dim[d]);
119 : }
120 : #endif
121 :
122 1716 : unsigned int idx = 0;
123 1716 : unsigned int d = 0;
124 :
125 : if constexpr (layout == LayoutType::LEFT)
126 2804 : (((idx +=
127 1144 : (d == 0 ? static_cast<unsigned int>(i) : static_cast<unsigned int>(i) * _dim.stride[d])),
128 : ++d),
129 : ...);
130 : else
131 1402 : (((idx += (d == inner - 1 ? static_cast<unsigned int>(i)
132 830 : : static_cast<unsigned int>(i) * _dim.stride[d])),
133 : ++d),
134 : ...);
135 :
136 1716 : return _data[idx];
137 : }
138 : #endif
139 :
140 : /**
141 : * The Kokkos jagged array class.
142 : * A Kokkos jagged array aids in treating jagged arrays conveniently and efficiently by using a
143 : * sequential data storage and providing multi-dimensional indexing using internal dope vectors.
144 : * A jagged array is divided into the inner and outer arrays. The outer array is the regular part of
145 : * a jagged array. Each entry of the outer array can hold an inner array, whose size can vary with
146 : * each other.
147 : * Calling create() will allocate the outer array, and inner arrays should be reserved one-by-one
148 : * through reserve(). Once the array structure is set, finalize() should be called. A finalized
149 : * array cannot be restructured unless it is reset completely by calling create() again.
150 : * The inner array returned by operator() or operator[] using the outer array index is held in a
151 : * temporary wrapper object. To avoid the overhead of temporary object creation, it is recommended
152 : * to store the object locally.
153 : * @tparam T The data type
154 : * @tparam inner The inner array dimension size
155 : * @tparam outer The outer array dimension size
156 : * @tparam index_type The array index type
157 : * @tparam layout The memory layout type
158 : */
159 : template <typename T,
160 : unsigned int inner,
161 : unsigned int outer,
162 : typename index_type = dof_id_type,
163 : LayoutType layout = LayoutType::LEFT>
164 : class JaggedArray
165 : {
166 : public:
167 : /**
168 : * Default constructor
169 : */
170 : JaggedArray() = default;
171 :
172 : #ifdef MOOSE_KOKKOS_SCOPE
173 : /**
174 : * Constructor
175 : */
176 : template <typename... size_type>
177 56 : JaggedArray(size_type... n)
178 56 : {
179 56 : create(n...);
180 56 : }
181 :
182 : /**
183 : * Allocate outer array
184 : * @param n The vector containing the size of each dimension for the outer array
185 : */
186 : void create(const std::vector<index_type> & n);
187 : /**
188 : * Allocate outer array
189 : * @param n The size of each dimension for the outer array
190 : */
191 : template <typename... size_type>
192 : void create(size_type... n);
193 : /**
194 : * Reserve inner array for an outer array entry
195 : * @param index The array containing the index for the outer array
196 : * @param dimension The array containing the size of each dimension for the inner array
197 : */
198 : void reserve(const std::array<uint64_t, outer> & index,
199 : const std::array<uint64_t, inner> & dimension);
200 : /**
201 : * Setup array structure
202 : */
203 : void finalize();
204 : /**
205 : * Copy data from host to device
206 : */
207 : void copyToDevice()
208 : {
209 : mooseAssert(_finalized, "KokkosJaggedArray not finalized.");
210 :
211 : _data.copyToDevice();
212 : }
213 : /**
214 : * Copy data from device to host
215 : */
216 54 : void copyToHost()
217 : {
218 : mooseAssert(_finalized, "KokkosJaggedArray not finalized.");
219 :
220 54 : _data.copyToHost();
221 54 : }
222 : /**
223 : * Copy data from host to device and deallocate host
224 : */
225 : void moveToDevice()
226 : {
227 : mooseAssert(_finalized, "KokkosJaggedArray not finalized.");
228 :
229 : _dims.moveToDevice();
230 : _offsets.moveToDevice();
231 : _data.moveToDevice();
232 : }
233 : /**
234 : * Copy data from device to host and deallocate device
235 : */
236 : void moveToHost()
237 : {
238 : mooseAssert(_finalized, "KokkosJaggedArray not finalized.");
239 :
240 : _dims.moveToHost();
241 : _offsets.moveToHost();
242 : _data.moveToHost();
243 : }
244 : /**
245 : * Get the underlying data array
246 : * @returns The data array
247 : */
248 : auto & array() { return _data; }
249 :
250 : /**
251 : * Get whether the array is finalized
252 : * @returns Whether the array is finalized
253 : */
254 : KOKKOS_FUNCTION bool isFinalized() const { return _finalized; }
255 : /**
256 : * Get whether the array was allocated on host
257 : * @returns Whether the array was allocated on host
258 : */
259 : KOKKOS_FUNCTION bool isHostAlloc() const { return _data.isHostAlloc(); }
260 : /**
261 : * Get whether the array was allocated on device
262 : * @returns Whether the array was allocated on device
263 : */
264 : KOKKOS_FUNCTION bool isDeviceAlloc() const { return _data.isDeviceAlloc(); }
265 : /**
266 : * Get the total data array size
267 : * @returns The total data array size
268 : */
269 : KOKKOS_FUNCTION index_type size() const { return _data.size(); }
270 : /**
271 : * Get the total outer array size
272 : * @returns The total outer array size
273 : */
274 : KOKKOS_FUNCTION index_type n() const { return _offsets.size(); }
275 : /**
276 : * Get the size of a dimension of the outer array
277 : * @param dim The dimension index
278 : * @returns The size of the dimension of the outer array
279 : */
280 : KOKKOS_FUNCTION index_type n(unsigned int dim) const { return _offsets.n(dim); }
281 : /**
282 : * Get an inner array
283 : * @param i The dimensionless outer array index
284 : * @returns The inner array wrapper object
285 : */
286 : KOKKOS_FUNCTION auto operator[](index_type i) const;
287 : /**
288 : * Get an inner array
289 : * @param i The index of each dimension
290 : * @returns The inner array wrapper object
291 : */
292 : template <typename... indices>
293 : KOKKOS_FUNCTION auto operator()(indices... i) const;
294 : #endif
295 :
296 : protected:
297 : /**
298 : * Sequential data array
299 : */
300 : Array1D<T, index_type> _data;
301 : /**
302 : * Dimension information of each inner array
303 : */
304 : Array<JaggedArrayInnerDim<inner>, outer, index_type> _dims;
305 : /**
306 : * Starting offset of each inner array into the sequential data array
307 : */
308 : Array<index_type, outer, index_type> _offsets;
309 : /**
310 : * Whether the array was finalized
311 : */
312 : bool _finalized = false;
313 : };
314 :
315 : #ifdef MOOSE_KOKKOS_SCOPE
316 : template <typename T,
317 : unsigned int inner,
318 : unsigned int outer,
319 : typename index_type,
320 : LayoutType layout>
321 : void
322 : JaggedArray<T, inner, outer, index_type, layout>::create(const std::vector<index_type> & n)
323 : {
324 : _data.destroy();
325 : _offsets.create(n);
326 : _dims.create(n);
327 :
328 : _finalized = false;
329 : }
330 :
331 : template <typename T,
332 : unsigned int inner,
333 : unsigned int outer,
334 : typename index_type,
335 : LayoutType layout>
336 : template <typename... size_type>
337 : void
338 56 : JaggedArray<T, inner, outer, index_type, layout>::create(size_type... n)
339 : {
340 56 : _data.destroy();
341 56 : _offsets.create(n...);
342 56 : _dims.create(n...);
343 :
344 56 : _finalized = false;
345 56 : }
346 :
347 : template <typename T,
348 : unsigned int inner,
349 : unsigned int outer,
350 : typename index_type,
351 : LayoutType layout>
352 : void
353 710 : JaggedArray<T, inner, outer, index_type, layout>::reserve(
354 : const std::array<uint64_t, outer> & index, const std::array<uint64_t, inner> & dimension)
355 : {
356 : mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");
357 :
358 710 : index_type idx = 0;
359 710 : index_type stride = 1;
360 :
361 2570 : for (unsigned int o = 0; o < outer; ++o)
362 : {
363 1860 : idx += index[o] * stride;
364 1860 : stride *= _offsets.n(o);
365 : }
366 :
367 2138 : for (unsigned int i = 0; i < inner; ++i)
368 1428 : _dims[idx].dim[i] = dimension[i];
369 :
370 710 : stride = 1;
371 :
372 : if constexpr (layout == LayoutType::LEFT)
373 1436 : for (unsigned int i = 0; i < inner; ++i)
374 : {
375 960 : _dims[idx].stride[i] = stride;
376 960 : stride *= dimension[i];
377 : }
378 : else
379 702 : for (int i = inner - 1; i >= 0; --i)
380 : {
381 468 : _dims[idx].stride[i] = stride;
382 468 : stride *= dimension[i];
383 : }
384 :
385 710 : _dims[idx].stride[inner] = stride;
386 710 : }
387 :
388 : template <typename T,
389 : unsigned int inner,
390 : unsigned int outer,
391 : typename index_type,
392 : LayoutType layout>
393 : void
394 56 : JaggedArray<T, inner, outer, index_type, layout>::finalize()
395 : {
396 : mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");
397 :
398 56 : index_type stride = 1;
399 :
400 812 : for (index_type o = 0; o < _offsets.size(); ++o)
401 : {
402 756 : stride = 1;
403 :
404 2322 : for (unsigned int i = 0; i < inner; ++i)
405 1566 : stride *= _dims[o][i];
406 :
407 756 : _offsets[o] = stride;
408 : }
409 :
410 56 : std::exclusive_scan(_offsets.begin(), _offsets.end(), _offsets.begin(), 0);
411 :
412 56 : _dims.copyToDevice();
413 56 : _offsets.copyToDevice();
414 :
415 : // Pad an extra element at the end to avoid accessing the bound in the following operators when
416 : // the last inner array has zero size
417 56 : _data.create(_offsets.last() + stride + 1);
418 :
419 56 : _finalized = true;
420 56 : }
421 :
422 : template <typename T,
423 : unsigned int inner,
424 : unsigned int outer,
425 : typename index_type,
426 : LayoutType layout>
427 : KOKKOS_FUNCTION auto
428 : JaggedArray<T, inner, outer, index_type, layout>::operator[](index_type i) const
429 : {
430 : auto data = &_data[_offsets[i]];
431 : const auto & dim = _dims[i];
432 :
433 : return JaggedArrayInnerData<T, inner, layout>(data, dim);
434 : }
435 :
436 : template <typename T,
437 : unsigned int inner,
438 : unsigned int outer,
439 : typename index_type,
440 : LayoutType layout>
441 : template <typename... indices>
442 : KOKKOS_FUNCTION auto
443 3799 : JaggedArray<T, inner, outer, index_type, layout>::operator()(indices... i) const
444 : {
445 3799 : auto data = &_data[_offsets(i...)];
446 3799 : const auto & dim = _dims(i...);
447 :
448 3799 : return JaggedArrayInnerData<T, inner, layout>(data, dim);
449 : }
450 : #endif
451 :
452 : } // namespace Moose::Kokkos
|
