[SCCP] Relax two-instruction range checks (#158495)

If we know x in R1, the range check `x in R2` can be relaxed into `x in
Union(R2, Inverse(R1))`. The latter one may be more efficient if we can
represent it with one icmp.
Fixes regressions introduced by
#156497.

Proof for `(X & -Pow2) == C -> (X - C) < Pow2`:
https://alive2.llvm.org/ce/z/HMgkuu

Compile-time impact:
https://llvm-compile-time-tracker.com/compare.php?from=ead4f3e271fdf6918aef2ede3a7134811147d276&to=bee3d902dd505cf9b11499ba4f230e4e8ae96b92&stat=instructions%3Au

Author: dtcxzyw

llvm/lib/Transforms/Utils/SCCPSolver.cpp

@@ -19,8 +19,10 @@
 #include "llvm/Analysis/ValueLattice.h"
 #include "llvm/Analysis/ValueLatticeUtils.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Casting.h"
@@ -284,6 +286,58 @@ static Value *simplifyInstruction(SCCPSolver &Solver,
     return Sub;
   }
 
+  // Relax range checks.
+  if (auto *ICmp = dyn_cast<ICmpInst>(&Inst)) {
+    Value *X;
+    auto MatchTwoInstructionExactRangeCheck =
+        [&]() -> std::optional<ConstantRange> {
+      const APInt *RHSC;
+      if (!match(ICmp->getOperand(1), m_APInt(RHSC)))
+        return std::nullopt;
+
+      Value *LHS = ICmp->getOperand(0);
+      ICmpInst::Predicate Pred = ICmp->getPredicate();
+      const APInt *Offset;
+      if (match(LHS, m_OneUse(m_AddLike(m_Value(X), m_APInt(Offset)))))
+        return ConstantRange::makeExactICmpRegion(Pred, *RHSC).sub(*Offset);
+      // Match icmp eq/ne X & NegPow2, C
+      if (ICmp->isEquality()) {
+        const APInt *Mask;
+        if (match(LHS, m_OneUse(m_And(m_Value(X), m_NegatedPower2(Mask)))) &&
+            RHSC->countr_zero() >= Mask->countr_zero()) {
+          ConstantRange CR(*RHSC, *RHSC - *Mask);
+          return Pred == ICmpInst::ICMP_EQ ? CR : CR.inverse();
+        }
+      }
+      return std::nullopt;
+    };
+
+    if (auto CR = MatchTwoInstructionExactRangeCheck()) {
+      ConstantRange LRange = GetRange(X);
+      // Early exit if we know nothing about X.
+      if (LRange.isFullSet())
+        return nullptr;
+      // We are allowed to refine the comparison to either true or false for out
+      // of range inputs. Here we refine the comparison to true, i.e. we relax
+      // the range check.
+      auto NewCR = CR->exactUnionWith(LRange.inverse());
+      // TODO: Check if we can narrow the range check to an equality test.
+      // E.g, for X in [0, 4), X - 3 u< 2 -> X == 3
+      if (!NewCR)
+        return nullptr;
+      ICmpInst::Predicate Pred;
+      APInt RHS;
+      // Check if we can represent NewCR as an icmp predicate.
+      if (NewCR->getEquivalentICmp(Pred, RHS)) {
+        IRBuilder<NoFolder> Builder(&Inst);
+        Value *NewICmp =
+            Builder.CreateICmp(Pred, X, ConstantInt::get(X->getType(), RHS));
+        InsertedValues.insert(NewICmp);
+        return NewICmp;
+      }
+    }
+  }
+
   return nullptr;
 }
 

llvm/test/Transforms/SCCP/relax-range-checks.ll

