Allow to access attributes of SpeciesInformation on the device

gyselax · Feb 16, 2024 · a3be632 · a3be632
1 parent d79fbfe
commit a3be632
Show file tree

Hide file tree

Showing 6 changed files with 124 additions and 101 deletions.
diff --git a/src/geometryXVx/rhs/collisions_utils.cpp b/src/geometryXVx/rhs/collisions_utils.cpp
@@ -40,23 +40,12 @@ void compute_nustar_profile(device_t<DSpanSpX> nustar_profile, double nustar0)
 {
     double const Lx = ddcHelper::total_interval_length(ddc::get_domain<IDimX>(nustar_profile));
 
-    auto masses_alloc = create_mirror_view_and_copy(
-            Kokkos::DefaultExecutionSpace(),
-            ddc::host_discrete_space<SpeciesInformation>().masses());
-    auto masses(masses_alloc.span_view());
-
-    auto charges_alloc = create_mirror_view_and_copy(
-            Kokkos::DefaultExecutionSpace(),
-            ddc::host_discrete_space<SpeciesInformation>().charges());
-    auto charges(charges_alloc.span_view());
-
     ddc::for_each(
             ddc::policies::parallel_device,
             nustar_profile.domain(),
             KOKKOS_LAMBDA(IndexSpX const ispx) {
-                double const coeff
-                        = Kokkos::sqrt(masses(ielec()) / masses(ddc::select<IDimSp>(ispx)))
-                          * Kokkos::pow(charges(ddc::select<IDimSp>(ispx)), 4) / Lx;
+                double const coeff = Kokkos::sqrt(mass(ielec()) / mass(ddc::select<IDimSp>(ispx)))
+                                     * Kokkos::pow(charge(ddc::select<IDimSp>(ispx)), 4) / Lx;
                 nustar_profile(ispx) = coeff * nustar0;
             });
 }

diff --git a/src/geometryXVx/rhs/krook_source_adaptive.cpp b/src/geometryXVx/rhs/krook_source_adaptive.cpp
@@ -30,24 +30,25 @@ KrookSourceAdaptive::KrookSourceAdaptive(
     , m_density(density)
     , m_temperature(temperature)
     , m_ftarget(gridvx)
+    , m_mask(gridx)
 {
     // mask that defines the region where the operator is active
+    DFieldX mask_host(gridx);
     switch (m_type) {
     case RhsType::Source:
         // the mask equals one in the interval [x_left, x_right]
-        m_mask = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Normal, false);
+        mask_host = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Normal, false);
         break;
     case RhsType::Sink:
         // the mask equals zero in the center of the plasma
-        m_mask = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Inverted, false);
+        mask_host = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Inverted, false);
         break;
     }
+    ddc::deepcopy(m_mask.span_view(), mask_host);
 
     // target distribution function
-    device_t<DFieldVx> ftarget_device(gridvx);
-    MaxwellianEquilibrium::compute_maxwellian(ftarget_device, m_density, m_temperature, 0.);
-    auto ftarget = ddc::create_mirror_view_and_copy(ftarget_device.span_view());
-    ddc::deepcopy(m_ftarget, ftarget);
+    MaxwellianEquilibrium::compute_maxwellian(m_ftarget.span_view(), m_density, m_temperature, 0.);
+    auto ftarget_host = ddc::create_mirror_view_and_copy(m_ftarget.span_view());
 
     switch (m_type) {
     case RhsType::Source:
@@ -56,17 +57,17 @@ KrookSourceAdaptive::KrookSourceAdaptive(
         ddc::expose_to_pdi("krook_source_adaptive_amplitude", m_amplitude);
         ddc::expose_to_pdi("krook_source_adaptive_density", m_density);
         ddc::expose_to_pdi("krook_source_adaptive_temperature", m_temperature);
-        ddc::expose_to_pdi("krook_source_adaptive_ftarget", m_ftarget);
-        ddc::expose_to_pdi("krook_source_adaptive_mask", m_mask);
+        ddc::expose_to_pdi("krook_source_adaptive_ftarget", ftarget_host);
+        ddc::expose_to_pdi("krook_source_adaptive_mask", mask_host);
         break;
     case RhsType::Sink:
         ddc::expose_to_pdi("krook_sink_adaptive_extent", m_extent);
         ddc::expose_to_pdi("krook_sink_adaptive_stiffness", m_stiffness);
         ddc::expose_to_pdi("krook_sink_adaptive_amplitude", m_amplitude);
         ddc::expose_to_pdi("krook_sink_adaptive_density", m_density);
         ddc::expose_to_pdi("krook_sink_adaptive_temperature", m_temperature);
-        ddc::expose_to_pdi("krook_sink_adaptive_ftarget", m_ftarget);
-        ddc::expose_to_pdi("krook_sink_adaptive_mask", m_mask);
+        ddc::expose_to_pdi("krook_sink_adaptive_ftarget", ftarget_host);
+        ddc::expose_to_pdi("krook_sink_adaptive_mask", mask_host);
         break;
     }
 }
