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         316 : 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         238 :       _factory.create<Moose::Kokkos::AuxKernel>(kernel_name, name, parameters, 0);
      22             : 
      23         316 :   if (kernel->isNodal())
      24         192 :     _kokkos_nodal_aux_storage.addObject(kernel, 0);
      25             :   else
      26         124 :     _kokkos_elemental_aux_storage.addObject(kernel, 0);
      27         316 : }
      28             : 
      29             : void
      30     1461310 : AuxiliarySystem::kokkosCompute(ExecFlagType type)
      31             : {
      32     4383930 :   TIME_SECTION("computeKokkosAuxKernel", 1);
      33             : 
      34     1461310 :   if (!_kokkos_elemental_aux_storage[type].size() && !_kokkos_nodal_aux_storage[type].size())
      35     1449580 :     return;
      36             : 
      37       11730 :   auto & systems = _fe_problem.getKokkosSystems();
      38             : 
      39             :   // Resolve dependencies
      40             : 
      41       11730 :   std::set<MooseVariableFieldBase *> needed_moose_vars;
      42       11730 :   std::set<TagID> needed_fe_var_vector_tags;
      43       11730 :   std::unordered_set<unsigned int> needed_mat_props;
      44             : 
      45       11730 :   auto solution_tag = _fe_problem.getVectorTagID("parallel_solution");
      46             : 
      47       71996 :   for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
      48       60266 :     if (tag._id != solution_tag)
      49       60266 :       needed_fe_var_vector_tags.insert(tag._id);
      50             : 
      51       11730 :   associateVectorToTag(solution(), solution_tag);
      52             : 
      53       23460 :   for (auto block : _fe_problem.mesh().meshSubdomains())
      54             :   {
      55       11730 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
      56             :                                                                             needed_moose_vars);
      57       11730 :     _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
      58             :         block, needed_fe_var_vector_tags);
      59             : 
      60       11730 :     _kokkos_elemental_aux_storage[type].updateBlockVariableDependency(block, needed_moose_vars);
      61       11730 :     _kokkos_elemental_aux_storage[type].updateBlockFEVariableCoupledVectorTagDependency(
      62             :         block, needed_fe_var_vector_tags);
      63       11730 :     _kokkos_elemental_aux_storage[type].updateBlockMatPropDependency(block, needed_mat_props);
      64             : 
      65       11730 :     _kokkos_nodal_aux_storage[type].updateBlockVariableDependency(block, needed_moose_vars);
      66       11730 :     _kokkos_nodal_aux_storage[type].updateBlockFEVariableCoupledVectorTagDependency(
      67             :         block, needed_fe_var_vector_tags);
      68             :   }
      69             : 
      70       56784 :   for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
      71             :   {
      72       45054 :     _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryVariableDependency(boundary,
      73             :                                                                                needed_moose_vars);
      74       45054 :     _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryFEVariableCoupledVectorTagDependency(
      75             :         boundary, needed_fe_var_vector_tags);
      76             : 
      77       45054 :     _kokkos_elemental_aux_storage[type].updateBoundaryVariableDependency(boundary,
      78             :                                                                          needed_moose_vars);
      79       45054 :     _kokkos_elemental_aux_storage[type].updateBoundaryFEVariableCoupledVectorTagDependency(
      80             :         boundary, needed_fe_var_vector_tags);
      81       45054 :     _kokkos_elemental_aux_storage[type].updateBoundaryMatPropDependency(boundary, needed_mat_props);
      82             : 
      83       45054 :     _kokkos_nodal_aux_storage[type].updateBoundaryVariableDependency(boundary, needed_moose_vars);
      84       45054 :     _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       11730 :   systems[number()].sync(solution_tag, Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
      91             : 
      92       35190 :   for (auto & system : systems)
      93             :   {
      94       23460 :     system.setActiveVariables(needed_moose_vars);
      95       23460 :     system.setActiveSolutionTags(needed_fe_var_vector_tags);
      96             : 
      97             :     {
      98       70380 :       TIME_SECTION("KokkosCopy", 1);
      99       23460 :       system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
     100       23460 :     }
     101             :     {
     102       70380 :       TIME_SECTION("KokkosReinit", 1);
     103       23460 :       system.reinit();
     104       23460 :     }
     105             :   }
     106             : 
     107       11730 :   systems.copyToDevice();
     108             : 
     109             :   {
     110       35190 :     TIME_SECTION("KokkosMaterial", 1);
     111             : 
     112             :     // Compute material properties
     113             : 
     114       11730 :     _fe_problem.prepareKokkosMaterials(needed_mat_props);
     115       11730 :     _fe_problem.reinitKokkosMaterials();
     116       11730 :   }
     117             : 
     118             :   {
     119       35190 :     TIME_SECTION("KokkosAuxKernel", 1);
     120             : 
     121             :     // Compute auxiliary kernels
     122             : 
     123       24607 :     for (auto nodal : _kokkos_nodal_aux_storage[type].getActiveObjects())
     124       12877 :       nodal->compute();
     125             : 
     126       15884 :     for (auto elemental : _kokkos_elemental_aux_storage[type].getActiveObjects())
     127        4154 :       elemental->compute();
     128       11730 :   }
     129             : 
     130             :   // Close and restore vectors
     131             : 
     132             :   {
     133       35190 :     TIME_SECTION("KokkosCopy", 1);
     134             : 
     135       35190 :     for (auto & system : systems)
     136       23460 :       system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     137             : 
     138       11730 :     systems[number()].sync(solution_tag, Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
     139             : 
     140       11730 :     _sys.update();
     141       11730 :   }
     142             : 
     143             :   // Clear
     144             : 
     145       35190 :   for (auto & system : systems)
     146             :   {
     147       23460 :     system.clearActiveVariables();
     148       23460 :     system.clearActiveSolutionTags();
     149             :   }
     150             : 
     151       11730 :   disassociateVectorFromTag(solution(), solution_tag);
     152     1461310 : }