Files
eigen/unsupported/Eigen/CXX11/src/Tensor/TensorFFT.h
Rasmus Munk Larsen e6b1020221 Adds a fast memcpy function to Eigen. This takes advantage of the following:
1. For small fixed sizes, the compiler generates inline code for memcpy, which is much faster.

2. My colleague eriche at googl dot com discovered that for large sizes, memmove is significantly faster than memcpy (at least on Linux with GCC or Clang). See benchmark numbers measured on a Haswell (HP Z440) workstation here: https://docs.google.com/a/google.com/spreadsheets/d/1jLs5bKzXwhpTySw65MhG1pZpsIwkszZqQTjwrd_n0ic/pubhtml This is of course surprising since memcpy is a less constrained version of memmove. This stackoverflow thread contains some speculation as to the causes: http://stackoverflow.com/questions/22793669/poor-memcpy-performance-on-linux

Below are numbers for copying and slicing tensors using the multithreaded TensorDevice. The numbers show significant improvements for memcpy of very small blocks and for memcpy of large blocks single threaded (we were already able to saturate memory bandwidth for >1 threads before on large blocks). The "slicingSmallPieces" benchmark also shows small consistent improvements, since memcpy cost is a fair portion of that particular computation.

The benchmarks operate on NxN matrices, and the names are of the form BM_$OP_${NUMTHREADS}T/${N}.

Measured improvements in wall clock time:

