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 : #include <iterator>
20 :
21 : #define usingKokkosArrayBaseMembers(T, dimension, index_type) \
22 : private: \
23 : using ArrayBase<T, dimension, index_type>::_n; \
24 : using ArrayBase<T, dimension, index_type>::_s; \
25 : using ArrayBase<T, dimension, index_type>::_d; \
26 : \
27 : public: \
28 : using typename ArrayBase<T, dimension, index_type>::signed_index_type; \
29 : using ArrayBase<T, dimension, index_type>::operator=
30 :
31 : namespace Moose::Kokkos
32 : {
33 :
34 : // This function simply calls ::Kokkos::kokkos_free, but it is separately defined in KokkosArray.K
35 : // because the Kokkos function cannot be directly seen by the host compiler
36 : void free(void * ptr);
37 :
38 : /**
39 : * The enumerator that dictates the memory copy direction
40 : */
41 : enum class MemcpyType
42 : {
43 : HOST_TO_HOST,
44 : HOST_TO_DEVICE,
45 : DEVICE_TO_HOST,
46 : DEVICE_TO_DEVICE
47 : };
48 :
49 : /**
50 : * The enumerator that dictates the memory layout
51 : */
52 : enum class LayoutType
53 : {
54 : LEFT,
55 : RIGHT
56 : };
57 :
58 : /**
59 : * The Kokkos array class
60 : */
61 : template <typename T,
62 : unsigned int dimension = 1,
63 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
64 : LayoutType layout = LayoutType::LEFT>
65 : class Array;
66 :
67 : /**
68 : * The type trait that determines if a template type is Kokkos array
69 : */
70 : ///@{
71 : template <typename>
72 : struct is_kokkos_array : std::false_type
73 : {
74 : };
75 :
76 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
77 : struct is_kokkos_array<Array<T, dimension, index_type, layout>> : std::true_type
78 : {
79 : };
80 : ///@}
81 :
82 : /**
83 : * The type trait that determines the default behavior of copy constructor and deepCopy()
84 : * If this type trait is set to true, the copy constructor will call deepCopy(),
85 : * and the deepCopy() method will copy-construct each entry.
86 : * If this type trait is set to false, the copy constructor will call shallowCopy(),
87 : * and the deepCopy() method will do a memory copy.
88 : */
89 : ///@{
90 : template <typename T>
91 : struct ArrayDeepCopy
92 : {
93 : static constexpr bool value = false;
94 : };
95 :
96 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
97 : struct ArrayDeepCopy<Array<T, dimension, index_type, layout>>
98 : {
99 : static constexpr bool value = ArrayDeepCopy<T>::value;
100 : };
101 : ///@}
102 :
103 : /**
104 : * The base class for Kokkos arrays
105 : */
106 : template <typename T, unsigned int dimension, typename index_type>
107 : class ArrayBase
108 : {
109 : static_assert(std::is_integral_v<index_type>, "Kokkos array index type must be an integral type");
110 : static_assert(std::is_unsigned_v<index_type>, "Kokkos array index type must be unsigned");
111 : static_assert(!std::is_same_v<bool, index_type>, "Kokkos array index type must not be bool");
112 :
113 : public:
114 : using unsigned_index_type = index_type;
115 : using signed_index_type = typename std::make_signed<index_type>::type;
116 :
117 : /**
118 : * Constructor
119 : * @param layout The memory layout type
120 : */
121 8616894 : ArrayBase(const LayoutType layout) : _layout(layout) {}
122 :
123 : #ifdef MOOSE_KOKKOS_SCOPE
124 : /**
125 : * Constructor
126 : * Initialize and allocate array with given dimensions
127 : * This allocates both host and device data
128 : * @param layout The memory layout type
129 : * @param n The size of each dimension
130 : */
131 : template <typename... size_type>
132 76 : ArrayBase(const LayoutType layout, size_type... n) : _layout(layout)
133 : {
134 76 : create(n...);
135 76 : }
136 : #endif
137 :
138 : /**
139 : * Copy constructor
140 : */
141 59346199 : ArrayBase(const ArrayBase<T, dimension, index_type> & array) : _layout(array._layout)
142 : {
143 : #ifndef MOOSE_KOKKOS_SCOPE
144 : static_assert(!ArrayDeepCopy<T>::value,
145 : "Kokkos array cannot be deep copied outside the Kokkos compilation scope");
146 : #endif
147 :
148 : if constexpr (ArrayDeepCopy<T>::value)
149 92288 : deepCopy(array);
150 : else
151 59253911 : shallowCopy(array);
152 59346199 : }
153 :
154 : /**
155 : * Destructor
156 : */
157 67701673 : ~ArrayBase() { destroy(); }
158 :
159 : /**
160 : * Free all data and reset
161 : */
162 : void destroy();
163 :
164 : /**
165 : * Shallow copy another Kokkos array
166 : * @param array The Kokkos array to be shallow copied
167 : */
168 : void shallowCopy(const ArrayBase<T, dimension, index_type> & array);
169 :
170 : /**
171 : * Get the reference count
172 : * @returns The reference count
173 : */
174 : unsigned int useCount() const { return _counter.use_count(); }
175 :
176 : /**
177 : * Get whether slot i is tracked as constructed
178 : *
179 : * For an array created with initialization (initialize = true) every slot is default-constructed
180 : * at allocation, so all slots are reported as constructed. For an array created without
181 : * initialization (initialize = false) only slots constructed through emplace() are tracked;
182 : * placement-new directly into the storage is not recorded here, so such a slot is reported as not
183 : * constructed and its destructor is skipped on free (a leak for types that own resources).
184 : * @param i The slot index
185 : * @returns true if the slot is tracked as constructed
186 : */
187 : bool isSlotConstructed(index_type i) const;
188 :
189 : template <bool is_const>
190 : class constructed_entry_range;
191 :
192 : /**
193 : * Iterator over tracked constructed entries.
194 : */
195 : template <bool is_const>
196 : class constructed_entry_iterator
197 : {
198 : public:
199 : using iterator_category = std::forward_iterator_tag;
200 : using value_type = T;
201 : using difference_type = std::ptrdiff_t;
202 : using pointer = std::conditional_t<is_const, const T *, T *>;
203 : using reference = std::conditional_t<is_const, const T &, T &>;
204 : using array_type = std::conditional_t<is_const, const ArrayBase, ArrayBase>;
205 :
206 676481 : reference operator*() const { return _array._host_data[_i]; }
207 : pointer operator->() const { return _array._host_data + _i; }
208 : constructed_entry_iterator & operator++();
209 : constructed_entry_iterator operator++(int);
210 : bool operator==(const constructed_entry_iterator & other) const;
211 : bool operator!=(const constructed_entry_iterator & other) const;
212 :
213 : private:
214 : friend class constructed_entry_range<is_const>;
215 :
216 : constructed_entry_iterator(array_type & array, index_type i);
217 :
218 : void advanceToConstructed();
219 :
220 : /**
221 : * Array whose constructed entries are being iterated.
222 : */
223 : array_type & _array;
224 : /**
225 : * Current slot index in _array.
226 : */
227 : index_type _i = 0;
228 : };
229 :
230 : /**
231 : * Range over tracked constructed entries.
232 : */
233 : template <bool is_const>
234 : class constructed_entry_range
235 : {
236 : public:
237 : using array_type = std::conditional_t<is_const, const ArrayBase, ArrayBase>;
238 : using iterator = constructed_entry_iterator<is_const>;
239 :
240 : explicit constructed_entry_range(array_type & array);
241 :
242 : iterator begin() const;
243 : iterator end() const;
244 :
245 : private:
246 : /**
247 : * Array that provides the constructed-entry iteration bounds.
248 : */
249 : array_type & _array;
250 : };
251 :
252 : /**
253 : * Get host-side range over entries tracked as constructed.
254 : * @returns Range that skips slots for which isSlotConstructed() is false
255 : */
256 : constructed_entry_range<false> constructedEntries();
257 :
258 : /**
259 : * Get host-side range over entries tracked as constructed.
260 : * @returns Range that skips slots for which isSlotConstructed() is false
261 : */
262 : constructed_entry_range<true> constructedEntries() const;
263 :
264 : #ifdef MOOSE_KOKKOS_SCOPE
265 : /**
266 : * Get whether the array was allocated either on host or device
267 : * @returns Whether the array was allocated either on host or device
268 : */
269 13226688 : KOKKOS_FUNCTION bool isAlloc() const { return _is_host_alloc || _is_device_alloc; }
270 : /**
271 : * Get whether the array was allocated on host
272 : * @returns Whether the array was allocated on host
273 : */
274 531 : KOKKOS_FUNCTION bool isHostAlloc() const { return _is_host_alloc; }
275 : /**
276 : * Get whether the array was allocated on device
277 : * @returns Whether the array was allocated on device
278 : */
279 1472 : KOKKOS_FUNCTION bool isDeviceAlloc() const { return _is_device_alloc; }
280 : /**
281 : * Get whether the host array was aliased
282 : * @returns Whether the host array was aliased
283 : */
284 : KOKKOS_FUNCTION bool isHostAlias() const { return _is_host_alias; }
285 : /**
286 : * Get whether the device array was aliased
287 : * @returns Whether the device array was aliased
288 : */
289 : KOKKOS_FUNCTION bool isDeviceAlias() const { return _is_device_alias; }
290 : /**
291 : * Get the total array size
292 : * @returns The total array size
293 : */
294 471192544 : KOKKOS_FUNCTION index_type size() const { return _size; }
295 : /**
296 : * Get the size of a dimension
297 : * @param dim The dimension index
298 : * @returns The size of the dimension
299 : */
300 3530 : KOKKOS_FUNCTION index_type n(unsigned int dim) const { return _n[dim]; }
301 : /**
302 : * Get the data pointer
303 : * @returns The pointer to the underlying data depending on the architecture this function is
304 : * being called on
305 : */
306 41629409 : KOKKOS_FUNCTION T * data() const
307 : {
308 41629409 : KOKKOS_IF_ON_HOST(return _host_data;)
309 :
310 41249679 : return _device_data;
311 : }
312 : /**
313 : * Get the first element
314 : * @returns The reference of the first element depending on the architecture this function is
315 : * being called on
316 : */
317 : KOKKOS_FUNCTION T & first() const
318 : {
319 : KOKKOS_IF_ON_HOST(return _host_data[0];)
320 :
321 : return _device_data[0];
322 : }
323 : /**
324 : * Get the last element
325 : * @returns The reference of the last element depending on the architecture this function is being
326 : * called on
327 : */
328 56 : KOKKOS_FUNCTION T & last() const
329 : {
330 56 : KOKKOS_IF_ON_HOST(return _host_data[_size - 1];)
331 :
332 0 : return _device_data[_size - 1];
333 : }
334 : /**
335 : * Get an array entry
336 : * @param i The dimensionless index
337 : * @returns The reference of the entry depending on the architecture this function is being called
338 : * on
339 : */
340 5275644048 : KOKKOS_FUNCTION T & operator[](index_type i) const
341 : {
342 : KOKKOS_ASSERT(i < _size);
343 :
344 5275644048 : KOKKOS_IF_ON_HOST(return _host_data[i];)
345 :
346 5217476807 : return _device_data[i];
347 : }
348 :
349 : /**
350 : * Get the host data pointer
351 : * @returns The pointer to the underlying host data
352 : */
353 1347071 : T * hostData() const { return _host_data; }
354 : /**
355 : * Get the device data pointer
356 : * @returns The pointer to the underlying device data
357 : */
358 1340213 : T * deviceData() const { return _device_data; }
359 : /**
360 : * Get the host unmanaged view
361 : * @returns The host unmanaged view
362 : */
363 : auto hostView() const
364 : {
365 : return ::Kokkos::View<T *, ::Kokkos::HostSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>>(
366 : _host_data, _size);
367 : }
368 : /**
369 : * Get the device unmanaged view
370 : * @returns The device unmanaged view
371 : */
372 596102 : auto deviceView() const
373 : {
374 596102 : return ::Kokkos::View<T *, MemSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>>(_device_data,
375 341948 : _size);
376 : }
377 : /**
378 : * Initialize array with given dimensions but do not allocate
379 : * @param n The size of each dimension
380 : */
381 : template <typename... size_type>
382 200 : void init(size_type... n)
383 : {
384 200 : createInternal<false, false, false>(n...);
385 200 : }
386 : /**
387 : * Allocate array on host and device
388 : * @tparam initialize Whether to initialize host data (calls default constructor)
389 : * @param n The vector containing the size of each dimension
390 : */
391 : template <bool initialize = true>
392 : void create(const std::vector<index_type> & n)
393 : {
394 : createInternal<true, true, initialize>(n);
395 : }
396 : /**
397 : * Allocate array on host and device
398 : * @tparam initialize Whether to initialize host data (calls default constructor)
399 : * @param n The size of each dimension
400 : */
401 : template <bool initialize = true, typename... size_type>
402 609810 : void create(size_type... n)
403 : {
404 609810 : createInternal<true, true, initialize>(n...);
405 609810 : }
406 : /**
407 : * Allocate array on host only
408 : * @tparam initialize Whether to initialize host data (calls default constructor)
409 : * @param n The vector containing the size of each dimension
410 : */
411 : template <bool initialize = true>
412 : void createHost(const std::vector<index_type> & n)
413 : {
414 : createInternal<true, false, initialize>(n);
415 : }
416 : /**
417 : * Allocate array on host only
418 : * @tparam initialize Whether to initialize host data (calls default constructor)
419 : * @param n The size of each dimension
420 : */
421 : template <bool initialize = true, typename... size_type>
422 24691 : void createHost(size_type... n)
423 : {
424 24691 : createInternal<true, false, initialize>(n...);
425 24691 : }
426 : /**
427 : * Allocate array on device only
428 : * @param n The vector containing the size of each dimension
429 : */
430 5120 : void createDevice(const std::vector<index_type> & n) { createInternal<false, true, false>(n); }
431 : /**
432 : * Allocate array on device only
433 : * @param n The size of each dimension
434 : */
435 : template <typename... size_type>
436 33896 : void createDevice(size_type... n)
437 : {
438 33896 : createInternal<false, true, false>(n...);
439 33896 : }
440 : /**
441 : * Point the host data to an external data instead of allocating it
442 : * @param ptr The pointer to the external host data
443 : */
444 : void aliasHost(T * ptr);
445 : /**
446 : * Point the device data to an external data instead of allocating it
447 : * @param ptr The pointer to the external device data
448 : */
449 : void aliasDevice(T * ptr);
450 : /**
451 : * Apply starting index offsets to each dimension
452 : * @param d The vector containing the offset of each dimension
453 : */
454 : void offset(const std::vector<signed_index_type> & d);
455 : /**
456 : * Apply starting index offsets to each dimension
457 : * @param d The offset of each dimension
458 : */
459 : template <typename... offset_type>
460 : void offset(offset_type... d);
461 : /**
462 : * Copy data from host to device
463 : */
464 : void copyToDevice();
465 : /**
466 : * Copy data from device to host
467 : */
468 : void copyToHost();
469 : /**
470 : * Copy data from an external data to this array
471 : * @param ptr The pointer to the external data
472 : * @param dir The copy direction
473 : * @param n The number of entries to copy
474 : * @param offset The starting offset of this array
475 : */
476 : void copyIn(const T * ptr, MemcpyType dir, index_type n, index_type offset = 0);
477 : /**
478 : * Copy data to an external data from this array
479 : * @param ptr The pointer to the external data
480 : * @param dir The copy direction
481 : * @param n The number of entries to copy
482 : * @param offset The starting offset of this array
483 : */
484 : void copyOut(T * ptr, MemcpyType dir, index_type n, index_type offset = 0);
485 : /**
486 : * Copy all the nested Kokkos arrays including self from host to device
487 : */
488 : void copyToDeviceNested();
489 : /**
490 : * Copy data from host to device and deallocate host
491 : * @param should_free_host Whether the host memory should be freed.
492 : * Host memory cannot be freed when there are shallow copies of this array that are still alive.
493 : * If \p should_free_host is true, and we cannot free for above reason, it will error.
494 : */
495 : void moveToDevice(bool should_free_host = true);
496 : /**
497 : * Copy data from device to host and deallocate device
498 : * @param should_free_device Whether the device memory should be freed.
499 : * Device memory cannot be freed when there are shallow copies of this array that are still alive.
500 : * If \p should_free_device is true, and we cannot free for above reason, it will error.
501 : */
502 : void moveToHost(bool should_free_device = true);
503 : /**
504 : * Deep copy another Kokkos array
505 : * If ArrayDeepCopy<T>::value is true, it will copy-construct each entry
506 : * If ArrayDeepCopy<T>::value is false, it will do a memory copy
507 : * @param array The Kokkos array to be deep copied
508 : */
509 : void deepCopy(const ArrayBase<T, dimension, index_type> & array);
510 : /**
511 : * Swap with another Kokkos array
512 : * @param array The Kokkos array to be swapped
513 : */
514 : void swap(ArrayBase<T, dimension, index_type> & array);
515 :
516 : /**
517 : * Assign a scalar value uniformly
518 : * @param scalar The scalar value to be assigned
519 : */
520 : auto & operator=(const T & scalar);
521 :
522 : /**
523 : * Array iterator
524 : */
525 : class iterator
526 : {
527 : public:
528 : using iterator_category = std::forward_iterator_tag;
529 : using value_type = T;
530 : using difference_type = std::ptrdiff_t;
531 : using pointer = T *;
532 : using reference = T &;
533 :
534 : KOKKOS_FUNCTION iterator() : it(nullptr) {}
535 4778399 : KOKKOS_FUNCTION explicit iterator(T * p) : it(p) {}
536 :
537 2591041 : KOKKOS_FUNCTION reference operator*() const { return *it; }
538 : KOKKOS_FUNCTION pointer operator->() const { return it; }
539 1122876 : KOKKOS_FUNCTION pointer operator&() const { return it; }
540 2590234 : KOKKOS_FUNCTION iterator & operator++()
541 : {
542 2590234 : ++it;
543 2590234 : return *this;
544 : }
545 807 : KOKKOS_FUNCTION iterator operator++(int)
546 : {
547 807 : iterator pre = *this;
548 807 : ++it;
549 807 : return pre;
550 : }
551 : KOKKOS_FUNCTION friend bool operator==(const iterator & a, const iterator & b)
552 : {
553 : return a.it == b.it;
554 : }
555 4417959 : KOKKOS_FUNCTION friend bool operator!=(const iterator & a, const iterator & b)
556 : {
557 4417959 : return a.it != b.it;
558 : }
559 :
560 : private:
561 : pointer it;
562 : };
563 :
564 : /**
565 : * Get the beginning iterator
566 : * @returns The beginning iterator
567 : */
568 2389236 : KOKKOS_FUNCTION iterator begin() const
569 : {
570 2389236 : KOKKOS_IF_ON_HOST(return iterator(_host_data);)
571 :
572 557590 : return iterator(_device_data);
573 : }
574 : /**
575 : * Get the end iterator
576 : * @returns The end iterator
577 : */
578 2389163 : KOKKOS_FUNCTION iterator end() const
579 : {
580 2389163 : KOKKOS_IF_ON_HOST(return iterator(_host_data + _size);)
581 :
582 557590 : return iterator(_device_data + _size);
583 : }
584 : #endif
585 :
586 : protected:
587 : /**
588 : * Size of each dimension
589 : */
590 : index_type _n[dimension] = {0};
591 : /**
592 : * Stride of each dimension
593 : */
594 : index_type _s[dimension] = {0};
595 : /**
596 : * Offset of each dimension
597 : */
598 : signed_index_type _d[dimension] = {0};
599 :
600 : #ifdef MOOSE_KOKKOS_SCOPE
601 : /**
602 : * The internal method to initialize and allocate this array
603 : * @tparam host Whether host data will be allocated
604 : * @tparam device Whether device data will be allocated
605 : * @tparam initialize Whether to initialize host data (calls default constructor)
606 : * @param n The size of each dimension
607 : */
608 : template <bool host, bool device, bool initialize, typename... size_type>
609 : void createInternal(size_type... n);
610 : /**
611 : * The internal method to initialize and allocate this array
612 : * @tparam host Whether host data will be allocated
613 : * @tparam device Whether device data will be allocated
614 : * @tparam initialize Whether to initialize host data (calls default constructor)
615 : * @param n The vector containing the size of each dimension
616 : */
617 : template <bool host, bool device, bool initialize>
618 : void createInternal(const std::vector<index_type> & n);
619 : /**
620 : * The internal method to initialize and allocate this array
621 : * @tparam initialize Whether to initialize host data (calls default constructor)
622 : * @param n The vector containing the size of each dimension
623 : * @param host The flag whether host data will be allocated
624 : * @param device The flag whether device data will be allocated
625 : */
626 : template <bool initialize>
627 : void createInternal(const std::vector<index_type> & n, bool host, bool device);
628 : /**
629 : * The internal method to perform a memory copy
630 : * @tparam TargetSpace The Kokkos memory space of target data
631 : * @tparam Sourcespace The Kokkos memory space of source data
632 : * @param target The pointer to the target data
633 : * @param source The pointer to the source data
634 : * @param n The number of entries to copy
635 : */
636 : template <typename TargetSpace, typename SourceSpace>
637 : void copyInternal(T * target, const T * source, index_type n);
638 : /**
639 : * Placement-new construct slot i from args, recording initialization
640 : * @param i The dimensionless slot index
641 : * @param args Arguments forwarded to T's constructor
642 : * @returns Reference to the constructed element
643 : */
644 : template <typename... Args>
645 : T & emplaceAt(index_type i, Args &&... args);
646 : #endif
647 :
648 : private:
649 : #ifdef MOOSE_KOKKOS_SCOPE
650 : /**
651 : * Allocate host data for an initialized array that has not allocated host data
652 : * @tparam initialize Whether to initialize host data (calls default constructor)
653 : */
654 : template <bool initialize>
655 : void allocHost();
656 : /**
657 : * Allocate device data for an initialized array that has not allocated device data
658 : */
659 : void allocDevice();
660 : #endif
661 : /**
662 : * Free host data
663 : */
664 : void freeHost();
665 : /**
666 : * Free device data
667 : */
668 : void freeDevice();
669 :
670 : /**
671 : * Reference counter
672 : */
673 : std::shared_ptr<unsigned int> _counter;
674 : /**
675 : * Non-null only for malloc-allocated arrays; tracks which slots have been emplace-constructed
676 : */
677 : std::shared_ptr<std::vector<bool>> _slots_constructed;
678 : /**
679 : * Flag whether array was initialized
680 : */
681 : bool _is_init = false;
682 : /**
683 : * Flag whether host data was allocated
684 : */
685 : bool _is_host_alloc = false;
686 : /**
687 : * Flag whether device data was allocated
688 : */
689 : bool _is_device_alloc = false;
690 : /**
691 : * Flag whether the host data points to an external data
692 : */
693 : bool _is_host_alias = false;
694 : /**
695 : * Flag whether the device data points to an external data
696 : */
697 : bool _is_device_alias = false;
698 : /**
699 : * Host data
700 : */
701 : T * _host_data = nullptr;
702 : /**
703 : * Device data
704 : */
705 : T * _device_data = nullptr;
706 : /**
707 : * Total size
708 : */
709 : index_type _size = 0;
710 : /**
711 : * Memory layout type
712 : */
713 : const LayoutType _layout;
714 : };
715 :
716 : template <typename T, unsigned int dimension, typename index_type>
717 : template <bool is_const>
718 577724 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::
719 407634 : constructed_entry_iterator(
720 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<
721 : is_const>::array_type & array,
722 407634 : index_type i)
723 577724 : : _array(array), _i(i)
724 : {
725 985358 : advanceToConstructed();
726 985358 : }
727 :
728 : template <typename T, unsigned int dimension, typename index_type>
729 : template <bool is_const>
730 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<is_const> &
731 676481 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator++()
732 : {
733 676481 : ++_i;
734 676481 : advanceToConstructed();
735 676481 : return *this;
736 : }
737 :
738 : template <typename T, unsigned int dimension, typename index_type>
739 : template <bool is_const>
740 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<is_const>
741 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator++(int)
742 : {
743 : constructed_entry_iterator pre = *this;
744 : ++(*this);
745 : return pre;
746 : }
747 :
748 : template <typename T, unsigned int dimension, typename index_type>
749 : template <bool is_const>
750 : bool
751 1169160 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator==(
752 : const constructed_entry_iterator & other) const
753 : {
754 1169160 : return &_array == &other._array && _i == other._i;
755 : }
756 :
757 : template <typename T, unsigned int dimension, typename index_type>
758 : template <bool is_const>
759 : bool
760 1169160 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator!=(
761 : const constructed_entry_iterator & other) const
762 : {
763 1169160 : return !(*this == other);
764 : }
765 :
766 : template <typename T, unsigned int dimension, typename index_type>
767 : template <bool is_const>
768 : void
769 1661839 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::advanceToConstructed()
770 : {
771 1975040 : while (_i < _array._size && !_array.isSlotConstructed(_i))
772 313201 : ++_i;
773 1661839 : }
774 :
775 : template <typename T, unsigned int dimension, typename index_type>
776 : template <bool is_const>
777 492679 : ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::constructed_entry_range(
778 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<
779 203817 : is_const>::array_type & array)
780 288862 : : _array(array)
781 : {
782 492679 : }
783 :
784 : template <typename T, unsigned int dimension, typename index_type>
785 : template <bool is_const>
786 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<is_const>::iterator
787 492679 : ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::begin() const
788 : {
789 492679 : return iterator(_array, 0);
790 : }
791 :
792 : template <typename T, unsigned int dimension, typename index_type>
793 : template <bool is_const>
794 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<is_const>::iterator
795 492679 : ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::end() const
796 : {
797 492679 : return iterator(_array, _array._size);
798 : }
799 :
800 : template <typename T, unsigned int dimension, typename index_type>
801 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<false>
802 492675 : ArrayBase<T, dimension, index_type>::constructedEntries()
803 : {
804 492675 : return constructed_entry_range<false>(*this);
805 : }
806 :
807 : template <typename T, unsigned int dimension, typename index_type>
808 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<true>
809 4 : ArrayBase<T, dimension, index_type>::constructedEntries() const
810 : {
811 4 : return constructed_entry_range<true>(*this);
812 : }
813 :
814 : template <typename T, unsigned int dimension, typename index_type>
815 : bool
816 8943799 : ArrayBase<T, dimension, index_type>::isSlotConstructed(index_type i) const
817 : {
818 : mooseAssert(i < _size, "isSlotConstructed index out of bounds");
819 :
820 8943799 : if (!_slots_constructed)
821 12 : return _is_host_alloc;
822 :
823 8943787 : return (*_slots_constructed)[i];
824 : }
825 :
826 : #ifdef MOOSE_KOKKOS_SCOPE
827 : template <typename T, unsigned int dimension, typename index_type>
828 : template <typename... Args>
829 : T &
830 102802 : ArrayBase<T, dimension, index_type>::emplaceAt(index_type i, Args &&... args)
831 : {
832 : mooseAssert(_is_host_alloc && _slots_constructed,
833 : "emplaceAt requires a malloc-allocated array (create<false>)");
834 : mooseAssert(i < _size, "emplaceAt index out of bounds");
835 :
836 102802 : new (_host_data + i) T(std::forward<Args>(args)...);
837 102802 : (*_slots_constructed)[i] = true;
838 :
839 102802 : return _host_data[i];
840 : }
841 : #endif // MOOSE_KOKKOS_SCOPE
842 :
843 : template <typename T, unsigned int dimension, typename index_type>
844 : void
845 2128154 : ArrayBase<T, dimension, index_type>::freeHost()
846 : {
847 2128154 : if (!_is_host_alloc)
848 62900 : return;
849 :
850 2065254 : if (_is_host_alias)
851 : {
852 11317 : _host_data = nullptr;
853 11317 : _is_host_alias = false;
854 : }
855 : else
856 : {
857 2053937 : if (!_slots_constructed)
858 : // Allocated by new
859 1262502 : delete[] _host_data;
860 : else
861 : {
862 : // Allocated by malloc
863 9429975 : for (const auto i : make_range(_size))
864 7954085 : if (isSlotConstructed(i))
865 102778 : _host_data[i].~T();
866 :
867 1475890 : std::free(_host_data);
868 :
869 1475890 : _slots_constructed.reset();
870 : }
871 : }
872 :
873 2065254 : _is_host_alloc = false;
874 : }
875 :
876 : template <typename T, unsigned int dimension, typename index_type>
877 : void
878 2103795 : ArrayBase<T, dimension, index_type>::freeDevice()
879 : {
880 2103795 : if (!_is_device_alloc)
881 24593 : return;
882 :
883 2079202 : if (_is_device_alias)
884 : {
885 92 : _device_data = nullptr;
886 92 : _is_device_alias = false;
887 : }
888 : else
889 2079110 : Moose::Kokkos::free(_device_data);
890 :
891 2079202 : _is_device_alloc = false;
892 : }
893 :
894 : template <typename T, unsigned int dimension, typename index_type>
895 : void
896 130241298 : ArrayBase<T, dimension, index_type>::destroy()
897 : {
898 130241298 : if (!_counter)
899 86171747 : return;
900 :
901 44069551 : if (_counter.use_count() > 1)
902 : {
903 41965756 : _host_data = nullptr;
904 41965756 : _device_data = nullptr;
905 : }
906 2103795 : else if (_counter.use_count() == 1)
907 : {
908 2103795 : freeHost();
909 2103795 : freeDevice();
910 : }
911 :
912 44069551 : _size = 0;
913 :
914 108200674 : for (unsigned int i = 0; i < dimension; ++i)
915 : {
916 64131123 : _n[i] = 0;
917 64131123 : _s[i] = 0;
918 64131123 : _d[i] = 0;
919 : }
920 :
921 44069551 : _is_init = false;
922 44069551 : _is_host_alloc = false;
923 44069551 : _is_device_alloc = false;
924 44069551 : _is_host_alias = false;
925 44069551 : _is_device_alias = false;
926 :
927 44069551 : _counter.reset();
928 44069551 : _slots_constructed.reset();
929 : }
930 :
931 : template <typename T, unsigned int dimension, typename index_type>
932 : void
933 61605120 : ArrayBase<T, dimension, index_type>::shallowCopy(const ArrayBase<T, dimension, index_type> & array)
934 : {
935 61605120 : if (_layout != array._layout)
936 0 : mooseError("Kokkos array error: cannot shallow copy arrays with different layouts.");
937 :
938 61605120 : destroy();
939 :
940 61605120 : _counter = array._counter;
941 61605120 : _slots_constructed = array._slots_constructed;
942 :
943 61605120 : _size = array._size;
944 :
945 154546701 : for (unsigned int i = 0; i < dimension; ++i)
946 : {
947 92941581 : _n[i] = array._n[i];
948 92941581 : _s[i] = array._s[i];
949 92941581 : _d[i] = array._d[i];
950 : }
951 :
952 61605120 : _is_init = array._is_init;
953 61605120 : _is_host_alloc = array._is_host_alloc;
954 61605120 : _is_device_alloc = array._is_device_alloc;
955 61605120 : _is_host_alias = array._is_host_alias;
956 61605120 : _is_device_alias = array._is_device_alias;
957 :
958 61605120 : _host_data = array._host_data;
959 61605120 : _device_data = array._device_data;
960 61605120 : }
961 :
962 : #ifdef MOOSE_KOKKOS_SCOPE
963 : template <typename T, unsigned int dimension, typename index_type>
964 : void
965 934629 : ArrayBase<T, dimension, index_type>::aliasHost(T * ptr)
966 : {
967 934629 : if (!_is_init)
968 0 : mooseError("Kokkos array error: attempted to alias host data before array initialization.");
969 :
970 934629 : if (_is_host_alloc && !_is_host_alias)
971 0 : mooseError("Kokkos array error: cannot alias host data because host data was not aliased.");
972 :
973 934629 : _host_data = ptr;
974 934629 : _is_host_alloc = true;
975 934629 : _is_host_alias = true;
976 934629 : }
977 :
978 : template <typename T, unsigned int dimension, typename index_type>
979 : void
980 428 : ArrayBase<T, dimension, index_type>::aliasDevice(T * ptr)
981 : {
982 428 : if (!_is_init)
983 0 : mooseError("Kokkos array error: attempted to alias device data before array initialization.");
984 :
985 428 : if (_is_device_alloc && !_is_device_alias)
986 0 : mooseError("Kokkos array error: cannot alias device data because device data was not aliased.");
987 :
988 428 : _device_data = ptr;
989 428 : _is_device_alloc = true;
990 428 : _is_device_alias = true;
991 428 : }
992 :
993 : template <typename T, unsigned int dimension, typename index_type>
994 : template <bool initialize>
995 : void
996 2057315 : ArrayBase<T, dimension, index_type>::allocHost()
997 : {
998 2057315 : if (_is_host_alloc)
999 0 : return;
1000 :
1001 : if constexpr (initialize)
1002 : {
1003 : static_assert(
1004 : std::is_default_constructible<T>::value,
1005 : "Data type is not default-constructible. Initialization argument should be set to false.");
1006 :
1007 3445555 : _host_data = new T[_size];
1008 : }
1009 : else
1010 : {
1011 1476330 : _host_data = static_cast<T *>(std::malloc(_size * sizeof(T)));
1012 1476330 : _slots_constructed = std::make_shared<std::vector<bool>>(_size, false);
1013 : }
1014 :
1015 2057315 : _is_host_alloc = true;
1016 : }
1017 :
1018 : template <typename T, unsigned int dimension, typename index_type>
1019 : void
1020 2082716 : ArrayBase<T, dimension, index_type>::allocDevice()
1021 : {
1022 2082716 : if (_is_device_alloc)
1023 0 : return;
1024 :
1025 2082716 : _device_data =
1026 1197867 : static_cast<T *>(::Kokkos::kokkos_malloc<ExecSpace::memory_space>(_size * sizeof(T)));
1027 :
1028 2082716 : _is_device_alloc = true;
1029 : }
1030 :
1031 : template <typename T, unsigned int dimension, typename index_type>
1032 : template <bool host, bool device, bool initialize>
1033 : void
1034 2107607 : ArrayBase<T, dimension, index_type>::createInternal(const std::vector<index_type> & n)
1035 : {
1036 2107607 : if (n.size() != dimension)
1037 0 : mooseError("Kokkos array error: the number of dimensions provided (",
1038 0 : n.size(),
1039 : ") must match the array dimension (",
1040 : dimension,
1041 : ").");
1042 :
1043 2107607 : if (_counter)
1044 1419 : destroy();
1045 :
1046 2107607 : _counter = std::make_shared<unsigned int>();
1047 :
1048 2107607 : uint64_t overflow_checker = 1;
1049 :
1050 2107607 : _size = 1;
1051 2107607 : _s[0] = 1;
1052 :
1053 4417816 : for (const auto i : make_range(dimension))
1054 : {
1055 2310209 : overflow_checker *= n[i];
1056 :
1057 2310209 : _n[i] = n[i];
1058 2310209 : _size *= n[i];
1059 : }
1060 :
1061 2107607 : if (overflow_checker > std::numeric_limits<index_type>::max())
1062 0 : mooseError("Kokkos array error: the dimensions provided (",
1063 : Moose::stringify(n),
1064 : ") has the total size of ",
1065 : overflow_checker,
1066 : " which exceeds the limit of ",
1067 : MooseUtils::prettyCppType<index_type>(),
1068 : ".");
1069 :
1070 2107607 : if (_layout == LayoutType::LEFT)
1071 : {
1072 2107597 : _s[0] = 1;
1073 :
1074 2310177 : for (unsigned int i = 1; i < dimension; ++i)
1075 202580 : _s[i] = _s[i - 1] * _n[i - 1];
1076 : }
1077 : else
1078 : {
1079 10 : _s[dimension - 1] = 1;
1080 :
1081 32 : for (int i = dimension - 2; i >= 0; --i)
1082 22 : _s[i] = _s[i + 1] * _n[i + 1];
1083 : }
1084 :
1085 : if constexpr (host)
1086 2057315 : allocHost<initialize>();
1087 :
1088 : if constexpr (device)
1089 2082716 : allocDevice();
1090 :
1091 2107607 : _is_init = true;
1092 2107607 : }
1093 :
1094 : template <typename T, unsigned int dimension, typename index_type>
1095 : template <bool initialize>
1096 : void
1097 93707 : ArrayBase<T, dimension, index_type>::createInternal(const std::vector<index_type> & n,
1098 : bool host,
1099 : bool device)
1100 : {
1101 93707 : if (host && device)
1102 92288 : createInternal<true, true, initialize>(n);
1103 1419 : else if (host && !device)
1104 0 : createInternal<true, false, initialize>(n);
1105 1419 : else if (!host && device)
1106 1419 : createInternal<false, true, initialize>(n);
1107 : else
1108 0 : createInternal<false, false, initialize>(n);
1109 93707 : }
1110 :
1111 : template <typename T, unsigned int dimension, typename index_type>
1112 : template <bool host, bool device, bool initialize, typename... size_type>
1113 : void
1114 668597 : ArrayBase<T, dimension, index_type>::createInternal(size_type... n)
1115 : {
1116 : static_assert((std::is_convertible<size_type, index_type>::value && ...),
1117 : "All arguments must be convertible to index_type");
1118 : static_assert(sizeof...(n) == dimension, "Number of arguments should match array dimension");
1119 :
1120 668597 : std::vector<index_type> dims;
1121 668597 : (dims.push_back(n), ...);
1122 :
1123 668597 : createInternal<host, device, initialize>(dims);
1124 668597 : }
1125 :
1126 : template <typename T, unsigned int dimension, typename index_type>
1127 : template <typename TargetSpace, typename SourceSpace>
1128 : void
1129 5449375 : ArrayBase<T, dimension, index_type>::copyInternal(T * target, const T * source, index_type n)
1130 : {
1131 5449375 : ::Kokkos::Impl::DeepCopy<TargetSpace, SourceSpace>(target, source, n * sizeof(T));
1132 5449375 : ::Kokkos::fence();
1133 5449375 : }
1134 :
1135 : template <typename T, unsigned int dimension, typename index_type>
1136 : void
1137 4654 : ArrayBase<T, dimension, index_type>::offset(const std::vector<signed_index_type> & d)
1138 : {
1139 4654 : if (d.size() > dimension)
1140 0 : mooseError("Kokkos array error: the number of offsets provided (",
1141 0 : d.size(),
1142 : ") cannot be larger than the array dimension (",
1143 : dimension,
1144 : ").");
1145 :
1146 9310 : for (const auto i : index_range(d))
1147 4656 : _d[i] = d[i];
1148 4654 : }
1149 :
1150 : template <typename T, unsigned int dimension, typename index_type>
1151 : template <typename... offset_type>
1152 : void
1153 4654 : ArrayBase<T, dimension, index_type>::offset(offset_type... d)
1154 : {
1155 : static_assert((std::is_convertible<offset_type, signed_index_type>::value && ...),
1156 : "All arguments must be convertible to signed_index_type");
1157 : static_assert(sizeof...(d) == dimension, "Number of arguments should match array dimension");
1158 :
1159 4654 : std::vector<signed_index_type> offsets;
1160 4654 : (offsets.push_back(d), ...);
1161 :
1162 4654 : offset(offsets);
1163 4654 : }
1164 :
1165 : template <typename T, unsigned int dimension, typename index_type>
1166 : void
1167 3376575 : ArrayBase<T, dimension, index_type>::copyToDevice()
1168 : {
1169 : // If host side memory is not allocated, do nothing
1170 3376575 : if (!_is_host_alloc)
1171 67303 : return;
1172 :
1173 : // If device side memory is not allocated,
1174 3309272 : if (!_is_device_alloc)
1175 : {
1176 0 : if (_counter.use_count() == 1)
1177 : // allocate memory if this array is not shared with other arrays
1178 0 : allocDevice();
1179 : else
1180 : // print error if this array is shared with other arrays
1181 0 : mooseError("Kokkos array error: cannot copy from host to device because device memory "
1182 : "was not allocated. Cannot allocate device memory for copy because the array is "
1183 : "being shared.");
1184 : }
1185 :
1186 : // Copy from host to device
1187 3309272 : copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data, _host_data, _size);
1188 : }
1189 :
1190 : template <typename T, unsigned int dimension, typename index_type>
1191 : void
1192 797106 : ArrayBase<T, dimension, index_type>::copyToHost()
1193 : {
1194 : // If device side memory is not allocated, do nothing
1195 797106 : if (!_is_device_alloc)
1196 0 : return;
1197 :
1198 : // If host side memory is not allocated,
1199 797106 : if (!_is_host_alloc)
1200 : {
1201 0 : if (_counter.use_count() == 1)
1202 : // allocate memory if this array is not shared with other arrays
1203 0 : allocHost<false>();
1204 : else
1205 : // print error if this array is shared with other arrays
1206 0 : mooseError("Kokkos array error: cannot copy from device to host because host memory "
1207 : "was not allocated. Cannot allocate host memory for copy because the array is "
1208 : "being shared.");
1209 : }
1210 :
1211 : // Copy from device to host
1212 797106 : copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data, _device_data, _size);
1213 : }
1214 :
1215 : template <typename T, unsigned int dimension, typename index_type>
1216 : void
1217 24359 : ArrayBase<T, dimension, index_type>::moveToDevice(bool should_free_host)
1218 : {
1219 : static_assert(!is_kokkos_array<T>::value,
1220 : "moveToDevice() not allowed for a nested array whose data type is another array.");
1221 :
1222 24359 : if (should_free_host && _counter.use_count() > 1)
1223 0 : mooseError("Kokkos array error: cannot move array from host to device because there is at "
1224 : "least one shallow copy of this array still alive.");
1225 :
1226 24359 : copyToDevice();
1227 :
1228 24359 : if (_counter.use_count() == 1)
1229 24359 : freeHost();
1230 24359 : }
1231 :
1232 : template <typename T, unsigned int dimension, typename index_type>
1233 : void
1234 : ArrayBase<T, dimension, index_type>::moveToHost(bool should_free_device)
1235 : {
1236 : if (should_free_device && _counter.use_count() > 1)
1237 : mooseError("Kokkos array error: cannot move array from device to host because there is at "
1238 : "least one shallow copy of this array still alive.");
1239 :
1240 : copyToHost();
1241 :
1242 : if (_counter.use_count() == 1)
1243 : freeDevice();
1244 : }
1245 :
1246 : template <typename T, unsigned int dimension, typename index_type>
1247 : void
1248 256 : ArrayBase<T, dimension, index_type>::copyIn(const T * ptr,
1249 : MemcpyType dir,
1250 : index_type n,
1251 : index_type offset)
1252 : {
1253 256 : if (n > _size)
1254 0 : mooseError("Kokkos array error: cannot copy in data larger than the array size.");
1255 :
1256 256 : if (offset > _size)
1257 0 : mooseError("Kokkos array error: offset cannot be larger than the array size.");
1258 :
1259 256 : if (dir == MemcpyType::HOST_TO_HOST)
1260 : {
1261 : // If host side memory is not allocated, print error
1262 0 : if (!_is_host_alloc)
1263 0 : mooseError(
1264 : "Kokkos array error: cannot copy in to the array because host memory was not allocated.");
1265 :
1266 : // Copy from host to host
1267 0 : copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(_host_data + offset, ptr, n);
1268 : }
1269 256 : else if (dir == MemcpyType::HOST_TO_DEVICE)
1270 : {
1271 : // If device side memory is not allocated, print error
1272 256 : if (!_is_device_alloc)
1273 0 : mooseError("Kokkos array error: cannot copy in to the array because device memory was not "
1274 : "allocated.");
1275 :
1276 : // Copy from host to device
1277 256 : copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data + offset, ptr, n);
1278 : }
1279 0 : else if (dir == MemcpyType::DEVICE_TO_HOST)
1280 : {
1281 : // If host side memory is not allocated, print error
1282 0 : if (!_is_host_alloc)
1283 0 : mooseError(
1284 : "Kokkos array error: cannot copy in to the array because host memory was not allocated.");
1285 :
1286 : // Copy from device to host
1287 0 : copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data + offset, ptr, n);
1288 : }
1289 0 : else if (dir == MemcpyType::DEVICE_TO_DEVICE)
1290 : {
1291 : // If device side memory is not allocated, print error
1292 0 : if (!_is_device_alloc)
1293 0 : mooseError("Kokkos array error: cannot copy in to the array because device memory was not "
1294 : "allocated.");
1295 :
1296 : // Copy from device to device
1297 0 : copyInternal<MemSpace, MemSpace>(_device_data + offset, ptr, n);
1298 : }
1299 256 : }
1300 :
1301 : template <typename T, unsigned int dimension, typename index_type>
1302 : void
1303 1139 : ArrayBase<T, dimension, index_type>::copyOut(T * ptr,
1304 : MemcpyType dir,
1305 : index_type n,
1306 : index_type offset)
1307 : {
1308 1139 : if (n > _size)
1309 0 : mooseError("Kokkos array error: cannot copy out data larger than the array size.");
1310 :
1311 1139 : if (offset > _size)
1312 0 : mooseError("Kokkos array error: offset cannot be larger than the array size.");
1313 :
1314 1139 : if (dir == MemcpyType::HOST_TO_HOST)
1315 : {
1316 : // If host side memory is not allocated, print error
1317 0 : if (!_is_host_alloc)
1318 0 : mooseError("Kokkos array error: cannot copy out from the array because host memory was not "
1319 : "allocated.");
1320 :
1321 : // Copy from host to host
1322 0 : copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
1323 : }
1324 1139 : else if (dir == MemcpyType::HOST_TO_DEVICE)
1325 : {
1326 : // If host side memory is not allocated, print error
1327 0 : if (!_is_host_alloc)
1328 0 : mooseError("Kokkos array error: cannot copy out from the array because host memory was not "
1329 : "allocated.");
1330 :
1331 : // Copy from host to device
1332 0 : copyInternal<MemSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
1333 : }
1334 1139 : else if (dir == MemcpyType::DEVICE_TO_HOST)
1335 : {
1336 : // If device side memory is not allocated, print error
1337 1139 : if (!_is_device_alloc)
1338 0 : mooseError("Kokkos array error: cannot copy out from the array because device memory was not "
1339 : "allocated.");
1340 :
1341 : // Copy from device to host
1342 1139 : copyInternal<::Kokkos::HostSpace, MemSpace>(ptr, _device_data + offset, n);
1343 : }
1344 0 : else if (dir == MemcpyType::DEVICE_TO_DEVICE)
1345 : {
1346 : // If device side memory is not allocated, print error
1347 0 : if (!_is_device_alloc)
1348 0 : mooseError("Kokkos array error: cannot copy out from the array because device memory was not "
1349 : "allocated.");
1350 :
1351 : // Copy from device to device
1352 0 : copyInternal<MemSpace, MemSpace>(ptr, _device_data + offset, n);
1353 : }
1354 1139 : }
1355 :
1356 : template <typename T>
1357 : void
1358 11451924 : copyToDeviceInner(T & /* data */)
1359 : {
1360 11451924 : }
1361 :
1362 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1363 : void
1364 266740 : copyToDeviceInner(Array<T, dimension, index_type, layout> & data)
1365 : {
1366 266740 : data.copyToDeviceNested();
1367 266740 : }
1368 :
1369 : template <typename T, unsigned int dimension, typename index_type>
1370 : void
1371 328152 : ArrayBase<T, dimension, index_type>::copyToDeviceNested()
1372 : {
1373 12046816 : for (index_type i = 0; i < _size; ++i)
1374 11718664 : copyToDeviceInner(_host_data[i]);
1375 :
1376 328152 : copyToDevice();
1377 328152 : }
1378 :
1379 : template <typename T, unsigned int dimension, typename index_type>
1380 : void
1381 93707 : ArrayBase<T, dimension, index_type>::deepCopy(const ArrayBase<T, dimension, index_type> & array)
1382 : {
1383 93707 : if (_layout != array._layout)
1384 0 : mooseError("Kokkos array error: cannot deep copy arrays with different layouts.");
1385 :
1386 92288 : if (ArrayDeepCopy<T>::value && !array._is_host_alloc)
1387 0 : mooseError(
1388 : "Kokkos array error: cannot deep copy using constructor from array without host data.");
1389 :
1390 281121 : std::vector<index_type> n(std::begin(array._n), std::end(array._n));
1391 :
1392 93707 : createInternal<false>(n, array._is_host_alloc, array._is_device_alloc);
1393 :
1394 : if constexpr (ArrayDeepCopy<T>::value)
1395 : {
1396 186357 : for (index_type i = 0; i < _size; ++i)
1397 94069 : emplaceAt(i, array._host_data[i]);
1398 :
1399 92288 : copyToDevice();
1400 : }
1401 : else
1402 : {
1403 1419 : if (_is_host_alloc)
1404 0 : std::memcpy(_host_data, array._host_data, _size * sizeof(T));
1405 :
1406 1419 : if (_is_device_alloc)
1407 1419 : copyInternal<MemSpace, MemSpace>(_device_data, array._device_data, _size);
1408 : }
1409 :
1410 187414 : for (unsigned int i = 0; i < dimension; ++i)
1411 : {
1412 93707 : _d[i] = array._d[i];
1413 93707 : _s[i] = array._s[i];
1414 : }
1415 93707 : }
1416 :
1417 : template <typename T, unsigned int dimension, typename index_type>
1418 : void
1419 2808 : ArrayBase<T, dimension, index_type>::swap(ArrayBase<T, dimension, index_type> & array)
1420 : {
1421 2808 : ArrayBase<T, dimension, index_type> clone(_layout);
1422 :
1423 2808 : clone.shallowCopy(*this);
1424 2808 : this->shallowCopy(array);
1425 2808 : array.shallowCopy(clone);
1426 2808 : }
1427 :
1428 : template <typename T, unsigned int dimension, typename index_type>
1429 : auto &
1430 596092 : ArrayBase<T, dimension, index_type>::operator=(const T & scalar)
1431 : {
1432 596092 : if (_is_host_alloc)
1433 596078 : std::fill_n(_host_data, _size, scalar);
1434 :
1435 596092 : if (_is_device_alloc)
1436 596092 : ::Kokkos::Experimental::fill_n(ExecSpace(), deviceView(), _size, scalar);
1437 :
1438 596092 : return *this;
1439 : }
1440 :
1441 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1442 : void
1443 1139 : dataStore(std::ostream & stream, Array<T, dimension, index_type, layout> & array, void * context)
1444 : {
1445 : using ::dataStore;
1446 :
1447 1139 : bool is_alloc = array.isAlloc();
1448 1139 : dataStore(stream, is_alloc, nullptr);
1449 :
1450 1139 : if (!is_alloc)
1451 0 : return;
1452 :
1453 1139 : std::string type = typeid(T).name();
1454 1139 : dataStore(stream, type, nullptr);
1455 :
1456 1139 : unsigned int dim = dimension;
1457 1139 : dataStore(stream, dim, nullptr);
1458 :
1459 2278 : for (unsigned int dim = 0; dim < dimension; ++dim)
1460 : {
1461 1139 : auto n = array.n(dim);
1462 1139 : dataStore(stream, n, nullptr);
1463 : }
1464 :
1465 1139 : if (array.isDeviceAlloc())
1466 : {
1467 : // We use malloc/free because we just want a memory copy
1468 : // If T is a Kokkos array and we use new/delete or vector to copy it out,
1469 : // the arrays will be destroyed on cleanup
1470 :
1471 1139 : T * data = static_cast<T *>(std::malloc(array.size() * sizeof(T)));
1472 :
1473 1139 : array.copyOut(data, MemcpyType::DEVICE_TO_HOST, array.size());
1474 :
1475 117441 : for (index_type i = 0; i < array.size(); ++i)
1476 116302 : dataStore(stream, data[i], context);
1477 :
1478 1139 : std::free(data);
1479 : }
1480 : else
1481 0 : for (auto & value : array)
1482 0 : dataStore(stream, value, context);
1483 1139 : }
1484 :
1485 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1486 : void
1487 531 : dataLoad(std::istream & stream, Array<T, dimension, index_type, layout> & array, void * context)
1488 : {
1489 : using ::dataLoad;
1490 :
1491 : bool is_alloc;
1492 531 : dataLoad(stream, is_alloc, nullptr);
1493 :
1494 531 : if (!is_alloc)
1495 0 : return;
1496 :
1497 531 : std::string from_type_name;
1498 531 : dataLoad(stream, from_type_name, nullptr);
1499 :
1500 531 : if (from_type_name != typeid(T).name())
1501 0 : mooseError("Kokkos array error: cannot load array because the stored array is of type '",
1502 : MooseUtils::prettyCppType(libMesh::demangle(from_type_name.c_str())),
1503 : "' but the loading array is of type '",
1504 : MooseUtils::prettyCppType(libMesh::demangle(typeid(T).name())),
1505 : "'.");
1506 :
1507 : unsigned int from_dimension;
1508 531 : dataLoad(stream, from_dimension, nullptr);
1509 :
1510 531 : if (from_dimension != dimension)
1511 0 : mooseError("Kokkos array error: cannot load array because the stored array is ",
1512 : from_dimension,
1513 : "D but the loading array is ",
1514 : dimension,
1515 : "D.");
1516 :
1517 1062 : std::vector<index_type> from_n(dimension);
1518 531 : std::vector<index_type> n(dimension);
1519 :
1520 1062 : for (unsigned int dim = 0; dim < dimension; ++dim)
1521 : {
1522 531 : dataLoad(stream, from_n[dim], nullptr);
1523 531 : n[dim] = array.n(dim);
1524 : }
1525 :
1526 531 : if (from_n != n)
1527 0 : mooseError("Kokkos array error: cannot load array because the stored array has dimensions (",
1528 : Moose::stringify(from_n),
1529 : ") but the loading array has dimensions (",
1530 : Moose::stringify(n),
1531 : ").");
1532 :
1533 531 : if (array.isHostAlloc())
1534 : {
1535 2706 : for (auto & value : array)
1536 2156 : dataLoad(stream, value, context);
1537 :
1538 275 : if (array.isDeviceAlloc())
1539 275 : array.copyToDevice();
1540 : }
1541 : else
1542 : {
1543 256 : std::vector<T> data(array.size());
1544 :
1545 51792 : for (auto & value : data)
1546 51536 : dataLoad(stream, value, context);
1547 :
1548 256 : array.copyIn(data.data(), MemcpyType::HOST_TO_DEVICE, array.size());
1549 256 : }
1550 531 : }
1551 : #endif
1552 :
1553 : /**
1554 : * The specialization of the Kokkos array class for each dimension.
1555 : * All array data that needs to be accessed on device in Kokkos objects should use this class.
1556 : * If the array is populated on host and is to be accessed on device, make sure to call
1557 : * copyToDevice() after populating data. For a nested Kokkos array, either copyToDeviceNested()
1558 : * should be called for the outermost array or copyToDevice() should be called for each instance of
1559 : * Kokkos array from the innermost to the outermost. Do not store this object as reference in your
1560 : * Kokkos object if it is used on device, because the reference refers to a host object and
1561 : * therefore is not accessible on device. If storing it as a reference is required, see
1562 : * ReferenceWrapper.
1563 : * @tparam T The data type
1564 : * @tparam dimension The array dimension size
1565 : * @tparam index_type The array index type
1566 : * @tparam layout The memory layout type
1567 : */
1568 : ///@{
1569 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1570 : class Array : public ArrayBase<T, dimension, index_type>
1571 : {
1572 : #ifdef MOOSE_KOKKOS_SCOPE
1573 : usingKokkosArrayBaseMembers(T, dimension, index_type);
1574 : #endif
1575 :
1576 : public:
1577 : /**
1578 : * Default constructor
1579 : */
1580 1836792 : Array() : ArrayBase<T, dimension, index_type>(layout) {}
1581 : /**
1582 : * Copy constructor
1583 : */
1584 28670403 : Array(const Array<T, dimension, index_type, layout> & array)
1585 16767477 : : ArrayBase<T, dimension, index_type>(array)
1586 : {
1587 28670403 : }
1588 : #ifdef MOOSE_KOKKOS_SCOPE
1589 : /**
1590 : * Constructor
1591 : * Initialize and allocate array with given dimensions
1592 : * This allocates both host and device data
1593 : * @param n The size of each dimension
1594 : */
1595 : template <typename... size_type>
1596 16 : Array(size_type... n) : ArrayBase<T, dimension, index_type>(layout, n...)
1597 : {
1598 16 : }
1599 : #endif
1600 :
1601 : /**
1602 : * Shallow copy another Kokkos array
1603 : * @param array The Kokkos array to be shallow copied
1604 : */
1605 : auto & operator=(const Array<T, dimension, index_type, layout> & array)
1606 : {
1607 : this->shallowCopy(array);
1608 :
1609 : return *this;
1610 : }
1611 :
1612 : #ifdef MOOSE_KOKKOS_SCOPE
1613 : /**
1614 : * Get an array entry
1615 : * @param i The index of each dimension
1616 : * @returns The reference of the entry depending on the architecture this function is being called
1617 : * on
1618 : */
1619 : template <typename... indices>
1620 1930892860 : KOKKOS_FUNCTION T & operator()(indices... i) const
1621 : {
1622 1930892860 : return this->operator[](linearIndex(i...));
1623 : }
1624 : /**
1625 : * Get an array entry using indices stored in an array
1626 : * @param idx The array storing the indices
1627 : * @returns The reference of the entry depending on the architecture this function is being
1628 : * called on
1629 : */
1630 : KOKKOS_FUNCTION T & operator()(const signed_index_type (&idx)[dimension]) const
1631 : {
1632 : return this->operator[](linearIndex(idx));
1633 : }
1634 : /**
1635 : * Placement-new construct an entry from args, recording initialization
1636 : * @param idx The array storing the indices
1637 : * @param args Arguments forwarded to T's constructor
1638 : * @returns Reference to the constructed element
1639 : */
1640 : template <typename indices, typename... Args>
1641 : T & emplace(const indices (&idx)[dimension], Args &&... args);
1642 : #endif
1643 :
1644 : private:
1645 : #ifdef MOOSE_KOKKOS_SCOPE
1646 : /**
1647 : * Get the dimensionless index for a multi-dimensional index
1648 : */
1649 : template <typename... indices>
1650 : KOKKOS_FUNCTION index_type linearIndex(indices... i) const;
1651 : /**
1652 : * Get the dimensionless index for indices stored in an array
1653 : */
1654 : ///@{
1655 : template <typename indices>
1656 8 : KOKKOS_FUNCTION index_type linearIndex(const indices (&idx)[dimension]) const
1657 : {
1658 8 : return linearIndexHelper(idx, std::make_integer_sequence<unsigned int, dimension>{});
1659 : }
1660 : template <typename indices, unsigned int... i>
1661 8 : KOKKOS_FUNCTION index_type linearIndexHelper(const indices (&idx)[dimension],
1662 : std::integer_sequence<unsigned int, i...>) const
1663 : {
1664 8 : return linearIndex(idx[i]...);
1665 : }
1666 : ///@}
1667 : #endif
1668 : };
1669 :
1670 : #ifdef MOOSE_KOKKOS_SCOPE
1671 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1672 : template <typename indices, typename... Args>
1673 : T &
1674 8 : Array<T, dimension, index_type, layout>::emplace(const indices (&idx)[dimension], Args &&... args)
1675 : {
1676 8 : return this->emplaceAt(linearIndex(idx), std::forward<Args>(args)...);
1677 : }
1678 :
1679 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1680 : template <typename... indices>
1681 : KOKKOS_FUNCTION index_type
1682 1930892868 : Array<T, dimension, index_type, layout>::linearIndex(indices... i) const
1683 : {
1684 : static_assert((std::is_convertible<indices, signed_index_type>::value && ...),
1685 : "All arguments must be convertible to signed_index_type");
1686 : static_assert(sizeof...(i) == dimension, "Number of arguments should match array dimension");
1687 :
1688 : #ifndef NDEBUG
1689 : {
1690 : signed_index_type idx[dimension] = {static_cast<signed_index_type>(i)...};
1691 :
1692 : for (unsigned int d = 0; d < sizeof...(i); ++d)
1693 : KOKKOS_ASSERT(idx[d] - _d[d] >= 0 && static_cast<index_type>(idx[d] - _d[d]) < _n[d]);
1694 : }
1695 : #endif
1696 :
1697 1930892868 : index_type idx = 0;
1698 1930892868 : unsigned int d = 0;
1699 :
1700 : if constexpr (layout == LayoutType::LEFT)
1701 2282956926 : (((idx += (d == 0 ? static_cast<signed_index_type>(i) - _d[d]
1702 2273942954 : : (static_cast<signed_index_type>(i) - _d[d]) * _s[d])),
1703 : ++d),
1704 : ...);
1705 : else
1706 14250 : (((idx += (d == dimension - 1 ? static_cast<signed_index_type>(i) - _d[d]
1707 11294 : : (static_cast<signed_index_type>(i) - _d[d]) * _s[d])),
1708 : ++d),
1709 : ...);
1710 :
1711 1930892868 : return idx;
1712 : }
1713 : #endif
1714 :
1715 : template <typename T, typename index_type>
1716 : class Array<T, 1, index_type, LayoutType::LEFT> : public ArrayBase<T, 1, index_type>
1717 : {
1718 : #ifdef MOOSE_KOKKOS_SCOPE
1719 : usingKokkosArrayBaseMembers(T, 1, index_type);
1720 : #endif
1721 :
1722 : public:
1723 : /**
1724 : * Default constructor
1725 : */
1726 6775666 : Array() : ArrayBase<T, 1, index_type>(LayoutType::LEFT) {}
1727 : /**
1728 : * Copy constructor
1729 : */
1730 30675796 : Array(const Array<T, 1, index_type, LayoutType::LEFT> & array)
1731 17995069 : : ArrayBase<T, 1, index_type>(array)
1732 : {
1733 30675796 : }
1734 : #ifdef MOOSE_KOKKOS_SCOPE
1735 : /**
1736 : * Constructor
1737 : * Initialize and allocate array with given size
1738 : * This allocates both host and device data
1739 : * @param n The array size
1740 : */
1741 60 : Array(index_type n) : ArrayBase<T, 1, index_type>(LayoutType::LEFT, n) {}
1742 : /**
1743 : * Constructor
1744 : * Initialize and allocate array by copying a standard vector variable
1745 : * This allocates and copies to both host and device data
1746 : * @param vector The standard vector variable to copy
1747 : */
1748 1628 : Array(const std::vector<T> & vector) : ArrayBase<T, 1, index_type>(LayoutType::LEFT)
1749 : {
1750 1628 : *this = vector;
1751 1628 : }
1752 : #endif
1753 :
1754 : /**
1755 : * Shallow copy another Kokkos array
1756 : * @param array The Kokkos array to be shallow copied
1757 : */
1758 2342785 : auto & operator=(const Array<T, 1, index_type, LayoutType::LEFT> & array)
1759 : {
1760 2342785 : this->shallowCopy(array);
1761 :
1762 2342785 : return *this;
1763 : }
1764 :
1765 : #ifdef MOOSE_KOKKOS_SCOPE
1766 : /**
1767 : * Copy a standard vector variable
1768 : * This re-initializes and re-allocates array with the size of the vector
1769 : * @tparam host Whether to allocate and copy to the host data
1770 : * @tparam device Whether to allocate and copy to the device data
1771 : * @param vector The standard vector variable to copy
1772 : */
1773 : template <bool host, bool device>
1774 1340183 : void copyVector(const std::vector<T> & vector)
1775 : {
1776 2690023 : this->template createInternal<host, device, false>({static_cast<index_type>(vector.size())});
1777 :
1778 : if (host)
1779 1330526 : std::memcpy(this->hostData(), vector.data(), this->size() * sizeof(T));
1780 :
1781 : if (device)
1782 1340183 : this->template copyInternal<MemSpace, ::Kokkos::HostSpace>(
1783 : this->deviceData(), vector.data(), this->size());
1784 1340183 : }
1785 : /**
1786 : * Copy a standard set variable
1787 : * This re-initializes and re-allocates array with the size of the set
1788 : * @tparam host Whether to allocate and copy to the host data
1789 : * @tparam device Whether to allocate and copy to the device data
1790 : * @param set The standard set variable to copy
1791 : */
1792 : template <bool host, bool device>
1793 1317363 : void copySet(const std::set<T> & set)
1794 : {
1795 1317363 : std::vector<T> vector(set.begin(), set.end());
1796 :
1797 1317363 : copyVector<host, device>(vector);
1798 1317363 : }
1799 :
1800 : /**
1801 : * Copy a standard vector variable
1802 : * This allocates and copies to both host and device data
1803 : * @param vector The standard vector variable to copy
1804 : */
1805 22820 : auto & operator=(const std::vector<T> & vector)
1806 : {
1807 22820 : copyVector<true, true>(vector);
1808 :
1809 22820 : return *this;
1810 : }
1811 : /**
1812 : * Copy a standard set variable
1813 : * This allocates and copies to both host and device data
1814 : * @param set The standard set variable to copy
1815 : */
1816 1307706 : auto & operator=(const std::set<T> & set)
1817 : {
1818 1307706 : copySet<true, true>(set);
1819 :
1820 1307706 : return *this;
1821 : }
1822 : /**
1823 : * Get an array entry
1824 : * @param i The array index
1825 : * @returns The reference of the entry depending on the architecture this function is being
1826 : * called on
1827 : */
1828 182599518 : KOKKOS_FUNCTION T & operator()(signed_index_type i) const
1829 : {
1830 182599518 : return this->operator[](linearIndex(i));
1831 : }
1832 : /**
1833 : * Placement-new construct an entry from args, recording initialization
1834 : * @param idx The array storing the index
1835 : * @param args Arguments forwarded to T's constructor
1836 : * @returns Reference to the constructed element
1837 : */
1838 : template <typename indices, typename... Args>
1839 : T & emplace(const indices (&idx)[1], Args &&... args);
1840 : /**
1841 : * Device BLAS operations
1842 : */
1843 : ///@{
1844 : /**
1845 : * Perform \p a * \p x \p op \p b * \p y and write the result to this array
1846 : * @param accumulate Whether to accumulate or overwrite the result
1847 : */
1848 : void axby(const T a,
1849 : const Array<T, 1, index_type, LayoutType::LEFT> & x,
1850 : const char op,
1851 : const T b,
1852 : const Array<T, 1, index_type, LayoutType::LEFT> & y,
1853 : const bool accumulate = false);
1854 : /**
1855 : * Scale \p x with \p a and write the result to this array
1856 : */
1857 : void scal(const T a, const Array<T, 1, index_type, LayoutType::LEFT> & x);
1858 : /**
1859 : * Scale this array with \p a
1860 : */
1861 : void scal(const T a);
1862 : /**
1863 : * Perform dot product between this array and \p x
1864 : */
1865 : T dot(const Array<T, 1, index_type, LayoutType::LEFT> & x);
1866 : /**
1867 : * Compute 2-norm of this array
1868 : */
1869 : T nrm2();
1870 : ///}@
1871 :
1872 : private:
1873 : /**
1874 : * Get the dimensionless index for an array index
1875 : */
1876 : KOKKOS_FUNCTION index_type linearIndex(signed_index_type i) const;
1877 : #endif
1878 : };
1879 : ///@}
1880 :
1881 : #ifdef MOOSE_KOKKOS_SCOPE
1882 : template <typename T, typename index_type>
1883 : template <typename indices, typename... Args>
1884 : T &
1885 8725 : Array<T, 1, index_type, LayoutType::LEFT>::emplace(const indices (&idx)[1], Args &&... args)
1886 : {
1887 8725 : return this->emplaceAt(linearIndex(idx[0]), std::forward<Args>(args)...);
1888 : }
1889 :
1890 : template <typename T, typename index_type>
1891 : KOKKOS_FUNCTION index_type
1892 182608243 : Array<T, 1, index_type, LayoutType::LEFT>::linearIndex(
1893 : typename Array<T, 1, index_type, LayoutType::LEFT>::signed_index_type i) const
1894 : {
1895 : KOKKOS_ASSERT(i - _d[0] >= 0 && static_cast<index_type>(i - _d[0]) < _n[0]);
1896 :
1897 182608243 : return i - _d[0];
1898 : }
1899 : #endif
1900 :
1901 : template <typename T, typename index_type = MOOSE_KOKKOS_INDEX_TYPE>
1902 : using Array1D = Array<T, 1, index_type, LayoutType::LEFT>;
1903 : template <typename T,
1904 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1905 : LayoutType layout = LayoutType::LEFT>
1906 : using Array2D = Array<T, 2, index_type, layout>;
1907 : template <typename T,
1908 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1909 : LayoutType layout = LayoutType::LEFT>
1910 : using Array3D = Array<T, 3, index_type, layout>;
1911 : template <typename T,
1912 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1913 : LayoutType layout = LayoutType::LEFT>
1914 : using Array4D = Array<T, 4, index_type, layout>;
1915 : template <typename T,
1916 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1917 : LayoutType layout = LayoutType::LEFT>
1918 : using Array5D = Array<T, 5, index_type, layout>;
1919 :
1920 : } // namespace Moose::Kokkos
|
