Resubmit "[CorrelatedValuePropagation] Fix prof branch_weights metadata handling for SwitchInst"

This reverts commit 5b32f60ec3.
The fix is in commit 4f9e68148b.

This patch fixes the CorrelatedValuePropagation pass to keep
prof branch_weights metadata of SwitchInst consistent.
It makes use of SwitchInstProfUpdateWrapper.
New tests are added.

Reviewed By: nikic
Differential Revision: https://reviews.llvm.org/D62126

llvm-svn: 362583

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

@@ -308,11 +308,11 @@ static bool processCmp(CmpInst *Cmp, LazyValueInfo *LVI) {
 /// that cannot fire no matter what the incoming edge can safely be removed. If
 /// a case fires on every incoming edge then the entire switch can be removed
 /// and replaced with a branch to the case destination.
-static bool processSwitch(SwitchInst *SI, LazyValueInfo *LVI,
+static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
                           DominatorTree *DT) {
   DomTreeUpdater DTU(*DT, DomTreeUpdater::UpdateStrategy::Lazy);
-  Value *Cond = SI->getCondition();
+  Value *Cond = I->getCondition();
-  BasicBlock *BB = SI->getParent();
+  BasicBlock *BB = I->getParent();
 
   // If the condition was defined in same block as the switch then LazyValueInfo
   // currently won't say anything useful about it, though in theory it could.
@@ -329,6 +329,10 @@ static bool processSwitch(SwitchInst *SI, LazyValueInfo *LVI,
   for (auto *Succ : successors(BB))
     SuccessorsCount[Succ]++;
 
+  { // Scope for SwitchInstProfUpdateWrapper. It must not live during
+    // ConstantFoldTerminator() as the underlying SwitchInst can be changed.
+    SwitchInstProfUpdateWrapper SI(*I);
+
     for (auto CI = SI->case_begin(), CE = SI->case_end(); CI != CE;) {
       ConstantInt *Case = CI->getCaseValue();
 
@@ -365,7 +369,7 @@ static bool processSwitch(SwitchInst *SI, LazyValueInfo *LVI,
         // This case never fires - remove it.
         BasicBlock *Succ = CI->getCaseSuccessor();
         Succ->removePredecessor(BB);
-      CI = SI->removeCase(CI);
+        CI = SI.removeCase(CI);
         CE = SI->case_end();
 
         // The condition can be modified by removePredecessor's PHI simplification
@@ -391,6 +395,7 @@ static bool processSwitch(SwitchInst *SI, LazyValueInfo *LVI,
       // Increment the case iterator since we didn't delete it.
       ++CI;
     }
+  }
 
   if (Changed)
     // If the switch has been simplified to the point where it can be replaced

llvm/test/Transforms/CorrelatedValuePropagation/profmd.ll

@@ -0,0 +1,119 @@
+; RUN: opt < %s -correlated-propagation -S | FileCheck %s
+
+; Removed several cases from switch.
+define i32 @switch1(i32 %s) {
+; CHECK-LABEL: @switch1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[S:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[NEGATIVE:%.*]], label [[OUT:%.*]]
+;
+entry:
+  %cmp = icmp slt i32 %s, 0
+  br i1 %cmp, label %negative, label %out
+
+negative:
+; CHECK:       negative:
+; CHECK-NEXT:    switch i32 [[S]], label [[OUT]] [
+; CHECK-NEXT:    i32 -2, label [[NEXT:%.*]]
+; CHECK-NEXT:    i32 -1, label [[NEXT]]
+  switch i32 %s, label %out [
+  i32 0, label %out
+  i32 1, label %out
+  i32 -1, label %next
+  i32 -2, label %next
+  i32 2, label %out
+  i32 3, label %out
+; CHECK-NEXT: !prof ![[MD0:[0-9]+]]
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6}
+
+out:
+  %p = phi i32 [ 1, %entry ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ]
+  ret i32 %p
+
+next:
+  %q = phi i32 [ 0, %negative ], [ 0, %negative ]
+  ret i32 %q
+}
+
+; Removed all cases from switch.
+define i32 @switch2(i32 %s) {
+; CHECK-LABEL: @switch2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[S:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[POSITIVE:%.*]], label [[OUT:%.*]]
+;
+entry:
+  %cmp = icmp sgt i32 %s, 0
+  br i1 %cmp, label %positive, label %out
+
+positive:
+  switch i32 %s, label %out [
+  i32 0, label %out
+  i32 -1, label %next
+  i32 -2, label %next
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2, i32 3}
+
+out:
+  %p = phi i32 [ -1, %entry ], [ 1, %positive ], [ 1, %positive ]
+  ret i32 %p
+
+next:
+  %q = phi i32 [ 0, %positive ], [ 0, %positive ]
+  ret i32 %q
+}
+
+; Change switch into conditional branch.
+define i32 @switch3(i32 %s) {
+; CHECK-LABEL: @switch3(
+;
+entry:
+  %cmp = icmp sgt i32 %s, 0
+  br i1 %cmp, label %positive, label %out
+
+positive:
+; CHECK:      positive:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 %s, 1
+; CHECK-NEXT:    br i1 [[CMP]], label [[NEXT:%.*]], label [[OUT:%.*]], !prof ![[MD1:[0-9]+]]
+  switch i32 %s, label %out [
+  i32 1, label %next
+  i32 -1, label %next
+  i32 -2, label %next
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2, i32 3}
+
+out:
+  %p = phi i32 [ -1, %entry ], [ 1, %positive ]
+  ret i32 %p
+
+next:
+  %q = phi i32 [ 0, %positive ], [ 0, %positive ], [ 0, %positive ]
+  ret i32 %q
+}
+
+; Removed all cases from switch.
+define i32 @switch4(i32 %s) {
+; CHECK-LABEL: @switch4(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[S:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[NEGATIVE:%.*]], label [[OUT:%.*]]
+;
+entry:
+  %cmp = icmp slt i32 %s, 0
+  br i1 %cmp, label %negative, label %out
+
+negative:
+; CHECK:       negative:
+; CHECK-NEXT:    br label %out
+  switch i32 %s, label %out [
+  i32 0, label %out
+  i32 1, label %out
+  i32 2, label %out
+  i32 3, label %out
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2, i32 3, i32 4}
+
+out:
+  %p = phi i32 [ 1, %entry ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ], [ -1, %negative ]
+  ret i32 %p
+}
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 99, i32 4, i32 3}
+; CHECK: ![[MD1]] = !{!"branch_weights", i32 1, i32 99}