[ConstraintElim] Eliminate exact right shifts in getConstraint

Author: dtcxzyw

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

@@ -718,6 +718,40 @@ ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
       Pred != CmpInst::ICMP_SLE && Pred != CmpInst::ICMP_SLT)
     return {};
 
+  // Check if we can eliminate right shifts.
+  auto ShiftOperands = [Pred](Value *&Op0, Value *&Op1) {
+    uint64_t ShAmtC;
+    Value *X;
+    if (!match(Op0, m_Exact(m_Shr(m_Value(X), m_ConstantInt(ShAmtC)))))
+      return false;
+    bool IsLShr = cast<BinaryOperator>(Op0)->getOpcode() == Instruction::LShr;
+    if (IsLShr && ICmpInst::isSigned(Pred))
+      return false;
+    // icmp pred (shr exact X, ShAmtC), (shr exact Y, ShAmtC) --> icmp pred X, Y
+    Value *Y;
+    if (match(Op1, m_Exact(m_Shr(m_Value(Y), m_SpecificInt(ShAmtC)))) &&
+        cast<BinaryOperator>(Op0)->getOpcode() ==
+            cast<BinaryOperator>(Op1)->getOpcode()) {
+      Op0 = X;
+      Op1 = Y;
+      return true;
+    }
+    // icmp pred (shr exact X, ShAmtC), C --> icmp pred X, (C << ShAmtC)
+    const APInt *RHSC;
+    if (!match(Op1, m_APInt(RHSC)))
+      return false;
+    bool Overflow = false;
+    APInt NewC = IsLShr ? RHSC->ushl_ov(ShAmtC, Overflow)
+                        : RHSC->sshl_ov(ShAmtC, Overflow);
+    if (Overflow)
+      return false;
+    Op0 = X;
+    Op1 = ConstantInt::get(Op1->getType(), NewC);
+    return true;
+  };
+  if (!ShiftOperands(Op0, Op1))
+    ShiftOperands(Op1, Op0);
+
   SmallVector<ConditionTy, 4> Preconditions;
   bool IsSigned = ForceSignedSystem || CmpInst::isSigned(Pred);
   auto &Value2Index = getValue2Index(IsSigned);

llvm/test/Transforms/ConstraintElimination/shr-exact.ll

@@ -8,8 +8,7 @@ define i1 @precond_icmp_ashr_and_rhsc(i64 %x) {
 ; CHECK-NEXT:    [[SHR:%.*]] = ashr exact i64 [[X]], 3
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ult i64 [[SHR]], 200
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[X]], 9223372036854775800
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %shr = ashr exact i64 %x, 3
@@ -27,8 +26,7 @@ define i1 @precond_icmp_ashr_and_ashr(i64 %x, i64 %y) {
 ; CHECK-NEXT:    [[SHRY:%.*]] = ashr exact i64 [[Y]], 3
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ult i64 [[SHRX]], [[SHRY]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[X]], [[Y]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %shrx = ashr exact i64 %x, 3
@@ -47,8 +45,7 @@ define i1 @precond_icmp_lshr_and_lshr(i64 %x, i64 %y) {
 ; CHECK-NEXT:    [[SHRY:%.*]] = lshr exact i64 [[Y]], 3
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ult i64 [[SHRX]], [[SHRY]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[X]], [[Y]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %shrx = lshr exact i64 %x, 3