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