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 : /**
35 : * Constructor
36 : * @param field_type The MOOSE variable field type
37 : * @param nodal Whether the residual object is a nodal residual object
38 : */
39 : ResidualObject(const InputParameters & parameters,
40 : Moose::VarFieldType field_type,
41 : bool nodal = false);
42 : /**
43 : * Copy constructor for parallel dispatch
44 : */
45 : ResidualObject(const ResidualObject & object);
46 :
47 : /**
48 : * Kokkos function tags
49 : */
50 : ///@{
51 : struct ResidualLoop
52 : {
53 : };
54 : struct JacobianLoop
55 : {
56 : };
57 : struct OffDiagJacobianLoop
58 : {
59 : };
60 : ///@}
61 :
62 41591 : virtual const MooseVariableBase & variable() const override { return _var; }
63 :
64 0 : virtual void computeOffDiagJacobian(unsigned int) override final
65 : {
66 0 : mooseError("computeOffDiagJacobian() is not used for Kokkos residual objects.");
67 : }
68 1860 : virtual void computeResidualAndJacobian() override
69 : {
70 1860 : computeResidual();
71 1860 : computeJacobian();
72 1860 : }
73 :
74 : protected:
75 : /**
76 : * Reference of the MOOSE variable
77 : */
78 : MooseVariableFieldBase & _var;
79 : /**
80 : * Kokkos variable
81 : */
82 : Variable _kokkos_var;
83 : /**
84 : * Kokkos thread object
85 : */
86 : Thread<> _thread;
87 : /**
88 : * Kokkos functor dispatchers
89 : */
90 : ///@{
91 : std::unique_ptr<DispatcherBase> _residual_dispatcher;
92 : std::unique_ptr<DispatcherBase> _jacobian_dispatcher;
93 : std::unique_ptr<DispatcherBase> _offdiag_jacobian_dispatcher;
94 : ///@}
95 :
96 : /**
97 : * TODO: Move to TransientInterface
98 : */
99 : ///@{
100 : /**
101 : * Time
102 : */
103 : Scalar<Real> _t;
104 : /**
105 : * Old time
106 : */
107 : Scalar<const Real> _t_old;
108 : /**
109 : * The number of the time step
110 : */
111 : Scalar<int> _t_step;
112 : /**
113 : * Time step size
114 : */
115 : Scalar<Real> _dt;
116 : /**
117 : * Size of the old time step
118 : */
119 : Scalar<Real> _dt_old;
120 : ///@}
121 :
122 : /**
123 : * Accumulate local elemental residual contribution to tagged vectors
124 : * @param local_re The local elemental residual contribution
125 : * @param elem The contiguous element ID
126 : * @param i The test function index
127 : * @param comp The variable component
128 : */
129 : KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
130 : const ContiguousElementID elem,
131 : const unsigned int i,
132 : const unsigned int comp = 0) const;
133 : /**
134 : * Accumulate or set local nodal residual contribution to tagged vectors
135 : * @param add Whether to add or set the local residual
136 : * @param local_re The local nodal residual contribution
137 : * @param node The contiguous node ID
138 : * @param comp The variable component
139 : */
140 : KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
141 : const Real local_re,
142 : const ContiguousNodeID node,
143 : const unsigned int comp = 0) const;
144 : /**
145 : * Accumulate local elemental Jacobian contribution to tagged matrices
146 : * @param local_ke The local elemental Jacobian contribution
147 : * @param elem The contiguous element ID
148 : * @param i The test function DOF index
149 : * @param j The trial function DOF index
150 : * @param jvar The variable number for column
151 : * @param comp The variable component
152 : */
153 : KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke,
154 : const ContiguousElementID elem,
155 : const unsigned int i,
156 : const unsigned int j,
157 : const unsigned int jvar,
158 : const unsigned int comp = 0) const;
159 : /**
160 : * Accumulate local elemental Jacobian contribution to tagged matrices using automatic
161 : * differentiation (AD)
162 : * @param local_ke The local elemental Jacobian contribution
163 : * @param datum The AssemblyDatum object of the current thread
164 : * @param i The test function DOF index
165 : * @param comp The variable component
166 : */
167 : KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
168 : const AssemblyDatum & datum,
169 : const unsigned int i,
170 : const unsigned int comp = 0) const;
171 : /**
172 : * Accumulate or set local nodal Jacobian contribution to tagged matrices
173 : * @param add Whether to add or set the local Jacobian
174 : * @param local_ke The local nodal Jacobian contribution
175 : * @param node The contiguous node ID
176 : * @param jvar The variable number for column
177 : * @param comp The variable component
178 : */
179 : KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
180 : const Real local_ke,
181 : const ContiguousNodeID node,
182 : const unsigned int jvar,
183 : const unsigned int comp = 0) const;
184 : /**
185 : * Accumulate or set local nodal Jacobian contribution to tagged matrices using automatic
186 : * differentiation (AD)
187 : * @param add Whether to add or set the local Jacobian
188 : * @param local_ke The local elemental Jacobian contribution
189 : * @param node The contiguous node ID
190 : * @param comp The variable component
191 : */
192 : KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
193 : const DNDerivativeType & local_ke,
194 : const ContiguousNodeID node,
195 : const unsigned int comp = 0) const;
196 :
197 : /**
198 : * The common loop structure template for computing elemental residual
199 : * @param datum The AssemblyDatum object of the current thread
200 : * @param body The quadrature point loop body
201 : */
202 : template <typename function>
203 : KOKKOS_FUNCTION void computeResidualInternal(AssemblyDatum & datum, function body) const;
204 : /**
205 : * The common loop structure template for computing elemental Jacobian
206 : * @param datum The AssemblyDatum object of the current thread
207 : * @param body The quadrature point loop body
208 : */
209 : template <typename function>
210 : KOKKOS_FUNCTION void computeJacobianInternal(AssemblyDatum & datum, function body) const;
211 :
212 : private:
213 : /**
214 : * Tags this object operates on
215 : */
216 : ///@{
217 : Array<TagID> _vector_tags;
218 : Array<TagID> _matrix_tags;
219 : ///@}
220 : };
221 :
222 : KOKKOS_FUNCTION inline void
223 17991826 : ResidualObject::accumulateTaggedElementalResidual(const Real local_re,
224 : const ContiguousElementID elem,
225 : const unsigned int i,
226 : const unsigned int comp) const
227 : {
228 17991826 : if (!local_re)
229 1524573 : return;
230 :
231 16467253 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
232 16467253 : auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
233 :
234 33741282 : for (unsigned int t = 0; t < _vector_tags.size(); ++t)
235 : {
236 17274029 : auto tag = _vector_tags[t];
237 :
238 17274029 : if (sys.isResidualTagActive(tag))
239 17154061 : ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), local_re);
240 : }
241 : }
242 :
243 : KOKKOS_FUNCTION inline void
244 1698340 : ResidualObject::accumulateTaggedNodalResidual(const bool add,
245 : const Real local_re,
246 : const ContiguousNodeID node,
247 : const unsigned int comp) const
248 : {
249 1698340 : if (!local_re && add)
250 660339 : return;
251 :
252 1038001 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
253 1038001 : auto dof = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
254 :
255 2422012 : for (unsigned int t = 0; t < _vector_tags.size(); ++t)
256 : {
257 1384011 : auto tag = _vector_tags[t];
258 :
259 1384011 : if (sys.isResidualTagActive(tag))
260 : {
261 1373751 : if (add)
262 921195 : sys.getVectorDofValue(dof, tag) += local_re;
263 : else
264 452556 : sys.getVectorDofValue(dof, tag) = local_re;
265 : }
266 : }
267 : }
268 :
269 : KOKKOS_FUNCTION inline void
270 9996452 : ResidualObject::accumulateTaggedElementalMatrix(const Real local_ke,
271 : const ContiguousElementID elem,
272 : const unsigned int i,
273 : const unsigned int j,
274 : const unsigned int jvar,
275 : const unsigned int comp) const
276 : {
277 9996452 : if (!local_ke)
278 27753 : return;
279 :
280 9968699 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
281 9968699 : auto row = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
282 9968699 : auto col = sys.getElemGlobalDofIndex(elem, j, jvar);
283 :
284 25324758 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
285 : {
286 15356059 : auto tag = _matrix_tags[t];
287 :
288 15356059 : if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))
289 9385509 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke);
290 : }
291 : }
292 :
293 : KOKKOS_FUNCTION inline void
294 1103456 : ResidualObject::accumulateTaggedElementalMatrix(const DNDerivativeType & local_ke,
295 : const AssemblyDatum & datum,
296 : const unsigned int i,
297 : const unsigned int comp) const
298 : {
299 1103456 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
300 1103456 : auto row = sys.getElemLocalDofIndex(datum.elem().id, i, _kokkos_var.var(comp));
301 :
302 2706412 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
303 : {
304 1602956 : auto tag = _matrix_tags[t];
305 :
306 1602956 : if (sys.isMatrixTagActive(tag) && !sys.hasNodalBCMatrixTag(row, tag))
307 4879930 : for (unsigned int j = 0; j < local_ke.size(); ++j)
308 : {
309 3891220 : auto col = local_ke.raw_index(j);
310 :
311 3891220 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke.raw_at(j));
312 : }
313 : }
314 1103456 : }
315 :
316 : KOKKOS_FUNCTION inline void
317 499839 : ResidualObject::accumulateTaggedNodalMatrix(const bool add,
318 : const Real local_ke,
319 : const ContiguousNodeID node,
320 : const unsigned int jvar,
321 : const unsigned int comp) const
322 : {
323 499839 : if (!local_ke && add)
324 157350 : return;
325 :
326 342489 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
327 342489 : auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
328 342489 : auto col = sys.getNodeGlobalDofIndex(node, jvar);
329 :
330 852745 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
331 : {
332 510256 : auto tag = _matrix_tags[t];
333 :
334 510256 : if (sys.isMatrixTagActive(tag))
335 : {
336 359977 : auto & matrix = sys.getMatrix(tag);
337 :
338 359977 : if (add)
339 301590 : matrix(row, col) += local_ke;
340 : else
341 : {
342 58387 : matrix.zero(row);
343 58387 : matrix(row, col) = local_ke;
344 : }
345 : }
346 : }
347 : }
348 :
349 : KOKKOS_FUNCTION inline void
350 33390 : ResidualObject::accumulateTaggedNodalMatrix(const bool add,
351 : const DNDerivativeType & local_ke,
352 : const ContiguousNodeID node,
353 : const unsigned int comp) const
354 : {
355 33390 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
356 33390 : auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
357 :
358 100170 : for (unsigned int t = 0; t < _matrix_tags.size(); ++t)
359 : {
360 66780 : auto tag = _matrix_tags[t];
361 66780 : auto & matrix = sys.getMatrix(tag);
362 :
363 66780 : if (sys.isMatrixTagActive(tag))
364 : {
365 33390 : if (!add)
366 33390 : matrix.zero(row);
367 :
368 66780 : for (unsigned int j = 0; j < local_ke.size(); ++j)
369 : {
370 33390 : auto col = local_ke.raw_index(j);
371 :
372 33390 : if (add)
373 0 : matrix(row, col) += local_ke.raw_at(j);
374 : else
375 33390 : matrix(row, col) = local_ke.raw_at(j);
376 : }
377 : }
378 : }
379 33390 : }
380 :
381 : template <typename function>
382 : KOKKOS_FUNCTION void
383 4221911 : ResidualObject::computeResidualInternal(AssemblyDatum & datum, function body) const
384 : {
385 : Real local_re[MAX_CACHED_DOF];
386 :
387 4221911 : unsigned int stride = MAX_CACHED_DOF * datum.num_local_threads();
388 4221911 : unsigned int num_batches = datum.n_dofs() / stride;
389 :
390 4221911 : if (datum.n_dofs() % stride)
391 4221911 : ++num_batches;
392 :
393 8443822 : for (unsigned int batch = 0; batch < num_batches; ++batch)
394 : {
395 4221911 : unsigned int ib = batch * stride;
396 4221911 : unsigned int ie = ::Kokkos::min(ib + stride, datum.n_dofs());
397 :
398 4221911 : const unsigned int n = ie - ib;
399 4221911 : const unsigned int d = n / datum.num_local_threads();
400 4221911 : const unsigned int m = n % datum.num_local_threads();
401 4221911 : const unsigned int t = datum.local_thread_id();
402 :
403 4221911 : ib += t * d + (t < m ? t : m);
404 4221911 : ie = ib + d + (t < m ? 1 : 0);
405 :
406 21031549 : for (unsigned int i = ib; i < ie; ++i)
407 16809638 : local_re[i - ib] = 0;
408 :
409 4221911 : body(local_re - ib, ib, ie);
410 :
411 21031549 : for (unsigned int i = ib; i < ie; ++i)
412 16809638 : accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
413 : }
414 4221911 : }
415 :
416 : template <typename function>
417 : KOKKOS_FUNCTION void
418 364217 : ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) const
419 : {
420 : Real local_ke[MAX_CACHED_DOF];
421 :
422 1717099 : for (unsigned int j = datum.local_thread_id(); j < datum.n_jdofs();
423 1352882 : j += datum.num_local_threads())
424 : {
425 1352882 : unsigned int num_batches = datum.n_idofs() / MAX_CACHED_DOF;
426 :
427 1352882 : if (datum.n_idofs() % MAX_CACHED_DOF)
428 1352882 : ++num_batches;
429 :
430 2705764 : for (unsigned int batch = 0; batch < num_batches; ++batch)
431 : {
432 1352882 : unsigned int ib = batch * MAX_CACHED_DOF;
433 1352882 : unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_idofs());
434 :
435 6843478 : for (unsigned int i = ib; i < ie; ++i)
436 5490596 : local_ke[i - ib] = 0;
437 :
438 1352882 : body(local_ke - ib, ib, ie, j);
439 :
440 6843478 : for (unsigned int i = ib; i < ie; ++i)
441 5490596 : accumulateTaggedElementalMatrix(local_ke[i - ib], datum.elem().id, i, j, datum.jvar());
442 : }
443 : }
444 364217 : }
445 :
446 : } // namespace Moose::Kokkos
|
