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 9491679 : 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 67690531 : 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 92315 : deepCopy(array);
150 : else
151 67598216 : shallowCopy(array);
152 67690531 : }
153 :
154 : /**
155 : * Destructor
156 : */
157 76904193 : ~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 704866 : 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 18268334 : 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 578063241 : 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 52735312 : KOKKOS_FUNCTION T * data() const
307 : {
308 52735312 : KOKKOS_IF_ON_HOST(return _host_data;)
309 :
310 52333713 : 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 6124871881 : KOKKOS_FUNCTION T & operator[](index_type i) const
341 : {
342 : KOKKOS_ASSERT(i < _size);
343 :
344 6124871881 : KOKKOS_IF_ON_HOST(return _host_data[i];)
345 :
346 6061725378 : return _device_data[i];
347 : }
348 :
349 : /**
350 : * Get the host data pointer
351 : * @returns The pointer to the underlying host data
352 : */
353 1425914 : T * hostData() const { return _host_data; }
354 : /**
355 : * Get the device data pointer
356 : * @returns The pointer to the underlying device data
357 : */
358 1418737 : 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 630996 : auto deviceView() const
373 : {
374 630996 : return ::Kokkos::View<T *, MemSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>>(_device_data,
375 359944 : _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 732884 : void create(size_type... n)
403 : {
404 732884 : createInternal<true, true, initialize>(n...);
405 732884 : }
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 27102 : void createHost(size_type... n)
423 : {
424 27102 : createInternal<true, false, initialize>(n...);
425 27102 : }
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 36970 : void createDevice(size_type... n)
437 : {
438 36970 : createInternal<false, true, false>(n...);
439 36970 : }
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 4980238 : KOKKOS_FUNCTION explicit iterator(T * p) : it(p) {}
536 :
537 2748813 : KOKKOS_FUNCTION reference operator*() const { return *it; }
538 : KOKKOS_FUNCTION pointer operator->() const { return it; }
539 1123492 : KOKKOS_FUNCTION pointer operator&() const { return it; }
540 2748006 : KOKKOS_FUNCTION iterator & operator++()
541 : {
542 2748006 : ++it;
543 2748006 : 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 4676281 : KOKKOS_FUNCTION friend bool operator!=(const iterator & a, const iterator & b)
556 : {
557 4676281 : 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 2490176 : KOKKOS_FUNCTION iterator begin() const
569 : {
570 2490176 : 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 2490062 : KOKKOS_FUNCTION iterator end() const
579 : {
580 2490062 : 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 607018 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::
719 434948 : constructed_entry_iterator(
720 : typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<
721 : is_const>::array_type & array,
722 434948 : index_type i)
723 607018 : : _array(array), _i(i)
724 : {
725 1041966 : advanceToConstructed();
726 1041966 : }
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 704866 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator++()
732 : {
733 704866 : ++_i;
734 704866 : advanceToConstructed();
735 704866 : 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 1225849 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator==(
752 : const constructed_entry_iterator & other) const
753 : {
754 1225849 : 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 1225849 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator!=(
761 : const constructed_entry_iterator & other) const
762 : {
763 1225849 : return !(*this == other);
764 : }
765 :
766 : template <typename T, unsigned int dimension, typename index_type>
767 : template <bool is_const>
768 : void
769 1746832 : ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::advanceToConstructed()
770 : {
771 2088364 : while (_i < _array._size && !_array.isSlotConstructed(_i))
772 341532 : ++_i;
773 1746832 : }
774 :
775 : template <typename T, unsigned int dimension, typename index_type>
776 : template <bool is_const>
777 520983 : 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 217474 : is_const>::array_type & array)
780 303509 : : _array(array)
781 : {
782 520983 : }
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 520983 : ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::begin() const
788 : {
789 520983 : 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 520983 : ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::end() const
796 : {
797 520983 : 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 520979 : ArrayBase<T, dimension, index_type>::constructedEntries()
803 : {
804 520979 : 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 14153955 : ArrayBase<T, dimension, index_type>::isSlotConstructed(index_type i) const
817 : {
818 : mooseAssert(i < _size, "isSlotConstructed index out of bounds");
819 :
820 14153955 : if (!_slots_constructed)
821 94108 : return _is_host_alloc;
822 :
823 14059847 : 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 102991 : 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 102991 : new (_host_data + i) T(std::forward<Args>(args)...);
837 102991 : (*_slots_constructed)[i] = true;
838 :
839 102991 : return _host_data[i];
840 : }
841 : #endif // MOOSE_KOKKOS_SCOPE
842 :
843 : template <typename T, unsigned int dimension, typename index_type>
844 : void
845 2336508 : ArrayBase<T, dimension, index_type>::freeHost()
846 : {
847 2336508 : if (!_is_host_alloc)
848 67326 : return;
849 :
850 2269182 : if (_is_host_alias)
851 : {
852 12005 : _host_data = nullptr;
853 12005 : _is_host_alias = false;
854 : }
855 : else
856 : {
857 2257177 : if (!_slots_constructed)
858 : // Allocated by new
859 1724098 : delete[] _host_data;
860 : else
861 : {
862 : // Allocated by malloc
863 14594892 : for (const auto i : make_range(_size))
864 13013429 : if (isSlotConstructed(i))
865 102967 : _host_data[i].~T();
866 :
867 1581463 : std::free(_host_data);
868 :
869 1581463 : _slots_constructed.reset();
870 : }
871 : }
872 :
873 2269182 : _is_host_alloc = false;
874 : }
875 :
876 : template <typename T, unsigned int dimension, typename index_type>
877 : void
878 2310699 : ArrayBase<T, dimension, index_type>::freeDevice()
879 : {
880 2310699 : if (!_is_device_alloc)
881 27004 : return;
882 :
883 2283695 : if (_is_device_alias)
884 : {
885 92 : _device_data = nullptr;
886 92 : _is_device_alias = false;
887 : }
888 : else
889 2283603 : Moose::Kokkos::free(_device_data);
890 :
891 2283695 : _is_device_alloc = false;
892 : }
893 :
894 : template <typename T, unsigned int dimension, typename index_type>
895 : void
896 147913500 : ArrayBase<T, dimension, index_type>::destroy()
897 : {
898 147913500 : if (!_counter)
899 97461469 : return;
900 :
901 50452031 : if (_counter.use_count() > 1)
902 : {
903 48141332 : _host_data = nullptr;
904 48141332 : _device_data = nullptr;
905 : }
906 2310699 : else if (_counter.use_count() == 1)
907 : {
908 2310699 : freeHost();
909 2310699 : freeDevice();
910 : }
911 :
912 50452031 : _size = 0;
913 :
914 127480088 : for (const auto i : make_range(dimension))
915 : {
916 77028057 : _n[i] = 0;
917 77028057 : _s[i] = 0;
918 77028057 : _d[i] = 0;
919 : }
920 :
921 50452031 : _is_init = false;
922 50452031 : _is_host_alloc = false;
923 50452031 : _is_device_alloc = false;
924 50452031 : _is_host_alias = false;
925 50452031 : _is_device_alias = false;
926 :
927 50452031 : _counter.reset();
928 50452031 : _slots_constructed.reset();
929 : }
930 :
931 : template <typename T, unsigned int dimension, typename index_type>
932 : void
933 70026199 : ArrayBase<T, dimension, index_type>::shallowCopy(const ArrayBase<T, dimension, index_type> & array)
934 : {
935 70026199 : if (_layout != array._layout)
936 0 : mooseError("Kokkos array error: cannot shallow copy arrays with different layouts.");
937 :
938 70026199 : destroy();
939 :
940 70026199 : _counter = array._counter;
941 70026199 : _slots_constructed = array._slots_constructed;
942 :
943 70026199 : _size = array._size;
944 :
945 178515986 : for (const auto i : make_range(dimension))
946 : {
947 108489787 : _n[i] = array._n[i];
948 108489787 : _s[i] = array._s[i];
949 108489787 : _d[i] = array._d[i];
950 : }
951 :
952 70026199 : _is_init = array._is_init;
953 70026199 : _is_host_alloc = array._is_host_alloc;
954 70026199 : _is_device_alloc = array._is_device_alloc;
955 70026199 : _is_host_alias = array._is_host_alias;
956 70026199 : _is_device_alias = array._is_device_alias;
957 :
958 70026199 : _host_data = array._host_data;
959 70026199 : _device_data = array._device_data;
960 70026199 : }
961 :
962 : #ifdef MOOSE_KOKKOS_SCOPE
963 : template <typename T, unsigned int dimension, typename index_type>
964 : void
965 977929 : ArrayBase<T, dimension, index_type>::aliasHost(T * ptr)
966 : {
967 977929 : if (!_is_init)
968 0 : mooseError("Kokkos array error: attempted to alias host data before array initialization.");
969 :
970 977929 : 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 977929 : _host_data = ptr;
974 977929 : _is_host_alloc = true;
975 977929 : _is_host_alias = true;
976 977929 : }
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 2260761 : ArrayBase<T, dimension, index_type>::allocHost()
997 : {
998 2260761 : 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 4025455 : _host_data = new T[_size];
1008 : }
1009 : else
1010 : {
1011 1581903 : _host_data = static_cast<T *>(std::malloc(_size * sizeof(T)));
1012 1581903 : _slots_constructed = std::make_shared<std::vector<bool>>(_size, false);
1013 : }
1014 :
1015 2260761 : _is_host_alloc = true;
1016 : }
1017 :
1018 : template <typename T, unsigned int dimension, typename index_type>
1019 : void
1020 2287415 : ArrayBase<T, dimension, index_type>::allocDevice()
1021 : {
1022 2287415 : if (_is_device_alloc)
1023 0 : return;
1024 :
1025 2287415 : _device_data =
1026 1305835 : static_cast<T *>(::Kokkos::kokkos_malloc<ExecSpace::memory_space>(_size * sizeof(T)));
1027 :
1028 2287415 : _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 2314717 : ArrayBase<T, dimension, index_type>::createInternal(const std::vector<index_type> & n)
1035 : {
1036 2314717 : 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 2314717 : if (_counter)
1044 1419 : destroy();
1045 :
1046 2314717 : _counter = std::make_shared<unsigned int>();
1047 :
1048 2314717 : uint64_t overflow_checker = 1;
1049 :
1050 2314717 : _size = 1;
1051 2314717 : _s[0] = 1;
1052 :
1053 4879987 : for (const auto i : make_range(dimension))
1054 : {
1055 2565270 : overflow_checker *= n[i];
1056 :
1057 2565270 : _n[i] = n[i];
1058 2565270 : _size *= n[i];
1059 : }
1060 :
1061 2314717 : 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 2314717 : if (_layout == LayoutType::LEFT)
1071 : {
1072 2314707 : _s[0] = 1;
1073 :
1074 2565238 : for (const auto i : make_range(1u, dimension))
1075 250531 : _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 2260761 : allocHost<initialize>();
1087 :
1088 : if constexpr (device)
1089 2287415 : allocDevice();
1090 :
1091 2314717 : _is_init = true;
1092 2314717 : }
1093 :
1094 : template <typename T, unsigned int dimension, typename index_type>
1095 : template <bool initialize>
1096 : void
1097 93734 : ArrayBase<T, dimension, index_type>::createInternal(const std::vector<index_type> & n,
1098 : bool host,
1099 : bool device)
1100 : {
1101 93734 : if (host && device)
1102 92315 : 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 93734 : }
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 797156 : 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 797156 : std::vector<index_type> dims;
1121 797156 : (dims.push_back(n), ...);
1122 :
1123 797156 : createInternal<host, device, initialize>(dims);
1124 797156 : }
1125 :
1126 : template <typename T, unsigned int dimension, typename index_type>
1127 : template <typename TargetSpace, typename SourceSpace>
1128 : void
1129 7286059 : ArrayBase<T, dimension, index_type>::copyInternal(T * target, const T * source, index_type n)
1130 : {
1131 7286059 : ::Kokkos::Impl::DeepCopy<TargetSpace, SourceSpace>(target, source, n * sizeof(T));
1132 7286059 : ::Kokkos::fence();
1133 7286059 : }
1134 :
1135 : template <typename T, unsigned int dimension, typename index_type>
1136 : void
1137 5008 : ArrayBase<T, dimension, index_type>::offset(const std::vector<signed_index_type> & d)
1138 : {
1139 5008 : 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 10018 : for (const auto i : index_range(d))
1147 5010 : _d[i] = d[i];
1148 5008 : }
1149 :
1150 : template <typename T, unsigned int dimension, typename index_type>
1151 : template <typename... offset_type>
1152 : void
1153 5008 : 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 5008 : std::vector<signed_index_type> offsets;
1160 5008 : (offsets.push_back(d), ...);
1161 :
1162 5008 : offset(offsets);
1163 5008 : }
1164 :
1165 : template <typename T, unsigned int dimension, typename index_type>
1166 : void
1167 5110953 : ArrayBase<T, dimension, index_type>::copyToDevice()
1168 : {
1169 : // If host side memory is not allocated, do nothing
1170 5110953 : if (!_is_host_alloc)
1171 88306 : return;
1172 :
1173 : // If device side memory is not allocated,
1174 5022647 : 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 5022647 : copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data, _host_data, _size);
1188 : }
1189 :
1190 : template <typename T, unsigned int dimension, typename index_type>
1191 : void
1192 841891 : ArrayBase<T, dimension, index_type>::copyToHost()
1193 : {
1194 : // If device side memory is not allocated, do nothing
1195 841891 : if (!_is_device_alloc)
1196 0 : return;
1197 :
1198 : // If host side memory is not allocated,
1199 841891 : 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 841891 : copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data, _device_data, _size);
1213 : }
1214 :
1215 : template <typename T, unsigned int dimension, typename index_type>
1216 : void
1217 25809 : 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 25809 : 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 25809 : copyToDevice();
1227 :
1228 25809 : if (_counter.use_count() == 1)
1229 25809 : freeHost();
1230 25809 : }
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 12121272 : copyToDeviceInner(T & /* data */)
1359 : {
1360 12121272 : }
1361 :
1362 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1363 : void
1364 299399 : copyToDeviceInner(Array<T, dimension, index_type, layout> & data)
1365 : {
1366 299399 : data.copyToDeviceNested();
1367 299399 : }
1368 :
1369 : template <typename T, unsigned int dimension, typename index_type>
1370 : void
1371 375187 : ArrayBase<T, dimension, index_type>::copyToDeviceNested()
1372 : {
1373 12795858 : for (const auto i : make_range(_size))
1374 12420671 : copyToDeviceInner(_host_data[i]);
1375 :
1376 375187 : copyToDevice();
1377 375187 : }
1378 :
1379 : template <typename T, unsigned int dimension, typename index_type>
1380 : void
1381 93734 : ArrayBase<T, dimension, index_type>::deepCopy(const ArrayBase<T, dimension, index_type> & array)
1382 : {
1383 93734 : if (_layout != array._layout)
1384 0 : mooseError("Kokkos array error: cannot deep copy arrays with different layouts.");
1385 :
1386 92315 : 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 281202 : std::vector<index_type> n(std::begin(array._n), std::end(array._n));
1391 :
1392 93734 : createInternal<false>(n, array._is_host_alloc, array._is_device_alloc);
1393 :
1394 93734 : if (_is_host_alloc)
1395 186411 : for (const auto i : make_range(_size))
1396 94096 : if (array.isSlotConstructed(i))
1397 94096 : emplaceAt(i, array._host_data[i]);
1398 :
1399 : if (ArrayDeepCopy<T>::value)
1400 92315 : copyToDevice();
1401 1419 : else if (_is_device_alloc)
1402 1419 : copyInternal<MemSpace, MemSpace>(_device_data, array._device_data, _size);
1403 :
1404 187468 : for (const auto i : make_range(dimension))
1405 : {
1406 93734 : _d[i] = array._d[i];
1407 93734 : _s[i] = array._s[i];
1408 : }
1409 93734 : }
1410 :
1411 : template <typename T, unsigned int dimension, typename index_type>
1412 : void
1413 2808 : ArrayBase<T, dimension, index_type>::swap(ArrayBase<T, dimension, index_type> & array)
1414 : {
1415 2808 : ArrayBase<T, dimension, index_type> clone(_layout);
1416 :
1417 2808 : clone.shallowCopy(*this);
1418 2808 : this->shallowCopy(array);
1419 2808 : array.shallowCopy(clone);
1420 2808 : }
1421 :
1422 : template <typename T, unsigned int dimension, typename index_type>
1423 : auto &
1424 630986 : ArrayBase<T, dimension, index_type>::operator=(const T & scalar)
1425 : {
1426 630986 : if (_is_host_alloc)
1427 630972 : std::fill_n(_host_data, _size, scalar);
1428 :
1429 630986 : if (_is_device_alloc)
1430 630986 : ::Kokkos::Experimental::fill_n(ExecSpace(), deviceView(), _size, scalar);
1431 :
1432 630986 : return *this;
1433 : }
1434 :
1435 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1436 : void
1437 1139 : dataStore(std::ostream & stream, Array<T, dimension, index_type, layout> & array, void * context)
1438 : {
1439 : using ::dataStore;
1440 :
1441 1139 : bool is_alloc = array.isAlloc();
1442 1139 : dataStore(stream, is_alloc, nullptr);
1443 :
1444 1139 : if (!is_alloc)
1445 0 : return;
1446 :
1447 1139 : std::string type = typeid(T).name();
1448 1139 : dataStore(stream, type, nullptr);
1449 :
1450 1139 : unsigned int dim = dimension;
1451 1139 : dataStore(stream, dim, nullptr);
1452 :
1453 2278 : for (const auto dim : make_range(dimension))
1454 : {
1455 1139 : auto n = array.n(dim);
1456 1139 : dataStore(stream, n, nullptr);
1457 : }
1458 :
1459 1139 : if (array.isDeviceAlloc())
1460 : {
1461 : // We use malloc/free because we just want a memory copy
1462 : // If T is a Kokkos array and we use new/delete or vector to copy it out,
1463 : // the arrays will be destroyed on cleanup
1464 :
1465 1139 : T * data = static_cast<T *>(std::malloc(array.size() * sizeof(T)));
1466 :
1467 1139 : array.copyOut(data, MemcpyType::DEVICE_TO_HOST, array.size());
1468 :
1469 117441 : for (const auto i : make_range(array.size()))
1470 116302 : dataStore(stream, data[i], context);
1471 :
1472 1139 : std::free(data);
1473 : }
1474 : else
1475 0 : for (auto & value : array)
1476 0 : dataStore(stream, value, context);
1477 1139 : }
1478 :
1479 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1480 : void
1481 531 : dataLoad(std::istream & stream, Array<T, dimension, index_type, layout> & array, void * context)
1482 : {
1483 : using ::dataLoad;
1484 :
1485 : bool is_alloc;
1486 531 : dataLoad(stream, is_alloc, nullptr);
1487 :
1488 531 : if (!is_alloc)
1489 0 : return;
1490 :
1491 531 : std::string from_type_name;
1492 531 : dataLoad(stream, from_type_name, nullptr);
1493 :
1494 531 : if (from_type_name != typeid(T).name())
1495 0 : mooseError("Kokkos array error: cannot load array because the stored array is of type '",
1496 : MooseUtils::prettyCppType(libMesh::demangle(from_type_name.c_str())),
1497 : "' but the loading array is of type '",
1498 : MooseUtils::prettyCppType(libMesh::demangle(typeid(T).name())),
1499 : "'.");
1500 :
1501 : unsigned int from_dimension;
1502 531 : dataLoad(stream, from_dimension, nullptr);
1503 :
1504 531 : if (from_dimension != dimension)
1505 0 : mooseError("Kokkos array error: cannot load array because the stored array is ",
1506 : from_dimension,
1507 : "D but the loading array is ",
1508 : dimension,
1509 : "D.");
1510 :
1511 1062 : std::vector<index_type> from_n(dimension);
1512 531 : std::vector<index_type> n(dimension);
1513 :
1514 1062 : for (const auto dim : make_range(dimension))
1515 : {
1516 531 : dataLoad(stream, from_n[dim], nullptr);
1517 531 : n[dim] = array.n(dim);
1518 : }
1519 :
1520 531 : if (from_n != n)
1521 0 : mooseError("Kokkos array error: cannot load array because the stored array has dimensions (",
1522 : Moose::stringify(from_n),
1523 : ") but the loading array has dimensions (",
1524 : Moose::stringify(n),
1525 : ").");
1526 :
1527 531 : if (array.isHostAlloc())
1528 : {
1529 2706 : for (auto & value : array)
1530 2156 : dataLoad(stream, value, context);
1531 :
1532 275 : if (array.isDeviceAlloc())
1533 275 : array.copyToDevice();
1534 : }
1535 : else
1536 : {
1537 256 : std::vector<T> data(array.size());
1538 :
1539 51792 : for (auto & value : data)
1540 51536 : dataLoad(stream, value, context);
1541 :
1542 256 : array.copyIn(data.data(), MemcpyType::HOST_TO_DEVICE, array.size());
1543 256 : }
1544 531 : }
1545 : #endif
1546 :
1547 : /**
1548 : * The specialization of the Kokkos array class for each dimension.
1549 : * All array data that needs to be accessed on device in Kokkos objects should use this class.
1550 : * If the array is populated on host and is to be accessed on device, make sure to call
1551 : * copyToDevice() after populating data. For a nested Kokkos array, either copyToDeviceNested()
1552 : * should be called for the outermost array or copyToDevice() should be called for each instance of
1553 : * Kokkos array from the innermost to the outermost. Do not store this object as reference in your
1554 : * Kokkos object if it is used on device, because the reference refers to a host object and
1555 : * therefore is not accessible on device. If storing it as a reference is required, see
1556 : * ReferenceWrapper.
1557 : * @tparam T The data type
1558 : * @tparam dimension The array dimension size
1559 : * @tparam index_type The array index type
1560 : * @tparam layout The memory layout type
1561 : */
1562 : ///@{
1563 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1564 : class Array : public ArrayBase<T, dimension, index_type>
1565 : {
1566 : #ifdef MOOSE_KOKKOS_SCOPE
1567 : usingKokkosArrayBaseMembers(T, dimension, index_type);
1568 : #endif
1569 :
1570 : public:
1571 : /**
1572 : * Default constructor
1573 : */
1574 2351087 : Array() : ArrayBase<T, dimension, index_type>(layout) {}
1575 : /**
1576 : * Copy constructor
1577 : */
1578 32867118 : Array(const Array<T, dimension, index_type, layout> & array)
1579 19109355 : : ArrayBase<T, dimension, index_type>(array)
1580 : {
1581 32867118 : }
1582 : #ifdef MOOSE_KOKKOS_SCOPE
1583 : /**
1584 : * Constructor
1585 : * Initialize and allocate array with given dimensions
1586 : * This allocates both host and device data
1587 : * @param n The size of each dimension
1588 : */
1589 : template <typename... size_type>
1590 16 : Array(size_type... n) : ArrayBase<T, dimension, index_type>(layout, n...)
1591 : {
1592 16 : }
1593 : #endif
1594 :
1595 : /**
1596 : * Shallow copy another Kokkos array
1597 : * @param array The Kokkos array to be shallow copied
1598 : */
1599 : auto & operator=(const Array<T, dimension, index_type, layout> & array)
1600 : {
1601 : this->shallowCopy(array);
1602 :
1603 : return *this;
1604 : }
1605 :
1606 : #ifdef MOOSE_KOKKOS_SCOPE
1607 : /**
1608 : * Get an array entry
1609 : * @param i The index of each dimension
1610 : * @returns The reference of the entry depending on the architecture this function is being called
1611 : * on
1612 : */
1613 : template <typename... indices>
1614 2225826259 : KOKKOS_FUNCTION T & operator()(indices... i) const
1615 : {
1616 2225826259 : return this->operator[](linearIndex(i...));
1617 : }
1618 : /**
1619 : * Get an array entry using indices stored in an array
1620 : * @param idx The array storing the indices
1621 : * @returns The reference of the entry depending on the architecture this function is being
1622 : * called on
1623 : */
1624 : KOKKOS_FUNCTION T & operator()(const signed_index_type (&idx)[dimension]) const
1625 : {
1626 : return this->operator[](linearIndex(idx));
1627 : }
1628 : /**
1629 : * Placement-new construct an entry from args, recording initialization
1630 : * @param idx The array storing the indices
1631 : * @param args Arguments forwarded to T's constructor
1632 : * @returns Reference to the constructed element
1633 : */
1634 : template <typename indices, typename... Args>
1635 : T & emplace(const indices (&idx)[dimension], Args &&... args);
1636 : #endif
1637 :
1638 : private:
1639 : #ifdef MOOSE_KOKKOS_SCOPE
1640 : /**
1641 : * Get the dimensionless index for a multi-dimensional index
1642 : */
1643 : template <typename... indices>
1644 : KOKKOS_FUNCTION index_type linearIndex(indices... i) const;
1645 : /**
1646 : * Get the dimensionless index for indices stored in an array
1647 : */
1648 : ///@{
1649 : template <typename indices>
1650 8 : KOKKOS_FUNCTION index_type linearIndex(const indices (&idx)[dimension]) const
1651 : {
1652 8 : return linearIndexHelper(idx, std::make_integer_sequence<unsigned int, dimension>{});
1653 : }
1654 : template <typename indices, unsigned int... i>
1655 8 : KOKKOS_FUNCTION index_type linearIndexHelper(const indices (&idx)[dimension],
1656 : std::integer_sequence<unsigned int, i...>) const
1657 : {
1658 8 : return linearIndex(idx[i]...);
1659 : }
1660 : ///@}
1661 : #endif
1662 : };
1663 :
1664 : #ifdef MOOSE_KOKKOS_SCOPE
1665 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1666 : template <typename indices, typename... Args>
1667 : T &
1668 8 : Array<T, dimension, index_type, layout>::emplace(const indices (&idx)[dimension], Args &&... args)
1669 : {
1670 8 : return this->emplaceAt(linearIndex(idx), std::forward<Args>(args)...);
1671 : }
1672 :
1673 : template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
1674 : template <typename... indices>
1675 : KOKKOS_FUNCTION index_type
1676 2225826267 : Array<T, dimension, index_type, layout>::linearIndex(indices... i) const
1677 : {
1678 : static_assert((std::is_convertible<indices, signed_index_type>::value && ...),
1679 : "All arguments must be convertible to signed_index_type");
1680 : static_assert(sizeof...(i) == dimension, "Number of arguments should match array dimension");
1681 :
1682 : #ifndef NDEBUG
1683 : {
1684 : signed_index_type idx[dimension] = {static_cast<signed_index_type>(i)...};
1685 :
1686 : for (unsigned int d = 0; d < sizeof...(i); ++d)
1687 : KOKKOS_ASSERT(idx[d] - _d[d] >= 0 && static_cast<index_type>(idx[d] - _d[d]) < _n[d]);
1688 : }
1689 : #endif
1690 :
1691 2225826267 : index_type idx = 0;
1692 2225826267 : unsigned int d = 0;
1693 :
1694 : if constexpr (layout == LayoutType::LEFT)
1695 2664943373 : (((idx += (d == 0 ? static_cast<signed_index_type>(i) - _d[d]
1696 2655517956 : : (static_cast<signed_index_type>(i) - _d[d]) * _s[d])),
1697 : ++d),
1698 : ...);
1699 : else
1700 14250 : (((idx += (d == dimension - 1 ? static_cast<signed_index_type>(i) - _d[d]
1701 11294 : : (static_cast<signed_index_type>(i) - _d[d]) * _s[d])),
1702 : ++d),
1703 : ...);
1704 :
1705 2225826267 : return idx;
1706 : }
1707 : #endif
1708 :
1709 : template <typename T, typename index_type>
1710 : class Array<T, 1, index_type, LayoutType::LEFT> : public ArrayBase<T, 1, index_type>
1711 : {
1712 : #ifdef MOOSE_KOKKOS_SCOPE
1713 : usingKokkosArrayBaseMembers(T, 1, index_type);
1714 : #endif
1715 :
1716 : public:
1717 : /**
1718 : * Default constructor
1719 : */
1720 7136156 : Array() : ArrayBase<T, 1, index_type>(LayoutType::LEFT) {}
1721 : /**
1722 : * Copy constructor
1723 : */
1724 34823413 : Array(const Array<T, 1, index_type, LayoutType::LEFT> & array)
1725 20311459 : : ArrayBase<T, 1, index_type>(array)
1726 : {
1727 34823413 : }
1728 : #ifdef MOOSE_KOKKOS_SCOPE
1729 : /**
1730 : * Constructor
1731 : * Initialize and allocate array with given size
1732 : * This allocates both host and device data
1733 : * @param n The array size
1734 : */
1735 60 : Array(index_type n) : ArrayBase<T, 1, index_type>(LayoutType::LEFT, n) {}
1736 : /**
1737 : * Constructor
1738 : * Initialize and allocate array by copying a standard vector variable
1739 : * This allocates and copies to both host and device data
1740 : * @param vector The standard vector variable to copy
1741 : */
1742 1628 : Array(const std::vector<T> & vector) : ArrayBase<T, 1, index_type>(LayoutType::LEFT)
1743 : {
1744 1628 : *this = vector;
1745 1628 : }
1746 : #endif
1747 :
1748 : /**
1749 : * Shallow copy another Kokkos array
1750 : * @param array The Kokkos array to be shallow copied
1751 : */
1752 2419559 : auto & operator=(const Array<T, 1, index_type, LayoutType::LEFT> & array)
1753 : {
1754 2419559 : this->shallowCopy(array);
1755 :
1756 2419559 : return *this;
1757 : }
1758 :
1759 : #ifdef MOOSE_KOKKOS_SCOPE
1760 : /**
1761 : * Copy a standard vector variable
1762 : * This re-initializes and re-allocates array with the size of the vector
1763 : * @tparam host Whether to allocate and copy to the host data
1764 : * @tparam device Whether to allocate and copy to the device data
1765 : * @param vector The standard vector variable to copy
1766 : */
1767 : template <bool host, bool device>
1768 1418707 : void copyVector(const std::vector<T> & vector)
1769 : {
1770 2847661 : this->template createInternal<host, device, false>({static_cast<index_type>(vector.size())});
1771 :
1772 : if (host)
1773 : {
1774 : if constexpr (std::is_trivially_copyable<T>())
1775 1408419 : std::memcpy(this->hostData(), vector.data(), this->size() * sizeof(T));
1776 : else
1777 41 : std::copy(vector.begin(), vector.end(), this->begin());
1778 : }
1779 :
1780 : if (device)
1781 1418707 : this->template copyInternal<MemSpace, ::Kokkos::HostSpace>(
1782 : this->deviceData(), vector.data(), this->size());
1783 1418707 : }
1784 : /**
1785 : * Copy a standard set variable
1786 : * This re-initializes and re-allocates array with the size of the set
1787 : * @tparam host Whether to allocate and copy to the host data
1788 : * @tparam device Whether to allocate and copy to the device data
1789 : * @param set The standard set variable to copy
1790 : */
1791 : template <bool host, bool device>
1792 1393958 : void copySet(const std::set<T> & set)
1793 : {
1794 1393958 : std::vector<T> vector(set.begin(), set.end());
1795 :
1796 1393958 : copyVector<host, device>(vector);
1797 1393958 : }
1798 :
1799 : /**
1800 : * Copy a standard vector variable
1801 : * This allocates and copies to both host and device data
1802 : * @param vector The standard vector variable to copy
1803 : */
1804 24749 : auto & operator=(const std::vector<T> & vector)
1805 : {
1806 24749 : copyVector<true, true>(vector);
1807 :
1808 24749 : return *this;
1809 : }
1810 : /**
1811 : * Copy a standard set variable
1812 : * This allocates and copies to both host and device data
1813 : * @param set The standard set variable to copy
1814 : */
1815 1383711 : auto & operator=(const std::set<T> & set)
1816 : {
1817 1383711 : copySet<true, true>(set);
1818 :
1819 1383711 : return *this;
1820 : }
1821 : /**
1822 : * Get an array entry
1823 : * @param i The array index
1824 : * @returns The reference of the entry depending on the architecture this function is being
1825 : * called on
1826 : */
1827 189176861 : KOKKOS_FUNCTION T & operator()(signed_index_type i) const
1828 : {
1829 189176861 : return this->operator[](linearIndex(i));
1830 : }
1831 : /**
1832 : * Placement-new construct an entry from args, recording initialization
1833 : * @param idx The array storing the index
1834 : * @param args Arguments forwarded to T's constructor
1835 : * @returns Reference to the constructed element
1836 : */
1837 : template <typename indices, typename... Args>
1838 : T & emplace(const indices (&idx)[1], Args &&... args);
1839 : /**
1840 : * Device BLAS operations
1841 : */
1842 : ///@{
1843 : /**
1844 : * Perform \p a * \p x \p op \p b * \p y and write the result to this array
1845 : * @param accumulate Whether to accumulate or overwrite the result
1846 : */
1847 : void axby(const T a,
1848 : const Array<T, 1, index_type, LayoutType::LEFT> & x,
1849 : const char op,
1850 : const T b,
1851 : const Array<T, 1, index_type, LayoutType::LEFT> & y,
1852 : const bool accumulate = false);
1853 : /**
1854 : * Scale \p x with \p a and write the result to this array
1855 : */
1856 : void scal(const T a, const Array<T, 1, index_type, LayoutType::LEFT> & x);
1857 : /**
1858 : * Scale this array with \p a
1859 : */
1860 : void scal(const T a);
1861 : /**
1862 : * Perform dot product between this array and \p x
1863 : */
1864 : T dot(const Array<T, 1, index_type, LayoutType::LEFT> & x);
1865 : /**
1866 : * Compute 2-norm of this array
1867 : */
1868 : T nrm2();
1869 : ///}@
1870 :
1871 : private:
1872 : /**
1873 : * Get the dimensionless index for an array index
1874 : */
1875 : KOKKOS_FUNCTION index_type linearIndex(signed_index_type i) const;
1876 : #endif
1877 : };
1878 : ///@}
1879 :
1880 : #ifdef MOOSE_KOKKOS_SCOPE
1881 : template <typename T, typename index_type>
1882 : template <typename indices, typename... Args>
1883 : T &
1884 8887 : Array<T, 1, index_type, LayoutType::LEFT>::emplace(const indices (&idx)[1], Args &&... args)
1885 : {
1886 8887 : return this->emplaceAt(linearIndex(idx[0]), std::forward<Args>(args)...);
1887 : }
1888 :
1889 : template <typename T, typename index_type>
1890 : KOKKOS_FUNCTION index_type
1891 189185748 : Array<T, 1, index_type, LayoutType::LEFT>::linearIndex(
1892 : typename Array<T, 1, index_type, LayoutType::LEFT>::signed_index_type i) const
1893 : {
1894 : KOKKOS_ASSERT(i - _d[0] >= 0 && static_cast<index_type>(i - _d[0]) < _n[0]);
1895 :
1896 189185748 : return i - _d[0];
1897 : }
1898 : #endif
1899 :
1900 : template <typename T, typename index_type = MOOSE_KOKKOS_INDEX_TYPE>
1901 : using Array1D = Array<T, 1, index_type, LayoutType::LEFT>;
1902 : template <typename T,
1903 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1904 : LayoutType layout = LayoutType::LEFT>
1905 : using Array2D = Array<T, 2, index_type, layout>;
1906 : template <typename T,
1907 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1908 : LayoutType layout = LayoutType::LEFT>
1909 : using Array3D = Array<T, 3, index_type, layout>;
1910 : template <typename T,
1911 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1912 : LayoutType layout = LayoutType::LEFT>
1913 : using Array4D = Array<T, 4, index_type, layout>;
1914 : template <typename T,
1915 : typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
1916 : LayoutType layout = LayoutType::LEFT>
1917 : using Array5D = Array<T, 5, index_type, layout>;
1918 :
1919 : } // namespace Moose::Kokkos
|
