[ValueTracking] Infer is-power-of-2 from assumptions. (#107745)

This patch tries to infer is-power-of-2 from assumptions. I don't see
that this kind of assumption exists in my dataset.
Related issue: rust-lang/rust#129795

Close #58996.

Author: dtcxzyw

llvm/lib/Analysis/ValueTracking.cpp

@@ -613,7 +613,7 @@ static bool isKnownNonZeroFromAssume(const Value *V, const SimplifyQuery &Q) {
     CmpInst::Predicate Pred;
     auto m_V = m_CombineOr(m_Specific(V), m_PtrToInt(m_Specific(V)));
     if (!match(I->getArgOperand(0), m_c_ICmp(Pred, m_V, m_Value(RHS))))
-      return false;
+      continue;
 
     if (cmpExcludesZero(Pred, RHS) && isValidAssumeForContext(I, Q.CxtI, Q.DT))
       return true;
@@ -2207,6 +2207,22 @@ static bool isPowerOfTwoRecurrence(const PHINode *PN, bool OrZero,
   }
 }
 
+/// Return true if we can infer that \p V is known to be a power of 2 from
+/// dominating condition \p Cond (e.g., ctpop(V) == 1).
+static bool isImpliedToBeAPowerOfTwoFromCond(const Value *V, bool OrZero,
+                                             const Value *Cond) {
+  ICmpInst::Predicate Pred;
+  const APInt *RHSC;
+  if (!match(Cond, m_ICmp(Pred, m_Intrinsic<Intrinsic::ctpop>(m_Specific(V)),
+                          m_APInt(RHSC))))
+    return false;
+  // ctpop(V) u< 2
+  if (OrZero && Pred == ICmpInst::ICMP_ULT && *RHSC == 2)
+    return true;
+  // ctpop(V) == 1
+  return Pred == ICmpInst::ICMP_EQ && *RHSC == 1;
+}
+
 /// Return true if the given value is known to have exactly one
 /// bit set when defined. For vectors return true if every element is known to
 /// be a power of two when defined. Supports values with integer or pointer
@@ -2222,6 +2238,18 @@ bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
   if (OrZero && V->getType()->getScalarSizeInBits() == 1)
     return true;
 
+  // Try to infer from assumptions.
+  if (Q.AC && Q.CxtI) {
+    for (auto &AssumeVH : Q.AC->assumptionsFor(V)) {
+      if (!AssumeVH)
+        continue;
+      CallInst *I = cast<CallInst>(AssumeVH);
+      if (isImpliedToBeAPowerOfTwoFromCond(V, OrZero, I->getArgOperand(0)) &&
+          isValidAssumeForContext(I, Q.CxtI, Q.DT))
+        return true;
+    }
+  }
+
   auto *I = dyn_cast<Instruction>(V);
   if (!I)
     return false;
@@ -9903,8 +9931,9 @@ void llvm::findValuesAffectedByCondition(
     } else if (match(V, m_ICmp(Pred, m_Value(A), m_Value(B)))) {
       AddCmpOperands(A, B);
 
+      bool HasRHSC = match(B, m_ConstantInt());
       if (ICmpInst::isEquality(Pred)) {
-        if (match(B, m_ConstantInt())) {
+        if (HasRHSC) {
           Value *Y;
           // (X & C) or (X | C) or (X ^ C).
           // (X << C) or (X >>_s C) or (X >>_u C).
@@ -9918,7 +9947,7 @@ void llvm::findValuesAffectedByCondition(
           }
         }
       } else {
-        if (match(B, m_ConstantInt())) {
+        if (HasRHSC) {
           // Handle (A + C1) u< C2, which is the canonical form of
           // A > C3 && A < C4.
           if (match(A, m_AddLike(m_Value(X), m_ConstantInt())))
@@ -9950,6 +9979,10 @@ void llvm::findValuesAffectedByCondition(
             InsertAffected(X);
         }
       }
+
+      if (IsAssume && HasRHSC &&
+          match(A, m_Intrinsic<Intrinsic::ctpop>(m_Value(X))))
+        AddAffected(X);
     } else if (match(Cond, m_FCmp(Pred, m_Value(A), m_Value(B)))) {
       AddCmpOperands(A, B);
 

llvm/test/Transforms/InstCombine/cttz.ll

@@ -276,3 +276,24 @@ define i32 @cttz_of_power_of_two_wrong_constant_2(i32 %x) {
   %r = call i32 @llvm.cttz.i32(i32 %add, i1 false)
   ret i32 %r
 }
+
+define i16 @cttz_assume(i16 %x) {
+; CHECK-LABEL: @cttz_assume(
+; CHECK-NEXT:    [[ADD:%.*]] = add i16 [[X:%.*]], 1
+; CHECK-NEXT:    [[COND0:%.*]] = icmp ult i16 [[ADD]], 10
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND0]])
+; CHECK-NEXT:    [[COND1:%.*]] = icmp ne i16 [[X]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND1]])
+; CHECK-NEXT:    [[CTTZ:%.*]] = call range(i16 0, 17) i16 @llvm.cttz.i16(i16 [[X]], i1 true)
+; CHECK-NEXT:    ret i16 [[CTTZ]]
+;
+  %add = add i16 %x, 1
+  %cond0 = icmp ult i16 %add, 10
+  call void @llvm.assume(i1 %cond0)
+
+  %cond1 = icmp ne i16 %x, 0
+  call void @llvm.assume(i1 %cond1)
+
+  %cttz = call i16 @llvm.cttz.i16(i16 %x, i1 false)
+  ret i16 %cttz
+}

llvm/test/Transforms/InstCombine/icmp.ll

@@ -5326,3 +5326,42 @@ define i1 @pr94897(i32 range(i32 -2147483648, 0) %x) {
   %cmp = icmp ugt i32 %shl, -50331648
   ret i1 %cmp
 }
+
+define i1 @icmp_and_inv_pow2_ne_0(i32 %A, i32 %B) {
+; CHECK-LABEL: @icmp_and_inv_pow2_ne_0(
+; CHECK-NEXT:    [[POPCNT:%.*]] = tail call range(i32 0, 33) i32 @llvm.ctpop.i32(i32 [[A:%.*]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[POPCNT]], 1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[A]], [[B:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[TMP1]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %popcnt = tail call i32 @llvm.ctpop.i32(i32 %A)
+  %cond = icmp eq i32 %popcnt, 1
+  call void @llvm.assume(i1 %cond)
+
+  %inv = xor i32 %B, -1
+  %and = and i32 %A, %inv
+  %cmp = icmp ne i32 %and, 0
+  ret i1 %cmp
+}
+
+define i1 @icmp_and_inv_pow2_or_zero_ne_0(i32 %A, i32 %B) {
+; CHECK-LABEL: @icmp_and_inv_pow2_or_zero_ne_0(
+; CHECK-NEXT:    [[POPCNT:%.*]] = tail call range(i32 0, 33) i32 @llvm.ctpop.i32(i32 [[A:%.*]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i32 [[POPCNT]], 2
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[INV:%.*]] = xor i32 [[B:%.*]], -1
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[A]], [[INV]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[AND]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %popcnt = tail call i32 @llvm.ctpop.i32(i32 %A)
+  %cond = icmp ult i32 %popcnt, 2
+  call void @llvm.assume(i1 %cond)
+
+  %inv = xor i32 %B, -1
+  %and = and i32 %A, %inv
+  %cmp = icmp ne i32 %and, 0
+  ret i1 %cmp
+}

llvm/test/Transforms/InstCombine/rem.ll

@@ -1041,3 +1041,35 @@ define <2 x i32> @PR62401(<2 x i1> %x, <2 x i32> %y) {
   %r = urem <2 x i32> %y, %sext.i1
   ret <2 x i32> %r
 }
+
+define i16 @rem_pow2_or_zero(i16 %x, i16 %y) {
+; CHECK-LABEL: @rem_pow2_or_zero(
+; CHECK-NEXT:    [[POPCNT:%.*]] = call range(i16 1, 17) i16 @llvm.ctpop.i16(i16 [[Y:%.*]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i16 [[POPCNT]], 2
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add i16 [[Y]], -1
+; CHECK-NEXT:    [[REM:%.*]] = and i16 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i16 [[REM]]
+;
+  %popcnt = call i16 @llvm.ctpop.i16(i16 %y)
+  %cond = icmp ult i16 %popcnt, 2
+  tail call void @llvm.assume(i1 %cond)
+  %rem = urem i16 %x, %y
+  ret i16 %rem
+}
+
+define i16 @rem_pow2(i16 %x, i16 %y) {
+; CHECK-LABEL: @rem_pow2(
+; CHECK-NEXT:    [[POPCNT:%.*]] = call range(i16 1, 17) i16 @llvm.ctpop.i16(i16 [[Y:%.*]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i16 [[POPCNT]], 1
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add i16 [[Y]], -1
+; CHECK-NEXT:    [[REM:%.*]] = and i16 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i16 [[REM]]
+;
+  %popcnt = call i16 @llvm.ctpop.i16(i16 %y)
+  %cond = icmp eq i16 %popcnt, 1
+  tail call void @llvm.assume(i1 %cond)
+  %rem = urem i16 %x, %y
+  ret i16 %rem
+}