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"

#define usingKokkosArrayBaseMembers(T, dimension)                                                  \
private:                                                                                           \
  using ArrayBase<T, dimension>::_n;                                                               \
  using ArrayBase<T, dimension>::_s;                                                               \
  using ArrayBase<T, dimension>::_d;                                                               \
                                                                                                   \
public:                                                                                            \
  using ArrayBase<T, dimension>::create;                                                           \
  using ArrayBase<T, dimension>::createHost;                                                       \
  using ArrayBase<T, dimension>::createDevice;                                                     \
  using ArrayBase<T, dimension>::offset;                                                           \
  using ArrayBase<T, dimension>::operator=

namespace Moose
{
namespace 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 MemcpyKind
{
  HOST_TO_HOST,
  HOST_TO_DEVICE,
  DEVICE_TO_HOST,
  DEVICE_TO_DEVICE
};

template <typename T, unsigned int dimension = 1>
class Array;

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

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

template <typename T, unsigned int dimension>
class ArrayBase
{
public:
  ArrayBase() = default;

  ArrayBase(const ArrayBase<T, dimension> & array)
  {
#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> & array);

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

#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 dof_id_type size() const { return _size; }
  KOKKOS_FUNCTION dof_id_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[](dof_id_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; }
  void create(const std::vector<dof_id_type> & n) { createInternal<true, true>(n); }
  void createHost(const std::vector<dof_id_type> & n) { createInternal<true, false>(n); }
  void createDevice(const std::vector<dof_id_type> & n) { createInternal<false, true>(n); }
  void aliasHost(T * ptr);
  void aliasDevice(T * ptr);
  void offset(const std::vector<dof_id_signed_type> & d);
  void copyToDevice();
  void copyToHost();
  void copyIn(const T * ptr, MemcpyKind dir, dof_id_type n, dof_id_type offset = 0);
  void copyOut(T * ptr, MemcpyKind dir, dof_id_type n, dof_id_type offset = 0);
  void copyToDeviceNested();
  void deepCopy(const ArrayBase<T, dimension> & array);
  void swap(ArrayBase<T, dimension> & array);

  auto & operator=(const T & scalar);

  class iterator
  {
  public:
    KOKKOS_FUNCTION iterator(T * it) : it(it) {}
    KOKKOS_FUNCTION bool operator==(const iterator & other) const { return it == other.it; }
    KOKKOS_FUNCTION bool operator!=(const iterator & other) const { return it != other.it; }
    KOKKOS_FUNCTION T & operator*() const { return *it; }
    KOKKOS_FUNCTION T * operator&() const { return it; }
    KOKKOS_FUNCTION iterator & operator++()
    {
      ++it;
      return *this;
    }
    KOKKOS_FUNCTION iterator operator++(int)
    {
      iterator pre = *this;
      ++it;
      return pre;
    }

  private:
    T * 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:
  dof_id_type _n[dimension] = {0};
  dof_id_type _s[dimension] = {0};
  dof_id_signed_type _d[dimension] = {0};

#ifdef MOOSE_KOKKOS_SCOPE

  template <bool host, bool device>
  void createInternal(const std::vector<dof_id_type> & n);
  void createInternal(const std::vector<dof_id_type> & n, bool host, bool device);
  template <typename TargetSpace, typename SourceSpace>
  void copyInternal(T * target, const T * source, dof_id_type n);
#endif

private:
#ifdef MOOSE_KOKKOS_SCOPE

  void allocHost();
  void allocDevice();
#endif

  std::shared_ptr<unsigned int> _counter;
  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;
  dof_id_type _size = 0;
};

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

  if (_counter.use_count() > 1)
  {
    _host_data = nullptr;
    _device_data = nullptr;
  }
  else if (_counter.use_count() == 1)
  {
    if (_is_host_alloc && !_is_host_alias)
    {
      if constexpr (std::is_default_constructible<T>::value)
        // Allocated by new
        delete[] _host_data;
      else
      {
        // Allocated by malloc
        for (dof_id_type i = 0; i < _size; ++i)
          _host_data[i].~T();

        std::free(_host_data);
      }
    }

    if (_is_device_alloc && !_is_device_alias)
      Moose::Kokkos::free(_device_data);
  }

  _size = 0;

  for (unsigned int i = 0; i < dimension; ++i)
  {
    _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();
}

template <typename T, unsigned int dimension>
void
ArrayBase<T, dimension>::shallowCopy(const ArrayBase<T, dimension> & array)
{
  destroy();

  _counter = array._counter;

  _size = array._size;

  for (unsigned int i = 0; i < dimension; ++i)
  {
    _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>
void
ArrayBase<T, dimension>::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>
void
ArrayBase<T, dimension>::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>
void
ArrayBase<T, dimension>::allocHost()
{
  if (_is_host_alloc)
    return;

  if constexpr (std::is_default_constructible<T>::value)
    _host_data = new T[_size];
  else
    _host_data = static_cast<T *>(std::malloc(_size * sizeof(T)));

  _is_host_alloc = true;
}

template <typename T, unsigned int dimension>
void
ArrayBase<T, dimension>::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>
template <bool host, bool device>
void
ArrayBase<T, dimension>::createInternal(const std::vector<dof_id_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>();

  _size = 1;
  _s[0] = 1;

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

    if (i)
      _s[i] = _s[i - 1] * _n[i - 1];
  }

  if constexpr (host)
    allocHost();

  if constexpr (device)
    allocDevice();

  _is_init = true;
}

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

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

template <typename T, unsigned int dimension>
void
ArrayBase<T, dimension>::offset(const std::vector<dof_id_signed_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>
void
ArrayBase<T, dimension>::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 side memory "
                 "was not allocated and array is being shared with other arrays.");
  }

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

template <typename T, unsigned int dimension>
void
ArrayBase<T, dimension>::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();
    else
      // print error if this array is shared with other arrays
      mooseError("Kokkos array error: cannot copy from device to host because host side memory "
                 "was not allocated and array is being shared with other arrays.");
  }

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

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

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

