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 1386 : System::System(SystemBase & system)
26 1054 : : MeshHolder(*system.mesh().getKokkosMesh()),
27 1054 : AssemblyHolder(system.feProblem().kokkosAssembly()),
28 1054 : _system(system),
29 1054 : _mesh(system.mesh()),
30 1054 : _dof_map(system.dofMap()),
31 1054 : _comm(system.dofMap().comm()),
32 1054 : _num_vars(system.system().n_vars()),
33 1054 : _num_local_dofs(system.system().n_local_dofs()),
34 4216 : _num_ghost_dofs(system.dofMap().get_send_list().size())
35 : {
36 1386 : setupVariables();
37 1386 : setupDofs();
38 :
39 1386 : if (dynamic_cast<NonlinearSystemBase *>(&_system))
40 : {
41 693 : setupSparsity();
42 693 : checkNodalBCs();
43 : }
44 :
45 1386 : _qp_solutions.create(MAX_TAG);
46 1386 : _qp_solutions_grad.create(MAX_TAG);
47 1386 : }
48 :
49 : void
50 1386 : System::setupVariables()
51 : {
52 1386 : auto & sys = _system.system();
53 :
54 1386 : _var_fe_types.create(_num_vars);
55 1386 : _var_subdomain_active.create(_num_vars, _mesh.meshSubdomains().size());
56 :
57 2568 : for (unsigned int var = 0; var < _num_vars; ++var)
58 : {
59 1182 : _var_fe_types[var] = kokkosAssembly().getFETypeID(sys.variable_type(var));
60 :
61 2426 : for (auto subdomain : _mesh.meshSubdomains())
62 1244 : _var_subdomain_active(var, kokkosMesh().getContiguousSubdomainID(subdomain)) =
63 940 : sys.variable(var).active_on_subdomain(subdomain);
64 : }
65 :
66 1386 : _var_fe_types.copyToDevice();
67 1386 : _var_subdomain_active.copyToDevice();
68 :
69 : // Set coupling
70 :
71 1386 : auto nl_system = dynamic_cast<NonlinearSystemBase *>(&_system);
72 :
73 1386 : if (nl_system)
74 : {
75 693 : _coupling.create(_num_vars);
76 :
77 693 : std::map<unsigned int, std::vector<unsigned int>> coupling;
78 :
79 693 : auto & ce = _system.feProblem().couplingEntries(0, nl_system->number());
80 :
81 1980 : for (const auto & [ivar, jvar] : ce)
82 1287 : if (ivar->number() != jvar->number())
83 378 : coupling[ivar->number()].push_back(jvar->number());
84 :
85 1602 : for (unsigned int var = 0; var < _num_vars; ++var)
86 909 : _coupling[var] = coupling[var];
87 :
88 693 : _coupling.copyToDevice();
89 693 : }
90 1386 : }
91 :
92 : void
93 1386 : System::setupDofs()
94 : {
95 1386 : auto & sys = _system.system();
96 :
97 1386 : auto num_elems = kokkosMesh().getContiguousElementMap().size();
98 1386 : auto num_nodes = kokkosMesh().getContiguousNodeMap().size();
99 :
100 1386 : _residual_tag_active.create(MAX_TAG);
101 1386 : _residual_tag_active = false;
102 :
103 1386 : _matrix_tag_active.create(MAX_TAG);
104 1386 : _matrix_tag_active = false;
105 :
106 1386 : _vectors.create(MAX_TAG);
107 1386 : _matrices.create(MAX_TAG);
108 :
109 1386 : _local_elem_dof_index.create(_num_vars);
110 1386 : _local_node_dof_index.create(_num_vars);
111 1386 : _local_to_global_dof_index.create(_num_local_dofs + _num_ghost_dofs);
112 :
113 1386 : _max_dofs_per_elem.create(_num_vars);
114 1386 : _max_dofs_per_elem = 0;
115 :
116 1386 : auto solution = dynamic_cast<PetscVector<Number> *>(sys.current_local_solution.get());
117 :
118 1386 : std::vector<dof_id_type> dof_indices;
119 :
120 2568 : for (unsigned int var = 0; var < _num_vars; ++var)
121 : {
122 57584 : for (auto & [elem, id] : kokkosMesh().getContiguousElementMap())
123 : {
124 56402 : _dof_map.dof_indices(elem, dof_indices, var);
125 :
126 96640 : _max_dofs_per_elem[var] =
127 40238 : std::max(_max_dofs_per_elem[var], static_cast<unsigned int>(dof_indices.size()));
128 : }
129 :
130 1182 : _local_elem_dof_index[var].create(num_elems, _max_dofs_per_elem[var]);
131 1182 : _local_elem_dof_index[var] = libMesh::DofObject::invalid_id;
132 :
133 1182 : _local_node_dof_index[var].create(num_nodes);
134 1182 : _local_node_dof_index[var] = libMesh::DofObject::invalid_id;
135 :
136 57584 : for (auto & [elem, id] : kokkosMesh().getContiguousElementMap())
137 : {
138 56402 : _dof_map.dof_indices(elem, dof_indices, var);
139 :
140 267440 : for (unsigned int dof = 0; dof < dof_indices.size(); ++dof)
141 : {
142 211038 : _local_elem_dof_index[var](id, dof) = solution->map_global_to_local_index(dof_indices[dof]);
143 211038 : _local_to_global_dof_index[_local_elem_dof_index[var](id, dof)] = dof_indices[dof];
144 : }
145 : }
146 :
147 71750 : for (auto & [node, id] : kokkosMesh().getContiguousNodeMap())
148 70568 : if (node->processor_id() == _comm.rank())
149 : {
150 69142 : _dof_map.dof_indices(node, dof_indices, var);
151 :
152 69142 : if (dof_indices.size())
153 67815 : _local_node_dof_index[var][id] = solution->map_global_to_local_index(dof_indices[0]);
154 : }
155 : }
156 :
157 1386 : _local_elem_dof_index.copyToDeviceNested();
158 1386 : _local_node_dof_index.copyToDeviceNested();
159 1386 : _local_to_global_dof_index.copyToDevice();
160 :
161 1386 : _max_dofs_per_elem.copyToDevice();
162 :
163 : // Setup DOF communication maps
164 :
165 1386 : auto num_procs = _comm.size();
166 :
167 2772 : std::vector<std::vector<dof_id_type>> send_list(num_procs);
168 1386 : std::vector<std::vector<dof_id_type>> recv_list(num_procs);
169 :
170 2805 : for (auto dof : _dof_map.get_send_list())
171 1419 : recv_list[_dof_map.dof_owner(dof)].push_back(dof);
172 :
173 3252 : for (processor_id_type proc = 0; proc < num_procs; proc++)
174 1866 : _comm.scatter(recv_list, send_list[proc], proc);
175 :
176 1386 : _local_comm_list.create(num_procs);
177 1386 : _ghost_comm_list.create(num_procs);
178 :
179 3252 : for (processor_id_type proc = 0; proc < num_procs; proc++)
180 : {
181 1866 : _local_comm_list[proc].create(send_list[proc].size());
182 1866 : _ghost_comm_list[proc].create(recv_list[proc].size());
183 :
184 3285 : for (dof_id_type i = 0; i < send_list[proc].size(); ++i)
185 1419 : _local_comm_list[proc][i] = solution->map_global_to_local_index(send_list[proc][i]);
186 :
187 3285 : for (dof_id_type i = 0; i < recv_list[proc].size(); ++i)
188 1419 : _ghost_comm_list[proc][i] = solution->map_global_to_local_index(recv_list[proc][i]);
189 : }
190 :
191 1386 : _local_comm_list.copyToDeviceNested();
192 1386 : _ghost_comm_list.copyToDeviceNested();
193 1386 : }
194 :
195 : void
196 693 : System::setupSparsity()
197 : {
198 693 : auto pattern = _dof_map.get_sparsity_pattern();
199 :
200 693 : std::vector<PetscInt> col_idx;
201 693 : std::vector<PetscInt> row_idx;
202 1386 : std::vector<PetscInt> row_ptr = {0};
203 :
204 54240 : for (dof_id_type r = _dof_map.first_dof(); r < _dof_map.end_dof(); ++r)
205 : {
206 53547 : auto & cols = pattern->get_sparsity_pattern().at(r - _dof_map.first_dof());
207 :
208 551150 : for (auto col : cols)
209 : {
210 497603 : col_idx.push_back(col);
211 497603 : row_idx.push_back(r);
212 : }
213 :
214 53547 : row_ptr.push_back(col_idx.size());
215 : }
216 :
217 693 : auto num_procs = _comm.size();
218 :
219 2772 : std::vector<std::vector<PetscInt>> send_cols(num_procs), send_num_cols(num_procs);
220 2079 : std::vector<std::vector<PetscInt>> recv_cols(num_procs), recv_num_cols(num_procs);
221 :
222 1626 : for (processor_id_type proc = 0; proc < num_procs; proc++)
223 2042 : for (auto r : _local_comm_list[proc])
224 : {
225 13657 : for (PetscInt c = row_ptr[r]; c < row_ptr[r + 1]; ++c)
226 12548 : send_cols[proc].push_back(col_idx[c]);
227 :
228 1109 : send_num_cols[proc].push_back(row_ptr[r + 1] - row_ptr[r]);
229 : }
230 :
231 1626 : for (processor_id_type proc = 0; proc < num_procs; proc++)
232 : {
233 933 : _comm.scatter(send_cols, recv_cols[proc], proc);
234 933 : _comm.scatter(send_num_cols, recv_num_cols[proc], proc);
235 : }
236 :
237 2079 : std::vector<PetscInt> col(num_procs), row(num_procs);
238 :
239 1802 : for (auto r : _dof_map.get_send_list())
240 : {
241 1109 : auto proc = _dof_map.dof_owner(r);
242 1109 : auto n_cols = recv_num_cols[proc][row[proc]++];
243 :
244 13657 : for (PetscInt c = 0; c < n_cols; ++c)
245 : {
246 12548 : col_idx.push_back(recv_cols[proc][col[proc]++]);
247 12548 : row_idx.push_back(r);
248 : }
249 :
250 1109 : row_ptr.push_back(col_idx.size());
251 : }
252 :
253 693 : _sparsity.col_idx = col_idx;
254 693 : _sparsity.row_idx = row_idx;
255 693 : _sparsity.row_ptr = row_ptr;
256 693 : }
257 :
258 : void
259 693 : System::checkNodalBCs()
260 : {
261 693 : auto & nl_system = static_cast<NonlinearSystemBase &>(_system);
262 :
263 693 : _nbc_residual_tag_dof.create(MAX_TAG);
264 693 : _nbc_matrix_tag_dof.create(MAX_TAG);
265 :
266 1894 : for (auto bc : nl_system.getKokkosNodalBCWarehouse().getActiveObjects())
267 : {
268 1201 : auto nbc = static_cast<NodalBCBase *>(bc.get());
269 :
270 1201 : std::set<TagID> extra_vector_tags;
271 1201 : std::set<TagID> extra_matrix_tags;
272 :
273 3603 : if (nbc->isParamValid("extra_vector_tags"))
274 384 : for (auto tag : nbc->getParam<std::vector<TagName>>("extra_vector_tags"))
275 96 : extra_vector_tags.insert(_system.feProblem().getVectorTagID(tag));
276 :
277 3603 : if (nbc->isParamValid("extra_matrix_tags"))
278 372 : for (auto tag : nbc->getParam<std::vector<TagName>>("extra_matrix_tags"))
279 120 : extra_matrix_tags.insert(_system.feProblem().getMatrixTagID(tag));
280 :
281 1201 : getNodalBCDofs(nbc, _nbc_dof);
282 :
283 1297 : for (auto tag : extra_vector_tags)
284 96 : getNodalBCDofs(nbc, _nbc_residual_tag_dof[tag]);
285 :
286 1321 : for (auto tag : extra_matrix_tags)
287 120 : getNodalBCDofs(nbc, _nbc_matrix_tag_dof[tag]);
288 1201 : }
289 :
290 693 : _nbc_dof.copyToDevice();
291 693 : _nbc_residual_tag_dof.copyToDeviceNested();
292 693 : _nbc_matrix_tag_dof.copyToDeviceNested();
293 693 : }
294 :
295 : void
296 1417 : System::getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs)
297 : {
298 1417 : auto var = nbc->variable().number();
299 :
300 1417 : if (!dofs.isAlloc())
301 : {
302 825 : dofs.create(_num_local_dofs + _num_ghost_dofs);
303 825 : dofs = false;
304 : }
305 :
306 9670 : for (auto node : nbc->getContiguousNodes())
307 9670 : dofs[_local_node_dof_index[var][node]] = true;
308 :
309 : // Let remote processes know about ghost DOFs associated with nodal BCs not to compute
310 : // residual on them
311 :
312 1417 : auto num_procs = _comm.size();
313 :
314 4251 : std::vector<std::vector<char>> send(num_procs), recv(num_procs);
315 :
316 3326 : for (processor_id_type proc = 0; proc < num_procs; proc++)
317 4214 : for (auto dof : _local_comm_list[proc])
318 2305 : send[proc].push_back(dofs[dof]);
319 :
320 3326 : for (processor_id_type proc = 0; proc < num_procs; proc++)
321 1909 : _comm.scatter(send, recv[proc], proc);
322 :
323 3326 : for (processor_id_type proc = 0; proc < num_procs; proc++)
324 4214 : for (dof_id_type i = 0; i < _ghost_comm_list[proc].size(); ++i)
325 2305 : dofs[_ghost_comm_list[proc][i]] = recv[proc][i];
326 1417 : }
327 :
328 : void
329 127148 : System::sync(const MemcpyKind dir)
330 : {
331 127148 : if (dir == MemcpyKind::HOST_TO_DEVICE)
332 : {
333 154045 : for (auto tag : _active_variable_tags)
334 : {
335 90471 : auto & vector = _system.getVector(tag);
336 :
337 90471 : _vectors[tag].create(vector, *this, false);
338 90471 : _vectors[tag].copyToDevice();
339 : }
340 :
341 145127 : for (auto tag : _active_residual_tags)
342 : {
343 81553 : auto & vector = _system.getVector(tag);
344 :
345 81553 : _vectors[tag].create(vector, *this, true);
346 81553 : _vectors[tag] = 0;
347 : }
348 :
349 68228 : for (auto tag : _active_matrix_tags)
350 : {
351 4654 : auto & matrix = _system.getMatrix(tag);
352 :
353 4654 : _matrices[tag].create(matrix, *this);
354 4654 : _matrices[tag] = 0;
355 : }
356 :
357 63574 : _vectors.copyToDevice();
358 63574 : _matrices.copyToDevice();
359 : }
360 63574 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
361 : {
362 154045 : for (auto tag : _active_variable_tags)
363 90471 : _vectors[tag].restore();
364 :
365 145127 : for (auto tag : _active_residual_tags)
366 81553 : _vectors[tag].close();
367 :
368 68228 : for (auto tag : _active_matrix_tags)
369 4654 : _matrices[tag].close();
370 : }
371 127148 : }
372 :
373 : void
374 7556 : System::sync(const std::set<TagID> & tags, const MemcpyKind dir)
375 : {
376 7556 : if (dir == MemcpyKind::HOST_TO_DEVICE)
377 : {
378 7576 : for (auto tag : tags)
379 : {
380 3798 : if (!_system.hasVector(tag))
381 1889 : continue;
382 :
383 1909 : auto & vector = _system.getVector(tag);
384 :
385 1909 : _vectors[tag].create(vector, *this, false);
386 1909 : _vectors[tag].copyToDevice();
387 : }
388 :
389 3778 : _vectors.copyToDevice();
390 : }
391 3778 : else if (dir == MemcpyKind::DEVICE_TO_HOST)
392 : {
393 7576 : for (auto tag : tags)
394 : {
395 3798 : if (!_system.hasVector(tag))
396 1889 : continue;
397 :
398 1909 : _vectors[tag].copyToHost();
399 1909 : _vectors[tag].restore();
400 : }
401 : }
402 7556 : }
403 :
404 : void
405 63574 : System::setActiveVariables(const std::set<MooseVariableFieldBase *> & vars)
406 : {
407 63574 : std::set<unsigned int> active_variables;
408 :
409 137330 : for (auto var : vars)
410 73756 : if (var->sys().number() == _system.number())
411 73756 : for (unsigned int i = 0; i < var->count(); ++i)
412 36878 : active_variables.insert(var->number() + i);
413 :
414 63574 : _active_variables = active_variables;
415 63574 : }
416 :
417 : void
418 63574 : System::setActiveVariableTags(const std::set<TagID> & tags)
419 : {
420 63574 : if (!_active_variables.size())
421 30153 : return;
422 :
423 33421 : std::set<TagID> active_variable_tags;
424 :
425 186266 : for (auto tag : tags)
426 152845 : if (_system.hasVector(tag))
427 90471 : active_variable_tags.insert(tag);
428 :
429 33421 : _active_variable_tags = active_variable_tags;
430 33421 : }
431 :
432 : void
433 27177 : System::setActiveResidualTags(const std::set<TagID> & tags)
434 : {
435 27177 : std::set<TagID> active_residual_tags;
436 :
437 108730 : for (auto tag : tags)
438 81553 : if (_system.hasVector(tag))
439 : {
440 81553 : active_residual_tags.insert(tag);
441 81553 : _residual_tag_active[tag] = true;
442 : }
443 :
444 27177 : _active_residual_tags = active_residual_tags;
445 :
446 27177 : _residual_tag_active.copyToDevice();
447 27177 : }
448 :
449 : void
450 4410 : System::setActiveMatrixTags(const std::set<TagID> & tags)
451 : {
452 4410 : std::set<TagID> active_matrix_tags;
453 :
454 13250 : for (auto tag : tags)
455 8840 : if (_system.hasMatrix(tag))
456 : {
457 4654 : active_matrix_tags.insert(tag);
458 4654 : _matrix_tag_active[tag] = true;
459 : }
460 :
461 4410 : _active_matrix_tags = active_matrix_tags;
462 :
463 4410 : _matrix_tag_active.copyToDevice();
464 4410 : }
465 :
466 : void
467 63574 : System::reinit()
468 : {
469 154045 : for (auto tag : _active_variable_tags)
470 : {
471 90471 : if (!_qp_solutions[tag].isAlloc())
472 1415 : _qp_solutions[tag].create(_mesh.meshSubdomains().size(), _num_vars);
473 :
474 90471 : if (!_qp_solutions_grad[tag].isAlloc())
475 1415 : _qp_solutions_grad[tag].create(_mesh.meshSubdomains().size(), _num_vars);
476 :
477 185642 : for (auto subdomain : _mesh.meshSubdomains())
478 : {
479 95171 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
480 :
481 197795 : for (auto var : _active_variables)
482 : {
483 102624 : if (!_var_subdomain_active(var, sid))
484 204 : continue;
485 :
486 102420 : if (!_qp_solutions[tag](sid, var).isAlloc())
487 1641 : _qp_solutions[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
488 :
489 102420 : if (!_qp_solutions_grad[tag](sid, var).isAlloc())
490 1641 : _qp_solutions_grad[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
491 : }
492 : }
493 :
494 90471 : _qp_solutions[tag].copyToDevice();
495 90471 : _qp_solutions_grad[tag].copyToDevice();
496 : }
497 :
498 63574 : _qp_solutions.copyToDevice();
499 63574 : _qp_solutions_grad.copyToDevice();
500 :
501 63574 : dof_id_type num_elems = kokkosMesh().getContiguousElementMap().size();
502 :
503 178748 : _thread.resize({kokkosAssembly().getMaxQpsPerElem(),
504 : num_elems,
505 51600 : _active_variables.size(),
506 51600 : _active_variable_tags.size()});
507 :
508 63574 : ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(0, _thread.size());
509 63574 : ::Kokkos::parallel_for(policy, *this);
510 63574 : ::Kokkos::fence();
511 63574 : }
512 :
513 : KOKKOS_FUNCTION void
514 18230688 : System::operator()(const ThreadID tid) const
515 : {
516 18230688 : auto qp = _thread(tid, 0);
517 18230688 : auto elem = _thread(tid, 1);
518 18230688 : auto var = _active_variables(_thread(tid, 2));
519 18230688 : auto tag = _active_variable_tags(_thread(tid, 3));
520 :
521 18230688 : auto info = kokkosMesh().getElementInfo(elem);
522 18230688 : auto sid = info.subdomain;
523 18230688 : auto elem_type = info.type;
524 :
525 18230688 : if (!_var_subdomain_active(var, sid))
526 1936 : return;
527 :
528 18228752 : auto fe_type = _var_fe_types[var];
529 18228752 : auto num_dofs = kokkosAssembly().getNumDofs(elem_type, fe_type);
530 18228752 : auto num_qps = kokkosAssembly().getNumQps(info);
531 18228752 : auto qp_offset = kokkosAssembly().getQpOffset(info);
532 :
533 18228752 : if (qp >= num_qps)
534 0 : return;
535 :
536 18228752 : auto & phi = kokkosAssembly().getPhi(sid, elem_type, fe_type);
537 18228752 : auto & grad_phi = kokkosAssembly().getGradPhi(sid, elem_type, fe_type);
538 18228752 : auto jacobian = kokkosAssembly().getJacobian(info, qp);
539 :
540 18228752 : Real value = 0;
541 18228752 : Real3 grad = 0;
542 :
543 96004640 : for (unsigned int i = 0; i < num_dofs; ++i)
544 : {
545 77775888 : auto vector = getVectorDofValue(getElemLocalDofIndex(elem, i, var), tag);
546 :
547 77775888 : value += vector * phi(i, qp);
548 77775888 : grad += vector * (jacobian * grad_phi(i, qp));
549 : }
550 :
551 18228752 : getVectorQpValue(info, qp_offset + qp, var, tag) = value;
552 18228752 : getVectorQpGrad(info, qp_offset + qp, var, tag) = grad;
553 : }
554 :
555 : } // namespace Kokkos
556 : } // namespace Moose
|
