[LoopPeeling] Better handling of branch weights for small values

[skatkov]

Summary

If the original exit weight for latch to exit is small the new
weights for copies of latch will be non-informative like
[0,1] if original exit weight is 1.

This CL tries to scale original exit weight to get reasonable values.

[reames]

lib/Transforms/Utils/LoopUnrollPeel.cpp

392

Are you based on a stale version? I don't have this function in this file on ToT.

402

This comment is concerning. This appears to be assigning meaning to the numeric value of the branch weights. The only meaning to the branch weights is in their *ratio*.

It looks like you're trying to extend a broken scheme, rather than fixing it. Please don't. :)

An alternative here would be to compute a header weight based on 1 + the estimated trip count, then decrement by 1 on each iteration. I think this would both be simpler, more stable, and likely get the same result on all of your existing test cases.

[skatkov]

lib/Transforms/Utils/LoopUnrollPeel.cpp

392

I'm basing on https://reviews.llvm.org/D63917

402

:) It is a tough topic :)

Let me try first to explain how I understand the current scheme works.

There is a latch. Say A is a weight to go to header and B is a weight to go to exit.
So A/(A+B) is a probability to go to the loop header.
The current scheme looks at it in a different angle. If during execution of the code we come to latch A+B times then A times we will go to the header and B times we go to the exit. It is evident that number of time we come to pre-header should be also B.
So it is similar to gathering a profile. So weight is a number of times we take the corresponding edge.

After the peeling we want to set the branch weights on latch and copies of the latch such as if we come to the first copy of the header B times then

1. Number of times we come to header is B (sum of all weights of edges coming to exit is B)
2. Number of times we take all copies of loop body is A
3. Each individual pair of weights for copy of latch is some reasonable distribution so that the probability to go to exit on the next iteration growth and almost 1 on the last iteration.

See comment for updateBranchWeights function where some invariants are described.

I do not say it is best idea but in this patch I tried to avoid modification it but a bit adjust to behaves as it was expected when the values for weights are small.

Now about your proposal.
As I understand you propose to set branch weights for i-th peeling iteration as (TC + 1 - i, 1), am I correct?
I'm fine with that but probably it is better to get feedback from the author of current approach.

According to git history it seems it is Michael Kuperstein, Michael do you see any problems with this modification/simplification?

[skatkov]

@mukper, Can you please comment on my last comment?

[skatkov]

lib/Transforms/Utils/LoopUnrollPeel.cpp

402

The patch implemented this idea is uploaded as https://reviews.llvm.org/D64235

[skatkov]

Abandon in a favor of https://reviews.llvm.org/D64235

lib/Transforms/Utils/LoopUnrollPeel.cpp

 static void updateBranchWeights(BasicBlock *Header, BranchInst *LatchBR,
                                 unsigned IterNumber, unsigned AvgIters,
                                 uint64_t &PeeledHeaderWeight) {
   // FIXME: Pick a more realistic distribution.
   // Currently the proportion of weight we assign to the fall-through
   // side of the branch drops linearly with the iteration number, and we use
   // a 0.9 fudge factor to make the drop-off less sharp...
   if (PeeledHeaderWeight) {
-    uint64_t FallThruWeight =
-        PeeledHeaderWeight * ((float)(AvgIters - IterNumber) / AvgIters * 0.9);
+    // Use ceil to get preference for in loop edge.
+    uint64_t FallThruWeight = ceil(
+        PeeledHeaderWeight * ((float)(AvgIters - IterNumber) / AvgIters * 0.9));
     uint64_t ExitWeight = PeeledHeaderWeight - FallThruWeight;
     PeeledHeaderWeight -= ExitWeight;
 
     unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
     MDBuilder MDB(LatchBR->getContext());
     MDNode *WeightNode =
         HeaderIdx ? MDB.createBranchWeights(ExitWeight, FallThruWeight)
                   : MDB.createBranchWeights(FallThruWeight, ExitWeight);
     LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
   }
 }
 
 /// Initialize the weights.
 ///
 /// \param Header The header block.
 /// \param LatchBR The latch branch.
 /// \param AvgIters The average number of iterations we expect the loop to have.
 /// \param[out] ExitWeight The weight of the edge from Latch to Exit block.
 /// \param[out] CurHeaderWeight The # of time the header is executed.
 static void initBranchWeights(BasicBlock *Header, BranchInst *LatchBR,
                               unsigned AvgIters, uint64_t &ExitWeight,
                               uint64_t &CurHeaderWeight) {
   uint64_t TrueWeight, FalseWeight;
   if (!LatchBR->extractProfMetadata(TrueWeight, FalseWeight))
     return;
   unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
   ExitWeight = HeaderIdx ? TrueWeight : FalseWeight;
   // The # of times the loop body executes is the sum of the exit block
   // weight and the # of times the backedges are taken.
   CurHeaderWeight = TrueWeight + FalseWeight;
+
+  // if ExitWeight is too small (for example due to normalization of weights
+  // then we will not be able to distribute this small weight between all
+  // peeled copies of branch. To get the distribution we'd like to see we can
+  // scale the backedge and exit weight to get enough value for ExitWeight.
+  // Taking into account that we know the multipliers used in
+  // updateBranchWeights in advance we can pre-compute what value of ExitWeight
+  // will be enough to have at least 1 for ExitWeight at the end.
+  // If AvgIters is greater then 10 we can have a problem with precision of
+  // float then just avoid any modification at all.
+  if (AvgIters > 10)
+    return;
+
+  float C = 1.0;
+  for (unsigned i = 0; i < AvgIters; i++)
+    C *= ((float)(AvgIters - i) / AvgIters * 0.9);
+  uint64_t MinExit = ceil(1 / C);
+  // If ExitWeight is enough big, no need to scale.
+  if (ExitWeight >= MinExit)
+    return;
+
+  uint64_t Multiplier = ceil((float)MinExit / ExitWeight);
+  // Handle overflow of CurHeaderWeight.
+  if (CurHeaderWeight * Multiplier >= 1u << 31)
+    return;
+
+  ExitWeight *= Multiplier;
+  CurHeaderWeight *= Multiplier;
 }
 
 /// Update the weights of original Latch block after peeling off all iterations.
 ///
 /// \param Header The header block.
 /// \param LatchBR The latch branch.
 /// \param ExitWeight The weight of the edge from Latch to Exit block.
 /// \param CurHeaderWeight The # of time the header is executed.

