1//===-------- MIRSampleProfile.cpp: MIRSampleFDO (For FSAFDO) -------------===//
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 provides the implementation of the MIRSampleProfile loader, mainly
10// for flow sensitive SampleFDO.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/CodeGen/MIRSampleProfile.h"
15#include "llvm/ADT/DenseMap.h"
16#include "llvm/ADT/DenseSet.h"
17#include "llvm/Analysis/BlockFrequencyInfoImpl.h"
18#include "llvm/CodeGen/MIRFSDiscriminatorOptions.h"
19#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
20#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
21#include "llvm/CodeGen/MachineDominators.h"
22#include "llvm/CodeGen/MachineInstr.h"
23#include "llvm/CodeGen/MachineLoopInfo.h"
24#include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
25#include "llvm/CodeGen/MachinePostDominators.h"
26#include "llvm/CodeGen/Passes.h"
27#include "llvm/IR/Function.h"
28#include "llvm/IR/PseudoProbe.h"
29#include "llvm/InitializePasses.h"
30#include "llvm/Support/CommandLine.h"
31#include "llvm/Support/Debug.h"
32#include "llvm/Support/VirtualFileSystem.h"
33#include "llvm/Support/raw_ostream.h"
34#include "llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h"
35#include "llvm/Transforms/Utils/SampleProfileLoaderBaseUtil.h"
36#include <optional>
37
38using namespace llvm;
39using namespace sampleprof;
40using namespace llvm::sampleprofutil;
41
42#define DEBUG_TYPE "fs-profile-loader"
43
44static cl::opt<bool> ShowFSBranchProb(
45 "show-fs-branchprob", cl::Hidden, cl::init(Val: false),
46 cl::desc("Print setting flow sensitive branch probabilities"));
47static cl::opt<unsigned> FSProfileDebugProbDiffThreshold(
48 "fs-profile-debug-prob-diff-threshold", cl::init(Val: 10),
49 cl::desc(
50 "Only show debug message if the branch probability is greater than "
51 "this value (in percentage)."));
52
53static cl::opt<unsigned> FSProfileDebugBWThreshold(
54 "fs-profile-debug-bw-threshold", cl::init(Val: 10000),
55 cl::desc("Only show debug message if the source branch weight is greater "
56 " than this value."));
57
58static cl::opt<bool> ViewBFIBefore("fs-viewbfi-before", cl::Hidden,
59 cl::init(Val: false),
60 cl::desc("View BFI before MIR loader"));
61static cl::opt<bool> ViewBFIAfter("fs-viewbfi-after", cl::Hidden,
62 cl::init(Val: false),
63 cl::desc("View BFI after MIR loader"));
64
65char MIRProfileLoaderPass::ID = 0;
66
67INITIALIZE_PASS_BEGIN(MIRProfileLoaderPass, DEBUG_TYPE,
68 "Load MIR Sample Profile",
69 /* cfg = */ false, /* is_analysis = */ false)
70INITIALIZE_PASS_DEPENDENCY(MachineBlockFrequencyInfoWrapperPass)
71INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
72INITIALIZE_PASS_DEPENDENCY(MachinePostDominatorTreeWrapperPass)
73INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
74INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
75INITIALIZE_PASS_END(MIRProfileLoaderPass, DEBUG_TYPE, "Load MIR Sample Profile",
76 /* cfg = */ false, /* is_analysis = */ false)
77
78char &llvm::MIRProfileLoaderPassID = MIRProfileLoaderPass::ID;
79
80FunctionPass *
81llvm::createMIRProfileLoaderPass(std::string File, std::string RemappingFile,
82 FSDiscriminatorPass P,
83 IntrusiveRefCntPtr<vfs::FileSystem> FS) {
84 return new MIRProfileLoaderPass(File, RemappingFile, P, std::move(FS));
85}
86
87namespace llvm {
88
89// Internal option used to control BFI display only after MBP pass.
90// Defined in CodeGen/MachineBlockFrequencyInfo.cpp:
91// -view-block-layout-with-bfi={none | fraction | integer | count}
92extern cl::opt<GVDAGType> ViewBlockLayoutWithBFI;
93
94// Command line option to specify the name of the function for CFG dump
95// Defined in Analysis/BlockFrequencyInfo.cpp: -view-bfi-func-name=
96extern cl::opt<std::string> ViewBlockFreqFuncName;
97
98std::optional<PseudoProbe> extractProbe(const MachineInstr &MI) {
99 if (MI.isPseudoProbe()) {
100 PseudoProbe Probe;
101 Probe.Id = MI.getOperand(i: 1).getImm();
102 Probe.Type = MI.getOperand(i: 2).getImm();
103 Probe.Attr = MI.getOperand(i: 3).getImm();
104 Probe.Factor = 1;
105 DILocation *DebugLoc = MI.getDebugLoc();
106 Probe.Discriminator = DebugLoc ? DebugLoc->getDiscriminator() : 0;
107 return Probe;
108 }
109
110 // Ignore callsite probes since they do not have FS discriminators.
111 return std::nullopt;
112}
113
114namespace afdo_detail {
115template <> struct IRTraits<MachineBasicBlock> {
116 using InstructionT = MachineInstr;
117 using BasicBlockT = MachineBasicBlock;
118 using FunctionT = MachineFunction;
119 using BlockFrequencyInfoT = MachineBlockFrequencyInfo;
120 using LoopT = MachineLoop;
121 using LoopInfoPtrT = MachineLoopInfo *;
122 using DominatorTreePtrT = MachineDominatorTree *;
123 using PostDominatorTreePtrT = MachinePostDominatorTree *;
124 using PostDominatorTreeT = MachinePostDominatorTree;
125 using OptRemarkEmitterT = MachineOptimizationRemarkEmitter;
126 using OptRemarkAnalysisT = MachineOptimizationRemarkAnalysis;
127 using PredRangeT =
128 iterator_range<SmallVectorImpl<MachineBasicBlock *>::iterator>;
129 using SuccRangeT =
130 iterator_range<SmallVectorImpl<MachineBasicBlock *>::iterator>;
131 static Function &getFunction(MachineFunction &F) { return F.getFunction(); }
132 static const MachineBasicBlock *getEntryBB(const MachineFunction *F) {
133 return GraphTraits<const MachineFunction *>::getEntryNode(F);
134 }
135 static PredRangeT getPredecessors(MachineBasicBlock *BB) {
136 return BB->predecessors();
137 }
138 static SuccRangeT getSuccessors(MachineBasicBlock *BB) {
139 return BB->successors();
140 }
141};
142} // namespace afdo_detail
143
144class MIRProfileLoader final
145 : public SampleProfileLoaderBaseImpl<MachineFunction> {
146public:
147 void setInitVals(MachineDominatorTree *MDT, MachinePostDominatorTree *MPDT,
148 MachineLoopInfo *MLI, MachineBlockFrequencyInfo *MBFI,
149 MachineOptimizationRemarkEmitter *MORE) {
150 DT = MDT;
151 PDT = MPDT;
152 LI = MLI;
153 BFI = MBFI;
154 ORE = MORE;
155 }
156 void setFSPass(FSDiscriminatorPass Pass) {
157 P = Pass;
158 LowBit = getFSPassBitBegin(P);
159 HighBit = getFSPassBitEnd(P);
160 assert(LowBit < HighBit && "HighBit needs to be greater than Lowbit");
161 }
162
163 MIRProfileLoader(StringRef Name, StringRef RemapName,
164 IntrusiveRefCntPtr<vfs::FileSystem> FS)
165 : SampleProfileLoaderBaseImpl(std::string(Name), std::string(RemapName),
166 std::move(FS)) {}
167
168 void setBranchProbs(MachineFunction &F);
169 bool runOnFunction(MachineFunction &F);
170 bool doInitialization(Module &M);
171 bool isValid() const { return ProfileIsValid; }
172
173protected:
174 friend class SampleCoverageTracker;
175
176 /// Hold the information of the basic block frequency.
177 MachineBlockFrequencyInfo *BFI;
178
179 /// PassNum is the sequence number this pass is called, start from 1.
180 FSDiscriminatorPass P;
181
182 // LowBit in the FS discriminator used by this instance. Note the number is
183 // 0-based. Base discrimnator use bit 0 to bit 11.
184 unsigned LowBit;
185 // HighwBit in the FS discriminator used by this instance. Note the number
186 // is 0-based.
187 unsigned HighBit;
188
189 bool ProfileIsValid = true;
190 ErrorOr<uint64_t> getInstWeight(const MachineInstr &MI) override {
191 if (FunctionSamples::ProfileIsProbeBased)
192 return getProbeWeight(Inst: MI);
193 if (ImprovedFSDiscriminator && MI.isMetaInstruction())
194 return std::error_code();
195 return getInstWeightImpl(Inst: MI);
196 }
197};
198
199template <>
200void SampleProfileLoaderBaseImpl<MachineFunction>::computeDominanceAndLoopInfo(
201 MachineFunction &F) {}
202
203void MIRProfileLoader::setBranchProbs(MachineFunction &F) {
204 LLVM_DEBUG(dbgs() << "\nPropagation complete. Setting branch probs\n");
205 for (auto &BI : F) {
206 MachineBasicBlock *BB = &BI;
207 if (BB->succ_size() < 2)
208 continue;
209 const MachineBasicBlock *EC = EquivalenceClass[BB];
210 uint64_t BBWeight = BlockWeights[EC];
211 uint64_t SumEdgeWeight = 0;
212 for (MachineBasicBlock *Succ : BB->successors()) {
213 Edge E = std::make_pair(x&: BB, y&: Succ);
214 SumEdgeWeight += EdgeWeights[E];
215 }
216
217 if (BBWeight != SumEdgeWeight) {
218 LLVM_DEBUG(dbgs() << "BBweight is not equal to SumEdgeWeight: BBWWeight="
219 << BBWeight << " SumEdgeWeight= " << SumEdgeWeight
220 << "\n");
221 BBWeight = SumEdgeWeight;
222 }
223 if (BBWeight == 0) {
224 LLVM_DEBUG(dbgs() << "SKIPPED. All branch weights are zero.\n");
225 continue;
226 }
227
228#ifndef NDEBUG
229 uint64_t BBWeightOrig = BBWeight;
230#endif
231 uint32_t MaxWeight = std::numeric_limits<uint32_t>::max();
232 uint32_t Factor = 1;
233 if (BBWeight > MaxWeight) {
234 Factor = BBWeight / MaxWeight + 1;
235 BBWeight /= Factor;
236 LLVM_DEBUG(dbgs() << "Scaling weights by " << Factor << "\n");
237 }
238
239 for (MachineBasicBlock::succ_iterator SI = BB->succ_begin(),
240 SE = BB->succ_end();
241 SI != SE; ++SI) {
242 MachineBasicBlock *Succ = *SI;
243 Edge E = std::make_pair(x&: BB, y&: Succ);
244 uint64_t EdgeWeight = EdgeWeights[E];
245 EdgeWeight /= Factor;
246
247 assert(BBWeight >= EdgeWeight &&
248 "BBweight is larger than EdgeWeight -- should not happen.\n");
249
250 BranchProbability OldProb = BB->getSuccProbability(Succ: SI);
251 BranchProbability NewProb(EdgeWeight, BBWeight);
252 if (OldProb == NewProb)
253 continue;
254 BB->setSuccProbability(I: SI, Prob: NewProb);
255#ifndef NDEBUG
256 if (!ShowFSBranchProb)
257 continue;
258 bool Show = false;
259 BranchProbability Diff;
260 if (OldProb > NewProb)
261 Diff = OldProb - NewProb;
262 else
263 Diff = NewProb - OldProb;
264 Show = (Diff >= BranchProbability(FSProfileDebugProbDiffThreshold, 100));
265 Show &= (BBWeightOrig >= FSProfileDebugBWThreshold);
266
267 auto DIL = BB->findBranchDebugLoc();
268 auto SuccDIL = Succ->findBranchDebugLoc();
269 if (Show) {
270 dbgs() << "Set branch fs prob: MBB (" << BB->getNumber() << " -> "
271 << Succ->getNumber() << "): ";
272 if (DIL)
273 dbgs() << DIL->getFilename() << ":" << DIL->getLine() << ":"
274 << DIL->getColumn();
275 if (SuccDIL)
276 dbgs() << "-->" << SuccDIL->getFilename() << ":" << SuccDIL->getLine()
277 << ":" << SuccDIL->getColumn();
278 dbgs() << " W=" << BBWeightOrig << " " << OldProb << " --> " << NewProb
279 << "\n";
280 }
281#endif
282 }
283 }
284}
285
286bool MIRProfileLoader::doInitialization(Module &M) {
287 auto &Ctx = M.getContext();
288
289 auto ReaderOrErr = sampleprof::SampleProfileReader::create(
290 Filename, C&: Ctx, FS&: *FS, P, RemapFilename: RemappingFilename);
291 if (std::error_code EC = ReaderOrErr.getError()) {
292 std::string Msg = "Could not open profile: " + EC.message();
293 Ctx.diagnose(DI: DiagnosticInfoSampleProfile(Filename, Msg));
294 return false;
295 }
296
297 Reader = std::move(ReaderOrErr.get());
298 Reader->setModule(&M);
299 ProfileIsValid = (Reader->read() == sampleprof_error::success);
300
301 // Load pseudo probe descriptors for probe-based function samples.
302 if (Reader->profileIsProbeBased()) {
303 ProbeManager = std::make_unique<PseudoProbeManager>(args&: M);
304 if (!ProbeManager->moduleIsProbed(M)) {
305 return false;
306 }
307 }
308
309 return true;
310}
311
312bool MIRProfileLoader::runOnFunction(MachineFunction &MF) {
313 // Do not load non-FS profiles. A line or probe can get a zero-valued
314 // discriminator at certain pass which could result in accidentally loading
315 // the corresponding base counter in the non-FS profile, while a non-zero
316 // discriminator would end up getting zero samples. This could in turn undo
317 // the sample distribution effort done by previous BFI maintenance and the
318 // probe distribution factor work for pseudo probes.
319 if (!Reader->profileIsFS())
320 return false;
321
322 Function &Func = MF.getFunction();
323 clearFunctionData(ResetDT: false);
324 Samples = Reader->getSamplesFor(F: Func);
325 if (!Samples || Samples->empty())
326 return false;
327
328 if (FunctionSamples::ProfileIsProbeBased) {
329 if (!ProbeManager->profileIsValid(F: MF.getFunction(), Samples: *Samples))
330 return false;
331 } else {
332 if (getFunctionLoc(F&: MF) == 0)
333 return false;
334 }
335
336 DenseSet<GlobalValue::GUID> InlinedGUIDs;
337 bool Changed = computeAndPropagateWeights(F&: MF, InlinedGUIDs);
338
339 // Set the new BPI, BFI.
340 setBranchProbs(MF);
341
342 return Changed;
343}
344
345} // namespace llvm
346
347MIRProfileLoaderPass::MIRProfileLoaderPass(
348 std::string FileName, std::string RemappingFileName, FSDiscriminatorPass P,
349 IntrusiveRefCntPtr<vfs::FileSystem> FS)
350 : MachineFunctionPass(ID), ProfileFileName(FileName), P(P) {
351 LowBit = getFSPassBitBegin(P);
352 HighBit = getFSPassBitEnd(P);
353
354 auto VFS = FS ? std::move(FS) : vfs::getRealFileSystem();
355 MIRSampleLoader = std::make_unique<MIRProfileLoader>(
356 args&: FileName, args&: RemappingFileName, args: std::move(VFS));
357 assert(LowBit < HighBit && "HighBit needs to be greater than Lowbit");
358}
359
360bool MIRProfileLoaderPass::runOnMachineFunction(MachineFunction &MF) {
361 if (!MIRSampleLoader->isValid())
362 return false;
363
364 LLVM_DEBUG(dbgs() << "MIRProfileLoader pass working on Func: "
365 << MF.getFunction().getName() << "\n");
366 MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
367 auto *MDT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
368 auto *MPDT =
369 &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
370
371 MF.RenumberBlocks();
372
373 MIRSampleLoader->setInitVals(
374 MDT, MPDT, MLI: &getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI,
375 MORE: &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE());
376
377 if (ViewBFIBefore && ViewBlockLayoutWithBFI != GVDT_None &&
378 (ViewBlockFreqFuncName.empty() ||
379 MF.getFunction().getName() == ViewBlockFreqFuncName)) {
380 MBFI->view(Name: "MIR_Prof_loader_b." + MF.getName(), isSimple: false);
381 }
382
383 bool Changed = MIRSampleLoader->runOnFunction(MF);
384 if (Changed)
385 MBFI->calculate(F: MF, MBPI: *MBFI->getMBPI(),
386 MLI: *&getAnalysis<MachineLoopInfoWrapperPass>().getLI());
387
388 if (ViewBFIAfter && ViewBlockLayoutWithBFI != GVDT_None &&
389 (ViewBlockFreqFuncName.empty() ||
390 MF.getFunction().getName() == ViewBlockFreqFuncName)) {
391 MBFI->view(Name: "MIR_prof_loader_a." + MF.getName(), isSimple: false);
392 }
393
394 return Changed;
395}
396
397bool MIRProfileLoaderPass::doInitialization(Module &M) {
398 LLVM_DEBUG(dbgs() << "MIRProfileLoader pass working on Module " << M.getName()
399 << "\n");
400
401 MIRSampleLoader->setFSPass(P);
402 return MIRSampleLoader->doInitialization(M);
403}
404
405void MIRProfileLoaderPass::getAnalysisUsage(AnalysisUsage &AU) const {
406 AU.setPreservesAll();
407 AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
408 AU.addRequired<MachineDominatorTreeWrapperPass>();
409 AU.addRequired<MachinePostDominatorTreeWrapperPass>();
410 AU.addRequiredTransitive<MachineLoopInfoWrapperPass>();
411 AU.addRequired<MachineOptimizationRemarkEmitterPass>();
412 MachineFunctionPass::getAnalysisUsage(AU);
413}
414