Skip to content

Commit

Permalink
Add a MemSpace template argument to the Polar spline classes
Browse files Browse the repository at this point in the history
Closes #371

See merge request gysela-developpers/gyselalibxx!725

--------------------------------------------
  • Loading branch information
Geoflow committed Oct 25, 2024
1 parent 14f5e30 commit 6af7eca
Show file tree
Hide file tree
Showing 12 changed files with 153 additions and 90 deletions.
12 changes: 7 additions & 5 deletions src/geometryRTheta/advection_field/advection_field_rp.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ class AdvectionFieldFinder
private:
Mapping const& m_mapping;

PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule> const
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule, Kokkos::HostSpace> const
m_polar_spline_evaluator;

SplineRThetaEvaluatorNullBound const m_spline_evaluator;
Expand Down Expand Up @@ -188,7 +188,7 @@ class AdvectionFieldFinder
* The advection field on the physical axis.
*/
void operator()(
SplinePolar& electrostatic_potential_coef,
host_t<SplinePolar>& electrostatic_potential_coef,
host_t<DVectorFieldRTheta<X, Y>> advection_field_xy) const
{
compute_advection_field_XY(
Expand Down Expand Up @@ -223,7 +223,8 @@ class AdvectionFieldFinder
Evaluator,
PolarSplineEvaluator<
PolarBSplinesRTheta,
ddc::NullExtrapolationRule>> && std::is_same_v<SplineType, SplinePolar>));
ddc::NullExtrapolationRule,
Kokkos::HostSpace>> && std::is_same_v<SplineType, host_t<SplinePolar>>));

IdxRangeRTheta const grid = get_idx_range(advection_field_xy);
host_t<DVectorFieldMemRTheta<X, Y>> electric_field(grid);
Expand Down Expand Up @@ -412,7 +413,7 @@ class AdvectionFieldFinder
* The advection field on the physical axis at the O-point.
*/
void operator()(
SplinePolar& electrostatic_potential_coef,
host_t<SplinePolar>& electrostatic_potential_coef,
host_t<DVectorFieldRTheta<R, Theta>> advection_field_rp,
CoordXY& advection_field_xy_center) const
{
Expand Down Expand Up @@ -453,7 +454,8 @@ class AdvectionFieldFinder
Evaluator,
PolarSplineEvaluator<
PolarBSplinesRTheta,
ddc::NullExtrapolationRule>> && std::is_same_v<SplineType, SplinePolar>));
ddc::NullExtrapolationRule,
Kokkos::HostSpace>> && std::is_same_v<SplineType, host_t<SplinePolar>>));

IdxRangeRTheta const grid_without_Opoint = get_idx_range(advection_field_rp);

Expand Down
7 changes: 4 additions & 3 deletions src/geometryRTheta/poisson/polarpoissonlikesolver.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,8 @@ class PolarSplineFEMPoissonLikeSolver

host_t<FieldMem<double, IdxRangeQuadratureRTheta>> int_volume;

PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule> m_polar_spline_evaluator;
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule, Kokkos::HostSpace>
m_polar_spline_evaluator;
std::unique_ptr<MatrixBatchCsr<Kokkos::DefaultExecutionSpace, MatrixBatchCsrSolver::CG>>
m_gko_matrix;
const int batch_idx {0};
Expand Down Expand Up @@ -653,7 +654,7 @@ class PolarSplineFEMPoissonLikeSolver
* The spline representation of the solution @f$\phi@f$.
*/
template <class RHSFunction>
void operator()(RHSFunction const& rhs, SplinePolar& spline) const
void operator()(RHSFunction const& rhs, host_t<SplinePolar>& spline) const
{
const int b_size = ddc::discrete_space<PolarBSplinesRTheta>().nbasis()
- ddc::discrete_space<BSplinesTheta_Polar>().nbasis();
Expand Down Expand Up @@ -806,7 +807,7 @@ class PolarSplineFEMPoissonLikeSolver
{
IdxRangeBSTheta_Polar polar_idx_range(
ddc::discrete_space<BSplinesTheta_Polar>().full_domain());
SplinePolar
host_t<SplinePolar>
spline(PolarBSplinesRTheta::singular_idx_range<PolarBSplinesRTheta>(),
IdxRangeBSRTheta(radial_bsplines, polar_idx_range));

Expand Down
4 changes: 2 additions & 2 deletions src/geometryRTheta/time_solver/bsl_predcorr.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -123,11 +123,11 @@ class BslPredCorrRTheta : public ITimeSolverRTheta
IdxStep<BSplinesR> {PolarBSplinesRTheta::continuity + 1}));
IdxRangeBSTheta polar_idx_range(ddc::discrete_space<BSplinesTheta>().full_domain());

