[llvm][NFC] Refactor code to use ProfDataUtils

[paulkirth]

Summary

In this patch we replace common code patterns with the use of utility
functions for dealing with profiling metadata. There should be no change
in functionality, as the existing checks should be preserved in all
cases.

[paulkirth]

fixups after rebase in stack

[paulkirth]

Fix wrong literal in expression

[bogner]

llvm/include/llvm/IR/Instruction.h

356

There's an argument for removing this entirely and updating the users to just use extractBranchWeights directly (possibly with a specialization of that function that makes the True/False case convenient). I'm not sure how strongly I feel about it, but all of the branch weight stuff being in one place seems kind of nice. WDYT?

361

Is this used anywhere? We can probably omit it until it is.

[paulkirth]

llvm/include/llvm/IR/Instruction.h

356

I think putting all the branch weights utilities in one place is a good idea. Replacing these users was something I was also thinking about, so maybe that's the way we should do it here.

I'll take a look at the users and see if there's a straightforward way to migrate them or if a specialization makes more sense.

361

Yeah, I was hesitant about adding the interface, but wanted to provide an option for switch statements. Since there's currently no users I'll remove the interface. Like you said once someone wants/needs it, It's really easy to add back.

[paulkirth]

llvm/lib/IR/Metadata.cpp

1521

This should either be removed, or also check for switch stmts. Well, anything branch weights may be attached to

[davidxl]

llvm/lib/IR/Metadata.cpp

1509–1510

missing space.

[paulkirth]

Respond to comments

• Update Interfaces in LLVM for to use extracBranchWeight instead of extractProfMetadata
• Fix typo
• remove unused API

[paulkirth]

I've left the extractProfMetadata interface to Instruction, since I'm unsure if we want to force downstream users to update their code. We could alternatively just leave the interface as is, and back out the places where we've update usage.

I'm waffling on that one, and whether those updates should be a separate change, or just done here. There's value in leaving the public interface to Instruction alone. We've already moved the logic into a central place, so I'm less convinced that the rest of the refactor is the best approach.

[davidxl]

Since there are other types of prof meta data other than branch weights, I think we should get rid of the extractProfMetaData interface (in the current form) and replace them with extractBranchWeights (as Justin suggested).

[bogner]

Let's go ahead and remove extractProfMetadata. Even if downstream users are depending on it the update is trivial, and we can always put it back temporarily if it turns out to be a bigger issue somehow.

[paulkirth]

Remove extractProfMetadata API from Instruction

[davidxl]

lgtm

[paulkirth]

Rebase

[thakis]

Looks like this doesn't build: http://45.33.8.238/linux/82503/step_4.txt

Please take a look, and revert for now if it takes a while to fix.

[paulkirth]

yeah, it looks like we missed something on this. I was already reverting.

[thakis]

Thanks!

Going forward, please build and run tests locally before committing, or at least wait for the precommit bots to complete.

[paulkirth]

I had tested locally, but you're right I should have waited for the bots to finish after rebaseing too. I'll be more careful in the future.

[paulkirth]

Fix missing pointer deref

[paulkirth]

Improve documentation strings, since we got a few warnings from buildbot

[paulkirth]

Rebase

llvm/include/llvm/IR/Instruction.h

public:
   AAMDNodes getAAMetadata() const;
 
   /// Sets the AA metadata on this instruction from the AAMDNodes structure.
   void setAAMetadata(const AAMDNodes &N);
 
   /// Retrieve the raw weight values of a conditional branch or select.
   /// Returns true on success with profile weights filled in.
   /// Returns false if no metadata or invalid metadata was found.
   bool extractProfMetadata(uint64_t &TrueVal, uint64_t &FalseVal) const;
 
