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 vector kernels on vector variables.
19 : */
20 : class VectorKernel : public KernelBase
21 : {
22 : public:
23 : static InputParameters validParams();
24 :
25 : /**
26 : * Constructor
27 : */
28 : VectorKernel(const InputParameters & parameters);
29 :
30 : /**
31 : * Dispatch residual calculation
32 : */
33 : virtual void computeResidual() override;
34 : /**
35 : * Dispatch diagonal and off-diagonal Jacobian calculation
36 : */
37 : virtual void computeJacobian() override;
38 :
39 : /**
40 : * Default methods to prevent compile errors even when these methods were not defined in the
41 : * derived class
42 : */
43 : ///@{
44 : /**
45 : * Compute diagonal Jacobian contribution on a quadrature point
46 : * @tparam Derived The object type
47 : * @param i The test function DOF index
48 : * @param j The trial function DOF index
49 : * @param qp The local quadrature point index
50 : * @param datum The AssemblyDatum object of the current thread
51 : * @returns The diagonal Jacobian contribution
52 : */
53 : template <typename Derived>
54 0 : KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
55 : const unsigned int /* j */,
56 : const unsigned int /* qp */,
57 : AssemblyDatum & /* datum */) const
58 : {
59 0 : ::Kokkos::abort("Default computeQpJacobian() should never be called. Make sure you properly "
60 : "redefined this method in your class without typos.");
61 :
62 : return 0;
63 : }
64 : /**
65 : * Compute off-diagonal Jacobian contribution on a quadrature point
66 : * @tparam Derived The object type
67 : * @param i The test function DOF index
68 : * @param j The trial function DOF index
69 : * @param jvar The variable number for column
70 : * @param qp The local quadrature point index
71 : * @param datum The AssemblyDatum object of the current thread
72 : * @returns The off-diagonal Jacobian contribution
73 : */
74 : template <typename Derived>
75 0 : KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
76 : const unsigned int /* j */,
77 : const unsigned int /* jvar */,
78 : const unsigned int /* qp */,
79 : AssemblyDatum & /* datum */) const
80 : {
81 0 : ::Kokkos::abort(
82 : "Default computeQpOffDiagJacobian() should never be called. Make sure you properly "
83 : "redefined this method in your class without typos.");
84 :
85 : return 0;
86 : }
87 : ///@}
88 :
89 : /**
90 : * Functions used to check if users have overriden the hook methods, whose calculations can be
91 : * skipped when not overriden
92 : * @returns The function pointer of the default hook method
93 : */
94 : ///@{
95 : template <typename Derived>
96 82798 : static auto defaultJacobian()
97 : {
98 82798 : return &VectorKernel::computeQpJacobian<Derived>;
99 : }
100 : template <typename Derived>
101 82798 : static auto defaultOffDiagJacobian()
102 : {
103 82798 : return &VectorKernel::computeQpOffDiagJacobian<Derived>;
104 : }
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 vector test function
153 : */
154 : const VectorVariableTestValue _test;
155 : /**
156 : * Gradient of the current vector test function
157 : */
158 : const VectorVariableTestGradient _grad_test;
159 : /**
160 : * Current vector shape function
161 : */
162 : const VectorVariablePhiValue _phi;
163 : /**
164 : * Gradient of the current vector shape function
165 : */
166 : const VectorVariablePhiGradient _grad_phi;
167 : /**
168 : * Current vector solution at quadrature points
169 : */
170 : const VectorVariableValue _u;
171 : /**
172 : * Gradient of the current vector solution at quadrature points
173 : */
174 : const VectorVariableGradient _grad_u;
175 : };
176 :
177 : template <typename Derived>
178 : KOKKOS_FUNCTION void
179 349006 : VectorKernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
180 : {
181 349006 : auto elem = kokkosBlockElementID(_thread(tid, 1));
182 :
183 698012 : AssemblyDatum datum(elem,
184 : libMesh::invalid_uint,
185 : kokkosAssembly(),
186 : kokkosSystems(),
187 349006 : _kokkos_var,
188 : _kokkos_var.var());
189 :
190 349006 : datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
191 :
192 349006 : kernel.computeResidualInternal(kernel, datum);
193 349006 : }
194 :
195 : template <typename Derived>
196 : KOKKOS_FUNCTION void
197 88264 : VectorKernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
198 : {
199 88264 : auto elem = kokkosBlockElementID(_thread(tid, 1));
200 :
201 176528 : AssemblyDatum datum(elem,
202 : libMesh::invalid_uint,
203 : kokkosAssembly(),
204 : kokkosSystems(),
205 88264 : _kokkos_var,
206 : _kokkos_var.var());
207 :
208 88264 : datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
209 :
210 88264 : kernel.computeJacobianInternal(kernel, datum);
211 88264 : }
212 :
213 : template <typename Derived>
214 : KOKKOS_FUNCTION void
215 45600 : VectorKernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
216 : {
217 45600 : auto elem = kokkosBlockElementID(_thread(tid, 2));
218 :
219 45600 : auto & sys = kokkosSystem(_kokkos_var.sys());
220 45600 : auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 1)];
221 :
222 45600 : if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
223 0 : return;
224 :
225 45600 : AssemblyDatum datum(
226 45600 : elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
227 :
228 45600 : datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
229 :
230 45600 : kernel.computeOffDiagJacobianInternal(kernel, datum);
231 : }
232 :
233 : template <typename Derived>
234 : KOKKOS_FUNCTION void
235 33964 : VectorKernel::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
236 : {
237 33964 : ResidualObject::computeResidualInternal(
238 : datum,
239 67928 : [&](Real * local_re, const unsigned int ib, const unsigned int ie)
240 : {
241 169820 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
242 1222704 : for (unsigned int i = ib; i < ie; ++i)
243 1086848 : local_re[i] += datum.JxW(qp) * kernel.template computeQpResidual<Derived>(i, qp, datum);
244 : });
245 33964 : }
246 :
247 : template <typename Derived>
248 : KOKKOS_FUNCTION void
249 21004 : VectorKernel::computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
250 : {
251 21004 : ResidualObject::computeJacobianInternal(
252 : datum,
253 336064 : [&](Real * local_ke, const unsigned int ib, const unsigned int ie, const unsigned int j)
254 : {
255 840160 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
256 6049152 : for (unsigned int i = ib; i < ie; ++i)
257 5377024 : local_ke[i] +=
258 5377024 : datum.JxW(qp) * kernel.template computeQpJacobian<Derived>(i, j, qp, datum);
259 : });
260 21004 : }
261 :
262 : template <typename Derived>
263 : KOKKOS_FUNCTION void
264 45600 : VectorKernel::computeOffDiagJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
265 : {
266 45600 : ResidualObject::computeJacobianInternal(
267 : datum,
268 638400 : [&](Real * local_ke, const unsigned int ib, const unsigned int ie, const unsigned int j)
269 : {
270 1596000 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
271 11491200 : for (unsigned int i = ib; i < ie; ++i)
272 10214400 : local_ke[i] += datum.JxW(qp) * kernel.template computeQpOffDiagJacobian<Derived>(
273 : i, j, datum.jvar(), qp, datum);
274 : });
275 45600 : }
276 :
277 : } // namespace Moose::Kokkos
|