@@ -0,0 +1,92 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=sccp -S | FileCheck %s
+
+define i1 @relax_range_check(i8 range(i8 0, 5) %x)  {
+; CHECK-LABEL: define i1 @relax_range_check(
+; CHECK-SAME: i8 range(i8 0, 5) [[X:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i8 [[X]], -3
+; CHECK-NEXT:    [[RET:%.*]] = icmp uge i8 [[X]], 3
+; CHECK-NEXT:    ret i1 [[RET]]
+;
+  %add = add i8 %x, -3
+  %ret = icmp ult i8 %add, 2
+  ret i1 %ret
+}
+
+define i1 @relax_range_check_highbits_check(i8 range(i8 2, 0) %x)  {
+; CHECK-LABEL: define i1 @relax_range_check_highbits_check(
+; CHECK-SAME: i8 range(i8 2, 0) [[X:%.*]]) {
+; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], -2
+; CHECK-NEXT:    [[RET:%.*]] = icmp ult i8 [[X]], 4
+; CHECK-NEXT:    ret i1 [[RET]]
+;
+  %and = and i8 %x, -2
+  %ret = icmp eq i8 %and, 2
+  ret i1 %ret
+}
+
+; Negative tests.
+
+define i1 @relax_range_check_one_instruction(i8 range(i8 0, 5) %x)  {
+; CHECK-LABEL: define i1 @relax_range_check_one_instruction(
+; CHECK-SAME: i8 range(i8 0, 5) [[X:%.*]]) {
+; CHECK-NEXT:    [[RET:%.*]] = icmp ult i8 [[X]], 2
+; CHECK-NEXT:    ret i1 [[RET]]
+;
+  %ret = icmp ult i8 %x, 2
+  ret i1 %ret
+}
+
+define i1 @relax_range_check_not_profitable(i8 range(i8 0, 6) %x)  {
+; CHECK-LABEL: define i1 @relax_range_check_not_profitable(
+; CHECK-SAME: i8 range(i8 0, 6) [[X:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i8 [[X]], -3
+; CHECK-NEXT:    [[RET:%.*]] = icmp ult i8 [[ADD]], 2
+; CHECK-NEXT:    ret i1 [[RET]]
+;
+  %add = add i8 %x, -3
+  %ret = icmp ult i8 %add, 2
+  ret i1 %ret
+}
+
+define i1 @relax_range_check_unknown_range(i64 %x)  {
+; CHECK-LABEL: define i1 @relax_range_check_unknown_range(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:    [[AND:%.*]] = and i64 [[X]], -67108864
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[AND]], 0
+; CHECK-NEXT:    ret i1 [[TMP1]]
+;
+  %and = and i64 %x, -67108864
+  %test = icmp eq i64 %and, 0
+  ret i1 %test
+}
+
+define i1 @relax_range_check_highbits_check_multiuse(i8 range(i8 2, 0) %x)  {
+; CHECK-LABEL: define i1 @relax_range_check_highbits_check_multiuse(
+; CHECK-SAME: i8 range(i8 2, 0) [[X:%.*]]) {
+; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], -2
+; CHECK-NEXT:    call void @use(i8 [[AND]])
+; CHECK-NEXT:    [[RET:%.*]] = icmp eq i8 [[AND]], 2
+; CHECK-NEXT:    ret i1 [[RET]]
+;
+  %and = and i8 %x, -2
+  call void @use(i8 %and)
+  %ret = icmp eq i8 %and, 2
+  ret i1 %ret
+}
+
+define i1 @relax_range_check_multiuse(i8 range(i8 0, 5) %x)  {
+; CHECK-LABEL: define i1 @relax_range_check_multiuse(
+; CHECK-SAME: i8 range(i8 0, 5) [[X:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i8 [[X]], -3
+; CHECK-NEXT:    call void @use(i8 [[ADD]])
+; CHECK-NEXT:    [[RET:%.*]] = icmp ult i8 [[ADD]], 2
+; CHECK-NEXT:    ret i1 [[RET]]
+;
+  %add = add i8 %x, -3
+  call void @use(i8 %add)
+  %ret = icmp ult i8 %add, 2
+  ret i1 %ret
+}
+
+declare void @use(i8)