Kokkos Functions System

Before reading this documentation, consider reading the following materials first for a better understanding of this documentation:

commentnote

Kokkos-MOOSE functions do not support automatic differention yet.

A Kokkos-MOOSE function can be created by subclassing Moose::Kokkos::FunctionBase (not Moose::Kokkos::Function), which should be registered with registerKokkosFunction() instead of registerMooseObject(). The signatures of the hook methods are defined as follows:


KOKKOS_FUNCTION Real value(Real t, Real3 p) const;
KOKKOS_FUNCTION Real3 vectorValue(Real t, Real3 p) const;
KOKKOS_FUNCTION Real3 gradient(Real t, Real3 p) const;
KOKKOS_FUNCTION Real3 curl(Real t, Real3 p) const;
KOKKOS_FUNCTION Real div(Real t, Real3 p) const;
KOKKOS_FUNCTION Real timeDerivative(Real t, Real3 p) const;
KOKKOS_FUNCTION Real timeIntegral(Real t1, Real t2, Real3 p) const;
KOKKOS_FUNCTION Real integral() const;
KOKKOS_FUNCTION Real average() const;

As in other Kokkos-MOOSE objects, they should be defined in the derived class as inlined public methods instead of virtual override. It is not mandatory to define each hook method in the derived class, but any hook method that was not defined in the derived class should not be called.

See the following source codes of KokkosPiecewiseConstant for an example of a function:

Listing 1: The KokkosPiecewiseConstant header file.


#pragma once

#include "KokkosPiecewiseTabularBase.h"
#include "KokkosUtils.h"

using Moose::Kokkos::Real3;

/**
 * Function which provides a piecewise constant interpolation of a provided (x,y) point data set.
 */
class KokkosPiecewiseConstant : public KokkosPiecewiseTabularBase
{
public:
  static InputParameters validParams();

  KokkosPiecewiseConstant(const InputParameters & parameters);

  KOKKOS_FUNCTION Real value(Real t, Real3 p) const;
  KOKKOS_FUNCTION Real integral() const;
  KOKKOS_FUNCTION Real average() const;

private:
  /// Enum for which direction to apply values
  const enum class Direction { LEFT, RIGHT, LEFT_INCLUSIVE, RIGHT_INCLUSIVE } _direction;
};

KOKKOS_FUNCTION inline Real
KokkosPiecewiseConstant::value(Real t, Real3 p) const
{
  using Moose::Kokkos::Utils::sign;

  const Real x = _has_axis ? p(_axis) : t;

  const auto len = functionSize();
  const Real tolerance = 1.0e-14;

  // endpoint cases
  if ((_direction == Direction::LEFT && x < (1 + tolerance * sign(domain(0))) * domain(0)) ||
      (_direction == Direction::RIGHT && x < (1 - tolerance * sign(domain(0))) * domain(0)) ||
      (_direction == Direction::LEFT_INCLUSIVE &&
       x < (1 - tolerance * sign(domain(0))) * domain(0)) ||
      (_direction == Direction::RIGHT_INCLUSIVE &&
       x < (1 + tolerance * sign(domain(0))) * domain(0)))
    return _scale_factor * range(0);
  else if ((_direction == Direction::LEFT &&
            x > (1 + tolerance * sign(domain(len - 1))) * domain(len - 1)) ||
           (_direction == Direction::RIGHT &&
            x > (1 - tolerance * sign(domain(len - 1))) * domain(len - 1)) ||
           (_direction == Direction::LEFT_INCLUSIVE &&
            x > (1 - tolerance * sign(domain(len - 1))) * domain(len - 1)) ||
           (_direction == Direction::RIGHT_INCLUSIVE &&
            x > (1 + tolerance * sign(domain(len - 1))) * domain(len - 1)))
    return _scale_factor * range(len - 1);

  for (unsigned int i = 1; i < len; ++i)
  {
    if (_direction == Direction::LEFT && x < (1 + tolerance * sign(domain(i))) * domain(i))
      return _scale_factor * range(i - 1);
    else if (_direction == Direction::LEFT_INCLUSIVE &&
             x < (1 - tolerance * sign(domain(i))) * domain(i))
      return _scale_factor * range(i - 1);
    else if ((_direction == Direction::RIGHT && x < (1 - tolerance * sign(domain(i))) * domain(i)))
      return _scale_factor * range(i);
    else if ((_direction == Direction::RIGHT_INCLUSIVE &&
              x < (1 + tolerance * sign(domain(i))) * domain(i)))
      return _scale_factor * range(i);
  }

  return 0.0;
}

