KokkosNodalKernel.h

Go to the documentation of this file.

//* 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 "KokkosNodalKernelBase.h"

namespace Moose
{
namespace Kokkos
{

template <typename Derived>
class NodalKernel : public NodalKernelBase
{
public:
  static InputParameters validParams();

  NodalKernel(const InputParameters & parameters);

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


  KOKKOS_FUNCTION Real computeQpJacobian(const ContiguousNodeID /* node */) const { return 0; }
  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* jvar */,
                                                const ContiguousNodeID /* node */) 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;

protected:
  const VariableNodalValue _u;

private:
  const bool _boundary_restricted;
  const bool _default_diag;
  const bool _default_offdiag;
};

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

template <typename Derived>
NodalKernel<Derived>::NodalKernel(const InputParameters & parameters)
  : NodalKernelBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD),
    _u(kokkosSystems(), _var),
    _boundary_restricted(boundaryRestricted()),
    _default_diag(&Derived::computeQpJacobian == &NodalKernel::computeQpJacobian),
    _default_offdiag(&Derived::computeQpOffDiagJacobian == &NodalKernel::computeQpOffDiagJacobian)
{
}

template <typename Derived>
void
NodalKernel<Derived>::computeResidual()
{
  ::Kokkos::RangePolicy<ResidualLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
      0, _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes());
  ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
  ::Kokkos::fence();
}

template <typename Derived>
void
NodalKernel<Derived>::computeJacobian()
{
  if (!_default_diag)
  {
    ::Kokkos::RangePolicy<JacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
        0, _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes());
    ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
    ::Kokkos::fence();
  }

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

    _thread.resize({sys.getCoupling(_kokkos_var.var()).size(),
                    _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes()});

    ::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
NodalKernel<Derived>::operator()(ResidualLoop, const ThreadID tid) const
{
  auto kernel = static_cast<const Derived *>(this);
  auto node = _boundary_restricted ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid);
  auto & sys = kokkosSystem(_kokkos_var.sys());

  if (!sys.isNodalDefined(node, _kokkos_var.var()))
    return;

  Real local_re = kernel->computeQpResidual(node);

  accumulateTaggedNodalResidual(true, local_re, node);
}

template <typename Derived>
KOKKOS_FUNCTION void
NodalKernel<Derived>::operator()(JacobianLoop, const ThreadID tid) const
{
  auto kernel = static_cast<const Derived *>(this);
  auto node = _boundary_restricted ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid);
  auto & sys = kokkosSystem(_kokkos_var.sys());

  if (!sys.isNodalDefined(node, _kokkos_var.var()))
    return;

  Real local_ke = kernel->computeQpJacobian(node);

  accumulateTaggedNodalMatrix(true, local_ke, node, _kokkos_var.var());
}

template <typename Derived>
KOKKOS_FUNCTION void
NodalKernel<Derived>::operator()(OffDiagJacobianLoop, const ThreadID tid) const
{
  auto kernel = static_cast<const Derived *>(this);
  auto node = _boundary_restricted ? kokkosBoundaryNodeID(_thread(tid, 1))
                                   : kokkosBlockNodeID(_thread(tid, 1));
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];

  if (!sys.isNodalDefined(node, _kokkos_var.var()))
    return;

  Real local_ke = kernel->computeQpOffDiagJacobian(jvar, node);

  accumulateTaggedNodalMatrix(true, local_ke, node, jvar);
}

} // namespace Kokkos
} // namespace Moose

#define usingKokkosNodalKernelMembers(T)                                                           \
  usingKokkosNodalKernelBaseMembers;                                                               \
                                                                                                   \
protected:                                                                                         \
  using Moose::Kokkos::NodalKernel<T>::_u;                                                         \
                                                                                                   \
public:                                                                                            \
  using Moose::Kokkos::NodalKernel<T>::operator()