[CorrelatedValuePropagation] Fix prof branch_weights metadata handling for SwitchInst

[yrouban]

Summary

This patch fixes the CorrelatedValuePropagation pass to keep prof branch_weights metadata of SwitchInst consistent.
It makes use of SwitchInstProfUpdateWrapper introduced in D62122.
A new test is added.

[nikic]

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

398

I think that something can go wrong here... Consider a case where we have a switch with branch weights, then we remove the first case (so that Changed=true in SwitchInstProfBranchWeightsWrapper), then we find that the second case is always true, such that this constant folding call will erase the switch, leaving behind a dangling reference in SwitchInstProfBranchWeightsWrapper, which will be accessed when the destructor runs at the end of the function.

[yrouban]

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

398

Good catch!
ConstantFoldTerminator() can change SwitchInst on its own.

[yrouban]

• renamed SwitchInstProfBranchWeightsWrapper to SwitchInstProfUpdateWrapper
• explicitly narrowed scope of the SwitchInstProfUpdateWrapper object so it does not overlap with ConstantFoldTerminator()
• added a test case where switch turns into conditional branch

[nikic]

This looks ok, but I wonder if it might make sense to have an explicit update() method for cases (like this) where we need to control the point where the SwitchInst is updated. Something like:

void update() {
  if (Changed)
    SI.setMetadata(LLVMContext::MD_prof, buildProfBranchWeightsMD());
  Changed = false;
}
~SwitchInstProfUpdateWrapper() {
  update();
}

I feel like this might be more obvious than extra scopes to control the destruction point.

[yrouban]

In D62126#1518460, @nikic wrote:

I considered this way and found it error prone. The wrapper stays alive after ConstantFoldTerminator() and developer may accidentally reuse it but the stored weights might be unsync with the underlying SwitchInst's weights changed by ConstantFoldTerminator().

[nikic]

LGTM

llvm/test/Transforms/CorrelatedValuePropagation/profmd.ll

50

These check lines should probably be positioned one block higher?

[asbirlea]

FYI, we're seeing a crash in llvm::SwitchInstProfUpdateWrapper::getProfBranchWeights() following this patch.
I'm trying to get a testcase available.

[nikic]

My guess would be that this is due to the assumption of valid branch_weights inside https://github.com/llvm-mirror/llvm/blob/2015c9266833ebae0e42a641252b2b120ef0f77a/lib/IR/Instructions.cpp#L3908, which is not a guarantee we have at this point.

[nikic]

Providing an obviously incorrect branch_weights for one of the added test cases is enough to cause an assertion failure:

define i32 @crash(i32 %s) {
entry:
  %cmp = icmp sgt i32 %s, 0
  br i1 %cmp, label %positive, label %out

positive:
  switch i32 %s, label %out [
  i32 1, label %next
  i32 -1, label %next
  i32 -2, label %next
  ], !prof !{!"branch_weights"} ; not enough elements

out:
  %p = phi i32 [ -1, %entry ], [ 1, %positive ]
  ret i32 %p

next:
  %q = phi i32 [ 0, %positive ], [ 0, %positive ], [ 0, %positive ]
  ret i32 %q
}

[nikic]

Reverted with rL361881 for now.

[asbirlea]

Thank you for the quick action!

[yrouban]

Thank you very much for the revert, the analysis and the test.
I'm going make the wrapper class safe by introducing invalid sticky state.

[yrouban]

In D62126#1520054, @nikic wrote:

Providing an obviously incorrect branch_weights for one of the added test cases is enough to cause an assertion failure:

I have just submitted a fix. Please see D62656. Once it is landed this patch can be landed unchanged.
I decided to not include the crash test case as it would be caught by verification proposed in D61179.

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

 
 /// Simplify a switch instruction by removing cases which can never fire. If the
 /// uselessness of a case could be determined locally then constant propagation
 /// would already have figured it out. Instead, walk the predecessors and
 /// statically evaluate cases based on information available on that edge. Cases
 /// 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) {
+  SwitchInstProfBranchWeightsWrapper SI(*I);
   DomTreeUpdater DTU(*DT, DomTreeUpdater::UpdateStrategy::Lazy);
   Value *Cond = SI->getCondition();
   BasicBlock *BB = SI->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.
   if (isa<Instruction>(Cond) && cast<Instruction>(Cond)->getParent() == BB)
     return false;
for (pred_iterator PI = PB; PI != PE; ++PI) {
         break;
       }
     }
 
     if (State == LazyValueInfo::False) {
       // 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
       // logic.
       Cond = SI->getCondition();
 
       ++NumDeadCases;
       Changed = true;
for (auto CI = SI->case_begin(), CE = SI->case_end(); CI != CE;) {
     // 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
     // by a branch then do so now.
     ConstantFoldTerminator(BB, /*DeleteDeadConditions = */ false,
                            /*TLI = */ nullptr, &DTU);
   return Changed;
 }
 
 // See if we can prove that the given overflow intrinsic will not overflow.
 static bool willNotOverflow(WithOverflowInst *WO, LazyValueInfo *LVI) {
   Value *RHS = WO->getRHS();
   ConstantRange RRange = LVI->getConstantRange(RHS, WO->getParent(), WO);
   ConstantRange NWRegion = ConstantRange::makeGuaranteedNoWrapRegion(

llvm/test/Transforms/CorrelatedValuePropagation/profmd.ll

+; RUN: opt < %s -correlated-propagation -S | FileCheck %s
+
+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
+}
+
+define i32 @switch2(i32 %s) {
+; CHECK-LABEL: @switch2(
+;
+entry:
+  %cmp = icmp sgt i32 %s, 0
+  br i1 %cmp, label %positive, label %out
+
+positive:
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[S:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[POSITIVE:%.*]], label [[OUT:%.*]]
+  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
+}
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 99, i32 4, i32 3}