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 3663 : 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 1675 : std::shared_ptr<Moose::Kokkos::KernelBase> kernel =
27 1988 : _factory.create<Moose::Kokkos::KernelBase>(kernel_name, name, parameters);
28 3661 : _kokkos_kernels.addObject(kernel, 0);
29 3661 : postAddResidualObject(*kernel);
30 3661 : }
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 4340 : NonlinearSystemBase::addKokkosBoundaryCondition(const std::string & bc_name,
46 : const std::string & name,
47 : InputParameters & parameters)
48 : {
49 : // Create the object
50 2000 : std::shared_ptr<Moose::Kokkos::BoundaryCondition> bc =
51 2340 : _factory.create<Moose::Kokkos::BoundaryCondition>(bc_name, name, parameters, 0);
52 4340 : postAddResidualObject(*bc);
53 :
54 : // Active BoundaryIDs for the object
55 4340 : const std::set<BoundaryID> & boundary_ids = bc->boundaryIDs();
56 4340 : auto bc_var = dynamic_cast<const MooseVariableFieldBase *>(&bc->variable());
57 4340 : _vars[0].addBoundaryVar(boundary_ids, bc_var);
58 :
59 : // Cast to the various types of BCs
60 2000 : std::shared_ptr<Moose::Kokkos::NodalBCBase> nbc =
61 2340 : std::dynamic_pointer_cast<Moose::Kokkos::NodalBCBase>(bc);
62 2000 : std::shared_ptr<Moose::Kokkos::IntegratedBCBase> ibc =
63 2340 : std::dynamic_pointer_cast<Moose::Kokkos::IntegratedBCBase>(bc);
64 :
65 : // NodalBCBase
66 4340 : if (nbc)
67 : {
68 3739 : 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 3739 : _kokkos_nodal_bcs.addObject(nbc);
76 3739 : _vars[0].addBoundaryVars(boundary_ids, nbc->getCoupledVars());
77 :
78 : // DirichletBCs that are preset
79 3739 : if (nbc->preset())
80 3267 : _kokkos_preset_nodal_bcs.addObject(nbc);
81 : }
82 :
83 : // IntegratedBCBase
84 601 : else if (ibc)
85 : {
86 601 : _kokkos_integrated_bcs.addObject(ibc);
87 601 : _vars[0].addBoundaryVars(boundary_ids, ibc->getCoupledVars());
88 : }
89 :
90 : else
91 0 : mooseError("Unknown Kokkos BoundaryCondition type for object named ", bc->name());
92 4340 : }
93 :
94 : void
95 4598 : NonlinearSystemBase::setKokkosInitialSolution()
96 : {
97 : // The preset Dirichlet BCs can have coupled variables from other systems, so we grab all systems
98 4598 : auto & systems = _fe_problem.getKokkosFESystems();
99 :
100 4598 : auto tag = _fe_problem.addVectorTag("parallel_solution", Moose::VECTOR_TAG_SOLUTION);
101 :
102 4598 : associateVectorToTag(solution(), tag);
103 :
104 9196 : std::set<TagID> needed_fe_var_vector_tags = {tag};
105 :
106 20788 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
107 16190 : _kokkos_preset_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
108 : boundary, needed_fe_var_vector_tags);
109 :
110 10712 : for (auto & system : systems.constructedEntries())
111 6114 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
112 :
113 4598 : systems.copyToDevice();
114 :
115 12366 : for (auto & nbc : _kokkos_preset_nodal_bcs.getActiveObjects())
116 7768 : std::static_pointer_cast<Moose::Kokkos::NodalBCBase>(nbc)->presetSolution(tag);
117 :
118 4598 : Kokkos::fence();
119 :
120 10712 : for (auto & system : systems.constructedEntries())
121 6114 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
122 :
123 4598 : disassociateVectorFromTag(solution(), tag);
124 4598 : }
125 :
126 : void
127 67208 : NonlinearSystemBase::computeKokkosResidual(const std::set<TagID> & tags)
128 : {
129 201624 : TIME_SECTION("computeKokkosResidual", 1);
130 :
131 : // Get warehouses
132 :
133 67208 : const auto & kernels = _kokkos_kernels.getVectorTagsObjectWarehouse(tags, 0);
134 67208 : const auto & nodal_kernels = _kokkos_nodal_kernels.getVectorTagsObjectWarehouse(tags, 0);
135 67208 : const auto & integrated_bcs = _kokkos_integrated_bcs.getVectorTagsObjectWarehouse(tags, 0);
136 :
137 67208 : if (!kernels.size() && !nodal_kernels.size() && !integrated_bcs.size())
138 0 : return;
139 :
140 : // Resolve dependencies
141 :
142 67208 : auto & systems = _fe_problem.getKokkosFESystems();
143 :
144 67208 : systems[number()].setActiveResidualTags(tags);
145 :
146 67208 : std::set<MooseVariableFieldBase *> needed_moose_vars;
147 67208 : std::set<TagID> needed_fe_var_vector_tags;
148 67208 : std::unordered_set<unsigned int> needed_mat_props;
149 :
150 353480 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
151 286272 : needed_fe_var_vector_tags.insert(tag._id);
152 :
153 67208 : kernels.updateVariableDependency(needed_moose_vars);
154 67208 : kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
155 67208 : kernels.updateMatPropDependency(needed_mat_props);
156 :
157 67208 : nodal_kernels.updateVariableDependency(needed_moose_vars);
158 67208 : nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
159 :
160 67208 : integrated_bcs.updateVariableDependency(needed_moose_vars);
161 67208 : integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
162 67208 : integrated_bcs.updateMatPropDependency(needed_mat_props);
163 :
164 67208 : 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 154822 : for (auto & system : systems.constructedEntries())
175 : {
176 87614 : system.setActiveVariables(needed_moose_vars);
177 87614 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
178 :
179 : {
180 262842 : TIME_SECTION("KokkosCopy", 1);
181 87614 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
182 87614 : }
183 : {
184 262842 : TIME_SECTION("KokkosReinit", 1);
185 87614 : system.reinit();
186 87614 : }
187 : }
188 :
189 67208 : systems.copyToDevice();
190 :
191 : {
192 201624 : TIME_SECTION("KokkosMaterial", 1);
193 :
194 : // Compute material properties
195 :
196 67208 : if (needed_mat_props.size())
197 : {
198 14417 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
199 14417 : _fe_problem.reinitKokkosMaterials();
200 : }
201 67208 : }
202 :
203 : {
204 201624 : TIME_SECTION("KokkosKernel", 1);
205 :
206 : // Compute kernels
207 :
208 198431 : for (auto & kernel : kernels.getActiveObjects())
209 131223 : kernel->computeResidual();
210 :
211 99097 : for (auto & nodal_kernel : nodal_kernels.getActiveObjects())
212 31889 : nodal_kernel->computeResidual();
213 :
214 93526 : for (auto & ibc : integrated_bcs.getActiveObjects())
215 26318 : ibc->computeResidual();
216 67208 : }
217 :
218 : // Close and restore vectors
219 :
220 : {
221 201624 : TIME_SECTION("KokkosClose", 1);
222 :
223 154822 : for (auto & system : systems.constructedEntries())
224 87614 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
225 67208 : }
226 :
227 : // Clear
228 :
229 67208 : systems[number()].clearActiveResidualTags();
230 :
231 154822 : for (auto & system : systems.constructedEntries())
232 : {
233 87614 : system.clearActiveVariables();
234 87614 : system.clearActiveSolutionTags();
235 : }
236 67208 : }
237 :
238 : void
239 73682 : NonlinearSystemBase::computeKokkosNodalBCsResidual(const std::set<TagID> & tags)
240 : {
241 221046 : TIME_SECTION("computeKokkosNodalBCsResidual", 1);
242 :
243 : // Get warehouses
244 :
245 73682 : const auto & nodal_bcs = _kokkos_nodal_bcs.getVectorTagsObjectWarehouse(tags, 0);
246 :
247 73682 : if (!nodal_bcs.hasActiveObjects())
248 6522 : return;
249 :
250 : // Resolve dependencies
251 :
252 67160 : auto & systems = _fe_problem.getKokkosFESystems();
253 :
254 67160 : systems[number()].setActiveResidualTags(tags);
255 :
256 67160 : std::set<TagID> needed_fe_var_vector_tags;
257 :
258 355339 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
259 288179 : needed_fe_var_vector_tags.insert(tag._id);
260 :
261 67160 : nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
262 :
263 : // Copy solution vectors and residuals
264 :
265 153378 : for (auto & system : systems.constructedEntries())
266 : {
267 86218 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
268 : {
269 258654 : TIME_SECTION("KokkosCopy", 1);
270 86218 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
271 86218 : }
272 : }
273 :
274 67160 : systems.copyToDevice();
275 :
276 : {
277 201480 : TIME_SECTION("KokkosKernel", 1);
278 :
279 : // Compute kernels
280 :
281 183876 : for (auto & nbc : nodal_bcs.getActiveObjects())
282 116716 : nbc->computeResidual();
283 67160 : }
284 :
285 : // Close and restore vectors
286 :
287 : {
288 201480 : TIME_SECTION("KokkosClose", 1);
289 :
290 153378 : for (auto & system : systems.constructedEntries())
291 86218 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
292 67160 : }
293 :
294 : // Clear
295 :
296 67160 : systems[number()].clearActiveResidualTags();
297 :
298 153378 : for (auto & system : systems.constructedEntries())
299 86218 : system.clearActiveSolutionTags();
300 73682 : }
301 :
302 : void
303 10381 : NonlinearSystemBase::computeKokkosJacobian(const std::set<TagID> & tags)
304 : {
305 31143 : TIME_SECTION("computeKokkosJacobian", 1);
306 :
307 : // Get warehouses
308 :
309 10381 : const auto & kernels = _kokkos_kernels.getMatrixTagsObjectWarehouse(tags, 0);
310 10381 : const auto & nodal_kernels = _kokkos_nodal_kernels.getMatrixTagsObjectWarehouse(tags, 0);
311 10381 : const auto & integrated_bcs = _kokkos_integrated_bcs.getMatrixTagsObjectWarehouse(tags, 0);
312 10381 : const auto & nodal_bcs = _kokkos_nodal_bcs.getMatrixTagsObjectWarehouse(tags, 0);
313 :
314 10381 : if (!kernels.hasActiveObjects() && !nodal_kernels.hasActiveObjects() &&
315 6080 : !integrated_bcs.hasActiveObjects() && !nodal_bcs.hasActiveObjects())
316 0 : return;
317 :
318 : // Resolve dependencies
319 :
320 10381 : auto & systems = _fe_problem.getKokkosFESystems();
321 :
322 10381 : systems[number()].setActiveMatrixTags(tags);
323 :
324 10381 : std::set<MooseVariableFieldBase *> needed_moose_vars;
325 10381 : std::set<TagID> needed_fe_var_vector_tags;
326 10381 : std::unordered_set<unsigned int> needed_mat_props;
327 :
328 56088 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
329 45707 : needed_fe_var_vector_tags.insert(tag._id);
330 :
331 10381 : kernels.updateVariableDependency(needed_moose_vars);
332 10381 : kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
333 10381 : kernels.updateMatPropDependency(needed_mat_props);
334 :
335 10381 : nodal_kernels.updateVariableDependency(needed_moose_vars);
336 10381 : nodal_kernels.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
337 :
338 10381 : integrated_bcs.updateVariableDependency(needed_moose_vars);
339 10381 : integrated_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
340 10381 : integrated_bcs.updateMatPropDependency(needed_mat_props);
341 :
342 10381 : nodal_bcs.updateVariableDependency(needed_moose_vars);
343 10381 : nodal_bcs.updateFEVariableCoupledVectorTagDependency(needed_fe_var_vector_tags);
344 :
345 10381 : 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 23870 : for (auto & system : systems.constructedEntries())
356 : {
357 13489 : system.setActiveVariables(needed_moose_vars);
358 13489 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
359 :
360 : {
361 40467 : TIME_SECTION("KokkosCopy", 1);
362 13489 : system.sync(Moose::Kokkos::MemcpyType::HOST_TO_DEVICE);
363 13489 : }
364 : {
365 40467 : TIME_SECTION("KokkosReinit", 1);
366 13489 : system.reinit();
367 13489 : }
368 : }
369 :
370 10381 : systems.copyToDevice();
371 :
372 : {
373 31143 : TIME_SECTION("KokkosMaterial", 1);
374 :
375 : // Compute material properties
376 :
377 10381 : if (needed_mat_props.size())
378 : {
379 1734 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
380 1734 : _fe_problem.reinitKokkosMaterials();
381 : }
382 10381 : }
383 :
384 : {
385 31143 : TIME_SECTION("KokkosKernel", 1);
386 :
387 : // Compute kernels
388 :
389 32749 : for (auto & kernel : kernels.getActiveObjects())
390 22368 : kernel->computeJacobian();
391 :
392 18575 : for (auto & nodal_kernel : nodal_kernels.getActiveObjects())
393 8194 : nodal_kernel->computeJacobian();
394 :
395 13551 : for (auto & ibc : integrated_bcs.getActiveObjects())
396 3170 : ibc->computeJacobian();
397 :
398 27128 : for (auto & nbc : nodal_bcs.getActiveObjects())
399 16747 : nbc->computeJacobian();
400 10381 : }
401 :
402 : // Close and restore vectors and matrices
403 :
404 : {
405 31143 : TIME_SECTION("KokkosClose", 1);
406 :
407 23870 : for (auto & system : systems.constructedEntries())
408 13489 : system.sync(Moose::Kokkos::MemcpyType::DEVICE_TO_HOST);
409 10381 : }
410 :
411 : // Clear
412 :
413 10381 : systems[number()].clearActiveMatrixTags();
414 :
415 23870 : for (auto & system : systems.constructedEntries())
416 : {
417 13489 : system.clearActiveVariables();
418 13489 : system.clearActiveSolutionTags();
419 : }
420 10381 : }
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 : }
|
