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             : #include "KokkosSystem.h"
      11             : 
      12             : #include "MooseMesh.h"
      13             : #include "FEProblemBase.h"
      14             : #include "AuxiliarySystem.h"
      15             : 
      16             : #include "libmesh/system.h"
      17             : 
      18             : namespace Moose::Kokkos
      19             : {
      20        6112 : System::System(SystemBase & system)
      21        1735 :   : PerfGraphInterface(system.feProblem().getMooseApp().perfGraph(), "KokkosSystem"),
      22        1735 :     MeshHolder(*system.mesh().getKokkosMesh()),
      23        1735 :     _system(system),
      24        3470 :     _mesh(system.mesh()),
      25        1735 :     _dof_map(system.dofMap()),
      26        1735 :     _comm(system.dofMap().comm()),
      27        1735 :     _num_vars(system.system().n_vars()),
      28        1735 :     _num_local_dofs(system.system().n_local_dofs()),
      29       10410 :     _num_ghost_dofs(system.dofMap().get_send_list().size())
      30             : {
      31        3194 :   setupVariables();
      32        3194 :   setupDofs();
      33             : 
      34        3194 :   if (!dynamic_cast<AuxiliarySystem *>(&_system))
      35        2240 :     setupSparsity();
      36        3194 : }
      37             : 
      38             : void
      39        3194 : System::setupVariables()
      40             : {
      41        3194 :   auto & sys = _system.system();
      42             : 
      43        3194 :   _var_subdomain_active.create(_num_vars, kokkosMesh().getNumSubdomains());
      44             : 
      45        7436 :   for (unsigned int var = 0; var < _num_vars; ++var)
      46        9561 :     for (auto subdomain : _mesh.meshSubdomains())
      47        5319 :       _var_subdomain_active(var, kokkosMesh().getContiguousSubdomainID(subdomain)) =
      48        2876 :           sys.variable(var).active_on_subdomain(subdomain);
      49             : 
      50        3194 :   _var_subdomain_active.copyToDevice();
      51        3194 : }
      52             : 
      53             : void
      54        3194 : System::setupDofs()
      55             : {
      56        3194 :   auto & sys = _system.system();
      57             : 
      58        3194 :   const auto num_total_elems = kokkosMesh().getNumLocalAndPossiblyOneNeighborLayerGhostElements();
      59             : 
      60        3194 :   _residual_tag_active.create(MAX_TAG);
      61        3194 :   _residual_tag_active = false;
      62             : 
      63        3194 :   _matrix_tag_active.create(MAX_TAG);
      64        3194 :   _matrix_tag_active = false;
      65             : 
      66        3194 :   _vectors.create(MAX_TAG);
      67        3194 :   _matrices.create(MAX_TAG);
      68             : 
      69        3194 :   _local_elem_dof_index.create(_num_vars);
      70        3194 :   _local_to_global_dof_index.create(_num_local_dofs + _num_ghost_dofs);
      71             : 
      72        3194 :   _max_dofs_per_elem.create(_num_vars);
      73        3194 :   _max_dofs_per_elem = 0;
      74             : 
      75        3194 :   auto solution = dynamic_cast<PetscVector<Number> *>(sys.current_local_solution.get());
      76             : 
      77             : #ifdef MOOSE_ENABLE_KOKKOS_GPU
      78             :   // Kokkos array thinks OpenMP clause is device code when using OpenMP backend
      79        1459 : #pragma omp parallel for
      80             : #endif
      81        4026 :   for (unsigned int var = 0; var < _num_vars; ++var)
      82             :   {
      83        2291 :     std::vector<dof_id_type> dof_indices;
      84             : 
      85      819399 :     for (const auto elem : _mesh.getMesh().active_local_element_ptr_range())
      86             :     {
      87      817108 :       _dof_map.dof_indices(elem, dof_indices, var);
      88             : 
      89     1634216 :       _max_dofs_per_elem[var] =
      90      817108 :           std::max(_max_dofs_per_elem[var], static_cast<unsigned int>(dof_indices.size()));
      91        2291 :     }
      92             : 
      93        2291 :     _local_elem_dof_index[var].create(_max_dofs_per_elem[var], num_total_elems);
      94        2291 :     _local_elem_dof_index[var] = libMesh::DofObject::invalid_id;
      95             : 
      96      819399 :     for (const auto elem : _mesh.getMesh().active_local_element_ptr_range())
      97             :     {
      98      817108 :       const auto id = kokkosMesh().getContiguousElementID(elem);
      99             : 
     100      817108 :       _dof_map.dof_indices(elem, dof_indices, var);
     101             : 
     102     4031254 :       for (unsigned int dof = 0; dof < dof_indices.size(); ++dof)
     103             :       {
     104     3214146 :         _local_elem_dof_index[var](dof, id) = solution->map_global_to_local_index(dof_indices[dof]);
     105     3214146 :         _local_to_global_dof_index[_local_elem_dof_index[var](dof, id)] = dof_indices[dof];
     106             :       }
     107        2291 :     }
     108             : 
     109        2331 :     for (auto & [ghost_elem, ghost_id] : kokkosMesh().getGhostElemIdMapping())
     110             :     {
     111          40 :       if (!_dof_map.is_evaluable(*ghost_elem, var))
     112          20 :         continue;
     113          20 :       _dof_map.dof_indices(ghost_elem, dof_indices, var);
     114          44 :       for (const auto dof : index_range(dof_indices))
     115             :       {
     116          24 :         const auto local_idx = solution->map_global_to_local_index(dof_indices[dof]);
     117          24 :         _local_elem_dof_index[var](dof, ghost_id) = local_idx;
     118          24 :         _local_to_global_dof_index[local_idx] = dof_indices[dof];
     119             :       }
     120             :     }
     121        2291 :   }
     122             : 
     123        3194 :   _local_elem_dof_index.copyToDeviceNested();
     124        3194 :   _local_to_global_dof_index.copyToDevice();
     125             : 
     126        3194 :   _max_dofs_per_elem.copyToDevice();
     127             : 
     128             :   // Setup DOF communication maps
     129             : 
     130        3194 :   auto num_procs = _comm.size();
     131             : 
     132        6388 :   std::vector<std::vector<dof_id_type>> send_list(num_procs);
     133        3194 :   std::vector<std::vector<dof_id_type>> recv_list(num_procs);
     134             : 
     135       16828 :   for (auto dof : _dof_map.get_send_list())
     136       13634 :     recv_list[_dof_map.dof_owner(dof)].push_back(dof);
     137             : 
     138        7868 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     139        4674 :     _comm.scatter(recv_list, send_list[proc], proc);
     140             : 
     141        3194 :   _local_comm_list.create(num_procs);
     142        3194 :   _ghost_comm_list.create(num_procs);
     143             : 
     144        7868 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     145             :   {
     146        4674 :     _local_comm_list[proc].create(send_list[proc].size());
     147        4674 :     _ghost_comm_list[proc].create(recv_list[proc].size());
     148             : 
     149       18308 :     for (dof_id_type i = 0; i < send_list[proc].size(); ++i)
     150       13634 :       _local_comm_list[proc][i] = solution->map_global_to_local_index(send_list[proc][i]);
     151             : 
     152       18308 :     for (dof_id_type i = 0; i < recv_list[proc].size(); ++i)
     153       13634 :       _ghost_comm_list[proc][i] = solution->map_global_to_local_index(recv_list[proc][i]);
     154             :   }
     155             : 
     156        3194 :   _local_comm_list.copyToDeviceNested();
     157        3194 :   _ghost_comm_list.copyToDeviceNested();
     158        3194 : }
     159             : 
     160             : void
     161        2240 : System::setupSparsity()
     162             : {
     163        2240 :   auto * const pattern = _dof_map.get_sparsity_pattern();
     164        2240 :   bool have_pattern = pattern;
     165             : 
     166             :   // Make sure to check whether the sparsity pattern graph returned by
     167             :   // pattern->get_sparsity_pattern() is non-empty
     168        2240 :   getComm().verify(have_pattern);
     169        2240 :   if (!have_pattern)
     170         182 :     return;
     171        2240 :   have_pattern = !pattern->get_sparsity_pattern().empty();
     172        2240 :   getComm().max(have_pattern);
     173        2240 :   if (!have_pattern)
     174         182 :     return;
     175             : 
     176        2058 :   std::vector<PetscInt> col_idx;
     177        2058 :   std::vector<PetscInt> row_idx;
     178        4116 :   std::vector<PetscInt> row_ptr = {0};
     179             : 
     180      140694 :   for (dof_id_type r = _dof_map.first_dof(); r < _dof_map.end_dof(); ++r)
     181             :   {
     182      138636 :     auto & cols = pattern->get_sparsity_pattern().at(r - _dof_map.first_dof());
     183             : 
     184     1271588 :     for (auto col : cols)
     185             :     {
     186     1132952 :       col_idx.push_back(col);
     187     1132952 :       row_idx.push_back(r);
     188             :     }
     189             : 
     190      138636 :     row_ptr.push_back(col_idx.size());
     191             :   }
     192             : 
     193        2058 :   auto num_procs = _comm.size();
     194             : 
     195        8232 :   std::vector<std::vector<PetscInt>> send_cols(num_procs), send_num_cols(num_procs);
     196        6174 :   std::vector<std::vector<PetscInt>> recv_cols(num_procs), recv_num_cols(num_procs);
     197             : 
     198        5048 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     199        6986 :     for (auto r : _local_comm_list[proc])
     200             :     {
     201       45140 :       for (PetscInt c = row_ptr[r]; c < row_ptr[r + 1]; ++c)
     202       41144 :         send_cols[proc].push_back(col_idx[c]);
     203             : 
     204        3996 :       send_num_cols[proc].push_back(row_ptr[r + 1] - row_ptr[r]);
     205             :     }
     206             : 
     207        5048 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     208             :   {
     209        2990 :     _comm.scatter(send_cols, recv_cols[proc], proc);
     210        2990 :     _comm.scatter(send_num_cols, recv_num_cols[proc], proc);
     211             :   }
     212             : 
     213        6174 :   std::vector<PetscInt> col(num_procs), row(num_procs);
     214             : 
     215        6054 :   for (auto r : _dof_map.get_send_list())
     216             :   {
     217        3996 :     auto proc = _dof_map.dof_owner(r);
     218        3996 :     auto n_cols = recv_num_cols[proc][row[proc]++];
     219             : 
     220       45140 :     for (PetscInt c = 0; c < n_cols; ++c)
     221             :     {
     222       41144 :       col_idx.push_back(recv_cols[proc][col[proc]++]);
     223       41144 :       row_idx.push_back(r);
     224             :     }
     225             : 
     226        3996 :     row_ptr.push_back(col_idx.size());
     227             :   }
     228             : 
     229        2058 :   _sparsity.col_idx = col_idx;
     230        2058 :   _sparsity.row_idx = row_idx;
     231        2058 :   _sparsity.row_ptr = row_ptr;
     232        2058 : }
     233             : 
     234             : void
     235      443748 : System::sync(const MemcpyType dir)
     236             : {
     237      443748 :   if (dir == MemcpyType::HOST_TO_DEVICE)
     238             :   {
     239      755847 :     for (auto tag : _active_solution_tags)
     240             :     {
     241      533973 :       auto & vector = _system.getVector(tag);
     242             : 
     243      533973 :       _vectors[tag].create(vector, *this, false);
     244      533973 :       _vectors[tag].copyToDevice();
     245             :     }
     246             : 
     247      601435 :     for (auto tag : _active_residual_tags)
     248             :     {
     249      379561 :       auto & vector = _system.getVector(tag);
     250             : 
     251      379561 :       _vectors[tag].create(vector, *this, true);
     252      379561 :       _vectors[tag].copyToDevice();
     253             :     }
     254             : 
     255      238433 :     for (auto tag : _active_matrix_tags)
     256             :     {
     257       16559 :       auto & matrix = _system.getMatrix(tag);
     258             : 
     259             :       {
     260       49677 :         TIME_SECTION("MatSetPreallocation", 1);
     261       16559 :         _matrices[tag].create(matrix, *this);
     262       16559 :       }
     263             : 
     264       16559 :       _matrices[tag] = 0;
     265             :     }
     266             : 
     267      221874 :     _vectors.copyToDevice();
     268      221874 :     _matrices.copyToDevice();
     269             :   }
     270      221874 :   else if (dir == MemcpyType::DEVICE_TO_HOST)
     271             :   {
     272      755847 :     for (auto tag : _active_solution_tags)
     273      533973 :       _vectors[tag].restore();
     274             : 
     275      601435 :     for (auto tag : _active_residual_tags)
     276      379561 :       _vectors[tag].close();
     277             : 
     278      238433 :     for (auto tag : _active_matrix_tags)
     279             :     {
     280       49677 :       TIME_SECTION("MatSetValues", 1);
     281       16559 :       _matrices[tag].close();
     282       16559 :     }
     283             :   }
     284      443748 : }
     285             : 
     286             : void
     287       45520 : System::sync(const std::set<TagID> & tags, const MemcpyType dir)
     288             : {
     289       45520 :   if (dir == MemcpyType::HOST_TO_DEVICE)
     290             :   {
     291       45535 :     for (auto tag : tags)
     292             :     {
     293       22775 :       if (!_system.hasVector(tag))
     294        1516 :         continue;
     295             : 
     296       21259 :       auto & vector = _system.getVector(tag);
     297             : 
     298       21259 :       _vectors[tag].create(vector, *this, false);
     299       21259 :       _vectors[tag].copyToDevice();
     300             :     }
     301             : 
     302       22760 :     _vectors.copyToDevice();
     303             :   }
     304       22760 :   else if (dir == MemcpyType::DEVICE_TO_HOST)
     305             :   {
     306       45535 :     for (auto tag : tags)
     307             :     {
     308       22775 :       if (!_system.hasVector(tag))
     309        1516 :         continue;
     310             : 
     311       21259 :       _vectors[tag].copyToHost();
     312       21259 :       _vectors[tag].restore();
     313             :     }
     314             :   }
     315       45520 : }
     316             : 
     317             : void
     318           0 : System::sync(const std::vector<TagID> & tags, const MemcpyType dir)
     319             : {
     320           0 :   sync(std::set<TagID>(tags.begin(), tags.end()), dir);
     321           0 : }
     322             : 
     323             : void
     324       33492 : System::sync(const TagID tag, const MemcpyType dir)
     325             : {
     326       66984 :   sync(std::set<TagID>{tag}, dir);
     327       33492 : }
     328             : 
     329             : void
     330      137101 : System::setActiveVariables(const std::set<MooseVariableFieldBase *> & vars)
     331             : {
     332      137101 :   std::set<unsigned int> active_variables;
     333             : 
     334      347078 :   for (auto var : vars)
     335      209977 :     if (var->sys().number() == _system.number())
     336      286872 :       for (unsigned int i = 0; i < var->count(); ++i)
     337      143436 :         active_variables.insert(var->number() + i);
     338             : 
     339      137101 :   _active_variables = active_variables;
     340      137101 : }
     341             : 
     342             : void
     343      221874 : System::setActiveSolutionTags(const std::set<TagID> & tags)
     344             : {
     345      221874 :   std::set<TagID> active_solution_tags;
     346             : 
     347     1162837 :   for (auto tag : tags)
     348      940963 :     if (_system.hasVector(tag))
     349      533973 :       active_solution_tags.insert(tag);
     350             : 
     351      221874 :   _active_solution_tags = active_solution_tags;
     352      221874 : }
     353             : 
     354             : void
     355      132713 : System::setActiveResidualTags(const std::set<TagID> & tags)
     356             : {
     357      132713 :   std::set<TagID> active_residual_tags;
     358             : 
     359      512274 :   for (auto tag : tags)
     360      379561 :     if (_system.hasVector(tag))
     361             :     {
     362      379561 :       active_residual_tags.insert(tag);
     363      379561 :       _residual_tag_active[tag] = true;
     364             :     }
     365             : 
     366      132713 :   _active_residual_tags = active_residual_tags;
     367             : 
     368      132713 :   _residual_tag_active.copyToDevice();
     369      132713 : }
     370             : 
     371             : void
     372       16197 : System::setActiveMatrixTags(const std::set<TagID> & tags)
     373             : {
     374       16197 :   std::set<TagID> active_matrix_tags;
     375             : 
     376       48163 :   for (auto tag : tags)
     377       31966 :     if (_system.hasMatrix(tag))
     378             :     {
     379       16559 :       active_matrix_tags.insert(tag);
     380       16559 :       _matrix_tag_active[tag] = true;
     381             :     }
     382             : 
     383       16197 :   _active_matrix_tags = active_matrix_tags;
     384             : 
     385       16197 :   _matrix_tag_active.copyToDevice();
     386       16197 : }
     387             : 
     388             : const libMesh::DofMap &
     389      379561 : System::getDofMap() const
     390             : {
     391      379561 :   return _system.dofMap();
     392             : }
     393             : 
     394             : } // namespace Moose::Kokkos