Kokkos AuxKernels System

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

AuxKernels System to understand the MOOSE auxiliary kernel system,
Getting Started with Kokkos-MOOSE to understand the programming practices for Kokkos-MOOSE,
Kokkos Kernels System to understand the common design pattern of objects in Kokkos-MOOSE.

note

Kokkos-MOOSE auxiliary kernels do not support boundary-restricted elemental kernels yet.

The Kokkos-MOOSE auxiliary kernels can be created by inheriting Moose::Kokkos::AuxKernel and registering them with registerKokkosAuxKernel(). There is no vector or array version at this moment. The signature of the computeValue() hook method is now as follows:


KOKKOS_FUNCTION Real computeValue(const unsigned int qp, AssemblyDatum & datum) const;

To reiterate, the hook method no longer uses virtual dispatch. Instead, it should be statically defined in the derived class as inlined public method.

Like the original MOOSE, the single interface is used for both elemental and nodal kernels, and whether the object is elemental or nodal can be queried with the isNodal() method on both CPU and GPU. Depending on whether the kernel is nodal or elemental, some data that can be queried through datum may be invalid. For instance, datum.node() will throw an invalid index in elemental kernels, and vice versa. Therefore, if you need to treat elemental and nodal variables separately, it is important to properly separate their code paths. For nodal kernels, the index qp is simply zero but still passed as an argument to have a unified interface, as explained in nodal boundary conditions.

The following example shows the implementation of KokkosVariableTimeIntegrationAux:

Listing 1: The KokkosVariableTimeIntegrationAux header file.


#pragma once

#include "KokkosAuxKernel.h"

/**
 * An AuxKernel that can be used to integrate a field variable in time
 * using a variety of different integration methods.  The result is
 * stored in another field variable.
 */
class KokkosVariableTimeIntegrationAux : public Moose::Kokkos::AuxKernel
{
public:
  static InputParameters validParams();

  KokkosVariableTimeIntegrationAux(const InputParameters & parameters);

  KOKKOS_FUNCTION Real computeValue(const unsigned int qp, AssemblyDatum & datum) const;

protected:
  KOKKOS_FUNCTION Real getIntegralValue(const unsigned int qp, AssemblyDatum & datum) const;

  Moose::Kokkos::Array<Moose::Kokkos::VariableValue> _coupled_vars;
  const Real _coef;
  const unsigned int _order;
  Moose::Kokkos::Array<Real> _integration_coef;

  /// The old variable value (zero if order == 3)
  const Moose::Kokkos::VariableValue _u_old;
  /// The older variable value (zero if order != 3)
  const Moose::Kokkos::VariableValue _u_older;
};

KOKKOS_FUNCTION inline Real
KokkosVariableTimeIntegrationAux::computeValue(const unsigned int qp, AssemblyDatum & datum) const
{
  Real integral = getIntegralValue(qp, datum);

  if (_order == 3)
    return _u_older(datum, qp) + _coef * integral;

  return _u_old(datum, qp) + _coef * integral;
}

KOKKOS_FUNCTION inline Real
KokkosVariableTimeIntegrationAux::getIntegralValue(const unsigned int qp,
                                                   AssemblyDatum & datum) const
{
  Real integral_value = 0.0;

  for (unsigned int i = 0; i < _order; ++i)
    integral_value += _integration_coef[i] * _coupled_vars[i](datum, qp) * _dt;

  /**
   * Subsequent timesteps may be unequal, so the standard Simpson rule
   * cannot be used. Use a different set of coefficients here.
   * J. McNAMEE, "A PROGRAM TO INTEGRATE A FUNCTION TABULATED AT
   * UNEQUAL INTERVALS," Internation Journal for Numerical Methods in
   * Engineering, Vol. 17, 217-279. (1981).
   */
  if (_order == 3 && _dt != _dt_old)
  {
    Real x0 = 0.0;
    Real x1 = _dt_old;
    Real x2 = _dt + _dt_old;
    Real y0 = _coupled_vars[2](datum, qp);
    Real y1 = _coupled_vars[1](datum, qp);
    Real y2 = _coupled_vars[0](datum, qp);
    Real term1 = (x2 - x0) * (y0 + (x2 - x0) * (y1 - y0) / (2.0 * (x1 - x0)));
    Real term2 = (2.0 * x2 * x2 - x0 * x2 - x0 * x0 + 3.0 * x0 * x1 - 3.0 * x1 * x2) / 6.0;
    Real term3 = (y2 - y1) / (x2 - x1) - (y1 - y0) / (x1 - x0);
    integral_value = term1 + term2 * term3;
  }

  return integral_value;
}

