[llvm] [LoopIdiom] Support 'shift until less-than' idiom (PR #95002)
Hari Limaye via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 1 06:06:54 PDT 2024
================
@@ -1528,6 +1563,107 @@ static PHINode *getRecurrenceVar(Value *VarX, Instruction *DefX,
return nullptr;
}
+/// Return true if the idiom is detected in the loop.
+///
+/// Additionally:
+/// 1) \p CntInst is set to the instruction Counting Leading Zeros (CTLZ)
+/// or nullptr if there is no such.
+/// 2) \p CntPhi is set to the corresponding phi node
+/// or nullptr if there is no such.
+/// 3) \p InitX is set to the value whose CTLZ could be used.
+/// 4) \p DefX is set to the instruction calculating Loop exit condition.
+/// 5) \p Threshold is set to the constant involved in the unsigned less-than
+/// comparison.
+///
+/// The core idiom we are trying to detect is:
+/// \code
+/// if (x0 < 2)
+/// goto loop-exit // the predcondition of the loop
+/// cnt0 = init-val
+/// do {
+/// x = phi (x0, x.next); //PhiX
+/// cnt = phi (cnt0, cnt.next)
+///
+/// cnt.next = cnt + 1;
+/// ...
+/// x.next = x >> 1; // DefX
+/// } while (x < 4)
+/// loop-exit:
+/// \endcode
+static bool detectShiftUntilLessThanIdiom(Loop *CurLoop, const DataLayout &DL,
+ Intrinsic::ID &IntrinID,
+ Value *&InitX, Instruction *&CntInst,
+ PHINode *&CntPhi, Instruction *&DefX,
+ unsigned &Threshold) {
+ BasicBlock *LoopEntry;
+
+ DefX = nullptr;
+ CntInst = nullptr;
+ CntPhi = nullptr;
+ LoopEntry = *(CurLoop->block_begin());
+
+ // step 1: Check if the loop-back branch is in desirable form.
+ if (Value *T = matchShiftULTCondition(
+ dyn_cast<BranchInst>(LoopEntry->getTerminator()), LoopEntry,
+ Threshold))
+ DefX = dyn_cast<Instruction>(T);
+ else
+ return false;
+
+ // step 2: Check the recurrence of variable X
+ if (!DefX || !isa<PHINode>(DefX))
+ return false;
+
+ PHINode *VarPhi = cast<PHINode>(DefX);
+ int Idx = VarPhi->getBasicBlockIndex(LoopEntry);
+ if (Idx == -1)
+ return false;
+
+ DefX = dyn_cast<Instruction>(VarPhi->getIncomingValue(Idx));
+ if (!DefX || DefX->getNumOperands() == 0 || DefX->getOperand(0) != VarPhi)
+ return false;
+
+ // step 3: detect instructions corresponding to "x.next = x >> 1"
+ if (DefX->getOpcode() != Instruction::LShr)
+ return false;
+
+ IntrinID = Intrinsic::ctlz;
+ ConstantInt *Shft = dyn_cast<ConstantInt>(DefX->getOperand(1));
+ if (!Shft || !Shft->isOne())
+ return false;
+
+ InitX = VarPhi->getIncomingValueForBlock(CurLoop->getLoopPreheader());
+
+ // step 4: Find the instruction which count the CTLZ: cnt.next = cnt + 1
+ // or cnt.next = cnt + -1.
+ // TODO: We can skip the step. If loop trip count is known (CTLZ),
+ // then all uses of "cnt.next" could be optimized to the trip count
+ // plus "cnt0". Currently it is not optimized.
+ // This step could be used to detect POPCNT instruction:
+ // cnt.next = cnt + (x.next & 1)
+ for (Instruction &Inst : llvm::make_range(
+ LoopEntry->getFirstNonPHI()->getIterator(), LoopEntry->end())) {
+ if (Inst.getOpcode() != Instruction::Add)
+ continue;
+
+ ConstantInt *Inc = dyn_cast<ConstantInt>(Inst.getOperand(1));
+ if (!Inc || (!Inc->isOne() && !Inc->isMinusOne()))
----------------
hazzlim wrote:
Good point - it seems to me we might as well (as the code for handling the decrement case already exists in `transformLoopToCountable`), so I've added some tests for this case.
https://github.com/llvm/llvm-project/pull/95002
More information about the llvm-commits
mailing list