Update edge weights properly when merging blocks in if-conversion.

[congh]

Summary

In if-conversion, there is a utility function MergeBlocks() that is used to merge blocks. However, when new edges are built in this function the edge weight is either not provided or not updated properly, leading to a modified CFG with incorrect edge weights. This patch corrects this issue. More information is provided in the comments.

[davidxl]

I'd like profile update to be isolated in its own helper function. Have some examples in the comments (like in the test case) -- the example should show cfg before and after. The parameters also need to be documented.

lib/CodeGen/IfConversion.cpp

1701 ↗

(On Diff #33655)

Should it be zero?

1736 ↗

(On Diff #33655)

Do you have example of this case?

[congh]

In D12513#237639, @davidxl wrote:

I'd like profile update to be isolated in its own helper function. Have some examples in the comments (like in the test case) -- the example should show cfg before and after. The parameters also need to be documented.

It is difficult to isolate profile update here as the profile data is needed when adding successors.

I have added a CFG example to the comment in this function. Some variables are also documented.

lib/CodeGen/IfConversion.cpp

1701 ↗

(On Diff #33655)

I made a mistake here: I should get To2FromWeight before the edge weight is set to zero. Need to switch this statement with the one above.

1736 ↗

(On Diff #33655)

Yes. The test case in test/CodeGen/ARM/ifcvt-iter-indbr.ll covers this case.

[congh]

Fixed a bug in the patch and added more comments.

[congh]

Ping?

[hans]

Looks ok to me. The code is pretty hard to follow, but I don't have any great ideas for making it easier.

lib/CodeGen/IfConversion.cpp

1689 ↗

(On Diff #33729)

I know this is old code, but while you're here, maybe change this to a SmallVector and rename to FromSuccs?

test/CodeGen/ARM/ifcvt-iter-indbr.ll

2 ↗

(On Diff #33729)

Should 2<&1 be 2>&1?

[congh]

In D12513#248922, @hans wrote:

Looks ok to me. The code is pretty hard to follow, but I don't have any great ideas for making it easier.

Thanks for the review! Yes, this part is not easy to understand but I also could not find a way to improve the overall readability.

lib/CodeGen/IfConversion.cpp

1689 ↗

(On Diff #33729)

Done. I found many uses of std::vector in this file. Do we need to change all of them to SmallVector?

test/CodeGen/ARM/ifcvt-iter-indbr.ll

2 ↗

(On Diff #33729)

Right.. corrected.

[hans]

lib/CodeGen/IfConversion.cpp

1689 ↗

(On Diff #33729)

You could look at them in a separate patch if you like, but it's probably not very important.

I figured this one would be nice to fix when you're editing the function already.

llvm/trunk/lib/CodeGen/IfConversion.cpp

 /// ToBBI.
 void IfConverter::MergeBlocks(BBInfo &ToBBI, BBInfo &FromBBI, bool AddEdges) {
   assert(!FromBBI.BB->hasAddressTaken() &&
          "Removing a BB whose address is taken!");
 
   ToBBI.BB->splice(ToBBI.BB->end(),
                    FromBBI.BB, FromBBI.BB->begin(), FromBBI.BB->end());
 
-  std::vector<MachineBasicBlock *> Succs(FromBBI.BB->succ_begin(),
+  SmallVector<MachineBasicBlock *, 4> FromSuccs(FromBBI.BB->succ_begin(),
                                                 FromBBI.BB->succ_end());
   MachineBasicBlock *NBB = getNextBlock(FromBBI.BB);
   MachineBasicBlock *FallThrough = FromBBI.HasFallThrough ? NBB : nullptr;
 
-  for (unsigned i = 0, e = Succs.size(); i != e; ++i) {
-    MachineBasicBlock *Succ = Succs[i];
+  // The edge weight from ToBBI.BB to FromBBI.BB, which is only needed when
+  // AddEdges is true and FromBBI.BB is a successor of ToBBI.BB.
+  uint32_t To2FromWeight = 0;
+  // WeightScale and SumWeight are for calculating successor probabilities of
+  // FromBBI.BB.
+  uint32_t WeightScale = 0;
+  uint32_t SumWeight = 0;
+  if (AddEdges && ToBBI.BB->isSuccessor(FromBBI.BB)) {
+    To2FromWeight = MBPI->getEdgeWeight(ToBBI.BB, FromBBI.BB);
+    // Set the edge weight from ToBBI.BB to FromBBI.BB to zero to avoid the edge
+    // weight being merged to other edges when this edge is removed later.
+    ToBBI.BB->setSuccWeight(
+        std::find(ToBBI.BB->succ_begin(), ToBBI.BB->succ_end(), FromBBI.BB), 0);
+    SumWeight = MBPI->getSumForBlock(FromBBI.BB, WeightScale);
+  }
+
+  for (unsigned i = 0, e = FromSuccs.size(); i != e; ++i) {
+    MachineBasicBlock *Succ = FromSuccs[i];
     // Fallthrough edge can't be transferred.
     if (Succ == FallThrough)
       continue;
+
+    uint32_t NewWeight = 0;
+    if (AddEdges) {
+      // Calculate the edge weight for the edge from ToBBI.BB to Succ, which is
+      // a portion of the edge weight from FromBBI.BB to Succ. The portion ratio
+      // is the edge probability from ToBBI.BB to FromBBI.BB (if FromBBI is a
+      // successor of ToBBI.BB. See comment below for excepion).
+      NewWeight = MBPI->getEdgeWeight(FromBBI.BB, Succ);
+
+      // To2FromWeight is 0 when FromBBI.BB is not a successor of ToBBI.BB. This
+      // only happens when if-converting a diamond CFG and FromBBI.BB is the
+      // tail BB.  In this case FromBBI.BB post-dominates ToBBI.BB and hence we
+      // could just use the weights on FromBBI.BB's out-edges when adding new
+      // successors.
+      if (To2FromWeight > 0) {
+        BranchProbability Prob(NewWeight / WeightScale, SumWeight);
+        NewWeight = Prob.scale(To2FromWeight);
+      }
+    }
+
     FromBBI.BB->removeSuccessor(Succ);
-    if (AddEdges && !ToBBI.BB->isSuccessor(Succ))
-      ToBBI.BB->addSuccessor(Succ);
+
+    if (AddEdges) {
+      // If the edge from ToBBI.BB to Succ already exists, update the weight of
+      // this edge by adding NewWeight to it. An example is shown below, in
+      // which A is ToBBI.BB and B is FromBBI.BB. In this case we don't have to
+      // set C as A's successor as it already is. We only need to update the
+      // edge weight on A->C. Note that B will not be immediately removed from
+      // A's successors. It is possible that B->D is not removed either if D is
+      // a fallthrough of B. Later the edge A->D (generated here) and B->D will
+      // be combined into one edge. To maintain correct edge weight of this
+      // combined edge, we need to set the edge weight of A->B to zero, which is
+      // already done above. The edge weight on A->D is calculated by scaling
+      // the original weight on A->B by the probability of B->D.
+      //
+      // Before ifcvt:      After ifcvt (assume B->D is kept):
+      //
+      //       A                A
+      //      /|               /|\
+      //     / B              / B|
+      //    | /|             |  ||
+      //    |/ |             |  |/
+      //    C  D             C  D
+      //
+      if (ToBBI.BB->isSuccessor(Succ))
+        ToBBI.BB->setSuccWeight(
+            std::find(ToBBI.BB->succ_begin(), ToBBI.BB->succ_end(), Succ),
+            MBPI->getEdgeWeight(ToBBI.BB, Succ) + NewWeight);
+      else
+        ToBBI.BB->addSuccessor(Succ, NewWeight);
+    }
   }
 
   // Now FromBBI always falls through to the next block!
   if (NBB && !FromBBI.BB->isSuccessor(NBB))
     FromBBI.BB->addSuccessor(NBB);
 
   ToBBI.Predicate.append(FromBBI.Predicate.begin(), FromBBI.Predicate.end());
   FromBBI.Predicate.clear();

llvm/trunk/test/CodeGen/ARM/ifcvt-iter-indbr.ll

 ; RUN: llc < %s -mtriple thumbv7s-apple-darwin  -asm-verbose=false | FileCheck %s
+; RUN: llc < %s -mtriple thumbv7s-apple-darwin  -asm-verbose=false -print-machineinstrs=if-converter 2>&1 | FileCheck --check-prefix=CHECK-WEIGHT %s
 
 declare i32 @foo(i32)
 declare i8* @bar(i32, i8*, i8*)
 
 ; Verify that we don't try to iteratively re-ifconvert a block with a
 ; (predicated) indirectbr terminator.
 ; If we do, we would ignore its fallthrough successor.
 

 ; CHECK-NEXT: LBB0_2:
 ; CHECK-NEXT:  movw r0, #1234
 ; CHECK-NEXT:  b [[FOOCALL:LBB[0-9_]+]]
 ; CHECK-NEXT: Ltmp
 ; CHECK-NEXT: LBB{{[0-9_]+}}:
 ; CHECK-NEXT:  movw r0, #4567
 ; CHECK-NEXT: [[FOOCALL]]:
 ; CHECK-NEXT:  blx _foo
+;
+; CHECK-WEIGHT: BB#0:
+; CHECK-WEIGHT: Successors according to CFG: BB#1(16) BB#2(8) BB#4(8)
+; CHECK-WEIGHT: BB#1:
+; CHECK-WEIGHT: Successors according to CFG: BB#2(24) BB#4(8)
 
 define i32 @test(i32 %a, i32 %a2, i32* %p, i32* %p2) {
 entry:
   %dst1 = call i8* @bar(i32 1, i8* blockaddress(@test, %bb1), i8* blockaddress(@test, %bb2))
   %dst2 = call i8* @bar(i32 2, i8* blockaddress(@test, %bb1), i8* blockaddress(@test, %bb2))
   %dst3 = call i8* @bar(i32 3, i8* blockaddress(@test, %bb1), i8* blockaddress(@test, %bb2))
   %cc1 = icmp eq i32 %a, 21
   br i1 %cc1, label %cc1t, label %cc1f

llvm/trunk/test/CodeGen/Hexagon/ifcvt-edge-weight.ll

+; RUN: llc -march=hexagon -mcpu=hexagonv5  -print-machineinstrs=if-converter %s -o /dev/null 2>&1 | FileCheck %s
+; Check that the edge weights are updated correctly after if-conversion.
+
+; CHECK: BB#3:
+; CHECK: Successors according to CFG: BB#2(10) BB#1(90)
+@a = external global i32
+@d = external global i32
+
+; In the following CFG, A,B,C,D will be if-converted into a single block.
+; Check if the edge weights on edges to E and F are maintained correctly.
+;
+;    A
+;   / \
+;  B   C
+;   \ /
+;    D
+;   / \
+;  E   F
+;
+define void @test1(i8 zeroext %la, i8 zeroext %lb) {
+entry:
+  %cmp0 = call i1 @pred()
+  br i1 %cmp0, label %if.else2, label %if.then0, !prof !1
+
+if.else2:
+  call void @bar(i32 2)
+  br label %if.end2
+
+if.end2:
+  call void @foo(i32 2)
+  br label %return
+
+if.end:
+  %storemerge = phi i32 [ %and, %if.else ], [ %shl, %if.then ]
+  store i32 %storemerge, i32* @a, align 4
+  %0 = load i32, i32* @d, align 4
+  %cmp2 = call i1 @pred()
+  br i1 %cmp2, label %if.end2, label %if.else2, !prof !2
+
+if.then0:
+  %cmp = icmp eq i8 %la, %lb
+  br i1 %cmp, label %if.then, label %if.else, !prof !1
+
+if.then:
+  %conv1 = zext i8 %la to i32
+  %shl = shl nuw nsw i32 %conv1, 16
+  br label %if.end
+
+if.else:
+  %and8 = and i8 %lb, %la
+  %and = zext i8 %and8 to i32
+  br label %if.end
+
+return:
+  call void @foo(i32 2)
+  ret void
+}
+
+declare void @foo(i32)
+declare void @bar(i32)
+declare i1 @pred()
+
+!1 = !{!"branch_weights", i32 80, i32 20}
+!2 = !{!"branch_weights", i32 10, i32 90}