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::Kokkos
21 : {
22 :
23 8666 : System::System(SystemBase & system)
24 2462 : : PerfGraphInterface(system.feProblem().getMooseApp().perfGraph(), "KokkosSystem"),
25 2462 : MeshHolder(*system.mesh().getKokkosMesh()),
26 2462 : AssemblyHolder(system.feProblem().kokkosAssembly()),
27 2462 : _system(system),
28 2462 : _mesh(system.mesh()),
29 2462 : _dof_map(system.dofMap()),
30 2462 : _comm(system.dofMap().comm()),
31 2462 : _num_vars(system.system().n_vars()),
32 2462 : _num_local_dofs(system.system().n_local_dofs()),
33 17234 : _num_ghost_dofs(system.dofMap().get_send_list().size())
34 : {
35 4530 : setupVariables();
36 4530 : setupDofs();
37 :
38 4530 : if (dynamic_cast<NonlinearSystemBase *>(&_system))
39 : {
40 2265 : setupSparsity();
41 2265 : setupNodalBCDofs();
42 : }
43 :
44 4530 : _qp_solutions.create(MAX_TAG);
45 4530 : _qp_solutions_grad.create(MAX_TAG);
46 4530 : }
47 :
48 : void
49 4530 : System::setupVariables()
50 : {
51 4530 : auto & sys = _system.system();
52 :
53 4530 : _var_fe_types.create(_num_vars);
54 4530 : _var_subdomain_active.create(_num_vars, kokkosMesh().getNumSubdomains());
55 :
56 8482 : for (unsigned int var = 0; var < _num_vars; ++var)
57 : {
58 3952 : _var_fe_types[var] = kokkosAssembly().getFETypeID(sys.variable_type(var));
59 :
60 8941 : for (auto subdomain : _mesh.meshSubdomains())
61 4989 : _var_subdomain_active(var, kokkosMesh().getContiguousSubdomainID(subdomain)) =
62 2707 : sys.variable(var).active_on_subdomain(subdomain);
63 : }
64 :
65 4530 : _var_fe_types.copyToDevice();
66 4530 : _var_subdomain_active.copyToDevice();
67 :
68 : // Set coupling
69 :
70 4530 : auto nl_system = dynamic_cast<NonlinearSystemBase *>(&_system);
71 :
72 4530 : if (nl_system)
73 : {
74 2265 : _coupling.create(_num_vars);
75 :
76 2265 : std::map<unsigned int, std::vector<unsigned int>> coupling;
77 :
78 2265 : auto & ce = _system.feProblem().couplingEntries(0, nl_system->number());
79 :
80 5479 : for (const auto & [ivar, jvar] : ce)
81 3214 : if (ivar->number() != jvar->number())
82 718 : coupling[ivar->number()].push_back(jvar->number());
83 :
84 4761 : for (unsigned int var = 0; var < _num_vars; ++var)
85 2496 : _coupling[var] = coupling[var];
86 :
87 2265 : _coupling.copyToDevice();
88 2265 : }
89 4530 : }
90 :
91 : void
92 4530 : System::setupDofs()
93 : {
94 4530 : auto & sys = _system.system();
95 :
96 4530 : auto num_elems = kokkosMesh().getNumLocalElements();
97 4530 : auto num_nodes = kokkosMesh().getNumLocalNodes();
98 :
99 4530 : _residual_tag_active.create(MAX_TAG);
100 4530 : _residual_tag_active = false;
101 :
102 4530 : _matrix_tag_active.create(MAX_TAG);
103 4530 : _matrix_tag_active = false;
104 :
105 4530 : _vectors.create(MAX_TAG);
106 4530 : _matrices.create(MAX_TAG);
107 :
108 4530 : _local_elem_dof_index.create(_num_vars);
109 4530 : _local_node_dof_index.create(_num_vars);
110 4530 : _local_to_global_dof_index.create(_num_local_dofs + _num_ghost_dofs);
111 :
112 4530 : _max_dofs_per_elem.create(_num_vars);
113 4530 : _max_dofs_per_elem = 0;
114 :
115 4530 : auto solution = dynamic_cast<PetscVector<Number> *>(sys.current_local_solution.get());
116 :
117 : #ifdef MOOSE_ENABLE_KOKKOS_GPU
118 : // Kokkos array thinks OpenMP clause is device code when using OpenMP backend
119 2068 : #pragma omp parallel for
120 : #endif
121 4604 : for (unsigned int var = 0; var < _num_vars; ++var)
122 : {
123 2142 : std::vector<dof_id_type> dof_indices;
124 :
125 802810 : for (const auto elem : _mesh.getMesh().active_local_element_ptr_range())
126 : {
127 800668 : _dof_map.dof_indices(elem, dof_indices, var);
128 :
129 1601336 : _max_dofs_per_elem[var] =
130 800668 : std::max(_max_dofs_per_elem[var], static_cast<unsigned int>(dof_indices.size()));
131 2142 : }
132 :
133 2142 : _local_elem_dof_index[var].create(num_elems, _max_dofs_per_elem[var]);
134 2142 : _local_elem_dof_index[var] = libMesh::DofObject::invalid_id;
135 :
136 2142 : _local_node_dof_index[var].create(num_nodes);
137 2142 : _local_node_dof_index[var] = libMesh::DofObject::invalid_id;
138 :
139 802810 : for (const auto elem : _mesh.getMesh().active_local_element_ptr_range())
140 : {
141 800668 : const auto id = kokkosMesh().getContiguousElementID(elem);
142 :
143 800668 : _dof_map.dof_indices(elem, dof_indices, var);
144 :
145 4003776 : for (unsigned int dof = 0; dof < dof_indices.size(); ++dof)
146 : {
147 3203108 : _local_elem_dof_index[var](id, dof) = solution->map_global_to_local_index(dof_indices[dof]);
148 3203108 : _local_to_global_dof_index[_local_elem_dof_index[var](id, dof)] = dof_indices[dof];
149 : }
150 2142 : }
151 :
152 876425 : for (const auto node : kokkosMesh().getLocalNodes())
153 874283 : if (node->processor_id() == _comm.rank())
154 : {
155 866542 : const auto id = kokkosMesh().getContiguousNodeID(node);
156 :
157 866542 : _dof_map.dof_indices(node, dof_indices, var);
158 :
159 866542 : if (dof_indices.size())
160 : {
161 840007 : for (unsigned int i = 1; i < dof_indices.size(); ++i)
162 0 : if (dof_indices[i] != dof_indices[i - 1] + 1)
163 0 : mooseError("Kokkos system error: a variable has multiple DOFs on a node, but the DOF "
164 : "indices are discontiguous. This is not supported.");
165 :
166 840007 : _local_node_dof_index[var][id] = solution->map_global_to_local_index(dof_indices[0]);
167 : }
168 : }
169 2142 : }
170 :
171 4530 : _local_elem_dof_index.copyToDeviceNested();
172 4530 : _local_node_dof_index.copyToDeviceNested();
173 4530 : _local_to_global_dof_index.copyToDevice();
174 :
175 4530 : _max_dofs_per_elem.copyToDevice();
176 :
177 : // Setup DOF communication maps
178 :
179 4530 : auto num_procs = _comm.size();
180 :
181 9060 : std::vector<std::vector<dof_id_type>> send_list(num_procs);
182 4530 : std::vector<std::vector<dof_id_type>> recv_list(num_procs);
183 :
184 18120 : for (auto dof : _dof_map.get_send_list())
185 13590 : recv_list[_dof_map.dof_owner(dof)].push_back(dof);
186 :
187 11276 : for (processor_id_type proc = 0; proc < num_procs; proc++)
188 6746 : _comm.scatter(recv_list, send_list[proc], proc);
189 :
190 4530 : _local_comm_list.create(num_procs);
191 4530 : _ghost_comm_list.create(num_procs);
192 :
193 11276 : for (processor_id_type proc = 0; proc < num_procs; proc++)
194 : {
195 6746 : _local_comm_list[proc].create(send_list[proc].size());
196 6746 : _ghost_comm_list[proc].create(recv_list[proc].size());
197 :
198 20336 : for (dof_id_type i = 0; i < send_list[proc].size(); ++i)
199 13590 : _local_comm_list[proc][i] = solution->map_global_to_local_index(send_list[proc][i]);
200 :
201 20336 : for (dof_id_type i = 0; i < recv_list[proc].size(); ++i)
202 13590 : _ghost_comm_list[proc][i] = solution->map_global_to_local_index(recv_list[proc][i]);
203 : }
204 :
205 4530 : _local_comm_list.copyToDeviceNested();
206 4530 : _ghost_comm_list.copyToDeviceNested();
207 4530 : }
208 :
209 : void
210 2265 : System::setupSparsity()
211 : {
212 2265 : auto pattern = _dof_map.get_sparsity_pattern();
213 :
214 2265 : if (!pattern)
215 0 : return;
216 :
217 2265 : std::vector<PetscInt> col_idx;
218 2265 : std::vector<PetscInt> row_idx;
219 4530 : std::vector<PetscInt> row_ptr = {0};
220 :
221 168901 : for (dof_id_type r = _dof_map.first_dof(); r < _dof_map.end_dof(); ++r)
222 : {
223 166636 : auto & cols = pattern->get_sparsity_pattern().at(r - _dof_map.first_dof());
224 :
225 1840262 : for (auto col : cols)
226 : {
227 1673626 : col_idx.push_back(col);
228 1673626 : row_idx.push_back(r);
229 : }
230 :
231 166636 : row_ptr.push_back(col_idx.size());
232 : }
233 :
234 2265 : auto num_procs = _comm.size();
235 :
236 9060 : std::vector<std::vector<PetscInt>> send_cols(num_procs), send_num_cols(num_procs);
237 6795 : std::vector<std::vector<PetscInt>> recv_cols(num_procs), recv_num_cols(num_procs);
238 :
239 5638 : for (processor_id_type proc = 0; proc < num_procs; proc++)
240 8495 : for (auto r : _local_comm_list[proc])
241 : {
242 62452 : for (PetscInt c = row_ptr[r]; c < row_ptr[r + 1]; ++c)
243 57330 : send_cols[proc].push_back(col_idx[c]);
244 :
245 5122 : send_num_cols[proc].push_back(row_ptr[r + 1] - row_ptr[r]);
246 : }
247 :
248 5638 : for (processor_id_type proc = 0; proc < num_procs; proc++)
249 : {
250 3373 : _comm.scatter(send_cols, recv_cols[proc], proc);
251 3373 : _comm.scatter(send_num_cols, recv_num_cols[proc], proc);
252 : }
253 :
254 6795 : std::vector<PetscInt> col(num_procs), row(num_procs);
255 :
256 7387 : for (auto r : _dof_map.get_send_list())
257 : {
258 5122 : auto proc = _dof_map.dof_owner(r);
259 5122 : auto n_cols = recv_num_cols[proc][row[proc]++];
260 :
261 62452 : for (PetscInt c = 0; c < n_cols; ++c)
262 : {
263 57330 : col_idx.push_back(recv_cols[proc][col[proc]++]);
264 57330 : row_idx.push_back(r);
265 : }
266 :
267 5122 : row_ptr.push_back(col_idx.size());
268 : }
269 :
270 2265 : _sparsity.col_idx = col_idx;
271 2265 : _sparsity.row_idx = row_idx;
272 2265 : _sparsity.row_ptr = row_ptr;
273 2265 : }
274 :
275 : void
276 2265 : System::setupNodalBCDofs()
277 : {
278 2265 : auto & nl_system = static_cast<NonlinearSystemBase &>(_system);
279 :
280 2265 : _nbc_matrix_tag_dof.create(MAX_TAG);
281 :
282 5674 : for (auto bc : nl_system.getKokkosNodalBCWarehouse().getActiveObjects())
283 : {
284 3409 : auto nbc = static_cast<NodalBCBase *>(bc.get());
285 :
286 3409 : auto matrix_tags = nbc->getMatrixTags({});
287 :
288 10414 : for (auto tag : matrix_tags)
289 7005 : getNodalBCDofs(nbc, _nbc_matrix_tag_dof[tag]);
290 3409 : }
291 :
292 2265 : _nbc_matrix_tag_dof.copyToDeviceNested();
293 2265 : }
294 :
295 : void
296 7005 : System::getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs)
297 : {
298 7005 : auto var = nbc->variable().number();
299 :
300 7005 : if (!dofs.isAlloc())
301 : {
302 3740 : dofs.create(_num_local_dofs + _num_ghost_dofs);
303 3740 : dofs = false;
304 : }
305 :
306 46189 : for (auto node : nbc->getContiguousNodes())
307 46189 : dofs[_local_node_dof_index[var][node]] = true;
308 :
309 : // Let remote processes know about ghost DOFs associated with nodal BCs not to contribute on them
310 :
311 7005 : auto num_procs = _comm.size();
312 :
313 21015 : std::vector<std::vector<char>> send(num_procs), recv(num_procs);
314 :
315 17426 : for (processor_id_type proc = 0; proc < num_procs; proc++)
316 25889 : for (auto dof : _local_comm_list[proc])
317 15468 : send[proc].push_back(dofs[dof]);
318 :
319 17426 : for (processor_id_type proc = 0; proc < num_procs; proc++)
320 10421 : _comm.scatter(send, recv[proc], proc);
321 :
322 17426 : for (processor_id_type proc = 0; proc < num_procs; proc++)
323 25889 : for (dof_id_type i = 0; i < _ghost_comm_list[proc].size(); ++i)
324 15468 : dofs[_ghost_comm_list[proc][i]] = recv[proc][i];
325 7005 : }
326 :
327 : void
328 639568 : System::sync(const MemcpyType dir)
329 : {
330 639568 : if (dir == MemcpyType::HOST_TO_DEVICE)
331 : {
332 1030772 : for (auto tag : _active_solution_tags)
333 : {
334 710988 : auto & vector = _system.getVector(tag);
335 :
336 710988 : _vectors[tag].create(vector, *this, false);
337 710988 : _vectors[tag].copyToDevice();
338 : }
339 :
340 696731 : for (auto tag : _active_residual_tags)
341 : {
342 376947 : auto & vector = _system.getVector(tag);
343 :
344 376947 : _vectors[tag].create(vector, *this, true);
345 376947 : _vectors[tag].copyToDevice();
346 : }
347 :
348 335937 : for (auto tag : _active_matrix_tags)
349 : {
350 16153 : auto & matrix = _system.getMatrix(tag);
351 :
352 : {
353 48459 : TIME_SECTION("MatSetPreallocation", 1);
354 16153 : _matrices[tag].create(matrix, *this);
355 16153 : }
356 :
357 16153 : _matrices[tag] = 0;
358 : }
359 :
360 319784 : _vectors.copyToDevice();
361 319784 : _matrices.copyToDevice();
362 : }
363 319784 : else if (dir == MemcpyType::DEVICE_TO_HOST)
364 : {
365 1030772 : for (auto tag : _active_solution_tags)
366 710988 : _vectors[tag].restore();
367 :
368 696731 : for (auto tag : _active_residual_tags)
369 376947 : _vectors[tag].close();
370 :
371 335937 : for (auto tag : _active_matrix_tags)
372 : {
373 48459 : TIME_SECTION("MatSetValues", 1);
374 16153 : _matrices[tag].close();
375 16153 : }
376 : }
377 639568 : }
378 :
379 : void
380 50946 : System::sync(const std::set<TagID> & tags, const MemcpyType dir)
381 : {
382 50946 : if (dir == MemcpyType::HOST_TO_DEVICE)
383 : {
384 50976 : for (auto tag : tags)
385 : {
386 25503 : if (!_system.hasVector(tag))
387 4368 : continue;
388 :
389 21135 : auto & vector = _system.getVector(tag);
390 :
391 21135 : _vectors[tag].create(vector, *this, false);
392 21135 : _vectors[tag].copyToDevice();
393 : }
394 :
395 25473 : _vectors.copyToDevice();
396 : }
397 25473 : else if (dir == MemcpyType::DEVICE_TO_HOST)
398 : {
399 50976 : for (auto tag : tags)
400 : {
401 25503 : if (!_system.hasVector(tag))
402 4368 : continue;
403 :
404 21135 : _vectors[tag].copyToHost();
405 21135 : _vectors[tag].restore();
406 : }
407 : }
408 50946 : }
409 :
410 : void
411 0 : System::sync(const std::vector<TagID> & tags, const MemcpyType dir)
412 : {
413 0 : sync(std::set<TagID>(tags.begin(), tags.end()), dir);
414 0 : }
415 :
416 : void
417 33474 : System::sync(const TagID tag, const MemcpyType dir)
418 : {
419 66948 : sync(std::set<TagID>{tag}, dir);
420 33474 : }
421 :
422 : void
423 190102 : System::setActiveVariables(const std::set<MooseVariableFieldBase *> & vars)
424 : {
425 190102 : std::set<unsigned int> active_variables;
426 :
427 475512 : for (auto var : vars)
428 285410 : if (var->sys().number() == _system.number())
429 285410 : for (unsigned int i = 0; i < var->count(); ++i)
430 142705 : active_variables.insert(var->number() + i);
431 :
432 190102 : _active_variables = active_variables;
433 190102 : }
434 :
435 : void
436 319784 : System::setActiveSolutionTags(const std::set<TagID> & tags)
437 : {
438 319784 : std::set<TagID> active_solution_tags;
439 :
440 1689934 : for (auto tag : tags)
441 1370150 : if (_system.hasVector(tag))
442 710988 : active_solution_tags.insert(tag);
443 :
444 319784 : _active_solution_tags = active_solution_tags;
445 319784 : }
446 :
447 : void
448 131257 : System::setActiveResidualTags(const std::set<TagID> & tags)
449 : {
450 131257 : std::set<TagID> active_residual_tags;
451 :
452 508204 : for (auto tag : tags)
453 376947 : if (_system.hasVector(tag))
454 : {
455 376947 : active_residual_tags.insert(tag);
456 376947 : _residual_tag_active[tag] = true;
457 : }
458 :
459 131257 : _active_residual_tags = active_residual_tags;
460 :
461 131257 : _residual_tag_active.copyToDevice();
462 131257 : }
463 :
464 : void
465 15791 : System::setActiveMatrixTags(const std::set<TagID> & tags)
466 : {
467 15791 : std::set<TagID> active_matrix_tags;
468 :
469 47237 : for (auto tag : tags)
470 31446 : if (_system.hasMatrix(tag))
471 : {
472 16153 : active_matrix_tags.insert(tag);
473 16153 : _matrix_tag_active[tag] = true;
474 : }
475 :
476 15791 : _active_matrix_tags = active_matrix_tags;
477 :
478 15791 : _matrix_tag_active.copyToDevice();
479 15791 : }
480 :
481 : void
482 190102 : System::reinit()
483 : {
484 615312 : for (auto tag : _active_solution_tags)
485 : {
486 425210 : if (!_qp_solutions[tag].isAlloc())
487 6477 : _qp_solutions[tag].create(kokkosMesh().getNumSubdomains(), _num_vars);
488 :
489 425210 : if (!_qp_solutions_grad[tag].isAlloc())
490 6477 : _qp_solutions_grad[tag].create(kokkosMesh().getNumSubdomains(), _num_vars);
491 :
492 873085 : for (auto subdomain : _mesh.meshSubdomains())
493 : {
494 447875 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
495 :
496 824077 : for (auto var : _active_variables)
497 : {
498 376202 : if (!_var_subdomain_active(var, sid))
499 626 : continue;
500 :
501 375576 : if (!_qp_solutions[tag](sid, var).isAlloc())
502 6547 : _qp_solutions[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
503 :
504 375576 : if (!_qp_solutions_grad[tag](sid, var).isAlloc())
505 6547 : _qp_solutions_grad[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
506 : }
507 : }
508 :
509 425210 : _qp_solutions[tag].copyToDevice();
510 425210 : _qp_solutions_grad[tag].copyToDevice();
511 : }
512 :
513 190102 : _qp_solutions.copyToDevice();
514 190102 : _qp_solutions_grad.copyToDevice();
515 :
516 190102 : dof_id_type num_elems = kokkosMesh().getNumLocalElements();
517 :
518 190102 : _thread.resize(kokkosAssembly().getMaxQpsPerElem(),
519 : num_elems,
520 : _active_variables.size(),
521 : _active_solution_tags.size());
522 :
523 190102 : ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(0, _thread.size());
524 190102 : ::Kokkos::parallel_for(policy, *this);
525 190102 : ::Kokkos::fence();
526 190102 : }
527 :
528 : KOKKOS_FUNCTION void
529 67877944 : System::operator()(const ThreadID tid) const
530 : {
531 67877944 : auto qp = _thread(tid, 0);
532 67877944 : auto elem = _thread(tid, 1);
533 67877944 : auto var = _active_variables(_thread(tid, 2));
534 67877944 : auto tag = _active_solution_tags(_thread(tid, 3));
535 :
536 67877944 : auto info = kokkosMesh().getElementInfo(elem);
537 67877944 : auto sid = info.subdomain;
538 67877944 : auto elem_type = info.type;
539 :
540 67877944 : if (!_var_subdomain_active(var, sid))
541 3872 : return;
542 :
543 67874072 : auto fe_type = _var_fe_types[var];
544 67874072 : auto num_dofs = kokkosAssembly().getNumDofs(elem_type, fe_type);
545 67874072 : auto num_qps = kokkosAssembly().getNumQps(info);
546 67874072 : auto qp_offset = kokkosAssembly().getQpOffset(info);
547 :
548 67874072 : if (qp >= num_qps)
549 0 : return;
550 :
551 67874072 : auto & phi = kokkosAssembly().getPhi(sid, elem_type, fe_type);
552 67874072 : auto & grad_phi = kokkosAssembly().getGradPhi(sid, elem_type, fe_type);
553 67874072 : auto jacobian = kokkosAssembly().getJacobian(info, qp);
554 :
555 67874072 : Real value = 0;
556 67874072 : Real3 grad = 0;
557 :
558 371598020 : for (unsigned int i = 0; i < num_dofs; ++i)
559 : {
560 303723948 : auto vector = getVectorDofValue(getElemLocalDofIndex(elem, i, var), tag);
561 :
562 303723948 : value += vector * phi(i, qp);
563 303723948 : grad += vector * (jacobian * grad_phi(i, qp));
564 : }
565 :
566 67874072 : getVectorQpValue(info, qp_offset + qp, var, tag) = value;
567 67874072 : getVectorQpGrad(info, qp_offset + qp, var, tag) = grad;
568 : }
569 :
570 : } // namespace Moose::Kokkos
|
