[llvm] [LICM] Collect direct loop children immediately (PR #226740)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Fabian Parzefall (pzfl)

<details>
<summary>Changes</summary>

Rather than collecting a vector with all nodes in the loop and filtering afterwards, proactively collect only nodes of that are directly contained within the target loop. This avoids collecting the same set of already processed blocks within each nested loop level.

---
Full diff: https://github.com/llvm/llvm-project/pull/226740.diff


3 Files Affected:

- (modified) llvm/include/llvm/Transforms/Utils/LoopUtils.h (+5-3) 
- (modified) llvm/lib/Transforms/Scalar/LICM.cpp (+1-5) 
- (modified) llvm/lib/Transforms/Utils/LoopUtils.cpp (+19-18) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Utils/LoopUtils.h b/llvm/include/llvm/Transforms/Utils/LoopUtils.h
index 74c549be35ddf..be6e36dbefa9c 100644
--- a/llvm/include/llvm/Transforms/Utils/LoopUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/LoopUtils.h
@@ -229,10 +229,12 @@ LLVM_ABI bool promoteLoopAccessesToScalars(
     MemorySSAUpdater &, ICFLoopSafetyInfo *, OptimizationRemarkEmitter *,
     bool AllowSpeculation, bool HasReadsOutsideSet);
 
-/// Does a BFS from a given node to all of its children inside a given loop.
-/// The returned vector of basic blocks includes the starting point.
+/// Does a DFS from a given node to all of its children inside a given loop.
+/// The returned vector of basic blocks includes only blocks contained in
+/// \p CurLoop and not contained in any subloops.
 LLVM_ABI SmallVector<BasicBlock *, 16>
-collectChildrenInLoop(DominatorTree *DT, DomTreeNode *N, const Loop *CurLoop);
+collectDirectChildrenInLoop(DominatorTree *DT, DomTreeNode *N,
+                            const LoopInfo *LI, const Loop *CurLoop);
 
 /// Returns the instructions that use values defined in the loop.
 LLVM_ABI SmallVector<Instruction *, 8> findDefsUsedOutsideOfLoop(Loop *L);
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index c4335754e6c81..cb6bc760ef8f3 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -566,14 +566,10 @@ bool llvm::sinkRegion(DomTreeNode *N, AAResults *AA, LoopInfo *LI,
   // before their children in the worklist and process the worklist in reverse
   // order.
   SmallVector<BasicBlock *, 16> Worklist =
-      collectChildrenInLoop(DT, N, CurLoop);
+      collectDirectChildrenInLoop(DT, N, LI, CurLoop);
 
   bool Changed = false;
   for (BasicBlock *BB : reverse(Worklist)) {
-    // subloop (which would already have been processed).
-    if (inSubLoop(BB, CurLoop, LI))
-      continue;
-
     for (BasicBlock::iterator II = BB->end(); II != BB->begin();) {
       Instruction &I = *--II;
 
diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
index 784c833152611..4d3488b0c6dfb 100644
--- a/llvm/lib/Transforms/Utils/LoopUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -481,27 +481,28 @@ TransformationMode llvm::hasLICMVersioningTransformation(const Loop *L) {
   return TM_Unspecified;
 }
 
-/// Does a BFS from a given node to all of its children inside a given loop.
-/// The returned vector of basic blocks includes the starting point.
-SmallVector<BasicBlock *, 16> llvm::collectChildrenInLoop(DominatorTree *DT,
-                                                          DomTreeNode *N,
-                                                          const Loop *CurLoop) {
-  SmallVector<BasicBlock *, 16> Worklist;
-  auto AddRegionToWorklist = [&](DomTreeNode *DTN) {
-    // Only include subregions in the top level loop.
-    BasicBlock *BB = DTN->getBlock();
-    if (CurLoop->contains(BB))
-      Worklist.push_back(DTN->getBlock());
-  };
-
-  AddRegionToWorklist(N);
+/// Does a DFS from a given node to all of its children inside a given loop.
+/// The returned vector of basic blocks includes only blocks contained in
+/// \p CurLoop and not contained in any subloops.
+SmallVector<BasicBlock *, 16>
+llvm::collectDirectChildrenInLoop(DominatorTree *DT, DomTreeNode *N,
+                                  const LoopInfo *LI, const Loop *CurLoop) {
+  SmallVector<BasicBlock *, 16> Result;
+  SmallVector<DomTreeNode *, 16> DFS;
+  DFS.push_back(N);
+  while (!DFS.empty()) {
+    DomTreeNode *DTN = DFS.pop_back_val();
+    for (DomTreeNode *Child : DTN->children())
+      if (CurLoop->contains(Child->getBlock()))
+        DFS.push_back(Child);
 
-  for (size_t I = 0; I < Worklist.size(); I++) {
-    for (DomTreeNode *Child : DT->getNode(Worklist[I])->children())
-      AddRegionToWorklist(Child);
+    BasicBlock *BB = DTN->getBlock();
+    // Filter blocks in subloops.
+    if (LI->getLoopFor(BB) == CurLoop)
+      Result.push_back(BB);
   }
 
-  return Worklist;
+  return Result;
 }
 
 bool llvm::isAlmostDeadIV(PHINode *PN, BasicBlock *LatchBlock, Value *Cond) {

``````````

</details>


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