Run on rmlarsen3.mtv (12 X 3501 MHz CPUs); 2017-01-20T11:26:31.493023454-08:00
CPU: Intel Haswell with HyperThreading (6 cores) dL1:32KB dL2:256KB dL3:15MB
Benchmark                          Base (ns)  New (ns) Improvement
------------------------------------------------------------------
BM_memcpy_1T/2                          3.48      2.39    +31.3%
BM_memcpy_1T/8                          12.3      6.51    +47.0%
BM_memcpy_1T/64                          371       383     -3.2%
BM_memcpy_1T/512                       66922     66720     +0.3%
BM_memcpy_1T/4k                      9892867   6849682    +30.8%
BM_memcpy_1T/5k                     14951099  10332856    +30.9%
BM_memcpy_2T/2                          3.50      2.46    +29.7%
BM_memcpy_2T/8                          12.3      7.66    +37.7%
BM_memcpy_2T/64                          371       376     -1.3%
BM_memcpy_2T/512                       66652     66788     -0.2%
BM_memcpy_2T/4k                      6145012   6117776     +0.4%
BM_memcpy_2T/5k                      9181478   9010942     +1.9%
BM_memcpy_4T/2                          3.47      2.47    +31.0%
BM_memcpy_4T/8                          12.3      6.67    +45.8
BM_memcpy_4T/64                          374       376     -0.5%
BM_memcpy_4T/512                       67833     68019     -0.3%
BM_memcpy_4T/4k                      5057425   5188253     -2.6%
BM_memcpy_4T/5k                      7555638   7779468     -3.0%
BM_memcpy_6T/2                          3.51      2.50    +28.8%
BM_memcpy_6T/8                          12.3      7.61    +38.1%
BM_memcpy_6T/64                          373       378     -1.3%
BM_memcpy_6T/512                       66871     66774     +0.1%
BM_memcpy_6T/4k                      5112975   5233502     -2.4%
BM_memcpy_6T/5k                      7614180   7772246     -2.1%
BM_memcpy_8T/2                          3.47      2.41    +30.5%
BM_memcpy_8T/8                          12.4      10.5    +15.3%
BM_memcpy_8T/64                          372       388     -4.3%
BM_memcpy_8T/512                       67373     66588     +1.2%
BM_memcpy_8T/4k                      5148462   5254897     -2.1%
BM_memcpy_8T/5k                      7660989   7799058     -1.8%
BM_memcpy_12T/2                         3.50      2.40    +31.4%
BM_memcpy_12T/8                         12.4      7.55    +39.1
BM_memcpy_12T/64                         374       378     -1.1%
BM_memcpy_12T/512                      67132     66683     +0.7%
BM_memcpy_12T/4k                     5185125   5292920     -2.1%
BM_memcpy_12T/5k                     7717284   7942684     -2.9%
BM_slicingSmallPieces_1T/2              47.3      47.5     +0.4%
BM_slicingSmallPieces_1T/8              53.6      52.3     +2.4%
BM_slicingSmallPieces_1T/64              491       476     +3.1%
BM_slicingSmallPieces_1T/512           21734     18814    +13.4%
BM_slicingSmallPieces_1T/4k           394660    396760     -0.5%
BM_slicingSmallPieces_1T/5k           218722    209244     +4.3%
BM_slicingSmallPieces_2T/2              80.7      79.9     +1.0%
BM_slicingSmallPieces_2T/8              54.2      53.1     +2.0
BM_slicingSmallPieces_2T/64              497       477     +4.0%
BM_slicingSmallPieces_2T/512           21732     18822    +13.4%
BM_slicingSmallPieces_2T/4k           392885    390490     +0.6%
BM_slicingSmallPieces_2T/5k           221988    208678     +6.0%
BM_slicingSmallPieces_4T/2              80.8      80.1     +0.9%
BM_slicingSmallPieces_4T/8              54.1      53.2     +1.7%
BM_slicingSmallPieces_4T/64              493       476     +3.4%
BM_slicingSmallPieces_4T/512           21702     18758    +13.6%
BM_slicingSmallPieces_4T/4k           393962    404023     -2.6%
BM_slicingSmallPieces_4T/5k           249667    211732    +15.2%
BM_slicingSmallPieces_6T/2              80.5      80.1     +0.5%
BM_slicingSmallPieces_6T/8              54.4      53.4     +1.8%
BM_slicingSmallPieces_6T/64              488       478     +2.0%
BM_slicingSmallPieces_6T/512           21719     18841    +13.3%
BM_slicingSmallPieces_6T/4k           394950    397583     -0.7%
BM_slicingSmallPieces_6T/5k           223080    210148     +5.8%
BM_slicingSmallPieces_8T/2              81.2      80.4     +1.0%
BM_slicingSmallPieces_8T/8              58.1      53.5     +7.9%
BM_slicingSmallPieces_8T/64              489       480     +1.8%
BM_slicingSmallPieces_8T/512           21586     18798    +12.9%
BM_slicingSmallPieces_8T/4k           394592    400165     -1.4%
BM_slicingSmallPieces_8T/5k           219688    208301     +5.2%
BM_slicingSmallPieces_12T/2             80.2      79.8     +0.7%
BM_slicingSmallPieces_12T/8             54.4      53.4     +1.8
BM_slicingSmallPieces_12T/64             488       476     +2.5%
BM_slicingSmallPieces_12T/512          21931     18831    +14.1%
BM_slicingSmallPieces_12T/4k          393962    396541     -0.7%
BM_slicingSmallPieces_12T/5k          218803    207965     +5.0%
2017-01-24 13:55:18 -08:00

