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 "KokkosKernelBase.h"
11 : #include "KokkosNodalKernelBase.h"
12 : #include "KokkosNodalBCBase.h"
13 : #include "KokkosIntegratedBCBase.h"
14 :
15 : #include "MaterialBase.h"
16 : #include "NonlinearSystemBase.h"
17 : #include "MooseVariableFieldBase.h"
18 : #include "FEProblemBase.h"
19 :
20 : void
21 3123 : NonlinearSystemBase::addKokkosKernel(const std::string & kernel_name,
22 : const std::string & name,
23 : InputParameters & parameters)
24 : {
25 : // Create the kernel object via the factory and add to warehouse
26 1416 : std::shared_ptr<Moose::Kokkos::KernelBase> kernel =
27 1707 : _factory.create<Moose::Kokkos::KernelBase>(kernel_name, name, parameters);
28 3121 : _kokkos_kernels.addObject(kernel, 0);
29 3121 : postAddResidualObject(*kernel);
30 3121 : }
31 :
32 : void
33 407 : NonlinearSystemBase::addKokkosNodalKernel(const std::string & kernel_name,
34 : const std::string & name,
35 : InputParameters & parameters)
36 : {
37 : // Create the kernel object via the factory and add to warehouse
38 192 : std::shared_ptr<Moose::Kokkos::NodalKernelBase> kernel =
39 215 : _factory.create<Moose::Kokkos::NodalKernelBase>(kernel_name, name, parameters);
40 407 : _kokkos_nodal_kernels.addObject(kernel, 0);
41 407 : postAddResidualObject(*kernel);
42 407 : }
43 :
44 : void
45 3978 : NonlinearSystemBase::addKokkosBoundaryCondition(const std::string & bc_name,
46 : const std::string & name,
47 : InputParameters & parameters)
48 : {
49 : // Create the object
50 1828 : std::shared_ptr<Moose::Kokkos::BoundaryCondition> bc =
51 2150 : _factory.create<Moose::Kokkos::BoundaryCondition>(bc_name, name, parameters, 0);
52 3978 : postAddResidualObject(*bc);
53 :
54 : // Active BoundaryIDs for the object
55 3978 : const std::set<BoundaryID> & boundary_ids = bc->boundaryIDs();
56 3978 : auto bc_var = dynamic_cast<const MooseVariableFieldBase *>(&bc->variable());
57 3978 : _vars[0].addBoundaryVar(boundary_ids, bc_var);
58 :
59 : // Cast to the various types of BCs
60 1828 : std::shared_ptr<Moose::Kokkos::NodalBCBase> nbc =
61 2150 : std::dynamic_pointer_cast<Moose::Kokkos::NodalBCBase>(bc);
62 1828 : std::shared_ptr<Moose::Kokkos::IntegratedBCBase> ibc =
63 2150 : std::dynamic_pointer_cast<Moose::Kokkos::IntegratedBCBase>(bc);
64 :
65 : // NodalBCBase
66 3978 : if (nbc)
67 : {
68 3411 : if (!nbc->variable().isNodal())
69 0 : mooseError("Trying to use nodal boundary condition '",
70 0 : nbc->name(),
71 : "' on a non-nodal variable '",
72 0 : nbc->variable().name(),
73 : "'.");
74 :
75 3411 : _kokkos_nodal_bcs.addObject(nbc);
76 3411 : _vars[0].addBoundaryVars(boundary_ids, nbc->getCoupledVars());
77 :
78 : // DirichletBCs that are preset
79 3411 : if (nbc->preset())
80 3121 : _kokkos_preset_nodal_bcs.addObject(nbc);
81 : }
82 :
83 : // IntegratedBCBase
84 567 : else if (ibc)
85 : {
86 567 : _kokkos_integrated_bcs.addObject(ibc);
87 567 : _vars[0].addBoundaryVars(boundary_ids, ibc->getCoupledVars());
88 : }
89 :
90 : else
91 0 : mooseError("Unknown Kokkos BoundaryCondition type for object named ", bc->name());
92 3978 : }
93 :
94 : void
95 4368 : NonlinearSystemBase::setKokkosInitialSolution()
96 : {
97 : // The preset Dirichlet BCs can have coupled variables from other systems, so we grab all systems
98 4368 : auto & systems = _fe_problem.getKokkosSystems();
99 :
100 4368 : auto tag = _fe_problem.addVectorTag("parallel_solution", Moose::VECTOR_TAG_SOLUTION);
101 :
102 4368 : associateVectorToTag(solution(), tag);
103 :
104 8736 : std::set<TagID> needed_fe_var_vector_tags = {tag};
105 :
106 19898 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
107 15530 : _kokkos_preset_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
108 : boundary, needed_fe_var_vector_tags);
109 :
110 13104 : for (auto & system : systems)
111 8736 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
112 :
113 4368 : systems.copyToDevice();
114 :
115 11737 : for (auto & nbc : _kokkos_preset_nodal_bcs.getActiveObjects())
116 7369 : std::static_pointer_cast<Moose::Kokkos::NodalBCBase>(nbc)->presetSolution(tag);
117 :
118 4368 : Kokkos::fence();
119 :
120 13104 : for (auto & system : systems)
121 8736 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
122 :
123 4368 : disassociateVectorFromTag(solution(), tag);
124 4368 : }
125 :
126 : void
127 59942 : NonlinearSystemBase::computeKokkosResidual(const std::set<TagID> & tags)
128 : {
129 179826 : TIME_SECTION("computeKokkosResidual", 1);
130 :
131 : // Get warehouses
132 :
133 59942 : const auto & kernels = _kokkos_kernels.getVectorTagsObjectWarehouse(tags, 0);
134 59942 : const auto & nodal_kernels = _kokkos_nodal_kernels.getVectorTagsObjectWarehouse(tags, 0);
135 59942 : const auto & integrated_bcs = _kokkos_integrated_bcs.getVectorTagsObjectWarehouse(tags, 0);
136 :
137 59942 : if (!kernels.size() && !nodal_kernels.size() && !integrated_bcs.size())
138 0 : return;
139 :
140 : // Resolve dependencies
141 :
142 59942 : auto & systems = _fe_problem.getKokkosSystems();
143 :
144 59942 : systems[number()].setActiveResidualTags(tags);
145 :
146 59942 : std::set<MooseVariableFieldBase *> needed_moose_vars;
147 59942 : std::set<TagID> needed_fe_var_vector_tags;
148 59942 : std::unordered_set<unsigned int> needed_mat_props;
149 :
150 323062 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
151 263120 : needed_fe_var_vector_tags.insert(tag._id);
152 :
153 59942 : kernels.updateVariableDependency(needed_moose_vars);
154 59942 : kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
155 59942 : kernels.updateMatPropDependency(needed_mat_props);
156 :
157 59942 : nodal_kernels.updateVariableDependency(needed_moose_vars);
158 59942 : nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
159 :
160 59942 : integrated_bcs.updateVariableDependency(needed_moose_vars);
161 59942 : integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
162 59942 : integrated_bcs.updateMatPropDependency(needed_mat_props);
163 :
164 59942 : if (needed_mat_props.size())
165 : {
166 14420 : _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
167 14420 : _fe_problem.getKokkosMaterialsWarehouse().updateFEVariableCoupledVectorTagDependency(
168 : needed_fe_var_vector_tags);
169 : }
170 :
171 : // Copy solution vectors and residuals and cache variable values at element quadature
172 : // points
173 :
174 179826 : for (auto & system : systems)
175 : {
176 119884 : system.setActiveVariables(needed_moose_vars);
177 119884 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
178 :
179 : {
180 359652 : TIME_SECTION("KokkosCopy", 1);
181 119884 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
182 119884 : }
183 : {
184 359652 : TIME_SECTION("KokkosReinit", 1);
185 119884 : system.reinit();
186 119884 : }
187 : }
188 :
189 59942 : systems.copyToDevice();
190 :
191 : {
192 179826 : TIME_SECTION("KokkosMaterial", 1);
193 :
194 : // Compute material properties
195 :
196 59942 : if (needed_mat_props.size())
197 : {
198 14420 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
199 14420 : _fe_problem.reinitKokkosMaterials();
200 : }
201 59942 : }
202 :
203 : {
204 179826 : TIME_SECTION("KokkosKernel", 1);
205 :
206 : // Compute kernels
207 :
208 175495 : for (auto & kernel : kernels.getActiveObjects())
209 115553 : kernel->computeResidual();
210 :
211 91831 : for (auto & nodal_kernel : nodal_kernels.getActiveObjects())
212 31889 : nodal_kernel->computeResidual();
213 :
214 86039 : for (auto & ibc : integrated_bcs.getActiveObjects())
215 26097 : ibc->computeResidual();
216 59942 : }
217 :
218 : // Close and restore vectors
219 :
220 : {
221 179826 : TIME_SECTION("KokkosClose", 1);
222 :
223 179826 : for (auto & system : systems)
224 119884 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
225 59942 : }
226 :
227 : // Clear
228 :
229 59942 : systems[number()].clearActiveResidualTags();
230 :
231 179826 : for (auto & system : systems)
232 : {
233 119884 : system.clearActiveVariables();
234 119884 : system.clearActiveSolutionTags();
235 : }
236 59942 : }
237 :
238 : void
239 66416 : NonlinearSystemBase::computeKokkosNodalBCsResidual(const std::set<TagID> & tags)
240 : {
241 199248 : TIME_SECTION("computeKokkosNodalBCsResidual", 1);
242 :
243 : // Get warehouses
244 :
245 66416 : const auto & nodal_bcs = _kokkos_nodal_bcs.getVectorTagsObjectWarehouse(tags, 0);
246 :
247 66416 : if (!nodal_bcs.hasActiveObjects())
248 1575 : return;
249 :
250 : // Resolve dependencies
251 :
252 64841 : auto & systems = _fe_problem.getKokkosSystems();
253 :
254 64841 : systems[number()].setActiveResidualTags(tags);
255 :
256 64841 : std::set<TagID> needed_fe_var_vector_tags;
257 :
258 349080 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
259 284239 : needed_fe_var_vector_tags.insert(tag._id);
260 :
261 64841 : nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
262 :
263 : // Copy solution vectors and residuals
264 :
265 194523 : for (auto & system : systems)
266 : {
267 129682 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
268 :
269 : {
270 389046 : TIME_SECTION("KokkosCopy", 1);
271 129682 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
272 129682 : }
273 : }
274 :
275 64841 : systems.copyToDevice();
276 :
277 : {
278 194523 : TIME_SECTION("KokkosKernel", 1);
279 :
280 : // Compute kernels
281 :
282 177305 : for (auto & nbc : nodal_bcs.getActiveObjects())
283 112464 : nbc->computeResidual();
284 64841 : }
285 :
286 : // Close and restore vectors
287 :
288 : {
289 194523 : TIME_SECTION("KokkosClose", 1);
290 :
291 194523 : for (auto & system : systems)
292 129682 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
293 64841 : }
294 :
295 : // Clear
296 :
297 64841 : systems[number()].clearActiveResidualTags();
298 :
299 194523 : for (auto & system : systems)
300 129682 : system.clearActiveSolutionTags();
301 66416 : }
302 :
303 : void
304 9317 : NonlinearSystemBase::computeKokkosJacobian(const std::set<TagID> & tags)
305 : {
306 27951 : TIME_SECTION("computeKokkosJacobian", 1);
307 :
308 : // Get warehouses
309 :
310 9317 : const auto & kernels = _kokkos_kernels.getMatrixTagsObjectWarehouse(tags, 0);
311 9317 : const auto & nodal_kernels = _kokkos_nodal_kernels.getMatrixTagsObjectWarehouse(tags, 0);
312 9317 : const auto & integrated_bcs = _kokkos_integrated_bcs.getMatrixTagsObjectWarehouse(tags, 0);
313 9317 : const auto & nodal_bcs = _kokkos_nodal_bcs.getMatrixTagsObjectWarehouse(tags, 0);
314 :
315 9317 : if (!kernels.hasActiveObjects() && !nodal_kernels.hasActiveObjects() &&
316 5514 : !integrated_bcs.hasActiveObjects() && !nodal_bcs.hasActiveObjects())
317 0 : return;
318 :
319 : // Resolve dependencies
320 :
321 9317 : auto & systems = _fe_problem.getKokkosSystems();
322 :
323 9317 : systems[number()].setActiveMatrixTags(tags);
324 :
325 9317 : std::set<MooseVariableFieldBase *> needed_moose_vars;
326 9317 : std::set<TagID> needed_fe_var_vector_tags;
327 9317 : std::unordered_set<unsigned int> needed_mat_props;
328 :
329 51264 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
330 41947 : needed_fe_var_vector_tags.insert(tag._id);
331 :
332 9317 : kernels.updateVariableDependency(needed_moose_vars);
333 9317 : kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
334 9317 : kernels.updateMatPropDependency(needed_mat_props);
335 :
336 9317 : nodal_kernels.updateVariableDependency(needed_moose_vars);
337 9317 : nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
338 :
339 9317 : integrated_bcs.updateVariableDependency(needed_moose_vars);
340 9317 : integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
341 9317 : integrated_bcs.updateMatPropDependency(needed_mat_props);
342 :
343 9317 : nodal_bcs.updateVariableDependency(needed_moose_vars);
344 9317 : nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
345 :
346 9317 : if (needed_mat_props.size())
347 : {
348 1734 : _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
349 1734 : _fe_problem.getKokkosMaterialsWarehouse().updateFEVariableCoupledVectorTagDependency(
350 : needed_fe_var_vector_tags);
351 : }
352 :
353 : // Copy solution vectors, initialize matrices, and cache variable values at element quadature
354 : // points
355 :
356 27951 : for (auto & system : systems)
357 : {
358 18634 : system.setActiveVariables(needed_moose_vars);
359 18634 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
360 :
361 : {
362 55902 : TIME_SECTION("KokkosCopy", 1);
363 18634 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
364 18634 : }
365 : {
366 55902 : TIME_SECTION("KokkosReinit", 1);
367 18634 : system.reinit();
368 18634 : }
369 : }
370 :
371 9317 : systems.copyToDevice();
372 :
373 : {
374 27951 : TIME_SECTION("KokkosMaterial", 1);
375 :
376 : // Compute material properties
377 :
378 9317 : if (needed_mat_props.size())
379 : {
380 1734 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
381 1734 : _fe_problem.reinitKokkosMaterials();
382 : }
383 9317 : }
384 :
385 : {
386 27951 : TIME_SECTION("KokkosKernel", 1);
387 :
388 : // Compute kernels
389 :
390 28937 : for (auto & kernel : kernels.getActiveObjects())
391 19620 : kernel->computeJacobian();
392 :
393 17511 : for (auto & nodal_kernel : nodal_kernels.getActiveObjects())
394 8194 : nodal_kernel->computeJacobian();
395 :
396 12332 : for (auto & ibc : integrated_bcs.getActiveObjects())
397 3015 : ibc->computeJacobian();
398 :
399 24795 : for (auto & nbc : nodal_bcs.getActiveObjects())
400 15478 : nbc->computeJacobian();
401 9317 : }
402 :
403 : // Close and restore vectors and matrices
404 :
405 : {
406 27951 : TIME_SECTION("KokkosClose", 1);
407 :
408 27951 : for (auto & system : systems)
409 18634 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
410 9317 : }
411 :
412 : // Clear
413 :
414 9317 : systems[number()].clearActiveMatrixTags();
415 :
416 27951 : for (auto & system : systems)
417 : {
418 18634 : system.clearActiveVariables();
419 18634 : system.clearActiveSolutionTags();
420 : }
421 9317 : }
422 :
423 : void
424 6474 : NonlinearSystemBase::computeKokkosResidualAndJacobian(const std::set<TagID> & vector_tags,
425 : const std::set<TagID> & matrix_tags)
426 : {
427 19422 : TIME_SECTION("computeKokkosResidualAndJacobian", 1);
428 :
429 6474 : if (!_kokkos_kernels.hasActiveObjects() && !_kokkos_nodal_kernels.hasActiveObjects() &&
430 3254 : !_kokkos_integrated_bcs.hasActiveObjects() && !_kokkos_nodal_bcs.hasActiveObjects())
431 0 : return;
432 :
433 : // Resolve dependencies
434 :
435 6474 : auto & systems = _fe_problem.getKokkosSystems();
436 :
437 6474 : systems[number()].setActiveResidualTags(vector_tags);
438 6474 : systems[number()].setActiveMatrixTags(matrix_tags);
439 :
440 6474 : std::set<MooseVariableFieldBase *> needed_moose_vars;
441 6474 : std::set<TagID> needed_fe_var_vector_tags;
442 6474 : std::unordered_set<unsigned int> needed_mat_props;
443 :
444 34650 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
445 28176 : needed_fe_var_vector_tags.insert(tag._id);
446 :
447 6474 : _kokkos_kernels.updateVariableDependency(needed_moose_vars);
448 6474 : _kokkos_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
449 6474 : _kokkos_kernels.updateMatPropDependency(needed_mat_props);
450 :
451 6474 : _kokkos_nodal_kernels.updateVariableDependency(needed_moose_vars);
452 6474 : _kokkos_nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
453 :
454 6474 : _kokkos_integrated_bcs.updateVariableDependency(needed_moose_vars);
455 6474 : _kokkos_integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
456 6474 : _kokkos_integrated_bcs.updateMatPropDependency(needed_mat_props);
457 :
458 6474 : _kokkos_nodal_bcs.updateVariableDependency(needed_moose_vars);
459 6474 : _kokkos_nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
460 :
461 6474 : if (needed_mat_props.size())
462 : {
463 0 : _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
464 0 : _fe_problem.getKokkosMaterialsWarehouse().updateFEVariableCoupledVectorTagDependency(
465 : needed_fe_var_vector_tags);
466 : }
467 :
468 : // Copy solution vectors and residuals, initialize matrices, and cache variable values at element
469 : // quadature points
470 :
471 19422 : for (auto & system : systems)
472 : {
473 12948 : system.setActiveVariables(needed_moose_vars);
474 12948 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
475 :
476 : {
477 38844 : TIME_SECTION("KokkosCopy", 1);
478 12948 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
479 12948 : }
480 : {
481 38844 : TIME_SECTION("KokkosReinit", 1);
482 12948 : system.reinit();
483 12948 : }
484 : }
485 :
486 6474 : systems.copyToDevice();
487 :
488 : {
489 19422 : TIME_SECTION("KokkosMaterial", 1);
490 :
491 : // Compute material properties
492 :
493 6474 : if (needed_mat_props.size())
494 : {
495 0 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
496 0 : _fe_problem.reinitKokkosMaterials();
497 : }
498 6474 : }
499 :
500 : {
501 19422 : TIME_SECTION("KokkosKernel", 1);
502 :
503 : // Compute kernels
504 :
505 22881 : for (auto & kernel : _kokkos_kernels.getActiveObjects())
506 16407 : kernel->computeResidualAndJacobian();
507 :
508 6474 : for (auto & nodal_kernel : _kokkos_nodal_kernels.getActiveObjects())
509 0 : nodal_kernel->computeResidualAndJacobian();
510 :
511 7686 : for (auto & ibc : _kokkos_integrated_bcs.getActiveObjects())
512 1212 : ibc->computeResidualAndJacobian();
513 :
514 : /// Nodal BC residuals are computed separately
515 20484 : for (auto & nbc : _kokkos_nodal_bcs.getActiveObjects())
516 14010 : nbc->computeJacobian();
517 6474 : }
518 :
519 : // Close and restore vectors and matrices
520 :
521 : {
522 19422 : TIME_SECTION("KokkosClose", 1);
523 :
524 19422 : for (auto & system : systems)
525 12948 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
526 6474 : }
527 :
528 : // Clear
529 :
530 6474 : systems[number()].clearActiveResidualTags();
531 6474 : systems[number()].clearActiveMatrixTags();
532 :
533 19422 : for (auto & system : systems)
534 : {
535 12948 : system.clearActiveVariables();
536 12948 : system.clearActiveSolutionTags();
537 : }
538 6474 : }
|
