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 "KokkosDatum.h"
13 :
14 : #include "MooseVariableFieldBase.h"
15 : #include "SystemBase.h"
16 :
17 : namespace Moose::Kokkos
18 : {
19 :
20 : /**
21 : * The Kokkos wrapper classes for MOOSE-like shape function access
22 : */
23 : ///@{
24 : class VariablePhiValue
25 : {
26 : public:
27 : /**
28 : * Get the current shape function
29 : * @param datum The AssemblyDatum object of the current thread
30 : * @param i The element-local DOF index
31 : * @param qp The local quadrature point index
32 : * @returns The shape function
33 : */
34 4906888 : KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
35 : {
36 4906888 : auto & elem = datum.elem();
37 4906888 : auto side = datum.side();
38 4906888 : auto fe = datum.jfe();
39 :
40 0 : return side == libMesh::invalid_uint
41 4906888 : ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
42 4906888 : : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
43 : }
44 : };
45 :
46 : class VariablePhiGradient
47 : {
48 : public:
49 : /**
50 : * Get the gradient of the current shape function
51 : * @param datum The AssemblyDatum object of the current thread
52 : * @param i The element-local DOF index
53 : * @param qp The local quadrature point index
54 : * @returns The gradient of the shape function
55 : */
56 25387624 : KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
57 : {
58 25387624 : auto & elem = datum.elem();
59 25387624 : auto side = datum.side();
60 25387624 : auto fe = datum.jfe();
61 :
62 0 : return datum.J(qp) *
63 : (side == libMesh::invalid_uint
64 25387624 : ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
65 50775248 : : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
66 : }
67 : };
68 :
69 : class VariableTestValue
70 : {
71 : public:
72 : /**
73 : * Get the current test function
74 : * @param datum The AssemblyDatum object of the current thread
75 : * @param i The element-local DOF index
76 : * @param qp The local quadrature point index
77 : * @returns The test function
78 : */
79 113036280 : KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
80 : {
81 113036280 : auto & elem = datum.elem();
82 113036280 : auto side = datum.side();
83 113036280 : auto fe = datum.ife();
84 :
85 0 : return side == libMesh::invalid_uint
86 113036280 : ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
87 113036280 : : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
88 : }
89 : };
90 :
91 : class VariableTestGradient
92 : {
93 : public:
94 : /**
95 : * Get the gradient of the current test function
96 : * @param datum The AssemblyDatum object of the current thread
97 : * @param i The element-local DOF index
98 : * @param qp The local quadrature point index
99 : * @returns The gradient of the test function
100 : */
101 68208188 : KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
102 : {
103 68208188 : auto & elem = datum.elem();
104 68208188 : auto side = datum.side();
105 68208188 : auto fe = datum.ife();
106 :
107 0 : return datum.J(qp) *
108 : (side == libMesh::invalid_uint
109 68208188 : ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
110 136416376 : : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
111 : }
112 : };
113 :
114 : using ADVariablePhiValue = VariablePhiValue;
115 : using ADVariablePhiGradient = VariablePhiGradient;
116 : using ADVariableTestValue = VariableTestValue;
117 : using ADVariableTestGradient = VariableTestGradient;
118 :
119 : ///@}
120 :
121 : /**
122 : * The Kokkos wrapper classes for MOOSE-like variable value access
123 : */
124 : ///@{
125 : template <bool is_ad>
126 : class VariableValueTempl
127 : {
128 : using real_type = std::conditional_t<is_ad, ADReal, Real>;
129 :
130 : public:
131 : /**
132 : * Default constructor
133 : */
134 8867 : VariableValueTempl() = default;
135 : /**
136 : * Constructor
137 : * @param var The Kokkos variable
138 : * @param dof Whether to get DOF values
139 : */
140 2898 : VariableValueTempl(Variable var, bool dof = false) : _var(var), _dof(dof) {}
141 : /**
142 : * Constructor
143 : * @param var The MOOSE variable
144 : * @param tag The vector tag name
145 : * @param dof Whether to get DOF values
146 : */
147 8872 : VariableValueTempl(const MooseVariableFieldBase & var,
148 : const TagName & tag = Moose::SOLUTION_TAG,
149 4040 : bool dof = false)
150 4832 : : _var(var, tag), _dof(dof)
151 : {
152 8872 : }
153 : /**
154 : * Constructor
155 : * @param vars The MOOSE variables
156 : * @param tag The vector tag name
157 : * @param dof Whether to get DOF values
158 : */
159 : ///@{
160 : VariableValueTempl(const std::vector<const MooseVariableFieldBase *> & vars,
161 : const TagName & tag = Moose::SOLUTION_TAG,
162 : bool dof = false)
163 : : _var(vars, tag), _dof(dof)
164 : {
165 : }
166 : VariableValueTempl(const std::vector<MooseVariableFieldBase *> & vars,
167 : const TagName & tag = Moose::SOLUTION_TAG,
168 : bool dof = false)
169 : : _var(vars, tag), _dof(dof)
170 : {
171 : }
172 : ///@}
173 :
174 : /**
175 : * Copy constructor for parallel dispatch
176 : */
177 : VariableValueTempl(const VariableValueTempl<is_ad> & object);
178 : /**
179 : * Copy assignment operator
180 : */
181 : VariableValueTempl<is_ad> & operator=(const VariableValueTempl<is_ad> & object);
182 :
183 : /**
184 : * Get whether the variable was coupled
185 : * @returns Whether the variable was coupled
186 : */
187 28032 : KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }
188 :
189 : /**
190 : * Get the current variable value
191 : * @param datum The Datum object of the current thread
192 : * @param idx The local quadrature point or DOF index
193 : * @param comp The variable component
194 : * @returns The variable value
195 : */
196 2592739 : KOKKOS_FUNCTION auto operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const
197 : {
198 2592739 : return get(datum, idx, comp);
199 : }
200 :
201 : /**
202 : * Get the current variable value
203 : * @param datum The AssemblyDatum object of the current thread
204 : * @param idx The local quadrature point or DOF index
205 : * @param comp The variable component
206 : * @returns The variable value
207 : */
208 : KOKKOS_FUNCTION auto
209 : operator()(AssemblyDatum & datum, unsigned int idx, unsigned int comp = 0) const;
210 :
211 : /**
212 : * Get the Kokkos variable
213 : * @returns The Kokkos variable
214 : */
215 460146 : KOKKOS_FUNCTION const Variable & variable() const { return _var; }
216 :
217 : private:
218 : /**
219 : * Get the current variable value
220 : * @param datum The Datum object of the current thread
221 : * @param idx The local quadrature point or DOF index
222 : * @param comp The variable component
223 : * @param seed The derivative seed (only meaningful for AD)
224 : * @returns The variable value
225 : */
226 : KOKKOS_FUNCTION auto
227 : get(Datum & datum, unsigned int idx, unsigned int comp = 0, Real seed = 0) const;
228 :
229 : /**
230 : * Coupled Kokkos variable
231 : */
232 : Variable _var;
233 : /**
234 : * Derivative seed of each component for AD
235 : */
236 : Array<Real> _seed;
237 : /**
238 : * Flag whether DOF values are requested
239 : */
240 : bool _dof = false;
241 : };
242 :
243 : template <bool is_ad>
244 589864 : VariableValueTempl<is_ad>::VariableValueTempl(const VariableValueTempl<is_ad> & object)
245 336230 : : _var(object._var), _seed(object._seed), _dof(object._dof)
246 : {
247 : if constexpr (is_ad)
248 59425 : if (_var.coupled())
249 : {
250 57533 : if (!_seed.isAlloc())
251 57533 : _seed.create(_var.components());
252 :
253 115066 : for (unsigned int comp = 0; comp < _var.components(); ++comp)
254 57533 : _seed[comp] =
255 29058 : _var.dot() ? _var.mooseVar(comp)->sys().duDotDu(_var.var(comp)) : (_var.old() ? 0 : 1);
256 :
257 57533 : _seed.copyToDevice();
258 : }
259 589864 : }
260 :
261 : template <bool is_ad>
262 : VariableValueTempl<is_ad> &
263 8867 : VariableValueTempl<is_ad>::operator=(const VariableValueTempl<is_ad> & object)
264 : {
265 8867 : _var = object._var;
266 8867 : _dof = object._dof;
267 :
268 8867 : return *this;
269 : }
270 :
271 : template <bool is_ad>
272 : KOKKOS_FUNCTION auto
273 14490855 : VariableValueTempl<is_ad>::operator()(AssemblyDatum & datum,
274 : unsigned int idx,
275 : unsigned int comp) const
276 : {
277 : if constexpr (is_ad)
278 : {
279 2369564 : Real seed =
280 2369564 : datum.do_derivatives() && _var.coupled() && _var.sys(comp) == datum.sys() ? _seed[comp] : 0;
281 :
282 2369564 : return get(datum, idx, comp, seed);
283 : }
284 : else
285 12121291 : return get(datum, idx, comp);
286 : }
287 :
288 : template <bool is_ad>
289 : KOKKOS_FUNCTION auto
290 17083594 : VariableValueTempl<is_ad>::get(Datum & datum,
291 : unsigned int idx,
292 : unsigned int comp,
293 : [[maybe_unused]] Real seed) const
294 : {
295 : KOKKOS_ASSERT(_var.initialized());
296 :
297 2369564 : real_type value;
298 :
299 17083594 : if (_var.coupled())
300 : {
301 17064380 : auto & sys = datum.system(_var.sys(comp));
302 17064380 : auto var = _var.var(comp);
303 17064380 : auto tag = _var.tag();
304 :
305 17064380 : if (_dof)
306 : {
307 : unsigned int dof;
308 :
309 4355794 : if (datum.isNodal())
310 : {
311 4355154 : auto node = datum.node();
312 4355154 : dof = sys.getNodeLocalDofIndex(node, 0, var);
313 : }
314 : else
315 : {
316 640 : auto elem = datum.elem().id;
317 640 : dof = sys.getElemLocalDofIndex(elem, idx, var);
318 : }
319 :
320 : if constexpr (is_ad)
321 68852 : value = sys.getVectorDofADValue(dof, tag, seed);
322 : else
323 4286942 : value = sys.getVectorDofValue(dof, tag);
324 : }
325 : else
326 : {
327 12708586 : auto & elem = datum.elem();
328 12708586 : auto side = datum.side();
329 :
330 : if constexpr (is_ad)
331 2285556 : value = side == libMesh::invalid_uint
332 4571112 : ? sys.getVectorQpADValue(elem, datum.qpOffset(), idx, var, tag, seed)
333 : : sys.getVectorQpADValueFace(elem, side, idx, var, tag, seed);
334 : else
335 10423030 : value = side == libMesh::invalid_uint
336 10423030 : ? sys.getVectorQpValue(elem, datum.qpOffset() + idx, var, tag)
337 221894 : : sys.getVectorQpValueFace(elem, side, idx, var, tag);
338 : }
339 : }
340 : else
341 19214 : value = _var.value(comp);
342 :
343 17083594 : return value;
344 0 : }
345 :
346 : template <bool is_ad>
347 : class VariableGradientTempl
348 : {
349 : using real3_type = std::conditional_t<is_ad, ADReal3, Real3>;
350 :
351 : public:
352 : /**
353 : * Default constructor
354 : */
355 : VariableGradientTempl() = default;
356 : /**
357 : * Constructor
358 : * @param var The Kokkos variable
359 : */
360 680 : VariableGradientTempl(Variable var) : _var(var) {}
361 : /**
362 : * Constructor
363 : * @param var The MOOSE variable
364 : * @param tag The vector tag name
365 : */
366 3705 : VariableGradientTempl(const MooseVariableFieldBase & var,
367 1657 : const TagName & tag = Moose::SOLUTION_TAG)
368 2048 : : _var(var, tag)
369 : {
370 3705 : }
371 : /**
372 : * Constructor
373 : * @param vars The MOOSE variables
374 : * @param tag The vector tag name
375 : */
376 : ///@{
377 : VariableGradientTempl(const std::vector<const MooseVariableFieldBase *> & vars,
378 : const TagName & tag = Moose::SOLUTION_TAG)
379 : : _var(vars, tag)
380 : {
381 : }
382 : VariableGradientTempl(const std::vector<MooseVariableFieldBase *> & vars,
383 : const TagName & tag = Moose::SOLUTION_TAG)
384 : : _var(vars, tag)
385 : {
386 : }
387 : ///@}
388 :
389 : /**
390 : * Copy constructor for parallel dispatch
391 : */
392 : VariableGradientTempl(const VariableGradientTempl<is_ad> & object);
393 : /**
394 : * Copy assignment operator
395 : */
396 : VariableGradientTempl<is_ad> & operator=(const VariableGradientTempl<is_ad> & object);
397 :
398 : /**
399 : * Get whether the variable was coupled
400 : * @returns Whether the variable was coupled
401 : */
402 : KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }
403 :
404 : /**
405 : * Get the current variable gradient
406 : * @param datum The Datum object of the current thread
407 : * @param qp The local quadrature point index
408 : * @param comp The variable component
409 : * @returns The variable gradient
410 : */
411 : KOKKOS_FUNCTION auto operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const
412 : {
413 : return get(datum, qp, comp);
414 : }
415 :
416 : /**
417 : * Get the current variable gradient
418 : * @param datum The AssemblyDatum object of the current thread
419 : * @param qp The local quadrature point index
420 : * @param comp The variable component
421 : * @returns The variable gradient
422 : */
423 : KOKKOS_FUNCTION auto
424 : operator()(AssemblyDatum & datum, unsigned int qp, unsigned int comp = 0) const;
425 :
426 : /**
427 : * Get the Kokkos variable
428 : * @returns The Kokkos variable
429 : */
430 : KOKKOS_FUNCTION const Variable & variable() const { return _var; }
431 :
432 : private:
433 : /**
434 : * Get the current variable gradient
435 : * @param datum The Datum object of the current thread
436 : * @param qp The local quadrature point index
437 : * @param comp The variable component
438 : * @param seed The derivative seed (only meaningful for AD)
439 : * @returns The variable gradient
440 : */
441 : KOKKOS_FUNCTION auto
442 : get(Datum & datum, unsigned int qp, unsigned int comp = 0, Real seed = 0) const;
443 :
444 : /**
445 : * Coupled Kokkos variable
446 : */
447 : Variable _var;
448 : /**
449 : * Derivative seed of each component for AD
450 : */
451 : Array<Real> _seed;
452 : };
453 :
454 : template <bool is_ad>
455 187636 : VariableGradientTempl<is_ad>::VariableGradientTempl(const VariableGradientTempl<is_ad> & object)
456 111310 : : _var(object._var), _seed(object._seed)
457 : {
458 : if constexpr (is_ad)
459 17931 : if (_var.coupled())
460 : {
461 17931 : if (!_seed.isAlloc())
462 17931 : _seed.create(_var.components());
463 :
464 35862 : for (unsigned int comp = 0; comp < _var.components(); ++comp)
465 17931 : _seed[comp] =
466 9068 : _var.dot() ? _var.mooseVar(comp)->sys().duDotDu(_var.var(comp)) : (_var.old() ? 0 : 1);
467 :
468 17931 : _seed.copyToDevice();
469 : }
470 187636 : }
471 :
472 : template <bool is_ad>
473 : VariableGradientTempl<is_ad> &
474 : VariableGradientTempl<is_ad>::operator=(const VariableGradientTempl<is_ad> & object)
475 : {
476 : _var = object._var;
477 :
478 : return *this;
479 : }
480 :
481 : template <bool is_ad>
482 : KOKKOS_FUNCTION auto
483 42819028 : VariableGradientTempl<is_ad>::operator()(AssemblyDatum & datum,
484 : unsigned int qp,
485 : unsigned int comp) const
486 : {
487 : if constexpr (is_ad)
488 : {
489 2369344 : Real seed =
490 2369344 : datum.do_derivatives() && _var.coupled() && _var.sys(comp) == datum.sys() ? _seed[comp] : 0;
491 :
492 2369344 : return get(datum, qp, comp, seed);
493 : }
494 : else
495 40449684 : return get(datum, qp, comp);
496 : }
497 :
498 : template <bool is_ad>
499 : KOKKOS_FUNCTION auto
500 42819028 : VariableGradientTempl<is_ad>::get(Datum & datum,
501 : unsigned int qp,
502 : unsigned int comp,
503 : [[maybe_unused]] Real seed) const
504 : {
505 : KOKKOS_ASSERT(_var.initialized());
506 :
507 42819028 : real3_type grad;
508 :
509 42819028 : if (_var.coupled())
510 : {
511 : KOKKOS_ASSERT(!datum.isNodal());
512 :
513 42819028 : auto & elem = datum.elem();
514 42819028 : auto side = datum.side();
515 :
516 : if constexpr (is_ad)
517 2369344 : grad =
518 : side == libMesh::invalid_uint
519 4738688 : ? datum.system(_var.sys(comp))
520 : .getVectorQpADGrad(
521 : elem, datum.J(qp), datum.qpOffset(), qp, _var.var(comp), _var.tag(), seed)
522 0 : : datum.system(_var.sys(comp))
523 : .getVectorQpADGradFace(
524 : elem, side, datum.J(qp), qp, _var.var(comp), _var.tag(), seed);
525 : else
526 40449684 : grad =
527 : side == libMesh::invalid_uint
528 80899368 : ? datum.system(_var.sys(comp))
529 40449684 : .getVectorQpGrad(elem, datum.qpOffset() + qp, _var.var(comp), _var.tag())
530 0 : : datum.system(_var.sys(comp))
531 : .getVectorQpGradFace(elem, side, datum.J(qp), qp, _var.var(comp), _var.tag());
532 : }
533 :
534 42819028 : return grad;
535 0 : }
536 :
537 : using VariableValue = VariableValueTempl<false>;
538 : using ADVariableValue = VariableValueTempl<true>;
539 : using VariableGradient = VariableGradientTempl<false>;
540 : using ADVariableGradient = VariableGradientTempl<true>;
541 :
542 : template <>
543 : struct ArrayDeepCopy<ADVariableValue>
544 : {
545 : static constexpr bool value = true;
546 : };
547 :
548 : template <>
549 : struct ArrayDeepCopy<ADVariableGradient>
550 : {
551 : static constexpr bool value = true;
552 : };
553 : ///@}
554 :
555 : } // namespace Moose::Kokkos
|
