Function TEST_F

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule_test.cc:69–123

Test of a single stream, where data dependencies fully determine the execution order.

// Test of a single stream, where data dependencies fully determine the
// execution order.
TEST_F(GpuHloScheduleTest, SequentialMatMul) {
  HloComputation::Builder builder("entry_computation");
  HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter(
      /*parameter_number=*/0, f32_2x2_, /*name=*/"x"));
  HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter(
      /*parameter_number=*/1, f32_2x2_, /*name=*/"y"));
  HloInstruction* z = builder.AddInstruction(HloInstruction::CreateParameter(
      /*parameter_number=*/2, f32_2x2_, /*name=*/"z"));
  HloInstruction* dot1 =
      builder.AddInstruction(CreateCanonicalDot(f32_2x2_, x, y));
  HloInstruction* dot2 =
      builder.AddInstruction(CreateCanonicalDot(f32_2x2_, dot1, z));

  auto module = CreateNewVerifiedModule();
  module->AddEntryComputation(builder.Build(dot2));

  std::unique_ptr<StreamAssignment> streams = AssignStreams(*module);
  EXPECT_EQ(streams->StreamNumberForHlo(*dot1),
            streams->StreamNumberForHlo(*dot2));

  auto schedule = BuildGpuHloSchedule(*module, *streams);
  // Remove parameters, which are unordered.
  EXPECT_EQ(RemoveHlo(schedule->ThunkLaunchOrder(), {x, y, z}),
            HloVec({dot1, dot2}));

  // Parameters x,y,z are mutually unordered, while dot1 and dot2 are
  // transitively ordered by operands.
  auto order = schedule->ConsumeHloOrdering();
  EXPECT_TRUE(order->ExecutesBefore(x, dot1));
  EXPECT_TRUE(order->ExecutesBefore(x, dot2));
  EXPECT_TRUE(order->ExecutesBefore(y, dot1));
  EXPECT_TRUE(order->ExecutesBefore(y, dot2));
  EXPECT_TRUE(order->ExecutesBefore(z, dot2));
  EXPECT_TRUE(order->ExecutesBefore(dot1, dot2));

  EXPECT_FALSE(order->ExecutesBefore(x, x));
  EXPECT_FALSE(order->ExecutesBefore(x, y));
  EXPECT_FALSE(order->ExecutesBefore(x, z));
  EXPECT_FALSE(order->ExecutesBefore(y, x));
  EXPECT_FALSE(order->ExecutesBefore(y, y));
  EXPECT_FALSE(order->ExecutesBefore(y, z));
  EXPECT_FALSE(order->ExecutesBefore(z, x));
  EXPECT_FALSE(order->ExecutesBefore(z, y));
  EXPECT_FALSE(order->ExecutesBefore(z, z));
  EXPECT_FALSE(order->ExecutesBefore(z, dot1));
  EXPECT_FALSE(order->ExecutesBefore(dot1, x));
  EXPECT_FALSE(order->ExecutesBefore(dot1, y));
  EXPECT_FALSE(order->ExecutesBefore(dot1, z));
  EXPECT_FALSE(order->ExecutesBefore(dot1, dot1));
  EXPECT_FALSE(order->ExecutesBefore(dot2, x));
  EXPECT_FALSE(order->ExecutesBefore(dot2, y));
  EXPECT_FALSE(order->ExecutesBefore(dot2, z));
  EXPECT_FALSE(order->ExecutesBefore(dot2, dot1));
  EXPECT_FALSE(order->ExecutesBefore(dot2, dot2));
}

// Test of a single stream, where data dependencies do not fully determine the
// execution order, but the stream assignment does.