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