[llvm] [ConstraintElim] Add range support for Constraint Elimination (PR #118458)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/118458

None

>From b8bae4478454bfdbaac042a1729b5d5dff057d7f Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Tue, 3 Dec 2024 17:54:42 +0800
Subject: [PATCH] [ConstraintElim] Add range support for Constraint Elimination

---
 .../Scalar/ConstraintElimination.cpp          | 64 ++++++++++++++++++-
 1 file changed, 62 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 4884c23f16e12a..b41255f1dbe7c6 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -1600,8 +1600,68 @@ void ConstraintInfo::addFact(CmpInst::Predicate Pred, Value *A, Value *B,
     DFSInStack.emplace_back(NumIn, NumOut, R.IsSigned,
                             std::move(ValuesToRelease));
 
-    if (!R.IsSigned) {
-      for (Value *V : NewVariables) {
+    // Add range information from attributes and metadata.
+    for (auto *V : NewVariables) {
+      if (V->getType()->getScalarSizeInBits() >= 64)
+        continue;
+      std::optional<ConstantRange> CR;
+      if (const auto *Arg = dyn_cast<Argument>(V)) {
+        CR = Arg->getRange();
+      } else if (const auto *CB = dyn_cast<CallBase>(V)) {
+        CR = CB->getRange();
+      } else if (const auto *I = dyn_cast<Instruction>(V)) {
+        if (auto *Range = I->getMetadata(LLVMContext::MD_range))
+          CR = getConstantRangeFromMetadata(*Range);
+      }
+      if (CR) {
+        if (R.IsSigned) {
+          int64_t MaxVal = CR->getSignedMax().getSExtValue();
+          int64_t MinVal = CR->getSignedMin().getSExtValue();
+          if (MaxVal != MaxConstraintValue) {
+            ConstraintTy VarPos(
+                SmallVector<int64_t, 8>(Value2Index.size() + 1, 0), false,
+                false, false);
+            VarPos.Coefficients[0] = MaxVal;
+            VarPos.Coefficients[Value2Index[V]] = 1;
+            CSToUse.addVariableRow(VarPos.Coefficients);
+            DFSInStack.emplace_back(NumIn, NumOut, R.IsSigned,
+                                    SmallVector<Value *, 2>());
+          }
+          if (MinVal != MinSignedConstraintValue) {
+            ConstraintTy VarPos(
+                SmallVector<int64_t, 8>(Value2Index.size() + 1, 0), false,
+                false, false);
+            VarPos.Coefficients[0] = MinVal;
+            VarPos.Coefficients[Value2Index[V]] = -1;
+            CSToUse.addVariableRow(VarPos.Coefficients);
+            DFSInStack.emplace_back(NumIn, NumOut, R.IsSigned,
+                                    SmallVector<Value *, 2>());
+          }
+        } else {
+          uint64_t MaxVal = CR->getUnsignedMax().getZExtValue();
+          uint64_t MinVal = CR->getUnsignedMin().getZExtValue();
+          if (MaxVal < static_cast<uint64_t>(MaxConstraintValue)) {
+            ConstraintTy VarPos(
+                SmallVector<int64_t, 8>(Value2Index.size() + 1, 0), false,
+                false, false);
+            VarPos.Coefficients[0] = MaxVal;
+            VarPos.Coefficients[Value2Index[V]] = 1;
+            CSToUse.addVariableRow(VarPos.Coefficients);
+            DFSInStack.emplace_back(NumIn, NumOut, R.IsSigned,
+                                    SmallVector<Value *, 2>());
+          }
+          MinVal =
+              std::min(MinVal, static_cast<uint64_t>(MaxConstraintValue - 1));
+          ConstraintTy VarPos(
+              SmallVector<int64_t, 8>(Value2Index.size() + 1, 0), false, false,
+              false);
+          VarPos.Coefficients[0] = MinVal;
+          VarPos.Coefficients[Value2Index[V]] = -1;
+          CSToUse.addVariableRow(VarPos.Coefficients);
+          DFSInStack.emplace_back(NumIn, NumOut, R.IsSigned,
+                                  SmallVector<Value *, 2>());
+        }
+      } else if (!R.IsSigned) {
         ConstraintTy VarPos(SmallVector<int64_t, 8>(Value2Index.size() + 1, 0),
                             false, false, false);
         VarPos.Coefficients[Value2Index[V]] = -1;