[IR] Improve SwitchInst API to support prof branch_weights

[yrouban]

Summary

This patch adds prof branch_weight as an optional parameter to add/remove and set successor methods.
This way preserves correctness of passes with respect to SwitchInst instructions' perf branch_weights metadata.

One of important assumption is that in a SwitchInst instruction we usually change its metadata first and then its cases. In add/remove methods the branch weights are not changed if their dimension is not equal to the number of successors because this means that the metadata has been changed by other means.

[yrouban]

llvm/lib/IR/Instructions.cpp

3818

must be uint64_t

[reames]

I'm not entirely sure about this interface. I see there's precedent in code for working with the branch weights when removing cases, so I see where you're coming from, but the rest of our interface (e.g. branches) don't have an analogous API.

I think that if we're going to add an API like this, it needs to be uniform across all the terminator instructios. Having it otherwise is just too confusing.

[yrouban]

In D60604#1464759, @reames wrote:

I think that if we're going to add an API like this, it needs to be uniform across all the terminator instructios. Having it otherwise is just too confusing.

Philip, as I see in the doc https://llvm.org/docs/BranchWeightMetadata.html#supported-instructions there are only 4 instructions allowed to have prof branch_weights: BranchInst, SwitchInst, IndirectBrInst and CallInst. CallInst is not a terminator instruction.
I'm proposing API changes for SwitchInst only. Very similar changes can be done to IndirectBrInst.
BranchInst and CallInst have fixed num of branch weight params so they do not need the changes in add/remove methods.
If you agree, I will prepare similar changes for all 4 instructions or just 2 (SwitchInst and IndirectBrInst).
The logic behind these changes is the following. I'd like to keep the prof branch_weights metadata to be kept correct through all passes that do not care about metadata and I do not want to change those passes. They add successors without specifying branch weights and remove successors disregarding the metadata consistency. On the other hand passes can explicitly specify the branch weights along with changed successors.

[yrouban]

rebased to the latest D60554 and formatted

[yrouban]

See 62122 and its children.

llvm/include/llvm/IR/Instructions.h

protected:
   friend class Instruction;
 
   SwitchInst *cloneImpl() const;
 
 public:
   // -2
   static const unsigned DefaultPseudoIndex = static_cast<unsigned>(~0L-1);
 
+  using CaseWeightOpt = Optional<uint32_t>;
+
   template <typename CaseHandleT> class CaseIteratorImpl;
 
   /// A handle to a particular switch case. It exposes a convenient interface
   /// to both the case value and the successor block.
   ///
   /// We define this as a template and instantiate it to form both a const and
   /// non-const handle.
   template <typename SwitchInstT, typename ConstantIntT, typename BasicBlockT>
public:
     /// Sets the new value for current case.
     void setValue(ConstantInt *V) {
       assert((unsigned)Index < SI->getNumCases() &&
              "Index out the number of cases.");
       SI->setOperand(2 + Index*2, reinterpret_cast<Value*>(V));
     }
 
     /// Sets the new successor for current case.