Listing 2: The KokkosVariableTimeIntegrationAux source file.


#include "KokkosVariableTimeIntegrationAux.h"

registerKokkosAuxKernel("MooseApp", KokkosVariableTimeIntegrationAux);

InputParameters
KokkosVariableTimeIntegrationAux::validParams()
{
  InputParameters params = AuxKernel::validParams();
  params.addClassDescription("Integrates a field variable in time.");
  params.addRequiredCoupledVar("variable_to_integrate", "The variable to be integrated");
  params.addParam<Real>("coefficient", 1.0, "A simple coefficient multiplying the integral");
  params.addParam<unsigned int>(
      "order", 2, "The order of global truncation error: midpoint=1, trapezoidal=2, Simpson=3");
  return params;
}

KokkosVariableTimeIntegrationAux::KokkosVariableTimeIntegrationAux(
    const InputParameters & parameters)
  : AuxKernel(parameters),
    _coef(getParam<Real>("coefficient")),
    _order(getParam<unsigned int>("order")),
    _u_old(_order != 3 ? uOld() : kokkosZeroValue()),
    _u_older(_order == 3 ? uOlder() : kokkosZeroValue())
{
  _integration_coef.create(_order);
  _coupled_vars.create(_order);

  switch (_order)
  {
    case 1:
      _integration_coef[0] = 1.0;
      _coupled_vars[0] = kokkosCoupledValue("variable_to_integrate");
      break;
    case 2:
      _integration_coef[0] = 0.5;
      _integration_coef[1] = 0.5;
      _coupled_vars[0] = kokkosCoupledValue("variable_to_integrate");
      _coupled_vars[1] = kokkosCoupledValueOld("variable_to_integrate");
      break;
    case 3:
      _integration_coef[0] = 1.0 / 3.0;
      _integration_coef[1] = 4.0 / 3.0;
      _integration_coef[2] = 1.0 / 3.0;
      _coupled_vars[0] = kokkosCoupledValue("variable_to_integrate");
      _coupled_vars[1] = kokkosCoupledValueOld("variable_to_integrate");
      _coupled_vars[2] = kokkosCoupledValueOlder("variable_to_integrate");
      break;
    default:
      mooseError("KokkosVariableTimeIntegrationAux: unknown time integration order specified");
  }

  _integration_coef.copyToDevice();
  _coupled_vars.copyToDevice();
}

Advanced Topics

The computational loops are defined in computeElementInternal() and computeNodeInternal() for elemental and nodal variables, respectively, and they call the hook method computeValue(). Depending on the type of operations, however, you may be able to write more efficient kernels by customizing the loops directly. For example, what KokkosCopyValueAux does is simply copy one auxiliary variable to another, and the shape function family and order are always identical between the two variables. In this case, you can directly loop over DOFs instead of quadrature points. See the following source code of KokkosCopyValueAux, where computeElementInternal() and computeNodeInternal() are redefined instead of computeValue():

Listing 3: The KokkosCopyValueAux header file.


#pragma once

#include "KokkosAuxKernel.h"

/**
 * Copies one variable onto an auxiliary variable
 */
class KokkosCopyValueAux : public Moose::Kokkos::AuxKernel
{
public:
  static InputParameters validParams();

  KokkosCopyValueAux(const InputParameters & parameters);

  template <typename Derived>
  KOKKOS_FUNCTION void computeElementInternal(const Derived & auxkernel,
                                              AssemblyDatum & datum) const;
  template <typename Derived>
  KOKKOS_FUNCTION void computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const;

protected:
  /// Variable used to specify state being copied
  unsigned short _state;

  /// The variable value to copy from
  const Moose::Kokkos::VariableValue _v;

  /// A reference to the variable to copy from
  const MooseVariable & _source_variable;
};

