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