[llvm][ir] Purge MD_prof custom accessors

[Dinistro]

Summary

This commit purges direct accesses to MD_prof metadata and replaces them
with the accessors provided from the utility file wherever possible.
This commit can be seen as the first step towards switching the branch weights to 64 bits.
See post here: https://discourse.llvm.org/t/extend-md-prof-branch-weights-metadata-from-32-to-64-bits/67492

[Dinistro]

Added the reviewers of the previous commits to ProfDataUtils.cpp.

[paulkirth]

Thanks for picking this up. When I implemented ProfDataUtils there were still several things I wanted to address but havent' had sufficient bandwidth to pick up again. Looks like you've got most of the ones that I recall so far.

I left some feedback inline. Overall I think this is a good change. I have a few nits, and concerns that we may be missing some important checks, but most are pretty minor.

The one thing I'd say that definitely needs to change IMO is the addition of the SwitchInst to the True/False extractBranchWeights API. It seems error prone and wrong to use in almost all situations. But I'm not opposed to the change if there is a compelling reason, but I still don't see much benefit to it w/in this patch, except for one place in Local.cpp.

llvm/include/llvm/IR/ProfDataUtils.h

76

This is incorrect. This overload is specifically for branch/select instructions where the number of successors is 2. There is a more general version that uses weights vector for all instruction types, including switches.

If this API isn't needed anywhere/in the code base anymore, then it can be removed IMO, but this comment shouldn't change.

llvm/lib/IR/Instructions.cpp

4606–4609

I don't think we do this check in ProfDatUtils, so we probably shouldn't remove it, should we? I know I've tripped this assert before when modifying some passes, so I think it's good one to keep.

llvm/lib/IR/ProfDataUtils.cpp

46

This constant is used to check the validity of branch weight metadata. Changing this to 2 implies that a MD_prof node w/ only the "branch_weight" moniker and a single weight are valid. I don't think that's true, so I don't think this should change unless we change how MD_prof is laid out as a whole.

144

This is wrong. you can't include a switch here. The overload above takes a SmallVectorImpl is the only appropriate API when dealing with a switch, since it can have an arbitrary number of arms.

The assertion here is to prevent API misuse. If used on a switch, this will always return false, so I don't think thats the right choice here. You get some utility by being able to use this in more places, but I think this will be easy to introduce subtle bugs. For instance, if a switch happens to only have 2 arms this will work when it arguably shouldn't.

My preference would be that everything settle on the SmallVectorImpl API, since that can correctly handle all cases. It isn't' as ergonomic in some of the code uses though, IIRC.

llvm/lib/Transforms/IPO/PartialInlining.cpp

719

llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp

578–579

Oh, this is a good change. The API use is much better at call sites.

Since the validation is basically done in ProfDataUtils, maybe rename this API to something more self descriptive?

588

nit: if we're changing the variable names, let's be consistent. I'd leave them as they were to limit the change, but if we're changing some, we should change them uniformly.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2524–2525

I think we discourage auto in cases like this, don't we?

See: https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

llvm/lib/Transforms/Scalar/LoopPredication.cpp

977–978

Should we also provide this API in ProfDataUtils? It's a validity check for branch weight nodes, and we already encapsulate most of it ...

llvm/lib/Transforms/Utils/Local.cpp

221–222

316–320

The overload will just get the MD node again anyway before calling the other API, so I'd suggest just using the MD node directly. Would you mind changing these and any others in this patch that use getBranchWeightMDNode prior to extractBranchWeights

[davidxl]

great cleanup overall.

llvm/lib/IR/ProfDataUtils.cpp

46

agree. What the reason to change it to 2?

144

Agree.

[Dinistro]

address reviewer comments

[Dinistro]

Minor improvement

[paulkirth]

llvm/lib/IR/ProfDataUtils.cpp

104–124

I'd consider inverting the logic between getValidBranchWeightMDNode() and hasValidBranchWeightMD()

llvm/lib/Transforms/Scalar/LoopPredication.cpp

992

I always forget the enumerate API exists ... what a massive improvement to readability.

llvm/lib/Transforms/Utils/Local.cpp

225

Should we track this improvement w/ an issue? If you don't want to that's fine, but I always lose track of these TODOs unless I file something in a bug tracker and ref it from my code. Totally up to you though.

316–319

uint64_t, maybe? I can never remember what is guaranteed to run before we scale down to 32-bit weights, but the size should at least be the same between the weights vector and DefWeight/CaseWeight, right?

