Kokkos Kernels System

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

• Kernels System to understand the MOOSE kernel system,

• Getting Started with Kokkos-MOOSE to understand the programming practices for Kokkos-MOOSE.

The Kokkos-MOOSE kernels are designed to resemble the original MOOSE kernels as much as possible for easier porting and adaptation. However, some differences still exist due to the fundamentally different programming paradigm between CPU and GPU. First, the way of subclassing the base MOOSE class to create your own kernel is different. In original MOOSE, you could simply inherit the base Kernel class:

class Diffusion : public Kernel
{
  ...
};

In Kokkos-MOOSE, however, due to the Curiously Recurring Template Pattern (CRTP) design adopted for static polymorphism, you should inherit a template base class Moose::Kokkos::Kernel with the template argument being the derived class type itself:

class KokkosDiffusion final : public Moose::Kokkos::Kernel<KokkosDiffusion>
{
  ...
};

In the above example, KokkosDiffusion cannot be inherited anymore. If you want to make a subclass that can be further inherited, the subclass should be a template class as well with the template argument being the derived class type, and the template argument should be relayed to the base class:

template <typename Derived>
class KokkosDiffusion : public Moose::Kokkos::Kernel<Derived>
{
  usingKokkosKernelMembers(Derived);
  ...
};

Due to the dependent names in C++, you cannot directly see the variables and functions of a base template class in a derived template class. Therefore, we provide usingKokkosKernelMembers() which is a convenient macro to bring them into the current scope. The argument of the macro should be the template argument.

However, by making KokkosDiffusion a template class, you can no longer have an instance of KokkosDiffusion itself. To achieve this, a dummy class that simply inherits KokkosDiffusion without adding capability should be separately defined and registered:

class KokkosDiffusionKernel final : public KokkosDiffusion<KokkosDiffusionKernel>
{
  ...
};

registerMooseObjectAliased("MooseApp", KokkosDiffusionKernel, "KokkosDiffusion");

The use of registerMooseObjectAliased(), which allows registering an object with a different name from the class name itself, would be optional, depending on how you name the template and dummy classes.

Second, the signatures of hook methods are different. In original MOOSE, the following virtual functions should or optionally have been overriden:

virtual Real computeQpResidual() override;
virtual Real computeQpJacobian() override;
virtual Real computeQpOffDiagJacobian(unsigned int jvar) override;

Indices such as _i, _j, and _qp were made available in those functions as member variables.

In Kokkos-MOOSE, the hook methods are no longer virtual functions. Instead, they are (re)defined in the derived class as public methods. Their signatures look like the following:

KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
                                       const unsigned int qp,
                                       ResidualDatum & datum) const;
KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int i,
                                       const unsigned int j,
                                       const unsigned int qp,
                                       ResidualDatum & datum) const;
KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int i,
                                              const unsigned int j,
                                              const unsigned int jvar,
                                              const unsigned int qp,
                                              ResidualDatum & datum) const;

Analogously to the original MOOSE, computeQpResidual() must be defined in the derived class, and the definition of computeQpJacobian() and computeQpOffDiagJacobian() are optional. The optional methods have default definitions in the base class, and redefining them in the derived class results in hiding the base class definitions. Beware not to misspell the function names when redefining the optional methods, as they will be silently ignored rather than throwing a compile error.

The indices can no longer be member variables due to the parallelization, and therefore they are now passed as function arguments. datum is a temporary object holding element-private data, and _q_point is available through datum as datum.q_point(qp). _current_elem does not have a direct replacement, as it is a pointer to the current libMesh element object that cannot be accessed on GPU. However, datum.elem() returns an object containing the information of the current element, and it can be used to retrieve mesh data from the Kokkos mesh object available through kokkosMesh(). Note that the mesh data available on GPU is currently very limited and will be continuously enhanced. Variable and shape objects now use a function call syntax instead of the array indexing syntax and require datum as an additional argument to the indices. Consequently, the following residual and Jacobian functions in Diffusion:

Real
Diffusion::computeQpResidual()
{
  return _grad_u[_qp] * _grad_test[_i][_qp];
}

Real
Diffusion::computeQpJacobian()
{
  return _grad_phi[_j][_qp] * _grad_test[_i][_qp];
}

