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 "KokkosSystem.h"
13 : #include "KokkosAssembly.h"
14 :
15 : class MooseMesh;
16 : class SystemBase;
17 :
18 : namespace Moose::Kokkos
19 : {
20 :
21 : class NodalBCBase;
22 :
23 : /**
24 : * The Kokkos FE system class. Each system in MOOSE with FE variables has a corresponding Kokkos FE
25 : * system.
26 : */
27 : class FESystem : public System, public AssemblyHolder
28 : {
29 : public:
30 : /**
31 : * Constructor for standalone use (pure FE simulations)
32 : * @param system The associated MOOSE system
33 : */
34 : FESystem(SystemBase & system);
35 :
36 : /**
37 : * Constructor for mixed FE+FV simulations, sharing device memory with an existing System
38 : * @param base The existing Kokkos System whose device memory to share
39 : * @param system The associated MOOSE system
40 : */
41 : FESystem(System & base, SystemBase & system);
42 :
43 : #ifdef MOOSE_KOKKOS_SCOPE
44 : ///@}
45 : /**
46 : * Allocate the quadrature point vectors for active variable and tags and cache
47 : * quadrature point values
48 : */
49 : void reinit();
50 :
51 : /**
52 : * Get the list of off-diagonal coupled variable numbers of a variable
53 : * @param var The variable number
54 : * @returns The list of off-diagonal coupled variable numbers
55 : */
56 256182 : KOKKOS_FUNCTION const auto & getCoupling(unsigned int var) const { return _coupling[var]; }
57 :
58 : /**
59 : * Check whether a local DOF index is associated with a nodal BC for an extra matrix tag
60 : * @param dof The local DOF index
61 : * @param tag The extra matrix tag
62 : * @returns Whether the local DOF index is covered by a nodal BC
63 : */
64 11463947 : KOKKOS_FUNCTION bool hasNodalBCMatrixTag(dof_id_type dof, TagID tag) const
65 : {
66 11463947 : return _nbc_matrix_tag_dof[tag].isAlloc() && _nbc_matrix_tag_dof[tag][dof];
67 : }
68 :
69 : /**
70 : * Get the FE type ID of a variable
71 : * @param var The variable number
72 : * @returns The FE type ID
73 : */
74 17433154 : KOKKOS_FUNCTION unsigned int getFETypeID(unsigned int var) const { return _var_fe_types[var]; }
75 :
76 : /**
77 : * Get the local DOF index of a variable for a node
78 : * @param node The contiguous node ID
79 : * @param i The node-local DOF index
80 : * @param var The variable number
81 : * @returns The local DOF index
82 : */
83 6826619 : KOKKOS_FUNCTION dof_id_type getNodeLocalDofIndex(ContiguousNodeID node,
84 : unsigned int i,
85 : unsigned int var) const
86 : {
87 6826619 : return _local_node_dof_index[var][node] + i;
88 : }
89 :
90 : /**
91 : * Get the global DOF index of a variable for a node
92 : * @param node The contiguous node ID
93 : * @param var The variable number
94 : * @returns The global DOF index
95 : */
96 342489 : KOKKOS_FUNCTION dof_id_type getNodeGlobalDofIndex(ContiguousNodeID node, unsigned int var) const
97 : {
98 342489 : return _local_to_global_dof_index[_local_node_dof_index[var][node]];
99 : }
100 :
101 : /**
102 : * Get whether a variable is defined on a node
103 : * @param node The contiguous node ID
104 : * @param var The variable number
105 : * @returns Whether the variable is defined on the node
106 : */
107 4919371 : KOKKOS_FUNCTION bool isNodalDefined(ContiguousNodeID node, unsigned int var) const
108 : {
109 4919371 : return _local_node_dof_index[var][node] != libMesh::DofObject::invalid_id;
110 : }
111 :
112 : /**
113 : * Get the DOF value of a tagged vector for automatic differentiation (AD)
114 : * @param dof The local DOF index
115 : * @param tag The vector tag
116 : * @param seed The derivative seed
117 : * @returns The DOF AD value with optional seed derivative
118 : */
119 : KOKKOS_FUNCTION ADReal getVectorDofADValue(const dof_id_type dof,
120 : const TagID tag,
121 : const Real seed) const;
122 : /**
123 : * Get the quadrature point value of a variable from a tagged vector
124 : * @param info The element information object
125 : * @param qp The subdomain-local flattened quadrature point index
126 : * @param var The variable number
127 : * @param tag The vector tag
128 : * @returns The quadrature value
129 : */
130 78077828 : KOKKOS_FUNCTION Real & getVectorQpValue(const ElementInfo info,
131 : const dof_id_type qp,
132 : const unsigned int var,
133 : const TagID tag) const
134 : {
135 78077828 : return _qp_solutions[tag](info.subdomain, var)[qp];
136 : }
137 : /**
138 : * Get the quadrature point value of a variable from a tagged vector for automatic differentiation
139 : * (AD)
140 : * @param info The element information object
141 : * @param offset The element's offset into the subdomain-local flattened quadrature point index
142 : * @param qp The local quadrature point index
143 : * @param var The variable number
144 : * @param tag The vector tag
145 : * @param seed The derivative seed
146 : * @returns The quadrature AD value
147 : */
148 : KOKKOS_FUNCTION ADReal getVectorQpADValue(const ElementInfo info,
149 : const dof_id_type offset,
150 : const dof_id_type qp,
151 : const unsigned int var,
152 : const TagID tag,
153 : const Real seed) const;
154 : /**
155 : * Get the quadrature point gradient of a variable from a tagged vector
156 : * @param info The element information object
157 : * @param qp The subdomain-local flattened quadrature point index
158 : * @param var The variable number
159 : * @param tag The vector tag
160 : * @returns The quadrature gradient
161 : */
162 108586136 : KOKKOS_FUNCTION Real3 & getVectorQpGrad(const ElementInfo info,
163 : const dof_id_type qp,
164 : const unsigned int var,
165 : const TagID tag) const
166 : {
167 108586136 : return _qp_solutions_grad[tag](info.subdomain, var)[qp];
168 : }
169 : /**
170 : * Get the quadrature point gradient of a variable from a tagged vector for automatic
171 : * differentiation (AD)
172 : * @param info The element information object
173 : * @param jacobian The inverse Jacobian matrix
174 : * @param offset The element's offset into the subdomain-local flattened quadrature point index
175 : * @param qp The local quadrature point index
176 : * @param var The variable number
177 : * @param tag The vector tag
178 : * @param seed The derivative seed
179 : * @returns The quadrature AD gradient
180 : */
181 : KOKKOS_FUNCTION ADReal3 getVectorQpADGrad(const ElementInfo info,
182 : const Real33 jacobian,
183 : const dof_id_type offset,
184 : const dof_id_type qp,
185 : const unsigned int var,
186 : const TagID tag,
187 : const Real seed) const;
188 : /**
189 : * Get the face quadrature point value of a variable from a tagged vector
190 : * @param info The element information object
191 : * @param side The side index
192 : * @param qp The local quadrature point index
193 : * @param var The variable number
194 : * @param tag The vector tag
195 : * @returns The face quadrature value
196 : */
197 : KOKKOS_FUNCTION Real getVectorQpValueFace(const ElementInfo info,
198 : const unsigned int side,
199 : const unsigned int qp,
200 : const unsigned int var,
201 : const TagID tag) const;
202 : /**
203 : * Get the face quadrature point value of a variable from a tagged vector for automatic
204 : * differentiation (AD)
205 : * @param info The element information object
206 : * @param side The side index
207 : * @param qp The local quadrature point index
208 : * @param var The variable number
209 : * @param tag The vector tag
210 : * @param seed The derivative seed
211 : * @returns The face quadrature AD value
212 : */
213 : KOKKOS_FUNCTION ADReal getVectorQpADValueFace(const ElementInfo info,
214 : const unsigned int side,
215 : const unsigned int qp,
216 : const unsigned int var,
217 : const TagID tag,
218 : const Real seed) const;
219 : /**
220 : * Get the face quadrature point gradient of a variable from a tagged vector
221 : * @param info The element information object
222 : * @param side The side index
223 : * @param jacobian The inverse Jacobian matrix
224 : * @param qp The local quadrature point index
225 : * @param var The variable number
226 : * @param tag The vector tag
227 : * @returns The face quadrature gradient
228 : */
229 : KOKKOS_FUNCTION Real3 getVectorQpGradFace(const ElementInfo info,
230 : const unsigned int side,
231 : const Real33 jacobian,
232 : const unsigned int qp,
233 : const unsigned int var,
234 : const TagID tag) const;
235 :
236 : /**
237 : * Get the face quadrature point gradient of a variable from a tagged vector for automatic
238 : * differentiation (AD)
239 : * @param info The element information object
240 : * @param side The side index
241 : * @param jacobian The inverse Jacobian matrix
242 : * @param qp The local quadrature point index
243 : * @param var The variable number
244 : * @param tag The vector tag
245 : * @param seed The derivative seed
246 : * @returns The face quadrature AD gradient
247 : */
248 : KOKKOS_FUNCTION ADReal3 getVectorQpADGradFace(const ElementInfo info,
249 : const unsigned int side,
250 : const Real33 jacobian,
251 : const unsigned int qp,
252 : const unsigned int var,
253 : const TagID tag,
254 : const Real seed) const;
255 :
256 : /**
257 : * Kokkos function for caching variable values on element quadrature points
258 : */
259 : KOKKOS_FUNCTION void operator()(const ThreadID tid) const;
260 : #endif
261 :
262 : private:
263 : /**
264 : * Setup variable data
265 : */
266 : void setupVariables();
267 :
268 : /**
269 : * Setup DOF data
270 : */
271 : void setupDofs();
272 :
273 : /**
274 : * Setup coupling data between variables
275 : */
276 : void setupCoupling();
277 :
278 : /**
279 : * Mark the DOFs covered by nodal BCs
280 : */
281 : void setupNodalBCDofs();
282 :
283 : /**
284 : * Get the list of DOFs covered by a nodal BC
285 : * @param nbc The Kokkos nodal BC object
286 : * @param dofs Local-plus-ghost DOF mask; entries are set true for DOFs covered by the nodal BC
287 : */
288 : void getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs);
289 :
290 : /**
291 : * Kokkos thread object
292 : */
293 : Thread<> _thread;
294 :
295 : /**
296 : * Cached elemental quadrature values and gradients
297 : */
298 : ///@{
299 : Array<Array2D<Array<Real>>> _qp_solutions;
300 : Array<Array2D<Array<Real3>>> _qp_solutions_grad;
301 : ///@}
302 :
303 : /**
304 : * Local nodal DOF indices of each variable
305 : */
306 : Array<Array<dof_id_type>> _local_node_dof_index;
307 :
308 : /**
309 : * FE type ID of each variable
310 : */
311 : Array<unsigned int> _var_fe_types;
312 :
313 : /**
314 : * Off-diagonal coupled variable numbers of each variable
315 : */
316 : Array<Array<unsigned int>> _coupling;
317 :
318 : /**
319 : * Per-matrix-tag local-plus-ghost DOF masks for nodal BC coverage
320 : */
321 : Array<Array<bool>> _nbc_matrix_tag_dof;
322 : };
323 :
324 : #ifdef MOOSE_KOKKOS_SCOPE
325 :
326 : KOKKOS_FUNCTION inline ADReal
327 17390940 : FESystem::getVectorDofADValue(const dof_id_type dof, TagID tag, const Real seed) const
328 : {
329 17390940 : ADReal value = _vectors[tag][dof];
330 :
331 17390940 : if (seed != 0)
332 17298214 : value.derivatives().insert(_local_to_global_dof_index[dof]) = seed;
333 :
334 17390940 : return value;
335 0 : }
336 :
337 : KOKKOS_FUNCTION inline ADReal
338 2270400 : FESystem::getVectorQpADValue(const ElementInfo info,
339 : const dof_id_type offset,
340 : const dof_id_type qp,
341 : const unsigned int var,
342 : const TagID tag,
343 : const Real seed) const
344 : {
345 2270400 : ADReal value = 0;
346 :
347 2270400 : if (seed == 0)
348 16800 : value = getVectorQpValue(info, offset + qp, var, tag);
349 : else
350 : {
351 2253600 : auto fe = _var_fe_types[var];
352 2253600 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
353 2253600 : auto & phi = kokkosAssembly().getPhi(info.subdomain, info.type, fe);
354 :
355 11268000 : for (unsigned int i = 0; i < n_dofs; ++i)
356 9014400 : value += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * phi(i, qp);
357 : }
358 :
359 2270400 : return value;
360 0 : }
361 :
362 : KOKKOS_FUNCTION inline ADReal3
363 2369344 : FESystem::getVectorQpADGrad(const ElementInfo info,
364 : const Real33 jacobian,
365 : const dof_id_type offset,
366 : const dof_id_type qp,
367 : const unsigned int var,
368 : const TagID tag,
369 : const Real seed) const
370 : {
371 2369344 : ADReal3 grad;
372 :
373 2369344 : if (seed == 0)
374 276560 : grad = getVectorQpGrad(info, offset + qp, var, tag);
375 : else
376 : {
377 2092784 : auto fe = _var_fe_types[var];
378 2092784 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
379 2092784 : auto & grad_phi = kokkosAssembly().getGradPhi(info.subdomain, info.type, fe);
380 :
381 10341680 : for (unsigned int i = 0; i < n_dofs; ++i)
382 8248896 : grad += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) *
383 16497792 : (jacobian * grad_phi(i, qp));
384 : }
385 :
386 2369344 : return grad;
387 0 : }
388 :
389 : KOKKOS_FUNCTION inline Real
390 221894 : FESystem::getVectorQpValueFace(const ElementInfo info,
391 : const unsigned int side,
392 : const unsigned int qp,
393 : const unsigned int var,
394 : const TagID tag) const
395 : {
396 221894 : auto fe = _var_fe_types[var];
397 221894 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
398 221894 : auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);
399 :
400 221894 : Real value = 0;
401 :
402 1120794 : for (unsigned int i = 0; i < n_dofs; ++i)
403 898900 : value += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * phi(i, qp);
404 :
405 221894 : return value;
406 : }
407 :
408 : KOKKOS_FUNCTION inline ADReal
409 15156 : FESystem::getVectorQpADValueFace(const ElementInfo info,
410 : const unsigned int side,
411 : const unsigned int qp,
412 : const unsigned int var,
413 : const TagID tag,
414 : const Real seed) const
415 : {
416 15156 : auto fe = _var_fe_types[var];
417 15156 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
418 15156 : auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);
419 :
420 15156 : ADReal value = 0;
421 :
422 73948 : for (unsigned int i = 0; i < n_dofs; ++i)
423 58792 : value += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * phi(i, qp);
424 :
425 15156 : return value;
426 0 : }
427 :
428 : KOKKOS_FUNCTION inline Real3
429 0 : FESystem::getVectorQpGradFace(const ElementInfo info,
430 : const unsigned int side,
431 : const Real33 jacobian,
432 : const unsigned int qp,
433 : const unsigned int var,
434 : const TagID tag) const
435 : {
436 0 : auto fe = _var_fe_types[var];
437 0 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
438 0 : auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);
439 :
440 0 : Real3 grad = 0;
441 :
442 0 : for (unsigned int i = 0; i < n_dofs; ++i)
443 0 : grad += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) *
444 0 : (jacobian * grad_phi(i, qp));
445 :
446 0 : return grad;
447 : }
448 :
449 : KOKKOS_FUNCTION inline ADReal3
450 0 : FESystem::getVectorQpADGradFace(const ElementInfo info,
451 : const unsigned int side,
452 : const Real33 jacobian,
453 : const unsigned int qp,
454 : const unsigned int var,
455 : const TagID tag,
456 : const Real seed) const
457 : {
458 0 : auto fe = _var_fe_types[var];
459 0 : auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
460 0 : auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);
461 :
462 0 : ADReal3 grad = ADReal(0);
463 :
464 0 : for (unsigned int i = 0; i < n_dofs; ++i)
465 0 : grad += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) *
466 0 : (jacobian * grad_phi(i, qp));
467 :
468 0 : return grad;
469 0 : }
470 : #endif
471 :
472 46648814 : MakeSystemHolder(FESystem);
473 : } // namespace Moose::Kokkos
|