[SimpleLoopUnswitch] Implement handling of prof branch_weights metadata for SwitchInst

[yrouban]

Summary

This patch implements a handling of prof branch_weights metadata for switch instructions.
New tests are created.

[reames]

The code change looks reasonable to me, but I'd like to let Chandler confirm since he's more familiar with the code.

You definitely need a few more tests.

[reames]

Two additional comments:

1. We appear to have the same lost profile problem for partial unswitching of branches.
2. When I wrote my early comment, I didn't realize this was sitting on top of proposed API changes. I'd suggest focusing on the earlier patches, then wan return to this one.

[yrouban]

In D60606#1464782, @reames wrote:

1. We appear to have the same lost profile problem for partial unswitching of branches.

Yes. But I believe that the fix would be unrelated to this patch. I do not see any way to fix branch_weights for partially unswitched branch instruction other than to drop the prof branch_weights metadata.

[yrouban]

added more test checks to trivial-unswitch.ll to cover all new and changed code

[yrouban]

• made use of SwitchInstProfBranchWeightsWrapper introduced in D62122.
• created new profmd tests instead of changing the existing tests

[yrouban]

@asbirlea, @fedor.sergeev, @mkazantsev, @chandlerc could you please review the patch or suggest somebody who can take a look?

[asbirlea]

I'd prefer to have Chandler approve this, since I'm not familiar with what you're trying to do here.
Some minor comments inline.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

746

Does this compile?

787–789

I find it somewhat confusing that this is making changes both through SI and SIW. But I don't really understand the SwitchInstProfBranchWeightsWrapper, so perhaps it's WAI.

[yrouban]

Rebased to the latest HEAD. Please review. Its sibling patches have been landed.

[reames]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

602–603

This bit appears unnecessarily complex. In particular, you can grab the default weight regardless of whether it's used or not later.

679

This is not idiomatic. Please leave NewSI as it was, then define a separate updater object which wraps that.

759

Hm, I think there's a potential reorg of the existing code which might make this easy to follow. Is there any reason we can't fully construct the switch, and then check if all successors are common? The checking early, transform late structure just seems oddly complex for no obvious reason?

llvm/test/Transforms/SimpleLoopUnswitch/basictest-profmd.ll

4 ↗

