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        6430 : System::System(SystemBase & system)
      21        1819 :   : PerfGraphInterface(system.feProblem().getMooseApp().perfGraph(), "KokkosSystem"),
      22        1819 :     MeshHolder(*system.mesh().getKokkosMesh()),
      23        1819 :     _system(system),
      24        3638 :     _mesh(system.mesh()),
      25        1819 :     _dof_map(system.dofMap()),
      26        1819 :     _comm(system.dofMap().comm()),
      27        1819 :     _num_vars(system.system().n_vars()),
      28        1819 :     _num_local_dofs(system.system().n_local_dofs()),
      29       10914 :     _num_ghost_dofs(system.dofMap().get_send_list().size())
      30             : {
      31        3356 :   setupVariables();
      32        3356 :   setupDofs();
      33             : 
      34        3356 :   if (!dynamic_cast<AuxiliarySystem *>(&_system))
      35        2402 :     setupSparsity();
      36        3356 : }
      37             : 
      38             : void
      39        3356 : System::setupVariables()
      40             : {
      41        3356 :   auto & sys = _system.system();
      42             : 
      43        3356 :   _var_subdomain_active.create(_num_vars, kokkosMesh().getNumSubdomains());
      44             : 
      45        7857 :   for (unsigned int var = 0; var < _num_vars; ++var)
      46       10108 :     for (auto subdomain : _mesh.meshSubdomains())
      47        5607 :       _var_subdomain_active(var, kokkosMesh().getContiguousSubdomainID(subdomain)) =
      48        3026 :           sys.variable(var).active_on_subdomain(subdomain);
      49             : 
      50        3356 :   _var_subdomain_active.copyToDevice();
      51        3356 : }
      52             : 
      53             : void
      54        3356 : System::setupDofs()
      55             : {
      56        3356 :   auto & sys = _system.system();
      57             : 
      58        3356 :   const auto num_total_elems = kokkosMesh().getNumLocalAndPossiblyOneNeighborLayerGhostElements();
      59             : 
      60        3356 :   _residual_tag_active.create(MAX_TAG);
      61        3356 :   _residual_tag_active = false;
      62             : 
      63        3356 :   _matrix_tag_active.create(MAX_TAG);
      64        3356 :   _matrix_tag_active = false;
      65             : 
      66        3356 :   _vectors.create(MAX_TAG);
      67        3356 :   _matrices.create(MAX_TAG);
      68             : 
      69        3356 :   _local_elem_dof_index.create(_num_vars);
      70        3356 :   _local_to_global_dof_index.create(_num_local_dofs + _num_ghost_dofs);
      71             : 
      72        3356 :   _max_dofs_per_elem.create(_num_vars);
      73        3356 :   _max_dofs_per_elem = 0;
      74             : 
      75        3356 :   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        1537 : #pragma omp parallel for
      80             : #endif
      81        4245 :   for (unsigned int var = 0; var < _num_vars; ++var)
      82             :   {
      83        2426 :     std::vector<dof_id_type> dof_indices;
      84             : 
      85      832234 :     for (const auto elem : _mesh.getMesh().active_local_element_ptr_range())
      86             :     {
      87      829808 :       _dof_map.dof_indices(elem, dof_indices, var);
      88             : 
      89     1659616 :       _max_dofs_per_elem[var] =
      90      829808 :           std::max(_max_dofs_per_elem[var], static_cast<unsigned int>(dof_indices.size()));
      91        2426 :     }
      92             : 
      93        2426 :     _local_elem_dof_index[var].create(_max_dofs_per_elem[var], num_total_elems);
      94        2426 :     _local_elem_dof_index[var] = libMesh::DofObject::invalid_id;
      95             : 
      96      832234 :     for (const auto elem : _mesh.getMesh().active_local_element_ptr_range())
      97             :     {
      98      829808 :       const auto id = kokkosMesh().getContiguousElementID(elem);
      99             : 
     100      829808 :       _dof_map.dof_indices(elem, dof_indices, var);
     101             : 
     102     4142159 :       for (unsigned int dof = 0; dof < dof_indices.size(); ++dof)
     103             :       {
     104     3312351 :         _local_elem_dof_index[var](dof, id) = solution->map_global_to_local_index(dof_indices[dof]);
     105     3312351 :         _local_to_global_dof_index[_local_elem_dof_index[var](dof, id)] = dof_indices[dof];
     106             :       }
     107        2426 :     }
     108             : 
     109        2466 :     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        2426 :   }
     122             : 
     123        3356 :   _local_elem_dof_index.copyToDeviceNested();
     124        3356 :   _local_to_global_dof_index.copyToDevice();
     125             : 
     126        3356 :   _max_dofs_per_elem.copyToDevice();
     127             : 
     128             :   // Setup DOF communication maps
     129             : 
     130        3356 :   auto num_procs = _comm.size();
     131             : 
     132        6712 :   std::vector<std::vector<dof_id_type>> send_list(num_procs);
     133        3356 :   std::vector<std::vector<dof_id_type>> recv_list(num_procs);
     134             : 
     135       18300 :   for (auto dof : _dof_map.get_send_list())
     136       14944 :     recv_list[_dof_map.dof_owner(dof)].push_back(dof);
     137             : 
     138        8280 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     139        4924 :     _comm.scatter(recv_list, send_list[proc], proc);
     140             : 
     141        3356 :   _local_comm_list.create(num_procs);
     142        3356 :   _ghost_comm_list.create(num_procs);
     143             : 
     144        8280 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     145             :   {
     146        4924 :     _local_comm_list[proc].create(send_list[proc].size());
     147        4924 :     _ghost_comm_list[proc].create(recv_list[proc].size());
     148             : 
     149       19868 :     for (dof_id_type i = 0; i < send_list[proc].size(); ++i)
     150       14944 :       _local_comm_list[proc][i] = solution->map_global_to_local_index(send_list[proc][i]);
     151             : 
     152       19868 :     for (dof_id_type i = 0; i < recv_list[proc].size(); ++i)
     153       14944 :       _ghost_comm_list[proc][i] = solution->map_global_to_local_index(recv_list[proc][i]);
     154             :   }
     155             : 
     156        3356 :   _local_comm_list.copyToDeviceNested();
     157        3356 :   _ghost_comm_list.copyToDeviceNested();
     158        3356 : }
     159             : 
     160             : void
     161        2402 : System::setupSparsity()
     162             : {
     163        2402 :   auto * const pattern = _dof_map.get_sparsity_pattern();
     164        2402 :   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        2402 :   getComm().verify(have_pattern);
     169        2402 :   if (!have_pattern)
     170         182 :     return;
     171        2402 :   have_pattern = !pattern->get_sparsity_pattern().empty();
     172        2402 :   getComm().max(have_pattern);
     173        2402 :   if (!have_pattern)
     174         182 :     return;
     175             : 
     176        2220 :   std::vector<PetscInt> col_idx;
     177        2220 :   std::vector<PetscInt> row_idx;
     178        4440 :   std::vector<PetscInt> row_ptr = {0};
     179             : 
     180      203968 :   for (dof_id_type r = _dof_map.first_dof(); r < _dof_map.end_dof(); ++r)
     181             :   {
     182      201748 :     auto & cols = pattern->get_sparsity_pattern().at(r - _dof_map.first_dof());
     183             : 
     184     3622608 :     for (auto col : cols)
     185             :     {
     186     3420860 :       col_idx.push_back(col);
     187     3420860 :       row_idx.push_back(r);
     188             :     }
     189             : 
     190      201748 :     row_ptr.push_back(col_idx.size());
     191             :   }
     192             : 
     193        2220 :   auto num_procs = _comm.size();
     194             : 
     195        8880 :   std::vector<std::vector<PetscInt>> send_cols(num_procs), send_num_cols(num_procs);
     196        6660 :   std::vector<std::vector<PetscInt>> recv_cols(num_procs), recv_num_cols(num_procs);
     197             : 
     198        5460 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     199        8546 :     for (auto r : _local_comm_list[proc])
     200             :     {
     201      103364 :       for (PetscInt c = row_ptr[r]; c < row_ptr[r + 1]; ++c)
     202       98058 :         send_cols[proc].push_back(col_idx[c]);
     203             : 
     204        5306 :       send_num_cols[proc].push_back(row_ptr[r + 1] - row_ptr[r]);
     205             :     }
     206             : 
     207        5460 :   for (processor_id_type proc = 0; proc < num_procs; proc++)
     208             :   {
     209        3240 :     _comm.scatter(send_cols, recv_cols[proc], proc);
     210        3240 :     _comm.scatter(send_num_cols, recv_num_cols[proc], proc);
     211             :   }
     212             : 
     213        6660 :   std::vector<PetscInt> col(num_procs), row(num_procs);
     214             : 
     215        7526 :   for (auto r : _dof_map.get_send_list())
     216             :   {
     217        5306 :     auto proc = _dof_map.dof_owner(r);
     218        5306 :     auto n_cols = recv_num_cols[proc][row[proc]++];
     219             : 
     220      103364 :     for (PetscInt c = 0; c < n_cols; ++c)
     221             :     {
     222       98058 :       col_idx.push_back(recv_cols[proc][col[proc]++]);
     223       98058 :       row_idx.push_back(r);
     224             :     }
     225             : 
     226        5306 :     row_ptr.push_back(col_idx.size());
     227             :   }
     228             : 
     229        2220 :   _sparsity.col_idx = col_idx;
     230        2220 :   _sparsity.row_idx = row_idx;
     231        2220 :   _sparsity.row_ptr = row_ptr;
     232        2220 : }
     233             : 
     234             : void
     235      462538 : System::sync(const MemcpyType dir)
     236             : {
     237      462538 :   if (dir == MemcpyType::HOST_TO_DEVICE)
     238             :   {
     239      786576 :     for (auto tag : _active_solution_tags)
     240             :     {
     241      555307 :       auto & vector = _system.getVector(tag);
     242             : 
     243      555307 :       _vectors[tag].create(vector, *this, false);
     244      555307 :       _vectors[tag].copyToDevice();
     245             :     }
     246             : 
     247      632699 :     for (auto tag : _active_residual_tags)
     248             :     {
     249      401430 :       auto & vector = _system.getVector(tag);
     250             : 
     251      401430 :       _vectors[tag].create(vector, *this, true);
     252      401430 :       _vectors[tag].copyToDevice();
     253             :     }
     254             : 
     255      248778 :     for (auto tag : _active_matrix_tags)
     256             :     {
     257       17509 :       auto & matrix = _system.getMatrix(tag);
     258             : 
     259             :       {
     260       52527 :         TIME_SECTION("MatSetPreallocation", 1);
     261       17509 :         _matrices[tag].create(matrix, *this);
     262       17509 :       }
     263             : 
     264       17509 :       _matrices[tag] = 0;
     265             :     }
     266             : 
     267      231269 :     _vectors.copyToDevice();
     268      231269 :     _matrices.copyToDevice();
     269             :   }
     270      231269 :   else if (dir == MemcpyType::DEVICE_TO_HOST)
     271             :   {
     272      786576 :     for (auto tag : _active_solution_tags)
     273      555307 :       _vectors[tag].restore();
     274             : 
     275      632699 :     for (auto tag : _active_residual_tags)
     276      401430 :       _vectors[tag].close();
     277             : 
     278      248778 :     for (auto tag : _active_matrix_tags)
     279             :     {
     280       52527 :       TIME_SECTION("MatSetValues", 1);
     281       17509 :       _matrices[tag].close();
     282       17509 :     }
     283             :   }
     284      462538 : }
     285             : 
     286             : void
     287       45714 : System::sync(const std::set<TagID> & tags, const MemcpyType dir)
     288             : {
     289       45714 :   if (dir == MemcpyType::HOST_TO_DEVICE)
     290             :   {
     291       45729 :     for (auto tag : tags)
     292             :     {
     293       22872 :       if (!_system.hasVector(tag))
     294        1516 :         continue;
     295             : 
     296       21356 :       auto & vector = _system.getVector(tag);
     297             : 
     298       21356 :       _vectors[tag].create(vector, *this, false);
     299       21356 :       _vectors[tag].copyToDevice();
     300             :     }
     301             : 
     302       22857 :     _vectors.copyToDevice();
     303             :   }
     304       22857 :   else if (dir == MemcpyType::DEVICE_TO_HOST)
     305             :   {
     306       45729 :     for (auto tag : tags)
     307             :     {
     308       22872 :       if (!_system.hasVector(tag))
     309        1516 :         continue;
     310             : 
     311       21356 :       _vectors[tag].copyToHost();
     312       21356 :       _vectors[tag].restore();
     313             :     }
     314             :   }
     315       45714 : }
     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       33486 : System::sync(const TagID tag, const MemcpyType dir)
     325             : {
     326       66972 :   sync(std::set<TagID>{tag}, dir);
     327       33486 : }
     328             : 
     329             : void
     330      144746 : System::setActiveVariables(const std::set<MooseVariableFieldBase *> & vars)
     331             : {
     332      144746 :   std::set<unsigned int> active_variables;
     333             : 
     334      363633 :   for (auto var : vars)
     335      218887 :     if (var->sys().number() == _system.number())
     336      304664 :       for (unsigned int i = 0; i < var->count(); ++i)
     337      152332 :         active_variables.insert(var->number() + i);
     338             : 
     339      144746 :   _active_variables = active_variables;
     340      144746 : }
     341             : 
     342             : void
     343      231269 : System::setActiveSolutionTags(const std::set<TagID> & tags)
     344             : {
     345      231269 :   std::set<TagID> active_solution_tags;
     346             : 
     347     1200564 :   for (auto tag : tags)
     348      969295 :     if (_system.hasVector(tag))
     349      555307 :       active_solution_tags.insert(tag);
     350             : 
     351      231269 :   _active_solution_tags = active_solution_tags;
     352      231269 : }
     353             : 
     354             : void
     355      141134 : System::setActiveResidualTags(const std::set<TagID> & tags)
     356             : {
     357      141134 :   std::set<TagID> active_residual_tags;
     358             : 
     359      542564 :   for (auto tag : tags)
     360      401430 :     if (_system.hasVector(tag))
     361             :     {
     362      401430 :       active_residual_tags.insert(tag);
     363      401430 :       _residual_tag_active[tag] = true;
     364             :     }
     365             : 
     366      141134 :   _active_residual_tags = active_residual_tags;
     367             : 
     368      141134 :   _residual_tag_active.copyToDevice();
     369      141134 : }
     370             : 
     371             : void
     372       17147 : System::setActiveMatrixTags(const std::set<TagID> & tags)
     373             : {
     374       17147 :   std::set<TagID> active_matrix_tags;
     375             : 
     376       51013 :   for (auto tag : tags)
     377       33866 :     if (_system.hasMatrix(tag))
     378             :     {
     379       17509 :       active_matrix_tags.insert(tag);
     380       17509 :       _matrix_tag_active[tag] = true;
     381             :     }
     382             : 
     383       17147 :   _active_matrix_tags = active_matrix_tags;
     384             : 
     385       17147 :   _matrix_tag_active.copyToDevice();
     386       17147 : }
     387             : 
     388             : } // namespace Moose::Kokkos