[XLA] Make HLO ordering module-scoped.

Add comparison of ordering of HLO instructions which are in different computations using the call graph. Previously, instructions in different computations were considered unordered. Ordering these instructions improves buffer liveness analysis and may enable better buffer sharing between values in different computations.
Change: 154592912

1 files changed, 77 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc
index 425bee601a..01b5fd9364 100644
--- a/tensorflow/compiler/xla/service/hlo_ordering_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc
@@ -78,6 +78,83 @@ TEST_F(HloOrderingTest, LastUseScheduledFirst) {
   EXPECT_TRUE(ordering.ExecutesBefore(add, negate));
 }
 
+TEST_F(HloOrderingTest, InstructionsInDifferentComputations) {
+  // Tests the ordering of instructions in different computations using the
+  // following HLO code:
+  //
+  // Entry computation:
+  //   %x = Call(A, {})
+  //   %y = Call(B, {%x})
+  //
+  // Computation A:
+  //   %a = Call(C, {})
+  //
+  // Computation B:
+  //   %b = Call(C, {})
+  //
+  // Computation C:
+  //   %c = Constant(42.0f)
+  //
+  // This results in a diamond-shaped callgraph.
+  HloModule module(TestName());
+  const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {});
+
+  auto builder_c = HloComputation::Builder("C");
+  HloInstruction* c = builder_c.AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(42.0f)));
+  HloComputation* computation_c =
+      module.AddEmbeddedComputation(builder_c.Build());
+
+  auto builder_b = HloComputation::Builder("B");
+  builder_b.AddInstruction(
+      HloInstruction::CreateParameter(0, scalar_shape, "param"));
+  HloInstruction* b = builder_b.AddInstruction(
+      HloInstruction::CreateCall(scalar_shape, {}, computation_c));
+  HloComputation* computation_b =
+      module.AddEmbeddedComputation(builder_b.Build());
+
+  auto builder_a = HloComputation::Builder("A");
+  HloInstruction* a = builder_a.AddInstruction(
+      HloInstruction::CreateCall(scalar_shape, {}, computation_c));
+  HloComputation* computation_a =
+      module.AddEmbeddedComputation(builder_a.Build());
+
+  auto builder = HloComputation::Builder(TestName());
+  HloInstruction* x = builder.AddInstruction(
+      HloInstruction::CreateCall(scalar_shape, {}, computation_a));
+  HloInstruction* y = builder.AddInstruction(
+      HloInstruction::CreateCall(scalar_shape, {x}, computation_b));
+  module.AddEntryComputation(builder.Build());
+
+  DependencyHloOrdering ordering(&module);
+  EXPECT_TRUE(ordering.ExecutesBefore(x, y));
+  EXPECT_FALSE(ordering.ExecutesBefore(y, x));
+
+  EXPECT_TRUE(ordering.ExecutesBefore(a, b));
+  EXPECT_FALSE(ordering.ExecutesBefore(b, a));
+
+  EXPECT_FALSE(ordering.ExecutesBefore(a, x));
+  EXPECT_TRUE(ordering.ExecutesBefore(a, y));
+  EXPECT_FALSE(ordering.ExecutesBefore(x, a));
+  EXPECT_FALSE(ordering.ExecutesBefore(y, a));
+
+  EXPECT_FALSE(ordering.ExecutesBefore(b, x));
+  EXPECT_FALSE(ordering.ExecutesBefore(b, y));
+  EXPECT_TRUE(ordering.ExecutesBefore(x, b));
+  EXPECT_FALSE(ordering.ExecutesBefore(y, b));
+
+  // Instruction 'c' is called from multiple callsites and should be unordered
+  // relative to all other instructions in the module.
+  EXPECT_FALSE(ordering.ExecutesBefore(c, a));
+  EXPECT_FALSE(ordering.ExecutesBefore(c, b));
+  EXPECT_FALSE(ordering.ExecutesBefore(c, x));
+  EXPECT_FALSE(ordering.ExecutesBefore(c, y));
+  EXPECT_FALSE(ordering.ExecutesBefore(a, c));
+  EXPECT_FALSE(ordering.ExecutesBefore(b, c));
+  EXPECT_FALSE(ordering.ExecutesBefore(x, c));
+  EXPECT_FALSE(ordering.ExecutesBefore(y, c));
+}
+
 }  // namespace
 
 }  // namespace xla