[llvm] [ConstraintElim] Optimize usub.sat intrinsic based on known constraints (PR #135744)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Iris (el-ev)

<details>
<summary>Changes</summary>

[ConstraintElim] Optimize usub.sat intrinsic based on known constraints

pre-commit test

update test

---
Full diff: https://github.com/llvm/llvm-project/pull/135744.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/ConstraintElimination.cpp (+29-3) 
- (modified) llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll (+37) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index ad41ea735f0fd..6608d65f797c5 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -1137,12 +1137,15 @@ void State::addInfoFor(BasicBlock &BB) {
           FactOrCheck::getCheck(DT.getNode(&BB), cast<CallInst>(&I)));
       // TODO: Check if it is possible to instead only added the min/max facts
       // when simplifying uses of the min/max intrinsics.
-      if (!isGuaranteedNotToBePoison(&I))
-        break;
+      if (isGuaranteedNotToBePoison(&I))
+        WorkList.push_back(FactOrCheck::getInstFact(DT.getNode(&BB), &I));
+      break;
+    case Intrinsic::usub_sat:
+      WorkList.push_back(
+          FactOrCheck::getCheck(DT.getNode(&BB), cast<CallInst>(&I)));
       [[fallthrough]];
     case Intrinsic::abs:
     case Intrinsic::uadd_sat:
-    case Intrinsic::usub_sat:
       WorkList.push_back(FactOrCheck::getInstFact(DT.getNode(&BB), &I));
       break;
     }
@@ -1519,6 +1522,24 @@ static bool checkAndReplaceCmp(CmpIntrinsic *I, ConstraintInfo &Info,
   return false;
 }
 
+static bool checkAndReplaceUSubSat(SaturatingInst *I, ConstraintInfo &Info,
+                                   SmallVectorImpl<Instruction *> &ToRemove) {
+  Value *LHS = I->getOperand(0);
+  Value *RHS = I->getOperand(1);
+  if (checkCondition(ICmpInst::ICMP_UGT, LHS, RHS, I, Info).value_or(false)) {
+    IRBuilder<> Builder(I->getParent(), I->getIterator());
+    I->replaceAllUsesWith(Builder.CreateSub(LHS, RHS));
+    ToRemove.push_back(I);
+    return true;
+  }
+  if (checkCondition(ICmpInst::ICMP_ULE, LHS, RHS, I, Info).value_or(false)) {
+    I->replaceAllUsesWith(ConstantInt::get(I->getType(), 0));
+    ToRemove.push_back(I);
+    return true;
+  }
+  return false;
+}
+
 static void
 removeEntryFromStack(const StackEntry &E, ConstraintInfo &Info,
                      Module *ReproducerModule,
@@ -1840,6 +1861,11 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
         Changed |= checkAndReplaceMinMax(MinMax, Info, ToRemove);
       } else if (auto *CmpIntr = dyn_cast<CmpIntrinsic>(Inst)) {
         Changed |= checkAndReplaceCmp(CmpIntr, Info, ToRemove);
+      } else if (auto *SatIntr = dyn_cast<SaturatingInst>(Inst)) {
+        if (SatIntr->getIntrinsicID() == Intrinsic::usub_sat)
+          Changed |= checkAndReplaceUSubSat(SatIntr, Info, ToRemove);
+        else
+          llvm_unreachable("Unexpected intrinsic.");
       }
       continue;
     }
diff --git a/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll b/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll
index b0db89dcfdab8..ea0d0d42eb081 100644
--- a/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll
+++ b/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll
@@ -29,6 +29,33 @@ define i1 @usub_sat_ule_lhs(i64 %a, i64 %b) {
   ret i1 %cmp
 }
 
+define i64 @usub_sat_when_lhs_ugt_rhs(i64 %a, i64 %b) {
+; CHECK-LABEL: define i64 @usub_sat_when_lhs_ugt_rhs(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[PRECOND:%.*]] = icmp ugt i64 [[A]], [[B]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    [[SUB_SAT:%.*]] = sub i64 [[A]], [[B]]
+; CHECK-NEXT:    ret i64 [[SUB_SAT]]
+;
+  %precond = icmp ugt i64 %a, %b
+  call void @llvm.assume(i1 %precond)
+  %sub.sat = call i64 @llvm.usub.sat.i64(i64 %a, i64 %b)
+  ret i64 %sub.sat
+}
+
+define i64 @usub_sat_when_lhs_ule_rhs(i64 %a, i64 %b) {
+; CHECK-LABEL: define i64 @usub_sat_when_lhs_ule_rhs(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[PRECOND:%.*]] = icmp ule i64 [[A]], [[B]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    ret i64 0
+;
+  %precond = icmp ule i64 %a, %b
+  call void @llvm.assume(i1 %precond)
+  %sub.sat = call i64 @llvm.usub.sat.i64(i64 %a, i64 %b)
+  ret i64 %sub.sat
+}
+
 ; Negative test
 define i1 @usub_sat_not_ule_rhs(i64 %a, i64 %b) {
 ; CHECK-LABEL: define i1 @usub_sat_not_ule_rhs(
@@ -41,3 +68,13 @@ define i1 @usub_sat_not_ule_rhs(i64 %a, i64 %b) {
   %cmp = icmp ule i64 %sub.sat, %b
   ret i1 %cmp
 }
+
+define i64 @usub_sat_without_precond(i64 %a, i64 %b) {
+; CHECK-LABEL: define i64 @usub_sat_without_precond(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[SUB_SAT:%.*]] = call i64 @llvm.usub.sat.i64(i64 [[A]], i64 [[B]])
+; CHECK-NEXT:    ret i64 [[SUB_SAT]]
+;
+  %sub.sat = call i64 @llvm.usub.sat.i64(i64 %a, i64 %b)
+  ret i64 %sub.sat
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/135744