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

[Xin Tong via Phabricator via llvm-commits]

trentxintong created this revision.
trentxintong added reviewers: mkuper, davidxl, anemet.
trentxintong added a subscriber: llvm-commits.

Regardless how the loop body weight is distributed, we should preserve
total loop body weight. i.e. we should have same weight reaching the body of the loop
or its duplicates in peeled and unpeeled case.


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,18 @@
 
   // 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. We
+    // also re-adjust ExitWeight so that the sum of the exit counts is the same
+    // as the entry count.
+    uint64_t BackEdgeWeight = 0;
+    if (ExitWeight < LoopBodyWeight)
+      BackEdgeWeight = LoopBodyWeight - ExitWeight;
+    else {
+      BackEdgeWeight = 1;
+      ExitWeight -= 1;
+    }
     MDBuilder MDB(LatchBR->getContext());
     MDNode *WeightNode =
         HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D28179.82736.patch
Type: text/x-patch
Size: 2704 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20161230/33fbcff3/attachment.bin>