Add DUCC FFT support

This commit is contained in:
Damiano Franzò
2025-05-12 17:56:02 +00:00
committed by Rasmus Munk Larsen
parent 6f1a143418
commit db85838ee2
9 changed files with 129 additions and 22 deletions

View File

@@ -0,0 +1,71 @@
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
namespace Eigen {
namespace internal {
template <typename _Scalar>
struct duccfft_impl {
using Scalar = _Scalar;
using Complex = std::complex<Scalar>;
using shape_t = ducc0::fmav_info::shape_t;
using stride_t = ducc0::fmav_info::stride_t;
inline void clear() {}
inline void fwd(Complex* dst, const Scalar* src, int nfft) {
const shape_t axes{0};
ducc0::cfmav<Scalar> m_in(src, shape_t{static_cast<size_t>(nfft)});
ducc0::vfmav<Complex> m_out(dst, shape_t{static_cast<size_t>(nfft) / 2 + 1});
ducc0::r2c(m_in, m_out, axes, /*forward=*/true, /*scale=*/static_cast<Scalar>(1));
}
inline void fwd(Complex* dst, const Complex* src, int nfft) {
const shape_t axes{0};
ducc0::cfmav<Complex> m_in(src, shape_t{static_cast<size_t>(nfft)});
ducc0::vfmav<Complex> m_out(dst, shape_t{static_cast<size_t>(nfft)});
ducc0::c2c(m_in, m_out, axes, /*forward=*/true, /*scale=*/static_cast<Scalar>(1));
}
inline void inv(Scalar* dst, const Complex* src, int nfft) {
const shape_t axes{0};
ducc0::cfmav<Complex> m_in(src, shape_t{static_cast<size_t>(nfft) / 2 + 1});
ducc0::vfmav<Scalar> m_out(dst, shape_t{static_cast<size_t>(nfft)});
ducc0::c2r(m_in, m_out, axes, /*forward=*/false, /*scale=*/static_cast<Scalar>(1));
}
inline void inv(Complex* dst, const Complex* src, int nfft) {
const shape_t axes{0};
ducc0::cfmav<Complex> m_in(src, shape_t{static_cast<size_t>(nfft)});
ducc0::vfmav<Complex> m_out(dst, shape_t{static_cast<size_t>(nfft)});
ducc0::c2c(m_in, m_out, axes, /*forward=*/false, /*scale=*/static_cast<Scalar>(1));
}
inline void fwd2(Complex* dst, const Complex* src, int nfft0, int nfft1) {
const shape_t axes{0, 1};
const shape_t in_shape{static_cast<size_t>(nfft0), static_cast<size_t>(nfft1)};
const shape_t out_shape{static_cast<size_t>(nfft0), static_cast<size_t>(nfft1)};
const stride_t stride{static_cast<ptrdiff_t>(nfft1), static_cast<ptrdiff_t>(1)};
ducc0::cfmav<Complex> m_in(src, in_shape, stride);
ducc0::vfmav<Complex> m_out(dst, out_shape, stride);
ducc0::c2c(m_in, m_out, axes, /*forward=*/true, /*scale=*/static_cast<Scalar>(1));
}
inline void inv2(Complex* dst, const Complex* src, int nfft0, int nfft1) {
const shape_t axes{0, 1};
const shape_t in_shape{static_cast<size_t>(nfft0), static_cast<size_t>(nfft1)};
const shape_t out_shape{static_cast<size_t>(nfft0), static_cast<size_t>(nfft1)};
const stride_t stride{static_cast<ptrdiff_t>(nfft1), static_cast<ptrdiff_t>(1)};
ducc0::cfmav<Complex> m_in(src, in_shape, stride);
ducc0::vfmav<Complex> m_out(dst, out_shape, stride);
ducc0::c2c(m_in, m_out, axes, /*forward=*/false, /*scale=*/static_cast<Scalar>(1));
}
};
} // namespace internal
} // namespace Eigen

View File

@@ -5,17 +5,16 @@
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
using namespace pocketfft;
using namespace pocketfft::detail;
namespace Eigen {
namespace internal {
template <typename _Scalar>
struct pocketfft_impl {
typedef _Scalar Scalar;
typedef std::complex<Scalar> Complex;
using Scalar = _Scalar;
using Complex = std::complex<Scalar>;
using shape_t = pocketfft::shape_t;
using stride_t = pocketfft::stride_t;
inline void clear() {}
@@ -24,14 +23,14 @@ struct pocketfft_impl {
const shape_t axes_{0};
const stride_t stride_in{sizeof(Scalar)};
const stride_t stride_out{sizeof(Complex)};
r2c(shape_, stride_in, stride_out, axes_, FORWARD, src, dst, static_cast<Scalar>(1));
pocketfft::r2c(shape_, stride_in, stride_out, axes_, pocketfft::FORWARD, src, dst, static_cast<Scalar>(1));
}
inline void fwd(Complex* dst, const Complex* src, int nfft) {
const shape_t shape_{static_cast<size_t>(nfft)};
const shape_t axes_{0};
const stride_t stride_{sizeof(Complex)};
c2c(shape_, stride_, stride_, axes_, FORWARD, src, dst, static_cast<Scalar>(1));
pocketfft::c2c(shape_, stride_, stride_, axes_, pocketfft::FORWARD, src, dst, static_cast<Scalar>(1));
}
inline void inv(Scalar* dst, const Complex* src, int nfft) {
@@ -39,28 +38,28 @@ struct pocketfft_impl {
const shape_t axes_{0};
const stride_t stride_in{sizeof(Complex)};
const stride_t stride_out{sizeof(Scalar)};
c2r(shape_, stride_in, stride_out, axes_, BACKWARD, src, dst, static_cast<Scalar>(1));
pocketfft::c2r(shape_, stride_in, stride_out, axes_, pocketfft::BACKWARD, src, dst, static_cast<Scalar>(1));
}
inline void inv(Complex* dst, const Complex* src, int nfft) {
const shape_t shape_{static_cast<size_t>(nfft)};
const shape_t axes_{0};
const stride_t stride_{sizeof(Complex)};
c2c(shape_, stride_, stride_, axes_, BACKWARD, src, dst, static_cast<Scalar>(1));
pocketfft::c2c(shape_, stride_, stride_, axes_, pocketfft::BACKWARD, src, dst, static_cast<Scalar>(1));
}
inline void fwd2(Complex* dst, const Complex* src, int nfft0, int nfft1) {
const shape_t shape_{static_cast<size_t>(nfft0), static_cast<size_t>(nfft1)};
const shape_t axes_{0, 1};
const stride_t stride_{static_cast<ptrdiff_t>(sizeof(Complex) * nfft1), static_cast<ptrdiff_t>(sizeof(Complex))};
c2c(shape_, stride_, stride_, axes_, FORWARD, src, dst, static_cast<Scalar>(1));
pocketfft::c2c(shape_, stride_, stride_, axes_, pocketfft::FORWARD, src, dst, static_cast<Scalar>(1));
}
inline void inv2(Complex* dst, const Complex* src, int nfft0, int nfft1) {
const shape_t shape_{static_cast<size_t>(nfft0), static_cast<size_t>(nfft1)};
const shape_t axes_{0, 1};
const stride_t stride_{static_cast<ptrdiff_t>(sizeof(Complex) * nfft1), static_cast<ptrdiff_t>(sizeof(Complex))};
c2c(shape_, stride_, stride_, axes_, BACKWARD, src, dst, static_cast<Scalar>(1));
pocketfft::c2c(shape_, stride_, stride_, axes_, pocketfft::BACKWARD, src, dst, static_cast<Scalar>(1));
}
};