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