From cd8e05e98feb5d981ad6786e95fe08af1e622dff Mon Sep 17 00:00:00 2001
From: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>
Date: Mon, 4 Nov 2024 17:48:05 +0900
Subject: [PATCH] Add inplace trasnform example

---
 examples/08_inplace_FFT/08_inplace_FFT.cpp | 88 ++++++++++++++++++++++
 examples/08_inplace_FFT/CMakeLists.txt     |  6 ++
 examples/CMakeLists.txt                    |  1 +
 3 files changed, 95 insertions(+)
 create mode 100644 examples/08_inplace_FFT/08_inplace_FFT.cpp
 create mode 100644 examples/08_inplace_FFT/CMakeLists.txt

diff --git a/examples/08_inplace_FFT/08_inplace_FFT.cpp b/examples/08_inplace_FFT/08_inplace_FFT.cpp
new file mode 100644
index 00000000..97e43e99
--- /dev/null
+++ b/examples/08_inplace_FFT/08_inplace_FFT.cpp
@@ -0,0 +1,88 @@
+// SPDX-FileCopyrightText: (C) The kokkos-fft development team, see COPYRIGHT.md file
+//
+// SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception
+
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Complex.hpp>
+#include <Kokkos_Random.hpp>
+#include <KokkosFFT.hpp>
+
+using execution_space = Kokkos::DefaultExecutionSpace;
+
+template <typename T>
+using RightView2D = Kokkos::View<T **, Kokkos::LayoutRight, execution_space>;
+
+template <std::size_t DIM>
+using axis_type = KokkosFFT::axis_type<DIM>;
+template <std::size_t DIM>
+using shape_type = KokkosFFT::shape_type<DIM>;
+
+int main(int argc, char *argv[]) {
+  Kokkos::initialize(argc, argv);
+  {
+    constexpr int n0 = 128, n1 = 128;
+    const Kokkos::complex<double> z(1.0, 1.0);
+
+    shape_type<2> shape;
+    shape[0] = n0;
+    shape[1] = n1;
+
+    // Forward and backward complex to complex transform
+    // Define a 2D complex view to handle data
+    RightView2D<Kokkos::complex<double>> xc2c("xc2c", n0, n1);
+
+    // Create an unmanaged view on the same data
+    RightView2D<Kokkos::complex<double>> xc2c_hat(xc2c.data(), n0, n1),
+        xc2c_inv(xc2c.data(), n0, n1);
+
+    // Fill the input view with random data
+    Kokkos::Random_XorShift64_Pool<> random_pool(12345);
+    execution_space exec;
+    Kokkos::fill_random(exec, xc2c, random_pool, z);
+
+    KokkosFFT::fft2(exec, xc2c, xc2c_hat);
+    KokkosFFT::ifft2(exec, xc2c_hat, xc2c_inv);
+
+    // Real to complex transform
+    // Define a 2D complex view to handle data
+    RightView2D<Kokkos::complex<double>> xr2c_hat("xr2c", n0, n1 / 2 + 1);
+
+    // Create an unmanaged view on the same data with the FFT shape,
+    // that is (n0, n1) -> (n0, n1/2+1) R2C transform
+    // The shape is incorrect from the view point of casting to real
+    // For casting, the shape should be (n0, (n0/2+1) * 2)
+    RightView2D<double> xr2c(reinterpret_cast<double *>(xr2c_hat.data()), n0,
+                             n1),
+        xr2c_padded(reinterpret_cast<double *>(xr2c_hat.data()), n0,
+                    (n0 / 2 + 1) * 2);
+
+    // Fill the input view with random data in real space through xr2c_padded
+    auto sub_xr2c_padded =
+        Kokkos::subview(xr2c_padded, Kokkos::ALL, Kokkos::make_pair(0, n1));
+    Kokkos::fill_random(exec, sub_xr2c_padded, random_pool, 1.0);
+
+    // Perform the real to complex transform
+    // [Important] You must use xr2c to define the FFT shape correctly
+    KokkosFFT::rfft2(exec, xr2c, xr2c_hat);
+
+    // Complex to real transform
+    // Define a 2D complex view to handle data
+    RightView2D<Kokkos::complex<double>> xc2r("xc2r", n0, n1 / 2 + 1);
+
+    // Create an unmanaged view on the same data with the FFT shape
+    RightView2D<double> xc2r_hat(reinterpret_cast<double *>(xc2r.data()), n0,
+                                 n1);
+
+    // Fill the input view with random data in complex space
+    Kokkos::fill_random(exec, xc2r, random_pool, z);
+
+    // Perform the complex to real transform
+    // [Important] You must use xc2r_hat to define the FFT shape correctly
+    KokkosFFT::irfft2(exec, xc2r, xc2r_hat);
+
+    Kokkos::fence();
+  }
+  Kokkos::finalize();
+
+  return 0;
+}
diff --git a/examples/08_inplace_FFT/CMakeLists.txt b/examples/08_inplace_FFT/CMakeLists.txt
new file mode 100644
index 00000000..20faadaf
--- /dev/null
+++ b/examples/08_inplace_FFT/CMakeLists.txt
@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: (C) The kokkos-fft development team, see COPYRIGHT.md file
+#
+# SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception
+
+add_executable(08_inplace_FFT 08_inplace_FFT.cpp)
+target_link_libraries(08_inplace_FFT PUBLIC KokkosFFT::fft)
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
index 919fef76..94b3c84a 100644
--- a/examples/CMakeLists.txt
+++ b/examples/CMakeLists.txt
@@ -9,3 +9,4 @@ add_subdirectory(04_batchedFFT)
 add_subdirectory(05_1DFFT_HOST_DEVICE)
 add_subdirectory(06_1DFFT_reuse_plans)
 add_subdirectory(07_unmanaged_views)
+add_subdirectory(08_inplace_FFT)