[llvm] 285e023 - ProfDataUtils: Add extractFromBranchWeightMD function; NFC

[Matthias Braun via llvm-commits]

Author: Matthias Braun
Date: 2023-09-11T10:38:06-07:00
New Revision: 285e0235f5f649e17eef8f0f25e1680c975dcfe4

URL: https://github.com/llvm/llvm-project/commit/285e0235f5f649e17eef8f0f25e1680c975dcfe4
DIFF: https://github.com/llvm/llvm-project/commit/285e0235f5f649e17eef8f0f25e1680c975dcfe4.diff

LOG: ProfDataUtils: Add extractFromBranchWeightMD function; NFC

Expose internal helper function as new `extractFromBranchWeightMD` API.

Differential revision: https://reviews.llvm.org/D157937

Added: 
    

Modified: 
    llvm/include/llvm/IR/ProfDataUtils.h
    llvm/lib/IR/ProfDataUtils.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/ProfDataUtils.h b/llvm/include/llvm/IR/ProfDataUtils.h
index 8019d0a3969b10e..b61199372de0de8 100644
--- a/llvm/include/llvm/IR/ProfDataUtils.h
+++ b/llvm/include/llvm/IR/ProfDataUtils.h
@@ -64,6 +64,11 @@ MDNode *getValidBranchWeightMDNode(const Instruction &I);
 bool extractBranchWeights(const MDNode *ProfileData,
                           SmallVectorImpl<uint32_t> &Weights);
 
+/// Faster version of extractBranchWeights() that skips checks and must only
+/// be called with "branch_weights" metadata nodes.
+void extractFromBranchWeightMD(const MDNode *ProfileData,
+                               SmallVectorImpl<uint32_t> &Weights);
+
 /// Extract branch weights attatched to an Instruction
 ///
 /// \param I The Instruction to extract weights from.

diff  --git a/llvm/lib/IR/ProfDataUtils.cpp b/llvm/lib/IR/ProfDataUtils.cpp
index e534368b05e4164..77b3c1cb95d686c 100644
--- a/llvm/lib/IR/ProfDataUtils.cpp
+++ b/llvm/lib/IR/ProfDataUtils.cpp
@@ -45,26 +45,6 @@ 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.");
-  Weights.resize(NOps - WeightsIdx);
-
-  for (unsigned Idx = WeightsIdx, E = NOps; Idx != E; ++Idx) {
-    ConstantInt *Weight =
-        mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
-    assert(Weight && "Malformed branch_weight in MD_prof node");
-    assert(Weight->getValue().getActiveBits() <= 32 &&
-           "Too many bits for uint32_t");
-    Weights[Idx - WeightsIdx] = Weight->getZExtValue();
-  }
-  return true;
-}
-
 // We may want to add support for other MD_prof types, so provide an abstraction
 // for checking the metadata type.
 bool isTargetMD(const MDNode *ProfData, const char *Name, unsigned MinOps) {
@@ -119,11 +99,30 @@ MDNode *getValidBranchWeightMDNode(const Instruction &I) {
   return nullptr;
 }
 
+void extractFromBranchWeightMD(const MDNode *ProfileData,
+                               SmallVectorImpl<uint32_t> &Weights) {
+  assert(isBranchWeightMD(ProfileData) && "wrong metadata");
+
+  unsigned NOps = ProfileData->getNumOperands();
+  assert(WeightsIdx < NOps && "Weights Index must be less than NOps.");
+  Weights.resize(NOps - WeightsIdx);
+
+  for (unsigned Idx = WeightsIdx, E = NOps; Idx != E; ++Idx) {
+    ConstantInt *Weight =
+        mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
+    assert(Weight && "Malformed branch_weight in MD_prof node");
+    assert(Weight->getValue().getActiveBits() <= 32 &&
+           "Too many bits for uint32_t");
+    Weights[Idx - WeightsIdx] = Weight->getZExtValue();
+  }
+}
+
 bool extractBranchWeights(const MDNode *ProfileData,
                           SmallVectorImpl<uint32_t> &Weights) {
   if (!isBranchWeightMD(ProfileData))
     return false;
-  return extractWeights(ProfileData, Weights);
+  extractFromBranchWeightMD(ProfileData, Weights);
+  return true;
 }
 
 bool extractBranchWeights(const Instruction &I,