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

[Fabian Parzefall via llvm-commits]

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

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.

>From a466a730201386a7fe45f1706a2e14eb65e0ee01 Mon Sep 17 00:00:00 2001
From: Fabian Parzefall <fabian.parzefall at mailbox.org>
Date: Sun, 20 Sep 2026 14:40:17 -0700
Subject: [PATCH] [LICM] Collect direct loop children immediately

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.
---
 .../include/llvm/Transforms/Utils/LoopUtils.h |  8 ++--
 llvm/lib/Transforms/Scalar/LICM.cpp           |  6 +--
 llvm/lib/Transforms/Utils/LoopUtils.cpp       | 37 ++++++++++---------
 3 files changed, 25 insertions(+), 26 deletions(-)

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) {