[ConstraintElim] Simplify cmp after uadd.sat/usub.sat (#135603)

- Closes #135557

2 files changed, 59 insertions, 1 deletions

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 456f508..ad41ea7 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -1141,6 +1141,8 @@ void State::addInfoFor(BasicBlock &BB) {
         break;
       [[fallthrough]];
     case Intrinsic::abs:
+    case Intrinsic::uadd_sat:
+    case Intrinsic::usub_sat:
       WorkList.push_back(FactOrCheck::getInstFact(DT.getNode(&BB), &I));
       break;
     }
@@ -1891,13 +1893,26 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
         AddFact(CmpInst::ICMP_SGE, CB.Inst, X);
         continue;
       }
-
       if (auto *MinMax = dyn_cast<MinMaxIntrinsic>(CB.Inst)) {
         Pred = ICmpInst::getNonStrictPredicate(MinMax->getPredicate());
         AddFact(Pred, MinMax, MinMax->getLHS());
         AddFact(Pred, MinMax, MinMax->getRHS());
         continue;
       }
+      if (auto *USatI = dyn_cast<SaturatingInst>(CB.Inst)) {
+        switch (USatI->getIntrinsicID()) {
+        default:
+          llvm_unreachable("Unexpected intrinsic.");
+        case Intrinsic::uadd_sat:
+          AddFact(ICmpInst::ICMP_UGE, USatI, USatI->getLHS());
+          AddFact(ICmpInst::ICMP_UGE, USatI, USatI->getRHS());
+          break;
+        case Intrinsic::usub_sat:
+          AddFact(ICmpInst::ICMP_ULE, USatI, USatI->getLHS());
+          break;
+        }
+        continue;
+      }
     }
 
     Value *A = nullptr, *B = nullptr;
diff --git a/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll b/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll
new file mode 100644
index 0000000..b0db89d
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll
@@ -0,0 +1,43 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=constraint-elimination -S %s | FileCheck %s
+
+declare i64 @llvm.uadd.sat.i64(i64, i64)
+declare i64 @llvm.usub.sat.i64(i64, i64)
+
+define i1 @uadd_sat_uge(i64 %a, i64 %b) {
+; CHECK-LABEL: define i1 @uadd_sat_uge(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[ADD_SAT:%.*]] = call i64 @llvm.uadd.sat.i64(i64 [[A]], i64 [[B]])
+; CHECK-NEXT:    [[CMP:%.*]] = and i1 true, true
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %add.sat = call i64 @llvm.uadd.sat.i64(i64 %a, i64 %b)
+  %cmp1 = icmp uge i64 %add.sat, %a
+  %cmp2 = icmp uge i64 %add.sat, %b
+  %cmp = and i1 %cmp1, %cmp2
+  ret i1 %cmp
+}
+
+define i1 @usub_sat_ule_lhs(i64 %a, i64 %b) {
+; CHECK-LABEL: define i1 @usub_sat_ule_lhs(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[SUB_SAT:%.*]] = call i64 @llvm.usub.sat.i64(i64 [[A]], i64 [[B]])
+; CHECK-NEXT:    ret i1 true
+;
+  %sub.sat = call i64 @llvm.usub.sat.i64(i64 %a, i64 %b)
+  %cmp = icmp ule i64 %sub.sat, %a
+  ret i1 %cmp
+}
+
+; Negative test
+define i1 @usub_sat_not_ule_rhs(i64 %a, i64 %b) {
+; CHECK-LABEL: define i1 @usub_sat_not_ule_rhs(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[SUB_SAT:%.*]] = call i64 @llvm.usub.sat.i64(i64 [[A]], i64 [[B]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ule i64 [[SUB_SAT]], [[B]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %sub.sat = call i64 @llvm.usub.sat.i64(i64 %a, i64 %b)
+  %cmp = icmp ule i64 %sub.sat, %b
+  ret i1 %cmp
+}