KokkosSystem.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 "KokkosTypes.h"
#include "KokkosVector.h"
#include "KokkosMatrix.h"
#include "KokkosAssembly.h"

#include "SystemBase.h"
#include "PerfGraphInterface.h"

#include "libmesh/communicator.h"

class MooseMesh;
class SystemBase;

namespace Moose::Kokkos
{

class NodalBCBase;

class System : public PerfGraphInterface, public MeshHolder, public AssemblyHolder
{
public:
  System(SystemBase & system);

#ifdef MOOSE_KOKKOS_SCOPE

  void sync(const MemcpyType dir);
  void sync(const std::set<TagID> & tags, const MemcpyType dir);
  void sync(const std::vector<TagID> & tags, const MemcpyType dir);
  void sync(const TagID tag, const MemcpyType dir);

  void reinit();
  void setActiveVariables(const std::set<MooseVariableFieldBase *> & vars);
  void setActiveSolutionTags(const std::set<TagID> & tags);
  void setActiveResidualTags(const std::set<TagID> & tags);
  void setActiveMatrixTags(const std::set<TagID> & tags);
  void clearActiveVariables() { _active_variables.destroy(); }
  void clearActiveSolutionTags() { _active_solution_tags.destroy(); }
  void clearActiveResidualTags()
  {
    _active_residual_tags.destroy();
    _residual_tag_active = false;
  }
  void clearActiveMatrixTags()
  {
    _active_matrix_tags.destroy();
    _matrix_tag_active = false;
  }
  auto & getSystem() { return _system; }
  const auto & getSystem() const { return _system; }

