Templated quadrature

gyselax · Nov 21, 2023 · 02e5f87 · 02e5f87
1 parent aedf5cc
commit 02e5f87
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 24 deletions.
diff --git a/src/quadrature/simpson_quadrature.hpp b/src/quadrature/simpson_quadrature.hpp
@@ -3,6 +3,7 @@
 #pragma once
 #include <ddc/ddc.hpp>
 
+#include <ddc_helper.hpp>
 
 /**
  * @brief Get the Simpson coefficients in 1D.
@@ -14,9 +15,12 @@
  *
  * @return The quadrature coefficients for the Simpson method defined on the provided domain.
  */
-template <class IDim>
-ddc::Chunk<double, ddc::DiscreteDomain<IDim>> simpson_quadrature_coefficients_1d(
-        ddc::DiscreteDomain<IDim> const& domain)
+template <class ExecSpace, class IDim>
+ddc::Chunk<
+        double,
+        ddc::DiscreteDomain<IDim>,
+        ddc::KokkosAllocator<double, typename ExecSpace::memory_space>>
+simpson_quadrature_coefficients_1d(ExecSpace const& space, ddc::DiscreteDomain<IDim> const& domain)
 {
     ddc::Chunk<double, ddc::DiscreteDomain<IDim>> coefficients(domain);
     ddc::DiscreteDomain<IDim> middle_domain
@@ -38,5 +42,16 @@ ddc::Chunk<double, ddc::DiscreteDomain<IDim>> simpson_quadrature_coefficients_1d
     }
 
 
-    return coefficients;
+    if constexpr (std::is_same_v<typename ExecSpace::memory_space, Kokkos::HostSpace>) {
+        return coefficients;
+    } else {
+        return ddc::create_mirror_and_copy(space, coefficients.span_view());
+    }
+}
+
+template <class IDim>
+ddc::Chunk<double, ddc::DiscreteDomain<IDim>, ddc::HostAllocator<double>>
+simpson_quadrature_coefficients_1d(ddc::DiscreteDomain<IDim> const& domain)
+{
+    return simpson_quadrature_coefficients_1d(Kokkos::DefaultHostExecutionSpace(), domain);
 }
diff --git a/src/utils/ddc_helper.hpp b/src/utils/ddc_helper.hpp
@@ -109,33 +109,38 @@ restrict_to_domain(
 namespace detail {
 
 /// \cond
-template <class>
-struct Device
+template <class, class>
+struct OnMemorySpace
 {
 };
 
-template <class ElementType, class SupportType, class Allocator>
-struct Device<ddc::Chunk<ElementType, SupportType, Allocator>>
+template <class NewMemorySpace, class ElementType, class SupportType, class Allocator>
+struct OnMemorySpace<NewMemorySpace, ddc::Chunk<ElementType, SupportType, Allocator>>
 {
-    using type = typename ddc::Chunk<
-            ElementType,
-            SupportType,
-            ddc::KokkosAllocator<ElementType, Kokkos::DefaultExecutionSpace::memory_space>>;
+    using type = typename ddc::
+            Chunk<ElementType, SupportType, ddc::KokkosAllocator<ElementType, NewMemorySpace>>;
 };
 
-template <class ElementType, class SupportType, class Layout, class MemorySpace>
-struct Device<ddc::ChunkSpan<ElementType, SupportType, Layout, MemorySpace>>
+template <
+        class NewMemorySpace,
+        class ElementType,
+        class SupportType,
+        class Layout,
+        class MemorySpace>
+struct OnMemorySpace<NewMemorySpace, ddc::ChunkSpan<ElementType, SupportType, Layout, MemorySpace>>
 {
-    using type = typename ddc::ChunkSpan<
-            ElementType,
-            SupportType,
-            Layout,
-            Kokkos::DefaultExecutionSpace::memory_space>;
+    using type = typename ddc::ChunkSpan<ElementType, SupportType, Layout, NewMemorySpace>;
 };
 /// \endcond
 
 } // namespace detail
 
 /// Alias template helper returning the "device" version of a `ddc::Chunk` or a `ddc::Chunkspan`
+template <class MemorySpace, class C>
+using on_memory_space_t = typename detail::OnMemorySpace<MemorySpace, C>::type;
+
+template <class C>
+using host_t = on_memory_space_t<Kokkos::HostSpace, C>;
+
 template <class C>
-using device_t = typename detail::Device<C>::type;
+using device_t = on_memory_space_t<Kokkos::DefaultExecutionSpace::memory_space, C>;
diff --git a/tests/geometryXVx/quadrature.cpp b/tests/geometryXVx/quadrature.cpp
@@ -100,14 +100,13 @@ double compute_error(int n_elems, Method meth)
     IDomainY const gridy(GrevillePointsY::get_domain());
 
     SplineYBuilder const builder_y(gridy);
-    std::function<ddc::Chunk<double, IDomainY>(IDomainY const&)> func;
+    DFieldY quadrature_coeffs;
     switch (meth) {
     case Method::TRAPEZ:
-        func = trapezoid_quadrature_coefficients<IDimY>;
+        quadrature_coeffs = trapezoid_quadrature_coefficients(gridy);
     case Method::SIMPSON:
-        func = simpson_quadrature_coefficients_1d<IDimY>;
+        quadrature_coeffs = simpson_quadrature_coefficients_1d(gridy);
     }
-    DFieldY const quadrature_coeffs = func(gridy);
     Quadrature<IDimY> const integrate(quadrature_coeffs);
 
     DFieldY values(gridy);