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 : * Get the underlying data array
261 : * @returns The data array
262 : */
263 : auto & array() { return _data; }
264 :
265 : /**
266 : * Get whether the array is finalized
267 : * @returns Whether the array is finalized
268 : */
269 : KOKKOS_FUNCTION bool isFinalized() const { return _finalized; }
270 : /**
271 : * Get the total data array size
272 : * @returns The total data array size
273 : */
274 : KOKKOS_FUNCTION dof_id_type size() const { return _data.size(); }
275 : /**
276 : * Get the total outer array size
277 : * @returns The total outer array size
278 : */
279 : KOKKOS_FUNCTION dof_id_type n() const { return _offsets.size(); }
280 : /**
281 : * Get the size of a dimension of the outer array
282 : * @param dim The dimension index
283 : * @returns The size of the dimension of the outer array
284 : */
285 : KOKKOS_FUNCTION dof_id_type n(unsigned int dim) const { return _offsets.n(dim); }
286 : /**
287 : * Get an inner array
288 : * @param i The dimensionless outer array index
289 : * @returns The inner array wrapper object
290 : */
291 : KOKKOS_FUNCTION auto operator[](dof_id_type i) const
292 : {
293 : auto data = &_data[_offsets[i]];
294 : const auto & dim = _dims[i];
295 :
296 : return JaggedArrayInnerData<T, inner>(data, dim);
297 : }
298 : #endif
299 :
300 : protected:
301 : /**
302 : * Sequential data array
303 : */
304 : Array<T> _data;
305 : /**
306 : * Dimension information of each inner array
307 : */
308 : Array<JaggedArrayInnerDim<inner>, outer> _dims;
309 : /**
310 : * Starting offset of each inner array into the sequential data array
311 : */
312 : Array<dof_id_type, outer> _offsets;
313 : /**
314 : * Whether the array was finalized
315 : */
316 : bool _finalized;
317 : };
318 :
319 : #define usingKokkosJaggedArrayBaseMembers(T, inner, outer) \
320 : private: \
321 : using JaggedArrayBase<T, inner, outer>::_data; \
322 : using JaggedArrayBase<T, inner, outer>::_dims; \
323 : using JaggedArrayBase<T, inner, outer>::_offsets
324 :
325 : #ifdef MOOSE_KOKKOS_SCOPE
326 : template <typename T, unsigned int inner, unsigned int outer>
327 : void
328 20 : JaggedArrayBase<T, inner, outer>::create(const std::vector<dof_id_type> & n)
329 : {
330 20 : _data.destroy();
331 20 : _offsets.create(n);
332 20 : _dims.create(n);
333 :
334 20 : _finalized = false;
335 20 : }
336 :
337 : template <typename T, unsigned int inner, unsigned int outer>
338 : void
339 242 : JaggedArrayBase<T, inner, outer>::reserve(const std::array<dof_id_type, outer> & index,
340 : const std::array<dof_id_type, inner> & dimension)
341 : {
342 : mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");
343 :
344 242 : dof_id_type idx = 0;
345 242 : dof_id_type stride = 1;
346 :
347 878 : for (unsigned int o = 0; o < outer; ++o)
348 : {
349 636 : idx += index[o] * stride;
350 636 : stride *= _offsets.n(o);
351 : }
352 :
353 242 : stride = 1;
354 :
355 734 : for (unsigned int i = 0; i < inner; ++i)
356 : {
357 492 : stride *= dimension[i];
358 :
359 492 : _dims[idx].dim[i] = dimension[i];
360 492 : _dims[idx].stride[i] = stride;
361 : }
362 242 : }
363 :
364 : template <typename T, unsigned int inner, unsigned int outer>
365 : void
366 20 : JaggedArrayBase<T, inner, outer>::finalize()
367 : {
368 : mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");
369 :
370 20 : dof_id_type stride = 1;
371 :
372 308 : for (dof_id_type o = 0; o < _offsets.size(); ++o)
373 : {
374 288 : stride = 1;
375 :
376 918 : for (unsigned int i = 0; i < inner; ++i)
377 630 : stride *= _dims[o][i];
378 :
379 288 : _offsets[o] = stride;
380 : }
381 :
382 20 : std::exclusive_scan(_offsets.begin(), _offsets.end(), _offsets.begin(), 0);
383 :
384 20 : _dims.copyToDevice();
385 20 : _offsets.copyToDevice();
386 20 : _data.create(_offsets.last() + stride);
387 :
388 20 : _finalized = true;
389 20 : }
390 : #endif
391 :
392 : /**
393 : * The specialization of the Kokkos jagged array class for each dimension.
394 : */
395 : ///{@
396 : template <typename T, unsigned int inner, unsigned int outer>
397 : class JaggedArray;
398 :
399 : template <typename T, unsigned int inner>
400 : class JaggedArray<T, inner, 1> : public JaggedArrayBase<T, inner, 1>
401 : {
402 : usingKokkosJaggedArrayBaseMembers(T, inner, 1);
403 :
404 : public:
405 : /**
406 : * Default constructor
407 : */
408 : JaggedArray() = default;
409 :
410 : #ifdef MOOSE_KOKKOS_SCOPE
411 : /**
412 : * Constructor
413 : * Allocate outer array with given dimensions
414 : * @param n0 The first dimension size for the outer array
415 : */
416 6 : JaggedArray(dof_id_type n0) { create(n0); }
417 : /**
418 : * Allocate outer array with given dimensions
419 : * @param n0 The first dimension size for the outer array
420 : */
421 18 : void create(dof_id_type n0) { JaggedArrayBase<T, inner, 1>::create({n0}); }
422 :
423 : /**
424 : * Get an inner array
425 : * @param i0 The first dimension outer array index
426 : * @returns The inner array wrapper object
427 : */
428 121 : KOKKOS_FUNCTION auto operator()(dof_id_type i0) const
429 : {
430 121 : auto data = &_data[_offsets(i0)];
431 121 : const auto & dim = _dims(i0);
432 :
433 121 : return JaggedArrayInnerData<T, inner>(data, dim);
434 : }
435 : #endif
436 : };
437 :
438 : template <typename T, unsigned int inner>
439 : class JaggedArray<T, inner, 2> : public JaggedArrayBase<T, inner, 2>
440 : {
441 : usingKokkosJaggedArrayBaseMembers(T, inner, 2);
442 :
443 : public:
444 : /**
445 : * Default constructor
446 : */
447 : JaggedArray() = default;
448 :
449 : #ifdef MOOSE_KOKKOS_SCOPE
450 : /**
451 : * Constructor
452 : * Allocate outer array with given dimensions
453 : * @param n0 The first dimension size for the outer array
454 : * @param n1 The second dimension size for the outer array
455 : */
456 6 : JaggedArray(dof_id_type n0, dof_id_type n1) { create(n0, n1); }
457 : /**
458 : * Allocate outer array with given dimensions
459 : * @param n0 The first dimension size for the outer array
460 : * @param n1 The second dimension size for the outer array
461 : */
462 18 : void create(dof_id_type n0, dof_id_type n1) { JaggedArrayBase<T, inner, 2>::create({n0, n1}); }
463 :
464 : /**
465 : * Get an inner array
466 : * @param i0 The first dimension outer array index
467 : * @param i1 The second dimension outer array index
468 : * @returns The inner array wrapper object
469 : */
470 303 : KOKKOS_FUNCTION auto operator()(dof_id_type i0, dof_id_type i1) const
471 : {
472 303 : auto data = &_data[_offsets(i0, i1)];
473 303 : const auto & dim = _dims(i0, i1);
474 :
475 303 : return JaggedArrayInnerData<T, inner>(data, dim);
476 : }
477 : #endif
478 : };
479 :
480 : template <typename T, unsigned int inner>
481 : class JaggedArray<T, inner, 3> : public JaggedArrayBase<T, inner, 3>
482 : {
483 : usingKokkosJaggedArrayBaseMembers(T, inner, 3);
484 :
485 : public:
486 : /**
487 : * Default constructor
488 : */
489 : JaggedArray() = default;
490 :
491 : #ifdef MOOSE_KOKKOS_SCOPE
492 : /**
493 : * Constructor
494 : * Allocate outer array with given dimensions
495 : * @param n0 The first dimension size for the outer array
496 : * @param n1 The second dimension size for the outer array
497 : * @param n2 The third dimension size for the outer array
498 : */
499 8 : JaggedArray(dof_id_type n0, dof_id_type n1, dof_id_type n2) { create(n0, n1, n2); }
500 : /**
501 : * Allocate outer array with given dimensions
502 : * @param n0 The first dimension size for the outer array
503 : * @param n1 The second dimension size for the outer array
504 : * @param n2 The third dimension size for the outer array
505 : */
506 8 : void create(dof_id_type n0, dof_id_type n1, dof_id_type n2)
507 : {
508 16 : JaggedArrayBase<T, inner, 3>::create({n0, n1, n2});
509 8 : }
510 :
511 : /**
512 : * Get an inner array
513 : * @param i0 The first dimension outer array index
514 : * @param i1 The second dimension outer array index
515 : * @param i2 The third dimension outer array index
516 : * @returns The inner array wrapper object
517 : */
518 853 : KOKKOS_FUNCTION auto operator()(dof_id_type i0, dof_id_type i1, dof_id_type i2) const
519 : {
520 853 : auto data = &_data[_offsets(i0, i1, i2)];
521 853 : const auto & dim = _dims(i0, i1, i2);
522 :
523 853 : return JaggedArrayInnerData<T, inner>(data, dim);
524 : }
525 : #endif
526 : };
527 : ///@}
528 :
529 : } // namespace Kokkos
530 : } // namespace Moose
|
