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 "KokkosFESystem.h"
11 : #include "KokkosNodalBCBase.h"
12 :
13 : #include "MooseMesh.h"
14 : #include "Assembly.h"
15 : #include "NonlinearSystemBase.h"
16 : #include "LinearSystem.h"
17 : #include "FEProblemBase.h"
18 :
19 : #include "libmesh/system.h"
20 :
21 : namespace Moose::Kokkos
22 : {
23 :
24 2988 : FESystem::FESystem(SystemBase & system)
25 1630 : : System(system), AssemblyHolder(system.feProblem().kokkosAssembly())
26 : {
27 2988 : setupVariables();
28 2988 : setupDofs();
29 2988 : setupCoupling();
30 :
31 2988 : if (dynamic_cast<NonlinearSystemBase *>(&_system))
32 2051 : setupNodalBCDofs();
33 :
34 2988 : _qp_solutions.create(MAX_TAG);
35 2988 : _qp_solutions_grad.create(MAX_TAG);
36 2988 : }
37 :
38 49 : FESystem::FESystem(System & base, SystemBase & system)
39 26 : : System(base), AssemblyHolder(system.feProblem().kokkosAssembly())
40 : {
41 49 : setupVariables();
42 49 : setupDofs();
43 :
44 49 : if (dynamic_cast<NonlinearSystemBase *>(&_system))
45 2 : setupNodalBCDofs();
46 :
47 49 : _qp_solutions.create(MAX_TAG);
48 49 : _qp_solutions_grad.create(MAX_TAG);
49 49 : }
50 :
51 : void
52 3037 : FESystem::setupVariables()
53 : {
54 3037 : auto & sys = _system.system();
55 :
56 3037 : _var_fe_types.create(_num_vars);
57 :
58 7102 : for (unsigned int var = 0; var < _num_vars; ++var)
59 4065 : _var_fe_types[var] = kokkosAssembly().getFETypeID(sys.variable_type(var));
60 :
61 3037 : _var_fe_types.copyToDevice();
62 3037 : }
63 :
64 : void
65 2988 : FESystem::setupCoupling()
66 : {
67 2988 : if (auto * const nl_system = dynamic_cast<NonlinearSystemBase *>(&_system))
68 : {
69 2051 : _coupling.create(_num_vars);
70 :
71 2051 : std::map<unsigned int, std::vector<unsigned int>> coupling;
72 :
73 2051 : auto & ce = _system.feProblem().couplingEntries(0, nl_system->number());
74 :
75 5280 : for (const auto & [ivar, jvar] : ce)
76 3229 : if (ivar->number() != jvar->number())
77 718 : coupling[ivar->number()].push_back(jvar->number());
78 :
79 4566 : for (const auto var : make_range(_num_vars))
80 2515 : _coupling[var] = coupling[var];
81 :
82 2051 : _coupling.copyToDevice();
83 2051 : }
84 2988 : }
85 :
86 : void
87 3037 : FESystem::setupDofs()
88 : {
89 3037 : auto & sys = _system.system();
90 :
91 3037 : auto num_nodes = kokkosMesh().getNumLocalNodes();
92 :
93 3037 : _local_node_dof_index.create(_num_vars);
94 :
95 1381 : auto * const solution =
96 1656 : libMesh::cast_ptr<PetscVector<Number> *>(sys.current_local_solution.get());
97 :
98 : #ifdef MOOSE_ENABLE_KOKKOS_GPU
99 : // Kokkos array thinks OpenMP clause is device code when using OpenMP backend
100 1381 : #pragma omp parallel for
101 : #endif
102 3858 : for (unsigned int var = 0; var < _num_vars; ++var)
103 : {
104 2202 : std::vector<dof_id_type> dof_indices;
105 :
106 2202 : _local_node_dof_index[var].create(num_nodes);
107 2202 : _local_node_dof_index[var] = libMesh::DofObject::invalid_id;
108 :
109 876993 : for (const auto node : kokkosMesh().getLocalNodes())
110 874791 : if (node->processor_id() == _comm.rank())
111 : {
112 867028 : const auto id = kokkosMesh().getContiguousNodeID(node);
113 :
114 867028 : _dof_map.dof_indices(node, dof_indices, var);
115 :
116 867028 : if (dof_indices.size())
117 : {
118 840061 : for (unsigned int i = 1; i < dof_indices.size(); ++i)
119 0 : if (dof_indices[i] != dof_indices[i - 1] + 1)
120 0 : mooseError("Kokkos system error: a variable has multiple DOFs on a node, but the DOF "
121 : "indices are discontiguous. This is not supported.");
122 :
123 840061 : _local_node_dof_index[var][id] = solution->map_global_to_local_index(dof_indices[0]);
124 : }
125 : }
126 2202 : }
127 :
128 3037 : _local_node_dof_index.copyToDeviceNested();
129 3037 : }
130 :
131 : void
132 2053 : FESystem::setupNodalBCDofs()
133 : {
134 2053 : auto & nl_system = static_cast<NonlinearSystemBase &>(_system);
135 :
136 2053 : _nbc_matrix_tag_dof.create(MAX_TAG);
137 :
138 5492 : for (auto bc : nl_system.getKokkosNodalBCWarehouse().getActiveObjects())
139 : {
140 3439 : auto nbc = static_cast<NodalBCBase *>(bc.get());
141 :
142 3439 : auto matrix_tags = nbc->getMatrixTags({});
143 :
144 10504 : for (auto tag : matrix_tags)
145 7065 : getNodalBCDofs(nbc, _nbc_matrix_tag_dof[tag]);
146 3439 : }
147 :
148 2053 : _nbc_matrix_tag_dof.copyToDeviceNested();
149 2053 : }
150 :
151 : void
152 7065 : FESystem::getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs)
153 : {
154 7065 : auto var = nbc->variable().number();
155 :
156 7065 : if (!dofs.isAlloc())
157 : {
158 3770 : dofs.create(_num_local_dofs + _num_ghost_dofs);
159 3770 : dofs = false;
160 : }
161 :
162 46293 : for (auto node : nbc->getContiguousNodes())
163 46293 : dofs[_local_node_dof_index[var][node]] = true;
164 :
165 : // Let remote processes know about ghost DOFs associated with nodal BCs not to contribute on them
166 :
167 7065 : auto num_procs = _comm.size();
168 :
169 21195 : std::vector<std::vector<char>> send(num_procs), recv(num_procs);
170 :
171 17578 : for (processor_id_type proc = 0; proc < num_procs; proc++)
172 26029 : for (auto dof : _local_comm_list[proc])
173 15516 : send[proc].push_back(dofs[dof]);
174 :
175 17578 : for (processor_id_type proc = 0; proc < num_procs; proc++)
176 10513 : _comm.scatter(send, recv[proc], proc);
177 :
178 17578 : for (processor_id_type proc = 0; proc < num_procs; proc++)
179 26029 : for (dof_id_type i = 0; i < _ghost_comm_list[proc].size(); ++i)
180 15516 : dofs[_ghost_comm_list[proc][i]] = recv[proc][i];
181 7065 : }
182 :
183 : void
184 137101 : FESystem::reinit()
185 : {
186 467345 : for (auto tag : _active_solution_tags)
187 : {
188 330244 : if (!_qp_solutions[tag].isAlloc())
189 4673 : _qp_solutions[tag].create(kokkosMesh().getNumSubdomains(), _num_vars);
190 :
191 330244 : if (!_qp_solutions_grad[tag].isAlloc())
192 4673 : _qp_solutions_grad[tag].create(kokkosMesh().getNumSubdomains(), _num_vars);
193 :
194 675630 : for (auto subdomain : _mesh.meshSubdomains())
195 : {
196 345386 : auto sid = kokkosMesh().getContiguousSubdomainID(subdomain);
197 :
198 722351 : for (auto var : _active_variables)
199 : {
200 376965 : if (!_var_subdomain_active(var, sid))
201 626 : continue;
202 :
203 376339 : if (!_qp_solutions[tag](sid, var).isAlloc())
204 6586 : _qp_solutions[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
205 :
206 376339 : if (!_qp_solutions_grad[tag](sid, var).isAlloc())
207 6586 : _qp_solutions_grad[tag](sid, var).createDevice(kokkosAssembly().getNumQps(sid));
208 : }
209 : }
210 :
211 330244 : _qp_solutions[tag].copyToDevice();
212 330244 : _qp_solutions_grad[tag].copyToDevice();
213 : }
214 :
215 137101 : _qp_solutions.copyToDevice();
216 137101 : _qp_solutions_grad.copyToDevice();
217 :
218 137101 : dof_id_type num_elems = kokkosMesh().getNumLocalElements();
219 :
220 137101 : _thread.resize(kokkosAssembly().getMaxQpsPerElem(),
221 : num_elems,
222 : _active_variables.size(),
223 : _active_solution_tags.size());
224 :
225 137101 : ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(0, _thread.size());
226 137101 : ::Kokkos::parallel_for(policy, *this);
227 137101 : ::Kokkos::fence();
228 137101 : }
229 :
230 : KOKKOS_FUNCTION void
231 67863764 : FESystem::operator()(const ThreadID tid) const
232 : {
233 67863764 : auto qp = _thread(tid, 0);
234 67863764 : auto elem = _thread(tid, 1);
235 67863764 : auto var = _active_variables(_thread(tid, 2));
236 67863764 : auto tag = _active_solution_tags(_thread(tid, 3));
237 :
238 67863764 : auto info = kokkosMesh().getElementInfo(elem);
239 67863764 : auto sid = info.subdomain;
240 67863764 : auto elem_type = info.type;
241 :
242 67863764 : if (!_var_subdomain_active(var, sid))
243 3872 : return;
244 :
245 67859892 : auto fe_type = _var_fe_types[var];
246 67859892 : auto num_dofs = kokkosAssembly().getNumDofs(elem_type, fe_type);
247 67859892 : auto num_qps = kokkosAssembly().getNumQps(info);
248 67859892 : auto qp_offset = kokkosAssembly().getQpOffset(info);
249 :
250 67859892 : if (qp >= num_qps)
251 0 : return;
252 :
253 67859892 : auto & phi = kokkosAssembly().getPhi(sid, elem_type, fe_type);
254 67859892 : auto & grad_phi = kokkosAssembly().getGradPhi(sid, elem_type, fe_type);
255 67859892 : auto jacobian = kokkosAssembly().getJacobian(info, qp);
256 :
257 67859892 : Real value = 0;
258 67859892 : Real3 grad = 0;
259 :
260 371507360 : for (unsigned int i = 0; i < num_dofs; ++i)
261 : {
262 303647468 : auto vector = getVectorDofValue(getElemLocalDofIndex(elem, i, var), tag);
263 :
264 303647468 : value += vector * phi(i, qp);
265 303647468 : grad += vector * (jacobian * grad_phi(i, qp));
266 : }
267 :
268 67859892 : getVectorQpValue(info, qp_offset + qp, var, tag) = value;
269 67859892 : getVectorQpGrad(info, qp_offset + qp, var, tag) = grad;
270 : }
271 :
272 : } // namespace Moose::Kokkos
|
