Fixed a bug that edge weights are not assigned correctly when lowering switch statement.

[congh]

Summary

This is a one-line-change patch that moves the update to UnhandledWeights to the correct position: it should be updated for all clusters instead of just range clusters.

[hans]

lgtm with some tweaks to the test.

test/CodeGen/X86/switch-edge-weight.ll

6 ↗

(On Diff #33281)

Please add a CHECK-LABEL and move this into or after the @test function.

The alloca, store and load, etc. seem unnecessary. Can the test just switch on the function parameter?

[congh]

test/CodeGen/X86/switch-edge-weight.ll

6 ↗

(On Diff #33281)

Done. I have simplified the function as suggested by you (previously I just copies it from another test case but I should look for a simplest one).

Thank you!

[congh]

Updated the test case.

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

if (I == W.LastCluster) {
       // For the last cluster, fall through to the default destination.
       Fallthrough = DefaultMBB;
     } else {
       Fallthrough = CurMF->CreateMachineBasicBlock(CurMBB->getBasicBlock());
       CurMF->insert(BBI, Fallthrough);
       // Put Cond in a virtual register to make it available from the new blocks.
       ExportFromCurrentBlock(Cond);
     }
+    UnhandledWeights -= I->Weight;
 
     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.
switch (I->Kind) {
           // Check I->Low <= Cond <= I->High.
           CC = ISD::SETLE;
           LHS = I->Low;
           MHS = Cond;
           RHS = I->High;
         }
 
         // The false weight is the sum of all unhandled cases.
-        UnhandledWeights -= I->Weight;
         CaseBlock CB(CC, LHS, RHS, MHS, I->MBB, Fallthrough, CurMBB, I->Weight,
                      UnhandledWeights);
 
         if (CurMBB == SwitchMBB)
           visitSwitchCase(CB, SwitchMBB);
         else
           SwitchCases.push_back(CB);
 

llvm/trunk/test/CodeGen/X86/switch-edge-weight.ll

+; RUN: llc -march=x86-64 -print-machineinstrs=expand-isel-pseudos %s -o /dev/null 2>&1 | FileCheck %s
+
+
+; CHECK-LABEL: test
+
+define void @test(i32 %x) nounwind {
+entry:
+  switch i32 %x, label %sw.default [
+    i32 54, label %sw.bb
+    i32 55, label %sw.bb
+    i32 56, label %sw.bb
+    i32 58, label %sw.bb
+    i32 67, label %sw.bb
+    i32 68, label %sw.bb
+    i32 134, label %sw.bb
+    i32 140, label %sw.bb
+  ], !prof !1
+
+sw.bb:
+  call void @foo(i32 0)
+  br label %sw.epilog
+
+sw.default:
+  call void @foo(i32 1)
+  br label %sw.epilog
+
+sw.epilog:
+  ret void
+
+; Check if weights are correctly assigned to edges generated from switch
+; statement.
+;
+; CHECK: BB#4:
+; CHECK: Successors according to CFG: BB#1(10) BB#6(10)
+; CHECK: BB#6:
+; CHECK: Successors according to CFG: BB#1(10) BB#2(10)
+}
+
+declare void @foo(i32)
+
+!1 = !{!"branch_weights", i32 10, i32 10, i32 10, i32 10, i32 10, i32 10, i32 10, i32 10, i32 10}