| 531 | } |
| 532 | |
| 533 | StatusOr<bool> ArCrsCombiner::RewriteGraph() { |
| 534 | if (all_reduce_map_.empty()) { |
| 535 | return false; |
| 536 | } |
| 537 | for (auto it : all_reduce_map_) { |
| 538 | auto pairs_vec = it.second; |
| 539 | for (auto pair : pairs_vec) { |
| 540 | auto all_reduce = pair.ar; |
| 541 | auto parent_computation = all_reduce->parent(); |
| 542 | auto channel_id = all_reduce->channel_id(); |
| 543 | auto prev = all_reduce->mutable_operand(0); |
| 544 | auto next = all_reduce->users()[0]; |
| 545 | TF_CHECK_OK(all_reduce->ReplaceUseWith(next, prev)); |
| 546 | TF_CHECK_OK(parent_computation->RemoveInstruction(all_reduce)); |
| 547 | while (!next->IsCrossReplicaAllReduce()) { |
| 548 | switch (next->opcode()) { |
| 549 | case HloOpcode::kBitcast: |
| 550 | case HloOpcode::kTranspose: |
| 551 | case HloOpcode::kReshape: |
| 552 | case HloOpcode::kConvert: |
| 553 | case HloOpcode::kMultiply: |
| 554 | break; |
| 555 | case HloOpcode::kAdd: |
| 556 | case HloOpcode::kSubtract: { |
| 557 | auto other_operand = (next->operands()[0] == prev) |
| 558 | ? next->operands()[1] |
| 559 | : next->operands()[0]; |
| 560 | // To move the AR past the addition/subtraction, we need to divide |
| 561 | // other_operand by the number of spatial partitions, except if |
| 562 | // other_operand is a cross-module AR, which can be eliminated. |
| 563 | if (other_operand->IsCrossModuleAllReduce() && |
| 564 | other_operand->user_count() == 1) { |
| 565 | TF_CHECK_OK(other_operand->ReplaceAllUsesWith( |
| 566 | other_operand->mutable_operand(0))); |
| 567 | } else { |
| 568 | auto shape = other_operand->shape(); |
| 569 | Literal lit(shape); |
| 570 | lit.PopulateWithValue<float>(num_spatial_partitions_); |
| 571 | auto divisor = parent_computation->AddInstruction( |
| 572 | HloInstruction::CreateConstant(lit.Clone())); |
| 573 | auto division = parent_computation->AddInstruction( |
| 574 | HloInstruction::CreateBinary(shape, HloOpcode::kDivide, |
| 575 | other_operand, divisor)); |
| 576 | TF_CHECK_OK(other_operand->ReplaceUseWith(next, division)); |
| 577 | } |
| 578 | break; |
| 579 | } |
| 580 | default: |
| 581 | LOG(FATAL) << "Unexpected instruction: " << next->ToShortString(); |
| 582 | } |
| 583 | prev = next; |
| 584 | next = next->users()[0]; |
| 585 | } |
| 586 | // The AllReduce and the CRS are combined to an all-core AllReduce. |
| 587 | // |
| 588 | // Note that we can just reuse the ReplicaGroup config of cross-replica |
| 589 | // all-reduce since we already checked that cross-partition all-reduce |
| 590 | // is always across all partitions (HasCombinableReplicaGroup). We need to |
nothing calls this directly
no test coverage detected