-    void setSuccessor(BasicBlock *S) {
-      SI->setSuccessor(getSuccessorIndex(), S);
+    void setSuccessor(BasicBlock *S, CaseWeightOpt W = None) {
+      SI->setSuccessor(getSuccessorIndex(), S, W);
     }
   };
 
   template <typename CaseHandleT>
   class CaseIteratorImpl
       : public iterator_facade_base<CaseIteratorImpl<CaseHandleT>,
                                     std::random_access_iterator_tag,
                                     CaseHandleT> {
public:
   // Accessor Methods for Switch stmt
   Value *getCondition() const { return getOperand(0); }
   void setCondition(Value *V) { setOperand(0, V); }
 
   BasicBlock *getDefaultDest() const {
     return cast<BasicBlock>(getOperand(1));
   }
 
-  void setDefaultDest(BasicBlock *DefaultCase) {
-    setOperand(1, reinterpret_cast<Value*>(DefaultCase));
+  CaseWeightOpt getDefaultCaseWeight() const { return getSuccessorWeight(0); }
+
+  void setDefaultDest(BasicBlock *DefaultCase, CaseWeightOpt W = None) {
+    setSuccessor(0, DefaultCase, W);
   }
 
   /// Return the number of 'cases' in this switch instruction, excluding the
   /// default case.
   unsigned getNumCases() const {
     return getNumOperands()/2 - 1;
   }
 
ConstantInt *findCaseDest(BasicBlock *BB) {
 
     return CI;
   }
 
   /// Add an entry to the switch instruction.
   /// Note:
   /// This action invalidates case_end(). Old case_end() iterator will
   /// point to the added case.
-  void addCase(ConstantInt *OnVal, BasicBlock *Dest);
+  void addCase(ConstantInt *OnVal, BasicBlock *Dest, CaseWeightOpt W = None);
 
   /// This method removes the specified case and its successor from the switch
   /// instruction. Note that this operation may reorder the remaining cases at
   /// index idx and above.
   /// Note:
   /// This action invalidates iterators for all cases following the one removed,
   /// including the case_end() iterator. It returns an iterator for the next
   /// case.
   CaseIt removeCase(CaseIt I, bool DropPerfMetadata = false);
 
   unsigned getNumSuccessors() const { return getNumOperands()/2; }
   BasicBlock *getSuccessor(unsigned idx) const {
     assert(idx < getNumSuccessors() &&"Successor idx out of range for switch!");
     return cast<BasicBlock>(getOperand(idx*2+1));
   }
-  void setSuccessor(unsigned idx, BasicBlock *NewSucc) {
-    assert(idx < getNumSuccessors() && "Successor # out of range for switch!");
-    setOperand(idx * 2 + 1, NewSucc);
-  }
+
+  CaseWeightOpt getSuccessorWeight(unsigned idx) const;
+
+  void setSuccessor(unsigned idx, BasicBlock *NewSucc, CaseWeightOpt W = None);
 
   // Methods for support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const Instruction *I) {
     return I->getOpcode() == Instruction::Switch;
   }
   static bool classof(const Value *V) {
     return isa<Instruction>(V) && classof(cast<Instruction>(V));
   }

llvm/lib/IR/Instructions.cpp

 static MDNode *getProfBranchWeightsMD(const SwitchInst *SI) {
   if (MDNode *ProfileData = SI->getMetadata(LLVMContext::MD_prof))
     if (auto *MDName = dyn_cast<MDString>(ProfileData->getOperand(0)))
       if (MDName->getString() == "branch_weights")
         return ProfileData;
   return nullptr;
 }
 
-/// addCase - Add an entry to the switch instruction...
-///
-void SwitchInst::addCase(ConstantInt *OnVal, BasicBlock *Dest) {
-  // Drop metdadata.
+// to append weight pass BI = SI->getNumSuccessors()
+static void setProfBranchWeight(SwitchInst *SI, unsigned BI, uint32_t BW) {
+  assert(BI <= SI->getNumSuccessors() && "case index is in range");
+  uint64_t SW = 0;
+  SmallVector<uint32_t, 8> Weights;
+
+  if (MDNode *ProfileData = getProfBranchWeightsMD(SI))
+    for (unsigned CI = 0, CE = std::max(BI + 1, SI->getNumSuccessors());
+         CI < CE; ++CI) {
+      uint32_t CW = (CI == BI) ? BW : mdconst::extract<ConstantInt>(
+                                          ProfileData->getOperand(CI + 1))
+                                          ->getValue()
+                                          .getZExtValue();
+      Weights.push_back(CW);
+      SW += CW;
+    }
+
+  if (!Weights.size() && BW) {
+    for (unsigned CI = 0, CE = std::max(BI + 1, SI->getNumSuccessors());
+         CI < CE; ++CI)
+      Weights.push_back(0);
+    Weights[BI] = BW;
+    SW = BW;
+  }
+
+  auto *NewProfileData =
+      !SW ? nullptr : MDBuilder(SI->getParent()->getContext())
+                          .createBranchWeights(Weights);
+  SI->setMetadata(LLVMContext::MD_prof, NewProfileData);
+}
+
+SwitchInst::CaseWeightOpt SwitchInst::getSuccessorWeight(unsigned idx) const {
   if (MDNode *ProfileData = getProfBranchWeightsMD(this))
-    if (ProfileData->getNumOperands() == getNumSuccessors() + 1)
-      setMetadata(LLVMContext::MD_prof, nullptr);
+    return mdconst::extract<ConstantInt>(ProfileData->getOperand(idx + 1))
+        ->getValue()
+        .getZExtValue();
 
+  return None;
+}
+
+void SwitchInst::setSuccessor(unsigned idx, BasicBlock *NewSucc,
+                              CaseWeightOpt W) {
+  assert(idx < getNumSuccessors() && "Successor # out of range for switch!");
+  if (W)
+    setProfBranchWeight(this, idx, *W);
+
+  unsigned OI = idx * 2 + 1;
+  setOperand(OI, NewSucc);
+}
+
+/// addCase - Add an entry to the switch instruction...
+///
+void SwitchInst::addCase(ConstantInt *OnVal, BasicBlock *Dest,
+                         CaseWeightOpt W) {
   unsigned NewCaseIdx = getNumCases();
   unsigned OpNo = getNumOperands();
   if (OpNo+2 > ReservedSpace)
     growOperands();  // Get more space!
   // Initialize some new operands.
   assert(OpNo+1 < ReservedSpace && "Growing didn't work!");
   setNumHungOffUseOperands(OpNo+2);
   CaseHandle Case(this, NewCaseIdx);
   Case.setValue(OnVal);
+
+  if (!W)
+    if (MDNode *ProfileData = getProfBranchWeightsMD(this))
+      if (ProfileData->getNumOperands() == getNumSuccessors() + 1) {
+        // prof branch_weights has been set properly, no need to touch them.
         Case.setSuccessor(Dest);
+        return;
+      }
+
+  // Pass some weight to make sure prof metadata will be fixed to have proper
+  // num of args.
+  Case.setSuccessor(Dest, W ? *W : 0);
 }
 
 /// removeCase - This method removes the specified case and its successor
 /// from the switch instruction.
 SwitchInst::CaseIt SwitchInst::removeCase(CaseIt I, bool DropPerfMetadata) {
   unsigned idx = I->getCaseIndex();
 
   assert(2 + idx*2 < getNumOperands() && "Case index out of range!!!");