KokkosAuxKernel.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 "KokkosDispatcher.h"
#include "KokkosVariableValue.h"
#include "KokkosMaterialPropertyValue.h"
#include "KokkosAssembly.h"
#include "KokkosFESystem.h"
 
#include "AuxKernelBase.h"
 
namespace Moose::Kokkos
{
 
class AuxKernel : public ::AuxKernelBase,
                  public MeshHolder,
                  public AssemblyHolder,
                  public FESystemHolder
{
public:
  static InputParameters validParams();
 
  AuxKernel(const InputParameters & parameters);
 
  AuxKernel(const AuxKernel & object);
 
  // Unused for Kokkos auxiliary kernels because all elements are computed in parallel
  virtual void subdomainSetup() override final {}
 
  virtual void compute() override;
 
  KOKKOS_FUNCTION bool isNodal() const { return _nodal; }
 
  struct ElementLoop
  {
  };
  struct NodeLoop
  {
  };
 
  template <typename Derived>
  KOKKOS_FUNCTION void operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const;
  template <typename Derived>
  KOKKOS_FUNCTION void operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const;
 
 
  template <typename Derived>
  KOKKOS_FUNCTION void computeElementInternal(const Derived & auxkernel,
                                              AssemblyDatum & datum) const;
  template <typename Derived>
  KOKKOS_FUNCTION void computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const;
 
protected:
  VariableValue uOld() const;
  VariableValue uOlder() const;
 
  KOKKOS_FUNCTION void setElementSolution(const Real * const values,
                                          const AssemblyDatum & datum,
                                          const unsigned int comp = 0) const;
  KOKKOS_FUNCTION void
  setNodeSolution(const Real value, const AssemblyDatum & datum, const unsigned int comp = 0) const;
 
  const bool _nodal;
 
  Variable _kokkos_var;
  std::unique_ptr<DispatcherBase> _element_dispatcher;
  std::unique_ptr<DispatcherBase> _node_dispatcher;
 
  const VariableTestValue _test;
  const VariableValue _u;
 
 
  Scalar<Real> _t;
  Scalar<const Real> _t_old;
  Scalar<int> _t_step;
  Scalar<Real> _dt;
  Scalar<Real> _dt_old;
 
private:
  void getKokkosMaterialPropertyHook(const std::string & prop_name_in,
                                     const unsigned int state) override final;
};
 
template <typename Derived>
KOKKOS_FUNCTION void
AuxKernel::operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const
{
  auto elem = kokkosBlockElementID(tid);
 
  AssemblyDatum datum(elem,
                      libMesh::invalid_uint,
                      kokkosAssembly(),
                      kokkosSystems(),
                      _kokkos_var,
                      _kokkos_var.var());
 
  auxkernel.computeElementInternal(auxkernel, datum);
}
 
template <typename Derived>
KOKKOS_FUNCTION void
AuxKernel::operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const
{
  auto node = _bnd ? 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());
 
  auxkernel.computeNodeInternal(auxkernel, datum);
}
 
template <typename Derived>
KOKKOS_FUNCTION void
AuxKernel::computeElementInternal(const Derived & auxkernel, AssemblyDatum & datum) const
{
  Real x[MAX_CACHED_DOF];
  Real b[MAX_CACHED_DOF];
  Real A[MAX_CACHED_DOF * MAX_CACHED_DOF];
 
  for (unsigned int i = 0; i < datum.n_dofs(); ++i)
  {
    x[i] = 0;
    b[i] = 0;
 
    for (unsigned int j = 0; j < datum.n_dofs(); ++j)
      A[j + datum.n_dofs() * i] = 0;
  }
 
  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
  {
    const auto value = auxkernel.template computeValue<Derived>(qp, datum);
 
    for (unsigned int i = 0; i < datum.n_dofs(); ++i)
    {
      const auto t = datum.JxW(qp) * _test(datum, i, qp);
 
      b[i] += t * value;
 
      for (unsigned int j = 0; j < datum.n_dofs(); ++j)
        A[j + datum.n_dofs() * i] += t * _test(datum, j, qp);
    }
  }
 
  if (datum.n_dofs() == 1)
    x[0] = b[0] / A[0];
  else
    Utils::choleskySolve(A, x, b, datum.n_dofs());
 
  setElementSolution(x, datum);
}
 
template <typename Derived>
KOKKOS_FUNCTION void
AuxKernel::computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const
{
  auto value = auxkernel.template computeValue<Derived>(0, datum);
 
  setNodeSolution(value, datum);
}
 
KOKKOS_FUNCTION inline void
AuxKernel::setElementSolution(const Real * const values,
                              const AssemblyDatum & datum,
                              const unsigned int comp) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  auto var_num = _kokkos_var.var(comp);
  auto tag = _kokkos_var.tag();
  auto elem = datum.elem().id;
 
  for (unsigned int i = 0; i < datum.n_dofs(); ++i)
    sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, i, var_num), tag) = values[i];
}
 
KOKKOS_FUNCTION inline void
AuxKernel::setNodeSolution(const Real value,
                           const AssemblyDatum & datum,
                           const unsigned int comp) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  auto var_num = _kokkos_var.var(comp);
  auto tag = _kokkos_var.tag();
  auto node = datum.node();
 
  sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, 0, var_num), tag) = value;
}
 
} // namespace Moose::Kokkos