KokkosKernelGrad.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 "KokkosKernel.h"

namespace Moose
{
namespace Kokkos
{

template <typename Derived>
class KernelGrad : public Kernel<Derived>
{
  usingKokkosKernelMembers(Derived);

public:
  static InputParameters validParams();

  KernelGrad(const InputParameters & parameters);


  KOKKOS_FUNCTION Real3 precomputeQpJacobian(const unsigned int /* j */,
                                             const unsigned int /* qp */,
                                             ResidualDatum & /* datum */) const
  {
    return Real3(0);
  }

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

protected:
  virtual bool defaultJacobian() const override
  {
    return &Derived::precomputeQpJacobian == &KernelGrad::precomputeQpJacobian;
  }
};

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

template <typename Derived>
KernelGrad<Derived>::KernelGrad(const InputParameters & parameters) : Kernel<Derived>(parameters)
{
}

template <typename Derived>
KOKKOS_FUNCTION void
KernelGrad<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();

          Real3 value = datum.JxW(qp) * kernel->precomputeQpResidual(qp, datum);

          for (unsigned int i = ib; i < ie; ++i)
            local_re[i] += value * _grad_test(datum, i, qp);
        }
      });
}

template <typename Derived>
KOKKOS_FUNCTION void
KernelGrad<Derived>::computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
{
  Real3 value;

  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();

          unsigned int j_old = libMesh::invalid_uint;

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

            if (j != j_old)
            {
              value = datum.JxW(qp) * kernel->precomputeQpJacobian(j, qp, datum);
              j_old = j;
            }

            local_ke[ij] += value * _grad_test(datum, i, qp);
          }
        }
      });
}

} // namespace Kokkos
} // namespace Moose

#define usingKokkosKernelGradMembers(T) usingKokkosKernelMembers(T)