Fix bug 23851: Preserve metadata for the unswitched branch in loop-unswitch

[weimingz]

Summary

Currently, when the branch is unswitched, the newly created branch does not copy the metadata (like branch_weights) from the old branch.

[weimingz]

ping?

[reames]

Can you upload a full context diff please? It's hard to assess what you're changing with the current limited context.

The general direction seems workable.

A couple of questions:

1. Are we sure the direction of the branches are the same? If not, we may need to invert the profile metadata.
2. Should we passing in the original terminator? Or just information about the metadata? Once I have context, I may have suggestions in this area.
3. You should white list the metadata. Not all metadata will neccessarily apply.

[weimingz]

full diff of the original patch

[weimingz]

HI Philip,

Thanks for reviewing and the suggestions. I will check the direction and white list the metadata.

Weiming

In D10458#189642, @reames wrote:

Can you upload a full context diff please? It's hard to assess what you're changing with the current limited context.

The general direction seems workable.

A couple of questions:

1. Are we sure the direction of the branches are the same? If not, we may need to invert the profile metadata.
2. Should we passing in the original terminator? Or just information about the metadata? Once I have context, I may have suggestions in this area.
3. You should white list the metadata. Not all metadata will neccessarily apply.

[weimingz]

1. White list the metadata to be copied. Currently, prof and dbg nodes are copied.
2. Handle the case if True/False is swapped in the new branch
3. I prefer passing "Instruction*" instead of "SmallVectorImpl< std::pair< unsigned, MDNode * >> &MDs" because 1) the latter is much longer 2) each call site needs to prepare the vector and call getAllMetadata().

[hfinkel]

lib/Transforms/Scalar/LoopUnswitch.cpp

693

Please name this

copyMetadata

the convention is not to capitalize the d.

720

I'd prefer there be a

// fallthrough

comment here.

[weimingz]

copyMetaData -> copyMetadata
and add comments for the fallthrough case

[hfinkel]

LGTM.

[weimingz]

Thank you, Hal!

[mcrosier]

Committed in r240378.

[mcrosier]

@reames: Can you please close this revision? I'm unable to do so because of your -1.

lib/Transforms/Scalar/LoopUnswitch.cpp

 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 #include <map>
void initLoopData() {
       loopHeader = currentLoop->getHeader();
       loopPreheader = currentLoop->getLoopPreheader();
     }
 
     /// 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 UnswitchIfProfitable(Value *LoopCond, Constant *Val);
+    bool UnswitchIfProfitable(Value *LoopCond, Constant *Val,
+                              TerminatorInst *TI = nullptr);
     void UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
-                                  BasicBlock *ExitBlock);
-    void UnswitchNontrivialCondition(Value *LIC, Constant *OnVal, Loop *L);
+                                  BasicBlock *ExitBlock, TerminatorInst *TI);
+    void UnswitchNontrivialCondition(Value *LIC, Constant *OnVal, Loop *L,
+                                     TerminatorInst *TI);
 
     void RewriteLoopBodyWithConditionConstant(Loop *L, Value *LIC,
                                               Constant *Val, bool isEqual);
 
     void EmitPreheaderBranchOnCondition(Value *LIC, Constant *Val,
                                         BasicBlock *TrueDest,
                                         BasicBlock *FalseDest,
-                                        Instruction *InsertPt);
+                                        Instruction *InsertPt,
+                                        TerminatorInst *TI);
 
     void SimplifyCode(std::vector<Instruction*> &Worklist, Loop *L);
     bool IsTrivialUnswitchCondition(Value *Cond, Constant **Val = nullptr,
                                     BasicBlock **LoopExit = nullptr);
 
   };
 }
 
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);
-        if (LoopCond && UnswitchIfProfitable(LoopCond,
-                                             ConstantInt::getTrue(Context))) {
+        if (LoopCond &&
+            UnswitchIfProfitable(LoopCond, ConstantInt::getTrue(Context), TI)) {
           ++NumBranches;
           return true;
         }
       }
     } else if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
       Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
                                              currentLoop, Changed);
       unsigned NumCases = SI->getNumCases();
for (BasicBlock::iterator I = Header->begin(), E = Header->end(); I != E; ++I)
     if (I->mayHaveSideEffects())
       return false;
   return true;
 }
 
 /// UnswitchIfProfitable - 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) {
