diff --git a/.github/workflows/ci-extended.yml b/.github/workflows/ci-extended.yml index 920cf3f0d2ca2..5bf810b830581 100644 --- a/.github/workflows/ci-extended.yml +++ b/.github/workflows/ci-extended.yml @@ -58,11 +58,14 @@ jobs: fi - name: Configure + env: + ENABLE_FFT: ${{ matrix.parallel == 'mpi' && 'ON' || 'OFF' }} run: | cmake -B build \ -DCMAKE_BUILD_TYPE=Release \ -DENABLE_ASAN=${{ steps.cmake-vars.outputs.enable_asan }} \ - -DMACHINE_VARIANT=${{ matrix.device }}-${{ matrix.parallel }} + -DMACHINE_VARIANT=${{ matrix.device }}-${{ matrix.parallel }} \ + -DPARTHENON_ENABLE_FFT=${ENABLE_FFT} - name: Build run: cmake --build build @@ -153,6 +156,8 @@ jobs: key: gold-standard - name: Configure + env: + ENABLE_FFT: ${{ matrix.parallel == 'mpi' && 'ON' || 'OFF' }} run: | git config --global --add safe.directory $(pwd) # The extra CMAKE_PREFIX_PATH seems to be required for current rocm6.2 and 7.2, see @@ -163,7 +168,8 @@ jobs: -DCMAKE_BUILD_TYPE=Release \ -DMACHINE_VARIANT=hip-${{ matrix.parallel }} \ -DCMAKE_PREFIX_PATH=/opt/rocm-7.2.4 \ - -DCMAKE_CXX_COMPILER=hipcc + -DCMAKE_CXX_COMPILER=hipcc \ + -DPARTHENON_ENABLE_FFT=${ENABLE_FFT} - name: Build run: cmake --build build diff --git a/CHANGELOG.md b/CHANGELOG.md index 81fb360013efe..49c029cfd7e4c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ## Current develop ### Added (new features/APIs/variables/...) +- [[PR 1403]](https://github.com/parthenon-hpc-lab/parthenon/pull/1403) Add interface for Fourier transforms on uniform meshes via heFFTe - [[PR 1408]](https://github.com/parthenon-hpc-lab/parthenon/pull/1408) Add GetAsUnresolvedString() method to ParameterInput - [[PR 1050]](https://github.com/parthenon-hpc-lab/parthenon/pull/1050) Add support for OpenPMD/ADIOS2 output (incl slices and coarsened dumps) - [[PR 1271]](https://github.com/parthenon-hpc-lab/parthenon/pull/1271) Add option to set a minimum number of teams for boundary communication kernels diff --git a/CMakeLists.txt b/CMakeLists.txt index 4c52bd93f6fe6..9f2f92eb290aa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -75,6 +75,9 @@ else() OFF) endif() +option(PARTHENON_ENABLE_FFT "Enable Parthenon FFT functionality" OFF) +option(PARTHENON_USE_SYSTEM_HEFFTE "Use system-installed heFFTe instead of fetching it" OFF) + # Check if we are using Parthenon as a subproject set(_parthenon_is_toplevel OFF) if (CMAKE_SOURCE_DIR STREQUAL PROJECT_SOURCE_DIR) @@ -492,7 +495,42 @@ set(CMAKE_INSTALL_INCLUDEDIR "${CMAKE_INSTALL_INCLUDEDIR}/parthenon") set(DOC_GEN_PATH "${CMAKE_SOURCE_DIR}/doc/sphinx/src/generated" CACHE STRING "Path to save generated data for docs.") +if (PARTHENON_ENABLE_FFT) + if(NOT ENABLE_MPI) + message(FATAL_ERROR + "PARTHENON_ENABLE_FFT requires MPI because the heFFTe backend uses " + "an MPI communicator. Reconfigure without PARTHENON_DISABLE_MPI=ON.") + endif() + if (PARTHENON_USE_SYSTEM_HEFFTE) + find_package(Heffte REQUIRED) + if (NOT Heffte_FOUND) + message(FATAL_ERROR "heFFTe requested to be used from the environment but it was " + "not found. Either append the path to the installed heFFTe lib via " + "Heffte_DIR or CMAKE_PREFIX_PATH, or use the version fetched by Parthenon by " + "keeping PARTHENON_USE_SYSTEM_HEFFTE=OFF") + endif() + else() + include(FetchContent) + set(Heffte_ENABLE_CUDA ${Kokkos_ENABLE_CUDA} CACHE BOOL "Enable CUDA for heFFTe" FORCE) + set(Heffte_ENABLE_ROCM ${Kokkos_ENABLE_HIP} CACHE BOOL "Enable ROCm for heFFTe" FORCE) + FetchContent_Declare(Heffte + GIT_REPOSITORY "https://github.com/icl-utk-edu/heffte.git" + GIT_TAG "4d8d459") # current (Jul26) GitHub `master` (from Nov25) + FetchContent_MakeAvailable(Heffte) + if (NOT TARGET Heffte::Heffte) + add_library(Heffte::Heffte ALIAS Heffte) + endif() + install(TARGETS Heffte EXPORT parthenonTargets) + endif() +endif() + add_subdirectory(src) + +if (PARTHENON_ENABLE_FFT) + target_link_libraries(parthenon PRIVATE Heffte::Heffte) + target_compile_definitions(parthenon PUBLIC PARTHENON_ENABLE_FFT) +endif() + add_subdirectory(example) add_subdirectory(benchmarks) diff --git a/README.md b/README.md index 1403d5a2b976d..0c353b2be2fcb 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ Parthenon -- a performance portable block-structured adaptive mesh refinement fr * OpenPMD and ADIOS2 (for outputs) * HDF5 (for outputs) * Ascent (for in situ visualization and analysis) +* heFFTe (for performing Fourier transforms on uniform meshes) ## Other diff --git a/doc/sphinx/src/chapters/advanced_features.rst b/doc/sphinx/src/chapters/advanced_features.rst index 863bd9d70e556..80992e928ed28 100644 --- a/doc/sphinx/src/chapters/advanced_features.rst +++ b/doc/sphinx/src/chapters/advanced_features.rst @@ -12,3 +12,4 @@ and solvers. ../interface/sparse ../integrators ../solvers + ../fourier_transforms \ No newline at end of file diff --git a/doc/sphinx/src/fourier_transforms.rst b/doc/sphinx/src/fourier_transforms.rst new file mode 100644 index 0000000000000..47165d50a3db0 --- /dev/null +++ b/doc/sphinx/src/fourier_transforms.rst @@ -0,0 +1,297 @@ +Fourier Transforms +================== + +Parthenon provides infrastructure for performing distributed Fast Fourier Transforms (FFTs) +on uniform meshes via the :cpp:class:`FFTManager` and :cpp:class:`UniformGridHelper` classes. +These are built on top of `heFFTe `_ and support +both CPU and GPU backends transparently. + +.. note:: + FFT functionality requires ``num_packs = 1`` in the input file, meaning all meshblocks + on a rank are packed into a single partition. This is required for the flat array indexing + used by heFFTe. Similarly, MPI is required (as it is a dependency of heFFTe). + +Overview +-------- + +The FFT infrastructure consists of two classes that work together: + +* :cpp:class:`FFTManager` — manages the FFT plan and performs forward/backward transforms +* :cpp:class:`UniformGridHelper` — provides mesh layout information and helper functions + for mapping between Parthenon's meshblock-based data layout and the flat arrays required + by heFFTe + +Both are accessible via the :cpp:class:`Mesh` object: + +.. code-block:: cpp + + auto fftManager = pmesh->GetFFTManager(); + auto uniformGridHelper = pmesh->GetUniformGridHelper(); + +Both classes use the :cpp:struct:`Box3D` struct to describe spatial extents: + +.. code-block:: cpp + + struct Box3D { + int low[3]; // lower bound in each dimension + int high[3]; // upper bound in each dimension + int size[3]; // size in each dimension: high - low + 1 + }; + +Normalization Convention +------------------------ + +The forward transform applies a :math:`1/N^3` normalization, and the backward transform +applies no normalization. This means the round-trip (forward followed by backward) recovers +the original field exactly, and Parseval's theorem reads: + +.. math:: + + \sum_{\mathbf{k}} |\hat{f}(\mathbf{k})|^2 = \frac{1}{N^3} \sum_{\mathbf{x}} |f(\mathbf{x})|^2 + +Physical wavenumbers are related to integer mode numbers by :math:`k_\mathrm{phys} = 2\pi k / L`, +assuming a periodic domain of size :math:`L`. + +Backends +-------- + +The backend is selected automatically at compile time based on the Kokkos execution space: + +* **GPU** (CUDA/HIP): uses heFFTe's GPU backend +* **CPU**: uses heFFTe's CPU backend (FFTW or MKL if available, otherwise stock) + +No code changes are required to switch between backends. + +Basic Usage +----------- + +The following example demonstrates the complete workflow for performing a forward and +backward FFT of a scalar field registered in Parthenon. + +Allocating arrays +~~~~~~~~~~~~~~~~~ + +FFT input and output arrays are standard Parthenon device arrays. The sizes are provided +by :cpp:class:`FFTManager`: + +.. code-block:: cpp + + const auto fft_size_inbox = fftManager->size_real_space_box(); + const auto fft_size_outbox = fftManager->size_fourier_space_box(); + + parthenon::ParArray1D input("input", fft_size_inbox); + parthenon::ParArray1D> output("output", fft_size_outbox); + parthenon::ParArray1D result("result", fft_size_inbox); + +Note that complex arrays must use Kokkos::complex, not std::complex, so that complex arithmetic is possible in Kokkos kernels. + +Gathering a field from the mesh +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +:cpp:func:`UniformGridHelper::GatherField` copies a single component of a named Parthenon +variable into a flat array suitable for FFT input: + +.. code-block:: cpp + + // Gather component 0 of "cons" into the input array + uniformGridHelper->GatherField("cons", 0, input); + +For derived quantities that require computation, +use a custom gather loop with :cpp:func:`UniformGridHelper::GetKernelHelper`: + +.. code-block:: cpp + + auto &md = pmesh->mesh_data.Get(); + auto cons = md->PackVariables(std::vector{"cons"}); + + auto &mbb = uniformGridHelper->MeshBlockBox; // interior cell bounds within a meshblock + + auto helper = uniformGridHelper->GetKernelHelper(); + + parthenon::par_for( + "GatherVelocity", 0, md->NumBlocks() - 1, + mbb.low[2], mbb.high[2], + mbb.low[1], mbb.high[1], + mbb.low[0], mbb.high[0], + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + const auto idx = helper.FlatIndex(b, k, j, i); + input(idx) = cons(b, 1, k, j, i) / cons(b, 0, k, j, i); + }); + +Performing the transforms +~~~~~~~~~~~~~~~~~~~~~~~~~ + +:cpp:func:`FFTManager::Forward` and :cpp:func:`FFTManager::Backward` operate on raw +device pointers: + +.. code-block:: cpp + + // Forward FFT (applies 1/N^3 normalization) + fftManager->Forward(input.data(), output.data()); + + // ... process output in Fourier space ... + + // Backward FFT (no normalization) + fftManager->Backward(output.data(), result.data()); + +Processing in Fourier space +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The local Fourier space box is accessible via :cpp:func:`FFTManager::fourier_space_box`. +Use :cpp:func:`FFTManager::GetKernelHelper` to obtain a device-copyable helper that +provides ``FourierFlatIndex`` and ``Wavevector``: + +.. code-block:: cpp + + auto fft_helper = fftManager->GetKernelHelper(); + auto outbox = fftManager->fourier_space_box(); + + parthenon::par_for( + "FourierSpaceKernel", + outbox.low[2], outbox.high[2], + outbox.low[1], outbox.high[1], + outbox.low[0], outbox.high[0], + KOKKOS_LAMBDA(const int kx3_idx, const int kx2_idx, const int kx1_idx) { + + const auto idx = fft_helper.FourierFlatIndex(kx3_idx, kx2_idx, kx1_idx); + + // integer wavevector components (negative frequencies unwrapped) + auto [kx3, kx2, kx1] = fft_helper.Wavevector(kx3_idx, kx2_idx, kx1_idx); + + // ... process output[idx] ... + }); + +.. note:: + The r2c transform only stores modes with :math:`k_{x1} \geq 0`. When computing + quantities like the power spectrum, modes with :math:`0 < k_{x1} < n_{x1}/2` must be + counted twice to account for Hermitian symmetry: + + .. code-block:: cpp + + const auto fac = ((kx1 > 0) && (2 * kx1 != nx1)) ? 2.0 : 1.0; + +Scattering a field back to the mesh +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +:cpp:func:`UniformGridHelper::ScatterField` copies a flat array back to a named +Parthenon variable on the mesh: + +.. code-block:: cpp + + parthenon::ParArray1D result("result", fft_size_inbox); + // ... fill result ... + uniformGridHelper->ScatterField(result, "my_derived_field", 0); + +The variable must be registered in the package before use (see :ref:`state`). + +API Reference +------------- + +Box3D +~~~~~ + +.. code-block:: cpp + + struct Box3D { + int low[3]; // lower index bound in each dimension + int high[3]; // upper index bound in each dimension + int size[3]; // size = high - low + 1 + }; + +FFTManager +~~~~~~~~~~ + +.. code-block:: cpp + + // Forward r2c FFT. Applies 1/N^3 normalization. + void Forward(const double* input, Kokkos::complex* output); + + // Backward c2r FFT. Applies no normalization. + void Backward(const Kokkos::complex* input, double* output); + + // Returns the local Fourier-space box (global Fourier indices) + Box3D fourier_space_box() const; + + // Returns the local real-space box (global cell indices) + Box3D real_space_box() const; + + // Total number of points in the local Fourier/real space box + std::size_t size_fourier_space_box() const; + std::size_t size_real_space_box() const; + + // Returns a device-copyable helper for use in Kokkos kernels. + // Capture by value in KOKKOS_LAMBDA. + KernelHelper GetKernelHelper() const; + +FFTManager::KernelHelper +~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + // Flat index into the local Fourier-space array + KOKKOS_INLINE_FUNCTION + std::int64_t FourierFlatIndex(const int k, const int j, const int i) const; + + // Flat index into the local real-space array + KOKKOS_INLINE_FUNCTION + std::int64_t RealFlatIndex(const int k, const int j, const int i) const; + + // Integer wavevector components (handles negative frequency unwrapping). + // For r2c transforms, kx >= 0 always. + // Returns {kx, ky, kz}. + KOKKOS_INLINE_FUNCTION + std::array Wavevector(const int k, const int j, const int i) const; + +UniformGridHelper +~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + // Gather a single component of a named variable into a flat FFT-ready array. + // output must be pre-allocated with size >= size_real_space_box() + void GatherField(const std::string &var_name, + const int var_index, + parthenon::ParArray1D &output); + + // Scatter a flat array back to a named variable on the mesh. + void ScatterField(const parthenon::ParArray1D &input, + const std::string &var_name, + const int var_index); + + // Returns a device-copyable helper for use in Kokkos kernels. + // Capture by value in KOKKOS_LAMBDA. + KernelHelper GetKernelHelper() const; + + // Local real-space box (global cell indices of this rank's domain) + Box3D LocalMeshBox; + + // Per-meshblock box (interior cell bounds within a single meshblock) + Box3D MeshBlockBox; + +UniformGridHelper::KernelHelper +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + // Flat index into the local real-space FFT array. + // Call from within a par_for loop over blocks and interior cells. + KOKKOS_INLINE_FUNCTION + std::int64_t FlatIndex(int b, int k, int j, int i) const; + +Limitations +----------- + +* Only uniform grids are supported. AMR is not compatible with the current FFT infrastructure. +* ``num_packs = 1`` is required (one partition per rank). +* Only cubic domains are fully supported for physical wavenumber calculations. + Non-cubic domains work for the FFT itself but wavenumber scaling must be handled manually. +* The r2c transform stores only modes with :math:`k_x \geq 0`, consistent with heFFTe's + default convention. +* Currently only 3D transforms are supported. + +See Also +-------- + +* :doc:`/src/interface/state` — registering variables for use with ``GatherField``/``ScatterField`` +* `heFFTe documentation `_ +* Fourier transform example: ``example/fourier_transform/`` diff --git a/example/CMakeLists.txt b/example/CMakeLists.txt index 1475ae8e1db29..d5d8cb86b83d8 100644 --- a/example/CMakeLists.txt +++ b/example/CMakeLists.txt @@ -24,3 +24,7 @@ add_subdirectory(poisson) add_subdirectory(poisson_gmg) add_subdirectory(diffusion) add_subdirectory(sparse_advection) +if(PARTHENON_ENABLE_FFT) + add_subdirectory(fourier_transform) +endif() + diff --git a/example/fourier_transform/CMakeLists.txt b/example/fourier_transform/CMakeLists.txt new file mode 100644 index 0000000000000..454e943259607 --- /dev/null +++ b/example/fourier_transform/CMakeLists.txt @@ -0,0 +1,15 @@ +#========================================================================================= +# Parthenon performance portable AMR framework +# Copyright(C) 2026 The Parthenon collaboration +# Licensed under the 3-clause BSD License, see LICENSE file for details +#======================================================================================== + +get_property(DRIVER_LIST GLOBAL PROPERTY DRIVERS_USED_IN_TESTS) +if(NOT PARTHENON_DISABLE_EXAMPLES AND PARTHENON_ENABLE_FFT) + add_executable( + fourier-transform + fourier_driver.cpp + fourier_driver.hpp + ) + target_link_libraries(fourier-transform PRIVATE Parthenon::parthenon) +endif() diff --git a/example/fourier_transform/fourier_driver.cpp b/example/fourier_transform/fourier_driver.cpp new file mode 100644 index 0000000000000..7e60fe8708c9b --- /dev/null +++ b/example/fourier_transform/fourier_driver.cpp @@ -0,0 +1,219 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#include +#include +#include +#include +#include + +#include + +#include "fourier_driver.hpp" +#include "utils/calc_spectrum.hpp" + +using namespace parthenon::driver::prelude; +using fourier_transform::FourierDriver; + +Packages_t ProcessPackages(std::unique_ptr &pin); +void FillTestField(MeshBlock *pmb, ParameterInput *pin); + +int main(int argc, char *argv[]) { + ParthenonManager pman; + pman.app_input->ProcessPackages = ProcessPackages; + pman.app_input->ProblemGenerator = FillTestField; + + auto manager_status = pman.ParthenonInitEnv(argc, argv); + if (manager_status == ParthenonStatus::complete) { + pman.ParthenonFinalize(); + return 0; + } + if (manager_status == ParthenonStatus::error) { + pman.ParthenonFinalize(); + return 1; + } + + pman.ParthenonInitPackagesAndMesh(); + { + FourierDriver driver(pman.pinput.get(), pman.app_input.get(), pman.pmesh.get()); + driver.Execute(); + } + pman.ParthenonFinalize(); + return 0; +} + +// Initialize a simple test field. Note that FFTs only work on a uniform grid, no AMR. +void FillTestField(MeshBlock *pmb, ParameterInput *pin) { + auto &mbd = pmb->meshblock_data.Get(); + auto field = mbd->Get("test_field").data; + auto vec_field = mbd->Get("test_vector_field").data; + + IndexRange ib = pmb->cellbounds.GetBoundsI(IndexDomain::interior); + IndexRange jb = pmb->cellbounds.GetBoundsJ(IndexDomain::interior); + IndexRange kb = pmb->cellbounds.GetBoundsK(IndexDomain::interior); + auto &coords = pmb->coords; + + pmb->par_for( + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int k, const int j, const int i) { + // smooth pattern with some "perturbations" + field(k, j, i) = Kokkos::sin(coords.Xc<1>(i)) * Kokkos::sin(coords.Xc<2>(j)) * + Kokkos::sin(coords.Xc<3>(k)) + + 0.1 * (k + j + i) / (kb.e + jb.e + ib.e); + vec_field(0, k, j, i) = + Kokkos::sin(coords.Xc<1>(i)) * Kokkos::sin(coords.Xc<2>(j)) * + Kokkos::sin(coords.Xc<3>(k)) + + 0.1 * (k + j + i) / (kb.e + jb.e + ib.e) + 1; // added mean component + vec_field(1, k, j, i) = Kokkos::sin(coords.Xc<1>(i)) * + Kokkos::cos(coords.Xc<2>(j)) * + Kokkos::sin(coords.Xc<3>(k)) + + 0.1 * (k + j + i) / (kb.e + jb.e + ib.e); + vec_field(2, k, j, i) = Kokkos::sin(coords.Xc<1>(i)) * + Kokkos::sin(coords.Xc<2>(j)) * + Kokkos::cos(coords.Xc<3>(k)) + + 0.1 * (k + j + i) / (kb.e + jb.e + ib.e); + }); +} + +Packages_t ProcessPackages(std::unique_ptr &pin) { + Packages_t packages; + + auto package = std::make_shared("fourier_transform"); + + // Register a scalar field for FFT round-trip test + parthenon::Metadata m({parthenon::Metadata::Cell, parthenon::Metadata::Derived, + parthenon::Metadata::OneCopy}); + package->AddField("test_field", m); + + // Register a vector field for FFT round-trip test + m = parthenon::Metadata({parthenon::Metadata::Cell, parthenon::Metadata::Derived, + parthenon::Metadata::OneCopy}, + std::vector({3})); + package->AddField("test_vector_field", m); + + packages.Add(package); + return packages; +} + +parthenon::DriverStatus FourierDriver::Execute() { + PreExecute(); + + auto &md = pmesh->mesh_data.Get(); + auto UniformGridHelper = + pmesh->GetUniformGridHelper(); // Helper class used to map block-local indices to a + // flat mesh index + auto FFTManager = pmesh->GetFFTManager(); // Class that holds and executes FFT plans + + // define input and output arrays for FFT: + parthenon::ParArray1D input("fft input", FFTManager->size_real_space_box()); + parthenon::ParArray1D> output( + "fft output", FFTManager->size_fourier_space_box()); + // also pre-allocate array for the recovered field after inverse FFT, to check + // round-trip accuracy: + parthenon::ParArray1D input_recovered("fft input recovered", + FFTManager->size_real_space_box()); + + auto test_field = md->PackVariables(std::vector{"test_field"}); + + // Gather block data into flat array for FFT input: + UniformGridHelper->GatherField("test_field", 0, input); + + // Perform forward FFT - applies 1/N^3 normalization: + FFTManager->Forward(input.data(), output.data()); + // Perform inverse FFT - applies no normalization: + FFTManager->Backward(output.data(), input_recovered.data()); + + // Check round-trip accuracy (get max difference across all points): + Real local_max_error = 0.0; + Kokkos::parallel_reduce( + "ComputeError", Kokkos::RangePolicy<>(0, FFTManager->size_real_space_box()), + KOKKOS_LAMBDA(const int idx, Real &max_err) { + Real diff = Kokkos::abs(input_recovered(idx) - input(idx)); + if (diff > max_err) max_err = diff; + }, + Kokkos::Max(local_max_error)); + + // Reduce across MPI ranks to get global maximum + Real max_error = local_max_error; +#ifdef MPI_PARALLEL + PARTHENON_MPI_CHECK(MPI_Allreduce(&local_max_error, &max_error, 1, MPI_PARTHENON_REAL, + MPI_MAX, MPI_COMM_WORLD)); +#endif + + if (parthenon::Globals::my_rank == 0) { + std::cout << "Max relative error after FFT round-trip: " << max_error << std::endl; + } + + // Now test the spectrum machinery + auto spectrum = + parthenon::utils::fft::CalcSpectrum(pmesh, "test_vector_field", {0, 1, 2}); + const auto spectrum_h = spectrum.GetHostMirrorAndCopy(); + + auto test_vector_field_pack = + md->PackVariables(std::vector{"test_vector_field"}); + IndexRange ib = md->GetBlockData(0)->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBlockData(0)->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBlockData(0)->GetBoundsK(IndexDomain::interior); + + auto mesh_size = pmesh->mesh_size; + const auto Nx = mesh_size.nx(parthenon::X1DIR); + const auto Ny = mesh_size.nx(parthenon::X2DIR); + const auto Nz = mesh_size.nx(parthenon::X3DIR); + + // Sanity checks (compare power in real space to spectral space power) + using parthenon::utils::fft::SpecReal; + Kokkos::Array sums{{0.0, 0.0, 0.0, 0.0}}; + Kokkos::parallel_reduce( + "fieldsqrd_sum", + Kokkos::MDRangePolicy>( + {0, kb.s, jb.s, ib.s}, + {test_vector_field_pack.GetDim(5), kb.e + 1, jb.e + 1, ib.e + 1}, + {1, 1, 1, ib.e + 1 - ib.s}), + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i, + SpecReal &sum_usqr, SpecReal &sum_u1, SpecReal &sum_u2, + SpecReal &sum_u3) { + const auto u1 = static_cast(test_vector_field_pack(b, 0, k, j, i)); + const auto u2 = static_cast(test_vector_field_pack(b, 1, k, j, i)); + const auto u3 = static_cast(test_vector_field_pack(b, 2, k, j, i)); + sum_u1 += u1; + sum_u2 += u2; + sum_u3 += u3; + sum_usqr += SQR(u1) + SQR(u2) + SQR(u3); + }, + sums[0], sums[1], sums[2], sums[3]); + +#ifdef MPI_PARALLEL + PARTHENON_MPI_CHECK( + MPI_Allreduce(MPI_IN_PLACE, sums.data(), 4, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD)); +#endif + const auto norm = + static_cast(Nx) * static_cast(Ny) * static_cast(Nz); + sums[0] /= norm; + sums[1] /= norm; + sums[2] /= norm; + sums[3] /= norm; + + // Sum power in spectrum + SpecReal spec_sum = 0.0; + for (int i = 0; i < static_cast(spectrum_h.extent(0)); i++) { + spec_sum += spectrum_h(i, 0); + } + if (parthenon::Globals::my_rank == 0) { + std::cout << "sum u^2=" << sums[0] << " sum uhat^2=" << spec_sum + << " ^2=" << SQR(sums[1]) + SQR(sums[2]) + SQR(sums[3]) + << " uhat(0)^2=" << spectrum_h(0, 0) << " sum u_1=" << sums[1] + << " sum u_2=" << sums[2] << " sum u_3=" << sums[3] << "\n"; + std::cout << std::format( + "Error in spectrum total power: {:.15e}\nError in spectrum mean: {:.15e}\n", + std::abs(sums[0] / spec_sum - 1.0), + std::abs((SQR(sums[1]) + SQR(sums[2]) + SQR(sums[3])) / spectrum_h(0, 0) - 1.0)); + } + + Driver::PostExecute(DriverStatus::complete); + return DriverStatus::complete; +} diff --git a/example/fourier_transform/fourier_driver.hpp b/example/fourier_transform/fourier_driver.hpp new file mode 100644 index 0000000000000..e6d2fbac5e07a --- /dev/null +++ b/example/fourier_transform/fourier_driver.hpp @@ -0,0 +1,45 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#ifndef EXAMPLE_FOURIER_TRANSFORM_FOURIER_DRIVER_HPP_ +#define EXAMPLE_FOURIER_TRANSFORM_FOURIER_DRIVER_HPP_ + +#include +#include + +#include + +namespace fourier_transform { +using namespace parthenon::driver::prelude; + +/** + * @brief Constructs a driver to demonstrate the use of Fourier transforms in Parthenon. + * The driver will compute the Fourier transform of a 3D array, compute the inverse + * Fourier transform to recover the original array, and compute the maximum error between + * the original and recovered arrays. + */ +class FourierDriver : public Driver { + public: + FourierDriver(ParameterInput *pin, ApplicationInput *app_in, Mesh *pm) + : Driver(pin, app_in, pm) { + InitializeOutputs(); + } + + /// MakeTaskList and MakeTasks aren't virtual routines on `Driver`, + // but each driver is expected to implement at least one of them. + /// TaskList MakeTaskList(MeshBlock *pmb); + template + TaskCollection MakeTaskCollection(T &blocks); + + /// `Execute` cycles until simulation completion. + DriverStatus Execute() override; +}; + +} // namespace fourier_transform + +#endif // EXAMPLE_FOURIER_TRANSFORM_FOURIER_DRIVER_HPP_ diff --git a/example/fourier_transform/parthinput.example b/example/fourier_transform/parthinput.example new file mode 100644 index 0000000000000..fe1bb0ba36b92 --- /dev/null +++ b/example/fourier_transform/parthinput.example @@ -0,0 +1,43 @@ +# ======================================================================================== +# Parthenon performance portable AMR framework +# Copyright(C) 2026 The Parthenon collaboration +# Licensed under the 3-clause BSD License, see LICENSE file for details +# ======================================================================================== + + +problem_id = fourier_transform_example + + +refinement = none +nghost = 3 + +nx1 = 16 +x1min = -2.0 +x1max = 2.0 +ix1_bc = periodic +ox1_bc = periodic + +nx2 = 16 +x2min = -2.0 +x2max = 2.0 +ix2_bc = periodic +ox2_bc = periodic + +nx3 = 16 +x3min = -2.0 +x3max = 2.0 +ix3_bc = periodic +ox3_bc = periodic + +# How many meshblocks to use in a premade default kernel. +# A value of <1 means use the whole mesh. +num_packs = 1 + + +nx1 = 8 +nx2 = 8 +nx3 = 8 + + +file_type = hdf5 +variables = in_or_out diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 2c6e292b7adff..93ce46a8c1962 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,4 +1,8 @@ #========================================================================================= +# Parthenon performance portable AMR framework +# Copyright(C) 2020-2026 The Parthenon collaboration +# Licensed under the 3-clause BSD License, see LICENSE file for detail +#========================================================================================= # (C) (or copyright) 2020-2026. Triad National Security, LLC. All rights reserved. # # This program was produced under U.S. Government contract 89233218CNA000001 for Los @@ -344,6 +348,20 @@ add_library(parthenon ) add_library(Parthenon::parthenon ALIAS parthenon) +if(PARTHENON_ENABLE_FFT) + target_sources(parthenon PRIVATE + utils/fft_manager.cpp + utils/fft_manager.hpp + utils/uniform_grid_helper.cpp + utils/uniform_grid_helper.hpp + utils/calc_spectrum.cpp + utils/calc_spectrum.hpp + outputs/spectrum.cpp + ) + target_link_libraries(parthenon PRIVATE Heffte::Heffte) + target_compile_definitions(parthenon PUBLIC PARTHENON_ENABLE_FFT) +endif() + set_target_properties(parthenon PROPERTIES SOVERSION ${parthenon_VERSION}) target_compile_features(parthenon PUBLIC cxx_std_20) diff --git a/src/mesh/mesh.hpp b/src/mesh/mesh.hpp index 95c79f8946efb..185a4e2102ceb 100644 --- a/src/mesh/mesh.hpp +++ b/src/mesh/mesh.hpp @@ -59,6 +59,11 @@ #include "utils/object_pool.hpp" #include "utils/partition_stl_containers.hpp" +#ifdef PARTHENON_ENABLE_FFT +#include "utils/fft_manager.hpp" +#include "utils/uniform_grid_helper.hpp" +#endif + namespace parthenon { // Forward declarations @@ -84,6 +89,7 @@ class Mesh { friend class HistoryOutput; friend class MeshBlock; friend class MeshRefinement; + friend class FFTManager; struct base_constructor_selector_t {}; Mesh(ParameterInput *pin, ApplicationInput *app_in, Packages_t &packages, @@ -93,6 +99,25 @@ class Mesh { hyper_rectangular_constructor_selector_t); public: +#ifdef PARTHENON_ENABLE_FFT + std::unique_ptr fft_manager; + + FFTManager *GetFFTManager() { + if (!fft_manager) { + fft_manager = std::make_unique(this); + } + return fft_manager.get(); + } + + std::unique_ptr uniform_grid_helper; + + UniformGridHelper *GetUniformGridHelper() { + if (!uniform_grid_helper) { + uniform_grid_helper = std::make_unique(this); + } + return uniform_grid_helper.get(); + } +#endif // 2x function overloads of ctor: normal and restarted simulation Mesh(ParameterInput *pin, ApplicationInput *app_in, Packages_t &packages, int test_flag = 0); diff --git a/src/outputs/outputs.cpp b/src/outputs/outputs.cpp index 1d35462dab64a..ba4d25c3092e2 100644 --- a/src/outputs/outputs.cpp +++ b/src/outputs/outputs.cpp @@ -361,6 +361,17 @@ Outputs::Outputs(Mesh *pm, ParameterInput *pin, SimTime *tm) { pnew_type = std::make_shared(op); } else if (op.file_type == "ascent") { pnew_type = std::make_shared(op); + } else if (op.file_type == "spectrum") { +#ifdef PARTHENON_ENABLE_FFT + pnew_type = std::make_shared(op); +#else + msg << "### FATAL ERROR in Outputs constructor" << std::endl + << "Executable not configured for Fourier transforms, but file format " + << "spectrum is requested in output block '" << op.block_name << "'. " + << "You can disable this block without deleting it by setting a dt < 0." + << std::endl; + PARTHENON_FAIL(msg); +#endif // ifdef PARTHENON_ENABLE_FFT } else if (op.file_type == "openpmd") { #ifdef PARTHENON_ENABLE_OPENPMD const auto backend_config = diff --git a/src/outputs/outputs.hpp b/src/outputs/outputs.hpp index 04cc85049b5b3..07c4f08b52598 100644 --- a/src/outputs/outputs.hpp +++ b/src/outputs/outputs.hpp @@ -115,6 +115,17 @@ using HstVec_list = std::vector; const char hist_param_key[] = "HistoryFunctions"; const char hist_vec_param_key[] = "HistoryVectorFunctions"; +//---------------------------------------------------------------------------------------- +//! \class SpectralOutput +// \brief derived OutputType class for Spectrum dumps + +class SpectralOutput : public OutputType { + public: + explicit SpectralOutput(const OutputParameters &oparams) : OutputType(oparams) {} + void WriteOutputFile(Mesh *pm, ParameterInput *pin, SimTime *tm, + const SignalHandler::OutputSignal signal) override; +}; + //---------------------------------------------------------------------------------------- //! \class HistoryFile // \brief derived OutputType class for history dumps diff --git a/src/outputs/spectrum.cpp b/src/outputs/spectrum.cpp new file mode 100644 index 0000000000000..2ac387cb7b4a4 --- /dev/null +++ b/src/outputs/spectrum.cpp @@ -0,0 +1,68 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#include +#include +#include +#include + +#include "globals.hpp" +#include "mesh/mesh.hpp" +#include "outputs/outputs.hpp" +#include "utils/calc_spectrum.hpp" +#include "utils/error_checking.hpp" + +namespace parthenon { + +//---------------------------------------------------------------------------------------- +//! \fn void SpectralOutput::WriteOutputFile() +// \brief Writes a spectrum output file + +void SpectralOutput::WriteOutputFile(Mesh *pm, ParameterInput *pin, SimTime *tm, + const SignalHandler::OutputSignal signal) { + const auto var_name = pin->GetString(output_params.block_name, "variable"); + const auto components = pin->GetVector(output_params.block_name, "components"); + const auto output_label = + pin->GetOrAddString(output_params.block_name, "output_label", var_name); + + auto spectra = utils::fft::CalcSpectrum(pm, var_name, components); + auto spectra_h = spectra.GetHostMirrorAndCopy(); + const auto num_bins = spectra_h.extent(0); + + if (parthenon::Globals::my_rank == 0) { + std::string suffix; + if (signal == SignalHandler::OutputSignal::now) { + suffix = "now"; + } else if (signal == SignalHandler::OutputSignal::final && + output_params.file_label_final) { + suffix = "final"; + } else { + suffix = std::format("{:0{}d}", output_params.file_number, + output_params.file_number_width); + } + + const std::string fname = std::format("{}.{}.{}.{}.spc", output_params.file_basename, + output_label, output_params.file_id, suffix); + + std::ofstream fout(fname); + if (!fout.is_open()) { + PARTHENON_FAIL("Could not open " + fname + " for writing"); + } + + fout << "# Bin val_sum K_sum Count\n"; + for (int i = 0; i < static_cast(num_bins); ++i) { + fout << std::format("{:d} {:.15e} {:.15e} {:.15e}\n", i, spectra_h(i, 0), + spectra_h(i, 1), spectra_h(i, 2)); + } + fout.close(); + } + + UpdateNextOutput_(pm, tm); +} + +} // namespace parthenon diff --git a/src/utils/calc_spectrum.cpp b/src/utils/calc_spectrum.cpp new file mode 100644 index 0000000000000..89e6412aa81f0 --- /dev/null +++ b/src/utils/calc_spectrum.cpp @@ -0,0 +1,149 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#include "utils/calc_spectrum.hpp" + +#include +#include +#include + +#include "defs.hpp" +#include "globals.hpp" +#include "mesh/mesh.hpp" +#include "parthenon_arrays.hpp" +#include "utils/error_checking.hpp" +#include "utils/fft_manager.hpp" +#include "utils/uniform_grid_helper.hpp" + +namespace parthenon { +namespace utils { +namespace fft { +parthenon::ParArray2D +CalcSpectrum(Mesh *pm, const parthenon::ParArray1D &input, const int n_comp) { + PARTHENON_REQUIRE_THROWS(pm != nullptr, "CalcSpectrum: mesh pointer must not be null"); + PARTHENON_REQUIRE_THROWS(n_comp > 0, "CalcSpectrum: n_comp must be positive"); + PARTHENON_REQUIRE_THROWS(pm->DefaultNumPartitions() == 1, + "Only num_packs=1 currently supported for heffte.") + + auto FFTManager = pm->GetFFTManager(); + const auto fft_size_inbox = FFTManager->size_real_space_box(); + const auto expected_input_size = static_cast(n_comp) * fft_size_inbox; + PARTHENON_REQUIRE_THROWS(input.size() == expected_input_size, + "CalcSpectrum: input array has size " + + std::to_string(input.size()) + + ", but expected n_comp * size_real_space_box() = " + + std::to_string(expected_input_size)); + + auto mesh_size = pm->mesh_size; + const auto nx = mesh_size.nx(X1DIR); + const auto ny = mesh_size.nx(X2DIR); + const auto nz = mesh_size.nx(X3DIR); + + parthenon::ParArray1D> output( + "fft output", n_comp * FFTManager->size_fourier_space_box()); + + for (int i = 0; i < n_comp; i++) { + FFTManager->Forward(input.data() + i * FFTManager->size_real_space_box(), + output.data() + i * FFTManager->size_fourier_space_box()); + } + + const auto k_max = std::sqrt(SQR(nx / 2) + SQR(ny / 2) + SQR(nz / 2)); + const auto num_bins = static_cast(std::ceil(k_max)) + 1; + + parthenon::ParArray2D spectra("spectra", num_bins, 3); + auto scatter_spectra = + Kokkos::Experimental::ScatterView( + spectra.KokkosView()); + + auto fb = FFTManager->fourier_space_box(); + + const auto fft_size_outbox = FFTManager->size_fourier_space_box(); + auto kernel_helper = FFTManager->GetKernelHelper(); + parthenon::par_for( + "CalcSpec", fb.low[2], fb.high[2], fb.low[1], fb.high[1], fb.low[0], fb.high[0], + KOKKOS_LAMBDA(const int k, const int j, const int i) { + auto k_vec = kernel_helper.Wavevector(k, j, i); + auto k_mag = std::sqrt(SQR(k_vec[0]) + SQR(k_vec[1]) + SQR(k_vec[2])); + auto k_mag_int = static_cast(std::floor(k_mag)); + const auto outidx = kernel_helper.FourierFlatIndex(k, j, i); + auto val = 0.0; + for (int n = 0; n < n_comp; n++) { + val += SQR(output[outidx + n * fft_size_outbox].real()) + + SQR(output[outidx + n * fft_size_outbox].imag()); + } + const auto fac = ((k_vec[2] > 0) && (2 * k_vec[2] != nx)) ? 2.0 : 1.0; + auto spec = scatter_spectra.access(); + spec(k_mag_int, 0) += fac * val; + spec(k_mag_int, 1) += fac * k_mag; + spec(k_mag_int, 2) += fac * 1.0; + }); + + Kokkos::Experimental::contribute(spectra.KokkosView(), scatter_spectra); + Kokkos::fence(); + +#ifdef MPI_PARALLEL + PARTHENON_REQUIRE_THROWS(sizeof(SpecReal) == sizeof(double), + "Need to fix comm data types manually."); + if (parthenon::Globals::my_rank == 0) { + PARTHENON_MPI_CHECK(MPI_Reduce(MPI_IN_PLACE, spectra.data(), spectra.size(), + MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD)); + } else { + PARTHENON_MPI_CHECK(MPI_Reduce(spectra.data(), spectra.data(), spectra.size(), + MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD)); + } +#endif + + return spectra; +} + +parthenon::ParArray2D CalcSpectrum(Mesh *pm, const std::string &var_name, + const std::vector &components) { + PARTHENON_REQUIRE_THROWS(pm != nullptr, "CalcSpectrum: mesh pointer must not be null"); + PARTHENON_REQUIRE_THROWS(!components.empty(), + "CalcSpectrum: at least one component is required"); + + auto &md = pm->mesh_data.Get(); + + IndexRange ib = md->GetBlockData(0)->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBlockData(0)->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBlockData(0)->GetBoundsK(IndexDomain::interior); + + auto vars = md->PackVariables(std::vector{var_name}); + + const int n_comp = components.size(); + auto FFTManager = pm->GetFFTManager(); + const auto fft_size_inbox = FFTManager->size_real_space_box(); + parthenon::ParArray1D input("fft input", n_comp * fft_size_inbox); + + parthenon::ParArray1D components_d("components", components.size()); + auto components_h = components_d.GetHostMirror(); + for (int n = 0; n < n_comp; n++) { + PARTHENON_REQUIRE_THROWS(components[n] >= 0 && components[n] < vars.GetDim(4), + "CalcSpectrum: component " + std::to_string(components[n]) + + " out of range for variable '" + var_name + "'"); + components_h(n) = components[n]; + } + components_d.DeepCopy(components_h); + + auto UniformGridHelper = pm->GetUniformGridHelper(); + auto helper = UniformGridHelper->GetKernelHelper(); + + par_for( + "Init FFT fields", 0, md->NumBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + const auto idx = helper.FlatIndex(b, k, j, i); + for (int n = 0; n < n_comp; n++) { + input(n * fft_size_inbox + idx) = vars(b, components_d(n), k, j, i); + } + }); + + return CalcSpectrum(pm, input, n_comp); +} +} // namespace fft +} // namespace utils +} // namespace parthenon diff --git a/src/utils/calc_spectrum.hpp b/src/utils/calc_spectrum.hpp new file mode 100644 index 0000000000000..fee5dde037585 --- /dev/null +++ b/src/utils/calc_spectrum.hpp @@ -0,0 +1,44 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#ifndef UTILS_CALC_SPECTRUM_HPP_ +#define UTILS_CALC_SPECTRUM_HPP_ + +#include +#include + +#include "parthenon_arrays.hpp" + +namespace parthenon { +class Mesh; +namespace utils { +namespace fft { + +// Keep spectra accumulations, MPI collectives, and output in double +// precision even when the application is built with Real=float. +using SpecReal = double; + +// Computes the shell-averaged power spectrum of the requested components of +// var_name on a uniform mesh. Returns a device-side array shaped [num_bins, 3]: +// col 0: power sum, col 1: wavenumber sum, col 2: bin count +// num_bins = ceil(k_max) + 1. An MPI_Reduce to rank 0 is performed internally, +// so only rank 0 holds meaningful data on return. +parthenon::ParArray2D CalcSpectrum(Mesh *pm, const std::string &var_name, + const std::vector &components); + +// Computes the shell-averaged power spectrum from a pre-gathered 1D real-space array. +// The input must contain n_comp contiguous fields, each of length +// FFTManager::size_real_space_box(). +parthenon::ParArray2D +CalcSpectrum(Mesh *pm, const parthenon::ParArray1D &input, int n_comp = 1); + +} // namespace fft +} // namespace utils +} // namespace parthenon + +#endif // UTILS_CALC_SPECTRUM_HPP_ diff --git a/src/utils/fft_manager.cpp b/src/utils/fft_manager.cpp new file mode 100644 index 0000000000000..74d9ff7c86266 --- /dev/null +++ b/src/utils/fft_manager.cpp @@ -0,0 +1,132 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#include +#include + +#include "fft_manager.hpp" +#include "heffte.h" +#include "mesh/mesh.hpp" + +namespace parthenon { + +struct FFTManager::Impl { +// @pgrete: Can the backend selection be made nicer? +#if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) + using BackendTag = heffte::backend::default_backend::type; +#elif defined(KOKKOS_ENABLE_SYCL) + static_assert(false, "heFFTe's SYCL backend is not yet tested with Parthenon. Please " + "test and enable this code."); +#else + using BackendTag = heffte::backend::default_backend::type; +#endif + + heffte::fft3d_r2c fft_plan; + ParArray1D> workspace_; + + Impl(const heffte::box3d<> &real_space_box, const heffte::box3d<> &fourier_space_box, + int r2c_direction, MPI_Comm comm, heffte::plan_options options) + : fft_plan(real_space_box, fourier_space_box, r2c_direction, comm, options), + workspace_("fft workspace", fft_plan.size_workspace()) {} +}; + +FFTManager::FFTManager(Mesh *mesh) : mesh_(mesh) { + auto UniformGridHelper = mesh_->GetUniformGridHelper(); + + auto mesh_size = mesh_->mesh_size; + nx1_ = mesh_size.nx(X1DIR); + nx2_ = mesh_size.nx(X2DIR); + nx3_ = mesh_size.nx(X3DIR); + + std::int64_t r2c_direction = 0; + + heffte::box3d<> real_indexes({0, 0, 0}, {nx1_ - 1, nx2_ - 1, nx3_ - 1}); + heffte::box3d<> complex_indexes({0, 0, 0}, {nx1_ / 2, nx2_ - 1, nx3_ - 1}); + + assert(real_indexes.r2c(r2c_direction) == complex_indexes); + + auto &mesh_start_idx = UniformGridHelper->local_mesh_box.low; + auto &mesh_end_idx = UniformGridHelper->local_mesh_box.high; + + const heffte::box3d<> real_space_box( + {mesh_start_idx[0], mesh_start_idx[1], mesh_start_idx[2]}, + {static_cast(mesh_end_idx[0]), static_cast(mesh_end_idx[1]), + static_cast(mesh_end_idx[2])}); + + std::array proc_grid = + heffte::proc_setup_min_surface(complex_indexes, parthenon::Globals::nranks); + + std::vector> complex_boxes = + heffte::split_world(complex_indexes, proc_grid); + + heffte::box3d<> const fourier_space_box = complex_boxes[parthenon::Globals::my_rank]; + + heffte::plan_options options = heffte::default_options(); + options.algorithm = heffte::reshape_algorithm::p2p_plined; + + impl_ = std::make_unique(real_space_box, fourier_space_box, r2c_direction, + MPI_COMM_WORLD, options); +} // FFTManager::FFTManager + +// ----------------------------- +// Forward / Backward +// ----------------------------- +void FFTManager::Forward(const Real *input, Kokkos::complex *output) { + impl_->fft_plan.forward( + // We have the interface use Kokkos:complex so that complex arithmetic in Kokkos + // kernels is possible, + // but heFFTe's interface uses std::complex, so we need to reinterpret_cast here. + input, reinterpret_cast *>(output), impl_->workspace_.data(), + heffte::scale::full); +} + +void FFTManager::Backward(const Kokkos::complex *input, Real *output) { + impl_->fft_plan.backward(reinterpret_cast *>(input), output, + impl_->workspace_.data(), heffte::scale::none); +} + +// ----------------------------- +// Boxes and sizes +// ----------------------------- +parthenon::Box3D FFTManager::fourier_space_box() const { + parthenon::Box3D box; + const auto &b = impl_->fft_plan.outbox(); + for (int i = 0; i < 3; i++) { + box.low[i] = b.low[i]; + box.high[i] = b.high[i]; + box.size[i] = b.high[i] - b.low[i] + 1; + } + return box; +} + +parthenon::Box3D FFTManager::real_space_box() const { + parthenon::Box3D box; + const auto &b = impl_->fft_plan.inbox(); + for (int i = 0; i < 3; i++) { + box.low[i] = b.low[i]; + box.high[i] = b.high[i]; + box.size[i] = b.high[i] - b.low[i] + 1; + } + return box; +} + +std::size_t FFTManager::size_fourier_space_box() const { + const auto box = fourier_space_box(); + return static_cast(box.size[0]) * static_cast(box.size[1]) * + static_cast(box.size[2]); +} + +std::size_t FFTManager::size_real_space_box() const { + const auto box = real_space_box(); + return static_cast(box.size[0]) * static_cast(box.size[1]) * + static_cast(box.size[2]); +} + +FFTManager::~FFTManager() = default; + +} // namespace parthenon diff --git a/src/utils/fft_manager.hpp b/src/utils/fft_manager.hpp new file mode 100644 index 0000000000000..35bd293b78249 --- /dev/null +++ b/src/utils/fft_manager.hpp @@ -0,0 +1,97 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#ifndef UTILS_FFT_MANAGER_HPP_ +#define UTILS_FFT_MANAGER_HPP_ + +#include +#include + +#include "parthenon_arrays.hpp" +#include "utils/uniform_grid_helper.hpp" + +namespace parthenon { + +class Mesh; + +class FFTManager { + friend class Mesh; + + public: + explicit FFTManager(Mesh *mesh); + ~FFTManager(); + + void Forward(const Real *input, Kokkos::complex *output); + + void Backward(const Kokkos::complex *input, Real *output); + + // ----------------------------- + // Box info + // ----------------------------- + + parthenon::Box3D fourier_space_box() const; + parthenon::Box3D real_space_box() const; + + std::size_t size_fourier_space_box() const; // total number of points + std::size_t size_real_space_box() const; + + // ----------------------------- + // Device-copyable kernel helper + // ----------------------------- + struct KernelHelper { + parthenon::Box3D fourier_box; + parthenon::Box3D real_box; + int nx1, nx2, nx3; + + // Flat index into the local Fourier-space array + KOKKOS_INLINE_FUNCTION + std::int64_t FourierFlatIndex(const int k, const int j, const int i) const { + return ((std::int64_t)(k - fourier_box.low[2]) * fourier_box.size[1] + + (j - fourier_box.low[1])) * + fourier_box.size[0] + + i - fourier_box.low[0]; + } + + // Flat index into the local real-space array + KOKKOS_INLINE_FUNCTION + std::int64_t RealFlatIndex(const int k, const int j, const int i) const { + return ((std::int64_t)(k - real_box.low[2]) * real_box.size[1] + + (j - real_box.low[1])) * + real_box.size[0] + + i - real_box.low[0]; + } + + // Integer wavevector components (handles negative frequencies) + // For r2c transforms, kx >= 0 always + KOKKOS_INLINE_FUNCTION + std::array Wavevector(const int k, const int j, const int i) const { + return {k <= nx3 / 2 ? k : k - nx3, j <= nx2 / 2 ? j : j - nx2, i}; + } + }; + + // Returns a device-copyable helper for use inside Kokkos kernels. + // Capture by value in KOKKOS_LAMBDA: + // auto helper = fftManager->GetKernelHelper(); + // par_for(..., KOKKOS_LAMBDA(...) { helper.FourierFlatIndex(...); }); + KernelHelper GetKernelHelper() const { + return {fourier_space_box(), real_space_box(), nx1_, nx2_, nx3_}; + } + + private: + struct Impl; // opaque implementation + std::unique_ptr impl_; // owns backend-specific data + + Mesh *mesh_; + + // Global mesh dimensions, stored during Initialize() + int nx1_ = 0, nx2_ = 0, nx3_ = 0; +}; + +} // namespace parthenon + +#endif // UTILS_FFT_MANAGER_HPP_ diff --git a/src/utils/uniform_grid_helper.cpp b/src/utils/uniform_grid_helper.cpp new file mode 100644 index 0000000000000..b0a734616dc48 --- /dev/null +++ b/src/utils/uniform_grid_helper.cpp @@ -0,0 +1,175 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#include +#include +#include +#include + +#include "mesh/mesh.hpp" +#include "utils/uniform_grid_helper.hpp" + +namespace parthenon { + +UniformGridHelper::UniformGridHelper(Mesh *mesh) : mesh_(mesh) { + loc_view = parthenon::ParArray2D("logical location of local blocks", + mesh_->GetNumMeshBlocksThisRank(), 3); + auto loc_view_h = loc_view.GetHostMirror(); + + const auto level = + mesh_->Forest().GetLegacyTreeLocation(mesh_->block_list[0]->loc).level(); + + std::array local_loc_min{ + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + }; + std::array local_loc_max{ + std::numeric_limits::min(), + std::numeric_limits::min(), + std::numeric_limits::min(), + }; + + // Set rank local min and max logical locations. + // Also check if all blocks are on the same level (we use this check instead of + // checking for refinement=none because AMR could have been used to dynamically refine + // a simulation. We just need to ensure that all blocks are on the same level to + // create an effective uniform grid.) + + // lkasselm: This is currently only checked once during initialization, + // but if the mesh is dynamically refined, this could be violated at later times. + // Needs to be checked at each call to Gather/Scatter. + + for (int b = 0; b < mesh_->GetNumMeshBlocksThisRank(); b++) { + auto pmb = mesh_->block_list[b]; + const auto loc = mesh_->Forest().GetLegacyTreeLocation(pmb->loc); + for (int i = 0; i <= 2; i++) { + local_loc_min.at(i) = std::min(loc.l(i), local_loc_min.at(i)); + local_loc_max.at(i) = std::max(loc.l(i), local_loc_max.at(i)); + loc_view_h(b, i) = loc.l(i); + } + PARTHENON_REQUIRE_THROWS( + loc.level() == level, + "uniform_grid_helper was initialized but not all blocks are on the same level."); + } + + // convert global logical locations to rank-local logical locs + for (int b = 0; b < mesh_->GetNumMeshBlocksThisRank(); b++) { + for (int i = 0; i <= 2; i++) { + loc_view_h(b, i) -= local_loc_min.at(i); + } + } + Kokkos::deep_copy(loc_view, loc_view_h); + + std::array local_nlocs{ + (local_loc_max.at(0) - local_loc_min.at(0)) + 1, + (local_loc_max.at(1) - local_loc_min.at(1)) + 1, + (local_loc_max.at(2) - local_loc_min.at(2)) + 1, + }; + const auto loc_max_vol = local_nlocs.at(0) * local_nlocs.at(1) * local_nlocs.at(2); + + PARTHENON_REQUIRE_THROWS( + loc_max_vol == mesh_->GetNumMeshBlocksThisRank(), + "Block coverage on rank cannot be matched to a contiguous " + "array, which is required for FFTs. Try a different amount of " + "ranks (one pack per rank, i.e. num_packs=1, will always work)."); + + const auto block_size_ = mesh_->GetDefaultBlockSize(); + mesh_block_box.size[0] = block_size_.nx(parthenon::X1DIR); + mesh_block_box.size[1] = block_size_.nx(parthenon::X2DIR); + mesh_block_box.size[2] = block_size_.nx(parthenon::X3DIR); + for (int i = 0; i < 3; i++) { + local_mesh_box.size[i] = local_nlocs[i] * mesh_block_box.size[i]; + local_mesh_box.low[i] = local_loc_min[i] * mesh_block_box.size[i]; + local_mesh_box.high[i] = local_mesh_box.low[i] + local_mesh_box.size[i] - 1; + } + + auto &md = mesh_->mesh_data.Get(); + IndexRange ib = md->GetBlockData(0)->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBlockData(0)->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBlockData(0)->GetBoundsK(IndexDomain::interior); + + mesh_block_box.low[0] = ib.s; + mesh_block_box.high[0] = ib.e; + mesh_block_box.low[1] = jb.s; + mesh_block_box.high[1] = jb.e; + mesh_block_box.low[2] = kb.s; + mesh_block_box.high[2] = kb.e; + + std::cout << "Initialized UniformGridHelper" << "\n"; + std::cout << "Mesh layout:" << "\n"; + std::cout << "Rank " << parthenon::Globals::my_rank << " local mesh box: low = (" + << local_mesh_box.low[0] << ", " << local_mesh_box.low[1] << ", " + << local_mesh_box.low[2] << "), high = (" << local_mesh_box.high[0] << ", " + << local_mesh_box.high[1] << ", " << local_mesh_box.high[2] << "), size = (" + << local_mesh_box.size[0] << ", " << local_mesh_box.size[1] << ", " + << local_mesh_box.size[2] << ")\n"; +} // UniformGridHelper::UniformGridHelper + +void UniformGridHelper::GatherField(const std::string &var_name, int var_index, + parthenon::ParArray1D &output) { + // Check that var_name and var_index correspond to a valid variable in the mesh data and + // that output array is large enough to hold the gathered data. + auto &md = mesh_->mesh_data.Get(); + auto vars = md->PackVariables(std::vector{var_name}); + PARTHENON_REQUIRE_THROWS(vars.GetDim(5) > 0, "GatherField: variable '" + var_name + + "' not found in mesh data"); + PARTHENON_REQUIRE_THROWS(var_index < vars.GetDim(4), "GatherField: var_index " + + std::to_string(var_index) + + " out of range"); + PARTHENON_REQUIRE_THROWS(output.size() >= local_mesh_box.size[0] * + local_mesh_box.size[1] * + local_mesh_box.size[2], + "GatherField: output array too small"); + + IndexRange ib = md->GetBlockData(0)->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBlockData(0)->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBlockData(0)->GetBoundsK(IndexDomain::interior); + + const auto vi = var_index; + auto helper = GetKernelHelper(); + + parthenon::par_for( + "UniformGridHelper::GatherField", 0, md->NumBlocks() - 1, kb.s, kb.e, jb.s, jb.e, + ib.s, ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + const auto idx = helper.FlatIndex(b, k, j, i); + output(idx) = vars(b, vi, k, j, i); + }); +} + +void UniformGridHelper::ScatterField(const parthenon::ParArray1D &input, + const std::string &var_name, int var_index) { + auto &md = mesh_->mesh_data.Get(); + auto vars = md->PackVariables(std::vector{var_name}); + PARTHENON_REQUIRE_THROWS(vars.GetDim(5) > 0, "ScatterField: variable '" + var_name + + "' not found in mesh data"); + PARTHENON_REQUIRE_THROWS(var_index < vars.GetDim(4), "ScatterField: var_index " + + std::to_string(var_index) + + " out of range"); + PARTHENON_REQUIRE_THROWS(input.size() >= local_mesh_box.size[0] * + local_mesh_box.size[1] * + local_mesh_box.size[2], + "ScatterField: input array too small"); + + IndexRange ib = md->GetBlockData(0)->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBlockData(0)->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBlockData(0)->GetBoundsK(IndexDomain::interior); + + const auto vi = var_index; + + auto helper = GetKernelHelper(); + + parthenon::par_for( + "UniformGridHelper::ScatterField", 0, md->NumBlocks() - 1, kb.s, kb.e, jb.s, jb.e, + ib.s, ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + const auto idx = helper.FlatIndex(b, k, j, i); + vars(b, vi, k, j, i) = input(idx); + }); +} + +} // namespace parthenon diff --git a/src/utils/uniform_grid_helper.hpp b/src/utils/uniform_grid_helper.hpp new file mode 100644 index 0000000000000..ebd79ff6a35c7 --- /dev/null +++ b/src/utils/uniform_grid_helper.hpp @@ -0,0 +1,76 @@ +//======================================================================================== +// Parthenon performance portable AMR framework +// Copyright(C) 2026 The Parthenon collaboration +// Licensed under the 3-clause BSD License, see LICENSE file for details +//======================================================================================== + +// This file was made in part with generative AI. + +#ifndef UTILS_UNIFORM_GRID_HELPER_HPP_ +#define UTILS_UNIFORM_GRID_HELPER_HPP_ + +#include +#include +#include +#include + +#include "parthenon_arrays.hpp" + +namespace parthenon { + +class Mesh; + +struct Box3D { + int low[3]; + int high[3]; + int size[3]; // size in each dimension: high - low + 1 +}; + +// Assuming a uniform grid, this class helps gather information about the grid layout +// across all meshblocks on a rank. +class UniformGridHelper { + public: + explicit UniformGridHelper(Mesh *mesh); + + Box3D mesh_block_box; + Box3D local_mesh_box; + + struct KernelHelper { + parthenon::ParArray2D loc_view; + Box3D mesh_block_box; + Box3D local_mesh_box; + + KOKKOS_INLINE_FUNCTION + std::int64_t FlatIndex(int b, int k, int j, int i) const { + const auto kk = k - mesh_block_box.low[2] + loc_view(b, 2) * mesh_block_box.size[2]; + const auto jj = j - mesh_block_box.low[1] + loc_view(b, 1) * mesh_block_box.size[1]; + const auto ii = i - mesh_block_box.low[0] + loc_view(b, 0) * mesh_block_box.size[0]; + return (std::int64_t)kk * local_mesh_box.size[1] * local_mesh_box.size[0] + + (std::int64_t)jj * local_mesh_box.size[0] + ii; + } + }; + + KernelHelper GetKernelHelper() const { + return {loc_view, mesh_block_box, local_mesh_box}; + } + + // Gathers a single component of a named variable from meshblocks + // into a contiguous 1D array suitable for FFT input. + // output must be pre-allocated with size >= size_real_space_box() + void GatherField(const std::string &var_name, const int var_index, + parthenon::ParArray1D &output); + + // Distributes a contiguous 1D array back to meshblocks. + // Inverse of GatherField. + void ScatterField(const parthenon::ParArray1D &input, const std::string &var_name, + const int var_index); + + private: + Mesh *mesh_; + parthenon::ParArray2D + loc_view; // logical location of local blocks; stored on device for use in kernels +}; + +} // namespace parthenon + +#endif // UTILS_UNIFORM_GRID_HELPER_HPP_ diff --git a/tst/regression/CMakeLists.txt b/tst/regression/CMakeLists.txt index d5e3c4877e07f..b03fd69c65d86 100644 --- a/tst/regression/CMakeLists.txt +++ b/tst/regression/CMakeLists.txt @@ -64,6 +64,15 @@ list(APPEND TEST_ARGS "--driver ${PROJECT_BINARY_DIR}/example/particle_leapfrog/ --num_steps 1") list(APPEND EXTRA_TEST_LABELS "") +# Fourier transform FFT round-trip accuracy test (only when HeFTe/FFT is enabled) +if(PARTHENON_ENABLE_FFT) + list(APPEND TEST_DIRS fourier_transform) + list(APPEND TEST_PROCS ${NUM_MPI_PROC_TESTING}) + list(APPEND TEST_ARGS "--driver ${PROJECT_BINARY_DIR}/example/fourier_transform/fourier-transform \ +--driver_input ${CMAKE_CURRENT_SOURCE_DIR}/test_suites/fourier_transform/parthinput.fourier_transform") + list(APPEND EXTRA_TEST_LABELS "") +endif() + if (ENABLE_HDF5) # h5py is needed for restart and hdf5 test diff --git a/tst/regression/test_suites/fourier_transform/__init__.py b/tst/regression/test_suites/fourier_transform/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/tst/regression/test_suites/fourier_transform/fourier_transform.py b/tst/regression/test_suites/fourier_transform/fourier_transform.py new file mode 100644 index 0000000000000..6924bc30f8bed --- /dev/null +++ b/tst/regression/test_suites/fourier_transform/fourier_transform.py @@ -0,0 +1,53 @@ +# ======================================================================================== +# Parthenon performance portable AMR framework +# Copyright(C) 2026 The Parthenon collaboration +# Licensed under the 3-clause BSD License, see LICENSE file for details +# ======================================================================================== + +import re +import sys +import utils.test_case + +sys.dont_write_bytecode = True + +# Maximum tolerated absolute error for a forward+inverse FFT round-trip on a +# 16^3 grid. +# Error tolerance currently set for double prec data and transforms. +ERROR_TOLERANCE = 1e-14 + + +class TestCase(utils.test_case.TestCaseAbs): + def Prepare(self, parameters, step): + parameters.coverage_status = "both" + return parameters + + def Analyse(self, parameters): + if not parameters.stdouts: + print("ERROR: no stdout captured from driver") + return False + + output = parameters.stdouts[0].decode() + + def parse_error(error_string): + match = re.search(rf"{error_string}:\s+([\d.eE+\-]+)", output) + if match is None: + print(f"ERROR: could not find '{error_string}' in output") + print("Driver output was:") + print(output) + return False + + max_error = float(match.group(1)) + print( + f"{error_string}: {max_error:.3e} (tolerance: {ERROR_TOLERANCE:.3e})" + ) + + if max_error > ERROR_TOLERANCE: + print("ERROR: exceeds tolerance") + return False + return True + + success = True + success &= parse_error("Max relative error after FFT round-trip") + success &= parse_error("Error in spectrum total power") + success &= parse_error("Error in spectrum mean") + return success diff --git a/tst/regression/test_suites/fourier_transform/parthinput.fourier_transform b/tst/regression/test_suites/fourier_transform/parthinput.fourier_transform new file mode 100644 index 0000000000000..399ab8ef496bd --- /dev/null +++ b/tst/regression/test_suites/fourier_transform/parthinput.fourier_transform @@ -0,0 +1,37 @@ +# ======================================================================================== +# Parthenon performance portable AMR framework +# Copyright(C) 2026 The Parthenon collaboration +# Licensed under the 3-clause BSD License, see LICENSE file for details +# ======================================================================================== + + +problem_id = fourier_transform + + +refinement = none +nghost = 3 + +nx1 = 16 +x1min = -2.0 +x1max = 2.0 +ix1_bc = periodic +ox1_bc = periodic + +nx2 = 16 +x2min = -2.0 +x2max = 2.0 +ix2_bc = periodic +ox2_bc = periodic + +nx3 = 16 +x3min = -2.0 +x3max = 2.0 +ix3_bc = periodic +ox3_bc = periodic + +num_packs = 1 + + +nx1 = 8 +nx2 = 8 +nx3 = 8 diff --git a/tst/unit/test_unit_sort.cpp b/tst/unit/test_unit_sort.cpp index eccc6cb6c0221..ded677f8a84ea 100644 --- a/tst/unit/test_unit_sort.cpp +++ b/tst/unit/test_unit_sort.cpp @@ -45,8 +45,6 @@ struct KeyComparator { }; TEST_CASE("Sorting", "[sort]") { -// No sort support on HIP at the moment -#if !defined(KOKKOS_ENABLE_HIP) GIVEN("An unordered list of integers") { ParArray1D data("Data to sort", N); @@ -94,5 +92,4 @@ TEST_CASE("Sorting", "[sort]") { REQUIRE(data_h(3).value_ == 4); REQUIRE(data_h(4).value_ == 5); } -#endif // !defined(KOKKOS_ENABLE_HIP) }