| 107 | // on outer dimensions. |
| 108 | template <int FFTRank, typename EigenDevice> |
| 109 | void EigenFftC2R(const EigenDevice& device, float* out, complex64* operand, |
| 110 | int64 input_batch, int64 fft_length0, int64 fft_length1, |
| 111 | int64 fft_length2) { |
| 112 | const std::array<int64, 3> fft_shape = { |
| 113 | {fft_length0, fft_length1, fft_length2}}; |
| 114 | |
| 115 | Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> in_dims; |
| 116 | in_dims[0] = input_batch; |
| 117 | Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> out_dims; |
| 118 | out_dims[0] = input_batch; |
| 119 | for (int i = 0; i < FFTRank; i++) { |
| 120 | in_dims[i + 1] = i == FFTRank - 1 ? fft_shape[i] / 2 + 1 : fft_shape[i]; |
| 121 | out_dims[i + 1] = fft_shape[i]; |
| 122 | } |
| 123 | const Eigen::TensorMap<Eigen::Tensor<complex64, FFTRank + 1, Eigen::RowMajor>, |
| 124 | Eigen::Aligned> |
| 125 | input(operand, in_dims); |
| 126 | Eigen::TensorMap<Eigen::Tensor<float, FFTRank + 1, Eigen::RowMajor>, |
| 127 | Eigen::Aligned> |
| 128 | output(out, out_dims); |
| 129 | |
| 130 | // Calculate the shape of the temporary tensor for the full FFT and the |
| 131 | // region we will slice from input given fft_shape. We slice input to |
| 132 | // fft_shape on its inner-most dimensions, except the last (which we |
| 133 | // slice to fft_shape[-1] / 2 + 1). |
| 134 | Eigen::Tensor<complex64, FFTRank + 1, Eigen::RowMajor> full_fft(out_dims); |
| 135 | |
| 136 | // Calculate the starting point and range of the source of |
| 137 | // negative frequency part. |
| 138 | auto neg_sizes = in_dims; |
| 139 | neg_sizes[FFTRank] = fft_shape[FFTRank - 1] - in_dims[FFTRank]; |
| 140 | Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> neg_target_indices; |
| 141 | neg_target_indices[FFTRank] = in_dims[FFTRank]; |
| 142 | |
| 143 | const Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> zero_start_indices; |
| 144 | Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> neg_start_indices; |
| 145 | neg_start_indices[FFTRank] = 1; |
| 146 | |
| 147 | full_fft.slice(zero_start_indices, in_dims).device(device) = input; |
| 148 | |
| 149 | // First, conduct IFFTs on outer dimensions. We save computation (and |
| 150 | // avoid touching uninitialized memory) by slicing full_fft to the |
| 151 | // subregion we wrote input to. |
| 152 | if (FFTRank > 1) { |
| 153 | const auto outer_axes = |
| 154 | Eigen::ArrayXi::LinSpaced(FFTRank - 1, 1, FFTRank - 1); |
| 155 | full_fft.slice(zero_start_indices, in_dims).device(device) = |
| 156 | full_fft.slice(zero_start_indices, in_dims) |
| 157 | .template fft<Eigen::BothParts, Eigen::FFT_REVERSE>(outer_axes); |
| 158 | } |
| 159 | |
| 160 | // Reconstruct the full FFT by appending reversed and conjugated |
| 161 | // spectrum as the negative frequency part. |
| 162 | Eigen::array<bool, FFTRank + 1> reverse_last_axis; |
| 163 | for (auto i = 0; i <= FFTRank; i++) { |
| 164 | reverse_last_axis[i] = i == FFTRank; |
| 165 | } |
| 166 | |