[llvm] 8cf1524 - [loop unroll] Fix `branch-weights` for unrolled loop.

[Mircea Trofin via llvm-commits]

Author: Mircea Trofin
Date: 2023-02-14T12:00:53-08:00
New Revision: 8cf1524cbc7f4d8dade005de69abcd07a9b5e32a

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

LOG: [loop unroll] Fix `branch-weights` for unrolled loop.

The branch weights of the unrolled loop need to be reduced by the
unroll factor.

Differential Revision: https://reviews.llvm.org/D143948

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/LoopUnroll.cpp
    llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll
    llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index e8f585b4a94dd..916c026928235 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -310,6 +310,9 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   const unsigned MaxTripCount = SE->getSmallConstantMaxTripCount(L);
   const bool MaxOrZero = SE->isBackedgeTakenCountMaxOrZero(L);
+  unsigned EstimatedLoopInvocationWeight = 0;
+  std::optional<unsigned> OriginalTripCount =
+      llvm::getLoopEstimatedTripCount(L, &EstimatedLoopInvocationWeight);
 
   // Effectively "DCE" unrolled iterations that are beyond the max tripcount
   // and will never be executed.
@@ -830,8 +833,16 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   Loop *OuterL = L->getParentLoop();
   // Update LoopInfo if the loop is completely removed.
-  if (CompletelyUnroll)
+  if (CompletelyUnroll) {
     LI->erase(L);
+    // We shouldn't try to use `L` anymore.
+    L = nullptr;
+  } else if (OriginalTripCount) {
+    // Update the trip count. Note that the remainder has already logic
+    // computing it in `UnrollRuntimeLoopRemainder`.
+    setLoopEstimatedTripCount(L, *OriginalTripCount / ULO.Count,
+                              EstimatedLoopInvocationWeight);
+  }
 
   // LoopInfo should not be valid, confirm that.
   if (UnrollVerifyLoopInfo)

diff  --git a/llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll b/llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll
index f3e3292021de4..00ddfc82feeb4 100644
--- a/llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll
+++ b/llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll
@@ -6,7 +6,7 @@
 ; CHECK: br i1 [[COND1:%.*]], label %for.end.loopexit.unr-lcssa.loopexit, label %for.body, !prof ![[#PROF:]], !llvm.loop ![[#LOOP:]]
 ; CHECK-LABEL: for.body.epil:
 ; CHECK: br i1 [[COND2:%.*]], label  %for.body.epil, label %for.end.loopexit.epilog-lcssa, !prof ![[#PROF2:]], !llvm.loop ![[#LOOP2:]]
-; CHECK: ![[#PROF]] = !{!"branch_weights", i32 1, i32 9999}
+; CHECK: ![[#PROF]] = !{!"branch_weights", i32 1, i32 2499}
 ; CHECK: ![[#PROF2]] = !{!"branch_weights", i32 3, i32 1}
 
 define i3 @test(ptr %a, i3 %n) {

diff  --git a/llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll b/llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll
index 8d39b175b2c3f..725b4cebe55f5 100644
--- a/llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll
+++ b/llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll
@@ -8,7 +8,9 @@
 ; CHECK: %mul.1 = mul
 ; CHECK: %mul.2 = mul
 ; CHECK: %mul.3 = mul
+; CHECK: br i1 %niter.ncmp.7, label %loop.end.unr-lcssa.loopexit, label %loop, !prof !1
 ; CHECK: loop.epil:
+; CHECK:   br i1 %epil.iter.cmp, label %loop.epil, label %loop.end.epilog-lcssa, !prof !2, !llvm.loop !3
 define i32 @bar_prof(ptr noalias nocapture readonly %src, i64 %c) !prof !1 {
 entry:
   br label %loop
@@ -57,3 +59,6 @@ loop.end:
 
 !1 = !{!"function_entry_count", i64 1}
 !2 = !{!"branch_weights", i32 1, i32 1000}
+
+; CHECK: !1 = !{!"branch_weights", i32 1, i32 124}
+; CHECK: !2 = !{!"branch_weights", i32 7, i32 1}
\ No newline at end of file