will look like the following in KokkosDiffusion:

KOKKOS_FUNCTION Real
KokkosDiffusion::computeQpResidual(const unsigned int i,
                                   const unsigned int qp,
                                   ResidualDatum & datum) const
{
  return _grad_u(datum, qp) * _grad_test(datum, i, qp);
}

KOKKOS_FUNCTION Real
KokkosDiffusion::computeQpJacobian(const unsigned int i,
                                   const unsigned int j,
                                   const unsigned int qp,
                                   ResidualDatum & datum) const
{
  return _grad_phi(datum, j, qp) * _grad_test(datum, i, qp);
}

See the following source codes of KokkosCoupledForce for another example of a kernel:

#pragma once

#include "KokkosKernel.h"

/**
 * Implements a source term proportional to the value of a coupled variable. Weak form: $(\\psi_i,
 * -\\sigma v)$.
 */
class KokkosCoupledForce final : public Moose::Kokkos::Kernel<KokkosCoupledForce>
{
public:
  static InputParameters validParams();

  KokkosCoupledForce(const InputParameters & parameters);

  KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
                                         const unsigned int qp,
                                         ResidualDatum & datum) const;

private:
  /// Coupled variable number
  const unsigned int _v_var;
  /// Coupled variable
  const Moose::Kokkos::VariableValue _v;
  /// Multiplier for the coupled force term
  const Real _coef;
};

KOKKOS_FUNCTION inline Real
KokkosCoupledForce::computeQpResidual(const unsigned int i,
                                      const unsigned int qp,
                                      ResidualDatum & datum) const
{
  return -_coef * _v(datum, qp) * _test(datum, i, qp);
}

#include "KokkosCoupledForce.h"

registerMooseObject("MooseApp", KokkosCoupledForce);

InputParameters
KokkosCoupledForce::validParams()
{
  InputParameters params = Kernel::validParams();

  params.addClassDescription("Implements a source term proportional to the value of a coupled "
                             "variable. Weak form: $(\\psi_i, -\\sigma v)$.");
  params.addRequiredCoupledVar("v", "The coupled variable which provides the force");
  params.addParam<Real>(
      "coef", 1.0, "Coefficent ($\\sigma$) multiplier for the coupled force term.");

  return params;
}

KokkosCoupledForce::KokkosCoupledForce(const InputParameters & parameters)
  : Kernel(parameters),
    _v_var(coupled("v")),
    _v(kokkosCoupledValue("v")),
    _coef(getParam<Real>("coef"))
{
  if (_var.number() == _v_var)
    paramError(
        "v", "Coupled variable 'v' needs to be different from 'variable' with KokkosCoupledForce.");
}

KOKKOS_FUNCTION Real precomputeQpResidual(const unsigned int qp,
                                          ResidualDatum & datum) const;
KOKKOS_FUNCTION Real precomputeQpJacobian(const unsigned int j,
                                          const unsigned int qp,
                                          ResidualDatum & datum) const;

KOKKOS_FUNCTION Real3 precomputeQpResidual(const unsigned int qp,
                                           ResidualDatum & datum) const;
KOKKOS_FUNCTION Real3 precomputeQpJacobian(const unsigned int j,
                                           const unsigned int qp,
                                           ResidualDatum & datum) const;

#pragma once

#include "KokkosKernelValue.h"

using Real3 = Moose::Kokkos::Real3;

class KokkosConvectionPrecompute final
  : public Moose::Kokkos::KernelValue<KokkosConvectionPrecompute>
{
public:
  static InputParameters validParams();

  KokkosConvectionPrecompute(const InputParameters & parameters);

  KOKKOS_FUNCTION Real precomputeQpResidual(const unsigned int qp, ResidualDatum & datum) const;
  KOKKOS_FUNCTION Real precomputeQpJacobian(const unsigned int j,
                                            const unsigned int qp,
                                            ResidualDatum & datum) const;

private:
  const Real3 _velocity;
};

KOKKOS_FUNCTION inline Real
KokkosConvectionPrecompute::precomputeQpResidual(const unsigned int qp, ResidualDatum & datum) const
{
  return _velocity * _grad_u(datum, qp);
}

