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 1469 : 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 320 : std::shared_ptr<Moose::Kokkos::KernelBase> kernel =
27 1149 : _factory.create<Moose::Kokkos::KernelBase>(kernel_name, name, parameters);
28 1467 : _kokkos_kernels.addObject(kernel, 0);
29 1467 : postAddResidualObject(*kernel);
30 1467 : }
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 1562 : NonlinearSystemBase::addKokkosBoundaryCondition(const std::string & bc_name,
46 : const std::string & name,
47 : InputParameters & parameters)
48 : {
49 : // Create the object
50 345 : std::shared_ptr<Moose::Kokkos::BoundaryCondition> bc =
51 1217 : _factory.create<Moose::Kokkos::BoundaryCondition>(bc_name, name, parameters, 0);
52 1562 : postAddResidualObject(*bc);
53 :
54 : // Active BoundaryIDs for the object
55 1562 : const std::set<BoundaryID> & boundary_ids = bc->boundaryIDs();
56 1562 : auto bc_var = dynamic_cast<const MooseVariableFieldBase *>(&bc->variable());
57 1562 : _vars[0].addBoundaryVar(boundary_ids, bc_var);
58 :
59 : // Cast to the various types of BCs
60 345 : std::shared_ptr<Moose::Kokkos::NodalBCBase> nbc =
61 1217 : std::dynamic_pointer_cast<Moose::Kokkos::NodalBCBase>(bc);
62 345 : std::shared_ptr<Moose::Kokkos::IntegratedBCBase> ibc =
63 1217 : std::dynamic_pointer_cast<Moose::Kokkos::IntegratedBCBase>(bc);
64 :
65 : // NodalBCBase
66 1562 : if (nbc)
67 : {
68 1293 : 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 1293 : _kokkos_nodal_bcs.addObject(nbc);
76 1293 : _vars[0].addBoundaryVars(boundary_ids, nbc->getCoupledVars());
77 :
78 : // DirichletBCs that are preset
79 1293 : if (nbc->preset())
80 1115 : _kokkos_preset_nodal_bcs.addObject(nbc);
81 : }
82 :
83 : // IntegratedBCBase
84 269 : else if (ibc)
85 : {
86 269 : _kokkos_integrated_bcs.addObject(ibc);
87 269 : _vars[0].addBoundaryVars(boundary_ids, ibc->getCoupledVars());
88 : }
89 :
90 : else
91 0 : mooseError("Unknown Kokkos BoundaryCondition type for object named ", bc->name());
92 1562 : }
93 :
94 : void
95 2476 : NonlinearSystemBase::setKokkosInitialSolution()
96 : {
97 : // The preset Dirichlet BCs can have coupled variables from other systems, so we grab all systems
98 2476 : auto & systems = _fe_problem.getKokkosSystems();
99 :
100 2476 : auto tag = _fe_problem.addVectorTag("parallel_solution", Moose::VECTOR_TAG_SOLUTION);
101 :
102 2476 : associateVectorToTag(solution(), tag);
103 :
104 4952 : std::set<TagID> needed_fe_var_vector_tags = {tag};
105 :
106 11147 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
107 8671 : _kokkos_preset_nodal_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
108 : boundary, needed_fe_var_vector_tags);
109 :
110 7428 : for (auto & system : systems)
111 4952 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
112 :
113 2476 : systems.copyToDevice();
114 :
115 6425 : for (auto nbc : _kokkos_preset_nodal_bcs.getActiveObjects())
116 3949 : std::static_pointer_cast<Moose::Kokkos::NodalBCBase>(nbc)->presetSolution(tag);
117 :
118 2476 : Kokkos::fence();
119 :
120 7428 : for (auto & system : systems)
121 4952 : system.sync(needed_fe_var_vector_tags, Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
122 :
123 2476 : disassociateVectorFromTag(solution(), tag);
124 2476 : }
125 :
126 : void
127 34353 : NonlinearSystemBase::computeKokkosResidual(const std::set<TagID> & tags)
128 : {
129 103059 : TIME_SECTION("computeKokkosResidual", 1);
130 :
131 34353 : auto & systems = _fe_problem.getKokkosSystems();
132 :
133 34353 : systems[number()].setActiveResidualTags(tags);
134 :
135 : // Resolve dependencies
136 :
137 34353 : std::set<MooseVariableFieldBase *> needed_moose_vars;
138 34353 : std::set<TagID> needed_fe_var_vector_tags;
139 34353 : std::unordered_set<unsigned int> needed_mat_props;
140 :
141 195739 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
142 161386 : needed_fe_var_vector_tags.insert(tag._id);
143 :
144 70334 : for (auto block : _fe_problem.mesh().meshSubdomains())
145 : {
146 35981 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
147 : needed_moose_vars);
148 35981 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
149 : block, needed_fe_var_vector_tags);
150 :
151 35981 : _kokkos_kernels.updateBlockVariableDependency(block, needed_moose_vars);
152 35981 : _kokkos_kernels.updateBlockFEVariableCoupledVectorTagDependency(block,
153 : needed_fe_var_vector_tags);
154 35981 : _kokkos_kernels.updateBlockMatPropDependency(block, needed_mat_props);
155 : }
156 :
157 162802 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
158 : {
159 128449 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryVariableDependency(boundary,
160 : needed_moose_vars);
161 128449 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryFEVariableCoupledVectorTagDependency(
162 : boundary, needed_fe_var_vector_tags);
163 :
164 128449 : _kokkos_integrated_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
165 128449 : _kokkos_integrated_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
166 : boundary, needed_fe_var_vector_tags);
167 128449 : _kokkos_integrated_bcs.updateBoundaryMatPropDependency(boundary, needed_mat_props);
168 :
169 128449 : _kokkos_nodal_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
170 128449 : _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 103059 : for (auto & system : systems)
177 : {
178 68706 : system.setActiveVariables(needed_moose_vars);
179 68706 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
180 :
181 : {
182 206118 : TIME_SECTION("KokkosCopy", 1);
183 68706 : system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
184 68706 : }
185 : {
186 206118 : TIME_SECTION("KokkosReinit", 1);
187 68706 : system.reinit();
188 68706 : }
189 : }
190 :
191 34353 : systems.copyToDevice();
192 :
193 : {
194 103059 : TIME_SECTION("KokkosMaterial", 1);
195 :
196 : // Compute material properties
197 :
198 34353 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
199 34353 : _fe_problem.reinitKokkosMaterials();
200 34353 : }
201 :
202 : {
203 103059 : TIME_SECTION("KokkosKernel", 1);
204 :
205 : // Compute kernels
206 :
207 103094 : for (auto kernel : _kokkos_kernels.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
208 68741 : kernel->computeResidual();
209 :
210 37871 : for (auto nodal_kernel :
211 75742 : _kokkos_nodal_kernels.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
212 20340 : nodal_kernel->computeResidual();
213 :
214 50232 : for (auto ibc : _kokkos_integrated_bcs.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
215 15879 : ibc->computeResidual();
216 :
217 88729 : for (auto nbc : _kokkos_nodal_bcs.getVectorTagsObjectWarehouse(tags, 0).getActiveObjects())
218 54376 : nbc->computeResidual();
219 34353 : }
220 :
221 : // Close and restore vectors
222 :
223 : {
224 103059 : TIME_SECTION("KokkosCopy", 1);
225 :
226 103059 : for (auto & system : systems)
227 68706 : system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
228 34353 : }
229 :
230 : // Clear
231 :
232 34353 : systems[number()].clearActiveResidualTags();
233 :
234 103059 : for (auto & system : systems)
235 : {
236 68706 : system.clearActiveVariables();
237 68706 : system.clearActiveSolutionTags();
238 : }
239 34353 : }
240 :
241 : void
242 5611 : NonlinearSystemBase::computeKokkosJacobian(const std::set<TagID> & tags)
243 : {
244 16833 : TIME_SECTION("computeKokkosJacobian", 1);
245 :
246 5611 : auto & systems = _fe_problem.getKokkosSystems();
247 :
248 5611 : systems[number()].setActiveMatrixTags(tags);
249 :
250 : // Resolve dependencies
251 :
252 5611 : std::set<MooseVariableFieldBase *> needed_moose_vars;
253 5611 : std::set<TagID> needed_fe_var_vector_tags;
254 5611 : std::unordered_set<unsigned int> needed_mat_props;
255 :
256 32345 : for (auto tag : _fe_problem.getVectorTags(Moose::VECTOR_TAG_SOLUTION))
257 26734 : needed_fe_var_vector_tags.insert(tag._id);
258 :
259 11484 : for (auto block : _fe_problem.mesh().meshSubdomains())
260 : {
261 5873 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockVariableDependency(block,
262 : needed_moose_vars);
263 5873 : _fe_problem.getKokkosMaterialsWarehouse().updateBlockFEVariableCoupledVectorTagDependency(
264 : block, needed_fe_var_vector_tags);
265 :
266 5873 : _kokkos_kernels.updateBlockVariableDependency(block, needed_moose_vars);
267 5873 : _kokkos_kernels.updateBlockFEVariableCoupledVectorTagDependency(block,
268 : needed_fe_var_vector_tags);
269 5873 : _kokkos_kernels.updateBlockMatPropDependency(block, needed_mat_props);
270 : }
271 :
272 24509 : for (auto boundary : _fe_problem.mesh().meshBoundaryIds())
273 : {
274 18898 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryVariableDependency(boundary,
275 : needed_moose_vars);
276 18898 : _fe_problem.getKokkosMaterialsWarehouse().updateBoundaryFEVariableCoupledVectorTagDependency(
277 : boundary, needed_fe_var_vector_tags);
278 :
279 18898 : _kokkos_integrated_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
280 18898 : _kokkos_integrated_bcs.updateBoundaryFEVariableCoupledVectorTagDependency(
281 : boundary, needed_fe_var_vector_tags);
282 18898 : _kokkos_integrated_bcs.updateBoundaryMatPropDependency(boundary, needed_mat_props);
283 :
284 18898 : _kokkos_nodal_bcs.updateBoundaryVariableDependency(boundary, needed_moose_vars);
285 18898 : _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 16833 : for (auto & system : systems)
292 : {
293 11222 : system.setActiveVariables(needed_moose_vars);
294 11222 : system.setActiveSolutionTags(needed_fe_var_vector_tags);
295 :
296 : {
297 33666 : TIME_SECTION("KokkosPreallocate", 1);
298 11222 : system.sync(Moose::Kokkos::MemcpyKind::HOST_TO_DEVICE);
299 11222 : }
300 : {
301 33666 : TIME_SECTION("KokkosReinit", 1);
302 11222 : system.reinit();
303 11222 : }
304 : }
305 :
306 5611 : systems.copyToDevice();
307 :
308 : {
309 16833 : TIME_SECTION("KokkosMaterial", 1);
310 :
311 : // Compute material properties
312 :
313 5611 : _fe_problem.prepareKokkosMaterials(needed_mat_props);
314 5611 : _fe_problem.reinitKokkosMaterials();
315 5611 : }
316 :
317 : {
318 16833 : TIME_SECTION("KokkosKernel", 1);
319 :
320 : // Compute kernels
321 :
322 17951 : for (auto kernel : _kokkos_kernels.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
323 12340 : kernel->computeJacobian();
324 :
325 6655 : for (auto nodal_kernel :
326 13800 : _kokkos_nodal_kernels.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
327 5424 : nodal_kernel->computeJacobian();
328 :
329 7513 : for (auto ibc : _kokkos_integrated_bcs.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
330 1902 : ibc->computeJacobian();
331 :
332 14795 : for (auto nbc : _kokkos_nodal_bcs.getMatrixTagsObjectWarehouse(tags, 0).getActiveObjects())
333 9184 : nbc->computeJacobian();
334 5611 : }
335 :
336 : // Close and restore vectors and matrices
337 :
338 : {
339 16833 : TIME_SECTION("KokkosCopy", 1);
340 :
341 16833 : for (auto & system : systems)
342 11222 : system.sync(Moose::Kokkos::MemcpyKind::DEVICE_TO_HOST);
343 5611 : }
344 :
345 : // Clear
346 :
347 5611 : systems[number()].clearActiveMatrixTags();
348 :
349 16833 : for (auto & system : systems)
350 : {
351 11222 : system.clearActiveVariables();
352 11222 : system.clearActiveSolutionTags();
353 : }
354 5611 : }
|
