From cc845e9de8c87c74d494f4e90e8fcf4fca264989 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n@gmail.com>
Date: Wed, 11 Jan 2023 10:18:46 +0100
Subject: [PATCH] [InstCombine] Handle assume(X & Pow2 != 0) in
 computeKnownBits()

If we know that X & Pow2 != 0, then the bit at that position is
known one.

Differential Revision: https://reviews.llvm.org/D140851
---
 llvm/lib/Analysis/AssumptionCache.cpp            | 25 +++++++++++++++---------
 llvm/lib/Analysis/ValueTracking.cpp              |  8 ++++++++
 llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll |  9 +++------
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Analysis/AssumptionCache.cpp b/llvm/lib/Analysis/AssumptionCache.cpp
index e7e476d..11796ef 100644
--- a/llvm/lib/Analysis/AssumptionCache.cpp
+++ b/llvm/lib/Analysis/AssumptionCache.cpp
@@ -105,7 +105,7 @@ findAffectedValues(CallBase *CI, TargetTransformInfo *TTI,
         if (match(V, m_BitwiseLogic(m_Value(A), m_Value(B)))) {
           AddAffected(A);
           AddAffected(B);
-        // (A << C) or (A >>_s C) or (A >>_u C) where C is some constant.
+          // (A << C) or (A >>_s C) or (A >>_u C) where C is some constant.
         } else if (match(V, m_Shift(m_Value(A), m_ConstantInt()))) {
           AddAffected(A);
         }
@@ -113,15 +113,22 @@ findAffectedValues(CallBase *CI, TargetTransformInfo *TTI,
 
       AddAffectedFromEq(A);
       AddAffectedFromEq(B);
+    } else if (Pred == ICmpInst::ICMP_NE) {
+      Value *X, *Y;
+      // Handle (a & b != 0). If a/b is a power of 2 we can use this
+      // information.
+      if (match(A, m_And(m_Value(X), m_Value(Y))) && match(B, m_Zero())) {
+        AddAffected(X);
+        AddAffected(Y);
+      }
+    } else if (Pred == ICmpInst::ICMP_ULT) {
+      Value *X;
+      // Handle (A + C1) u< C2, which is the canonical form of A > C3 && A < C4,
+      // and recognized by LVI at least.
+      if (match(A, m_Add(m_Value(X), m_ConstantInt())) &&
+          match(B, m_ConstantInt()))
+        AddAffected(X);
     }
-
-    Value *X;
-    // Handle (A + C1) u< C2, which is the canonical form of A > C3 && A < C4,
-    // and recognized by LVI at least.
-    if (Pred == ICmpInst::ICMP_ULT &&
-        match(A, m_Add(m_Value(X), m_ConstantInt())) &&
-        match(B, m_ConstantInt()))
-      AddAffected(X);
   }
 
   if (TTI) {
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 7c37b89..e49183d 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -940,6 +940,14 @@ static void computeKnownBitsFromAssume(const Value *V, KnownBits &Known,
           Known.Zero.setHighBits(RHSKnown.countMinLeadingZeros());
       }
       break;
+    case ICmpInst::ICMP_NE: {
+      // assume (v & b != 0) where b is a power of 2
+      const APInt *BPow2;
+      if (match(Cmp, m_ICmp(Pred, m_c_And(m_V, m_Power2(BPow2)), m_Zero())) &&
+          isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+        Known.One |= BPow2->zextOrTrunc(BitWidth);
+      }
+    } break;
     }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll b/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll
index 47c1985..224ea3c 100644
--- a/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-ne-pow2.ll
@@ -9,8 +9,7 @@ define i32 @pow2_32_assume(i32 %x) {
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[X:%.*]], 4
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[AND]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[AND2:%.*]] = and i32 [[X]], 4
-; CHECK-NEXT:    ret i32 [[AND2]]
+; CHECK-NEXT:    ret i32 4
 ;
   %and = and i32 %x, 4
   %cmp = icmp ne i32 %and, 0
@@ -39,8 +38,7 @@ define i64 @pow2_64_assume(i64 %x) {
 ; CHECK-NEXT:    [[AND:%.*]] = and i64 [[X:%.*]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[AND]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[OR:%.*]] = or i64 [[X]], 1
-; CHECK-NEXT:    ret i64 [[OR]]
+; CHECK-NEXT:    ret i64 [[X]]
 ;
   %and = and i64 %x, 1
   %cmp = icmp ne i64 %and, 0
@@ -69,8 +67,7 @@ define i16 @pow2_16_assume(i16 %x) {
 ; CHECK-NEXT:    [[AND:%.*]] = and i16 [[X:%.*]], 16384
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i16 [[AND]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    [[AND2:%.*]] = and i16 [[X]], 16384
-; CHECK-NEXT:    ret i16 [[AND2]]
+; CHECK-NEXT:    ret i16 16384
 ;
   %and = and i16 %x, 16384
   %cmp = icmp eq i16 %and, 16384
-- 
2.7.4