+  /// Retrieve the raw weight values of a conditional branch, select, or switch.
+  /// Returns true on success with profile weights filled in.
+  /// Returns false if no metadata or invalid metadata was found.
+  bool extractProfMetadata(SmallVector<uint32_t> &Weights) const;
+
   /// Retrieve total raw weight values of a branch.
   /// Returns true on success with profile total weights filled in.
   /// Returns false if no metadata was found.
   bool extractProfTotalWeight(uint64_t &TotalVal) const;
 
   /// Set the debug location information for this instruction.
   void setDebugLoc(DebugLoc Loc) { DbgLoc = std::move(Loc); }
 

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
 
   // 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;
-  Weights.reserve(TI->getNumSuccessors());
-  for (unsigned I = 1, E = WeightsNode->getNumOperands(); I != E; ++I) {
-    ConstantInt *Weight =
-        mdconst::dyn_extract<ConstantInt>(WeightsNode->getOperand(I));
-    if (!Weight)
-      return false;
-    assert(Weight->getValue().getActiveBits() <= 32 &&
-           "Too many bits for uint32_t");
-    Weights.push_back(Weight->getZExtValue());
-    WeightSum += Weights.back();
+
+  extractBranchWeights(*TI, 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 - 1));
+    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 - 1);
+      UnreachableIdxs.push_back(I);
     else
-      ReachableIdxs.push_back(I - 1);
+      ReachableIdxs.push_back(I);
   }
   assert(Weights.size() == TI->getNumSuccessors() && "Checked above");
 
   // If the sum of weights does not fit in 32 bits, scale every weight down
   // accordingly.
   uint64_t ScalingFactor =
       (WeightSum > UINT32_MAX) ? WeightSum / UINT32_MAX + 1 : 1;
 

llvm/lib/IR/Metadata.cpp

 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalObject.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/TrackingMDRef.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include <algorithm>
 #include <cassert>
void Instruction::getAllMetadataImpl(
   Value::getAllMetadata(Result);
 }
 
 bool Instruction::extractProfMetadata(uint64_t &TrueVal,
                                       uint64_t &FalseVal) const {
   assert(
       (getOpcode() == Instruction::Br || getOpcode() == Instruction::Select) &&
       "Looking for branch weights on something besides branch or select");
-
-  auto *ProfileData = getMetadata(LLVMContext::MD_prof);
-  if (!ProfileData || ProfileData->getNumOperands() != 3)
-    return false;
-
-  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
-  if (!ProfDataName || !ProfDataName->getString().equals("branch_weights"))
+  SmallVector<uint32_t>Weights;
+  if(!extractBranchWeights(*this, Weights) || Weights.size() != 3)
     return false;
 
-  auto *CITrue = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(1));
-  auto *CIFalse = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(2));
-  if (!CITrue || !CIFalse)
-    return false;
-
-  TrueVal = CITrue->getValue().getZExtValue();
-  FalseVal = CIFalse->getValue().getZExtValue();
+  TrueVal = Weights[0];
+  FalseVal = Weights[1];
 
   return true;
 }
 
+bool Instruction::extractProfMetadata(SmallVector<uint32_t> &Weights) const {
+  assert(
+      (getOpcode() == Instruction::Br || getOpcode() == Instruction::Select) &&
+      "Looking for branch weights on something besides branch or select");
+  return extractBranchWeights(*this, Weights);
+}
+
+
 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");
 
