Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout > Class Template Reference

The Kokkos jagged array class. More...

#include <KokkosJaggedArray.h>

Public Member Functions
	JaggedArray ()=default
	Default constructor. More...
template<typename... size_type>
	JaggedArray (size_type... n)
	Constructor. More...
void	create (const std::vector< index_type > &n)
	Allocate outer array. More...
template<typename... size_type>
void	create (size_type... n)
	Allocate outer array. More...
void	reserve (const std::array< uint64_t, outer > &index, const std::array< uint64_t, inner > &dimension)
	Reserve inner array for an outer array entry. More...
void	finalize ()
	Setup array structure. More...
void	copyToDevice ()
	Copy data from host to device. More...
void	copyToHost ()
	Copy data from device to host. More...
void	moveToDevice ()
	Copy data from host to device and deallocate host. More...
void	moveToHost ()
	Copy data from device to host and deallocate device. More...
auto &	array ()
	Get the underlying data array. More...
KOKKOS_FUNCTION bool	isFinalized () const
	Get whether the array is finalized. More...
KOKKOS_FUNCTION bool	isHostAlloc () const
	Get whether the array was allocated on host. More...
KOKKOS_FUNCTION bool	isDeviceAlloc () const
	Get whether the array was allocated on device. More...
KOKKOS_FUNCTION index_type	size () const
	Get the total data array size. More...
KOKKOS_FUNCTION index_type	n () const
	Get the total outer array size. More...
KOKKOS_FUNCTION index_type	n (unsigned int dim) const
	Get the size of a dimension of the outer array. More...
KOKKOS_FUNCTION auto	operator[] (index_type i) const
	Get an inner array. More...
template<typename... indices>
KOKKOS_FUNCTION auto	operator() (indices... i) const
	Get an inner array. More...

Protected Attributes
Array1D< T, index_type >	_data
	Sequential data array. More...
Array< JaggedArrayInnerDim< inner >, outer, index_type >	_dims
	Dimension information of each inner array. More...
Array< index_type, outer, index_type >	_offsets
	Starting offset of each inner array into the sequential data array. More...
bool	_finalized = false
	Whether the array was finalized. More...

Detailed Description

template<typename T, unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>
class Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >

The Kokkos jagged array class.

A Kokkos jagged array aids in treating jagged arrays conveniently and efficiently by using a sequential data storage and providing multi-dimensional indexing using internal dope vectors. A jagged array is divided into the inner and outer arrays. The outer array is the regular part of a jagged array. Each entry of the outer array can hold an inner array, whose size can vary with each other. Calling create() will allocate the outer array, and inner arrays should be reserved one-by-one through reserve(). Once the array structure is set, finalize() should be called. A finalized array cannot be restructured unless it is reset completely by calling create() again. The inner array returned by operator() or operator[] using the outer array index is held in a temporary wrapper object. To avoid the overhead of temporary object creation, it is recommended to store the object locally.

Template Parameters

T	The data type
inner	The inner array dimension size
outer	The outer array dimension size
index_type	The array index type
layout	The memory layout type

Definition at line 164 of file KokkosJaggedArray.h.

Constructor & Destructor Documentation

JaggedArray() [1/2]

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::JaggedArray ( )

default

Default constructor.

JaggedArray() [2/2]

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

template<typename... size_type>

Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::JaggedArray ( size_type...  n )

inline

Constructor.

Definition at line 177 of file KokkosJaggedArray.h.

  {
    create(n...);
  }

Member Function Documentation

array()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

auto& Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::array ( )

inline

Get the underlying data array.

Returns

The data array

Definition at line 248 of file KokkosJaggedArray.h.

{ return _data; }

copyToDevice()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::copyToDevice ( )

inline

Copy data from host to device.

Definition at line 207 of file KokkosJaggedArray.h.

  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _data.copyToDevice();
  }

copyToHost()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::copyToHost ( )

inline

Copy data from device to host.

Definition at line 216 of file KokkosJaggedArray.h.

  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _data.copyToHost();
  }

create() [1/2]

template<typename T , unsigned int inner, unsigned int outer, typename index_type , LayoutType layout>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::create

(

const std::vector< index_type > &

n

)

Allocate outer array.

Parameters

n	The vector containing the size of each dimension for the outer array

Definition at line 322 of file KokkosJaggedArray.h.

Referenced by Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::JaggedArray().

{
  _data.destroy();
  _offsets.create(n);
  _dims.create(n);

  _finalized = false;
}

create() [2/2]

template<typename T , unsigned int inner, unsigned int outer, typename index_type , LayoutType layout>

template<typename... size_type>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::create

(

size_type...

n

)

Allocate outer array.

Parameters

n	The size of each dimension for the outer array

Definition at line 338 of file KokkosJaggedArray.h.

{
  _data.destroy();
  _offsets.create(n...);
  _dims.create(n...);

  _finalized = false;
}

finalize()

template<typename T , unsigned int inner, unsigned int outer, typename index_type , LayoutType layout>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::finalize

(

)

Setup array structure.

Definition at line 394 of file KokkosJaggedArray.h.

