Skip to content

Commit

Permalink
Port poisson (r,theta) weak integrals on gpu
Browse files Browse the repository at this point in the history 
Closes #370

In addition to the changes described in the issue, the following modifications were necessary to be able to compute weak integrals on gpu:

* Add device builder and evaluators aliases in geometry file : `SplineRThetaBuilder_device, SplineRThetaEvaluatorConstBound_device, SplineRThetaEvaluatorNullBound_device`
* Change the interpolator used in the poisson solver class template parameters to : `SplineRThetaEvaluatorNullBound_device`
* Modify the constructor to take device allocated structures `eg: coeff_alpha, coeff_beta`
* The integral volume associated with each point used in the quadrature scheme is now filled in a function : `fill_int_volume`
* FEM matrix filling is splitted into 3 functions
  * `compute_singular_elements`
  * `compute_overlapping_singular_elements`
  * `compute_stencil_elements`
* Propagate necessary changes for instantiation and `operator()` call in poisson solver test file.
* Propagate necessary changes for instantiation and `operator()` call in `BslPredCorrRTheta, BslExplicitPredCorrRTheta, BslImplicitPredCorrRTheta` classes
* Propagate necessary changes for instantiation and `operator()` call in `diocotron `and `vortex_merger` simulations.

See merge request gysela-developpers/gyselalibxx!773

--------------------------------------------

Co-authored-by: Emily Bourne <emily.bourne@epfl.ch>
  • Loading branch information
@AbdelhadiKara @EmilyBourne
AbdelhadiKara and EmilyBourne committed Jan 9, 2025
1 parent 4022e06 commit 860495e
Show file tree
Hide file tree
Showing 13 changed files with 730 additions and 286 deletions.
3 changes: 2 additions & 1 deletion ci_tools/gyselalib_static_analysis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -602,7 +602,8 @@ def check_exec_space_usage(file):
scope = scopes[scope_id]

if config.data[current_idx]['str'] == 'KOKKOS_LAMBDA':
ast_parent = config.data_xml.find(f".token[@id='{config.data[current_idx-1]['astParent']}']").attrib
ast_parent_id = config.data[current_idx-1].get('astParent', config.data[current_idx-1]['id'])
ast_parent = config.data_xml.find(f".token[@id='{ast_parent_id}']").attrib
parallel_func_start_idx = config.data.index(ast_parent)
args = [t['str'] for t in config.data[parallel_func_start_idx+1:current_idx]]
if 'DefaultHostExecutionSpace' in args:
Expand Down
85 changes: 53 additions & 32 deletions simulations/geometryRTheta/diocotron/diocotron.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,10 @@ using PoissonSolver = PolarSplineFEMPoissonLikeSolver<
GridR,
GridTheta,
PolarBSplinesRTheta,
SplineRThetaEvaluatorNullBound_host>;
using DiscreteMappingBuilder = DiscreteToCartesianBuilder<
SplineRThetaEvaluatorNullBound>;
using DiscreteMappingBuilder
= DiscreteToCartesianBuilder<X, Y, SplineRThetaBuilder, SplineRThetaEvaluatorConstBound>;
using DiscreteMappingBuilder_host = DiscreteToCartesianBuilder<
X,
Y,
SplineRThetaBuilder_host,
Expand Down Expand Up @@ -104,31 +106,46 @@ int main(int argc, char** argv)


// OPERATORS ======================================================================================
SplineRThetaBuilder_host const builder(mesh_rp);
SplineRThetaBuilder const builder(mesh_rp);
SplineRThetaBuilder_host const builder_host(mesh_rp);


// --- Define the mapping. ------------------------------------------------------------------------
ddc::ConstantExtrapolationRule<R, Theta> boundary_condition_r_left(
ddc::coordinate(mesh_r.front()));
ddc::ConstantExtrapolationRule<R, Theta> boundary_condition_r_right(
ddc::coordinate(mesh_r.back()));

