From 29389f298c175607caa60b07a564b5e90c157336 Mon Sep 17 00:00:00 2001
From: yasahi-hpc <57478230+yasahi-hpc@users.noreply.github.com>
Date: Tue, 5 Nov 2024 10:01:09 +0100
Subject: [PATCH] Add inplace transform example (#190)

* Add inplace trasnform example

* fix: in_inplace is unsed for C2C

* fix: inplace example based on reviews

* remove unused types

---------

Co-authored-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>
---
 common/src/KokkosFFT_Extents.hpp           |  2 +-
 examples/08_inplace_FFT/08_inplace_FFT.cpp | 75 ++++++++++++++++++++++
 examples/08_inplace_FFT/CMakeLists.txt     |  6 ++
 examples/CMakeLists.txt                    |  1 +
 4 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 examples/08_inplace_FFT/08_inplace_FFT.cpp
 create mode 100644 examples/08_inplace_FFT/CMakeLists.txt

diff --git a/common/src/KokkosFFT_Extents.hpp b/common/src/KokkosFFT_Extents.hpp
index 95a4dc34..e3bea237 100644
--- a/common/src/KokkosFFT_Extents.hpp
+++ b/common/src/KokkosFFT_Extents.hpp
@@ -33,7 +33,7 @@ namespace Impl {
 template <typename InViewType, typename OutViewType, std::size_t DIM = 1>
 auto get_extents(const InViewType& in, const OutViewType& out,
                  axis_type<DIM> axes, shape_type<DIM> shape = {},
-                 bool is_inplace = false) {
+                 [[maybe_unused]] bool is_inplace = false) {
   using in_value_type     = typename InViewType::non_const_value_type;
   using out_value_type    = typename OutViewType::non_const_value_type;
   using array_layout_type = typename InViewType::array_layout;
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..7e17904d
--- /dev/null
+++ b/examples/08_inplace_FFT/08_inplace_FFT.cpp
@@ -0,0 +1,75 @@
+// 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>;
+
+int main(int argc, char *argv[]) {
+  Kokkos::initialize(argc, argv);
+  {
+    const int n0 = 128, n1 = 128;
+    const Kokkos::complex<double> z(1.0, 1.0);
+
+    // Forward and backward complex to complex transform
+    // Define a 2D complex view to handle data
+    RightView2D<Kokkos::complex<double>> xc2c("xc2c", 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);
+    KokkosFFT::ifft2(exec, xc2c, xc2c);
+
+    // Real to complex transform
+    // Define a 2D complex view to handle data
+    RightView2D<Kokkos::complex<double>> xr2c_hat("xr2c", n0, n1 / 2 + 1);
+
+    // Create unmanaged views 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);
+
+    exec.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)