mirror of
https://gitlab.com/libeigen/eigen.git
synced 2026-04-10 11:34:33 +08:00
249 lines
8.7 KiB
C++
249 lines
8.7 KiB
C++
// 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/.
|
|
|
|
#include <mkl_dfti.h>
|
|
|
|
// IWYU pragma: private
|
|
#include "./InternalHeaderCheck.h"
|
|
|
|
#include <complex>
|
|
#include <memory>
|
|
|
|
namespace Eigen {
|
|
namespace internal {
|
|
namespace imklfft {
|
|
|
|
#define RUN_OR_ASSERT(EXPR, ERROR_MSG) \
|
|
{ \
|
|
MKL_LONG status = (EXPR); \
|
|
eigen_assert(status == DFTI_NO_ERROR && (ERROR_MSG)); \
|
|
};
|
|
|
|
inline MKL_Complex16* complex_cast(const std::complex<double>* p) {
|
|
return const_cast<MKL_Complex16*>(reinterpret_cast<const MKL_Complex16*>(p));
|
|
}
|
|
|
|
inline MKL_Complex8* complex_cast(const std::complex<float>* p) {
|
|
return const_cast<MKL_Complex8*>(reinterpret_cast<const MKL_Complex8*>(p));
|
|
}
|
|
|
|
/*
|
|
* Parameters:
|
|
* precision: enum, Precision of the transform: DFTI_SINGLE or DFTI_DOUBLE.
|
|
* forward_domain: enum, Forward domain of the transform: DFTI_COMPLEX or
|
|
* DFTI_REAL. dimension: MKL_LONG Dimension of the transform. sizes: MKL_LONG if
|
|
* dimension = 1.Length of the transform for a one-dimensional transform. sizes:
|
|
* Array of type MKL_LONG otherwise. Lengths of each dimension for a
|
|
* multi-dimensional transform.
|
|
*/
|
|
inline void configure_descriptor(std::shared_ptr<DFTI_DESCRIPTOR>& handl, enum DFTI_CONFIG_VALUE precision,
|
|
enum DFTI_CONFIG_VALUE forward_domain, MKL_LONG dimension, MKL_LONG* sizes) {
|
|
eigen_assert(dimension == 1 || dimension == 2 && "Transformation dimension must be less than 3.");
|
|
|
|
DFTI_DESCRIPTOR_HANDLE res = nullptr;
|
|
if (dimension == 1) {
|
|
RUN_OR_ASSERT(DftiCreateDescriptor(&res, precision, forward_domain, dimension, *sizes),
|
|
"DftiCreateDescriptor failed.")
|
|
handl.reset(res, [](DFTI_DESCRIPTOR_HANDLE handle) { DftiFreeDescriptor(&handle); });
|
|
if (forward_domain == DFTI_REAL) {
|
|
// Set CCE storage
|
|
RUN_OR_ASSERT(DftiSetValue(handl.get(), DFTI_CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX),
|
|
"DftiSetValue failed.")
|
|
}
|
|
} else {
|
|
RUN_OR_ASSERT(DftiCreateDescriptor(&res, precision, DFTI_COMPLEX, dimension, sizes), "DftiCreateDescriptor failed.")
|
|
handl.reset(res, [](DFTI_DESCRIPTOR_HANDLE handle) { DftiFreeDescriptor(&handle); });
|
|
}
|
|
|
|
RUN_OR_ASSERT(DftiSetValue(handl.get(), DFTI_PLACEMENT, DFTI_NOT_INPLACE), "DftiSetValue failed.")
|
|
RUN_OR_ASSERT(DftiCommitDescriptor(handl.get()), "DftiCommitDescriptor failed.")
|
|
}
|
|
|
|
template <typename T>
|
|
struct plan {};
|
|
|
|
template <>
|
|
struct plan<float> {
|
|
typedef float scalar_type;
|
|
typedef MKL_Complex8 complex_type;
|
|
|
|
std::shared_ptr<DFTI_DESCRIPTOR> m_plan;
|
|
|
|
plan() = default;
|
|
|
|
enum DFTI_CONFIG_VALUE precision = DFTI_SINGLE;
|
|
|
|
inline void forward(complex_type* dst, complex_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
|
|
}
|
|
|
|
inline void inverse(complex_type* dst, complex_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
|
|
}
|
|
|
|
inline void forward(complex_type* dst, scalar_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
|
|
}
|
|
|
|
inline void inverse(scalar_type* dst, complex_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
|
|
}
|
|
|
|
inline void forward2(complex_type* dst, complex_type* src, int n0, int n1) {
|
|
if (m_plan == 0) {
|
|
MKL_LONG sizes[2] = {n0, n1};
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
|
|
}
|
|
|
|
inline void inverse2(complex_type* dst, complex_type* src, int n0, int n1) {
|
|
if (m_plan == 0) {
|
|
MKL_LONG sizes[2] = {n0, n1};
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct plan<double> {
|
|
typedef double scalar_type;
|
|
typedef MKL_Complex16 complex_type;
|
|
|
|
std::shared_ptr<DFTI_DESCRIPTOR> m_plan;
|
|
|
|
plan() = default;
|
|
|
|
enum DFTI_CONFIG_VALUE precision = DFTI_DOUBLE;
|
|
|
|
inline void forward(complex_type* dst, complex_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
|
|
}
|
|
|
|
inline void inverse(complex_type* dst, complex_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
|
|
}
|
|
|
|
inline void forward(complex_type* dst, scalar_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
|
|
}
|
|
|
|
inline void inverse(scalar_type* dst, complex_type* src, MKL_LONG nfft) {
|
|
if (m_plan == 0) {
|
|
configure_descriptor(m_plan, precision, DFTI_REAL, 1, &nfft);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
|
|
}
|
|
|
|
inline void forward2(complex_type* dst, complex_type* src, int n0, int n1) {
|
|
if (m_plan == 0) {
|
|
MKL_LONG sizes[2] = {n0, n1};
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeForward(m_plan.get(), src, dst), "DftiComputeForward failed.")
|
|
}
|
|
|
|
inline void inverse2(complex_type* dst, complex_type* src, int n0, int n1) {
|
|
if (m_plan == 0) {
|
|
MKL_LONG sizes[2] = {n0, n1};
|
|
configure_descriptor(m_plan, precision, DFTI_COMPLEX, 2, sizes);
|
|
}
|
|
RUN_OR_ASSERT(DftiComputeBackward(m_plan.get(), src, dst), "DftiComputeBackward failed.")
|
|
}
|
|
};
|
|
|
|
template <typename Scalar_>
|
|
struct imklfft_impl {
|
|
typedef Scalar_ Scalar;
|
|
typedef std::complex<Scalar> Complex;
|
|
|
|
inline void clear() { m_plans.clear(); }
|
|
|
|
// complex-to-complex forward FFT
|
|
inline void fwd(Complex* dst, const Complex* src, int nfft) {
|
|
MKL_LONG size = nfft;
|
|
get_plan(nfft, dst, src).forward(complex_cast(dst), complex_cast(src), size);
|
|
}
|
|
|
|
// real-to-complex forward FFT
|
|
inline void fwd(Complex* dst, const Scalar* src, int nfft) {
|
|
MKL_LONG size = nfft;
|
|
get_plan(nfft, dst, src).forward(complex_cast(dst), const_cast<Scalar*>(src), nfft);
|
|
}
|
|
|
|
// 2-d complex-to-complex
|
|
inline void fwd2(Complex* dst, const Complex* src, int n0, int n1) {
|
|
get_plan(n0, n1, dst, src).forward2(complex_cast(dst), complex_cast(src), n0, n1);
|
|
}
|
|
|
|
// inverse complex-to-complex
|
|
inline void inv(Complex* dst, const Complex* src, int nfft) {
|
|
MKL_LONG size = nfft;
|
|
get_plan(nfft, dst, src).inverse(complex_cast(dst), complex_cast(src), nfft);
|
|
}
|
|
|
|
// half-complex to scalar
|
|
inline void inv(Scalar* dst, const Complex* src, int nfft) {
|
|
MKL_LONG size = nfft;
|
|
get_plan(nfft, dst, src).inverse(const_cast<Scalar*>(dst), complex_cast(src), nfft);
|
|
}
|
|
|
|
// 2-d complex-to-complex
|
|
inline void inv2(Complex* dst, const Complex* src, int n0, int n1) {
|
|
get_plan(n0, n1, dst, src).inverse2(complex_cast(dst), complex_cast(src), n0, n1);
|
|
}
|
|
|
|
private:
|
|
std::map<int64_t, plan<Scalar>> m_plans;
|
|
|
|
inline plan<Scalar>& get_plan(int nfft, void* dst, const void* src) {
|
|
int inplace = dst == src ? 1 : 0;
|
|
int aligned = ((reinterpret_cast<size_t>(src) & 15) | (reinterpret_cast<size_t>(dst) & 15)) == 0 ? 1 : 0;
|
|
int64_t key = ((nfft << 2) | (inplace << 1) | aligned) << 1;
|
|
|
|
// Create element if key does not exist.
|
|
return m_plans[key];
|
|
}
|
|
|
|
inline plan<Scalar>& get_plan(int n0, int n1, void* dst, const void* src) {
|
|
int inplace = (dst == src) ? 1 : 0;
|
|
int aligned = ((reinterpret_cast<size_t>(src) & 15) | (reinterpret_cast<size_t>(dst) & 15)) == 0 ? 1 : 0;
|
|
int64_t key = (((((int64_t)n0) << 31) | (n1 << 2) | (inplace << 1) | aligned) << 1) + 1;
|
|
|
|
// Create element if key does not exist.
|
|
return m_plans[key];
|
|
}
|
|
};
|
|
|
|
#undef RUN_OR_ASSERT
|
|
|
|
} // namespace imklfft
|
|
} // namespace internal
|
|
} // namespace Eigen
|