SplineRThetaEvaluatorConstBound_host spline_evaluator_extrapol(
SplineRThetaEvaluatorConstBound spline_evaluator_extrapol(
boundary_condition_r_left,
boundary_condition_r_right,
ddc::PeriodicExtrapolationRule<Theta>(),
ddc::PeriodicExtrapolationRule<Theta>());

SplineRThetaEvaluatorConstBound_host spline_evaluator_extrapol_host(
boundary_condition_r_left,
boundary_condition_r_right,
ddc::PeriodicExtrapolationRule<Theta>(),
ddc::PeriodicExtrapolationRule<Theta>());

const LogicalToPhysicalMapping to_physical_mapping;
DiscreteMappingBuilder const discrete_mapping_builder(
Kokkos::DefaultHostExecutionSpace(),
Kokkos::DefaultExecutionSpace(),
to_physical_mapping,
builder,
spline_evaluator_extrapol);
DiscreteMappingBuilder_host const discrete_mapping_builder_host(
Kokkos::DefaultHostExecutionSpace(),
to_physical_mapping,
builder_host,
spline_evaluator_extrapol_host);
DiscreteToCartesian const discrete_mapping = discrete_mapping_builder();
DiscreteToCartesian const discrete_mapping_host = discrete_mapping_builder_host();

ddc::init_discrete_space<PolarBSplinesRTheta>(discrete_mapping);

IdxRangeBSRTheta const idx_range_bsplinesRTheta = get_spline_idx_range(builder);
ddc::init_discrete_space<PolarBSplinesRTheta>(discrete_mapping_host);

IdxRangeBSRTheta const idx_range_bsplinesRTheta = get_spline_idx_range(builder_host);


// --- Time integration method --------------------------------------------------------------------
Expand Down Expand Up @@ -160,39 +177,42 @@ int main(int argc, char** argv)
// --- Advection operator -------------------------------------------------------------------------
ddc::NullExtrapolationRule r_extrapolation_rule;
ddc::PeriodicExtrapolationRule<Theta> p_extrapolation_rule;
SplineRThetaEvaluatorNullBound_host spline_evaluator(
SplineRThetaEvaluatorNullBound spline_evaluator(
r_extrapolation_rule,
r_extrapolation_rule,
p_extrapolation_rule,
p_extrapolation_rule);
SplineRThetaEvaluatorNullBound_host spline_evaluator_host(
r_extrapolation_rule,
r_extrapolation_rule,
p_extrapolation_rule,
p_extrapolation_rule);

PreallocatableSplineInterpolatorRTheta interpolator(builder, spline_evaluator);
PreallocatableSplineInterpolatorRTheta interpolator(builder_host, spline_evaluator_host);

AdvectionDomain advection_domain(to_physical_mapping);

SplineFootFinder find_feet(
time_stepper,
advection_domain,
to_physical_mapping,
builder,
spline_evaluator_extrapol);
builder_host,
spline_evaluator_extrapol_host);

BslAdvectionRTheta advection_operator(interpolator, find_feet, to_physical_mapping);



// --- Poisson solver -----------------------------------------------------------------------------
// Coefficients alpha and beta of the Poisson equation:
host_t<DFieldMemRTheta> coeff_alpha(mesh_rp);
host_t<DFieldMemRTheta> coeff_beta(mesh_rp);
DFieldMemRTheta coeff_alpha(mesh_rp);
DFieldMemRTheta coeff_beta(mesh_rp);
ddc::parallel_fill(coeff_alpha, -1);
ddc::parallel_fill(coeff_beta, 0);

ddc::for_each(mesh_rp, [&](IdxRTheta const irp) {
coeff_alpha(irp) = -1.0;
coeff_beta(irp) = 0.0;
});

host_t<Spline2DMem> coeff_alpha_spline(idx_range_bsplinesRTheta);
host_t<Spline2DMem> coeff_beta_spline(idx_range_bsplinesRTheta);
Spline2DMem coeff_alpha_spline(idx_range_bsplinesRTheta);
Spline2DMem coeff_beta_spline(idx_range_bsplinesRTheta);

builder(get_field(coeff_alpha_spline), get_const_field(coeff_alpha));
builder(get_field(coeff_beta_spline), get_const_field(coeff_beta));
Expand All @@ -208,29 +228,29 @@ int main(int argc, char** argv)
BslPredCorrRTheta predcorr_operator(
to_physical_mapping,
advection_operator,
builder,
spline_evaluator,
builder_host,
spline_evaluator_host,
poisson_solver);
#elif defined(EXPLICIT_PREDCORR)
BslExplicitPredCorrRTheta predcorr_operator(
advection_domain,
to_physical_mapping,
advection_operator,
mesh_rp,
builder,
spline_evaluator,
builder_host,
spline_evaluator_host,
poisson_solver,
spline_evaluator_extrapol);
spline_evaluator_extrapol_host);
#elif defined(IMPLICIT_PREDCORR)
BslImplicitPredCorrRTheta predcorr_operator(
advection_domain,
to_physical_mapping,
advection_operator,
mesh_rp,
builder,
spline_evaluator,
builder_host,
spline_evaluator_host,
poisson_solver,
spline_evaluator_extrapol);
spline_evaluator_extrapol_host);
#endif

