| 583 | } |
| 584 | |
| 585 | Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { |
| 586 | HloInstruction *lhs, *rhs; |
| 587 | CHECK(Match(add, m::Add(m::Op(&lhs), m::Op(&rhs)))); |
| 588 | |
| 589 | // A + 0 => A |
| 590 | VLOG(10) << "trying transform [A + 0 => A]: " << add->ToString(); |
| 591 | if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(add, lhs)) { |
| 592 | return Status::OK(); |
| 593 | } |
| 594 | // 0 + A => A |
| 595 | VLOG(10) << "trying transform [0 + A => A]: " << add->ToString(); |
| 596 | if (IsAll(lhs, 0) && ReplaceInstructionIfSameShape(add, rhs)) { |
| 597 | return Status::OK(); |
| 598 | } |
| 599 | |
| 600 | // Canonicalization: Put constants on the right. This makes the reassociation |
| 601 | // rules below simpler. |
| 602 | VLOG(10) << "trying transform [Const + A => A + Const]"; |
| 603 | if (Match(add, m::Add(m::Constant(), m::NonConstant()))) { |
| 604 | return ReplaceWithNewInstruction( |
| 605 | add, |
| 606 | HloInstruction::CreateBinary(add->shape(), HloOpcode::kAdd, rhs, lhs)); |
| 607 | } |
| 608 | |
| 609 | // Reassociate to allow constant folding. |
| 610 | // |
| 611 | // Note: This is not general. For example, we won't reassociate |
| 612 | // |
| 613 | // (A + C1) + (B + C2) => A + B + (C1 + C2). |
| 614 | // |
| 615 | VLOG(10) << "trying transform [(A + C1) + C2 => A + (C1 + C2)]"; |
| 616 | HloInstruction *a, *c1, *c2; |
| 617 | if (Match(add, m::Add(m::Add(m::NonConstant(&a), m::Constant(&c1)), |
| 618 | m::Constant(&c2))) || |
| 619 | Match(add, m::Add(m::Add(m::NonConstant(&a), |
| 620 | m::Broadcast(m::ConstantScalar(&c1))), |
| 621 | m::Broadcast(m::ConstantScalar(&c2))))) { |
| 622 | TF_ASSIGN_OR_RETURN(auto* sum_of_constants, |
| 623 | MakeBinaryHlo(HloOpcode::kAdd, c1, c2)); |
| 624 | if (ShapeUtil::IsScalar(sum_of_constants->shape()) && |
| 625 | !ShapeUtil::IsScalar(add->shape())) { |
| 626 | sum_of_constants = computation_->AddInstruction( |
| 627 | HloInstruction::CreateBroadcast(add->shape(), sum_of_constants, {})); |
| 628 | } |
| 629 | return ReplaceWithNewInstruction( |
| 630 | add, HloInstruction::CreateBinary(add->shape(), HloOpcode::kAdd, a, |
| 631 | sum_of_constants)); |
| 632 | } |
| 633 | |
| 634 | // Convert add with fullshape into add with partial shape when a |
| 635 | // portion of add is effective: |
| 636 | // zero (fullshape) rhs (partialshape) |
| 637 | // . | | |
| 638 | // . lhs . dynamic_update_slice (fullshape) |
| 639 | // . | | |
| 640 | // Add (fullshape) |
| 641 | // |
| 642 | // to: |
no test coverage detected