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 "KokkosKernelBase.h"
13 :
14 : namespace Moose
15 : {
16 : namespace Kokkos
17 : {
18 :
19 : /**
20 : * The base class for a user to derive their own Kokkos kernels.
21 : *
22 : * The user should define computeQpResidual(), computeQpJacobian(), and computeQpOffDiagJacobian()
23 : * as inlined public methods in their derived class (not virtual override). The signature of
24 : * computeQpResidual() expected to be defined in the derived class is as follows:
25 : *
26 : * @param i The element-local DOF index
27 : * @param qp The local quadrature point index
28 : * @param datum The AssemblyDatum object of the current thread
29 : * @returns The residual contribution
30 : *
31 : * KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
32 : * const unsigned int qp,
33 : * AssemblyDatum & datum) const;
34 : *
35 : * The signatures of computeQpJacobian() and computeOffDiagQpJacobian() can be found in the code
36 : * below, and their definition in the derived class is optional. If they are defined in the derived
37 : * class, they will hide the default definitions in the base class.
38 : */
39 : class Kernel : public KernelBase
40 : {
41 : public:
42 : static InputParameters validParams();
43 :
44 : /**
45 : * Constructor
46 : */
47 : Kernel(const InputParameters & parameters);
48 :
49 : /**
50 : * Dispatch residual calculation
51 : */
52 : virtual void computeResidual() override;
53 : /**
54 : * Dispatch diagonal and off-diagonal Jacobian calculation
55 : */
56 : virtual void computeJacobian() override;
57 :
58 : /**
59 : * Default methods to prevent compile errors even when these methods were not defined in the
60 : * derived class
61 : */
62 : ///@{
63 : /**
64 : * Compute diagonal Jacobian contribution on a quadrature point
65 : * @param i The test function DOF index
66 : * @param j The trial function DOF index
67 : * @param qp The local quadrature point index
68 : * @param datum The AssemblyDatum object of the current thread
69 : * @returns The diagonal Jacobian contribution
70 : */
71 0 : KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
72 : const unsigned int /* j */,
73 : const unsigned int /* qp */,
74 : AssemblyDatum & /* datum */) const
75 : {
76 0 : return 0;
77 : }
78 : /**
79 : * Compute off-diagonal Jacobian contribution on a quadrature point
80 : * @param i The test function DOF index
81 : * @param j The trial function DOF index
82 : * @param jvar The variable number for column
83 : * @param qp The local quadrature point index
84 : * @param datum The AssemblyDatum object of the current thread
85 : * @returns The off-diagonal Jacobian contribution
86 : */
87 0 : KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
88 : const unsigned int /* j */,
89 : const unsigned int /* jvar */,
90 : const unsigned int /* qp */,
91 : AssemblyDatum & /* datum */) const
92 : {
93 0 : return 0;
94 : }
95 : /**
96 : * Get the function pointer of the default computeQpJacobian()
97 : * @returns The function pointer
98 : */
99 489352 : static auto defaultJacobian() { return &Kernel::computeQpJacobian; }
100 : /**
101 : * Get the function pointer of the default computeQpOffDiagJacobian()
102 : * @returns The function pointer
103 : */
104 564610 : static auto defaultOffDiagJacobian() { return &Kernel::computeQpOffDiagJacobian; }
105 : ///@}
106 :
107 : /**
108 : * Shims for hook methods that can be leveraged to implement static polymorphism
109 : */
110 : ///{@
111 : template <typename Derived>
112 65637808 : KOKKOS_FUNCTION Real computeQpResidualShim(const Derived & kernel,
113 : const unsigned int i,
114 : const unsigned int qp,
115 : AssemblyDatum & datum) const
116 : {
117 65637808 : return kernel.computeQpResidual(i, qp, datum);
118 : }
119 : template <typename Derived>
120 42536864 : KOKKOS_FUNCTION Real computeQpJacobianShim(const Derived & kernel,
121 : const unsigned int i,
122 : const unsigned int j,
123 : const unsigned int qp,
124 : AssemblyDatum & datum) const
125 : {
126 42536864 : return kernel.computeQpJacobian(i, j, qp, datum);
127 : }
128 : template <typename Derived>
129 80000 : KOKKOS_FUNCTION Real computeQpOffDiagJacobianShim(const Derived & kernel,
130 : const unsigned int i,
131 : const unsigned int j,
132 : const unsigned int jvar,
133 : const unsigned int qp,
134 : AssemblyDatum & datum) const
135 : {
136 80000 : return kernel.computeQpOffDiagJacobian(i, j, jvar, qp, datum);
137 : }
138 : ///@}
139 :
140 : /**
141 : * The parallel computation entry functions called by Kokkos
142 : */
143 : ///@{
144 : template <typename Derived>
145 : KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
146 : template <typename Derived>
147 : KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
148 : template <typename Derived>
149 : KOKKOS_FUNCTION void
150 : operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
151 : ///@}
152 :
153 : /**
154 : * The parallel computation bodies that can be customized in the derived class by defining
155 : * them in the derived class with the same signature.
156 : * Make sure to define them as inlined public methods if to be defined in the derived class.
157 : */
158 : ///@{
159 : /**
160 : * Compute residual
161 : * @param kernel The kernel object of the final derived type
162 : * @param datum The AssemblyDatum object of the current thread
163 : */
164 : template <typename Derived>
165 : KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const;
166 : /**
167 : * Compute diagonal Jacobian
168 : * @param kernel The kernel object of the final derived type
169 : * @param datum The AssemblyDatum object of the current thread
170 : */
171 : template <typename Derived>
172 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const;
173 : /**
174 : * Compute off-diagonal Jacobian
175 : * @param kernel The kernel object of the final derived type
176 : * @param datum The AssemblyDatum object of the current thread
177 : */
178 : template <typename Derived>
179 : KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
180 : AssemblyDatum & datum) const;
181 : ///@}
182 :
183 : protected:
184 : /**
185 : * Current test function
186 : */
187 : const VariableTestValue _test;
188 : /**
189 : * Gradient of the current test function
190 : */
191 : const VariableTestGradient _grad_test;
192 : /**
193 : * Current shape function
194 : */
195 : const VariablePhiValue _phi;
196 : /**
197 : * Gradient of the current shape function
198 : */
199 : const VariablePhiGradient _grad_phi;
200 : /**
201 : * Current solution at quadrature points
202 : */
203 : const VariableValue _u;
204 : /**
205 : * Gradient of the current solution at quadrature points
206 : */
207 : const VariableGradient _grad_u;
208 : };
209 :
210 : template <typename Derived>
211 : KOKKOS_FUNCTION void
212 3552908 : Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
213 : {
214 3552908 : auto elem = kokkosBlockElementID(tid);
215 :
216 7105816 : AssemblyDatum datum(elem,
217 : libMesh::invalid_uint,
218 : kokkosAssembly(),
219 : kokkosSystems(),
220 3552908 : _kokkos_var,
221 : _kokkos_var.var());
222 :
223 3552908 : kernel.computeResidualInternal(kernel, datum);
224 3552908 : }
225 :
226 : template <typename Derived>
227 : KOKKOS_FUNCTION void
228 592515 : Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
229 : {
230 592515 : auto elem = kokkosBlockElementID(tid);
231 :
232 1185030 : AssemblyDatum datum(elem,
233 : libMesh::invalid_uint,
234 : kokkosAssembly(),
235 : kokkosSystems(),
236 592515 : _kokkos_var,
237 : _kokkos_var.var());
238 :
239 592515 : kernel.computeJacobianInternal(kernel, datum);
240 592515 : }
241 :
242 : template <typename Derived>
243 : KOKKOS_FUNCTION void
244 1250 : Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
245 : {
246 1250 : auto elem = kokkosBlockElementID(_thread(tid, 1));
247 :
248 1250 : auto & sys = kokkosSystem(_kokkos_var.sys());
249 1250 : auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
250 :
251 1250 : if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
252 0 : return;
253 :
254 1250 : AssemblyDatum datum(
255 1250 : elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
256 :
257 1250 : kernel.computeOffDiagJacobianInternal(kernel, datum);
258 : }
259 :
260 : template <typename Derived>
261 : KOKKOS_FUNCTION void
262 2006052 : Kernel::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
263 : {
264 2006052 : ResidualObject::computeResidualInternal(
265 : datum,
266 4012104 : [&](Real * local_re, const unsigned int ib, const unsigned int ie)
267 : {
268 10642928 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
269 : {
270 8636876 : datum.reinit();
271 :
272 44452044 : for (unsigned int i = ib; i < ie; ++i)
273 35815168 : local_re[i] += datum.JxW(qp) * kernel.computeQpResidualShim(kernel, i, qp, datum);
274 : }
275 : });
276 2006052 : }
277 :
278 : template <typename Derived>
279 : KOKKOS_FUNCTION void
280 314998 : Kernel::computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
281 : {
282 314998 : ResidualObject::computeJacobianInternal(
283 : datum,
284 677612 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
285 : {
286 1759582 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
287 : {
288 1420776 : datum.reinit();
289 :
290 23861352 : for (unsigned int ij = ijb; ij < ije; ++ij)
291 : {
292 22440576 : unsigned int i = ij % datum.n_jdofs();
293 22440576 : unsigned int j = ij / datum.n_jdofs();
294 :
295 22440576 : local_ke[ij] += datum.JxW(qp) * kernel.computeQpJacobianShim(kernel, i, j, qp, datum);
296 : }
297 : }
298 : });
299 314998 : }
300 :
301 : template <typename Derived>
302 : KOKKOS_FUNCTION void
303 1250 : Kernel::computeOffDiagJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
304 : {
305 1250 : ResidualObject::computeJacobianInternal(
306 : datum,
307 2500 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
308 : {
309 6250 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
310 : {
311 5000 : datum.reinit();
312 :
313 85000 : for (unsigned int ij = ijb; ij < ije; ++ij)
314 : {
315 80000 : unsigned int i = ij % datum.n_jdofs();
316 80000 : unsigned int j = ij / datum.n_jdofs();
317 :
318 80000 : local_ke[ij] += datum.JxW(qp) * kernel.computeQpOffDiagJacobianShim(
319 : kernel, i, j, datum.jvar(), qp, datum);
320 : }
321 : }
322 : });
323 1250 : }
324 :
325 : } // namespace Kokkos
326 : } // namespace Moose
|