  if (dir == MemcpyKind::HOST_TO_HOST)
  {
    // If host side memory is not allocated, do nothing
    if (!_is_host_alloc)
      return;

    // Copy from host to host
    copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(_host_data + offset, ptr, n);
  }
  else if (dir == MemcpyKind::HOST_TO_DEVICE)
  {
    // If device side memory is not allocated, do nothing
    if (!_is_device_alloc)
      return;

    // Copy from host to device
    copyInternal<MemSpace, ::Kokkos::HostSpace>(_device_data + offset, ptr, n);
  }
  else if (dir == MemcpyKind::DEVICE_TO_HOST)
  {
    // If host side memory is not allocated, do nothing
    if (!_is_host_alloc)
      return;

    // Copy from device to host
    copyInternal<::Kokkos::HostSpace, MemSpace>(_host_data + offset, ptr, n);
  }
  else if (dir == MemcpyKind::DEVICE_TO_DEVICE)
  {
    // If device side memory is not allocated, do nothing
    if (!_is_device_alloc)
      return;

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

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

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

  if (dir == MemcpyKind::HOST_TO_HOST)
  {
    // If host side memory is not allocated, do nothing
    if (!_is_host_alloc)
      return;

    // Copy from host to host
    copyInternal<::Kokkos::HostSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
  }
  else if (dir == MemcpyKind::HOST_TO_DEVICE)
  {
    // If host side memory is not allocated, do nothing
    if (!_is_host_alloc)
      return;

    // Copy from host to device
    copyInternal<MemSpace, ::Kokkos::HostSpace>(ptr, _host_data + offset, n);
  }
  else if (dir == MemcpyKind::DEVICE_TO_HOST)
  {
    // If device side memory is not allocated, do nothing
    if (!_is_device_alloc)
      return;

    // Copy from device to host
    copyInternal<::Kokkos::HostSpace, MemSpace>(ptr, _device_data + offset, n);
  }
  else if (dir == MemcpyKind::DEVICE_TO_DEVICE)
  {
    // If device side memory is not allocated, do nothing
    if (!_is_device_alloc)
      return;

    // 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>
void
copyToDeviceInner(Array<T, dimension> & data)
{
  data.copyToDeviceNested();
}

template <typename T, unsigned int dimension>
void
ArrayBase<T, dimension>::copyToDeviceNested()
{
  for (unsigned int i = 0; i < _size; ++i)
    copyToDeviceInner(_host_data[i]);

  copyToDevice();
}

template <typename T, unsigned int dimension>
void
ArrayBase<T, dimension>::deepCopy(const ArrayBase<T, dimension> & array)
{
  if (ArrayDeepCopy<T>::value && !array._is_host_alloc)
    mooseError(
        "Kokkos array error: cannot deep copy using constructor from array without host data.");

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

  createInternal(n, array._is_host_alloc, array._is_device_alloc);

  if constexpr (ArrayDeepCopy<T>::value)
  {
    for (dof_id_type i = 0; i < _size; ++i)
      new (_host_data + i) T(array._host_data[i]);

    copyToDevice();
  }
  else
  {
    if (_is_host_alloc)
      std::memcpy(_host_data, array._host_data, _size * sizeof(T));

    if (_is_device_alloc)
      copyInternal<MemSpace, MemSpace>(_device_data, array._device_data, _size);
  }

  for (unsigned int i = 0; i < dimension; ++i)
  {
    _d[i] = array._d[i];
    _s[i] = array._s[i];
  }
}

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

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

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

  if (_is_device_alloc)
  {
    ::Kokkos::View<T *, MemSpace, ::Kokkos::MemoryTraits<::Kokkos::Unmanaged>> data(_device_data,
                                                                                    _size);
    ::Kokkos::Experimental::fill_n(ExecSpace(), data, _size, scalar);
  }

  return *this;
}

template <typename T, unsigned int dimension>
void
dataStore(std::ostream & stream, Array<T, dimension> & 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 (unsigned int dim = 0; dim < dimension; ++dim)
  {
    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, MemcpyKind::DEVICE_TO_HOST, array.size());

    for (dof_id_type i = 0; i < array.size(); ++i)
      dataStore(stream, data[i], context);

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

template <typename T, unsigned int dimension>
void
dataLoad(std::istream & stream, Array<T, dimension> & 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 an 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 an array because the stored array is ",
               from_dimension,
               "D but the loading array is ",
               dimension,
               "D.");

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

  for (unsigned int dim = 0; dim < dimension; ++dim)
  {
    dataLoad(stream, from_n[dim], nullptr);
    n[dim] = array.n(dim);
  }

  if (from_n != n)
    mooseError("Kokkos array error: cannot load an 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(), MemcpyKind::HOST_TO_DEVICE, array.size());
  }
}
#endif

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

public:
  Array() = default;
  Array(const Array<T, 1> & array) : ArrayBase<T, 1>(array) {}
  auto & operator=(const Array<T, 1> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  Array(dof_id_type n0) { create(n0); }
  Array(const std::vector<T> & vector) { *this = vector; }

  void init(dof_id_type n0) { this->template createInternal<false, false>({n0}); }
  void create(dof_id_type n0) { this->template createInternal<true, true>({n0}); }
  void createHost(dof_id_type n0) { this->template createInternal<true, false>({n0}); }
  void createDevice(dof_id_type n0) { this->template createInternal<false, true>({n0}); }
  void offset(dof_id_signed_type d0) { ArrayBase<T, 1>::offset({d0}); }

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

    if (host)
      std::memcpy(this->hostData(), vector.data(), this->size() * sizeof(T));

    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()(dof_id_signed_type i0) const
  {
    KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);

    return this->operator[](i0 - _d[0]);
  }


  template <typename U = T>
  typename std::enable_if<std::is_same<U, Real>::value, void>::type
  axby(const U a,
       const Array<U, 1> & x,
       const char op,
       const U b,
       const Array<U, 1> & y,
       const bool accumulate = false);
  template <typename U = T>
  typename std::enable_if<std::is_same<U, Real>::value, void>::type scal(const U a,
                                                                         const Array<U, 1> & x);
  template <typename U = T>
  typename std::enable_if<std::is_same<U, Real>::value, void>::type scal(const U a);
  template <typename U = T>
  typename std::enable_if<std::is_same<U, Real>::value, Real>::type dot(const Array<U, 1> & x);
  template <typename U = T>
  typename std::enable_if<std::is_same<U, Real>::value, Real>::type nrm2();
#endif
};

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

public:
  Array() = default;
  Array(const Array<T, 2> & array) : ArrayBase<T, 2>(array) {}
  auto & operator=(const Array<T, 2> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  Array(dof_id_type n0, dof_id_type n1) { create(n0, n1); }

  void init(dof_id_type n0, dof_id_type n1)
  {
    this->template createInternal<false, false>({n0, n1});
  }
  void create(dof_id_type n0, dof_id_type n1)
  {
    this->template createInternal<true, true>({n0, n1});
  }
  void createHost(dof_id_type n0, dof_id_type n1)
  {
    this->template createInternal<true, false>({n0, n1});
  }
  void createDevice(dof_id_type n0, dof_id_type n1)
  {
    this->template createInternal<false, true>({n0, n1});
  }
  void offset(dof_id_signed_type d0, dof_id_signed_type d1) { ArrayBase<T, 2>::offset({d0, d1}); }

  KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0, dof_id_signed_type i1) const
  {
    KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
    KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);

