[llvm] [LoopRotate] Preserve block frequencies for rotated loops (PR #202219)

[Alok Kumar Sharma via llvm-commits]

alokkrsharma wrote:

> @jdoerfert @MatzeB @arsenm @mtrofin
> 
> I appreciate this PR's effort to fix profile corruption in LoopRotate and to drive the conversation forward. However, in my opinion, `updateBranchWeights` in `LoopRotationUtils.cpp` is already more complex than it needs to be, and this PR makes it more complex. I think there is a simpler way that we have already discussed for other passes.
> 
> ### Preserve branch weights to preserve BFI
> LoopRotate is one of the passes where I think we should implement [[RFC] Fix Loop Transformations to Preserve Block Frequencies](https://discourse.llvm.org/t/rfc-fix-loop-transformations-to-preserve-block-frequencies/85785), which was previously reviewed and accepted by the community. The premise is simple: when possible, transformations should maintain branch weights consistent with the original loop for the sake of preserving BFI.
> 
> In the case of LoopRotate, the original loop header acts a guard for every loop iteration, so its branch weights encode the probability that the iteration will execute once it is reached. Otherwise, the loop exits. Splitting the header into a preheader (the guard of the first iteration) and a latch (the guard of every remaining iteration) does not change the probability that any iteration will execute once it is reached. That is, merely copying the branch weights should be sufficient to preserve BFI.
> 
> The above simplification should hold even for multi-exit loops: the extra exits within a loop iteration do not change the probability of entering each loop iteration to begin with once it is reached. So, the simplification would have a second win: it fixes multi-exit loops.
> 
> ### Set estimated trip counts
> Like LoopPeel and LoopUnroll before the above RFC, LoopRotate currently tries to encode some sense of trip counts in branch weights. That is the source of its `updateBranchWeights` complexity, and preserving that approach is the source of this PR's complexity, as I understand it.
> 
> Based on a few experiments I have tried in the past, LoopRotate currently manages to do so in a way that does not corrupt the aggregate loop body BFI in the single-exit case. That is why LoopRotate has not been a high priority for me. First, it reduces the latch probability to reduce the estimated trip count there by 1 (e.g., if the original header is reached 11 times, then the loop body executes 10 times). Then, it compensates by increasing the new preheader probability, thus preserving BFI.
> 
> I know of no general way to apply LoopRotate's approach in other passes, like LoopPeel or LoopUnroll, where previous attempts to encode trip counts in branch weights corrupted BFI. The above RFC explains their situation and offers a simpler and more universal solution: stop conflating trip counts and branch weights, and instead encode the estimated trip count in separate metadata. So LoopRotate would need to do that too, but that would still be significantly simpler than what its `updateBranchWeights` does now.
> 
> ### Compiler contradicts uniform loop probability
> But what do we do if the compiler proves the first iteration is unconditional (`!HasConditionalPreHeader` in the implementation)? This is not the first time we have seen the compiler prove the original uniform loop probability is incorrect for a specific iteration. As in other cases, to preserve aggregate loop body BFI, we can adjust the remaining latch probability to compensate.
> 
> For single-exit loops, that is a small bit of code, and PR #168250 already landed to solve a nearly identical issue for LoopPeel. There is a more complex PR series for LoopUnroll, ending at PR #182405, but that is far more than is needed here.
> 
> For multi-exit loops, the solution will be more challenging because we then need to calculate the probability of reaching the latch from the header. I think that calculation is already encapsulated in BFI, so one approach is to import BFI as this PR does.
> 
> ### Summary
> Following the above RFC would make LoopRotate more consistent with our current goals in other passes, like LoopPeel and LoopUnroll. It would also make LoopRotate simpler than it is even now while additionally fixing multi-exit loops except in the multi-exit `!HasConditionalPreHeader` case. For that case, we likely need to import BFI as this PR does, but the result should still be simpler than this PR's approach.

@jdenny-ornl — thanks again for the review. I replaced the earlier patch with an RFC-aligned version: preserve header branch weights when the guard stays conditional, move post-rotation trip-count updates into loop metadata, and use local BFI only for folded-guard multi-exit latch weights. The BFI-to-branch-weight rescaling path and new pass flags are removed. Let me know if this matches what you had in mind.

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