Assign weights to edges to jump table when lowering switch statement.

[congh]

Summary

Currently, when lowering switch statement and a new basic block is built for jump table, the edge to this new block is not assigned with a correct weight. This patch collects the edge weight from all its successors and assign this sum of weights to the edge (and also the other fall-through edge). Test cases are adjusted accordingly.

[davidxl]

Can we add more unit test cases (with annotated profile data) to cover profile update for various switch lowering strategies?

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1901 ↗

(On Diff #32598)

Same bug occurs here too. The weight to MBB should be the sum of the ExtraWeight of all cases of 'B'. See visitBitTestCase.

[congh]

In D12166#229003, @davidxl wrote:

Can we add more unit test cases (with annotated profile data) to cover profile update for various switch lowering strategies?

OK. I will investigate on switch lowering and try to make a test case covering all cases with profile data check.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1901 ↗

(On Diff #32598)

You are right. I will also fix this issue in this patch.

[congh]

Update the patch by also adding forgotten edge weight when adding bit test header for lowering switch statement.

[davidxl]

This looks fine to me in general, but other reviewers or Hans may need to take a look.

(The additional test cases can be added as a follow up).

[congh]

In D12166#230433, @davidxl wrote:

This looks fine to me in general, but other reviewers or Hans may need to take a look.

(The additional test cases can be added as a follow up).

OK, thanks! I have added Hans as a reviewer.

[hans]

This looks fine to me, but it would be really nice to get tests here.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1900 ↗

(On Diff #32778)

Does it get the DefaultWeight right here? (Not a criticism, just wondering if this ever worked.)

8019 ↗

(On Diff #32778)

I think the sum of all the branch weights on a branch (in this case a SwitchInst) will not overflow an uint32_t. Since we're adding up weights from the same switch it should be safe to use uint32_t for JumpWeight and skip the ScaleWeights below.

[congh]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1900 ↗

(On Diff #32778)

Yes. getEdgeWeight() will retrieve edge weight through BPI from BB instead of MBB.

8019 ↗

(On Diff #32778)

You are right. I will replace uint64_t with uint32_t and remove ScaleWeights() here. Thanks for pointing it out!

[congh]

Change the type of weight from uint64_t to uint32_t as there won't be overflow in calculation.

Add a test case that checks if the edge weight to the jump table is correct.

[congh]

Fixed a small issue in the newly added test case.

[llvm-commits]

Hi Hans

Could you please take a look at this patch and see if it is ok to be
checked in? Thank you!

Cong

thanks,
Cong

• msg-3525-207.txt162 BDownload

[hans]

Yes, lgtm. Sorry for the delay.

test/CodeGen/X86/switch-jump-table.ll

58 ↗

(On Diff #33137)

This isn't really checking the probabilities though, but the block placement. Maybe that's the best we can do? :-/

Maybe use a more descriptive prefix than CHECK2? CHECK-PLACEMENT?

[congh]

test/CodeGen/X86/switch-jump-table.ll

58 ↗

(On Diff #33137)

As this fix is for lowering IR to MC so I think the best pass that we should indicate here is expand-isel-pseudos. I have updated it. I also changed CHECK2 into CHECK-JT-WEIGHT to make it more descriptive. Thanks!

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h

struct BitTestCase {
     MachineBasicBlock *ThisBB;
     MachineBasicBlock *TargetBB;
     uint32_t ExtraWeight;
   };
 
   typedef SmallVector<BitTestCase, 3> BitTestInfo;
 
   struct BitTestBlock {
-    BitTestBlock(APInt F, APInt R, const Value* SV,
-                 unsigned Rg, MVT RgVT, bool E, bool CR,
-                 MachineBasicBlock* P, MachineBasicBlock* D,
-                 BitTestInfo C):
-      First(F), Range(R), SValue(SV), Reg(Rg), RegVT(RgVT), Emitted(E),
-      ContiguousRange(CR), Parent(P), Default(D), Cases(std::move(C)) { }
+    BitTestBlock(APInt F, APInt R, const Value *SV, unsigned Rg, MVT RgVT,
+                 bool E, bool CR, MachineBasicBlock *P, MachineBasicBlock *D,
+                 BitTestInfo C, uint32_t W)
+        : First(F), Range(R), SValue(SV), Reg(Rg), RegVT(RgVT), Emitted(E),
+          ContiguousRange(CR), Parent(P), Default(D), Cases(std::move(C)),
+          Weight(W) {}
     APInt First;
     APInt Range;
     const Value *SValue;
     unsigned Reg;
     MVT RegVT;
     bool Emitted;
     bool ContiguousRange;
     MachineBasicBlock *Parent;
     MachineBasicBlock *Default;
     BitTestInfo Cases;
+    uint32_t Weight;
   };
 
   /// Minimum jump table density, in percent.
   enum { MinJumpTableDensity = 40 };
 
   /// Check whether a range of clusters is dense enough for a jump table.
   bool isDense(const CaseClusterVector &Clusters, unsigned *TotalCases,
                unsigned First, unsigned Last);

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::visitBitTestHeader(BitTestBlock &B,
   }
 
   B.RegVT = VT.getSimpleVT();
   B.Reg = FuncInfo.CreateReg(B.RegVT);
   SDValue CopyTo = DAG.getCopyToReg(getControlRoot(), dl, B.Reg, Sub);
 
   MachineBasicBlock* MBB = B.Cases[0].ThisBB;
 
-  addSuccessorWithWeight(SwitchBB, B.Default);
-  addSuccessorWithWeight(SwitchBB, MBB);
+  uint32_t DefaultWeight = getEdgeWeight(SwitchBB, B.Default);
+  addSuccessorWithWeight(SwitchBB, B.Default, DefaultWeight);
+  addSuccessorWithWeight(SwitchBB, MBB, B.Weight);
 
   SDValue BrRange = DAG.getNode(ISD::BRCOND, dl,
                                 MVT::Other, CopyTo, RangeCmp,
                                 DAG.getBasicBlock(B.Default));
 
   // Avoid emitting unnecessary branches to the next block.
   if (MBB != NextBlock(SwitchBB))
     BrRange = DAG.getNode(ISD::BR, dl, MVT::Other, BrRange,
bool SelectionDAGBuilder::buildBitTests(CaseClusterVector &Clusters,
 
   for (auto &CB : CBV) {
     MachineBasicBlock *BitTestBB =
         FuncInfo.MF->CreateMachineBasicBlock(SI->getParent());
     BTI.push_back(BitTestCase(CB.Mask, BitTestBB, CB.BB, CB.ExtraWeight));
   }
   BitTestCases.emplace_back(std::move(LowBound), std::move(CmpRange),
                             SI->getCondition(), -1U, MVT::Other, false,
-                            ContiguousRange, nullptr, nullptr, std::move(BTI));
+                            ContiguousRange, nullptr, nullptr, std::move(BTI),
+                            TotalWeight);
 
   BTCluster = CaseCluster::bitTests(Clusters[First].Low, Clusters[Last].High,
                                     BitTestCases.size() - 1, TotalWeight);
   return true;
 }
 
 void SelectionDAGBuilder::findBitTestClusters(CaseClusterVector &Clusters,
                                               const SwitchInst *SI) {
switch (I->Kind) {
       case CC_JumpTable: {
         // FIXME: Optimize away range check based on pivot comparisons.
         JumpTableHeader *JTH = &JTCases[I->JTCasesIndex].first;
         JumpTable *JT = &JTCases[I->JTCasesIndex].second;
 
         // The jump block hasn't been inserted yet; insert it here.
         MachineBasicBlock *JumpMBB = JT->MBB;
         CurMF->insert(BBI, JumpMBB);
-        addSuccessorWithWeight(CurMBB, Fallthrough);
-        addSuccessorWithWeight(CurMBB, JumpMBB);
+
+        // Collect the sum of weights of outgoing edges from JumpMBB, which will
+        // be the edge weight on CurMBB->JumpMBB.
+        uint32_t JumpWeight = 0;
+        for (auto Succ : JumpMBB->successors())
+          JumpWeight += getEdgeWeight(JumpMBB, Succ);
+        uint32_t FallthruWeight = getEdgeWeight(CurMBB, Fallthrough);
+
+        addSuccessorWithWeight(CurMBB, Fallthrough, FallthruWeight);
+        addSuccessorWithWeight(CurMBB, JumpMBB, JumpWeight);
 
         // The jump table header will be inserted in our current block, do the
         // range check, and fall through to our fallthrough block.
         JTH->HeaderBB = CurMBB;
         JT->Default = Fallthrough; // FIXME: Move Default to JumpTableHeader.
 
         // If we're in the right place, emit the jump table header right now.
         if (CurMBB == SwitchMBB) {

llvm/trunk/test/CodeGen/ARM/jump-table-islands.ll

 ; RUN: llc -mtriple=armv7-apple-ios8.0 -o - %s | FileCheck %s
 
 %BigInt = type i5500
 
 define %BigInt @test_moved_jumptable(i1 %tst, i32 %sw, %BigInt %l) {
 ; CHECK-LABEL: test_moved_jumptable:
 
 ; CHECK:   adr {{r[0-9]+}}, [[JUMP_TABLE:LJTI[0-9]+_[0-9]+]]
 ; CHECK:   b [[SKIP_TABLE:LBB[0-9]+_[0-9]+]]
 
 ; CHECK: [[JUMP_TABLE]]:
 ; CHECK:   .data_region jt32
 ; CHECK:   .long LBB{{[0-9]+_[0-9]+}}-[[JUMP_TABLE]]
 
 ; CHECK: [[SKIP_TABLE]]:
-; CHECK:   add pc, {{r[0-9]+}}, {{r[0-9]+}}
+; CHECK:   add pc, lr, {{r[0-9]+}}
   br i1 %tst, label %simple, label %complex
 
 simple:
   br label %end
 
 complex:
   switch i32 %sw, label %simple [ i32 0, label %other
                                   i32 1, label %third

llvm/trunk/test/CodeGen/Mips/nacl-align.ll

 
 
 ; CHECK-LABEL:       test1:
 
 ; CHECK:             .align  4
 ; CHECK-NEXT:    ${{BB[0-9]+_[0-9]+}}:
 ; CHECK-NEXT:        jr      $ra
 ; CHECK-NEXT:        addiu   $2, $zero, 111
-; CHECK-NEXT:        .align  4
 ; CHECK-NEXT:    ${{BB[0-9]+_[0-9]+}}:
 ; CHECK-NEXT:        jr      $ra
-; CHECK-NEXT:        addiu   $2, $zero, 222
+; CHECK-NEXT:        addiu   $2, $zero, 555
 ; CHECK-NEXT:        .align  4
 ; CHECK-NEXT:    ${{BB[0-9]+_[0-9]+}}:
 ; CHECK-NEXT:        jr      $ra
-; CHECK-NEXT:        addiu   $2, $zero, 333
+; CHECK-NEXT:        addiu   $2, $zero, 222
 ; CHECK-NEXT:        .align  4
 ; CHECK-NEXT:    ${{BB[0-9]+_[0-9]+}}:
 ; CHECK-NEXT:        jr      $ra
-; CHECK-NEXT:        addiu   $2, $zero, 444
+; CHECK-NEXT:        addiu   $2, $zero, 333
 
 }
 
 
 ; This test tests that a block whose address is taken is bundle-aligned in NaCl.
 
 @bb_array = constant [2 x i8*] [i8* blockaddress(@test2, %bb1),
                                 i8* blockaddress(@test2, %bb2)], align 4

llvm/trunk/test/CodeGen/X86/switch-jump-table.ll

-; RUN: llc -mtriple=i686-pc-gnu-linux < %s | FileCheck %s
+; RUN: llc -mtriple=i686-pc-gnu-linux < %s | FileCheck %s -check-prefix=CHECK
+; RUN: llc -mtriple=i686-pc-gnu-linux -print-machineinstrs=expand-isel-pseudos %s -o /dev/null 2>&1 | FileCheck %s -check-prefix=CHECK-JT-WEIGHT
 
 
 ; An unreachable default destination is replaced with the most popular case label.
 
-define void @sum2(i32 %x, i32* %to) {
-; CHECK-LABEL: sum2:
+define void @foo(i32 %x, i32* %to) {
+; CHECK-LABEL: foo:
 ; CHECK: movl 4(%esp), [[REG:%e[a-z]{2}]]
 ; CHECK: cmpl $3, [[REG]]
-; CHECK: jbe .LBB0_1
+; CHECK: ja .LBB0_6
+; CHECK-NEXT: # BB#1:
+; CHECK-NEXT: jmpl *.LJTI0_0(,[[REG]],4)
 ; CHECK: movl $4
 ; CHECK: retl
-; CHECK-LABEL: .LBB0_1:
-; CHECK-NEXT: jmpl *.LJTI0_0(,[[REG]],4)
 
 entry:
   switch i32 %x, label %default [
     i32 0, label %bb0
     i32 1, label %bb1
     i32 2, label %bb2
     i32 3, label %bb3
     i32 4, label %bb4
default:
   unreachable
 
 ; The jump table has four entries.
 ; CHECK-LABEL: .LJTI0_0:
 ; CHECK-NEXT: .long  .LBB0_2
 ; CHECK-NEXT: .long  .LBB0_3
 ; CHECK-NEXT: .long  .LBB0_4
 ; CHECK-NEXT: .long  .LBB0_5
-; CHECK-NOT: .long
 }
+
+; Check if branch probabilities are correctly assigned to the jump table.
+
+define void @bar(i32 %x, i32* %to) {
+; CHECK-JT-WEIGHT-LABEL: bar:
+; CHECK-JT-WEIGHT: Successors according to CFG: BB#6(16) BB#8(96)
+; CHECK-JT-WEIGHT: Successors according to CFG: BB#1(16) BB#2(16) BB#3(16) BB#4(16) BB#5(32)
+
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+    i32 5, label %bb4
+  ], !prof !1
+bb0:
+  store i32 0, i32* %to
+  br label %exit
+bb1:
+  store i32 1, i32* %to
+  br label %exit
+bb2:
+  store i32 2, i32* %to
+  br label %exit
+bb3:
+  store i32 3, i32* %to
+  br label %exit
+bb4:
+  store i32 4, i32* %to
+  br label %exit
+default:
+  store i32 5, i32* %to
+  br label %exit
+exit:
+  ret void
+}
+
+!1 = !{!"branch_weights", i32 16, i32 16, i32 16, i32 16, i32 16, i32 16, i32 16}

llvm/trunk/test/CodeGen/X86/switch.ll

 bb3: tail call void @g(i32 3) br label %return
 bb4: tail call void @g(i32 4) br label %return
 return: ret void
 
 ; Cases 0-5 could be lowered with two bit tests,
 ; but with 6-8, the whole switch is suitable for a jump table.
 ; CHECK-LABEL: jt_is_better
 ; CHECK: cmpl $8
-; CHECK: jbe
+; CHECK: ja
 ; CHECK: jmpq *.LJTI
 }
 
 
 define void @bt_is_better(i32 %x) {
 entry:
   switch i32 %x, label %return [
     i32 0, label %bb0

llvm/trunk/test/CodeGen/X86/x86-shrink-wrapping.ll

 ; CHECK: movq 8(%rdi), %rdi
 ; CHECK-NEXT: movzwl (%rdi), %e[[BF_LOAD2:[a-z]+]]
 ; CHECK-NEXT: leal -54(%r[[BF_LOAD2]]), [[TMP:%e[a-z]+]]
 ; CHECK-NEXT: cmpl $14, [[TMP]]
 ; CHECK-NEXT: ja [[LOR_LHS_FALSE:LBB[0-9_]+]]
 ;
 ; CHECK: movl $24599, [[TMP2:%e[a-z]+]]
 ; CHECK-NEXT: btl [[TMP]], [[TMP2]]
-; CHECK-NEXT: jb [[CLEANUP]]
+; CHECK-NEXT: jae [[LOR_LHS_FALSE:LBB[0-9_]+]]
+;
+; CHECK: [[CLEANUP]]: ## %cleanup
+; DISABLE: popq
+; CHECK-NEXT: retq
 ;
 ; CHECK: [[LOR_LHS_FALSE]]: ## %lor.lhs.false
 ; CHECK: cmpl $134, %e[[BF_LOAD2]]
 ; CHECK-NEXT: je [[CLEANUP]]
 ;
 ; CHECK: cmpl $140, %e[[BF_LOAD2]]
 ; CHECK-NEXT: je [[CLEANUP]]
 ;
 ; ENABLE: pushq
 ; CHECK: callq _find_temp_slot_from_address
 ; CHECK-NEXT: testq   %rax, %rax
 ;
 ; The adjustment must use LEA here (or be moved above the test).
 ; ENABLE-NEXT: leaq 8(%rsp), %rsp
 ;
 ; CHECK-NEXT: je [[CLEANUP]]
 ;
 ; CHECK: movb $1, 57(%rax)
-;
-; CHECK: [[CLEANUP]]: ## %cleanup
-; DISABLE: popq
-; CHECK-NEXT: retq
 define void @useLEA(%struct.rtx_def* readonly %x) {
 entry:
   %cmp = icmp eq %struct.rtx_def* %x, null
   br i1 %cmp, label %cleanup, label %if.end
 
 if.end:                                           ; preds = %entry
   %tmp = getelementptr inbounds %struct.rtx_def, %struct.rtx_def* %x, i64 0, i32 0
   %bf.load = load i32, i32* %tmp, align 8