+bool LoopUnswitch::UnswitchIfProfitable(Value *LoopCond, Constant *Val,
+                                        TerminatorInst *TI) {
   Function *F = loopHeader->getParent();
   Constant *CondVal = nullptr;
   BasicBlock *ExitBlock = nullptr;
 
   if (IsTrivialUnswitchCondition(LoopCond, &CondVal, &ExitBlock)) {
     // If the condition is trivial, always unswitch. There is no code growth
     // for this case.
-    UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, ExitBlock);
+    UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, ExitBlock, TI);
     return true;
   }
 
   // Check to see if it would be profitable to unswitch current loop.
 
   // Do not do non-trivial unswitch while optimizing for size.
   if (OptimizeForSize || F->hasFnAttribute(Attribute::OptimizeForSize))
     return false;
 
-  UnswitchNontrivialCondition(LoopCond, Val, currentLoop);
+  UnswitchNontrivialCondition(LoopCond, Val, currentLoop, TI);
   return true;
 }
 
 /// CloneLoop - 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();
   LPM->insertLoop(New, PL);
 
   // Add all of the blocks in L to the new loop.
   for (Loop::block_iterator I = L->block_begin(), E = L->block_end();
        I != E; ++I)
     if (LI->getLoopFor(*I) == L)
       New->addBasicBlockToLoop(cast<BasicBlock>(VM[*I]), *LI);
 
   // Add all of the subloops to the new loop.
   for (Loop::iterator I = L->begin(), E = L->end(); I != E; ++I)
     CloneLoop(*I, New, VM, LI, LPM);
 
   return New;
 }
 
+static void copyMetadata(Instruction *DstInst, const Instruction *SrcInst,
+                         bool Swapped) {
+  if (!SrcInst || !SrcInst->hasMetadata())
+    return;
+
+  SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+  SrcInst->getAllMetadata(MDs);
+  for (auto &MD : MDs) {
+    switch (MD.first) {
+    default:
+      break;
+    case LLVMContext::MD_prof:
+      if (Swapped && MD.second->getNumOperands() == 3 &&
+          isa<MDString>(MD.second->getOperand(0))) {
+        MDString *MDName = cast<MDString>(MD.second->getOperand(0));
+        if (MDName->getString() == "branch_weights") {
+          auto *ValT = cast_or_null<ConstantAsMetadata>(
+                           MD.second->getOperand(1))->getValue();
+          auto *ValF = cast_or_null<ConstantAsMetadata>(
+                           MD.second->getOperand(2))->getValue();
+          assert(ValT && ValF && "Invalid Operands of branch_weights");
+          auto NewMD =
+              MDBuilder(DstInst->getParent()->getContext())
+                  .createBranchWeights(cast<ConstantInt>(ValF)->getZExtValue(),
+                                       cast<ConstantInt>(ValT)->getZExtValue());
+          MD.second = NewMD;
+        }
+      }
+      // fallthrough.
+    case LLVMContext::MD_dbg:
+      DstInst->setMetadata(MD.first, MD.second);
+    }
+  }
+}
+
 /// EmitPreheaderBranchOnCondition - Emit a conditional branch on two values
 /// if LIC == Val, branch to TrueDst, otherwise branch to FalseDest.  Insert the
 /// code immediately before InsertPt.
 void LoopUnswitch::EmitPreheaderBranchOnCondition(Value *LIC, Constant *Val,
                                                   BasicBlock *TrueDest,
                                                   BasicBlock *FalseDest,
-                                                  Instruction *InsertPt) {
+                                                  Instruction *InsertPt,
+                                                  TerminatorInst *TI) {
   // Insert a conditional branch on LIC to the two preheaders.  The original
   // code is the true version and the new code is the false version.
   Value *BranchVal = LIC;
+  bool Swapped = false;
   if (!isa<ConstantInt>(Val) ||
       Val->getType() != Type::getInt1Ty(LIC->getContext()))
     BranchVal = new ICmpInst(InsertPt, ICmpInst::ICMP_EQ, LIC, Val);
-  else if (Val != ConstantInt::getTrue(Val->getContext()))
+  else if (Val != ConstantInt::getTrue(Val->getContext())) {
     // We want to enter the new loop when the condition is true.
     std::swap(TrueDest, FalseDest);
+    Swapped = true;
+  }
 
   // Insert the new branch.
   BranchInst *BI = BranchInst::Create(TrueDest, FalseDest, BranchVal, InsertPt);
+  copyMetadata(BI, TI, Swapped);
 
   // If either edge is critical, split it. This helps preserve LoopSimplify
   // form for enclosing loops.
   auto Options = CriticalEdgeSplittingOptions(DT, LI).setPreserveLCSSA();
   SplitCriticalEdge(BI, 0, Options);
   SplitCriticalEdge(BI, 1, Options);
 }
 
 /// UnswitchTrivialCondition - Given a loop that has a trivial unswitchable
 /// condition in it (a cond branch from its header block to its latch block,
 /// where the path through the loop that doesn't execute its body has no
 /// side-effects), unswitch it.  This doesn't involve any code duplication, just
 /// moving the conditional branch outside of the loop and updating loop info.