@@ -86,11 +87,11 @@ void KrookSourceAdaptive::get_amplitudes(
     }
     IndexSp iion(iion_opt.value());
     IDomainVx const gridvx = allfdistribu.domain<IDimVx>();
-    DFieldVx const quadrature_coeffs(trapezoid_quadrature_coefficients(gridvx));
+    DFieldVx const quadrature_coeffs_host(trapezoid_quadrature_coefficients(gridvx));
     auto quadrature_coeffs_alloc = ddc::create_mirror_view_and_copy(
             Kokkos::DefaultExecutionSpace(),
-            quadrature_coeffs.span_view());
-    auto quadrature_coeffs_device = quadrature_coeffs_alloc.span_view();
+            quadrature_coeffs_host.span_view());
+    auto quadrature_coeffs = quadrature_coeffs_alloc.span_view();
 
     auto const& amplitude = m_amplitude;
     auto const& density = m_density;
@@ -103,71 +104,50 @@ void KrookSourceAdaptive::get_amplitudes(
                 double density_electron = 0.;
 
                 for (IndexVx ivx : allfdistribu.domain<IDimVx>()) {
-                    density_ion += quadrature_coeffs_device(ivx) * allfdistribu(iion, ix, ivx);
-                    density_electron
-                            += quadrature_coeffs_device(ivx) * allfdistribu(ielec(), ix, ivx);
+                    density_ion += quadrature_coeffs(ivx) * allfdistribu(iion, ix, ivx);
+                    density_electron += quadrature_coeffs(ivx) * allfdistribu(ielec(), ix, ivx);
                 }
                 amplitudes(IndexSpX(ielec(), ix))
                         = amplitude * (density_ion - density) / (density_electron - density);
             });
 }
 
 void KrookSourceAdaptive::get_derivative(
-        DSpanSpXVx df,
-        DViewSpXVx allfdistribu,
-        DViewSpXVx allfdistribu_start) const
+        device_t<DSpanSpXVx> df,
+        device_t<DViewSpXVx> allfdistribu,
+        device_t<DViewSpXVx> allfdistribu_start) const
 {
-    auto df_device_alloc
-            = ddc::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), df.span_view());
-    auto df_device = df_device_alloc.span_view();
+    IDomainSpX grid_sp_x(allfdistribu.domain<IDimSp, IDimX>());
 
-    auto allfdistribu_device_alloc = ddc::
-            create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), allfdistribu.span_view());
-    auto allfdistribu_device = allfdistribu_device_alloc.span_view();
+    device_t<DFieldSpX> amplitudes_alloc(grid_sp_x);
+    auto amplitudes = amplitudes_alloc.span_view();
+    get_amplitudes(amplitudes, allfdistribu);
 
