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_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; }

  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 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 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 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<dof_id_type> _local_to_global_dof_index;

  Array<unsigned int> _max_dofs_per_elem;

  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