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 "KokkosKernelBase.h"
      11             : #include "KokkosNodalKernelBase.h"
      12             : #include "KokkosNodalBCBase.h"
      13             : #include "KokkosIntegratedBCBase.h"
      14             : 
      15             : #include "MaterialBase.h"
      16             : #include "NonlinearSystemBase.h"
      17             : #include "MooseVariableFieldBase.h"
      18             : #include "FEProblemBase.h"
      19             : 
      20             : void
      21        1299 : NonlinearSystemBase::addKokkosKernel(const std::string & kernel_name,
      22             :                                      const std::string & name,
      23             :                                      InputParameters & parameters)
      24             : {
      25             :   // Create the kernel object via the factory and add to warehouse
      26         305 :   std::shared_ptr<Moose::Kokkos::KernelBase> kernel =
      27         994 :       _factory.create<Moose::Kokkos::KernelBase>(kernel_name, name, parameters);
      28        1299 :   _kokkos_kernels.addObject(kernel, 0);
      29        1299 :   postAddResidualObject(*kernel);
      30        1299 : }
      31             : 
      32             : void
      33         272 : NonlinearSystemBase::addKokkosNodalKernel(const std::string & kernel_name,
      34             :                                           const std::string & name,
      35             :                                           InputParameters & parameters)
      36             : {
      37             :   // Create the kernel object via the factory and add to warehouse
      38          57 :   std::shared_ptr<Moose::Kokkos::NodalKernelBase> kernel =
      39         215 :       _factory.create<Moose::Kokkos::NodalKernelBase>(kernel_name, name, parameters);
      40         272 :   _kokkos_nodal_kernels.addObject(kernel, 0);
      41         272 :   postAddResidualObject(*kernel);
      42         272 : }
      43             : 
      44             : void
      45        1402 : NonlinearSystemBase::addKokkosBoundaryCondition(const std::string & bc_name,
      46             :                                                 const std::string & name,
      47             :                                                 InputParameters & parameters)
      48             : {
      49             :   // Create the object
      50         333 :   std::shared_ptr<Moose::Kokkos::BoundaryCondition> bc =
      51        1069 :       _factory.create<Moose::Kokkos::BoundaryCondition>(bc_name, name, parameters, 0);
      52        1402 :   postAddResidualObject(*bc);
      53             : 
      54             :   // Active BoundaryIDs for the object
      55        1402 :   const std::set<BoundaryID> & boundary_ids = bc->boundaryIDs();
      56        1402 :   auto bc_var = dynamic_cast<const MooseVariableFieldBase *>(&bc->variable());
      57        1402 :   _vars[0].addBoundaryVar(boundary_ids, bc_var);
      58             : 
      59             :   // Cast to the various types of BCs
      60         333 :   std::shared_ptr<Moose::Kokkos::NodalBCBase> nbc =
      61        1069 :       std::dynamic_pointer_cast<Moose::Kokkos::NodalBCBase>(bc);
      62         333 :   std::shared_ptr<Moose::Kokkos::IntegratedBCBase> ibc =
      63        1069 :       std::dynamic_pointer_cast<Moose::Kokkos::IntegratedBCBase>(bc);
      64             : 
      65             :   // NodalBCBase
      66        1402 :   if (nbc)
      67             :   {
      68        1201 :     if (!nbc->variable().isNodal())
      69           0 :       mooseError("Trying to use nodal boundary condition '",
      70           0 :                  nbc->name(),
      71             :                  "' on a non-nodal variable '",
      72           0 :                  nbc->variable().name(),
      73             :                  "'.");
      74             : 
      75        1201 :     _kokkos_nodal_bcs.addObject(nbc);
      76        1201 :     _vars[0].addBoundaryVars(boundary_ids, nbc->getCoupledVars());
      77             : 
      78             :     // DirichletBCs that are preset
      79        1201 :     if (nbc->preset())
      80        1023 :       _kokkos_preset_nodal_bcs.addObject(nbc);
      81             :   }
      82             : 
      83             :   // IntegratedBCBase
      84         201 :   else if (ibc)
      85             :   {
      86         201 :     _kokkos_integrated_bcs.addObject(ibc);
      87         201 :     _vars[0].addBoundaryVars(boundary_ids, ibc->getCoupledVars());
      88             :   }
      89             : 
      90             :   else
      91           0 :     mooseError("Unknown Kokkos BoundaryCondition type for object named ", bc->name());
      92        1402 : }
      93             : 
      94             : void
      95        1889 : NonlinearSystemBase::setKokkosInitialSolution()
      96             : {
      97             :   // The preset Dirichlet BCs can have coupled variables from other systems, so we grab all systems
      98        1889 :   auto & systems = _fe_problem.getKokkosSystems();
      99             : 
     100        1889 :   auto tag = _fe_problem.addVectorTag("parallel_solution", Moose::VECTOR_TAG_SOLUTION);
     101             : 
     102        1889 :   associateVectorToTag(solution(), tag);
     103             : 
     104        3778 :   std::set<TagID> needed_fe_var_vector_tags = {tag};
     105             : 
     106        8212 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
     107        6323 :     _kokkos_preset_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
     108             :         boundary, needed_fe_var_vector_tags);
     109             : 
     110        5667 :   for (auto & system : systems)
     111        3778 :     system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
     112             : 
     113        1889 :   systems.copyToDevice();
     114             : 
     115        5251 :   for (auto nbc : _kokkos_preset_nodal_bcs.getActiveObjects())
     116        3362 :     std::static_pointer_cast<Moose::Kokkos::NodalBCBase>(nbc)->presetSolution(tag);
     117             : 
     118        1889 :   Kokkos::fence();
     119             : 
     120        5667 :   for (auto & system : systems)
     121        3778 :     system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     122             : 
     123        1889 :   disassociateVectorFromTag(solution(), tag);
     124        1889 : }
     125             : 
     126             : void
     127       27177 : NonlinearSystemBase::computeKokkosResidual(const std::set<TagID> & tags)
     128             : {
     129       81531 :   TIME_SECTION("computeKokkosResidual", 1);
     130             : 
     131       27177 :   auto & systems = _fe_problem.getKokkosSystems();
     132             : 
     133       27177 :   systems[number()].setActiveResidualTags(tags);
     134             : 
     135             :   // Resolve dependencies
     136             : 
     137       27177 :   std::set<MooseVariableFieldBase *> needed_moose_vars;
     138       27177 :   std::set<TagID> needed_fe_var_vector_tags;
     139       27177 :   std::unordered_set<unsigned int> needed_mat_props;
     140             : 
     141      152436 :   for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
     142      125259 :     needed_fe_var_vector_tags.insert(tag._id);
     143             : 
     144       55986 :   for (auto block : _fe_problem.mesh().meshSubdomains())
     145             :   {
     146       28809 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
     147             :                                                                             needed_moose_vars);
     148       28809 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
     149             :         block, needed_fe_var_vector_tags);
     150             : 
     151       28809 :     _kokkos_kernels.updateBlockVariableDependency(block, needed_moose_vars);
     152       28809 :     _kokkos_kernels.updateBlockFEVariableCoupledVectorTagDependency(block,
     153             :                                                                     needed_fe_var_vector_tags);
     154       28809 :     _kokkos_kernels.updateBlockMatPropDependency(block, needed_mat_props);
     155             :   }
     156             : 
     157      126654 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
     158             :   {
     159       99477 :     _kokkos_integrated_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
     160       99477 :     _kokkos_integrated_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
     161             :         boundary, needed_fe_var_vector_tags);
     162             : 
     163       99477 :     _kokkos_nodal_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
     164       99477 :     _kokkos_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(boundary,
     165             :                                                                          needed_fe_var_vector_tags);
     166             :   }
     167             : 
     168      126654 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
     169       99477 :     _kokkos_integrated_bcs.updateBoundaryMatPropDependency(boundary, needed_mat_props);
     170             : 
     171             :   // Copy data and cache variable values at element quadature points
     172             : 
     173       81531 :   for (auto & system : systems)
     174             :   {
     175       54354 :     system.setActiveVariables(needed_moose_vars);
     176       54354 :     system.setActiveVariableTags(needed_fe_var_vector_tags);
     177             : 
     178             :     {
     179      163062 :       TIME_SECTION("KokkosCopy", 1);
     180       54354 :       system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
     181       54354 :     }
     182             :     {
     183      163062 :       TIME_SECTION("KokkosReinit", 1);
     184       54354 :       system.reinit();
     185       54354 :     }
     186             :   }
     187             : 
     188       27177 :   systems.copyToDevice();
     189             : 
     190             :   {
     191       81531 :     TIME_SECTION("KokkosMaterial", 1);
     192             : 
     193             :     // Compute material properties
     194             : 
     195       27177 :     _fe_problem.prepareKokkosMaterials(needed_mat_props);
     196       27177 :     _fe_problem.reinitKokkosMaterials();
     197       27177 :   }
     198             : 
     199             :   {
     200       81531 :     TIME_SECTION("KokkosKernel", 1);
     201             : 
     202             :     // Compute kernels
     203             : 
     204       81438 :     for (auto kernel : _kokkos_kernels.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
     205       54261 :       kernel->computeResidual();
     206             : 
     207       30695 :     for (auto nodal_kernel :
     208       61018 :          _kokkos_nodal_kernels.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
     209       20340 :       nodal_kernel->computeResidual();
     210             : 
     211       35513 :     for (auto ibc : _kokkos_integrated_bcs.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
     212        8336 :       ibc->computeResidual();
     213             : 
     214       74806 :     for (auto nbc : _kokkos_nodal_bcs.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
     215       47629 :       nbc->computeResidual();
     216       27177 :   }
     217             : 
     218             :   // Close and restore vectors
     219             : 
     220             :   {
     221       81531 :     TIME_SECTION("KokkosCopy", 1);
     222             : 
     223       81531 :     for (auto & system : systems)
     224       54354 :       system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     225       27177 :   }
     226             : 
     227             :   // Clear
     228             : 
     229       27177 :   systems[number()].clearActiveResidualTags();
     230             : 
     231       81531 :   for (auto & system : systems)
     232             :   {
     233       54354 :     system.clearActiveVariables();
     234       54354 :     system.clearActiveVariableTags();
     235             :   }
     236       27177 : }
     237             : 
     238             : void
     239        4410 : NonlinearSystemBase::computeKokkosJacobian(const std::set<TagID> & tags)
     240             : {
     241       13230 :   TIME_SECTION("computeKokkosJacobian", 1);
     242             : 
     243        4410 :   auto & systems = _fe_problem.getKokkosSystems();
     244             : 
     245        4410 :   systems[number()].setActiveMatrixTags(tags);
     246             : 
     247             :   // Resolve dependencies
     248             : 
     249        4410 :   std::set<MooseVariableFieldBase *> needed_moose_vars;
     250        4410 :   std::set<TagID> needed_fe_var_vector_tags;
     251        4410 :   std::unordered_set<unsigned int> needed_mat_props;
     252             : 
     253       25086 :   for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
     254       20676 :     needed_fe_var_vector_tags.insert(tag._id);
     255             : 
     256        9082 :   for (auto block : _fe_problem.mesh().meshSubdomains())
     257             :   {
     258        4672 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
     259             :                                                                             needed_moose_vars);
     260        4672 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
     261             :         block, needed_fe_var_vector_tags);
     262             : 
     263        4672 :     _kokkos_kernels.updateBlockVariableDependency(block, needed_moose_vars);
     264        4672 :     _kokkos_kernels.updateBlockFEVariableCoupledVectorTagDependency(block,
     265             :                                                                     needed_fe_var_vector_tags);
     266        4672 :     _kokkos_kernels.updateBlockMatPropDependency(block, needed_mat_props);
     267             :   }
     268             : 
     269       18496 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
     270             :   {
     271       14086 :     _kokkos_integrated_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
     272       14086 :     _kokkos_integrated_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
     273             :         boundary, needed_fe_var_vector_tags);
     274             : 
     275       14086 :     _kokkos_nodal_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
     276       14086 :     _kokkos_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(boundary,
     277             :                                                                          needed_fe_var_vector_tags);
     278             :   }
     279             : 
     280       18496 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
     281       14086 :     _kokkos_integrated_bcs.updateBoundaryMatPropDependency(boundary, needed_mat_props);
     282             : 
     283             :   // Copy data and cache variable values at element quadature points
     284             : 
     285       13230 :   for (auto & system : systems)
     286             :   {
     287        8820 :     system.setActiveVariables(needed_moose_vars);
     288        8820 :     system.setActiveVariableTags(needed_fe_var_vector_tags);
     289             : 
     290             :     {
     291       26460 :       TIME_SECTION("KokkosPreallocate", 1);
     292        8820 :       system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
     293        8820 :     }
     294             :     {
     295       26460 :       TIME_SECTION("KokkosReinit", 1);
     296        8820 :       system.reinit();
     297        8820 :     }
     298             :   }
     299             : 
     300        4410 :   systems.copyToDevice();
     301             : 
     302             :   {
     303       13230 :     TIME_SECTION("KokkosMaterial", 1);
     304             : 
     305             :     // Compute material properties
     306             : 
     307        4410 :     _fe_problem.prepareKokkosMaterials(needed_mat_props);
     308        4410 :     _fe_problem.reinitKokkosMaterials();
     309        4410 :   }
     310             : 
     311             :   {
     312       13230 :     TIME_SECTION("KokkosKernel", 1);
     313             : 
     314             :     // Compute kernels
     315             : 
     316       14375 :     for (auto kernel : _kokkos_kernels.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
     317        9965 :       kernel->computeJacobian();
     318             : 
     319        5454 :     for (auto nodal_kernel :
     320       11482 :          _kokkos_nodal_kernels.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
     321        5424 :       nodal_kernel->computeJacobian();
     322             : 
     323        5140 :     for (auto ibc : _kokkos_integrated_bcs.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
     324         730 :       ibc->computeJacobian();
     325             : 
     326       12424 :     for (auto nbc : _kokkos_nodal_bcs.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
     327        8014 :       nbc->computeJacobian();
     328        4410 :   }
     329             : 
     330             :   // Close and restore vectors and matrices
     331             : 
     332             :   {
     333       13230 :     TIME_SECTION("KokkosCopy", 1);
     334             : 
     335       13230 :     for (auto & system : systems)
     336        8820 :       system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     337        4410 :   }
     338             : 
     339             :   // Clear
     340             : 
     341        4410 :   systems[number()].clearActiveMatrixTags();
     342             : 
     343       13230 :   for (auto & system : systems)
     344             :   {
     345        8820 :     system.clearActiveVariables();
     346        8820 :     system.clearActiveVariableTags();
     347             :   }
     348        4410 : }