KokkosADKernel.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::Kokkos
{
 
class ADKernel : public KernelBase
{
public:
  static InputParameters validParams();
 
  ADKernel(const InputParameters & parameters);
 
  virtual void computeResidual() override;
  virtual void computeJacobian() override;
  virtual void computeResidualAndJacobian() override;
 
  template <typename Derived>
  KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
 
  template <typename Derived>
  KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const;
 
protected:
  virtual void dispatch();
 
  const ADVariableTestValue _test;
  const ADVariableTestGradient _grad_test;
  const ADVariablePhiValue _phi;
  const ADVariablePhiGradient _grad_phi;
  const ADVariableValue _u;
  const ADVariableGradient _grad_u;
  bool _computing_residual = false;
  bool _computing_jacobian = false;
};
 
template <typename Derived>
KOKKOS_FUNCTION void
ADKernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
{
  auto elem = kokkosBlockElementID(_thread(tid, 1));
 
  AssemblyDatum datum(elem,
                      libMesh::invalid_uint,
                      kokkosAssembly(),
                      kokkosSystems(),
                      _kokkos_var,
                      _kokkos_var.var());
 
  datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
 
  datum.do_derivatives(_computing_jacobian);
 
  kernel.computeResidualInternal(kernel, datum);
}
 
template <typename Derived>
KOKKOS_FUNCTION void
ADKernel::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
{
  for (unsigned int i = datum.local_thread_id(); i < datum.n_dofs(); i += datum.num_local_threads())
  {
    ADReal local_re = 0;
 
    for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
      local_re += datum.JxW(qp) * kernel.template computeQpResidual<Derived>(i, qp, datum);
 
    if (_computing_residual)
      accumulateTaggedElementalResidual(local_re.value(), datum.elem().id, i);
    if (_computing_jacobian)
      accumulateTaggedElementalMatrix(local_re.derivatives(), datum, i);
  }
}
 
} // namespace Moose::Kokkos