Skip to content

[llvm][profdata][NFC] Support 64-bit weights in ProfDataUtils #86607

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions llvm/include/llvm/IR/ProfDataUtils.h
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,14 @@ 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);
/// be called with "branch_weights" metadata nodes. Supports uint32_t.
void extractFromBranchWeightMD32(const MDNode *ProfileData,
SmallVectorImpl<uint32_t> &Weights);

/// Faster version of extractBranchWeights() that skips checks and must only
/// be called with "branch_weights" metadata nodes. Supports uint64_t.
void extractFromBranchWeightMD64(const MDNode *ProfileData,
SmallVectorImpl<uint64_t> &Weights);

/// Extract branch weights attatched to an Instruction
///
Expand Down
42 changes: 27 additions & 15 deletions llvm/lib/IR/ProfDataUtils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,26 @@ bool isTargetMD(const MDNode *ProfData, const char *Name, unsigned MinOps) {
return ProfDataName->getString().equals(Name);
}

template <typename T,
typename = typename std::enable_if<std::is_arithmetic_v<T>>>
static void extractFromBranchWeightMD(const MDNode *ProfileData,
SmallVectorImpl<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();
}
}

} // namespace

namespace llvm {
Expand Down Expand Up @@ -100,22 +120,14 @@ 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);
void extractFromBranchWeightMD32(const MDNode *ProfileData,
SmallVectorImpl<uint32_t> &Weights) {
extractFromBranchWeightMD(ProfileData, Weights);
}

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();
}
void extractFromBranchWeightMD64(const MDNode *ProfileData,
SmallVectorImpl<uint64_t> &Weights) {
extractFromBranchWeightMD(ProfileData, Weights);
}

bool extractBranchWeights(const MDNode *ProfileData,
Expand Down
2 changes: 1 addition & 1 deletion llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ static void updateBranchWeights(BranchInst &PreHeaderBI, BranchInst &LoopBI,
return;

SmallVector<uint32_t, 2> Weights;
extractFromBranchWeightMD(WeightMD, Weights);
extractFromBranchWeightMD32(WeightMD, Weights);
if (Weights.size() != 2)
return;
uint32_t OrigLoopExitWeight = Weights[0];
Expand Down
7 changes: 2 additions & 5 deletions llvm/lib/Transforms/Utils/SimplifyCFG.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1066,11 +1066,8 @@ static int ConstantIntSortPredicate(ConstantInt *const *P1,
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) {
ConstantInt *CI = mdconst::extract<ConstantInt>(MD->getOperand(i));
Weights.push_back(CI->getValue().getZExtValue());
}
assert(MD && "Invalid branch-weight metadata");
extractFromBranchWeightMD64(MD, Weights);

// If TI is a conditional eq, the default case is the false case,
// and the corresponding branch-weight data is at index 2. We swap the
Expand Down
Loading