KokkosArray.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

#ifdef MOOSE_KOKKOS_SCOPE
#include "KokkosHeader.h"
#endif

#include "Conversion.h"
#include "DataIO.h"

#include <iterator>

#define usingKokkosArrayBaseMembers(T, dimension, index_type)                                      \
private:                                                                                           \
  using ArrayBase<T, dimension, index_type>::_n;                                                   \
  using ArrayBase<T, dimension, index_type>::_s;                                                   \
  using ArrayBase<T, dimension, index_type>::_d;                                                   \
                                                                                                   \
public:                                                                                            \
  using typename ArrayBase<T, dimension, index_type>::signed_index_type;                           \
  using ArrayBase<T, dimension, index_type>::operator=

namespace Moose::Kokkos
{

// This function simply calls ::Kokkos::kokkos_free, but it is separately defined in KokkosArray.K
// because the Kokkos function cannot be directly seen by the host compiler
void free(void * ptr);

enum class MemcpyType
{
  HOST_TO_HOST,
  HOST_TO_DEVICE,
  DEVICE_TO_HOST,
  DEVICE_TO_DEVICE
};

enum class LayoutType
{
  LEFT,
  RIGHT
};

template <typename T,
          unsigned int dimension = 1,
          typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
          LayoutType layout = LayoutType::LEFT>
class Array;

template <typename>
struct is_kokkos_array : std::false_type
{
};

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
struct is_kokkos_array<Array<T, dimension, index_type, layout>> : std::true_type
{
};

template <typename T>
struct ArrayDeepCopy
{
  static constexpr bool value = false;
};

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
struct ArrayDeepCopy<Array<T, dimension, index_type, layout>>
{
  static constexpr bool value = ArrayDeepCopy<T>::value;
};

template <typename T, unsigned int dimension, typename index_type>
class ArrayBase
{
  static_assert(std::is_integral_v<index_type>, "Kokkos array index type must be an integral type");
  static_assert(std::is_unsigned_v<index_type>, "Kokkos array index type must be unsigned");
  static_assert(!std::is_same_v<bool, index_type>, "Kokkos array index type must not be bool");

public:
  using unsigned_index_type = index_type;
  using signed_index_type = typename std::make_signed<index_type>::type;

  ArrayBase(const LayoutType layout) : _layout(layout) {}

#ifdef MOOSE_KOKKOS_SCOPE

  template <typename... size_type>
  ArrayBase(const LayoutType layout, size_type... n) : _layout(layout)
  {
    create(n...);
  }
#endif

  ArrayBase(const ArrayBase<T, dimension, index_type> & array) : _layout(array._layout)
  {
#ifndef MOOSE_KOKKOS_SCOPE
    static_assert(!ArrayDeepCopy<T>::value,
                  "Kokkos array cannot be deep copied outside the Kokkos compilation scope");
#endif

    if constexpr (ArrayDeepCopy<T>::value)
      deepCopy(array);
    else
      shallowCopy(array);
  }

  ~ArrayBase() { destroy(); }

  void destroy();

  void shallowCopy(const ArrayBase<T, dimension, index_type> & array);

  unsigned int useCount() const { return _counter.use_count(); }

  bool isSlotConstructed(index_type i) const;

  template <bool is_const>
  class constructed_entry_range;

  template <bool is_const>
  class constructed_entry_iterator
  {
  public:
    using iterator_category = std::forward_iterator_tag;
    using value_type = T;
    using difference_type = std::ptrdiff_t;
    using pointer = std::conditional_t<is_const, const T *, T *>;
    using reference = std::conditional_t<is_const, const T &, T &>;
    using array_type = std::conditional_t<is_const, const ArrayBase, ArrayBase>;

    reference operator*() const { return _array._host_data[_i]; }
    pointer operator->() const { return _array._host_data + _i; }
    constructed_entry_iterator & operator++();
    constructed_entry_iterator operator++(int);
    bool operator==(const constructed_entry_iterator & other) const;
    bool operator!=(const constructed_entry_iterator & other) const;

  private:
    friend class constructed_entry_range<is_const>;

    constructed_entry_iterator(array_type & array, index_type i);

    void advanceToConstructed();

    array_type & _array;
    index_type _i = 0;
  };

  template <bool is_const>
  class constructed_entry_range
  {
  public:
    using array_type = std::conditional_t<is_const, const ArrayBase, ArrayBase>;
    using iterator = constructed_entry_iterator<is_const>;

    explicit constructed_entry_range(array_type & array);

    iterator begin() const;
    iterator end() const;

  private:
    array_type & _array;
  };

  constructed_entry_range<false> constructedEntries();

  constructed_entry_range<true> constructedEntries() const;

#ifdef MOOSE_KOKKOS_SCOPE

