From d485e71aeaba3a8158a1c6347966484c04fef246 Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Fri, 8 Dec 2023 00:30:37 +0900 Subject: [PATCH] Add simple examples and reference numpy scripts --- CMakeLists.txt | 8 +++- examples/01_1DFFT/01_1DFFT.cpp | 43 ++++++++++++++++++++++ examples/01_1DFFT/CMakeLists.txt | 2 + examples/01_1DFFT/numpy_1DFFT.py | 17 +++++++++ examples/02_2DFFT/02_2DFFT.cpp | 43 ++++++++++++++++++++++ examples/02_2DFFT/CMakeLists.txt | 2 + examples/02_2DFFT/numpy_2DFFT.py | 17 +++++++++ examples/03_NDFFT/03_NDFFT.cpp | 43 ++++++++++++++++++++++ examples/03_NDFFT/CMakeLists.txt | 2 + examples/03_NDFFT/numpy_NDFFT.py | 17 +++++++++ examples/04_batchedFFT/04_batchedFFT.cpp | 43 ++++++++++++++++++++++ examples/04_batchedFFT/CMakeLists.txt | 2 + examples/04_batchedFFT/numpy_batchedFFT.py | 17 +++++++++ examples/CMakeLists.txt | 4 ++ 14 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 examples/01_1DFFT/01_1DFFT.cpp create mode 100644 examples/01_1DFFT/CMakeLists.txt create mode 100644 examples/01_1DFFT/numpy_1DFFT.py create mode 100644 examples/02_2DFFT/02_2DFFT.cpp create mode 100644 examples/02_2DFFT/CMakeLists.txt create mode 100644 examples/02_2DFFT/numpy_2DFFT.py create mode 100644 examples/03_NDFFT/03_NDFFT.cpp create mode 100644 examples/03_NDFFT/CMakeLists.txt create mode 100644 examples/03_NDFFT/numpy_NDFFT.py create mode 100644 examples/04_batchedFFT/04_batchedFFT.cpp create mode 100644 examples/04_batchedFFT/CMakeLists.txt create mode 100644 examples/04_batchedFFT/numpy_batchedFFT.py create mode 100644 examples/CMakeLists.txt diff --git a/CMakeLists.txt b/CMakeLists.txt index 297002b0..11f2d40f 100644 --- a/CMakeLists.txt +++ b/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) \ No newline at end of file +add_subdirectory(fft) +if(BUILD_EXAMPLES) + add_subdirectory(examples) +endif() \ No newline at end of file diff --git a/examples/01_1DFFT/01_1DFFT.cpp b/examples/01_1DFFT/01_1DFFT.cpp new file mode 100644 index 00000000..3169f9eb --- /dev/null +++ b/examples/01_1DFFT/01_1DFFT.cpp @@ -0,0 +1,43 @@ +#include +#include +#include +#include + +using execution_space = Kokkos::DefaultExecutionSpace; +template using View1D = Kokkos::View; + +int main( int argc, char* argv[] ) { + Kokkos::initialize( argc, argv ); + { + constexpr int n0 = 128; + const Kokkos::complex I(1.0, 1.0); + + // 1D C2C FFT (Forward and Backward) + View1D > xc2c("xc2c", n0); + View1D > xc2c_hat("xc2c_hat", n0); + View1D > 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 xr2c("xr2c", n0); + View1D > xr2c_hat("xr2c_hat", n0/2+1); + Kokkos::fill_random(xr2c, random_pool, 1); + + KokkosFFT::rfft(xr2c, xr2c_hat); + + // 1D C2R FFT + View1D > xc2r("xr2c_hat", n0/2+1); + View1D xc2r_hat("xc2r", n0); + Kokkos::fill_random(xc2r, random_pool, I); + + KokkosFFT::irfft(xc2r, xc2r_hat); + } + Kokkos::finalize(); + + return 0; +} \ No newline at end of file diff --git a/examples/01_1DFFT/CMakeLists.txt b/examples/01_1DFFT/CMakeLists.txt new file mode 100644 index 00000000..19aa79a5 --- /dev/null +++ b/examples/01_1DFFT/CMakeLists.txt @@ -0,0 +1,2 @@ +add_executable(01_1DFFT 01_1DFFT.cpp) +target_link_libraries(01_1DFFT PUBLIC Kokkos::fft) \ No newline at end of file diff --git a/examples/01_1DFFT/numpy_1DFFT.py b/examples/01_1DFFT/numpy_1DFFT.py new file mode 100644 index 00000000..274978f5 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/examples/02_2DFFT/02_2DFFT.cpp b/examples/02_2DFFT/02_2DFFT.cpp new file mode 100644 index 00000000..0faeced7 --- /dev/null +++ b/examples/02_2DFFT/02_2DFFT.cpp @@ -0,0 +1,43 @@ +#include +#include +#include +#include + +using execution_space = Kokkos::DefaultExecutionSpace; +template using View2D = Kokkos::View; + +int main( int argc, char* argv[] ) { + Kokkos::initialize( argc, argv ); + { + constexpr int n0 = 128, n1 = 128; + const Kokkos::complex I(1.0, 1.0); + + // 2D C2C FFT (Forward and Backward) + View2D > xc2c("xc2c", n0, n1); + View2D > xc2c_hat("xc2c_hat", n0, n1); + View2D > 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 xr2c("xr2c", n0, n1); + View2D > xr2c_hat("xr2c_hat", n0, n1/2+1); + Kokkos::fill_random(xr2c, random_pool, 1); + + KokkosFFT::rfft2(xr2c, xr2c_hat); + + // 2D C2R FFT + View2D > xc2r("xr2c_hat", n0, n1/2+1); + View2D xc2r_hat("xc2r", n0, n1); + Kokkos::fill_random(xc2r, random_pool, I); + + KokkosFFT::irfft2(xc2r, xc2r_hat); + } + Kokkos::finalize(); + + return 0; +} \ No newline at end of file diff --git a/examples/02_2DFFT/CMakeLists.txt b/examples/02_2DFFT/CMakeLists.txt new file mode 100644 index 00000000..18c7d096 --- /dev/null +++ b/examples/02_2DFFT/CMakeLists.txt @@ -0,0 +1,2 @@ +add_executable(02_2DFFT 02_2DFFT.cpp) +target_link_libraries(02_2DFFT PUBLIC Kokkos::fft) \ No newline at end of file diff --git a/examples/02_2DFFT/numpy_2DFFT.py b/examples/02_2DFFT/numpy_2DFFT.py new file mode 100644 index 00000000..50d4169d --- /dev/null +++ b/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) \ No newline at end of file diff --git a/examples/03_NDFFT/03_NDFFT.cpp b/examples/03_NDFFT/03_NDFFT.cpp new file mode 100644 index 00000000..991e67c3 --- /dev/null +++ b/examples/03_NDFFT/03_NDFFT.cpp @@ -0,0 +1,43 @@ +#include +#include +#include +#include + +using execution_space = Kokkos::DefaultExecutionSpace; +template using View3D = Kokkos::View; + +int main( int argc, char* argv[] ) { + Kokkos::initialize( argc, argv ); + { + constexpr int n0 = 128, n1 = 128, n2 = 16; + const Kokkos::complex I(1.0, 1.0); + + // 3D C2C FFT (Forward and Backward) + View3D > xc2c("xc2c", n0, n1, n2); + View3D > xc2c_hat("xc2c_hat", n0, n1, n2); + View3D > 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 xr2c("xr2c", n0, n1, n2); + View3D > 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 > xc2r("xr2c_hat", n0, n1, n2/2+1); + View3D xc2r_hat("xc2r", n0, n1, n2); + Kokkos::fill_random(xc2r, random_pool, I); + + KokkosFFT::irfftn(xc2r, xc2r_hat); + } + Kokkos::finalize(); + + return 0; +} \ No newline at end of file diff --git a/examples/03_NDFFT/CMakeLists.txt b/examples/03_NDFFT/CMakeLists.txt new file mode 100644 index 00000000..0dac1239 --- /dev/null +++ b/examples/03_NDFFT/CMakeLists.txt @@ -0,0 +1,2 @@ +add_executable(03_NDFFT 03_NDFFT.cpp) +target_link_libraries(03_NDFFT PUBLIC Kokkos::fft) \ No newline at end of file diff --git a/examples/03_NDFFT/numpy_NDFFT.py b/examples/03_NDFFT/numpy_NDFFT.py new file mode 100644 index 00000000..dfd93ef8 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/examples/04_batchedFFT/04_batchedFFT.cpp b/examples/04_batchedFFT/04_batchedFFT.cpp new file mode 100644 index 00000000..3554ab96 --- /dev/null +++ b/examples/04_batchedFFT/04_batchedFFT.cpp @@ -0,0 +1,43 @@ +#include +#include +#include +#include + +using execution_space = Kokkos::DefaultExecutionSpace; +template using View3D = Kokkos::View; + +int main( int argc, char* argv[] ) { + Kokkos::initialize( argc, argv ); + { + constexpr int n0 = 128, n1 = 128, n2 = 16; + const Kokkos::complex I(1.0, 1.0); + + // 1D batched C2C FFT (Forward and Backward) + View3D > xc2c("xc2c", n0, n1, n2); + View3D > xc2c_hat("xc2c_hat", n0, n1, n2); + View3D > 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 xr2c("xr2c", n0, n1, n2); + View3D > 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 > xc2r("xr2c_hat", n0, n1, n2/2+1); + View3D 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; +} \ No newline at end of file diff --git a/examples/04_batchedFFT/CMakeLists.txt b/examples/04_batchedFFT/CMakeLists.txt new file mode 100644 index 00000000..5fe03b30 --- /dev/null +++ b/examples/04_batchedFFT/CMakeLists.txt @@ -0,0 +1,2 @@ +add_executable(04_batchedFFT 04_batchedFFT.cpp) +target_link_libraries(04_batchedFFT PUBLIC Kokkos::fft) \ No newline at end of file diff --git a/examples/04_batchedFFT/numpy_batchedFFT.py b/examples/04_batchedFFT/numpy_batchedFFT.py new file mode 100644 index 00000000..06a486c0 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt new file mode 100644 index 00000000..bc70f031 --- /dev/null +++ b/examples/CMakeLists.txt @@ -0,0 +1,4 @@ +add_subdirectory(01_1DFFT) +add_subdirectory(02_2DFFT) +add_subdirectory(03_NDFFT) +add_subdirectory(04_batchedFFT) \ No newline at end of file