SplinePolar electrostatic_potential_coef(
host_t<SplinePolar> electrostatic_potential_coef(
PolarBSplinesRTheta::singular_idx_range<PolarBSplinesRTheta>(),
IdxRangeBSRTheta(radial_bsplines, polar_idx_range));
ddc::NullExtrapolationRule extrapolation_rule;
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule>
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule, Kokkos::HostSpace>
polar_spline_evaluator(extrapolation_rule);


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -158,12 +158,12 @@ class BslExplicitPredCorrRTheta : public ITimeSolverRTheta
// --- Electrostatic potential (phi). -------------------------------------------------------------
host_t<DFieldMemRTheta> electrical_potential(grid);

SplinePolar electrostatic_potential_coef(
host_t<SplinePolar> electrostatic_potential_coef(
PolarBSplinesRTheta::singular_idx_range<PolarBSplinesRTheta>(),
IdxRangeBSRTheta(radial_bsplines, polar_idx_range));

ddc::NullExtrapolationRule extrapolation_rule;
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule>
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule, Kokkos::HostSpace>
polar_spline_evaluator(extrapolation_rule);

// --- For the computation of advection field from the electrostatic potential (phi): -------------
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -158,12 +158,12 @@ class BslImplicitPredCorrRTheta : public ITimeSolverRTheta
// --- Electrostatic potential (phi). -------------------------------------------------------------
host_t<DFieldMemRTheta> electrical_potential(grid);

SplinePolar electrostatic_potential_coef(
host_t<SplinePolar> electrostatic_potential_coef(
PolarBSplinesRTheta::singular_idx_range<PolarBSplinesRTheta>(),
IdxRangeBSRTheta(radial_bsplines, polar_idx_range));

ddc::NullExtrapolationRule extrapolation_rule;
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule>
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule, Kokkos::HostSpace>
polar_spline_evaluator(extrapolation_rule);

// --- For the computation of advection field from the electrostatic potential (phi): -------------
Expand Down
18 changes: 10 additions & 8 deletions src/utils/ddc_alias_inline_functions.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -82,17 +82,19 @@ inline constexpr bool
template <class ElementType, class IdxRangeType, class MemoryType>
inline constexpr bool enable_mem_type<FieldMem<ElementType, IdxRangeType, MemoryType>> = true;

template <class PolarBSplines>
inline constexpr bool enable_data_access_methods<PolarSpline<PolarBSplines>> = true;
template <class PolarBSplines, class MemorySpace>
inline constexpr bool enable_data_access_methods<PolarSpline<PolarBSplines, MemorySpace>> = true;

template <class PolarBSplines>
inline constexpr bool enable_data_access_methods<PolarSplineView<PolarBSplines>> = true;
template <class PolarBSplines, class MemorySpace>
inline constexpr bool
enable_data_access_methods<PolarSplineView<PolarBSplines, MemorySpace>> = true;

template <class PolarBSplines>
inline constexpr bool enable_data_access_methods<PolarSplineSpan<PolarBSplines>> = true;
template <class PolarBSplines, class MemorySpace>
inline constexpr bool
enable_data_access_methods<PolarSplineSpan<PolarBSplines, MemorySpace>> = true;

template <class PolarBSplines>
inline constexpr bool enable_mem_type<PolarSpline<PolarBSplines>> = true;
template <class PolarBSplines, class MemorySpace>
inline constexpr bool enable_mem_type<PolarSpline<PolarBSplines, MemorySpace>> = true;

template <typename Type>
static constexpr bool has_idx_range_v = detail::HasIdxRange<Type>::value;
Expand Down
16 changes: 16 additions & 0 deletions src/utils/ddc_helper.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@

#include <ddc/ddc.hpp>

#include <sll/polar_spline.hpp>

#include "ddc_aliases.hpp"
#include "directional_tag.hpp"
#include "transpose.hpp"
Expand Down Expand Up @@ -313,6 +315,20 @@ struct OnMemorySpace<
using type = VectorField<ElementType, SupportType, NDTag, NewMemorySpace, Layout>;
};

/**
* @brief Get a new `PolarSpline` type with the same parametrisation
* except in the memory space which is set to NewMemorySpace.
* @tparam NewMemorySpace The new memory space.
* @tparam PolarSplineType Type of the B-splines.
* @tparam MemorySpace The original memory space of the chunk of the VectorFieldMem.
* @see VectorField
*/
template <class NewMemorySpace, class PolarSplineType, class MemorySpace>
struct OnMemorySpace<NewMemorySpace, PolarSpline<PolarSplineType, MemorySpace>>
{
using type = PolarSpline<PolarSplineType, NewMemorySpace>;
};


template <template <class Tag> class Templ, class TypeSeq>
struct ApplyTemplateToTypeSeq;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,8 @@ TEST(AdvectionFieldRThetaComputation, TestAdvectionFieldFinder)


ddc::NullExtrapolationRule r_extrapolation_rule;
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule> polar_spline_evaluator(
r_extrapolation_rule);
PolarSplineEvaluator<PolarBSplinesRTheta, ddc::NullExtrapolationRule, Kokkos::HostSpace>
polar_spline_evaluator(r_extrapolation_rule);

// --- Define the mapping. ------------------------------------------------------------------------
const Mapping mapping;
Expand Down
38 changes: 23 additions & 15 deletions vendor/sll/include/sll/polar_bsplines.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -490,7 +490,7 @@ class PolarBSplines
*
* @param[out] int_vals The integrals of the basis splines.
*/
void integrals(PolarSplineSpan<DDim> int_vals) const;
void integrals(PolarSplineSpan<DDim, MemorySpace> int_vals) const;

/**
* Get the total number of basis functions.
Expand Down Expand Up @@ -662,7 +662,7 @@ ddc::DiscreteElement<BSplinesR, BSplinesTheta> PolarBSplines<BSplinesR, BSplines
template <class BSplinesR, class BSplinesTheta, int C>
template <class DDim, class MemorySpace>
void PolarBSplines<BSplinesR, BSplinesTheta, C>::Impl<DDim, MemorySpace>::integrals(
PolarSplineSpan<DDim> int_vals) const
PolarSplineSpan<DDim, MemorySpace> int_vals) const
{
auto r_bspl_space = ddc::discrete_space<BSplinesR>();
auto theta_bspl_space = ddc::discrete_space<BSplinesTheta>();
Expand All @@ -676,21 +676,28 @@ void PolarBSplines<BSplinesR, BSplinesTheta, C>::Impl<DDim, MemorySpace>::integr
|| int_vals.spline_coef.domain().template extent<BSplinesTheta>()
== theta_bspl_space.size());

ddc::Chunk<double, ddc::DiscreteDomain<BSplinesR>> r_integrals_alloc(
r_bspl_space.full_domain().take_first(
ddc::Chunk<double, ddc::DiscreteDomain<BSplinesR>, ddc::KokkosAllocator<double, MemorySpace>>
r_integrals_alloc(r_bspl_space.full_domain().take_first(
ddc::DiscreteVector<BSplinesR> {r_bspl_space.nbasis()}));
ddc::Chunk<double, ddc::DiscreteDomain<BSplinesTheta>> theta_integrals_alloc(
theta_bspl_space.full_domain().take_first(
ddc::Chunk<
double,
ddc::DiscreteDomain<BSplinesTheta>,
ddc::KokkosAllocator<double, MemorySpace>>
theta_integrals_alloc(theta_bspl_space.full_domain().take_first(
ddc::DiscreteVector<BSplinesTheta> {theta_bspl_space.size()}));
ddc::ChunkSpan r_integrals = r_integrals_alloc.span_view();
ddc::ChunkSpan theta_integrals = theta_integrals_alloc.span_view();

ddc::integrals(Kokkos::DefaultHostExecutionSpace(), r_integrals);
ddc::integrals(Kokkos::DefaultHostExecutionSpace(), theta_integrals);

ddc::for_each(singular_idx_range<DDim>(), [&](auto k) {
int_vals.singular_spline_coef(k) = ddc::transform_reduce(
ddc::select<BSplinesR, BSplinesTheta>(m_singular_basis_elements.domain()),
auto singular_domain
= ddc::select<BSplinesR, BSplinesTheta>(m_singular_basis_elements.domain());
ddc::ChunkSpan singular_spline_integrals = int_vals.singular_spline_coef.span_view();

ddc::for_each(singular_idx_range<DDim>(), [&](ddc::DiscreteElement<DDim> k) {
singular_spline_integrals(k) = ddc::transform_reduce(
singular_domain,
0.0,
ddc::reducer::sum<double>(),
[&](tensor_product_index_type const idx) {
Expand All @@ -702,19 +709,20 @@ void PolarBSplines<BSplinesR, BSplinesTheta, C>::Impl<DDim, MemorySpace>::integr

ddc::DiscreteDomain<BSplinesR> r_tensor_product_dom(
ddc::select<BSplinesR>(int_vals.spline_coef.domain()));

tensor_product_idx_range_type
tensor_bspline_idx_range(r_tensor_product_dom, theta_integrals.domain());
ddc::ChunkSpan spline_integrals = int_vals.spline_coef.span_view();

ddc::for_each(tensor_bspline_idx_range, [&](auto idx) {
int_vals.spline_coef(idx) = r_integrals(ddc::select<BSplinesR>(idx))
* theta_integrals(ddc::select<BSplinesTheta>(idx));
ddc::for_each(tensor_bspline_idx_range, [&](tensor_product_index_type idx) {
spline_integrals(idx) = r_integrals(ddc::select<BSplinesR>(idx))
* theta_integrals(ddc::select<BSplinesTheta>(idx));
});

if (int_vals.spline_coef.domain().template extent<BSplinesTheta>() == theta_bspl_space.size()) {
ddc::DiscreteDomain<BSplinesTheta> periodic_points(theta_integrals.domain().take_last(
ddc::DiscreteVector<BSplinesTheta> {BSplinesTheta::degree()}));
tensor_product_idx_range_type repeat_idx_range(r_tensor_product_dom, periodic_points);
ddc::for_each(repeat_idx_range, [&](auto idx) { int_vals.spline_coef(idx) = 0.0; });
ddc::for_each(repeat_idx_range, [&](tensor_product_index_type idx) {
spline_integrals(idx) = 0.0;
});
}
}
Loading

0 comments on commit 6af7eca

Please sign in to comment.