[LoopUnswitch] Split unswitch quota based on branch weight

[reames]

Summary

This patch teaches LoopUnswitch to split unswitch quota to old loop and newly generated loop based on branch weight.

Consider the following example

for (...)
   if (cond1) 
     dummy1()
   else
     dummy2()
   if (cond2)
     dummy3()
   else
     break

This loop can be unswitched twice based on cond1 and cond2, ending with 4 loops if we have enough quota. However, if we dont have enough quota, we should spend most of them on the hot branch based on branch weight. In this example, if the branch to dummy1() is hotter than dummy2(), we should give the quota to loop with dummy1() and end up with:

if (cond1)
  if (cond2)
    for (...)
      dummy1()
      dummy3()
  else
    for (...)
      dummy1()
      break
else
  for (...) 
    dummy2()
    if (cond2)
      dummy3()
    else
      break

If we don't spilt quota based on branch weight, the colder branch will get unswitched and not help performance.

[davidxl]

lib/Transforms/Scalar/LoopUnswitch.cpp

378

spill --> split?

384

TotalWeight = OldLoopWeight + NewLoopWeight;

... = (Quota * NewLoopWeight + TotalWeight/2)/TotalWeight;

for the right rounding.

571

Why not using BPI interfaces?

[reames]

lib/Transforms/Scalar/LoopUnswitch.cpp

187

This should probably be a cl::opt. One should be 100-TheOther.

573

This isn't quite right when the LoopCond isn't BI->getCondition. The branch's weight doesn't tell us anything about LoopConds likelyhood of failing.

580

Wait, what?! Oh, these aren't normalized. I'd probably do the normalization here rather than at the use site. It'd be more clear.

test/Transforms/LoopUnswitch/unswitch-with-branch-weight.ll

18

This comment (both here and in the change description) is really confusing. You're not choosing between cond1 and cond2 at all. You're choosing which *version* of cond2 might get unswitched once the unswitch of cond1 is done.

45

This test case is really hard to follow. Please write this out in full so that the unswitched loop structure can be seen.

[reames]

Marking as "Request Changes" since previous review comments don't seem to have been addressed. Just getting it more clearly marked in my review queue.

lib/Transforms/Scalar/LoopUnswitch.cpp