KOKKOS_FUNCTION inline Real
KokkosPiecewiseConstant::integral() const
{
  const auto len = functionSize();

  unsigned int offset = 0;

  if (_direction == Direction::RIGHT || _direction == Direction::RIGHT_INCLUSIVE)
    offset = 1;

  Real sum = 0;

  for (unsigned int i = 0; i < len - 1; ++i)
    sum += range(i + offset) * (domain(i + 1) - domain(i));

  return _scale_factor * sum;
}

KOKKOS_FUNCTION inline Real
KokkosPiecewiseConstant::average() const
{
  return integral() / (domain(functionSize() - 1) - domain(0));
}
(framework/include/kokkos/functions/KokkosPiecewiseConstant.h)

Listing 2: The KokkosPiecewiseConstant source file.


#include "KokkosPiecewiseConstant.h"

registerKokkosFunction("MooseApp", KokkosPiecewiseConstant);

InputParameters
KokkosPiecewiseConstant::validParams()
{
  InputParameters params = KokkosPiecewiseTabularBase::validParams();
  MooseEnum direction("LEFT RIGHT LEFT_INCLUSIVE RIGHT_INCLUSIVE", "LEFT");
  params.addParam<MooseEnum>(
      "direction", direction, "Direction to look to find value: " + direction.getRawNames());
  params.addClassDescription("Defines piece-wise constant data using a set of x-y data pairs");
  return params;
}

KokkosPiecewiseConstant::KokkosPiecewiseConstant(const InputParameters & parameters)
  : KokkosPiecewiseTabularBase(parameters),
    _direction(getParam<MooseEnum>("direction").getEnum<Direction>())
{
}
(framework/src/kokkos/functions/KokkosPiecewiseConstant.K)

Functions can be acquired in your object by calling getKokkosFunction<T>, where T should be your function type. Namely, the actual type of the function should be known in advance. This is the limitation of the current Kokkos-MOOSE function implementation which is being addressed, and more discussions can be found below. It is recommended to store the acquired function in a ReferenceWrapper instance. Otherwise, it should be stored as a copy, and it becomes your responsibility to maintain synchronization between the original function and the copy.

Use of Dynamic Polymorphism

Unlike most of other Kokkos-MOOSE objects, functions are pluggable objects that are to be used in other objects. Therefore, it is desired to use functions type-agnostically in your object so that any type of function can be plugged in. Such design requires dynamic polymorphism, which is represented by virtual functions. Although using virtual functions on GPU can sacrifice performance due to the inability of code inlining and vtable lookup and is still advised to avoid if possible, it can be useful for avoiding code duplications and improving code maintainability.

In fact, Kokkos-MOOSE already has the code path for acquiring functions as a type-agnostic abstract type and using virtual dispatch to evaluate actual functions. There is another non-template API getKokkosFunction which returns an abstract type Moose::Kokkos::Function. It is not the base class of your function but a wrapper class that contains your function and provides shims to call the function methods.

However, that code path is currently blocked for GPU backends, and you will hit a runtime error when you attempt to use it. While our target GPU backends (CUDA, HIP, and Intel SYCL) support virtual functions on GPU, it requires the relocatable device code (RDC) option to be enabled during compilation, which the current MOOSE build system is lacking. Kokkos imposes a restriction to have a consistent RDC configuration across the whole software stack, but PETSc, from which MOOSE is currently acquring Kokkos, requires significant changes in its build system to enable RDC. Namely, the current code is functional under the assumption that the MOOSE software stack can be built with the RDC option, but it involves a major rework from our upstream dependency. The PETSc and MOOSE team are looking into this issue, and it is expected to be resolved in a foreseeable future.

Available Objects