Kokkos-FFT

Kokkos-FFT implements local interfaces between Kokkos and de facto standard FFT libraries, including FFTW, cufft, hipfft (rocfft), and oneMKL. "Local" means not using MPI, or running within a single MPI process without knowing about MPI. We are inclined to implement the numpy.fft-like interfaces adapted for Kokkos. A key concept is that "As easy as numpy, as fast as vendor libraries". Accordingly, our API follows the API by numpy.fft with minor differences. A fft library dedicated to Kokkos Device backend (e.g. cufft for CUDA backend) is automatically used. If something is wrong with runtime values (say View extents), it will raise runtime errors (C++ std::runtime_error). See documentations for more information.

Here is an example for 1D real to complex transform with rfft in Kokkos-FFT.

#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>;
constexpr int n = 4;

View1D<double> x("x", n);
View1D<Kokkos::complex<double> > x_hat("x_hat", n/2+1);

Kokkos::Random_XorShift64_Pool<> random_pool(12345);
Kokkos::fill_random(x, random_pool, 1);
Kokkos::fence();

KokkosFFT::rfft(execution_space(), x, x_hat);

This is equivalent to the following python code.

import numpy as np
x = np.random.rand(4)
x_hat = np.fft.rfft(x)

There are two major differences: execution_space argument and output value (x_hat) is an argument of API (not returned value from API). As imagined, Kokkos-FFT only accepts Kokkos Views as input data. The accessibilities of Views from execution_space are statically checked (compilation errors if not accessible).

Depending on a View dimension, it automatically uses the batched plans as follows

#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>;
constexpr int n0 = 4, n1 = 8;

View2D<double> x("x", n0, n1);
View2D<Kokkos::complex<double> > x_hat("x_hat", n0, n1/2+1);

Kokkos::Random_XorShift64_Pool<> random_pool(12345);
Kokkos::fill_random(x, random_pool, 1);
Kokkos::fence();

// FFT along -1 axis and batched along 0th axis
int axis = -1;
KokkosFFT::rfft(execution_space(), x, x_hat, KokkosFFT::Normalization::backward, axis);

This is equivalent to

import numpy as np
x = np.random.rand(4, 8)
x_hat = np.fft.rfft(x, axis=-1)

In this example, the 1D batched rfft over 2D View along axis -1 is executed. Some basic examples are found in examples.

Using Kokkos-FFT

For the moment, there are two ways to use Kokkos-FFT: including as a subdirectory in CMake project or installing as a library. First of all, you need to clone this repo.

git clone --recursive https://github.com/kokkos/kokkos-fft.git

Prerequisites

To use Kokkos-FFT, we need the following:

• CMake 3.22+
• Kokkos 4.7+
• gcc 10.4.0+ (CPUs)
• IntelLLVM 2024.2.1+ (CPUs, Intel GPUs)
• nvcc 12.2.0+ (NVIDIA GPUs)
• rocm 6.3.0+ (AMD GPUs)

Warning

A compatible C++ compiler that supports at least C++20 is necessary

CMake

Since Kokkos-FFT is a header-only library, it is enough to simply add as a subdirectory. It is assumed that kokkos and Kokkos-FFT are placed under <project_directory>/tpls.

Here is an example to use Kokkos-FFT in the following CMake project.

---/
 |
 └──<project_directory>/
    |--tpls
    |    |--kokkos/
    |    └──kokkos-fft/
    |--CMakeLists.txt
    └──hello.cpp

The CMakeLists.txt would be

cmake_minimum_required(VERSION 3.23)
project(kokkos-fft-as-subdirectory LANGUAGES CXX)

add_subdirectory(tpls/kokkos)
add_subdirectory(tpls/kokkos-fft)

add_executable(hello-kokkos-fft hello.cpp)
target_link_libraries(hello-kokkos-fft PUBLIC Kokkos::kokkos KokkosFFT::fft)

For compilation, we basically rely on the CMake options for Kokkos. For example, the compile options for A100 GPU is as follows.

cmake -B build \
      -DCMAKE_CXX_COMPILER=g++ \
      -DCMAKE_BUILD_TYPE=Release \
      -DKokkos_ENABLE_CUDA=ON \
      -DKokkos_ARCH_AMPERE80=ON
cmake --build build -j 8

This way, all the functionalities are executed on A100 GPUs. For installation, details are provided in the documentation.

Spack

Kokkos-FFT can also be installed with spack. For example, the recipe for H100 GPU with cufft is as follows.

git clone --depth=2 --branch=v1.1.0 https://github.com/spack/spack.git
source spack/share/spack/setup-env.sh # For bash

spack install kokkos-fft device_backend=cufft ^kokkos +cuda +wrapper cuda_arch=90

We have two main parameters to configure Spack:

• host_backend: Enable device backend FFT library (fftw-serial or fftw-openmp)
• device_backend: Enable device backend FFT library (cufft, hipfft, or onemkl)

To enable the device library, you need to install kokkos with the corresponding Kokkos backend. See this page for detail.

Support

Find us on Slack channel or open a GitHub issue.

Contributing

Please see this page for details on how to contribute.

Citing Kokkos-FFT

Please see this page.

LICENSE

Kokkos-FFT is mainly distributed under either the MIT license, or at your option, the Apache-2.0 license with LLVM exception. More strictly, FindFFTW.cmake file is provided under BSD-3-Clause license and reuse.yml file is provided under CC0-1.0 license. Licenses are automatically confirmed with REUSE compliance check in CI (see this page for detail).