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 3138 : 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 1423 : std::shared_ptr<Moose::Kokkos::KernelBase> kernel =
27 1715 : _factory.create<Moose::Kokkos::KernelBase>(kernel_name, name, parameters);
28 3136 : _kokkos_kernels.addObject(kernel, 0);
29 3136 : postAddResidualObject(*kernel);
30 3136 : }
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 4008 : NonlinearSystemBase::addKokkosBoundaryCondition(const std::string & bc_name,
46 : const std::string & name,
47 : InputParameters & parameters)
48 : {
49 : // Create the object
50 1842 : std::shared_ptr<Moose::Kokkos::BoundaryCondition> bc =
51 2166 : _factory.create<Moose::Kokkos::BoundaryCondition>(bc_name, name, parameters, 0);
52 4008 : postAddResidualObject(*bc);
53 :
54 : // Active BoundaryIDs for the object
55 4008 : const std::set<BoundaryID> & boundary_ids = bc->boundaryIDs();
56 4008 : auto bc_var = dynamic_cast<const MooseVariableFieldBase *>(&bc->variable());
57 4008 : _vars[0].addBoundaryVar(boundary_ids, bc_var);
58 :
59 : // Cast to the various types of BCs
60 1842 : std::shared_ptr<Moose::Kokkos::NodalBCBase> nbc =
61 2166 : std::dynamic_pointer_cast<Moose::Kokkos::NodalBCBase>(bc);
62 1842 : std::shared_ptr<Moose::Kokkos::IntegratedBCBase> ibc =
63 2166 : std::dynamic_pointer_cast<Moose::Kokkos::IntegratedBCBase>(bc);
64 :
65 : // NodalBCBase
66 4008 : if (nbc)
67 : {
68 3441 : 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 3441 : _kokkos_nodal_bcs.addObject(nbc);
76 3441 : _vars[0].addBoundaryVars(boundary_ids, nbc->getCoupledVars());
77 :
78 : // DirichletBCs that are preset
79 3441 : if (nbc->preset())
80 3151 : _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 4008 : }
93 :
94 : void
95 4498 : NonlinearSystemBase::setKokkosInitialSolution()
96 : {
97 : // The preset Dirichlet BCs can have coupled variables from other systems, so we grab all systems
98 4498 : auto & systems = _fe_problem.getKokkosFESystems();
99 :
100 4498 : auto tag = _fe_problem.addVectorTag("parallel_solution", Moose::VECTOR_TAG_SOLUTION);
101 :
102 4498 : associateVectorToTag(solution(), tag);
103 :
104 8996 : std::set<TagID> needed_fe_var_vector_tags = {tag};
105 :
106 20288 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
107 15790 : _kokkos_preset_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
108 : boundary, needed_fe_var_vector_tags);
109 :
110 10512 : for (auto & system : systems.constructedEntries())
111 6014 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
112 :
113 4498 : systems.copyToDevice();
114 :
115 12127 : for (auto & nbc : _kokkos_preset_nodal_bcs.getActiveObjects())
116 7629 : std::static_pointer_cast<Moose::Kokkos::NodalBCBase>(nbc)->presetSolution(tag);
117 :
118 4498 : Kokkos::fence();
119 :
120 10512 : for (auto & system : systems.constructedEntries())
121 6014 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
122 :
123 4498 : disassociateVectorFromTag(solution(), tag);
124 4498 : }
125 :
126 : void
127 60524 : NonlinearSystemBase::computeKokkosResidual(const std::set<TagID> & tags)
128 : {
129 181572 : TIME_SECTION("computeKokkosResidual", 1);
130 :
131 : // Get warehouses
132 :
133 60524 : const auto & kernels = _kokkos_kernels.getVectorTagsObjectWarehouse(tags, 0);
134 60524 : const auto & nodal_kernels = _kokkos_nodal_kernels.getVectorTagsObjectWarehouse(tags, 0);
135 60524 : const auto & integrated_bcs = _kokkos_integrated_bcs.getVectorTagsObjectWarehouse(tags, 0);
136 :
137 60524 : if (!kernels.size() && !nodal_kernels.size() && !integrated_bcs.size())
138 0 : return;
139 :
140 : // Resolve dependencies
141 :
142 60524 : auto & systems = _fe_problem.getKokkosFESystems();
143 :
144 60524 : systems[number()].setActiveResidualTags(tags);
145 :
146 60524 : std::set<MooseVariableFieldBase *> needed_moose_vars;
147 60524 : std::set<TagID> needed_fe_var_vector_tags;
148 60524 : std::unordered_set<unsigned int> needed_mat_props;
149 :
150 324799 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
151 264275 : needed_fe_var_vector_tags.insert(tag._id);
152 :
153 60524 : kernels.updateVariableDependency(needed_moose_vars);
154 60524 : kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
155 60524 : kernels.updateMatPropDependency(needed_mat_props);
156 :
157 60524 : nodal_kernels.updateVariableDependency(needed_moose_vars);
158 60524 : nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
159 :
160 60524 : integrated_bcs.updateVariableDependency(needed_moose_vars);
161 60524 : integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
162 60524 : integrated_bcs.updateMatPropDependency(needed_mat_props);
163 :
164 60524 : if (needed_mat_props.size())
165 : {
166 14417 : _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
167 14417 : _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 141441 : for (auto & system : systems.constructedEntries())
175 : {
176 80917 : system.setActiveVariables(needed_moose_vars);
177 80917 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
178 :
179 : {
180 242751 : TIME_SECTION("KokkosCopy", 1);
181 80917 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
182 80917 : }
183 : {
184 242751 : TIME_SECTION("KokkosReinit", 1);
185 80917 : system.reinit();
186 80917 : }
187 : }
188 :
189 60524 : systems.copyToDevice();
190 :
191 : {
192 181572 : TIME_SECTION("KokkosMaterial", 1);
193 :
194 : // Compute material properties
195 :
196 60524 : if (needed_mat_props.size())
197 : {
198 14417 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
199 14417 : _fe_problem.reinitKokkosMaterials();
200 : }
201 60524 : }
202 :
203 : {
204 181572 : TIME_SECTION("KokkosKernel", 1);
205 :
206 : // Compute kernels
207 :
208 176656 : for (auto & kernel : kernels.getActiveObjects())
209 116132 : kernel->computeResidual();
210 :
211 92413 : for (auto & nodal_kernel : nodal_kernels.getActiveObjects())
212 31889 : nodal_kernel->computeResidual();
213 :
214 86617 : for (auto & ibc : integrated_bcs.getActiveObjects())
215 26093 : ibc->computeResidual();
216 60524 : }
217 :
218 : // Close and restore vectors
219 :
220 : {
221 181572 : TIME_SECTION("KokkosClose", 1);
222 :
223 141441 : for (auto & system : systems.constructedEntries())
224 80917 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
225 60524 : }
226 :
227 : // Clear
228 :
229 60524 : systems[number()].clearActiveResidualTags();
230 :
231 141441 : for (auto & system : systems.constructedEntries())
232 : {
233 80917 : system.clearActiveVariables();
234 80917 : system.clearActiveSolutionTags();
235 : }
236 60524 : }
237 :
238 : void
239 66998 : NonlinearSystemBase::computeKokkosNodalBCsResidual(const std::set<TagID> & tags)
240 : {
241 200994 : TIME_SECTION("computeKokkosNodalBCsResidual", 1);
242 :
243 : // Get warehouses
244 :
245 66998 : const auto & nodal_bcs = _kokkos_nodal_bcs.getVectorTagsObjectWarehouse(tags, 0);
246 :
247 66998 : if (!nodal_bcs.hasActiveObjects())
248 1575 : return;
249 :
250 : // Resolve dependencies
251 :
252 65423 : auto & systems = _fe_problem.getKokkosFESystems();
253 :
254 65423 : systems[number()].setActiveResidualTags(tags);
255 :
256 65423 : std::set<TagID> needed_fe_var_vector_tags;
257 :
258 350817 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
259 285394 : needed_fe_var_vector_tags.insert(tag._id);
260 :
261 65423 : nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
262 :
263 : // Copy solution vectors and residuals
264 :
265 149891 : for (auto & system : systems.constructedEntries())
266 : {
267 84468 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
268 : {
269 253404 : TIME_SECTION("KokkosCopy", 1);
270 84468 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
271 84468 : }
272 : }
273 :
274 65423 : systems.copyToDevice();
275 :
276 : {
277 196269 : TIME_SECTION("KokkosKernel", 1);
278 :
279 : // Compute kernels
280 :
281 179055 : for (auto & nbc : nodal_bcs.getActiveObjects())
282 113632 : nbc->computeResidual();
283 65423 : }
284 :
285 : // Close and restore vectors
286 :
287 : {
288 196269 : TIME_SECTION("KokkosClose", 1);
289 :
290 149891 : for (auto & system : systems.constructedEntries())
291 84468 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
292 65423 : }
293 :
294 : // Clear
295 :
296 65423 : systems[number()].clearActiveResidualTags();
297 :
298 149891 : for (auto & system : systems.constructedEntries())
299 84468 : system.clearActiveSolutionTags();
300 66998 : }
301 :
302 : void
303 9431 : NonlinearSystemBase::computeKokkosJacobian(const std::set<TagID> & tags)
304 : {
305 28293 : TIME_SECTION("computeKokkosJacobian", 1);
306 :
307 : // Get warehouses
308 :
309 9431 : const auto & kernels = _kokkos_kernels.getMatrixTagsObjectWarehouse(tags, 0);
310 9431 : const auto & nodal_kernels = _kokkos_nodal_kernels.getMatrixTagsObjectWarehouse(tags, 0);
311 9431 : const auto & integrated_bcs = _kokkos_integrated_bcs.getMatrixTagsObjectWarehouse(tags, 0);
312 9431 : const auto & nodal_bcs = _kokkos_nodal_bcs.getMatrixTagsObjectWarehouse(tags, 0);
313 :
314 9431 : if (!kernels.hasActiveObjects() && !nodal_kernels.hasActiveObjects() &&
315 5572 : !integrated_bcs.hasActiveObjects() && !nodal_bcs.hasActiveObjects())
316 0 : return;
317 :
318 : // Resolve dependencies
319 :
320 9431 : auto & systems = _fe_problem.getKokkosFESystems();
321 :
322 9431 : systems[number()].setActiveMatrixTags(tags);
323 :
324 9431 : std::set<MooseVariableFieldBase *> needed_moose_vars;
325 9431 : std::set<TagID> needed_fe_var_vector_tags;
326 9431 : std::unordered_set<unsigned int> needed_mat_props;
327 :
328 51606 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
329 42175 : needed_fe_var_vector_tags.insert(tag._id);
330 :
331 9431 : kernels.updateVariableDependency(needed_moose_vars);
332 9431 : kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
333 9431 : kernels.updateMatPropDependency(needed_mat_props);
334 :
335 9431 : nodal_kernels.updateVariableDependency(needed_moose_vars);
336 9431 : nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
337 :
338 9431 : integrated_bcs.updateVariableDependency(needed_moose_vars);
339 9431 : integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
340 9431 : integrated_bcs.updateMatPropDependency(needed_mat_props);
341 :
342 9431 : nodal_bcs.updateVariableDependency(needed_moose_vars);
343 9431 : nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
344 :
345 9431 : if (needed_mat_props.size())
346 : {
347 1734 : _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
348 1734 : _fe_problem.getKokkosMaterialsWarehouse().updateFEVariableCoupledVectorTagDependency(
349 : needed_fe_var_vector_tags);
350 : }
351 :
352 : // Copy solution vectors, initialize matrices, and cache variable values at element quadature
353 : // points
354 :
355 21966 : for (auto & system : systems.constructedEntries())
356 : {
357 12535 : system.setActiveVariables(needed_moose_vars);
358 12535 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
359 :
360 : {
361 37605 : TIME_SECTION("KokkosCopy", 1);
362 12535 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
363 12535 : }
364 : {
365 37605 : TIME_SECTION("KokkosReinit", 1);
366 12535 : system.reinit();
367 12535 : }
368 : }
369 :
370 9431 : systems.copyToDevice();
371 :
372 : {
373 28293 : TIME_SECTION("KokkosMaterial", 1);
374 :
375 : // Compute material properties
376 :
377 9431 : if (needed_mat_props.size())
378 : {
379 1734 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
380 1734 : _fe_problem.reinitKokkosMaterials();
381 : }
382 9431 : }
383 :
384 : {
385 28293 : TIME_SECTION("KokkosKernel", 1);
386 :
387 : // Compute kernels
388 :
389 29165 : for (auto & kernel : kernels.getActiveObjects())
390 19734 : kernel->computeJacobian();
391 :
392 17625 : for (auto & nodal_kernel : nodal_kernels.getActiveObjects())
393 8194 : nodal_kernel->computeJacobian();
394 :
395 12446 : for (auto & ibc : integrated_bcs.getActiveObjects())
396 3015 : ibc->computeJacobian();
397 :
398 25137 : for (auto & nbc : nodal_bcs.getActiveObjects())
399 15706 : nbc->computeJacobian();
400 9431 : }
401 :
402 : // Close and restore vectors and matrices
403 :
404 : {
405 28293 : TIME_SECTION("KokkosClose", 1);
406 :
407 21966 : for (auto & system : systems.constructedEntries())
408 12535 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
409 9431 : }
410 :
411 : // Clear
412 :
413 9431 : systems[number()].clearActiveMatrixTags();
414 :
415 21966 : for (auto & system : systems.constructedEntries())
416 : {
417 12535 : system.clearActiveVariables();
418 12535 : system.clearActiveSolutionTags();
419 : }
420 9431 : }
421 :
422 : void
423 6474 : NonlinearSystemBase::computeKokkosResidualAndJacobian(const std::set<TagID> & vector_tags,
424 : const std::set<TagID> & matrix_tags)
425 : {
426 19422 : TIME_SECTION("computeKokkosResidualAndJacobian", 1);
427 :
428 6474 : if (!_kokkos_kernels.hasActiveObjects() && !_kokkos_nodal_kernels.hasActiveObjects() &&
429 3254 : !_kokkos_integrated_bcs.hasActiveObjects() && !_kokkos_nodal_bcs.hasActiveObjects())
430 0 : return;
431 :
432 : // Resolve dependencies
433 :
434 6474 : auto & systems = _fe_problem.getKokkosFESystems();
435 :
436 6474 : systems[number()].setActiveResidualTags(vector_tags);
437 6474 : systems[number()].setActiveMatrixTags(matrix_tags);
438 :
439 6474 : std::set<MooseVariableFieldBase *> needed_moose_vars;
440 6474 : std::set<TagID> needed_fe_var_vector_tags;
441 6474 : std::unordered_set<unsigned int> needed_mat_props;
442 :
443 34650 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
444 28176 : needed_fe_var_vector_tags.insert(tag._id);
445 :
446 6474 : _kokkos_kernels.updateVariableDependency(needed_moose_vars);
447 6474 : _kokkos_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
448 6474 : _kokkos_kernels.updateMatPropDependency(needed_mat_props);
449 :
450 6474 : _kokkos_nodal_kernels.updateVariableDependency(needed_moose_vars);
451 6474 : _kokkos_nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
452 :
453 6474 : _kokkos_integrated_bcs.updateVariableDependency(needed_moose_vars);
454 6474 : _kokkos_integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
455 6474 : _kokkos_integrated_bcs.updateMatPropDependency(needed_mat_props);
456 :
457 6474 : _kokkos_nodal_bcs.updateVariableDependency(needed_moose_vars);
458 6474 : _kokkos_nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
459 :
460 6474 : if (needed_mat_props.size())
461 : {
462 0 : _fe_problem.getKokkosMaterialsWarehouse().updateVariableDependency(needed_moose_vars);
463 0 : _fe_problem.getKokkosMaterialsWarehouse().updateFEVariableCoupledVectorTagDependency(
464 : needed_fe_var_vector_tags);
465 : }
466 :
467 : // Copy solution vectors and residuals, initialize matrices, and cache variable values at element
468 : // quadature points
469 :
470 12993 : for (auto & system : systems.constructedEntries())
471 : {
472 6519 : system.setActiveVariables(needed_moose_vars);
473 6519 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
474 :
475 : {
476 19557 : TIME_SECTION("KokkosCopy", 1);
477 6519 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
478 6519 : }
479 : {
480 19557 : TIME_SECTION("KokkosReinit", 1);
481 6519 : system.reinit();
482 6519 : }
483 : }
484 :
485 6474 : systems.copyToDevice();
486 :
487 : {
488 19422 : TIME_SECTION("KokkosMaterial", 1);
489 :
490 : // Compute material properties
491 :
492 6474 : if (needed_mat_props.size())
493 : {
494 0 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
495 0 : _fe_problem.reinitKokkosMaterials();
496 : }
497 6474 : }
498 :
499 : {
500 19422 : TIME_SECTION("KokkosKernel", 1);
501 :
502 : // Compute kernels
503 :
504 22881 : for (auto & kernel : _kokkos_kernels.getActiveObjects())
505 16407 : kernel->computeResidualAndJacobian();
506 :
507 6474 : for (auto & nodal_kernel : _kokkos_nodal_kernels.getActiveObjects())
508 0 : nodal_kernel->computeResidualAndJacobian();
509 :
510 7686 : for (auto & ibc : _kokkos_integrated_bcs.getActiveObjects())
511 1212 : ibc->computeResidualAndJacobian();
512 :
513 : /// Nodal BC residuals are computed separately
514 20484 : for (auto & nbc : _kokkos_nodal_bcs.getActiveObjects())
515 14010 : nbc->computeJacobian();
516 6474 : }
517 :
518 : // Close and restore vectors and matrices
519 :
520 : {
521 19422 : TIME_SECTION("KokkosClose", 1);
522 :
523 12993 : for (auto & system : systems.constructedEntries())
524 6519 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
525 6474 : }
526 :
527 : // Clear
528 :
529 6474 : systems[number()].clearActiveResidualTags();
530 6474 : systems[number()].clearActiveMatrixTags();
531 :
532 12993 : for (auto & system : systems.constructedEntries())
533 : {
534 6519 : system.clearActiveVariables();
535 6519 : system.clearActiveSolutionTags();
536 : }
537 6474 : }
|
