Partial GPU port of FEM PoissonXVx

PoissonXVx is partially ported:

* Fem(Non)PeriodicPoissonSolver::operator() involves GPU data.
* solve_matrix_system() takes GPU data as arguments but deepcopy them on CPU (still solved using LAPACK).
* Splines evaluation and derivation (electric field computation) are performed on GPU.
* Old ElectricField operator was already unused, it is deprecated and removed.

docs/CODING_STANDARD.md

@@ -33,6 +33,7 @@
 * we don't rely on case to distinguish between variables
 * there are two types of DDC objects representing a multidimensional array : `ddc::Chunk` (which possesses the data) and `ddc::ChunkSpan` (which does not own the data but can be captured by `KOKKOS_LAMBDA`). We suffix `Chunk` with `_alloc` if both variables are needed locally.
 * if a variable is mirrored between host (CPU) and device (GPU) memories, the variable representing data on host is `_host` suffixed
+* capturing classes members through `KOKKOS_LAMBDA` or `KOKKOS_CLASS_LAMBDA` may be complicated, we often need to copy-by-reference the member to a local variable, which must be `_proxy` suffixed
 
 ## Style
 * We use the style specified by the `.clang-format` file using clang-format 10

src/geometryXVx/geometry/geometry.hpp

@@ -144,17 +144,17 @@ using SplineVxEvaluator = ddc::SplineEvaluator<
         IDimX,
         IDimVx>;
 using SplineXBuilder_1d = ddc::SplineBuilder<
-        Kokkos::DefaultHostExecutionSpace,
-        Kokkos::DefaultHostExecutionSpace::memory_space,
+        Kokkos::DefaultExecutionSpace,
+        Kokkos::DefaultExecutionSpace::memory_space,
         BSplinesX,
         IDimX,
         SplineXBoundary,
         SplineXBoundary,
         ddc::SplineSolver::GINKGO,
         IDimX>;
 using SplineXEvaluator_1d = ddc::SplineEvaluator<
-        Kokkos::DefaultHostExecutionSpace,
-        Kokkos::DefaultHostExecutionSpace::memory_space,
+        Kokkos::DefaultExecutionSpace,
+        Kokkos::DefaultExecutionSpace::memory_space,
         BSplinesX,
         IDimX,
 #ifdef PERIODIC_RDIMX
@@ -297,6 +297,9 @@ using SpanSpXVx = device_t<ddc::ChunkSpan<ElementType, IDomainSpXVx>>;
 template <class ElementType>
 using SpanSpVx = device_t<ddc::ChunkSpan<ElementType, IDomainSpVx>>;
 
+template <class ElementType>
+using BSSpanX = device_t<ddc::ChunkSpan<ElementType, BSDomainX>>;
+
 
 
 template <class DomainType>
@@ -315,6 +318,7 @@ using DSpanSpXVx = SpanSpXVx<double>;
 
 using DSpanSpVx = SpanSpVx<double>;
 
+using DBSSpanX = BSSpanX<double>;
 
 
 template <class ElementType>

src/geometryXVx/poisson/CMakeLists.txt

@@ -4,7 +4,6 @@ foreach(GEOMETRY_VARIANT IN LISTS GEOMETRY_XVx_VARIANTS_LIST)
 
 add_library("poisson_${GEOMETRY_VARIANT}" STATIC
     chargedensitycalculator.cpp
-    electricfield.cpp
     nullpoissonsolver.cpp
 )
 

src/geometryXVx/poisson/femnonperiodicpoissonsolver.cpp

