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 411971 : static auto defaultJacobian() { return &Kernel::computeQpJacobian; }
100 : /**
101 : * Get the function pointer of the default computeQpOffDiagJacobian()
102 : * @returns The function pointer
103 : */
104 486843 : static auto defaultOffDiagJacobian() { return &Kernel::computeQpOffDiagJacobian; }
105 : ///@}
106 :
107 : /**
108 : * The parallel computation entry functions called by Kokkos
109 : */
110 : ///@{
111 : template <typename Derived>
112 : KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
113 : template <typename Derived>
114 : KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
115 : template <typename Derived>
116 : KOKKOS_FUNCTION void
117 : operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
118 : ///@}
119 :
120 : /**
121 : * The parallel computation bodies that can be customized in the derived class by defining
122 : * them in the derived class with the same signature.
123 : * Make sure to define them as inlined public methods if to be defined in the derived class.
124 : */
125 : ///@{
126 : /**
127 : * Compute residual
128 : * @param kernel The kernel object of the final derived type
129 : * @param datum The AssemblyDatum object of the current thread
130 : */
131 : template <typename Derived>
132 : KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const;
133 : /**
134 : * Compute diagonal Jacobian
135 : * @param kernel The kernel object of the final derived type
136 : * @param datum The AssemblyDatum object of the current thread
137 : */
138 : template <typename Derived>
139 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const;
140 : /**
141 : * Compute off-diagonal Jacobian
142 : * @param kernel The kernel object of the final derived type
143 : * @param datum The AssemblyDatum object of the current thread
144 : */
145 : template <typename Derived>
146 : KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
147 : AssemblyDatum & datum) const;
148 : ///@}
149 :
150 : protected:
151 : /**
152 : * Current test function
153 : */
154 : const VariableTestValue _test;
155 : /**
156 : * Gradient of the current test function
157 : */
158 : const VariableTestGradient _grad_test;
159 : /**
160 : * Current shape function
161 : */
162 : const VariablePhiValue _phi;
163 : /**
164 : * Gradient of the current shape function
165 : */
166 : const VariablePhiGradient _grad_phi;
167 : /**
168 : * Current solution at quadrature points
169 : */
170 : const VariableValue _u;
171 : /**
172 : * Gradient of the current solution at quadrature points
173 : */
174 : const VariableGradient _grad_u;
175 : };
176 :
177 : template <typename Derived>
178 : KOKKOS_FUNCTION void
179 2449508 : Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
180 : {
181 2449508 : auto elem = kokkosBlockElementID(tid);
182 :
183 4899016 : AssemblyDatum datum(elem,
184 : libMesh::invalid_uint,
185 : kokkosAssembly(),
186 : kokkosSystems(),
187 2449508 : _kokkos_var,
188 : _kokkos_var.var());
189 :
190 2449508 : kernel.computeResidualInternal(kernel, datum);
191 2449508 : }
192 :
193 : template <typename Derived>
194 : KOKKOS_FUNCTION void
195 421715 : Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
196 : {
197 421715 : auto elem = kokkosBlockElementID(tid);
198 :
199 843430 : AssemblyDatum datum(elem,
200 : libMesh::invalid_uint,
201 : kokkosAssembly(),
202 : kokkosSystems(),
203 421715 : _kokkos_var,
204 : _kokkos_var.var());
205 :
206 421715 : kernel.computeJacobianInternal(kernel, datum);
207 421715 : }
208 :
209 : template <typename Derived>
210 : KOKKOS_FUNCTION void
211 1250 : Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
212 : {
213 1250 : auto elem = kokkosBlockElementID(_thread(tid, 1));
214 :
215 1250 : auto & sys = kokkosSystem(_kokkos_var.sys());
216 1250 : auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
217 :
218 1250 : if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
219 0 : return;
220 :
221 1250 : AssemblyDatum datum(
222 1250 : elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
223 :
224 1250 : kernel.computeOffDiagJacobianInternal(kernel, datum);
225 : }
226 :
227 : template <typename Derived>
228 : KOKKOS_FUNCTION void
229 1454352 : Kernel::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
230 : {
231 1454352 : ResidualObject::computeResidualInternal(
232 : datum,
233 2908704 : [&](Real * local_re, const unsigned int ib, const unsigned int ie)
234 : {
235 7884428 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
236 : {
237 6430076 : datum.reinit();
238 :
239 33418044 : for (unsigned int i = ib; i < ie; ++i)
240 26987968 : local_re[i] += datum.JxW(qp) * kernel.computeQpResidual(i, qp, datum);
241 : }
242 : });
243 1454352 : }
244 :
245 : template <typename Derived>
246 : KOKKOS_FUNCTION void
247 229598 : Kernel::computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
248 : {
249 229598 : ResidualObject::computeJacobianInternal(
250 : datum,
251 506812 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
252 : {
253 1332582 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
254 : {
255 1079176 : datum.reinit();
256 :
257 18054152 : for (unsigned int ij = ijb; ij < ije; ++ij)
258 : {
259 16974976 : unsigned int i = ij % datum.n_jdofs();
260 16974976 : unsigned int j = ij / datum.n_jdofs();
261 :
262 16974976 : local_ke[ij] += datum.JxW(qp) * kernel.computeQpJacobian(i, j, qp, datum);
263 : }
264 : }
265 : });
266 229598 : }
267 :
268 : template <typename Derived>
269 : KOKKOS_FUNCTION void
270 1250 : Kernel::computeOffDiagJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
271 : {
272 1250 : ResidualObject::computeJacobianInternal(
273 : datum,
274 2500 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
275 : {
276 6250 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
277 : {
278 5000 : datum.reinit();
279 :
280 85000 : for (unsigned int ij = ijb; ij < ije; ++ij)
281 : {
282 80000 : unsigned int i = ij % datum.n_jdofs();
283 80000 : unsigned int j = ij / datum.n_jdofs();
284 :
285 80000 : local_ke[ij] +=
286 80000 : datum.JxW(qp) * kernel.computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
287 : }
288 : }
289 : });
290 1250 : }
291 :
292 : } // namespace Kokkos
293 : } // namespace Moose
|
