[LVI] Infer value ranges from ctpop results

Fixes #115751

Author: zsrkmyn

llvm/lib/Analysis/LazyValueInfo.cpp

@@ -386,6 +386,9 @@ class LazyValueInfoImpl {
 
   void solve();
 
+  ValueLatticeElement getValueFromICmpCtpop(ICmpInst::Predicate Pred,
+                                            Value *RHS);
+
   // For the following methods, if UseBlockValue is true, the function may
   // push additional values to the worklist and return nullopt. If
   // UseBlockValue is false, it will never return nullopt.
@@ -1159,6 +1162,29 @@ getRangeViaSLT(CmpInst::Predicate Pred, APInt RHS,
   return std::nullopt;
 }
 
+/// Get value range for a "ctpop(Val) Pred RHS" condition.
+ValueLatticeElement
+LazyValueInfoImpl::getValueFromICmpCtpop(ICmpInst::Predicate Pred, Value *RHS) {
+  unsigned BitWidth = RHS->getType()->getPrimitiveSizeInBits();
+
+  auto RHSConst = dyn_cast<ConstantInt>(RHS);
+  if (!RHSConst)
+    return ValueLatticeElement::getOverdefined();
+
+  auto &RHSVal = RHSConst->getValue();
+
+  ConstantRange ResValRange = ConstantRange::makeExactICmpRegion(Pred, RHSVal);
+
+  unsigned ResMin = ResValRange.getUnsignedMin().getLimitedValue(BitWidth);
+  unsigned ResMax = ResValRange.getUnsignedMax().getLimitedValue(BitWidth);
+
+  APInt ValMin, ValMax;
+  APInt AllOnes = APInt::getAllOnes(BitWidth);
+  ValMin = AllOnes.lshr(BitWidth - ResMin);
+  ValMax = AllOnes.shl(BitWidth - ResMax);
+  return ValueLatticeElement::getRange(ConstantRange{ValMin, ValMax + 1});
+}
+
 std::optional<ValueLatticeElement> LazyValueInfoImpl::getValueFromICmpCondition(
     Value *Val, ICmpInst *ICI, bool isTrueDest, bool UseBlockValue) {
   Value *LHS = ICI->getOperand(0);
@@ -1192,6 +1218,9 @@ std::optional<ValueLatticeElement> LazyValueInfoImpl::getValueFromICmpCondition(
     return getValueFromSimpleICmpCondition(SwappedPred, LHS, Offset, ICI,
                                            UseBlockValue);
 
+  if (match(LHS, m_Intrinsic<Intrinsic::ctpop>(m_Specific(Val))))
+    return getValueFromICmpCtpop(EdgePred, RHS);
+
   const APInt *Mask, *C;
   if (match(LHS, m_And(m_Specific(Val), m_APInt(Mask))) &&
       match(RHS, m_APInt(C))) {

llvm/test/Transforms/CorrelatedValuePropagation/ctpop-range.ll

@@ -0,0 +1,102 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=correlated-propagation %s | FileCheck %s
+
+declare void @use(i1)
+
+define void @ctpop(i8 %v) {
+; CHECK-LABEL: define void @ctpop(
+; CHECK-SAME: i8 [[V:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[RES:%.*]] = call range(i8 0, 9) i8 @llvm.ctpop.i8(i8 [[V]])
+; CHECK-NEXT:    [[C0_0:%.*]] = icmp samesign uge i8 [[RES]], 3
+; CHECK-NEXT:    [[C0_1:%.*]] = icmp samesign ule i8 [[RES]], 7
+; CHECK-NEXT:    [[C0:%.*]] = and i1 [[C0_0]], [[C0_1]]
+; CHECK-NEXT:    br i1 [[C0]], label %[[RANGE_3_8:.*]], label %[[TEST2:.*]]
+; CHECK:       [[RANGE_3_8]]:
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    call void @use(i1 false)
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    call void @use(i1 false)
+; CHECK-NEXT:    ret void
+; CHECK:       [[TEST2]]:
+; CHECK-NEXT:    [[C2_0:%.*]] = icmp samesign uge i8 [[RES]], 1
+; CHECK-NEXT:    [[C2_1:%.*]] = icmp samesign ule i8 [[RES]], 4
+; CHECK-NEXT:    [[C2:%.*]] = and i1 [[C2_0]], [[C2_1]]
+; CHECK-NEXT:    br i1 [[C2]], label %[[RANGE_1_5:.*]], label %[[TEST3:.*]]
+; CHECK:       [[RANGE_1_5]]:
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    call void @use(i1 false)
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    call void @use(i1 false)
+; CHECK-NEXT:    ret void
+; CHECK:       [[TEST3]]:
+; CHECK-NEXT:    [[C3:%.*]] = icmp samesign uge i8 [[RES]], 8
+; CHECK-NEXT:    br i1 [[C3]], label %[[RANGE_8_9:.*]], label %[[TEST4:.*]]
+; CHECK:       [[RANGE_8_9]]:
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    ret void
+; CHECK:       [[TEST4]]:
+; CHECK-NEXT:    [[C4:%.*]] = icmp eq i8 [[RES]], 0
+; CHECK-NEXT:    br i1 [[C4]], label %[[RANGE_0_1:.*]], label %[[ED:.*]]
+; CHECK:       [[RANGE_0_1]]:
+; CHECK-NEXT:    call void @use(i1 true)
+; CHECK-NEXT:    ret void
+; CHECK:       [[ED]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %res = call range(i8 0, 9) i8 @llvm.ctpop.i8(i8 %v)
+  %c0.0 = icmp uge i8 %res, 3
+  %c0.1 = icmp ule i8 %res, 7
+  %c0 = and i1 %c0.0, %c0.1
+  br i1 %c0, label %range.3.8, label %test2
+
+range.3.8:
+  %cmp0 = icmp uge i8 %v, 7
+  call void @use(i1 %cmp0) ; true
+  %cmp1 = icmp ult i8 %v, 7
+  call void @use(i1 %cmp1) ; false
+  %cmp2 = icmp ule i8 %v, 254
+  call void @use(i1 %cmp2) ; true
+  %cmp3 = icmp ugt i8 %v, 254
+  call void @use(i1 %cmp3) ; false
+  ret void
+
+test2:
+  %c2.0 = icmp uge i8 %res, 1
+  %c2.1 = icmp ule i8 %res, 4
+  %c2 = and i1 %c2.0, %c2.1
+  br i1 %c2, label %range.1.5, label %test3
+
+range.1.5:
+  %cmp8 = icmp uge i8 %v, 1
+  call void @use(i1 %cmp8) ; true
+  %cmp9 = icmp ult i8 %v, 1
+  call void @use(i1 %cmp9) ; false
+  %cmp10 = icmp ule i8 %v, 240
+  call void @use(i1 %cmp10) ; true
+  %cmp11 = icmp ugt i8 %v, 240
+  call void @use(i1 %cmp11) ; false
+  ret void
+
+test3:
+  %c3 = icmp uge i8 %res, 8
+  br i1 %c3, label %range.8.9, label %test4
+
+range.8.9:
+  %cmp4 = icmp eq i8 %v, -1
+  call void @use(i1 %cmp4) ; true
+  ret void
+
+test4:
+  %c4 = icmp eq i8 %res, 0
+  br i1 %c4, label %range.0.1, label %ed
+
+range.0.1:
+  %cmp5 = icmp eq i8 %v, 0
+  call void @use(i1 %cmp5) ; true
+  ret void
+
+ed:
+  ret void
+}