Make 1D sweeps along each transform axis.
| 985 | |
| 986 | // Make 1D sweeps along each transform axis. |
| 987 | void Sweep(int64 fft_rank, FftType fft_type, |
| 988 | const absl::Span<const int64> fft_lengths, |
| 989 | const absl::Span<const int64> fft_strides, |
| 990 | absl::Span<complex128> data, absl::Span<complex128> buffer) { |
| 991 | const bool inverse = fft_type == FftType::IFFT || fft_type == FftType::IRFFT; |
| 992 | const bool input_is_truncated = fft_type == FftType::IRFFT; |
| 993 | const bool output_is_truncated = fft_type == FftType::RFFT; |
| 994 | |
| 995 | // Recursively visit each column of the data along the sweep_axis. Calculate |
| 996 | // linearized index of that column's first element and the stride, then invoke |
| 997 | // 1D transform. |
| 998 | // For RFFT, avoid calculating unused output values: first, compute only |
| 999 | // (length_x / 2) + 1 values along the X axis, then limit the X coordinate to |
| 1000 | // [0 ... (length / 2)] during the sweeps along other axes. Similarly, for |
| 1001 | // IRFFT sweep along higher dimensions first, while keeping the X coordinate |
| 1002 | // in the [0 ... (length / 2)] range, then re-create negative frequencies |
| 1003 | // omitted in the input and perform the full-length transform along the X axis |
| 1004 | // in the last sweep. |
| 1005 | std::function<void(int64, int64, int64)> sweep = [&](int64 sweep_axis, |
| 1006 | int64 axis, |
| 1007 | int64 start) { |
| 1008 | if (axis < 0) { |
| 1009 | // Base case: invoke 1D transform. |
| 1010 | const int64 length = fft_lengths[sweep_axis]; |
| 1011 | const int64 stride = fft_strides[sweep_axis]; |
| 1012 | const bool expand_input = input_is_truncated && sweep_axis == 0; |
| 1013 | const bool contract_oputput = output_is_truncated && sweep_axis == 0; |
| 1014 | Dft1D(length, start, stride, inverse, contract_oputput, expand_input, |
| 1015 | data, buffer); |
| 1016 | } else if (axis == sweep_axis) { |
| 1017 | // Visit only the elements with coordinate 0 along the sweep axis. |
| 1018 | sweep(sweep_axis, axis - 1, start); |
| 1019 | } else { |
| 1020 | const int64 length = fft_lengths[axis]; |
| 1021 | const bool is_truncated = input_is_truncated || output_is_truncated; |
| 1022 | const int64 ub = is_truncated && axis == 0 ? (length / 2) + 1 : length; |
| 1023 | for (int64 i = 0; i < ub; i++) { |
| 1024 | sweep(sweep_axis, axis - 1, start + i * fft_strides[axis]); |
| 1025 | } |
| 1026 | } |
| 1027 | }; |
| 1028 | if (input_is_truncated) { |
| 1029 | // Sweep along the X axis last for IRFFT. |
| 1030 | for (int64 sweep_axis = fft_rank - 1; sweep_axis >= 0; sweep_axis--) { |
| 1031 | sweep(sweep_axis, fft_rank - 1, 0); |
| 1032 | } |
| 1033 | } else { |
| 1034 | // Sweep along the X axis first for RFFT. The order does not matter for FFT |
| 1035 | // and IFFT types; handle them here as well. |
| 1036 | for (int64 sweep_axis = 0; sweep_axis < fft_rank; sweep_axis++) { |
| 1037 | sweep(sweep_axis, fft_rank - 1, 0); |
| 1038 | } |
| 1039 | } |
| 1040 | } |
| 1041 | |
| 1042 | // These templates convert the data from the input data type to the type used in |
| 1043 | // calculations and then to the output data type. They are intended to be used |