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