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 "KokkosSystem.h"

#include "AuxKernelBase.h"

namespace Moose
{
namespace Kokkos
{

class AuxKernel : public ::AuxKernelBase,
                  public MeshHolder,
                  public AssemblyHolder,
                  public SystemHolder
{
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 Real computeValueShim(const Derived & auxkernel,
                                        const unsigned int qp,
                                        AssemblyDatum & datum) const
  {
    return auxkernel.computeValue(qp, datum);
  }

  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.computeValueShim(auxkernel, qp, datum);

    datum.reinit();

    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.computeValueShim(auxkernel, 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 Kokkos
} // namespace Moose