template <typename Derived>
KOKKOS_FUNCTION void
KokkosCopyValueAux::computeElementInternal(const Derived & /* auxkernel */,
                                           AssemblyDatum & datum) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto var = _kokkos_var.var();
  auto tag = _kokkos_var.tag();
  auto elem = datum.elem().id;

  for (unsigned int i = 0; i < datum.n_dofs(); ++i)
    sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, i, var), tag) = _v(datum, i);
}

template <typename Derived>
KOKKOS_FUNCTION void
KokkosCopyValueAux::computeNodeInternal(const Derived & /* auxkernel */,
                                        AssemblyDatum & datum) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto var = _kokkos_var.var();
  auto tag = _kokkos_var.tag();
  auto node = datum.node();

  sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, 0, var), tag) = _v(datum, 0);
}

Listing 4: The KokkosCopyValueAux source file.


#include "KokkosCopyValueAux.h"

registerKokkosAuxKernel("MooseApp", KokkosCopyValueAux);

InputParameters
KokkosCopyValueAux::validParams()
{
  InputParameters params = AuxKernel::validParams();
  params.addClassDescription("Returns the specified variable as an auxiliary variable with a "
                             "simple copy of the variable values.");
  params.addRequiredCoupledVar("source", "Variable to take the value of.");
  MooseEnum stateEnum("CURRENT=0 OLD=1 OLDER=2", "CURRENT");
  params.addParam<MooseEnum>(
      "state",
      stateEnum,
      "This parameter specifies the state being copied. CURRENT=0 OLD=1 OLDER=2. Copying an older "
      "state allows access to previous solution information if necessary.");
  return params;
}

KokkosCopyValueAux::KokkosCopyValueAux(const InputParameters & parameters)
  : AuxKernel(parameters),
    _state(getParam<MooseEnum>("state")),
    _v(_state == 0   ? kokkosCoupledDofValue("source")
       : _state == 1 ? kokkosCoupledDofValueOld("source")
                     : kokkosCoupledDofValueOlder("source")),
    _source_variable(*getVar("source", 0))
{
  if (_var.feType().family != _source_variable.feType().family)
    paramError("variable",
               "Source (" + Moose::stringify(_source_variable.feType().family) + ") and target (" +
                   Moose::stringify(_var.feType().family) +
                   ") variable have different families. You may use a ProjectionAux "
                   "instead of the CopyValueAux");
  if (_var.feType().order != _source_variable.feType().order)
    paramError("variable",
               "Source (" + Moose::stringify(_source_variable.feType().order) + ") and target (" +
                   Moose::stringify(_var.feType().order) +
                   ") variable are of different orders. You may use a ProjectionAux "
                   "instead of the CopyValueAux");
}

Available Objects

Moose App
KokkosConstantAuxCreates a constant field in the domain.
KokkosCopyValueAuxReturns the specified variable as an auxiliary variable with a simple copy of the variable values.
KokkosMaterialRealAuxOutputs element volume-averaged material properties
KokkosTagVectorAuxExtract DOF values from a tagged vector into an AuxVariable
KokkosVariableTimeIntegrationAuxIntegrates a field variable in time.

(framework/include/kokkos/auxkernels/KokkosVariableTimeIntegrationAux.h)

// This file is part of the MOOSE framework
// https://mooseframework.inl.gov
//
// 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

#include "KokkosAuxKernel.h"

/**
 * An AuxKernel that can be used to integrate a field variable in time
 * using a variety of different integration methods.  The result is
 * stored in another field variable.
 */
class KokkosVariableTimeIntegrationAux : public Moose::Kokkos::AuxKernel
{
public:
  static InputParameters validParams();

  KokkosVariableTimeIntegrationAux(const InputParameters & parameters);

  KOKKOS_FUNCTION Real computeValue(const unsigned int qp, AssemblyDatum & datum) const;

protected:
  KOKKOS_FUNCTION Real getIntegralValue(const unsigned int qp, AssemblyDatum & datum) const;

  Moose::Kokkos::Array<Moose::Kokkos::VariableValue> _coupled_vars;
  const Real _coef;
  const unsigned int _order;
  Moose::Kokkos::Array<Real> _integration_coef;

