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 "KokkosDispatcher.h"
13 : #include "KokkosFESystem.h"
14 : #include "KokkosVariableValue.h"
15 : #include "KokkosMaterialPropertyValue.h"
16 :
17 : #include "MooseVariableBase.h"
18 : #include "ResidualObject.h"
19 :
20 : namespace Moose::Kokkos
21 : {
22 :
23 : /**
24 : * The base class for Kokkos residual objects
25 : */
26 : class ResidualObject : public ::ResidualObject,
27 : public MeshHolder,
28 : public AssemblyHolder,
29 : public FESystemHolder
30 : {
31 : public:
32 : static InputParameters validParams();
33 :
34 : /// Whether this object's hooks factor out the test function
35 : static constexpr bool use_precompute_hooks = false;
36 :
37 : /**
38 : * Constructor
39 : * @param field_type The MOOSE variable field type
40 : * @param nodal Whether the residual object is a nodal residual object
41 : */
42 : ResidualObject(const InputParameters & parameters,
43 : Moose::VarFieldType field_type,
44 : bool nodal = false);
45 : /**
46 : * Copy constructor for parallel dispatch
47 : */
48 : ResidualObject(const ResidualObject & object);
49 :
50 : /**
51 : * Kokkos function tags
52 : */
53 : ///@{
54 : struct ResidualLoop
55 : {
56 : };
57 : struct JacobianLoop
58 : {
59 : };
60 : struct OffDiagJacobianLoop
61 : {
62 : };
63 : ///@}
64 :
65 45961 : virtual const MooseVariableBase & variable() const override { return _var; }
66 :
67 0 : virtual void computeOffDiagJacobian(unsigned int) override final
68 : {
69 0 : mooseError("computeOffDiagJacobian() is not used for Kokkos residual objects.");
70 : }
71 1860 : virtual void computeResidualAndJacobian() override
72 : {
73 1860 : computeResidual();
74 1860 : computeJacobian();
75 1860 : }
76 :
77 : protected:
78 : /**
79 : * Reference of the MOOSE variable
80 : */
81 : MooseVariableFieldBase & _var;
82 : /**
83 : * Kokkos variable
84 : */
85 : Variable _kokkos_var;
86 : /**
87 : * Kokkos thread object
88 : */
89 : Thread<> _thread;
90 : /**
91 : * Kokkos functor dispatchers
92 : */
93 : ///@{
94 : std::unique_ptr<DispatcherBase> _residual_dispatcher;
95 : std::unique_ptr<DispatcherBase> _jacobian_dispatcher;
96 : std::unique_ptr<DispatcherBase> _offdiag_jacobian_dispatcher;
97 : ///@}
98 :
99 : /**
100 : * Mesh dimension
101 : */
102 : const unsigned int _dimension;
103 :
104 : /**
105 : * TODO: Move to TransientInterface
106 : */
107 : ///@{
108 : /**
109 : * Time
110 : */
111 : Scalar<Real> _t;
112 : /**
113 : * Old time
114 : */
115 : Scalar<const Real> _t_old;
116 : /**
117 : * The number of the time step
118 : */
119 : Scalar<int> _t_step;
120 : /**
121 : * Time step size
122 : */
123 : Scalar<Real> _dt;
124 : /**
125 : * Size of the old time step
126 : */
127 : Scalar<Real> _dt_old;
128 : ///@}
129 :
130 : /**
131 : * Accumulate local elemental residual contribution to tagged vectors
132 : * @param local_re The local elemental residual contribution
133 : * @param elem The contiguous element ID
134 : * @param i The test function index
135 : * @param comp The variable component
136 : */
137 : KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
138 : const ContiguousElementID elem,
139 : const unsigned int i,
140 : const unsigned int comp = 0) const;
141 : /**
142 : * Accumulate or set local nodal residual contribution to tagged vectors
143 : * @param add Whether to add or set the local residual
144 : * @param local_re The local nodal residual contribution
145 : * @param node The contiguous node ID
146 : * @param comp The variable component
147 : */
148 : KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
149 : const Real local_re,
150 : const ContiguousNodeID node,
151 : const unsigned int comp = 0) const;
152 : /**
153 : * Accumulate or set local nodal residual contribution to tagged vectors for vector FE variables
154 : * @param add Whether to add or set the local residual
155 : * @param local_re The local nodal residual contribution
156 : * @param node The contiguous node ID
157 : * @param i The vector component/node-local DOF index
158 : */
159 : KOKKOS_FUNCTION void accumulateTaggedVectorNodalResidual(const bool add,
160 : const Real local_re,
161 : const ContiguousNodeID node,
162 : const unsigned int i) const;
163 : /**
164 : * Accumulate local elemental Jacobian contribution to tagged matrices
165 : * @param local_ke The local elemental Jacobian contribution
166 : * @param elem The contiguous element ID
167 : * @param i The test function DOF index
168 : * @param j The trial function DOF index
169 : * @param jvar The variable number for column
170 : * @param comp The variable component
171 : */
172 : KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke,
173 : const ContiguousElementID elem,
174 : const unsigned int i,
175 : const unsigned int j,
176 : const unsigned int jvar,
177 : const unsigned int comp = 0) const;
178 : /**
179 : * Accumulate local elemental Jacobian contribution to tagged matrices using automatic
180 : * differentiation (AD)
181 : * @param local_ke The local elemental Jacobian contribution
182 : * @param datum The AssemblyDatum object of the current thread
183 : * @param i The test function DOF index
184 : * @param comp The variable component
185 : */
186 : KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
187 : const AssemblyDatum & datum,
188 : const unsigned int i,
189 : const unsigned int comp = 0) const;
190 : /**
191 : * Accumulate or set local nodal Jacobian contribution to tagged matrices
192 : * @param add Whether to add or set the local Jacobian
193 : * @param local_ke The local nodal Jacobian contribution
194 : * @param node The contiguous node ID
195 : * @param jvar The variable number for column
196 : * @param comp The variable component
197 : */
198 : KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
199 : const Real local_ke,
200 : const ContiguousNodeID node,
201 : const unsigned int jvar,
202 : const unsigned int comp = 0) const;
203 : /**
204 : * Accumulate or set local nodal Jacobian contribution to tagged matrices for vector FE variables
205 : * @param add Whether to add or set the local Jacobian
206 : * @param local_ke The local nodal Jacobian contribution
207 : * @param node The contiguous node ID
208 : * @param i The vector component/node-local DOF index for the row
209 : * @param j The node-local DOF index for the column
210 : * @param jvar The variable number for column
211 : */
212 : KOKKOS_FUNCTION void accumulateTaggedVectorNodalMatrix(const bool add,
213 : const Real local_ke,
214 : const ContiguousNodeID node,
215 : const unsigned int i,
216 : const unsigned int j,
217 : const unsigned int jvar) const;
218 : /**
219 : * Accumulate or set local nodal Jacobian contribution to tagged matrices using automatic
220 : * differentiation (AD)
221 : * @param add Whether to add or set the local Jacobian
222 : * @param local_ke The local elemental Jacobian contribution
223 : * @param node The contiguous node ID
224 : * @param comp The variable component
225 : */
226 : KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
227 : const DNDerivativeType & local_ke,
228 : const ContiguousNodeID node,
229 : const unsigned int comp = 0) const;
230 :
231 : /**
232 : * The common loop structure template for computing elemental residual
233 : * @param datum The AssemblyDatum object of the current thread
234 : * @param body The quadrature point loop body
235 : */
236 : template <typename function>
237 : KOKKOS_FUNCTION void computeResidualInternal(AssemblyDatum & datum, function body) const;
238 : /**
239 : * The common loop structure template for computing elemental Jacobian
240 : * @param datum The AssemblyDatum object of the current thread
241 : * @param body The quadrature point loop body
242 : */
243 : template <typename function>
244 : KOKKOS_FUNCTION void computeJacobianInternal(AssemblyDatum & datum, function body) const;
245 :
246 : private:
247 : /**
248 : * Tags this object operates on
249 : */
250 : ///@{
251 : Array<TagID> _vector_tags;
252 : Array<TagID> _matrix_tags;
253 : ///@}
254 : };
255 :
256 : KOKKOS_FUNCTION inline void
257 22849738 : ResidualObject::accumulateTaggedElementalResidual(const Real local_re,
258 : const ContiguousElementID elem,
259 : const unsigned int i,
260 : const unsigned int comp) const
261 : {
262 22849738 : if (!local_re)
263 1937092 : return;
264 :
265 20912646 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
266 20912646 : auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
267 :
268 42632068 : for (unsigned int t = 0; t < _vector_tags.size(); ++t)
269 : {
270 21719422 : auto tag = _vector_tags[t];
271 :
272 21719422 : if (sys.isResidualTagActive(tag))
273 21599454 : ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), local_re);
274 : }
275 : }
276 :
277 : KOKKOS_FUNCTION inline void
278 1703968 : ResidualObject::accumulateTaggedNodalResidual(const bool add,
279 : const Real local_re,
280 : const ContiguousNodeID node,
281 : const unsigned int comp) const
282 : {
283 1703968 : if (!local_re && add)
284 659958 : return;
285 :
286 1044010 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
287 1044010 : auto dof = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
288 :
289 2434030 : for (unsigned int t = 0; t < _vector_tags.size(); ++t)
290 : {
291 1390020 : auto tag = _vector_tags[t];
292 :
293 1390020 : if (sys.isResidualTagActive(tag))
294 : {
295 1379760 : if (add)
296 921576 : sys.getVectorDofValue(dof, tag) += local_re;
297 : else
298 458184 : sys.getVectorDofValue(dof, tag) = local_re;
299 : }
300 : }
301 : }
302 :
303 : KOKKOS_FUNCTION inline void
304 108116 : ResidualObject::accumulateTaggedVectorNodalResidual(const bool add,
305 : const Real local_re,
306 : const ContiguousNodeID node,
307 : const unsigned int i) const
308 : {
309 108116 : if (!local_re && add)
310 0 : return;
311 :
312 108116 : auto & sys = kokkosSystem(_kokkos_var.sys());
313 108116 : auto dof = sys.getNodeLocalDofIndex(node, i, _kokkos_var.var());
314 :
315 216232 : for (unsigned int t = 0; t < _vector_tags.size(); ++t)
316 : {
317 108116 : auto tag = _vector_tags[t];
318 :
319 108116 : if (sys.isResidualTagActive(tag))
320 : {
321 108116 : if (add)
322 0 : sys.getVectorDofValue(dof, tag) += local_re;
323 : else
324 108116 : sys.getVectorDofValue(dof, tag) = local_re;
325 : }
326 : }
327 : }
328 :
329 : KOKKOS_FUNCTION inline void
330 19724200 : ResidualObject::accumulateTaggedElementalMatrix(const Real local_ke,
331 : const ContiguousElementID elem,
332 : const unsigned int i,
333 : const unsigned int j,
334 : const unsigned int jvar,
335 : const unsigned int comp) const
336 : {
337 19724200 : if (!local_ke)
338 5917810 : return;
339 :
340 13806390 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
341 13806390 : auto row = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
342 13806390 : auto col = sys.getElemGlobalDofIndex(elem, j, jvar);
343 :
344 33786828 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
345 : {
346 19980438 : auto tag = _matrix_tags[t];
347 :
348 19980438 : if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))
349 13067680 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke);
350 : }
351 : }
352 :
353 : KOKKOS_FUNCTION inline void
354 1103456 : ResidualObject::accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
355 : const AssemblyDatum & datum,
356 : const unsigned int i,
357 : const unsigned int comp) const
358 : {
359 1103456 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
360 1103456 : auto row = sys.getElemLocalDofIndex(datum.elem().id, i, _kokkos_var.var(comp));
361 :
362 2706412 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
363 : {
364 1602956 : auto tag = _matrix_tags[t];
365 :
366 1602956 : if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))
367 4879930 : for (unsigned int j = 0; j < local_ke.size(); ++j)
368 : {
369 3891220 : auto col = local_ke.raw_index(j);
370 :
371 3891220 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke.raw_at(j));
372 : }
373 : }
374 1103456 : }
375 :
376 : KOKKOS_FUNCTION inline void
377 502650 : ResidualObject::accumulateTaggedNodalMatrix(const bool add,
378 : const Real local_ke,
379 : const ContiguousNodeID node,
380 : const unsigned int jvar,
381 : const unsigned int comp) const
382 : {
383 502650 : if (!local_ke && add)
384 157350 : return;
385 :
386 345300 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
387 345300 : auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
388 345300 : auto col = sys.getNodeGlobalDofIndex(node, 0, jvar);
389 :
390 861178 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
391 : {
392 515878 : auto tag = _matrix_tags[t];
393 :
394 515878 : if (sys.isMatrixTagActive(tag))
395 : {
396 362788 : auto & matrix = sys.getMatrix(tag);
397 :
398 362788 : if (add)
399 301590 : matrix(row, col) += local_ke;
400 : else
401 : {
402 61198 : matrix.zero(row);
403 61198 : matrix(row, col) = local_ke;
404 : }
405 : }
406 : }
407 : }
408 :
409 : KOKKOS_FUNCTION inline void
410 9696 : ResidualObject::accumulateTaggedVectorNodalMatrix(const bool add,
411 : const Real local_ke,
412 : const ContiguousNodeID node,
413 : const unsigned int i,
414 : const unsigned int j,
415 : const unsigned int jvar) const
416 : {
417 9696 : if (!local_ke && add)
418 0 : return;
419 :
420 9696 : auto & sys = kokkosSystem(_kokkos_var.sys());
421 9696 : auto row = sys.getNodeLocalDofIndex(node, i, _kokkos_var.var());
422 9696 : auto col = sys.getNodeGlobalDofIndex(node, j, jvar);
423 :
424 29088 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
425 : {
426 19392 : auto tag = _matrix_tags[t];
427 :
428 19392 : if (sys.isMatrixTagActive(tag))
429 : {
430 9696 : auto & matrix = sys.getMatrix(tag);
431 :
432 9696 : if (add)
433 0 : matrix(row, col) += local_ke;
434 : else
435 : {
436 9696 : matrix.zero(row);
437 9696 : matrix(row, col) = local_ke;
438 : }
439 : }
440 : }
441 : }
442 :
443 : KOKKOS_FUNCTION inline void
444 33390 : ResidualObject::accumulateTaggedNodalMatrix(const bool add,
445 : const DNDerivativeType & local_ke,
446 : const ContiguousNodeID node,
447 : const unsigned int comp) const
448 : {
449 33390 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
450 33390 : auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
451 :
452 100170 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
453 : {
454 66780 : auto tag = _matrix_tags[t];
455 66780 : auto & matrix = sys.getMatrix(tag);
456 :
457 66780 : if (sys.isMatrixTagActive(tag))
458 : {
459 33390 : if (!add)
460 33390 : matrix.zero(row);
461 :
462 66780 : for (unsigned int j = 0; j < local_ke.size(); ++j)
463 : {
464 33390 : auto col = local_ke.raw_index(j);
465 :
466 33390 : if (add)
467 0 : matrix(row, col) += local_ke.raw_at(j);
468 : else
469 33390 : matrix(row, col) = local_ke.raw_at(j);
470 : }
471 : }
472 : }
473 33390 : }
474 :
475 : template <typename function>
476 : KOKKOS_FUNCTION void
477 4640363 : ResidualObject::computeResidualInternal(AssemblyDatum & datum, function body) const
478 : {
479 : Real local_re[MAX_CACHED_DOF];
480 :
481 4640363 : unsigned int stride = MAX_CACHED_DOF * datum.num_local_threads();
482 4640363 : unsigned int num_batches = datum.n_dofs() / stride;
483 :
484 4640363 : if (datum.n_dofs() % stride)
485 4640363 : ++num_batches;
486 :
487 9280726 : for (unsigned int batch = 0; batch < num_batches; ++batch)
488 : {
489 4640363 : unsigned int ib = batch * stride;
490 4640363 : unsigned int ie = ::Kokkos::min(ib + stride, datum.n_dofs());
491 :
492 4640363 : const unsigned int n = ie - ib;
493 4640363 : const unsigned int d = n / datum.num_local_threads();
494 4640363 : const unsigned int m = n % datum.num_local_threads();
495 4640363 : const unsigned int t = datum.local_thread_id();
496 :
497 4640363 : ib += t * d + (t < m ? t : m);
498 4640363 : ie = ib + d + (t < m ? 1 : 0);
499 :
500 26308241 : for (unsigned int i = ib; i < ie; ++i)
501 21667878 : local_re[i - ib] = 0;
502 :
503 4640363 : body(local_re - ib, ib, ie);
504 :
505 26308241 : for (unsigned int i = ib; i < ie; ++i)
506 21667878 : accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
507 : }
508 4640363 : }
509 :
510 : template <typename function>
511 : KOKKOS_FUNCTION void
512 510130 : ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) const
513 : {
514 : Real local_ke[MAX_CACHED_DOF];
515 :
516 2910470 : for (unsigned int j = datum.local_thread_id(); j < datum.n_jdofs();
517 2400340 : j += datum.num_local_threads())
518 : {
519 2400340 : unsigned int num_batches = datum.n_idofs() / MAX_CACHED_DOF;
520 :
521 2400340 : if (datum.n_idofs() % MAX_CACHED_DOF)
522 2400340 : ++num_batches;
523 :
524 4800680 : for (unsigned int batch = 0; batch < num_batches; ++batch)
525 : {
526 2400340 : unsigned int ib = batch * MAX_CACHED_DOF;
527 2400340 : unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_idofs());
528 :
529 16045308 : for (unsigned int i = ib; i < ie; ++i)
530 13644968 : local_ke[i - ib] = 0;
531 :
532 2400340 : body(local_ke - ib, ib, ie, j);
533 :
534 16045308 : for (unsigned int i = ib; i < ie; ++i)
535 13644968 : accumulateTaggedElementalMatrix(local_ke[i - ib], datum.elem().id, i, j, datum.jvar());
536 : }
537 : }
538 510130 : }
539 :
540 : } // namespace Moose::Kokkos
|
