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