test/Transforms/LoopUnroll/peel-loop-pgo-small-weigth.ll

+; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll 2>&1 | FileCheck %s
+; RUN: opt < %s -S -debug-only=loop-unroll -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+; Make sure we use the profile information correctly to peel-off 3 iterations
+; from the loop, and update the branch weights for the peeled loop properly.
+
+; CHECK: Loop Unroll: F[basic]
+; CHECK: PEELING loop %for.body with iteration count 4!
+
+; CHECK-LABEL: @basic
+; CHECK: br i1 %{{.*}}, label %[[NEXT0:.*]], label %for.cond.for.end_crit_edge, !prof !1
+; CHECK: [[NEXT0]]:
+; CHECK: br i1 %{{.*}}, label %[[NEXT1:.*]], label %for.cond.for.end_crit_edge, !prof !2
+; CHECK: [[NEXT1]]:
+; CHECK: br i1 %{{.*}}, label %[[NEXT2:.*]], label %for.cond.for.end_crit_edge, !prof !3
+; CHECK: [[NEXT2]]:
+; CHECK: br i1 %{{.*}}, label %[[NEXT3:.*]], label %for.cond.for.end_crit_edge, !prof !4
+; CHECK: [[NEXT3]]:
+; CHECK: br i1 %{{.*}}, label %for.body, label %{{.*}}, !prof !5
+
+define void @basic(i32* %p, i32 %k) #0 !prof !0 {
+entry:
+  %cmp3 = icmp slt i32 0, %k
+  br i1 %cmp3, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.body
+  %i.05 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.body ]
+  %p.addr.04 = phi i32* [ %p, %for.body.lr.ph ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i32, i32* %p.addr.04, i32 1
+  store i32 %i.05, i32* %p.addr.04, align 4
+  %inc = add nsw i32 %i.05, 1
+  %cmp = icmp slt i32 %inc, %k
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge, !prof !1
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  ret void
+}
+
+attributes #0 = { nounwind }
+
+!0 = !{!"function_entry_count", i64 1}
+!1 = !{!"branch_weights", i32 4, i32 1}
+
+;CHECK: !1 = !{!"branch_weights", i32 16, i32 1}
+;CHECK: !2 = !{!"branch_weights", i32 11, i32 5}
+;CHECK: !3 = !{!"branch_weights", i32 5, i32 6}
+;CHECK: !4 = !{!"branch_weights", i32 2, i32 3}
+;CHECK: !5 = !{!"branch_weights", i32 34, i32 2}
+

test/Transforms/LoopUnroll/peel-loop-pgo.ll

 !10 = !{!"DetailedSummary", !11}
 !11 = !{!12, !13, !14}
 !12 = !{i32 10000, i64 3, i32 2}
 !13 = !{i32 999000, i64 1, i32 10}
 !14 = !{i32 999999, i64 1, i32 10}
 !15 = !{!"function_entry_count", i64 1}
 !16 = !{!"branch_weights", i32 3001, i32 1001}
 
-;CHECK: !15 = !{!"branch_weights", i32 900, i32 101}
-;CHECK: !16 = !{!"branch_weights", i32 540, i32 360}
-;CHECK: !17 = !{!"branch_weights", i32 162, i32 378}
-;CHECK: !18 = !{!"branch_weights", i32 1399, i32 162}
+;CHECK: !15 = !{!"branch_weights", i32 901, i32 100}
+;CHECK: !16 = !{!"branch_weights", i32 541, i32 360}
+;CHECK: !17 = !{!"branch_weights", i32 163, i32 378}
+;CHECK: !18 = !{!"branch_weights", i32 1396, i32 163}