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::Kokkos
21 : {
22 :
23 : /**
24 : * The base class for a user to derive their own Kokkos auxiliary kernels.
25 : *
26 : * The user should define computeValue() as inlined public method in their derived class (not
27 : * virtual override). The signature of computeValue() expected to be defined in the derived
28 : * class is as follows:
29 : *
30 : * @tparam Derived The object type
31 : * @param qp The local quadrature point index
32 : * @param datum The AssemblyDatum object of the current thread
33 : * @returns The value at the quadrature point
34 : *
35 : * template <typename Derived>
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 1002 : 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 233800 : AuxKernel::operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const
209 : {
210 233800 : auto elem = kokkosBlockElementID(tid);
211 :
212 467600 : AssemblyDatum datum(elem,
213 : libMesh::invalid_uint,
214 : kokkosAssembly(),
215 : kokkosSystems(),
216 233800 : _kokkos_var,
217 : _kokkos_var.var());
218 :
219 233800 : auxkernel.computeElementInternal(auxkernel, datum);
220 233800 : }
221 :
222 : template <typename Derived>
223 : KOKKOS_FUNCTION void
224 2534856 : AuxKernel::operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const
225 : {
226 2534856 : auto node = _bnd ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid);
227 2534856 : auto & sys = kokkosSystem(_kokkos_var.sys());
228 :
229 2534856 : if (!sys.isNodalDefined(node, _kokkos_var.var()))
230 1505280 : return;
231 :
232 1029576 : AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
233 :
234 1029576 : auxkernel.computeNodeInternal(auxkernel, datum);
235 : }
236 :
237 : template <typename Derived>
238 : KOKKOS_FUNCTION void
239 233200 : 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 1165356 : for (unsigned int i = 0; i < datum.n_dofs(); ++i)
246 : {
247 932156 : x[i] = 0;
248 932156 : b[i] = 0;
249 :
250 6051136 : for (unsigned int j = 0; j < datum.n_dofs(); ++j)
251 5118980 : A[j + datum.n_dofs() * i] = 0;
252 : }
253 :
254 2159520 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
255 : {
256 1926320 : const auto value = auxkernel.template computeValue<Derived>(qp, datum);
257 :
258 10139788 : for (unsigned int i = 0; i < datum.n_dofs(); ++i)
259 : {
260 8213468 : const auto t = datum.JxW(qp) * _test(datum, i, qp);
261 :
262 8213468 : b[i] += t * value;
263 :
264 54094528 : for (unsigned int j = 0; j < datum.n_dofs(); ++j)
265 45881060 : A[j + datum.n_dofs() * i] += t * _test(datum, j, qp);
266 : }
267 : }
268 :
269 233200 : if (datum.n_dofs() == 1)
270 92948 : x[0] = b[0] / A[0];
271 : else
272 140252 : Utils::choleskySolve(A, x, b, datum.n_dofs());
273 :
274 233200 : setElementSolution(x, datum);
275 233200 : }
276 :
277 : template <typename Derived>
278 : KOKKOS_FUNCTION void
279 1028650 : AuxKernel::computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const
280 : {
281 1028650 : auto value = auxkernel.template computeValue<Derived>(0, datum);
282 :
283 1028650 : setNodeSolution(value, datum);
284 1028650 : }
285 :
286 : KOKKOS_FUNCTION inline void
287 233200 : AuxKernel::setElementSolution(const Real * const values,
288 : const AssemblyDatum & datum,
289 : const unsigned int comp) const
290 : {
291 233200 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
292 233200 : auto var_num = _kokkos_var.var(comp);
293 233200 : auto tag = _kokkos_var.tag();
294 233200 : auto elem = datum.elem().id;
295 :
296 1165356 : for (unsigned int i = 0; i < datum.n_dofs(); ++i)
297 932156 : sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, i, var_num), tag) = values[i];
298 233200 : }
299 :
300 : KOKKOS_FUNCTION inline void
301 1028650 : AuxKernel::setNodeSolution(const Real value,
302 : const AssemblyDatum & datum,
303 : const unsigned int comp) const
304 : {
305 1028650 : auto & sys = kokkosSystem(_kokkos_var.sys(comp));
306 1028650 : auto var_num = _kokkos_var.var(comp);
307 1028650 : auto tag = _kokkos_var.tag();
308 1028650 : auto node = datum.node();
309 :
310 1028650 : sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, 0, var_num), tag) = value;
311 1028650 : }
312 :
313 : } // namespace Moose::Kokkos
|