[Dinistro]

address additional comments

[Dinistro]

llvm/lib/IR/Instructions.cpp

4606–4609

I changed this back to have the llvm_unreachable. The newly added getValidBranchWeightMDNode could be used here but that would not cause an assertion violation.

llvm/lib/IR/ProfDataUtils.cpp

46

I changed this because calls can also have branch weights: https://llvm.org/docs/BranchWeightMetadata.html#callinst
It turns out that branch weights of calls different from other branch weights and can thus either way not be handled with the current state of this utility, so I reverted this part.

llvm/lib/Transforms/Utils/Local.cpp

225

Good idea, I tracked this and another problem with the subsequent swap in here: https://github.com/llvm/llvm-project/issues/59956

316–319

The extractBranchWeights function expects an SmallVectorImp<uint32_t>, so I now changed the Weight variables to uint32_t as well.

[paulkirth]

This is LGTM from my perspective, as long as the 32-bit usage in Local.cpp is safe. Thanks for taking this up. I think it's a big improvement.

llvm/lib/Transforms/Utils/Local.cpp

316–319

I think this is fine but we may want to double check that this only gets called after scaling down to 32-bit. Otherwise, maybe we should have a uint64_t version too, or make this a templated function in the .h, so the you can choose at the callsite. That can be a separate change though, as long as 32-bit is safe.

@davidxl do you know if this gets called after scaling down to 32-bit?

[davidxl]

lgtm

llvm/lib/Transforms/Utils/Local.cpp

316–319

IRPGO scales the count before setting the edge weights. Frontend based instrumentation does the same. Static branch weights have small values so they are fine too.

[Dinistro]

Thank you very much for the reviews. Could one of you commit this revision for me as I do not yet have commit access?

[paulkirth]

@Dinistro I can do that for you, but looking at your history, I'd say you should just request access. This is especially true, since I think you plan some follow up work w/ 64-bit weights, correct? https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access. Usually people get access w/in an hour or so, but w/in a day is the expectation.

But let me know, and I can land this for you. But just FYI, if I do so I'll be the Author unfortunately, but I will attribute you in the commit summary. For some reason the committer doesn't get emails from bots when breakages happen, and only the author does. So to avoid that I've just started doing it that way. It's very imperfect, and not my preference, but it's not easy to catch breakages otherwise.

