AnyFFT: 1D Support (#5112)

* AnyFFT 1D: FFTW * AnyFFT 1D: CuFFT
ECP-WarpX · Aug 7, 2024 · 547794d · 547794d
1 parent fe36717
commit 547794d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/Source/ablastr/math/fft/WrapCuFFT.cpp b/Source/ablastr/math/fft/WrapCuFFT.cpp
@@ -42,8 +42,11 @@ namespace ablastr::math::anyfft
             } else if (dim == 2) {
                 result = cufftPlan2d(
                     &(fft_plan.m_plan), real_size[1], real_size[0], VendorR2C);
+            } else if (dim == 1) {
+                result = cufftPlan1d(
+                    &(fft_plan.m_plan), real_size[0], VendorR2C, 1);
             } else {
-                ABLASTR_ABORT_WITH_MESSAGE("only dim=2 and dim=3 have been implemented");
+                ABLASTR_ABORT_WITH_MESSAGE("only dim=1 and dim=2 and dim=3 have been implemented");
             }
         } else {
             if (dim == 3) {
@@ -52,6 +55,10 @@ namespace ablastr::math::anyfft
             } else if (dim == 2) {
                 result = cufftPlan2d(
                     &(fft_plan.m_plan), real_size[1], real_size[0], VendorC2R);
+            } else if (dim == 1) {
+                int batch = 2;
+                result = cufftPlan1d(
+                    &(fft_plan.m_plan), real_size[0], VendorC2R, 1);
             } else {
                 ABLASTR_ABORT_WITH_MESSAGE("only dim=2 and dim=3 have been implemented");
             }

diff --git a/Source/ablastr/math/fft/WrapFFTW.cpp b/Source/ablastr/math/fft/WrapFFTW.cpp
@@ -25,11 +25,15 @@ namespace ablastr::math::anyfft
     const auto VendorCreatePlanC2R3D = fftwf_plan_dft_c2r_3d;
     const auto VendorCreatePlanR2C2D = fftwf_plan_dft_r2c_2d;
     const auto VendorCreatePlanC2R2D = fftwf_plan_dft_c2r_2d;
+    const auto VendorCreatePlanR2C1D = fftwf_plan_dft_r2c_1d;
+    const auto VendorCreatePlanC2R1D = fftwf_plan_dft_c2r_1d;
 #else
     const auto VendorCreatePlanR2C3D = fftw_plan_dft_r2c_3d;
     const auto VendorCreatePlanC2R3D = fftw_plan_dft_c2r_3d;
     const auto VendorCreatePlanR2C2D = fftw_plan_dft_r2c_2d;
     const auto VendorCreatePlanC2R2D = fftw_plan_dft_c2r_2d;
+    const auto VendorCreatePlanR2C1D = fftw_plan_dft_r2c_1d;
+    const auto VendorCreatePlanC2R1D = fftw_plan_dft_c2r_1d;
 #endif
 
     FFTplan CreatePlan(const amrex::IntVect& real_size, amrex::Real * const real_array,
@@ -56,9 +60,12 @@ namespace ablastr::math::anyfft
             } else if (dim == 2) {
                 fft_plan.m_plan = VendorCreatePlanR2C2D(
                     real_size[1], real_size[0], real_array, complex_array, FFTW_ESTIMATE);
+            } else if (dim == 1) {
+                fft_plan.m_plan = VendorCreatePlanR2C1D(
+                    real_size[0], real_array, complex_array, FFTW_ESTIMATE);
             } else {
                 ABLASTR_ABORT_WITH_MESSAGE(
-                    "only dim=2 and dim=3 have been implemented");
+                    "only dim=1 and dim=2 and dim=3 have been implemented");
             }
         } else if (dir == direction::C2R){
             if (dim == 3) {
@@ -67,9 +74,12 @@ namespace ablastr::math::anyfft
             } else if (dim == 2) {
                 fft_plan.m_plan = VendorCreatePlanC2R2D(
                     real_size[1], real_size[0], complex_array, real_array, FFTW_ESTIMATE);
+            } else if (dim == 1) {
+                fft_plan.m_plan = VendorCreatePlanC2R1D(
+                    real_size[0], complex_array, real_array, FFTW_ESTIMATE);
             } else {
                 ABLASTR_ABORT_WITH_MESSAGE(
-                    "only dim=2 and dim=3 have been implemented. Should be easy to add dim=1.");
+                    "only dim=1 and dim=2 and dim=3 have been implemented.");
             }
         }