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 : #ifdef MOOSE_KOKKOS_SCOPE
13 : #include "KokkosHeader.h"
14 : #endif
15 :
16 : #include "Conversion.h"
17 : #include "DataIO.h"
18 :
19 : #define usingKokkosArrayBaseMembers(T, dimension) \
20 : private: \
21 : using ArrayBase<T, dimension>::_n; \
22 : using ArrayBase<T, dimension>::_s; \
23 : using ArrayBase<T, dimension>::_d; \
24 : \
25 : public: \
26 : using ArrayBase<T, dimension>::create; \
27 : using ArrayBase<T, dimension>::createHost; \
28 : using ArrayBase<T, dimension>::createDevice; \
29 : using ArrayBase<T, dimension>::offset; \
30 : using ArrayBase<T, dimension>::operator=
31 :
32 : namespace Moose
33 : {
34 : namespace Kokkos
35 : {
36 :
37 : // This function simply calls ::Kokkos::kokkos_free, but it is separately defined in KokkosArray.K
38 : // because the Kokkos function cannot be directly seen by the host compiler
39 : void free(void * ptr);
40 :
41 : /**
42 : * The enumerator that dictates the memory copy direction
43 : */
44 : enum class MemcpyKind
45 : {
46 : HOST_TO_HOST,
47 : HOST_TO_DEVICE,
48 : DEVICE_TO_HOST,
49 : DEVICE_TO_DEVICE
50 : };
51 :
52 : /**
53 : * The Kokkos array class
54 : */
55 : template <typename T, unsigned int dimension = 1>
56 : class Array;
57 :
58 : /**
59 : * The type trait that determines the default behavior of copy constructor and deepCopy()
60 : * If this type trait is set to true, the copy constructor will call deepCopy(),
61 : * and the deepCopy() method will copy-construct each entry.
62 : * If this type trait is set to false, the copy constructor will call shallowCopy(),
63 : * and the deepCopy() method will do a memory copy.
64 : */
65 : ///@{
66 : template <typename T>
67 : struct ArrayDeepCopy
68 : {
69 : static const bool value = false;
70 : };
71 :
72 : template <typename T, unsigned int dimension>
73 : struct ArrayDeepCopy<Array<T, dimension>>
74 : {
75 : static const bool value = ArrayDeepCopy<T>::value;
76 : };
77 : ///@}
78 :
79 : /**
80 : * The base class for Kokkos arrays
81 : */
82 : template <typename T, unsigned int dimension>
83 : class ArrayBase
84 : {
85 : public:
86 : /**
87 : * Default constructor
88 : */
89 4955530 : ArrayBase() = default;
90 :
91 : /**
92 : * Copy constructor
93 : */
94 18329948 : ArrayBase(const ArrayBase<T, dimension> & array)
95 18329948 : {
96 : #ifndef MOOSE_KOKKOS_SCOPE
97 : static_assert(!ArrayDeepCopy<T>::value,
98 : "Kokkos array cannot be deep copied outside the Kokkos compilation scope");
99 : #endif
100 :
101 : if constexpr (ArrayDeepCopy<T>::value)
102 59238 : deepCopy(array);
103 : else
104 18270710 : shallowCopy(array);
105 18329948 : }
106 :
107 : /**
108 : * Destructor
109 : */
110 23025016 : ~ArrayBase() { destroy(); }
111 :
112 : /**
113 : * Free all data and reset
114 : */
115 : void destroy();
116 :
117 : /**
118 : * Shallow copy another Kokkos array
119 : * @param array The Kokkos array to be shallow copied
120 : */
121 : void shallowCopy(const ArrayBase<T, dimension> & array);
122 :
123 : /**
124 : * Get the reference count
125 : * @returns The reference count
126 : */
127 : unsigned int useCount() const { return _counter.use_count(); }
128 :
129 : #ifdef MOOSE_KOKKOS_SCOPE
130 : /**
131 : * Get whether the array was allocated either on host or device
132 : * @returns Whether the array was allocated either on host or device
133 : */
134 16968672 : KOKKOS_FUNCTION bool isAlloc() const { return _is_host_alloc || _is_device_alloc; }
135 : /**
136 : * Get whether the array was allocated on host
137 : * @returns Whether the array was allocated on host
138 : */
139 466 : KOKKOS_FUNCTION bool isHostAlloc() const { return _is_host_alloc; }
140 : /**
141 : * Get whether the array was allocated on device
142 : * @returns Whether the array was allocated on device
143 : */
144 1232 : KOKKOS_FUNCTION bool isDeviceAlloc() const { return _is_device_alloc; }
145 : /**
146 : * Get whether the host array was aliased
147 : * @returns Whether the host array was aliased
148 : */
149 : KOKKOS_FUNCTION bool isHostAlias() const { return _is_host_alias; }
150 : /**
151 : * Get whether the device array was aliased
152 : * @returns Whether the device array was aliased
153 : */
154 : KOKKOS_FUNCTION bool isDeviceAlias() const { return _is_device_alias; }
155 : /**
156 : * Get the total array size
157 : * @returns The total array size
158 : */
159 261761263 : KOKKOS_FUNCTION dof_id_type size() const { return _size; }
160 : /**
161 : * Get the size of a dimension
162 : * @param dim The dimension index
163 : * @returns The size of the dimension
164 : */
165 2036 : KOKKOS_FUNCTION dof_id_type n(unsigned int dim) const { return _n[dim]; }
166 : /**
167 : * Get the data pointer
168 : * @returns The pointer to the underlying data depending on the architecture this function is
169 : * being called on
170 : */
171 18269104 : KOKKOS_FUNCTION T * data() const
172 : {
173 18269104 : KOKKOS_IF_ON_HOST(return _host_data;)
174 :
175 18188796 : return _device_data;
176 : }
177 : /**
178 : * Get the first element
179 : * @returns The reference of the first element depending on the architecture this function is
180 : * being called on
181 : */
182 : KOKKOS_FUNCTION T & first() const
183 : {
184 : KOKKOS_IF_ON_HOST(return _host_data[0];)
185 :
186 : return _device_data[0];
187 : }
188 : /**
189 : * Get the last element
190 : * @returns The reference of the last element depending on the architecture this function is being
191 : * called on
192 : */
193 20 : KOKKOS_FUNCTION T & last() const
194 : {
195 20 : KOKKOS_IF_ON_HOST(return _host_data[_size - 1];)
196 :
197 0 : return _device_data[_size - 1];
198 : }
199 : /**
200 : * Get an array entry
201 : * @param i The dimensionless index
202 : * @returns The reference of the entry depending on the architecture this function is being called
203 : * on
204 : */
205 3260218075 : KOKKOS_FUNCTION T & operator[](dof_id_type i) const
206 : {
207 : KOKKOS_ASSERT(i < _size);
208 :
209 3260218075 : KOKKOS_IF_ON_HOST(return _host_data[i];)
210 :
211 3254127903 : return _device_data[i];
212 : }
213 :
214 : /**
215 : * Get the host data pointer
216 : * @returns The pointer to the underlying host data
217 : */
218 178135 : T * hostData() const { return _host_data; }
219 : /**
220 : * Get the device data pointer
221 : * @returns The pointer to the underlying device data
222 : */
223 173490 : T * deviceData() const { return _device_data; }
224 : /**
225 : * Allocate array on host and device
226 : * @param n The vector containing the size of each dimension
227 : */
228 40 : void create(const std::vector<dof_id_type> & n) { createInternal<true, true>(n); }
229 : /**
230 : * Allocate array on host only
231 : * @param n The vector containing the size of each dimension
232 : */
233 : void createHost(const std::vector<dof_id_type> & n) { createInternal<true, false>(n); }
234 : /**
235 : * Allocate array on device only
236 : * @param n The vector containing the size of each dimension
237 : */
238 1506 : void createDevice(const std::vector<dof_id_type> & n) { createInternal<false, true>(n); }
239 : /**
240 : * Point the host data to an external data instead of allocating it
241 : * @param ptr The pointer to the external host data
242 : */
243 : void aliasHost(T * ptr);
244 : /**
245 : * Point the device data to an external data instead of allocating it
246 : * @param ptr The pointer to the external device data
247 : */
248 : void aliasDevice(T * ptr);
249 : /**
250 : * Apply starting index offsets to each dimension
251 : * @param d The vector containing the offset of each dimension
252 : */
253 : void offset(const std::vector<dof_id_signed_type> & d);
254 : /**
255 : * Copy data from host to device
256 : */
257 : void copyToDevice();
258 : /**
259 : * Copy data from device to host
260 : */
261 : void copyToHost();
262 : /**
263 : * Copy data from an external data to this array
264 : * @param ptr The pointer to the external data
265 : * @param dir The copy direction
266 : * @param n The number of entries to copy
267 : * @param offset The starting offset of this array
268 : */
269 : void copyIn(const T * ptr, MemcpyKind dir, dof_id_type n, dof_id_type offset = 0);
270 : /**
271 : * Copy data to an external data from this array
272 : * @param ptr The pointer to the external data
273 : * @param dir The copy direction
274 : * @param n The number of entries to copy
275 : * @param offset The starting offset of this array
276 : */
277 : void copyOut(T * ptr, MemcpyKind dir, dof_id_type n, dof_id_type offset = 0);
278 : /**
279 : * Copy all the nested Kokkos arrays including self from host to device
280 : */
281 : void copyToDeviceNested();
282 : /**
283 : * Deep copy another Kokkos array
284 : * If ArrayDeepCopy<T>::value is true, it will copy-construct each entry
285 : * If ArrayDeepCopy<T>::value is false, it will do a memory copy
286 : * @param array The Kokkos array to be deep copied
287 : */
288 : void deepCopy(const ArrayBase<T, dimension> & array);
289 : /**
290 : * Swap with another Kokkos array
291 : * @param array The Kokkos array to be swapped
292 : */
293 : void swap(ArrayBase<T, dimension> & array);
294 :
295 : /**
296 : * Assign a scalar value uniformly
297 : * @param scalar The scalar value to be assigned
298 : */
299 : auto & operator=(const T & scalar);
300 :
301 : /**
302 : * Array iterator
303 : */
304 : class iterator
305 : {
306 : public:
307 2689910 : KOKKOS_FUNCTION iterator(T * it) : it(it) {}
308 : KOKKOS_FUNCTION bool operator==(const iterator & other) const { return it == other.it; }
309 1695814 : KOKKOS_FUNCTION bool operator!=(const iterator & other) const { return it != other.it; }
310 910107 : KOKKOS_FUNCTION T & operator*() const { return *it; }
311 1117816 : KOKKOS_FUNCTION T * operator&() const { return it; }
312 909783 : KOKKOS_FUNCTION iterator & operator++()
313 : {
314 909783 : ++it;
315 909783 : return *this;
316 : }
317 324 : KOKKOS_FUNCTION iterator operator++(int)
318 : {
319 324 : iterator pre = *this;
320 324 : ++it;
321 324 : return pre;
322 : }
323 :
324 : private:
325 : T * it;
326 : };
327 :
328 : /**
329 : * Get the beginning iterator
330 : * @returns The beginning iterator
331 : */
332 1344971 : KOKKOS_FUNCTION iterator begin() const
333 : {
334 1344971 : KOKKOS_IF_ON_HOST(return iterator(_host_data);)
335 :
336 557590 : return iterator(_device_data);
337 : }
338 : /**
339 : * Get the end iterator
340 : * @returns The end iterator
341 : */
342 1344939 : KOKKOS_FUNCTION iterator end() const
343 : {
344 1344939 : KOKKOS_IF_ON_HOST(return iterator(_host_data + _size);)
345 :
346 557590 : return iterator(_device_data + _size);
347 : }
348 : #endif
349 :
350 : protected:
351 : /**
352 : * Size of each dimension
353 : */
354 : dof_id_type _n[dimension] = {0};
355 : /**
356 : * Stride of each dimension
357 : */
358 : dof_id_type _s[dimension] = {0};
359 : /**
360 : * Offset of each dimension
361 : */
362 : dof_id_signed_type _d[dimension] = {0};
363 :
364 : #ifdef MOOSE_KOKKOS_SCOPE
365 : /**
366 : * The internal method to initialize and allocate this array
367 : * @tparam host Whether host data will be allocated
368 : * @tparam device Whether device data will be allocated
369 : * @param n The vector containing the size of each dimension
370 : */
371 : template <bool host, bool device>
372 : void createInternal(const std::vector<dof_id_type> & n);
373 : /**
374 : * The internal method to initialize and allocate this array
375 : * @param n The vector containing the size of each dimension
376 : * @param host The flag whether host data will be allocated
377 : * @param device The flag whether device data will be allocated
378 : */
379 : void createInternal(const std::vector<dof_id_type> & n, bool host, bool device);
380 : /**
381 : * The internal method to perform a memory copy
382 : * @tparam TargetSpace The Kokkos memory space of target data
383 : * @tparam Sourcespace The Kokkos memory space of source data
384 : * @param target The pointer to the target data
385 : * @param source The pointer to the source data
386 : * @param n The number of entries to copy
387 : */
388 : template <typename TargetSpace, typename SourceSpace>
389 : void copyInternal(T * target, const T * source, dof_id_type n);
390 : #endif
391 :
392 : private:
393 : #ifdef MOOSE_KOKKOS_SCOPE
394 : /**
395 : * Allocate host data for an initialized array that has not allocated host data
396 : */
397 : void allocHost();
398 : /**
399 : * Allocate device data for an initialized array that has not allocated device data
400 : */
401 : void allocDevice();
402 : #endif
403 :
404 : /**
405 : * Reference counter
406 : */
407 : std::shared_ptr<unsigned int> _counter;
408 : /**
409 : * Flag whether array was initialized
410 : */
411 : bool _is_init = false;
412 : /**
413 : * Flag whether host data was allocated
414 : */
415 : bool _is_host_alloc = false;
416 : /**
417 : * Flag whether device data was allocated
418 : */
419 : bool _is_device_alloc = false;
420 : /**
421 : * Flag whether the host data points to an external data
422 : */
423 : bool _is_host_alias = false;
424 : /**
425 : * Flag whether the device data points to an external data
426 : */
427 : bool _is_device_alias = false;
428 : /**
429 : * Host data
430 : */
431 : T * _host_data = nullptr;
432 : /**
433 : * Device data
434 : */
435 : T * _device_data = nullptr;
436 : /**
437 : * Total size
438 : */
439 : dof_id_type _size = 0;
440 : };
441 :
442 : template <typename T, unsigned int dimension>
443 : void
444 42707093 : ArrayBase<T, dimension>::destroy()
445 : {
446 42707093 : if (!_counter)
447 25326859 : return;
448 :
449 17380234 : if (_counter.use_count() > 1)
450 : {
451 17004176 : _host_data = nullptr;
452 17004176 : _device_data = nullptr;
453 : }
454 376058 : else if (_counter.use_count() == 1)
455 : {
456 376058 : if (_is_host_alloc && !_is_host_alias)
457 : {
458 : if constexpr (std::is_default_constructible<T>::value)
459 : // Allocated by new
460 698158 : delete[] _host_data;
461 : else
462 : {
463 : // Allocated by malloc
464 2175 : for (dof_id_type i = 0; i < _size; ++i)
465 1450 : _host_data[i].~T();
466 :
467 725 : std::free(_host_data);
468 : }
469 : }
470 :
471 376058 : if (_is_device_alloc && !_is_device_alias)
472 376046 : Moose::Kokkos::free(_device_data);
473 : }
474 :
475 17380234 : _size = 0;
476 :
477 43705133 : for (unsigned int i = 0; i < dimension; ++i)
478 : {
479 26324899 : _n[i] = 0;
480 26324899 : _s[i] = 0;
481 26324899 : _d[i] = 0;
482 : }
483 :
484 17380234 : _is_init = false;
485 17380234 : _is_host_alloc = false;
486 17380234 : _is_device_alloc = false;
487 17380234 : _is_host_alias = false;
488 17380234 : _is_device_alias = false;
489 :
490 17380234 : _counter.reset();
491 : }
492 :
493 : template <typename T, unsigned int dimension>
494 : void
495 19297830 : ArrayBase<T, dimension>::shallowCopy(const ArrayBase<T, dimension> & array)
496 : {
497 19297830 : destroy();
498 :
499 19297830 : _counter = array._counter;
500 :
501 19297830 : _size = array._size;
502 :
503 47753681 : for (unsigned int i = 0; i < dimension; ++i)
504 : {
505 28455851 : _n[i] = array._n[i];
506 28455851 : _s[i] = array._s[i];
507 28455851 : _d[i] = array._d[i];
508 : }
509 :
510 19297830 : _is_init = array._is_init;
511 19297830 : _is_host_alloc = array._is_host_alloc;
512 19297830 : _is_device_alloc = array._is_device_alloc;
513 19297830 : _is_host_alias = array._is_host_alias;
514 19297830 : _is_device_alias = array._is_device_alias;
515 :
516 19297830 : _host_data = array._host_data;
517 19297830 : _device_data = array._device_data;
518 19297830 : }
519 :
520 : #ifdef MOOSE_KOKKOS_SCOPE
521 : template <typename T, unsigned int dimension>
522 : void
523 206187 : ArrayBase<T, dimension>::aliasHost(T * ptr)
524 : {
525 206187 : if (!_is_init)
526 0 : mooseError("Kokkos array error: attempted to alias host data before array initialization.");
527 :
528 206187 : if (_is_host_alloc && !_is_host_alias)
529 0 : mooseError("Kokkos array error: cannot alias host data because host data was not aliased.");
530 :
531 206187 : _host_data = ptr;
532 206187 : _is_host_alloc = true;
533 206187 : _is_host_alias = true;
534 206187 : }
535 :
536 : template <typename T, unsigned int dimension>
537 : void
538 32 : ArrayBase<T, dimension>::aliasDevice(T * ptr)
539 : {
540 32 : if (!_is_init)
541 0 : mooseError("Kokkos array error: attempted to alias device data before array initialization.");
542 :
543 32 : if (_is_device_alloc && !_is_device_alias)
544 0 : mooseError("Kokkos array error: cannot alias device data because device data was not aliased.");
545 :
546 32 : _device_data = ptr;
547 32 : _is_device_alloc = true;
548 32 : _is_device_alias = true;
549 32 : }
550 :
551 : template <typename T, unsigned int dimension>
552 : void
553 366338 : ArrayBase<T, dimension>::allocHost()
554 : {
555 366338 : if (_is_host_alloc)
556 0 : return;
557 :
558 : if constexpr (std::is_default_constructible<T>::value)
559 997385 : _host_data = new T[_size];
560 : else
561 741 : _host_data = static_cast<T *>(std::malloc(_size * sizeof(T)));
562 :
563 366338 : _is_host_alloc = true;
564 : }
565 :
566 : template <typename T, unsigned int dimension>
567 : void
568 379755 : ArrayBase<T, dimension>::allocDevice()
569 : {
570 379755 : if (_is_device_alloc)
571 0 : return;
572 :
573 379755 : _device_data =
574 301781 : static_cast<T *>(::Kokkos::kokkos_malloc<ExecSpace::memory_space>(_size * sizeof(T)));
575 :
576 379755 : _is_device_alloc = true;
577 : }
578 :
579 : template <typename T, unsigned int dimension>
580 : template <bool host, bool device>
581 : void
582 379767 : ArrayBase<T, dimension>::createInternal(const std::vector<dof_id_type> & n)
583 : {
584 379767 : if (n.size() != dimension)
585 0 : mooseError("Kokkos array error: the number of dimensions provided (",
586 0 : n.size(),
587 : ") must match the array dimension (",
588 : dimension,
589 : ").");
590 :
591 379767 : if (_counter)
592 7613 : destroy();
593 :
594 379767 : _counter = std::make_shared<unsigned int>();
595 :
596 379767 : _size = 1;
597 379767 : _s[0] = 1;
598 :
599 809825 : for (const auto i : make_range(dimension))
600 : {
601 430058 : _n[i] = n[i];
602 430058 : _size *= n[i];
603 :
604 430058 : if (i)
605 50291 : _s[i] = _s[i - 1] * _n[i - 1];
606 : }
607 :
608 : if constexpr (host)
609 366338 : allocHost();
610 :
611 : if constexpr (device)
612 379755 : allocDevice();
613 :
614 379767 : _is_init = true;
615 379767 : }
616 :
617 : template <typename T, unsigned int dimension>
618 : void
619 60182 : ArrayBase<T, dimension>::createInternal(const std::vector<dof_id_type> & n, bool host, bool device)
620 : {
621 60182 : if (host && device)
622 59238 : createInternal<true, true>(n);
623 944 : else if (host && !device)
624 0 : createInternal<true, false>(n);
625 944 : else if (!host && device)
626 944 : createInternal<false, true>(n);
627 : else
628 0 : createInternal<false, false>(n);
629 60182 : }
630 :
631 : template <typename T, unsigned int dimension>
632 : template <typename TargetSpace, typename SourceSpace>
633 : void
634 1279840 : ArrayBase<T, dimension>::copyInternal(T * target, const T * source, dof_id_type n)
635 : {
636 1279840 : ::Kokkos::Impl::DeepCopy<TargetSpace, SourceSpace>(target, source, n * sizeof(T));
637 1279840 : ::Kokkos::fence();
638 1279840 : }
639 :
640 : template <typename T, unsigned int dimension>
641 : void
642 1698 : ArrayBase<T, dimension>::offset(const std::vector<dof_id_signed_type> & d)
643 : {
644 1698 : if (d.size() > dimension)
645 0 : mooseError("Kokkos array error: the number of offsets provided (",
646 0 : d.size(),
647 : ") cannot be larger than the array dimension (",
648 : dimension,
649 : ").");
650 :
651 3396 : for (const auto i : index_range(d))
652 1698 : _d[i] = d[i];
653 1698 : }
654 :
655 : template <typename T, unsigned int dimension>
656 : void
657 976613 : ArrayBase<T, dimension>::copyToDevice()
658 : {
659 : // If host side memory is not allocated, do nothing
660 976613 : if (!_is_host_alloc)
661 46737 : return;
662 :
663 : // If device side memory is not allocated,
664 929876 : if (!_is_device_alloc)
665 : {
666 0 : if (_counter.use_count() == 1)
667 : // allocate memory if this array is not shared with other arrays
668 0 : allocDevice();
669 : else
670 : // print error if this array is shared with other arrays
671 0 : mooseError("Kokkos array error: cannot copy from host to device because device side memory "
672 : "was not allocated and array is being shared with other arrays.");
673 : }
674 :
675 : // Copy from host to device
676 929876 : copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data, _host_data, _size);
677 : }
678 :
679 : template <typename T, unsigned int dimension>
680 : void
681 174400 : ArrayBase<T, dimension>::copyToHost()
682 : {
683 : // If device side memory is not allocated, do nothing
684 174400 : if (!_is_device_alloc)
685 0 : return;
686 :
687 : // If host side memory is not allocated,
688 174400 : if (!_is_host_alloc)
689 : {
690 0 : if (_counter.use_count() == 1)
691 : // allocate memory if this array is not shared with other arrays
692 0 : allocHost();
693 : else
694 : // print error if this array is shared with other arrays
695 0 : mooseError("Kokkos array error: cannot copy from device to host because host side memory "
696 : "was not allocated and array is being shared with other arrays.");
697 : }
698 :
699 : // Copy from device to host
700 174400 : copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data, _device_data, _size);
701 : }
702 :
703 : template <typename T, unsigned int dimension>
704 : void
705 226 : ArrayBase<T, dimension>::copyIn(const T * ptr, MemcpyKind dir, dof_id_type n, dof_id_type offset)
706 : {
707 226 : if (n > _size)
708 0 : mooseError("Kokkos array error: cannot copyin data larger than the array size.");
709 :
710 226 : if (offset > _size)
711 0 : mooseError("Kokkos array error: offset cannot be larger than the array size.");
712 :
713 226 : if (dir == MemcpyKind::HOST_TO_HOST)
714 : {
715 : // If host side memory is not allocated, do nothing
716 0 : if (!_is_host_alloc)
717 0 : return;
718 :
719 : // Copy from host to host
720 0 : copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(_host_data + offset, ptr, n);
721 : }
722 226 : else if (dir == MemcpyKind::HOST_TO_DEVICE)
723 : {
724 : // If device side memory is not allocated, do nothing
725 226 : if (!_is_device_alloc)
726 0 : return;
727 :
728 : // Copy from host to device
729 226 : copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data + offset, ptr, n);
730 : }
731 0 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
732 : {
733 : // If host side memory is not allocated, do nothing
734 0 : if (!_is_host_alloc)
735 0 : return;
736 :
737 : // Copy from device to host
738 0 : copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data + offset, ptr, n);
739 : }
740 0 : else if (dir == MemcpyKind::DEVICE_TO_DEVICE)
741 : {
742 : // If device side memory is not allocated, do nothing
743 0 : if (!_is_device_alloc)
744 0 : return;
745 :
746 : // Copy from device to device
747 0 : copyInternal<MemSpace, MemSpace>(_device_data + offset, ptr, n);
748 : }
749 : }
750 :
751 : template <typename T, unsigned int dimension>
752 : void
753 934 : ArrayBase<T, dimension>::copyOut(T * ptr, MemcpyKind dir, dof_id_type n, dof_id_type offset)
754 : {
755 934 : if (n > _size)
756 0 : mooseError("Kokkos array error: cannot copyout data larger than the array size.");
757 :
758 934 : if (offset > _size)
759 0 : mooseError("Kokkos array error: offset cannot be larger than the array size.");
760 :
761 934 : if (dir == MemcpyKind::HOST_TO_HOST)
762 : {
763 : // If host side memory is not allocated, do nothing
764 0 : if (!_is_host_alloc)
765 0 : return;
766 :
767 : // Copy from host to host
768 0 : copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
769 : }
770 934 : else if (dir == MemcpyKind::HOST_TO_DEVICE)
771 : {
772 : // If host side memory is not allocated, do nothing
773 0 : if (!_is_host_alloc)
774 0 : return;
775 :
776 : // Copy from host to device
777 0 : copyInternal<MemSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
778 : }
779 934 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
780 : {
781 : // If device side memory is not allocated, do nothing
782 934 : if (!_is_device_alloc)
783 0 : return;
784 :
785 : // Copy from device to host
786 934 : copyInternal<::Kokkos::HostSpace, MemSpace>(ptr, _device_data + offset, n);
787 : }
788 0 : else if (dir == MemcpyKind::DEVICE_TO_DEVICE)
789 : {
790 : // If device side memory is not allocated, do nothing
791 0 : if (!_is_device_alloc)
792 0 : return;
793 :
794 : // Copy from device to device
795 0 : copyInternal<MemSpace, MemSpace>(ptr, _device_data + offset, n);
796 : }
797 : }
798 :
799 : template <typename T>
800 : void
801 1145246 : copyToDeviceInner(T & /* data */)
802 : {
803 1145246 : }
804 :
805 : template <typename T, unsigned int dimension>
806 : void
807 91255 : copyToDeviceInner(Array<T, dimension> & data)
808 : {
809 91255 : data.copyToDeviceNested();
810 91255 : }
811 :
812 : template <typename T, unsigned int dimension>
813 : void
814 112750 : ArrayBase<T, dimension>::copyToDeviceNested()
815 : {
816 1349251 : for (unsigned int i = 0; i < _size; ++i)
817 1236501 : copyToDeviceInner(_host_data[i]);
818 :
819 112750 : copyToDevice();
820 112750 : }
821 :
822 : template <typename T, unsigned int dimension>
823 : void
824 60182 : ArrayBase<T, dimension>::deepCopy(const ArrayBase<T, dimension> & array)
825 : {
826 59238 : if (ArrayDeepCopy<T>::value && !array._is_host_alloc)
827 0 : mooseError(
828 : "Kokkos array error: cannot deep copy using constructor from array without host data.");
829 :
830 180546 : std::vector<dof_id_type> n(std::begin(array._n), std::end(array._n));
831 :
832 60182 : createInternal(n, array._is_host_alloc, array._is_device_alloc);
833 :
834 : if constexpr (ArrayDeepCopy<T>::value)
835 : {
836 118662 : for (dof_id_type i = 0; i < _size; ++i)
837 59424 : new (_host_data + i) T(array._host_data[i]);
838 :
839 59238 : copyToDevice();
840 : }
841 : else
842 : {
843 944 : if (_is_host_alloc)
844 0 : std::memcpy(_host_data, array._host_data, _size * sizeof(T));
845 :
846 944 : if (_is_device_alloc)
847 944 : copyInternal<MemSpace, MemSpace>(_device_data, array._device_data, _size);
848 : }
849 :
850 120364 : for (unsigned int i = 0; i < dimension; ++i)
851 : {
852 60182 : _d[i] = array._d[i];
853 60182 : _s[i] = array._s[i];
854 : }
855 60182 : }
856 :
857 : template <typename T, unsigned int dimension>
858 : void
859 1868 : ArrayBase<T, dimension>::swap(ArrayBase<T, dimension> & array)
860 : {
861 1868 : ArrayBase<T, dimension> clone;
862 :
863 1868 : clone.shallowCopy(*this);
864 1868 : this->shallowCopy(array);
865 1868 : array.shallowCopy(clone);
866 1868 : }
867 :
868 : template <typename T, unsigned int dimension>
869 : auto &
870 215922 : ArrayBase<T, dimension>::operator=(const T & scalar)
871 : {
872 215922 : if (_is_host_alloc)
873 215890 : std::fill_n(_host_data, _size, scalar);
874 :
875 215922 : if (_is_device_alloc)
876 : {
877 215922 : ::Kokkos::View<T *, MemSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>> data(_device_data,
878 : _size);
879 215922 : ::Kokkos::Experimental::fill_n(ExecSpace(), data, _size, scalar);
880 215922 : }
881 :
882 215922 : return *this;
883 : }
884 :
885 : template <typename T, unsigned int dimension>
886 : void
887 934 : dataStore(std::ostream & stream, Array<T, dimension> & array, void * context)
888 : {
889 : using ::dataStore;
890 :
891 934 : bool is_alloc = array.isAlloc();
892 934 : dataStore(stream, is_alloc, nullptr);
893 :
894 934 : if (!is_alloc)
895 0 : return;
896 :
897 934 : std::string type = typeid(T).name();
898 934 : dataStore(stream, type, nullptr);
899 :
900 934 : unsigned int dim = dimension;
901 934 : dataStore(stream, dim, nullptr);
902 :
903 1868 : for (unsigned int dim = 0; dim < dimension; ++dim)
904 : {
905 934 : auto n = array.n(dim);
906 934 : dataStore(stream, n, nullptr);
907 : }
908 :
909 934 : if (array.isDeviceAlloc())
910 : {
911 : // We use malloc/free because we just want a memory copy
912 : // If T is a Kokkos array and we use new/delete or vector to copy it out,
913 : // the arrays will be destroyed on cleanup
914 :
915 934 : T * data = static_cast<T *>(std::malloc(array.size() * sizeof(T)));
916 :
917 934 : array.copyOut(data, MemcpyKind::DEVICE_TO_HOST, array.size());
918 :
919 97366 : for (dof_id_type i = 0; i < array.size(); ++i)
920 96432 : dataStore(stream, data[i], context);
921 :
922 934 : std::free(data);
923 : }
924 : else
925 0 : for (auto & value : array)
926 0 : dataStore(stream, value, context);
927 934 : }
928 :
929 : template <typename T, unsigned int dimension>
930 : void
931 466 : dataLoad(std::istream & stream, Array<T, dimension> & array, void * context)
932 : {
933 : using ::dataLoad;
934 :
935 : bool is_alloc;
936 466 : dataLoad(stream, is_alloc, nullptr);
937 :
938 466 : if (!is_alloc)
939 0 : return;
940 :
941 466 : std::string from_type_name;
942 466 : dataLoad(stream, from_type_name, nullptr);
943 :
944 466 : if (from_type_name != typeid(T).name())
945 0 : mooseError("Kokkos array error: cannot load an array because the stored array is of type '",
946 : MooseUtils::prettyCppType(libMesh::demangle(from_type_name.c_str())),
947 : "' but the loading array is of type '",
948 : MooseUtils::prettyCppType(libMesh::demangle(typeid(T).name())),
949 : "'.");
950 :
951 : unsigned int from_dimension;
952 466 : dataLoad(stream, from_dimension, nullptr);
953 :
954 466 : if (from_dimension != dimension)
955 0 : mooseError("Kokkos array error: cannot load an array because the stored array is ",
956 : from_dimension,
957 : "D but the loading array is ",
958 : dimension,
959 : "D.");
960 :
961 932 : std::vector<dof_id_type> from_n(dimension);
962 466 : std::vector<dof_id_type> n(dimension);
963 :
964 932 : for (unsigned int dim = 0; dim < dimension; ++dim)
965 : {
966 466 : dataLoad(stream, from_n[dim], nullptr);
967 466 : n[dim] = array.n(dim);
968 : }
969 :
970 466 : if (from_n != n)
971 0 : mooseError("Kokkos array error: cannot load an array because the stored array has dimensions (",
972 : Moose::stringify(from_n),
973 : ") but the loading array has dimensions (",
974 : Moose::stringify(n),
975 : ").");
976 :
977 466 : if (array.isHostAlloc())
978 : {
979 2106 : for (auto & value : array)
980 1626 : dataLoad(stream, value, context);
981 :
982 240 : if (array.isDeviceAlloc())
983 240 : array.copyToDevice();
984 : }
985 : else
986 : {
987 226 : std::vector<T> data(array.size());
988 :
989 46002 : for (auto & value : data)
990 45776 : dataLoad(stream, value, context);
991 :
992 226 : array.copyIn(data.data(), MemcpyKind::HOST_TO_DEVICE, array.size());
993 226 : }
994 466 : }
995 : #endif
996 :
997 : /**
998 : * The specialization of the Kokkos array class for each dimension.
999 : * All array data that needs to be accessed on device in Kokkos objects should use this class.
1000 : * If the array is populated on host and is to be accessed on device, make sure to call
1001 : * copyToDevice() after populating data. For a nested Kokkos array, either copyToDeviceNested()
1002 : * should be called for the outermost array or copyToDevice() should be called for each instance of
1003 : * Kokkos array from the innermost to the outermost. Do not store this object as reference in your
1004 : * Kokkos object if it is used on device, because the reference refers to a host object and
1005 : * therefore is not accessible on device. If storing it as a reference is required, see
1006 : * ReferenceWrapper.
1007 : */
1008 : ///@{
1009 : template <typename T>
1010 : class Array<T, 1> : public ArrayBase<T, 1>
1011 : {
1012 : #ifdef MOOSE_KOKKOS_SCOPE
1013 : usingKokkosArrayBaseMembers(T, 1);
1014 : #endif
1015 :
1016 : public:
1017 : /**
1018 : * Default constructor
1019 : */
1020 3966626 : Array() = default;
1021 : /**
1022 : * Copy constructor
1023 : */
1024 11057996 : Array(const Array<T, 1> & array) : ArrayBase<T, 1>(array) {}
1025 : /**
1026 : * Shallow copy another Kokkos array
1027 : * @param array The Kokkos array to be shallow copied
1028 : */
1029 1021516 : auto & operator=(const Array<T, 1> & array)
1030 : {
1031 1021516 : this->shallowCopy(array);
1032 :
1033 1021516 : return *this;
1034 : }
1035 :
1036 : #ifdef MOOSE_KOKKOS_SCOPE
1037 : /**
1038 : * Constructor
1039 : * Initialize and allocate array with given dimensions
1040 : * This allocates both host and device data
1041 : * @param n0 The first dimension size
1042 : */
1043 44 : Array(dof_id_type n0) { create(n0); }
1044 : /**
1045 : * Constructor
1046 : * Initialize and allocate array by copying a standard vector variable
1047 : * This allocates and copies to both host and device data
1048 : * @param vector The standard vector variable to copy
1049 : */
1050 414 : Array(const std::vector<T> & vector) { *this = vector; }
1051 :
1052 : /**
1053 : * Initialize array with given dimensions but do not allocate
1054 : * @param n0 The first dimension size
1055 : */
1056 36 : void init(dof_id_type n0) { this->template createInternal<false, false>({n0}); }
1057 : /**
1058 : * Initialize and allocate array with given dimensions on host and device
1059 : * @param n0 The first dimension size
1060 : */
1061 263631 : void create(dof_id_type n0) { this->template createInternal<true, true>({n0}); }
1062 : /**
1063 : * Initialize and allocate array with given dimensions on host only
1064 : * @param n0 The first dimension size
1065 : */
1066 0 : void createHost(dof_id_type n0) { this->template createInternal<true, false>({n0}); }
1067 : /**
1068 : * Initialize and allocate array with given dimensions on device only
1069 : * @param n0 The first dimension size
1070 : */
1071 32901 : void createDevice(dof_id_type n0) { this->template createInternal<false, true>({n0}); }
1072 : /**
1073 : * Set starting index offsets
1074 : * @param d0 The first dimension offset
1075 : */
1076 5094 : void offset(dof_id_signed_type d0) { ArrayBase<T, 1>::offset({d0}); }
1077 :
1078 : /**
1079 : * Copy a standard vector variable
1080 : * This re-initializes and re-allocates array with the size of the vector
1081 : * @tparam host Whether to allocate and copy to the host data
1082 : * @tparam device Whether to allocate and copy to the device data
1083 : * @param vector The standard vector variable to copy
1084 : */
1085 : template <bool host, bool device>
1086 173460 : void copyVector(const std::vector<T> & vector)
1087 : {
1088 346920 : this->template createInternal<host, device>({static_cast<dof_id_type>(vector.size())});
1089 :
1090 : if (host)
1091 173460 : std::memcpy(this->hostData(), vector.data(), this->size() * sizeof(T));
1092 :
1093 : if (device)
1094 173460 : this->template copyInternal<MemSpace, ::Kokkos::HostSpace>(
1095 : this->deviceData(), vector.data(), this->size());
1096 173460 : }
1097 : /**
1098 : * Copy a standard set variable
1099 : * This re-initializes and re-allocates array with the size of the set
1100 : * @tparam host Whether to allocate and copy to the host data
1101 : * @tparam device Whether to allocate and copy to the device data
1102 : * @param set The standard set variable to copy
1103 : */
1104 : template <bool host, bool device>
1105 169604 : void copySet(const std::set<T> & set)
1106 : {
1107 169604 : std::vector<T> vector(set.begin(), set.end());
1108 :
1109 169604 : copyVector<host, device>(vector);
1110 169604 : }
1111 :
1112 : /**
1113 : * Copy a standard vector variable
1114 : * This allocates and copies to both host and device data
1115 : * @param vector The standard vector variable to copy
1116 : */
1117 3856 : auto & operator=(const std::vector<T> & vector)
1118 : {
1119 3856 : copyVector<true, true>(vector);
1120 :
1121 3856 : return *this;
1122 : }
1123 : /**
1124 : * Copy a standard set variable
1125 : * This allocates and copies to both host and device data
1126 : * @param set The standard set variable to copy
1127 : */
1128 169604 : auto & operator=(const std::set<T> & set)
1129 : {
1130 169604 : copySet<true, true>(set);
1131 :
1132 169604 : return *this;
1133 : }
1134 :
1135 : /**
1136 : * Get an array entry
1137 : * @param i0 The first dimension index
1138 : * @returns The reference of the entry depending on the architecture this function is being
1139 : * called on
1140 : */
1141 127135679 : KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0) const
1142 : {
1143 : KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
1144 :
1145 127135679 : return this->operator[](i0 - _d[0]);
1146 : }
1147 :
1148 : /**
1149 : * Device BLAS operations
1150 : */
1151 : ///@{
1152 : /**
1153 : * Perform \p a * \p x \p op \p b * \p y and write the result to this array
1154 : * @param accumulate Whether to accumulate or overwrite the result
1155 : */
1156 : template <typename U = T>
1157 : typename std::enable_if<std::is_same<U, Real>::value, void>::type
1158 : axby(const U a,
1159 : const Array<U, 1> & x,
1160 : const char op,
1161 : const U b,
1162 : const Array<U, 1> & y,
1163 : const bool accumulate = false);
1164 : /**
1165 : * Scale \p x with \p a and write the result to this array
1166 : */
1167 : template <typename U = T>
1168 : typename std::enable_if<std::is_same<U, Real>::value, void>::type scal(const U a,
1169 : const Array<U, 1> & x);
1170 : /**
1171 : * Scale this array with \p a
1172 : */
1173 : template <typename U = T>
1174 : typename std::enable_if<std::is_same<U, Real>::value, void>::type scal(const U a);
1175 : /**
1176 : * Perform dot product between this array and \p x
1177 : */
1178 : template <typename U = T>
1179 : typename std::enable_if<std::is_same<U, Real>::value, Real>::type dot(const Array<U, 1> & x);
1180 : /**
1181 : * Compute 2-norm of this array
1182 : */
1183 : template <typename U = T>
1184 : typename std::enable_if<std::is_same<U, Real>::value, Real>::type nrm2();
1185 : ///}@
1186 : #endif
1187 : };
1188 :
1189 : template <typename T>
1190 : class Array<T, 2> : public ArrayBase<T, 2>
1191 : {
1192 : #ifdef MOOSE_KOKKOS_SCOPE
1193 : usingKokkosArrayBaseMembers(T, 2);
1194 : #endif
1195 :
1196 : public:
1197 : /**
1198 : * Default constructor
1199 : */
1200 768738 : Array() = default;
1201 : /**
1202 : * Copy constructor
1203 : */
1204 5385883 : Array(const Array<T, 2> & array) : ArrayBase<T, 2>(array) {}
1205 : /**
1206 : * Shallow copy another Kokkos array
1207 : * @param array The Kokkos array to be shallow copied
1208 : */
1209 : auto & operator=(const Array<T, 2> & array)
1210 : {
1211 : this->shallowCopy(array);
1212 :
1213 : return *this;
1214 : }
1215 :
1216 : #ifdef MOOSE_KOKKOS_SCOPE
1217 : /**
1218 : * Constructor
1219 : * Initialize and allocate array with given dimensions
1220 : * This allocates both host and device data
1221 : * @param n0 The first dimension size
1222 : * @param n1 The second dimension size
1223 : */
1224 : Array(dof_id_type n0, dof_id_type n1) { create(n0, n1); }
1225 :
1226 : /**
1227 : * Initialize array with given dimensions but do not allocate
1228 : * @param n0 The first dimension size
1229 : * @param n1 The second dimension size
1230 : */
1231 : void init(dof_id_type n0, dof_id_type n1)
1232 : {
1233 : this->template createInternal<false, false>({n0, n1});
1234 : }
1235 : /**
1236 : * Initialize and allocate array with given dimensions on host and device
1237 : * @param n0 The first dimension size
1238 : * @param n1 The second dimension size
1239 : */
1240 41271 : void create(dof_id_type n0, dof_id_type n1)
1241 : {
1242 82542 : this->template createInternal<true, true>({n0, n1});
1243 41271 : }
1244 : /**
1245 : * Initialize and allocate array with given dimensions on host only
1246 : * @param n0 The first dimension size
1247 : * @param n1 The second dimension size
1248 : */
1249 : void createHost(dof_id_type n0, dof_id_type n1)
1250 : {
1251 : this->template createInternal<true, false>({n0, n1});
1252 : }
1253 : /**
1254 : * Initialize and allocate array with given dimensions on device only
1255 : * @param n0 The first dimension size
1256 : * @param n1 The second dimension size
1257 : */
1258 : void createDevice(dof_id_type n0, dof_id_type n1)
1259 : {
1260 : this->template createInternal<false, true>({n0, n1});
1261 : }
1262 : /**
1263 : * Set starting index offsets
1264 : * @param d0 The first dimension offset
1265 : * @param d1 The second dimension offset
1266 : */
1267 : void offset(dof_id_signed_type d0, dof_id_signed_type d1) { ArrayBase<T, 2>::offset({d0, d1}); }
1268 :
1269 : /**
1270 : * Get an array entry
1271 : * @param i0 The first dimension index
1272 : * @param i1 The second dimension index
1273 : * @returns The reference of the entry depending on the architecture this function is being called
1274 : * on
1275 : */
1276 1037600309 : KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0, dof_id_signed_type i1) const
1277 : {
1278 : KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
1279 : KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
1280 :
1281 1037600309 : return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1]);
1282 : }
1283 : #endif
1284 : };
1285 :
1286 : template <typename T>
1287 : class Array<T, 3> : public ArrayBase<T, 3>
1288 : {
1289 : #ifdef MOOSE_KOKKOS_SCOPE
1290 : usingKokkosArrayBaseMembers(T, 3);
1291 : #endif
1292 :
1293 : public:
1294 : /**
1295 : * Default constructor
1296 : */
1297 219708 : Array() = default;
1298 : /**
1299 : * Copy constructor
1300 : */
1301 1886069 : Array(const Array<T, 3> & array) : ArrayBase<T, 3>(array) {}
1302 : /**
1303 : * Shallow copy another Kokkos array
1304 : * @param array The Kokkos array to be shallow copied
1305 : */
1306 : auto & operator=(const Array<T, 3> & array)
1307 : {
1308 : this->shallowCopy(array);
1309 :
1310 : return *this;
1311 : }
1312 :
1313 : #ifdef MOOSE_KOKKOS_SCOPE
1314 : /**
1315 : * Constructor
1316 : * Initialize and allocate array with given dimensions
1317 : * This allocates both host and device data
1318 : * @param n0 The first dimension size
1319 : * @param n1 The second dimension size
1320 : * @param n2 The third dimension size
1321 : */
1322 : Array(dof_id_type n0, dof_id_type n1, dof_id_type n2) { create(n0, n1, n2); }
1323 :
1324 : /**
1325 : * Initialize array with given dimensions but do not allocate
1326 : * @param n0 The first dimension size
1327 : * @param n1 The second dimension size
1328 : * @param n2 The third dimension size
1329 : */
1330 : void init(dof_id_type n0, dof_id_type n1, dof_id_type n2)
1331 : {
1332 : this->template createInternal<false, false>({n0, n1, n2});
1333 : }
1334 : /**
1335 : * Initialize and allocate array with given dimensions on host and device
1336 : * @param n0 The first dimension size
1337 : * @param n1 The second dimension size
1338 : * @param n2 The third dimension size
1339 : */
1340 4452 : void create(dof_id_type n0, dof_id_type n1, dof_id_type n2)
1341 : {
1342 8904 : this->template createInternal<true, true>({n0, n1, n2});
1343 4452 : }
1344 : /**
1345 : * Initialize and allocate array with given dimensions on host only
1346 : * @param n0 The first dimension size
1347 : * @param n1 The second dimension size
1348 : * @param n2 The third dimension size
1349 : */
1350 : void createHost(dof_id_type n0, dof_id_type n1, dof_id_type n2)
1351 : {
1352 : this->template createInternal<true, false>({n0, n1, n2});
1353 : }
1354 : /**
1355 : * Initialize and allocate array with given dimensions on device only
1356 : * @param n0 The first dimension size
1357 : * @param n1 The second dimension size
1358 : * @param n2 The third dimension size
1359 : */
1360 : void createDevice(dof_id_type n0, dof_id_type n1, dof_id_type n2)
1361 : {
1362 : this->template createInternal<false, true>({n0, n1, n2});
1363 : }
1364 : /**
1365 : * Set starting index offsets
1366 : * @param d0 The first dimension offset
1367 : * @param d1 The second dimension offset
1368 : * @param d2 The third dimension offset
1369 : */
1370 : void offset(dof_id_signed_type d0, dof_id_signed_type d1, dof_id_signed_type d2)
1371 : {
1372 : ArrayBase<T, 3>::offset({d0, d1, d2});
1373 : }
1374 :
1375 : /**
1376 : * Get an array entry
1377 : * @param i0 The first dimension index
1378 : * @param i1 The second dimension index
1379 : * @param i2 The third dimension index
1380 : * @returns The reference of the entry depending on the architecture this function is being
1381 : * called on
1382 : */
1383 : KOKKOS_FUNCTION T &
1384 251559952 : operator()(dof_id_signed_type i0, dof_id_signed_type i1, dof_id_signed_type i2) const
1385 : {
1386 : KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
1387 : KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
1388 : KOKKOS_ASSERT(i2 - _d[2] >= 0 && static_cast<dof_id_type>(i2 - _d[2]) < _n[2]);
1389 :
1390 251559952 : return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1] + (i2 - _d[2]) * _s[2]);
1391 : }
1392 : #endif
1393 : };
1394 :
1395 : template <typename T>
1396 : class Array<T, 4> : public ArrayBase<T, 4>
1397 : {
1398 : #ifdef MOOSE_KOKKOS_SCOPE
1399 : usingKokkosArrayBaseMembers(T, 4);
1400 : #endif
1401 :
1402 : public:
1403 : /**
1404 : * Default constructor
1405 : */
1406 0 : Array() = default;
1407 : /**
1408 : * Copy constructor
1409 : */
1410 : Array(const Array<T, 4> & array) : ArrayBase<T, 4>(array) {}
1411 : /**
1412 : * Shallow copy another Kokkos array
1413 : * @param array The Kokkos array to be shallow copied
1414 : */
1415 : auto & operator=(const Array<T, 4> & array)
1416 : {
1417 : this->shallowCopy(array);
1418 :
1419 : return *this;
1420 : }
1421 :
1422 : #ifdef MOOSE_KOKKOS_SCOPE
1423 : /**
1424 : * Constructor
1425 : * Initialize and allocate array with given dimensions
1426 : * This allocates both host and device data
1427 : * @param n0 The first dimension size
1428 : * @param n1 The second dimension size
1429 : * @param n2 The third dimension size
1430 : * @param n3 The fourth dimension size
1431 : */
1432 : Array(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3) { create(n0, n1, n2, n3); }
1433 :
1434 : /**
1435 : * Initialize array with given dimensions but do not allocate
1436 : * @param n0 The first dimension size
1437 : * @param n1 The second dimension size
1438 : * @param n2 The third dimension size
1439 : * @param n3 The fourth dimension size
1440 : */
1441 : void init(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
1442 : {
1443 : this->template createInternal<false, false>({n0, n1, n2, n3});
1444 : }
1445 : /**
1446 : * Initialize and allocate array with given dimensions on host and device
1447 : * @param n0 The first dimension size
1448 : * @param n1 The second dimension size
1449 : * @param n2 The third dimension size
1450 : * @param n3 The fourth dimension size
1451 : */
1452 : void create(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
1453 : {
1454 : this->template createInternal<true, true>({n0, n1, n2, n3});
1455 : }
1456 : /**
1457 : * Initialize and allocate array with given dimensions on host only
1458 : * @param n0 The first dimension size
1459 : * @param n1 The second dimension size
1460 : * @param n2 The third dimension size
1461 : * @param n3 The fourth dimension size
1462 : */
1463 : void createHost(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
1464 : {
1465 : this->template createInternal<true, false>({n0, n1, n2, n3});
1466 : }
1467 : /**
1468 : * Initialize and allocate array with given dimensions on device only
1469 : * @param n0 The first dimension size
1470 : * @param n1 The second dimension size
1471 : * @param n2 The third dimension size
1472 : * @param n3 The fourth dimension size
1473 : */
1474 : void createDevice(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
1475 : {
1476 : this->template createInternal<false, true>({n0, n1, n2, n3});
1477 : }
1478 : /**
1479 : * Set starting index offsets
1480 : * @param d0 The first dimension offset
1481 : * @param d1 The second dimension offset
1482 : * @param d2 The third dimension offset
1483 : * @param d3 The fourth dimension offset
1484 : */
1485 : void
1486 : offset(dof_id_signed_type d0, dof_id_signed_type d1, dof_id_signed_type d2, dof_id_signed_type d3)
1487 : {
1488 : ArrayBase<T, 4>::offset({d0, d1, d2, d3});
1489 : }
1490 :
1491 : /**
1492 : * Get an array entry
1493 : * @param i0 The first dimension index
1494 : * @param i1 The second dimension index
1495 : * @param i2 The third dimension index
1496 : * @param i3 The fourth dimension index
1497 : * @returns The reference of the entry depending on the architecture this function is being called
1498 : * on
1499 : */
1500 : KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0,
1501 : dof_id_signed_type i1,
1502 : dof_id_signed_type i2,
1503 : dof_id_signed_type i3) const
1504 : {
1505 : KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
1506 : KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
1507 : KOKKOS_ASSERT(i2 - _d[2] >= 0 && static_cast<dof_id_type>(i2 - _d[2]) < _n[2]);
1508 : KOKKOS_ASSERT(i3 - _d[3] >= 0 && static_cast<dof_id_type>(i3 - _d[3]) < _n[3]);
1509 :
1510 : return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1] + (i2 - _d[2]) * _s[2] +
1511 : (i3 - _d[3]) * _s[3]);
1512 : }
1513 : #endif
1514 : };
1515 :
1516 : template <typename T>
1517 : class Array<T, 5> : public ArrayBase<T, 5>
1518 : {
1519 : #ifdef MOOSE_KOKKOS_SCOPE
1520 : usingKokkosArrayBaseMembers(T, 5);
1521 : #endif
1522 :
1523 : public:
1524 : /**
1525 : * Default constructor
1526 : */
1527 0 : Array() = default;
1528 : /**
1529 : * Copy constructor
1530 : */
1531 : Array(const Array<T, 5> & array) : ArrayBase<T, 5>(array) {}
1532 : /**
1533 : * Shallow copy another Kokkos array
1534 : * @param array The Kokkos array to be shallow copied
1535 : */
1536 : auto & operator=(const Array<T, 5> & array)
1537 : {
1538 : this->shallowCopy(array);
1539 :
1540 : return *this;
1541 : }
1542 :
1543 : #ifdef MOOSE_KOKKOS_SCOPE
1544 : /**
1545 : * Constructor
1546 : * Initialize and allocate array with given dimensions
1547 : * This allocates both host and device data
1548 : * @param n0 The first dimension size
1549 : * @param n1 The second dimension size
1550 : * @param n2 The third dimension size
1551 : * @param n3 The fourth dimension size
1552 : * @param n4 The fifth dimension size
1553 : */
1554 : Array(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
1555 : {
1556 : create(n0, n1, n2, n3, n4);
1557 : }
1558 :
1559 : /**
1560 : * Initialize array with given dimensions but do not allocate
1561 : * @param n0 The first dimension size
1562 : * @param n1 The second dimension size
1563 : * @param n2 The third dimension size
1564 : * @param n3 The fourth dimension size
1565 : * @param n4 The fifth dimension size
1566 : */
1567 : void init(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
1568 : {
1569 : this->template createInternal<false, false>({n0, n1, n2, n3, n4});
1570 : }
1571 : /**
1572 : * Initialize and allocate array with given dimensions on host and device
1573 : * @param n0 The first dimension size
1574 : * @param n1 The second dimension size
1575 : * @param n2 The third dimension size
1576 : * @param n3 The fourth dimension size
1577 : * @param n4 The fifth dimension size
1578 : */
1579 : void create(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
1580 : {
1581 : this->template createInternal<true, true>({n0, n1, n2, n3, n4});
1582 : }
1583 : /**
1584 : * Initialize and allocate array with given dimensions on host only
1585 : * @param n0 The first dimension size
1586 : * @param n1 The second dimension size
1587 : * @param n2 The third dimension size
1588 : * @param n3 The fourth dimension size
1589 : * @param n4 The fifth dimension size
1590 : */
1591 : void createHost(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
1592 : {
1593 : this->template createInternal<true, false>({n0, n1, n2, n3, n4});
1594 : }
1595 : /**
1596 : * Initialize and allocate array with given dimensions on device only
1597 : * @param n0 The first dimension size
1598 : * @param n1 The second dimension size
1599 : * @param n2 The third dimension size
1600 : * @param n3 The fourth dimension size
1601 : * @param n4 The fifth dimension size
1602 : */
1603 : void createDevice(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
1604 : {
1605 : this->template createInternal<false, true>({n0, n1, n2, n3, n4});
1606 : }
1607 : /**
1608 : * Set starting index offsets
1609 : * @param d0 The first dimension offset
1610 : * @param d1 The second dimension offset
1611 : * @param d2 The third dimension offset
1612 : * @param d3 The fourth dimension offset
1613 : * @param d4 The fifth dimension offset
1614 : */
1615 : void offset(dof_id_signed_type d0,
1616 : dof_id_signed_type d1,
1617 : dof_id_signed_type d2,
1618 : dof_id_signed_type d3,
1619 : dof_id_signed_type d4)
1620 : {
1621 : ArrayBase<T, 5>::offset({d0, d1, d2, d3, d4});
1622 : }
1623 :
1624 : /**
1625 : * Get an array entry
1626 : * @param i0 The first dimension index
1627 : * @param i1 The second dimension index
1628 : * @param i2 The third dimension index
1629 : * @param i3 The fourth dimension index
1630 : * @param i4 The fifth dimension index
1631 : * @returns The reference of the entry depending on the architecture this function is being called
1632 : * on
1633 : */
1634 : KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0,
1635 : dof_id_signed_type i1,
1636 : dof_id_signed_type i2,
1637 : dof_id_signed_type i3,
1638 : dof_id_signed_type i4) const
1639 : {
1640 : KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
1641 : KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
1642 : KOKKOS_ASSERT(i2 - _d[2] >= 0 && static_cast<dof_id_type>(i2 - _d[2]) < _n[2]);
1643 : KOKKOS_ASSERT(i3 - _d[3] >= 0 && static_cast<dof_id_type>(i3 - _d[3]) < _n[3]);
1644 : KOKKOS_ASSERT(i4 - _d[4] >= 0 && static_cast<dof_id_type>(i4 - _d[4]) < _n[4]);
1645 :
1646 : return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1] + (i2 - _d[2]) * _s[2] +
1647 : (i3 - _d[3]) * _s[3] + (i4 - _d[4]) * _s[4]);
1648 : }
1649 : #endif
1650 : };
1651 : ///@}
1652 :
1653 : template <typename T>
1654 : using Array1D = Array<T, 1>;
1655 : template <typename T>
1656 : using Array2D = Array<T, 2>;
1657 : template <typename T>
1658 : using Array3D = Array<T, 3>;
1659 : template <typename T>
1660 : using Array4D = Array<T, 4>;
1661 : template <typename T>
1662 : using Array5D = Array<T, 5>;
1663 :
1664 : } // namespace Kokkos
1665 : } // namespace Moose
|
