1 //===- ProfDataUtils.cpp - Utility functions for MD_prof Metadata ---------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements utilities for working with Profiling Metadata.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "llvm/IR/ProfDataUtils.h"
14 #include "llvm/ADT/SmallVector.h"
15 #include "llvm/ADT/Twine.h"
16 #include "llvm/IR/Constants.h"
17 #include "llvm/IR/Function.h"
18 #include "llvm/IR/Instructions.h"
19 #include "llvm/IR/LLVMContext.h"
20 #include "llvm/IR/Metadata.h"
21 #include "llvm/Support/BranchProbability.h"
22 #include "llvm/Support/CommandLine.h"
23 
24 using namespace llvm;
25 
26 namespace {
27 
28 // MD_prof nodes have the following layout
29 //
30 // In general:
31 // { String name,         Array of i32   }
32 //
33 // In terms of Types:
34 // { MDString,            [i32, i32, ...]}
35 //
36 // Concretely for Branch Weights
37 // { "branch_weights",    [i32 1, i32 10000]}
38 //
39 // We maintain some constants here to ensure that we access the branch weights
40 // correctly, and can change the behavior in the future if the layout changes
41 
42 // The index at which the weights vector starts
43 constexpr unsigned WeightsIdx = 1;
44 
45 // the minimum number of operands for MD_prof nodes with branch weights
46 constexpr unsigned MinBWOps = 3;
47 
48 bool extractWeights(const MDNode *ProfileData,
49                     SmallVectorImpl<uint32_t> &Weights) {
50   // Assume preconditions are already met (i.e. this is valid metadata)
51   assert(ProfileData && "ProfileData was nullptr in extractWeights");
52   unsigned NOps = ProfileData->getNumOperands();
53 
54   assert(WeightsIdx < NOps && "Weights Index must be less than NOps.");
55   Weights.resize(NOps - WeightsIdx);
56 
57   for (unsigned Idx = WeightsIdx, E = NOps; Idx != E; ++Idx) {
58     ConstantInt *Weight =
59         mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
60     assert(Weight && "Malformed branch_weight in MD_prof node");
61     assert(Weight->getValue().getActiveBits() <= 32 &&
62            "Too many bits for uint32_t");
63     Weights[Idx - WeightsIdx] = Weight->getZExtValue();
64   }
65   return true;
66 }
67 
68 // We may want to add support for other MD_prof types, so provide an abstraction
69 // for checking the metadata type.
70 bool isTargetMD(const MDNode *ProfData, const char *Name, unsigned MinOps) {
71   // TODO: This routine may be simplified if MD_prof used an enum instead of a
72   // string to differentiate the types of MD_prof nodes.
73   if (!ProfData || !Name || MinOps < 2)
74     return false;
75 
76   unsigned NOps = ProfData->getNumOperands();
77   if (NOps < MinOps)
78     return false;
79 
80   auto *ProfDataName = dyn_cast<MDString>(ProfData->getOperand(0));
81   if (!ProfDataName)
82     return false;
83 
84   return ProfDataName->getString().equals(Name);
85 }
86 
87 } // namespace
88 
89 namespace llvm {
90 
91 bool hasProfMD(const Instruction &I) {
92   return nullptr != I.getMetadata(LLVMContext::MD_prof);
93 }
94 
95 bool isBranchWeightMD(const MDNode *ProfileData) {
96   return isTargetMD(ProfileData, "branch_weights", MinBWOps);
97 }
98 
99 bool hasBranchWeightMD(const Instruction &I) {
100   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
101   return isBranchWeightMD(ProfileData);
102 }
103 
104 bool hasValidBranchWeightMD(const Instruction &I) {
105   return getValidBranchWeightMDNode(I);
106 }
107 
108 MDNode *getBranchWeightMDNode(const Instruction &I) {
109   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
110   if (!isBranchWeightMD(ProfileData))
111     return nullptr;
112   return ProfileData;
113 }
114 
115 MDNode *getValidBranchWeightMDNode(const Instruction &I) {
116   auto *ProfileData = getBranchWeightMDNode(I);
117   if (ProfileData && ProfileData->getNumOperands() == 1 + I.getNumSuccessors())
118     return ProfileData;
119   return nullptr;
120 }
121 
122 bool extractBranchWeights(const MDNode *ProfileData,
123                           SmallVectorImpl<uint32_t> &Weights) {
124   if (!isBranchWeightMD(ProfileData))
125     return false;
126   return extractWeights(ProfileData, Weights);
127 }
128 
129 bool extractBranchWeights(const Instruction &I,
130                           SmallVectorImpl<uint32_t> &Weights) {
131   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
132   return extractBranchWeights(ProfileData, Weights);
133 }
134 
135 bool extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
136                           uint64_t &FalseVal) {
137   assert((I.getOpcode() == Instruction::Br ||
138           I.getOpcode() == Instruction::Select) &&
139          "Looking for branch weights on something besides branch, select, or "
140          "switch");
141 
142   SmallVector<uint32_t, 2> Weights;
143   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
144   if (!extractBranchWeights(ProfileData, Weights))
145     return false;
146 
147   if (Weights.size() > 2)
148     return false;
149 
150   TrueVal = Weights[0];
151   FalseVal = Weights[1];
152   return true;
153 }
154 
155 bool extractProfTotalWeight(const MDNode *ProfileData, uint64_t &TotalVal) {
156   TotalVal = 0;
157   if (!ProfileData)
158     return false;
159 
160   auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
161   if (!ProfDataName)
162     return false;
163 
164   if (ProfDataName->getString().equals("branch_weights")) {
165     for (unsigned Idx = 1; Idx < ProfileData->getNumOperands(); Idx++) {
166       auto *V = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
167       assert(V && "Malformed branch_weight in MD_prof node");
168       TotalVal += V->getValue().getZExtValue();
169     }
170     return true;
171   }
172 
173   if (ProfDataName->getString().equals("VP") &&
174       ProfileData->getNumOperands() > 3) {
175     TotalVal = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(2))
176                    ->getValue()
177                    .getZExtValue();
178     return true;
179   }
180   return false;
181 }
182 
183 bool extractProfTotalWeight(const Instruction &I, uint64_t &TotalVal) {
184   return extractProfTotalWeight(I.getMetadata(LLVMContext::MD_prof), TotalVal);
185 }
186 
187 } // namespace llvm
188