[ValueTracking] Infer is-power-of-2 from assumptions.

[dtcxzyw]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/107745.diff

3 Files Affected:

• (modified) llvm/lib/Analysis/ValueTracking.cpp (+37-2)
• (modified) llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll (+1-1)
• (modified) llvm/test/Transforms/InstCombine/rem.ll (+16)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 3a0ec99ee5ea1e..6951f65006fa52 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -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,20 @@ 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 (!isValidAssumeForContext(I, Q.CxtI, Q.DT))
+        continue;
+
+      if (isImpliedToBeAPowerOfTwoFromCond(V, OrZero, I->getArgOperand(0)))
+        return true;
+    }
+  }
+
   auto *I = dyn_cast<Instruction>(V);
   if (!I)
     return false;
@@ -9903,8 +9933,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 +9949,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 +9981,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);
 
diff --git a/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll b/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll
index 618f5d641dc1ab..ffc6aef2aafda4 100644
--- a/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll
@@ -306,7 +306,7 @@ define i32 @pow2_32_nonconst_assume(i32 %x, i32 %y) {
 
 define i32 @pow2_32_gtnonconst_assume(i32 %x, i32 %y) {
 ; CHECK-LABEL: @pow2_32_gtnonconst_assume(
-; CHECK-NEXT:    [[CTPOP:%.*]] = call range(i32 1, 33) i32 @llvm.ctpop.i32(i32 [[Y:%.*]])
+; CHECK-NEXT:    [[CTPOP:%.*]] = call range(i32 0, 33) i32 @llvm.ctpop.i32(i32 [[Y:%.*]])
 ; CHECK-NEXT:    [[YP2:%.*]] = icmp eq i32 [[CTPOP]], 1
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[YP2]])
 ; CHECK-NEXT:    [[YGT:%.*]] = icmp ugt i32 [[Y]], [[X:%.*]]
diff --git a/llvm/test/Transforms/InstCombine/rem.ll b/llvm/test/Transforms/InstCombine/rem.ll
index 05ff214f91b8ce..314e85ae1b42d8 100644
--- a/llvm/test/Transforms/InstCombine/rem.ll
+++ b/llvm/test/Transforms/InstCombine/rem.ll
@@ -1041,3 +1041,19 @@ 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(i16 %x, i16 %y) {
+; CHECK-LABEL: @rem_pow2(
+; CHECK-NEXT:    [[POPCNT:%.*]] = tail 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 = tail call i16 @llvm.ctpop.i16(i16 %y)
+  %cond = icmp sle i16 %popcnt, 1
+  tail call void @llvm.assume(i1 %cond)
+  %rem = urem i16 %x, %y
+  ret i16 %rem
+}

[nikic]

This looks fine, but we should have tests for both the pow2 and pow2-or-zero cases. Do you plan to also extend this to the dominating condition case?

[nikic]

LGTM

[goldsteinn]

LGTM