  KOKKOS_FUNCTION bool isAlloc() const { return _is_host_alloc || _is_device_alloc; }
  KOKKOS_FUNCTION bool isHostAlloc() const { return _is_host_alloc; }
  KOKKOS_FUNCTION bool isDeviceAlloc() const { return _is_device_alloc; }
  KOKKOS_FUNCTION bool isHostAlias() const { return _is_host_alias; }
  KOKKOS_FUNCTION bool isDeviceAlias() const { return _is_device_alias; }
  KOKKOS_FUNCTION index_type size() const { return _size; }
  KOKKOS_FUNCTION index_type n(unsigned int dim) const { return _n[dim]; }
  KOKKOS_FUNCTION T * data() const
  {
    KOKKOS_IF_ON_HOST(return _host_data;)

    return _device_data;
  }
  KOKKOS_FUNCTION T & first() const
  {
    KOKKOS_IF_ON_HOST(return _host_data[0];)

    return _device_data[0];
  }
  KOKKOS_FUNCTION T & last() const
  {
    KOKKOS_IF_ON_HOST(return _host_data[_size - 1];)

    return _device_data[_size - 1];
  }
  KOKKOS_FUNCTION T & operator[](index_type i) const
  {
    KOKKOS_ASSERT(i < _size);

    KOKKOS_IF_ON_HOST(return _host_data[i];)

    return _device_data[i];
  }

  T * hostData() const { return _host_data; }
  T * deviceData() const { return _device_data; }
  auto hostView() const
  {
    return ::Kokkos::View<T *, ::Kokkos::HostSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>>(
        _host_data, _size);
  }
  auto deviceView() const
  {
    return ::Kokkos::View<T *, MemSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>>(_device_data,
                                                                                      _size);
  }
  template <typename... size_type>
  void init(size_type... n)
  {
    createInternal<false, false, false>(n...);
  }
  template <bool initialize = true>
  void create(const std::vector<index_type> & n)
  {
    createInternal<true, true, initialize>(n);
  }
  template <bool initialize = true, typename... size_type>
  void create(size_type... n)
  {
    createInternal<true, true, initialize>(n...);
  }
  template <bool initialize = true>
  void createHost(const std::vector<index_type> & n)
  {
    createInternal<true, false, initialize>(n);
  }
  template <bool initialize = true, typename... size_type>
  void createHost(size_type... n)
  {
    createInternal<true, false, initialize>(n...);
  }
  void createDevice(const std::vector<index_type> & n) { createInternal<false, true, false>(n); }
  template <typename... size_type>
  void createDevice(size_type... n)
  {
    createInternal<false, true, false>(n...);
  }
  void aliasHost(T * ptr);
  void aliasDevice(T * ptr);
  void offset(const std::vector<signed_index_type> & d);
  template <typename... offset_type>
  void offset(offset_type... d);
  void copyToDevice();
  void copyToHost();
  void copyIn(const T * ptr, MemcpyType dir, index_type n, index_type offset = 0);
  void copyOut(T * ptr, MemcpyType dir, index_type n, index_type offset = 0);
  void copyToDeviceNested();
  void moveToDevice(bool should_free_host = true);
  void moveToHost(bool should_free_device = true);
  void deepCopy(const ArrayBase<T, dimension, index_type> & array);
  void swap(ArrayBase<T, dimension, index_type> & array);

  auto & operator=(const T & scalar);

  class iterator
  {
  public:
    using iterator_category = std::forward_iterator_tag;
    using value_type = T;
    using difference_type = std::ptrdiff_t;
    using pointer = T *;
    using reference = T &;

    KOKKOS_FUNCTION iterator() : it(nullptr) {}
    KOKKOS_FUNCTION explicit iterator(T * p) : it(p) {}

    KOKKOS_FUNCTION reference operator*() const { return *it; }
    KOKKOS_FUNCTION pointer operator->() const { return it; }
    KOKKOS_FUNCTION pointer operator&() const { return it; }
    KOKKOS_FUNCTION iterator & operator++()
    {
      ++it;
      return *this;
    }
    KOKKOS_FUNCTION iterator operator++(int)
    {
      iterator pre = *this;
      ++it;
      return pre;
    }
    KOKKOS_FUNCTION friend bool operator==(const iterator & a, const iterator & b)
    {
      return a.it == b.it;
    }
    KOKKOS_FUNCTION friend bool operator!=(const iterator & a, const iterator & b)
    {
      return a.it != b.it;
    }

  private:
    pointer it;
  };

  KOKKOS_FUNCTION iterator begin() const
  {
    KOKKOS_IF_ON_HOST(return iterator(_host_data);)

    return iterator(_device_data);
  }
  KOKKOS_FUNCTION iterator end() const
  {
    KOKKOS_IF_ON_HOST(return iterator(_host_data + _size);)

    return iterator(_device_data + _size);
  }
#endif

protected:
  index_type _n[dimension] = {0};
  index_type _s[dimension] = {0};
  signed_index_type _d[dimension] = {0};

#ifdef MOOSE_KOKKOS_SCOPE