@@ -30,7 +30,7 @@ NUBSplineXEvaluator_1d jit_build_nubsplinesx(SplineXEvaluator_1d const& spline_x
     // Boundary values are never evaluated
     return NUBSplineXEvaluator_1d(
             ddc::replace_dim_of<typename SplineXEvaluator_1d::bsplines_type, NUBSplinesX>(
-                    spline_x_evaluator.spline_domain(),
+                    spline_x_evaluator.bsplines_domain(),
                     ddc::discrete_space<NUBSplinesX>().full_domain()),
             ddc::NullExtrapolationRule(),
             ddc::NullExtrapolationRule());
@@ -51,7 +51,7 @@ FemNonPeriodicPoissonSolver::FemNonPeriodicPoissonSolver(
     , m_compute_rho(compute_rho)
     , m_nbasis(ddc::discrete_space<BSplinesX>().nbasis())
     , m_ncells(ddc::discrete_space<BSplinesX>().ncells())
-    , m_quad_coef(ddc::DiscreteDomain<QMeshX>(
+    , m_quad_coef_alloc(ddc::DiscreteDomain<QMeshX>(
               ddc::DiscreteElement<QMeshX>(0),
               ddc::DiscreteVector<QMeshX>(s_npts_gauss * m_ncells)))
 {
@@ -67,10 +67,12 @@ FemNonPeriodicPoissonSolver::FemNonPeriodicPoissonSolver(
 
     // Calculate the integration coefficients
     GaussLegendre<QDimX> const gl(s_npts_gauss);
-    std::vector<ddc::Coordinate<QDimX>> eval_pts_data(m_quad_coef.domain().size());
+    std::vector<ddc::Coordinate<QDimX>> eval_pts_data(m_quad_coef_alloc.domain().size());
     ddc::ChunkSpan<ddc::Coordinate<QDimX>, ddc::DiscreteDomain<QMeshX>> const
-            eval_pts(eval_pts_data.data(), m_quad_coef.domain());
-    gl.compute_points_and_weights_on_mesh(eval_pts, m_quad_coef.span_view(), knots.span_cview());
+            eval_pts(eval_pts_data.data(), m_quad_coef_alloc.domain());
+    auto quad_coef_host = ddc::create_mirror_and_copy(m_quad_coef_alloc.span_view());
+    gl.compute_points_and_weights_on_mesh(eval_pts, quad_coef_host.span_view(), knots.span_cview());
+    ddc::deepcopy(m_quad_coef_alloc, quad_coef_host);
 
     ddc::init_discrete_space<QMeshX>(eval_pts_data);
 
@@ -94,9 +96,11 @@ void FemNonPeriodicPoissonSolver::build_matrix()
     m_fem_matrix = Matrix::
             make_new_banded(matrix_size, n_lower_diags, n_lower_diags, positive_definite_symmetric);
 
+    auto quad_coef_host = ddc::create_mirror_and_copy(m_quad_coef_alloc.span_cview());
+
     // Fill the banded part of the matrix
     std::array<double, s_degree + 1> derivs;
-    ddc::for_each(m_quad_coef.domain(), [&](ddc::DiscreteElement<QMeshX> const ix) {
+    ddc::for_each(m_quad_coef_alloc.domain(), [&](ddc::DiscreteElement<QMeshX> const ix) {
         ddc::Coordinate<RDimX> const coord = coord_from_quad_point(ddc::coordinate(ix));
         ddc::DiscreteElement<NUBSplinesX> const jmin
                 = ddc::discrete_space<NUBSplinesX>().eval_deriv(derivs, coord);
@@ -108,7 +112,7 @@ void FemNonPeriodicPoissonSolver::build_matrix()
                 if (j_idx != -1 && j_idx != matrix_size && k_idx != -1 && k_idx != matrix_size) {
                     double a_jk = m_fem_matrix->get_element(j_idx, k_idx);
                     // Update element
-                    a_jk += derivs[j] * derivs[k] * m_quad_coef(ix);
+                    a_jk += derivs[j] * derivs[k] * quad_coef_host(ix);
 
                     m_fem_matrix->set_element(j_idx, k_idx, a_jk);
                 }
@@ -125,36 +129,46 @@ void FemNonPeriodicPoissonSolver::build_matrix()
 //               Solve the Poisson equation
 //---------------------------------------------------------------------------
 void FemNonPeriodicPoissonSolver::solve_matrix_system(
-        ddc::ChunkSpan<double, NUBSDomainX> const phi_spline_coef,
-        ddc::ChunkSpan<double, BSDomainX> const rho_spline_coef) const
+        DNUBSSpanX const phi_spline_coef,
+        DBSViewX const rho_spline_coef) const
 {
-    std::array<double, s_degree + 1> values;
-
-    for (int i(0); i < m_nbasis; ++i) {
-        phi_spline_coef(ddc::DiscreteElement<NUBSplinesX>(i)) = 0.0;
-    }
+    ddc::fill(phi_spline_coef, 0.0);
 
     int const rhs_size = m_nbasis - 2;
-    DSpan1D const phi_rhs(phi_spline_coef.data_handle() + 1, rhs_size);
+    int const nbasis_proxy = m_nbasis;
+    SplineXEvaluator_1d spline_x_evaluator_proxy = m_spline_x_evaluator;
+    ddc::ChunkSpan quad_coef = m_quad_coef_alloc.span_view();
 
     // Fill phi_rhs(i) with \int rho(x) b_i(x) dx
     // Rk: phi_rhs no longer contains spline coefficients, but is the
     //     RHS of the matrix equation
-    ddc::for_each(m_quad_coef.domain(), [&](ddc::DiscreteElement<QMeshX> const ix) {
-        ddc::Coordinate<RDimX> const coord = coord_from_quad_point(ddc::coordinate(ix));
-        ddc::DiscreteElement<BSplinesX> const jmin
-                = ddc::discrete_space<BSplinesX>().eval_basis(values, coord);
-        double const rho_val = m_spline_x_evaluator(coord, rho_spline_coef.span_cview());
-        for (int j = 0; j < s_degree + 1; ++j) {
-            int const j_idx = (jmin.uid() + j) % m_nbasis - 1;
-            if (j_idx != -1 && j_idx != rhs_size) {
-                phi_rhs(j_idx) += rho_val * values[j] * m_quad_coef(ix);
-            }
-        }
-    });
+    ddc::for_each(
+            ddc::policies::parallel_device,
+            m_quad_coef_alloc.domain(),
+            KOKKOS_LAMBDA(ddc::DiscreteElement<QMeshX> const ix) {
+                ddc::Coordinate<RDimX> const coord = coord_from_quad_point(ddc::coordinate(ix));
+                std::array<double, s_degree + 1> values;
+                ddc::DiscreteElement<BSplinesX> const jmin
+                        = ddc::discrete_space<BSplinesX>().eval_basis(values, coord);
+                double const rho_val
+                        = spline_x_evaluator_proxy(coord, rho_spline_coef.span_cview());
+                for (int j = 0; j < s_degree + 1; ++j) {
+                    int const j_idx = (jmin.uid() + j) % nbasis_proxy;
+                    if (j_idx != 0 && j_idx != rhs_size + 1) {
+                        Kokkos::atomic_fetch_add(
+                                &phi_spline_coef(ddc::DiscreteElement<NUBSplinesX>(j_idx)),
+                                rho_val * values[j] * quad_coef(ix));
+                    }
+                }
+            });
+
+    auto phi_spline_coef_host = ddc::create_mirror_and_copy(phi_spline_coef);
+    DSpan1D const phi_rhs_host(phi_spline_coef_host.data_handle() + 1, rhs_size);
 
     // Solve the matrix equation to find the spline coefficients of phi
-    m_fem_matrix->solve_inplace(phi_rhs);
+    m_fem_matrix->solve_inplace(phi_rhs_host);
+
+    ddc::deepcopy(phi_spline_coef, phi_spline_coef_host);
 }
 
 
@@ -180,32 +194,32 @@ void FemNonPeriodicPoissonSolver::operator()(
         DViewSpXVx allfdistribu) const
 {
     Kokkos::Profiling::pushRegion("PoissonSolver");
-    auto electrostatic_potential_host = ddc::create_mirror_and_copy(electrostatic_potential);
-    auto electric_field_host = ddc::create_mirror_and_copy(electric_field);
     assert(electrostatic_potential.domain() == ddc::get_domain<IDimX>(allfdistribu));
     IDomainX const dom_x = electrostatic_potential.domain();
 
     // Compute the RHS of the Poisson equation
-    host_t<DFieldX> rho_host(dom_x);
     DFieldX rho(dom_x);
     m_compute_rho(rho, allfdistribu);
-    ddc::deepcopy(rho_host, rho);
 
     //
-    ddc::Chunk<double, BSDomainX> rho_spline_coef(m_spline_x_builder.spline_domain());
-    m_spline_x_builder(rho_spline_coef.span_view(), rho_host.span_cview());
-    ddc::Chunk<double, NUBSDomainX> phi_spline_coef(
-            ddc::discrete_space<NUBSplinesX>().full_domain());
-    solve_matrix_system(phi_spline_coef, rho_spline_coef);
+    ddc::Chunk rho_spline_coef(m_spline_x_builder.spline_domain(), ddc::DeviceAllocator<double>());
+    m_spline_x_builder(rho_spline_coef.span_view(), rho.span_cview());
+    ddc::Chunk phi_spline_coef(
+            ddc::discrete_space<NUBSplinesX>().full_domain(),
+            ddc::DeviceAllocator<double>());
+    solve_matrix_system(phi_spline_coef, rho_spline_coef.span_cview());
 
     //
-    ddc::for_each(dom_x, [&](IndexX const ix) {
-        electrostatic_potential_host(ix)
-                = m_spline_x_nu_evaluator(ddc::coordinate(ix), phi_spline_coef.span_cview());
-        electric_field_host(ix)
-                = -m_spline_x_nu_evaluator.deriv(ddc::coordinate(ix), phi_spline_coef.span_cview());
-    });
-    ddc::deepcopy(electrostatic_potential, electrostatic_potential_host);
-    ddc::deepcopy(electric_field, electric_field_host);
+    NUBSplineXEvaluator_1d spline_x_nu_evaluator_proxy = m_spline_x_nu_evaluator;
+    ddc::ChunkSpan phi_spline_coef_view = phi_spline_coef.span_cview();
+    ddc::for_each(
+            ddc::policies::parallel_device,
+            dom_x,
+            KOKKOS_LAMBDA(IndexX const ix) {
+                electrostatic_potential(ix)
+                        = spline_x_nu_evaluator_proxy(ddc::coordinate(ix), phi_spline_coef_view);
+                electric_field(ix) = -spline_x_nu_evaluator_proxy
+                                              .deriv(ddc::coordinate(ix), phi_spline_coef_view);
+            });
     Kokkos::Profiling::popRegion();
 }

src/geometryXVx/poisson/femnonperiodicpoissonsolver.hpp

@@ -15,8 +15,8 @@ using NUBSDomainX = ddc::DiscreteDomain<NUBSplinesX>;
 using UBSplinesX = ddc::UniformBSplines<RDimX, BSDegreeX>;
 
 using NUBSplineXEvaluator_1d = ddc::SplineEvaluator<
-        Kokkos::DefaultHostExecutionSpace,
-        Kokkos::DefaultHostExecutionSpace::memory_space,
+        Kokkos::DefaultExecutionSpace,
+        Kokkos::DefaultExecutionSpace::memory_space,
         NUBSplinesX,
         IDimX,
         ddc::NullExtrapolationRule,
@@ -46,6 +46,8 @@ class FemNonPeriodicPoissonSolver : public IPoissonSolver
 
 private:
     using QMeshX = ddc::NonUniformPointSampling<QDimX>;
+    using DQFieldX = device_t<ddc::Chunk<double, ddc::DiscreteDomain<QMeshX>>>;
+    using DNUBSSpanX = device_t<ddc::ChunkSpan<double, NUBSDomainX>>;
 
 private:
     // Spline degree in x direction
@@ -69,17 +71,19 @@ class FemNonPeriodicPoissonSolver : public IPoissonSolver
     // Number of cells in x direction
     int m_ncells;
 
-    ddc::Chunk<double, ddc::DiscreteDomain<QMeshX>> m_quad_coef;
+    DQFieldX m_quad_coef_alloc;
 
     std::unique_ptr<Matrix> m_fem_matrix;
 
 private:
-    static ddc::Coordinate<QDimX> quad_point_from_coord(ddc::Coordinate<RDimX> const& coord)
+    static KOKKOS_FUNCTION ddc::Coordinate<QDimX> quad_point_from_coord(
+            ddc::Coordinate<RDimX> const& coord)
     {
         return ddc::Coordinate<QDimX>(ddc::get<RDimX>(coord));
     }
 
-    static ddc::Coordinate<RDimX> coord_from_quad_point(ddc::Coordinate<QDimX> const& coord)
+    static KOKKOS_FUNCTION ddc::Coordinate<RDimX> coord_from_quad_point(
+            ddc::Coordinate<QDimX> const& coord)
     {
         return ddc::Coordinate<RDimX>(ddc::get<QDimX>(coord));
     }
@@ -110,7 +114,12 @@ class FemNonPeriodicPoissonSolver : public IPoissonSolver
 private:
     void build_matrix();
 
-    void solve_matrix_system(
-            ddc::ChunkSpan<double, NUBSDomainX> phi_spline_coef,
-            ddc::ChunkSpan<double, BSDomainX> rho_spline_coef) const;
+public:
+    /**
+     * [SHOULD BE PRIVATE (Kokkos limitation)]
+     *
+     * @param[out] phi_spline_coef
+     * @param[in] rho_spline_coef
+     */
+    void solve_matrix_system(DNUBSSpanX phi_spline_coef, DBSViewX rho_spline_coef) const;
 };