| 663 | } |
| 664 | |
| 665 | Status HloCostAnalysis::HandleFft(const HloInstruction* fft) { |
| 666 | auto real_shape = |
| 667 | fft->operand(0)->shape().IsTuple() |
| 668 | ? ShapeUtil::GetTupleElementShape(fft->operand(0)->shape(), 0) |
| 669 | : fft->operand(0)->shape(); |
| 670 | constexpr int kFmaPerComplexMul = 4; |
| 671 | int64 log_factors = 1; |
| 672 | for (int64 dim : fft->fft_length()) { |
| 673 | log_factors *= tensorflow::Log2Floor(dim); |
| 674 | } |
| 675 | current_properties_[kFlopsKey] = kFmaFlops * kFmaPerComplexMul * log_factors * |
| 676 | ShapeUtil::ElementsIn(real_shape); |
| 677 | return Status::OK(); |
| 678 | } |
| 679 | |
| 680 | Status HloCostAnalysis::HandleTriangularSolve(const HloInstruction* hlo) { |
| 681 | // Half of operand 0 is read. |