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
{
namespace Kokkos
{

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

#ifdef MOOSE_KOKKOS_SCOPE

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

#ifdef MOOSE_KOKKOS_SCOPE

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

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

    return _data[i];
  }

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

#define usingKokkosJaggedArrayInnerDataBaseMembers(T, inner)                                       \
private:                                                                                           \
  using JaggedArrayInnerDataBase<T, inner>::_data;                                                 \
  using JaggedArrayInnerDataBase<T, inner>::_dim

template <typename T, unsigned int inner>
class JaggedArrayInnerData;

template <typename T>
class JaggedArrayInnerData<T, 1> : public JaggedArrayInnerDataBase<T, 1>
{
  usingKokkosJaggedArrayInnerDataBaseMembers(T, 1);

public:
  KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<1> dim)
    : JaggedArrayInnerDataBase<T, 1>(data, dim)
  {
  }

  KOKKOS_FUNCTION T & operator()(dof_id_type i0) const
  {
    KOKKOS_ASSERT(i0 < _dim[0]);

    return _data[i0];
  }
};

template <typename T>
class JaggedArrayInnerData<T, 2> : public JaggedArrayInnerDataBase<T, 2>
{
  usingKokkosJaggedArrayInnerDataBaseMembers(T, 2);

public:
  KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<2> dim)
    : JaggedArrayInnerDataBase<T, 2>(data, dim)
  {
  }

  KOKKOS_FUNCTION T & operator()(dof_id_type i0, dof_id_type i1) const
  {
    KOKKOS_ASSERT(i0 < _dim[0]);
    KOKKOS_ASSERT(i1 < _dim[1]);

    return _data[i0 + _dim.stride[0] * i1];
  }
};

template <typename T>
class JaggedArrayInnerData<T, 3> : public JaggedArrayInnerDataBase<T, 3>
{
  usingKokkosJaggedArrayInnerDataBaseMembers(T, 3);

public:
  KOKKOS_FUNCTION JaggedArrayInnerData(T * data, JaggedArrayInnerDim<3> dim)
    : JaggedArrayInnerDataBase<T, 3>(data, dim)
  {
  }

  KOKKOS_FUNCTION T & operator()(dof_id_type i0, dof_id_type i1, dof_id_type i2) const
  {
    KOKKOS_ASSERT(i0 < _dim[0]);
    KOKKOS_ASSERT(i1 < _dim[1]);
    KOKKOS_ASSERT(i2 < _dim[2]);

    return _data[i0 + _dim.stride[0] * i1 + _dim.stride[1] * i2];
  }
};
#endif

template <typename T, unsigned int inner, unsigned int outer>
class JaggedArrayBase
{
#ifdef MOOSE_KOKKOS_SCOPE
public:
  void create(const std::vector<dof_id_type> & n);
  void reserve(const std::array<dof_id_type, outer> & index,
               const std::array<dof_id_type, inner> & dimension);
  void finalize();
  void copyToDevice()
  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

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

    _data.copyToHost();
  }
  auto & array() { return _data; }

  KOKKOS_FUNCTION bool isFinalized() const { return _finalized; }
  KOKKOS_FUNCTION dof_id_type size() const { return _data.size(); }
  KOKKOS_FUNCTION dof_id_type n() const { return _offsets.size(); }
  KOKKOS_FUNCTION dof_id_type n(unsigned int dim) const { return _offsets.n(dim); }
  KOKKOS_FUNCTION auto operator[](dof_id_type i) const
  {
    auto data = &_data[_offsets[i]];
    const auto & dim = _dims[i];

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

protected:
  Array<T> _data;
  Array<JaggedArrayInnerDim<inner>, outer> _dims;
  Array<dof_id_type, outer> _offsets;
  bool _finalized;
};

#define usingKokkosJaggedArrayBaseMembers(T, inner, outer)                                         \
private:                                                                                           \
  using JaggedArrayBase<T, inner, outer>::_data;                                                   \
  using JaggedArrayBase<T, inner, outer>::_dims;                                                   \
  using JaggedArrayBase<T, inner, outer>::_offsets

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

  _finalized = false;
}

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

  dof_id_type idx = 0;
  dof_id_type stride = 1;

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

  stride = 1;

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

    _dims[idx].dim[i] = dimension[i];
    _dims[idx].stride[i] = stride;
  }
}

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

  dof_id_type stride = 1;

  for (dof_id_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();
  _data.create(_offsets.last() + stride);

  _finalized = true;
}
#endif

template <typename T, unsigned int inner, unsigned int outer>
class JaggedArray;

template <typename T, unsigned int inner>
class JaggedArray<T, inner, 1> : public JaggedArrayBase<T, inner, 1>
{
  usingKokkosJaggedArrayBaseMembers(T, inner, 1);

public:
  JaggedArray() = default;

#ifdef MOOSE_KOKKOS_SCOPE

  JaggedArray(dof_id_type n0) { create(n0); }
  void create(dof_id_type n0) { JaggedArrayBase<T, inner, 1>::create({n0}); }

  KOKKOS_FUNCTION auto operator()(dof_id_type i0) const
  {
    auto data = &_data[_offsets(i0)];
    const auto & dim = _dims(i0);

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

template <typename T, unsigned int inner>
class JaggedArray<T, inner, 2> : public JaggedArrayBase<T, inner, 2>
{
  usingKokkosJaggedArrayBaseMembers(T, inner, 2);

public:
  JaggedArray() = default;

#ifdef MOOSE_KOKKOS_SCOPE

  JaggedArray(dof_id_type n0, dof_id_type n1) { create(n0, n1); }
  void create(dof_id_type n0, dof_id_type n1) { JaggedArrayBase<T, inner, 2>::create({n0, n1}); }

  KOKKOS_FUNCTION auto operator()(dof_id_type i0, dof_id_type i1) const
  {
    auto data = &_data[_offsets(i0, i1)];
    const auto & dim = _dims(i0, i1);

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

template <typename T, unsigned int inner>
class JaggedArray<T, inner, 3> : public JaggedArrayBase<T, inner, 3>
{
  usingKokkosJaggedArrayBaseMembers(T, inner, 3);

public:
  JaggedArray() = default;

#ifdef MOOSE_KOKKOS_SCOPE

  JaggedArray(dof_id_type n0, dof_id_type n1, dof_id_type n2) { create(n0, n1, n2); }
  void create(dof_id_type n0, dof_id_type n1, dof_id_type n2)
  {
    JaggedArrayBase<T, inner, 3>::create({n0, n1, n2});
  }

  KOKKOS_FUNCTION auto operator()(dof_id_type i0, dof_id_type i1, dof_id_type i2) const
  {
    auto data = &_data[_offsets(i0, i1, i2)];
    const auto & dim = _dims(i0, i1, i2);

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

} // namespace Kokkos
} // namespace Moose