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
{

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 unsigned int /* qp */,
                                         AssemblyDatum & /* datum */) const
  {
    return 0;
  }
  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* jvar */,
                                                const unsigned int /* qp */,
                                                AssemblyDatum & /* datum */) const
  {
    return 0;
  }
  static auto defaultJacobian() { return &NodalKernel::computeQpJacobian; }
  static auto defaultOffDiagJacobian() { return &NodalKernel::computeQpOffDiagJacobian; }

  template <typename Derived>
  KOKKOS_FUNCTION Real computeQpResidualShim(const Derived & kernel,
                                             const unsigned int qp,
                                             AssemblyDatum & datum) const
  {
    return kernel.computeQpResidual(qp, datum);
  }
  template <typename Derived>
  KOKKOS_FUNCTION Real computeQpJacobianShim(const Derived & kernel,
                                             const unsigned int qp,
                                             AssemblyDatum & datum) const
  {
    return kernel.computeQpJacobian(qp, datum);
  }
  template <typename Derived>
  KOKKOS_FUNCTION Real computeQpOffDiagJacobianShim(const Derived & kernel,
                                                    const unsigned int jvar,
                                                    const unsigned int qp,
                                                    AssemblyDatum & datum) const
  {
    return kernel.computeQpOffDiagJacobian(jvar, qp, datum);
  }

  template <typename Derived>
  KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
  template <typename Derived>
  KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
  template <typename Derived>
  KOKKOS_FUNCTION void
  operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;

protected:
  const VariableValue _u;

private:
  const bool _boundary_restricted;
};

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

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

  AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());

  Real local_re = kernel.computeQpResidualShim(kernel, 0, datum);

  accumulateTaggedNodalResidual(true, local_re, node);
}

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

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

  AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());

  Real local_ke = kernel.computeQpJacobianShim(kernel, 0, datum);

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

template <typename Derived>
KOKKOS_FUNCTION void
NodalKernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
{
  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;

  AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);

  Real local_ke = kernel.computeQpOffDiagJacobianShim(kernel, jvar, 0, datum);

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

} // namespace Kokkos
} // namespace Moose