  /// The old variable value (zero if order == 3)
  const Moose::Kokkos::VariableValue _u_old;
  /// The older variable value (zero if order != 3)
  const Moose::Kokkos::VariableValue _u_older;
};

KOKKOS_FUNCTION inline Real
KokkosVariableTimeIntegrationAux::computeValue(const unsigned int qp, AssemblyDatum & datum) const
{
  Real integral = getIntegralValue(qp, datum);

  if (_order == 3)
    return _u_older(datum, qp) + _coef * integral;

  return _u_old(datum, qp) + _coef * integral;
}

KOKKOS_FUNCTION inline Real
KokkosVariableTimeIntegrationAux::getIntegralValue(const unsigned int qp,
                                                   AssemblyDatum & datum) const
{
  Real integral_value = 0.0;

  for (unsigned int i = 0; i < _order; ++i)
    integral_value += _integration_coef[i] * _coupled_vars[i](datum, qp) * _dt;

  /**
   * Subsequent timesteps may be unequal, so the standard Simpson rule
   * cannot be used. Use a different set of coefficients here.
   * J. McNAMEE, "A PROGRAM TO INTEGRATE A FUNCTION TABULATED AT
   * UNEQUAL INTERVALS," Internation Journal for Numerical Methods in
   * Engineering, Vol. 17, 217-279. (1981).
   */
  if (_order == 3 && _dt != _dt_old)
  {
    Real x0 = 0.0;
    Real x1 = _dt_old;
    Real x2 = _dt + _dt_old;
    Real y0 = _coupled_vars[2](datum, qp);
    Real y1 = _coupled_vars[1](datum, qp);
    Real y2 = _coupled_vars[0](datum, qp);
    Real term1 = (x2 - x0) * (y0 + (x2 - x0) * (y1 - y0) / (2.0 * (x1 - x0)));
    Real term2 = (2.0 * x2 * x2 - x0 * x2 - x0 * x0 + 3.0 * x0 * x1 - 3.0 * x1 * x2) / 6.0;
    Real term3 = (y2 - y1) / (x2 - x1) - (y1 - y0) / (x1 - x0);
    integral_value = term1 + term2 * term3;
  }

  return integral_value;
}

(framework/src/kokkos/auxkernels/KokkosVariableTimeIntegrationAux.K)

// This file is part of the MOOSE framework
// https://mooseframework.inl.gov
//
// 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

#include "KokkosVariableTimeIntegrationAux.h"

registerKokkosAuxKernel("MooseApp", KokkosVariableTimeIntegrationAux);

InputParameters
KokkosVariableTimeIntegrationAux::validParams()
{
  InputParameters params = AuxKernel::validParams();
  params.addClassDescription("Integrates a field variable in time.");
  params.addRequiredCoupledVar("variable_to_integrate", "The variable to be integrated");
  params.addParam<Real>("coefficient", 1.0, "A simple coefficient multiplying the integral");
  params.addParam<unsigned int>(
      "order", 2, "The order of global truncation error: midpoint=1, trapezoidal=2, Simpson=3");
  return params;
}

KokkosVariableTimeIntegrationAux::KokkosVariableTimeIntegrationAux(
    const InputParameters & parameters)
  : AuxKernel(parameters),
    _coef(getParam<Real>("coefficient")),
    _order(getParam<unsigned int>("order")),
    _u_old(_order != 3 ? uOld() : kokkosZeroValue()),
    _u_older(_order == 3 ? uOlder() : kokkosZeroValue())
{
  _integration_coef.create(_order);
  _coupled_vars.create(_order);

  switch (_order)
  {
    case 1:
      _integration_coef[0] = 1.0;
      _coupled_vars[0] = kokkosCoupledValue("variable_to_integrate");
      break;
    case 2:
      _integration_coef[0] = 0.5;
      _integration_coef[1] = 0.5;
      _coupled_vars[0] = kokkosCoupledValue("variable_to_integrate");
      _coupled_vars[1] = kokkosCoupledValueOld("variable_to_integrate");
      break;
    case 3:
      _integration_coef[0] = 1.0 / 3.0;
      _integration_coef[1] = 4.0 / 3.0;
      _integration_coef[2] = 1.0 / 3.0;
      _coupled_vars[0] = kokkosCoupledValue("variable_to_integrate");
      _coupled_vars[1] = kokkosCoupledValueOld("variable_to_integrate");
      _coupled_vars[2] = kokkosCoupledValueOlder("variable_to_integrate");
      break;
    default:
      mooseError("KokkosVariableTimeIntegrationAux: unknown time integration order specified");
  }

  _integration_coef.copyToDevice();
  _coupled_vars.copyToDevice();
}