-  TotalVal = 0;
-  auto *ProfileData = getMetadata(LLVMContext::MD_prof);
-  if (!ProfileData)
-    return false;
-
-  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
-  if (!ProfDataName)
-    return false;
-
-  if (ProfDataName->getString().equals("branch_weights")) {
-    TotalVal = 0;
-    for (unsigned i = 1; i < ProfileData->getNumOperands(); i++) {
-      auto *V = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(i));
-      if (!V)
-        return false;
-      TotalVal += V->getValue().getZExtValue();
-    }
-    return true;
-  } else if (ProfDataName->getString().equals("VP") &&
-             ProfileData->getNumOperands() > 3) {
-    TotalVal = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(2))
-                   ->getValue()
-                   .getZExtValue();
-    return true;
-  }
-  return false;
+ return extractTotalBranchWeights(getMetadata(LLVMContext::MD_prof), 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/Transforms/Utils/MisExpect.cpp

 #include "llvm/Transforms/Utils/MisExpect.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <cstdint>
 #include <functional>
 #include <numeric>
 
void emitMisexpectDiagnostic(Instruction *I, LLVMContext &Ctx,
   ORE.emit(OptimizationRemark(DEBUG_TYPE, "misexpect", Cond) << RemStr.str());
 }
 
 } // namespace
 
 namespace llvm {
 namespace misexpect {
 
-// Helper function to extract branch weights into a vector
-Optional<SmallVector<uint32_t, 4>> extractWeights(Instruction *I,
-                                                  LLVMContext &Ctx) {
-  assert(I && "MisExpect::extractWeights given invalid pointer");
-
-  auto *ProfileData = I->getMetadata(LLVMContext::MD_prof);
-  if (!ProfileData)
-    return None;
-
-  unsigned NOps = ProfileData->getNumOperands();
-  if (NOps < 3)
-    return None;
-
-  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
-  if (!ProfDataName || !ProfDataName->getString().equals("branch_weights"))
-    return None;
-
-  SmallVector<uint32_t, 4> Weights(NOps - 1);
-  for (unsigned Idx = 1; Idx < NOps; Idx++) {
-    ConstantInt *Value =
-        mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
-    uint32_t V = Value->getZExtValue();
-    Weights[Idx - 1] = V;
-  }
-
-  return Weights;
-}
-
 // TODO: when clang allows c++17, use std::clamp instead
 uint32_t clamp(uint64_t value, uint32_t low, uint32_t hi) {
   if (value > hi)
     return hi;
   if (value < low)
     return low;
   return value;
 }
void verifyMisExpect(Instruction &I, ArrayRef<uint32_t> RealWeights,
   // When the profile weight is below the threshold, we emit the diagnostic
   if (ProfiledWeight < ScaledThreshold)
     emitMisexpectDiagnostic(&I, I.getContext(), ProfiledWeight,
                             RealWeightsTotal);
 }
 
 void checkBackendInstrumentation(Instruction &I,
                                  const ArrayRef<uint32_t> RealWeights) {
-  auto ExpectedWeightsOpt = extractWeights(&I, I.getContext());
-  if (!ExpectedWeightsOpt)
+  SmallVector<uint32_t> ExpectedWeights;
+  if (!extractBranchWeights(I, ExpectedWeights))
     return;
-  auto ExpectedWeights = ExpectedWeightsOpt.getValue();
   verifyMisExpect(I, RealWeights, ExpectedWeights);
 }
 
 void checkFrontendInstrumentation(Instruction &I,
                                   const ArrayRef<uint32_t> ExpectedWeights) {
-  auto RealWeightsOpt = extractWeights(&I, I.getContext());
-  if (!RealWeightsOpt)
+  SmallVector<uint32_t> RealWeights;
+  if (!extractBranchWeights(I, RealWeights))
     return;
-  auto RealWeights = RealWeightsOpt.getValue();
   verifyMisExpect(I, RealWeights, ExpectedWeights);
 }
 
 void checkExpectAnnotations(Instruction &I,
                             const ArrayRef<uint32_t> ExistingWeights,
                             bool IsFrontendInstr) {
   if (IsFrontendInstr) {
     checkFrontendInstrumentation(I, ExistingWeights);
   } else {
     checkBackendInstrumentation(I, ExistingWeights);
   }
 }
 
 } // namespace misexpect
 } // namespace llvm
 #undef DEBUG_TYPE

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/IR/Operator.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/BranchProbability.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
static int ConstantIntSortPredicate(ConstantInt *const *P1,
                                     ConstantInt *const *P2) {
   const ConstantInt *LHS = *P1;
   const ConstantInt *RHS = *P2;
   if (LHS == RHS)
     return 0;
   return LHS->getValue().ult(RHS->getValue()) ? 1 : -1;
 }
 
