| 43 | }; |
| 44 | |
| 45 | TEST_F(StreamAssignmentTest, SequentialMatMul) { |
| 46 | HloComputation::Builder builder("entry_computation"); |
| 47 | HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter( |
| 48 | /*parameter_number=*/0, f32_2x2_, /*name=*/"x")); |
| 49 | HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter( |
| 50 | /*parameter_number=*/1, f32_2x2_, /*name=*/"y")); |
| 51 | HloInstruction* z = builder.AddInstruction(HloInstruction::CreateParameter( |
| 52 | /*parameter_number=*/2, f32_2x2_, /*name=*/"z")); |
| 53 | HloInstruction* dot1 = |
| 54 | builder.AddInstruction(CreateCanonicalDot(f32_2x2_, x, y)); |
| 55 | HloInstruction* dot2 = |
| 56 | builder.AddInstruction(CreateCanonicalDot(f32_2x2_, dot1, z)); |
| 57 | |
| 58 | auto module = CreateNewVerifiedModule(); |
| 59 | module->AddEntryComputation(builder.Build(dot2)); |
| 60 | |
| 61 | std::unique_ptr<StreamAssignment> assignment = AssignStreams(*module); |
| 62 | EXPECT_EQ(assignment->StreamNumberForHlo(*dot1), |
| 63 | assignment->StreamNumberForHlo(*dot2)); |
| 64 | } |
| 65 | |
| 66 | TEST_F(StreamAssignmentTest, ConcurrentMatMul) { |
| 67 | HloComputation::Builder builder("entry_computation"); |
nothing calls this directly
no test coverage detected