KOKKOS_FUNCTION inline Real
KokkosConvectionPrecompute::precomputeQpJacobian(const unsigned int j,
                                                 const unsigned int qp,
                                                 ResidualDatum & datum) const
{
  return _velocity * _grad_phi(datum, j, qp);
}

#pragma once

#include "KokkosKernelGrad.h"

using Real3 = Moose::Kokkos::Real3;

class KokkosDiffusionPrecompute final : public Moose::Kokkos::KernelGrad<KokkosDiffusionPrecompute>
{
public:
  static InputParameters validParams();

  KokkosDiffusionPrecompute(const InputParameters & parameters);

  KOKKOS_FUNCTION Real3 precomputeQpResidual(const unsigned int qp, ResidualDatum & datum) const;
  KOKKOS_FUNCTION Real3 precomputeQpJacobian(const unsigned int j,
                                             const unsigned int qp,
                                             ResidualDatum & datum) const;
};

KOKKOS_FUNCTION inline Real3
KokkosDiffusionPrecompute::precomputeQpResidual(const unsigned int qp, ResidualDatum & datum) const
{
  // Note we do not multiply by the gradient of the test function. That is done in the parent
  // class
  return _grad_u(datum, qp);
}

KOKKOS_FUNCTION inline Real3
KokkosDiffusionPrecompute::precomputeQpJacobian(const unsigned int j,
                                                const unsigned int qp,
                                                ResidualDatum & datum) const
{
  // Note we do not multiply by the gradient of the test function. That is done in the parent
  // class
  return _grad_phi(datum, j, qp);
}

return _test[_i][_qp] * _u_dot[_qp];

return _test(datum, i, qp) * _u_dot(datum, qp);

return _test[_i][_qp] * _phi[_j][_qp] * _du_dot_du[_qp];

return _test(datum, i, qp) * _phi(datum, j, qp) * _du_dot_du;

#pragma once

#include "KokkosKernel.h"

/**
 * This calculates the time derivative for a coupled variable
 **/
class KokkosCoupledTimeDerivative final : public Moose::Kokkos::Kernel<KokkosCoupledTimeDerivative>
{
public:
  static InputParameters validParams();

  KokkosCoupledTimeDerivative(const InputParameters & parameters);

  KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
                                         const unsigned int qp,
                                         ResidualDatum & datum) const;
  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int i,
                                                const unsigned int j,
                                                const unsigned int jvar,
                                                const unsigned int qp,
                                                ResidualDatum & datum) const;

protected:
  const Moose::Kokkos::VariableValue _v_dot;
  const Moose::Kokkos::Scalar<const Real> _dv_dot;
  const unsigned int _v_var;
};

KOKKOS_FUNCTION inline Real
KokkosCoupledTimeDerivative::computeQpResidual(const unsigned int i,
                                               const unsigned int qp,
                                               ResidualDatum & datum) const
{
  return _test(datum, i, qp) * _v_dot(datum, qp);
}

KOKKOS_FUNCTION inline Real
KokkosCoupledTimeDerivative::computeQpOffDiagJacobian(const unsigned int i,
                                                      const unsigned int j,
                                                      const unsigned int jvar,
                                                      const unsigned int qp,
                                                      ResidualDatum & datum) const
{
  if (jvar == _v_var)
    return _test(datum, i, qp) * _phi(datum, j, qp) * _dv_dot;
  else
    return 0;
}

#include "KokkosCoupledTimeDerivative.h"

registerMooseObject("MooseApp", KokkosCoupledTimeDerivative);

InputParameters
KokkosCoupledTimeDerivative::validParams()
{
  InputParameters params = Kernel::validParams();
  params.addClassDescription("Time derivative Kernel that acts on a coupled variable. Weak form: "
                             "$(\\psi_i, \\frac{\\partial v_h}{\\partial t})$.");
  params.addRequiredCoupledVar("v", "Coupled variable");
  return params;
}

KokkosCoupledTimeDerivative::KokkosCoupledTimeDerivative(const InputParameters & parameters)
  : Kernel(parameters),
    _v_dot(kokkosCoupledDot("v")),
    _dv_dot(kokkosCoupledDotDu("v")),
    _v_var(coupled("v"))
{
}