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 "KokkosTypes.h"
11 : #include "KokkosFunction.h"
12 : #include "KokkosUserObject.h"
13 : #include "KokkosPostprocessor.h"
14 : #include "KokkosVectorPostprocessor.h"
15 : #include "KokkosReporter.h"
16 :
17 : #include "FEProblemBase.h"
18 : #include "DisplacedProblem.h"
19 : #include "NonlinearSystemBase.h"
20 : #include "AuxiliarySystem.h"
21 : #include "MaterialBase.h"
22 : #include "MooseVariableDependencyInterface.h"
23 : #include "MaterialPropertyInterface.h"
24 : #include "Coupleable.h"
25 :
26 : Moose::Kokkos::System &
27 0 : FEProblemBase::getKokkosSystem(const unsigned int sys_num)
28 : {
29 0 : return _kokkos_systems[sys_num];
30 : }
31 :
32 : const Moose::Kokkos::System &
33 0 : FEProblemBase::getKokkosSystem(const unsigned int sys_num) const
34 : {
35 0 : return _kokkos_systems[sys_num];
36 : }
37 :
38 : void
39 3123 : FEProblemBase::addKokkosKernel(const std::string & kernel_name,
40 : const std::string & name,
41 : InputParameters & parameters)
42 : {
43 : parallel_object_only();
44 :
45 6246 : const auto nl_sys_num = determineSolverSystem(parameters.varName("variable", name), true).second;
46 3123 : if (!isSolverSystemNonlinear(nl_sys_num))
47 0 : mooseError("You are trying to add a Kernel to a linear variable/system, which is not "
48 : "supported at the moment!");
49 :
50 4539 : setResidualObjectParamsAndLog(
51 1707 : kernel_name, name, parameters, nl_sys_num, "KokkosKernel", _reinit_displaced_elem);
52 :
53 3123 : _nl[nl_sys_num]->addKokkosKernel(kernel_name, name, parameters);
54 :
55 3121 : _has_kokkos_objects = true;
56 3121 : _has_kokkos_residual_objects = true;
57 3121 : }
58 :
59 : void
60 407 : FEProblemBase::addKokkosNodalKernel(const std::string & kernel_name,
61 : const std::string & name,
62 : InputParameters & parameters)
63 : {
64 : parallel_object_only();
65 :
66 814 : const auto nl_sys_num = determineSolverSystem(parameters.varName("variable", name), true).second;
67 :
68 599 : setResidualObjectParamsAndLog(
69 215 : kernel_name, name, parameters, nl_sys_num, "KokkosNodalKernel", _reinit_displaced_elem);
70 :
71 407 : _nl[nl_sys_num]->addKokkosNodalKernel(kernel_name, name, parameters);
72 :
73 407 : _has_kokkos_objects = true;
74 407 : _has_kokkos_residual_objects = true;
75 407 : }
76 :
77 : void
78 3978 : FEProblemBase::addKokkosBoundaryCondition(const std::string & bc_name,
79 : const std::string & name,
80 : InputParameters & parameters)
81 : {
82 : parallel_object_only();
83 :
84 7956 : const auto nl_sys_num = determineSolverSystem(parameters.varName("variable", name), true).second;
85 3978 : if (!isSolverSystemNonlinear(nl_sys_num))
86 0 : mooseError(
87 : "You are trying to add a BoundaryCondition to a linear variable/system, which is not "
88 : "supported at the moment!");
89 :
90 5806 : setResidualObjectParamsAndLog(
91 2150 : bc_name, name, parameters, nl_sys_num, "KokkosBoundaryCondition", _reinit_displaced_face);
92 :
93 3978 : _nl[nl_sys_num]->addKokkosBoundaryCondition(bc_name, name, parameters);
94 :
95 3978 : _has_kokkos_objects = true;
96 3978 : _has_kokkos_residual_objects = true;
97 3978 : }
98 :
99 : void
100 621 : FEProblemBase::addKokkosAuxKernel(const std::string & kernel_name,
101 : const std::string & name,
102 : InputParameters & parameters)
103 : {
104 : parallel_object_only();
105 :
106 621 : setAuxKernelParamsAndLog(kernel_name, name, parameters, "KokkosAuxKernel");
107 :
108 621 : _aux->addKokkosKernel(kernel_name, name, parameters);
109 :
110 621 : _has_kokkos_objects = true;
111 621 : }
112 :
113 : void
114 226 : FEProblemBase::addKokkosFunction(const std::string & type,
115 : const std::string & name,
116 : InputParameters & parameters)
117 : {
118 : parallel_object_only();
119 :
120 452 : parameters.set<SubProblem *>("_subproblem") = this;
121 :
122 24 : std::shared_ptr<Moose::FunctionBase> func =
123 202 : _factory.create<Moose::FunctionBase>(type, name, parameters);
124 226 : logAdd("KokkosFunction", name, type, parameters);
125 226 : _kokkos_functions.addObject(func);
126 :
127 226 : _has_kokkos_objects = true;
128 226 : }
129 :
130 : bool
131 179 : FEProblemBase::hasKokkosFunction(const std::string & name) const
132 : {
133 179 : return _kokkos_functions.hasActiveObject(name);
134 : }
135 :
136 : Moose::Kokkos::Function
137 117 : FEProblemBase::getKokkosFunction(const std::string & name)
138 : {
139 : #ifdef MOOSE_ENABLE_KOKKOS_GPU
140 0 : mooseError("Retrieving a Kokkos function as abstract type is currently not supported for GPU.");
141 : #endif
142 :
143 117 : const auto & function = getKokkosFunction<Moose::FunctionBase>(name);
144 :
145 : std::shared_ptr<Moose::Kokkos::FunctionWrapperHostBase> wrapper =
146 117 : Moose::Kokkos::FunctorRegistry::buildFunction(&function, function.type());
147 :
148 234 : return Moose::Kokkos::Function(wrapper);
149 117 : }
150 :
151 : void
152 1022 : FEProblemBase::addKokkosPostprocessor(const std::string & pp_name,
153 : const std::string & name,
154 : InputParameters & parameters)
155 : {
156 1022 : checkUserObjectNameCollision(name, "KokkosPostprocessor");
157 :
158 1022 : addKokkosUserObject(pp_name, name, parameters);
159 1022 : }
160 :
161 : void
162 169 : FEProblemBase::addKokkosVectorPostprocessor(const std::string & pp_name,
163 : const std::string & name,
164 : InputParameters & parameters)
165 : {
166 169 : checkUserObjectNameCollision(name, "KokkosVectorPostprocessor");
167 :
168 169 : addKokkosUserObject(pp_name, name, parameters);
169 169 : }
170 :
171 : void
172 34 : FEProblemBase::addKokkosReporter(const std::string & pp_name,
173 : const std::string & name,
174 : InputParameters & parameters)
175 : {
176 34 : checkUserObjectNameCollision(name, "KokkosReporter");
177 :
178 34 : addKokkosUserObject(pp_name, name, parameters);
179 34 : }
180 :
181 : void
182 1225 : FEProblemBase::addKokkosUserObject(const std::string & user_object_name,
183 : const std::string & name,
184 : InputParameters & parameters)
185 : {
186 : parallel_object_only();
187 :
188 : // Add the _subproblem and _sys parameters depending on use_displaced_mesh
189 1225 : addObjectParamsHelper(parameters, name);
190 :
191 : // Create the UserObject
192 552 : std::shared_ptr<UserObjectBase> user_object =
193 673 : _factory.create<UserObjectBase>(user_object_name, name, parameters);
194 1225 : logAdd("KokkosUserObject", name, user_object_name, parameters);
195 :
196 1225 : theWarehouse().add(user_object);
197 :
198 1225 : _has_kokkos_objects = true;
199 1225 : }
200 :
201 : bool
202 55950 : FEProblemBase::hasKokkosUserObject(const std::string & name) const
203 : {
204 55950 : std::vector<UserObjectBase *> objs;
205 55950 : theWarehouse()
206 109088 : .query()
207 54544 : .condition<AttribSystem>("KokkosUserObject")
208 54544 : .condition<AttribName>(name)
209 54544 : .queryInto(objs);
210 111900 : return !objs.empty();
211 55950 : }
212 :
213 : void
214 942 : FEProblemBase::addKokkosMaterial(const std::string & mat_name,
215 : const std::string & name,
216 : InputParameters & parameters)
217 : {
218 1884 : addMaterialHelper({&_kokkos_materials}, mat_name, name, parameters);
219 :
220 934 : _has_kokkos_objects = true;
221 934 : }
222 :
223 : MaterialData &
224 24171 : FEProblemBase::getKokkosMaterialData(Moose::MaterialDataType type, const MooseObject * object) const
225 : {
226 24171 : switch (type)
227 : {
228 13265 : case Moose::BLOCK_MATERIAL_DATA:
229 13265 : if (object)
230 5393 : _kokkos_material_props.addConsumer(type, object);
231 13265 : return _kokkos_material_props.getMaterialData(0);
232 913 : case Moose::NEIGHBOR_MATERIAL_DATA:
233 913 : if (object)
234 913 : _kokkos_neighbor_material_props.addConsumer(type, object);
235 913 : return _kokkos_neighbor_material_props.getMaterialData(0);
236 9993 : case Moose::BOUNDARY_MATERIAL_DATA:
237 : case Moose::FACE_MATERIAL_DATA:
238 : case Moose::INTERFACE_MATERIAL_DATA:
239 9993 : if (object)
240 1743 : _kokkos_bnd_material_props.addConsumer(type, object);
241 9993 : return _kokkos_bnd_material_props.getMaterialData(0);
242 : }
243 :
244 0 : mooseError("FEProblemBase::getKokkosMaterialData(): Invalid MaterialDataType ", type);
245 : }
246 :
247 : const std::set<const MooseObject *> &
248 2265 : FEProblemBase::getKokkosMaterialPropertyStorageConsumers(Moose::MaterialDataType type) const
249 : {
250 2265 : switch (type)
251 : {
252 0 : case Moose::BLOCK_MATERIAL_DATA:
253 0 : return _kokkos_material_props.getConsumers(type);
254 0 : case Moose::NEIGHBOR_MATERIAL_DATA:
255 0 : return _kokkos_neighbor_material_props.getConsumers(type);
256 2265 : case Moose::BOUNDARY_MATERIAL_DATA:
257 : case Moose::FACE_MATERIAL_DATA:
258 : case Moose::INTERFACE_MATERIAL_DATA:
259 2265 : return _kokkos_bnd_material_props.getConsumers(type);
260 : }
261 :
262 0 : mooseError(
263 : "FEProblemBase::getKokkosMaterialPropertyStorageConsumers(): Invalid MaterialDataType ",
264 : type);
265 : }
266 :
267 : void
268 2265 : FEProblemBase::initKokkos()
269 : {
270 11325 : TIME_SECTION("InitializeKokkos", 1, "Initializing Kokkos");
271 :
272 : // Error on unsupported options
273 :
274 2265 : if (haveDisplaced())
275 0 : mooseError("Kokkos does not support displaced mesh yet.");
276 :
277 2265 : if (adaptivity().isOn())
278 0 : mooseError("Kokkos does not support adaptivity yet.");
279 :
280 : // Initialize Kokkos mesh
281 2265 : _mesh.getKokkosMesh()->update();
282 :
283 : // Initialize Kokkos assembly
284 2265 : _kokkos_assembly.init();
285 :
286 : // Call functions that require Kokkos mesh to be initialized first
287 21077 : for (auto hook : _kokkos_mesh_initialization_hooks)
288 18812 : hook();
289 :
290 : // Initialize Kokkos systems
291 :
292 2265 : unsigned int max_system_number = 0;
293 :
294 4530 : for (unsigned int s = 0; s < numNonlinearSystems(); ++s)
295 2265 : max_system_number = std::max(max_system_number, getNonlinearSystemBase(s).number());
296 :
297 2265 : max_system_number = std::max(max_system_number, getAuxiliarySystem().number());
298 :
299 : // Kokkos system does not have a default constructor, so this simply allocates buffer to which we
300 : // should do placement new later
301 2265 : _kokkos_systems.create<false>(max_system_number + 1);
302 :
303 4530 : for (unsigned int s = 0; s < numNonlinearSystems(); ++s)
304 : {
305 2265 : auto & nl = getNonlinearSystemBase(s);
306 2265 : new (&_kokkos_systems[nl.number()]) Moose::Kokkos::System(nl);
307 : }
308 :
309 2265 : auto & aux = getAuxiliarySystem();
310 2265 : new (&_kokkos_systems[aux.number()]) Moose::Kokkos::System(aux);
311 :
312 : // Initialize Kokkos material properties
313 :
314 2265 : _kokkos_material_props.allocateKokkosProperties();
315 2265 : _kokkos_bnd_material_props.allocateKokkosProperties();
316 2265 : _kokkos_neighbor_material_props.allocateKokkosProperties();
317 2265 : }
318 :
319 : void
320 781 : FEProblemBase::initKokkosStatefulProps()
321 : {
322 : // Resolve dependencies
323 :
324 781 : std::set<MooseVariableFieldBase *> needed_moose_vars;
325 781 : std::set<TagID> needed_fe_var_vector_tags;
326 :
327 781 : _kokkos_materials.updateVariableDependency(needed_moose_vars);
328 781 : _kokkos_materials.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
329 :
330 : // Copy data and preallocate quadature point solution vectors
331 :
332 2343 : for (auto & system : _kokkos_systems)
333 : {
334 1562 : system.setActiveVariables(needed_moose_vars);
335 1562 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
336 :
337 1562 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
338 1562 : system.reinit();
339 : }
340 :
341 781 : _kokkos_systems.copyToDevice();
342 :
343 : // Initialize stateful properties
344 :
345 1910 : for (auto & material : _kokkos_materials.getActiveObjects())
346 1129 : if (!material->hasRestoredProperties() && !material->boundaryRestricted())
347 1060 : material->initStatefulProperties(0);
348 :
349 1878 : for (auto & material : _kokkos_materials[Moose::FACE_MATERIAL_DATA].getActiveObjects())
350 1097 : if (!material->hasRestoredProperties())
351 1058 : material->initStatefulProperties(0);
352 :
353 1878 : for (auto & material : _kokkos_materials[Moose::NEIGHBOR_MATERIAL_DATA].getActiveObjects())
354 1097 : if (!material->hasRestoredProperties())
355 1060 : material->initStatefulProperties(0);
356 :
357 1910 : for (auto & material : _kokkos_materials.getActiveObjects())
358 1129 : if (!material->hasRestoredProperties() && material->boundaryRestricted())
359 30 : material->initStatefulProperties(0);
360 :
361 : // Copy to old and older states
362 :
363 781 : _kokkos_material_props.copy();
364 781 : _kokkos_bnd_material_props.copy();
365 781 : _kokkos_neighbor_material_props.copy();
366 :
367 : // Clear
368 :
369 2343 : for (auto & system : _kokkos_systems)
370 : {
371 1562 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
372 :
373 1562 : system.clearActiveVariables();
374 1562 : system.clearActiveSolutionTags();
375 : }
376 781 : }
377 :
378 : void
379 23146 : FEProblemBase::prepareKokkosMaterials(
380 : const std::unordered_set<unsigned int> & consumer_needed_mat_props)
381 : {
382 23146 : auto needed_mat_props = consumer_needed_mat_props;
383 :
384 23146 : _kokkos_materials.updateMatPropDependency(needed_mat_props, 0, true);
385 :
386 23146 : setActiveMaterialProperties(needed_mat_props, 0);
387 23146 : }
388 :
389 : void
390 23146 : FEProblemBase::reinitKokkosMaterials()
391 : {
392 23146 : if (hasActiveMaterialProperties(0))
393 : {
394 52273 : for (auto & material : _kokkos_materials.getActiveObjects())
395 29127 : if (material->hasActiveProperties() && !material->boundaryRestricted())
396 27550 : material->computeProperties();
397 :
398 50921 : for (auto & material : _kokkos_materials[Moose::FACE_MATERIAL_DATA].getActiveObjects())
399 27775 : if (material->hasActiveProperties())
400 27550 : material->computeProperties();
401 :
402 50921 : for (auto & material : _kokkos_materials[Moose::NEIGHBOR_MATERIAL_DATA].getActiveObjects())
403 27775 : if (material->hasActiveProperties())
404 27550 : material->computeProperties();
405 :
406 52273 : for (auto & material : _kokkos_materials.getActiveObjects())
407 29127 : if (material->hasActiveProperties() && material->boundaryRestricted())
408 1352 : material->computeProperties();
409 : }
410 23146 : }
411 :
412 : void
413 209388 : FEProblemBase::computeKokkosUserObjectsInternal(const ExecFlagType & type,
414 : TheWarehouse::Query & query)
415 : {
416 : try
417 : {
418 1046940 : TIME_SECTION("computeKokkosUserObjects", 1, "Computing Kokkos User Objects");
419 :
420 : // Query everything first to obtain a list of execution groups
421 :
422 209388 : std::vector<Moose::Kokkos::UserObject *> uos;
423 209388 : query.clone().queryIntoUnsorted(uos);
424 :
425 209388 : if (!uos.size())
426 207588 : return;
427 :
428 : // Resolve dependencies
429 :
430 1800 : std::set<MooseVariableFEBase *> needed_moose_vars;
431 1800 : std::unordered_set<unsigned int> needed_mat_props;
432 1800 : std::set<TagID> needed_fe_var_vector_tags;
433 :
434 4249 : for (const auto uo : uos)
435 : {
436 2449 : auto v_uo = dynamic_cast<MooseVariableDependencyInterface *>(uo);
437 2449 : if (v_uo)
438 : {
439 2359 : const auto & v_deps = v_uo->getMooseVariableDependencies();
440 2359 : needed_moose_vars.insert(v_deps.begin(), v_deps.end());
441 : }
442 :
443 2449 : auto m_uo = dynamic_cast<MaterialPropertyInterface *>(uo);
444 2449 : if (m_uo)
445 : {
446 1650 : auto & m_deps = m_uo->getMatPropDependencies();
447 1650 : needed_mat_props.insert(m_deps.begin(), m_deps.end());
448 : }
449 :
450 2449 : auto c_uo = dynamic_cast<Coupleable *>(uo);
451 2449 : if (c_uo)
452 : {
453 2359 : const auto & tag_deps = c_uo->getFEVariableCoupleableVectorTags();
454 2359 : needed_fe_var_vector_tags.insert(tag_deps.begin(), tag_deps.end());
455 : }
456 : }
457 :
458 1800 : if (needed_mat_props.size())
459 : {
460 147 : _kokkos_materials.updateVariableDependency(needed_moose_vars);
461 147 : _kokkos_materials.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
462 : }
463 :
464 : // Copy vectors and cache variable values at element quadature points
465 :
466 5400 : for (auto & system : _kokkos_systems)
467 : {
468 3600 : system.setActiveVariables(needed_moose_vars);
469 3600 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
470 :
471 : {
472 10800 : TIME_SECTION("KokkosCopy", 1);
473 3600 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
474 3600 : }
475 : {
476 10800 : TIME_SECTION("KokkosReinit", 1);
477 3600 : system.reinit();
478 3600 : }
479 : }
480 :
481 1800 : _kokkos_systems.copyToDevice();
482 :
483 : {
484 5400 : TIME_SECTION("KokkosMaterial", 1);
485 :
486 : // Compute material properties
487 :
488 1800 : if (needed_mat_props.size())
489 : {
490 147 : prepareKokkosMaterials(needed_mat_props);
491 147 : reinitKokkosMaterials();
492 : }
493 1800 : }
494 :
495 : {
496 5400 : TIME_SECTION("KokkosUserObject", 1);
497 :
498 : // Compute user objects
499 :
500 1800 : if (type == EXEC_LINEAR)
501 0 : for (auto obj : uos)
502 0 : obj->residualSetup();
503 1800 : else if (type == EXEC_NONLINEAR)
504 0 : for (auto obj : uos)
505 0 : obj->jacobianSetup();
506 :
507 4249 : for (auto obj : uos)
508 2449 : obj->initialize();
509 :
510 4249 : for (auto obj : uos)
511 2449 : obj->compute();
512 :
513 1800 : kokkosJoinAndFinalize(uos);
514 1800 : }
515 :
516 : // Close and restore vectors
517 :
518 : {
519 5400 : TIME_SECTION("KokkosCopy", 1);
520 :
521 5400 : for (auto & system : _kokkos_systems)
522 3600 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
523 1800 : }
524 :
525 : // Clear
526 :
527 5400 : for (auto & system : _kokkos_systems)
528 : {
529 3600 : system.clearActiveVariables();
530 3600 : system.clearActiveSolutionTags();
531 : }
532 416976 : }
533 0 : catch (...)
534 : {
535 0 : handleException("computeKokkosUserObjectsInternal");
536 0 : }
537 : }
538 :
539 : void
540 1800 : FEProblemBase::kokkosJoinAndFinalize(const std::vector<Moose::Kokkos::UserObject *> & userobjs)
541 : {
542 4249 : for (auto obj : userobjs)
543 : {
544 2449 : obj->finalize();
545 :
546 2449 : auto pp = dynamic_cast<const Moose::Kokkos::Postprocessor *>(obj);
547 2449 : if (pp)
548 : {
549 2256 : _reporter_data.finalize(obj->name());
550 2256 : setPostprocessorValueByName(obj->name(), pp->getValue());
551 : }
552 :
553 2449 : auto vpp = dynamic_cast<Moose::Kokkos::VectorPostprocessor *>(obj);
554 2449 : if (vpp)
555 163 : _reporter_data.finalize(obj->name());
556 :
557 2449 : auto reporter = dynamic_cast<Moose::Kokkos::Reporter *>(obj);
558 2449 : if (reporter)
559 30 : _reporter_data.finalize(obj->name());
560 : }
561 1800 : }
|