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