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 16309 : 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 14526916 : ResidualObject::accumulateTaggedElementalResidual(const Real local_re,
208 : const ContiguousElementID elem,
209 : const unsigned int i,
210 : const unsigned int comp) const
211 : {
212 14526916 : if (!local_re)
213 930832 : return;
214 :
215 13596084 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
216 13596084 : auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
217 :
218 27879648 : for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
219 : {
220 14283564 : auto tag = _vector_tags[t];
221 :
222 14283564 : if (sys.isResidualTagActive(tag))
223 : {
224 : bool has_nodal_bc =
225 14282892 : _is_extra_vector_tag[t] ? sys.hasNodalBCResidualTag(dof, tag) : sys.hasNodalBC(dof);
226 :
227 14282892 : if (!has_nodal_bc)
228 13095404 : ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), local_re);
229 : }
230 : }
231 : }
232 :
233 : KOKKOS_FUNCTION inline void
234 1596618 : ResidualObject::accumulateTaggedNodalResidual(const bool add,
235 : const Real local_re,
236 : const ContiguousNodeID node,
237 : const unsigned int comp) const
238 : {
239 1596618 : if (!local_re)
240 987134 : 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 8861328 : 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 8861328 : if (!local_ke)
268 27690 : return;
269 :
270 8833638 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
271 8833638 : auto row = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
272 8833638 : auto col = sys.getElemGlobalDofIndex(elem, j, jvar);
273 :
274 22565932 : for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
275 : {
276 13732294 : auto tag = _matrix_tags[t];
277 :
278 13732294 : if (sys.isMatrixTagActive(tag))
279 : {
280 : bool has_nodal_bc =
281 9225430 : _is_extra_matrix_tag[t] ? sys.hasNodalBCMatrixTag(row, tag) : sys.hasNodalBC(row);
282 :
283 9225430 : if (!has_nodal_bc)
284 8316600 : ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke);
285 : }
286 : }
287 : }
288 :
289 : KOKKOS_FUNCTION inline void
290 491748 : 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 491748 : if (!local_ke)
297 157350 : return;
298 :
299 334398 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
300 334398 : auto row = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp));
301 334398 : auto col = sys.getNodeGlobalDofIndex(node, jvar);
302 :
303 777098 : for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
304 : {
305 442700 : auto tag = _matrix_tags[t];
306 :
307 442700 : if (sys.isMatrixTagActive(tag))
308 : {
309 351886 : auto & matrix = sys.getMatrix(tag);
310 :
311 351886 : if (add)
312 301590 : matrix(row, col) += local_ke;
313 : else
314 : {
315 50296 : matrix.zero(row);
316 50296 : matrix(row, col) = local_ke;
317 : }
318 : }
319 : }
320 : }
321 :
322 : template <typename function>
323 : KOKKOS_FUNCTION void
324 3651328 : ResidualObject::computeResidualInternal(AssemblyDatum & datum, function body) const
325 : {
326 : Real local_re[MAX_CACHED_DOF];
327 :
328 3651328 : unsigned int num_batches = datum.n_dofs() / MAX_CACHED_DOF;
329 :
330 3651328 : if (datum.n_dofs() % MAX_CACHED_DOF)
331 3651328 : ++num_batches;
332 :
333 7302656 : for (unsigned int batch = 0; batch < num_batches; ++batch)
334 : {
335 3651328 : unsigned int ib = batch * MAX_CACHED_DOF;
336 3651328 : unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_dofs());
337 :
338 18178244 : for (unsigned int i = ib; i < ie; ++i)
339 14526916 : local_re[i - ib] = 0;
340 :
341 3651328 : body(local_re - ib, ib, ie);
342 :
343 18178244 : for (unsigned int i = ib; i < ie; ++i)
344 14526916 : accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
345 : }
346 3651328 : }
347 :
348 : template <typename function>
349 : KOKKOS_FUNCTION void
350 317656 : ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) const
351 : {
352 : Real local_ke[MAX_CACHED_DOF];
353 :
354 317656 : unsigned int num_batches = datum.n_idofs() * datum.n_jdofs() / MAX_CACHED_DOF;
355 :
356 317656 : if ((datum.n_idofs() * datum.n_jdofs()) % MAX_CACHED_DOF)
357 317656 : ++num_batches;
358 :
359 659120 : for (unsigned int batch = 0; batch < num_batches; ++batch)
360 : {
361 341464 : unsigned int ijb = batch * MAX_CACHED_DOF;
362 341464 : unsigned int ije = ::Kokkos::min(ijb + MAX_CACHED_DOF, datum.n_idofs() * datum.n_jdofs());
363 :
364 5087336 : for (unsigned int ij = ijb; ij < ije; ++ij)
365 4745872 : local_ke[ij - ijb] = 0;
366 :
367 341464 : body(local_ke - ijb, ijb, ije);
368 :
369 5087336 : for (unsigned int ij = ijb; ij < ije; ++ij)
370 : {
371 4745872 : unsigned int i = ij % datum.n_jdofs();
372 4745872 : unsigned int j = ij / datum.n_jdofs();
373 :
374 4745872 : accumulateTaggedElementalMatrix(local_ke[ij - ijb], datum.elem().id, i, j, datum.jvar());
375 : }
376 : }
377 317656 : }
378 :
379 : } // namespace Kokkos
380 : } // namespace Moose
|
