[AIE] Add inner-loop versioning to pipeline small-trip-count loops by mludevid · Pull Request #1078 · Xilinx/llvm-aie

mludevid · 2026-07-10T15:18:11Z

This PR adds inner-loop versioning for AIE targets so the post-pipeliner can software-pipeline inner loops even when their statically known minimum trip count is too small. A loop marked with the aie-loop-versioning hint is split into two runtime-guarded copies: a pipelined fast copy and a verbatim, un-pipelined slow copy. A trip-count check chooses between them at runtime, so small trip counts stay correct while large ones get the pipelined schedule.

andcarminati · 2026-07-13T10:23:08Z

Could we merge the first commit in advance?

F-Stuckmann · 2026-07-13T11:58:30Z

+// Reaching definition of Reg seen from Use, scanning backwards into
+// predecessors -- the threshold pseudo may sit in a dominating block.
+static MachineInstr *
+findReachingDef(MachineBasicBlock &BB, MachineBasicBlock::reverse_iterator Use,


Do we have to reimplement reachingDef? I have a suspicion we may not need to reimplement it

I'm not sure what we could reuse for this. If you have a specific suggestion please let me know. I'm not aware of any.

Let's say that this is a simplified version for when you know that you have a single definition that dominates uses. General reaching defs for non-ssa code require a bit more expensive fixed-point dataflow analysis that could return you a set of definitions. Maybe findUniqueReachingDef?

The fact that we cannot hoist the intrinsic here (properties of the instruction), your search will stop in the same block of the instruction that uses the register discovered by findGuardConditionReg. Basically, def and use will be in the same block, could we aggressively simplify? I feel we have a complex code to deal with a simple problem.

andcarminati · 2026-07-13T12:13:55Z

+  bool runOnFunction(Function &F) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    // We deliberately do not force LoopSimplify/LCSSA: that canonicalization


Maybe this could be moved to an earlier position, as this is part of the contract of the pass.

Where to? The file header? Personally I think it makes sense to keep it here.

The comments only, to help to describe the pass pre and post conditions. Please no header files.

F-Stuckmann · 2026-07-21T11:43:19Z

+  return Intrinsic::not_intrinsic;
+}
+
+void dropLoopMetadata(Loop &L, ArrayRef<StringRef> Keys) {


CHECK: we are dropping the Key metadata, correct? You may want to rename the argument to clarify the intention.

F-Stuckmann · 2026-07-21T11:48:16Z

 }