  template <bool host, bool device, bool initialize, typename... size_type>
  void createInternal(size_type... n);
  template <bool host, bool device, bool initialize>
  void createInternal(const std::vector<index_type> & n);
  template <bool initialize>
  void createInternal(const std::vector<index_type> & n, bool host, bool device);
  template <typename TargetSpace, typename SourceSpace>
  void copyInternal(T * target, const T * source, index_type n);
  template <typename... Args>
  T & emplaceAt(index_type i, Args &&... args);
#endif

private:
#ifdef MOOSE_KOKKOS_SCOPE

  template <bool initialize>
  void allocHost();
  void allocDevice();
#endif

  void freeHost();
  void freeDevice();

  std::shared_ptr<unsigned int> _counter;
  std::shared_ptr<std::vector<bool>> _slots_constructed;
  bool _is_init = false;
  bool _is_host_alloc = false;
  bool _is_device_alloc = false;
  bool _is_host_alias = false;
  bool _is_device_alias = false;
  T * _host_data = nullptr;
  T * _device_data = nullptr;
  index_type _size = 0;
  const LayoutType _layout;
};

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::
    constructed_entry_iterator(
        typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<
            is_const>::array_type & array,
        index_type i)
  : _array(array), _i(i)
{
  advanceToConstructed();
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<is_const> &
ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator++()
{
  ++_i;
  advanceToConstructed();
  return *this;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
typename ArrayBase<T, dimension, index_type>::template constructed_entry_iterator<is_const>
ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator++(int)
{
  constructed_entry_iterator pre = *this;
  ++(*this);
  return pre;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
bool
ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator==(
    const constructed_entry_iterator & other) const
{
  return &_array == &other._array && _i == other._i;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
bool
ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::operator!=(
    const constructed_entry_iterator & other) const
{
  return !(*this == other);
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
void
ArrayBase<T, dimension, index_type>::constructed_entry_iterator<is_const>::advanceToConstructed()
{
  while (_i < _array._size && !_array.isSlotConstructed(_i))
    ++_i;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::constructed_entry_range(
    typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<
        is_const>::array_type & array)
  : _array(array)
{
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<is_const>::iterator
ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::begin() const
{
  return iterator(_array, 0);
}

template <typename T, unsigned int dimension, typename index_type>
template <bool is_const>
typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<is_const>::iterator
ArrayBase<T, dimension, index_type>::constructed_entry_range<is_const>::end() const
{
  return iterator(_array, _array._size);
}

template <typename T, unsigned int dimension, typename index_type>
typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<false>
ArrayBase<T, dimension, index_type>::constructedEntries()
{
  return constructed_entry_range<false>(*this);
}

template <typename T, unsigned int dimension, typename index_type>
typename ArrayBase<T, dimension, index_type>::template constructed_entry_range<true>
ArrayBase<T, dimension, index_type>::constructedEntries() const
{
  return constructed_entry_range<true>(*this);
}

template <typename T, unsigned int dimension, typename index_type>
bool
ArrayBase<T, dimension, index_type>::isSlotConstructed(index_type i) const
{
  mooseAssert(i < _size, "isSlotConstructed index out of bounds");

  if (!_slots_constructed)
    return _is_host_alloc;

  return (*_slots_constructed)[i];
}

#ifdef MOOSE_KOKKOS_SCOPE
template <typename T, unsigned int dimension, typename index_type>
template <typename... Args>
T &
ArrayBase<T, dimension, index_type>::emplaceAt(index_type i, Args &&... args)
{
  mooseAssert(_is_host_alloc && _slots_constructed,
              "emplaceAt requires a malloc-allocated array (create<false>)");
  mooseAssert(i < _size, "emplaceAt index out of bounds");

  new (_host_data + i) T(std::forward<Args>(args)...);
  (*_slots_constructed)[i] = true;

  return _host_data[i];
}
#endif // MOOSE_KOKKOS_SCOPE

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::freeHost()
{
  if (!_is_host_alloc)
    return;

  if (_is_host_alias)
  {
    _host_data = nullptr;
    _is_host_alias = false;
  }
  else
  {
    if (!_slots_constructed)
      // Allocated by new
      delete[] _host_data;
    else
    {
      // Allocated by malloc
      for (const auto i : make_range(_size))
        if (isSlotConstructed(i))
          _host_data[i].~T();

      std::free(_host_data);

      _slots_constructed.reset();
    }
  }

  _is_host_alloc = false;
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::freeDevice()
{
  if (!_is_device_alloc)
    return;

  if (_is_device_alias)
  {
    _device_data = nullptr;
    _is_device_alias = false;
  }
  else
    Moose::Kokkos::free(_device_data);

  _is_device_alloc = false;
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::destroy()
{
  if (!_counter)
    return;

  if (_counter.use_count() > 1)
  {
    _host_data = nullptr;
    _device_data = nullptr;
  }
  else if (_counter.use_count() == 1)
  {
    freeHost();
    freeDevice();
  }

  _size = 0;

  for (const auto i : make_range(dimension))
  {
    _n[i] = 0;
    _s[i] = 0;
    _d[i] = 0;
  }

  _is_init = false;
  _is_host_alloc = false;
  _is_device_alloc = false;
  _is_host_alias = false;
  _is_device_alias = false;

  _counter.reset();
  _slots_constructed.reset();
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::shallowCopy(const ArrayBase<T, dimension, index_type> & array)
{
  if (_layout != array._layout)
    mooseError("Kokkos array error: cannot shallow copy arrays with different layouts.");

  destroy();

  _counter = array._counter;
  _slots_constructed = array._slots_constructed;

  _size = array._size;

  for (const auto i : make_range(dimension))
  {
    _n[i] = array._n[i];
    _s[i] = array._s[i];
    _d[i] = array._d[i];
  }

  _is_init = array._is_init;
  _is_host_alloc = array._is_host_alloc;
  _is_device_alloc = array._is_device_alloc;
  _is_host_alias = array._is_host_alias;
  _is_device_alias = array._is_device_alias;

  _host_data = array._host_data;
  _device_data = array._device_data;
}

#ifdef MOOSE_KOKKOS_SCOPE
template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::aliasHost(T * ptr)
{
  if (!_is_init)
    mooseError("Kokkos array error: attempted to alias host data before array initialization.");

  if (_is_host_alloc && !_is_host_alias)
    mooseError("Kokkos array error: cannot alias host data because host data was not aliased.");

  _host_data = ptr;
  _is_host_alloc = true;
  _is_host_alias = true;
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::aliasDevice(T * ptr)
{
  if (!_is_init)
    mooseError("Kokkos array error: attempted to alias device data before array initialization.");

  if (_is_device_alloc && !_is_device_alias)
    mooseError("Kokkos array error: cannot alias device data because device data was not aliased.");

  _device_data = ptr;
  _is_device_alloc = true;
  _is_device_alias = true;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool initialize>
void
ArrayBase<T, dimension, index_type>::allocHost()
{
  if (_is_host_alloc)
    return;

  if constexpr (initialize)
  {
    static_assert(
        std::is_default_constructible<T>::value,
        "Data type is not default-constructible. Initialization argument should be set to false.");

    _host_data = new T[_size];
  }
  else
  {
    _host_data = static_cast<T *>(std::malloc(_size * sizeof(T)));
    _slots_constructed = std::make_shared<std::vector<bool>>(_size, false);
  }

  _is_host_alloc = true;
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::allocDevice()
{
  if (_is_device_alloc)
    return;

  _device_data =
      static_cast<T *>(::Kokkos::kokkos_malloc<ExecSpace::memory_space>(_size * sizeof(T)));

  _is_device_alloc = true;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool host, bool device, bool initialize>
void
ArrayBase<T, dimension, index_type>::createInternal(const std::vector<index_type> & n)
{
  if (n.size() != dimension)
    mooseError("Kokkos array error: the number of dimensions provided (",
               n.size(),
               ") must match the array dimension (",
               dimension,
               ").");

  if (_counter)
    destroy();

  _counter = std::make_shared<unsigned int>();

  uint64_t overflow_checker = 1;

  _size = 1;
  _s[0] = 1;

  for (const auto i : make_range(dimension))
  {
    overflow_checker *= n[i];

    _n[i] = n[i];
    _size *= n[i];
  }

  if (overflow_checker > std::numeric_limits<index_type>::max())
    mooseError("Kokkos array error: the dimensions provided (",
               Moose::stringify(n),
               ") has the total size of ",
               overflow_checker,
               " which exceeds the limit of ",
               MooseUtils::prettyCppType<index_type>(),
               ".");

  if (_layout == LayoutType::LEFT)
  {
    _s[0] = 1;

    for (const auto i : make_range(1u, dimension))
      _s[i] = _s[i - 1] * _n[i - 1];
  }
  else
  {
    _s[dimension - 1] = 1;

    for (int i = dimension - 2; i >= 0; --i)
      _s[i] = _s[i + 1] * _n[i + 1];
  }

  if constexpr (host)
    allocHost<initialize>();

  if constexpr (device)
    allocDevice();

  _is_init = true;
}

template <typename T, unsigned int dimension, typename index_type>
template <bool initialize>
void
ArrayBase<T, dimension, index_type>::createInternal(const std::vector<index_type> & n,
                                                    bool host,
                                                    bool device)
{
  if (host && device)
    createInternal<true, true, initialize>(n);
  else if (host && !device)
    createInternal<true, false, initialize>(n);
  else if (!host && device)
    createInternal<false, true, initialize>(n);
  else
    createInternal<false, false, initialize>(n);
}

template <typename T, unsigned int dimension, typename index_type>
template <bool host, bool device, bool initialize, typename... size_type>
void
ArrayBase<T, dimension, index_type>::createInternal(size_type... n)
{
  static_assert((std::is_convertible<size_type, index_type>::value && ...),
                "All arguments must be convertible to index_type");
  static_assert(sizeof...(n) == dimension, "Number of arguments should match array dimension");

  std::vector<index_type> dims;
  (dims.push_back(n), ...);

  createInternal<host, device, initialize>(dims);
}

template <typename T, unsigned int dimension, typename index_type>
template <typename TargetSpace, typename SourceSpace>
void
ArrayBase<T, dimension, index_type>::copyInternal(T * target, const T * source, index_type n)
{
  ::Kokkos::Impl::DeepCopy<TargetSpace, SourceSpace>(target, source, n * sizeof(T));
  ::Kokkos::fence();
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::offset(const std::vector<signed_index_type> & d)
{
  if (d.size() > dimension)
    mooseError("Kokkos array error: the number of offsets provided (",
               d.size(),
               ") cannot be larger than the array dimension (",
               dimension,
               ").");

  for (const auto i : index_range(d))
    _d[i] = d[i];
}

template <typename T, unsigned int dimension, typename index_type>
template <typename... offset_type>
void
ArrayBase<T, dimension, index_type>::offset(offset_type... d)
{
  static_assert((std::is_convertible<offset_type, signed_index_type>::value && ...),
                "All arguments must be convertible to signed_index_type");
  static_assert(sizeof...(d) == dimension, "Number of arguments should match array dimension");

  std::vector<signed_index_type> offsets;
  (offsets.push_back(d), ...);

  offset(offsets);
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::copyToDevice()
{
  // If host side memory is not allocated, do nothing
  if (!_is_host_alloc)
    return;

  // If device side memory is not allocated,
  if (!_is_device_alloc)
  {
    if (_counter.use_count() == 1)
      // allocate memory if this array is not shared with other arrays
      allocDevice();
    else
      // print error if this array is shared with other arrays
      mooseError("Kokkos array error: cannot copy from host to device because device memory "
                 "was not allocated. Cannot allocate device memory for copy because the array is "
                 "being shared.");
  }

  // Copy from host to device
  copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data, _host_data, _size);
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::copyToHost()
{
  // If device side memory is not allocated, do nothing
  if (!_is_device_alloc)
    return;

  // If host side memory is not allocated,
  if (!_is_host_alloc)
  {
    if (_counter.use_count() == 1)
      // allocate memory if this array is not shared with other arrays
      allocHost<false>();
    else
      // print error if this array is shared with other arrays
      mooseError("Kokkos array error: cannot copy from device to host because host memory "
                 "was not allocated. Cannot allocate host memory for copy because the array is "
                 "being shared.");
  }

  // Copy from device to host
  copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data, _device_data, _size);
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::moveToDevice(bool should_free_host)
{
  static_assert(!is_kokkos_array<T>::value,
                "moveToDevice() not allowed for a nested array whose data type is another array.");

  if (should_free_host && _counter.use_count() > 1)
    mooseError("Kokkos array error: cannot move array from host to device because there is at "
               "least one shallow copy of this array still alive.");

  copyToDevice();

  if (_counter.use_count() == 1)
    freeHost();
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::moveToHost(bool should_free_device)
{
  if (should_free_device && _counter.use_count() > 1)
    mooseError("Kokkos array error: cannot move array from device to host because there is at "
               "least one shallow copy of this array still alive.");

  copyToHost();

  if (_counter.use_count() == 1)
    freeDevice();
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::copyIn(const T * ptr,
                                            MemcpyType dir,
                                            index_type n,
                                            index_type offset)
{
  if (n > _size)
    mooseError("Kokkos array error: cannot copy in data larger than the array size.");

  if (offset > _size)
    mooseError("Kokkos array error: offset cannot be larger than the array size.");

  if (dir == MemcpyType::HOST_TO_HOST)
  {
    // If host side memory is not allocated, print error
    if (!_is_host_alloc)
      mooseError(
          "Kokkos array error: cannot copy in to the array because host memory was not allocated.");

    // Copy from host to host
    copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(_host_data + offset, ptr, n);
  }
  else if (dir == MemcpyType::HOST_TO_DEVICE)
  {
    // If device side memory is not allocated, print error
    if (!_is_device_alloc)
      mooseError("Kokkos array error: cannot copy in to the array because device memory was not "
                 "allocated.");

    // Copy from host to device
    copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data + offset, ptr, n);
  }
  else if (dir == MemcpyType::DEVICE_TO_HOST)
  {
    // If host side memory is not allocated, print error
    if (!_is_host_alloc)
      mooseError(
          "Kokkos array error: cannot copy in to the array because host memory was not allocated.");

    // Copy from device to host
    copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data + offset, ptr, n);
  }
  else if (dir == MemcpyType::DEVICE_TO_DEVICE)
  {
    // If device side memory is not allocated, print error
    if (!_is_device_alloc)
      mooseError("Kokkos array error: cannot copy in to the array because device memory was not "
                 "allocated.");

    // Copy from device to device
    copyInternal<MemSpace, MemSpace>(_device_data + offset, ptr, n);
  }
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::copyOut(T * ptr,
                                             MemcpyType dir,
                                             index_type n,
                                             index_type offset)
{
  if (n > _size)
    mooseError("Kokkos array error: cannot copy out data larger than the array size.");

  if (offset > _size)
    mooseError("Kokkos array error: offset cannot be larger than the array size.");

  if (dir == MemcpyType::HOST_TO_HOST)
  {
    // If host side memory is not allocated, print error
    if (!_is_host_alloc)
      mooseError("Kokkos array error: cannot copy out from the array because host memory was not "
                 "allocated.");

    // Copy from host to host
    copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
  }
  else if (dir == MemcpyType::HOST_TO_DEVICE)
  {
    // If host side memory is not allocated, print error
    if (!_is_host_alloc)
      mooseError("Kokkos array error: cannot copy out from the array because host memory was not "
                 "allocated.");

    // Copy from host to device
    copyInternal<MemSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
  }
  else if (dir == MemcpyType::DEVICE_TO_HOST)
  {
    // If device side memory is not allocated, print error
    if (!_is_device_alloc)
      mooseError("Kokkos array error: cannot copy out from the array because device memory was not "
                 "allocated.");

    // Copy from device to host
    copyInternal<::Kokkos::HostSpace, MemSpace>(ptr, _device_data + offset, n);
  }
  else if (dir == MemcpyType::DEVICE_TO_DEVICE)
  {
    // If device side memory is not allocated, print error
    if (!_is_device_alloc)
      mooseError("Kokkos array error: cannot copy out from the array because device memory was not "
                 "allocated.");

    // Copy from device to device
    copyInternal<MemSpace, MemSpace>(ptr, _device_data + offset, n);
  }
}

template <typename T>
void
copyToDeviceInner(T & /* data */)
{
}

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
void
copyToDeviceInner(Array<T, dimension, index_type, layout> & data)
{
  data.copyToDeviceNested();
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::copyToDeviceNested()
{
  for (const auto i : make_range(_size))
    copyToDeviceInner(_host_data[i]);

  copyToDevice();
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::deepCopy(const ArrayBase<T, dimension, index_type> & array)
{
  if (_layout != array._layout)
    mooseError("Kokkos array error: cannot deep copy arrays with different layouts.");

  if (ArrayDeepCopy<T>::value && !array._is_host_alloc)
    mooseError(
        "Kokkos array error: cannot deep copy using constructor from array without host data.");

  std::vector<index_type> n(std::begin(array._n), std::end(array._n));

  createInternal<false>(n, array._is_host_alloc, array._is_device_alloc);

  if (_is_host_alloc)
    for (const auto i : make_range(_size))
      if (array.isSlotConstructed(i))
        emplaceAt(i, array._host_data[i]);

  if (ArrayDeepCopy<T>::value)
    copyToDevice();
  else if (_is_device_alloc)
    copyInternal<MemSpace, MemSpace>(_device_data, array._device_data, _size);

  for (const auto i : make_range(dimension))
  {
    _d[i] = array._d[i];
    _s[i] = array._s[i];
  }
}

template <typename T, unsigned int dimension, typename index_type>
void
ArrayBase<T, dimension, index_type>::swap(ArrayBase<T, dimension, index_type> & array)
{
  ArrayBase<T, dimension, index_type> clone(_layout);

  clone.shallowCopy(*this);
  this->shallowCopy(array);
  array.shallowCopy(clone);
}

template <typename T, unsigned int dimension, typename index_type>
auto &
ArrayBase<T, dimension, index_type>::operator=(const T & scalar)
{
  if (_is_host_alloc)
    std::fill_n(_host_data, _size, scalar);

  if (_is_device_alloc)
    ::Kokkos::Experimental::fill_n(ExecSpace(), deviceView(), _size, scalar);

  return *this;
}

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
void
dataStore(std::ostream & stream, Array<T, dimension, index_type, layout> & array, void * context)
{
  using ::dataStore;

  bool is_alloc = array.isAlloc();
  dataStore(stream, is_alloc, nullptr);

  if (!is_alloc)
    return;

  std::string type = typeid(T).name();
  dataStore(stream, type, nullptr);

  unsigned int dim = dimension;
  dataStore(stream, dim, nullptr);

  for (const auto dim : make_range(dimension))
  {
    auto n = array.n(dim);
    dataStore(stream, n, nullptr);
  }

  if (array.isDeviceAlloc())
  {
    // We use malloc/free because we just want a memory copy
    // If T is a Kokkos array and we use new/delete or vector to copy it out,
    // the arrays will be destroyed on cleanup

    T * data = static_cast<T *>(std::malloc(array.size() * sizeof(T)));

    array.copyOut(data, MemcpyType::DEVICE_TO_HOST, array.size());

    for (const auto i : make_range(array.size()))
      dataStore(stream, data[i], context);

    std::free(data);
  }
  else
    for (auto & value : array)
      dataStore(stream, value, context);
}

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
void
dataLoad(std::istream & stream, Array<T, dimension, index_type, layout> & array, void * context)
{
  using ::dataLoad;

  bool is_alloc;
  dataLoad(stream, is_alloc, nullptr);

  if (!is_alloc)
    return;

  std::string from_type_name;
  dataLoad(stream, from_type_name, nullptr);

  if (from_type_name != typeid(T).name())
    mooseError("Kokkos array error: cannot load array because the stored array is of type '",
               MooseUtils::prettyCppType(libMesh::demangle(from_type_name.c_str())),
               "' but the loading array is of type '",
               MooseUtils::prettyCppType(libMesh::demangle(typeid(T).name())),
               "'.");

  unsigned int from_dimension;
  dataLoad(stream, from_dimension, nullptr);

  if (from_dimension != dimension)
    mooseError("Kokkos array error: cannot load array because the stored array is ",
               from_dimension,
               "D but the loading array is ",
               dimension,
               "D.");

  std::vector<index_type> from_n(dimension);
  std::vector<index_type> n(dimension);

  for (const auto dim : make_range(dimension))
  {
    dataLoad(stream, from_n[dim], nullptr);
    n[dim] = array.n(dim);
  }

  if (from_n != n)
    mooseError("Kokkos array error: cannot load array because the stored array has dimensions (",
               Moose::stringify(from_n),
               ") but the loading array has dimensions (",
               Moose::stringify(n),
               ").");

  if (array.isHostAlloc())
  {
    for (auto & value : array)
      dataLoad(stream, value, context);

    if (array.isDeviceAlloc())
      array.copyToDevice();
  }
  else
  {
    std::vector<T> data(array.size());

    for (auto & value : data)
      dataLoad(stream, value, context);

    array.copyIn(data.data(), MemcpyType::HOST_TO_DEVICE, array.size());
  }
}
#endif

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
class Array : public ArrayBase<T, dimension, index_type>
{
#ifdef MOOSE_KOKKOS_SCOPE
  usingKokkosArrayBaseMembers(T, dimension, index_type);
#endif

public:
  Array() : ArrayBase<T, dimension, index_type>(layout) {}
  Array(const Array<T, dimension, index_type, layout> & array)
    : ArrayBase<T, dimension, index_type>(array)
  {
  }
#ifdef MOOSE_KOKKOS_SCOPE

  template <typename... size_type>
  Array(size_type... n) : ArrayBase<T, dimension, index_type>(layout, n...)
  {
  }
#endif

  auto & operator=(const Array<T, dimension, index_type, layout> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  template <typename... indices>
  KOKKOS_FUNCTION T & operator()(indices... i) const
  {
    return this->operator[](linearIndex(i...));
  }
  KOKKOS_FUNCTION T & operator()(const signed_index_type (&idx)[dimension]) const
  {
    return this->operator[](linearIndex(idx));
  }
  template <typename indices, typename... Args>
  T & emplace(const indices (&idx)[dimension], Args &&... args);
#endif

private:
#ifdef MOOSE_KOKKOS_SCOPE

  template <typename... indices>
  KOKKOS_FUNCTION index_type linearIndex(indices... i) const;
  template <typename indices>
  KOKKOS_FUNCTION index_type linearIndex(const indices (&idx)[dimension]) const
  {
    return linearIndexHelper(idx, std::make_integer_sequence<unsigned int, dimension>{});
  }
  template <typename indices, unsigned int... i>
  KOKKOS_FUNCTION index_type linearIndexHelper(const indices (&idx)[dimension],
                                               std::integer_sequence<unsigned int, i...>) const
  {
    return linearIndex(idx[i]...);
  }
#endif
};

#ifdef MOOSE_KOKKOS_SCOPE
template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
template <typename indices, typename... Args>
T &
Array<T, dimension, index_type, layout>::emplace(const indices (&idx)[dimension], Args &&... args)
{
  return this->emplaceAt(linearIndex(idx), std::forward<Args>(args)...);
}

template <typename T, unsigned int dimension, typename index_type, LayoutType layout>
template <typename... indices>
KOKKOS_FUNCTION index_type
Array<T, dimension, index_type, layout>::linearIndex(indices... i) const
{
  static_assert((std::is_convertible<indices, signed_index_type>::value && ...),
                "All arguments must be convertible to signed_index_type");
  static_assert(sizeof...(i) == dimension, "Number of arguments should match array dimension");

#ifndef NDEBUG
  {
    signed_index_type idx[dimension] = {static_cast<signed_index_type>(i)...};

    for (unsigned int d = 0; d < sizeof...(i); ++d)
      KOKKOS_ASSERT(idx[d] - _d[d] >= 0 && static_cast<index_type>(idx[d] - _d[d]) < _n[d]);
  }
#endif

  index_type idx = 0;
  unsigned int d = 0;

  if constexpr (layout == LayoutType::LEFT)
    (((idx += (d == 0 ? static_cast<signed_index_type>(i) - _d[d]
                      : (static_cast<signed_index_type>(i) - _d[d]) * _s[d])),
      ++d),
     ...);
  else
    (((idx += (d == dimension - 1 ? static_cast<signed_index_type>(i) - _d[d]
                                  : (static_cast<signed_index_type>(i) - _d[d]) * _s[d])),
      ++d),
     ...);

  return idx;
}
#endif

template <typename T, typename index_type>
class Array<T, 1, index_type, LayoutType::LEFT> : public ArrayBase<T, 1, index_type>
{
#ifdef MOOSE_KOKKOS_SCOPE
  usingKokkosArrayBaseMembers(T, 1, index_type);
#endif

public:
  Array() : ArrayBase<T, 1, index_type>(LayoutType::LEFT) {}
  Array(const Array<T, 1, index_type, LayoutType::LEFT> & array)
    : ArrayBase<T, 1, index_type>(array)
  {
  }
#ifdef MOOSE_KOKKOS_SCOPE

  Array(index_type n) : ArrayBase<T, 1, index_type>(LayoutType::LEFT, n) {}
  Array(const std::vector<T> & vector) : ArrayBase<T, 1, index_type>(LayoutType::LEFT)
  {
    *this = vector;
  }
#endif

  auto & operator=(const Array<T, 1, index_type, LayoutType::LEFT> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  template <bool host, bool device>
  void copyVector(const std::vector<T> & vector)
  {
    this->template createInternal<host, device, false>({static_cast<index_type>(vector.size())});

    if (host)
    {
      if constexpr (std::is_trivially_copyable<T>())
        std::memcpy(this->hostData(), vector.data(), this->size() * sizeof(T));
      else
        std::copy(vector.begin(), vector.end(), this->begin());
    }

    if (device)
      this->template copyInternal<MemSpace, ::Kokkos::HostSpace>(
          this->deviceData(), vector.data(), this->size());
  }
  template <bool host, bool device>
  void copySet(const std::set<T> & set)
  {
    std::vector<T> vector(set.begin(), set.end());

    copyVector<host, device>(vector);
  }

  auto & operator=(const std::vector<T> & vector)
  {
    copyVector<true, true>(vector);

    return *this;
  }
  auto & operator=(const std::set<T> & set)
  {
    copySet<true, true>(set);

    return *this;
  }
  KOKKOS_FUNCTION T & operator()(signed_index_type i) const
  {
    return this->operator[](linearIndex(i));
  }
  template <typename indices, typename... Args>
  T & emplace(const indices (&idx)[1], Args &&... args);

  void axby(const T a,
            const Array<T, 1, index_type, LayoutType::LEFT> & x,
            const char op,
            const T b,
            const Array<T, 1, index_type, LayoutType::LEFT> & y,
            const bool accumulate = false);
  void scal(const T a, const Array<T, 1, index_type, LayoutType::LEFT> & x);
  void scal(const T a);
  T dot(const Array<T, 1, index_type, LayoutType::LEFT> & x);
  T nrm2();

private:
  KOKKOS_FUNCTION index_type linearIndex(signed_index_type i) const;
#endif
};

#ifdef MOOSE_KOKKOS_SCOPE
template <typename T, typename index_type>
template <typename indices, typename... Args>
T &
Array<T, 1, index_type, LayoutType::LEFT>::emplace(const indices (&idx)[1], Args &&... args)
{
  return this->emplaceAt(linearIndex(idx[0]), std::forward<Args>(args)...);
}

template <typename T, typename index_type>
KOKKOS_FUNCTION index_type
Array<T, 1, index_type, LayoutType::LEFT>::linearIndex(
    typename Array<T, 1, index_type, LayoutType::LEFT>::signed_index_type i) const
{
  KOKKOS_ASSERT(i - _d[0] >= 0 && static_cast<index_type>(i - _d[0]) < _n[0]);

  return i - _d[0];
}
#endif

template <typename T, typename index_type = MOOSE_KOKKOS_INDEX_TYPE>
using Array1D = Array<T, 1, index_type, LayoutType::LEFT>;
template <typename T,
          typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
          LayoutType layout = LayoutType::LEFT>
using Array2D = Array<T, 2, index_type, layout>;
template <typename T,
          typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
          LayoutType layout = LayoutType::LEFT>
using Array3D = Array<T, 3, index_type, layout>;
template <typename T,
          typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
          LayoutType layout = LayoutType::LEFT>
using Array4D = Array<T, 4, index_type, layout>;
template <typename T,
          typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
          LayoutType layout = LayoutType::LEFT>
using Array5D = Array<T, 5, index_type, layout>;

} // namespace Moose::Kokkos