[ConstraintElim] Simplify cmp after uadd.sat/usub.sat

[el-ev]

• Closes [ConstraintElim] Failure to eliminate cmp of usub.sat #135557

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Iris (el-ev)

Changes

• Closes #135557

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/ConstraintElimination.cpp (+16-1)
• (added) llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll (+36)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 456f5086309cf..ea03b6e5d5570 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 *SatI = dyn_cast<SaturatingInst>(CB.Inst)) {
+        switch (SatI->getIntrinsicID()) {
+        default:
+          continue;
+        case Intrinsic::uadd_sat:
+          Pred = ICmpInst::ICMP_UGE;
+          break;
+        case Intrinsic::usub_sat:
+          Pred = ICmpInst::ICMP_ULE;
+          break;
+        }
+        AddFact(Pred, SatI, SatI->getLHS());
+        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 0000000000000..ec30795a3e95a
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/uadd-usub-sat.ll
@@ -0,0 +1,36 @@
+; 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:    [[PRECOND:%.*]] = icmp ugt i64 [[A]], [[B]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    [[ADD_SAT:%.*]] = call i64 @llvm.uadd.sat.i64(i64 [[A]], i64 1)
+; CHECK-NEXT:    ret i1 true
+;
+  %precond = icmp ugt i64 %a, %b
+  call void @llvm.assume(i1 %precond)
+  %add.sat = call i64 @llvm.uadd.sat.i64(i64 %a, i64 1)
+  %cmp = icmp ugt i64 %add.sat, %b
+  ret i1 %cmp
+}
+
+
+define i1 @usub_sat_ule(i64 %a, i64 %b) {
+; CHECK-LABEL: define i1 @usub_sat_ule(
+; CHECK-SAME: i64 [[A:%.*]], i64 [[B:%.*]]) {
+; CHECK-NEXT:    [[PRECOND:%.*]] = icmp ult i64 [[A]], [[B]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    [[SUB_SAT:%.*]] = call i64 @llvm.usub.sat.i64(i64 [[A]], i64 1)
+; CHECK-NEXT:    ret i1 true
+;
+  %precond = icmp ult i64 %a, %b
+  call void @llvm.assume(i1 %precond)
+  %sub.sat = call i64 @llvm.usub.sat.i64(i64 %a, i64 1)
+  %cmp = icmp ult i64 %sub.sat, %b
+  ret i1 %cmp
+}

[dtcxzyw]

LGTM. Thank you!

[fhahn]

Thanks for the improvement!

There still are a number of improvements to simplifying intrinsics, e.g. add_with_overflow is still missing, while we already support ssub_with_overflow.

[nikic]

We could also simplify usub.sat to sub if LHS >= RHS. This patch just adds fact, but we could also simplify.

[el-ev]

@fhahn @nikic I would like to handle them in some followup patches.

[aeubanks]

I believe this is causing a miscompile:

$ cat /tmp/a.ll
define i1 @f(i64 %conv6, i64 %str.coerce, ptr %conv) {
  %a = load i32, ptr %conv, align 4
  %cmp = icmp sgt i32 %a, -1
  %conv2 = zext nneg i32 %a to i64
  %add = add i64 %str.coerce, %conv6
  %spec.select = call i64 @llvm.usub.sat.i64(i64 %conv2, i64 %add)
  ret i1 %cmp
}

declare i64 @llvm.usub.sat.i64(i64, i64)

$ opt -p constraint-elimination -S /tmp/a.ll
define i1 @f(i64 %conv6, i64 %str.coerce, ptr %conv) {
  %a = load i32, ptr %conv, align 4
  %conv2 = zext nneg i32 %a to i64
  %add = add i64 %str.coerce, %conv6
  %spec.select = call i64 @llvm.usub.sat.i64(i64 %conv2, i64 %add)
  ret i1 true
}

; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare i64 @llvm.usub.sat.i64(i64, i64) #0

[el-ev]

I believe this is causing a miscompile:

$ cat /tmp/a.ll
define i1 @f(i64 %conv6, i64 %str.coerce, ptr %conv) {
  %a = load i32, ptr %conv, align 4
  %cmp = icmp sgt i32 %a, -1
  %conv2 = zext nneg i32 %a to i64
  %add = add i64 %str.coerce, %conv6
  %spec.select = call i64 @llvm.usub.sat.i64(i64 %conv2, i64 %add)
  ret i1 %cmp
}

declare i64 @llvm.usub.sat.i64(i64, i64)

$ opt -p constraint-elimination -S /tmp/a.ll
define i1 @f(i64 %conv6, i64 %str.coerce, ptr %conv) {
  %a = load i32, ptr %conv, align 4
  %conv2 = zext nneg i32 %a to i64
  %add = add i64 %str.coerce, %conv6
  %spec.select = call i64 @llvm.usub.sat.i64(i64 %conv2, i64 %add)
  ret i1 true
}

; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare i64 @llvm.usub.sat.i64(i64, i64) #0

The nneg flag on zext is providing some assumption that %a >=0 holds true

[aeubanks]

I believe it's not UB if %a >= 0 is not true, the later instructions just end up being poison

[el-ev]

I'm unsure what is the correct behavior, but zext nneg does add a constraint here.
https://github.com/llvm/llvm-project/blob/31ddaef8d18d643ff4c343d03ddfe2edae7d22a2/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp#L527C1-L529C33
cc @dtcxzyw

[nikic]

@el-ev The problem is the missing check for

llvm-project/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

Lines 1140 to 1141 in 31ddaef

	if (!isGuaranteedNotToBePoison(&I))
	break;

for usub_sat.

[el-ev]

@nikic Thanks. Could you please take a look at #135744?

[nikic]

@el-ev Would it be possible to fix this separately from adding support for simplifying usub.sat? These are not really related...

[aeubanks]

I've reverted this to keep head green, please reland with the fix