Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://www.mooseframework.org
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #pragma once
      11             : 
      12             : #include "KokkosTypes.h"
      13             : #include "KokkosVector.h"
      14             : #include "KokkosMatrix.h"
      15             : #include "PerfGraphInterface.h"
      16             : #include "KokkosMesh.h"
      17             : 
      18             : #include "libmesh/communicator.h"
      19             : 
      20             : class MooseMesh;
      21             : class SystemBase;
      22             : 
      23             : namespace Moose::Kokkos
      24             : {
      25             : 
      26             : /**
      27             :  * The Kokkos base system class. Each system in MOOSE has a corresponding Kokkos
      28             :  * base system.
      29             :  */
      30             : class System : public PerfGraphInterface, public MeshHolder
      31             : {
      32             : public:
      33             :   /**
      34             :    * Constructor
      35             :    * @param system The associated MOOSE system
      36             :    */
      37             :   System(SystemBase & system);
      38             : 
      39             :   /**
      40             :    * Defaulted copy constructor
      41             :    * Used by FESystem in mixed FE+FV simulations.
      42             :    */
      43      137150 :   System(const System & src) = default;
      44             : 
      45             :   /**
      46             :    * CSR format sparsity data
      47             :    */
      48             :   struct Sparsity
      49             :   {
      50             :     Array<PetscInt> col_idx;
      51             :     Array<PetscInt> row_idx;
      52             :     Array<PetscInt> row_ptr;
      53             :   };
      54             : 
      55             : #ifdef MOOSE_KOKKOS_SCOPE
      56             :   /**
      57             :    * Synchronize the active tagged vectors and matrices between host and device
      58             :    * @param dir Copy direction
      59             :    */
      60             :   void sync(const MemcpyType dir);
      61             :   /**
      62             :    * Synchronize the specified tagged vectors between host and device
      63             :    * @param tags The vector tags
      64             :    * @param dir Copy direction
      65             :    */
      66             :   ///@{
      67             :   void sync(const std::set<TagID> & tags, const MemcpyType dir);
      68             :   void sync(const std::vector<TagID> & tags, const MemcpyType dir);
      69             :   void sync(const TagID tag, const MemcpyType dir);
      70             :   ///@}
      71             : 
      72             :   /**
      73             :    * Set the active variables
      74             :    * @param vars The active MOOSE variables
      75             :    */
      76             :   void setActiveVariables(const std::set<MooseVariableFieldBase *> & vars);
      77             : 
      78             :   /**
      79             :    * Set the active solution tags
      80             :    * @param tags The active solution tags
      81             :    */
      82             :   void setActiveSolutionTags(const std::set<TagID> & tags);
      83             : 
      84             :   /**
      85             :    * Set the active residual tags
      86             :    * @param tags The active residual tags
      87             :    */
      88             :   void setActiveResidualTags(const std::set<TagID> & tags);
      89             : 
      90             :   /**
      91             :    * Set the active matrix tags
      92             :    * @param vars The active matrix tags
      93             :    */
      94             :   void setActiveMatrixTags(const std::set<TagID> & tags);
      95             : 
      96             :   /**
      97             :    * Clear the cached active variables
      98             :    */
      99      137406 :   void clearActiveVariables() { _active_variables.destroy(); }
     100             : 
     101             :   /**
     102             :    * Clear the cached active solution tags
     103             :    */
     104      221874 :   void clearActiveSolutionTags() { _active_solution_tags.destroy(); }
     105             : 
     106             :   /**
     107             :    * Clear the cached active residual tags
     108             :    */
     109      132713 :   void clearActiveResidualTags()
     110             :   {
     111      132713 :     _active_residual_tags.destroy();
     112      132713 :     _residual_tag_active = false;
     113      132713 :   }
     114             : 
     115             :   /**
     116             :    * Clear the cached active matrix tags
     117             :    */
     118       16197 :   void clearActiveMatrixTags()
     119             :   {
     120       16197 :     _active_matrix_tags.destroy();
     121       16197 :     _matrix_tag_active = false;
     122       16197 :   }
     123             : 
     124             :   /**
     125             :    * Get the MOOSE system
     126             :    * @returns The MOOSE system
     127             :    */
     128             :   ///@{
     129             :   SystemBase & getSystem() { return _system; }
     130             :   const SystemBase & getSystem() const { return _system; }
     131             :   ///@}
     132             : 
     133             :   /**
     134             :    * Get the libMesh DOF map
     135             :    * @returns The libMesh DOF map
     136             :    */
     137             :   const libMesh::DofMap & getDofMap() const;
     138             : 
     139             :   /**
     140             :    * Get the libMesh communicator
     141             :    * @returns The libMesh communicator
     142             :    */
     143        4480 :   const Parallel::Communicator & getComm() const { return _comm; }
     144             : 
     145             :   /**
     146             :    * Get the list of local DOF indices to communicate
     147             :    * @returns The list of local DOF indices to communicate
     148             :    */
     149             :   const Array<Array<dof_id_type>> & getLocalCommList() const { return _local_comm_list; }
     150             : 
     151             :   /**
     152             :    * Get the list of ghost DOF indices to communicate
     153             :    * @returns The list of ghost DOF indices to communicate
     154             :    */
     155             :   const Array<Array<dof_id_type>> & getGhostCommList() const { return _ghost_comm_list; }
     156             : 
     157             :   /**
     158             :    * Get the sparisty pattern data
     159             :    * @returns The sparisty pattern data
     160             :    */
     161        1924 :   const Sparsity & getSparsity() const { return _sparsity; }
     162             : 
     163             :   /**
     164             :    * Check whether a variable is active on a subdomain
     165             :    * @param var The variable number
     166             :    * @param subdomain The contiguous subdomain ID
     167             :    * @returns Whether the variable is active
     168             :    */
     169      492790 :   KOKKOS_FUNCTION bool isVariableActive(unsigned int var, ContiguousSubdomainID subdomain) const
     170             :   {
     171      492790 :     return _var_subdomain_active(var, subdomain);
     172             :   }
     173             : 
     174             :   /**
     175             :    * Check whether a residual tag is active
     176             :    * @param tag The residual tag
     177             :    * @returns Whether the residual tag is active
     178             :    */
     179    18670026 :   KOKKOS_FUNCTION bool isResidualTagActive(TagID tag) const { return _residual_tag_active[tag]; }
     180             : 
     181             :   /**
     182             :    * Check whether a matrix tag is active
     183             :    * @param tag The matrix tag
     184             :    * @returns Whether the matrix tag is active
     185             :    */
     186    17615848 :   KOKKOS_FUNCTION bool isMatrixTagActive(TagID tag) const { return _matrix_tag_active[tag]; }
     187             : 
     188             :   /**
     189             :    * Get the number of local DOFs
     190             :    * @returns The number of local DOFs
     191             :    */
     192       15579 :   KOKKOS_FUNCTION dof_id_type getNumLocalDofs() const { return _num_local_dofs; }
     193             : 
     194             :   /**
     195             :    * Get the number of ghost DOFs
     196             :    * @returns The number of ghost DOFs
     197             :    */
     198      388812 :   KOKKOS_FUNCTION dof_id_type getNumGhostDofs() const { return _num_ghost_dofs; }
     199             : 
     200             :   /**
     201             :    * Get the local DOF index of a variable for an element
     202             :    * @param elem The contiguous element ID
     203             :    * @param i The element-local DOF index
     204             :    * @param var The variable number
     205             :    * @returns The local DOF index
     206             :    */
     207   350392564 :   KOKKOS_FUNCTION dof_id_type getElemLocalDofIndex(ContiguousElementID elem,
     208             :                                                    unsigned int i,
     209             :                                                    unsigned int var) const
     210             :   {
     211   350392564 :     return _local_elem_dof_index[var](i, elem);
     212             :   }
     213             : 
     214             :   /**
     215             :    * Get the global DOF index of a variable for an element
     216             :    * @param elem The contiguous element ID
     217             :    * @param i The element-local DOF index
     218             :    * @param var The variable number
     219             :    * @returns The global DOF index
     220             :    */
     221    10050013 :   KOKKOS_FUNCTION dof_id_type getElemGlobalDofIndex(ContiguousElementID elem,
     222             :                                                     unsigned int i,
     223             :                                                     unsigned int var) const
     224             :   {
     225    10050013 :     return _local_to_global_dof_index[_local_elem_dof_index[var](i, elem)];
     226             :   }
     227             : 
     228             :   /**
     229             :    * Get the global DOF index of a local DOF index
     230             :    * @param dof The local DOF index
     231             :    * @returns The global DOF index
     232             :    */
     233             :   KOKKOS_FUNCTION dof_id_type localToGlobalDofIndex(dof_id_type dof) const
     234             :   {
     235             :     return _local_to_global_dof_index[dof];
     236             :   }
     237             : 
     238             :   /**
     239             :    * Get a tagged Kokkos vector
     240             :    * @param tag The vector tag
     241             :    * @returns The Kokkos vector
     242             :    */
     243             :   KOKKOS_FUNCTION Vector & getVector(TagID tag) const { return _vectors[tag]; }
     244             : 
     245             :   /**
     246             :    * Get a tagged Kokkos matrix
     247             :    * @param tag The matrix tag
     248             :    * @returns The Kokkos matrix
     249             :    */
     250      426757 :   KOKKOS_FUNCTION Matrix & getMatrix(TagID tag) const { return _matrices[tag]; }
     251             : 
     252             :   /**
     253             :    * Get the DOF value of a tagged vector
     254             :    * @param dof The local DOF index
     255             :    * @param tag The vector tag
     256             :    * @returns The DOF value
     257             :    */
     258   329361089 :   KOKKOS_FUNCTION Real & getVectorDofValue(const dof_id_type dof, const TagID tag) const
     259             :   {
     260   329361089 :     return _vectors[tag][dof];
     261             :   }
     262             : 
     263             :   /**
     264             :    * Get an entry from a tagged matrix
     265             :    * @param row The local row index
     266             :    * @param col The global column index
     267             :    * @param tag The matrix tag
     268             :    * @returns The entry from the tagged matrix
     269             :    */
     270    13356526 :   KOKKOS_FUNCTION Real & getMatrixValue(dof_id_type row, dof_id_type col, TagID tag) const
     271             :   {
     272    13356526 :     return _matrices[tag](row, col);
     273             :   }
     274             : 
     275             : #endif
     276             : 
     277             : protected:
     278             :   /**
     279             :    * Reference of the MOOSE system
     280             :    */
     281             :   SystemBase & _system;
     282             : 
     283             :   /**
     284             :    * Reference of the MOOSE mesh
     285             :    */
     286             :   const MooseMesh & _mesh;
     287             : 
     288             :   /**
     289             :    * Reference of the libMesh DOF map
     290             :    */
     291             :   const libMesh::DofMap & _dof_map;
     292             : 
     293             :   /**
     294             :    * Reference of the libMesh communicator
     295             :    */
     296             :   const Parallel::Communicator & _comm;
     297             : 
     298             :   /**
     299             :    * Number of variables
     300             :    */
     301             :   const unsigned int _num_vars;
     302             : 
     303             :   /**
     304             :    * Number of local DOFs
     305             :    */
     306             :   const dof_id_type _num_local_dofs;
     307             : 
     308             :   /**
     309             :    * Number of ghost DOFs
     310             :    */
     311             :   const dof_id_type _num_ghost_dofs;
     312             : 
     313             :   /**
     314             :    * Kokkos vectors and matrices on device
     315             :    */
     316             :   ///@{
     317             :   Array<Vector> _vectors;
     318             :   Array<Matrix> _matrices;
     319             :   ///@}
     320             : 
     321             :   /**
     322             :    * Local element DOF indices of each variable
     323             :    */
     324             :   Array<Array2D<dof_id_type>> _local_elem_dof_index;
     325             : 
     326             :   /**
     327             :    * Map from local DOF index to global DOF index
     328             :    */
     329             :   Array<dof_id_type> _local_to_global_dof_index;
     330             : 
     331             :   /**
     332             :    * Maximum number of DOFs per element for each variable
     333             :    */
     334             :   Array<unsigned int> _max_dofs_per_elem;
     335             : 
     336             :   /**
     337             :    * Whether each variable is active on subdomains
     338             :    */
     339             :   Array2D<bool> _var_subdomain_active;
     340             : 
     341             :   /**
     342             :    * List of active variable numbers
     343             :    */
     344             :   Array<unsigned int> _active_variables;
     345             : 
     346             :   /**
     347             :    * List of active tags
     348             :    */
     349             :   ///@{
     350             :   Array<TagID> _active_solution_tags;
     351             :   Array<TagID> _active_residual_tags;
     352             :   Array<TagID> _active_matrix_tags;
     353             :   ///@}
     354             : 
     355             :   /**
     356             :    * Flag whether each tag is active
     357             :    */
     358             :   ///@{
     359             :   Array<bool> _residual_tag_active;
     360             :   Array<bool> _matrix_tag_active;
     361             :   ///@}
     362             : 
     363             :   /**
     364             :    * List of DOFs to send and receive
     365             :    */
     366             :   ///@{
     367             :   Array<Array<dof_id_type>> _local_comm_list;
     368             :   Array<Array<dof_id_type>> _ghost_comm_list;
     369             :   ///@}
     370             : 
     371             : private:
     372             :   /**
     373             :    * Setup variable data
     374             :    */
     375             :   void setupVariables();
     376             : 
     377             :   /**
     378             :    * Setup DOF data
     379             :    */
     380             :   void setupDofs();
     381             : 
     382             :   /**
     383             :    * Setup sparsity data
     384             :    */
     385             :   void setupSparsity();
     386             : 
     387             :   /**
     388             :    * Matrix sparsity pattern data
     389             :    */
     390             :   Sparsity _sparsity;
     391             : };
     392             : 
     393             : #ifdef MOOSE_KOKKOS_SCOPE
     394             : #define MakeSystemHolderMethods(SystemTypeName)                                                    \
     395             :   KOKKOS_FUNCTION const Array<SystemTypeName> & kokkosSystems() const                              \
     396             :   {                                                                                                \
     397             :     KOKKOS_IF_ON_HOST(return _systems_host;)                                                       \
     398             :     return _systems_device;                                                                        \
     399             :   }                                                                                                \
     400             :   Array<SystemTypeName> & kokkosSystems() { return _systems_host; }                                \
     401             :   KOKKOS_FUNCTION const SystemTypeName & kokkosSystem(unsigned int sys) const                      \
     402             :   {                                                                                                \
     403             :     KOKKOS_IF_ON_HOST(return _systems_host[sys];)                                                  \
     404             :     return _systems_device[sys];                                                                   \
     405             :   }                                                                                                \
     406             :   SystemTypeName & kokkosSystem(unsigned int sys) { return _systems_host[sys]; }
     407             : #else
     408             : #define MakeSystemHolderMethods(SystemTypeName)
     409             : #endif
     410             : 
     411             : /**
     412             :  * The Kokkos interface that holds the host reference of the Kokkos systems and copies it to device
     413             :  * during parallel dispatch.
     414             :  * Maintains synchronization between host and device Kokkos systems and provides access to the
     415             :  * appropriate Kokkos systems depending on the architecture.
     416             :  */
     417             : // clang-format off
     418             : #define MakeSystemHolder(SystemTypeName)                                                           \
     419             :   class SystemTypeName##Holder                                                                     \
     420             :   {                                                                                                \
     421             :   public:                                                                                          \
     422             :     SystemTypeName##Holder(Array<SystemTypeName> & systems)                                        \
     423             :       : _systems_host(systems), _systems_device(systems)                                           \
     424             :     {                                                                                              \
     425             :     }                                                                                              \
     426             :     SystemTypeName##Holder(const SystemTypeName##Holder & holder)                                  \
     427             :       : _systems_host(holder._systems_host), _systems_device(holder._systems_host)                 \
     428             :     {                                                                                              \
     429             :     }                                                                                              \
     430             :     MakeSystemHolderMethods(SystemTypeName)                                                        \
     431             :   private:                                                                                         \
     432             :     Array<SystemTypeName> & _systems_host;                                                         \
     433             :     const Array<SystemTypeName> _systems_device;                                                   \
     434             :   }
     435             : // clang-format on
     436             : 
     437      149217 : MakeSystemHolder(System);
     438             : } // namespace Moose::Kokkos