Add simple examples and reference numpy scripts

CMakeLists.txt

@@ -4,6 +4,9 @@ project(kokkos-fft LANGUAGES CXX)
 # Add cmake helpers for FFTW
 list(INSERT CMAKE_MODULE_PATH 0 "${CMAKE_SOURCE_DIR}/cmake")
 
+# Options
+option(BUILD_EXAMPLES "Build kokkos-fft examples" ON)
+
 find_package(Kokkos CONFIG)
 if(NOT kokkos_FOUND)
     add_subdirectory(tpls/kokkos)
@@ -19,4 +22,7 @@ if(BUILD_TESTING)
 endif()
 
 add_subdirectory(common)
-add_subdirectory(fft)
+add_subdirectory(fft)
+if(BUILD_EXAMPLES)
+    add_subdirectory(examples)
+endif()

examples/01_1DFFT/01_1DFFT.cpp

@@ -0,0 +1,43 @@
+#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);
+
+    KokkosFFT::fft(xc2c, xc2c_hat);
+    KokkosFFT::ifft(xc2c_hat, xc2c_inv);
+
+    // 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::rfft(xr2c, xr2c_hat);
+
+    // 1D C2R FFT
+    View1D<Kokkos::complex<double> > xc2r("xr2c_hat", n0/2+1);
+    View1D<double> xc2r_hat("xc2r", n0);
+    Kokkos::fill_random(xc2r, random_pool, I);
+
+    KokkosFFT::irfft(xc2r, xc2r_hat);
+  }
+  Kokkos::finalize();
+
+  return 0;
+}

examples/01_1DFFT/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_executable(01_1DFFT 01_1DFFT.cpp)
+target_link_libraries(01_1DFFT PUBLIC Kokkos::fft)

examples/01_1DFFT/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)

examples/02_2DFFT/02_2DFFT.cpp

@@ -0,0 +1,43 @@
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Complex.hpp>
+#include <Kokkos_Random.hpp>
+#include <KokkosFFT.hpp>
+
+using execution_space = Kokkos::DefaultExecutionSpace;
+template <typename T> using View2D = Kokkos::View<T**, execution_space>;
+
+int main( int argc, char* argv[] ) {
+  Kokkos::initialize( argc, argv );
+  {
+    constexpr int n0 = 128, n1 = 128;
+    const Kokkos::complex<double> I(1.0, 1.0);
+
+    // 2D C2C FFT (Forward and Backward)
+    View2D<Kokkos::complex<double> > xc2c("xc2c", n0, n1);
+    View2D<Kokkos::complex<double> > xc2c_hat("xc2c_hat", n0, n1);
+    View2D<Kokkos::complex<double> > xc2c_inv("xc2c_inv", n0, n1);
+
+    Kokkos::Random_XorShift64_Pool<> random_pool(12345);
+    Kokkos::fill_random(xc2c, random_pool, I);
+
+    KokkosFFT::fft2(xc2c, xc2c_hat);
+    KokkosFFT::ifft2(xc2c_hat, xc2c_inv);
+
+    // 2D R2C FFT
+    View2D<double> xr2c("xr2c", n0, n1);
+    View2D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0, n1/2+1);
+    Kokkos::fill_random(xr2c, random_pool, 1);
+
+    KokkosFFT::rfft2(xr2c, xr2c_hat);
+
+    // 2D C2R FFT
+    View2D<Kokkos::complex<double> > xc2r("xr2c_hat", n0, n1/2+1);
+    View2D<double> xc2r_hat("xc2r", n0, n1);
+    Kokkos::fill_random(xc2r, random_pool, I);
+
+    KokkosFFT::irfft2(xc2r, xc2r_hat);
+  }
+  Kokkos::finalize();
+
+  return 0;
+}

examples/02_2DFFT/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_executable(02_2DFFT 02_2DFFT.cpp)
+target_link_libraries(02_2DFFT PUBLIC Kokkos::fft)

examples/02_2DFFT/numpy_2DFFT.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+if __name__ == '__main__':
+    n0, n1 = 128, 128
+
+    # 2D C2C FFT (Forward and Backward)
+    xc2c = np.random.rand(n0, n1) + 1j * np.random.rand(n0, n1)
+    xc2c_hat = np.fft.fft2(xc2c)
+    xc2c_inv = np.fft.ifft2(xc2c_hat)
+
+    # 2D R2C FFT
+    xr2c = np.random.rand(n0, n1)
+    xr2c_hat = np.fft.rfft2(xr2c)
+
+    # 2D C2R FFT
+    xc2r = np.random.rand(n0, n1//2+1)
+    xc2r_hat = np.fft.irfft2(xc2r)

examples/03_NDFFT/03_NDFFT.cpp

@@ -0,0 +1,43 @@
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Complex.hpp>
+#include <Kokkos_Random.hpp>
+#include <KokkosFFT.hpp>
+
+using execution_space = Kokkos::DefaultExecutionSpace;
+template <typename T> using View3D = Kokkos::View<T***, execution_space>;
+
+int main( int argc, char* argv[] ) {
+  Kokkos::initialize( argc, argv );
+  {
+    constexpr int n0 = 128, n1 = 128, n2 = 16;
+    const Kokkos::complex<double> I(1.0, 1.0);
+
+    // 3D C2C FFT (Forward and Backward)
+    View3D<Kokkos::complex<double> > xc2c("xc2c", n0, n1, n2);
+    View3D<Kokkos::complex<double> > xc2c_hat("xc2c_hat", n0, n1, n2);
+    View3D<Kokkos::complex<double> > xc2c_inv("xc2c_inv", n0, n1, n2);
+
+    Kokkos::Random_XorShift64_Pool<> random_pool(12345);
+    Kokkos::fill_random(xc2c, random_pool, I);
+
+    KokkosFFT::fftn(xc2c, xc2c_hat);
+    KokkosFFT::ifftn(xc2c_hat, xc2c_inv);
+
+    // 3D R2C FFT
+    View3D<double> xr2c("xr2c", n0, n1, n2);
+    View3D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0, n1, n2/2+1);
+    Kokkos::fill_random(xr2c, random_pool, 1);
+
+    KokkosFFT::rfftn(xr2c, xr2c_hat);
+
+    // 3D C2R FFT
+    View3D<Kokkos::complex<double> > xc2r("xr2c_hat", n0, n1, n2/2+1);
+    View3D<double> xc2r_hat("xc2r", n0, n1, n2);
+    Kokkos::fill_random(xc2r, random_pool, I);
+
+    KokkosFFT::irfftn(xc2r, xc2r_hat);
+  }
+  Kokkos::finalize();
+
+  return 0;
+}

