KokkosThread.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 "MooseUtils.h"
#include "Conversion.h"

namespace Moose::Kokkos
{

using ThreadID = MOOSE_KOKKOS_INDEX_TYPE;

template <typename thread_id_type = ThreadID, unsigned int max_dimension = 4>
class Thread
{
  static_assert(std::is_integral_v<thread_id_type>,
                "Kokkos thread index type must be an integral type");
  static_assert(std::is_unsigned_v<thread_id_type>, "Kokkos thread index type must be unsigned");
  static_assert(!std::is_same_v<bool, thread_id_type>, "Kokkos thread index type must not be bool");

public:
  using id_type = thread_id_type;

#ifdef MOOSE_KOKKOS_SCOPE

  template <typename... size_type>
  void resize(size_type... sizes);
  KOKKOS_FUNCTION thread_id_type size() const { return _size; }
  KOKKOS_FUNCTION thread_id_type size(const unsigned int dim) const
  {
    KOKKOS_ASSERT(dim < _dim);

    return _dims[dim];
  }
  KOKKOS_FUNCTION thread_id_type operator()(const thread_id_type tid, const unsigned int dim) const
  {
    KOKKOS_ASSERT(dim < _dim);

    return (tid / _strides[dim]) % _dims[dim];
  }
#endif

protected:
  thread_id_type _size = 0;
  unsigned int _dim = 0;
  thread_id_type _dims[max_dimension];
  thread_id_type _strides[max_dimension];
};

#ifdef MOOSE_KOKKOS_SCOPE
template <typename thread_id_type, unsigned int max_dimension>
template <typename... size_type>
void
Thread<thread_id_type, max_dimension>::resize(size_type... sizes)
{
  static_assert((std::is_convertible<size_type, thread_id_type>::value && ...),
                "All arguments must be convertible to thread_id_type");
  static_assert(sizeof...(sizes) <= max_dimension, "Number of arguments exceeds maximum dimension");

  std::vector<thread_id_type> dims;
  (dims.push_back(sizes), ...);

  uint64_t size = 1;
  _dim = dims.size();

  for (unsigned int dim = 0; dim < _dim; ++dim)
  {
    _dims[dim] = dims[dim];
    _strides[dim] = dim ? _strides[dim - 1] * dims[dim - 1] : 1;
    size *= dims[dim];
  }

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

  _size = size;
}
#endif

} // namespace Moose::Kokkos