KokkosResidualObject.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 "KokkosFESystem.h"
#include "KokkosVariableValue.h"
#include "KokkosMaterialPropertyValue.h"
 
#include "MooseVariableBase.h"
#include "ResidualObject.h"
 
namespace Moose::Kokkos
{
 
class ResidualObject : public ::ResidualObject,
                       public MeshHolder,
                       public AssemblyHolder,
                       public FESystemHolder
{
public:
  static InputParameters validParams();
 
  static constexpr bool uses_precompute_hooks = false;
 
  static constexpr bool supports_scalar_jacobian = false;
 
  ResidualObject(const InputParameters & parameters,
                 Moose::VarFieldType field_type,
                 bool nodal = false);
  ResidualObject(const ResidualObject & object);
 
  struct ResidualLoop
  {
  };
  struct JacobianLoop
  {
  };
  struct OffDiagJacobianLoop
  {
  };
  struct OffDiagJacobianScalarLoop
  {
  };
 
  virtual const MooseVariableBase & variable() const override { return _var; }
 
  virtual void computeOffDiagJacobian(unsigned int) override final
  {
    mooseError("computeOffDiagJacobian() is not used for Kokkos residual objects.");
  }
  virtual void computeResidualAndJacobian() override
  {
    computeResidual();
    computeJacobian();
  }
 
protected:
  MooseVariableFieldBase & _var;
  Variable _kokkos_var;
  Thread<> _thread;
  std::unique_ptr<DispatcherBase> _residual_dispatcher;
  std::unique_ptr<DispatcherBase> _jacobian_dispatcher;
  std::unique_ptr<DispatcherBase> _offdiag_jacobian_dispatcher;
  std::unique_ptr<DispatcherBase> _offdiag_jacobian_scalar_dispatcher;
 
  const unsigned int _dimension;
 
 
  Scalar<Real> _t;
  Scalar<const Real> _t_old;
  Scalar<int> _t_step;
  Scalar<Real> _dt;
  Scalar<Real> _dt_old;
 
  KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
                                                         const ContiguousElementID elem,
                                                         const unsigned int i,
                                                         const unsigned int comp = 0) const;
  KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
                                                     const Real local_re,
                                                     const ContiguousNodeID node,
                                                     const unsigned int comp = 0) const;
  KOKKOS_FUNCTION void accumulateTaggedVectorNodalResidual(const bool add,
                                                           const Real local_re,
                                                           const ContiguousNodeID node,
                                                           const unsigned int i) const;
  KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke,
                                                       const ContiguousElementID elem,
                                                       const unsigned int i,
                                                       const unsigned int j,
                                                       const unsigned int jvar,
                                                       const unsigned int comp = 0) const;
  KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
                                                       const AssemblyDatum & datum,
                                                       const unsigned int i,
                                                       const unsigned int comp = 0) const;
  KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
                                                   const Real local_ke,
                                                   const ContiguousNodeID node,
                                                   const unsigned int jvar,
                                                   const unsigned int comp = 0) const;
  KOKKOS_FUNCTION void accumulateTaggedVectorNodalMatrix(const bool add,
                                                         const Real local_ke,
                                                         const ContiguousNodeID node,
                                                         const unsigned int i,
                                                         const unsigned int j,
                                                         const unsigned int jvar) const;
  KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
                                                   const DNDerivativeType & local_ke,
                                                   const ContiguousNodeID node,
                                                   const unsigned int comp = 0) const;
 
  template <typename function>
  KOKKOS_FUNCTION void computeResidualInternal(AssemblyDatum & datum, function body) const;
  template <typename function>
  KOKKOS_FUNCTION void computeJacobianInternal(AssemblyDatum & datum, function body) const;
 
