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

[Mircea Trofin via Phabricator via llvm-commits]

mtrofin updated this revision to Diff 497399.
mtrofin added a comment.

feedback


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D143948/new/

https://reviews.llvm.org/D143948

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


Index: llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll
===================================================================
--- llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll
+++ 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
Index: llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll
===================================================================
--- llvm/test/Transforms/LoopUnroll/runtime-loop-branchweight.ll
+++ 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) {
Index: llvm/lib/Transforms/Utils/LoopUnroll.cpp
===================================================================
--- llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ 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<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 @@
 
   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)


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D143948.497399.patch
Type: text/x-patch
Size: 2969 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230214/352ac551/attachment-0001.bin>