(On Diff #205507)

If this doesn't unswitch, what is it testing?

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

6 ↗

(On Diff #205507)

Given the control flow, you need at least one test case where we unswitch the default only, one where we unswitch cases only, and one where we do a combination of both.

[yrouban]

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

6 ↗

(On Diff #205507)

This advice revealed one case I missed to set weight for the default switch case. There were no branch in the original code for this case.

[yrouban]

Addressed most comments. The main changes are the following:

• added more cases to trivial-unswitch-profmd.ll
• brushed up basictest-profmd.ll
• added NewSI default case weight if the default successor was not an exit

[reames]

LGTM. The flow through here is still confusing, but it's a lot better than it was. I'm fine seeing this landed, and then I may take a shot at an NFC to clarify what's going on.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
   SmallVector<int, 4> ExitCaseIndices;
   for (auto Case : SI.cases()) {
     auto *SuccBB = Case.getCaseSuccessor();
     if (!L.contains(SuccBB) &&
         areLoopExitPHIsLoopInvariant(L, *ParentBB, *SuccBB))
       ExitCaseIndices.push_back(Case.getCaseIndex());
   }
   BasicBlock *DefaultExitBB = nullptr;
+  SwitchInst::CaseWeightOpt DefaultExitWeight;
   if (!L.contains(SI.getDefaultDest()) &&
       areLoopExitPHIsLoopInvariant(L, *ParentBB, *SI.getDefaultDest()) &&
-      !isa<UnreachableInst>(SI.getDefaultDest()->getTerminator()))
+      !isa<UnreachableInst>(SI.getDefaultDest()->getTerminator())) {
     DefaultExitBB = SI.getDefaultDest();
-  else if (ExitCaseIndices.empty())
+    DefaultExitWeight = SI.getDefaultCaseWeight();
+  } else if (ExitCaseIndices.empty())
     return false;
 
   LLVM_DEBUG(dbgs() << "    unswitching trivial switch...\n");
 
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
 
   // We may need to invalidate SCEVs for the outermost loop reached by any of
   // the exits.
   Loop *OuterL = &L;
 
   if (DefaultExitBB) {
     // Clear out the default destination temporarily to allow accurate
     // predecessor lists to be examined below.
     SI.setDefaultDest(nullptr);
     // Check the loop containing this exit.
     Loop *ExitL = LI.getLoopFor(DefaultExitBB);
     if (!ExitL || ExitL->contains(OuterL))
       OuterL = ExitL;
   }
 
   // Store the exit cases into a separate data structure and remove them from
   // the switch.
-  SmallVector<std::pair<ConstantInt *, BasicBlock *>, 4> ExitCases;
+  SmallVector<std::tuple<ConstantInt *, BasicBlock *, SwitchInst::CaseWeightOpt>, 4> ExitCases;
   ExitCases.reserve(ExitCaseIndices.size());
   // We walk the case indices backwards so that we remove the last case first
   // and don't disrupt the earlier indices.
   for (unsigned Index : reverse(ExitCaseIndices)) {
     auto CaseI = SI.case_begin() + Index;
     // Compute the outer loop from this exit.
     Loop *ExitL = LI.getLoopFor(CaseI->getCaseSuccessor());
     if (!ExitL || ExitL->contains(OuterL))
       OuterL = ExitL;
     // Save the value of this case.
-    ExitCases.push_back({CaseI->getCaseValue(), CaseI->getCaseSuccessor()});
+    ExitCases.emplace_back(CaseI->getCaseValue(),
+                           CaseI->getCaseSuccessor(),
+                           SI.getSuccessorWeight(CaseI->getSuccessorIndex()));
     // Delete the unswitched cases.
     SI.removeCase(CaseI);
   }
 
   if (SE) {
     if (OuterL)
       SE->forgetLoop(OuterL);
     else
static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
 
   // Split the preheader, so that we know that there is a safe place to insert
   // the switch.
   BasicBlock *OldPH = L.getLoopPreheader();
   BasicBlock *NewPH = SplitEdge(OldPH, L.getHeader(), &DT, &LI, MSSAU);
   OldPH->getTerminator()->eraseFromParent();
 
   // Now add the unswitched switch.
   auto *NewSI = SwitchInst::Create(LoopCond, NewPH, ExitCases.size(), OldPH);
 
   // Rewrite the IR for the unswitched basic blocks. This requires two steps.
   // First, we split any exit blocks with remaining in-loop predecessors. Then
   // we update the PHIs in one of two ways depending on if there was a split.
   // We walk in reverse so that we split in the same order as the cases
   // appeared. This is purely for convenience of reading the resulting IR, but
   // it doesn't cost anything really.
   SmallPtrSet<BasicBlock *, 2> UnswitchedExitBBs;
if (pred_empty(DefaultExitBB)) {
       rewritePHINodesForExitAndUnswitchedBlocks(*DefaultExitBB, *SplitBB,
                                                 *ParentBB, *OldPH,
                                                 /*FullUnswitch*/ true);
       DefaultExitBB = SplitExitBBMap[DefaultExitBB] = SplitBB;
     }
   }
   // Note that we must use a reference in the for loop so that we update the
   // container.
-  for (auto &CasePair : reverse(ExitCases)) {
+  for (auto &ExitCase : reverse(ExitCases)) {
     // Grab a reference to the exit block in the pair so that we can update it.
-    BasicBlock *ExitBB = CasePair.second;
+    BasicBlock *ExitBB = std::get<1>(ExitCase);
 
     // If this case is the last edge into the exit block, we can simply reuse it
     // as it will no longer be a loop exit. No mapping necessary.
     if (pred_empty(ExitBB)) {
       // Only rewrite once.
       if (UnswitchedExitBBs.insert(ExitBB).second)
         rewritePHINodesForUnswitchedExitBlock(*ExitBB, *ParentBB, *OldPH);
       continue;
     }
 
     // Otherwise we need to split the exit block so that we retain an exit
     // block from the loop and a target for the unswitched condition.
     BasicBlock *&SplitExitBB = SplitExitBBMap[ExitBB];
     if (!SplitExitBB) {
       // If this is the first time we see this, do the split and remember it.
       SplitExitBB = SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI, MSSAU);
       rewritePHINodesForExitAndUnswitchedBlocks(*ExitBB, *SplitExitBB,
                                                 *ParentBB, *OldPH,
                                                 /*FullUnswitch*/ true);
     }
     // Update the case pair to point to the split block.
-    CasePair.second = SplitExitBB;
+    std::get<1>(ExitCase) = SplitExitBB;
   }
 
   // Now add the unswitched cases. We do this in reverse order as we built them
   // in reverse order.
-  for (auto CasePair : reverse(ExitCases)) {
-    ConstantInt *CaseVal = CasePair.first;
-    BasicBlock *UnswitchedBB = CasePair.second;
+  for (auto &ExitCase : reverse(ExitCases)) {
+    ConstantInt *CaseVal = std::get<0>(ExitCase);
+    BasicBlock *UnswitchedBB = std::get<1>(ExitCase);
 
-    NewSI->addCase(CaseVal, UnswitchedBB);
+    NewSI->addCase(CaseVal, UnswitchedBB, std::get<2>(ExitCase));
   }
 
   // If the default was unswitched, re-point it and add explicit cases for
   // entering the loop.
   if (DefaultExitBB) {
-    NewSI->setDefaultDest(DefaultExitBB);
+    NewSI->setDefaultDest(DefaultExitBB, DefaultExitWeight);
 
     // We removed all the exit cases, so we just copy the cases to the
     // unswitched switch.
-    for (auto Case : SI.cases())
-      NewSI->addCase(Case.getCaseValue(), NewPH);
+    for (const auto &Case : SI.cases())
+      NewSI->addCase(Case.getCaseValue(), NewPH,
+                     SI.getSuccessorWeight(Case.getSuccessorIndex()));
   }
 
   // If we ended up with a common successor for every path through the switch
   // after unswitching, rewrite it to an unconditional branch to make it easy
   // to recognize. Otherwise we potentially have to recognize the default case
   // pointing at unreachable and other complexity.
   if (CommonSuccBB) {
     BasicBlock *BB = SI.getParent();
     // We may have had multiple edges to this common successor block, so remove
     // them as predecessors. We skip the first one, either the default or the
     // actual first case.
     bool SkippedFirst = DefaultExitBB == nullptr;
     for (auto Case : SI.cases()) {
       assert(Case.getCaseSuccessor() == CommonSuccBB &&
              "Non-common successor!");
static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
   } else if (DefaultExitBB) {
     assert(SI.getNumCases() > 0 &&
            "If we had no cases we'd have a common successor!");
     // Move the last case to the default successor. This is valid as if the
     // default got unswitched it cannot be reached. This has the advantage of
     // being simple and keeping the number of edges from this switch to
     // successors the same, and avoiding any PHI update complexity.
     auto LastCaseI = std::prev(SI.case_end());
-    SI.setDefaultDest(LastCaseI->getCaseSuccessor());
+
+    SI.setDefaultDest(LastCaseI->getCaseSuccessor(),
+                      SI.getSuccessorWeight(LastCaseI->getSuccessorIndex()));
     SI.removeCase(LastCaseI);
   }
 
   // Walk the unswitched exit blocks and the unswitched split blocks and update
   // the dominator tree based on the CFG edits. While we are walking unordered
   // containers here, the API for applyUpdates takes an unordered list of
   // updates and requires them to not contain duplicates.
   SmallVector<DominatorTree::UpdateType, 4> DTUpdates;

llvm/test/Transforms/SimpleLoopUnswitch/basictest.ll

 ; CHECK: !prof ![[MD0:[0-9]+]]
 loop_begin:
 
   %var_val = load i32, i32* %var
 
   switch i32 %c, label %default [
       i32 1, label %inc
       i32 2, label %dec
-  ], !prof !{!"branch_weights", i32 1, i32 1, i32 1}
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2}
 
 inc:
   call void @incf() noreturn nounwind
   br label %loop_begin
 dec:
 ; CHECK: call void @decf()
 ; CHECK-NOT: call void @decf()
   call void @decf() noreturn nounwind noduplicate

 for.end:
   ret void
 }
 
 declare void @incf() noreturn
 declare void @decf() noreturn
 declare void @conv() convergent
 
-; CHECK: ![[MD0]] = !{!"branch_weights", i32 1, i32 1}
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 2, i32 1}

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

 loop_begin:
   %var_val = load i32, i32* %var
   switch i32 %cond2, label %loop_exit2 [
     i32 0, label %loop0
     i32 1, label %loop1
     i32 13, label %loop_exit1
     i32 2, label %loop2
     i32 42, label %loop_exit3
-  ], !prof !{!"branch_weights", i32 1, i32 1, i32 1, i32 1, i32 1, i32 1}
+  ], !prof !{!"branch_weights", i32 99, i32 1, i32 2, i32 3, i32 4, i32 5}
 ; CHECK:       loop_begin:
 ; CHECK-NEXT:    load
 ; CHECK-NEXT:    switch i32 %cond2, label %loop2 [
 ; CHECK-NEXT:      i32 0, label %loop0
 ; CHECK-NEXT:      i32 1, label %loop1
 ; CHECK-NEXT:    ], !prof ![[MD0:[0-9]+]]
 
 loop0:

 ; CHECK-NEXT:    br label %header
 
 loopexit:
   ret void
 ; CHECK:       loopexit:
 ; CHECK-NEXT:    ret
 }
 
-; CHECK: ![[MD0]] = !{!"branch_weights", i32 1, i32 1, i32 1}
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 4, i32 1, i32 2}