KokkosKernel.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

#include "KokkosKernelBase.h"

namespace Moose
{
namespace Kokkos
{

template <typename Derived>
class Kernel : public KernelBase
{
public:
  static InputParameters validParams();

  Kernel(const InputParameters & parameters);

  virtual void computeResidual() override;
  virtual void computeJacobian() override;


  KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
                                         const unsigned int /* j */,
                                         const unsigned int /* qp */,
                                         ResidualDatum & /* datum */) const
  {
    return 0;
  }
  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
                                                const unsigned int /* j */,
                                                const unsigned int /* jvar */,
                                                const unsigned int /* qp */,
                                                ResidualDatum & /* datum */) const
  {
    return 0;
  }

  KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid) const;
  KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid) const;
  KOKKOS_FUNCTION void operator()(OffDiagJacobianLoop, const ThreadID tid) const;


  KOKKOS_FUNCTION void computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const;
  KOKKOS_FUNCTION void computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const;
  KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived * kernel,
                                                      ResidualDatum & datum) const;

protected:
  const VariableTestValue _test;
  const VariableTestGradient _grad_test;
  const VariablePhiValue _phi;
  const VariablePhiGradient _grad_phi;
  const VariableValue _u;
  const VariableGradient _grad_u;

  virtual bool defaultJacobian() const
  {
    return &Derived::computeQpJacobian == &Kernel::computeQpJacobian;
  }
  virtual bool defaultOffDiagJacobian() const
  {
    return &Derived::computeQpOffDiagJacobian == &Kernel::computeQpOffDiagJacobian;
  }
};

template <typename Derived>
InputParameters
Kernel<Derived>::validParams()
{
  InputParameters params = KernelBase::validParams();
  return params;
}

template <typename Derived>
Kernel<Derived>::Kernel(const InputParameters & parameters)
  : KernelBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD),
    _test(),
    _grad_test(),
    _phi(),
    _grad_phi(),
    _u(_var),
    _grad_u(_var)
{
  addMooseVariableDependency(&_var);
}

template <typename Derived>
void
Kernel<Derived>::computeResidual()
{
  ::Kokkos::RangePolicy<ResidualLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
      0, numKokkosBlockElements());
  ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
  ::Kokkos::fence();
}

template <typename Derived>
void
Kernel<Derived>::computeJacobian()
{
  if (!defaultJacobian())
  {
    ::Kokkos::RangePolicy<JacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
        0, numKokkosBlockElements());
    ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
    ::Kokkos::fence();
  }

  if (!defaultOffDiagJacobian())
  {
    auto & sys = kokkosSystem(_kokkos_var.sys());

    _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBlockElements()});

    ::Kokkos::RangePolicy<OffDiagJacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
        0, _thread.size());
    ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
    ::Kokkos::fence();
  }
}

template <typename Derived>
KOKKOS_FUNCTION void
Kernel<Derived>::operator()(ResidualLoop, const ThreadID tid) const
{
  auto kernel = static_cast<const Derived *>(this);
  auto elem = kokkosBlockElementID(tid);

  ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());

  kernel->computeResidualInternal(kernel, datum);
}

template <typename Derived>
KOKKOS_FUNCTION void
Kernel<Derived>::operator()(JacobianLoop, const ThreadID tid) const
{
  auto kernel = static_cast<const Derived *>(this);
  auto elem = kokkosBlockElementID(tid);

  ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());

  kernel->computeJacobianInternal(kernel, datum);
}

template <typename Derived>
KOKKOS_FUNCTION void
Kernel<Derived>::operator()(OffDiagJacobianLoop, const ThreadID tid) const
{
  auto kernel = static_cast<const Derived *>(this);
  auto elem = kokkosBlockElementID(_thread(tid, 1));

  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];

  if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
    return;

  ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);

  kernel->computeOffDiagJacobianInternal(kernel, datum);
}

template <typename Derived>
KOKKOS_FUNCTION void
Kernel<Derived>::computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const
{
  ResidualObject::computeResidualInternal(
      datum,
      [&](Real * local_re, const unsigned int ib, const unsigned int ie)
      {
        for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
        {
          datum.reinit();

          for (unsigned int i = ib; i < ie; ++i)
            local_re[i] += datum.JxW(qp) * kernel->computeQpResidual(i, qp, datum);
        }
      });
}

template <typename Derived>
KOKKOS_FUNCTION void
Kernel<Derived>::computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
{
  ResidualObject::computeJacobianInternal(
      datum,
      [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
      {
        for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
        {
          datum.reinit();

          for (unsigned int ij = ijb; ij < ije; ++ij)
          {
            unsigned int i = ij % datum.n_jdofs();
            unsigned int j = ij / datum.n_jdofs();

            local_ke[ij] += datum.JxW(qp) * kernel->computeQpJacobian(i, j, qp, datum);
          }
        }
      });
}

template <typename Derived>
KOKKOS_FUNCTION void
Kernel<Derived>::computeOffDiagJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
{
  ResidualObject::computeJacobianInternal(
      datum,
      [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
      {
        for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
        {
          datum.reinit();

          for (unsigned int ij = ijb; ij < ije; ++ij)
          {
            unsigned int i = ij % datum.n_jdofs();
            unsigned int j = ij / datum.n_jdofs();

            local_ke[ij] +=
                datum.JxW(qp) * kernel->computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
          }
        }
      });
}

} // namespace Kokkos
} // namespace Moose

#define usingKokkosKernelMembers(T)                                                                \
  usingKokkosKernelBaseMembers;                                                                    \
                                                                                                   \
protected:                                                                                         \
  using Moose::Kokkos::Kernel<T>::_test;                                                           \
  using Moose::Kokkos::Kernel<T>::_grad_test;                                                      \
  using Moose::Kokkos::Kernel<T>::_phi;                                                            \
  using Moose::Kokkos::Kernel<T>::_grad_phi;                                                       \
  using Moose::Kokkos::Kernel<T>::_u;                                                              \
  using Moose::Kokkos::Kernel<T>::_grad_u;                                                         \
                                                                                                   \
public:                                                                                            \
  using Moose::Kokkos::Kernel<T>::operator()