[ValueTracking] Infer is-power-of-2 from dominating conditions

llvm/include/llvm/Analysis/ValueTracking.h

@@ -119,6 +119,9 @@ bool isKnownToBeAPowerOfTwo(const Value *V, const DataLayout &DL,
                             const DominatorTree *DT = nullptr,
                             bool UseInstrInfo = true);
 
+bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
+                            const SimplifyQuery &Q);
+
 bool isOnlyUsedInZeroComparison(const Instruction *CxtI);
 
 bool isOnlyUsedInZeroEqualityComparison(const Instruction *CxtI);

llvm/include/llvm/Transforms/InstCombine/InstCombiner.h

@@ -450,7 +450,8 @@ class LLVM_LIBRARY_VISIBILITY InstCombiner {
   bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero = false,
                               unsigned Depth = 0,
                               const Instruction *CxtI = nullptr) {
-    return llvm::isKnownToBeAPowerOfTwo(V, DL, OrZero, Depth, &AC, CxtI, &DT);
+    return llvm::isKnownToBeAPowerOfTwo(V, OrZero, Depth,
+                                        SQ.getWithInstruction(CxtI));
   }
 
   bool MaskedValueIsZero(const Value *V, const APInt &Mask, unsigned Depth = 0,

llvm/lib/Analysis/ValueTracking.cpp

@@ -265,9 +265,6 @@ bool llvm::isOnlyUsedInZeroEqualityComparison(const Instruction *I) {
   });
 }
 
