Skip to content

Commit b8608bb

Browse files
odiazibodiazib
authored
TChem: Using Impl::Aerosol_RHS instead of problem to compute RHS. (#113)
* TChem: Using Impl::Aerosol_RHS instead of problem to compute RHS. * save
1 parent 76290be commit b8608bb

File tree

1 file changed

+58
-54
lines changed

1 file changed

+58
-54
lines changed

‎src/examples/TChem_AerosolChemistry_RHSs.cpp‎

Lines changed: 58 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ int main(int argc, char *argv[]) {
6161
"use_cloned_samples", "If true, one state vector will be cloned.", &use_cloned_samples);
6262
opts.set_option<bool>("do_jac", "Evaluate Jacobian matrix", &do_jac);
6363
opts.set_option<bool>("do_rhs", "Evaluate RHS", &do_rhs);
64-
64+
6565

6666
const bool r_parse = opts.parse(argc, argv);
6767
if (r_parse)
@@ -192,12 +192,48 @@ int main(int argc, char *argv[]) {
192192

193193
using range_type = Kokkos::pair<ordinal_type, ordinal_type>;
194194
const ordinal_type level = 1;
195-
const ordinal_type per_team_extent = problem_type::getWorkSpaceSize(kmcd,amcd)
196-
+ number_of_equations;
195+
ordinal_type per_team_extent=0.0;
196+
if (do_jac){
197+
per_team_extent = problem_type::getWorkSpaceSize(kmcd,amcd);
198+
}
199+
200+
if (do_rhs){
201+
per_team_extent
202+
= TChem::Impl::Aerosol_RHS<real_type, device_type>::getWorkSpaceSize(kmcd, amcd);
203+
}
197204
const ordinal_type per_team_scratch =
198205
TChem::Scratch<real_type_1d_view_type>::shmem_size(per_team_extent);
199206
policy.set_scratch_size(level, Kokkos::PerTeam(per_team_scratch));
200207

208+
real_type_2d_view_type y2d("y2d", nBatch, number_of_equations);
209+
210+
Kokkos::parallel_for(
211+
"fill_y", Kokkos::RangePolicy<TChem::exec_space>(0, nBatch),
212+
KOKKOS_LAMBDA(const int i) {
213+
const ordinal_type n_active_gas_species = kmcd.nSpec - kmcd.nConstSpec;
214+
const real_type_1d_view_type state_at_i =
215+
Kokkos::subview(state, i, Kokkos::ALL());
216+
const ordinal_type total_n_species = kmcd.nSpec + amcd.nParticles*amcd.nSpec;
217+
TChem::Impl::StateVector<real_type_1d_view_type> sv_at_i(total_n_species, state_at_i);
218+
const real_type_1d_view_type Ys = sv_at_i.MassFractions();
219+
220+
const auto activeYs = Kokkos::subview(Ys, range_type(0, n_active_gas_species));
221+
const real_type_1d_view_type partYs = Kokkos::subview(Ys, range_type(kmcd.nSpec, total_n_species));
222+
223+
for (ordinal_type j=0;j<n_active_gas_species;++j){
224+
y2d(i, j) = activeYs(j);
225+
}
226+
227+
for (ordinal_type j=n_active_gas_species;j<total_n_species- kmcd.nConstSpec;++j)
228+
{
229+
y2d(i, j) = partYs(j-n_active_gas_species);
230+
}
231+
232+
});
233+
234+
235+
236+
201237
Kokkos::Timer timer;
202238
FILE *fout_times = fopen(outputFileTimes.c_str(), "w");
203239
fprintf(fout_times, "{\n");
@@ -223,6 +259,9 @@ int main(int argc, char *argv[]) {
223259
const real_type_1d_view_type state_at_i =
224260
Kokkos::subview(state, i, Kokkos::ALL());
225261

262+
const real_type_1d_view_type vals =
263+
Kokkos::subview(y2d, i, Kokkos::ALL());
264+
226265
const real_type_1d_view_type number_conc_at_i =
227266
Kokkos::subview(num_concentration, i, Kokkos::ALL());
228267
TChem::Scratch<real_type_1d_view_type> work(member.team_scratch(level),
@@ -235,42 +274,19 @@ int main(int argc, char *argv[]) {
235274
const real_type pressure = sv_at_i.Pressure();
236275
const real_type_1d_view_type Ys = sv_at_i.MassFractions();
237276
const ordinal_type n_active_gas_species = kmcd.nSpec - kmcd.nConstSpec;
238-
const real_type_1d_view_type activeYs = Kokkos::subview(Ys,
239-
range_type(0, n_active_gas_species));
240277
const real_type_1d_view_type constYs = Kokkos::subview(Ys,
241-
range_type(n_active_gas_species, kmcd.nSpec));
242-
const real_type_1d_view_type partYs = Kokkos::subview(Ys, range_type(kmcd.nSpec, total_n_species));
243-
244-
real_type_1d_view_type vals(wptr, m);
245-
wptr += m;
278+
range_type(n_active_gas_species, kmcd.nSpec));
246279

247280
/// problem workspace
248281
/// problem setup
249-
const ordinal_type problem_workspace_size = problem_type::getWorkSpaceSize(kmcd,amcd);
282+
const ordinal_type problem_workspace_size
283+
= TChem::Impl::Aerosol_RHS<real_type, device_type>::getWorkSpaceSize(kmcd, amcd);
250284
auto pw = real_type_1d_view_type(wptr, problem_workspace_size);
251285
wptr +=problem_workspace_size;
252-
problem_type problem;
253-
problem._kmcd = kmcd;
254-
problem._amcd = amcd;
255-
256-
/// initialize problem
257-
// problem._fac = fac_at_i;
258-
problem._work = pw;
259-
problem._temperature = temperature;
260-
problem._pressure = pressure;
261-
problem._const_concentration = constYs;
262-
problem._number_conc = number_conc_at_i;
263-
// active gas species
264-
for (ordinal_type i=0;i<n_active_gas_species;++i){
265-
vals(i) = activeYs(i);
266-
}
267-
268-
for (ordinal_type i=n_active_gas_species;i<total_n_species- kmcd.nConstSpec;++i)
269-
{
270-
vals(i) = partYs(i-n_active_gas_species);
271-
}
272-
273-
problem.computeFunction(member,vals,rhs_at_i);
286+
TChem::Impl::Aerosol_RHS<real_type, device_type>
287+
::team_invoke(member, temperature, pressure,
288+
number_conc_at_i, vals,
289+
constYs, rhs_at_i, pw, kmcd, amcd);
274290
});
275291

276292
Kokkos::Profiling::popRegion();
@@ -282,11 +298,11 @@ int main(int argc, char *argv[]) {
282298

283299
if (do_jac){
284300
fprintf(fout_times, "}, \n ");// reaction rates
285-
}
301+
}
286302
else {
287303
fprintf(fout_times, "} \n ");// reaction rates
288304
}
289-
305+
290306
}
291307

292308
if (do_jac){
@@ -314,6 +330,9 @@ int main(int argc, char *argv[]) {
314330
const real_type_1d_view_type fac_at_i =
315331
Kokkos::subview(fac, i, Kokkos::ALL());
316332

333+
const real_type_1d_view_type vals =
334+
Kokkos::subview(y2d, i, Kokkos::ALL());
335+
317336
const real_type_1d_view_type number_conc_at_i =
318337
Kokkos::subview(num_concentration, i, Kokkos::ALL());
319338
TChem::Scratch<real_type_1d_view_type> work(member.team_scratch(level),
@@ -326,15 +345,8 @@ int main(int argc, char *argv[]) {
326345
const real_type pressure = sv_at_i.Pressure();
327346
const real_type_1d_view_type Ys = sv_at_i.MassFractions();
328347
const ordinal_type n_active_gas_species = kmcd.nSpec - kmcd.nConstSpec;
329-
const auto activeYs = real_type_1d_view_type(Ys.data(),
330-
n_active_gas_species );
331-
const auto constYs = real_type_1d_view_type(Ys.data(),
332-
+n_active_gas_species, kmcd.nSpec );
333-
const real_type_1d_view_type partYs = Kokkos::subview(Ys, range_type(kmcd.nSpec, total_n_species));
334-
335-
real_type_1d_view_type vals(wptr, m);
336-
wptr += m;
337-
348+
const real_type_1d_view_type constYs = Kokkos::subview(Ys,
349+
range_type(n_active_gas_species, kmcd.nSpec));
338350
/// problem workspace
339351
/// problem setup
340352
const ordinal_type problem_workspace_size = problem_type::getWorkSpaceSize(kmcd,amcd);
@@ -352,14 +364,6 @@ int main(int argc, char *argv[]) {
352364
problem._const_concentration= constYs;
353365
problem._number_conc =number_conc_at_i;
354366
// active gas species
355-
for (ordinal_type i=0;i<n_active_gas_species;++i){
356-
vals(i) = activeYs(i);
357-
}
358-
359-
for (ordinal_type i=n_active_gas_species;i<total_n_species- kmcd.nConstSpec;++i)
360-
{
361-
vals(i) = partYs(i-n_active_gas_species);
362-
}
363367
problem.computeNumericalJacobian(member,vals,jacobian_at_i);
364368
});
365369
Kokkos::Profiling::popRegion();
@@ -373,8 +377,8 @@ int main(int argc, char *argv[]) {
373377

374378
if (verbose & do_rhs) {
375379

376-
377-
if (use_cloned_samples) {
380+
381+
if (use_cloned_samples) {
378382
auto rhs_at_0 = Kokkos::subview(rhs, 0, Kokkos::ALL());
379383
auto rhs_host_at_0 = Kokkos::create_mirror_view(rhs_at_0);
380384
Kokkos::deep_copy(rhs_host_at_0, rhs_at_0);

0 commit comments

Comments
 (0)