    return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1]);
  }
#endif
};

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

public:
  Array() = default;
  Array(const Array<T, 3> & array) : ArrayBase<T, 3>(array) {}
  auto & operator=(const Array<T, 3> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  Array(dof_id_type n0, dof_id_type n1, dof_id_type n2) { create(n0, n1, n2); }

  void init(dof_id_type n0, dof_id_type n1, dof_id_type n2)
  {
    this->template createInternal<false, false>({n0, n1, n2});
  }
  void create(dof_id_type n0, dof_id_type n1, dof_id_type n2)
  {
    this->template createInternal<true, true>({n0, n1, n2});
  }
  void createHost(dof_id_type n0, dof_id_type n1, dof_id_type n2)
  {
    this->template createInternal<true, false>({n0, n1, n2});
  }
  void createDevice(dof_id_type n0, dof_id_type n1, dof_id_type n2)
  {
    this->template createInternal<false, true>({n0, n1, n2});
  }
  void offset(dof_id_signed_type d0, dof_id_signed_type d1, dof_id_signed_type d2)
  {
    ArrayBase<T, 3>::offset({d0, d1, d2});
  }

  KOKKOS_FUNCTION T &
  operator()(dof_id_signed_type i0, dof_id_signed_type i1, dof_id_signed_type i2) const
  {
    KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
    KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
    KOKKOS_ASSERT(i2 - _d[2] >= 0 && static_cast<dof_id_type>(i2 - _d[2]) < _n[2]);

    return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1] + (i2 - _d[2]) * _s[2]);
  }
