[llvm] 8f1efc2 - [Transforms][Utils] Preserve branch weights in LoopSplitUtils (#213626)

[via llvm-commits]

Author: Ashutosh Nema
Date: 2026-08-03T10:25:12Z
New Revision: 8f1efc26e47f6a3c7b2c623b5de44b99529e3a02

URL: https://github.com/llvm/llvm-project/commit/8f1efc26e47f6a3c7b2c623b5de44b99529e3a02
DIFF: https://github.com/llvm/llvm-project/commit/8f1efc26e47f6a3c7b2c623b5de44b99529e3a02.diff

LOG: [Transforms][Utils] Preserve branch weights in LoopSplitUtils (#213626)

Carry the original latch's branch weights onto the clamped latch, and
mark the newly created partition-guard branches as having unknown
weights so profile-tracking passes are not misled.

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/LoopSplitUtils.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/LoopSplitUtils.cpp b/llvm/lib/Transforms/Utils/LoopSplitUtils.cpp
index 1c8ea08ac80c6..db07e3f571f66 100644
--- a/llvm/lib/Transforms/Utils/LoopSplitUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopSplitUtils.cpp
@@ -60,6 +60,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
@@ -506,7 +507,16 @@ static void rewriteLatch(Loop *PL, Value *IndOp, Value *SelEnd,
                                                /*Inclusive=*/!LatchComparesPHI);
   Value *NewCmp = B.CreateICmp(Pred, IndOp, Bound, "itr.chk");
   B.SetInsertPoint(Term);
-  B.CreateCondBr(NewCmp, PL->getHeader(), Exit);
+  auto *NewBr = B.CreateCondBr(NewCmp, PL->getHeader(), Exit);
+  // Carry the original latch's weights over, mapping by which successor stayed
+  // in the loop.
+  uint64_t TrueW, FalseW;
+  if (extractBranchWeights(*Term, TrueW, FalseW)) {
+    bool Succ0InLoop = PL->contains(Term->getSuccessor(0));
+    setFittedBranchWeights(
+        *NewBr, {Succ0InLoop ? TrueW : FalseW, Succ0InLoop ? FalseW : TrueW},
+        /*IsExpected=*/false);
+  }
   Term->eraseFromParent();
   if (Cmp->use_empty())
     Cmp->eraseFromParent();
@@ -551,7 +561,10 @@ void LoopSplitUtils::chainPartitions(SplitState &S) {
       B.CreateBr(P.Preheader);
     } else {
       Value *Enter = B.CreateICmp(GuardPred, P.StartVal, P.SelEnd, "itr.chk");
-      B.CreateCondBr(Enter, P.Preheader, MergeAfter);
+      auto *GuardBr = B.CreateCondBr(Enter, P.Preheader, MergeAfter);
+      // New control flow with no source profile; record the weights as unknown
+      // so profile-tracking passes are not misled.
+      setExplicitlyUnknownBranchWeightsIfProfiled(*GuardBr, DEBUG_TYPE);
     }
     GuardTerm->eraseFromParent();