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
20 : {
21 : namespace Kokkos
22 : {
23 :
24 : /**
25 : * The base class for Kokkos residual objects
26 : */
27 : class ResidualObject : public ::ResidualObject,
28 : public MeshHolder,
29 : public AssemblyHolder,
30 : public SystemHolder
31 : {
32 : public:
33 : static InputParameters validParams();
34 :
35 : /**
36 : * Constructor
37 : * @param field_type The MOOSE variable field type
38 : * @param nodal Whether the residual object is a nodal residual object
39 : */
40 : ResidualObject(const InputParameters & parameters,
41 : Moose::VarFieldType field_type,
42 : bool nodal = false);
43 : /**
44 : * Copy constructor for parallel dispatch
45 : */
46 : ResidualObject(const ResidualObject & object);
47 :
48 : /**
49 : * Kokkos function tags
50 : */
51 : ///@{
52 : struct ResidualLoop
53 : {
54 : };
55 : struct JacobianLoop
56 : {
57 : };
58 : struct OffDiagJacobianLoop
59 : {
60 : };
61 : ///@}
62 :
63 15043 : virtual const MooseVariableBase & variable() const override { return _var; }
64 :
65 0 : virtual void computeOffDiagJacobian(unsigned int) override final
66 : {
67 0 : mooseError("computeOffDiagJacobian() is not used for Kokkos residual objects.");
68 : }
69 0 : virtual void computeResidualAndJacobian() override final
70 : {
71 0 : computeResidual();
72 0 : computeJacobian();
73 0 : }
74 :
75 : protected:
76 : /**
77 : * Reference of the MOOSE variable
78 : */
79 : MooseVariableFieldBase & _var;
80 : /**
81 : * Kokkos variable
82 : */
83 : Variable _kokkos_var;
84 : /**
85 : * Kokkos thread object
86 : */
87 : Thread _thread;
88 : /**
89 : * Kokkos functor dispatchers
90 : */
91 : ///@{
92 : std::unique_ptr<DispatcherBase> _residual_dispatcher;
93 : std::unique_ptr<DispatcherBase> _jacobian_dispatcher;
94 : std::unique_ptr<DispatcherBase> _offdiag_jacobian_dispatcher;
95 : ///@}
96 :
97 : /**
98 : * TODO: Move to TransientInterface
99 : */
100 : ///@{
101 : /**
102 : * Time
103 : */
104 : Scalar<Real> _t;
105 : /**
106 : * Old time
107 : */
108 : Scalar<const Real> _t_old;
109 : /**
110 : * The number of the time step
111 : */
112 : Scalar<int> _t_step;
113 : /**
114 : * Time step size
115 : */
116 : Scalar<Real> _dt;
117 : /**
118 : * Size of the old time step
119 : */
120 : Scalar<Real> _dt_old;
121 : ///@}
122 :
123 : /**
124 : * Accumulate local elemental residual contribution to tagged vectors
125 : * @param local_re The local elemental residual contribution
126 : * @param elem The contiguous element ID
127 : * @param i The test function index
128 : * @param comp The variable component
129 : */
130 : KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
131 : const ContiguousElementID elem,
132 : const unsigned int i,
133 : const unsigned int comp = 0) const;
134 : /**
135 : * Accumulate or set local nodal residual contribution to tagged vectors
136 : * @param add The flag whether to add or set the local residual
137 : * @param local_re The local nodal residual contribution
138 : * @param node The contiguous node ID
139 : * @param comp The variable component
140 : */
141 : KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
142 : const Real local_re,
143 : const ContiguousNodeID node,
144 : const unsigned int comp = 0) const;
145 : /**
146 : * Accumulate local elemental Jacobian contribution to tagged matrices
147 : * @param local_ke The local elemental Jacobian contribution
148 : * @param elem The contiguous element ID
149 : * @param i The test function DOF index
150 : * @param j The trial function DOF index
151 : * @param jvar The variable number for column
152 : * @param comp The variable component
153 : */
154 : KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke,
155 : const ContiguousElementID elem,
156 : const unsigned int i,
157 : const unsigned int j,
158 : const unsigned int jvar,
159 : const unsigned int comp = 0) const;
160 : /**
161 : * Accumulate or set local nodal Jacobian contribution to tagged matrices
162 : * @param add The flag whether to add or set the local residual
163 : * @param local_ke The local nodal Jacobian contribution
164 : * @param node The contiguous node ID
165 : * @param jvar The variable number for column
166 : * @param comp The variable component
167 : */
168 : KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
169 : const Real local_ke,
170 : const ContiguousNodeID node,
171 : const unsigned int jvar,
172 : const unsigned int comp = 0) const;
173 :
174 : /**
175 : * The common loop structure template for computing elemental residual
176 : * @param datum The AssemblyDatum object of the current thread
177 : * @param body The quadrature point loop body
178 : */
179 : template <typename function>
180 : KOKKOS_FUNCTION void computeResidualInternal(AssemblyDatum & datum, function body) const;
181 : /**
182 : * The common loop structure template for computing elemental Jacobian
183 : * @param datum The AssemblyDatum object of the current thread
184 : * @param body The quadrature point loop body
185 : */
186 : template <typename function>
187 : KOKKOS_FUNCTION void computeJacobianInternal(AssemblyDatum & datum, function body) const;
188 :
189 : private:
190 : /**
191 : * Tags this object operates on
192 : */
193 : ///@{
194 : Array<TagID> _vector_tags;
195 : Array<TagID> _matrix_tags;
196 : ///@}
197 : /**
198 : * Flag whether each tag
199 : */
200 : ///@{
201 : Array<bool> _is_extra_vector_tag;
202 : Array<bool> _is_extra_matrix_tag;
203 : ///@}
204 : };
205 :
206 : KOKKOS_FUNCTION inline void
207 9892636 : ResidualObject::accumulateTaggedElementalResidual(const Real local_re,
208 : const ContiguousElementID elem,
209 : const unsigned int i,
210 : const unsigned int comp) const
211 : {
212 9892636 : if (!local_re)
213 632892 : return;
214 :
215 9259744 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
216 9259744 : auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
217 :
218 19206968 : for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
219 : {
220 9947224 : auto tag = _vector_tags[t];
221 :
222 9947224 : if (sys.isResidualTagActive(tag))
223 : {
224 : bool has_nodal_bc =
225 9946552 : _is_extra_vector_tag[t] ? sys.hasNodalBCResidualTag(dof, tag) : sys.hasNodalBC(dof);
226 :
227 9946552 : if (!has_nodal_bc)
228 8970364 : ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), local_re);
229 : }
230 : }
231 : }
232 :
233 : KOKKOS_FUNCTION inline void
234 1535931 : ResidualObject::accumulateTaggedNodalResidual(const bool add,
235 : const Real local_re,
236 : const ContiguousNodeID node,
237 : const unsigned int comp) const
238 : {
239 1535931 : if (!local_re)
240 926447 : return;
241 :
242 609484 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
243 609484 : auto dof = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
244 :
245 1535092 : for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
246 : {
247 925608 : auto tag = _vector_tags[t];
248 :
249 925608 : if (sys.isResidualTagActive(tag))
250 : {
251 925572 : if (add)
252 922299 : sys.getVectorDofValue(dof, tag) += local_re;
253 : else
254 3273 : sys.getVectorDofValue(dof, tag) = local_re;
255 : }
256 : }
257 : }
258 :
259 : KOKKOS_FUNCTION inline void
260 6128528 : ResidualObject::accumulateTaggedElementalMatrix(const Real local_ke,
261 : const ContiguousElementID elem,
262 : const unsigned int i,
263 : const unsigned int j,
264 : const unsigned int jvar,
265 : const unsigned int comp) const
266 : {
267 6128528 : if (!local_ke)
268 27690 : return;
269 :
270 6100838 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
271 6100838 : auto row = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
272 6100838 : auto col = sys.getElemGlobalDofIndex(elem, j, jvar);
273 :
274 15733932 : for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
275 : {
276 9633094 : auto tag = _matrix_tags[t];
277 :
278 9633094 : if (sys.isMatrixTagActive(tag))
279 : {
280 : bool has_nodal_bc =
281 6492630 : _is_extra_matrix_tag[t] ? sys.hasNodalBCMatrixTag(row, tag) : sys.hasNodalBC(row);
282 :
283 6492630 : if (!has_nodal_bc)
284 5720440 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke);
285 : }
286 : }
287 : }
288 :
289 : KOKKOS_FUNCTION inline void
290 482354 : ResidualObject::accumulateTaggedNodalMatrix(const bool add,
291 : const Real local_ke,
292 : const ContiguousNodeID node,
293 : const unsigned int jvar,
294 : const unsigned int comp) const
295 : {
296 482354 : if (!local_ke)
297 157350 : return;
298 :
299 325004 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
300 325004 : auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
301 325004 : auto col = sys.getNodeGlobalDofIndex(node, jvar);
302 :
303 758310 : for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
304 : {
305 433306 : auto tag = _matrix_tags[t];
306 :
307 433306 : if (sys.isMatrixTagActive(tag))
308 : {
309 342492 : auto & matrix = sys.getMatrix(tag);
310 :
311 342492 : if (add)
312 301590 : matrix(row, col) += local_ke;
313 : else
314 : {
315 40902 : matrix.zero(row);
316 40902 : matrix(row, col) = local_ke;
317 : }
318 : }
319 : }
320 : }
321 :
322 : template <typename function>
323 : KOKKOS_FUNCTION void
324 2492758 : ResidualObject::computeResidualInternal(AssemblyDatum & datum, function body) const
325 : {
326 : Real local_re[MAX_CACHED_DOF];
327 :
328 2492758 : unsigned int num_batches = datum.n_dofs() / MAX_CACHED_DOF;
329 :
330 2492758 : if (datum.n_dofs() % MAX_CACHED_DOF)
331 2492758 : ++num_batches;
332 :
333 4985516 : for (unsigned int batch = 0; batch < num_batches; ++batch)
334 : {
335 2492758 : unsigned int ib = batch * MAX_CACHED_DOF;
336 2492758 : unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_dofs());
337 :
338 12385394 : for (unsigned int i = ib; i < ie; ++i)
339 9892636 : local_re[i - ib] = 0;
340 :
341 2492758 : body(local_re - ib, ib, ie);
342 :
343 12385394 : for (unsigned int i = ib; i < ie; ++i)
344 9892636 : accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
345 : }
346 2492758 : }
347 :
348 : template <typename function>
349 : KOKKOS_FUNCTION void
350 232256 : ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) const
351 : {
352 : Real local_ke[MAX_CACHED_DOF];
353 :
354 232256 : unsigned int num_batches = datum.n_idofs() * datum.n_jdofs() / MAX_CACHED_DOF;
355 :
356 232256 : if ((datum.n_idofs() * datum.n_jdofs()) % MAX_CACHED_DOF)
357 232256 : ++num_batches;
358 :
359 488320 : for (unsigned int batch = 0; batch < num_batches; ++batch)
360 : {
361 256064 : unsigned int ijb = batch * MAX_CACHED_DOF;
362 256064 : unsigned int ije = ::Kokkos::min(ijb + MAX_CACHED_DOF, datum.n_idofs() * datum.n_jdofs());
363 :
364 3635536 : for (unsigned int ij = ijb; ij < ije; ++ij)
365 3379472 : local_ke[ij - ijb] = 0;
366 :
367 256064 : body(local_ke - ijb, ijb, ije);
368 :
369 3635536 : for (unsigned int ij = ijb; ij < ije; ++ij)
370 : {
371 3379472 : unsigned int i = ij % datum.n_jdofs();
372 3379472 : unsigned int j = ij / datum.n_jdofs();
373 :
374 3379472 : accumulateTaggedElementalMatrix(local_ke[ij - ijb], datum.elem().id, i, j, datum.jvar());
375 : }
376 : }
377 232256 : }
378 :
379 : } // namespace Kokkos
380 : } // namespace Moose
|
