Refactor metadata copying/swapping code to make it sharable

[davidxl]

Summary

No functional change is intended.

[davidxl]

Updated the patch according to review comments.

[davidxl]

Addressed review feedbacks.

include/llvm/IR/Instruction.h

   /// instruction providing the most-general result.
   void getAAMetadata(AAMDNodes &N, bool Merge = false) const;
 
   /// Set the metadata of the specified kind to the specified node. This updates
   /// or replaces metadata if already present, or removes it if Node is null.
   void setMetadata(unsigned KindID, MDNode *Node);
   void setMetadata(StringRef Kind, MDNode *Node);
 
+  /// Copy metadata from \p SrcInst to this instruction.
+  void copyMetadata(const Instruction &SrcInst);
+
+  /// If the instruction has "branch_weights" MD_prof metadata and the MDNode
+  /// has three operands (including name string), swap the order of the
+  /// metadata.
+  void swapProfMetadata();
+
+  /// Preserve only meta data specified in WL and strip the rest. For instance,
+  /// to keep MD_prof data, WL[MD_prof] must be non-zero.
+  void stripAllMetadataExcept(ArrayRef<uint8_t> WL);
+
   /// Drop all unknown metadata except for debug locations.
   /// @{
   /// Passes are required to drop metadata they don't understand. This is a
   /// convenience method for passes to do so.
   void dropUnknownNonDebugMetadata(ArrayRef<unsigned> KnownIDs);
   void dropUnknownNonDebugMetadata() {
     return dropUnknownNonDebugMetadata(None);
   }

include/llvm/IR/LLVMContext.h

enum {
     MD_dereferenceable = 12,          // "dereferenceable"
     MD_dereferenceable_or_null = 13,  // "dereferenceable_or_null"
     MD_make_implicit = 14,            // "make.implicit"
     MD_unpredictable = 15,            // "unpredictable"
     MD_invariant_group = 16,          // "invariant.group"
     MD_align = 17,                    // "align"
     MD_loop = 18,                     // "llvm.loop"
     MD_type = 19,                     // "type"
+    MD_last                           // End marker
   };
 
   /// Known operand bundle tag IDs, which always have the same value.  All
   /// operand bundle tags that LLVM has special knowledge of are listed here.
   /// Additionally, this scheme allows LLVM to efficiently check for specific
   /// operand bundle tags without comparing strings.
   enum {
     OB_deopt = 0,         // "deopt"

lib/IR/Instruction.cpp

 bool Instruction::isNilpotent(unsigned Opcode) {
   return Opcode == Xor;
 }
 
 Instruction *Instruction::cloneImpl() const {
   llvm_unreachable("Subclass of Instruction failed to implement cloneImpl");
 }
 
+void Instruction::swapProfMetadata() {
+  MDNode *ProfileData = getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData || ProfileData->getNumOperands() != 3 ||
+      !isa<MDString>(ProfileData->getOperand(0)))
+    return;
+
+  MDString *MDName = cast<MDString>(ProfileData->getOperand(0));
+  if (MDName->getString() != "branch_weights")
+    return;
+
+  // The first operand is the name. Fetch them backwards and build a new one.
+  Metadata *Ops[] = {ProfileData->getOperand(0), ProfileData->getOperand(2),
+                     ProfileData->getOperand(1)};
+  setMetadata(LLVMContext::MD_prof,
+              MDNode::get(ProfileData->getContext(), Ops));
+}
+
+void Instruction::copyMetadata(const Instruction &SrcInst) {
+  if (!SrcInst.hasMetadata())
+    return;
+
+  // Otherwise, enumerate and copy over metadata from the old instruction to the
+  // new one.
+  SmallVector<std::pair<unsigned, MDNode *>, 4> TheMDs;
+  SrcInst.getAllMetadataOtherThanDebugLoc(TheMDs);
+  for (const auto &MD : TheMDs)
+    setMetadata(MD.first, MD.second);
+
+  setDebugLoc(SrcInst.getDebugLoc());
+  return;
+}
+
+void Instruction::stripAllMetadataExcept(ArrayRef<uint8_t> WL) {
+  if (!hasMetadata())
+    return;
+
+  SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+  getAllMetadata(MDs);
+  for (auto &MD : MDs) {
+    if (!WL[MD.first])
+      setMetadata(MD.first, nullptr);
+  }
+}
+
 Instruction *Instruction::clone() const {
   Instruction *New = nullptr;
   switch (getOpcode()) {
   default:
     llvm_unreachable("Unhandled Opcode.");
 #define HANDLE_INST(num, opc, clas)                                            \
   case Instruction::opc:                                                       \
     New = cast<clas>(this)->cloneImpl();                                       \
     break;
 #include "llvm/IR/Instruction.def"
 #undef HANDLE_INST
   }
 
   New->SubclassOptionalData = SubclassOptionalData;
-  if (!hasMetadata())
-    return New;
-
-  // Otherwise, enumerate and copy over metadata from the old instruction to the
-  // new one.
-  SmallVector<std::pair<unsigned, MDNode *>, 4> TheMDs;
-  getAllMetadataOtherThanDebugLoc(TheMDs);
-  for (const auto &MD : TheMDs)
-    New->setMetadata(MD.first, MD.second);
-
-  New->setDebugLoc(getDebugLoc());
+  New->copyMetadata(*this);
   return New;
 }

lib/IR/Instructions.cpp

 
 void BranchInst::swapSuccessors() {
   assert(isConditional() &&
          "Cannot swap successors of an unconditional branch");
   Op<-1>().swap(Op<-2>());
 
   // Update profile metadata if present and it matches our structural
   // expectations.
-  MDNode *ProfileData = getMetadata(LLVMContext::MD_prof);
-  if (!ProfileData || ProfileData->getNumOperands() != 3)
-    return;
-
-  // The first operand is the name. Fetch them backwards and build a new one.
-  Metadata *Ops[] = {ProfileData->getOperand(0), ProfileData->getOperand(2),
-                     ProfileData->getOperand(1)};
-  setMetadata(LLVMContext::MD_prof,
-              MDNode::get(ProfileData->getContext(), Ops));
+  swapProfMetadata();
 }
 
 BasicBlock *BranchInst::getSuccessorV(unsigned idx) const {
   return getSuccessor(idx);
 }
 unsigned BranchInst::getNumSuccessorsV() const {
   return getNumSuccessors();
 }

lib/Transforms/Scalar/LoopUnswitch.cpp

 
   // Add all of the subloops to the new loop.
   for (Loop *I : *L)
     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_make_implicit:
-    case LLVMContext::MD_dbg:
-      DstInst->setMetadata(MD.first, MD.second);
-    }
-  }
-}
-
 /// 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,
                                                   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())) {
     // 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 (TI) {
+    BI->copyMetadata(*TI);
+    if (Swapped)
+      BI->swapProfMetadata();
+    // FIXME: check why white list is needed here:
+    std::vector<uint8_t> WL(LLVMContext::MD_last, 0);
+    WL[LLVMContext::MD_dbg] = 1;
+    WL[LLVMContext::MD_prof] = 1;
+    WL[LLVMContext::MD_make_implicit] = 1;
+    BI->stripAllMetadataExcept(WL);
+  }
 
   // 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);
 }