Merge pull request #19 from CExA-project/reuse-plan

Reuse plan
kokkos · Jan 8, 2024 · f1ee91f · f1ee91f
2 parents 66b979a + 7270941
commit f1ee91f
Show file tree

Hide file tree

Showing 8 changed files with 945 additions and 36 deletions.
diff --git a/.github/workflows/cmake.yml b/.github/workflows/cmake.yml
@@ -82,11 +82,15 @@ jobs:
     runs-on: ubuntu-latest
 
     env:
-      backends: HIP HIP_HOST_DEVICE
+      #backends: HIP HIP_HOST_DEVICE
       architecture: VEGA90A
       CMAKE_CXX_COMPILER: hipcc
       container: amd_env
 
+    strategy:
+      matrix:
+        backend: [ {name: HIP, option: ""}, {name: HIP_HOST_DEVICE, option: "-DKokkosFFT_ENABLE_HOST_AND_DEVICE=ON"} ]
+
     steps:
     - name: Free Disk Space (Ubuntu)
       uses: jlumbroso/[email protected]
@@ -107,22 +111,25 @@ jobs:
       # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make.
       # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type
       run: |
-        docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_HIP \
+        docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_${{matrix.backend.name}} \
             -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} -DCMAKE_CXX_COMPILER=${{env.CMAKE_CXX_COMPILER}} \
-            -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON
+            -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON ${{matrix.backend.option}}
 
-    - name: Configure CMake for HIP backend with HOST and DEVICE option
-      # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make.
-      # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type
-      run: |
-        docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_HIP_HOST_DEVICE \
-            -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} -DCMAKE_CXX_COMPILER=${{env.CMAKE_CXX_COMPILER}} \
-            -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON \
-            -DKokkosFFT_ENABLE_HOST_AND_DEVICE=ON
+    #- name: Configure CMake for HIP backend with HOST and DEVICE option
+    #  # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make.
+    #  # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type
+    #  run: |
+    #    docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_HIP_HOST_DEVICE \
+    #        -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} -DCMAKE_CXX_COMPILER=${{env.CMAKE_CXX_COMPILER}} \
+    #        -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON \
+    #        -DKokkosFFT_ENABLE_HOST_AND_DEVICE=ON
 
     - name: Build
       # Build your program with the given configuration
       run: |
