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 1482 : System::System(SystemBase & system)
26 1126 : : MeshHolder(*system.mesh().getKokkosMesh()),
27 1126 : AssemblyHolder(system.feProblem().kokkosAssembly()),
28 1126 : _system(system),
29 1126 : _mesh(system.mesh()),
30 1126 : _dof_map(system.dofMap()),
31 1126 : _comm(system.dofMap().comm()),
32 1126 : _num_vars(system.system().n_vars()),
33 1126 : _num_local_dofs(system.system().n_local_dofs()),
34 4504 : _num_ghost_dofs(system.dofMap().get_send_list().size())
35 : {
36 1482 : setupVariables();
37 1482 : setupDofs();
38 :
39 1482 : if (dynamic_cast<NonlinearSystemBase *>(&_system))
40 : {
41 741 : setupSparsity();
42 741 : checkNodalBCs();
43 : }
44 :
45 1482 : _qp_solutions.create(MAX_TAG);
46 1482 : _qp_solutions_grad.create(MAX_TAG);
47 1482 : }
48 :
49 : void
50 1482 : System::setupVariables()
51 : {
52 1482 : auto & sys = _system.system();
53 :
54 1482 : _var_fe_types.create(_num_vars);
55 1482 : _var_subdomain_active.create(_num_vars, _mesh.meshSubdomains().size());
56 :
57 2906 : for (unsigned int var = 0; var < _num_vars; ++var)
58 : {
59 1424 : _var_fe_types[var] = kokkosAssembly().getFETypeID(sys.variable_type(var));
60 :
61 2910 : for (auto subdomain : _mesh.meshSubdomains())
62 1486 : _var_subdomain_active(var, kokkosMesh().getContiguousSubdomainID(subdomain)) =
63 1124 : sys.variable(var).active_on_subdomain(subdomain);
64 : }
65 :
66 1482 : _var_fe_types.copyToDevice();
67 1482 : _var_subdomain_active.copyToDevice();
68 :
69 : // Set coupling
70 :
71 1482 : auto nl_system = dynamic_cast<NonlinearSystemBase *>(&_system);
72 :
73 1482 : if (nl_system)
74 : {
75 741 : _coupling.create(_num_vars);
76 :
77 741 : std::map<unsigned int, std::vector<unsigned int>> coupling;
78 :
79 741 : auto & ce = _system.feProblem().couplingEntries(0, nl_system->number());
80 :
81 2076 : for (const auto & [ivar, jvar] : ce)
82 1335 : if (ivar->number() != jvar->number())
83 378 : coupling[ivar->number()].push_back(jvar->number());
84 :
85 1698 : for (unsigned int var = 0; var < _num_vars; ++var)
86 957 : _coupling[var] = coupling[var];
87 :
88 741 : _coupling.copyToDevice();
89 741 : }
90 1482 : }
91 :
92 : void
93 1482 : System::setupDofs()
94 : {
95 1482 : auto & sys = _system.system();
96 :
97 1482 : auto num_elems = kokkosMesh().getContiguousElementMap().size();
98 1482 : auto num_nodes = kokkosMesh().getContiguousNodeMap().size();
99 :
100 1482 : _residual_tag_active.create(MAX_TAG);
101 1482 : _residual_tag_active = false;
102 :
103 1482 : _matrix_tag_active.create(MAX_TAG);
104 1482 : _matrix_tag_active = false;
105 :
106 1482 : _vectors.create(MAX_TAG);
107 1482 : _matrices.create(MAX_TAG);
108 :
109 1482 : _local_elem_dof_index.create(_num_vars);
110 1482 : _local_node_dof_index.create(_num_vars);
111 1482 : _local_to_global_dof_index.create(_num_local_dofs + _num_ghost_dofs);
112 :
113 1482 : _max_dofs_per_elem.create(_num_vars);
114 1482 : _max_dofs_per_elem = 0;
115 :
116 1482 : auto solution = dynamic_cast<PetscVector<Number> *>(sys.current_local_solution.get());
117 :
118 1482 : std::vector<dof_id_type> dof_indices;
119 :
120 2906 : for (unsigned int var = 0; var < _num_vars; ++var)
121 : {
122 73976 : for (auto & [elem, id] : kokkosMesh().getContiguousElementMap())
123 : {
124 72552 : _dof_map.dof_indices(elem, dof_indices, var);
125 :
126 124355 : _max_dofs_per_elem[var] =
127 51803 : std::max(_max_dofs_per_elem[var], static_cast<unsigned int>(dof_indices.size()));
128 : }
129 :
130 1424 : _local_elem_dof_index[var].create(num_elems, _max_dofs_per_elem[var]);
131 1424 : _local_elem_dof_index[var] = libMesh::DofObject::invalid_id;
132 :
133 1424 : _local_node_dof_index[var].create(num_nodes);
134 1424 : _local_node_dof_index[var] = libMesh::DofObject::invalid_id;
135 :
136 73976 : for (auto & [elem, id] : kokkosMesh().getContiguousElementMap())
137 : {
138 72552 : _dof_map.dof_indices(elem, dof_indices, var);
139 :
140 349830 : for (unsigned int dof = 0; dof < dof_indices.size(); ++dof)
141 : {
142 277278 : _local_elem_dof_index[var](id, dof) = solution->map_global_to_local_index(dof_indices[dof]);
143 277278 : _local_to_global_dof_index[_local_elem_dof_index[var](id, dof)] = dof_indices[dof];
144 : }
145 : }
146 :
147 106418 : for (auto & [node, id] : kokkosMesh().getContiguousNodeMap())
148 104994 : if (node->processor_id() == _comm.rank())
149 : {
150 102994 : _dof_map.dof_indices(node, dof_indices, var);
151 :
152 102994 : if (dof_indices.size())
153 : {
154 80815 : 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 80815 : _local_node_dof_index[var][id] = solution->map_global_to_local_index(dof_indices[0]);
160 : }
161 : }
162 : }
163 :
164 1482 : _local_elem_dof_index.copyToDeviceNested();
165 1482 : _local_node_dof_index.copyToDeviceNested();
166 1482 : _local_to_global_dof_index.copyToDevice();
167 :
168 1482 : _max_dofs_per_elem.copyToDevice();
169 :
170 : // Setup DOF communication maps
171 :
172 1482 : auto num_procs = _comm.size();
173 :
174 2964 : std::vector<std::vector<dof_id_type>> send_list(num_procs);
175 1482 : std::vector<std::vector<dof_id_type>> recv_list(num_procs);
176 :
177 3115 : for (auto dof : _dof_map.get_send_list())
178 1633 : recv_list[_dof_map.dof_owner(dof)].push_back(dof);
179 :
180 3476 : for (processor_id_type proc = 0; proc < num_procs; proc++)
181 1994 : _comm.scatter(recv_list, send_list[proc], proc);
182 :
183 1482 : _local_comm_list.create(num_procs);
184 1482 : _ghost_comm_list.create(num_procs);
185 :
186 3476 : for (processor_id_type proc = 0; proc < num_procs; proc++)
187 : {
188 1994 : _local_comm_list[proc].create(send_list[proc].size());
189 1994 : _ghost_comm_list[proc].create(recv_list[proc].size());
190 :
191 3627 : for (dof_id_type i = 0; i < send_list[proc].size(); ++i)
192 1633 : _local_comm_list[proc][i] = solution->map_global_to_local_index(send_list[proc][i]);
193 :
194 3627 : for (dof_id_type i = 0; i < recv_list[proc].size(); ++i)
195 1633 : _ghost_comm_list[proc][i] = solution->map_global_to_local_index(recv_list[proc][i]);
196 : }
197 :
198 1482 : _local_comm_list.copyToDeviceNested();
199 1482 : _ghost_comm_list.copyToDeviceNested();
200 1482 : }
201 :
202 : void
203 741 : System::setupSparsity()
204 : {
205 741 : auto pattern = _dof_map.get_sparsity_pattern();
206 :
207 741 : std::vector<PetscInt> col_idx;
208 741 : std::vector<PetscInt> row_idx;
209 1482 : std::vector<PetscInt> row_ptr = {0};
210 :
211 61158 : for (dof_id_type r = _dof_map.first_dof(); r < _dof_map.end_dof(); ++r)
212 : {
213 60417 : auto & cols = pattern->get_sparsity_pattern().at(r - _dof_map.first_dof());
214 :
215 643010 : for (auto col : cols)
216 : {
217 582593 : col_idx.push_back(col);
218 582593 : row_idx.push_back(r);
219 : }
220 :
221 60417 : row_ptr.push_back(col_idx.size());
222 : }
223 :
224 741 : auto num_procs = _comm.size();
225 :
226 2964 : std::vector<std::vector<PetscInt>> send_cols(num_procs), send_num_cols(num_procs);
227 2223 : std::vector<std::vector<PetscInt>> recv_cols(num_procs), recv_num_cols(num_procs);
228 :
229 1738 : for (processor_id_type proc = 0; proc < num_procs; proc++)
230 2200 : for (auto r : _local_comm_list[proc])
231 : {
232 15049 : for (PetscInt c = row_ptr[r]; c < row_ptr[r + 1]; ++c)
233 13846 : send_cols[proc].push_back(col_idx[c]);
234 :
235 1203 : send_num_cols[proc].push_back(row_ptr[r + 1] - row_ptr[r]);
236 : }
237 :
238 1738 : for (processor_id_type proc = 0; proc < num_procs; proc++)
239 : {
240 997 : _comm.scatter(send_cols, recv_cols[proc], proc);
241 997 : _comm.scatter(send_num_cols, recv_num_cols[proc], proc);
242 : }
243 :
244 2223 : std::vector<PetscInt> col(num_procs), row(num_procs);
245 :
246 1944 : for (auto r : _dof_map.get_send_list())
247 : {
248 1203 : auto proc = _dof_map.dof_owner(r);
249 1203 : auto n_cols = recv_num_cols[proc][row[proc]++];
250 :
251 15049 : for (PetscInt c = 0; c < n_cols; ++c)
252 : {
253 13846 : col_idx.push_back(recv_cols[proc][col[proc]++]);
254 13846 : row_idx.push_back(r);
255 : }
256 :
257 1203 : row_ptr.push_back(col_idx.size());
258 : }
259 :
260 741 : _sparsity.col_idx = col_idx;
261 741 : _sparsity.row_idx = row_idx;
262 741 : _sparsity.row_ptr = row_ptr;
263 741 : }
264 :
265 : void
266 741 : System::checkNodalBCs()
267 : {
268 741 : auto & nl_system = static_cast<NonlinearSystemBase &>(_system);
269 :
270 741 : _nbc_residual_tag_dof.create(MAX_TAG);
271 741 : _nbc_matrix_tag_dof.create(MAX_TAG);
272 :
273 1966 : for (auto bc : nl_system.getKokkosNodalBCWarehouse().getActiveObjects())
274 : {
275 1225 : auto nbc = static_cast<NodalBCBase *>(bc.get());
276 :
277 1225 : std::set<TagID> extra_vector_tags;
278 1225 : std::set<TagID> extra_matrix_tags;
279 :
280 3675 : 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 3675 : 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 1225 : getNodalBCDofs(nbc, _nbc_dof);
289 :
290 1321 : for (auto tag : extra_vector_tags)
291 96 : getNodalBCDofs(nbc, _nbc_residual_tag_dof[tag]);
292 :
293 1345 : for (auto tag : extra_matrix_tags)
294 120 : getNodalBCDofs(nbc, _nbc_matrix_tag_dof[tag]);
295 1225 : }
296 :
297 741 : _nbc_dof.copyToDevice();
298 741 : _nbc_residual_tag_dof.copyToDeviceNested();
299 741 : _nbc_matrix_tag_dof.copyToDeviceNested();
300 741 : }
301 :
302 : void
303 1441 : System::getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs)
304 : {
305 1441 : auto var = nbc->variable().number();
306 :
307 1441 : if (!dofs.isAlloc())
308 : {
309 837 : dofs.create(_num_local_dofs + _num_ghost_dofs);
310 837 : dofs = false;
311 : }
312 :
313 9914 : for (auto node : nbc->getContiguousNodes())
314 9914 : 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 1441 : auto num_procs = _comm.size();
320 :
321 4323 : std::vector<std::vector<char>> send(num_procs), recv(num_procs);
322 :
323 3382 : for (processor_id_type proc = 0; proc < num_procs; proc++)
324 4294 : for (auto dof : _local_comm_list[proc])
325 2353 : send[proc].push_back(dofs[dof]);
326 :
327 3382 : for (processor_id_type proc = 0; proc < num_procs; proc++)
328 1941 : _comm.scatter(send, recv[proc], proc);
329 :
330 3382 : for (processor_id_type proc = 0; proc < num_procs; proc++)
331 4294 : for (dof_id_type i = 0; i < _ghost_comm_list[proc].size(); ++i)
332 2353 : dofs[_ghost_comm_list[proc][i]] = recv[proc][i];
333 1441 : }
334 :
335 : void
336 154788 : System::sync(const MemcpyKind dir)
337 : {
338 154788 : if (dir == MemcpyKind::HOST_TO_DEVICE)
339 : {
340 194202 : for (auto tag : _active_solution_tags)
341 : {
342 116808 : auto & vector = _system.getVector(tag);
343 :
344 116808 : _vectors[tag].create(vector, *this, false);
345 116808 : _vectors[tag].copyToDevice();
346 : }
347 :
348 157702 : for (auto tag : _active_residual_tags)
349 : {
350 80308 : auto & vector = _system.getVector(tag);
351 :
352 80308 : _vectors[tag].create(vector, *this, true);
353 80308 : _vectors[tag] = 0;
354 : }
355 :
356 82073 : for (auto tag : _active_matrix_tags)
357 : {
358 4679 : auto & matrix = _system.getMatrix(tag);
359 :
360 4679 : _matrices[tag].create(matrix, *this);
361 4679 : _matrices[tag] = 0;
362 : }
363 :
364 77394 : _vectors.copyToDevice();
365 77394 : _matrices.copyToDevice();
366 : }
367 77394 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
368 : {
369 194202 : for (auto tag : _active_solution_tags)
370 116808 : _vectors[tag].restore();
371 :
372 157702 : for (auto tag : _active_residual_tags)
373 80308 : _vectors[tag].close();
374 :
375 82073 : for (auto tag : _active_matrix_tags)
376 4679 : _matrices[tag].close();
377 : }
378 154788 : }
379 :
380 : void
381 21944 : System::sync(const std::set<TagID> & tags, const MemcpyKind dir)
382 : {
383 21944 : if (dir == MemcpyKind::HOST_TO_DEVICE)
384 : {
385 21964 : for (auto tag : tags)
386 : {
387 10992 : if (!_system.hasVector(tag))
388 1889 : continue;
389 :
390 9103 : auto & vector = _system.getVector(tag);
391 :
392 9103 : _vectors[tag].create(vector, *this, false);
393 9103 : _vectors[tag].copyToDevice();
394 : }
395 :
396 10972 : _vectors.copyToDevice();
397 : }
398 10972 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
399 : {
400 21964 : for (auto tag : tags)
401 : {
402 10992 : if (!_system.hasVector(tag))
403 1889 : continue;
404 :
405 9103 : _vectors[tag].copyToHost();
406 9103 : _vectors[tag].restore();
407 : }
408 : }
409 21944 : }
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 14388 : System::sync(const TagID tag, const MemcpyKind dir)
419 : {
420 28776 : sync(std::set<TagID>{tag}, dir);
421 14388 : }
422 :
423 : void
424 77394 : System::setActiveVariables(const std::set<MooseVariableFieldBase *> & vars)
425 : {
426 77394 : std::set<unsigned int> active_variables;
427 :
428 180740 : for (auto var : vars)
429 103346 : if (var->sys().number() == _system.number())
430 103346 : for (unsigned int i = 0; i < var->count(); ++i)
431 51673 : active_variables.insert(var->number() + i);
432 :
433 77394 : _active_variables = active_variables;
434 77394 : }
435 :
436 : void
437 77394 : System::setActiveSolutionTags(const std::set<TagID> & tags)
438 : {
439 77394 : if (!_active_variables.size())
440 34489 : return;
441 :
442 42905 : std::set<TagID> active_solution_tags;
443 :
444 237772 : for (auto tag : tags)
445 194867 : if (_system.hasVector(tag))
446 116808 : active_solution_tags.insert(tag);
447 :
448 42905 : _active_solution_tags = active_solution_tags;
449 42905 : }
450 :
451 : void
452 26756 : System::setActiveResidualTags(const std::set<TagID> & tags)
453 : {
454 26756 : std::set<TagID> active_residual_tags;
455 :
456 107064 : for (auto tag : tags)
457 80308 : if (_system.hasVector(tag))
458 : {
459 80308 : active_residual_tags.insert(tag);
460 80308 : _residual_tag_active[tag] = true;
461 : }
462 :
463 26756 : _active_residual_tags = active_residual_tags;
464 :
465 26756 : _residual_tag_active.copyToDevice();
466 26756 : }
467 :
468 : void
469 4437 : System::setActiveMatrixTags(const std::set<TagID> & tags)
470 : {
471 4437 : std::set<TagID> active_matrix_tags;
472 :
473 13300 : for (auto tag : tags)
474 8863 : if (_system.hasMatrix(tag))
475 : {
476 4679 : active_matrix_tags.insert(tag);
477 4679 : _matrix_tag_active[tag] = true;
478 : }
479 :
480 4437 : _active_matrix_tags = active_matrix_tags;
481 :
482 4437 : _matrix_tag_active.copyToDevice();
483 4437 : }
484 :
485 : void
486 77394 : System::reinit()
487 : {
488 194202 : for (auto tag : _active_solution_tags)
489 : {
490 116808 : if (!_qp_solutions[tag].isAlloc())
491 1831 : _qp_solutions[tag].create(_mesh.meshSubdomains().size(), _num_vars);
492 :
493 116808 : if (!_qp_solutions_grad[tag].isAlloc())
494 1831 : _qp_solutions_grad[tag].create(_mesh.meshSubdomains().size(), _num_vars);
495 :
496 238312 : for (auto subdomain : _mesh.meshSubdomains())
497 : {
498 121504 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
499 :
500 266265 : for (auto var : _active_variables)
501 : {
502 144761 : if (!_var_subdomain_active(var, sid))
503 204 : continue;
504 :
505 144557 : if (!_qp_solutions[tag](sid, var).isAlloc())
506 2197 : _qp_solutions[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
507 :
508 144557 : if (!_qp_solutions_grad[tag](sid, var).isAlloc())
509 2197 : _qp_solutions_grad[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
510 : }
511 : }
512 :
513 116808 : _qp_solutions[tag].copyToDevice();
514 116808 : _qp_solutions_grad[tag].copyToDevice();
515 : }
516 :
517 77394 : _qp_solutions.copyToDevice();
518 77394 : _qp_solutions_grad.copyToDevice();
519 :
520 77394 : dof_id_type num_elems = kokkosMesh().getContiguousElementMap().size();
521 :
522 218380 : _thread.resize({kokkosAssembly().getMaxQpsPerElem(),
523 : num_elems,
524 63592 : _active_variables.size(),
525 63592 : _active_solution_tags.size()});
526 :
527 77394 : ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(0, _thread.size());
528 77394 : ::Kokkos::parallel_for(policy, *this);
529 77394 : ::Kokkos::fence();
530 77394 : }
531 :
532 : KOKKOS_FUNCTION void
533 42835772 : System::operator()(const ThreadID tid) const
534 : {
535 42835772 : auto qp = _thread(tid, 0);
536 42835772 : auto elem = _thread(tid, 1);
537 42835772 : auto var = _active_variables(_thread(tid, 2));
538 42835772 : auto tag = _active_solution_tags(_thread(tid, 3));
539 :
540 42835772 : auto info = kokkosMesh().getElementInfo(elem);
541 42835772 : auto sid = info.subdomain;
542 42835772 : auto elem_type = info.type;
543 :
544 42835772 : if (!_var_subdomain_active(var, sid))
545 1936 : return;
546 :
547 42833836 : auto fe_type = _var_fe_types[var];
548 42833836 : auto num_dofs = kokkosAssembly().getNumDofs(elem_type, fe_type);
549 42833836 : auto num_qps = kokkosAssembly().getNumQps(info);
550 42833836 : auto qp_offset = kokkosAssembly().getQpOffset(info);
551 :
552 42833836 : if (qp >= num_qps)
553 0 : return;
554 :
555 42833836 : auto & phi = kokkosAssembly().getPhi(sid, elem_type, fe_type);
556 42833836 : auto & grad_phi = kokkosAssembly().getGradPhi(sid, elem_type, fe_type);
557 42833836 : auto jacobian = kokkosAssembly().getJacobian(info, qp);
558 :
559 42833836 : Real value = 0;
560 42833836 : Real3 grad = 0;
561 :
562 248688012 : for (unsigned int i = 0; i < num_dofs; ++i)
563 : {
564 205854176 : auto vector = getVectorDofValue(getElemLocalDofIndex(elem, i, var), tag);
565 :
566 205854176 : value += vector * phi(i, qp);
567 205854176 : grad += vector * (jacobian * grad_phi(i, qp));
568 : }
569 :
570 42833836 : getVectorQpValue(info, qp_offset + qp, var, tag) = value;
571 42833836 : getVectorQpGrad(info, qp_offset + qp, var, tag) = grad;
572 : }
573 :
574 : } // namespace Kokkos
575 : } // namespace Moose
|
