[SimpleLoopUnswitch] Discard stale prof branch_weights for partial unswitched branches

[yrouban]

Summary

In https://reviews.llvm.org/D60606#1464782 Philip Reames pointed out that partial unswitched branches lose their profile information.
This patch fixes SimpleLoopUnswitch pass to make it happen: to not reuse the stale perf branch_weights if any but explicitly discard them for partially unswitched branches.
Couple of tests are fixed to check the changes both for trivial and non-trivial unswitch.

[reames]

The change description and your tests make it hard to determine what your intended result is. Please fix.

Please update your change description to explain *why* stripping metadata is the best answer. A small example might help. Or alternatively, turn the small examples into (new) test cases and reference them from the proposed commit message.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

490

I believe the same problem exists for full unswitch as well, and thus this change should not be conditional.

As an example, consider:
loop {

if (A) break
if (B) break;

}

Where we're unswitching B, but not A for some reason. The frequency on the new B branch could be quite different depending on whether A and B are correlated.

2013

I'd suggest moving this to right after the branch successor rewriting.

Also, there is a bug here if this function is called w/BI null and SI non null.

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll

2703

Please separate and commit your test changes, then rebase over them. At the moment, I can't actually tell what's changing vs what are changes to the tests.

note that you should make sure to have tests both with and without profile weights. You don't need to (and should not) fully duplicate, but at least a few of each are required.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

if (ExitDirection)
                  Instruction::Or &&
              "Must have an `or` of `i1`s for the condition!");
     else
       assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
                  Instruction::And &&
              "Must have an `and` of `i1`s for the condition!");
     buildPartialUnswitchConditionalBranch(*OldPH, Invariants, ExitDirection,
                                           *UnswitchedBB, *NewPH);