examples/03_NDFFT/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_executable(03_NDFFT 03_NDFFT.cpp)
+target_link_libraries(03_NDFFT PUBLIC Kokkos::fft)

examples/03_NDFFT/numpy_NDFFT.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+if __name__ == '__main__':
+    n0, n1, n2 = 128, 128, 16
+
+    # 3D C2C FFT (Forward and Backward)
+    xc2c = np.random.rand(n0, n1, n2) + 1j * np.random.rand(n0, n1, n2)
+    xc2c_hat = np.fft.fftn(xc2c)
+    xc2c_inv = np.fft.ifftn(xc2c_hat)
+
+    # 3D R2C FFT
+    xr2c = np.random.rand(n0, n1, n2)
+    xr2c_hat = np.fft.rfftn(xr2c)
+
+    # 3D C2R FFT
+    xc2r = np.random.rand(n0, n1, n2//2+1)
+    xc2r_hat = np.fft.irfftn(xc2r)

examples/04_batchedFFT/04_batchedFFT.cpp

@@ -0,0 +1,43 @@
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Complex.hpp>
+#include <Kokkos_Random.hpp>
+#include <KokkosFFT.hpp>
+
+using execution_space = Kokkos::DefaultExecutionSpace;
+template <typename T> using View3D = Kokkos::View<T***, execution_space>;
+
+int main( int argc, char* argv[] ) {
+  Kokkos::initialize( argc, argv );
+  {
+    constexpr int n0 = 128, n1 = 128, n2 = 16;
+    const Kokkos::complex<double> I(1.0, 1.0);
+
+    // 1D batched C2C FFT (Forward and Backward)
+    View3D<Kokkos::complex<double> > xc2c("xc2c", n0, n1, n2);
+    View3D<Kokkos::complex<double> > xc2c_hat("xc2c_hat", n0, n1, n2);
+    View3D<Kokkos::complex<double> > xc2c_inv("xc2c_inv", n0, n1, n2);
+
+    Kokkos::Random_XorShift64_Pool<> random_pool(12345);
+    Kokkos::fill_random(xc2c, random_pool, I);
+
+    KokkosFFT::fft(xc2c, xc2c_hat, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/-1);
+    KokkosFFT::ifft(xc2c_hat, xc2c_inv, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/-1);
+
+    // 1D batched R2C FFT
+    View3D<double> xr2c("xr2c", n0, n1, n2);
+    View3D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0, n1, n2/2+1);
+    Kokkos::fill_random(xr2c, random_pool, 1);
+
+    KokkosFFT::rfft(xr2c, xr2c_hat, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/-1);
+
+    // 1D batched C2R FFT
+    View3D<Kokkos::complex<double> > xc2r("xr2c_hat", n0, n1, n2/2+1);
+    View3D<double> xc2r_hat("xc2r", n0, n1, n2);
+    Kokkos::fill_random(xc2r, random_pool, I);
+
+    KokkosFFT::irfft(xc2r, xc2r_hat, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/-1);
+  }
+  Kokkos::finalize();
+
+  return 0;
+}

examples/04_batchedFFT/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_executable(04_batchedFFT 04_batchedFFT.cpp)
+target_link_libraries(04_batchedFFT PUBLIC Kokkos::fft)

examples/04_batchedFFT/numpy_batchedFFT.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+if __name__ == '__main__':
+    n0, n1, n2 = 128, 128, 16
+
+    # 1D batched C2C FFT (Forward and Backward)
+    xc2c = np.random.rand(n0, n1, n2) + 1j * np.random.rand(n0, n1, n2)
+    xc2c_hat = np.fft.fft(xc2c, axis=-1)
+    xc2c_inv = np.fft.ifft(xc2c_hat, axis=-1)
+
+    # 1D batched R2C FFT
+    xr2c = np.random.rand(n0, n1, n2)
+    xr2c_hat = np.fft.rfft(xr2c, axis=-1)
+
+    # 1D batched C2R FFT
+    xc2r = np.random.rand(n0, n1, n2//2+1)
+    xc2r_hat = np.fft.irfft(xc2r, axis=-1)

examples/CMakeLists.txt

@@ -0,0 +1,4 @@
+add_subdirectory(01_1DFFT)
+add_subdirectory(02_2DFFT)
+add_subdirectory(03_NDFFT)
+add_subdirectory(04_batchedFFT)