Merge branch 'ebourne-273-fft-template-params' into 'main'

Improve FFT template parametrisation

Closes #273

See merge request gysela-developpers/gyselalibxx!747

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

src/geometryRTheta/advection_field/advection_field_rp.hpp

@@ -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;
@@ -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(
@@ -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);
@@ -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
     {
@@ -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);
 

src/geometryRTheta/poisson/polarpoissonlikesolver.hpp

@@ -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};
@@ -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();
@@ -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));
 

src/geometryRTheta/time_solver/bsl_predcorr.hpp

@@ -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);
 
 

src/geometryRTheta/time_solver/bsl_predcorr_second_order_explicit.hpp

@@ -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): -------------

src/geometryRTheta/time_solver/bsl_predcorr_second_order_implicit.hpp

@@ -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): -------------

src/utils/ddc_alias_inline_functions.hpp

@@ -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;

src/utils/ddc_helper.hpp

@@ -7,6 +7,8 @@
 
 #include <ddc/ddc.hpp>
 
+#include <sll/polar_spline.hpp>
+
 #include "ddc_aliases.hpp"
 #include "directional_tag.hpp"
 #include "transpose.hpp"
@@ -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;

tests/geometryRTheta/advection_field_rp/advection_field_gtest.cpp

@@ -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;

vendor/sll/include/sll/polar_bsplines.hpp

@@ -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.
@@ -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>();
@@ -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) {
@@ -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;
+        });
     }
 }