public:
     // Returns true if another unswitching could be done within the cost
     // threshold.
     bool CostAllowsUnswitching();
 
     // Clone all loop-unswitch related loop properties.
     // Redistribute unswitching quotas.
     // Note, that new loop data is stored inside the VMap.
     void cloneData(const Loop *NewLoop, const Loop *OldLoop,
-                   const ValueToValueMapTy &VMap);
+                   const ValueToValueMapTy &VMap, uint64_t NewLoopWeight,
+                   uint64_t OldLoopWeight);
   };
 
   class LoopUnswitch : public LoopPass {
     LoopInfo *LI;  // Loop information
     LPPassManager *LPM;
     AssumptionCache *AC;
 
     // Used to check if second loop needs processing after
     // RewriteLoopBodyWithConditionConstant rewrites first loop.
     std::vector<Loop*> LoopProcessWorklist;
 
     LUAnalysisCache BranchesInfo;
 
     bool EnabledPGO;
 
     // BFI and ColdEntryFreq are only used when PGO and
     // LoopUnswitchWithBlockFrequency are enabled.
     BlockFrequencyInfo BFI;
     BlockFrequency ColdEntryFreq;
 
+    // After unswitch, the remaining quota should be split
+    // to old loop and new loop based on branch weight. If
+    // no branch weight is available, use the default below
+    // to split evenly.
+    uint64_t DefaultOldLoopWeight = 50;
+    uint64_t DefaultNewLoopWeight = 50;
+
     bool OptimizeForSize;
     bool redoLoop;
 
     Loop *currentLoop;
     DominatorTree *DT;
     BasicBlock *loopHeader;
     BasicBlock *loopPreheader;
 
private:
     /// Split all of the edges from inside the loop to their exit blocks.
     /// Update the appropriate Phi nodes as we do so.
     void SplitExitEdges(Loop *L,
                         const SmallVectorImpl<BasicBlock *> &ExitBlocks);
 
     bool TryTrivialLoopUnswitch(bool &Changed);
 
     bool UnswitchIfProfitable(Value *LoopCond, Constant *Val,
+                              uint64_t NewLoopWeight, uint64_t OldLoopWeight,
                               TerminatorInst *TI = nullptr);
     void UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
                                   BasicBlock *ExitBlock, TerminatorInst *TI);
     void UnswitchNontrivialCondition(Value *LIC, Constant *OnVal, Loop *L,
+                                     uint64_t NewLoopWeight,
+                                     uint64_t OldLoopWeight,
                                      TerminatorInst *TI);
 
     void RewriteLoopBodyWithConditionConstant(Loop *L, Value *LIC,
                                               Constant *Val, bool isEqual);
 
     void EmitPreheaderBranchOnCondition(Value *LIC, Constant *Val,
                                         BasicBlock *TrueDest,
                                         BasicBlock *FalseDest,

 bool LUAnalysisCache::CostAllowsUnswitching() {
   return CurrentLoopProperties->CanBeUnswitchedCount > 0;
 }
 
 // Clone all loop-unswitch related loop properties.
 // Redistribute unswitching quotas.
 // Note, that new loop data is stored inside the VMap.
 void LUAnalysisCache::cloneData(const Loop *NewLoop, const Loop *OldLoop,
-                                const ValueToValueMapTy &VMap) {
+                                const ValueToValueMapTy &VMap,
+                                uint64_t NewLoopWeight,
+                                uint64_t OldLoopWeight) {
 
   LoopProperties &NewLoopProps = LoopsProperties[NewLoop];
   LoopProperties &OldLoopProps = *CurrentLoopProperties;
   UnswitchedValsMap &Insts = OldLoopProps.UnswitchedVals;
 
   // Reallocate "can-be-unswitched quota"
 
   --OldLoopProps.CanBeUnswitchedCount;
   ++OldLoopProps.WasUnswitchedCount;
   NewLoopProps.WasUnswitchedCount = 0;
   unsigned Quota = OldLoopProps.CanBeUnswitchedCount;
-  NewLoopProps.CanBeUnswitchedCount = Quota / 2;
-  OldLoopProps.CanBeUnswitchedCount = Quota - Quota / 2;
+
+  // Spill quota based on branch weight. Notice that integer
+  // division will round to floor. Therefore to make sure the
+  // hotter branch always get a higher quota, add one if
+  // NewLoopWeight is greater OldLoopWeight (if they are equal,
+  // OldLoop get the extra count which matches the original
+  // behavior without branch weight).
+  NewLoopProps.CanBeUnswitchedCount =
+      Quota * NewLoopWeight / (NewLoopWeight + OldLoopWeight) +
+      NewLoopWeight > OldLoopWeight ? 1 : 0;
+  OldLoopProps.CanBeUnswitchedCount =
+      Quota - NewLoopProps.CanBeUnswitchedCount;
 
   NewLoopProps.SizeEstimation = OldLoopProps.SizeEstimation;
 
   // Clone unswitched values info:
   // for new loop switches we clone info about values that was
   // already unswitched and has redundant successors.
   for (UnswitchedValsIt I = Insts.begin(); I != Insts.end(); ++I) {
     const SwitchInst *OldInst = I->first;
for (Loop::block_iterator I = currentLoop->block_begin(),
     if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
       // If this isn't branching on an invariant condition, we can't unswitch
       // it.
       if (BI->isConditional()) {
         // See if this, or some part of it, is loop invariant.  If so, we can
         // unswitch on it if we desire.
         Value *LoopCond = FindLIVLoopCondition(BI->getCondition(),
                                                currentLoop, Changed);
+
+        // Get branch weights from metadata. Use default if metadata
+        // is not available.
+        uint64_t NewLoopWeight = DefaultNewLoopWeight,
+                 OldLoopWeight = DefaultOldLoopWeight;
+        MDNode *MD = BI->getMetadata(LLVMContext::MD_prof);
+        if (MD && MD->getOperand(0)) {
+          if (MDString *MDS = dyn_cast<MDString>(MD->getOperand(0))) {
+            if (MDS->getString().equals("branch_weights")) {
+
+              ConstantInt *CI = mdconst::extract<ConstantInt>(MD->getOperand(1));
+              NewLoopWeight = CI->getValue().getZExtValue();
+              CI = mdconst::extract<ConstantInt>(MD->getOperand(2));
+              OldLoopWeight = CI->getValue().getZExtValue();
+            }
+          }
+        }
+
         if (LoopCond &&
-            UnswitchIfProfitable(LoopCond, ConstantInt::getTrue(Context), TI)) {
+            UnswitchIfProfitable(LoopCond, ConstantInt::getTrue(Context),
+                                 NewLoopWeight, OldLoopWeight, TI)) {
           ++NumBranches;
           return true;
         }
       }
     } else if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
       Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
                                              currentLoop, Changed);
       unsigned NumCases = SI->getNumCases();
if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
             UnswitchVal = UnswitchValCandidate;
             break;
           }
         }
 
         if (!UnswitchVal)
           continue;
 
