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

Yingwei Zheng via llvm-commits llvm-commits at lists.llvm.org
Sun Sep 8 18:49:57 PDT 2024


https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/107745

>From 638fdb7b9e1bee02cfa4ab213f7080c7093d8c42 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 9 Sep 2024 09:35:42 +0800
Subject: [PATCH 1/2] [ValueTracking] Add pre-commit tests. NFC.

---
 llvm/test/Transforms/InstCombine/cttz.ll | 21 +++++++++++++
 llvm/test/Transforms/InstCombine/icmp.ll | 40 ++++++++++++++++++++++++
 llvm/test/Transforms/InstCombine/rem.ll  | 30 ++++++++++++++++++
 3 files changed, 91 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/cttz.ll b/llvm/test/Transforms/InstCombine/cttz.ll
index e106faf9cb38f0..ad7b864cb7c4a2 100644
--- a/llvm/test/Transforms/InstCombine/cttz.ll
+++ 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 false)
+; 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
+}
diff --git a/llvm/test/Transforms/InstCombine/icmp.ll b/llvm/test/Transforms/InstCombine/icmp.ll
index e492055fea8b8d..392312261985f6 100644
--- a/llvm/test/Transforms/InstCombine/icmp.ll
+++ b/llvm/test/Transforms/InstCombine/icmp.ll
@@ -5326,3 +5326,43 @@ 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:    [[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 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
+}
diff --git a/llvm/test/Transforms/InstCombine/rem.ll b/llvm/test/Transforms/InstCombine/rem.ll
index 05ff214f91b8ce..f76b4dd85eacd0 100644
--- a/llvm/test/Transforms/InstCombine/rem.ll
+++ b/llvm/test/Transforms/InstCombine/rem.ll
@@ -1041,3 +1041,33 @@ 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:    [[REM:%.*]] = urem i16 [[X:%.*]], [[Y]]
+; 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:    [[REM:%.*]] = urem i16 [[X:%.*]], [[Y]]
+; 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
+}

>From dc99b690cdded18108409b67045b1e8f34f0122c Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 9 Sep 2024 09:49:29 +0800
Subject: [PATCH 2/2] [ValueTracking] Infer is-power-of-2 from assumptions.

---
 llvm/lib/Analysis/ValueTracking.cpp      | 39 ++++++++++++++++++++++--
 llvm/test/Transforms/InstCombine/cttz.ll |  2 +-
 llvm/test/Transforms/InstCombine/icmp.ll |  5 ++-
 llvm/test/Transforms/InstCombine/rem.ll  |  6 ++--
 4 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 3a0ec99ee5ea1e..e9fc9bbe8c4936 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/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);
 
diff --git a/llvm/test/Transforms/InstCombine/cttz.ll b/llvm/test/Transforms/InstCombine/cttz.ll
index ad7b864cb7c4a2..cb0bc59ae79958 100644
--- a/llvm/test/Transforms/InstCombine/cttz.ll
+++ b/llvm/test/Transforms/InstCombine/cttz.ll
@@ -284,7 +284,7 @@ define i16 @cttz_assume(i16 %x) {
 ; 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 false)
+; 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
diff --git a/llvm/test/Transforms/InstCombine/icmp.ll b/llvm/test/Transforms/InstCombine/icmp.ll
index 392312261985f6..ecf21b8a42cf50 100644
--- a/llvm/test/Transforms/InstCombine/icmp.ll
+++ b/llvm/test/Transforms/InstCombine/icmp.ll
@@ -5332,9 +5332,8 @@ define i1 @icmp_and_inv_pow2_ne_0(i32 %A, i32 %B) {
 ; 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:    [[INV:%.*]] = xor i32 [[B:%.*]], -1
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[A]], [[INV]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[AND]], 0
+; 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)
diff --git a/llvm/test/Transforms/InstCombine/rem.ll b/llvm/test/Transforms/InstCombine/rem.ll
index f76b4dd85eacd0..9d2a947d6b45c9 100644
--- a/llvm/test/Transforms/InstCombine/rem.ll
+++ b/llvm/test/Transforms/InstCombine/rem.ll
@@ -1047,7 +1047,8 @@ define i16 @rem_pow2_or_zero(i16 %x, i16 %y) {
 ; 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:    [[REM:%.*]] = urem i16 [[X:%.*]], [[Y]]
+; 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)
@@ -1062,7 +1063,8 @@ define i16 @rem_pow2(i16 %x, i16 %y) {
 ; 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:    [[REM:%.*]] = urem i16 [[X:%.*]], [[Y]]
+; 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)



More information about the llvm-commits mailing list