-        for backend in ${{ env.backends }}; do
-          docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake --build build_${backend} --config ${{env.BUILD_TYPE}} -j 2
-        done
+        docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake --build build_${{matrix.backend.name}} --config ${{env.BUILD_TYPE}} -j 2
+
+      #run: |
+      #  for backend in ${{ env.backends }}; do
+      #    docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake --build build_${backend} --config ${{env.BUILD_TYPE}} -j 2
+      #  done
diff --git a/examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp b/examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp
@@ -0,0 +1,49 @@
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Complex.hpp>
+#include <Kokkos_Random.hpp>
+#include <KokkosFFT.hpp>
+
+using execution_space = Kokkos::DefaultExecutionSpace;
+template <typename T> using View1D = Kokkos::View<T*, execution_space>;
+
+int main( int argc, char* argv[] ) {
+  Kokkos::initialize( argc, argv );
+  {
+    constexpr int n0 = 128;
+    const Kokkos::complex<double> I(1.0, 1.0);
+
+    // 1D C2C FFT (Forward and Backward)
+    View1D<Kokkos::complex<double> > xc2c("xc2c", n0);
+    View1D<Kokkos::complex<double> > xc2c_hat("xc2c_hat", n0);
+    View1D<Kokkos::complex<double> > xc2c_inv("xc2c_inv", n0);
+
+    Kokkos::Random_XorShift64_Pool<> random_pool(12345);
+    Kokkos::fill_random(xc2c, random_pool, I);
+
+    int axis = -1;
+    KokkosFFT::Impl::Plan fft_plan(execution_space(), xc2c, xc2c_hat, KokkosFFT::Impl::Direction::Forward, axis);
+    KokkosFFT::fft(execution_space(), xc2c, xc2c_hat, fft_plan);
+
+    KokkosFFT::Impl::Plan ifft_plan(execution_space(), xc2c_hat, xc2c_inv, KokkosFFT::Impl::Direction::Backward, axis);
+    KokkosFFT::ifft(execution_space(), xc2c_hat, xc2c_inv, ifft_plan);
+
+    // 1D R2C FFT
+    View1D<double> xr2c("xr2c", n0);
+    View1D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0/2+1);
+    Kokkos::fill_random(xr2c, random_pool, 1);
+
+    KokkosFFT::Impl::Plan rfft_plan(execution_space(), xr2c, xr2c_hat, KokkosFFT::Impl::Direction::Forward, axis);
+    KokkosFFT::rfft(execution_space(), xr2c, xr2c_hat, rfft_plan);
+
+    // 1D C2R FFT
+    View1D<Kokkos::complex<double> > xc2r("xc2r_hat", n0/2+1);
+    View1D<double> xc2r_hat("xc2r", n0);
+    Kokkos::fill_random(xc2r, random_pool, I);
+
+    KokkosFFT::Impl::Plan irfft_plan(execution_space(), xc2r, xc2r_hat, KokkosFFT::Impl::Direction::Backward, axis);
+    KokkosFFT::irfft(execution_space(), xc2r, xc2r_hat, irfft_plan);
+  }
+  Kokkos::finalize();
+
+  return 0;
+}
diff --git a/examples/06_1DFFT_reuse_plans/CMakeLists.txt b/examples/06_1DFFT_reuse_plans/CMakeLists.txt
@@ -0,0 +1,2 @@
+add_executable(06_1DFFT_reuse_plans 06_1DFFT_reuse_plans.cpp)
+target_link_libraries(06_1DFFT_reuse_plans PUBLIC Kokkos::fft)
diff --git a/examples/06_1DFFT_reuse_plans/numpy_1DFFT.py b/examples/06_1DFFT_reuse_plans/numpy_1DFFT.py
@@ -0,0 +1,17 @@
+import numpy as np
+
+if __name__ == '__main__':
+    n0 = 128
+
+    # 1D C2C FFT (Forward and Backward)
+    xc2c = np.random.rand(n0) + 1j * np.random.rand(n0)
+    xc2c_hat = np.fft.fft(xc2c)
+    xc2c_inv = np.fft.ifft(xc2c_hat)
+
+    # 1D R2C FFT
+    xr2c = np.random.rand(n0)
+    xr2c_hat = np.fft.rfft(xr2c)
+
+    # 1D C2R FFT
+    xc2r = np.random.rand(n0//2+1)
+    xc2r_hat = np.fft.irfft(xc2r)
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -2,6 +2,5 @@ add_subdirectory(01_1DFFT)
 add_subdirectory(02_2DFFT)
 add_subdirectory(03_NDFFT)
 add_subdirectory(04_batchedFFT)
-if(KokkosFFT_ENABLE_HOST_AND_DEVICE)
-  add_subdirectory(05_1DFFT_HOST_DEVICE)
-endif()
+add_subdirectory(05_1DFFT_HOST_DEVICE)
+add_subdirectory(06_1DFFT_reuse_plans)
diff --git a/fft/src/KokkosFFT_Plans.hpp b/fft/src/KokkosFFT_Plans.hpp
@@ -33,6 +33,7 @@ namespace KokkosFFT {
 namespace Impl {
   template <typename ExecutionSpace, typename InViewType, typename OutViewType, std::size_t DIM = 1>
   class Plan {
+    using execSpace = ExecutionSpace;
     using in_value_type = typename InViewType::non_const_value_type;
     using out_value_type = typename OutViewType::non_const_value_type;
     using float_type = KokkosFFT::Impl::real_type_t<in_value_type>;
@@ -46,24 +47,27 @@ namespace Impl {
     fft_size_type m_fft_size;
     map_type m_map, m_map_inv;
     bool m_is_transpose_needed;
+    axis_type<DIM> m_axes;
+    KokkosFFT::Impl::Direction m_direction;
 
     // Only used when transpose needed
     nonConstInViewType m_in_T;
     nonConstOutViewType m_out_T;
 
   public:
-    explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, int axis) : m_fft_size(1), m_is_transpose_needed(false) {
+    explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, int axis) : m_fft_size(1), m_is_transpose_needed(false), m_direction(direction) {
       static_assert(Kokkos::is_view<InViewType>::value,
                     "KokkosFFT::Plan: InViewType is not a Kokkos::View.");
       static_assert(Kokkos::is_view<OutViewType>::value,
                     "KokkosFFT::Plan: OutViewType is not a Kokkos::View.");
 
+      m_axes = {axis};
       std::tie(m_map, m_map_inv) = KokkosFFT::Impl::get_map_axes(in, axis);
       m_is_transpose_needed = KokkosFFT::Impl::is_transpose_needed(m_map);
-      m_fft_size = KokkosFFT::Impl::_create(exec_space, m_plan, in, out, direction, axis_type<1>{axis});
+      m_fft_size = KokkosFFT::Impl::_create(exec_space, m_plan, in, out, direction, m_axes);
     }
 
-    explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, axis_type<DIM> axes) : m_fft_size(1), m_is_transpose_needed(false) {
+    explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, axis_type<DIM> axes) : m_fft_size(1), m_is_transpose_needed(false), m_direction(direction), m_axes(axes) {
       static_assert(Kokkos::is_view<InViewType>::value,
                     "KokkosFFT::Plan: InViewType is not a Kokkos::View.");
       static_assert(Kokkos::is_view<OutViewType>::value,
@@ -78,6 +82,29 @@ namespace Impl {
       _destroy<ExecutionSpace, float_type>(m_plan);
     }
 
+    template <typename ExecutionSpace2, typename InViewType2, typename OutViewType2>
+    void good(KokkosFFT::Impl::Direction direction, axis_type<DIM> axes) const {
+      static_assert(std::is_same_v<ExecutionSpace2, execSpace>,
+                    "KokkosFFT::Plan: is_good: Execution spaces for plan and execution are inconsistent.");
+
+      using nonConstInViewType2 = std::remove_cv_t<InViewType2>;
+      using nonConstOutViewType2 = std::remove_cv_t<OutViewType2>;
+      static_assert(std::is_same_v<nonConstInViewType2, nonConstInViewType>,
+                    "KokkosFFT::Plan: is_good: InViewType for plan and execution are inconsistent.");
+      static_assert(std::is_same_v<nonConstOutViewType2, nonConstOutViewType>,
+                    "KokkosFFT::Plan: is_good: OutViewType for plan and execution are inconsistent.");
+
+      if(direction != m_direction) {
+        throw std::runtime_error("KokkosFFT::Impl::Plan::good: direction for plan and execution are inconsistent.");
+      }
+
+      if(axes != m_axes) {
+        throw std::runtime_error("KokkosFFT::Impl::Plan::good: axes for plan and execution are inconsistent.");
+      }
+
+      // [TO DO] Check view extents
+    }
+
     fft_plan_type plan() const { return m_plan; }
     fft_size_type fft_size() const { return m_fft_size; }
     bool is_transpose_needed() const { return m_is_transpose_needed; }
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		add_executable(06_1DFFT_reuse_plans 06_1DFFT_reuse_plans.cpp)
		target_link_libraries(06_1DFFT_reuse_plans PUBLIC Kokkos::fft)