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 "KokkosAuxKernel.h"
      11             : 
      12             : #include "AuxiliarySystem.h"
      13             : #include "MaterialBase.h"
      14             : 
      15             : void
      16         280 : AuxiliarySystem::addKokkosKernel(const std::string & kernel_name,
      17             :                                  const std::string & name,
      18             :                                  InputParameters & parameters)
      19             : {
      20          78 :   std::shared_ptr<Moose::Kokkos::AuxKernel> kernel =
      21         202 :       _factory.create<Moose::Kokkos::AuxKernel>(kernel_name, name, parameters, 0);
      22             : 
      23         280 :   if (kernel->isNodal())
      24         156 :     _kokkos_nodal_aux_storage.addObject(kernel, 0);
      25             :   else
      26         124 :     _kokkos_elemental_aux_storage.addObject(kernel, 0);
      27         280 : }
      28             : 
      29             : void
      30     1453629 : AuxiliarySystem::kokkosCompute(ExecFlagType type)
      31             : {
      32     4360887 :   TIME_SECTION("computeKokkosAuxKernel", 1);
      33             : 
      34     1453629 :   if (!_kokkos_elemental_aux_storage[type].size() && !_kokkos_nodal_aux_storage[type].size())
      35     1446435 :     return;
      36             : 
      37        7194 :   auto & systems = _fe_problem.getKokkosSystems();
      38             : 
      39             :   // Resolve dependencies
      40             : 
      41        7194 :   std::set<MooseVariableFieldBase *> needed_moose_vars;
      42        7194 :   std::set<TagID> needed_fe_var_vector_tags;
      43        7194 :   std::unordered_set<unsigned int> needed_mat_props;
      44             : 
      45        7194 :   auto solution_tag = _fe_problem.getVectorTagID("parallel_solution");
      46             : 
      47       44780 :   for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
      48       37586 :     if (tag._id != solution_tag)
      49       37586 :       needed_fe_var_vector_tags.insert(tag._id);
      50             : 
      51        7194 :   associateVectorToTag(solution(), solution_tag);
      52             : 
      53       14388 :   for (auto block : _fe_problem.mesh().meshSubdomains())
      54             :   {
      55        7194 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
      56             :                                                                             needed_moose_vars);
      57        7194 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
      58             :         block, needed_fe_var_vector_tags);
      59             : 
      60        7194 :     _kokkos_elemental_aux_storage[type].updateBlockVariableDependency(block, needed_moose_vars);
      61        7194 :     _kokkos_elemental_aux_storage[type].updateBlockFEVariableCoupledVectorTagDependency(
      62             :         block, needed_fe_var_vector_tags);
      63        7194 :     _kokkos_elemental_aux_storage[type].updateBlockMatPropDependency(block, needed_mat_props);
      64             : 
      65        7194 :     _kokkos_nodal_aux_storage[type].updateBlockVariableDependency(block, needed_moose_vars);
      66        7194 :     _kokkos_nodal_aux_storage[type].updateBlockFEVariableCoupledVectorTagDependency(
      67             :         block, needed_fe_var_vector_tags);
      68             :   }
      69             : 
      70       34104 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
      71             :   {
      72       26910 :     _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryVariableDependency(boundary,
      73             :                                                                                needed_moose_vars);
      74       26910 :     _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryFEVariableCoupledVectorTagDependency(
      75             :         boundary, needed_fe_var_vector_tags);
      76             : 
      77       26910 :     _kokkos_elemental_aux_storage[type].updateBoundaryVariableDependency(boundary,
      78             :                                                                          needed_moose_vars);
      79       26910 :     _kokkos_elemental_aux_storage[type].updateBoundaryFEVariableCoupledVectorTagDependency(
      80             :         boundary, needed_fe_var_vector_tags);
      81       26910 :     _kokkos_elemental_aux_storage[type].updateBoundaryMatPropDependency(boundary, needed_mat_props);
      82             : 
      83       26910 :     _kokkos_nodal_aux_storage[type].updateBoundaryVariableDependency(boundary, needed_moose_vars);
      84       26910 :     _kokkos_nodal_aux_storage[type].updateBoundaryFEVariableCoupledVectorTagDependency(
      85             :         boundary, needed_fe_var_vector_tags);
      86             :   }
      87             : 
      88             :   // Copy data and cache variable values at element quadature points
      89             : 
      90        7194 :   systems[number()].sync(solution_tag, Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
      91             : 
      92       21582 :   for (auto & system : systems)
      93             :   {
      94       14388 :     system.setActiveVariables(needed_moose_vars);
      95       14388 :     system.setActiveSolutionTags(needed_fe_var_vector_tags);
      96             : 
      97             :     {
      98       43164 :       TIME_SECTION("KokkosCopy", 1);
      99       14388 :       system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
     100       14388 :     }
     101             :     {
     102       43164 :       TIME_SECTION("KokkosReinit", 1);
     103       14388 :       system.reinit();
     104       14388 :     }
     105             :   }
     106             : 
     107        7194 :   systems.copyToDevice();
     108             : 
     109             :   {
     110       21582 :     TIME_SECTION("KokkosMaterial", 1);
     111             : 
     112             :     // Compute material properties
     113             : 
     114        7194 :     _fe_problem.prepareKokkosMaterials(needed_mat_props);
     115        7194 :     _fe_problem.reinitKokkosMaterials();
     116        7194 :   }
     117             : 
     118             :   {
     119       21582 :     TIME_SECTION("KokkosAuxKernel", 1);
     120             : 
     121             :     // Compute auxiliary kernels
     122             : 
     123       15535 :     for (auto nodal : _kokkos_nodal_aux_storage[type].getActiveObjects())
     124        8341 :       nodal->compute();
     125             : 
     126       11348 :     for (auto elemental : _kokkos_elemental_aux_storage[type].getActiveObjects())
     127        4154 :       elemental->compute();
     128        7194 :   }
     129             : 
     130             :   // Close and restore vectors
     131             : 
     132             :   {
     133       21582 :     TIME_SECTION("KokkosCopy", 1);
     134             : 
     135       21582 :     for (auto & system : systems)
     136       14388 :       system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     137             : 
     138        7194 :     systems[number()].sync(solution_tag, Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     139             : 
     140        7194 :     _sys.update();
     141        7194 :   }
     142             : 
     143             :   // Clear
     144             : 
     145       21582 :   for (auto & system : systems)
     146             :   {
     147       14388 :     system.clearActiveVariables();
     148       14388 :     system.clearActiveSolutionTags();
     149             :   }
     150             : 
     151        7194 :   disassociateVectorFromTag(solution(), solution_tag);
     152     1453629 : }