| 695 | } |
| 696 | |
| 697 | Status HloEvaluator::HandleCompare(HloInstruction* compare) { |
| 698 | ComparisonDirection direction = compare->comparison_direction(); |
| 699 | auto lhs = compare->operand(0); |
| 700 | auto rhs = compare->operand(1); |
| 701 | DCHECK(ShapeUtil::SameDimensions(compare->shape(), rhs->shape()) && |
| 702 | ShapeUtil::SameDimensions(lhs->shape(), rhs->shape())); |
| 703 | |
| 704 | TF_RET_CHECK(lhs->shape().element_type() == rhs->shape().element_type()); |
| 705 | |
| 706 | const Literal& lhs_literal = GetEvaluatedLiteralFor(lhs); |
| 707 | const Literal& rhs_literal = GetEvaluatedLiteralFor(rhs); |
| 708 | |
| 709 | // Note here we switch on the operand's type. |
| 710 | switch (lhs->shape().element_type()) { |
| 711 | case PRED: { |
| 712 | TF_ASSIGN_OR_RETURN( |
| 713 | evaluated_[compare], |
| 714 | Compare<bool>(compare->shape(), direction, lhs_literal, rhs_literal)); |
| 715 | } break; |
| 716 | case U8: { |
| 717 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 718 | Compare<uint8>(compare->shape(), direction, |
| 719 | lhs_literal, rhs_literal)); |
| 720 | } break; |
| 721 | case U16: { |
| 722 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 723 | Compare<uint16>(compare->shape(), direction, |
| 724 | lhs_literal, rhs_literal)); |
| 725 | } break; |
| 726 | case U32: { |
| 727 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 728 | Compare<uint32>(compare->shape(), direction, |
| 729 | lhs_literal, rhs_literal)); |
| 730 | } break; |
| 731 | case U64: { |
| 732 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 733 | Compare<uint64>(compare->shape(), direction, |
| 734 | lhs_literal, rhs_literal)); |
| 735 | } break; |
| 736 | case S8: { |
| 737 | TF_ASSIGN_OR_RETURN( |
| 738 | evaluated_[compare], |
| 739 | Compare<int8>(compare->shape(), direction, lhs_literal, rhs_literal)); |
| 740 | } break; |
| 741 | case S16: { |
| 742 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 743 | Compare<int16>(compare->shape(), direction, |
| 744 | lhs_literal, rhs_literal)); |
| 745 | } break; |
| 746 | case S32: { |
| 747 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 748 | Compare<int32>(compare->shape(), direction, |
| 749 | lhs_literal, rhs_literal)); |
| 750 | } break; |
| 751 | case S64: { |
| 752 | TF_ASSIGN_OR_RETURN(evaluated_[compare], |
| 753 | Compare<int64>(compare->shape(), direction, |
| 754 | lhs_literal, rhs_literal)); |
no test coverage detected