// ================================================================================================
Expand Down Expand Up @@ -296,20 +316,21 @@ int main(int argc, char** argv)
});

// Compute phi equilibrium phi_eq from Poisson solver. ***********
host_t<DFieldMemRTheta> phi_eq(mesh_rp);
DFieldMemRTheta phi_eq(mesh_rp);
host_t<DFieldMemRTheta> phi_eq_host(mesh_rp);
host_t<Spline2DMem> rho_coef_eq(idx_range_bsplinesRTheta);
builder(get_field(rho_coef_eq), get_const_field(rho_eq));
PoissonLikeRHSFunction poisson_rhs_eq(get_const_field(rho_coef_eq), spline_evaluator);
builder_host(get_field(rho_coef_eq), get_const_field(rho_eq));
PoissonLikeRHSFunction poisson_rhs_eq(get_const_field(rho_coef_eq), spline_evaluator_host);
poisson_solver(poisson_rhs_eq, get_field(phi_eq));

ddc::parallel_deepcopy(phi_eq, phi_eq_host);

// --- Save initial data --------------------------------------------------------------------------
ddc::PdiEvent("initialization")
.with("x_coords", coords_x)
.and_with("y_coords", coords_y)
.and_with("jacobian", jacobian)
.and_with("density_eq", rho_eq)
.and_with("electrical_potential_eq", phi_eq);
.and_with("electrical_potential_eq", phi_eq_host);


// ================================================================================================
Expand Down
77 changes: 50 additions & 27 deletions simulations/geometryRTheta/vortex_merger/vortex_merger.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,10 @@ using PoissonSolver = PolarSplineFEMPoissonLikeSolver<
GridR,
GridTheta,
PolarBSplinesRTheta,
SplineRThetaEvaluatorNullBound_host>;
using DiscreteMappingBuilder = DiscreteToCartesianBuilder<
SplineRThetaEvaluatorNullBound>;
using DiscreteMappingBuilder
= DiscreteToCartesianBuilder<X, Y, SplineRThetaBuilder, SplineRThetaEvaluatorConstBound>;
using DiscreteMappingBuilder_host = DiscreteToCartesianBuilder<
X,
Y,
SplineRThetaBuilder_host,
Expand Down Expand Up @@ -99,31 +101,43 @@ int main(int argc, char** argv)


// OPERATORS ======================================================================================
SplineRThetaBuilder_host const builder(grid);
SplineRThetaBuilder const builder(grid);
SplineRThetaBuilder_host const builder_host(grid);

// --- Define the mapping. ------------------------------------------------------------------------
ddc::ConstantExtrapolationRule<R, Theta> boundary_condition_r_left(
ddc::coordinate(mesh_r.front()));
ddc::ConstantExtrapolationRule<R, Theta> boundary_condition_r_right(
ddc::coordinate(mesh_r.back()));

