Current section

Files

Jump to
mob_dev priv cpp_nif nx_eigen_fft_eigen.cpp
Raw

priv/cpp_nif/nx_eigen_fft_eigen.cpp

// nx_eigen_fft_eigen.cpp — Eigen-backed FFT implementation of the
// nx_eigen_fft.h C contract.
//
// We ship this in mob_dev (not in upstream nx_eigen) so the on-device
// path through `mix mob.enable nxeigen` gets `Nx.fft` working without
// a separate FFTW cross-compile. Eigen's FFT module ships kissfft as
// the default backend — header-only, embedded in the Eigen tarball at
// unsupported/Eigen/FFT — so this file plus the Eigen headers gives
// us a complete pluggable FFT with no external library dep.
//
// ## Contract delta vs FFTW
//
// nx_eigen_fft.h spec:
// * Forward unnormalised: X[k] = Σ x[n] · exp(-2πi nk/N)
// * Inverse unnormalised: X̃[k] = Σ x[n] · exp(+2πi nk/N)
// * Caller divides the IDFT by n.
//
// Eigen's `FFT<>` defaults match this for `fwd` (kissfft is unscaled)
// but its `inv` divides by N unless `Unscaled` flag is set. We set
// the flag on every transform so the contract holds exactly.
//
// ## Layout
//
// The C contract is interleaved real/imag floats:
// [re_0, im_0, re_1, im_1, ..., re_{n-1}, im_{n-1}]
// `std::complex<float>` and `std::complex<double>` are guaranteed by
// C++ to have this exact layout (sizeof = 2 × sizeof(scalar), no
// padding, real then imag). Bitwise-compatible reinterpret_cast.
//
// ## Performance note
//
// Eigen's default kissfft backend is roughly 2× slower than FFTW for
// large transforms but typically microseconds for audio-sized buffers
// (≤8192). If a future workload measures a real bottleneck, swap to
// FFTW via a parallel `nx_eigen_fft_fftw.cpp` cross-compile — the
// pluggable contract was designed for exactly that.
#include "nx_eigen_fft.h"
#include <complex>
#include <unsupported/Eigen/FFT>
namespace {
template <typename Scalar>
int fft_forward(const Scalar *in, Scalar *out, int n) {
if (n <= 0) {
return 1;
}
Eigen::FFT<Scalar> fft;
fft.SetFlag(Eigen::FFT<Scalar>::Unscaled);
const auto *src = reinterpret_cast<const std::complex<Scalar> *>(in);
auto *dst = reinterpret_cast<std::complex<Scalar> *>(out);
// Eigen's `fwd(dst, src, n)` accepts overlapping buffers (kissfft
// copies internally), so the contract's in-place allowance holds.
fft.fwd(dst, src, static_cast<Eigen::Index>(n));
return 0;
}
template <typename Scalar>
int fft_inverse(const Scalar *in, Scalar *out, int n) {
if (n <= 0) {
return 1;
}
Eigen::FFT<Scalar> fft;
fft.SetFlag(Eigen::FFT<Scalar>::Unscaled);
const auto *src = reinterpret_cast<const std::complex<Scalar> *>(in);
auto *dst = reinterpret_cast<std::complex<Scalar> *>(out);
fft.inv(dst, src, static_cast<Eigen::Index>(n));
return 0;
}
} // namespace
extern "C" {
int nx_eigen_fft_forward_f32(const float *in, float *out, int n) {
return fft_forward<float>(in, out, n);
}
int nx_eigen_fft_forward_f64(const double *in, double *out, int n) {
return fft_forward<double>(in, out, n);
}
int nx_eigen_fft_inverse_f32(const float *in, float *out, int n) {
return fft_inverse<float>(in, out, n);
}
int nx_eigen_fft_inverse_f64(const double *in, double *out, int n) {
return fft_inverse<double>(in, out, n);
}
} // extern "C"