#endif
};

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

public:
  Array() = default;
  Array(const Array<T, 4> & array) : ArrayBase<T, 4>(array) {}
  auto & operator=(const Array<T, 4> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  Array(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3) { create(n0, n1, n2, n3); }

  void init(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
  {
    this->template createInternal<false, false>({n0, n1, n2, n3});
  }
  void create(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
  {
    this->template createInternal<true, true>({n0, n1, n2, n3});
  }
  void createHost(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
  {
    this->template createInternal<true, false>({n0, n1, n2, n3});
  }
  void createDevice(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3)
  {
    this->template createInternal<false, true>({n0, n1, n2, n3});
  }
  void
  offset(dof_id_signed_type d0, dof_id_signed_type d1, dof_id_signed_type d2, dof_id_signed_type d3)
  {
    ArrayBase<T, 4>::offset({d0, d1, d2, d3});
  }

  KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0,
                                 dof_id_signed_type i1,
                                 dof_id_signed_type i2,
                                 dof_id_signed_type i3) const
  {
    KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
    KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
    KOKKOS_ASSERT(i2 - _d[2] >= 0 && static_cast<dof_id_type>(i2 - _d[2]) < _n[2]);
    KOKKOS_ASSERT(i3 - _d[3] >= 0 && static_cast<dof_id_type>(i3 - _d[3]) < _n[3]);

    return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1] + (i2 - _d[2]) * _s[2] +
                            (i3 - _d[3]) * _s[3]);
  }
#endif
};

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