-static bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
-                                   const SimplifyQuery &Q);
-
 bool llvm::isKnownToBeAPowerOfTwo(const Value *V, const DataLayout &DL,
                                   bool OrZero, unsigned Depth,
                                   AssumptionCache *AC, const Instruction *CxtI,
@@ -2210,12 +2207,15 @@ 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) {
+                                             const Value *Cond,
+                                             bool CondIsTrue) {
   ICmpInst::Predicate Pred;
   const APInt *RHSC;
   if (!match(Cond, m_ICmp(Pred, m_Intrinsic<Intrinsic::ctpop>(m_Specific(V)),
                           m_APInt(RHSC))))
     return false;
+  if (!CondIsTrue)
+    Pred = ICmpInst::getInversePredicate(Pred);
   // ctpop(V) u< 2
   if (OrZero && Pred == ICmpInst::ICMP_ULT && *RHSC == 2)
     return true;
@@ -2227,8 +2227,8 @@ static bool isImpliedToBeAPowerOfTwoFromCond(const Value *V, bool OrZero,
 /// 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
 /// types and vectors of integers.
-bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
-                            const SimplifyQuery &Q) {
+bool llvm::isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
+                                  const SimplifyQuery &Q) {
   assert(Depth <= MaxAnalysisRecursionDepth && "Limit Search Depth");
 
   if (isa<Constant>(V))
@@ -2244,12 +2244,32 @@ bool isKnownToBeAPowerOfTwo(const Value *V, bool OrZero, unsigned Depth,
       if (!AssumeVH)
         continue;
       CallInst *I = cast<CallInst>(AssumeVH);
-      if (isImpliedToBeAPowerOfTwoFromCond(V, OrZero, I->getArgOperand(0)) &&
+      if (isImpliedToBeAPowerOfTwoFromCond(V, OrZero, I->getArgOperand(0),
+                                           /*CondIsTrue=*/true) &&
           isValidAssumeForContext(I, Q.CxtI, Q.DT))
         return true;
     }
   }
 
+  // Handle dominating conditions.
+  if (Q.DC && Q.CxtI && Q.DT) {
+    for (BranchInst *BI : Q.DC->conditionsFor(V)) {
+      Value *Cond = BI->getCondition();
+
+      BasicBlockEdge Edge0(BI->getParent(), BI->getSuccessor(0));
+      if (isImpliedToBeAPowerOfTwoFromCond(V, OrZero, Cond,
+                                           /*CondIsTrue=*/true) &&
+          Q.DT->dominates(Edge0, Q.CxtI->getParent()))
+        return true;
+
+      BasicBlockEdge Edge1(BI->getParent(), BI->getSuccessor(1));
+      if (isImpliedToBeAPowerOfTwoFromCond(V, OrZero, Cond,
+                                           /*CondIsTrue=*/false) &&
+          Q.DT->dominates(Edge1, Q.CxtI->getParent()))
+        return true;
+    }
+  }
+
   auto *I = dyn_cast<Instruction>(V);
   if (!I)
     return false;
@@ -9980,8 +10000,7 @@ void llvm::findValuesAffectedByCondition(
         }
       }
 
-      if (IsAssume && HasRHSC &&
-          match(A, m_Intrinsic<Intrinsic::ctpop>(m_Value(X))))
+      if (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/rem.ll

@@ -1073,3 +1073,106 @@ define i16 @rem_pow2(i16 %x, i16 %y) {
   %rem = urem i16 %x, %y
   ret i16 %rem
 }
+
+define i64 @rem_pow2_domcond(i64 %a, i64 %b) {
+; CHECK-LABEL: @rem_pow2_domcond(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CPOP:%.*]] = call range(i64 0, 65) i64 @llvm.ctpop.i64(i64 [[B:%.*]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[CPOP]], 1
+; CHECK-NEXT:    br i1 [[COND]], label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[B]], -1
+; CHECK-NEXT:    [[REM:%.*]] = and i64 [[A:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret i64 [[REM]]
+; CHECK:       bb2:
+; CHECK-NEXT:    ret i64 0
+;
+start:
+  %cpop = call i64 @llvm.ctpop.i64(i64 %b)
+  %cond = icmp eq i64 %cpop, 1
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  %rem = urem i64 %a, %b
+  ret i64 %rem
+
+bb2:
+  ret i64 0
+}
+
+define i64 @rem_pow2_domcond_in_else(i64 %a, i64 %b) {
+; CHECK-LABEL: @rem_pow2_domcond_in_else(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CPOP:%.*]] = call range(i64 0, 65) i64 @llvm.ctpop.i64(i64 [[B:%.*]])
+; CHECK-NEXT:    [[COND_NOT:%.*]] = icmp eq i64 [[CPOP]], 1
+; CHECK-NEXT:    br i1 [[COND_NOT]], label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[B]], -1
+; CHECK-NEXT:    [[REM:%.*]] = and i64 [[A:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret i64 [[REM]]
+; CHECK:       bb2:
+; CHECK-NEXT:    ret i64 0
+;
+start:
+  %cpop = call i64 @llvm.ctpop.i64(i64 %b)
+  %cond = icmp ne i64 %cpop, 1
+  br i1 %cond, label %bb2, label %bb1
+
+bb1:
+  %rem = urem i64 %a, %b
+  ret i64 %rem
+
+bb2:
+  ret i64 0
+}
+
+define i64 @rem_pow2_or_zero_domcond(i64 %a, i64 %b) {
+; CHECK-LABEL: @rem_pow2_or_zero_domcond(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CPOP:%.*]] = call range(i64 0, 65) i64 @llvm.ctpop.i64(i64 [[B:%.*]])
+; CHECK-NEXT:    [[COND:%.*]] = icmp ult i64 [[CPOP]], 2
+; CHECK-NEXT:    br i1 [[COND]], label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[B]], -1
+; CHECK-NEXT:    [[REM:%.*]] = and i64 [[A:%.*]], [[TMP0]]
+; CHECK-NEXT:    ret i64 [[REM]]
+; CHECK:       bb2:
+; CHECK-NEXT:    ret i64 0
+;
+start:
+  %cpop = call i64 @llvm.ctpop.i64(i64 %b)
+  %cond = icmp ult i64 %cpop, 2
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  %rem = urem i64 %a, %b
+  ret i64 %rem
+
+bb2:
+  ret i64 0
+}
+
+define i64 @rem_pow2_non_domcond(i64 %a, i64 %b) {
+; CHECK-LABEL: @rem_pow2_non_domcond(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[CPOP:%.*]] = call range(i64 0, 65) i64 @llvm.ctpop.i64(i64 [[B:%.*]])
+; CHECK-NEXT:    [[COND_NOT:%.*]] = icmp eq i64 [[CPOP]], 1
+; CHECK-NEXT:    br i1 [[COND_NOT]], label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[REM:%.*]] = urem i64 [[A:%.*]], [[B]]
+; CHECK-NEXT:    ret i64 [[REM]]
+; CHECK:       bb2:
+; CHECK-NEXT:    br label [[BB1]]
+;
+start:
+  %cpop = call i64 @llvm.ctpop.i64(i64 %b)
+  %cond = icmp ne i64 %cpop, 1
+  br i1 %cond, label %bb2, label %bb1
+
+bb1:
+  %rem = urem i64 %a, %b
+  ret i64 %rem
+
+bb2:
+  br label %bb1
+}