-    auto allfdistribu_start_device_alloc = ddc::create_mirror_view_and_copy(
-            Kokkos::DefaultExecutionSpace(),
-            allfdistribu_start.span_view());
-    auto allfdistribu_start_device = allfdistribu_start_device_alloc.span_view();
-
-    device_t<DFieldSpX> amplitudes_device_alloc(
-            ddc::get_domain<IDimSp, IDimX>(allfdistribu_device));
-    auto amplitudes_device = amplitudes_device_alloc.span_view();
-    get_amplitudes(amplitudes_device, allfdistribu_device);
-
-    auto mask_device_alloc
-            = ddc::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), m_mask.span_view());
-    auto mask_device = mask_device_alloc.span_view();
-
-    auto ftarget_device_alloc = ddc::
-            create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), m_ftarget.span_view());
-    auto ftarget_device = ftarget_device_alloc.span_view();
+    auto ftarget(m_ftarget.span_view());
+    auto mask(m_mask.span_view());
 
     ddc::for_each(
             ddc::policies::parallel_device,
-            allfdistribu_device.domain(),
+            allfdistribu.domain(),
             KOKKOS_LAMBDA(IndexSpXVx const ispxvx) {
                 IndexSp isp(ddc::select<IDimSp>(ispxvx));
                 IndexX ix(ddc::select<IDimX>(ispxvx));
                 IndexVx ivx(ddc::select<IDimVx>(ispxvx));
-                df_device(ispxvx) = -mask_device(ix) * amplitudes_device(isp, ix)
-                                    * (allfdistribu_start_device(ispxvx) - ftarget_device(ivx));
+                df(ispxvx) = -mask(ix) * amplitudes(isp, ix)
+                             * (allfdistribu_start(ispxvx) - ftarget(ivx));
             });
-    ddc::deepcopy(df, df_device);
 }
 
 device_t<DSpanSpXVx> KrookSourceAdaptive::operator()(
-        device_t<DSpanSpXVx> const allfdistribu_device,
+        device_t<DSpanSpXVx> const allfdistribu,
         double const dt) const
 {
     Kokkos::Profiling::pushRegion("KrookSource");
-    auto allfdistribu_alloc = ddc::create_mirror_view_and_copy(allfdistribu_device);
-    ddc::ChunkSpan allfdistribu = allfdistribu_alloc.span_view();
-    RK2<DFieldSpXVx> timestepper(allfdistribu.domain());
-    timestepper.update(allfdistribu, dt, [&](DSpanSpXVx df, DViewSpXVx f) {
+    RK2<device_t<DFieldSpXVx>> timestepper(allfdistribu.domain());
+
+    timestepper.update(allfdistribu, dt, [&](device_t<DSpanSpXVx> df, device_t<DViewSpXVx> f) {
         get_derivative(df, f, allfdistribu);
     });
-    ddc::deepcopy(allfdistribu_device, allfdistribu);
     Kokkos::Profiling::popRegion();
-    return allfdistribu_device;
+    return allfdistribu;
 }
diff --git a/src/geometryXVx/rhs/krook_source_adaptive.hpp b/src/geometryXVx/rhs/krook_source_adaptive.hpp
@@ -30,8 +30,8 @@ class KrookSourceAdaptive : public IRightHandSide
     double m_amplitude;
     double m_density;
     double m_temperature;
-    DFieldX m_mask;
-    DFieldVx m_ftarget;
+    device_t<DFieldX> m_mask;
+    device_t<DFieldVx> m_ftarget;
 
 public:
     /**
@@ -100,5 +100,6 @@ class KrookSourceAdaptive : public IRightHandSide
      * @param[in] f The distribution function.
      * @param[in] f0 An optional parameter.
      */
-    void get_derivative(DSpanSpXVx df, DViewSpXVx f, DViewSpXVx f0) const;
+    void get_derivative(device_t<DSpanSpXVx> df, device_t<DViewSpXVx> f, device_t<DViewSpXVx> f0)
+            const;
 };
diff --git a/src/geometryXVx/rhs/krook_source_constant.cpp b/src/geometryXVx/rhs/krook_source_constant.cpp
@@ -24,23 +24,25 @@ KrookSourceConstant::KrookSourceConstant(
     , m_density(density)
     , m_temperature(temperature)
     , m_ftarget(gridvx)
+    , m_mask(gridx)
 {
     // mask that defines the region where the operator is active
+    DFieldX mask_host(gridx);
     switch (m_type) {
     case RhsType::Source:
         // the mask equals one in the interval [x_left, x_right]
-        m_mask = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Normal, false);
+        mask_host = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Normal, false);
         break;
     case RhsType::Sink:
         // the mask equals zero in the center of the plasma
-        m_mask = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Inverted, false);
+        mask_host = mask_tanh(gridx, m_extent, m_stiffness, MaskType::Inverted, false);
         break;
     }
+    ddc::deepcopy(m_mask.span_view(), mask_host);
+
     // target distribution function