public:
  Array() = default;
  Array(const Array<T, 5> & array) : ArrayBase<T, 5>(array) {}
  auto & operator=(const Array<T, 5> & array)
  {
    this->shallowCopy(array);

    return *this;
  }

#ifdef MOOSE_KOKKOS_SCOPE

  Array(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
  {
    create(n0, n1, n2, n3, n4);
  }

  void init(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
  {
    this->template createInternal<false, false>({n0, n1, n2, n3, n4});
  }
  void create(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
  {
    this->template createInternal<true, true>({n0, n1, n2, n3, n4});
  }
  void createHost(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
  {
    this->template createInternal<true, false>({n0, n1, n2, n3, n4});
  }
  void createDevice(dof_id_type n0, dof_id_type n1, dof_id_type n2, dof_id_type n3, dof_id_type n4)
  {
    this->template createInternal<false, true>({n0, n1, n2, n3, n4});
  }
  void offset(dof_id_signed_type d0,
              dof_id_signed_type d1,
              dof_id_signed_type d2,
              dof_id_signed_type d3,
              dof_id_signed_type d4)
  {
    ArrayBase<T, 5>::offset({d0, d1, d2, d3, d4});
  }

  KOKKOS_FUNCTION T & operator()(dof_id_signed_type i0,
                                 dof_id_signed_type i1,
                                 dof_id_signed_type i2,
                                 dof_id_signed_type i3,
                                 dof_id_signed_type i4) const
  {
    KOKKOS_ASSERT(i0 - _d[0] >= 0 && static_cast<dof_id_type>(i0 - _d[0]) < _n[0]);
    KOKKOS_ASSERT(i1 - _d[1] >= 0 && static_cast<dof_id_type>(i1 - _d[1]) < _n[1]);
    KOKKOS_ASSERT(i2 - _d[2] >= 0 && static_cast<dof_id_type>(i2 - _d[2]) < _n[2]);
    KOKKOS_ASSERT(i3 - _d[3] >= 0 && static_cast<dof_id_type>(i3 - _d[3]) < _n[3]);
    KOKKOS_ASSERT(i4 - _d[4] >= 0 && static_cast<dof_id_type>(i4 - _d[4]) < _n[4]);

    return this->operator[](i0 - _d[0] + (i1 - _d[1]) * _s[1] + (i2 - _d[2]) * _s[2] +
                            (i3 - _d[3]) * _s[3] + (i4 - _d[4]) * _s[4]);
  }
#endif
};

template <typename T>
using Array1D = Array<T, 1>;
template <typename T>
using Array2D = Array<T, 2>;
template <typename T>
using Array3D = Array<T, 3>;
template <typename T>
using Array4D = Array<T, 4>;
template <typename T>
using Array5D = Array<T, 5>;

} // namespace Kokkos
} // namespace Moose