1//===-- InlineAsmPrepare - Prepare inline asm for code gen ----------------===//
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 pass lowers callbrs in LLVM IR in order to to assist SelectionDAG's
10// codegen.
11//
12// In particular, this pass assists in inserting register copies for the output
13// values of a callbr along the edges leading to the indirect target blocks.
14// Though the output SSA value is defined by the callbr instruction itself in
15// the IR representation, the value cannot be copied to the appropriate virtual
16// registers prior to jumping to an indirect label, since the jump occurs
17// within the user-provided assembly blob.
18//
19// Instead, those copies must occur separately at the beginning of each
20// indirect target. That requires that we create a separate SSA definition in
21// each of them (via llvm.callbr.landingpad), and may require splitting
22// critical edges so we have a location to place the intrinsic. Finally, we
23// remap users of the original callbr output SSA value to instead point to the
24// appropriate llvm.callbr.landingpad value.
25//
26// Ideally, this could be done inside SelectionDAG, or in the
27// MachineInstruction representation, without the use of an IR-level intrinsic.
28// But, within the current framework, it’s simpler to implement as an IR pass.
29// (If support for callbr in GlobalISel is implemented, it’s worth considering
30// whether this is still required.)
31//
32//===----------------------------------------------------------------------===//
33
34#include "llvm/CodeGen/InlineAsmPrepare.h"
35#include "llvm/ADT/ArrayRef.h"
36#include "llvm/ADT/SmallPtrSet.h"
37#include "llvm/ADT/SmallVector.h"
38#include "llvm/ADT/iterator.h"
39#include "llvm/Analysis/CFG.h"
40#include "llvm/CodeGen/Passes.h"
41#include "llvm/IR/BasicBlock.h"
42#include "llvm/IR/Dominators.h"
43#include "llvm/IR/Function.h"
44#include "llvm/IR/IRBuilder.h"
45#include "llvm/IR/Instructions.h"
46#include "llvm/IR/IntrinsicInst.h"
47#include "llvm/IR/Intrinsics.h"
48#include "llvm/InitializePasses.h"
49#include "llvm/Pass.h"
50#include "llvm/Transforms/Utils/BasicBlockUtils.h"
51#include "llvm/Transforms/Utils/SSAUpdater.h"
52
53using namespace llvm;
54
55#define DEBUG_TYPE "inline-asm-prepare"
56
57static bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT);
58static bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs,
59 DominatorTree &DT);
60static void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,
61 SSAUpdater &SSAUpdate);
62static SmallVector<CallBrInst *, 2> FindCallBrs(Function &F);
63
64namespace {
65
66class InlineAsmPrepare : public FunctionPass {
67public:
68 InlineAsmPrepare() : FunctionPass(ID) {}
69 void getAnalysisUsage(AnalysisUsage &AU) const override;
70 bool runOnFunction(Function &F) override;
71 static char ID;
72};
73
74} // end anonymous namespace
75
76PreservedAnalyses InlineAsmPreparePass::run(Function &F,
77 FunctionAnalysisManager &FAM) {
78 bool Changed = false;
79 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(F);
80
81 if (CBRs.empty())
82 return PreservedAnalyses::all();
83
84 auto &DT = FAM.getResult<DominatorTreeAnalysis>(IR&: F);
85
86 Changed |= SplitCriticalEdges(CBRs, DT);
87 Changed |= InsertIntrinsicCalls(CBRs, DT);
88
89 if (!Changed)
90 return PreservedAnalyses::all();
91 PreservedAnalyses PA;
92 PA.preserve<DominatorTreeAnalysis>();
93 return PA;
94}
95
96char InlineAsmPrepare::ID = 0;
97INITIALIZE_PASS_BEGIN(InlineAsmPrepare, "inline-asm-prepare",
98 "Prepare inline asm insts", false, false)
99INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
100INITIALIZE_PASS_END(InlineAsmPrepare, "inline-asm-prepare",
101 "Prepare inline asm insts", false, false)
102
103FunctionPass *llvm::createInlineAsmPreparePass() {
104 return new InlineAsmPrepare();
105}
106
107void InlineAsmPrepare::getAnalysisUsage(AnalysisUsage &AU) const {
108 AU.addPreserved<DominatorTreeWrapperPass>();
109}
110
111SmallVector<CallBrInst *, 2> FindCallBrs(Function &F) {
112 SmallVector<CallBrInst *, 2> CBRs;
113 for (BasicBlock &BB : F)
114 if (auto *CBR = dyn_cast<CallBrInst>(Val: BB.getTerminator()))
115 if (!CBR->getType()->isVoidTy() && !CBR->use_empty())
116 CBRs.push_back(Elt: CBR);
117 return CBRs;
118}
119
120bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) {
121 bool Changed = false;
122 CriticalEdgeSplittingOptions Options(&DT);
123 Options.setMergeIdenticalEdges();
124
125 // The indirect destination might be duplicated between another parameter...
126 // %0 = callbr ... [label %x, label %x]
127 // ...hence MergeIdenticalEdges and AllowIndentical edges, but we don't need
128 // to split the default destination if it's duplicated between an indirect
129 // destination...
130 // %1 = callbr ... to label %x [label %x]
131 // ...hence starting at 1 and checking against successor 0 (aka the default
132 // destination).
133 for (CallBrInst *CBR : CBRs)
134 for (unsigned i = 1, e = CBR->getNumSuccessors(); i != e; ++i)
135 if (CBR->getSuccessor(i) == CBR->getSuccessor(i: 0) ||
136 isCriticalEdge(TI: CBR, SuccNum: i, /*AllowIdenticalEdges*/ true))
137 if (SplitKnownCriticalEdge(TI: CBR, SuccNum: i, Options))
138 Changed = true;
139 return Changed;
140}
141
142bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) {
143 bool Changed = false;
144 SmallPtrSet<const BasicBlock *, 4> Visited;
145 IRBuilder<> Builder(CBRs[0]->getContext());
146 for (CallBrInst *CBR : CBRs) {
147 if (!CBR->getNumIndirectDests())
148 continue;
149
150 SSAUpdater SSAUpdate;
151 SSAUpdate.Initialize(Ty: CBR->getType(), Name: CBR->getName());
152 SSAUpdate.AddAvailableValue(BB: CBR->getParent(), V: CBR);
153 SSAUpdate.AddAvailableValue(BB: CBR->getDefaultDest(), V: CBR);
154
155 for (BasicBlock *IndDest : CBR->getIndirectDests()) {
156 if (!Visited.insert(Ptr: IndDest).second)
157 continue;
158 Builder.SetInsertPoint(&*IndDest->begin());
159 CallInst *Intrinsic = Builder.CreateIntrinsic(
160 RetTy: CBR->getType(), ID: Intrinsic::callbr_landingpad, Args: {CBR});
161 SSAUpdate.AddAvailableValue(BB: IndDest, V: Intrinsic);
162 UpdateSSA(DT, CBR, Intrinsic, SSAUpdate);
163 Changed = true;
164 }
165 }
166 return Changed;
167}
168
169static bool IsInSameBasicBlock(const Use &U, const BasicBlock *BB) {
170 const auto *I = dyn_cast<Instruction>(Val: U.getUser());
171 return I && I->getParent() == BB;
172}
173
174#ifndef NDEBUG
175static void PrintDebugDomInfo(const DominatorTree &DT, const Use &U,
176 const BasicBlock *BB, bool IsDefaultDest) {
177 if (!isa<Instruction>(U.getUser()))
178 return;
179 LLVM_DEBUG(dbgs() << "Use: " << *U.getUser() << ", in block "
180 << cast<Instruction>(U.getUser())->getParent()->getName()
181 << ", is " << (DT.dominates(BB, U) ? "" : "NOT ")
182 << "dominated by " << BB->getName() << " ("
183 << (IsDefaultDest ? "in" : "") << "direct)\n");
184}
185#endif
186
187void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,
188 SSAUpdater &SSAUpdate) {
189
190 SmallPtrSet<Use *, 4> Visited;
191 BasicBlock *DefaultDest = CBR->getDefaultDest();
192 BasicBlock *LandingPad = Intrinsic->getParent();
193
194 SmallVector<Use *, 4> Uses(make_pointer_range(Range: CBR->uses()));
195 for (Use *U : Uses) {
196 if (!Visited.insert(Ptr: U).second)
197 continue;
198
199#ifndef NDEBUG
200 PrintDebugDomInfo(DT, *U, LandingPad, /*IsDefaultDest*/ false);
201 PrintDebugDomInfo(DT, *U, DefaultDest, /*IsDefaultDest*/ true);
202#endif
203
204 // Don't rewrite the use in the newly inserted intrinsic.
205 if (const auto *II = dyn_cast<IntrinsicInst>(Val: U->getUser()))
206 if (II->getIntrinsicID() == Intrinsic::callbr_landingpad)
207 continue;
208
209 // If the Use is in the same BasicBlock as the Intrinsic call, replace
210 // the Use with the value of the Intrinsic call.
211 if (IsInSameBasicBlock(U: *U, BB: LandingPad)) {
212 U->set(Intrinsic);
213 continue;
214 }
215
216 // If the Use is dominated by the default dest, do not touch it.
217 if (DT.dominates(BB: DefaultDest, U: *U))
218 continue;
219
220 SSAUpdate.RewriteUse(U&: *U);
221 }
222}
223
224bool InlineAsmPrepare::runOnFunction(Function &F) {
225 bool Changed = false;
226 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(F);
227
228 if (CBRs.empty())
229 return Changed;
230
231 // It's highly likely that most programs do not contain CallBrInsts. Follow a
232 // similar pattern from SafeStackLegacyPass::runOnFunction to reuse previous
233 // domtree analysis if available, otherwise compute it lazily. This avoids
234 // forcing Dominator Tree Construction at -O0 for programs that likely do not
235 // contain CallBrInsts. It does pessimize programs with callbr at higher
236 // optimization levels, as the DominatorTree created here is not reused by
237 // subsequent passes.
238 DominatorTree *DT;
239 std::optional<DominatorTree> LazilyComputedDomTree;
240 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
241 DT = &DTWP->getDomTree();
242 else {
243 LazilyComputedDomTree.emplace(args&: F);
244 DT = &*LazilyComputedDomTree;
245 }
246
247 if (SplitCriticalEdges(CBRs, DT&: *DT))
248 Changed = true;
249
250 if (InsertIntrinsicCalls(CBRs, DT&: *DT))
251 Changed = true;
252
253 return Changed;
254}
255