diff options
author | 2017-03-01 17:37:09 -0800 | |
---|---|---|
committer | 2017-03-01 17:48:18 -0800 | |
commit | af2c7253bb1f9d135ad9b0c6a271741205ab57fd (patch) | |
tree | 5bc2c03e568e73b419531679dce6a86d61064a72 /tensorflow/compiler/xla/service/hlo_cost_analysis.cc | |
parent | 5ce1f684d2bb6e9220f68b9541d33342b5452918 (diff) |
[XLA] Add support for profiling multiple computations
While we are here, add support for getting the cost analysis for call HLOs.
Change: 148952748
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_cost_analysis.cc')
-rw-r--r-- | tensorflow/compiler/xla/service/hlo_cost_analysis.cc | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index 41be997584..8fe1897e75 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -369,7 +369,13 @@ Status HloCostAnalysis::HandleFusion(HloInstruction* fusion) { Status HloCostAnalysis::HandleCall( HloInstruction* call, tensorflow::gtl::ArraySlice<HloInstruction*> operands, HloComputation* computation) { - return Unimplemented("call"); + HloCostAnalysis computation_visitor(shape_size_); + TF_RETURN_IF_ERROR(computation->Accept(&computation_visitor)); + + current_flop_count_ = computation_visitor.flop_count(); + current_transcendental_count_ = computation_visitor.transcendental_count(); + current_bytes_accessed_ = computation_visitor.bytes_accessed(); + return Status::OK(); } Status HloCostAnalysis::HandleCustomCall( |