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
15 : {
16 : namespace Kokkos
17 : {
18 :
19 : /**
20 : * A simple object holding the dimension information of an inner array.
21 : * @tparam inner The inner array dimension size
22 : */
23 : template <unsigned int inner>
24 : struct JaggedArrayInnerDim
25 : {
26 : /**
27 : * Size of each dimension
28 : */
29 : dof_id_type dim[inner] = {0};
30 : /**
31 : * Stride of each dimension
32 : */
33 : dof_id_type stride[inner] = {0};
34 :
35 : #ifdef MOOSE_KOKKOS_SCOPE
36 : /**
37 : * Get the size of a dimension
38 : * @param i The dimension index
39 : * @returns The size of the dimension
40 : */
41 1959 : KOKKOS_FUNCTION dof_id_type operator[](unsigned int i) const { return dim[i]; }
42 : #endif
43 : };
44 :
45 : #ifdef MOOSE_KOKKOS_SCOPE
46 : /**
47 : * The base class of the inner array wrapper. This object provides access to an inner array with
48 : * local multi-dimensional indexing and is created as a temporary object by the jagged array class.
49 : * @tparam T The data type
50 : * @tparam inner The inner array dimension size
51 : */
52 : template <typename T, unsigned int inner>
53 : class JaggedArrayInnerDataBase
54 : {
55 : public:
56 : /**
57 : * Constructor
58 : * @param data The pointer to the inner array data
59 : * @param dim The inner array dimension information
60 : */
61 1277 : KOKKOS_FUNCTION JaggedArrayInnerDataBase(T * data, JaggedArrayInnerDim<inner> dim)
62 697 : : _data(data), _dim(dim)
63 : {
64 1277 : }
65 :
66 : /**
67 : * Get the total inner array size
68 : * @returns The total inner array size
69 : */
70 4 : KOKKOS_FUNCTION dof_id_type size() const { return _dim.stride[inner - 1]; }
71 : /**
72 : * Get the size of a dimension of the inner array
73 : * @param dim The dimension index
74 : * @returns The size of the dimension of the inner array
75 : */
76 1329 : KOKKOS_FUNCTION dof_id_type n(unsigned int dim) const { return _dim[dim]; }
77 : /**
78 : * Get an array entry
79 : * @param i The dimensionless inner array index
80 : */
81 : KOKKOS_FUNCTION T & operator[](dof_id_type i) const
82 : {
83 : KOKKOS_ASSERT(i < _dim.stride[inner - 1]);
84 :
85 : return _data[i];
86 : }
87 :
88 : protected:
89 : /**
90 : * Pointer to the inner array data
91 : */
92 : T * _data;
93 : /**
94 : * Inner array dimension information
95 : */
96 : JaggedArrayInnerDim<inner> _dim;
97 : };
98 :
99 : #define usingKokkosJaggedArrayInnerDataBaseMembers(T, inner) \
100 : private: \
101 : using JaggedArrayInnerDataBase<T, inner>::_data; \
102 : using JaggedArrayInnerDataBase<T, inner>::_dim
103 :
104 : /**
105 : * The specialization of the inner array wrapper class for each dimension.
106 : */
107 : ///{@
108 : template <typename T, unsigned int inner>
109 : class JaggedArrayInnerData;
110 :
111 : template <typename T>
112 : class JaggedArrayInnerData<T, 1> : public JaggedArrayInnerDataBase<T, 1>
113 : {
114 : usingKokkosJaggedArrayInnerDataBaseMembers(T, 1);
115 :
116 : public:
117 : /**
118 : * Constructor
119 : * @param data The pointer to the inner array data
120 : * @param dim The inner array dimension information
121 : */
122 195 : KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<1> dim)
123 117 : : JaggedArrayInnerDataBase<T, 1>(data, dim)
124 : {
125 195 : }
126 :
127 : /**
128 : * Get an array entry
129 : * @param i0 The first dimension inner array index
130 : */
131 234 : KOKKOS_FUNCTION T & operator()(dof_id_type i0) const
132 : {
133 : KOKKOS_ASSERT(i0 < _dim[0]);
134 :
135 234 : return _data[i0];
136 : }
137 : };
138 :
139 : template <typename T>
140 : class JaggedArrayInnerData<T, 2> : public JaggedArrayInnerDataBase<T, 2>
141 : {
142 : usingKokkosJaggedArrayInnerDataBaseMembers(T, 2);
143 :
144 : public:
145 : /**
146 : * Constructor
147 : * @param data The pointer to the inner array data
148 : * @param dim The inner array dimension information
149 : */
150 355 : KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<2> dim)
151 197 : : JaggedArrayInnerDataBase<T, 2>(data, dim)
152 : {
153 355 : }
154 :
155 : /**
156 : * Get an array entry
157 : * @param i0 The first dimension inner array index
158 : * @param i1 The second dimension inner array index
159 : */
160 474 : KOKKOS_FUNCTION T & operator()(dof_id_type i0, dof_id_type i1) const
161 : {
162 : KOKKOS_ASSERT(i0 < _dim[0]);
163 : KOKKOS_ASSERT(i1 < _dim[1]);
164 :
165 474 : return _data[i0 + _dim.stride[0] * i1];
166 : }
167 : };
168 :
169 : template <typename T>
170 : class JaggedArrayInnerData<T, 3> : public JaggedArrayInnerDataBase<T, 3>
171 : {
172 : usingKokkosJaggedArrayInnerDataBaseMembers(T, 3);
173 :
174 : public:
175 : /**
176 : * Constructor
177 : * @param data The pointer to the inner array data
178 : * @param dim The inner array dimension information
179 : */
180 727 : KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<3> dim)
181 383 : : JaggedArrayInnerDataBase<T, 3>(data, dim)
182 : {
183 727 : }
184 :
185 : /**
186 : * Get an array entry
187 : * @param i0 The first dimension inner array index
188 : * @param i1 The second dimension inner array index
189 : * @param i2 The third dimension inner array index
190 : */
191 1008 : KOKKOS_FUNCTION T & operator()(dof_id_type i0, dof_id_type i1, dof_id_type i2) const
192 : {
193 : KOKKOS_ASSERT(i0 < _dim[0]);
194 : KOKKOS_ASSERT(i1 < _dim[1]);
195 : KOKKOS_ASSERT(i2 < _dim[2]);
196 :
197 1008 : return _data[i0 + _dim.stride[0] * i1 + _dim.stride[1] * i2];
198 : }
199 : };
200 : ///@}
201 : #endif
202 :
203 : /**
204 : * The base class for Kokkos jagged arrays.
205 : * A Kokkos jagged array aids in treating jagged arrays conveniently and efficiently by using a
206 : * sequential data storage and providing multi-dimensional indexing using internal dope vectors.
207 : * A jagged array is divided into the inner and outer arrays. The outer array is the regular part of
208 : * a jagged array. Each entry of the outer array can hold an inner array, whose size can vary with
209 : * each other.
210 : * Calling create() will allocate the outer array, and inner arrays should be reserved one-by-one
211 : * through reserve(). Once the array structure is set, finalize() should be called. A finalized
212 : * array cannot be restructured unless it is reset completely by calling create() again.
213 : * The inner array returned by operator() or operator[] using the outer array index is held in a
214 : * temporary wrapper object. To avoid the overhead of temporary object creation, it is recommended
215 : * to store the object locally.
216 : * @tparam T The data type
217 : * @tparam inner The inner array dimension size
218 : * @tparam outer The outer array dimension size
219 : */
220 : template <typename T, unsigned int inner, unsigned int outer>
221 : class JaggedArrayBase
222 : {
223 : #ifdef MOOSE_KOKKOS_SCOPE
224 : public:
225 : /**
226 : * Allocate outer array
227 : * @param n The vector containing the size of each dimension for the outer array
228 : */
229 : void create(const std::vector<dof_id_type> & n);
230 : /**
231 : * Reserve inner array for an outer array entry
232 : * @param index The array containing the index for the outer array
233 : * @param dimension The array containing the size of each dimension for the inner array
234 : */
235 : void reserve(const std::array<dof_id_type, outer> & index,
236 : const std::array<dof_id_type, inner> & dimension);
237 : /**
238 : * Setup array structure
239 : */
240 : void finalize();
241 : /**
242 : * Copy data from host to device
243 : */
244 : void copyToDevice()
245 : {
246 : mooseAssert(_finalized, "KokkosJaggedArray not finalized.");
247 :
248 : _data.copyToDevice();
249 : }
250 : /**
251 : * Copy data from device to host
252 : */
253 18 : void copyToHost()
254 : {
255 : mooseAssert(_finalized, "KokkosJaggedArray not finalized.");
256 :
257 18 : _data.copyToHost();
258 18 : }
259 :
260 : /**
261 : * Get whether the array is finalized
262 : * @returns Whether the array is finalized
263 : */
264 : KOKKOS_FUNCTION bool isFinalized() const { return _finalized; }
265 : /**
266 : * Get the total data array size
267 : * @returns The total data array size
268 : */
269 : KOKKOS_FUNCTION dof_id_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 dof_id_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 dof_id_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[](dof_id_type i) const
287 : {
288 : auto data = &_data[_offsets[i]];
289 : const auto & dim = _dims[i];
290 :
291 : return JaggedArrayInnerData<T, inner>(data, dim);
292 : }
293 : #endif
294 :
295 : protected:
296 : /**
297 : * Sequential data array
298 : */
299 : Array<T> _data;
300 : /**
301 : * Dimension information of each inner array
302 : */
303 : Array<JaggedArrayInnerDim<inner>, outer> _dims;
304 : /**
305 : * Starting offset of each inner array into the sequential data array
306 : */
307 : Array<dof_id_type, outer> _offsets;
308 : /**
309 : * Whether the array was finalized
310 : */
311 : bool _finalized;
312 : };
313 :
314 : #define usingKokkosJaggedArrayBaseMembers(T, inner, outer) \
315 : private: \
316 : using JaggedArrayBase<T, inner, outer>::_data; \
317 : using JaggedArrayBase<T, inner, outer>::_dims; \
318 : using JaggedArrayBase<T, inner, outer>::_offsets
319 :
320 : #ifdef MOOSE_KOKKOS_SCOPE
321 : template <typename T, unsigned int inner, unsigned int outer>
322 : void
323 20 : JaggedArrayBase<T, inner, outer>::create(const std::vector<dof_id_type> & n)
324 : {
325 20 : _data.destroy();
326 20 : _offsets.create(n);
327 20 : _dims.create(n);
328 :
329 20 : _finalized = false;
330 20 : }
331 :
332 : template <typename T, unsigned int inner, unsigned int outer>
333 : void
334 242 : JaggedArrayBase<T, inner, outer>::reserve(const std::array<dof_id_type, outer> & index,
335 : const std::array<dof_id_type, inner> & dimension)
336 : {
337 : mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");
338 :
339 242 : dof_id_type idx = 0;
340 242 : dof_id_type stride = 1;
341 :
342 878 : for (unsigned int o = 0; o < outer; ++o)
343 : {
344 636 : idx += index[o] * stride;
345 636 : stride *= _offsets.n(o);
346 : }
347 :
348 242 : stride = 1;
349 :
350 734 : for (unsigned int i = 0; i < inner; ++i)
351 : {
352 492 : stride *= dimension[i];
353 :
354 492 : _dims[idx].dim[i] = dimension[i];
355 492 : _dims[idx].stride[i] = stride;
356 : }
357 242 : }
358 :
359 : template <typename T, unsigned int inner, unsigned int outer>
360 : void
361 20 : JaggedArrayBase<T, inner, outer>::finalize()
362 : {
363 : mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");
364 :
365 20 : dof_id_type stride = 1;
366 :
367 308 : for (dof_id_type o = 0; o < _offsets.size(); ++o)
368 : {
369 288 : stride = 1;
370 :
371 918 : for (unsigned int i = 0; i < inner; ++i)
372 630 : stride *= _dims[o][i];
373 :
374 288 : _offsets[o] = stride;
375 : }
376 :
377 20 : std::exclusive_scan(_offsets.begin(), _offsets.end(), _offsets.begin(), 0);
378 :
379 20 : _dims.copyToDevice();
380 20 : _offsets.copyToDevice();
381 20 : _data.create(_offsets.last() + stride);
382 :
383 20 : _finalized = true;
384 20 : }
385 : #endif
386 :
387 : /**
388 : * The specialization of the Kokkos jagged array class for each dimension.
389 : */
390 : ///{@
391 : template <typename T, unsigned int inner, unsigned int outer>
392 : class JaggedArray;
393 :
394 : template <typename T, unsigned int inner>
395 : class JaggedArray<T, inner, 1> : public JaggedArrayBase<T, inner, 1>
396 : {
397 : usingKokkosJaggedArrayBaseMembers(T, inner, 1);
398 :
399 : public:
400 : /**
401 : * Default constructor
402 : */
403 : JaggedArray() = default;
404 :
405 : #ifdef MOOSE_KOKKOS_SCOPE
406 : /**
407 : * Constructor
408 : * Allocate outer array with given dimensions
409 : * @param n0 The first dimension size for the outer array
410 : */
411 6 : JaggedArray(dof_id_type n0) { create(n0); }
412 : /**
413 : * Allocate outer array with given dimensions
414 : * @param n0 The first dimension size for the outer array
415 : */
416 18 : void create(dof_id_type n0) { JaggedArrayBase<T, inner, 1>::create({n0}); }
417 :
418 : /**
419 : * Get an inner array
420 : * @param i0 The first dimension outer array index
421 : * @returns The inner array wrapper object
422 : */
423 121 : KOKKOS_FUNCTION auto operator()(dof_id_type i0) const
424 : {
425 121 : auto data = &_data[_offsets(i0)];
426 121 : const auto & dim = _dims(i0);
427 :
428 121 : return JaggedArrayInnerData<T, inner>(data, dim);
429 : }
430 : #endif
431 : };
432 :
433 : template <typename T, unsigned int inner>
434 : class JaggedArray<T, inner, 2> : public JaggedArrayBase<T, inner, 2>
435 : {
436 : usingKokkosJaggedArrayBaseMembers(T, inner, 2);
437 :
438 : public:
439 : /**
440 : * Default constructor
441 : */
442 : JaggedArray() = default;
443 :
444 : #ifdef MOOSE_KOKKOS_SCOPE
445 : /**
446 : * Constructor
447 : * Allocate outer array with given dimensions
448 : * @param n0 The first dimension size for the outer array
449 : * @param n1 The second dimension size for the outer array
450 : */
451 6 : JaggedArray(dof_id_type n0, dof_id_type n1) { create(n0, n1); }
452 : /**
453 : * Allocate outer array with given dimensions
454 : * @param n0 The first dimension size for the outer array
455 : * @param n1 The second dimension size for the outer array
456 : */
457 18 : void create(dof_id_type n0, dof_id_type n1) { JaggedArrayBase<T, inner, 2>::create({n0, n1}); }
458 :
459 : /**
460 : * Get an inner array
461 : * @param i0 The first dimension outer array index
462 : * @param i1 The second dimension outer array index
463 : * @returns The inner array wrapper object
464 : */
465 303 : KOKKOS_FUNCTION auto operator()(dof_id_type i0, dof_id_type i1) const
466 : {
467 303 : auto data = &_data[_offsets(i0, i1)];
468 303 : const auto & dim = _dims(i0, i1);
469 :
470 303 : return JaggedArrayInnerData<T, inner>(data, dim);
471 : }
472 : #endif
473 : };
474 :
475 : template <typename T, unsigned int inner>
476 : class JaggedArray<T, inner, 3> : public JaggedArrayBase<T, inner, 3>
477 : {
478 : usingKokkosJaggedArrayBaseMembers(T, inner, 3);
479 :
480 : public:
481 : /**
482 : * Default constructor
483 : */
484 : JaggedArray() = default;
485 :
486 : #ifdef MOOSE_KOKKOS_SCOPE
487 : /**
488 : * Constructor
489 : * Allocate outer array with given dimensions
490 : * @param n0 The first dimension size for the outer array
491 : * @param n1 The second dimension size for the outer array
492 : * @param n2 The third dimension size for the outer array
493 : */
494 8 : JaggedArray(dof_id_type n0, dof_id_type n1, dof_id_type n2) { create(n0, n1, n2); }
495 : /**
496 : * Allocate outer array with given dimensions
497 : * @param n0 The first dimension size for the outer array
498 : * @param n1 The second dimension size for the outer array
499 : * @param n2 The third dimension size for the outer array
500 : */
501 8 : void create(dof_id_type n0, dof_id_type n1, dof_id_type n2)
502 : {
503 16 : JaggedArrayBase<T, inner, 3>::create({n0, n1, n2});
504 8 : }
505 :
506 : /**
507 : * Get an inner array
508 : * @param i0 The first dimension outer array index
509 : * @param i1 The second dimension outer array index
510 : * @param i2 The third dimension outer array index
511 : * @returns The inner array wrapper object
512 : */
513 853 : KOKKOS_FUNCTION auto operator()(dof_id_type i0, dof_id_type i1, dof_id_type i2) const
514 : {
515 853 : auto data = &_data[_offsets(i0, i1, i2)];
516 853 : const auto & dim = _dims(i0, i1, i2);
517 :
518 853 : return JaggedArrayInnerData<T, inner>(data, dim);
519 : }
520 : #endif
521 : };
522 : ///@}
523 :
524 : } // namespace Kokkos
525 : } // namespace Moose
|