+bool AIEBaseInstrInfo::isLoopVersionThresholdDef(const MachineInstr &MI) const {
+  auto Opcode = getLoopVersionThresholdOpcode();


nit: you can cost the opcode

andcarminati · 2026-07-21T16:08:03Z

  if (!isInt<10>(MinTripCount))
-    report_fatal_error(
-        "AIE loop versioning: threshold does not fit the scalar-move immediate");
+    report_fatal_error("AIE loop versioning: threshold does not fit the "


Just a matter of taste ;-). This fatal is in TII, but the message resembles something emitted by the loop versioning pass.

andcarminati · 2026-07-21T16:24:31Z

+  // A fall-through preheader may be empty, so pick a DebugLoc defensively.
+  const DebugLoc GuardLoc =
+      Preheader->empty() ? DebugLoc() : Preheader->front().getDebugLoc();
+  auto bailWithMsg = [&](StringRef Msg) {


nit BailWithMsg

andcarminati · 2026-07-21T16:33:39Z

+// Reaching definition of Reg seen from Use, scanning backwards into
+// predecessors -- the threshold pseudo may sit in a dominating block.
+static MachineInstr *
+findReachingDef(MachineBasicBlock &BB, MachineBasicBlock::reverse_iterator Use,


This could live in a utility file.

andcarminati · 2026-07-22T07:15:06Z

+    for (const MachineOperand &MO : MI.operands()) {
+      if (MO.isMBB())
+        TargetsBlock = true;
+      else if (MO.isReg() && MO.isUse() && MO.getReg().isPhysical())


Why checking if it is physical?

andcarminati · 2026-07-22T07:28:04Z

+      else if (MO.isReg() && MO.isUse() && MO.getReg().isPhysical())
+        RegUse = MO.getReg();
+    }
+    if (TargetsBlock && RegUse.isValid())


I would restructure this function in such a way we stop the search if we hit this if. Or simply a return here. Alternatively, we are looking to the second last instruction (before the barrier), why do we need a search? Also, If we keep the search for simplicity, why not a reverse traversal of the MBB?

andcarminati · 2026-07-22T07:29:49Z

+  return nullptr;
+}
+
+MachineBasicBlock *PostPipeliner::findVersionGuard() const {


nit: AIELoopUtils with a nice comment?

andcarminati · 2026-07-22T08:52:51Z

+;
+; (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates
+;
+; End-to-end: a versioned loop that the post-pipeliner can schedule gets its


We may use two sets of folders: loop versioning exclusive folder (IR tests - we need to create) and end-to-end (this folder exists).

andcarminati · 2026-07-22T08:57:55Z

+  addExitPHIs(HighLoop, ExitBlock, VMap, DefsUsedOutside);
+  formDedicatedExitBlocks(HighLoop, &DT, &LI, nullptr, /*PreserveLCSSA=*/true);
+  formDedicatedExitBlocks(&L, &DT, &LI, nullptr, /*PreserveLCSSA=*/true);
+


Could we extract all the code below as a member function related to metadata handling, like updateLoopsMetadata? You can use more descriptive methods for the previous steps also. This is a suggestion to improve readability - if you agree, of course.

andcarminati · 2026-07-22T09:04:00Z

+  // not re-version either. Mark the high-trip-count copy as versioned (carrying
+  // the request's value) so the postpipeliner recognizes it as the pipelined
+  // copy; the low copy stays the verbatim un-pipelined fallback.
+  int64_t HintValue = *AIELoopUtils::getLoopHintInt(


andcarminati · 2026-07-22T09:04:59Z

+}
+
+bool AIELoopVersioner::isProfitable() const {
+  // The runtime guard compares the trip count in i32, and the pipelined


Could we have a test rejected by this function?

andcarminati · 2026-07-22T09:12:03Z

+  Value *Threshold =
+      Builder.CreateIntrinsic(ThresholdIID, {}, {Builder.getInt32(-1)},
+                              /*FMFSource=*/nullptr, "lver.threshold");
+  if (TripCount->getType() != I32Ty)


We guarantee to not be greater than 32, so we may have 16bit or 8bit, right? Could we have a test?

andcarminati · 2026-07-22T09:15:45Z

+  Value *TakeLow = Builder.CreateICmpULT(TripCount, Threshold, "lver.low");
+  GuardBB->setName(Header->getName() + ".lver.guard");
+
+  SmallVector<Instruction *, 8> DefsUsedOutside = findDefsUsedOutsideOfLoop(&L);


Is it possible to move this a bit more close to the user?

andcarminati · 2026-07-22T09:48:23Z

+        break;
+    if (PN)
+      continue;
+    PN = PHINode::Create(Inst->getType(), 2, Inst->getName() + ".lver");


We are enforcing LCSSA everywhere. Here we catch a possible LCSSA violation, how is this possible? If we don't have a phi node and we need to create a new one we are definitely out of LCSSA. I think we could assert the property and simplify the code.

I guess this works as well:

void AIELoopVersioner::addExitPHIs( Loop *ClonedLoop, BasicBlock *ExitBlock, ValueToValueMapTy &VMap, ArrayRef<Instruction *> DefsUsedOutside) const { for (PHINode &PN : ExitBlock->phis()) { assert(PN.getNumIncomingValues() == 1 && PN.getIncomingBlock(0) == L.getExitingBlock() && "expected single-edge LCSSA phi"); Value *OrigValue = PN.getIncomingValue(0); Value *ClonedValue = VMap.lookup(OrigValue); PN.addIncoming(ClonedValue ? ClonedValue : OrigValue, ClonedLoop->getExitingBlock()); } }

I can't find a counter example.

andcarminati · 2026-07-22T11:40:55Z

+
+  DT.changeImmediateDominator(ExitBlock, GuardBB);
+  addExitPHIs(HighLoop, ExitBlock, VMap, DefsUsedOutside);
+  formDedicatedExitBlocks(HighLoop, &DT, &LI, nullptr, /*PreserveLCSSA=*/true);


Also, I guess we don't need to call formDedicatedExitBlocks here, simplifyLoop called previously guarantees this. If you see a potential need, we could create a test.

F-Stuckmann · 2026-07-24T09:59:39Z

+}
+
+void dropLoopMetadata(Loop &L, ArrayRef<StringRef> Keys) {
+  MDNode *LoopID = L.getLoopID();


F-Stuckmann · 2026-07-24T10:00:30Z

+void dropLoopMetadata(Loop &L, ArrayRef<StringRef> Keys) {
+  MDNode *LoopID = L.getLoopID();
+  if (!LoopID)
+    return;


nit: we could also early exit if Keys is empty

F-Stuckmann · 2026-07-24T13:39:00Z

+  // Collect enabled loops before mutating, so cloning does not disturb
+  // iteration.
+  SmallVector<Loop *, 4> Candidates;
+  for (Loop *L : LI.getLoopsInPreorder())


is the order relevant?

F-Stuckmann · 2026-07-24T13:41:07Z

+  // Read-only profitability gate first, so an unprofitable loop is rejected
+  // without any IR mutation.
+  if (!isProfitable()) {
+    LLVM_DEBUG(dbgs() << "  Not profitable to version\n");


nit: please mention the loop that is not profitable.
We may also want to emit optimiazation remarks, about loopVersion success/skipping.

F-Stuckmann · 2026-07-24T13:45:17Z

+                   ArrayRef<Instruction *> DefsUsedOutside) const;
+  /// Remove the versioning request hint so the postpipeliner leaves \p Loop
+  /// alone.
+  void stripVersioningHint(Loop &Loop) const;


nit: Why do we strip loop version hints? The postpipeliner specifically only checks if we have versioned metadata and not the versioning request. I don't think removing the loop version request is necessary if we didn't loopversion.

F-Stuckmann · 2026-07-24T14:32:53Z

+
+  BasicBlock *ExitBlock = canonicalizeAndCheckStructure();
+  if (!ExitBlock) {
+    // Versioning bailed but the request hint is still on the loop. Consume it


nit: Instead of a comment, could you just print the debug message?

mludevid requested review from F-Stuckmann, SagarMaheshwari99, abhinay-anubola, abnikant, andcarminati, katerynamuts, khallouh, konstantinschwarz, martien-de-jong, niwinanto and stephenneuendorffer as code owners July 10, 2026 15:18

andcarminati reviewed Jul 13, 2026

View reviewed changes

Comment thread llvm/lib/Target/AIE/Utils/AIELoopUtils.h

andcarminati reviewed Jul 13, 2026

View reviewed changes

Comment thread llvm/lib/Target/AIE/aie2ps/AIE2PSInstrInfo.cpp

andcarminati reviewed Jul 13, 2026

View reviewed changes

Comment thread llvm/lib/Target/AIE/AIEInnerLoopVersioning.cpp Outdated

andcarminati reviewed Jul 13, 2026

View reviewed changes

Comment thread llvm/lib/Target/AIE/AIEInnerLoopVersioning.cpp Outdated

andcarminati reviewed Jul 13, 2026

View reviewed changes

Comment thread llvm/lib/Target/AIE/AIEInnerLoopVersioning.cpp Outdated