diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp --- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp +++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp @@ -310,6 +310,9 @@ const unsigned MaxTripCount = SE->getSmallConstantMaxTripCount(L); const bool MaxOrZero = SE->isBackedgeTakenCountMaxOrZero(L); + unsigned EstimatedLoopInvocationWeight = 0; + std::optional OriginalTripCount = + llvm::getLoopEstimatedTripCount(L, &EstimatedLoopInvocationWeight); // Effectively "DCE" unrolled iterations that are beyond the max tripcount // and will never be executed. @@ -830,8 +833,16 @@ 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 --- 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 --- 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 @@ !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