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"

/**
 * 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);

  using Real3 = Moose::Kokkos::Real3;

  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();
  constexpr 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. The non-templated getKokkosFunction() code path is hypothetically functional under the assumption that the MOOSE software stack can be built with the RDC option, but this hypothetical will only become a reality after a PETSc build system rework. The PETSc and MOOSE team are looking into this issue, and it is expected to be resolved in the foreseeable future.

Available Objects

  • Moose App
  • ADParsedFunctionFunction created by parsing a string
  • ADPiecewiseLinearLinearly interpolates between pairs of x-y data
  • Axisymmetric2D3DSolutionFunctionFunction for reading a 2D axisymmetric solution from file and mapping it to a 3D Cartesian model
  • BicubicSplineFunctionDefine a bicubic spline function from interpolated data defined by input parameters.
  • CoarsenedPiecewiseLinearPerform a point reduction of the tabulated data upon initialization, then evaluate using a linear interpolation.
  • CompositeFunctionMultiplies an arbitrary set of functions together
  • ConstantFunctionA function that returns a constant value as defined by an input parameter.
  • ImageFunctionFunction with values sampled from an image or image stack.
  • KokkosConstantFunctionA function that returns a constant value as defined by an input parameter.
  • KokkosPiecewiseConstantDefines piece-wise constant data using a set of x-y data pairs
  • LinearCombinationFunctionReturns the linear combination of the functions
  • MFEMParsedFunctionParses scalar function of position, time and scalar problem coefficients (including scalar variables).
  • ParsedFunctionFunction created by parsing a string
  • ParsedGradFunctionDefines a function and its gradient using input file parameters.
  • ParsedVectorFunctionReturns a vector function based on string descriptions for each component.
  • PeriodicFunctionProvides a periodic function by repeating a user-supplied base function in time and/or any of the three Cartesian coordinate directions
  • PiecewiseBilinearInterpolates values from a csv file
  • PiecewiseConstantDefines piece-wise constant data using a set of x-y data pairs
  • PiecewiseConstantFromCSVUses data read from CSV to assign values
  • PiecewiseLinearLinearly interpolates between pairs of x-y data
  • PiecewiseLinearFromVectorPostprocessorProvides piecewise linear interpolation of from two columns of a VectorPostprocessor
  • PiecewiseMultiConstantPiecewiseMulticonstant performs constant interpolation on 1D, 2D, 3D or 4D data. The data_file specifies the axes directions and the function values. If a point lies outside the data range, the appropriate end value is used.
  • PiecewiseMulticonstantPiecewiseMulticonstant performs constant interpolation on 1D, 2D, 3D or 4D data. The data_file specifies the axes directions and the function values. If a point lies outside the data range, the appropriate end value is used.
  • PiecewiseMultilinearPiecewiseMultilinear performs linear interpolation on 1D, 2D, 3D or 4D data. The data_file specifies the axes directions and the function values. If a point lies outside the data range, the appropriate end value is used.
  • SolutionFunctionFunction for reading a solution from file.
  • SplineFunctionDefine a spline function from interpolated data defined by input parameters.
  • VectorPostprocessorFunctionProvides piecewise linear interpolation of from two columns of a VectorPostprocessor
  • Reactor App
  • MultiControlDrumFunctionA function that returns an absorber fraction for multiple control drums application.
  • Porous Flow App
  • MovingPlanarFrontThis function defines the position of a moving front. The front is an infinite plane with normal pointing from start_posn to end_posn. The front's distance from start_posn is defined by 'distance', so if the 'distance' function is time dependent, the front's position will change with time. Roughly speaking, the function returns true_value for points lying in between start_posn and start_posn + distance. Precisely speaking, two planes are constructed, both with normal pointing from start_posn to end_posn. The first plane passes through start_posn; the second plane passes through end_posn. Given a point p and time t, this function returns false_value if ANY of the following are true: (a) t<activation_time; (b) t>=deactivation_time; (c) p is 'behind' start_posn (ie, p lies on one side of the start_posn plane and end_posn lies on the other side); (d) p is 'ahead' of the front (ie, p lies one one side of the front and start_posn lies on the other side); (e) the distance between p and the front is greater than active_length. Otherwise, the point is 'in the active zone' and the function returns true_value.
  • Fluid Properties App
  • SaturationDensityFunctionComputes saturation density from temperature function
  • SaturationPressureFunctionComputes saturation pressure from temperature function and 2-phase fluid properties object
  • SaturationTemperatureFunctionComputes saturation temperature from pressure function and 2-phase fluid properties object
  • TwoPhaseNCGPartialPressureFunctionComputes a property from a TwoPhaseNCGPartialPressureFluidProperties object.
  • Stochastic Tools App
  • ScaledAbsDifferenceDRLRewardFunctionEvaluates a scaled absolute difference reward function for a process which is controlled by a Deep Reinforcement Learning based surrogate.
  • Thermal Hydraulics App
  • CircularAreaHydraulicDiameterFunctionComputes hydraulic diameter for a circular area from its area function
  • CosineHumpFunctionComputes a cosine hump of a user-specified width and height
  • CosineTransitionFunctionComputes a cosine transtition of a user-specified width between two values
  • CubicTransitionFunctionComputes a cubic polynomial transition between two functions
  • GeneralizedCircumferenceComputes a generalized circumference from a function providing the area.
  • PiecewiseFunctionFunction which provides a piecewise representation of arbitrary functions
  • TimeRampFunctionRamps up to a value from another value over time.
  • Optimization App
  • NearestReporterCoordinatesFunctionThis Function finds the nearest point in the specified vectors of coordinates and returns the values specified in the vector of values at the index of the nearest point. All the vectors must be specified using either vector postprocessors or reporter vectors. This function interpolates linearly in time with transient data.
  • ParameterMeshFunctionOptimization function with parameters represented by a mesh and finite-element shape functions.
  • ParsedOptimizationFunctionFunction used for optimization that uses a parsed expression with parameter dependence.
  • Level Set App
  • LevelSetOlssonBubbleImplementation of 'bubble' ranging from 0 to 1.
  • LevelSetOlssonPlaneImplementation of a level set function to represent a plane.
  • LevelSetOlssonVortexA function for creating vortex velocity fields for level set equation benchmark problems.
  • Phase Field App
  • FourierNoiseGenerate noise from a fourier series
  • Functional Expansion Tools App
  • FunctionSeriesThis function uses a convolution of functional series (functional expansion or FX) to create a 1D, 2D, or 3D function