-void LoopUnswitch::UnswitchTrivialCondition(Loop *L, Value *Cond,
-                                            Constant *Val,
-                                            BasicBlock *ExitBlock) {
+void LoopUnswitch::UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
+                                            BasicBlock *ExitBlock,
+                                            TerminatorInst *TI) {
   DEBUG(dbgs() << "loop-unswitch: Trivial-Unswitch loop %"
                << loopHeader->getName() << " [" << L->getBlocks().size()
-        << " blocks] in Function " << L->getHeader()->getParent()->getName()
-        << " on cond: " << *Val << " == " << *Cond << "\n");
+               << " blocks] in Function "
+               << L->getHeader()->getParent()->getName() << " on cond: " << *Val
+               << " == " << *Cond << "\n");
 
   // First step, split the preheader, so that we know that there is a safe place
   // to insert the conditional branch.  We will change loopPreheader to have a
   // conditional branch on Cond.
   BasicBlock *NewPH = SplitEdge(loopPreheader, loopHeader, DT, LI);
 
   // Now that we have a place to insert the conditional branch, create a place
   // to branch to: this is the exit block out of the loop that we should
   // short-circuit to.
 
   // Split this block now, so that the loop maintains its exit block, and so
   // that the jump from the preheader can execute the contents of the exit block
   // without actually branching to it (the exit block should be dominated by the
   // loop header, not the preheader).
   assert(!L->contains(ExitBlock) && "Exit block is in the loop?");
   BasicBlock *NewExit = SplitBlock(ExitBlock, ExitBlock->begin(), DT, LI);
 
   // Okay, now we have a position to branch from and a position to branch to,
   // insert the new conditional branch.
   EmitPreheaderBranchOnCondition(Cond, Val, NewExit, NewPH,
-                                 loopPreheader->getTerminator());
+                                 loopPreheader->getTerminator(), TI);
   LPM->deleteSimpleAnalysisValue(loopPreheader->getTerminator(), L);
   loopPreheader->getTerminator()->eraseFromParent();
 
   // We need to reprocess this loop, it could be unswitched again.
   redoLoop = true;
 
   // Now that we know that the loop is never entered when this condition is a
   // particular value, rewrite the loop with this info.  We know that this will
SplitBlockPredecessors(ExitBlock, Preds, ".us-lcssa",
                            /*PreserveLCSSA*/ true);
   }
 }
 
 /// UnswitchNontrivialCondition - 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) {
+                                               Loop *L, 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 (ScalarEvolution *SE = getAnalysisIfAvailable<ScalarEvolution>())
     SE->forgetLoop(L);
