[llvm] [LoopVectorize] Enable vectorization of uncountable loops with invariant bound loads via speculative hoisting and runtime versioning. (PR #195027)

[Ashutosh Nema via llvm-commits]

================
@@ -8163,6 +8164,201 @@ static void connectEpilogueVectorLoop(VPlan &EpiPlan, Loop *L,
       Phi.eraseFromParent();
 }
 
+// Is profitable to generate runtime check dominated version of the loop? which
+// has speculatively hoisted the loop bound load.
+static bool isSpeculativeBoundVersioningProfitable() { return true; }
+
+/*
+
+              preheader
+                |
+              rtcheck-block
+                /  \
+               /    \
+        for.preheader    ver.preheader
+            |               |
+         for.body       ver.for.body
+
+*/
+static Loop *tryToVersionLoopForInvariantBoundLoad(
+    Loop *L, LoadInst *BoundLoad, const DataLayout &DL, DominatorTree &DT,
+    LoopInfo &LI, AAResults &AA, PredicatedScalarEvolution &PSE,
+    AssumptionCache *AC) {
+  LoadInst *InnerLoad = BoundLoad;
+  if (!InnerLoad)
+    return nullptr;
+  auto *PreHeader = L->getLoopPreheader();
+  if (!PreHeader)
+    return nullptr;
+
+  Value *BoundPtrOl = InnerLoad->getPointerOperand();
+  if (auto *BoundPtrI = dyn_cast<Instruction>(BoundPtrOl))
+    if (!DT.properlyDominates(BoundPtrI->getParent(), PreHeader))
+      return nullptr;
+
+  // All the relevent base pointers of participating stores should dominate both
+  // loops.
+  LoopAccessInfoManager PreFlightLAIs(*PSE.getSE(), AA, DT, LI, nullptr,
+                                      nullptr, AC);
+  const LoopAccessInfo &PreLAI = PreFlightLAIs.getInfo(*L);
+  const auto &RtPtrChecking = *PreLAI.getRuntimePointerChecking();
+
+  for (const auto &Check : RtPtrChecking.getChecks()) {
+    for (const RuntimeCheckingPtrGroup *Group : {Check.first, Check.second}) {
+      for (unsigned Idx : Group->Members) {
+        Value *Ptr = RtPtrChecking.Pointers[Idx].PointerValue;
+        Value *Base = getUnderlyingObject(Ptr);
+        if (auto *BaseI = dyn_cast<Instruction>(Base))
+          if (!DT.properlyDominates(BaseI->getParent(), PreHeader)) {
+            return nullptr;
+          }
+      }
+    }
+  }
+
+  Function *ParentF = PreHeader->getParent();
+  LLVMContext &Ctx = ParentF->getContext();
+
+  BasicBlock *RtCheckBB =
+      BasicBlock::Create(Ctx, "ivbound.rtcheck", ParentF, PreHeader);
+
+  // Temporary terminator for SCEVExapnder to have valid insert point.
+  IRBuilder<> TmpBuilder(RtCheckBB);
+  Instruction *TmpTerm = TmpBuilder.CreateUnreachable();
+
+  // Re-wire the CFG.
+  // Redirect the predecessors of preheader the new rtcheck block.
+  SmallVector<BasicBlock *, 4> PreHeaderPreds(predecessors(PreHeader));
+  for (auto *Pred : PreHeaderPreds)
+    Pred->getTerminator()->replaceUsesOfWith(PreHeader, RtCheckBB);
+
+  BasicBlock *OldPreheadIdom = DT.getNode(PreHeader)->getIDom()->getBlock();
+  DT.addNewBlock(RtCheckBB, OldPreheadIdom);
+  DT.changeImmediateDominator(PreHeader, RtCheckBB);
+
+  // Fix the preheader phis via rtcblock phis.
+  for (PHINode &PN : make_early_inc_range(PreHeader->phis())) {
+    PHINode *RtcPN =
+        PHINode::Create(PN.getType(), PreHeaderPreds.size(),
+                        PN.getName() + ".rtcphi", RtCheckBB->begin());
+    for (unsigned i = 0; i < PN.getNumIncomingValues(); i++) {
+      BasicBlock *IBB = PN.getIncomingBlock(i);
+      assert(is_contained(PreHeaderPreds, IBB) &&
+             "Incoming Basic Block not part of predecessors list!");
+      RtcPN->addIncoming(PN.getIncomingValue(i), IBB);
+    }
+    PN.replaceAllUsesWith(RtcPN);
+    PN.eraseFromParent();
+  }
+
+  // Clone the loop including preheader.
+  ValueToValueMapTy VMap;
+  SmallVector<BasicBlock *, 4> NewBlocks;
+  Loop *VerLoop = cloneLoopWithPreheader(PreHeader, RtCheckBB, L, VMap, ".lver",
+                                         &LI, &DT, NewBlocks);
+
+  remapInstructionsInBlocks(NewBlocks, VMap);
+
+  BasicBlock *VerPreHeader = cast<BasicBlock>(VMap[PreHeader]);
+  BasicBlock *VerLatch = cast<BasicBlock>(VMap[L->getLoopLatch()]);
+
+  if (Loop *ParLoop = L->getParentLoop())
+    ParLoop->addBasicBlockToLoop(RtCheckBB, LI);
+
+  VerLatch->getTerminator()->setMetadata(LLVMContext::MD_loop, nullptr);
+
+  LoadInst *HoistLoad = cast<LoadInst>(InnerLoad->clone());
+  HoistLoad->setName(InnerLoad->getName() + ".speculatively.hoisted");
+  // Put this hoisted load in RtcheckBB, since SCEVExapnder needs it.
+  HoistLoad->insertBefore(TmpTerm->getIterator());
+  Instruction *ClonedInnerLoad = cast<Instruction>(VMap[InnerLoad]);
+  ClonedInnerLoad->replaceAllUsesWith(HoistLoad);
+  ClonedInnerLoad->eraseFromParent();
+
+  // fix lcssa phis.
+  SmallVector<BasicBlock *, 4> ExitBlocks;
+  L->getExitBlocks(ExitBlocks);
+  for (BasicBlock *ExBB : ExitBlocks) {
+    DT.changeImmediateDominator(ExBB, RtCheckBB);
+    for (PHINode &Phi : ExBB->phis()) {
+      int idx = Phi.getBasicBlockIndex(L->getLoopLatch());
+      if (idx < 0)
+        continue;
+      Value *OrigVal = Phi.getIncomingValue(idx);
+      Value *NewVal =
+          VMap.count(OrigVal) ? cast<Value>(VMap[OrigVal]) : OrigVal;
+      Phi.addIncoming(NewVal, VerLatch);
+    }
+  }
+
+  formDedicatedExitBlocks(L, &DT, &LI, nullptr, true);
+  formDedicatedExitBlocks(VerLoop, &DT, &LI, nullptr, true);
+
+  // RTC condition generation using SCEV.
+  auto &SE = *PSE.getSE();
+  LoopAccessInfoManager NewLAIs(SE, AA, DT, LI, nullptr, nullptr, AC);
+  const LoopAccessInfo &VerLAI = NewLAIs.getInfo(*VerLoop);
+
+  // Generate LAA based ptr runtime checks.
+  SCEVExpander MemExp(*VerLAI.getRuntimePointerChecking()->getSE(),
+                      "ivbound.mem.check");
+  Value *MemRTC = nullptr;
+  if (VerLAI.getRuntimePointerChecking()->Need)
----------------
nema-ashutosh wrote:

We always expect the Need to insert a runtime check should be TRUE here, as we already performed versioning.

It should be an assertion, instead of if condition.

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