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::Kokkos
15 : {
16 :
17 : /**
18 : * The base class for a user to derive their own Kokkos kernels.
19 : *
20 : * The user should define computeQpResidual(), computeQpJacobian(), and computeQpOffDiagJacobian()
21 : * as inlined public methods in their derived class (not virtual override). The signature of
22 : * computeQpResidual() expected to be defined in the derived class is as follows:
23 : *
24 : * @tparam Derived The object type
25 : * @param i The element-local DOF index
26 : * @param qp The local quadrature point index
27 : * @param datum The AssemblyDatum object of the current thread
28 : * @returns The residual contribution
29 : *
30 : * template <typename Derived>
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 computeQpOffDiagJacobian() 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 : * @tparam Derived The object type
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 AssemblyDatum object of the current thread
70 : * @returns The diagonal Jacobian contribution
71 : */
72 : template <typename Derived>
73 0 : KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
74 : const unsigned int /* j */,
75 : const unsigned int /* qp */,
76 : AssemblyDatum & /* datum */) const
77 : {
78 0 : ::Kokkos::abort("Default computeQpJacobian() should never be called. Make sure you properly "
79 : "redefined this method in your class without typos.");
80 :
81 : return 0;
82 : }
83 : /**
84 : * Compute off-diagonal Jacobian contribution on a quadrature point
85 : * @tparam Derived The object type
86 : * @param i The test function DOF index
87 : * @param j The trial function DOF index
88 : * @param jvar The variable number for column
89 : * @param qp The local quadrature point index
90 : * @param datum The AssemblyDatum object of the current thread
91 : * @returns The off-diagonal Jacobian contribution
92 : */
93 : template <typename Derived>
94 0 : KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
95 : const unsigned int /* j */,
96 : const unsigned int /* jvar */,
97 : const unsigned int /* qp */,
98 : AssemblyDatum & /* datum */) const
99 : {
100 0 : ::Kokkos::abort(
101 : "Default computeQpOffDiagJacobian() should never be called. Make sure you properly "
102 : "redefined this method in your class without typos.");
103 :
104 : return 0;
105 : }
106 : ///@}
107 :
108 : /**
109 : * Functions used to check if users have overriden the hook methods, whose calculations can be
110 : * skipped when not overriden
111 : * @returns The function pointer of the default hook method
112 : */
113 : ///@{
114 : template <typename Derived>
115 363020 : static auto defaultJacobian()
116 : {
117 363020 : return &Kernel::computeQpJacobian<Derived>;
118 : }
119 : template <typename Derived>
120 403375 : static auto defaultOffDiagJacobian()
121 : {
122 403375 : return &Kernel::computeQpOffDiagJacobian<Derived>;
123 : }
124 : ///@}
125 :
126 : /**
127 : * The parallel computation entry functions called by Kokkos
128 : */
129 : ///@{
130 : template <typename Derived>
131 : KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
132 : template <typename Derived>
133 : KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
134 : template <typename Derived>
135 : KOKKOS_FUNCTION void
136 : operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
137 : ///@}
138 :
139 : /**
140 : * The parallel computation bodies that can be customized in the derived class by defining
141 : * them in the derived class with the same signature.
142 : * Make sure to define them as inlined public methods if to be defined in the derived class.
143 : */
144 : ///@{
145 : /**
146 : * Compute residual
147 : * @param kernel The kernel object of the final derived type
148 : * @param datum The AssemblyDatum object of the current thread
149 : */
150 : template <typename Derived>
151 : KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const;
152 : /**
153 : * Compute diagonal Jacobian
154 : * @param kernel The kernel object of the final derived type
155 : * @param datum The AssemblyDatum object of the current thread
156 : */
157 : template <typename Derived>
158 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const;
159 : /**
160 : * Compute off-diagonal Jacobian
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 computeOffDiagJacobianInternal(const Derived & kernel,
166 : AssemblyDatum & datum) const;
167 : ///@}
168 :
169 : protected:
170 : /**
171 : * Current test function
172 : */
173 : const VariableTestValue _test;
174 : /**
175 : * Gradient of the current test function
176 : */
177 : const VariableTestGradient _grad_test;
178 : /**
179 : * Current shape function
180 : */
181 : const VariablePhiValue _phi;
182 : /**
183 : * Gradient of the current shape function
184 : */
185 : const VariablePhiGradient _grad_phi;
186 : /**
187 : * Current solution at quadrature points
188 : */
189 : const VariableValue _u;
190 : /**
191 : * Gradient of the current solution at quadrature points
192 : */
193 : const VariableGradient _grad_u;
194 : };
195 :
196 : template <typename Derived>
197 : KOKKOS_FUNCTION void
198 16420992 : Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
199 : {
200 16420992 : auto elem = kokkosBlockElementID(_thread(tid, 1));
201 :
202 32841984 : AssemblyDatum datum(elem,
203 : libMesh::invalid_uint,
204 : kokkosAssembly(),
205 : kokkosSystems(),
206 16420992 : _kokkos_var,
207 : _kokkos_var.var());
208 :
209 16420992 : datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
210 :
211 16420992 : kernel.computeResidualInternal(kernel, datum);
212 16420992 : }
213 :
214 : template <typename Derived>
215 : KOKKOS_FUNCTION void
216 2653084 : Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
217 : {
218 2653084 : auto elem = kokkosBlockElementID(_thread(tid, 1));
219 :
220 5306168 : AssemblyDatum datum(elem,
221 : libMesh::invalid_uint,
222 : kokkosAssembly(),
223 : kokkosSystems(),
224 2653084 : _kokkos_var,
225 : _kokkos_var.var());
226 :
227 2653084 : datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
228 :
229 2653084 : kernel.computeJacobianInternal(kernel, datum);
230 2653084 : }
231 :
232 : template <typename Derived>
233 : KOKKOS_FUNCTION void
234 5000 : Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
235 : {
236 5000 : auto elem = kokkosBlockElementID(_thread(tid, 2));
237 :
238 5000 : auto & sys = kokkosSystem(_kokkos_var.sys());
239 5000 : auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 1)];
240 :
241 5000 : if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
242 0 : return;
243 :
244 5000 : AssemblyDatum datum(
245 5000 : elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
246 :
247 5000 : datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
248 :
249 5000 : kernel.computeOffDiagJacobianInternal(kernel, datum);
250 : }
251 :
252 : template <typename Derived>
253 : KOKKOS_FUNCTION void
254 8957428 : Kernel::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
255 : {
256 8957428 : ResidualObject::computeResidualInternal(
257 : datum,
258 17914856 : [&](Real * local_re, const unsigned int ib, const unsigned int ie)
259 : {
260 47972900 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
261 79682064 : for (unsigned int i = ib; i < ie; ++i)
262 40666592 : local_re[i] += datum.JxW(qp) * kernel.template computeQpResidual<Derived>(i, qp, datum);
263 : });
264 8957428 : }
265 :
266 : template <typename Derived>
267 : KOKKOS_FUNCTION void
268 1445416 : Kernel::computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
269 : {
270 1445416 : ResidualObject::computeJacobianInternal(
271 : datum,
272 2683760 : [&](Real * local_ke, const unsigned int ib, const unsigned int ie, const unsigned int j)
273 : {
274 7085368 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
275 31173696 : for (unsigned int i = ib; i < ie; ++i)
276 25430208 : local_ke[i] +=
277 25430208 : datum.JxW(qp) * kernel.template computeQpJacobian<Derived>(i, j, qp, datum);
278 : });
279 1445416 : }
280 :
281 : template <typename Derived>
282 : KOKKOS_FUNCTION void
283 0 : Kernel::computeOffDiagJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
284 : {
285 0 : ResidualObject::computeJacobianInternal(
286 : datum,
287 0 : [&](Real * local_ke, const unsigned int ib, const unsigned int ie, const unsigned int j)
288 : {
289 0 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
290 0 : for (unsigned int i = ib; i < ie; ++i)
291 0 : local_ke[i] += datum.JxW(qp) * kernel.template computeQpOffDiagJacobian<Derived>(
292 : i, j, datum.jvar(), qp, datum);
293 : });
294 0 : }
295 :
296 : } // namespace Moose::Kokkos
|
