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 "LinearSystem.h"
      11             : 
      12             : #include "Attributes.h"
      13             : #include "KokkosLinearFVElementalKernel.h"
      14             : #include "KokkosLinearFVFluxKernel.h"
      15             : #include "KokkosLinearFVBoundaryCondition.h"
      16             : #include "MooseUtils.h"
      17             : #include "MaterialPropertyInterface.h"
      18             : #include "MaterialBase.h"
      19             : #include "FEProblemBase.h"
      20             : 
      21             : #include "libmesh/equation_systems.h"
      22             : 
      23             : #include <set>
      24             : #include <unordered_set>
      25             : 
      26             : void
      27         400 : LinearSystem::addKokkosKernel(const std::string & kernel_name,
      28             :                               const std::string & name,
      29             :                               InputParameters & parameters)
      30             : {
      31         400 :   auto kernel = _factory.create<Moose::Kokkos::LinearFVKernel>(kernel_name, name, parameters);
      32             : 
      33         718 :   if (!std::dynamic_pointer_cast<Moose::Kokkos::LinearFVElementalKernel>(kernel) &&
      34         318 :       !std::dynamic_pointer_cast<Moose::Kokkos::LinearFVFluxKernel>(kernel))
      35           0 :     mooseError("Unknown Kokkos LinearFVKernel type for object named ", kernel->name());
      36             : 
      37         400 :   _fe_problem.theWarehouse().add(kernel);
      38         400 : }
      39             : 
      40             : void
      41         339 : LinearSystem::addKokkosBoundaryCondition(const std::string & bc_name,
      42             :                                          const std::string & name,
      43             :                                          InputParameters & parameters)
      44             : {
      45         339 :   auto bc = _factory.create<Moose::Kokkos::LinearFVBoundaryCondition>(bc_name, name, parameters);
      46         339 :   _fe_problem.theWarehouse().add(bc);
      47             : 
      48         339 :   const auto & boundary_ids = bc->boundaryIDs();
      49         339 :   const auto * const bc_var = bc->variable().mooseVar();
      50         339 :   _vars[0].addBoundaryVar(boundary_ids, bc_var);
      51         339 : }
      52             : 
      53             : void
      54         898 : LinearSystem::initialSetupKokkosLinearFV()
      55             : {
      56             :   // Kokkos objects are currently instantiated only for thread 0 (single copy per object).
      57             :   // All three queries below use AttribThread(0) for this reason. If Kokkos objects are ever
      58             :   // allowed to be multi-threaded, this function will need to loop over threads and each kernel's
      59             :   // initialSetup will work correctly via its own _tid member.
      60         898 :   std::vector<Moose::Kokkos::LinearFVElementalKernel *> elemental_kernels;
      61         898 :   _fe_problem.theWarehouse()
      62         809 :       .query()
      63         809 :       .template condition<AttribSysNum>(number())
      64         809 :       .template condition<AttribSystem>("LinearFVElementalKernel")
      65         809 :       .template condition<AttribKokkos>(true)
      66        1618 :       .template condition<AttribThread>(0)
      67         809 :       .queryInto(elemental_kernels);
      68             : 
      69        1069 :   for (auto * kernel : elemental_kernels)
      70         171 :     kernel->initialSetup();
      71             : 
      72         898 :   std::vector<Moose::Kokkos::LinearFVBoundaryCondition *> bcs;
      73         898 :   _fe_problem.theWarehouse()
      74         809 :       .query()
      75         809 :       .template condition<AttribSysNum>(number())
      76         809 :       .template condition<AttribSystem>("LinearFVBoundaryCondition")
      77         809 :       .template condition<AttribKokkos>(true)
      78        1618 :       .template condition<AttribThread>(0)
      79         809 :       .queryInto(bcs);
      80             : 
      81        1237 :   for (auto * bc : bcs)
      82         339 :     bc->initialSetup();
      83             : 
      84         898 :   std::vector<Moose::Kokkos::LinearFVFluxKernel *> flux_kernels;
      85         898 :   _fe_problem.theWarehouse()
      86         809 :       .query()
      87         809 :       .template condition<AttribSysNum>(number())
      88         809 :       .template condition<AttribSystem>("LinearFVFluxKernel")
      89         809 :       .template condition<AttribKokkos>(true)
      90        1618 :       .template condition<AttribThread>(0)
      91         809 :       .queryInto(flux_kernels);
      92             : 
      93        1127 :   for (auto * kernel : flux_kernels)
      94         229 :     kernel->initialSetup();
      95         898 : }
      96             : 
      97             : void
      98         292 : LinearSystem::computeKokkosLinearSystem(const std::set<TagID> & vector_tags,
      99             :                                         const std::set<TagID> & matrix_tags)
     100             : {
     101         876 :   TIME_SECTION("computeKokkosLinearSystem", 1);
     102             : 
     103         292 :   std::vector<Moose::Kokkos::LinearFVElementalKernel *> elemental_for_vector_tags;
     104         292 :   _fe_problem.theWarehouse()
     105         152 :       .query()
     106         152 :       .template condition<AttribSysNum>(number())
     107         152 :       .template condition<AttribSystem>("LinearFVElementalKernel")
     108         152 :       .template condition<AttribKokkos>(true)
     109         304 :       .template condition<AttribThread>(0)
     110         152 :       .template condition<AttribVectorTags>(vector_tags)
     111         152 :       .queryInto(elemental_for_vector_tags);
     112             : 
     113         292 :   std::vector<Moose::Kokkos::LinearFVElementalKernel *> elemental_for_matrix_tags;
     114         292 :   _fe_problem.theWarehouse()
     115         152 :       .query()
     116         152 :       .template condition<AttribSysNum>(number())
     117         152 :       .template condition<AttribSystem>("LinearFVElementalKernel")
     118         152 :       .template condition<AttribKokkos>(true)
     119         304 :       .template condition<AttribThread>(0)
     120         152 :       .template condition<AttribMatrixTags>(matrix_tags)
     121         152 :       .queryInto(elemental_for_matrix_tags);
     122             : 
     123         292 :   std::vector<Moose::Kokkos::LinearFVFluxKernel *> flux_for_vector_tags;
     124         292 :   _fe_problem.theWarehouse()
     125         152 :       .query()
     126         152 :       .template condition<AttribSysNum>(number())
     127         152 :       .template condition<AttribSystem>("LinearFVFluxKernel")
     128         152 :       .template condition<AttribKokkos>(true)
     129         304 :       .template condition<AttribThread>(0)
     130         152 :       .template condition<AttribVectorTags>(vector_tags)
     131         152 :       .queryInto(flux_for_vector_tags);
     132             : 
     133         292 :   std::vector<Moose::Kokkos::LinearFVFluxKernel *> flux_for_matrix_tags;
     134         292 :   _fe_problem.theWarehouse()
     135         152 :       .query()
     136         152 :       .template condition<AttribSysNum>(number())
     137         152 :       .template condition<AttribSystem>("LinearFVFluxKernel")
     138         152 :       .template condition<AttribKokkos>(true)
     139         304 :       .template condition<AttribThread>(0)
     140         152 :       .template condition<AttribMatrixTags>(matrix_tags)
     141         152 :       .queryInto(flux_for_matrix_tags);
     142             : 
     143         292 :   std::vector<Moose::Kokkos::LinearFVBoundaryCondition *> bcs;
     144         292 :   _fe_problem.theWarehouse()
     145         152 :       .query()
     146         152 :       .template condition<AttribSysNum>(number())
     147         152 :       .template condition<AttribSystem>("LinearFVBoundaryCondition")
     148         152 :       .template condition<AttribKokkos>(true)
     149         304 :       .template condition<AttribThread>(0)
     150         152 :       .queryInto(bcs);
     151             : 
     152         292 :   std::vector<Moose::Kokkos::LinearFVElementalKernel *> elemental_kernels;
     153         292 :   MooseUtils::getUnion(elemental_for_vector_tags, elemental_for_matrix_tags, elemental_kernels);
     154             : 
     155         292 :   std::vector<Moose::Kokkos::LinearFVFluxKernel *> flux_kernels;
     156         292 :   MooseUtils::getUnion(flux_for_vector_tags, flux_for_matrix_tags, flux_kernels);
     157             : 
     158         292 :   if (elemental_kernels.empty() && flux_kernels.empty())
     159           0 :     return;
     160             : 
     161         292 :   auto & our_kokkos_system = _fe_problem.getKokkosSystem(number());
     162         292 :   our_kokkos_system.setActiveResidualTags(vector_tags);
     163         292 :   our_kokkos_system.setActiveMatrixTags(matrix_tags);
     164             : 
     165         292 :   std::set<TagID> needed_solution_vector_tags;
     166             : 
     167         727 :   for (const auto & tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
     168         435 :     needed_solution_vector_tags.insert(tag._id);
     169             : 
     170         292 :   auto & kokkos_systems = _fe_problem.getKokkosSystems();
     171         597 :   for (auto & system : kokkos_systems.constructedEntries())
     172             :   {
     173         305 :     system.setActiveSolutionTags(needed_solution_vector_tags);
     174         305 :     system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
     175             :   }
     176             : 
     177         292 :   kokkos_systems.copyToDevice();
     178             : 
     179         292 :   std::unordered_set<unsigned int> value_bc_vars;
     180         292 :   std::unordered_set<unsigned int> normal_gradient_bc_vars;
     181             : 
     182         632 :   for (const auto kernel : flux_kernels)
     183             :   {
     184         340 :     if (kernel->needsBoundaryValueData())
     185          28 :       value_bc_vars.insert(kernel->variable().var());
     186         340 :     if (kernel->needsBoundaryNormalGradientData())
     187         312 :       normal_gradient_bc_vars.insert(kernel->variable().var());
     188             :   }
     189             : 
     190         855 :   for (auto bc : bcs)
     191             :   {
     192         563 :     const auto bc_var = bc->variable().var();
     193         563 :     if (value_bc_vars.count(bc_var))
     194          56 :       bc->computeBoundaryValueData();
     195         563 :     if (normal_gradient_bc_vars.count(bc_var))
     196         563 :       bc->computeBoundaryNormalGradientData();
     197             :   }
     198             : 
     199         461 :   for (auto kernel : elemental_for_matrix_tags)
     200         169 :     kernel->computeMatrix();
     201         632 :   for (auto kernel : flux_for_matrix_tags)
     202         340 :     kernel->computeMatrix();
     203         461 :   for (auto kernel : elemental_for_vector_tags)
     204         169 :     kernel->computeRightHandSide();
     205         632 :   for (auto kernel : flux_for_vector_tags)
     206         340 :     kernel->computeRightHandSide();
     207             : 
     208         597 :   for (auto & system : kokkos_systems.constructedEntries())
     209         305 :     system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
     210             : 
     211         292 :   our_kokkos_system.clearActiveResidualTags();
     212         292 :   our_kokkos_system.clearActiveMatrixTags();
     213             : 
     214         597 :   for (auto & system : kokkos_systems.constructedEntries())
     215             :   {
     216         305 :     system.clearActiveVariables();
     217         305 :     system.clearActiveSolutionTags();
     218             :   }
     219         292 : }