| 40 | using AllReduceSimplifierTest = HloTestBase; |
| 41 | |
| 42 | TEST_F(AllReduceSimplifierTest, ReplicatedParameters) { |
| 43 | const char* kModuleStr = R"( |
| 44 | HloModule m |
| 45 | |
| 46 | sum { |
| 47 | a = f32[] parameter(0) |
| 48 | b = f32[] parameter(1) |
| 49 | ROOT add.2 = f32[] add(a, b) |
| 50 | } |
| 51 | |
| 52 | max { |
| 53 | a.1 = f32[] parameter(0) |
| 54 | b.1 = f32[] parameter(1) |
| 55 | ROOT max = f32[] maximum(a.1, b.1) |
| 56 | } |
| 57 | |
| 58 | min { |
| 59 | a.2 = f32[] parameter(0) |
| 60 | b.2 = f32[] parameter(1) |
| 61 | ROOT min = f32[] minimum(a.2, b.2) |
| 62 | } |
| 63 | |
| 64 | sum.1 { |
| 65 | a.3 = f32[] parameter(0) |
| 66 | b.3 = f32[] parameter(1) |
| 67 | ROOT add.1 = f32[] add(a.3, b.3) |
| 68 | } |
| 69 | |
| 70 | test { |
| 71 | p0 = f32[8,16] parameter(0), parameter_replication={true} |
| 72 | p1 = f32[8,16] parameter(1), parameter_replication={false} |
| 73 | p2 = f32[] parameter(2), parameter_replication={true} |
| 74 | all-reduce = f32[8,16] all-reduce(p0), replica_groups={}, to_apply=sum |
| 75 | all-reduce.1 = f32[8,16] all-reduce(p0), replica_groups={}, to_apply=max |
| 76 | all-reduce.2 = f32[8,16] all-reduce(p1), replica_groups={}, to_apply=min |
| 77 | all-reduce.3 = f32[] all-reduce(p2), replica_groups={}, to_apply=sum.1 |
| 78 | ROOT tuple = (f32[8,16], f32[8,16], f32[8,16], f32[]) tuple(all-reduce, all-reduce.1, all-reduce.2, all-reduce.3) |
| 79 | } |
| 80 | )"; |
| 81 | TF_ASSERT_OK_AND_ASSIGN(auto module, |
| 82 | ParseAndReturnVerifiedModule(kModuleStr)); |
| 83 | AllReduceSimplifier simplifier(/*replica_count=*/8); |
| 84 | ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); |
| 85 | EXPECT_THAT( |
| 86 | module->entry_computation()->root_instruction(), |
| 87 | GmockMatch(m::Tuple( |
| 88 | m::MultiplyAnyOrder(m::Parameter(0), |
| 89 | m::Broadcast(m::Convert(m::ConstantScalar(8)))), |
| 90 | m::Parameter(0), m::AllReduce(m::Parameter(1)), |
| 91 | m::MultiplyAnyOrder(m::Parameter(2), |
| 92 | m::Convert(m::ConstantScalar(8)))))); |
| 93 | } |
| 94 | |
| 95 | TEST_F(AllReduceSimplifierTest, AllReduceAfterAllReduce) { |
| 96 | const char* kModuleStr = R"( |
nothing calls this directly
no test coverage detected