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