[XLA] Several fixes to HLO reachability analysis.

(1) Account for control dependencies in reachability.
(2) Invert sense of reachability. We draw our HLO graphs with arrows from producers to consumers so it makes more sense for reachability to be defined along the direction of these edges.
(3) Rename ComputeTransitiveOperands to ComputeReachability.

PiperOrigin-RevId: 160366307

1 files changed, 10 insertions, 5 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index ff76cc7bf6..4b33947371 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -587,14 +587,19 @@ void HloComputation::ReachabilityMap::SetReachableAndTransitiveClosure(
 }
 
 std::unique_ptr<HloComputation::ReachabilityMap>
-HloComputation::ComputeTransitiveOperands() const {
-  const auto all = MakeInstructionPostOrder();
+HloComputation::ComputeReachability() const {
+  const std::list<HloInstruction*> all = MakeInstructionPostOrder();
   auto result = MakeUnique<HloComputation::ReachabilityMap>(all);
 
-  // Fill in the dependency bit matrix
-  for (const auto* hlo : all) {
+  // Fill in the dependency bit matrix. Iterate in reverse topological order.
+  for (auto it = all.rbegin(); it != all.rend(); ++it) {
+    const HloInstruction* hlo = *it;
+    result->SetReachable(hlo, hlo);
     for (const HloInstruction* operand : hlo->operands()) {
-      result->SetReachableAndTransitiveClosure(hlo, operand);
+      result->SetReachableAndTransitiveClosure(operand, hlo);
+    }
+    for (const HloInstruction* pred : hlo->control_predecessors()) {
+      result->SetReachableAndTransitiveClosure(pred, hlo);
     }
   }
   return result;