for (BasicBlock::iterator I = NewBlocks[i]->begin(),
       RemapInstruction(I, VMap,RF_NoModuleLevelChanges|RF_IgnoreMissingEntries);
 
   // Rewrite the original preheader to select between versions of the loop.
   BranchInst *OldBR = cast<BranchInst>(loopPreheader->getTerminator());
   assert(OldBR->isUnconditional() && OldBR->getSuccessor(0) == LoopBlocks[0] &&
          "Preheader splitting did not work correctly!");
 
   // Emit the new branch that selects between the two versions of this loop.
-  EmitPreheaderBranchOnCondition(LIC, Val, NewBlocks[0], LoopBlocks[0], OldBR);
+  EmitPreheaderBranchOnCondition(LIC, Val, NewBlocks[0], LoopBlocks[0], OldBR,
+                                 TI);
   LPM->deleteSimpleAnalysisValue(OldBR, L);
   OldBR->eraseFromParent();
 
   LoopProcessWorklist.push_back(NewLoop);
   redoLoop = true;
 
   // Keep a WeakVH holding onto LIC.  If the first call to RewriteLoopBody
   // deletes the instruction (for example by simplifying a PHI that feeds into

test/Transforms/LoopUnswitch/2015-06-17-Metadata.ll

+;RUN: opt  -loop-unswitch -simplifycfg -S < %s | FileCheck %s
+
+define i32 @foo(i32 %a, i32 %b) {
+;CHECK-LABEL: foo
+entry:
+  br label %for.body.lr.ph
+
+for.body.lr.ph:                                   ; preds = %entry
+  %cmp0 = icmp sgt i32 %b, 0
+  br i1 %cmp0, label %for.body, label %for.cond.cleanup
+
+for.body:                                         ; preds = %for.inc, %for.body.lr.ph
+  %inc.i = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
+  %mul.i = phi i32 [ 3, %for.body.lr.ph ], [ %mul.p, %for.inc ]
+  %add.i = phi i32 [ %a, %for.body.lr.ph ], [ %add.p, %for.inc ]
+  %cmp1 = icmp eq i32 %a, 12345
+  br i1 %cmp1, label %if.then, label %if.else, !prof !0
+; CHECK: %cmp1 = icmp eq i32 %a, 12345
+; CHECK-NEXT: br i1 %cmp1, label %if.then.us, label %if.else, !prof !0
+if.then:                                          ; preds = %for.body
+; CHECK: if.then.us:
+; CHECK: add nsw i32 %{{.*}}, 123
+; CHECK: %exitcond.us = icmp eq i32 %inc.us, %b
+; CHECK: br i1 %exitcond.us, label %for.cond.cleanup, label %if.then.us
+  %add = add nsw i32 %add.i, 123
+  br label %for.inc
+
+if.else:                                          ; preds = %for.body
+  %mul = mul nsw i32 %mul.i, %b
+  br label %for.inc
+; CHECK: if.else:
+; CHECK: %mul = mul nsw i32 %mul.i, %b
+; CHECK: %inc = add nuw nsw i32 %inc.i, 1
+; CHECK: %exitcond = icmp eq i32 %inc, %b
+; CHECK: br i1 %exitcond, label %for.cond.cleanup, label %if.else
+for.inc:                                          ; preds = %if.then, %if.else
+  %mul.p = phi i32 [ %b, %if.then ], [ %mul, %if.else ]
+  %add.p = phi i32 [ %add, %if.then ], [ %a, %if.else ]
+  %inc = add nuw nsw i32 %inc.i, 1
+  %exitcond = icmp eq i32 %inc, %b
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.inc, %for.body.lr.ph
+  %t2 = phi i32 [ %b, %for.body.lr.ph ], [ %mul.p, %for.inc ]
+  %t1 = phi i32 [ %a, %for.body.lr.ph ], [ %add.p, %for.inc ]
+  %add3 = add nsw i32 %t2, %t1
+  ret i32 %add3
+}
+
+define void @foo_swapped(i32 %a, i32 %b) {
+;CHECK-LABEL: foo_swapped
+entry:
+  br label %for.body
+;CHECK: entry:
+;CHECK-NEXT: %cmp1 = icmp eq i32 1, 2
+;CHECK-NEXT: br i1 %cmp1, label %for.body, label %for.cond.cleanup.split, !prof !1
+;CHECK: for.body:
+for.body:                                         ; preds = %for.inc, %entry
+  %inc.i = phi i32 [ 0, %entry ], [ %inc, %if.then ]
+  %add.i = phi i32 [ 100, %entry ], [ %add, %if.then ]
+  %inc = add nuw nsw i32 %inc.i, 1
+  %cmp1 = icmp eq i32 1, 2
+  br i1 %cmp1, label %if.then, label  %for.cond.cleanup, !prof !0
+
+if.then:                                          ; preds = %for.body
+  %add = add nsw i32 %a, %add.i
+
+  %exitcond = icmp eq i32 %inc, %b
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.inc, %for.body.lr.ph, %for.body
+  ret void
+}
+!0 = !{!"branch_weights", i32 64, i32 4}
+
+;CHECK: !0 = !{!"branch_weights", i32 64, i32 4}
+;CHECK: !1 = !{!"branch_weights", i32 4, i32 64}