-        if (UnswitchIfProfitable(LoopCond, UnswitchVal)) {
+        if (UnswitchIfProfitable(LoopCond, UnswitchVal, DefaultNewLoopWeight,
+                                 DefaultOldLoopWeight)) {
           ++NumSwitches;
           return true;
         }
       }
     }
 
     // Scan the instructions to check for unswitchable values.
     for (BasicBlock::iterator BBI = (*I)->begin(), E = (*I)->end();
          BBI != E; ++BBI)
       if (SelectInst *SI = dyn_cast<SelectInst>(BBI)) {
         Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
                                                currentLoop, Changed);
         if (LoopCond && UnswitchIfProfitable(LoopCond,
-                                             ConstantInt::getTrue(Context))) {
+                                             ConstantInt::getTrue(Context),
+                                             DefaultNewLoopWeight,
+                                             DefaultOldLoopWeight)) {
           ++NumSelects;
           return true;
         }
       }
   }
   return Changed;
 }
 
if (isTrivialLoopExitBlockHelper(L, BB, ExitBB, Visited))
     return ExitBB;
   return nullptr;
 }
 
 /// We have found that we can unswitch currentLoop when LoopCond == Val to
 /// simplify the loop.  If we decide that this is profitable,
 /// unswitch the loop, reprocess the pieces, then return true.
 bool LoopUnswitch::UnswitchIfProfitable(Value *LoopCond, Constant *Val,
+                                        uint64_t NewLoopWeight,
+                                        uint64_t OldLoopWeight,
                                         TerminatorInst *TI) {
   // Check to see if it would be profitable to unswitch current loop.
   if (!BranchesInfo.CostAllowsUnswitching()) {
     DEBUG(dbgs() << "NOT unswitching loop %"
                  << currentLoop->getHeader()->getName()
                  << " at non-trivial condition '" << *Val
                  << "' == " << *LoopCond << "\n"
                  << ". Cost too high.\n");
     return false;
   }
 
-  UnswitchNontrivialCondition(LoopCond, Val, currentLoop, TI);
+  UnswitchNontrivialCondition(LoopCond, Val, currentLoop, NewLoopWeight,
+                              OldLoopWeight, TI);
   return true;
 }
 
 /// Recursively clone the specified loop and all of its children,
 /// mapping the blocks with the specified map.
 static Loop *CloneLoop(Loop *L, Loop *PL, ValueToValueMapTy &VM,
                        LoopInfo *LI, LPPassManager *LPM) {
   Loop *New = new Loop();
SplitBlockPredecessors(ExitBlock, Preds, ".us-lcssa", DT, LI,
                            /*PreserveLCSSA*/ true);
   }
 }
 
 /// We determined that the loop is profitable to unswitch when LIC equal Val.
 /// Split it into loop versions and test the condition outside of either loop.
 /// Return the loops created as Out1/Out2.
 void LoopUnswitch::UnswitchNontrivialCondition(Value *LIC, Constant *Val,
-                                               Loop *L, TerminatorInst *TI) {
+                                               Loop *L, uint64_t NewLoopWeight,
+                                               uint64_t OldLoopWeight,
+                                               TerminatorInst *TI) {
   Function *F = loopHeader->getParent();
   DEBUG(dbgs() << "loop-unswitch: Unswitching loop %"
         << loopHeader->getName() << " [" << L->getBlocks().size()
         << " blocks] in Function " << F->getName()
         << " when '" << *Val << "' == " << *LIC << "\n");
 
   if (auto *SEWP = getAnalysisIfAvailable<ScalarEvolutionWrapperPass>())
     SEWP->getSE().forgetLoop(L);
void LoopUnswitch::UnswitchNontrivialCondition(Value *LIC, Constant *Val,
   // whole function's cache.
   AC->clear();
 
   // Now we create the new Loop object for the versioned loop.
   Loop *NewLoop = CloneLoop(L, L->getParentLoop(), VMap, LI, LPM);
 
   // Recalculate unswitching quota, inherit simplified switches info for NewBB,
   // Probably clone more loop-unswitch related loop properties.
-  BranchesInfo.cloneData(NewLoop, L, VMap);
+  BranchesInfo.cloneData(NewLoop, L, VMap, NewLoopWeight, OldLoopWeight);
 
   Loop *ParentLoop = L->getParentLoop();
   if (ParentLoop) {
     // Make sure to add the cloned preheader and exit blocks to the parent loop
     // as well.
     ParentLoop->addBasicBlockToLoop(NewBlocks[0], *LI);
   }
 

test/Transforms/LoopUnswitch/unswitch-with-branch-weight.ll

+; RUN: opt < %s -loop-unswitch -loop-unswitch-threshold=16 -simplifycfg -S < %s 2>&1 | FileCheck %s
+
+
+;; This test is verify that loop unswitch takes advantage of branch weight 
+;; to spill unswitch quota to old loop and newly generated loop. The test
+;; is equivalent to the following code:
+;;  for (...)
+;;    if (cond1) 
+;;      dummy1()
+;;    else
+;;      dummy2()
+;;    if (cond2)
+;;      dummy3()
+;;    else
+;;      break
+;;
+;; This can be unswitched twice based on cond1 and cond2, ending with 4
+;; loops if we have enough quota. However, if we dont have enough quota,
+;; we should spend most of them on the hot branch based on branch weight.
+;; In this test, the branch to dummy1() is hotter than dummy2(), so we 
+;; give more quota to that branch and end up like:
+;; 
+;;  if (cond1)
+;;    if (cond2)
+;;      for (...)
+;;        dummy1()
+;;        dummy3()
+;;    else
+;;      for (...)
+;;        dummy1()
+;;        break
+;;  else
+;;    for (...) 
+;;      dummy2()
+;;      if (cond2)
+;;        dummy3()
+;;      else
+;;        break
+;;
+;; If we don't spill quota based on branch weight, the colder branch
+;; will get unswitched.
+
+define i32 @test(i1 %cond1, i1 %cond2) {
+
+; CHECK:      call void @dummy1() #0
+; CHECK-NEXT: call void @dummy3() #0
+
+; CHECK:      call void @dummy2() #0
+; CHECK-NEXT: br i1 %cond2
+
+  br label %loop_begin
+
+loop_begin:  
+  br i1 %cond1, label %loop_br1, label %loop_br2, !prof !0
+  
+loop_br1:
+  call void @dummy1() nounwind
+  br label %loop_continue
+loop_br2:
+  call void @dummy2() nounwind
+  br label %loop_continue
+
+loop_continue:
+  br i1 %cond2, label %loop_br3, label %loop_exit
+
+loop_br3:
+  call void @dummy3() nounwind
+  br label %loop_begin
+
+loop_exit:
+  ret i32 0
+}
+
+declare void @dummy1() 
+declare void @dummy2() 
+declare void @dummy3() 
+
+!0 = !{!"branch_weights", i32 1000, i32 1}