652 lines
23 KiB
C++

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2015 Jianwei Cui <thucjw@gmail.com>
//
// 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/.
#ifndef EIGEN_CXX11_TENSOR_TENSOR_FFT_H
#define EIGEN_CXX11_TENSOR_TENSOR_FFT_H
// This code requires the ability to initialize arrays of constant
// values directly inside a class.
#if __cplusplus >= 201103L || EIGEN_COMP_MSVC >= 1900
namespace Eigen {
/** \class TensorFFT
* \ingroup CXX11_Tensor_Module
*
* \brief Tensor FFT class.
*
* TODO:
* Vectorize the Cooley Tukey and the Bluestein algorithm
* Add support for multithreaded evaluation
* Improve the performance on GPU
*/
template <bool NeedUprade> struct MakeComplex {
template <typename T>
EIGEN_DEVICE_FUNC
T operator() (const T& val) const { return val; }
};
template <> struct MakeComplex<true> {
template <typename T>
EIGEN_DEVICE_FUNC
std::complex<T> operator() (const T& val) const { return std::complex<T>(val, 0); }
};
template <> struct MakeComplex<false> {
template <typename T>
EIGEN_DEVICE_FUNC
std::complex<T> operator() (const std::complex<T>& val) const { return val; }
};
template <int ResultType> struct PartOf {
template <typename T> T operator() (const T& val) const { return val; }
};
template <> struct PartOf<RealPart> {
template <typename T> T operator() (const std::complex<T>& val) const { return val.real(); }
};
template <> struct PartOf<ImagPart> {
template <typename T> T operator() (const std::complex<T>& val) const { return val.imag(); }
};
namespace internal {
template <typename FFT, typename XprType, int FFTResultType, int FFTDir>
struct traits<TensorFFTOp<FFT, XprType, FFTResultType, FFTDir> > : public traits<XprType> {
typedef traits<XprType> XprTraits;
typedef typename NumTraits<typename XprTraits::Scalar>::Real RealScalar;
typedef typename std::complex<RealScalar> ComplexScalar;
typedef typename XprTraits::Scalar InputScalar;
typedef typename conditional<FFTResultType == RealPart || FFTResultType == ImagPart, RealScalar, ComplexScalar>::type OutputScalar;
typedef typename XprTraits::StorageKind StorageKind;
typedef typename XprTraits::Index Index;
typedef typename XprType::Nested Nested;
typedef typename remove_reference<Nested>::type _Nested;
static const int NumDimensions = XprTraits::NumDimensions;
static const int Layout = XprTraits::Layout;
};
template <typename FFT, typename XprType, int FFTResultType, int FFTDirection>
struct eval<TensorFFTOp<FFT, XprType, FFTResultType, FFTDirection>, Eigen::Dense> {
typedef const TensorFFTOp<FFT, XprType, FFTResultType, FFTDirection>& type;
};
template <typename FFT, typename XprType, int FFTResultType, int FFTDirection>
struct nested<TensorFFTOp<FFT, XprType, FFTResultType, FFTDirection>, 1, typename eval<TensorFFTOp<FFT, XprType, FFTResultType, FFTDirection> >::type> {
typedef TensorFFTOp<FFT, XprType, FFTResultType, FFTDirection> type;
};
} // end namespace internal
template <typename FFT, typename XprType, int FFTResultType, int FFTDir>
class TensorFFTOp : public TensorBase<TensorFFTOp<FFT, XprType, FFTResultType, FFTDir>, ReadOnlyAccessors> {
public:
typedef typename Eigen::internal::traits<TensorFFTOp>::Scalar Scalar;
typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
typedef typename std::complex<RealScalar> ComplexScalar;
typedef typename internal::conditional<FFTResultType == RealPart || FFTResultType == ImagPart, RealScalar, ComplexScalar>::type OutputScalar;
typedef OutputScalar CoeffReturnType;
typedef typename Eigen::internal::nested<TensorFFTOp>::type Nested;
typedef typename Eigen::internal::traits<TensorFFTOp>::StorageKind StorageKind;
typedef typename Eigen::internal::traits<TensorFFTOp>::Index Index;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorFFTOp(const XprType& expr, const FFT& fft)
: m_xpr(expr), m_fft(fft) {}
EIGEN_DEVICE_FUNC
const FFT& fft() const { return m_fft; }
EIGEN_DEVICE_FUNC
const typename internal::remove_all<typename XprType::Nested>::type& expression() const {
return m_xpr;
}
protected:
typename XprType::Nested m_xpr;
const FFT m_fft;
};
// Eval as rvalue
template <typename FFT, typename ArgType, typename Device, int FFTResultType, int FFTDir>
struct TensorEvaluator<const TensorFFTOp<FFT, ArgType, FFTResultType, FFTDir>, Device> {
typedef TensorFFTOp<FFT, ArgType, FFTResultType, FFTDir> XprType;
typedef typename XprType::Index Index;
static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
typedef DSizes<Index, NumDims> Dimensions;
typedef typename XprType::Scalar Scalar;
typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
typedef typename std::complex<RealScalar> ComplexScalar;
typedef typename TensorEvaluator<ArgType, Device>::Dimensions InputDimensions;
typedef internal::traits<XprType> XprTraits;
typedef typename XprTraits::Scalar InputScalar;
typedef typename internal::conditional<FFTResultType == RealPart || FFTResultType == ImagPart, RealScalar, ComplexScalar>::type OutputScalar;
typedef OutputScalar CoeffReturnType;
typedef typename PacketType<OutputScalar, Device>::type PacketReturnType;
static const int PacketSize = internal::unpacket_traits<PacketReturnType>::size;
enum {
IsAligned = false,
PacketAccess = true,
BlockAccess = false,
Layout = TensorEvaluator<ArgType, Device>::Layout,
CoordAccess = false,
RawAccess = false
};
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device) : m_fft(op.fft()), m_impl(op.expression(), device), m_data(NULL), m_device(device) {
const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims = m_impl.dimensions();
for (int i = 0; i < NumDims; ++i) {
eigen_assert(input_dims[i] > 0);
m_dimensions[i] = input_dims[i];
}
if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
m_strides[0] = 1;
for (int i = 1; i < NumDims; ++i) {
m_strides[i] = m_strides[i - 1] * m_dimensions[i - 1];
}
} else {
m_strides[NumDims - 1] = 1;
for (int i = NumDims - 2; i >= 0; --i) {
m_strides[i] = m_strides[i + 1] * m_dimensions[i + 1];
}
}
m_size = m_dimensions.TotalSize();
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const {
return m_dimensions;
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(OutputScalar* data) {
m_impl.evalSubExprsIfNeeded(NULL);
if (data) {
evalToBuf(data);
return false;
} else {
m_data = (CoeffReturnType*)m_device.allocate(sizeof(CoeffReturnType) * m_size);
evalToBuf(m_data);
return true;
}
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
if (m_data) {
m_device.deallocate(m_data);
m_data = NULL;
}
m_impl.cleanup();
}
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE CoeffReturnType coeff(Index index) const {
return m_data[index];
}
template <int LoadMode>
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE PacketReturnType
packet(Index index) const {
return internal::ploadt<PacketReturnType, LoadMode>(m_data + index);
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost
costPerCoeff(bool vectorized) const {
return TensorOpCost(sizeof(CoeffReturnType), 0, 0, vectorized, PacketSize);
}
EIGEN_DEVICE_FUNC Scalar* data() const { return m_data; }
private:
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalToBuf(OutputScalar* data) {
const bool write_to_out = internal::is_same<OutputScalar, ComplexScalar>::value;
ComplexScalar* buf = write_to_out ? (ComplexScalar*)data : (ComplexScalar*)m_device.allocate(sizeof(ComplexScalar) * m_size);
for (Index i = 0; i < m_size; ++i) {
buf[i] = MakeComplex<internal::is_same<InputScalar, RealScalar>::value>()(m_impl.coeff(i));
}
for (size_t i = 0; i < m_fft.size(); ++i) {
Index dim = m_fft[i];
eigen_assert(dim >= 0 && dim < NumDims);
Index line_len = m_dimensions[dim];
eigen_assert(line_len >= 1);
ComplexScalar* line_buf = (ComplexScalar*)m_device.allocate(sizeof(ComplexScalar) * line_len);
const bool is_power_of_two = isPowerOfTwo(line_len);
const Index good_composite = is_power_of_two ? 0 : findGoodComposite(line_len);
const Index log_len = is_power_of_two ? getLog2(line_len) : getLog2(good_composite);
ComplexScalar* a = is_power_of_two ? NULL : (ComplexScalar*)m_device.allocate(sizeof(ComplexScalar) * good_composite);
ComplexScalar* b = is_power_of_two ? NULL : (ComplexScalar*)m_device.allocate(sizeof(ComplexScalar) * good_composite);
ComplexScalar* pos_j_base_powered = is_power_of_two ? NULL : (ComplexScalar*)m_device.allocate(sizeof(ComplexScalar) * (line_len + 1));
if (!is_power_of_two) {
// Compute twiddle factors
// t_n = exp(sqrt(-1) * pi * n^2 / line_len)
// for n = 0, 1,..., line_len-1.
// For n > 2 we use the recurrence t_n = t_{n-1}^2 / t_{n-2} * t_1^2
pos_j_base_powered[0] = ComplexScalar(1, 0);
if (line_len > 1) {
const RealScalar pi_over_len(EIGEN_PI / line_len);
const ComplexScalar pos_j_base = ComplexScalar(
std::cos(pi_over_len), std::sin(pi_over_len));
pos_j_base_powered[1] = pos_j_base;
if (line_len > 2) {
const ComplexScalar pos_j_base_sq = pos_j_base * pos_j_base;
for (int j = 2; j < line_len + 1; ++j) {
pos_j_base_powered[j] = pos_j_base_powered[j - 1] *
pos_j_base_powered[j - 1] /
pos_j_base_powered[j - 2] * pos_j_base_sq;
}
}
}
}
for (Index partial_index = 0; partial_index < m_size / line_len; ++partial_index) {
const Index base_offset = getBaseOffsetFromIndex(partial_index, dim);
// get data into line_buf
const Index stride = m_strides[dim];
if (stride == 1) {
m_device.memcpy(line_buf, &buf[base_offset], line_len*sizeof(ComplexScalar));
} else {
Index offset = base_offset;
for (int j = 0; j < line_len; ++j, offset += stride) {
line_buf[j] = buf[offset];
}
}
// processs the line
if (is_power_of_two) {
processDataLineCooleyTukey(line_buf, line_len, log_len);
}
else {
processDataLineBluestein(line_buf, line_len, good_composite, log_len, a, b, pos_j_base_powered);
}
// write back
if (FFTDir == FFT_FORWARD && stride == 1) {
m_device.memcpy(&buf[base_offset], line_buf, line_len*sizeof(ComplexScalar));
} else {
Index offset = base_offset;
const ComplexScalar div_factor = ComplexScalar(1.0 / line_len, 0);
for (int j = 0; j < line_len; ++j, offset += stride) {
buf[offset] = (FFTDir == FFT_FORWARD) ? line_buf[j] : line_buf[j] * div_factor;
}
}
}
m_device.deallocate(line_buf);
if (!is_power_of_two) {
m_device.deallocate(a);
m_device.deallocate(b);
m_device.deallocate(pos_j_base_powered);
}
}
if(!write_to_out) {
for (Index i = 0; i < m_size; ++i) {
data[i] = PartOf<FFTResultType>()(buf[i]);
}
m_device.deallocate(buf);
}
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static bool isPowerOfTwo(Index x) {
eigen_assert(x > 0);
return !(x & (x - 1));
}
// The composite number for padding, used in Bluestein's FFT algorithm
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Index findGoodComposite(Index n) {
Index i = 2;
while (i < 2 * n - 1) i *= 2;
return i;
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static Index getLog2(Index m) {
Index log2m = 0;
while (m >>= 1) log2m++;
return log2m;
}
// Call Cooley Tukey algorithm directly, data length must be power of 2
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void processDataLineCooleyTukey(ComplexScalar* line_buf, Index line_len, Index log_len) {
eigen_assert(isPowerOfTwo(line_len));
scramble_FFT(line_buf, line_len);
compute_1D_Butterfly<FFTDir>(line_buf, line_len, log_len);
}
// Call Bluestein's FFT algorithm, m is a good composite number greater than (2 * n - 1), used as the padding length
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void processDataLineBluestein(ComplexScalar* line_buf, Index line_len, Index good_composite, Index log_len, ComplexScalar* a, ComplexScalar* b, const ComplexScalar* pos_j_base_powered) {
Index n = line_len;
Index m = good_composite;
ComplexScalar* data = line_buf;
for (Index i = 0; i < n; ++i) {
if(FFTDir == FFT_FORWARD) {
a[i] = data[i] * numext::conj(pos_j_base_powered[i]);
}
else {
a[i] = data[i] * pos_j_base_powered[i];
}
}
for (Index i = n; i < m; ++i) {
a[i] = ComplexScalar(0, 0);
}
for (Index i = 0; i < n; ++i) {
if(FFTDir == FFT_FORWARD) {
b[i] = pos_j_base_powered[i];
}
else {
b[i] = numext::conj(pos_j_base_powered[i]);
}
}
for (Index i = n; i < m - n; ++i) {
b[i] = ComplexScalar(0, 0);
}
for (Index i = m - n; i < m; ++i) {
if(FFTDir == FFT_FORWARD) {
b[i] = pos_j_base_powered[m-i];
}
else {
b[i] = numext::conj(pos_j_base_powered[m-i]);
}
}
scramble_FFT(a, m);
compute_1D_Butterfly<FFT_FORWARD>(a, m, log_len);
scramble_FFT(b, m);
compute_1D_Butterfly<FFT_FORWARD>(b, m, log_len);
for (Index i = 0; i < m; ++i) {
a[i] *= b[i];
}
scramble_FFT(a, m);
compute_1D_Butterfly<FFT_REVERSE>(a, m, log_len);
//Do the scaling after ifft
for (Index i = 0; i < m; ++i) {
a[i] /= m;
}
for (Index i = 0; i < n; ++i) {
if(FFTDir == FFT_FORWARD) {
data[i] = a[i] * numext::conj(pos_j_base_powered[i]);
}
else {
data[i] = a[i] * pos_j_base_powered[i];
}
}
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void scramble_FFT(ComplexScalar* data, Index n) {
eigen_assert(isPowerOfTwo(n));
Index j = 1;
for (Index i = 1; i < n; ++i){
if (j > i) {
std::swap(data[j-1], data[i-1]);
}
Index m = n >> 1;
while (m >= 2 && j > m) {
j -= m;
m >>= 1;
}
j += m;
}
}
template <int Dir>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void butterfly_2(ComplexScalar* data) {
ComplexScalar tmp = data[1];
data[1] = data[0] - data[1];
data[0] += tmp;
}
template <int Dir>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void butterfly_4(ComplexScalar* data) {
ComplexScalar tmp[4];
tmp[0] = data[0] + data[1];
tmp[1] = data[0] - data[1];
tmp[2] = data[2] + data[3];
if (Dir == FFT_FORWARD) {
tmp[3] = ComplexScalar(0.0, -1.0) * (data[2] - data[3]);
} else {
tmp[3] = ComplexScalar(0.0, 1.0) * (data[2] - data[3]);
}
data[0] = tmp[0] + tmp[2];
data[1] = tmp[1] + tmp[3];
data[2] = tmp[0] - tmp[2];
data[3] = tmp[1] - tmp[3];
}
template <int Dir>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void butterfly_8(ComplexScalar* data) {
ComplexScalar tmp_1[8];
ComplexScalar tmp_2[8];
tmp_1[0] = data[0] + data[1];
tmp_1[1] = data[0] - data[1];
tmp_1[2] = data[2] + data[3];
if (Dir == FFT_FORWARD) {
tmp_1[3] = (data[2] - data[3]) * ComplexScalar(0, -1);
} else {
tmp_1[3] = (data[2] - data[3]) * ComplexScalar(0, 1);
}
tmp_1[4] = data[4] + data[5];
tmp_1[5] = data[4] - data[5];
tmp_1[6] = data[6] + data[7];
if (Dir == FFT_FORWARD) {
tmp_1[7] = (data[6] - data[7]) * ComplexScalar(0, -1);
} else {
tmp_1[7] = (data[6] - data[7]) * ComplexScalar(0, 1);
}
tmp_2[0] = tmp_1[0] + tmp_1[2];
tmp_2[1] = tmp_1[1] + tmp_1[3];
tmp_2[2] = tmp_1[0] - tmp_1[2];
tmp_2[3] = tmp_1[1] - tmp_1[3];
tmp_2[4] = tmp_1[4] + tmp_1[6];
// SQRT2DIV2 = sqrt(2)/2
#define SQRT2DIV2 0.7071067811865476
if (Dir == FFT_FORWARD) {
tmp_2[5] = (tmp_1[5] + tmp_1[7]) * ComplexScalar(SQRT2DIV2, -SQRT2DIV2);
tmp_2[6] = (tmp_1[4] - tmp_1[6]) * ComplexScalar(0, -1);
tmp_2[7] = (tmp_1[5] - tmp_1[7]) * ComplexScalar(-SQRT2DIV2, -SQRT2DIV2);
} else {
tmp_2[5] = (tmp_1[5] + tmp_1[7]) * ComplexScalar(SQRT2DIV2, SQRT2DIV2);
tmp_2[6] = (tmp_1[4] - tmp_1[6]) * ComplexScalar(0, 1);
tmp_2[7] = (tmp_1[5] - tmp_1[7]) * ComplexScalar(-SQRT2DIV2, SQRT2DIV2);
}
data[0] = tmp_2[0] + tmp_2[4];
data[1] = tmp_2[1] + tmp_2[5];
data[2] = tmp_2[2] + tmp_2[6];
data[3] = tmp_2[3] + tmp_2[7];
data[4] = tmp_2[0] - tmp_2[4];
data[5] = tmp_2[1] - tmp_2[5];
data[6] = tmp_2[2] - tmp_2[6];
data[7] = tmp_2[3] - tmp_2[7];
}
template <int Dir>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void butterfly_1D_merge(
ComplexScalar* data, Index n, Index n_power_of_2) {
// Original code:
// RealScalar wtemp = std::sin(M_PI/n);
// RealScalar wpi = -std::sin(2 * M_PI/n);
const RealScalar wtemp = m_sin_PI_div_n_LUT[n_power_of_2];
const RealScalar wpi = (Dir == FFT_FORWARD)
? m_minus_sin_2_PI_div_n_LUT[n_power_of_2]
: -m_minus_sin_2_PI_div_n_LUT[n_power_of_2];
const ComplexScalar wp(wtemp, wpi);
const ComplexScalar wp_one = wp + ComplexScalar(1, 0);
const ComplexScalar wp_one_2 = wp_one * wp_one;
const ComplexScalar wp_one_3 = wp_one_2 * wp_one;
const ComplexScalar wp_one_4 = wp_one_3 * wp_one;
const Index n2 = n / 2;
ComplexScalar w(1.0, 0.0);
for (Index i = 0; i < n2; i += 4) {
ComplexScalar temp0(data[i + n2] * w);
ComplexScalar temp1(data[i + 1 + n2] * w * wp_one);
ComplexScalar temp2(data[i + 2 + n2] * w * wp_one_2);
ComplexScalar temp3(data[i + 3 + n2] * w * wp_one_3);
w = w * wp_one_4;
data[i + n2] = data[i] - temp0;
data[i] += temp0;
data[i + 1 + n2] = data[i + 1] - temp1;
data[i + 1] += temp1;
data[i + 2 + n2] = data[i + 2] - temp2;
data[i + 2] += temp2;
data[i + 3 + n2] = data[i + 3] - temp3;
data[i + 3] += temp3;
}
}
template <int Dir>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void compute_1D_Butterfly(
ComplexScalar* data, Index n, Index n_power_of_2) {
eigen_assert(isPowerOfTwo(n));
if (n > 8) {
compute_1D_Butterfly<Dir>(data, n / 2, n_power_of_2 - 1);
compute_1D_Butterfly<Dir>(data + n / 2, n / 2, n_power_of_2 - 1);
butterfly_1D_merge<Dir>(data, n, n_power_of_2);
} else if (n == 8) {
butterfly_8<Dir>(data);
} else if (n == 4) {
butterfly_4<Dir>(data);
} else if (n == 2) {
butterfly_2<Dir>(data);
}
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index getBaseOffsetFromIndex(Index index, Index omitted_dim) const {
Index result = 0;
if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
for (int i = NumDims - 1; i > omitted_dim; --i) {
const Index partial_m_stride = m_strides[i] / m_dimensions[omitted_dim];
const Index idx = index / partial_m_stride;
index -= idx * partial_m_stride;
result += idx * m_strides[i];
}
result += index;
}
else {
for (Index i = 0; i < omitted_dim; ++i) {
const Index partial_m_stride = m_strides[i] / m_dimensions[omitted_dim];
const Index idx = index / partial_m_stride;
index -= idx * partial_m_stride;
result += idx * m_strides[i];
}
result += index;
}
// Value of index_coords[omitted_dim] is not determined to this step
return result;
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index getIndexFromOffset(Index base, Index omitted_dim, Index offset) const {
Index result = base + offset * m_strides[omitted_dim] ;
return result;
}
protected:
Index m_size;
const FFT& m_fft;
Dimensions m_dimensions;
array<Index, NumDims> m_strides;
TensorEvaluator<ArgType, Device> m_impl;
CoeffReturnType* m_data;
const Device& m_device;
// This will support a maximum FFT size of 2^32 for each dimension
// m_sin_PI_div_n_LUT[i] = (-2) * std::sin(M_PI / std::pow(2,i)) ^ 2;
const RealScalar m_sin_PI_div_n_LUT[32] = {
RealScalar(0.0),
RealScalar(-2),
RealScalar(-0.999999999999999),
RealScalar(-0.292893218813453),
RealScalar(-0.0761204674887130),
RealScalar(-0.0192147195967696),
RealScalar(-0.00481527332780311),
RealScalar(-0.00120454379482761),
RealScalar(-3.01181303795779e-04),
RealScalar(-7.52981608554592e-05),
RealScalar(-1.88247173988574e-05),
RealScalar(-4.70619042382852e-06),
RealScalar(-1.17654829809007e-06),
RealScalar(-2.94137117780840e-07),
RealScalar(-7.35342821488550e-08),
RealScalar(-1.83835707061916e-08),
RealScalar(-4.59589268710903e-09),
RealScalar(-1.14897317243732e-09),
RealScalar(-2.87243293150586e-10),
RealScalar( -7.18108232902250e-11),
RealScalar(-1.79527058227174e-11),
RealScalar(-4.48817645568941e-12),
RealScalar(-1.12204411392298e-12),
RealScalar(-2.80511028480785e-13),
RealScalar(-7.01277571201985e-14),
RealScalar(-1.75319392800498e-14),
RealScalar(-4.38298482001247e-15),
RealScalar(-1.09574620500312e-15),
RealScalar(-2.73936551250781e-16),
RealScalar(-6.84841378126949e-17),
RealScalar(-1.71210344531737e-17),
RealScalar(-4.28025861329343e-18)
};
// m_minus_sin_2_PI_div_n_LUT[i] = -std::sin(2 * M_PI / std::pow(2,i));
const RealScalar m_minus_sin_2_PI_div_n_LUT[32] = {
RealScalar(0.0),
RealScalar(0.0),
RealScalar(-1.00000000000000e+00),
RealScalar(-7.07106781186547e-01),
RealScalar(-3.82683432365090e-01),
RealScalar(-1.95090322016128e-01),
RealScalar(-9.80171403295606e-02),
RealScalar(-4.90676743274180e-02),
RealScalar(-2.45412285229123e-02),
RealScalar(-1.22715382857199e-02),
RealScalar(-6.13588464915448e-03),
RealScalar(-3.06795676296598e-03),
RealScalar(-1.53398018628477e-03),
RealScalar(-7.66990318742704e-04),
RealScalar(-3.83495187571396e-04),
RealScalar(-1.91747597310703e-04),
RealScalar(-9.58737990959773e-05),
RealScalar(-4.79368996030669e-05),
RealScalar(-2.39684498084182e-05),
RealScalar(-1.19842249050697e-05),
RealScalar(-5.99211245264243e-06),
RealScalar(-2.99605622633466e-06),
RealScalar(-1.49802811316901e-06),
RealScalar(-7.49014056584716e-07),
RealScalar(-3.74507028292384e-07),
RealScalar(-1.87253514146195e-07),
RealScalar(-9.36267570730981e-08),
RealScalar(-4.68133785365491e-08),
RealScalar(-2.34066892682746e-08),
RealScalar(-1.17033446341373e-08),
RealScalar(-5.85167231706864e-09),
RealScalar(-2.92583615853432e-09)
};
};
} // end namespace Eigen
#endif // EIGEN_HAS_CONSTEXPR
#endif // EIGEN_CXX11_TENSOR_TENSOR_FFT_H