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 "PerfGraphInterface.h"
#include "KokkosMesh.h"
 
#include "libmesh/communicator.h"
 
class MooseMesh;
class SystemBase;
 
namespace Moose::Kokkos
{
 
class System : public PerfGraphInterface, public MeshHolder
{
public:
  System(SystemBase & system);
 
  System(const System & src) = default;
 
  struct Sparsity
  {
    Array<PetscInt> col_idx;
    Array<PetscInt> row_ptr;
  };
 
#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 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;
  }
 
  SystemBase & getSystem() { return _system; }
  const SystemBase & getSystem() const { return _system; }
 
  const libMesh::DofMap & getDofMap() const { return _dof_map; }
 
  const Parallel::Communicator & getComm() const { return _comm; }
 
  const Array<Array<dof_id_type>> & getLocalCommList() const { return _local_comm_list; }
 
  const Array<Array<dof_id_type>> & getGhostCommList() const { return _ghost_comm_list; }
 
  const Sparsity & getSparsity() const { return _sparsity; }
 
  void buildCooIndices(std::vector<PetscInt> & coo_i, std::vector<PetscInt> & coo_j) const;
 
  KOKKOS_FUNCTION bool isScalarVariable(unsigned int var) const { return _var_is_scalar[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 dof_id_type getNumLocalDofs() const { return _num_local_dofs; }
 
  KOKKOS_FUNCTION dof_id_type getNumGhostDofs() const { return _num_ghost_dofs; }
 
  KOKKOS_FUNCTION unsigned int getNumScalarDofs(unsigned int var) const
  {
    return _scalar_dof_index[var].size();
  }
 
  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 getScalarLocalDofIndex(unsigned int i, unsigned int var) const
  {
    return _scalar_dof_index[var][i];
  }
 
  KOKKOS_FUNCTION dof_id_type getElemGlobalDofIndex(ContiguousElementID elem,
                                                    unsigned int i,
                                                    unsigned int var) const
  {
    return _local_to_global_dof_index[getElemLocalDofIndex(elem, i, var)];
  }
 
  KOKKOS_FUNCTION dof_id_type getScalarGlobalDofIndex(unsigned int i, unsigned int var) const
  {
    return _local_to_global_dof_index[getScalarLocalDofIndex(i, var)];
  }
 
  KOKKOS_FUNCTION dof_id_type localToGlobalDofIndex(dof_id_type dof) const
  {
    return _local_to_global_dof_index[dof];
  }
 
  KOKKOS_FUNCTION Vector & getVector(TagID tag) const { return _vectors[tag]; }
 
  KOKKOS_FUNCTION Matrix & 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 Real & getMatrixValue(dof_id_type row, dof_id_type col, TagID tag) const
  {
    return _matrices[tag](row, col);
  }
 
#endif
 
protected:
  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<dof_id_type>> _local_elem_dof_index;
 
  Array<Array<dof_id_type>> _scalar_dof_index;
 
  Array<dof_id_type> _local_to_global_dof_index;
 
  Array<unsigned int> _max_dofs_per_elem;
 
  Array<bool> _var_is_scalar;
 
  Array2D<bool> _var_subdomain_active;
 
  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<dof_id_type>> _local_comm_list;
  Array<Array<dof_id_type>> _ghost_comm_list;
 
private:
  void setupVariables();
 
  void setupDofs();
 
  void setupSparsity();
 
  Sparsity _sparsity;
};
 
#ifdef MOOSE_KOKKOS_SCOPE
#define MakeSystemHolderMethods(SystemTypeName)                                                    \
  KOKKOS_FUNCTION const Array<SystemTypeName> & kokkosSystems() const                              \
  {                                                                                                \
    KOKKOS_IF_ON_HOST(return _systems_host;)                                                       \
    return _systems_device;                                                                        \
  }                                                                                                \
  Array<SystemTypeName> & kokkosSystems() { return _systems_host; }                                \
  KOKKOS_FUNCTION const SystemTypeName & kokkosSystem(unsigned int sys) const                      \
  {                                                                                                \
    KOKKOS_IF_ON_HOST(return _systems_host[sys];)                                                  \
    return _systems_device[sys];                                                                   \
  }                                                                                                \
  SystemTypeName & kokkosSystem(unsigned int sys) { return _systems_host[sys]; }
#else
#define MakeSystemHolderMethods(SystemTypeName)
#endif
 
// clang-format off
#define MakeSystemHolder(SystemTypeName)                                                           \
  class SystemTypeName##Holder                                                                     \
  {                                                                                                \
  public:                                                                                          \
    SystemTypeName##Holder(Array<SystemTypeName> & systems)                                        \
      : _systems_host(systems), _systems_device(systems)                                           \
    {                                                                                              \
    }                                                                                              \
    SystemTypeName##Holder(const SystemTypeName##Holder & holder)                                  \
      : _systems_host(holder._systems_host), _systems_device(holder._systems_host)                 \
    {                                                                                              \
    }                                                                                              \
    MakeSystemHolderMethods(SystemTypeName)                                                        \
  private:                                                                                         \
    Array<SystemTypeName> & _systems_host;                                                         \
    const Array<SystemTypeName> _systems_device;                                                   \
  }
// clang-format on
 
MakeSystemHolder(System);
} // namespace Moose::Kokkos