Don't do two map lookups when one will do; NFC

PiperOrigin-RevId: 159001841

1 files changed, 6 insertions, 10 deletions

diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc b/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc
index 5b29686100..b9a496be43 100644
--- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc
+++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc
@@ -51,22 +51,18 @@ void DfsHloVisitor::SetVisited(const HloInstruction& instruction) {
 }
 
 bool DfsHloVisitor::IsVisiting(const HloInstruction& instruction) {
-  if (visit_state_.count(&instruction) == 0) {
-    return false;
-  }
-  return visit_state_[&instruction] == VisitState::kVisiting;
+  auto it = visit_state_.find(&instruction);
+  return it != visit_state_.end() && it->second == VisitState::kVisiting;
 }
 
 bool DfsHloVisitor::DidVisit(const HloInstruction& instruction) {
-  if (visit_state_.count(&instruction) == 0) {
-    return false;
-  }
-  return visit_state_[&instruction] == VisitState::kVisited;
+  auto it = visit_state_.find(&instruction);
+  return it != visit_state_.end() && it->second == VisitState::kVisited;
 }
 
 bool DfsHloVisitor::NotVisited(const HloInstruction& instruction) {
-  return visit_state_.count(&instruction) == 0 ||
-         visit_state_[&instruction] == VisitState::kNotVisited;
+  auto it = visit_state_.find(&instruction);
+  return it == visit_state_.end() || it->second == VisitState::kNotVisited;
 }
 
 Status DfsHloVisitor::Preprocess(HloInstruction* hlo) { return Status::OK(); }