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 "KokkosTypes.h"
13 :
14 : #include "MooseMesh.h"
15 :
16 : #include "libmesh/elem_range.h"
17 : #include "libmesh/fe_base.h"
18 : #include "libmesh/fe_type.h"
19 :
20 : class FEProblemBase;
21 :
22 : namespace Moose
23 : {
24 : namespace Kokkos
25 : {
26 :
27 : /**
28 : * The Kokkos assembly class
29 : */
30 : class Assembly : public MeshHolder
31 : {
32 : public:
33 : /**
34 : * Constructor
35 : * @param problem The MOOSE problem
36 : */
37 : Assembly(FEProblemBase & problem);
38 : /**
39 : * Initialize assembly
40 : */
41 : void init();
42 :
43 : #ifdef MOOSE_KOKKOS_SCOPE
44 : /**
45 : * Get the FE type ID
46 : * @param type The libMesh FEType object
47 : * @returns The FE type ID
48 : */
49 1424 : unsigned int getFETypeID(FEType type) const { return libmesh_map_find(_fe_type_map, type); }
50 : /**
51 : * Get the mesh dimension
52 : * @returns The mesh dimension
53 : */
54 0 : KOKKOS_FUNCTION unsigned int getDimension() const { return _dimension; }
55 : /**
56 : * Get the maximum number of quadrature points per element in the current partition
57 : * @returns The maximum number of quadrature points per element
58 : */
59 77394 : KOKKOS_FUNCTION unsigned int getMaxQpsPerElem() const { return _max_qps_per_elem; }
60 : /**
61 : * Get the total number of elemental quadrature points in a subdomain
62 : * @param subdomain The contiguous subdomain ID
63 : * @returns The number of quadrature points
64 : */
65 4959 : KOKKOS_FUNCTION dof_id_type getNumQps(ContiguousSubdomainID subdomain) const
66 : {
67 4959 : return _n_subdomain_qps[subdomain];
68 : }
69 : /**
70 : * Get the number of quadrature points of an element
71 : * @param info The element information object
72 : * @returns The number of quadrature points
73 : */
74 46944259 : KOKKOS_FUNCTION unsigned int getNumQps(ElementInfo info) const { return _n_qps[info.id]; }
75 : /**
76 : * Get the total number of facial quadrature points in a subdomain
77 : * Note: this number does not represent the real number of facial quadrature points but only
78 : * counts the facial quadrature points that need global caching, such as face material properties
79 : * @param subdomain The contiguous subdomain ID
80 : * @returns The number of quadrature points
81 : */
82 1169 : KOKKOS_FUNCTION dof_id_type getNumFaceQps(ContiguousSubdomainID subdomain) const
83 : {
84 1169 : return _n_subdomain_qps_face[subdomain];
85 : }
86 : /**
87 : * Get the number of quadrature points of a side of an element
88 : * @param info The element information object
89 : * @param side The side index
90 : * @returns The number of quadrature points
91 : */
92 121790 : KOKKOS_FUNCTION unsigned int getNumFaceQps(ElementInfo info, unsigned int side) const
93 : {
94 121790 : return _n_qps_face(side, info.id);
95 : }
96 : /**
97 : * Get the starting offset of quadrature points of an element into the global quadrature point
98 : * index
99 : * @param info The element information object
100 : * @returns The starting offset
101 : */
102 134970399 : KOKKOS_FUNCTION dof_id_type getQpOffset(ElementInfo info) const { return _qp_offset[info.id]; }
103 : /**
104 : * Get the starting offset of quadrature points of a side of an element into the global quadrature
105 : * point index
106 : * @param info The element information object
107 : * @param side The side index
108 : * @returns The starting offset
109 : */
110 121790 : KOKKOS_FUNCTION dof_id_type getQpFaceOffset(ElementInfo info, unsigned int side) const
111 : {
112 121790 : return _qp_offset_face(side, info.id);
113 : }
114 : /**
115 : * Get the number of DOFs of a FE type for an element type
116 : * @param elem_type The element type ID
117 : * @param fe_type The FE type ID
118 : * @returns The number of DOFs
119 : */
120 49264550 : KOKKOS_FUNCTION unsigned int getNumDofs(unsigned int elem_type, unsigned int fe_type) const
121 : {
122 49264550 : return _n_dofs(elem_type, fe_type);
123 : }
124 : /**
125 : * Get the shape functions of a FE type for an element type and subdomain
126 : * @param subdomain The contiguous subdomain ID
127 : * @param elem_type The element type ID
128 : * @param fe_type The FE type ID
129 : * @returns The shape functions at quadrature points
130 : */
131 : KOKKOS_FUNCTION const auto &
132 147323516 : getPhi(ContiguousSubdomainID subdomain, unsigned int elem_type, unsigned int fe_type) const
133 : {
134 147323516 : return _phi(subdomain, elem_type, fe_type);
135 : }
136 : /**
137 : * Get the face shape functions of a FE type for an element type and subdomain
138 : * @param subdomain The contiguous subdomain ID
139 : * @param elem_type The element type ID
140 : * @param fe_type The FE type ID
141 : * @returns The shape functions of all sides at quadrature points
142 : */
143 : KOKKOS_FUNCTION const auto &
144 666240 : getPhiFace(ContiguousSubdomainID subdomain, unsigned int elem_type, unsigned int fe_type) const
145 : {
146 666240 : return _phi_face(subdomain, elem_type, fe_type);
147 : }
148 : /**
149 : * Get the gradient of shape functions of a FE type for an element type and subdomain
150 : * @param subdomain The contiguous subdomain ID
151 : * @param elem_type The element type ID
152 : * @param fe_type The FE type ID
153 : * @returns The gradient of shape functions at quadrature points
154 : */
155 : KOKKOS_FUNCTION const auto &
156 102653020 : getGradPhi(ContiguousSubdomainID subdomain, unsigned int elem_type, unsigned int fe_type) const
157 : {
158 102653020 : return _grad_phi(subdomain, elem_type, fe_type);
159 : }
160 : /**
161 : * Get the gradient of face shape functions of a FE type for an element type and subdomain
162 : * @param subdomain The contiguous subdomain ID
163 : * @param elem_type The element type ID
164 : * @param fe_type The FE type ID
165 : * @returns The gradient of shape functions of all sides at quadrature points
166 : */
167 0 : KOKKOS_FUNCTION const auto & getGradPhiFace(ContiguousSubdomainID subdomain,
168 : unsigned int elem_type,
169 : unsigned int fe_type) const
170 : {
171 0 : return _grad_phi_face(subdomain, elem_type, fe_type);
172 : }
173 : /**
174 : * Get the inverse of Jacobian matrix of an element quadrature point
175 : * @param info The element information object
176 : * @param qp The local quadrature point index
177 : * @returns The inverse of Jacobian matrix
178 : */
179 57938676 : KOKKOS_FUNCTION Real33 getJacobian(ElementInfo info, unsigned int qp) const
180 : {
181 57938676 : return _jacobian[info.subdomain][getQpOffset(info) + qp];
182 : }
183 : /**
184 : * Get the transformed Jacobian weight of an element quadrature point
185 : * @param info The element information object
186 : * @param qp The local quadrature point index
187 : * @returns The inverse of Jacobian matrix
188 : */
189 15104840 : KOKKOS_FUNCTION Real getJxW(ElementInfo info, unsigned int qp) const
190 : {
191 15104840 : return _jxw[info.subdomain][getQpOffset(info) + qp];
192 : }
193 : /**
194 : * Get the coordinate of an element quadrature point
195 : * @param info The element information object
196 : * @param qp The local quadrature point index
197 : * @returns The inverse of Jacobian matrix
198 : */
199 15104840 : KOKKOS_FUNCTION Real3 getQPoint(ElementInfo info, unsigned int qp) const
200 : {
201 15104840 : return _xyz[info.subdomain][getQpOffset(info) + qp];
202 : }
203 :
204 : /**
205 : * Get the coordinate transform factor for a point in a subdomain
206 : * @param subdomain The contiguous subdomain ID
207 : * @param point The point coordinate
208 : * @returns The coordinate transform factor
209 : */
210 : KOKKOS_FUNCTION Real coordTransformFactor(const ContiguousSubdomainID subdomain,
211 : const Real3 point) const;
212 : /**
213 : * Compute physical transformation data for an element
214 : * @param info The element information object
215 : * @param qp The local quadrature point index
216 : * @param jacobian The pointer to store the inverse of Jacobian matrix
217 : * @param JxW The pointer to store transformed Jacobian weight
218 : * @param q_points The pointer to store physical quadrature point coordinate
219 : */
220 : KOKKOS_FUNCTION void computePhysicalMap(const ElementInfo info,
221 : const unsigned int qp,
222 : Real33 * const jacobian,
223 : Real * const JxW,
224 : Real3 * const q_points) const;
225 : /**
226 : * Compute physical transformation data for a side
227 : * @param info The element information object
228 : * @param side The side index
229 : * @param qp The local quadrature point index
230 : * @param jacobian The pointer to store the inverse of Jacobian matrix
231 : * @param JxW The pointer to store transformed Jacobian weight
232 : * @param q_points The pointer to store physical quadrature point coordinate
233 : */
234 : KOKKOS_FUNCTION void computePhysicalMap(const ElementInfo info,
235 : const unsigned int side,
236 : const unsigned int qp,
237 : Real33 * const jacobian,
238 : Real * const JxW,
239 : Real3 * const q_points) const;
240 :
241 : /**
242 : * Kokkos function for caching physical maps on element quadrature points
243 : */
244 : KOKKOS_FUNCTION void operator()(const ThreadID tid) const;
245 :
246 : /**
247 : * Get the list of boundaries to cache face material properties
248 : * @returns The list of boundaries
249 : */
250 563 : const auto & getMaterialBoundaries() const { return _material_boundaries; }
251 : #endif
252 :
253 : private:
254 : /**
255 : * Initialize quadrature data
256 : */
257 : void initQuadrature();
258 : /**
259 : * Initialize shape data
260 : */
261 : void initShape();
262 : /**
263 : * Cache physical maps on element quadrature points
264 : */
265 : void cachePhysicalMap();
266 :
267 : /**
268 : * Reference of the MOOSE problem
269 : */
270 : FEProblemBase & _problem;
271 : /**
272 : * Reference of the MOOSE mesh
273 : */
274 : MooseMesh & _mesh;
275 : /**
276 : * FE type ID map
277 : */
278 : std::map<FEType, unsigned int> _fe_type_map;
279 :
280 : /**
281 : * Mesh dimension
282 : */
283 : const unsigned int _dimension;
284 : /**
285 : * Coordinate system type of each subdomain
286 : */
287 : Array<Moose::CoordinateSystemType> _coord_type;
288 : /**
289 : * Radial coordinate index in cylindrical coordinate system
290 : */
291 : unsigned int _rz_radial_coord = libMesh::invalid_uint;
292 : /**
293 : * General axisymmetric axis of each subdomain in cylindrical coordinate system
294 : */
295 : Array<Pair<Real3, Real3>> _rz_axis;
296 :
297 : /**
298 : * Starting offset into the global quadrature point index
299 : */
300 : ///@{
301 : Array<dof_id_type> _qp_offset;
302 : Array2D<dof_id_type> _qp_offset_face;
303 : ///@}
304 : /**
305 : * Number of quadrature points
306 : */
307 : ///@{
308 : Array<unsigned int> _n_qps;
309 : Array2D<unsigned int> _n_qps_face;
310 :
311 : unsigned int _max_qps_per_elem = 0;
312 :
313 : Array<dof_id_type> _n_subdomain_qps;
314 : Array<dof_id_type> _n_subdomain_qps_face;
315 : ///@}
316 : /**
317 : * Quadrature points and weights for reference elements
318 : */
319 : ///@{
320 : Array2D<Array<Real3>> _q_points;
321 : Array2D<Array<Array<Real3>>> _q_points_face;
322 : Array2D<Array<Real>> _weights;
323 : Array2D<Array<Array<Real>>> _weights_face;
324 : ///@}
325 : /**
326 : * Shape functions for reference elements
327 : */
328 : ///@{
329 : Array3D<Array2D<Real>> _phi;
330 : Array3D<Array<Array2D<Real>>> _phi_face;
331 : Array3D<Array2D<Real3>> _grad_phi;
332 : Array3D<Array<Array2D<Real3>>> _grad_phi_face;
333 : Array2D<unsigned int> _n_dofs;
334 : ///@}
335 : /**
336 : * Shape functions for computing reference-to-physical maps
337 : */
338 : ///@{
339 : Array2D<Array2D<Real>> _map_phi;
340 : Array2D<Array<Array2D<Real>>> _map_phi_face;
341 : Array2D<Array<Array2D<Real>>> _map_psi_face;
342 : Array2D<Array2D<Real3>> _map_grad_phi;
343 : Array2D<Array<Array2D<Real3>>> _map_grad_phi_face;
344 : Array2D<Array<Array2D<Real3>>> _map_grad_psi_face;
345 : ///@}
346 : /**
347 : * Cached physical maps on element quadrature points
348 : */
349 : ///@{
350 : Array<Array<Real33>> _jacobian;
351 : Array<Array<Real>> _jxw;
352 : Array<Array<Real3>> _xyz;
353 : ///@}
354 :
355 : /**
356 : * Boundaries to cache face material properties
357 : */
358 : std::set<BoundaryID> _material_boundaries;
359 : };
360 :
361 : #ifdef MOOSE_KOKKOS_SCOPE
362 : KOKKOS_FUNCTION inline Real
363 228180 : Assembly::coordTransformFactor(const ContiguousSubdomainID subdomain, const Real3 point) const
364 : {
365 228180 : switch (_coord_type[subdomain])
366 : {
367 225240 : case Moose::COORD_XYZ:
368 225240 : return 1;
369 2800 : case Moose::COORD_RZ:
370 2800 : if (_rz_radial_coord == libMesh::invalid_uint)
371 0 : return 2 * M_PI *
372 0 : (point - _rz_axis[subdomain].first).cross_product(_rz_axis[subdomain].second).norm();
373 : else
374 2800 : return 2 * M_PI * point(_rz_radial_coord);
375 140 : case Moose::COORD_RSPHERICAL:
376 140 : return 4 * M_PI * point(0) * point(0);
377 0 : default:
378 0 : return 0;
379 : }
380 : }
381 :
382 : KOKKOS_FUNCTION inline void
383 122216 : Assembly::computePhysicalMap(const ElementInfo info,
384 : const unsigned int qp,
385 : Real33 * const jacobian,
386 : Real * const JxW,
387 : Real3 * const q_points) const
388 : {
389 122216 : auto sid = info.subdomain;
390 122216 : auto eid = info.id;
391 122216 : auto elem_type = info.type;
392 122216 : auto num_nodes = kokkosMesh().getNumNodes(elem_type);
393 :
394 122216 : auto & phi = _map_phi(sid, elem_type);
395 122216 : auto & grad_phi = _map_grad_phi(sid, elem_type);
396 :
397 122216 : Real33 J;
398 122216 : Real3 xyz;
399 :
400 673852 : for (unsigned int node = 0; node < num_nodes; ++node)
401 : {
402 551636 : auto points = kokkosMesh().getNodePoint(kokkosMesh().getContiguousNodeID(eid, node));
403 :
404 551636 : if (jacobian || JxW)
405 551636 : J += grad_phi(node, qp).cartesian_product(points);
406 :
407 551636 : xyz += phi(node, qp) * points;
408 : }
409 :
410 122216 : if (jacobian)
411 122216 : *jacobian = J.inverse(_dimension);
412 122216 : if (JxW)
413 122216 : *JxW =
414 122216 : J.determinant(_dimension) * _weights(sid, elem_type)[qp] * coordTransformFactor(sid, xyz);
415 122216 : if (q_points)
416 122216 : *q_points = xyz;
417 122216 : }
418 :
419 : KOKKOS_FUNCTION inline void
420 105964 : Assembly::computePhysicalMap(const ElementInfo info,
421 : const unsigned int side,
422 : const unsigned int qp,
423 : Real33 * const jacobian,
424 : Real * const JxW,
425 : Real3 * const q_points) const
426 : {
427 105964 : auto sid = info.subdomain;
428 105964 : auto eid = info.id;
429 105964 : auto elem_type = info.type;
430 105964 : auto num_nodes = kokkosMesh().getNumNodes(elem_type);
431 105964 : auto num_side_nodes = kokkosMesh().getNumNodes(elem_type, side);
432 :
433 105964 : auto & phi = _map_phi_face(sid, elem_type)(side);
434 105964 : auto & grad_phi = _map_grad_phi_face(sid, elem_type)(side);
435 :
436 105964 : Real33 J;
437 105964 : Real3 xyz;
438 :
439 528768 : for (unsigned int node = 0; node < num_nodes; ++node)
440 : {
441 422804 : auto points = kokkosMesh().getNodePoint(kokkosMesh().getContiguousNodeID(eid, node));
442 :
443 422804 : if (jacobian)
444 422804 : J += grad_phi(node, qp).cartesian_product(points);
445 :
446 422804 : if (JxW || q_points)
447 422804 : xyz += phi(node, qp) * points;
448 : }
449 :
450 105964 : if (jacobian)
451 105964 : *jacobian = J.inverse(_dimension);
452 105964 : if (q_points)
453 105964 : *q_points = xyz;
454 :
455 105964 : if (JxW)
456 : {
457 105964 : Real33 J;
458 :
459 105964 : auto & grad_psi = _map_grad_psi_face(sid, elem_type)(side);
460 :
461 317366 : for (unsigned int node = 0; node < num_side_nodes; ++node)
462 : {
463 211402 : auto points = kokkosMesh().getNodePoint(kokkosMesh().getContiguousNodeID(eid, side, node));
464 :
465 211402 : J += grad_psi(node, qp).cartesian_product(points);
466 : }
467 :
468 105964 : *JxW = ::Kokkos::sqrt((J * J.transpose()).determinant(_dimension - 1)) *
469 105964 : _weights_face(sid, elem_type)[side][qp] * coordTransformFactor(sid, xyz);
470 : }
471 105964 : }
472 : #endif
473 :
474 : /**
475 : * The Kokkos interface that holds the host reference of the Kokkos assembly and copies it to device
476 : * during parallel dispatch.
477 : * Maintains synchronization between host and device Kokkos assemblies and provides access to the
478 : * appropriate Kokkos assembly depending on the architecture.
479 : */
480 : class AssemblyHolder
481 : {
482 : public:
483 : /**
484 : * Constructor
485 : * @param assembly The Kokkos assembly
486 : */
487 5747 : AssemblyHolder(const Assembly & assembly) : _assembly_host(assembly), _assembly_device(assembly)
488 : {
489 5747 : }
490 : /**
491 : * Copy constructor
492 : */
493 300793 : AssemblyHolder(const AssemblyHolder & holder)
494 247294 : : _assembly_host(holder._assembly_host), _assembly_device(holder._assembly_host)
495 : {
496 300793 : }
497 :
498 : #ifdef MOOSE_KOKKOS_SCOPE
499 : /**
500 : * Get the const reference of the Kokkos assembly
501 : * @returns The const reference of the Kokkos assembly depending on the architecture this function
502 : * is being called on
503 : */
504 264135718 : KOKKOS_FUNCTION const Assembly & kokkosAssembly() const
505 : {
506 264135718 : KOKKOS_IF_ON_HOST(return _assembly_host;)
507 :
508 264051943 : return _assembly_device;
509 : }
510 : #endif
511 :
512 : private:
513 : /**
514 : * Host reference of the Kokkos assembly
515 : */
516 : const Assembly & _assembly_host;
517 : /**
518 : * Device copy of the Kokkos assembly
519 : */
520 : const Assembly _assembly_device;
521 : };
522 :
523 : } // namespace Kokkos
524 : } // namespace Moose
|