  const auto & getDofMap() const { return _system.dofMap(); }
  const auto & getComm() const { return _comm; }
  const auto & getLocalCommList() const { return _local_comm_list; }
  const auto & getGhostCommList() const { return _ghost_comm_list; }
  const auto & getSparsity() const { return _sparsity; }
  KOKKOS_FUNCTION const auto & getCoupling(unsigned int var) const { return _coupling[var]; }
  KOKKOS_FUNCTION bool isVariableActive(unsigned int var, ContiguousSubdomainID subdomain) const
  {
    return _var_subdomain_active(var, subdomain);
  }
  KOKKOS_FUNCTION bool isResidualTagActive(TagID tag) const { return _residual_tag_active[tag]; }
  KOKKOS_FUNCTION bool isMatrixTagActive(TagID tag) const { return _matrix_tag_active[tag]; }
  KOKKOS_FUNCTION bool hasNodalBCMatrixTag(dof_id_type dof, TagID tag) const
  {
    return _nbc_matrix_tag_dof[tag].isAlloc() && _nbc_matrix_tag_dof[tag][dof];
  }
  KOKKOS_FUNCTION unsigned int getFETypeID(unsigned int var) const { return _var_fe_types[var]; }
  KOKKOS_FUNCTION dof_id_type getNumLocalDofs() const { return _num_local_dofs; }
  KOKKOS_FUNCTION dof_id_type getNumGhostDofs() const { return _num_ghost_dofs; }
  KOKKOS_FUNCTION unsigned int getMaxDofsPerElem(unsigned int var) const
  {
    return _max_dofs_per_elem[var];
  }
  KOKKOS_FUNCTION dof_id_type getElemLocalDofIndex(ContiguousElementID elem,
                                                   unsigned int i,
                                                   unsigned int var) const
  {
    return _local_elem_dof_index[var](i, elem);
  }
  KOKKOS_FUNCTION dof_id_type getNodeLocalDofIndex(ContiguousNodeID node,
                                                   unsigned int i,
                                                   unsigned int var) const
  {
    return _local_node_dof_index[var][node] + i;
  }
  KOKKOS_FUNCTION dof_id_type getElemGlobalDofIndex(ContiguousElementID elem,
                                                    unsigned int i,
                                                    unsigned int var) const
  {
    return _local_to_global_dof_index[_local_elem_dof_index[var](i, elem)];
  }
  KOKKOS_FUNCTION dof_id_type getNodeGlobalDofIndex(ContiguousNodeID node, unsigned int var) const
  {
    return _local_to_global_dof_index[_local_node_dof_index[var][node]];
  }
  KOKKOS_FUNCTION dof_id_type localToGlobalDofIndex(dof_id_type dof) const
  {
    return _local_to_global_dof_index[dof];
  }
  KOKKOS_FUNCTION bool isNodalDefined(ContiguousNodeID node, unsigned int var) const
  {
    return _local_node_dof_index[var][node] != libMesh::DofObject::invalid_id;
  }
  KOKKOS_FUNCTION auto & getVector(TagID tag) const { return _vectors[tag]; }
  KOKKOS_FUNCTION auto & getMatrix(TagID tag) const { return _matrices[tag]; }
  KOKKOS_FUNCTION Real & getVectorDofValue(const dof_id_type dof, const TagID tag) const
  {
    return _vectors[tag][dof];
  }
  KOKKOS_FUNCTION ADReal getVectorDofADValue(const dof_id_type dof,
                                             const TagID tag,
                                             const Real seed) const;
  KOKKOS_FUNCTION Real & getVectorQpValue(const ElementInfo info,
                                          const dof_id_type qp,
                                          const unsigned int var,
                                          const TagID tag) const
  {
    return _qp_solutions[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION ADReal getVectorQpADValue(const ElementInfo info,
                                            const dof_id_type offset,
                                            const dof_id_type qp,
                                            const unsigned int var,
                                            const TagID tag,
                                            const Real seed) const;
  KOKKOS_FUNCTION Real3 & getVectorQpGrad(const ElementInfo info,
                                          const dof_id_type qp,
                                          const unsigned int var,
                                          const TagID tag) const
  {
    return _qp_solutions_grad[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION ADReal3 getVectorQpADGrad(const ElementInfo info,
                                            const Real33 jacobian,
                                            const dof_id_type offset,
                                            const dof_id_type qp,
                                            const unsigned int var,
                                            const TagID tag,
                                            const Real seed) const;
  KOKKOS_FUNCTION Real getVectorQpValueFace(const ElementInfo info,
                                            const unsigned int side,
                                            const unsigned int qp,
                                            const unsigned int var,
                                            const TagID tag) const;
  KOKKOS_FUNCTION ADReal getVectorQpADValueFace(const ElementInfo info,
                                                const unsigned int side,
                                                const unsigned int qp,
                                                const unsigned int var,
                                                const TagID tag,
                                                const Real seed) const;
  KOKKOS_FUNCTION Real3 getVectorQpGradFace(const ElementInfo info,
                                            const unsigned int side,
                                            const Real33 jacobian,
                                            const unsigned int qp,
                                            const unsigned int var,
                                            const TagID tag) const;
  KOKKOS_FUNCTION ADReal3 getVectorQpADGradFace(const ElementInfo info,
                                                const unsigned int side,
                                                const Real33 jacobian,
                                                const unsigned int qp,
                                                const unsigned int var,
                                                const TagID tag,
                                                const Real seed) const;
  KOKKOS_FUNCTION Real & getMatrixValue(dof_id_type row, dof_id_type col, TagID tag) const
  {
    return _matrices[tag](row, col);
  }

  KOKKOS_FUNCTION void operator()(const ThreadID tid) const;
#endif

  struct Sparsity
  {
    Array<PetscInt> col_idx;
    Array<PetscInt> row_idx;
    Array<PetscInt> row_ptr;
  };

private:
  void setupVariables();
  void setupDofs();
  void setupSparsity();
  void setupNodalBCDofs();
  void getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs);

  Thread<> _thread;
  SystemBase & _system;
  const MooseMesh & _mesh;
  const libMesh::DofMap & _dof_map;
  const Parallel::Communicator & _comm;
  const unsigned int _num_vars;
  const dof_id_type _num_local_dofs;
  const dof_id_type _num_ghost_dofs;

  Array<Vector> _vectors;
  Array<Matrix> _matrices;

  Array<Array2D<Array<Real>>> _qp_solutions;
  Array<Array2D<Array<Real3>>> _qp_solutions_grad;

  Array<Array2D<dof_id_type>> _local_elem_dof_index;
  Array<Array<dof_id_type>> _local_node_dof_index;

  Array<dof_id_type> _local_to_global_dof_index;
  Array<unsigned int> _max_dofs_per_elem;
  Array<unsigned int> _var_fe_types;
  Array2D<bool> _var_subdomain_active;
  Array<Array<unsigned int>> _coupling;

  Array<unsigned int> _active_variables;
  Array<TagID> _active_solution_tags;
  Array<TagID> _active_residual_tags;
  Array<TagID> _active_matrix_tags;

  Array<bool> _residual_tag_active;
  Array<bool> _matrix_tag_active;

  Array<Array<bool>> _nbc_matrix_tag_dof;
  Array<Array<dof_id_type>> _local_comm_list;
  Array<Array<dof_id_type>> _ghost_comm_list;

  Sparsity _sparsity;
};

#ifdef MOOSE_KOKKOS_SCOPE
KOKKOS_FUNCTION inline ADReal
System::getVectorDofADValue(const dof_id_type dof, TagID tag, const Real seed) const
{
  ADReal value = _vectors[tag][dof];

  if (seed != 0)
    value.derivatives().insert(_local_to_global_dof_index[dof]) = seed;

  return value;
}

KOKKOS_FUNCTION inline ADReal
System::getVectorQpADValue(const ElementInfo info,
                           const dof_id_type offset,
                           const dof_id_type qp,
                           const unsigned int var,
                           const TagID tag,
                           const Real seed) const
{
  ADReal value = 0;

  if (seed == 0)
    value = getVectorQpValue(info, offset + qp, var, tag);
  else
  {
    auto fe = _var_fe_types[var];
    auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
    auto & phi = kokkosAssembly().getPhi(info.subdomain, info.type, fe);

    for (unsigned int i = 0; i < n_dofs; ++i)
      value += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * phi(i, qp);
  }

  return value;
}

KOKKOS_FUNCTION inline ADReal3
System::getVectorQpADGrad(const ElementInfo info,
                          const Real33 jacobian,
                          const dof_id_type offset,
                          const dof_id_type qp,
                          const unsigned int var,
                          const TagID tag,
                          const Real seed) const
{
  ADReal3 grad;

  if (seed == 0)
    grad = getVectorQpGrad(info, offset + qp, var, tag);
  else
  {
    auto fe = _var_fe_types[var];
    auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
    auto & grad_phi = kokkosAssembly().getGradPhi(info.subdomain, info.type, fe);

    for (unsigned int i = 0; i < n_dofs; ++i)
      grad += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) *
              (jacobian * grad_phi(i, qp));
  }

  return grad;
}

KOKKOS_FUNCTION inline Real
System::getVectorQpValueFace(const ElementInfo info,
                             const unsigned int side,
                             const unsigned int qp,
                             const unsigned int var,
                             const TagID tag) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);

