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 polymorphic design of the original MOOSE is reproduced statically by leveraging the Curiously
23 : * Recurring Template Pattern (CRTP), a programming idiom that involves a class template inheriting
24 : * from a template instantiation of itself. When the user derives their Kokkos object from this
25 : * class, the inheritance structure will look like:
26 : *
27 : * class UserKernel final : public Moose::Kokkos::Kernel<UserKernel>
28 : *
29 : * It is important to note that the template argument should point to the last derived class.
30 : * Therefore, if the user wants to define a derived class that can be further inherited, the derived
31 : * class should be a class template as well. Otherwise, it is recommended to mark the derived class
32 : * as final to prevent its inheritence by mistake.
33 : *
34 : * The user is expected to define computeQpResidual(), computeQpJacobian(), and
35 : * computeQpOffDiagJacobian() as inlined public methods in their derived class (not virtual
36 : * override). The signature of computeQpResidual() expected to be defined in the derived class is as
37 : * follows:
38 : *
39 : * @param i The element-local DOF index
40 : * @param qp The local quadrature point index
41 : * @param datum The ResidualDatum object of the current thread
42 : * @returns The residual contribution
43 : *
44 : * KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
45 : * const unsigned int qp,
46 : * ResidualDatum & datum) const;
47 : *
48 : * The signatures of computeQpJacobian() and computeOffDiagQpJacobian() can be found in the code
49 : * below, and their definition in the derived class is optional. If they are defined in the derived
50 : * class, they will hide the default definitions in the base class.
51 : */
52 : template <typename Derived>
53 : class Kernel : public KernelBase
54 : {
55 : public:
56 : static InputParameters validParams();
57 :
58 : /**
59 : * Constructor
60 : */
61 : Kernel(const InputParameters & parameters);
62 :
63 : /**
64 : * Dispatch residual calculation
65 : */
66 : virtual void computeResidual() override;
67 : /**
68 : * Dispatch diagonal and off-diagonal Jacobian calculation
69 : */
70 : virtual void computeJacobian() override;
71 :
72 : /**
73 : * Default methods to prevent compile errors even when these methods were not defined in the
74 : * derived class
75 : */
76 : ///@{
77 : /**
78 : * Compute diagonal Jacobian contribution on a quadrature point
79 : * @param i The test function DOF index
80 : * @param j The trial function DOF index
81 : * @param qp The local quadrature point index
82 : * @param datum The ResidualDatum object of the current thread
83 : * @returns The diagonal Jacobian contribution
84 : */
85 0 : KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
86 : const unsigned int /* j */,
87 : const unsigned int /* qp */,
88 : ResidualDatum & /* datum */) const
89 : {
90 0 : return 0;
91 : }
92 : /**
93 : * Compute off-diagonal Jacobian contribution on a quadrature point
94 : * @param i The test function DOF index
95 : * @param j The trial function DOF index
96 : * @param jvar The variable number for column
97 : * @param qp The local quadrature point index
98 : * @param datum The ResidualDatum object of the current thread
99 : * @returns The off-diagonal Jacobian contribution
100 : */
101 0 : KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
102 : const unsigned int /* j */,
103 : const unsigned int /* jvar */,
104 : const unsigned int /* qp */,
105 : ResidualDatum & /* datum */) const
106 : {
107 0 : return 0;
108 : }
109 : ///@}
110 :
111 : /**
112 : * The parallel computation entry functions called by Kokkos
113 : */
114 : ///@{
115 : KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid) const;
116 : KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid) const;
117 : KOKKOS_FUNCTION void operator()(OffDiagJacobianLoop, const ThreadID tid) 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 ResidualDatum object of the current thread
130 : */
131 : KOKKOS_FUNCTION void computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const;
132 : /**
133 : * Compute diagonal Jacobian
134 : * @param kernel The kernel object of the final derived type
135 : * @param datum The ResidualDatum object of the current thread
136 : */
137 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const;
138 : /**
139 : * Compute off-diagonal Jacobian
140 : * @param kernel The kernel object of the final derived type
141 : * @param datum The ResidualDatum object of the current thread
142 : */
143 : KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived * kernel,
144 : ResidualDatum & datum) const;
145 : ///@}
146 :
147 : protected:
148 : /**
149 : * Current test function
150 : */
151 : const VariableTestValue _test;
152 : /**
153 : * Gradient of the current test function
154 : */
155 : const VariableTestGradient _grad_test;
156 : /**
157 : * Current shape function
158 : */
159 : const VariablePhiValue _phi;
160 : /**
161 : * Gradient of the current shape function
162 : */
163 : const VariablePhiGradient _grad_phi;
164 : /**
165 : * Current solution at quadrature points
166 : */
167 : const VariableValue _u;
168 : /**
169 : * Gradient of the current solution at quadrature points
170 : */
171 : const VariableGradient _grad_u;
172 :
173 : /**
174 : * Get whether computeQpJacobian() was not defined in the derived class
175 : * @returns Whether computeQpJacobian() was not defined in the derived class
176 : */
177 9945 : virtual bool defaultJacobian() const
178 : {
179 9945 : return &Derived::computeQpJacobian == &Kernel::computeQpJacobian;
180 : }
181 : /**
182 : * Get whether computeQpOffDiagJacobian() was not defined in the derived class
183 : * @returns Whether computeQpOffDiagJacobian() was not defined in the derived class
184 : */
185 9965 : virtual bool defaultOffDiagJacobian() const
186 : {
187 9965 : return &Derived::computeQpOffDiagJacobian == &Kernel::computeQpOffDiagJacobian;
188 : }
189 : };
190 :
191 : template <typename Derived>
192 : InputParameters
193 121022 : Kernel<Derived>::validParams()
194 : {
195 121022 : InputParameters params = KernelBase::validParams();
196 121022 : return params;
197 : }
198 :
199 : template <typename Derived>
200 1299 : Kernel<Derived>::Kernel(const InputParameters & parameters)
201 : : KernelBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD),
202 : _test(),
203 : _grad_test(),
204 : _phi(),
205 : _grad_phi(),
206 994 : _u(_var),
207 994 : _grad_u(_var)
208 : {
209 1299 : addMooseVariableDependency(&_var);
210 1299 : }
211 :
212 : template <typename Derived>
213 : void
214 54261 : Kernel<Derived>::computeResidual()
215 : {
216 54261 : ::Kokkos::RangePolicy<ResidualLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
217 : 0, numKokkosBlockElements());
218 54261 : ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
219 54261 : ::Kokkos::fence();
220 54261 : }
221 :
222 : template <typename Derived>
223 : void
224 9965 : Kernel<Derived>::computeJacobian()
225 : {
226 9965 : if (!defaultJacobian())
227 : {
228 8234 : ::Kokkos::RangePolicy<JacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
229 : 0, numKokkosBlockElements());
230 8234 : ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
231 8234 : ::Kokkos::fence();
232 8234 : }
233 :
234 9965 : if (!defaultOffDiagJacobian())
235 : {
236 48 : auto & sys = kokkosSystem(_kokkos_var.sys());
237 :
238 96 : _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBlockElements()});
239 :
240 48 : ::Kokkos::RangePolicy<OffDiagJacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
241 : 0, _thread.size());
242 48 : ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
243 48 : ::Kokkos::fence();
244 48 : }
245 9965 : }
246 :
247 : template <typename Derived>
248 : KOKKOS_FUNCTION void
249 2449070 : Kernel<Derived>::operator()(ResidualLoop, const ThreadID tid) const
250 : {
251 2449070 : auto kernel = static_cast<const Derived *>(this);
252 2449070 : auto elem = kokkosBlockElementID(tid);
253 :
254 2449070 : ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
255 :
256 2449070 : kernel->computeResidualInternal(kernel, datum);
257 2449070 : }
258 :
259 : template <typename Derived>
260 : KOKKOS_FUNCTION void
261 421696 : Kernel<Derived>::operator()(JacobianLoop, const ThreadID tid) const
262 : {
263 421696 : auto kernel = static_cast<const Derived *>(this);
264 421696 : auto elem = kokkosBlockElementID(tid);
265 :
266 421696 : ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
267 :
268 421696 : kernel->computeJacobianInternal(kernel, datum);
269 421696 : }
270 :
271 : template <typename Derived>
272 : KOKKOS_FUNCTION void
273 1250 : Kernel<Derived>::operator()(OffDiagJacobianLoop, const ThreadID tid) const
274 : {
275 1250 : auto kernel = static_cast<const Derived *>(this);
276 1250 : auto elem = kokkosBlockElementID(_thread(tid, 1));
277 :
278 1250 : auto & sys = kokkosSystem(_kokkos_var.sys());
279 1250 : auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
280 :
281 1250 : if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
282 0 : return;
283 :
284 1250 : ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
285 :
286 1250 : kernel->computeOffDiagJacobianInternal(kernel, datum);
287 : }
288 :
289 : template <typename Derived>
290 : KOKKOS_FUNCTION void
291 1454133 : Kernel<Derived>::computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const
292 : {
293 1454133 : ResidualObject::computeResidualInternal(
294 : datum,
295 2908266 : [&](Real * local_re, const unsigned int ib, const unsigned int ie)
296 : {
297 7118833 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
298 : {
299 5664700 : datum.reinit();
300 :
301 29591164 : for (unsigned int i = ib; i < ie; ++i)
302 23926464 : local_re[i] += datum.JxW(qp) * kernel->computeQpResidual(i, qp, datum);
303 : }
304 : });
305 1454133 : }
306 :
307 : template <typename Derived>
308 : KOKKOS_FUNCTION void
309 229598 : Kernel<Derived>::computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
310 : {
311 229598 : ResidualObject::computeJacobianInternal(
312 : datum,
313 506812 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
314 : {
315 1258582 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
316 : {
317 1005176 : datum.reinit();
318 :
319 16796152 : for (unsigned int ij = ijb; ij < ije; ++ij)
320 : {
321 15790976 : unsigned int i = ij % datum.n_jdofs();
322 15790976 : unsigned int j = ij / datum.n_jdofs();
323 :
324 15790976 : local_ke[ij] += datum.JxW(qp) * kernel->computeQpJacobian(i, j, qp, datum);
325 : }
326 : }
327 : });
328 229598 : }
329 :
330 : template <typename Derived>
331 : KOKKOS_FUNCTION void
332 1250 : Kernel<Derived>::computeOffDiagJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
333 : {
334 1250 : ResidualObject::computeJacobianInternal(
335 : datum,
336 2500 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
337 : {
338 6250 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
339 : {
340 5000 : datum.reinit();
341 :
342 85000 : for (unsigned int ij = ijb; ij < ije; ++ij)
343 : {
344 80000 : unsigned int i = ij % datum.n_jdofs();
345 80000 : unsigned int j = ij / datum.n_jdofs();
346 :
347 80000 : local_ke[ij] +=
348 80000 : datum.JxW(qp) * kernel->computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
349 : }
350 : }
351 : });
352 1250 : }
353 :
354 : } // namespace Kokkos
355 : } // namespace Moose
356 :
357 : #define usingKokkosKernelMembers(T) \
358 : usingKokkosKernelBaseMembers; \
359 : \
360 : protected: \
361 : using Moose::Kokkos::Kernel<T>::_test; \
362 : using Moose::Kokkos::Kernel<T>::_grad_test; \
363 : using Moose::Kokkos::Kernel<T>::_phi; \
364 : using Moose::Kokkos::Kernel<T>::_grad_phi; \
365 : using Moose::Kokkos::Kernel<T>::_u; \
366 : using Moose::Kokkos::Kernel<T>::_grad_u; \
367 : \
368 : public: \
369 : using Moose::Kokkos::Kernel<T>::operator()
|