-static inline bool HasBranchWeights(const Instruction *I) {
-  MDNode *ProfMD = I->getMetadata(LLVMContext::MD_prof);
-  if (ProfMD && ProfMD->getOperand(0))
-    if (MDString *MDS = dyn_cast<MDString>(ProfMD->getOperand(0)))
-      return MDS->getString().equals("branch_weights");
-
-  return false;
-}
-
 /// Get Weights of a given terminator, the default weight is at the front
 /// of the vector. If TI is a conditional eq, we need to swap the branch-weight
 /// metadata.
 static void GetBranchWeights(Instruction *TI,
                              SmallVectorImpl<uint64_t> &Weights) {
   MDNode *MD = TI->getMetadata(LLVMContext::MD_prof);
   assert(MD);
   for (unsigned i = 1, e = MD->getNumOperands(); i < e; ++i) {
bool SimplifyCFGOpt::PerformValueComparisonIntoPredecessorFolding(
 
   // Based on whether the default edge from PTI goes to BB or not, fill in
   // PredCases and PredDefault with the new switch cases we would like to
   // build.
   SmallMapVector<BasicBlock *, int, 8> NewSuccessors;
 
   // Update the branch weight metadata along the way
   SmallVector<uint64_t, 8> Weights;
-  bool PredHasWeights = HasBranchWeights(PTI);
-  bool SuccHasWeights = HasBranchWeights(TI);
+  bool PredHasWeights = hasBranchWeightMD(*PTI);
+  bool SuccHasWeights = hasBranchWeightMD(*TI);
 
   if (PredHasWeights) {
     GetBranchWeights(PTI, Weights);
     // branch-weight metadata is inconsistent here.
     if (Weights.size() != 1 + PredCases.size())
       PredHasWeights = SuccHasWeights = false;
   } else if (SuccHasWeights)
     // If there are no predecessor weights but there are successor weights,
bool SimplifyCFGOpt::SimplifySwitchOnSelect(SwitchInst *SI,
   // Find the relevant condition and destinations.
   Value *Condition = Select->getCondition();
   BasicBlock *TrueBB = SI->findCaseValue(TrueVal)->getCaseSuccessor();
   BasicBlock *FalseBB = SI->findCaseValue(FalseVal)->getCaseSuccessor();
 
   // Get weight for TrueBB and FalseBB.
   uint32_t TrueWeight = 0, FalseWeight = 0;
   SmallVector<uint64_t, 8> Weights;
-  bool HasWeights = HasBranchWeights(SI);
+  bool HasWeights = hasBranchWeightMD(*SI);
   if (HasWeights) {
     GetBranchWeights(SI, Weights);
     if (Weights.size() == 1 + SI->getNumCases()) {
       TrueWeight =
           (uint32_t)Weights[SI->findCaseValue(TrueVal)->getSuccessorIndex()];
       FalseWeight =
           (uint32_t)Weights[SI->findCaseValue(FalseVal)->getSuccessorIndex()];
     }
bool SimplifyCFGOpt::TurnSwitchRangeIntoICmp(SwitchInst *SI,
   // If NumCases overflowed, then all possible values jump to the successor.
   if (NumCases->isNullValue() && !ContiguousCases->empty())
     Cmp = ConstantInt::getTrue(SI->getContext());
   else
     Cmp = Builder.CreateICmpULT(Sub, NumCases, "switch");
   BranchInst *NewBI = Builder.CreateCondBr(Cmp, ContiguousDest, OtherDest);
 
   // Update weight for the newly-created conditional branch.
-  if (HasBranchWeights(SI)) {
+  if (hasBranchWeightMD(*SI)) {
     SmallVector<uint64_t, 8> Weights;
     GetBranchWeights(SI, Weights);
     if (Weights.size() == 1 + SI->getNumCases()) {
       uint64_t TrueWeight = 0;
       uint64_t FalseWeight = 0;
       for (size_t I = 0, E = Weights.size(); I != E; ++I) {
         if (SI->getSuccessor(I) == ContiguousDest)
           TrueWeight += Weights[I];