KokkosJaggedArray.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

#include "KokkosArray.h"

namespace Moose::Kokkos
{

template <unsigned int inner>
struct JaggedArrayInnerDim
{
  unsigned int dim[inner] = {0};
  unsigned int stride[inner + 1] = {0};

#ifdef MOOSE_KOKKOS_SCOPE

  KOKKOS_FUNCTION unsigned int operator[](unsigned int i) const { return dim[i]; }
#endif
};

#ifdef MOOSE_KOKKOS_SCOPE

template <typename T, unsigned int inner, LayoutType layout>
class JaggedArrayInnerData
{
public:
  KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<inner> dim)
    : _data(data), _dim(dim)
  {
  }

  KOKKOS_FUNCTION unsigned int size() const { return _dim.stride[inner]; }
  KOKKOS_FUNCTION unsigned int n(unsigned int dim) const { return _dim[dim]; }
  KOKKOS_FUNCTION T & operator[](unsigned int i) const
  {
    KOKKOS_ASSERT(i < _dim.stride[inner]);

    return _data[i];
  }
  template <typename... indices>
  KOKKOS_FUNCTION T & operator()(indices... i) const;

protected:
  T * _data;
  JaggedArrayInnerDim<inner> _dim;
};

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

#ifndef NDEBUG
  {
    unsigned int idx[inner] = {static_cast<unsigned int>(i)...};

    for (unsigned int d = 0; d < sizeof...(i); ++d)
      KOKKOS_ASSERT(idx[d] < _dim[d]);
  }
#endif

  unsigned int idx = 0;
  unsigned int d = 0;

  if constexpr (layout == LayoutType::LEFT)
    (((idx +=
       (d == 0 ? static_cast<unsigned int>(i) : static_cast<unsigned int>(i) * _dim.stride[d])),
      ++d),
     ...);
  else
    (((idx += (d == inner - 1 ? static_cast<unsigned int>(i)
                              : static_cast<unsigned int>(i) * _dim.stride[d])),
      ++d),
     ...);

  return _data[idx];
}
#endif

template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type = MOOSE_KOKKOS_INDEX_TYPE,
          LayoutType layout = LayoutType::LEFT>
class JaggedArray
{
public:
  JaggedArray() = default;

#ifdef MOOSE_KOKKOS_SCOPE

  template <typename... size_type>
  JaggedArray(size_type... n)
  {
    create(n...);
  }

  void create(const std::vector<index_type> & n);
  template <typename... size_type>
  void create(size_type... n);
  void reserve(const std::array<uint64_t, outer> & index,
               const std::array<uint64_t, inner> & dimension);
  void finalize();
  void copyToDevice()
  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _data.copyToDevice();
  }
  void copyToHost()
  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _data.copyToHost();
  }
  void moveToDevice()
  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _dims.moveToDevice();
    _offsets.moveToDevice();
    _data.moveToDevice();
  }
  void moveToHost()
  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _dims.moveToHost();
    _offsets.moveToHost();
    _data.moveToHost();
  }
  auto & array() { return _data; }

  KOKKOS_FUNCTION bool isFinalized() const { return _finalized; }
  KOKKOS_FUNCTION bool isHostAlloc() const { return _data.isHostAlloc(); }
  KOKKOS_FUNCTION bool isDeviceAlloc() const { return _data.isDeviceAlloc(); }
  KOKKOS_FUNCTION index_type size() const { return _data.size(); }
  KOKKOS_FUNCTION index_type n() const { return _offsets.size(); }
  KOKKOS_FUNCTION index_type n(unsigned int dim) const { return _offsets.n(dim); }
  KOKKOS_FUNCTION auto operator[](index_type i) const;
  template <typename... indices>
  KOKKOS_FUNCTION auto operator()(indices... i) const;
#endif

protected:
  Array1D<T, index_type> _data;
  Array<JaggedArrayInnerDim<inner>, outer, index_type> _dims;
  Array<index_type, outer, index_type> _offsets;
  bool _finalized = false;
};

#ifdef MOOSE_KOKKOS_SCOPE
template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type,
          LayoutType layout>
void
JaggedArray<T, inner, outer, index_type, layout>::create(const std::vector<index_type> & n)
{
  _data.destroy();
  _offsets.create(n);
  _dims.create(n);

  _finalized = false;
}

template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type,
          LayoutType layout>
template <typename... size_type>
void
JaggedArray<T, inner, outer, index_type, layout>::create(size_type... n)
{
  _data.destroy();
  _offsets.create(n...);
  _dims.create(n...);

  _finalized = false;
}

template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type,
          LayoutType layout>
void
JaggedArray<T, inner, outer, index_type, layout>::reserve(
    const std::array<uint64_t, outer> & index, const std::array<uint64_t, inner> & dimension)
{
  mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");

  index_type idx = 0;
  index_type stride = 1;

  for (unsigned int o = 0; o < outer; ++o)
  {
    idx += index[o] * stride;
    stride *= _offsets.n(o);
  }

  for (unsigned int i = 0; i < inner; ++i)
    _dims[idx].dim[i] = dimension[i];

  stride = 1;

  if constexpr (layout == LayoutType::LEFT)
    for (unsigned int i = 0; i < inner; ++i)
    {
      _dims[idx].stride[i] = stride;
      stride *= dimension[i];
    }
  else
    for (int i = inner - 1; i >= 0; --i)
    {
      _dims[idx].stride[i] = stride;
      stride *= dimension[i];
    }

  _dims[idx].stride[inner] = stride;
}

template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type,
          LayoutType layout>
void
JaggedArray<T, inner, outer, index_type, layout>::finalize()
{
  mooseAssert(!_finalized, "KokkosJaggedArray already finalized.");

  index_type stride = 1;

  for (index_type o = 0; o < _offsets.size(); ++o)
  {
    stride = 1;

    for (unsigned int i = 0; i < inner; ++i)
      stride *= _dims[o][i];

    _offsets[o] = stride;
  }

  std::exclusive_scan(_offsets.begin(), _offsets.end(), _offsets.begin(), 0);

  _dims.copyToDevice();
  _offsets.copyToDevice();

  // Pad an extra element at the end to avoid accessing the bound in the following operators when
  // the last inner array has zero size
  _data.create(_offsets.last() + stride + 1);

  _finalized = true;
}

template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type,
          LayoutType layout>
KOKKOS_FUNCTION auto
JaggedArray<T, inner, outer, index_type, layout>::operator[](index_type i) const
{
  auto data = &_data[_offsets[i]];
  const auto & dim = _dims[i];

  return JaggedArrayInnerData<T, inner, layout>(data, dim);
}

template <typename T,
          unsigned int inner,
          unsigned int outer,
          typename index_type,
          LayoutType layout>
template <typename... indices>
KOKKOS_FUNCTION auto
JaggedArray<T, inner, outer, index_type, layout>::operator()(indices... i) const
{
  auto data = &_data[_offsets(i...)];
  const auto & dim = _dims(i...);

  return JaggedArrayInnerData<T, inner, layout>(data, dim);
}
#endif

} // namespace Moose::Kokkos