Line data Source code
1 : //* This file is part of the MOOSE framework
2 : //* https://mooseframework.inl.gov
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 "KokkosDispatcher.h"
13 : #include "KokkosVariableValue.h"
14 : #include "KokkosMaterialPropertyValue.h"
15 : #include "KokkosAssembly.h"
16 : #include "KokkosSystem.h"
17 :
18 : #include "AuxKernelBase.h"
19 :
20 : namespace Moose
21 : {
22 : namespace Kokkos
23 : {
24 :
25 : /**
26 : * The base class for a user to derive their own Kokkos auxiliary kernels.
27 : *
28 : * The user should define computeValue() as inlined public method in their derived class (not
29 : * virtual override). The signature of computeValue() expected to be defined in the derived
30 : * class is as follows:
31 : *
32 : * @param qp The local quadrature point index
33 : * @param datum The AssemblyDatum object of the current thread
34 : * @returns The value at the quadrature point
35 : *
36 : * KOKKOS_FUNCTION Real computeValue(const unsigned int qp, AssemblyDatum & datum) const;
37 : */
38 : class AuxKernel : public ::AuxKernelBase,
39 : public MeshHolder,
40 : public AssemblyHolder,
41 : public SystemHolder
42 : {
43 : public:
44 : static InputParameters validParams();
45 :
46 : /**
47 : * Constructor
48 : */
49 : AuxKernel(const InputParameters & parameters);
50 :
51 : /**
52 : * Copy constructor for parallel dispatch
53 : */
54 : AuxKernel(const AuxKernel & object);
55 :
56 : // Unused for Kokkos auxiliary kernels because all elements are computed in parallel
57 0 : virtual void subdomainSetup() override final {}
58 :
59 : /**
60 : * Dispatch calculation
61 : */
62 : virtual void compute() override;
63 :
64 : /**
65 : * Get whether this auxiliary kernel is nodal
66 : * @returns Whether this auxiliary kernel is nodal
67 : */
68 488 : KOKKOS_FUNCTION bool isNodal() const { return _nodal; }
69 :
70 : /**
71 : * Kokkos function tags
72 : */
73 : ///@{
74 : struct ElementLoop
75 : {
76 : };
77 : struct NodeLoop
78 : {
79 : };
80 : ///@}
81 :
82 : /**
83 : * The parallel computation entry functions called by Kokkos
84 : */
85 : ///@{
86 : template <typename Derived>
87 : KOKKOS_FUNCTION void operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const;
88 : template <typename Derived>
89 : KOKKOS_FUNCTION void operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const;
90 : ///@}
91 :
92 : /**
93 : * The parallel computation bodies that can be customized in the derived class by defining
94 : * them in the derived class with the same signature.
95 : * Make sure to define them as inlined public methods if to be defined in the derived class.
96 : */
97 : ///@{
98 : /**
99 : * Compute an element
100 : * @param kernel The auxiliary kernel object of the final derived type
101 : * @param datum The AssemblyDatum object of the current thread
102 : */
103 : template <typename Derived>
104 : KOKKOS_FUNCTION void computeElementInternal(const Derived & auxkernel,
105 : AssemblyDatum & datum) const;
106 : /**
107 : * Compute a node
108 : * @param kernel The auxiliary kernel object of the final derived type
109 : * @param datum The AssemblyDatum object of the current thread
110 : */
111 : template <typename Derived>
112 : KOKKOS_FUNCTION void computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const;
113 : ///@}
114 :
115 : protected:
116 : /**
117 : * Retrieve the old value of the variable that this kernel operates on
118 : * @returns The old variable value object
119 : */
120 : VariableValue uOld() const;
121 : /**
122 : * Retrieve the older value of the variable that this kernel operates on
123 : * @returns The older variable value object
124 : */
125 : VariableValue uOlder() const;
126 :
127 : /**
128 : * Set element values to the auxiliary solution vector
129 : * @param values The array containing the solution values of the element
130 : * @param datum The AssemblyDatum object of the current thread
131 : * @param comp The variable component
132 : */
133 : KOKKOS_FUNCTION void setElementSolution(const Real * const values,
134 : const AssemblyDatum & datum,
135 : const unsigned int comp = 0) const;
136 : /**
137 : * Set node value to the auxiliary solution vector
138 : * @param values The node solution value
139 : * @param datum The AssemblyDatum object of the current thread
140 : * @param comp The variable component
141 : */
142 : KOKKOS_FUNCTION void
143 : setNodeSolution(const Real value, const AssemblyDatum & datum, const unsigned int comp = 0) const;
144 :
145 : /**
146 : * Flag whether this kernel is nodal
147 : */
148 : const bool _nodal;
149 :
150 : /**
151 : * Kokkos variable
152 : */
153 : Variable _kokkos_var;
154 : /**
155 : * Kokkos functor dispatchers
156 : */
157 : ///@{
158 : std::unique_ptr<DispatcherBase> _element_dispatcher;
159 : std::unique_ptr<DispatcherBase> _node_dispatcher;
160 : ///@}
161 :
162 : /**
163 : * Current test function
164 : */
165 : const VariableTestValue _test;
166 : /**
167 : * Current solution
168 : */
169 : const VariableValue _u;
170 :
171 : /**
172 : * TODO: Move to TransientInterface
173 : */
174 : ///@{
175 : /**
176 : * Time
177 : */
178 : Scalar<Real> _t;
179 : /**
180 : * Old time
181 : */
182 : Scalar<const Real> _t_old;
183 : /**
184 : * The number of the time step
185 : */
186 : Scalar<int> _t_step;
187 : /**
188 : * Time step size
189 : */
190 : Scalar<Real> _dt;
191 : /**
192 : * Size of the old time step
193 : */
194 : Scalar<Real> _dt_old;
195 : ///@}
196 :
197 : private:
198 : /**
199 : * Override of the MaterialPropertyInterface function to throw on material property request for
200 : * nodal kernels
201 : */
202 : void getKokkosMaterialPropertyHook(const std::string & prop_name_in,
203 : const unsigned int state) override final;
204 : };
205 :
206 : template <typename Derived>
207 : KOKKOS_FUNCTION void
208 210176 : AuxKernel::operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const
209 : {
210 210176 : auto elem = kokkosBlockElementID(tid);
211 :
212 420352 : AssemblyDatum datum(elem,
213 : libMesh::invalid_uint,
214 : kokkosAssembly(),
215 : kokkosSystems(),
216 210176 : _kokkos_var,
217 : _kokkos_var.var());
218 :
219 210176 : auxkernel.computeElementInternal(auxkernel, datum);
220 210176 : }
221 :
222 : template <typename Derived>
223 : KOKKOS_FUNCTION void
224 2081260 : AuxKernel::operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const
225 : {
226 2081260 : auto node = _bnd ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid);
227 2081260 : auto & sys = kokkosSystem(_kokkos_var.sys());
228 :
229 2081260 : if (!sys.isNodalDefined(node, _kokkos_var.var()))
230 1505280 : return;
231 :
232 575980 : AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
233 :
234 575980 : auxkernel.computeNodeInternal(auxkernel, datum);
235 : }
236 :
237 : template <typename Derived>
238 : KOKKOS_FUNCTION void
239 209576 : AuxKernel::computeElementInternal(const Derived & auxkernel, AssemblyDatum & datum) const
240 : {
241 : Real x[MAX_CACHED_DOF];
242 : Real b[MAX_CACHED_DOF];
243 : Real A[MAX_CACHED_DOF * MAX_CACHED_DOF];
244 :
245 1104152 : for (unsigned int i = 0; i < datum.n_dofs(); ++i)
246 : {
247 894576 : x[i] = 0;
248 894576 : b[i] = 0;
249 :
250 5899152 : for (unsigned int j = 0; j < datum.n_dofs(); ++j)
251 5004576 : A[j + datum.n_dofs() * i] = 0;
252 : }
253 :
254 2023184 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
255 : {
256 1813608 : const auto value = auxkernel.computeValue(qp, datum);
257 :
258 1813608 : datum.reinit();
259 :
260 9792216 : for (unsigned int i = 0; i < datum.n_dofs(); ++i)
261 : {
262 7978608 : const auto t = datum.JxW(qp) * _test(datum, i, qp);
263 :
264 7978608 : b[i] += t * value;
265 :
266 52947216 : for (unsigned int j = 0; j < datum.n_dofs(); ++j)
267 44968608 : A[j + datum.n_dofs() * i] += t * _test(datum, j, qp);
268 : }
269 : }
270 :
271 209576 : if (datum.n_dofs() == 1)
272 72576 : x[0] = b[0] / A[0];
273 : else
274 137000 : Utils::choleskySolve(A, x, b, datum.n_dofs());
275 :
276 209576 : setElementSolution(x, datum);
277 209576 : }
278 :
279 : template <typename Derived>
280 : KOKKOS_FUNCTION void
281 575054 : AuxKernel::computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const
282 : {
283 575054 : auto value = auxkernel.computeValue(0, datum);
284 :
285 575054 : setNodeSolution(value, datum);
286 575054 : }
287 :
288 : KOKKOS_FUNCTION inline void
289 209576 : AuxKernel::setElementSolution(const Real * const values,
290 : const AssemblyDatum & datum,
291 : const unsigned int comp) const
292 : {
293 209576 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
294 209576 : auto var_num = _kokkos_var.var(comp);
295 209576 : auto tag = _kokkos_var.tag();
296 209576 : auto elem = datum.elem().id;
297 :
298 1104152 : for (unsigned int i = 0; i < datum.n_dofs(); ++i)
299 894576 : sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, i, var_num), tag) = values[i];
300 209576 : }
301 :
302 : KOKKOS_FUNCTION inline void
303 575054 : AuxKernel::setNodeSolution(const Real value,
304 : const AssemblyDatum & datum,
305 : const unsigned int comp) const
306 : {
307 575054 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
308 575054 : auto var_num = _kokkos_var.var(comp);
309 575054 : auto tag = _kokkos_var.tag();
310 575054 : auto node = datum.node();
311 :
312 575054 : sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, 0, var_num), tag) = value;
313 575054 : }
314 :
315 : } // namespace Kokkos
316 : } // namespace Moose
|