SplineRThetaEvaluatorConstBound_host spline_evaluator_extrapol(
SplineRThetaEvaluatorConstBound spline_evaluator_extrapol(
boundary_condition_r_left,
boundary_condition_r_right,
ddc::PeriodicExtrapolationRule<Theta>(),
ddc::PeriodicExtrapolationRule<Theta>());
SplineRThetaEvaluatorConstBound_host spline_evaluator_extrapol_host(
boundary_condition_r_left,
boundary_condition_r_right,
ddc::PeriodicExtrapolationRule<Theta>(),
ddc::PeriodicExtrapolationRule<Theta>());

const LogicalToPhysicalMapping to_physical_mapping;
DiscreteMappingBuilder const discrete_mapping_builder(
Kokkos::DefaultHostExecutionSpace(),
Kokkos::DefaultExecutionSpace(),
to_physical_mapping,
builder,
spline_evaluator_extrapol);
DiscreteMappingBuilder_host const discrete_mapping_builder_host(
Kokkos::DefaultHostExecutionSpace(),
to_physical_mapping,
builder_host,
spline_evaluator_extrapol_host);
DiscreteToCartesian const discrete_mapping = discrete_mapping_builder();
DiscreteToCartesian const discrete_mapping_host = discrete_mapping_builder_host();

ddc::init_discrete_space<PolarBSplinesRTheta>(discrete_mapping);
ddc::init_discrete_space<PolarBSplinesRTheta>(discrete_mapping_host);

IdxRangeBSRTheta const idx_range_bsplinesRTheta = get_spline_idx_range(builder);
IdxRangeBSRTheta const idx_range_bsplinesRTheta = get_spline_idx_range(builder_host);


// --- Time integration method --------------------------------------------------------------------
Expand All @@ -135,39 +149,42 @@ int main(int argc, char** argv)
// --- Advection operator -------------------------------------------------------------------------
ddc::NullExtrapolationRule r_extrapolation_rule;
ddc::PeriodicExtrapolationRule<Theta> p_extrapolation_rule;
SplineRThetaEvaluatorNullBound_host spline_evaluator(
SplineRThetaEvaluatorNullBound spline_evaluator(
r_extrapolation_rule,
r_extrapolation_rule,
p_extrapolation_rule,
p_extrapolation_rule);
SplineRThetaEvaluatorNullBound_host spline_evaluator_host(
r_extrapolation_rule,
r_extrapolation_rule,
p_extrapolation_rule,
p_extrapolation_rule);

PreallocatableSplineInterpolatorRTheta interpolator(builder, spline_evaluator);
PreallocatableSplineInterpolatorRTheta interpolator(builder_host, spline_evaluator_host);

AdvectionDomain advection_domain(to_physical_mapping);

SplineFootFinder find_feet(
time_stepper,
advection_domain,
to_physical_mapping,
builder,
spline_evaluator_extrapol);
builder_host,
spline_evaluator_extrapol_host);

BslAdvectionRTheta advection_operator(interpolator, find_feet, to_physical_mapping);



// --- Poisson solver -----------------------------------------------------------------------------
// Coefficients alpha and beta of the Poisson equation:
host_t<DFieldMemRTheta> coeff_alpha(grid);
host_t<DFieldMemRTheta> coeff_beta(grid);
DFieldMemRTheta coeff_alpha(grid);
DFieldMemRTheta coeff_beta(grid);

ddc::for_each(grid, [&](IdxRTheta const irp) {
coeff_alpha(irp) = -1.0;
coeff_beta(irp) = 0.0;
});
ddc::parallel_fill(coeff_alpha, -1);
ddc::parallel_fill(coeff_beta, 0);

host_t<Spline2DMem> coeff_alpha_spline(idx_range_bsplinesRTheta);
host_t<Spline2DMem> coeff_beta_spline(idx_range_bsplinesRTheta);
Spline2DMem coeff_alpha_spline(idx_range_bsplinesRTheta);
Spline2DMem coeff_beta_spline(idx_range_bsplinesRTheta);

builder(get_field(coeff_alpha_spline), get_const_field(coeff_alpha));
builder(get_field(coeff_beta_spline), get_const_field(coeff_beta));
Expand All @@ -184,10 +201,10 @@ int main(int argc, char** argv)
to_physical_mapping,
advection_operator,
grid,
builder,
spline_evaluator,
builder_host,
spline_evaluator_host,
poisson_solver,
spline_evaluator_extrapol);
spline_evaluator_extrapol_host);



Expand Down Expand Up @@ -243,8 +260,12 @@ int main(int argc, char** argv)
double const phi_max(1.);


VortexMergerEquilibria
equilibrium(to_physical_mapping, grid, builder, spline_evaluator, poisson_solver);
VortexMergerEquilibria equilibrium(
to_physical_mapping,
grid,
builder_host,
spline_evaluator_host,
poisson_solver);
std::function<double(double const)> const function = [&](double const x) { return x * x; };
host_t<DFieldMemRTheta> rho_eq(grid);
equilibrium.set_equilibrium(get_field(rho_eq), function, phi_max, tau);
Expand All @@ -264,11 +285,13 @@ int main(int argc, char** argv)


// Compute phi equilibrium phi_eq from Poisson solver. ***********
host_t<DFieldMemRTheta> phi_eq(grid);
DFieldMemRTheta phi_eq(grid);
host_t<DFieldMemRTheta> phi_eq_host(grid);
host_t<Spline2DMem> rho_coef_eq(idx_range_bsplinesRTheta);
builder(get_field(rho_coef_eq), get_const_field(rho_eq));
PoissonLikeRHSFunction poisson_rhs_eq(get_const_field(rho_coef_eq), spline_evaluator);
builder_host(get_field(rho_coef_eq), get_const_field(rho_eq));
PoissonLikeRHSFunction poisson_rhs_eq(get_const_field(rho_coef_eq), spline_evaluator_host);
poisson_solver(poisson_rhs_eq, get_field(phi_eq));
ddc::parallel_deepcopy(phi_eq, phi_eq_host);


// --- Save initial data --------------------------------------------------------------------------
Expand Down
Loading

0 comments on commit 860495e

Please sign in to comment.