[llvm] LoopIdiomRecognize: detect and convert powi idiom (PR #72650)

[Ramkumar Ramachandra via llvm-commits]

================
@@ -925,6 +930,111 @@ bool LoopIdiomRecognize::processLoopMemSet(MemSetInst *MSI,
                                  BECount, IsNegStride, /*IsLoopMemset=*/true);
 }
 
+static CallInst *createPowiIntrinsic(IRBuilder<> &IRBuilder, Value *Base,
+                                     Value *Exp, const DebugLoc &DL) {
+  Value *Ops[] = {Base, Exp};
+  Type *Tys[] = {Base->getType(), Exp->getType()};
+
+  Module *M = IRBuilder.GetInsertBlock()->getParent()->getParent();
+  Function *Func = Intrinsic::getDeclaration(M, Intrinsic::powi, Tys);
+  CallInst *CI = IRBuilder.CreateCall(Func, Ops);
+  CI->setDebugLoc(DL);
+  return CI;
+}
+
+// Checks that the Phi is an fmul fast with a loop-invariant operand, and
+// returns the the fmul instruction.
+static Instruction *detectPowiIdiom(PHINode *Phi, BasicBlock *PH,
+                                    BasicBlock *Latch, Loop *CurLoop) {
+  LLVM_DEBUG(dbgs() << DEBUG_TYPE " Performing powi idiom detection\n");
+
+  // The phi must have two incoming values (one from the preheader, and another
+  // from the latch), it must have one use (which we will subsequently check is
+  // an fmul fast instruction), and it must be a floating-point type.
+  if (Phi->getNumIncomingValues() != 2 || !Phi->hasOneUse() ||
+      Phi->getBasicBlockIndex(PH) < 0 || Phi->getBasicBlockIndex(Latch) < 0 ||
+      !Phi->getType()->isFloatingPointTy()) {
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE " Unable to operate on this PHI node\n");
+    return nullptr;
+  }
+
+  // Further, check that the incoming value from the preheader is 1.0.
+  auto *ConstFP = dyn_cast<ConstantFP>(Phi->getIncomingValueForBlock(PH));
+  if (!ConstFP || !ConstFP->isExactlyValue(1.0)) {
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE " Initial value comparison failed\n");
+    return nullptr;
+  }
+
+  auto *I = cast<Instruction>(Phi->use_begin()->getUser());
+  Value *Op1, *Op2;
+  if (!match(I, m_FMul(m_Value(Op1), m_Value(Op2))) || !I->isFast()) {
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE " fmul-fast test failed\n");
+    return nullptr;
+  }
+  Value *Base = Op1 == Phi ? Op2 : Op1;
+  if (CurLoop->isLoopInvariant(Base))
+    return I;
+  else
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE " Base is not loop-invariant\n");
+  return nullptr;
+}
+
+/// Detect the powi idiom, and convert it to an intrinsic.
+bool LoopIdiomRecognize::processLoopPowi(BasicBlock *BB, const SCEV *BECount) {
+  // We only process loops where the IV is, at most, a C int type, since this is
+  // used for the exponent of the powi.
+  PHINode *IV = CurLoop->getInductionVariable(*SE);
+  if (!IV || IV->getType()->getScalarSizeInBits() > 32)
+    return false;
+
+  // If the loop doesn't have a valid preheader and latch, give up now.
+  BasicBlock *PH = CurLoop->getLoopPreheader();
+  BasicBlock *Latch = CurLoop->getLoopLatch();
+  if (!PH || !Latch)
+    return false;
+
+  // Collect all phis that are not the induction phi.
+  SmallVector<PHINode *, 1> Phis;
+  for (Instruction &I : *BB) {
+    if (auto *Phi = dyn_cast<PHINode>(&I))
+      if (Phi != IV)
+        Phis.push_back(Phi);
+  }
+
+  // Find the Phi corresponding to the powi idiom, amongst all phis except the
+  // induction phi.
+  for (PHINode *Phi : Phis) {
+    if (Instruction *FMul = detectPowiIdiom(Phi, PH, Latch, CurLoop)) {
+      // Find the trip count, and expand the SCEV to find the exponent of the
+      // powi.
+      IRBuilder<> Builder(PH->getTerminator());
+      SCEVExpander Expander(*SE, *DL, "loop-idiom");
+      SCEVExpanderCleaner ExpCleaner(Expander);
+      Type *ExpTy = IV->getType();
+      const SCEV *TripCount =
+          SE->getTripCountFromExitCount(BECount, ExpTy, CurLoop);
+      if (!Expander.isSafeToExpand(TripCount)) {
+        LLVM_DEBUG(dbgs() << DEBUG_TYPE " Trip count not safe to expand\n");
+        return false;
+      }
+      Value *Exp =
+          Expander.expandCodeFor(TripCount, ExpTy, PH->getTerminator());
+
+      // Insert the powi intrinsic, and replace its uses outside the block.
+      const DebugLoc &Loc = FMul->getDebugLoc();
+      Value *Base = isa<PHINode>(FMul->getOperand(1)) ? FMul->getOperand(2)
+                                                      : FMul->getOperand(1);
----------------
artagnon wrote:

Oops! Thanks for catching this bug!

https://github.com/llvm/llvm-project/pull/72650