[llvm] [LoopSafetyInfo] Compute block colors lazily (NFC) (PR #221954)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/221954

Computing the block colors is fairly expensive, but only rarely actually needed. We can delay computation to first use.

Fixes https://github.com/llvm/llvm-project/issues/221162.

>From 362379a2beb6851e7cb330f1298d9f25196fe998 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Tue, 8 Sep 2026 12:45:40 +0200
Subject: [PATCH] [LoopSafetyInfo] Compute block colors lazily (NFC)

Computing the block colors is fairly expensive, but only rarely
actually needed. We can delay computation to first use.
---
 llvm/include/llvm/Analysis/MustExecute.h |  8 ++++----
 llvm/lib/Analysis/MustExecute.cpp        | 20 +++++++++++++-------
 llvm/lib/Transforms/Scalar/LICM.cpp      |  5 +++--
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MustExecute.h b/llvm/include/llvm/Analysis/MustExecute.h
index 7e64d0622a62a..d22c3fc1dfb15 100644
--- a/llvm/include/llvm/Analysis/MustExecute.h
+++ b/llvm/include/llvm/Analysis/MustExecute.h
@@ -54,7 +54,7 @@ class raw_ostream;
 /// changes.
 class LoopSafetyInfo {
   // Used to update funclet bundle operands.
-  DenseMap<BasicBlock *, ColorVector> BlockColors;
+  mutable std::optional<DenseMap<BasicBlock *, ColorVector>> BlockColors;
 
   // Cache whether (the start of) this block is guaranteed to execute if the
   // loop is entered.
@@ -63,12 +63,12 @@ class LoopSafetyInfo {
   bool allLoopPathsLeadToBlockImpl(const BasicBlock *BB,
                                    const DominatorTree *DT) const;
 
+  /// Computes block colors.
+  void computeBlockColors() const;
+
 protected:
   const Loop *CurLoop;
 
-  /// Computes block colors.
-  LLVM_ABI void computeBlockColors();
-
 public:
   /// Returns block colors map that is used to update funclet operand bundles.
   LLVM_ABI const DenseMap<BasicBlock *, ColorVector> &getBlockColors() const;
diff --git a/llvm/lib/Analysis/MustExecute.cpp b/llvm/lib/Analysis/MustExecute.cpp
index 2462ef716dd52..463477ce4aef1 100644
--- a/llvm/lib/Analysis/MustExecute.cpp
+++ b/llvm/lib/Analysis/MustExecute.cpp
@@ -28,12 +28,17 @@ using namespace llvm;
 
 const DenseMap<BasicBlock *, ColorVector> &
 LoopSafetyInfo::getBlockColors() const {
-  return BlockColors;
+  computeBlockColors();
+  return *BlockColors;
 }
 
 void LoopSafetyInfo::copyColors(BasicBlock *New, BasicBlock *Old) {
-  ColorVector &ColorsForNewBlock = BlockColors[New];
-  ColorVector &ColorsForOldBlock = BlockColors[Old];
+  // Nothing to update if colors have not been computed yet.
+  if (!BlockColors)
+    return;
+
+  ColorVector &ColorsForNewBlock = (*BlockColors)[New];
+  ColorVector &ColorsForOldBlock = (*BlockColors)[Old];
   ColorsForNewBlock = ColorsForOldBlock;
 }
 
@@ -62,8 +67,6 @@ void SimpleLoopSafetyInfo::computeLoopSafetyInfo() {
     if (MayThrow)
       break;
   }
-
-  computeBlockColors();
 }
 
 bool ICFLoopSafetyInfo::blockMayThrow(const BasicBlock *BB) const {
@@ -85,7 +88,6 @@ void ICFLoopSafetyInfo::computeLoopSafetyInfo() {
       MayThrow = true;
       break;
     }
-  computeBlockColors();
 }
 
 void ICFLoopSafetyInfo::insertInstructionTo(const Instruction *Inst,
@@ -99,7 +101,11 @@ void ICFLoopSafetyInfo::removeInstruction(const Instruction *Inst) {
   MW.removeInstruction(Inst);
 }
 
-void LoopSafetyInfo::computeBlockColors() {
+void LoopSafetyInfo::computeBlockColors() const {
+  if (BlockColors)
+    return;
+  BlockColors.emplace();
+
   // Compute funclet colors if we might sink/hoist in a function with a funclet
   // personality routine.
   Function *Fn = CurLoop->getHeader()->getParent();
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index bc12aca7c1086..e4bcad4ac731d 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -1427,7 +1427,6 @@ static bool isNotUsedOrFoldableInLoop(const Instruction &I, const Loop *CurLoop,
                                       const LoopSafetyInfo *SafetyInfo,
                                       TargetTransformInfo *TTI,
                                       bool &FoldableInLoop, bool LoopNestMode) {
-  const auto &BlockColors = SafetyInfo->getBlockColors();
   bool IsFoldable = isFoldableInLoop(I, CurLoop, TTI);
   for (const User *U : I.users()) {
     const Instruction *UI = cast<Instruction>(U);
@@ -1439,10 +1438,12 @@ static bool isNotUsedOrFoldableInLoop(const Instruction &I, const Loop *CurLoop,
 
       // We need to sink a callsite to a unique funclet.  Avoid sinking if the
       // phi use is too muddled.
-      if (isa<CallInst>(I))
+      if (isa<CallInst>(I)) {
+        const auto &BlockColors = SafetyInfo->getBlockColors();
         if (!BlockColors.empty() &&
             BlockColors.find(const_cast<BasicBlock *>(BB))->second.size() != 1)
           return false;
+      }
 
       if (LoopNestMode) {
         while (isa<PHINode>(UI) && UI->hasOneUser() &&