Push isCtpopFast check down, merge with adjustPower2Test

Author: s-barannikov

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -474,7 +474,7 @@ class CodeGenPrepare {
   bool optimizeURem(Instruction *Rem);
   bool combineToUSubWithOverflow(CmpInst *Cmp, ModifyDT &ModifiedDT);
   bool combineToUAddWithOverflow(CmpInst *Cmp, ModifyDT &ModifiedDT);
-  bool unfoldPow2Test(CmpInst *Cmp);
+  bool unfoldPowerOf2Test(CmpInst *Cmp);
   void verifyBFIUpdates(Function &F);
   bool _run(Function &F);
 };
@@ -1764,7 +1764,9 @@ bool CodeGenPrepare::combineToUSubWithOverflow(CmpInst *Cmp,
 }
 
 // Decanonicalizes icmp+ctpop power-of-two test if ctpop is slow.
-bool CodeGenPrepare::unfoldPow2Test(CmpInst *Cmp) {
+// The same transformation exists in DAG combiner, but we repeat it here because
+// DAG builder can break the pattern by moving icmp into a successor block.
+bool CodeGenPrepare::unfoldPowerOf2Test(CmpInst *Cmp) {
   CmpPredicate Pred;
   Value *X;
   const APInt *C;
@@ -1774,48 +1776,60 @@ bool CodeGenPrepare::unfoldPow2Test(CmpInst *Cmp) {
                          m_APIntAllowPoison(C))))
     return false;
 
-  // This transformation increases the number of instructions, don't do it if
-  // ctpop is fast.
-  Type *OpTy = X->getType();
-  if (TLI->isCtpopFast(TLI->getValueType(*DL, OpTy)))
+  // We're only interested in "is power of 2 [or zero]" patterns.
+  bool IsStrictlyPowerOf2Test = ICmpInst::isEquality(Pred) && *C == 1;
+  bool IsPowerOf2OrZeroTest = (Pred == CmpInst::ICMP_ULT && *C == 2) ||
+                              (Pred == CmpInst::ICMP_UGT && *C == 1);
+  if (!IsStrictlyPowerOf2Test && !IsPowerOf2OrZeroTest)
     return false;
 
-  // ctpop(x) u< 2 -> (x & (x - 1)) == 0
-  // ctpop(x) u> 1 -> (x & (x - 1)) != 0
-  // Also handles ctpop(x) == 1 and ctpop(x) != 1 if ctpop(x) is known non-zero.
-  if ((Pred == CmpInst::ICMP_ULT && *C == 2) ||
-      (Pred == CmpInst::ICMP_UGT && *C == 1) ||
-      (ICmpInst::isEquality(Pred) && *C == 1 &&
-       isKnownNonZero(Cmp->getOperand(0), *DL))) {
+  // Some targets have better codegen for `ctpop(x) u</u>= 2/1`than for
+  // `ctpop(x) ==/!= 1`. If ctpop is fast, only try changing the comparison,
+  // and otherwise expand ctpop into a few simple instructions.
+  Type *OpTy = X->getType();
+  if (TLI->isCtpopFast(TLI->getValueType(*DL, OpTy))) {
+    // Look for `ctpop(x) ==/!= 1`, where `ctpop(x)` is known to be non-zero.
+    if (!IsStrictlyPowerOf2Test || !isKnownNonZero(Cmp->getOperand(0), *DL))
+      return false;
+
+    // ctpop(x) == 1 -> ctpop(x) u< 2
+    // ctpop(x) != 1 -> ctpop(x) u> 1
+    if (Pred == ICmpInst::ICMP_EQ) {
+      Cmp->setOperand(1, ConstantInt::get(OpTy, 2));
+      Cmp->setPredicate(ICmpInst::ICMP_ULT);
+    } else {
+      Cmp->setPredicate(ICmpInst::ICMP_UGT);
+    }
+    return true;
+  }
+
+  Value *NewCmp;
+  if (IsPowerOf2OrZeroTest ||
+      (IsStrictlyPowerOf2Test && isKnownNonZero(Cmp->getOperand(0), *DL))) {
+    // ctpop(x) u< 2 -> (x & (x - 1)) == 0
+    // ctpop(x) u> 1 -> (x & (x - 1)) != 0
     IRBuilder<> Builder(Cmp);
     Value *Sub = Builder.CreateAdd(X, Constant::getAllOnesValue(OpTy));
     Value *And = Builder.CreateAnd(X, Sub);
     CmpInst::Predicate NewPred =
         (Pred == CmpInst::ICMP_ULT || Pred == CmpInst::ICMP_EQ)
             ? CmpInst::ICMP_EQ
             : CmpInst::ICMP_NE;
-    Value *NewCmp =
-        Builder.CreateICmp(NewPred, And, ConstantInt::getNullValue(OpTy));
-    Cmp->replaceAllUsesWith(NewCmp);
-    RecursivelyDeleteTriviallyDeadInstructions(Cmp);
-    return true;
-  }
-
-  // ctpop(x) == 1 -> (x ^ (x - 1)) u> (x - 1)
-  // ctpop(x) != 1 -> (x ^ (x - 1)) u<= (x - 1)
-  if (ICmpInst::isEquality(Pred) && *C == 1) {
+    NewCmp = Builder.CreateICmp(NewPred, And, ConstantInt::getNullValue(OpTy));
+  } else {
+    // ctpop(x) == 1 -> (x ^ (x - 1)) u> (x - 1)
+    // ctpop(x) != 1 -> (x ^ (x - 1)) u<= (x - 1)
     IRBuilder<> Builder(Cmp);
     Value *Sub = Builder.CreateAdd(X, Constant::getAllOnesValue(OpTy));
     Value *Xor = Builder.CreateXor(X, Sub);
     CmpInst::Predicate NewPred =
         Pred == CmpInst::ICMP_EQ ? CmpInst::ICMP_UGT : CmpInst::ICMP_ULE;
-    Value *NewCmp = Builder.CreateICmp(NewPred, Xor, Sub);
-    Cmp->replaceAllUsesWith(NewCmp);
-    RecursivelyDeleteTriviallyDeadInstructions(Cmp);
-    return true;
+    NewCmp = Builder.CreateICmp(NewPred, Xor, Sub);
   }
 
-  return false;
+  Cmp->replaceAllUsesWith(NewCmp);
+  RecursivelyDeleteTriviallyDeadInstructions(Cmp);
+  return true;
 }
 
 /// Sink the given CmpInst into user blocks to reduce the number of virtual