(framework/include/kokkos/auxkernels/KokkosCopyValueAux.h)

// This file is part of the MOOSE framework
// https://mooseframework.inl.gov
//
// 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

#include "KokkosAuxKernel.h"

/**
 * Copies one variable onto an auxiliary variable
 */
class KokkosCopyValueAux : public Moose::Kokkos::AuxKernel
{
public:
  static InputParameters validParams();

  KokkosCopyValueAux(const InputParameters & parameters);

  template <typename Derived>
  KOKKOS_FUNCTION void computeElementInternal(const Derived & auxkernel,
                                              AssemblyDatum & datum) const;
  template <typename Derived>
  KOKKOS_FUNCTION void computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const;

protected:
  /// Variable used to specify state being copied
  unsigned short _state;

  /// The variable value to copy from
  const Moose::Kokkos::VariableValue _v;

  /// A reference to the variable to copy from
  const MooseVariable & _source_variable;
};

template <typename Derived>
KOKKOS_FUNCTION void
KokkosCopyValueAux::computeElementInternal(const Derived & /* auxkernel */,
                                           AssemblyDatum & datum) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto var = _kokkos_var.var();
  auto tag = _kokkos_var.tag();
  auto elem = datum.elem().id;

  for (unsigned int i = 0; i < datum.n_dofs(); ++i)
    sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, i, var), tag) = _v(datum, i);
}

template <typename Derived>
KOKKOS_FUNCTION void
KokkosCopyValueAux::computeNodeInternal(const Derived & /* auxkernel */,
                                        AssemblyDatum & datum) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto var = _kokkos_var.var();
  auto tag = _kokkos_var.tag();
  auto node = datum.node();

  sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, 0, var), tag) = _v(datum, 0);
}

(framework/src/kokkos/auxkernels/KokkosCopyValueAux.K)

// This file is part of the MOOSE framework
// https://mooseframework.inl.gov
//
// 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

#include "KokkosCopyValueAux.h"

registerKokkosAuxKernel("MooseApp", KokkosCopyValueAux);

InputParameters
KokkosCopyValueAux::validParams()
{
  InputParameters params = AuxKernel::validParams();
  params.addClassDescription("Returns the specified variable as an auxiliary variable with a "
                             "simple copy of the variable values.");
  params.addRequiredCoupledVar("source", "Variable to take the value of.");
  MooseEnum stateEnum("CURRENT=0 OLD=1 OLDER=2", "CURRENT");
  params.addParam<MooseEnum>(
      "state",
      stateEnum,
      "This parameter specifies the state being copied. CURRENT=0 OLD=1 OLDER=2. Copying an older "
      "state allows access to previous solution information if necessary.");
  return params;
}

KokkosCopyValueAux::KokkosCopyValueAux(const InputParameters & parameters)
  : AuxKernel(parameters),
    _state(getParam<MooseEnum>("state")),
    _v(_state == 0   ? kokkosCoupledDofValue("source")
       : _state == 1 ? kokkosCoupledDofValueOld("source")
                     : kokkosCoupledDofValueOlder("source")),
    _source_variable(*getVar("source", 0))
{
  if (_var.feType().family != _source_variable.feType().family)
    paramError("variable",
               "Source (" + Moose::stringify(_source_variable.feType().family) + ") and target (" +
                   Moose::stringify(_var.feType().family) +
                   ") variable have different families. You may use a ProjectionAux "
                   "instead of the CopyValueAux");
  if (_var.feType().order != _source_variable.feType().order)
    paramError("variable",
               "Source (" + Moose::stringify(_source_variable.feType().order) + ") and target (" +
                   Moose::stringify(_var.feType().order) +
                   ") variable are of different orders. You may use a ProjectionAux "
                   "instead of the CopyValueAux");
}

Advanced Topics
Available Objects