private:
  Array<TagID> _vector_tags;
  Array<TagID> _matrix_tags;
};
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedElementalResidual(const Real local_re,
                                                  const ContiguousElementID elem,
                                                  const unsigned int i,
                                                  const unsigned int comp) const
{
  if (!local_re)
    return;
 
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  const auto var = _kokkos_var.var(comp);
  auto dof = sys.getElemLocalDofIndex(elem, i, var);
  const auto scaled_local_re = local_re * sys.getVariableScalingFactor(var);
 
  for (unsigned int t = 0; t < _vector_tags.size(); ++t)
  {
    auto tag = _vector_tags[t];
 
    if (sys.isResidualTagActive(tag))
      ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), scaled_local_re);
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedNodalResidual(const bool add,
                                              const Real local_re,
                                              const ContiguousNodeID node,
                                              const unsigned int comp) const
{
  if (!local_re && add)
    return;
 
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  auto dof = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
 
  for (unsigned int t = 0; t < _vector_tags.size(); ++t)
  {
    auto tag = _vector_tags[t];
 
    if (sys.isResidualTagActive(tag))
    {
      if (add)
        sys.getVectorDofValue(dof, tag) += local_re;
      else
        sys.getVectorDofValue(dof, tag) = local_re;
    }
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedVectorNodalResidual(const bool add,
                                                    const Real local_re,
                                                    const ContiguousNodeID node,
                                                    const unsigned int i) const
{
  if (!local_re && add)
    return;
 
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto dof = sys.getNodeLocalDofIndex(node, i, _kokkos_var.var());
 
  for (unsigned int t = 0; t < _vector_tags.size(); ++t)
  {
    auto tag = _vector_tags[t];
 
    if (sys.isResidualTagActive(tag))
    {
      if (add)
        sys.getVectorDofValue(dof, tag) += local_re;
      else
        sys.getVectorDofValue(dof, tag) = local_re;
    }
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedElementalMatrix(const Real local_ke,
                                                const ContiguousElementID elem,
                                                const unsigned int i,
                                                const unsigned int j,
                                                const unsigned int jvar,
                                                const unsigned int comp) const
{
  if (!local_ke)
    return;
 
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  const auto ivar = _kokkos_var.var(comp);
  auto row = sys.getElemLocalDofIndex(elem, i, ivar);
  auto col = sys.isScalarVariable(jvar) ? sys.getScalarGlobalDofIndex(j, jvar)
                                        : sys.getElemGlobalDofIndex(elem, j, jvar);
  const auto scaled_local_ke = local_ke * sys.getVariableScalingFactor(ivar);
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
 
    if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))
      ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), scaled_local_ke);
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
                                                const AssemblyDatum & datum,
                                                const unsigned int i,
                                                const unsigned int comp) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  const auto ivar = _kokkos_var.var(comp);
  auto row = sys.getElemLocalDofIndex(datum.elem().id, i, ivar);
  const auto scaling_factor = sys.getVariableScalingFactor(ivar);
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
 
    if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))
      for (unsigned int j = 0; j < local_ke.size(); ++j)
      {
        auto col = local_ke.raw_index(j);
 
        ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag),
                             local_ke.raw_at(j) * scaling_factor);
      }
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedNodalMatrix(const bool add,
                                            const Real local_ke,
                                            const ContiguousNodeID node,
                                            const unsigned int jvar,
                                            const unsigned int comp) const
{
  if (!local_ke && add)
    return;
 
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
  auto col = sys.getNodeGlobalDofIndex(node, 0, jvar);
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
 
    if (sys.isMatrixTagActive(tag))
    {
      auto & matrix = sys.getMatrix(tag);
 
      if (add)
        matrix(row, col) += local_ke;
      else
      {
        matrix.zero(row);
        matrix(row, col) = local_ke;
      }
    }
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedVectorNodalMatrix(const bool add,
                                                  const Real local_ke,
                                                  const ContiguousNodeID node,
                                                  const unsigned int i,
                                                  const unsigned int j,
                                                  const unsigned int jvar) const
{
  if (!local_ke && add)
    return;
 
  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto row = sys.getNodeLocalDofIndex(node, i, _kokkos_var.var());
  auto col = sys.getNodeGlobalDofIndex(node, j, jvar);
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
 
    if (sys.isMatrixTagActive(tag))
    {
      auto & matrix = sys.getMatrix(tag);
 
      if (add)
        matrix(row, col) += local_ke;
      else
      {
        matrix.zero(row);
        matrix(row, col) = local_ke;
      }
    }
  }
}
 
KOKKOS_FUNCTION inline void
ResidualObject::accumulateTaggedNodalMatrix(const bool add,
                                            const DNDerivativeType & local_ke,
                                            const ContiguousNodeID node,
                                            const unsigned int comp) const
{
  auto & sys = kokkosSystem(_kokkos_var.sys(comp));
  auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
    auto & matrix = sys.getMatrix(tag);
 
    if (sys.isMatrixTagActive(tag))
    {
      if (!add)
        matrix.zero(row);
 
      for (unsigned int j = 0; j < local_ke.size(); ++j)
      {
        auto col = local_ke.raw_index(j);
 
        if (add)
          matrix(row, col) += local_ke.raw_at(j);
        else
          matrix(row, col) = local_ke.raw_at(j);
      }
    }
  }
}
 
template <typename function>
KOKKOS_FUNCTION void
ResidualObject::computeResidualInternal(AssemblyDatum & datum, function body) const
{
  Real local_re[MAX_CACHED_DOF];
 
  unsigned int stride = MAX_CACHED_DOF * datum.num_local_threads();
  unsigned int num_batches = datum.n_dofs() / stride;
 
  if (datum.n_dofs() % stride)
    ++num_batches;
 
  for (unsigned int batch = 0; batch < num_batches; ++batch)
  {
    unsigned int ib = batch * stride;
    unsigned int ie = ::Kokkos::min(ib + stride, datum.n_dofs());
 
    const unsigned int n = ie - ib;
    const unsigned int d = n / datum.num_local_threads();
    const unsigned int m = n % datum.num_local_threads();
    const unsigned int t = datum.local_thread_id();
 
    ib += t * d + (t < m ? t : m);
    ie = ib + d + (t < m ? 1 : 0);
 
    for (unsigned int i = ib; i < ie; ++i)
      local_re[i - ib] = 0;
 
    body(local_re - ib, ib, ie);
 
    for (unsigned int i = ib; i < ie; ++i)
      accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
  }
}
 
template <typename function>
KOKKOS_FUNCTION void
ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) const
{
  Real local_ke[MAX_CACHED_DOF];
 
  for (unsigned int j = datum.local_thread_id(); j < datum.n_jdofs();
       j += datum.num_local_threads())
  {
    unsigned int num_batches = datum.n_idofs() / MAX_CACHED_DOF;
 
    if (datum.n_idofs() % MAX_CACHED_DOF)
      ++num_batches;
 
    for (unsigned int batch = 0; batch < num_batches; ++batch)
    {
      unsigned int ib = batch * MAX_CACHED_DOF;
      unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_idofs());
 
      for (unsigned int i = ib; i < ie; ++i)
        local_ke[i - ib] = 0;
 
      body(local_ke - ib, ib, ie, j);
 
      for (unsigned int i = ib; i < ie; ++i)
        accumulateTaggedElementalMatrix(local_ke[i - ib], datum.elem().id, i, j, datum.jvar());
    }
  }
}
 
} // namespace Moose::Kokkos