@@ -2204,31 +2218,6 @@ bool CodeGenPrepare::optimizeURem(Instruction *Rem) {
   return false;
 }
 
-/// Some targets have better codegen for `ctpop(X) u< 2` than `ctpop(X) == 1`.
-/// This function converts `ctpop(X) ==/!= 1` into `ctpop(X) u</u> 2/1` if the
-/// result cannot be zero.
-static bool adjustIsPower2Test(CmpInst *Cmp, const TargetLowering &TLI,
-                               const TargetTransformInfo &TTI,
-                               const DataLayout &DL) {
-  CmpPredicate Pred;
-  if (!match(Cmp, m_ICmp(Pred, m_Intrinsic<Intrinsic::ctpop>(), m_One())))
-    return false;
-  if (!ICmpInst::isEquality(Pred))
-    return false;
-  auto *II = cast<IntrinsicInst>(Cmp->getOperand(0));
-
-  if (isKnownNonZero(II, DL)) {
-    if (Pred == ICmpInst::ICMP_EQ) {
-      Cmp->setOperand(1, ConstantInt::get(II->getType(), 2));
-      Cmp->setPredicate(ICmpInst::ICMP_ULT);
-    } else {
-      Cmp->setPredicate(ICmpInst::ICMP_UGT);
-    }
-    return true;
-  }
-  return false;
-}
-
 bool CodeGenPrepare::optimizeCmp(CmpInst *Cmp, ModifyDT &ModifiedDT) {
   if (sinkCmpExpression(Cmp, *TLI))
     return true;
@@ -2239,7 +2228,7 @@ bool CodeGenPrepare::optimizeCmp(CmpInst *Cmp, ModifyDT &ModifiedDT) {
   if (combineToUSubWithOverflow(Cmp, ModifiedDT))
     return true;
 
-  if (unfoldPow2Test(Cmp))
+  if (unfoldPowerOf2Test(Cmp))
     return true;
 
   if (foldICmpWithDominatingICmp(Cmp, *TLI))
@@ -2251,9 +2240,6 @@ bool CodeGenPrepare::optimizeCmp(CmpInst *Cmp, ModifyDT &ModifiedDT) {
   if (foldFCmpToFPClassTest(Cmp, *TLI, *DL))
     return true;
 
-  if (adjustIsPower2Test(Cmp, *TLI, *TTI, *DL))
-    return true;
-
   return false;
 }