  Real value = 0;

  for (unsigned int i = 0; i < n_dofs; ++i)
    value += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * phi(i, qp);

  return value;
}

KOKKOS_FUNCTION inline ADReal
System::getVectorQpADValueFace(const ElementInfo info,
                               const unsigned int side,
                               const unsigned int qp,
                               const unsigned int var,
                               const TagID tag,
                               const Real seed) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);

  ADReal value = 0;

  for (unsigned int i = 0; i < n_dofs; ++i)
    value += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * phi(i, qp);

  return value;
}

KOKKOS_FUNCTION inline Real3
System::getVectorQpGradFace(const ElementInfo info,
                            const unsigned int side,
                            const Real33 jacobian,
                            const unsigned int qp,
                            const unsigned int var,
                            const TagID tag) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);

  Real3 grad = 0;

  for (unsigned int i = 0; i < n_dofs; ++i)
    grad += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) *
            (jacobian * grad_phi(i, qp));

  return grad;
}

KOKKOS_FUNCTION inline ADReal3
System::getVectorQpADGradFace(const ElementInfo info,
                              const unsigned int side,
                              const Real33 jacobian,
                              const unsigned int qp,
                              const unsigned int var,
                              const TagID tag,
                              const Real seed) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);

  ADReal3 grad = ADReal(0);

  for (unsigned int i = 0; i < n_dofs; ++i)
    grad += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) *
            (jacobian * grad_phi(i, qp));

  return grad;
}
#endif

class SystemHolder
{
public:
  SystemHolder(Array<System> & systems) : _systems_host(systems), _systems_device(systems) {}
  SystemHolder(const SystemHolder & holder)
    : _systems_host(holder._systems_host), _systems_device(holder._systems_host)
  {
  }

#ifdef MOOSE_KOKKOS_SCOPE

  KOKKOS_FUNCTION const Array<System> & kokkosSystems() const
  {
    KOKKOS_IF_ON_HOST(return _systems_host;)

    return _systems_device;
  }
  Array<System> & kokkosSystems() { return _systems_host; }
  KOKKOS_FUNCTION const System & kokkosSystem(unsigned int sys) const
  {
    KOKKOS_IF_ON_HOST(return _systems_host[sys];)

    return _systems_device[sys];
  }
  System & kokkosSystem(unsigned int sys) { return _systems_host[sys]; }
#endif

private:
  Array<System> & _systems_host;
  const Array<System> _systems_device;
};

} // namespace Moose::Kokkos