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 "KokkosVariableValue.h"
13 : #include "KokkosMaterialPropertyValue.h"
14 :
15 : #include "MooseVariableBase.h"
16 : #include "ResidualObject.h"
17 :
18 : namespace Moose
19 : {
20 : namespace 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 SystemHolder
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 14815 : 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 0 : virtual void computeResidualAndJacobian() override final
69 : {
70 0 : computeResidual();
71 0 : computeJacobian();
72 0 : }
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 : /**
89 : * TODO: Move to TransientInterface
90 : */
91 : ///@{
92 : /**
93 : * Time
94 : */
95 : Scalar<Real> _t;
96 : /**
97 : * Old time
98 : */
99 : Scalar<const Real> _t_old;
100 : /**
101 : * The number of the time step
102 : */
103 : Scalar<int> _t_step;
104 : /**
105 : * Time step size
106 : */
107 : Scalar<Real> _dt;
108 : /**
109 : * Size of the old time step
110 : */
111 : Scalar<Real> _dt_old;
112 : ///@}
113 :
114 : /**
115 : * Accumulate local elemental residual contribution to tagged vectors
116 : * @param local_re The local elemental residual contribution
117 : * @param elem The contiguous element ID
118 : * @param i The test function index
119 : * @param comp The variable component
120 : */
121 : KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
122 : const ContiguousElementID elem,
123 : const unsigned int i,
124 : const unsigned int comp = 0) const;
125 : /**
126 : * Accumulate or set local nodal residual contribution to tagged vectors
127 : * @param add The flag whether to add or set the local residual
128 : * @param local_re The local nodal residual contribution
129 : * @param node The contiguous node ID
130 : * @param comp The variable component
131 : */
132 : KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
133 : const Real local_re,
134 : const ContiguousNodeID node,
135 : const unsigned int comp = 0) const;
136 : /**
137 : * Accumulate local elemental Jacobian contribution to tagged matrices
138 : * @param local_ke The local elemental Jacobian contribution
139 : * @param elem The contiguous element ID
140 : * @param i The test function DOF index
141 : * @param j The trial function DOF index
142 : * @param jvar The variable number for column
143 : * @param comp The variable component
144 : */
145 : KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke,
146 : const ContiguousElementID elem,
147 : const unsigned int i,
148 : const unsigned int j,
149 : const unsigned int jvar,
150 : const unsigned int comp = 0) const;
151 : /**
152 : * Accumulate or set local nodal Jacobian contribution to tagged matrices
153 : * @param add The flag whether to add or set the local residual
154 : * @param local_ke The local nodal Jacobian contribution
155 : * @param node The contiguous node ID
156 : * @param jvar The variable number for column
157 : * @param comp The Variable component
158 : */
159 : KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
160 : const Real local_ke,
161 : const ContiguousNodeID node,
162 : const unsigned int jvar,
163 : const unsigned int comp = 0) const;
164 :
165 : /**
166 : * The common loop structure template for computing elemental residual
167 : * @param datum The ResidualDatum object of the current thread
168 : * @param body The quadrature-point loop body
169 : */
170 : template <typename function>
171 : KOKKOS_FUNCTION void computeResidualInternal(ResidualDatum & datum, function body) const;
172 : /**
173 : * The common loop structure template for computing elemental Jacobian
174 : * @param datum The ResidualDatum object of the current thread
175 : * @param body The quadrature-point loop body
176 : */
177 : template <typename function>
178 : KOKKOS_FUNCTION void computeJacobianInternal(ResidualDatum & datum, function body) const;
179 :
180 : private:
181 : /**
182 : * Tags this object operates on
183 : */
184 : ///@{
185 : Array<TagID> _vector_tags;
186 : Array<TagID> _matrix_tags;
187 : ///@}
188 : /**
189 : * Flag whether each tag
190 : */
191 : ///@{
192 : Array<bool> _is_extra_vector_tag;
193 : Array<bool> _is_extra_matrix_tag;
194 : ///@}
195 : };
196 :
197 : KOKKOS_FUNCTION inline void
198 9890804 : ResidualObject::accumulateTaggedElementalResidual(const Real local_re,
199 : const ContiguousElementID elem,
200 : const unsigned int i,
201 : const unsigned int comp) const
202 : {
203 9890804 : if (!local_re)
204 632776 : return;
205 :
206 9258028 : auto & sys = kokkosSystem(_kokkos_var.sys());
207 9258028 : auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
208 :
209 19203536 : for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
210 : {
211 9945508 : auto tag = _vector_tags[t];
212 :
213 9945508 : if (sys.isResidualTagActive(tag))
214 : {
215 : bool has_nodal_bc =
216 9944836 : _is_extra_vector_tag[t] ? sys.hasNodalBCResidualTag(dof, tag) : sys.hasNodalBC(dof);
217 :
218 9944836 : if (!has_nodal_bc)
219 8968728 : ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), local_re);
220 : }
221 : }
222 : }
223 :
224 : KOKKOS_FUNCTION inline void
225 1535909 : ResidualObject::accumulateTaggedNodalResidual(const bool add,
226 : const Real local_re,
227 : const ContiguousNodeID node,
228 : const unsigned int comp) const
229 : {
230 1535909 : if (!local_re)
231 926425 : return;
232 :
233 609484 : auto & sys = kokkosSystem(_kokkos_var.sys());
234 609484 : auto dof = sys.getNodeLocalDofIndex(node, _kokkos_var.var(comp));
235 :
236 1535092 : for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
237 : {
238 925608 : auto tag = _vector_tags[t];
239 :
240 925608 : if (sys.isResidualTagActive(tag))
241 : {
242 925572 : if (add)
243 922299 : sys.getVectorDofValue(dof, tag) += local_re;
244 : else
245 3273 : sys.getVectorDofValue(dof, tag) = local_re;
246 : }
247 : }
248 : }
249 :
250 : KOKKOS_FUNCTION inline void
251 6128224 : ResidualObject::accumulateTaggedElementalMatrix(const Real local_ke,
252 : const ContiguousElementID elem,
253 : const unsigned int i,
254 : const unsigned int j,
255 : const unsigned int jvar,
256 : const unsigned int comp) const
257 : {
258 6128224 : if (!local_ke)
259 27690 : return;
260 :
261 6100534 : auto & sys = kokkosSystem(_kokkos_var.sys());
262 6100534 : auto row = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
263 6100534 : auto col = sys.getElemGlobalDofIndex(elem, j, jvar);
264 :
265 15733020 : for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
266 : {
267 9632486 : auto tag = _matrix_tags[t];
268 :
269 9632486 : if (sys.isMatrixTagActive(tag))
270 : {
271 : bool has_nodal_bc =
272 6492326 : _is_extra_matrix_tag[t] ? sys.hasNodalBCMatrixTag(row, tag) : sys.hasNodalBC(row);
273 :
274 6492326 : if (!has_nodal_bc)
275 5720136 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke);
276 : }
277 : }
278 : }
279 :
280 : KOKKOS_FUNCTION inline void
281 482354 : ResidualObject::accumulateTaggedNodalMatrix(const bool add,
282 : const Real local_ke,
283 : const ContiguousNodeID node,
284 : const unsigned int jvar,
285 : const unsigned int comp) const
286 : {
287 482354 : if (!local_ke)
288 157350 : return;
289 :
290 325004 : auto & sys = kokkosSystem(_kokkos_var.sys());
291 325004 : auto row = sys.getNodeLocalDofIndex(node, _kokkos_var.var(comp));
292 325004 : auto col = sys.getNodeGlobalDofIndex(node, jvar);
293 :
294 758310 : for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
295 : {
296 433306 : auto tag = _matrix_tags[t];
297 :
298 433306 : if (sys.isMatrixTagActive(tag))
299 : {
300 342492 : auto & matrix = sys.getMatrix(tag);
301 :
302 342492 : if (add)
303 301590 : matrix(row, col) += local_ke;
304 : else
305 : {
306 40902 : matrix.zero(row);
307 40902 : matrix(row, col) = local_ke;
308 : }
309 : }
310 : }
311 : }
312 :
313 : template <typename function>
314 : KOKKOS_FUNCTION void
315 2492300 : ResidualObject::computeResidualInternal(ResidualDatum & datum, function body) const
316 : {
317 : Real local_re[MAX_CACHED_DOF];
318 :
319 2492300 : unsigned int num_batches = datum.n_dofs() / MAX_CACHED_DOF;
320 :
321 2492300 : if (datum.n_dofs() % MAX_CACHED_DOF)
322 2492300 : ++num_batches;
323 :
324 4984600 : for (unsigned int batch = 0; batch < num_batches; ++batch)
325 : {
326 2492300 : unsigned int ib = batch * MAX_CACHED_DOF;
327 2492300 : unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_dofs());
328 :
329 12383104 : for (unsigned int i = ib; i < ie; ++i)
330 9890804 : local_re[i - ib] = 0;
331 :
332 2492300 : body(local_re - ib, ib, ie);
333 :
334 12383104 : for (unsigned int i = ib; i < ie; ++i)
335 9890804 : accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
336 : }
337 2492300 : }
338 :
339 : template <typename function>
340 : KOKKOS_FUNCTION void
341 232256 : ResidualObject::computeJacobianInternal(ResidualDatum & datum, function body) const
342 : {
343 : Real local_ke[MAX_CACHED_DOF];
344 :
345 232256 : unsigned int num_batches = datum.n_idofs() * datum.n_jdofs() / MAX_CACHED_DOF;
346 :
347 232256 : if ((datum.n_idofs() * datum.n_jdofs()) % MAX_CACHED_DOF)
348 232256 : ++num_batches;
349 :
350 488320 : for (unsigned int batch = 0; batch < num_batches; ++batch)
351 : {
352 256064 : unsigned int ijb = batch * MAX_CACHED_DOF;
353 256064 : unsigned int ije = ::Kokkos::min(ijb + MAX_CACHED_DOF, datum.n_idofs() * datum.n_jdofs());
354 :
355 3635536 : for (unsigned int ij = ijb; ij < ije; ++ij)
356 3379472 : local_ke[ij - ijb] = 0;
357 :
358 256064 : body(local_ke - ijb, ijb, ije);
359 :
360 3635536 : for (unsigned int ij = ijb; ij < ije; ++ij)
361 : {
362 3379472 : unsigned int i = ij % datum.n_jdofs();
363 3379472 : unsigned int j = ij / datum.n_jdofs();
364 :
365 3379472 : accumulateTaggedElementalMatrix(local_ke[ij - ijb], datum.elem().id, i, j, datum.jvar());
366 : }
367 : }
368 232256 : }
369 :
370 : } // namespace Kokkos
371 : } // namespace Moose
372 :
373 : #define usingKokkosResidualObjectMembers \
374 : public: \
375 : usingPostprocessorInterfaceMembers; \
376 : \
377 : protected: \
378 : using Moose::Kokkos::ResidualObject::kokkosAssembly; \
379 : using Moose::Kokkos::ResidualObject::kokkosSystems; \
380 : using Moose::Kokkos::ResidualObject::kokkosSystem; \
381 : using Moose::Kokkos::ResidualObject::accumulateTaggedElementalResidual; \
382 : using Moose::Kokkos::ResidualObject::accumulateTaggedNodalResidual; \
383 : using Moose::Kokkos::ResidualObject::accumulateTaggedElementalMatrix; \
384 : using Moose::Kokkos::ResidualObject::accumulateTaggedNodalMatrix; \
385 : using Moose::Kokkos::ResidualObject::_var; \
386 : using Moose::Kokkos::ResidualObject::_kokkos_var; \
387 : using Moose::Kokkos::ResidualObject::_thread; \
388 : using Moose::Kokkos::ResidualObject::_t; \
389 : using Moose::Kokkos::ResidualObject::_t_old; \
390 : using Moose::Kokkos::ResidualObject::_t_step; \
391 : using Moose::Kokkos::ResidualObject::_dt; \
392 : using Moose::Kokkos::ResidualObject::_dt_old; \
393 : \
394 : public: \
395 : using Moose::Kokkos::ResidualObject::ResidualLoop; \
396 : using Moose::Kokkos::ResidualObject::JacobianLoop; \
397 : using Moose::Kokkos::ResidualObject::OffDiagJacobianLoop
|