-    device_t<DFieldVx> ftarget_device(gridvx);
-    MaxwellianEquilibrium::compute_maxwellian(ftarget_device, m_density, m_temperature, 0.);
-    auto ftarget = ddc::create_mirror_view_and_copy(ftarget_device.span_view());
-    ddc::deepcopy(m_ftarget, ftarget);
+    MaxwellianEquilibrium::compute_maxwellian(m_ftarget.span_view(), m_density, m_temperature, 0.);
+    auto ftarget_host = ddc::create_mirror_view_and_copy(m_ftarget.span_view());
 
     switch (m_type) {
     case RhsType::Source:
@@ -49,17 +51,17 @@ KrookSourceConstant::KrookSourceConstant(
         ddc::expose_to_pdi("krook_source_constant_amplitude", m_amplitude);
         ddc::expose_to_pdi("krook_source_constant_density", m_density);
         ddc::expose_to_pdi("krook_source_constant_temperature", m_temperature);
-        ddc::expose_to_pdi("krook_source_constant_ftarget", m_ftarget);
-        ddc::expose_to_pdi("krook_source_constant_mask", m_mask);
+        ddc::expose_to_pdi("krook_source_constant_ftarget", ftarget_host);
+        ddc::expose_to_pdi("krook_source_constant_mask", mask_host);
         break;
     case RhsType::Sink:
         ddc::expose_to_pdi("krook_sink_constant_extent", m_extent);
         ddc::expose_to_pdi("krook_sink_constant_stiffness", m_stiffness);
         ddc::expose_to_pdi("krook_sink_constant_amplitude", m_amplitude);
         ddc::expose_to_pdi("krook_sink_constant_density", m_density);
         ddc::expose_to_pdi("krook_sink_constant_temperature", m_temperature);
-        ddc::expose_to_pdi("krook_sink_constant_ftarget", m_ftarget);
-        ddc::expose_to_pdi("krook_sink_constant_mask", m_mask);
+        ddc::expose_to_pdi("krook_sink_constant_ftarget", ftarget_host);
+        ddc::expose_to_pdi("krook_sink_constant_mask", mask_host);
         break;
     }
 }
@@ -69,26 +71,20 @@ device_t<DSpanSpXVx> KrookSourceConstant::operator()(
         double const dt) const
 {
     Kokkos::Profiling::pushRegion("KrookSource");
-    auto ftarget_device_alloc = ddc::
-            create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), m_ftarget.span_view());
-    auto ftarget_device = ftarget_device_alloc.span_view();
-
-    auto mask_device_alloc
-            = ddc::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), m_mask.span_view());
-    auto mask_device = mask_device_alloc.span_view();
 
+    auto ftarget = m_ftarget.span_view();
+    auto mask = m_mask.span_view();
     auto const& amplitude = m_amplitude;
 
     ddc::for_each(
             ddc::policies::parallel_device,
             allfdistribu.domain(),
             KOKKOS_LAMBDA(IndexSpXVx const ispxvx) {
                 allfdistribu(ispxvx)
-                        = ftarget_device(ddc::select<IDimVx>(ispxvx))
-                          + (allfdistribu(ispxvx) - ftarget_device(ddc::select<IDimVx>(ispxvx)))
+                        = ftarget(ddc::select<IDimVx>(ispxvx))
+                          + (allfdistribu(ispxvx) - ftarget(ddc::select<IDimVx>(ispxvx)))
                                     * Kokkos::exp(
-                                            -amplitude * mask_device(ddc::select<IDimX>(ispxvx))
-                                            * dt);
+                                            -amplitude * mask(ddc::select<IDimX>(ispxvx)) * dt);
             });
 
     Kokkos::Profiling::popRegion();

diff --git a/src/geometryXVx/rhs/krook_source_constant.hpp b/src/geometryXVx/rhs/krook_source_constant.hpp
@@ -28,8 +28,8 @@ class KrookSourceConstant : public IRightHandSide
     double m_amplitude;
     double m_density;
     double m_temperature;
-    DFieldX m_mask;
-    DFieldVx m_ftarget;
+    device_t<DFieldX> m_mask;
+    device_t<DFieldVx> m_ftarget;
 
 public:
     /**