[Transforms][Utils] Preserve branch weights in LoopSplitUtils

[nema-ashutosh]

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.

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

Author: Ashutosh Nema (nema-ashutosh)

Changes

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.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Utils/LoopSplitUtils.cpp (+16-2)

diff --git a/llvm/lib/Transforms/Utils/LoopSplitUtils.cpp b/llvm/lib/Transforms/Utils/LoopSplitUtils.cpp
index 1c8ea08ac80c6..24295f6867dc6 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,17 @@ 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 branch weights onto the clamped latch, matching
+  // by which original successor stayed in the loop (the "keep iterating" edge).
+  SmallVector<uint32_t, 2> Weights;
+  if (extractBranchWeights(*Term, Weights)) {
+    bool Succ0InLoop = PL->contains(Term->getSuccessor(0));
+    setBranchWeights(*NewBr,
+                     {Succ0InLoop ? Weights[0] : Weights[1],
+                      Succ0InLoop ? Weights[1] : Weights[0]},
+                     /*IsExpected=*/false);
+  }
   Term->eraseFromParent();
   if (Cmp->use_empty())
     Cmp->eraseFromParent();
@@ -551,7 +562,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();
 

[nema-ashutosh]

@artagnon - please review, its a fix for #205995 build breakage

[artagnon]

Could use this overload:

bool llvm::extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
                                uint64_t &FalseVal)

There are only ever two weights?

[nema-ashutosh]

Sure, switched to the extractBranchWeights overload, and also moved to setFittedBranchWeights to avoid the narrowing cast. hope its OK !

[artagnon]

Merging. By the way, if you can expose your email on GitHub, that would be good -- we generally don't prefer github-noreply emails in commit logs?

[artagnon]

Would also be good to add a test for this, now that the urgent fix has been merged?

[nema-ashutosh]

sure, thanks !

[artagnon]

LGTM, with suggestion to use other overload, thanks.

[mtrofin]

You're missing unit tests for this. A simple way to add them could be adding !prof to the subset of your current tests from #205995 that covers these cases.

[mtrofin]

can the probability associated with GuardPred be derived from the existing probabilities?

another possibility may be to derive it from the BFI of the loop basic blocks pre-split.

[artagnon]

You're missing unit tests for this. A simple way to add them could be adding !prof to the subset of your current tests from #205995 that covers these cases.

It was an urgent fix as buildbots were failing, and the test has been added as a follow-up: #213647.

[fhahn]

It was an urgent fix as buildbots were failing, and the test has been added as a follow-up: #213647.

In that case, reverting is usually a safe option as well