llvm/include/llvm/Analysis/CFGPrinter.h

 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/HeatUtils.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/Support/DOTGraphTraits.h"
 #include "llvm/Support/FormatVariadic.h"
 
 namespace llvm {
 template <class GraphType> struct GraphTraits;
 class CFGViewerPass : public PassInfoMixin<CFGViewerPass> {
 public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
std::string getEdgeAttributes(const BasicBlock *Node, const_succ_iterator I,
     // profile count (due to scaling).
 
     uint64_t Freq = CFGInfo->getFreq(Node);
     std::string Attrs = formatv("label=\"W:{0}\" penwidth={1}",
                                 (uint64_t)(Freq * WeightPercent), Width);
     if (Attrs.size())
       return Attrs;
 
-    MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
+    MDNode *WeightsNode = getBranchWeightMDNode(*TI);
     if (!WeightsNode)
       return "";
 
-    MDString *MDName = cast<MDString>(WeightsNode->getOperand(0));
-    if (MDName->getString() != "branch_weights")
-      return "";
-
     OpNo = I.getSuccessorIndex() + 1;
     if (OpNo >= WeightsNode->getNumOperands())
       return "";
     ConstantInt *Weight =
         mdconst::dyn_extract<ConstantInt>(WeightsNode->getOperand(OpNo));
     if (!Weight)
       return "";
     return ("label=\"W:" + std::to_string(Weight->getZExtValue()) +

llvm/include/llvm/IR/Instructions.h

 /// A wrapper class to simplify modification of SwitchInst cases along with
 /// their prof branch_weights metadata.
 class SwitchInstProfUpdateWrapper {
   SwitchInst &SI;
   std::optional<SmallVector<uint32_t, 8>> Weights;
   bool Changed = false;
 
 protected:
-  static MDNode *getProfBranchWeightsMD(const SwitchInst &SI);
-
   MDNode *buildProfBranchWeightsMD();
 
   void init();
 
 public:
   using CaseWeightOpt = std::optional<uint32_t>;
   SwitchInst *operator->() { return &SI; }
   SwitchInst &operator*() { return SI; }

llvm/include/llvm/IR/ProfDataUtils.h

 
 /// Checks if an instructions has Branch Weight Metadata
 ///
 /// \param I The instruction to check
 /// \returns True if I has an MD_prof node containing Branch Weights. False
 /// otherwise.
 bool hasBranchWeightMD(const Instruction &I);
 
+/// Checks if an instructions has valid Branch Weight Metadata
+///
+/// \param I The instruction to check
+/// \returns True if I has an MD_prof node containing valid Branch Weights,
+/// i.e., one weight for each successor. False otherwise.
+bool hasValidBranchWeightMD(const Instruction &I);
+
+/// Get the branch weights metadata node
+///
+/// \param I The Instruction to get the weights from.
+/// \returns A pointer to I's branch weights metadata node, if it exists.
+/// Nullptr otherwise.
+MDNode *getBranchWeightMDNode(const Instruction &I);
+
+/// Get the valid branch weights metadata node
+///
+/// \param I The Instruction to get the weights from.
+/// \returns A pointer to I's valid branch weights metadata node, if it exists.
+/// Nullptr otherwise.
+MDNode *getValidBranchWeightMDNode(const Instruction &I);
+
 /// Extract branch weights from MD_prof metadata
 ///
 /// \param ProfileData A pointer to an MDNode.
 /// \param [out] Weights An output vector to fill with branch weights
 /// \returns True if weights were extracted, False otherwise. When false Weights
 /// will be cleared.
 bool extractBranchWeights(const MDNode *ProfileData,
                           SmallVectorImpl<uint32_t> &Weights);
 
 /// Extract branch weights attatched to an Instruction
 ///
 /// \param I The Instruction to extract weights from.
 /// \param [out] Weights An output vector to fill with branch weights
 /// \returns True if weights were extracted, False otherwise. When false Weights
 /// will be cleared.
 bool extractBranchWeights(const Instruction &I,
                           SmallVectorImpl<uint32_t> &Weights);
 
 /// Extract branch weights from a conditional branch or select Instruction.
 ///
 /// \param I The instruction to extract branch weights from.
 /// \param [out] TrueVal will contain the branch weight for the True branch
 /// \param [out] FalseVal will contain the branch weight for the False branch
 /// \returns True on success with profile weights filled in. False if no
 /// metadata or invalid metadata was found.
 bool extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
                           uint64_t &FalseVal);
 
 /// Retrieve the total of all weights from MD_prof data.
 ///
 /// \param ProfileData The profile data to extract the total weight from
 /// \param [out] TotalWeights input variable to fill with total weights
 /// \returns True on success with profile total weights filled in. False if no
 /// metadata was found.
 bool extractProfTotalWeight(const MDNode *ProfileData, uint64_t &TotalWeights);
 
+/// Retrieve the total of all weights from an instruction.
+///
+/// \param I The instruction to extract the total weight from
+/// \param [out] TotalWeights input variable to fill with total weights
+/// \returns True on success with profile total weights filled in. False if no
+/// metadata was found.
+bool extractProfTotalWeight(const Instruction &I, uint64_t &TotalWeights);
+
 } // namespace llvm
 #endif

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 // set to min of metadata and unreachable heuristic.
 bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
   const Instruction *TI = BB->getTerminator();
   assert(TI->getNumSuccessors() > 1 && "expected more than one successor!");
   if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) || isa<IndirectBrInst>(TI) ||
         isa<InvokeInst>(TI) || isa<CallBrInst>(TI)))
     return false;
 
-  MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
+  MDNode *WeightsNode = getValidBranchWeightMDNode(*TI);
   if (!WeightsNode)
     return false;
 
   // Check that the number of successors is manageable.
   assert(TI->getNumSuccessors() < UINT32_MAX && "Too many successors");
 
-  // Ensure there are weights for all of the successors. Note that the first
-  // operand to the metadata node is a name, not a weight.
-  if (WeightsNode->getNumOperands() != TI->getNumSuccessors() + 1)
-    return false;
-
   // Build up the final weights that will be used in a temporary buffer.
   // Compute the sum of all weights to later decide whether they need to
   // be scaled to fit in 32 bits.
   uint64_t WeightSum = 0;
   SmallVector<uint32_t, 2> Weights;
   SmallVector<unsigned, 2> UnreachableIdxs;
   SmallVector<unsigned, 2> ReachableIdxs;
 
