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 ResidualDatum 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 : * ResidualDatum & 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 :
40 : class Kernel : public KernelBase
41 : {
42 : public:
43 : static InputParameters validParams();
44 :
45 : /**
46 : * Constructor
47 : */
48 : Kernel(const InputParameters & parameters);
49 :
50 : /**
51 : * Dispatch residual calculation
52 : */
53 : virtual void computeResidual() override;
54 : /**
55 : * Dispatch diagonal and off-diagonal Jacobian calculation
56 : */
57 : virtual void computeJacobian() override;
58 :
59 : /**
60 : * Default methods to prevent compile errors even when these methods were not defined in the
61 : * derived class
62 : */
63 : ///@{
64 : /**
65 : * Compute diagonal Jacobian contribution on a quadrature point
66 : * @param i The test function DOF index
67 : * @param j The trial function DOF index
68 : * @param qp The local quadrature point index
69 : * @param datum The ResidualDatum object of the current thread
70 : * @returns The diagonal Jacobian contribution
71 : */
72 0 : KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
73 : const unsigned int /* j */,
74 : const unsigned int /* qp */,
75 : ResidualDatum & /* datum */) const
76 : {
77 0 : return 0;
78 : }
79 : /**
80 : * Compute off-diagonal Jacobian contribution on a quadrature point
81 : * @param i The test function DOF index
82 : * @param j The trial function DOF index
83 : * @param jvar The variable number for column
84 : * @param qp The local quadrature point index
85 : * @param datum The ResidualDatum object of the current thread
86 : * @returns The off-diagonal Jacobian contribution
87 : */
88 0 : KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
89 : const unsigned int /* j */,
90 : const unsigned int /* jvar */,
91 : const unsigned int /* qp */,
92 : ResidualDatum & /* datum */) const
93 : {
94 0 : return 0;
95 : }
96 : /**
97 : * Get the function pointer of the default computeQpJacobian()
98 : * @returns The function pointer
99 : */
100 411146 : static auto defaultJacobian() { return &Kernel::computeQpJacobian; }
101 : /**
102 : * Get the function pointer of the default computeQpOffDiagJacobian()
103 : * @returns The function pointer
104 : */
105 485868 : static auto defaultOffDiagJacobian() { return &Kernel::computeQpOffDiagJacobian; }
106 : ///@}
107 :
108 : /**
109 : * The parallel computation entry functions called by Kokkos
110 : */
111 : ///@{
112 : template <typename Derived>
113 : KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
114 : template <typename Derived>
115 : KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
116 : template <typename Derived>
117 : KOKKOS_FUNCTION void
118 : operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
119 : ///@}
120 :
121 : /**
122 : * The parallel computation bodies that can be customized in the derived class by defining
123 : * them in the derived class with the same signature.
124 : * Make sure to define them as inlined public methods if to be defined in the derived class.
125 : */
126 : ///@{
127 : /**
128 : * Compute residual
129 : * @param kernel The kernel object of the final derived type
130 : * @param datum The ResidualDatum object of the current thread
131 : */
132 : template <typename Derived>
133 : KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, ResidualDatum & datum) const;
134 : /**
135 : * Compute diagonal Jacobian
136 : * @param kernel The kernel object of the final derived type
137 : * @param datum The ResidualDatum object of the current thread
138 : */
139 : template <typename Derived>
140 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, ResidualDatum & datum) const;
141 : /**
142 : * Compute off-diagonal Jacobian
143 : * @param kernel The kernel object of the final derived type
144 : * @param datum The ResidualDatum object of the current thread
145 : */
146 : template <typename Derived>
147 : KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
148 : ResidualDatum & datum) const;
149 : ///@}
150 :
151 : protected:
152 : /**
153 : * Current test function
154 : */
155 : const VariableTestValue _test;
156 : /**
157 : * Gradient of the current test function
158 : */
159 : const VariableTestGradient _grad_test;
160 : /**
161 : * Current shape function
162 : */
163 : const VariablePhiValue _phi;
164 : /**
165 : * Gradient of the current shape function
166 : */
167 : const VariablePhiGradient _grad_phi;
168 : /**
169 : * Current solution at quadrature points
170 : */
171 : const VariableValue _u;
172 : /**
173 : * Gradient of the current solution at quadrature points
174 : */
175 : const VariableGradient _grad_u;
176 : };
177 :
178 : template <typename Derived>
179 : KOKKOS_FUNCTION void
180 2449070 : Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
181 : {
182 2449070 : auto elem = kokkosBlockElementID(tid);
183 :
184 4898140 : ResidualDatum datum(elem,
185 : libMesh::invalid_uint,
186 : kokkosAssembly(),
187 : kokkosSystems(),
188 2449070 : _kokkos_var,
189 : _kokkos_var.var());
190 :
191 2449070 : kernel.computeResidualInternal(kernel, datum);
192 2449070 : }
193 :
194 : template <typename Derived>
195 : KOKKOS_FUNCTION void
196 421696 : Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
197 : {
198 421696 : auto elem = kokkosBlockElementID(tid);
199 :
200 843392 : ResidualDatum datum(elem,
201 : libMesh::invalid_uint,
202 : kokkosAssembly(),
203 : kokkosSystems(),
204 421696 : _kokkos_var,
205 : _kokkos_var.var());
206 :
207 421696 : kernel.computeJacobianInternal(kernel, datum);
208 421696 : }
209 :
210 : template <typename Derived>
211 : KOKKOS_FUNCTION void
212 1250 : Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
213 : {
214 1250 : auto elem = kokkosBlockElementID(_thread(tid, 1));
215 :
216 1250 : auto & sys = kokkosSystem(_kokkos_var.sys());
217 1250 : auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
218 :
219 1250 : if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
220 0 : return;
221 :
222 1250 : ResidualDatum datum(
223 1250 : elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
224 :
225 1250 : kernel.computeOffDiagJacobianInternal(kernel, datum);
226 : }
227 :
228 : template <typename Derived>
229 : KOKKOS_FUNCTION void
230 1454133 : Kernel::computeResidualInternal(const Derived & kernel, ResidualDatum & datum) const
231 : {
232 1454133 : ResidualObject::computeResidualInternal(
233 : datum,
234 2908266 : [&](Real * local_re, const unsigned int ib, const unsigned int ie)
235 : {
236 7118833 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
237 : {
238 5664700 : datum.reinit();
239 :
240 29591164 : for (unsigned int i = ib; i < ie; ++i)
241 23926464 : local_re[i] += datum.JxW(qp) * kernel.computeQpResidual(i, qp, datum);
242 : }
243 : });
244 1454133 : }
245 :
246 : template <typename Derived>
247 : KOKKOS_FUNCTION void
248 229598 : Kernel::computeJacobianInternal(const Derived & kernel, ResidualDatum & datum) const
249 : {
250 229598 : ResidualObject::computeJacobianInternal(
251 : datum,
252 506812 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
253 : {
254 1258582 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
255 : {
256 1005176 : datum.reinit();
257 :
258 16796152 : for (unsigned int ij = ijb; ij < ije; ++ij)
259 : {
260 15790976 : unsigned int i = ij % datum.n_jdofs();
261 15790976 : unsigned int j = ij / datum.n_jdofs();
262 :
263 15790976 : local_ke[ij] += datum.JxW(qp) * kernel.computeQpJacobian(i, j, qp, datum);
264 : }
265 : }
266 : });
267 229598 : }
268 :
269 : template <typename Derived>
270 : KOKKOS_FUNCTION void
271 1250 : Kernel::computeOffDiagJacobianInternal(const Derived & kernel, ResidualDatum & datum) const
272 : {
273 1250 : ResidualObject::computeJacobianInternal(
274 : datum,
275 2500 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
276 : {
277 6250 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
278 : {
279 5000 : datum.reinit();
280 :
281 85000 : for (unsigned int ij = ijb; ij < ije; ++ij)
282 : {
283 80000 : unsigned int i = ij % datum.n_jdofs();
284 80000 : unsigned int j = ij / datum.n_jdofs();
285 :
286 80000 : local_ke[ij] +=
287 80000 : datum.JxW(qp) * kernel.computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
288 : }
289 : }
290 : });
291 1250 : }
292 :
293 : } // namespace Kokkos
294 : } // namespace Moose
|
