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         621 : AuxiliarySystem::addKokkosKernel(const std::string & kernel_name,
      17             :                                  const std::string & name,
      18             :                                  InputParameters & parameters)
      19             : {
      20         286 :   std::shared_ptr<Moose::Kokkos::AuxKernel> kernel =
      21         335 :       _factory.create<Moose::Kokkos::AuxKernel>(kernel_name, name, parameters, 0);
      22             : 
      23         621 :   if (kernel->isNodal())
      24         340 :     _kokkos_nodal_aux_storage.addObject(kernel, 0);
      25             :   else
      26         281 :     _kokkos_elemental_aux_storage.addObject(kernel, 0);
      27         621 : }
      28             : 
      29             : void
      30     1491715 : AuxiliarySystem::kokkosCompute(ExecFlagType type)
      31             : {
      32     4475145 :   TIME_SECTION("computeKokkosAuxKernel", 3);
      33             : 
      34     4423445 :   if (!_kokkos_elemental_aux_storage[type].hasActiveObjects() &&
      35     2931730 :       !_kokkos_nodal_aux_storage[type].hasActiveObjects())
      36     1474984 :     return;
      37             : 
      38       16731 :   auto & systems = _fe_problem.getKokkosSystems();
      39             : 
      40             :   // Resolve dependencies
      41             : 
      42       16731 :   std::set<MooseVariableFieldBase *> needed_moose_vars;
      43       16731 :   std::set<TagID> needed_fe_var_vector_tags;
      44       16731 :   std::unordered_set<unsigned int> needed_mat_props;
      45             : 
      46       16731 :   auto solution_tag = _fe_problem.getVectorTagID("parallel_solution");
      47             : 
      48       99156 :   for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
      49       82425 :     if (tag._id != solution_tag)
      50       82425 :       needed_fe_var_vector_tags.insert(tag._id);
      51             : 
      52       16731 :   associateVectorToTag(solution(), solution_tag);
      53             : 
      54       16731 :   _kokkos_elemental_aux_storage[type].updateVariableDependency(needed_moose_vars);
      55       16731 :   _kokkos_elemental_aux_storage[type].updateFEVariableCoupledVectorTagDependency(
      56             :       needed_fe_var_vector_tags);
      57       16731 :   _kokkos_elemental_aux_storage[type].updateMatPropDependency(needed_mat_props);
      58             : 
      59       16731 :   _kokkos_nodal_aux_storage[type].updateVariableDependency(needed_moose_vars);
      60       16731 :   _kokkos_nodal_aux_storage[type].updateFEVariableCoupledVectorTagDependency(
      61             :       needed_fe_var_vector_tags);
      62             : 
      63       16731 :   if (needed_mat_props.size())
      64             :   {
      65        6839 :     _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
      66        6839 :     _fe_problem.getKokkosMaterialsWarehouse().updateFEVariableCoupledVectorTagDependency(
      67             :         needed_fe_var_vector_tags);
      68             :   }
      69             : 
      70             :   // Copy vectors and cache variable values at element quadature points
      71             : 
      72       16731 :   systems[number()].sync(solution_tag, Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
      73             : 
      74       50193 :   for (auto & system : systems)
      75             :   {
      76       33462 :     system.setActiveVariables(needed_moose_vars);
      77       33462 :     system.setActiveSolutionTags(needed_fe_var_vector_tags);
      78             : 
      79             :     {
      80      100386 :       TIME_SECTION("KokkosCopy", 1);
      81       33462 :       system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
      82       33462 :     }
      83             :     {
      84      100386 :       TIME_SECTION("KokkosReinit", 1);
      85       33462 :       system.reinit();
      86       33462 :     }
      87             :   }
      88             : 
      89       16731 :   systems.copyToDevice();
      90             : 
      91             :   {
      92       50193 :     TIME_SECTION("KokkosMaterial", 1);
      93             : 
      94             :     // Compute material properties
      95             : 
      96       16731 :     if (needed_mat_props.size())
      97             :     {
      98        6839 :       _fe_problem.prepareKokkosMaterials(needed_mat_props);
      99        6839 :       _fe_problem.reinitKokkosMaterials();
     100             :     }
     101       16731 :   }
     102             : 
     103             :   {
     104       50193 :     TIME_SECTION("KokkosAuxKernel", 1);
     105             : 
     106             :     // Compute auxiliary kernels
     107             : 
     108       34621 :     for (auto & nodal : _kokkos_nodal_aux_storage[type].getActiveObjects())
     109       17890 :       nodal->compute();
     110             : 
     111       23763 :     for (auto & elemental : _kokkos_elemental_aux_storage[type].getActiveObjects())
     112        7032 :       elemental->compute();
     113       16731 :   }
     114             : 
     115             :   // Close and restore vectors
     116             : 
     117             :   {
     118       50193 :     TIME_SECTION("KokkosCopy", 1);
     119             : 
     120       50193 :     for (auto & system : systems)
     121       33462 :       system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
     122             : 
     123       16731 :     systems[number()].sync(solution_tag, Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
     124             : 
     125       16731 :     _sys.update();
     126       16731 :   }
     127             : 
     128             :   // Clear
     129             : 
     130       50193 :   for (auto & system : systems)
     131             :   {
     132       33462 :     system.clearActiveVariables();
     133       33462 :     system.clearActiveSolutionTags();
     134             :   }
     135             : 
     136       16731 :   disassociateVectorFromTag(solution(), solution_tag);
     137     1491715 : }