+    // If there is any prof branch_weights metadata then discard them because
+    // they are no longer valid for the split condition.
+    BI.setMetadata(LLVMContext::MD_prof, nullptr);
   }
 
   // Update the dominator tree with the added edge.
   DT.insertEdge(OldPH, UnswitchedBB);
 
   // After the dominator tree was updated with the added edge, update MemorySSA
   // if available.
   if (MSSAU) {
static void unswitchNontrivialInvariants(
   // between the unswitched versions, and we will have a new preheader for the
   // original loop.
   BasicBlock *SplitBB = L.getLoopPreheader();
   BasicBlock *LoopPH = SplitEdge(SplitBB, L.getHeader(), &DT, &LI, MSSAU);
 
   // Keep track of the dominator tree updates needed.
   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
 
+  // Before cloning for partial unswitch if there is any prof branch_weights
+  // metadata then discard them because they are no longer valid for the split
+  // condition.
+  if (!FullUnswitch)
+    BI->setMetadata(LLVMContext::MD_prof, nullptr);
+
   // Clone the loop for each unswitched successor.
   SmallVector<std::unique_ptr<ValueToValueMapTy>, 4> VMaps;
   VMaps.reserve(UnswitchedSuccBBs.size());
   SmallDenseMap<BasicBlock *, BasicBlock *, 4> ClonedPHs;
   for (auto *SuccBB : UnswitchedSuccBBs) {
     VMaps.emplace_back(new ValueToValueMapTy());
     ClonedPHs[SuccBB] = buildClonedLoopBlocks(
         L, LoopPH, SplitBB, ExitBlocks, ParentBB, SuccBB, RetainedSuccBB,

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll

loop_exit:
   ret i32 0
 ; CHECK:       loop_exit:
 ; CHECK-NEXT:    ret
 }
 
 ; Non-trivial partial loop unswitching of an invariant input to an 'or'.
 define i32 @test25(i1* %ptr, i1 %cond) {
 ; CHECK-LABEL: @test25(
 entry:
   br label %loop_begin
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 %cond, label %entry.split.us, label %entry.split
 
 loop_begin:
   %v1 = load i1, i1* %ptr
   %cond_or = or i1 %v1, %cond
-  br i1 %cond_or, label %loop_a, label %loop_b
+  br i1 %cond_or, label %loop_a, label %loop_b, !prof !{!"branch_weights", i32 1, i32 2}
 
 loop_a:
   call i32 @a()
   br label %latch
 ; The 'loop_a' unswitched loop.
 ;
 ; CHECK:       entry.split.us:
 ; CHECK-NEXT:    br label %loop_begin.us
 ;
 ; CHECK:       loop_begin.us:
 ; CHECK-NEXT:    %[[V1_US:.*]] = load i1, i1* %ptr
 ; CHECK-NEXT:    %[[OR_US:.*]] = or i1 %[[V1_US]], true
 ; CHECK-NEXT:    br label %loop_a.us
 ;
 ; CHECK:       loop_a.us:
 ; CHECK-NEXT:    call i32 @a()
 ; CHECK-NEXT:    br label %latch.us
 ;
 ; CHECK:       latch.us:
 ; CHECK-NEXT:    %[[V2_US:.*]] = load i1, i1* %ptr
-; CHECK-NEXT:    br i1 %[[V2_US]], label %loop_begin.us, label %loop_exit.split.us
+; CHECK-NEXT:    br i1 %[[V2_US]], label %loop_begin.us, label %loop_exit.split.us, !prof ![[MD0:[0-9]+]]
 ;
 ; CHECK:       loop_exit.split.us:
 ; CHECK-NEXT:    br label %loop_exit
 
 loop_b:
   call i32 @b()
   br label %latch
 ; The original loop.
 ;
 ; CHECK:       entry.split:
 ; CHECK-NEXT:    br label %loop_begin
 ;
 ; CHECK:       loop_begin:
 ; CHECK-NEXT:    %[[V1:.*]] = load i1, i1* %ptr
 ; CHECK-NEXT:    %[[OR:.*]] = or i1 %[[V1]], false
-; CHECK-NEXT:    br i1 %[[OR]], label %loop_a, label %loop_b
+; CHECK-NEXT:    br i1 %[[OR]], label %loop_a, label %loop_b{{$}}
 ;
 ; CHECK:       loop_a:
 ; CHECK-NEXT:    call i32 @a()
 ; CHECK-NEXT:    br label %latch
 ;
 ; CHECK:       loop_b:
 ; CHECK-NEXT:    call i32 @b()
 ; CHECK-NEXT:    br label %latch
 
 latch:
   %v2 = load i1, i1* %ptr
-  br i1 %v2, label %loop_begin, label %loop_exit
+  br i1 %v2, label %loop_begin, label %loop_exit, !prof !{!"branch_weights", i32 3, i32 4}
 ; CHECK:       latch:
 ; CHECK-NEXT:    %[[V2:.*]] = load i1, i1* %ptr
-; CHECK-NEXT:    br i1 %[[V2]], label %loop_begin, label %loop_exit.split
+; CHECK-NEXT:    br i1 %[[V2]], label %loop_begin, label %loop_exit.split, !prof ![[MD0:[0-9]+]]
 
 loop_exit:
   ret i32 0
 ; CHECK:       loop_exit.split:
 ; CHECK-NEXT:    br label %loop_exit
 ;
 ; CHECK:       loop_exit:
 ; CHECK-NEXT:    ret
outer.latch:
   %outer.cond = call i1 @cond()
   br i1 %outer.cond, label %outer.header, label %exit
 
 exit:
   ret void
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 }
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 3, i32 4}

llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

 ; CHECK-NEXT:    ret void
 }
 
 define i32 @test_partial_condition_unswitch_and(i32* %var, i1 %cond1, i1 %cond2) {
 ; CHECK-LABEL: @test_partial_condition_unswitch_and(
 entry:
   br label %loop_begin
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    br i1 %cond1, label %entry.split, label %loop_exit.split
+; CHECK-NEXT:    br i1 %cond1, label %entry.split, label %loop_exit.split, !prof ![[MD0:[0-9]+]]
 ;
 ; CHECK:       entry.split:
-; CHECK-NEXT:    br i1 %cond2, label %entry.split.split, label %loop_exit
+; CHECK-NEXT:    br i1 %cond2, label %entry.split.split, label %loop_exit.split1{{$}}
 ;
 ; CHECK:       entry.split.split:
 ; CHECK-NEXT:    br label %loop_begin
 
 loop_begin:
-  br i1 %cond1, label %continue, label %loop_exit
+  br i1 %cond1, label %continue, label %loop_exit, !prof !{!"branch_weights", i32 1, i32 2}
 ; CHECK:       loop_begin:
 ; CHECK-NEXT:    br label %continue
 
 continue:
   %var_val = load i32, i32* %var
   %var_cond = trunc i32 %var_val to i1
   %cond_and = and i1 %var_cond, %cond2
-  br i1 %cond_and, label %do_something, label %loop_exit
+  br i1 %cond_and, label %do_something, label %loop_exit, !prof !{!"branch_weights", i32 3, i32 4}
 ; CHECK:       continue:
 ; CHECK-NEXT:    %[[VAR:.*]] = load i32
 ; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
 ; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 %[[VAR_COND]], true
-; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit
+; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit{{$}}
 
 do_something:
   call void @some_func() noreturn nounwind
   br label %loop_begin
 ; CHECK:       do_something:
 ; CHECK-NEXT:    call
 ; CHECK-NEXT:    br label %loop_begin
 

 ; CHECK-NEXT:    call void @sink(i32 %[[PHI]])
 ; CHECK-NEXT:    br label %header
 
 loopexit:
   ret void
 ; CHECK:       loopexit:
 ; CHECK-NEXT:    ret
 }
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 1, i32 2}