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 "libmesh/communicator.h"

class MooseMesh;
class SystemBase;

namespace Moose
{
namespace Kokkos
{

class NodalBCBase;

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

#ifdef MOOSE_KOKKOS_SCOPE

  void sync(const MemcpyKind dir);
  void sync(const std::set<TagID> & tags, const MemcpyKind dir);
  void sync(const std::vector<TagID> & tags, const MemcpyKind dir);
  void sync(const TagID tag, const MemcpyKind 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 hasNodalBC(dof_id_type dof) const
  {
    return _nbc_dof.isAlloc() && _nbc_dof[dof];
  }
  KOKKOS_FUNCTION bool hasNodalBCResidualTag(dof_id_type dof, TagID tag) const
  {
    return _nbc_residual_tag_dof[tag].isAlloc() && _nbc_residual_tag_dof[tag][dof];
  }
  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](elem, i);
  }
  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](elem, i)];
  }
  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(dof_id_type dof, TagID tag) const
  {
    return _vectors[tag][dof];
  }
  KOKKOS_FUNCTION Real &
  getVectorQpValue(ElementInfo info, dof_id_type qp, unsigned int var, TagID tag) const
  {
    return _qp_solutions[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION Real3 &
  getVectorQpGrad(ElementInfo info, dof_id_type qp, unsigned int var, TagID tag) const
  {
    return _qp_solutions_grad[tag](info.subdomain, var)[qp];
  }
  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 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 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 checkNodalBCs();
  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<bool> _nbc_dof;
  Array<Array<bool>> _nbc_residual_tag_dof;
  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 Real
System::getVectorQpValueFace(
    ElementInfo info, unsigned int side, unsigned int qp, unsigned int var, 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 Real3
System::getVectorQpGradFace(ElementInfo info,
                            unsigned int side,
                            Real33 jacobian,
                            unsigned int qp,
                            unsigned int var,
                            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;
}
#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 Kokkos
} // namespace Moose