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