-  extractBranchWeights(*TI, Weights);
+  extractBranchWeights(WeightsNode, Weights);
   for (unsigned I = 0, E = Weights.size(); I != E; ++I) {
     WeightSum += Weights[I];
     const LoopBlock SrcLoopBB = getLoopBlock(BB);
     const LoopBlock DstLoopBB = getLoopBlock(TI->getSuccessor(I));
     auto EstimatedWeight = getEstimatedEdgeWeight({SrcLoopBB, DstLoopBB});
     if (EstimatedWeight &&
         *EstimatedWeight <= static_cast<uint32_t>(BlockExecWeight::UNREACHABLE))
       UnreachableIdxs.push_back(I);

llvm/lib/IR/Instruction.cpp

 
 #include "llvm/IR/Instruction.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Operator.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/Type.h"
 using namespace llvm;
 
 Instruction::Instruction(Type *ty, unsigned it, Use *Ops, unsigned NumOps,
                          Instruction *InsertBefore)
   : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr) {
 
   // If requested, insert this instruction into a basic block...
if (getSuccessor(Idx) == OldBB)
       setSuccessor(Idx, NewBB);
 }
 
 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")
+  MDNode *ProfileData = getBranchWeightMDNode(*this);
+  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));
 }

llvm/lib/IR/Instructions.cpp

 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/ModRef.h"
 #include "llvm/Support/TypeSize.h"

 void SwitchInst::growOperands() {
   unsigned e = getNumOperands();
   unsigned NumOps = e*3;
 
   ReservedSpace = NumOps;
   growHungoffUses(ReservedSpace);
 }
 
-MDNode *
-SwitchInstProfUpdateWrapper::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;
-}
-
 MDNode *SwitchInstProfUpdateWrapper::buildProfBranchWeightsMD() {
   assert(Changed && "called only if metadata has changed");
 
   if (!Weights)
     return nullptr;
 
   assert(SI.getNumSuccessors() == Weights->size() &&
          "num of prof branch_weights must accord with num of successors");
 
   bool AllZeroes = all_of(*Weights, [](uint32_t W) { return W == 0; });
 
   if (AllZeroes || Weights->size() < 2)
     return nullptr;
 
   return MDBuilder(SI.getParent()->getContext()).createBranchWeights(*Weights);
 }
 
 void SwitchInstProfUpdateWrapper::init() {
-  MDNode *ProfileData = getProfBranchWeightsMD(SI);
+  MDNode *ProfileData = getBranchWeightMDNode(SI);
   if (!ProfileData)
     return;
 
   if (ProfileData->getNumOperands() != SI.getNumSuccessors() + 1) {
     llvm_unreachable("number of prof branch_weights metadata operands does "
                      "not correspond to number of succesors");
   }
 
   SmallVector<uint32_t, 8> Weights;
-  for (unsigned CI = 1, CE = SI.getNumSuccessors(); CI <= CE; ++CI) {
-    ConstantInt *C = mdconst::extract<ConstantInt>(ProfileData->getOperand(CI));
-    uint32_t CW = C->getValue().getZExtValue();
-    Weights.push_back(CW);
-  }
+  if (!extractBranchWeights(ProfileData, Weights))
+    return;
   this->Weights = std::move(Weights);
 }
 
 SwitchInst::CaseIt
 SwitchInstProfUpdateWrapper::removeCase(SwitchInst::CaseIt I) {
   if (Weights) {
     assert(SI.getNumSuccessors() == Weights->size() &&
            "num of prof branch_weights must accord with num of successors");
if (*W != OldW) {
       OldW = *W;
     }
   }
 }
 
 SwitchInstProfUpdateWrapper::CaseWeightOpt
 SwitchInstProfUpdateWrapper::getSuccessorWeight(const SwitchInst &SI,
                                                 unsigned idx) {
-  if (MDNode *ProfileData = getProfBranchWeightsMD(SI))
+  if (MDNode *ProfileData = getBranchWeightMDNode(SI))
     if (ProfileData->getNumOperands() == SI.getNumSuccessors() + 1)
       return mdconst::extract<ConstantInt>(ProfileData->getOperand(idx + 1))
           ->getValue()
           .getZExtValue();
 
   return std::nullopt;
 }
 

