[PATCH] D28179: Make sure total loop body weight is preserved in loop peeling

[Xin Tong via Phabricator via llvm-commits]

trentxintong updated this revision to Diff 82787.
trentxintong added a comment.

Rework BE count for the remainder loop a bit.


https://reviews.llvm.org/D28179

Files:
  lib/Transforms/Utils/LoopUnrollPeel.cpp
  test/Transforms/LoopUnroll/peel-loop-pgo.ll


Index: test/Transforms/LoopUnroll/peel-loop-pgo.ll
===================================================================
--- test/Transforms/LoopUnroll/peel-loop-pgo.ll
+++ test/Transforms/LoopUnroll/peel-loop-pgo.ll
@@ -43,5 +43,5 @@
 ;CHECK: !1 = !{!"branch_weights", i32 900, i32 101}
 ;CHECK: !2 = !{!"branch_weights", i32 540, i32 360}
 ;CHECK: !3 = !{!"branch_weights", i32 162, i32 378}
-;CHECK: !4 = !{!"branch_weights", i32 560, i32 162}
+;CHECK: !4 = !{!"branch_weights", i32 1399, i32 162}
 
Index: lib/Transforms/Utils/LoopUnrollPeel.cpp
===================================================================
--- lib/Transforms/Utils/LoopUnrollPeel.cpp
+++ lib/Transforms/Utils/LoopUnrollPeel.cpp
@@ -335,10 +335,12 @@
   unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
 
   uint64_t TrueWeight, FalseWeight;
-  uint64_t ExitWeight = 0, BackEdgeWeight = 0;
+  uint64_t ExitWeight = 0, LoopBodyWeight = 0;
   if (LatchBR->extractProfMetadata(TrueWeight, FalseWeight)) {
     ExitWeight = HeaderIdx ? TrueWeight : FalseWeight;
-    BackEdgeWeight = HeaderIdx ? FalseWeight : TrueWeight;
+    // The # of times the loop body executes is the sum of the exit block
+    // weight and the # of times the backedges are taken.
+    LoopBodyWeight = TrueWeight + FalseWeight;
   }
 
   // For each peeled-off iteration, make a copy of the loop.
@@ -351,10 +353,10 @@
     // the current peeled-off static iteration uses up.
     // FIXME: due to the way the distribution is constructed, we need a
     // guard here to make sure we don't end up with non-positive weights.
-    if (ExitWeight < BackEdgeWeight)
-      BackEdgeWeight -= ExitWeight;
+    if (ExitWeight < LoopBodyWeight)
+      LoopBodyWeight -= ExitWeight;
     else
-      BackEdgeWeight = 1;
+      LoopBodyWeight = 1;
 
     cloneLoopBlocks(L, Iter, InsertTop, InsertBot, Exit,
                     NewBlocks, LoopBlocks, VMap, LVMap, LI);
@@ -388,6 +390,14 @@
 
   // Adjust the branch weights on the loop exit.
   if (ExitWeight) {
+    // The backedge count is the difference of remaining loop body weight and
+    // loop body entry count. If the remaining loop body weight is smaller than
+    // this loop body entry weight, we mark the loop backedge weight as 1.
+    uint64_t BackEdgeWeight = 0;
+    if (ExitWeight < LoopBodyWeight)
+      BackEdgeWeight = LoopBodyWeight - ExitWeight;
+    else
+      BackEdgeWeight = 1;
     MDBuilder MDB(LatchBR->getContext());
     MDNode *WeightNode =
         HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D28179.82787.patch
Type: text/x-patch
Size: 2559 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20170101/dcbfceee/attachment.bin>