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
  {
    if (_vars.size() > 1)
      mooseError("variable() was called for a multi-variable Kokkos residual object. Use "
                 "variables() instead.");
 
    return *_vars[0];
  }
 
  virtual std::vector<const MooseVariableBase *> variables() const override
  {
    return {_vars.begin(), _vars.end()};
  }
 
  virtual void computeOffDiagJacobian(unsigned int) override final
  {
    mooseError("computeOffDiagJacobian() is not used for Kokkos residual objects.");
  }
  virtual void computeResidualAndJacobian() override
  {
    computeResidual();
    computeJacobian();
  }
 
protected:
  const std::vector<MooseVariableFieldBase *> _vars;
  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) const;
  KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
                                                     const Real local_re,
                                                     const ContiguousNodeID node,
                                                     const unsigned int comp) 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) const;
  KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
                                                       const AssemblyDatum & datum,
                                                       const unsigned int i,
                                                       const unsigned int comp) const;
  KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
                                                   const Real local_ke,
                                                   const ContiguousNodeID node,
                                                   const unsigned int jvar,
                                                   const unsigned int comp) 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) 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:
  std::vector<MooseVariableFieldBase *> getMooseVariables(Moose::VarFieldType field_type);
 
  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));
  auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
  auto constraints = sys.getLocalConstraints(dof);
 
  for (unsigned int t = 0; t < _vector_tags.size(); ++t)
  {
    auto tag = _vector_tags[t];
 
    if (sys.isResidualTagActive(tag))
    {
      if (constraints.size())
        for (unsigned int c = 0; c < constraints.size(); ++c)
        {
          auto [constraint_dof, weight] = constraints[c];
 
          ::Kokkos::atomic_add(&sys.getVectorDofValue(constraint_dof, tag), weight * local_re);
        }
      else
        ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), 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.getScalarLocalDofIndex(j, jvar)
                                        : sys.getElemLocalDofIndex(elem, j, jvar);
  auto col_global = sys.localToGlobalDofIndex(col);
  auto row_constraints = sys.getLocalConstraints(row);
  auto col_constraints = sys.getLocalConstraints(col);
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
 
    if (sys.isMatrixTagActive(tag))
    {
      for (unsigned int rc = 0; rc < ::Kokkos::max(1u, row_constraints.size()); ++rc)
      {
        auto constraint_row = row_constraints.size() ? row_constraints[rc].first : row;
        auto weight_row = row_constraints.size() ? row_constraints[rc].second : 1.0;
 
        if (sys.hasNodalBCMatrixTag(constraint_row, tag))
          continue;
 
        for (unsigned int cc = 0; cc < ::Kokkos::max(1u, col_constraints.size()); ++cc)
        {
          auto constraint_col = col_constraints.size()
                                    ? sys.localToGlobalDofIndex(col_constraints[cc].first)
                                    : col_global;
          auto weight_col = col_constraints.size() ? col_constraints[cc].second : 1.0;
 
          ::Kokkos::atomic_add(&sys.getMatrixValue(constraint_row, constraint_col, tag),
                               weight_row * weight_col * 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);
  auto row_constraints = sys.getLocalConstraints(row);
 
  for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
  {
    auto tag = _matrix_tags[t];
 
    if (sys.isMatrixTagActive(tag))
    {
      for (unsigned int rc = 0; rc < ::Kokkos::max(1u, row_constraints.size()); ++rc)
      {
        auto constraint_row = row_constraints.size() ? row_constraints[rc].first : row;
        auto weight_row = row_constraints.size() ? row_constraints[rc].second : 1.0;
 
        if (sys.hasNodalBCMatrixTag(constraint_row, tag))
          continue;
 
        for (unsigned int j = 0; j < local_ke.size(); ++j)
        {
          auto col = local_ke.raw_index(j);
          auto col_constraints = sys.getLocalConstraints(col);
 
          for (unsigned int cc = 0; cc < ::Kokkos::max(1u, col_constraints.size()); ++cc)
          {
            auto constraint_col =
                sys.localToGlobalDofIndex(col_constraints.size() ? col_constraints[cc].first : col);
            auto weight_col = col_constraints.size() ? col_constraints[cc].second : 1.0;
 
            ::Kokkos::atomic_add(&sys.getMatrixValue(constraint_row, constraint_col, tag),
                                 weight_row * weight_col * local_ke.raw_at(j));
          }
        }
      }
    }
  }
}
 
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 = sys.localToGlobalDofIndex(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];
 
  const auto scaling_factor =
      kokkosSystem(_kokkos_var.sys()).getVariableScalingFactor(_kokkos_var.var());
 
  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] * scaling_factor, datum.elem().id, i, datum.comp());
  }
}
 
template <typename function>
KOKKOS_FUNCTION void
ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) const
{
  Real local_ke[MAX_CACHED_DOF];
 
  const auto scaling_factor =
      kokkosSystem(_kokkos_var.sys()).getVariableScalingFactor(_kokkos_var.var());
 
  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] * scaling_factor, datum.elem().id, i, j, datum.jvar(), datum.comp());
    }
  }
}
 
} // namespace Moose::Kokkos