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 "KokkosKernelBase.h"
11 : #include "KokkosNodalKernelBase.h"
12 : #include "KokkosNodalBCBase.h"
13 : #include "KokkosIntegratedBCBase.h"
14 :
15 : #include "MaterialBase.h"
16 : #include "NonlinearSystemBase.h"
17 : #include "MooseVariableFieldBase.h"
18 : #include "FEProblemBase.h"
19 :
20 : void
21 1335 : NonlinearSystemBase::addKokkosKernel(const std::string & kernel_name,
22 : const std::string & name,
23 : InputParameters & parameters)
24 : {
25 : // Create the kernel object via the factory and add to warehouse
26 314 : std::shared_ptr<Moose::Kokkos::KernelBase> kernel =
27 1021 : _factory.create<Moose::Kokkos::KernelBase>(kernel_name, name, parameters);
28 1335 : _kokkos_kernels.addObject(kernel, 0);
29 1335 : postAddResidualObject(*kernel);
30 1335 : }
31 :
32 : void
33 272 : NonlinearSystemBase::addKokkosNodalKernel(const std::string & kernel_name,
34 : const std::string & name,
35 : InputParameters & parameters)
36 : {
37 : // Create the kernel object via the factory and add to warehouse
38 57 : std::shared_ptr<Moose::Kokkos::NodalKernelBase> kernel =
39 215 : _factory.create<Moose::Kokkos::NodalKernelBase>(kernel_name, name, parameters);
40 272 : _kokkos_nodal_kernels.addObject(kernel, 0);
41 272 : postAddResidualObject(*kernel);
42 272 : }
43 :
44 : void
45 1426 : NonlinearSystemBase::addKokkosBoundaryCondition(const std::string & bc_name,
46 : const std::string & name,
47 : InputParameters & parameters)
48 : {
49 : // Create the object
50 339 : std::shared_ptr<Moose::Kokkos::BoundaryCondition> bc =
51 1087 : _factory.create<Moose::Kokkos::BoundaryCondition>(bc_name, name, parameters, 0);
52 1426 : postAddResidualObject(*bc);
53 :
54 : // Active BoundaryIDs for the object
55 1426 : const std::set<BoundaryID> & boundary_ids = bc->boundaryIDs();
56 1426 : auto bc_var = dynamic_cast<const MooseVariableFieldBase *>(&bc->variable());
57 1426 : _vars[0].addBoundaryVar(boundary_ids, bc_var);
58 :
59 : // Cast to the various types of BCs
60 339 : std::shared_ptr<Moose::Kokkos::NodalBCBase> nbc =
61 1087 : std::dynamic_pointer_cast<Moose::Kokkos::NodalBCBase>(bc);
62 339 : std::shared_ptr<Moose::Kokkos::IntegratedBCBase> ibc =
63 1087 : std::dynamic_pointer_cast<Moose::Kokkos::IntegratedBCBase>(bc);
64 :
65 : // NodalBCBase
66 1426 : if (nbc)
67 : {
68 1225 : if (!nbc->variable().isNodal())
69 0 : mooseError("Trying to use nodal boundary condition '",
70 0 : nbc->name(),
71 : "' on a non-nodal variable '",
72 0 : nbc->variable().name(),
73 : "'.");
74 :
75 1225 : _kokkos_nodal_bcs.addObject(nbc);
76 1225 : _vars[0].addBoundaryVars(boundary_ids, nbc->getCoupledVars());
77 :
78 : // DirichletBCs that are preset
79 1225 : if (nbc->preset())
80 1047 : _kokkos_preset_nodal_bcs.addObject(nbc);
81 : }
82 :
83 : // IntegratedBCBase
84 201 : else if (ibc)
85 : {
86 201 : _kokkos_integrated_bcs.addObject(ibc);
87 201 : _vars[0].addBoundaryVars(boundary_ids, ibc->getCoupledVars());
88 : }
89 :
90 : else
91 0 : mooseError("Unknown Kokkos BoundaryCondition type for object named ", bc->name());
92 1426 : }
93 :
94 : void
95 1889 : NonlinearSystemBase::setKokkosInitialSolution()
96 : {
97 : // The preset Dirichlet BCs can have coupled variables from other systems, so we grab all systems
98 1889 : auto & systems = _fe_problem.getKokkosSystems();
99 :
100 1889 : auto tag = _fe_problem.addVectorTag("parallel_solution", Moose::VECTOR_TAG_SOLUTION);
101 :
102 1889 : associateVectorToTag(solution(), tag);
103 :
104 3778 : std::set<TagID> needed_fe_var_vector_tags = {tag};
105 :
106 8212 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
107 6323 : _kokkos_preset_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
108 : boundary, needed_fe_var_vector_tags);
109 :
110 5667 : for (auto & system : systems)
111 3778 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
112 :
113 1889 : systems.copyToDevice();
114 :
115 5251 : for (auto nbc : _kokkos_preset_nodal_bcs.getActiveObjects())
116 3362 : std::static_pointer_cast<Moose::Kokkos::NodalBCBase>(nbc)->presetSolution(tag);
117 :
118 1889 : Kokkos::fence();
119 :
120 5667 : for (auto & system : systems)
121 3778 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
122 :
123 1889 : disassociateVectorFromTag(solution(), tag);
124 1889 : }
125 :
126 : void
127 26756 : NonlinearSystemBase::computeKokkosResidual(const std::set<TagID> & tags)
128 : {
129 80268 : TIME_SECTION("computeKokkosResidual", 1);
130 :
131 26756 : auto & systems = _fe_problem.getKokkosSystems();
132 :
133 26756 : systems[number()].setActiveResidualTags(tags);
134 :
135 : // Resolve dependencies
136 :
137 26756 : std::set<MooseVariableFieldBase *> needed_moose_vars;
138 26756 : std::set<TagID> needed_fe_var_vector_tags;
139 26756 : std::unordered_set<unsigned int> needed_mat_props;
140 :
141 150157 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
142 123401 : needed_fe_var_vector_tags.insert(tag._id);
143 :
144 55140 : for (auto block : _fe_problem.mesh().meshSubdomains())
145 : {
146 28384 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
147 : needed_moose_vars);
148 28384 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
149 : block, needed_fe_var_vector_tags);
150 :
151 28384 : _kokkos_kernels.updateBlockVariableDependency(block, needed_moose_vars);
152 28384 : _kokkos_kernels.updateBlockFEVariableCoupledVectorTagDependency(block,
153 : needed_fe_var_vector_tags);
154 28384 : _kokkos_kernels.updateBlockMatPropDependency(block, needed_mat_props);
155 : }
156 :
157 124817 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
158 : {
159 98061 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryVariableDependency(boundary,
160 : needed_moose_vars);
161 98061 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryFEVariableCoupledVectorTagDependency(
162 : boundary, needed_fe_var_vector_tags);
163 :
164 98061 : _kokkos_integrated_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
165 98061 : _kokkos_integrated_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
166 : boundary, needed_fe_var_vector_tags);
167 98061 : _kokkos_integrated_bcs.updateBoundaryMatPropDependency(boundary, needed_mat_props);
168 :
169 98061 : _kokkos_nodal_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
170 98061 : _kokkos_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(boundary,
171 : needed_fe_var_vector_tags);
172 : }
173 :
174 : // Copy data and cache variable values at element quadature points
175 :
176 80268 : for (auto & system : systems)
177 : {
178 53512 : system.setActiveVariables(needed_moose_vars);
179 53512 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
180 :
181 : {
182 160536 : TIME_SECTION("KokkosCopy", 1);
183 53512 : system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
184 53512 : }
185 : {
186 160536 : TIME_SECTION("KokkosReinit", 1);
187 53512 : system.reinit();
188 53512 : }
189 : }
190 :
191 26756 : systems.copyToDevice();
192 :
193 : {
194 80268 : TIME_SECTION("KokkosMaterial", 1);
195 :
196 : // Compute material properties
197 :
198 26756 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
199 26756 : _fe_problem.reinitKokkosMaterials();
200 26756 : }
201 :
202 : {
203 80268 : TIME_SECTION("KokkosKernel", 1);
204 :
205 : // Compute kernels
206 :
207 80303 : for (auto kernel : _kokkos_kernels.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
208 53547 : kernel->computeResidual();
209 :
210 30274 : for (auto nodal_kernel :
211 61068 : _kokkos_nodal_kernels.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
212 20340 : nodal_kernel->computeResidual();
213 :
214 35038 : for (auto ibc : _kokkos_integrated_bcs.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
215 8282 : ibc->computeResidual();
216 :
217 73535 : for (auto nbc : _kokkos_nodal_bcs.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
218 46779 : nbc->computeResidual();
219 26756 : }
220 :
221 : // Close and restore vectors
222 :
223 : {
224 80268 : TIME_SECTION("KokkosCopy", 1);
225 :
226 80268 : for (auto & system : systems)
227 53512 : system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
228 26756 : }
229 :
230 : // Clear
231 :
232 26756 : systems[number()].clearActiveResidualTags();
233 :
234 80268 : for (auto & system : systems)
235 : {
236 53512 : system.clearActiveVariables();
237 53512 : system.clearActiveSolutionTags();
238 : }
239 26756 : }
240 :
241 : void
242 4437 : NonlinearSystemBase::computeKokkosJacobian(const std::set<TagID> & tags)
243 : {
244 13311 : TIME_SECTION("computeKokkosJacobian", 1);
245 :
246 4437 : auto & systems = _fe_problem.getKokkosSystems();
247 :
248 4437 : systems[number()].setActiveMatrixTags(tags);
249 :
250 : // Resolve dependencies
251 :
252 4437 : std::set<MooseVariableFieldBase *> needed_moose_vars;
253 4437 : std::set<TagID> needed_fe_var_vector_tags;
254 4437 : std::unordered_set<unsigned int> needed_mat_props;
255 :
256 25301 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
257 20864 : needed_fe_var_vector_tags.insert(tag._id);
258 :
259 9136 : for (auto block : _fe_problem.mesh().meshSubdomains())
260 : {
261 4699 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
262 : needed_moose_vars);
263 4699 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
264 : block, needed_fe_var_vector_tags);
265 :
266 4699 : _kokkos_kernels.updateBlockVariableDependency(block, needed_moose_vars);
267 4699 : _kokkos_kernels.updateBlockFEVariableCoupledVectorTagDependency(block,
268 : needed_fe_var_vector_tags);
269 4699 : _kokkos_kernels.updateBlockMatPropDependency(block, needed_mat_props);
270 : }
271 :
272 18639 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
273 : {
274 14202 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryVariableDependency(boundary,
275 : needed_moose_vars);
276 14202 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryFEVariableCoupledVectorTagDependency(
277 : boundary, needed_fe_var_vector_tags);
278 :
279 14202 : _kokkos_integrated_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
280 14202 : _kokkos_integrated_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
281 : boundary, needed_fe_var_vector_tags);
282 14202 : _kokkos_integrated_bcs.updateBoundaryMatPropDependency(boundary, needed_mat_props);
283 :
284 14202 : _kokkos_nodal_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
285 14202 : _kokkos_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(boundary,
286 : needed_fe_var_vector_tags);
287 : }
288 :
289 : // Copy data and cache variable values at element quadature points
290 :
291 13311 : for (auto & system : systems)
292 : {
293 8874 : system.setActiveVariables(needed_moose_vars);
294 8874 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
295 :
296 : {
297 26622 : TIME_SECTION("KokkosPreallocate", 1);
298 8874 : system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
299 8874 : }
300 : {
301 26622 : TIME_SECTION("KokkosReinit", 1);
302 8874 : system.reinit();
303 8874 : }
304 : }
305 :
306 4437 : systems.copyToDevice();
307 :
308 : {
309 13311 : TIME_SECTION("KokkosMaterial", 1);
310 :
311 : // Compute material properties
312 :
313 4437 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
314 4437 : _fe_problem.reinitKokkosMaterials();
315 4437 : }
316 :
317 : {
318 13311 : TIME_SECTION("KokkosKernel", 1);
319 :
320 : // Compute kernels
321 :
322 14429 : for (auto kernel : _kokkos_kernels.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
323 9992 : kernel->computeJacobian();
324 :
325 5481 : for (auto nodal_kernel :
326 11532 : _kokkos_nodal_kernels.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
327 5424 : nodal_kernel->computeJacobian();
328 :
329 5165 : for (auto ibc : _kokkos_integrated_bcs.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
330 728 : ibc->computeJacobian();
331 :
332 12447 : for (auto nbc : _kokkos_nodal_bcs.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
333 8010 : nbc->computeJacobian();
334 4437 : }
335 :
336 : // Close and restore vectors and matrices
337 :
338 : {
339 13311 : TIME_SECTION("KokkosCopy", 1);
340 :
341 13311 : for (auto & system : systems)
342 8874 : system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
343 4437 : }
344 :
345 : // Clear
346 :
347 4437 : systems[number()].clearActiveMatrixTags();
348 :
349 13311 : for (auto & system : systems)
350 : {
351 8874 : system.clearActiveVariables();
352 8874 : system.clearActiveSolutionTags();
353 : }
354 4437 : }
|
