Method HandleFft

tensorflow/compiler/xla/service/hlo_evaluator.cc:1362–1434

Flexible implementation of the discrete Fourier transform. All transform types (FFT, IFFT, RFFT, and IRFFT) are supported, as well as the arbitrary rank and length of each dimension of the transform, and arbitrary layouts of the input and output literals. The input literal in operand 0 provides input data, which must be complex64 for FFT, IFFT, IRFFT transforms and float for RFFT. The transform i

// O(N0*N1*...*Nn*(log(N0)+log(N1)+...+log(Nn)) in the best case.
//
Status HloEvaluator::HandleFft(HloInstruction* fft) {
  const FftType fft_type = fft->fft_type();
  std::vector<int64> fft_lengths = fft->fft_length();
  const int64 fft_rank = fft_lengths.size();
  const Literal& input_literal = GetEvaluatedLiteralFor(fft->operand(0));
  const Shape& input_shape = input_literal.shape();
  const Shape& output_shape = fft->shape();
  Literal output_literal = Literal::CreateFromShape(output_shape);

  // Make fft_lengths[0] the minor-most dimension.
  absl::c_reverse(fft_lengths);

  TF_RETURN_IF_ERROR(CheckParameters(input_shape, output_shape, fft_rank,
                                     fft_type, fft_lengths));

  const auto fft_strides = ComputeStrides(fft_lengths);

  // Working set size.
  const int64 fft_size = fft_strides[fft_rank];

  if (fft_size > 0) {
    // Linearized working data set.
    std::vector<complex128> data(fft_size);

    // Temporary buffer allocated once and used in 1D sweeps. For dimension
    // length values that are powers of 2, the buffer should be twice as large.
    int64 buffer_size = 0;
    for (auto len : fft_lengths) {
      int64 size = IsPowerOfTwo(static_cast<uint64>(len)) ? len * 2 : len;
      buffer_size = std::max(buffer_size, size);
    }
    std::vector<complex128> buffer(buffer_size);

    // Sizes of each axis of input and output literals.
    const auto input_lengths = GetDimensionLengths(input_literal);
    const auto output_lengths = GetDimensionLengths(output_literal);

    // Strides for generating linearized indices into multidimensional arrays.
    const auto input_strides = ComputeStrides(input_lengths, input_literal);
    const auto output_strides = ComputeStrides(output_lengths, output_literal);

    // Visit all elements in the dimensions with ranks above the FFT rank. For
    // each such element invoke the transform. Use separate indices for the
    // input and the output to allow different layouts.
    auto base_case = [&](int64 axis, int64 output_index, int64 input_index,
                         bool within_src_bounds) {
      if (axis == fft_rank - 1) {
        // Base case: copy the data from the input literal, apply the
        // transform, and copy the result to the output literal.
        CHECK(within_src_bounds);
        bool input_is_zero =
            CopyDataFromInput(input_literal, input_index, fft_rank, fft_type,
                              fft_size, fft_lengths, fft_strides, input_lengths,
                              input_strides, absl::MakeSpan(data));
        if (!input_is_zero) {
          // Make 1D sweeps along each transform axis.
          Sweep(fft_rank, fft_type, fft_lengths, fft_strides,
                absl::MakeSpan(data), absl::MakeSpan(buffer));