{
  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;
}

isDeviceAlloc()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

KOKKOS_FUNCTION bool Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::isDeviceAlloc ( ) const

inline

Get whether the array was allocated on device.

Returns

Whether the array was allocated on device

Definition at line 264 of file KokkosJaggedArray.h.

{ return _data.isDeviceAlloc(); }

isFinalized()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

KOKKOS_FUNCTION bool Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::isFinalized ( ) const

inline

Get whether the array is finalized.

Returns

Whether the array is finalized

Definition at line 254 of file KokkosJaggedArray.h.

{ return _finalized; }

isHostAlloc()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

KOKKOS_FUNCTION bool Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::isHostAlloc ( ) const

inline

Get whether the array was allocated on host.

Returns

Whether the array was allocated on host

Definition at line 259 of file KokkosJaggedArray.h.

{ return _data.isHostAlloc(); }

moveToDevice()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToDevice ( )

inline

Copy data from host to device and deallocate host.

Definition at line 225 of file KokkosJaggedArray.h.

  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _dims.moveToDevice();
    _offsets.moveToDevice();
    _data.moveToDevice();
  }

moveToHost()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToHost ( )

inline

Copy data from device to host and deallocate device.

Definition at line 236 of file KokkosJaggedArray.h.

  {
    mooseAssert(_finalized, "KokkosJaggedArray not finalized.");

    _dims.moveToHost();
    _offsets.moveToHost();
    _data.moveToHost();
  }

n() [1/2]

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

KOKKOS_FUNCTION index_type Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::n ( ) const

inline

Get the total outer array size.

Returns

The total outer array size

Definition at line 274 of file KokkosJaggedArray.h.

Referenced by Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::JaggedArray().

{ return _offsets.size(); }

n() [2/2]

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

KOKKOS_FUNCTION index_type Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::n ( unsigned int  dim ) const

inline

Get the size of a dimension of the outer array.

Parameters

dim	The dimension index

Returns

The size of the dimension of the outer array

Definition at line 280 of file KokkosJaggedArray.h.

{ return _offsets.n(dim); }

operator()()

template<typename T , unsigned int inner, unsigned int outer, typename index_type , LayoutType layout>

template<typename... indices>

KOKKOS_FUNCTION auto Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::operator()

(

indices...

i

)

const

Get an inner array.

Parameters

i	The index of each dimension

Returns

The inner array wrapper object

Definition at line 443 of file KokkosJaggedArray.h.

{
  auto data = &_data[_offsets(i...)];
  const auto & dim = _dims(i...);

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

operator[]()

template<typename T , unsigned int inner, unsigned int outer, typename index_type , LayoutType layout>

KOKKOS_FUNCTION auto Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::operator[]

(

index_type

i

)

const

Get an inner array.

Parameters

i	The dimensionless outer array index

Returns

The inner array wrapper object

Definition at line 428 of file KokkosJaggedArray.h.

{
  auto data = &_data[_offsets[i]];
  const auto & dim = _dims[i];

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

reserve()

template<typename T , unsigned int inner, unsigned int outer, typename index_type , LayoutType layout>

void Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::reserve	(	const std::array< uint64_t, outer > &	index,
		const std::array< uint64_t, inner > &	dimension
	)

Reserve inner array for an outer array entry.

Parameters

index	The array containing the index for the outer array
dimension	The array containing the size of each dimension for the inner array

Definition at line 353 of file KokkosJaggedArray.h.

{
  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;
}

size()

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

KOKKOS_FUNCTION index_type Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::size ( ) const

inline

Get the total data array size.

Returns

The total data array size

Definition at line 269 of file KokkosJaggedArray.h.

{ return _data.size(); }

Member Data Documentation

_data

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

Array1D<T, index_type> Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::_data

protected

Sequential data array.

Definition at line 300 of file KokkosJaggedArray.h.

Referenced by Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::array(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::copyToDevice(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::copyToHost(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::isDeviceAlloc(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::isHostAlloc(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToDevice(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToHost(), and Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::size().

_dims

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

Array<JaggedArrayInnerDim<inner>, outer, index_type> Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::_dims

protected

Dimension information of each inner array.

Definition at line 304 of file KokkosJaggedArray.h.

Referenced by Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToDevice(), and Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToHost().

_finalized

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

bool Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::_finalized = false

protected

Whether the array was finalized.

Definition at line 312 of file KokkosJaggedArray.h.

Referenced by Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::copyToDevice(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::copyToHost(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::isFinalized(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToDevice(), and Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToHost().

_offsets

template<typename T , unsigned int inner, unsigned int outer, typename index_type = MOOSE_KOKKOS_INDEX_TYPE, LayoutType layout = LayoutType::LEFT>

Array<index_type, outer, index_type> Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::_offsets

protected

Starting offset of each inner array into the sequential data array.

Definition at line 308 of file KokkosJaggedArray.h.

Referenced by Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToDevice(), Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::moveToHost(), and Moose::Kokkos::JaggedArray< T, inner, outer, index_type, layout >::n().

---

The documentation for this class was generated from the following file:

• include/kokkos/base/KokkosJaggedArray.h