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 : #include "KokkosAssembly.h"
11 :
12 : #include "MooseMesh.h"
13 : #include "FEProblemBase.h"
14 : #include "NonlinearSystemBase.h"
15 : #include "AuxiliarySystem.h"
16 : #include "Assembly.h"
17 : #include "BoundaryRestrictable.h"
18 :
19 : #include "libmesh/fe_interface.h"
20 : #include "libmesh/reference_elem.h"
21 :
22 : namespace Moose::Kokkos
23 : {
24 :
25 46593 : Assembly::Assembly(FEProblemBase & problem)
26 44492 : : MeshHolder(*problem.mesh().getKokkosMesh()),
27 44492 : _problem(problem),
28 44492 : _mesh(problem.mesh()),
29 133476 : _dimension(_mesh.dimension())
30 : {
31 46593 : }
32 :
33 : void
34 2265 : Assembly::init()
35 : {
36 : // Cache mesh information
37 :
38 2265 : const auto num_subdomains = kokkosMesh().getNumSubdomains();
39 :
40 2265 : _coord_type.create(num_subdomains);
41 :
42 5211 : for (auto subdomain : _mesh.meshSubdomains())
43 2946 : _coord_type[kokkosMesh().getContiguousSubdomainID(subdomain)] = _mesh.getCoordSystem(subdomain);
44 :
45 2265 : _coord_type.copyToDevice();
46 :
47 2265 : if (_mesh.usingGeneralAxisymmetricCoordAxes())
48 : {
49 0 : _rz_axis.create(num_subdomains);
50 :
51 0 : for (auto subdomain : _mesh.meshSubdomains())
52 0 : _rz_axis[kokkosMesh().getContiguousSubdomainID(subdomain)] =
53 0 : _mesh.getGeneralAxisymmetricCoordAxis(subdomain);
54 :
55 0 : _rz_axis.copyToDevice();
56 : }
57 : else
58 2265 : _rz_radial_coord = _mesh.getAxisymmetricRadialCoord();
59 :
60 : // Initialize quadrature and shape data
61 :
62 2265 : initQuadrature();
63 2265 : initShape();
64 2265 : cachePhysicalMap();
65 2265 : }
66 :
67 : void
68 2265 : Assembly::initQuadrature()
69 : {
70 2265 : const auto num_subdomains = kokkosMesh().getNumSubdomains();
71 2265 : const auto num_elem_types = kokkosMesh().getNumLocalElementTypes();
72 :
73 2265 : _q_points.create(num_subdomains, num_elem_types);
74 2265 : _q_points_face.create(num_subdomains, num_elem_types);
75 2265 : _weights.create(num_subdomains, num_elem_types);
76 2265 : _weights_face.create(num_subdomains, num_elem_types);
77 :
78 : // Find boundaries where material properties should be computed
79 :
80 1034 : auto & boundary_objects =
81 1231 : _problem.getKokkosMaterialPropertyStorageConsumers(Moose::BOUNDARY_MATERIAL_DATA);
82 :
83 3085 : for (auto object : boundary_objects)
84 : {
85 820 : auto boundary_restriction = dynamic_cast<const BoundaryRestrictable *>(object);
86 :
87 820 : if (boundary_restriction)
88 1691 : for (auto boundary : boundary_restriction->boundaryIDs())
89 871 : _material_boundaries.insert(boundary);
90 : else
91 0 : mooseError("Kokkos assembly error: ", object->name(), " is not boundary-restricted.");
92 : }
93 :
94 : // Cache quadrature data
95 :
96 2265 : std::map<SubdomainID, std::map<ElemType, unsigned int>> n_qps;
97 2265 : std::map<SubdomainID, std::map<ElemType, std::vector<unsigned int>>> n_qps_face;
98 :
99 2265 : _max_qps_per_elem = 0;
100 :
101 5211 : for (auto subdomain : _mesh.meshSubdomains())
102 : {
103 2946 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
104 :
105 2946 : auto & assembly = _problem.assembly(0, 0);
106 2946 : auto qrule = assembly.writeableQRule(_dimension, subdomain, {});
107 2946 : auto qrule_face = assembly.writeableQRuleFace(_dimension, subdomain, {});
108 :
109 5888 : for (auto & [elem_type, elem_type_id] : kokkosMesh().getElementTypeMap())
110 : {
111 2942 : auto elem = &libMesh::ReferenceElem::get(elem_type);
112 :
113 2942 : _q_points_face(sid, elem_type_id).create(elem->n_sides());
114 2942 : _weights_face(sid, elem_type_id).create(elem->n_sides());
115 :
116 : // Cache volume quadrature of each reference element
117 :
118 2942 : qrule->init(*elem, /* p-level */ 0);
119 2942 : n_qps[subdomain][elem_type] = qrule->n_points();
120 :
121 2942 : _q_points(sid, elem_type_id).create(qrule->n_points());
122 2942 : _weights(sid, elem_type_id).create(qrule->n_points());
123 :
124 14342 : for (unsigned int qp = 0; qp < qrule->n_points(); ++qp)
125 : {
126 11400 : _q_points(sid, elem_type_id)[qp] = qrule->qp(qp);
127 11400 : _weights(sid, elem_type_id)[qp] = qrule->w(qp);
128 : }
129 :
130 : // Cache face quadrature of each reference element
131 :
132 14348 : for (unsigned int side = 0; side < elem->n_sides(); ++side)
133 : {
134 11406 : qrule_face->init(*elem->side_ptr(side), /* p-level */ 0);
135 11406 : n_qps_face[subdomain][elem_type].push_back(qrule_face->n_points());
136 :
137 11406 : _q_points_face(sid, elem_type_id)[side].create(qrule_face->n_points());
138 11406 : _weights_face(sid, elem_type_id)[side].create(qrule_face->n_points());
139 :
140 34058 : for (unsigned int qp = 0; qp < qrule_face->n_points(); ++qp)
141 : {
142 22652 : _q_points_face(sid, elem_type_id)[side][qp] = qrule_face->qp(qp);
143 22652 : _weights_face(sid, elem_type_id)[side][qp] = qrule_face->w(qp);
144 : }
145 : }
146 :
147 2942 : _max_qps_per_elem = std::max(_max_qps_per_elem, n_qps[subdomain][elem_type]);
148 : }
149 : }
150 :
151 2265 : const auto num_elems = _mesh.nActiveLocalElem();
152 :
153 2265 : _n_subdomain_qps.create(num_subdomains);
154 2265 : _n_subdomain_qps_face.create(num_subdomains);
155 2265 : _n_subdomain_qps = 0;
156 2265 : _n_subdomain_qps_face = 0;
157 :
158 2265 : _n_qps.create(num_elems);
159 2265 : _n_qps_face.create(_mesh.getMaxSidesPerElem(), num_elems);
160 2265 : _n_qps_face = 0;
161 :
162 2265 : _qp_offset.create(num_elems);
163 2265 : _qp_offset_face.create(_mesh.getMaxSidesPerElem(), num_elems);
164 2265 : _qp_offset_face = libMesh::DofObject::invalid_id;
165 :
166 2265 : _elem_face_property_idx.create(_mesh.getMaxSidesPerElem(), num_elems);
167 2265 : _elem_face_property_idx = libMesh::DofObject::invalid_id;
168 :
169 2265 : _n_elem_face_properties.create(num_subdomains);
170 2265 : _n_elem_face_properties = 0;
171 :
172 646184 : for (auto elem : *_mesh.getActiveLocalElementRange())
173 : {
174 643919 : auto eid = kokkosMesh().getContiguousElementID(elem);
175 643919 : auto sid = kokkosMesh().getContiguousSubdomainID(elem->subdomain_id());
176 :
177 643919 : _n_qps[eid] = n_qps[elem->subdomain_id()][elem->type()];
178 643919 : _qp_offset[eid] = _n_subdomain_qps[sid];
179 643919 : _n_subdomain_qps[sid] += _n_qps[eid];
180 :
181 3214903 : for (unsigned int side = 0; side < elem->n_sides(); ++side)
182 2570984 : _n_qps_face(side, eid) = n_qps_face[elem->subdomain_id()][elem->type()][side];
183 : }
184 :
185 3034 : for (auto boundary : _material_boundaries)
186 5326 : for (auto elem_id : _mesh.getBoundaryActiveSemiLocalElemIds(boundary))
187 : {
188 4557 : auto elem = _mesh.elemPtr(elem_id);
189 :
190 4557 : if (elem->processor_id() == _problem.processor_id())
191 : {
192 3867 : auto sid = kokkosMesh().getContiguousSubdomainID(elem->subdomain_id());
193 3867 : auto eid = kokkosMesh().getContiguousElementID(elem);
194 3867 : auto side = _mesh.sideWithBoundaryID(elem, boundary);
195 :
196 3867 : _qp_offset_face(side, eid) = _n_subdomain_qps_face[sid];
197 3867 : _n_subdomain_qps_face[sid] += _n_qps_face(side, eid);
198 :
199 3867 : _elem_face_property_idx(side, eid) = _n_elem_face_properties[sid];
200 3867 : ++_n_elem_face_properties[sid];
201 : }
202 769 : }
203 :
204 2265 : _q_points.copyToDeviceNested();
205 2265 : _q_points_face.copyToDeviceNested();
206 2265 : _weights.copyToDeviceNested();
207 2265 : _weights_face.copyToDeviceNested();
208 :
209 2265 : _n_qps.copyToDevice();
210 2265 : _n_qps_face.copyToDevice();
211 2265 : _n_subdomain_qps.copyToDevice();
212 2265 : _n_subdomain_qps_face.copyToDevice();
213 2265 : _qp_offset.copyToDevice();
214 2265 : _qp_offset_face.copyToDevice();
215 :
216 2265 : _elem_face_property_idx.copyToDevice();
217 2265 : _n_elem_face_properties.copyToDevice();
218 2265 : }
219 :
220 : void
221 2265 : Assembly::initShape()
222 : {
223 : // Generate the list of unique FE types
224 :
225 2265 : std::set<FEType> fe_types;
226 :
227 4530 : auto getFETypes = [&](::System & system)
228 : {
229 8482 : for (unsigned int var = 0; var < system.n_vars(); var++)
230 3952 : fe_types.insert(system.variable_type(var));
231 6795 : };
232 :
233 4530 : for (unsigned int nl = 0; nl < _problem.numNonlinearSystems(); ++nl)
234 2265 : getFETypes(_problem.getNonlinearSystemBase(nl).system());
235 :
236 2265 : getFETypes(_problem.getAuxiliarySystem().system());
237 :
238 2265 : _fe_type_map.clear();
239 :
240 4878 : for (auto & fet : fe_types)
241 2613 : _fe_type_map[fet] = _fe_type_map.size();
242 :
243 : // Cache reference shape data
244 :
245 2265 : const auto num_subdomains = kokkosMesh().getNumSubdomains();
246 2265 : const auto num_elem_types = kokkosMesh().getNumLocalElementTypes();
247 :
248 2265 : _phi.create(num_subdomains, num_elem_types, _fe_type_map.size());
249 2265 : _phi_face.create(num_subdomains, num_elem_types, _fe_type_map.size());
250 2265 : _grad_phi.create(num_subdomains, num_elem_types, _fe_type_map.size());
251 2265 : _grad_phi_face.create(num_subdomains, num_elem_types, _fe_type_map.size());
252 :
253 2265 : _map_phi.create(num_subdomains, num_elem_types);
254 2265 : _map_phi_face.create(num_subdomains, num_elem_types);
255 2265 : _map_psi_face.create(num_subdomains, num_elem_types);
256 2265 : _map_grad_phi.create(num_subdomains, num_elem_types);
257 2265 : _map_grad_phi_face.create(num_subdomains, num_elem_types);
258 2265 : _map_grad_psi_face.create(num_subdomains, num_elem_types);
259 :
260 2265 : _normal_dx_dxi.create(num_subdomains, num_elem_types);
261 2265 : _normal_dx_deta.create(num_subdomains, num_elem_types);
262 :
263 2265 : _n_dofs.create(num_elem_types, _fe_type_map.size());
264 2265 : _n_dofs = 0;
265 :
266 5211 : for (auto subdomain : _mesh.meshSubdomains())
267 : {
268 2946 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
269 :
270 2946 : auto & assembly = _problem.assembly(0, 0);
271 2946 : auto qrule = assembly.writeableQRule(_dimension, subdomain, {});
272 2946 : auto qrule_face = assembly.writeableQRuleFace(_dimension, subdomain, {});
273 :
274 6444 : for (auto & [fe_type, fe_type_id] : _fe_type_map)
275 : {
276 3498 : std::unique_ptr<FEBase> fe(FEBase::build(_dimension, fe_type));
277 3498 : std::unique_ptr<FEBase> fe_face(FEBase::build(_dimension, fe_type));
278 :
279 3498 : fe->attach_quadrature_rule(qrule);
280 3498 : fe_face->attach_quadrature_rule(qrule_face);
281 :
282 6992 : for (auto & [elem_type, elem_type_id] : kokkosMesh().getElementTypeMap())
283 : {
284 3494 : auto elem = &libMesh::ReferenceElem::get(elem_type);
285 :
286 3494 : auto & phi = fe->get_phi();
287 3494 : auto & grad_phi = fe->get_dphi();
288 :
289 3494 : fe->reinit(elem);
290 :
291 3494 : _n_dofs(elem_type_id, fe_type_id) = phi.size();
292 :
293 3494 : _phi(sid, elem_type_id, fe_type_id).create(phi.size(), qrule->n_points());
294 3494 : _grad_phi(sid, elem_type_id, fe_type_id).create(phi.size(), qrule->n_points());
295 :
296 16545 : for (unsigned int i = 0; i < phi.size(); ++i)
297 71881 : for (unsigned int qp = 0; qp < qrule->n_points(); ++qp)
298 : {
299 58830 : _phi(sid, elem_type_id, fe_type_id)(i, qp) = phi[i][qp];
300 58830 : _grad_phi(sid, elem_type_id, fe_type_id)(i, qp) = grad_phi[i][qp];
301 : }
302 :
303 3494 : _phi_face(sid, elem_type_id, fe_type_id).create(elem->n_sides());
304 3494 : _grad_phi_face(sid, elem_type_id, fe_type_id).create(elem->n_sides());
305 :
306 17216 : for (unsigned int side = 0; side < elem->n_sides(); ++side)
307 : {
308 13722 : auto & phi = fe_face->get_phi();
309 13722 : auto & grad_phi = fe_face->get_dphi();
310 :
311 13722 : fe_face->reinit(elem, side);
312 :
313 13722 : _phi_face(sid, elem_type_id, fe_type_id)(side).create(phi.size(), qrule_face->n_points());
314 13722 : _grad_phi_face(sid, elem_type_id, fe_type_id)(side).create(phi.size(),
315 : qrule_face->n_points());
316 :
317 66432 : for (unsigned int i = 0; i < phi.size(); ++i)
318 168420 : for (unsigned int qp = 0; qp < qrule_face->n_points(); ++qp)
319 : {
320 115710 : _phi_face(sid, elem_type_id, fe_type_id)(side)(i, qp) = phi[i][qp];
321 115710 : _grad_phi_face(sid, elem_type_id, fe_type_id)(side)(i, qp) = grad_phi[i][qp];
322 : }
323 : }
324 : }
325 3498 : }
326 :
327 5888 : for (auto & [elem_type, elem_type_id] : kokkosMesh().getElementTypeMap())
328 : {
329 2942 : auto elem = &libMesh::ReferenceElem::get(elem_type);
330 2942 : auto fe_type = FEType(elem->default_order(), LAGRANGE);
331 2942 : auto shape_deriv_ptr = libMesh::FEInterface::shape_deriv_function(fe_type, elem);
332 :
333 2942 : std::unique_ptr<FEBase> fe(FEBase::build(_dimension, fe_type));
334 2942 : std::unique_ptr<FEBase> fe_face(FEBase::build(_dimension, fe_type));
335 :
336 2942 : fe->attach_quadrature_rule(qrule);
337 2942 : fe_face->attach_quadrature_rule(qrule_face);
338 :
339 2942 : auto & phi = fe->get_phi();
340 2942 : auto & grad_phi = fe->get_dphi();
341 :
342 2942 : fe->reinit(elem);
343 :
344 2942 : _map_phi(sid, elem_type_id).create(phi.size(), qrule->n_points());
345 2942 : _map_grad_phi(sid, elem_type_id).create(phi.size(), qrule->n_points());
346 :
347 14943 : for (unsigned int i = 0; i < phi.size(); ++i)
348 63082 : for (unsigned int qp = 0; qp < qrule->n_points(); ++qp)
349 : {
350 51081 : _map_phi(sid, elem_type_id)(i, qp) = phi[i][qp];
351 51081 : _map_grad_phi(sid, elem_type_id)(i, qp) = grad_phi[i][qp];
352 : }
353 :
354 2942 : _map_phi_face(sid, elem_type_id).create(elem->n_sides());
355 2942 : _map_grad_phi_face(sid, elem_type_id).create(elem->n_sides());
356 2942 : _map_psi_face(sid, elem_type_id).create(elem->n_sides());
357 2942 : _map_grad_psi_face(sid, elem_type_id).create(elem->n_sides());
358 :
359 2942 : _normal_dx_dxi(sid, elem_type_id).create(elem->n_sides());
360 2942 : _normal_dx_deta(sid, elem_type_id).create(elem->n_sides());
361 :
362 14348 : for (unsigned int side = 0; side < elem->n_sides(); ++side)
363 : {
364 11406 : auto side_ptr = elem->build_side_ptr(side);
365 :
366 11406 : auto & phi = fe_face->get_phi();
367 11406 : auto & grad_phi = fe_face->get_dphi();
368 :
369 11406 : fe_face->reinit(elem, side);
370 :
371 11406 : _map_phi_face(sid, elem_type_id)(side).create(phi.size(), qrule_face->n_points());
372 11406 : _map_grad_phi_face(sid, elem_type_id)(side).create(phi.size(), qrule_face->n_points());
373 :
374 59706 : for (unsigned int i = 0; i < phi.size(); ++i)
375 151104 : for (unsigned int qp = 0; qp < qrule_face->n_points(); ++qp)
376 : {
377 102804 : _map_phi_face(sid, elem_type_id)(side)(i, qp) = phi[i][qp];
378 102804 : _map_grad_phi_face(sid, elem_type_id)(side)(i, qp) = grad_phi[i][qp];
379 : }
380 :
381 11406 : auto & psi = fe_face->get_fe_map().get_psi();
382 11406 : auto & dpsidxi = fe_face->get_fe_map().get_dpsidxi();
383 11406 : auto & dpsideta = fe_face->get_fe_map().get_dpsideta();
384 :
385 11406 : _map_psi_face(sid, elem_type_id)(side).create(psi.size(), qrule_face->n_points());
386 11406 : _map_grad_psi_face(sid, elem_type_id)(side).create(psi.size(), qrule_face->n_points());
387 :
388 35046 : for (unsigned int i = 0; i < psi.size(); ++i)
389 73512 : for (unsigned int qp = 0; qp < qrule_face->n_points(); ++qp)
390 : {
391 49872 : _map_psi_face(sid, elem_type_id)(side)(i, qp) = psi[i][qp];
392 49872 : if (elem->dim() > 1)
393 49340 : _map_grad_psi_face(sid, elem_type_id)(side)(i, qp)(0) = dpsidxi[i][qp];
394 49872 : if (elem->dim() > 2)
395 8160 : _map_grad_psi_face(sid, elem_type_id)(side)(i, qp)(1) = dpsideta[i][qp];
396 : }
397 :
398 11406 : _normal_dx_dxi(sid, elem_type_id)(side).create(phi.size(), qrule_face->n_points());
399 11406 : _normal_dx_deta(sid, elem_type_id)(side).create(phi.size(), qrule_face->n_points());
400 :
401 34058 : for (unsigned int qp = 0; qp < qrule_face->n_points(); ++qp)
402 : {
403 22652 : Point reference_point;
404 :
405 22652 : if (_dimension == 1)
406 532 : reference_point = side ? Point(1) : Point(-1);
407 22120 : else if (_dimension == 2)
408 61260 : for (unsigned int i = 0; i < psi.size(); ++i)
409 41180 : reference_point.add_scaled(side_ptr->point(i), psi[i][qp]);
410 :
411 125456 : for (unsigned int i = 0; i < phi.size(); ++i)
412 : {
413 102804 : if (_dimension < 3)
414 86484 : _normal_dx_dxi(sid, elem_type_id)(side)(i, qp) =
415 46908 : shape_deriv_ptr(fe_type, elem, i, 0, reference_point, false);
416 102804 : if (_dimension == 2)
417 85420 : _normal_dx_deta(sid, elem_type_id)(side)(i, qp) =
418 46348 : shape_deriv_ptr(fe_type, elem, i, 1, reference_point, false);
419 : }
420 : }
421 11406 : }
422 2942 : }
423 : }
424 :
425 2265 : _phi.copyToDeviceNested();
426 2265 : _phi_face.copyToDeviceNested();
427 2265 : _grad_phi.copyToDeviceNested();
428 2265 : _grad_phi_face.copyToDeviceNested();
429 :
430 2265 : _map_phi.copyToDeviceNested();
431 2265 : _map_phi_face.copyToDeviceNested();
432 2265 : _map_psi_face.copyToDeviceNested();
433 2265 : _map_grad_phi.copyToDeviceNested();
434 2265 : _map_grad_phi_face.copyToDeviceNested();
435 2265 : _map_grad_psi_face.copyToDeviceNested();
436 :
437 2265 : _normal_dx_dxi.copyToDeviceNested();
438 2265 : _normal_dx_deta.copyToDeviceNested();
439 :
440 2265 : _n_dofs.copyToDevice();
441 2265 : }
442 :
443 : void
444 2265 : Assembly::cachePhysicalMap()
445 : {
446 2265 : const auto num_subdomains = kokkosMesh().getNumSubdomains();
447 2265 : const auto num_elems = kokkosMesh().getNumLocalElements();
448 :
449 2265 : _jacobian.create(num_subdomains);
450 2265 : _jxw.create(num_subdomains);
451 2265 : _xyz.create(num_subdomains);
452 :
453 5211 : for (auto subdomain : _mesh.meshSubdomains())
454 : {
455 2946 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
456 :
457 2946 : _jacobian[sid].createDevice(_n_subdomain_qps[sid]);
458 2946 : _jxw[sid].createDevice(_n_subdomain_qps[sid]);
459 2946 : _xyz[sid].createDevice(_n_subdomain_qps[sid]);
460 : }
461 :
462 2265 : _jacobian.copyToDeviceNested();
463 2265 : _jxw.copyToDeviceNested();
464 2265 : _xyz.copyToDeviceNested();
465 :
466 2265 : ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(0, num_elems);
467 2265 : ::Kokkos::parallel_for(policy, *this);
468 2265 : ::Kokkos::fence();
469 2265 : }
470 :
471 : KOKKOS_FUNCTION void
472 385485 : Assembly::operator()(const ThreadID tid) const
473 : {
474 385485 : auto info = kokkosMesh().getElementInfo(tid);
475 385485 : auto offset = getQpOffset(info);
476 :
477 385485 : auto jacobian = &_jacobian[info.subdomain][offset];
478 385485 : auto jxw = &_jxw[info.subdomain][offset];
479 385485 : auto xyz = &_xyz[info.subdomain][offset];
480 :
481 1005001 : for (unsigned int qp = 0; qp < getNumQps(info); ++qp)
482 619516 : computePhysicalMap(info, qp, &jacobian[qp], &jxw[qp], &xyz[qp]);
483 385485 : }
484 :
485 : } // namespace Moose::Kokkos
|