llvm/lib/IR/Metadata.cpp

 bool Instruction::extractProfTotalWeight(uint64_t &TotalVal) const {
   assert(
       (getOpcode() == Instruction::Br || getOpcode() == Instruction::Select ||
        getOpcode() == Instruction::Call || getOpcode() == Instruction::Invoke ||
        getOpcode() == Instruction::IndirectBr ||
        getOpcode() == Instruction::Switch) &&
       "Looking for branch weights on something besides branch");
 
-  return ::extractProfTotalWeight(getMetadata(LLVMContext::MD_prof), TotalVal);
+  return ::extractProfTotalWeight(*this, TotalVal);
 }
 
 void GlobalObject::copyMetadata(const GlobalObject *Other, unsigned Offset) {
   SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
   Other->getAllMetadata(MDs);
   for (auto &MD : MDs) {
     // We need to adjust the type metadata offset.
     if (Offset != 0 && MD.first == LLVMContext::MD_type) {

llvm/lib/IR/ProfDataUtils.cpp

 //
 // We maintain some constants here to ensure that we access the branch weights
 // correctly, and can change the behavior in the future if the layout changes
 
 // The index at which the weights vector starts
 constexpr unsigned WeightsIdx = 1;
 
 // the minimum number of operands for MD_prof nodes with branch weights
 constexpr unsigned MinBWOps = 3;
 
 bool extractWeights(const MDNode *ProfileData,
                     SmallVectorImpl<uint32_t> &Weights) {
   // Assume preconditions are already met (i.e. this is valid metadata)
   assert(ProfileData && "ProfileData was nullptr in extractWeights");
   unsigned NOps = ProfileData->getNumOperands();
 
   assert(WeightsIdx < NOps && "Weights Index must be less than NOps.");
bool isBranchWeightMD(const MDNode *ProfileData) {
   return isTargetMD(ProfileData, "branch_weights", MinBWOps);
 }
 
 bool hasBranchWeightMD(const Instruction &I) {
   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
   return isBranchWeightMD(ProfileData);
 }
 
+bool hasValidBranchWeightMD(const Instruction &I) {
+  auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  if (!isBranchWeightMD(ProfileData))
+    return false;
+  if (ProfileData && ProfileData->getNumOperands() == 1 + I.getNumSuccessors())
+    return true;
+  return false;
+}
+
+MDNode *getBranchWeightMDNode(const Instruction &I) {
+  auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  if (!isBranchWeightMD(ProfileData))
+    return nullptr;
+  return ProfileData;
+}
+
+MDNode *getValidBranchWeightMDNode(const Instruction &I) {
+  if (!hasValidBranchWeightMD(I))
+    return nullptr;
+  return I.getMetadata(LLVMContext::MD_prof);
+}
+
 bool extractBranchWeights(const MDNode *ProfileData,
                           SmallVectorImpl<uint32_t> &Weights) {
   if (!isBranchWeightMD(ProfileData))
     return false;
   return extractWeights(ProfileData, Weights);
 }
 
 bool extractBranchWeights(const Instruction &I,
                           SmallVectorImpl<uint32_t> &Weights) {
   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
   return extractBranchWeights(ProfileData, Weights);
 }
 
 bool extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
                           uint64_t &FalseVal) {
   assert((I.getOpcode() == Instruction::Br ||
           I.getOpcode() == Instruction::Select) &&
-         "Looking for branch weights on something besides branch or select");
+         "Looking for branch weights on something besides branch, select, or "
+         "switch");
 
   SmallVector<uint32_t, 2> Weights;
   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
   if (!extractBranchWeights(ProfileData, Weights))
     return false;
 
   if (Weights.size() > 2)
     return false;
if (ProfDataName->getString().equals("VP") &&
     TotalVal = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(2))
                    ->getValue()
                    .getZExtValue();
     return true;
   }
   return false;
 }
 
+bool extractProfTotalWeight(const Instruction &I, uint64_t &TotalVal) {
+  return extractProfTotalWeight(I.getMetadata(LLVMContext::MD_prof), TotalVal);
+}
+
 } // namespace llvm

