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 : #pragma once
11 :
12 : #include "KokkosTypes.h"
13 : #include "KokkosVector.h"
14 : #include "KokkosMatrix.h"
15 : #include "KokkosAssembly.h"
16 :
17 : #include "SystemBase.h"
18 :
19 : #include "libmesh/communicator.h"
20 :
21 : class MooseMesh;
22 : class SystemBase;
23 :
24 : namespace Moose
25 : {
26 : namespace Kokkos
27 : {
28 :
29 : class NodalBCBase;
30 :
31 : /**
32 : * The Kokkos system class. Each nonlinear and auxiliary system in MOOSE has a corresponding Kokkos
33 : * system.
34 : */
35 : class System : public MeshHolder, public AssemblyHolder
36 : {
37 : public:
38 : /**
39 : * Constructor
40 : * @param system The associated MOOSE system
41 : */
42 : System(SystemBase & system);
43 :
44 : #ifdef MOOSE_KOKKOS_SCOPE
45 : /**
46 : * Synchronize the active tagged vectors and matrices between host and device
47 : * @param dir Copy direction
48 : */
49 : void sync(const MemcpyKind dir);
50 : /**
51 : * Synchronize the specified tagged vectors between host and device
52 : * @param tags The vector tags
53 : * @param dir Copy direction
54 : */
55 : void sync(const std::set<TagID> & tags, const MemcpyKind dir);
56 : /**
57 : * Allocate the quadrature point solution vectors for active variable and tags and cache
58 : * quadrature point values
59 : */
60 : void reinit();
61 : /**
62 : * Set the active variables
63 : * @param vars The active MOOSE variables
64 : */
65 : void setActiveVariables(const std::set<MooseVariableFieldBase *> & vars);
66 : /**
67 : * Set the active solution tags
68 : * @param tags The active solution tags
69 : */
70 : void setActiveVariableTags(const std::set<TagID> & tags);
71 : /**
72 : * Set the active residual tags
73 : * @param tags The active residual tags
74 : */
75 : void setActiveResidualTags(const std::set<TagID> & tags);
76 : /**
77 : * Set the active matrix tags
78 : * @param vars The active matrix tags
79 : */
80 : void setActiveMatrixTags(const std::set<TagID> & tags);
81 : /**
82 : * Clear the cached active variables
83 : */
84 63574 : void clearActiveVariables() { _active_variables.destroy(); }
85 : /**
86 : * Clear the cached active solution tags
87 : */
88 63574 : void clearActiveVariableTags() { _active_variable_tags.destroy(); }
89 : /**
90 : * Clear the cached active residual tags
91 : */
92 27177 : void clearActiveResidualTags()
93 : {
94 27177 : _active_residual_tags.destroy();
95 27177 : _residual_tag_active = false;
96 27177 : }
97 : /**
98 : * Clear the cached active matrix tags
99 : */
100 4410 : void clearActiveMatrixTags()
101 : {
102 4410 : _active_matrix_tags.destroy();
103 4410 : _matrix_tag_active = false;
104 4410 : }
105 : /**
106 : * Get the MOOSE system
107 : * @returns The MOOSE system
108 : */
109 : const auto & getSystem() const { return _system; }
110 : /**
111 : * Get the libMesh DOF map
112 : * @returns The libMesh DOF map
113 : */
114 81553 : const auto & getDofMap() const { return _system.dofMap(); }
115 : /**
116 : * Get the libMesh communicator
117 : * @returns The libMesh communicator
118 : */
119 : const auto & getComm() const { return _comm; }
120 : /**
121 : * Get the list of local DOF indices to communicate
122 : * @returns The list of local DOF indices to communicate
123 : */
124 : const auto & getLocalCommList() const { return _local_comm_list; }
125 : /**
126 : * Get the list of ghost DOF indices to communicate
127 : * @returns The list of ghost DOF indices to communicate
128 : */
129 : const auto & getGhostCommList() const { return _ghost_comm_list; }
130 : /**
131 : * Get the sparisty pattern data
132 : * @returns The sparisty pattern data
133 : */
134 647 : const auto & getSparsity() const { return _sparsity; }
135 : /**
136 : * Get the list of off-diagonal coupled variable numbers of a variable
137 : * @param var The variable number
138 : * @returns The list of off-diagonal coupled variable numbers
139 : */
140 254340 : KOKKOS_FUNCTION const auto & getCoupling(unsigned int var) const { return _coupling[var]; }
141 : /**
142 : * Check whether a variable is active on a subdomain
143 : * @param var The variable number
144 : * @param subdomain The contiguous subdomain ID
145 : * @returns Whether the variable is active
146 : */
147 1270 : KOKKOS_FUNCTION bool isVariableActive(unsigned int var, ContiguousSubdomainID subdomain) const
148 : {
149 1270 : return _var_subdomain_active(var, subdomain);
150 : }
151 : /**
152 : * Check whether a residual tag is active
153 : * @param tag The residual tag
154 : * @returns Whether the residual tag is active
155 : */
156 10871116 : KOKKOS_FUNCTION bool isResidualTagActive(TagID tag) const { return _residual_tag_active[tag]; }
157 : /**
158 : * Check whether a matrix tag is active
159 : * @param tag The matrix tag
160 : * @returns Whether the matrix tag is active
161 : */
162 10065792 : KOKKOS_FUNCTION bool isMatrixTagActive(TagID tag) const { return _matrix_tag_active[tag]; }
163 : /**
164 : * Check whether a local DOF index is covered by a nodal BC
165 : * @param dof The local DOF index
166 : * @returns Whether the local DOF index is covered by a nodal BC
167 : */
168 14966722 : KOKKOS_FUNCTION bool hasNodalBC(dof_id_type dof) const
169 : {
170 14966722 : return _nbc_dof.isAlloc() && _nbc_dof[dof];
171 : }
172 : /**
173 : * Check whether a local DOF index is covered by a nodal BC for an extra residual tag
174 : * @param dof The local DOF index
175 : * @param tag The extra residual tag
176 : * @returns Whether the local DOF index is covered by a nodal BC
177 : */
178 687240 : KOKKOS_FUNCTION bool hasNodalBCResidualTag(dof_id_type dof, TagID tag) const
179 : {
180 687240 : return _nbc_residual_tag_dof[tag].isAlloc() && _nbc_residual_tag_dof[tag][dof];
181 : }
182 : /**
183 : * Check whether a local DOF index is associated with a nodal BC for an extra matrix tag
184 : * @param dof The local DOF index
185 : * @param tag The extra matrix tag
186 : * @returns Whether the local DOF index is covered by a nodal BC
187 : */
188 783200 : KOKKOS_FUNCTION bool hasNodalBCMatrixTag(dof_id_type dof, TagID tag) const
189 : {
190 783200 : return _nbc_matrix_tag_dof[tag].isAlloc() && _nbc_matrix_tag_dof[tag][dof];
191 : }
192 : /**
193 : * Get the FE type ID of a variable
194 : * @param var The variable number
195 : * @returns The FE type ID
196 : */
197 5833212 : KOKKOS_FUNCTION unsigned int getFETypeID(unsigned int var) const { return _var_fe_types[var]; }
198 : /**
199 : * Get the number of local DOFs
200 : * @returns The number of local DOFs
201 : */
202 5266 : KOKKOS_FUNCTION dof_id_type getNumLocalDofs() const { return _num_local_dofs; }
203 : /**
204 : * Get the number of ghost DOFs
205 : * @returns The number of ghost DOFs
206 : */
207 84640 : KOKKOS_FUNCTION dof_id_type getNumGhostDofs() const { return _num_ghost_dofs; }
208 : /**
209 : * Get the maximum number of DOFs per element of a variable
210 : * This number is local to each process
211 : * @param var The variable number
212 : * @returns The maximum number of DOFs per element
213 : */
214 : KOKKOS_FUNCTION unsigned int getMaxDofsPerElem(unsigned int var) const
215 : {
216 : return _max_dofs_per_elem[var];
217 : }
218 : /**
219 : * Get the local DOF index of a variable for an element
220 : * @param elem The contiguous element ID
221 : * @param i The element-local DOF index
222 : * @param var The variable number
223 : * @returns The local DOF index
224 : */
225 93837402 : KOKKOS_FUNCTION dof_id_type getElemLocalDofIndex(ContiguousElementID elem,
226 : unsigned int i,
227 : unsigned int var) const
228 : {
229 93837402 : return _local_elem_dof_index[var](elem, i);
230 : }
231 : /**
232 : * Get the local DOF index of a variable for a node
233 : * @param node The contiguous node ID
234 : * @param var The variable number
235 : * @returns The local DOF index
236 : */
237 2673822 : KOKKOS_FUNCTION dof_id_type getNodeLocalDofIndex(ContiguousNodeID node, unsigned int var) const
238 : {
239 2673822 : return _local_node_dof_index[var][node];
240 : }
241 : /**
242 : * Get the global DOF index of a variable for an element
243 : * @param elem The contiguous element ID
244 : * @param i The element-local DOF index
245 : * @param var The variable number
246 : * @returns The global DOF index
247 : */
248 6100534 : KOKKOS_FUNCTION dof_id_type getElemGlobalDofIndex(ContiguousElementID elem,
249 : unsigned int i,
250 : unsigned int var) const
251 : {
252 6100534 : return _local_to_global_dof_index[_local_elem_dof_index[var](elem, i)];
253 : }
254 : /**
255 : * Get the global DOF index of a variable for a node
256 : * @param node The contiguous node ID
257 : * @param var The variable number
258 : * @returns The global DOF index
259 : */
260 325004 : KOKKOS_FUNCTION dof_id_type getNodeGlobalDofIndex(ContiguousNodeID node, unsigned int var) const
261 : {
262 325004 : return _local_to_global_dof_index[_local_node_dof_index[var][node]];
263 : }
264 : /**
265 : * Get the global DOF index of a local DOF index
266 : * @param dof The local DOF index
267 : * @returns The global DOF index
268 : */
269 : KOKKOS_FUNCTION dof_id_type localToGlobalDofIndex(dof_id_type dof) const
270 : {
271 : return _local_to_global_dof_index[dof];
272 : }
273 : /**
274 : * Get whether a variable is defined on a node
275 : * @param node The contiguous node ID
276 : * @param var The variable number
277 : * @returns Whether the variable is defined on the node
278 : */
279 2018263 : KOKKOS_FUNCTION bool isNodalDefined(ContiguousNodeID node, unsigned int var) const
280 : {
281 2018263 : return _local_node_dof_index[var][node] != libMesh::DofObject::invalid_id;
282 : }
283 : /**
284 : * Get a tagged Kokkos vector
285 : * @param tag The vector tag
286 : * @returns The Kokkos vector
287 : */
288 : KOKKOS_FUNCTION auto & getVector(TagID tag) const { return _vectors[tag]; }
289 : /**
290 : * Get a tagged Kokkos matrix
291 : * @param tag The matrix tag
292 : * @returns The Kokkos matrix
293 : */
294 342492 : KOKKOS_FUNCTION auto & getMatrix(TagID tag) const { return _matrices[tag]; }
295 : /**
296 : * Get the DOF value of a tagged vector
297 : * @param dof The local DOF index
298 : * @param tag The vector tag
299 : * @returns The DOF value
300 : */
301 90112474 : KOKKOS_FUNCTION Real & getVectorDofValue(dof_id_type dof, TagID tag) const
302 : {
303 90112474 : return _vectors[tag][dof];
304 : }
305 : /**
306 : * Get the quadrature value of a variable from a tagged vector
307 : * @param info The element information object
308 : * @param qp The global quadrature point index
309 : * @param var The variable number
310 : * @param tag The vector tag
311 : * @returns The quadrature value
312 : */
313 : KOKKOS_FUNCTION Real &
314 36804496 : getVectorQpValue(ElementInfo info, dof_id_type qp, unsigned int var, TagID tag) const
315 : {
316 36804496 : return _qp_solutions[tag](info.subdomain, var)[qp];
317 : }
318 : /**
319 : * Get the quadrature gradient of a variable from a tagged vector
320 : * @param info The element information object
321 : * @param qp The global quadrature point index
322 : * @param var The variable number
323 : * @param tag The vector tag
324 : * @returns The quadrature gradient
325 : */
326 : KOKKOS_FUNCTION Real3 &
327 41123696 : getVectorQpGrad(ElementInfo info, dof_id_type qp, unsigned int var, TagID tag) const
328 : {
329 41123696 : return _qp_solutions_grad[tag](info.subdomain, var)[qp];
330 : }
331 : /**
332 : * Get the face quadrature value of a variable from a tagged vector
333 : * @param info The element information object
334 : * @param side The side index
335 : * @param qp The local quadrature point index
336 : * @param var The vriable number
337 : * @param tag The vector tag
338 : * @returns The face quadrature value
339 : */
340 : KOKKOS_FUNCTION Real getVectorQpValueFace(const ElementInfo info,
341 : const unsigned int side,
342 : const unsigned int qp,
343 : const unsigned int var,
344 : const TagID tag) const;
345 : /**
346 : * Get the face quadrature gradient of a variable from a tagged vector
347 : * @param info The element information object
348 : * @param side The side index
349 : * @param jacobian The inverse Jacobian matrix
350 : * @param qp The local quadrature point index
351 : * @param var The variable number
352 : * @param tag The vector tag
353 : * @returns The face quadrature gradient
354 : */
355 : KOKKOS_FUNCTION Real3 getVectorQpGradFace(const ElementInfo info,
356 : const unsigned int side,
357 : const Real33 jacobian,
358 : const unsigned int qp,
359 : const unsigned int var,
360 : const TagID tag) const;
361 : /**
362 : * Get an entry from a tagged matrix
363 : * @param row The local row index
364 : * @param col The global column index
365 : * @param tag The matrix tag
366 : * @returns The entry from the tagged matrix
367 : */
368 5720136 : KOKKOS_FUNCTION Real & getMatrixValue(dof_id_type row, dof_id_type col, TagID tag) const
369 : {
370 5720136 : return _matrices[tag](row, col);
371 : }
372 :
373 : /**
374 : * Kokkos function for caching variable values on element quadrature points
375 : */
376 : KOKKOS_FUNCTION void operator()(const ThreadID tid) const;
377 : #endif
378 :
379 : /**
380 : * CSR format sparsity data
381 : */
382 : struct Sparsity
383 : {
384 : Array<PetscInt> col_idx;
385 : Array<PetscInt> row_idx;
386 : Array<PetscInt> row_ptr;
387 : };
388 :
389 : private:
390 : /**
391 : * Setup variable data
392 : */
393 : void setupVariables();
394 : /**
395 : * Setup DOF data
396 : */
397 : void setupDofs();
398 : /**
399 : * Setup sparsity data
400 : */
401 : void setupSparsity();
402 : /**
403 : * Check if the DOFs are covered by nodal BCs
404 : */
405 : void checkNodalBCs();
406 : /**
407 : * Get the list of DOFs covered by a nodal BC
408 : * @param nbc The Kokkos nodal BC object
409 : * @param dofs The flag whether each DOF is covered by the nodal BC
410 : */
411 : void getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs);
412 :
413 : /**
414 : * Kokkos thread object
415 : */
416 : Thread _thread;
417 : /**
418 : * Reference of the MOOSE system
419 : */
420 : SystemBase & _system;
421 : /**
422 : * Reference of the MOOSE mesh
423 : */
424 : const MooseMesh & _mesh;
425 : /**
426 : * Reference of the libMesh DOF map
427 : */
428 : const libMesh::DofMap & _dof_map;
429 : /**
430 : * Reference of the libMesh communicator
431 : */
432 : const Parallel::Communicator & _comm;
433 : /**
434 : * Number of variables
435 : */
436 : const unsigned int _num_vars;
437 : /**
438 : * Number of local DOFs
439 : */
440 : const dof_id_type _num_local_dofs;
441 : /**
442 : * Number of ghost DOFs
443 : */
444 : const dof_id_type _num_ghost_dofs;
445 :
446 : /**
447 : * Kokkos vectors and matrices on device
448 : */
449 : ///@{
450 : Array<Vector> _vectors;
451 : Array<Matrix> _matrices;
452 : ///@}
453 : /**
454 : * Cached elemental quadrature values and gradients
455 : */
456 : ///@{
457 : Array<Array2D<Array<Real>>> _qp_solutions;
458 : Array<Array2D<Array<Real3>>> _qp_solutions_grad;
459 : ///@}
460 : /**
461 : * Local DOF index of each variable
462 : */
463 : ///@{
464 : Array<Array2D<dof_id_type>> _local_elem_dof_index;
465 : Array<Array<dof_id_type>> _local_node_dof_index;
466 : ///@}
467 : /**
468 : * Map from local DOF index to global DOF index
469 : */
470 : Array<dof_id_type> _local_to_global_dof_index;
471 : /**
472 : * Maximum number of DOFs per element for each variable
473 : */
474 : Array<unsigned int> _max_dofs_per_elem;
475 : /**
476 : * FE type ID of each variable
477 : */
478 : Array<unsigned int> _var_fe_types;
479 : /**
480 : * Whether each variable is active on subdomains
481 : */
482 : Array2D<bool> _var_subdomain_active;
483 : /**
484 : * Off-diagonal coupled variable numbers of each variable
485 : */
486 : Array<Array<unsigned int>> _coupling;
487 :
488 : /**
489 : * List of active variable numbers
490 : */
491 : Array<unsigned int> _active_variables;
492 : /**
493 : * List of active tags
494 : */
495 : ///@{
496 : Array<TagID> _active_variable_tags;
497 : Array<TagID> _active_residual_tags;
498 : Array<TagID> _active_matrix_tags;
499 : ///@}
500 : /**
501 : * Flag whether each tag is active
502 : */
503 : ///@{
504 : Array<bool> _residual_tag_active;
505 : Array<bool> _matrix_tag_active;
506 : ///@}
507 : /**
508 : * Flag whether each DOF is covered by a nodal BC
509 : */
510 : Array<bool> _nbc_dof;
511 : /**
512 : * Flag whether each DOF is covered by a nodal BC for extra residual tags
513 : */
514 : ///@{
515 : Array<Array<bool>> _nbc_residual_tag_dof;
516 : Array<Array<bool>> _nbc_matrix_tag_dof;
517 : ///@}
518 : /**
519 : * List of DOFs to send and receive
520 : */
521 : ///@{
522 : Array<Array<dof_id_type>> _local_comm_list;
523 : Array<Array<dof_id_type>> _ghost_comm_list;
524 : ///@}
525 :
526 : /**
527 : * Matrix sparsity pattern data
528 : */
529 : Sparsity _sparsity;
530 : };
531 :
532 : #ifdef MOOSE_KOKKOS_SCOPE
533 : KOKKOS_FUNCTION inline Real
534 176196 : System::getVectorQpValueFace(
535 : ElementInfo info, unsigned int side, unsigned int qp, unsigned int var, TagID tag) const
536 : {
537 176196 : auto fe = _var_fe_types[var];
538 176196 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
539 176196 : auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);
540 :
541 176196 : Real value = 0;
542 :
543 879148 : for (unsigned int i = 0; i < n_dofs; ++i)
544 702952 : value += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * phi(i, qp);
545 :
546 176196 : return value;
547 : }
548 : KOKKOS_FUNCTION inline Real3
549 0 : System::getVectorQpGradFace(ElementInfo info,
550 : unsigned int side,
551 : Real33 jacobian,
552 : unsigned int qp,
553 : unsigned int var,
554 : TagID tag) const
555 : {
556 0 : auto fe = _var_fe_types[var];
557 0 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
558 0 : auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);
559 :
560 0 : Real3 grad = 0;
561 :
562 0 : for (unsigned int i = 0; i < n_dofs; ++i)
563 0 : grad += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) *
564 0 : (jacobian * grad_phi(i, qp));
565 :
566 0 : return grad;
567 : }
568 : #endif
569 :
570 : /**
571 : * The Kokkos interface that holds the host reference of the Kokkos systems and copies it to device
572 : * during parallel dispatch.
573 : * Maintains synchronization between host and device Kokkos systems and provides access to the
574 : * appropriate Kokkos systems depending on the architecture.
575 : */
576 : class SystemHolder
577 : {
578 : public:
579 : /**
580 : * Constructor
581 : * @param systems The Kokkos systems
582 : */
583 5618 : SystemHolder(Array<System> & systems) : _systems_host(systems), _systems_device(systems) {}
584 : /**
585 : * Copy constructor
586 : */
587 297665 : SystemHolder(const SystemHolder & holder)
588 242544 : : _systems_host(holder._systems_host), _systems_device(holder._systems_host)
589 : {
590 297665 : }
591 :
592 : #ifdef MOOSE_KOKKOS_SCOPE
593 : /**
594 : * Get the const reference of the Kokkos systems
595 : * @returns The const reference of the Kokkos systems depending on the architecture this function
596 : * is being called on
597 : */
598 3628874 : KOKKOS_FUNCTION const Array<System> & kokkosSystems() const
599 : {
600 3628874 : KOKKOS_IF_ON_HOST(return _systems_host;)
601 :
602 3628874 : return _systems_device;
603 : }
604 : /**
605 : * Get the writeable host reference of the Kokkos systems
606 : * @returns The writeable host reference of the Kokkos systems
607 : */
608 1509 : Array<System> & kokkosSystems() { return _systems_host; }
609 : /**
610 : * Get the const reference of a Kokkos system
611 : * @param sys The system number
612 : * @returns The const reference of the Kokkos system depending on the architecture this function
613 : * is being called on
614 : */
615 21773176 : KOKKOS_FUNCTION const System & kokkosSystem(unsigned int sys) const
616 : {
617 21773176 : KOKKOS_IF_ON_HOST(return _systems_host[sys];)
618 :
619 21773176 : return _systems_device[sys];
620 : }
621 : /**
622 : * Get the writeable reference of a Kokkos system
623 : * @param sys The system number
624 : * @returns The writeable host reference of the Kokkos system
625 : */
626 3556 : System & kokkosSystem(unsigned int sys) { return _systems_host[sys]; }
627 : #endif
628 :
629 : private:
630 : /**
631 : * Host reference of the Kokkos systems
632 : */
633 : Array<System> & _systems_host;
634 : /**
635 : * Device copy of the Kokkos systems
636 : */
637 : const Array<System> _systems_device;
638 : };
639 :
640 : } // namespace Kokkos
641 : } // namespace Moose
|
