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 "KokkosSystem.h"
11 : #include "KokkosNodalBCBase.h"
12 :
13 : #include "MooseMesh.h"
14 : #include "Assembly.h"
15 : #include "NonlinearSystemBase.h"
16 : #include "FEProblemBase.h"
17 :
18 : #include "libmesh/system.h"
19 :
20 : namespace Moose
21 : {
22 : namespace Kokkos
23 : {
24 :
25 1614 : System::System(SystemBase & system)
26 1252 : : MeshHolder(*system.mesh().getKokkosMesh()),
27 1252 : AssemblyHolder(system.feProblem().kokkosAssembly()),
28 1252 : _system(system),
29 1252 : _mesh(system.mesh()),
30 1252 : _dof_map(system.dofMap()),
31 1252 : _comm(system.dofMap().comm()),
32 1252 : _num_vars(system.system().n_vars()),
33 1252 : _num_local_dofs(system.system().n_local_dofs()),
34 5008 : _num_ghost_dofs(system.dofMap().get_send_list().size())
35 : {
36 1614 : setupVariables();
37 1614 : setupDofs();
38 :
39 1614 : if (dynamic_cast<NonlinearSystemBase *>(&_system))
40 : {
41 807 : setupSparsity();
42 807 : checkNodalBCs();
43 : }
44 :
45 1614 : _qp_solutions.create(MAX_TAG);
46 1614 : _qp_solutions_grad.create(MAX_TAG);
47 1614 : }
48 :
49 : void
50 1614 : System::setupVariables()
51 : {
52 1614 : auto & sys = _system.system();
53 :
54 1614 : _var_fe_types.create(_num_vars);
55 1614 : _var_subdomain_active.create(_num_vars, _mesh.meshSubdomains().size());
56 :
57 3140 : for (unsigned int var = 0; var < _num_vars; ++var)
58 : {
59 1526 : _var_fe_types[var] = kokkosAssembly().getFETypeID(sys.variable_type(var));
60 :
61 3114 : for (auto subdomain : _mesh.meshSubdomains())
62 1588 : _var_subdomain_active(var, kokkosMesh().getContiguousSubdomainID(subdomain)) =
63 1223 : sys.variable(var).active_on_subdomain(subdomain);
64 : }
65 :
66 1614 : _var_fe_types.copyToDevice();
67 1614 : _var_subdomain_active.copyToDevice();
68 :
69 : // Set coupling
70 :
71 1614 : auto nl_system = dynamic_cast<NonlinearSystemBase *>(&_system);
72 :
73 1614 : if (nl_system)
74 : {
75 807 : _coupling.create(_num_vars);
76 :
77 807 : std::map<unsigned int, std::vector<unsigned int>> coupling;
78 :
79 807 : auto & ce = _system.feProblem().couplingEntries(0, nl_system->number());
80 :
81 2208 : for (const auto & [ivar, jvar] : ce)
82 1401 : if (ivar->number() != jvar->number())
83 378 : coupling[ivar->number()].push_back(jvar->number());
84 :
85 1830 : for (unsigned int var = 0; var < _num_vars; ++var)
86 1023 : _coupling[var] = coupling[var];
87 :
88 807 : _coupling.copyToDevice();
89 807 : }
90 1614 : }
91 :
92 : void
93 1614 : System::setupDofs()
94 : {
95 1614 : auto & sys = _system.system();
96 :
97 1614 : auto num_elems = kokkosMesh().getContiguousElementMap().size();
98 1614 : auto num_nodes = kokkosMesh().getContiguousNodeMap().size();
99 :
100 1614 : _residual_tag_active.create(MAX_TAG);
101 1614 : _residual_tag_active = false;
102 :
103 1614 : _matrix_tag_active.create(MAX_TAG);
104 1614 : _matrix_tag_active = false;
105 :
106 1614 : _vectors.create(MAX_TAG);
107 1614 : _matrices.create(MAX_TAG);
108 :
109 1614 : _local_elem_dof_index.create(_num_vars);
110 1614 : _local_node_dof_index.create(_num_vars);
111 1614 : _local_to_global_dof_index.create(_num_local_dofs + _num_ghost_dofs);
112 :
113 1614 : _max_dofs_per_elem.create(_num_vars);
114 1614 : _max_dofs_per_elem = 0;
115 :
116 1614 : auto solution = dynamic_cast<PetscVector<Number> *>(sys.current_local_solution.get());
117 :
118 1614 : std::vector<dof_id_type> dof_indices;
119 :
120 3140 : for (unsigned int var = 0; var < _num_vars; ++var)
121 : {
122 82078 : for (auto & [elem, id] : kokkosMesh().getContiguousElementMap())
123 : {
124 80552 : _dof_map.dof_indices(elem, dof_indices, var);
125 :
126 140055 : _max_dofs_per_elem[var] =
127 59503 : std::max(_max_dofs_per_elem[var], static_cast<unsigned int>(dof_indices.size()));
128 : }
129 :
130 1526 : _local_elem_dof_index[var].create(num_elems, _max_dofs_per_elem[var]);
131 1526 : _local_elem_dof_index[var] = libMesh::DofObject::invalid_id;
132 :
133 1526 : _local_node_dof_index[var].create(num_nodes);
134 1526 : _local_node_dof_index[var] = libMesh::DofObject::invalid_id;
135 :
136 82078 : for (auto & [elem, id] : kokkosMesh().getContiguousElementMap())
137 : {
138 80552 : _dof_map.dof_indices(elem, dof_indices, var);
139 :
140 389830 : for (unsigned int dof = 0; dof < dof_indices.size(); ++dof)
141 : {
142 309278 : _local_elem_dof_index[var](id, dof) = solution->map_global_to_local_index(dof_indices[dof]);
143 309278 : _local_to_global_dof_index[_local_elem_dof_index[var](id, dof)] = dof_indices[dof];
144 : }
145 : }
146 :
147 116464 : for (auto & [node, id] : kokkosMesh().getContiguousNodeMap())
148 114938 : if (node->processor_id() == _comm.rank())
149 : {
150 112674 : _dof_map.dof_indices(node, dof_indices, var);
151 :
152 112674 : if (dof_indices.size())
153 : {
154 90495 : for (unsigned int i = 1; i < dof_indices.size(); ++i)
155 0 : if (dof_indices[i] != dof_indices[i - 1] + 1)
156 0 : mooseError("Kokkos system error: a variable has multiple DOFs on a node, but the DOF "
157 : "indices are discontiguous. This is not supported.");
158 :
159 90495 : _local_node_dof_index[var][id] = solution->map_global_to_local_index(dof_indices[0]);
160 : }
161 : }
162 : }
163 :
164 1614 : _local_elem_dof_index.copyToDeviceNested();
165 1614 : _local_node_dof_index.copyToDeviceNested();
166 1614 : _local_to_global_dof_index.copyToDevice();
167 :
168 1614 : _max_dofs_per_elem.copyToDevice();
169 :
170 : // Setup DOF communication maps
171 :
172 1614 : auto num_procs = _comm.size();
173 :
174 3228 : std::vector<std::vector<dof_id_type>> send_list(num_procs);
175 1614 : std::vector<std::vector<dof_id_type>> recv_list(num_procs);
176 :
177 3511 : for (auto dof : _dof_map.get_send_list())
178 1897 : recv_list[_dof_map.dof_owner(dof)].push_back(dof);
179 :
180 3796 : for (processor_id_type proc = 0; proc < num_procs; proc++)
181 2182 : _comm.scatter(recv_list, send_list[proc], proc);
182 :
183 1614 : _local_comm_list.create(num_procs);
184 1614 : _ghost_comm_list.create(num_procs);
185 :
186 3796 : for (processor_id_type proc = 0; proc < num_procs; proc++)
187 : {
188 2182 : _local_comm_list[proc].create(send_list[proc].size());
189 2182 : _ghost_comm_list[proc].create(recv_list[proc].size());
190 :
191 4079 : for (dof_id_type i = 0; i < send_list[proc].size(); ++i)
192 1897 : _local_comm_list[proc][i] = solution->map_global_to_local_index(send_list[proc][i]);
193 :
194 4079 : for (dof_id_type i = 0; i < recv_list[proc].size(); ++i)
195 1897 : _ghost_comm_list[proc][i] = solution->map_global_to_local_index(recv_list[proc][i]);
196 : }
197 :
198 1614 : _local_comm_list.copyToDeviceNested();
199 1614 : _ghost_comm_list.copyToDeviceNested();
200 1614 : }
201 :
202 : void
203 807 : System::setupSparsity()
204 : {
205 807 : auto pattern = _dof_map.get_sparsity_pattern();
206 :
207 807 : std::vector<PetscInt> col_idx;
208 807 : std::vector<PetscInt> row_idx;
209 1614 : std::vector<PetscInt> row_ptr = {0};
210 :
211 67516 : for (dof_id_type r = _dof_map.first_dof(); r < _dof_map.end_dof(); ++r)
212 : {
213 66709 : auto & cols = pattern->get_sparsity_pattern().at(r - _dof_map.first_dof());
214 :
215 699274 : for (auto col : cols)
216 : {
217 632565 : col_idx.push_back(col);
218 632565 : row_idx.push_back(r);
219 : }
220 :
221 66709 : row_ptr.push_back(col_idx.size());
222 : }
223 :
224 807 : auto num_procs = _comm.size();
225 :
226 3228 : std::vector<std::vector<PetscInt>> send_cols(num_procs), send_num_cols(num_procs);
227 2421 : std::vector<std::vector<PetscInt>> recv_cols(num_procs), recv_num_cols(num_procs);
228 :
229 1898 : for (processor_id_type proc = 0; proc < num_procs; proc++)
230 2462 : for (auto r : _local_comm_list[proc])
231 : {
232 16645 : for (PetscInt c = row_ptr[r]; c < row_ptr[r + 1]; ++c)
233 15274 : send_cols[proc].push_back(col_idx[c]);
234 :
235 1371 : send_num_cols[proc].push_back(row_ptr[r + 1] - row_ptr[r]);
236 : }
237 :
238 1898 : for (processor_id_type proc = 0; proc < num_procs; proc++)
239 : {
240 1091 : _comm.scatter(send_cols, recv_cols[proc], proc);
241 1091 : _comm.scatter(send_num_cols, recv_num_cols[proc], proc);
242 : }
243 :
244 2421 : std::vector<PetscInt> col(num_procs), row(num_procs);
245 :
246 2178 : for (auto r : _dof_map.get_send_list())
247 : {
248 1371 : auto proc = _dof_map.dof_owner(r);
249 1371 : auto n_cols = recv_num_cols[proc][row[proc]++];
250 :
251 16645 : for (PetscInt c = 0; c < n_cols; ++c)
252 : {
253 15274 : col_idx.push_back(recv_cols[proc][col[proc]++]);
254 15274 : row_idx.push_back(r);
255 : }
256 :
257 1371 : row_ptr.push_back(col_idx.size());
258 : }
259 :
260 807 : _sparsity.col_idx = col_idx;
261 807 : _sparsity.row_idx = row_idx;
262 807 : _sparsity.row_ptr = row_ptr;
263 807 : }
264 :
265 : void
266 807 : System::checkNodalBCs()
267 : {
268 807 : auto & nl_system = static_cast<NonlinearSystemBase &>(_system);
269 :
270 807 : _nbc_residual_tag_dof.create(MAX_TAG);
271 807 : _nbc_matrix_tag_dof.create(MAX_TAG);
272 :
273 2098 : for (auto bc : nl_system.getKokkosNodalBCWarehouse().getActiveObjects())
274 : {
275 1291 : auto nbc = static_cast<NodalBCBase *>(bc.get());
276 :
277 1291 : std::set<TagID> extra_vector_tags;
278 1291 : std::set<TagID> extra_matrix_tags;
279 :
280 3873 : if (nbc->isParamValid("extra_vector_tags"))
281 384 : for (auto tag : nbc->getParam<std::vector<TagName>>("extra_vector_tags"))
282 96 : extra_vector_tags.insert(_system.feProblem().getVectorTagID(tag));
283 :
284 3873 : if (nbc->isParamValid("extra_matrix_tags"))
285 372 : for (auto tag : nbc->getParam<std::vector<TagName>>("extra_matrix_tags"))
286 120 : extra_matrix_tags.insert(_system.feProblem().getMatrixTagID(tag));
287 :
288 1291 : getNodalBCDofs(nbc, _nbc_dof);
289 :
290 1387 : for (auto tag : extra_vector_tags)
291 96 : getNodalBCDofs(nbc, _nbc_residual_tag_dof[tag]);
292 :
293 1411 : for (auto tag : extra_matrix_tags)
294 120 : getNodalBCDofs(nbc, _nbc_matrix_tag_dof[tag]);
295 1291 : }
296 :
297 807 : _nbc_dof.copyToDevice();
298 807 : _nbc_residual_tag_dof.copyToDeviceNested();
299 807 : _nbc_matrix_tag_dof.copyToDeviceNested();
300 807 : }
301 :
302 : void
303 1507 : System::getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs)
304 : {
305 1507 : auto var = nbc->variable().number();
306 :
307 1507 : if (!dofs.isAlloc())
308 : {
309 903 : dofs.create(_num_local_dofs + _num_ghost_dofs);
310 903 : dofs = false;
311 : }
312 :
313 10552 : for (auto node : nbc->getContiguousNodes())
314 10552 : dofs[_local_node_dof_index[var][node]] = true;
315 :
316 : // Let remote processes know about ghost DOFs associated with nodal BCs not to compute
317 : // residual on them
318 :
319 1507 : auto num_procs = _comm.size();
320 :
321 4521 : std::vector<std::vector<char>> send(num_procs), recv(num_procs);
322 :
323 3542 : for (processor_id_type proc = 0; proc < num_procs; proc++)
324 4556 : for (auto dof : _local_comm_list[proc])
325 2521 : send[proc].push_back(dofs[dof]);
326 :
327 3542 : for (processor_id_type proc = 0; proc < num_procs; proc++)
328 2035 : _comm.scatter(send, recv[proc], proc);
329 :
330 3542 : for (processor_id_type proc = 0; proc < num_procs; proc++)
331 4556 : for (dof_id_type i = 0; i < _ghost_comm_list[proc].size(); ++i)
332 2521 : dofs[_ghost_comm_list[proc][i]] = recv[proc][i];
333 1507 : }
334 :
335 : void
336 208016 : System::sync(const MemcpyKind dir)
337 : {
338 208016 : if (dir == MemcpyKind::HOST_TO_DEVICE)
339 : {
340 256201 : for (auto tag : _active_solution_tags)
341 : {
342 152193 : auto & vector = _system.getVector(tag);
343 :
344 152193 : _vectors[tag].create(vector, *this, false);
345 152193 : _vectors[tag].copyToDevice();
346 : }
347 :
348 207107 : for (auto tag : _active_residual_tags)
349 : {
350 103099 : auto & vector = _system.getVector(tag);
351 :
352 103099 : _vectors[tag].create(vector, *this, true);
353 103099 : _vectors[tag] = 0;
354 : }
355 :
356 109861 : for (auto tag : _active_matrix_tags)
357 : {
358 5853 : auto & matrix = _system.getMatrix(tag);
359 :
360 5853 : _matrices[tag].create(matrix, *this);
361 5853 : _matrices[tag] = 0;
362 : }
363 :
364 104008 : _vectors.copyToDevice();
365 104008 : _matrices.copyToDevice();
366 : }
367 104008 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
368 : {
369 256201 : for (auto tag : _active_solution_tags)
370 152193 : _vectors[tag].restore();
371 :
372 207107 : for (auto tag : _active_residual_tags)
373 103099 : _vectors[tag].close();
374 :
375 109861 : for (auto tag : _active_matrix_tags)
376 5853 : _matrices[tag].close();
377 : }
378 208016 : }
379 :
380 : void
381 33364 : System::sync(const std::set<TagID> & tags, const MemcpyKind dir)
382 : {
383 33364 : if (dir == MemcpyKind::HOST_TO_DEVICE)
384 : {
385 33384 : for (auto tag : tags)
386 : {
387 16702 : if (!_system.hasVector(tag))
388 2476 : continue;
389 :
390 14226 : auto & vector = _system.getVector(tag);
391 :
392 14226 : _vectors[tag].create(vector, *this, false);
393 14226 : _vectors[tag].copyToDevice();
394 : }
395 :
396 16682 : _vectors.copyToDevice();
397 : }
398 16682 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
399 : {
400 33384 : for (auto tag : tags)
401 : {
402 16702 : if (!_system.hasVector(tag))
403 2476 : continue;
404 :
405 14226 : _vectors[tag].copyToHost();
406 14226 : _vectors[tag].restore();
407 : }
408 : }
409 33364 : }
410 :
411 : void
412 0 : System::sync(const std::vector<TagID> & tags, const MemcpyKind dir)
413 : {
414 0 : sync(std::set<TagID>(tags.begin(), tags.end()), dir);
415 0 : }
416 :
417 : void
418 23460 : System::sync(const TagID tag, const MemcpyKind dir)
419 : {
420 46920 : sync(std::set<TagID>{tag}, dir);
421 23460 : }
422 :
423 : void
424 104008 : System::setActiveVariables(const std::set<MooseVariableFieldBase *> & vars)
425 : {
426 104008 : std::set<unsigned int> active_variables;
427 :
428 233968 : for (auto var : vars)
429 129960 : if (var->sys().number() == _system.number())
430 129960 : for (unsigned int i = 0; i < var->count(); ++i)
431 64980 : active_variables.insert(var->number() + i);
432 :
433 104008 : _active_variables = active_variables;
434 104008 : }
435 :
436 : void
437 104008 : System::setActiveSolutionTags(const std::set<TagID> & tags)
438 : {
439 104008 : if (!_active_variables.size())
440 47796 : return;
441 :
442 56212 : std::set<TagID> active_solution_tags;
443 :
444 313078 : for (auto tag : tags)
445 256866 : if (_system.hasVector(tag))
446 152193 : active_solution_tags.insert(tag);
447 :
448 56212 : _active_solution_tags = active_solution_tags;
449 56212 : }
450 :
451 : void
452 34353 : System::setActiveResidualTags(const std::set<TagID> & tags)
453 : {
454 34353 : std::set<TagID> active_residual_tags;
455 :
456 137452 : for (auto tag : tags)
457 103099 : if (_system.hasVector(tag))
458 : {
459 103099 : active_residual_tags.insert(tag);
460 103099 : _residual_tag_active[tag] = true;
461 : }
462 :
463 34353 : _active_residual_tags = active_residual_tags;
464 :
465 34353 : _residual_tag_active.copyToDevice();
466 34353 : }
467 :
468 : void
469 5611 : System::setActiveMatrixTags(const std::set<TagID> & tags)
470 : {
471 5611 : std::set<TagID> active_matrix_tags;
472 :
473 16822 : for (auto tag : tags)
474 11211 : if (_system.hasMatrix(tag))
475 : {
476 5853 : active_matrix_tags.insert(tag);
477 5853 : _matrix_tag_active[tag] = true;
478 : }
479 :
480 5611 : _active_matrix_tags = active_matrix_tags;
481 :
482 5611 : _matrix_tag_active.copyToDevice();
483 5611 : }
484 :
485 : void
486 104008 : System::reinit()
487 : {
488 256201 : for (auto tag : _active_solution_tags)
489 : {
490 152193 : if (!_qp_solutions[tag].isAlloc())
491 2101 : _qp_solutions[tag].create(_mesh.meshSubdomains().size(), _num_vars);
492 :
493 152193 : if (!_qp_solutions_grad[tag].isAlloc())
494 2101 : _qp_solutions_grad[tag].create(_mesh.meshSubdomains().size(), _num_vars);
495 :
496 309082 : for (auto subdomain : _mesh.meshSubdomains())
497 : {
498 156889 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
499 :
500 337035 : for (auto var : _active_variables)
501 : {
502 180146 : if (!_var_subdomain_active(var, sid))
503 204 : continue;
504 :
505 179942 : if (!_qp_solutions[tag](sid, var).isAlloc())
506 2467 : _qp_solutions[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
507 :
508 179942 : if (!_qp_solutions_grad[tag](sid, var).isAlloc())
509 2467 : _qp_solutions_grad[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
510 : }
511 : }
512 :
513 152193 : _qp_solutions[tag].copyToDevice();
514 152193 : _qp_solutions_grad[tag].copyToDevice();
515 : }
516 :
517 104008 : _qp_solutions.copyToDevice();
518 104008 : _qp_solutions_grad.copyToDevice();
519 :
520 104008 : dof_id_type num_elems = kokkosMesh().getContiguousElementMap().size();
521 :
522 297622 : _thread.resize({kokkosAssembly().getMaxQpsPerElem(),
523 : num_elems,
524 89606 : _active_variables.size(),
525 89606 : _active_solution_tags.size()});
526 :
527 104008 : ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(0, _thread.size());
528 104008 : ::Kokkos::parallel_for(policy, *this);
529 104008 : ::Kokkos::fence();
530 104008 : }
531 :
532 : KOKKOS_FUNCTION void
533 53213772 : System::operator()(const ThreadID tid) const
534 : {
535 53213772 : auto qp = _thread(tid, 0);
536 53213772 : auto elem = _thread(tid, 1);
537 53213772 : auto var = _active_variables(_thread(tid, 2));
538 53213772 : auto tag = _active_solution_tags(_thread(tid, 3));
539 :
540 53213772 : auto info = kokkosMesh().getElementInfo(elem);
541 53213772 : auto sid = info.subdomain;
542 53213772 : auto elem_type = info.type;
543 :
544 53213772 : if (!_var_subdomain_active(var, sid))
545 1936 : return;
546 :
547 53211836 : auto fe_type = _var_fe_types[var];
548 53211836 : auto num_dofs = kokkosAssembly().getNumDofs(elem_type, fe_type);
549 53211836 : auto num_qps = kokkosAssembly().getNumQps(info);
550 53211836 : auto qp_offset = kokkosAssembly().getQpOffset(info);
551 :
552 53211836 : if (qp >= num_qps)
553 0 : return;
554 :
555 53211836 : auto & phi = kokkosAssembly().getPhi(sid, elem_type, fe_type);
556 53211836 : auto & grad_phi = kokkosAssembly().getGradPhi(sid, elem_type, fe_type);
557 53211836 : auto jacobian = kokkosAssembly().getJacobian(info, qp);
558 :
559 53211836 : Real value = 0;
560 53211836 : Real3 grad = 0;
561 :
562 300578012 : for (unsigned int i = 0; i < num_dofs; ++i)
563 : {
564 247366176 : auto vector = getVectorDofValue(getElemLocalDofIndex(elem, i, var), tag);
565 :
566 247366176 : value += vector * phi(i, qp);
567 247366176 : grad += vector * (jacobian * grad_phi(i, qp));
568 : }
569 :
570 53211836 : getVectorQpValue(info, qp_offset + qp, var, tag) = value;
571 53211836 : getVectorQpGrad(info, qp_offset + qp, var, tag) = grad;
572 : }
573 :
574 : } // namespace Kokkos
575 : } // namespace Moose
|
