Function Fft1D

tensorflow/compiler/xla/service/hlo_evaluator.cc:874–925

Non-recursive implementation of the Cooley-Tukey radix-2 decimation in time. Performs 1D FFT transform for the lengths, which are powers of 2. Runs in O(length * log(length)) time. Uses the same parameters as the naive implementation above, except that the preallocated buffer must be at least twice as big as the length of the transform, because the buffer is used to hold both input and output valu

// hold both input and output values for each stage of the transform.
//
void Fft1D(int64 length, int64 start, int64 stride, bool inverse,
           bool contract_output, bool expand_input, absl::Span<complex128> data,
           absl::Span<complex128> buffer) {
  CHECK(IsPowerOfTwo(static_cast<uint64>(length)));
  const bool input_is_zero =
      GatherToBuffer(data, length, start, stride, expand_input, buffer);

  if (!input_is_zero) {
    auto generate_twiddles = [](int64 length, bool inverse) {
      std::vector<complex128> twiddles;
      // Need only half the twiddles.
      for (int64 k = 0; k < length / 2; k++) {
        twiddles.push_back(Twiddle(k, length, inverse));
      }
      return twiddles;
    };

    // Indices into the parts of the buffer used for input and output values.
    int64 in_base = length;
    int64 out_base = 0;

    // At each stage, we "split" the input data into num_blocks, with block_size
    // values in each block.
    for (int64 num_blocks = 1; num_blocks < length; num_blocks *= 2) {
      // Swap input and output parts of the buffer.
      std::swap(in_base, out_base);
      auto twiddles = generate_twiddles(num_blocks * 2, inverse);
      const int64 block_size = length / num_blocks;
      const int64 next_iteration_block_size = block_size / 2;
      for (int64 block = 0; block < num_blocks; block++) {
        const int64 in_offset = in_base + block * block_size;
        const int64 out_offset = out_base + block * next_iteration_block_size;
        // For each (even, odd) pair of values in the block, calculate two
        // output values as even + twiddle * odd and even - twiddle * odd.
        for (int64 pair = 0; pair < block_size / 2; pair++) {
          const complex128 even = buffer[in_offset + pair];
          const complex128 odd = buffer[in_offset + block_size / 2 + pair];
          const complex128 twiddled_odd = twiddles[block] * odd;
          buffer[out_offset + pair] = even + twiddled_odd;
          buffer[out_offset + length / 2 + pair] = even - twiddled_odd;
        }
      }
    }
    // Copy computed result back to data.
    const int64 ub = contract_output ? length / 2 + 1 : length;
    for (int64 k = 0; k < ub; k++) {
      complex128 value = buffer[out_base + k];
      data[start + k * stride] =
          inverse ? value / complex128(length, 0.0) : value;
    }
  }
}

// Determine, which implementation of 1D transform to use and call it.
void Dft1D(int64 length, int64 start, int64 stride, bool inverse,