[llvm] [ConstraintElim] Add facts implied by llvm.abs (PR #73189)

[Alexander Shaposhnikov via llvm-commits]

https://github.com/alexander-shaposhnikov updated https://github.com/llvm/llvm-project/pull/73189

>From c13819ccea048833d24214221cc52d0b52eaf4a9 Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <ashaposhnikov at google.com>
Date: Sun, 3 Dec 2023 13:42:45 +0000
Subject: [PATCH] [ConstraintElim] Add facts implied by llvm.abs

---
 .../Scalar/ConstraintElimination.cpp          | 52 +++++++++++--------
 .../Transforms/ConstraintElimination/abs.ll   | 39 ++++++++++++++
 2 files changed, 70 insertions(+), 21 deletions(-)
 create mode 100644 llvm/test/Transforms/ConstraintElimination/abs.ll

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 7aadd810c1da3..f0d18f5aa0fd3 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -984,32 +984,37 @@ void State::addInfoFor(BasicBlock &BB) {
       continue;
     }
 
-    if (match(&I, m_Intrinsic<Intrinsic::ssub_with_overflow>())) {
+    auto *II = dyn_cast<IntrinsicInst>(&I);
+    Intrinsic::ID ID = II ? II->getIntrinsicID() : Intrinsic::not_intrinsic;
+    switch (ID) {
+    case Intrinsic::assume:
+      Value *A, *B;
+      CmpInst::Predicate Pred;
+      if (match(I.getOperand(0), m_ICmp(Pred, m_Value(A), m_Value(B)))) {
+        if (GuaranteedToExecute) {
+          // The assume is guaranteed to execute when BB is entered, hence Cond
+          // holds on entry to BB.
+          WorkList.emplace_back(FactOrCheck::getConditionFact(
+              DT.getNode(I.getParent()), Pred, A, B));
+        } else {
+          WorkList.emplace_back(
+              FactOrCheck::getInstFact(DT.getNode(I.getParent()), &I));
+        }
+      }
+      break;
+    case Intrinsic::ssub_with_overflow:
       WorkList.push_back(
           FactOrCheck::getCheck(DT.getNode(&BB), cast<CallInst>(&I)));
-      continue;
-    }
-
-    if (isa<MinMaxIntrinsic>(&I)) {
+      break;
+    case Intrinsic::abs:
+    case Intrinsic::umin:
+    case Intrinsic::umax:
+    case Intrinsic::smin:
+    case Intrinsic::smax:
       WorkList.push_back(FactOrCheck::getInstFact(DT.getNode(&BB), &I));
-      continue;
+      break;
     }
 
-    Value *A, *B;
-    CmpInst::Predicate Pred;
-    // For now, just handle assumes with a single compare as condition.
-    if (match(&I, m_Intrinsic<Intrinsic::assume>(
-                      m_ICmp(Pred, m_Value(A), m_Value(B))))) {
-      if (GuaranteedToExecute) {
-        // The assume is guaranteed to execute when BB is entered, hence Cond
-        // holds on entry to BB.
-        WorkList.emplace_back(FactOrCheck::getConditionFact(
-            DT.getNode(I.getParent()), Pred, A, B));
-      } else {
-        WorkList.emplace_back(
-            FactOrCheck::getInstFact(DT.getNode(I.getParent()), &I));
-      }
-    }
     GuaranteedToExecute &= isGuaranteedToTransferExecutionToSuccessor(&I);
   }
 
@@ -1629,6 +1634,11 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
 
     ICmpInst::Predicate Pred;
     if (!CB.isConditionFact()) {
+      if (Value *X; match(CB.Inst, m_Intrinsic<Intrinsic::abs>(m_Value(X)))) {
+        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());
diff --git a/llvm/test/Transforms/ConstraintElimination/abs.ll b/llvm/test/Transforms/ConstraintElimination/abs.ll
new file mode 100644
index 0000000000000..37fd00535d4b1
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/abs.ll
@@ -0,0 +1,39 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -passes=constraint-elimination -S %s | FileCheck %s
+
+define i1 @abs_int_min_is_not_poison(i32 %arg) {
+; CHECK-LABEL: define i1 @abs_int_min_is_not_poison(
+; CHECK-SAME: i32 [[ARG:%.*]]) {
+; CHECK-NEXT:    [[ABS:%.*]] = tail call i32 @llvm.abs.i32(i32 [[ARG]], i1 false)
+; CHECK-NEXT:    ret i1 true
+;
+  %abs = tail call i32 @llvm.abs.i32(i32 %arg, i1 false)
+  %cmp = icmp sge i32 %abs, %arg
+  ret i1 %cmp
+}
+
+define i1 @abs_int_min_is_poison(i32 %arg) {
+; CHECK-LABEL: define i1 @abs_int_min_is_poison(
+; CHECK-SAME: i32 [[ARG:%.*]]) {
+; CHECK-NEXT:    [[ABS:%.*]] = tail call i32 @llvm.abs.i32(i32 [[ARG]], i1 true)
+; CHECK-NEXT:    ret i1 true
+;
+  %abs = tail call i32 @llvm.abs.i32(i32 %arg, i1 true)
+  %cmp = icmp sge i32 %abs, %arg
+  ret i1 %cmp
+}
+
+define i1 @abs_plus_one(i32 %arg) {
+; CHECK-LABEL: define i1 @abs_plus_one(
+; CHECK-SAME: i32 [[ARG:%.*]]) {
+; CHECK-NEXT:    [[ABS:%.*]] = tail call i32 @llvm.abs.i32(i32 [[ARG]], i1 true)
+; CHECK-NEXT:    [[ABS_PLUS_ONE:%.*]] = add nsw i32 [[ABS]], 1
+; CHECK-NEXT:    ret i1 true
+;
+  %abs = tail call i32 @llvm.abs.i32(i32 %arg, i1 true)
+  %abs_plus_one = add nsw i32 %abs, 1
+  %cmp = icmp sge i32 %abs_plus_one, %arg
+  ret i1 %cmp
+}
+
+declare i32 @llvm.abs.i32(i32, i1 immarg)