llvm/lib/Transforms/IPO/PartialInlining.cpp

 static bool hasProfileData(const Function &F, const FunctionOutliningInfo &OI) {
   if (F.hasProfileData())
     return true;
   // Now check if any of the entry block has MD_prof data:
   for (auto *E : OI.Entries) {
     BranchInst *BR = dyn_cast<BranchInst>(E->getTerminator());
     if (!BR || BR->isUnconditional())
       continue;
-    uint64_t T, F;
-    if (extractBranchWeights(*BR, T, F))
+    if (hasBranchWeightMD(*BR))
       return true;
   }
   return false;
 }
 
 BranchProbability PartialInlinerImpl::getOutliningCallBBRelativeFreq(
     FunctionCloner &Cloner) const {
   BasicBlock *OutliningCallBB = Cloner.OutlinedFunctions.back().second;

llvm/lib/Transforms/Instrumentation/ControlHeightReduction.cpp

 #include "llvm/Analysis/RegionIterator.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
if (isHoistable(I, DT)) {
       }
     }
     Visited[I] = false;
     return false;
   }
   // Non-instructions are considered hoistable.
   return true;
 }
 
-// Returns true and sets the true probability and false probability of an
-// MD_prof metadata if it's well-formed.
-static bool checkMDProf(MDNode *MD, BranchProbability &TrueProb,
+// Constructs the true and false branch probabilities if the the instruction has
+// valid branch weights. Returns true when this was successful, false otherwise.
+static bool extractBranchProbabilities(Instruction *I,
+                                       BranchProbability &TrueProb,
                                        BranchProbability &FalseProb) {
-  if (!MD) return false;
-  MDString *MDName = cast<MDString>(MD->getOperand(0));
-  if (MDName->getString() != "branch_weights" ||
-      MD->getNumOperands() != 3)
+  uint64_t TrueWeight;
+  uint64_t FalseWeight;
+  if (!extractBranchWeights(*I, TrueWeight, FalseWeight))
     return false;
-  ConstantInt *TrueWeight = mdconst::extract<ConstantInt>(MD->getOperand(1));
-  ConstantInt *FalseWeight = mdconst::extract<ConstantInt>(MD->getOperand(2));
-  if (!TrueWeight || !FalseWeight)
-    return false;
-  uint64_t TrueWt = TrueWeight->getValue().getZExtValue();
-  uint64_t FalseWt = FalseWeight->getValue().getZExtValue();
-  uint64_t SumWt = TrueWt + FalseWt;
+  uint64_t SumWeight = TrueWeight + FalseWeight;
 
-  assert(SumWt >= TrueWt && SumWt >= FalseWt &&
+  assert(SumWeight >= TrueWeight && SumWeight >= FalseWeight &&
          "Overflow calculating branch probabilities.");
 
   // Guard against 0-to-0 branch weights to avoid a division-by-zero crash.
-  if (SumWt == 0)
+  if (SumWeight == 0)
     return false;
 
-  TrueProb = BranchProbability::getBranchProbability(TrueWt, SumWt);
-  FalseProb = BranchProbability::getBranchProbability(FalseWt, SumWt);
+  TrueProb = BranchProbability::getBranchProbability(TrueWeight, SumWeight);
+  FalseProb = BranchProbability::getBranchProbability(FalseWeight, SumWeight);
   return true;
 }
 
 static BranchProbability getCHRBiasThreshold() {
   return BranchProbability::getBranchProbability(
       static_cast<uint64_t>(CHRBiasThreshold * 1000000), 1000000);
 }
 

 // branch of the region is biased.
 static bool checkBiasedBranch(BranchInst *BI, Region *R,
                               DenseSet<Region *> &TrueBiasedRegionsGlobal,
                               DenseSet<Region *> &FalseBiasedRegionsGlobal,
                               DenseMap<Region *, BranchProbability> &BranchBiasMap) {
   if (!BI->isConditional())
     return false;
   BranchProbability ThenProb, ElseProb;
-  if (!checkMDProf(BI->getMetadata(LLVMContext::MD_prof),
-                   ThenProb, ElseProb))
+  if (!extractBranchProbabilities(BI, ThenProb, ElseProb))
     return false;
   BasicBlock *IfThen = BI->getSuccessor(0);
   BasicBlock *IfElse = BI->getSuccessor(1);
   assert((IfThen == R->getExit() || IfElse == R->getExit()) &&
          IfThen != IfElse &&
          "Invariant from findScopes");
   if (IfThen == R->getExit()) {
     // Swap them so that IfThen/ThenProb means going into the conditional code

 // Returns true and insert a select into the right biased set and the map if the
 // select is biased.
 static bool checkBiasedSelect(
     SelectInst *SI, Region *R,
     DenseSet<SelectInst *> &TrueBiasedSelectsGlobal,
     DenseSet<SelectInst *> &FalseBiasedSelectsGlobal,
     DenseMap<SelectInst *, BranchProbability> &SelectBiasMap) {
   BranchProbability TrueProb, FalseProb;
-  if (!checkMDProf(SI->getMetadata(LLVMContext::MD_prof),
-                   TrueProb, FalseProb))
+  if (!extractBranchProbabilities(SI, TrueProb, FalseProb))
     return false;
   CHR_DEBUG(dbgs() << "SI " << *SI << " ");
   CHR_DEBUG(dbgs() << "TrueProb " << TrueProb << " ");
   CHR_DEBUG(dbgs() << "FalseProb " << FalseProb << "\n");
   return checkBias(SI, TrueProb, FalseProb,
                    TrueBiasedSelectsGlobal, FalseBiasedSelectsGlobal,
                    SelectBiasMap);
 }

llvm/lib/Transforms/Scalar/JumpThreading.cpp

BasicBlock *JumpThreadingPass::splitBlockPreds(BasicBlock *BB,
   }
 
   DTU->applyUpdatesPermissive(Updates);
   return NewBBs[0];
 }
 
 bool JumpThreadingPass::doesBlockHaveProfileData(BasicBlock *BB) {
   const Instruction *TI = BB->getTerminator();
   assert(TI->getNumSuccessors() > 1 && "not a split");
-
-  MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
-  if (!WeightsNode)
-    return false;
-
-  MDString *MDName = cast<MDString>(WeightsNode->getOperand(0));
-  if (MDName->getString() != "branch_weights")
-    return false;
-
-  // Ensure there are weights for all of the successors. Note that the first
-  // operand to the metadata node is a name, not a weight.
-  return WeightsNode->getNumOperands() == TI->getNumSuccessors() + 1;
+  return hasValidBranchWeightMD(*TI);
 }
 
 /// Update the block frequency of BB and branch weight and the metadata on the
 /// edge BB->SuccBB. This is done by scaling the weight of BB->SuccBB by 1 -
 /// Freq(PredBB->BB) / Freq(BB->SuccBB).
 void JumpThreadingPass::updateBlockFreqAndEdgeWeight(BasicBlock *PredBB,
                                                      BasicBlock *BB,
                                                      BasicBlock *NewBB,

llvm/lib/Transforms/Scalar/LoopPredication.cpp

 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/GuardUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
bool LoopPredication::isLoopProfitableToPredicate() {
   // being complete for an entire function.
 
   // If the latch exits into a deoptimize or an unreachable block, do not
   // predicate on that latch check.
   auto *LatchExitBlock = LatchTerm->getSuccessor(LatchBrExitIdx);
   if (isa<UnreachableInst>(LatchTerm) ||
       LatchExitBlock->getTerminatingDeoptimizeCall())
     return false;
 
-  auto IsValidProfileData = [](MDNode *ProfileData, const Instruction *Term) {
-    if (!ProfileData || !ProfileData->getOperand(0))
-      return false;
-    if (MDString *MDS = dyn_cast<MDString>(ProfileData->getOperand(0)))
-      if (!MDS->getString().equals("branch_weights"))
-        return false;
-    if (ProfileData->getNumOperands() != 1 + Term->getNumSuccessors())
-      return false;
-    return true;
-  };
-  MDNode *LatchProfileData = LatchTerm->getMetadata(LLVMContext::MD_prof);
   // Latch terminator has no valid profile data, so nothing to check
   // profitability on.
-  if (!IsValidProfileData(LatchProfileData, LatchTerm))
+  if (!hasValidBranchWeightMD(*LatchTerm))
     return true;
 
   auto ComputeBranchProbability =
       [&](const BasicBlock *ExitingBlock,
           const BasicBlock *ExitBlock) -> BranchProbability {
     auto *Term = ExitingBlock->getTerminator();
-    MDNode *ProfileData = Term->getMetadata(LLVMContext::MD_prof);
     unsigned NumSucc = Term->getNumSuccessors();
-    if (IsValidProfileData(ProfileData, Term)) {
-      uint64_t Numerator = 0, Denominator = 0, ProfVal = 0;
-      for (unsigned i = 0; i < NumSucc; i++) {
-        ConstantInt *CI =
-            mdconst::extract<ConstantInt>(ProfileData->getOperand(i + 1));
-        ProfVal = CI->getValue().getZExtValue();
+    if (MDNode *ProfileData = getValidBranchWeightMDNode(*Term)) {
+      SmallVector<uint32_t> Weights;
+      extractBranchWeights(ProfileData, Weights);
+      uint64_t Numerator = 0, Denominator = 0;
+      for (auto [i, Weight] : llvm::enumerate(Weights)) {
         if (Term->getSuccessor(i) == ExitBlock)
-          Numerator += ProfVal;
-        Denominator += ProfVal;
+          Numerator += Weight;
+        Denominator += Weight;
       }
       return BranchProbability::getBranchProbability(Numerator, Denominator);
     } else {
       assert(LatchBlock != ExitingBlock &&
              "Latch term should always have profile data!");
       // No profile data, so we choose the weight as 1/num_of_succ(Src)
       return BranchProbability::getBranchProbability(1, NumSucc);
     }

llvm/lib/Transforms/Utils/Local.cpp

 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsWebAssembly.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
for (auto i = SI->case_begin(), e = SI->case_end(); i != e;) {
       if (i->getCaseValue() == CI) {
         TheOnlyDest = i->getCaseSuccessor();
         break;
       }
 
       // Check to see if this branch is going to the same place as the default
       // dest.  If so, eliminate it as an explicit compare.
       if (i->getCaseSuccessor() == DefaultDest) {
-        MDNode *MD = SI->getMetadata(LLVMContext::MD_prof);
+        MDNode *MD = getBranchWeightMDNode(*SI);
         unsigned NCases = SI->getNumCases();
         // Fold the case metadata into the default if there will be any branches
         // left, unless the metadata doesn't match the switch.
         if (NCases > 1 && MD && MD->getNumOperands() == 2 + NCases) {
           // Collect branch weights into a vector.
           SmallVector<uint32_t, 8> Weights;
-          for (unsigned MD_i = 1, MD_e = MD->getNumOperands(); MD_i < MD_e;
-               ++MD_i) {
-            auto *CI = mdconst::extract<ConstantInt>(MD->getOperand(MD_i));
-            Weights.push_back(CI->getValue().getZExtValue());
-          }
+          extractBranchWeights(MD, Weights);
+
           // Merge weight of this case to the default weight.
           unsigned idx = i->getCaseIndex();
+          // TODO: no overflow check
           Weights[0] += Weights[idx+1];
+          // TODO: that's incorrect, no?
           // Remove weight for this case.
           std::swap(Weights[idx+1], Weights.back());
           Weights.pop_back();
           SI->setMetadata(LLVMContext::MD_prof,
                           MDBuilder(BB->getContext()).
                           createBranchWeights(Weights));
         }
         // Remove this entry.
if (SI->getNumCases() == 1) {
       auto FirstCase = *SI->case_begin();
       Value *Cond = Builder.CreateICmpEQ(SI->getCondition(),
           FirstCase.getCaseValue(), "cond");
 
       // Insert the new branch.
       BranchInst *NewBr = Builder.CreateCondBr(Cond,
                                                FirstCase.getCaseSuccessor(),
                                                SI->getDefaultDest());
-      MDNode *MD = SI->getMetadata(LLVMContext::MD_prof);
-      if (MD && MD->getNumOperands() == 3) {
-        ConstantInt *SICase =
-            mdconst::dyn_extract<ConstantInt>(MD->getOperand(2));
-        ConstantInt *SIDef =
-            mdconst::dyn_extract<ConstantInt>(MD->getOperand(1));
-        assert(SICase && SIDef);
+      SmallVector<uint32_t> Weights;
+      if (extractBranchWeights(*SI, Weights) && Weights.size() == 2) {
+        uint64_t DefWeight = Weights[0];
+        uint64_t CaseWeight = Weights[1];
         // The TrueWeight should be the weight for the single case of SI.
         NewBr->setMetadata(LLVMContext::MD_prof,
-                        MDBuilder(BB->getContext()).
-                        createBranchWeights(SICase->getValue().getZExtValue(),
-                                            SIDef->getValue().getZExtValue()));
+                           MDBuilder(BB->getContext())
+                               .createBranchWeights(CaseWeight, DefWeight));
       }
 
       // Update make.implicit metadata to the newly-created conditional branch.
       MDNode *MakeImplicitMD = SI->getMetadata(LLVMContext::MD_make_implicit);
       if (MakeImplicitMD)
         NewBr->setMetadata(LLVMContext::MD_make_implicit, MakeImplicitMD);
 
       // Delete the old switch.