aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/compiler/xla/service/hlo_cost_analysis.cc
diff options
context:
space:
mode:
authorGravatar David Majnemer <majnemer@google.com>2017-03-01 17:37:09 -0800
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2017-03-01 17:48:18 -0800
commitaf2c7253bb1f9d135ad9b0c6a271741205ab57fd (patch)
tree5bc2c03e568e73b419531679dce6a86d61064a72 /tensorflow/compiler/xla/service/hlo_cost_analysis.cc
parent5ce1f684d2bb6e9220f68b9541d33342b5452918 (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.cc8
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(