[SimplifyCFG] Fix prof branch_weights MD while removing unreachable switch cases

[yrouban]

Summary

SimplifyCFG has a bug that results in inconsistent prof branch_weights metadata if unreachable switch cases are removed.
This patch fixes this bug by making use of the newly introduced SwitchInstProfBranchWeightsWrapper class (see D62122).
A new test is created.

[yrouban]

@davidxl, @eraman, @reames please review.

[reames]

LGTM w/two required changes.

The second is that the updater itself should check for the right number of metadata nodes and do nothing if already corrupt (as mentioned in another review thread).

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

4217 ↗

(On Diff #200455)

Naming here is non-idiomatic. Please leave SI for the SwitchInst and use SU or something for the updater class.

4219 ↗

(On Diff #200455)

Not related to your change, but there's no reason a switch can't have several cases leading to the same unreachable block. A separate change to handle that would be nice if you wanted.

llvm/trunk/lib/Transforms/Utils/SimplifyCFG.cpp

if (auto *BI = dyn_cast<BranchInst>(TI)) {
         } else if (BI->getSuccessor(1) == BB) {
           Builder.CreateAssumption(Cond);
           Builder.CreateBr(BI->getSuccessor(0));
           EraseTerminatorAndDCECond(BI);
           Changed = true;
         }
       }
     } else if (auto *SI = dyn_cast<SwitchInst>(TI)) {
-      for (auto i = SI->case_begin(), e = SI->case_end(); i != e;) {
+      SwitchInstProfUpdateWrapper SU(*SI);
+      for (auto i = SU->case_begin(), e = SU->case_end(); i != e;) {
         if (i->getCaseSuccessor() != BB) {
           ++i;
           continue;
         }
-        BB->removePredecessor(SI->getParent());
-        i = SI->removeCase(i);
-        e = SI->case_end();
+        BB->removePredecessor(SU->getParent());
+        i = SU.removeCase(i);
+        e = SU->case_end();
         Changed = true;
       }
     } else if (auto *II = dyn_cast<InvokeInst>(TI)) {
       if (II->getUnwindDest() == BB) {
         removeUnwindEdge(TI->getParent());
         Changed = true;
       }
     } else if (auto *CSI = dyn_cast<CatchSwitchInst>(TI)) {

llvm/trunk/test/Transforms/SimplifyCFG/switch-profmd.ll

+; RUN: opt -S -simplifycfg < %s | FileCheck %s
+
+declare i32 @f(i32 %val)
+
+; Check that eliminating cases with unreachable branches keeps
+; prof branch_weights metadata consistent with switch instruction.
+define i32 @test_switch_to_unreachable(i32 %val) {
+; CHECK-LABEL: test_switch_to_unreachable
+; CHECK: switch
+; CHECK-NOT: i32 0, label %on0
+  switch i32 %val, label %otherwise [
+    i32 0, label %on0
+    i32 1, label %on1
+    i32 2, label %on2
+  ], !prof !{!"branch_weights", i32 99, i32 0, i32 1, i32 2}
+; CHECK: !prof ![[MD0:[0-9]+]]
+
+otherwise:
+  %result = call i32 @f(i32 -1)
+  ret i32 %result
+
+on0:
+  unreachable
+  ret i32 125
+
+on1:
+  %result1 = call i32 @f(i32 -2)
+  ret i32 %result1
+
+on2:
+  %result2 = call i32 @f(i32 -3)
+  ret i32 %result2
+}
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 99, i32 2, i32 1}