1//===- CallSiteSplitting.cpp ----------------------------------------------===//
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 a transformation that tries to split a call-site to pass
10// more constrained arguments if its argument is predicated in the control flow
11// so that we can expose better context to the later passes (e.g, inliner, jump
12// threading, or IPA-CP based function cloning, etc.).
13// As of now we support two cases :
14//
15// 1) Try to a split call-site with constrained arguments, if any constraints
16// on any argument can be found by following the single predecessors of the
17// all site's predecessors. Currently this pass only handles call-sites with 2
18// predecessors. For example, in the code below, we try to split the call-site
19// since we can predicate the argument(ptr) based on the OR condition.
20//
21// Split from :
22// if (!ptr || c)
23// callee(ptr);
24// to :
25// if (!ptr)
26// callee(null) // set the known constant value
27// else if (c)
28// callee(nonnull ptr) // set non-null attribute in the argument
29//
30// 2) We can also split a call-site based on constant incoming values of a PHI
31// For example,
32// from :
33// Header:
34// %c = icmp eq i32 %i1, %i2
35// br i1 %c, label %Tail, label %TBB
36// TBB:
37// br label Tail%
38// Tail:
39// %p = phi i32 [ 0, %Header], [ 1, %TBB]
40// call void @bar(i32 %p)
41// to
42// Header:
43// %c = icmp eq i32 %i1, %i2
44// br i1 %c, label %Tail-split0, label %TBB
45// TBB:
46// br label %Tail-split1
47// Tail-split0:
48// call void @bar(i32 0)
49// br label %Tail
50// Tail-split1:
51// call void @bar(i32 1)
52// br label %Tail
53// Tail:
54// %p = phi i32 [ 0, %Tail-split0 ], [ 1, %Tail-split1 ]
55//
56//===----------------------------------------------------------------------===//
57
58#include "llvm/Transforms/Scalar/CallSiteSplitting.h"
59#include "llvm/ADT/Statistic.h"
60#include "llvm/Analysis/DomTreeUpdater.h"
61#include "llvm/Analysis/TargetLibraryInfo.h"
62#include "llvm/Analysis/TargetTransformInfo.h"
63#include "llvm/IR/IntrinsicInst.h"
64#include "llvm/IR/PatternMatch.h"
65#include "llvm/Support/CommandLine.h"
66#include "llvm/Support/Debug.h"
67#include "llvm/Transforms/Utils/Cloning.h"
68#include "llvm/Transforms/Utils/Local.h"
69
70using namespace llvm;
71using namespace PatternMatch;
72
73#define DEBUG_TYPE "callsite-splitting"
74
75STATISTIC(NumCallSiteSplit, "Number of call-site split");
76
77/// Only allow instructions before a call, if their CodeSize cost is below
78/// DuplicationThreshold. Those instructions need to be duplicated in all
79/// split blocks.
80static cl::opt<unsigned>
81 DuplicationThreshold("callsite-splitting-duplication-threshold", cl::Hidden,
82 cl::desc("Only allow instructions before a call, if "
83 "their cost is below DuplicationThreshold"),
84 cl::init(Val: 5));
85
86static void addNonNullAttribute(CallBase &CB, Value *Op) {
87 unsigned ArgNo = 0;
88 for (auto &I : CB.args()) {
89 if (&*I == Op)
90 CB.addParamAttr(ArgNo, Kind: Attribute::NonNull);
91 ++ArgNo;
92 }
93}
94
95static void setConstantInArgument(CallBase &CB, Value *Op,
96 Constant *ConstValue) {
97 unsigned ArgNo = 0;
98 for (auto &I : CB.args()) {
99 if (&*I == Op) {
100 // It is possible we have already added the non-null attribute to the
101 // parameter by using an earlier constraining condition.
102 CB.removeParamAttr(ArgNo, Kind: Attribute::NonNull);
103 CB.setArgOperand(i: ArgNo, v: ConstValue);
104 }
105 ++ArgNo;
106 }
107}
108
109static bool isCondRelevantToAnyCallArgument(ICmpInst *Cmp, CallBase &CB) {
110 assert(isa<Constant>(Cmp->getOperand(1)) && "Expected a constant operand.");
111 Value *Op0 = Cmp->getOperand(i_nocapture: 0);
112 unsigned ArgNo = 0;
113 for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I, ++ArgNo) {
114 // Don't consider constant or arguments that are already known non-null.
115 if (isa<Constant>(Val: *I) || CB.paramHasAttr(ArgNo, Kind: Attribute::NonNull))
116 continue;
117
118 if (*I == Op0)
119 return true;
120 }
121 return false;
122}
123
124using ConditionTy = std::pair<ICmpInst *, unsigned>;
125using ConditionsTy = SmallVector<ConditionTy, 2>;
126
127/// If From has a conditional jump to To, add the condition to Conditions,
128/// if it is relevant to any argument at CB.
129static void recordCondition(CallBase &CB, BasicBlock *From, BasicBlock *To,
130 ConditionsTy &Conditions) {
131 auto *BI = dyn_cast<BranchInst>(Val: From->getTerminator());
132 if (!BI || !BI->isConditional())
133 return;
134
135 CmpPredicate Pred;
136 Value *Cond = BI->getCondition();
137 if (!match(V: Cond, P: m_ICmp(Pred, L: m_Value(), R: m_Constant())))
138 return;
139
140 ICmpInst *Cmp = cast<ICmpInst>(Val: Cond);
141 if (Pred == ICmpInst::ICMP_EQ || Pred == ICmpInst::ICMP_NE)
142 if (isCondRelevantToAnyCallArgument(Cmp, CB))
143 Conditions.push_back(Elt: {Cmp, From->getTerminator()->getSuccessor(Idx: 0) == To
144 ? Pred
145 : Cmp->getInverseCmpPredicate()});
146}
147
148/// Record ICmp conditions relevant to any argument in CB following Pred's
149/// single predecessors. If there are conflicting conditions along a path, like
150/// x == 1 and x == 0, the first condition will be used. We stop once we reach
151/// an edge to StopAt.
152static void recordConditions(CallBase &CB, BasicBlock *Pred,
153 ConditionsTy &Conditions, BasicBlock *StopAt) {
154 BasicBlock *From = Pred;
155 BasicBlock *To = Pred;
156 SmallPtrSet<BasicBlock *, 4> Visited;
157 while (To != StopAt && !Visited.count(Ptr: From->getSinglePredecessor()) &&
158 (From = From->getSinglePredecessor())) {
159 recordCondition(CB, From, To, Conditions);
160 Visited.insert(Ptr: From);
161 To = From;
162 }
163}
164
165static void addConditions(CallBase &CB, const ConditionsTy &Conditions) {
166 for (const auto &Cond : Conditions) {
167 Value *Arg = Cond.first->getOperand(i_nocapture: 0);
168 Constant *ConstVal = cast<Constant>(Val: Cond.first->getOperand(i_nocapture: 1));
169 if (Cond.second == ICmpInst::ICMP_EQ)
170 setConstantInArgument(CB, Op: Arg, ConstValue: ConstVal);
171 else if (ConstVal->getType()->isPointerTy() && ConstVal->isNullValue()) {
172 assert(Cond.second == ICmpInst::ICMP_NE);
173 addNonNullAttribute(CB, Op: Arg);
174 }
175 }
176}
177
178static SmallVector<BasicBlock *, 2> getTwoPredecessors(BasicBlock *BB) {
179 SmallVector<BasicBlock *, 2> Preds(predecessors((BB)));
180 assert(Preds.size() == 2 && "Expected exactly 2 predecessors!");
181 return Preds;
182}
183
184static bool canSplitCallSite(CallBase &CB, TargetTransformInfo &TTI) {
185 if (CB.isConvergent() || CB.cannotDuplicate())
186 return false;
187
188 // FIXME: As of now we handle only CallInst. InvokeInst could be handled
189 // without too much effort.
190 if (!isa<CallInst>(Val: CB))
191 return false;
192
193 BasicBlock *CallSiteBB = CB.getParent();
194 // Need 2 predecessors and cannot split an edge from an IndirectBrInst.
195 SmallVector<BasicBlock *, 2> Preds(predecessors(BB: CallSiteBB));
196 if (Preds.size() != 2 || isa<IndirectBrInst>(Val: Preds[0]->getTerminator()) ||
197 isa<IndirectBrInst>(Val: Preds[1]->getTerminator()))
198 return false;
199
200 // BasicBlock::canSplitPredecessors is more aggressive, so checking for
201 // BasicBlock::isEHPad as well.
202 if (!CallSiteBB->canSplitPredecessors() || CallSiteBB->isEHPad())
203 return false;
204
205 // Allow splitting a call-site only when the CodeSize cost of the
206 // instructions before the call is less then DuplicationThreshold. The
207 // instructions before the call will be duplicated in the split blocks and
208 // corresponding uses will be updated.
209 InstructionCost Cost = 0;
210 for (auto &InstBeforeCall :
211 llvm::make_range(x: CallSiteBB->begin(), y: CB.getIterator())) {
212 Cost += TTI.getInstructionCost(U: &InstBeforeCall,
213 CostKind: TargetTransformInfo::TCK_CodeSize);
214 if (Cost >= DuplicationThreshold)
215 return false;
216 }
217
218 return true;
219}
220
221static Instruction *
222cloneInstForMustTail(Instruction *I, BasicBlock::iterator Before, Value *V) {
223 Instruction *Copy = I->clone();
224 Copy->setName(I->getName());
225 Copy->insertBefore(InsertPos: Before);
226 if (V)
227 Copy->setOperand(i: 0, Val: V);
228 return Copy;
229}
230
231/// Copy mandatory `musttail` return sequence that follows original `CI`, and
232/// link it up to `NewCI` value instead:
233///
234/// * (optional) `bitcast NewCI to ...`
235/// * `ret bitcast or NewCI`
236///
237/// Insert this sequence right before `SplitBB`'s terminator, which will be
238/// cleaned up later in `splitCallSite` below.
239static void copyMustTailReturn(BasicBlock *SplitBB, Instruction *CI,
240 Instruction *NewCI) {
241 bool IsVoid = SplitBB->getParent()->getReturnType()->isVoidTy();
242 auto II = std::next(x: CI->getIterator());
243
244 BitCastInst* BCI = dyn_cast<BitCastInst>(Val: &*II);
245 if (BCI)
246 ++II;
247
248 ReturnInst* RI = dyn_cast<ReturnInst>(Val: &*II);
249 assert(RI && "`musttail` call must be followed by `ret` instruction");
250
251 Instruction *TI = SplitBB->getTerminator();
252 Value *V = NewCI;
253 if (BCI)
254 V = cloneInstForMustTail(I: BCI, Before: TI->getIterator(), V);
255 cloneInstForMustTail(I: RI, Before: TI->getIterator(), V: IsVoid ? nullptr : V);
256
257 // FIXME: remove TI here, `DuplicateInstructionsInSplitBetween` has a bug
258 // that prevents doing this now.
259}
260
261/// For each (predecessor, conditions from predecessors) pair, it will split the
262/// basic block containing the call site, hook it up to the predecessor and
263/// replace the call instruction with new call instructions, which contain
264/// constraints based on the conditions from their predecessors.
265/// For example, in the IR below with an OR condition, the call-site can
266/// be split. In this case, Preds for Tail is [(Header, a == null),
267/// (TBB, a != null, b == null)]. Tail is replaced by 2 split blocks, containing
268/// CallInst1, which has constraints based on the conditions from Head and
269/// CallInst2, which has constraints based on the conditions coming from TBB.
270///
271/// From :
272///
273/// Header:
274/// %c = icmp eq i32* %a, null
275/// br i1 %c %Tail, %TBB
276/// TBB:
277/// %c2 = icmp eq i32* %b, null
278/// br i1 %c %Tail, %End
279/// Tail:
280/// %ca = call i1 @callee (i32* %a, i32* %b)
281///
282/// to :
283///
284/// Header: // PredBB1 is Header
285/// %c = icmp eq i32* %a, null
286/// br i1 %c %Tail-split1, %TBB
287/// TBB: // PredBB2 is TBB
288/// %c2 = icmp eq i32* %b, null
289/// br i1 %c %Tail-split2, %End
290/// Tail-split1:
291/// %ca1 = call @callee (i32* null, i32* %b) // CallInst1
292/// br %Tail
293/// Tail-split2:
294/// %ca2 = call @callee (i32* nonnull %a, i32* null) // CallInst2
295/// br %Tail
296/// Tail:
297/// %p = phi i1 [%ca1, %Tail-split1],[%ca2, %Tail-split2]
298///
299/// Note that in case any arguments at the call-site are constrained by its
300/// predecessors, new call-sites with more constrained arguments will be
301/// created in createCallSitesOnPredicatedArgument().
302static void splitCallSite(CallBase &CB,
303 ArrayRef<std::pair<BasicBlock *, ConditionsTy>> Preds,
304 DomTreeUpdater &DTU) {
305 BasicBlock *TailBB = CB.getParent();
306 bool IsMustTailCall = CB.isMustTailCall();
307
308 PHINode *CallPN = nullptr;
309
310 // `musttail` calls must be followed by optional `bitcast`, and `ret`. The
311 // split blocks will be terminated right after that so there're no users for
312 // this phi in a `TailBB`.
313 if (!IsMustTailCall && !CB.use_empty()) {
314 CallPN = PHINode::Create(Ty: CB.getType(), NumReservedValues: Preds.size(), NameStr: "phi.call");
315 CallPN->setDebugLoc(CB.getDebugLoc());
316 }
317
318 LLVM_DEBUG(dbgs() << "split call-site : " << CB << " into \n");
319
320 assert(Preds.size() == 2 && "The ValueToValueMaps array has size 2.");
321 // ValueToValueMapTy is neither copy nor moveable, so we use a simple array
322 // here.
323 ValueToValueMapTy ValueToValueMaps[2];
324 for (unsigned i = 0; i < Preds.size(); i++) {
325 BasicBlock *PredBB = Preds[i].first;
326 BasicBlock *SplitBlock = DuplicateInstructionsInSplitBetween(
327 BB: TailBB, PredBB, StopAt: &*std::next(x: CB.getIterator()), ValueMapping&: ValueToValueMaps[i],
328 DTU);
329 assert(SplitBlock && "Unexpected new basic block split.");
330
331 auto *NewCI =
332 cast<CallBase>(Val: &*std::prev(x: SplitBlock->getTerminator()->getIterator()));
333 addConditions(CB&: *NewCI, Conditions: Preds[i].second);
334
335 // Handle PHIs used as arguments in the call-site.
336 for (PHINode &PN : TailBB->phis()) {
337 unsigned ArgNo = 0;
338 for (auto &CI : CB.args()) {
339 if (&*CI == &PN) {
340 NewCI->setArgOperand(i: ArgNo, v: PN.getIncomingValueForBlock(BB: SplitBlock));
341 }
342 ++ArgNo;
343 }
344 }
345 LLVM_DEBUG(dbgs() << " " << *NewCI << " in " << SplitBlock->getName()
346 << "\n");
347 if (CallPN)
348 CallPN->addIncoming(V: NewCI, BB: SplitBlock);
349
350 // Clone and place bitcast and return instructions before `TI`
351 if (IsMustTailCall)
352 copyMustTailReturn(SplitBB: SplitBlock, CI: &CB, NewCI);
353 }
354
355 NumCallSiteSplit++;
356
357 // FIXME: remove TI in `copyMustTailReturn`
358 if (IsMustTailCall) {
359 // Remove superfluous `br` terminators from the end of the Split blocks
360 // NOTE: Removing terminator removes the SplitBlock from the TailBB's
361 // predecessors. Therefore we must get complete list of Splits before
362 // attempting removal.
363 SmallVector<BasicBlock *, 2> Splits(predecessors(BB: (TailBB)));
364 assert(Splits.size() == 2 && "Expected exactly 2 splits!");
365 for (BasicBlock *BB : Splits) {
366 BB->getTerminator()->eraseFromParent();
367 DTU.applyUpdatesPermissive(Updates: {{DominatorTree::Delete, BB, TailBB}});
368 }
369
370 // Erase the tail block once done with musttail patching
371 DTU.deleteBB(DelBB: TailBB);
372 return;
373 }
374
375 BasicBlock::iterator OriginalBegin = TailBB->begin();
376 // Replace users of the original call with a PHI mering call-sites split.
377 if (CallPN) {
378 CallPN->insertBefore(BB&: *TailBB, InsertPos: OriginalBegin);
379 CB.replaceAllUsesWith(V: CallPN);
380 }
381
382 // Remove instructions moved to split blocks from TailBB, from the duplicated
383 // call instruction to the beginning of the basic block. If an instruction
384 // has any uses, add a new PHI node to combine the values coming from the
385 // split blocks. The new PHI nodes are placed before the first original
386 // instruction, so we do not end up deleting them. By using reverse-order, we
387 // do not introduce unnecessary PHI nodes for def-use chains from the call
388 // instruction to the beginning of the block.
389 auto I = CB.getReverseIterator();
390 Instruction *OriginalBeginInst = &*OriginalBegin;
391 while (I != TailBB->rend()) {
392 Instruction *CurrentI = &*I++;
393 if (!CurrentI->use_empty()) {
394 // If an existing PHI has users after the call, there is no need to create
395 // a new one.
396 if (isa<PHINode>(Val: CurrentI))
397 continue;
398 PHINode *NewPN = PHINode::Create(Ty: CurrentI->getType(), NumReservedValues: Preds.size());
399 NewPN->setDebugLoc(CurrentI->getDebugLoc());
400 for (auto &Mapping : ValueToValueMaps) {
401 Value *V = Mapping[CurrentI];
402 NewPN->addIncoming(V, BB: cast<Instruction>(Val: V)->getParent());
403 }
404 NewPN->insertBefore(BB&: *TailBB, InsertPos: TailBB->begin());
405 CurrentI->replaceAllUsesWith(V: NewPN);
406 }
407 CurrentI->dropDbgRecords();
408 CurrentI->eraseFromParent();
409 // We are done once we handled the first original instruction in TailBB.
410 if (CurrentI == OriginalBeginInst)
411 break;
412 }
413}
414
415// Return true if the call-site has an argument which is a PHI with only
416// constant incoming values.
417static bool isPredicatedOnPHI(CallBase &CB) {
418 BasicBlock *Parent = CB.getParent();
419 if (&CB != &*Parent->getFirstNonPHIOrDbg())
420 return false;
421
422 for (auto &PN : Parent->phis()) {
423 for (auto &Arg : CB.args()) {
424 if (&*Arg != &PN)
425 continue;
426 assert(PN.getNumIncomingValues() == 2 &&
427 "Unexpected number of incoming values");
428 if (PN.getIncomingBlock(i: 0) == PN.getIncomingBlock(i: 1))
429 return false;
430 if (PN.getIncomingValue(i: 0) == PN.getIncomingValue(i: 1))
431 continue;
432 if (isa<Constant>(Val: PN.getIncomingValue(i: 0)) &&
433 isa<Constant>(Val: PN.getIncomingValue(i: 1)))
434 return true;
435 }
436 }
437 return false;
438}
439
440using PredsWithCondsTy = SmallVector<std::pair<BasicBlock *, ConditionsTy>, 2>;
441
442// Check if any of the arguments in CS are predicated on a PHI node and return
443// the set of predecessors we should use for splitting.
444static PredsWithCondsTy shouldSplitOnPHIPredicatedArgument(CallBase &CB) {
445 if (!isPredicatedOnPHI(CB))
446 return {};
447
448 auto Preds = getTwoPredecessors(BB: CB.getParent());
449 return {{Preds[0], {}}, {Preds[1], {}}};
450}
451
452// Checks if any of the arguments in CS are predicated in a predecessor and
453// returns a list of predecessors with the conditions that hold on their edges
454// to CS.
455static PredsWithCondsTy shouldSplitOnPredicatedArgument(CallBase &CB,
456 DomTreeUpdater &DTU) {
457 auto Preds = getTwoPredecessors(BB: CB.getParent());
458 if (Preds[0] == Preds[1])
459 return {};
460
461 // We can stop recording conditions once we reached the immediate dominator
462 // for the block containing the call site. Conditions in predecessors of the
463 // that node will be the same for all paths to the call site and splitting
464 // is not beneficial.
465 assert(DTU.hasDomTree() && "We need a DTU with a valid DT!");
466 auto *CSDTNode = DTU.getDomTree().getNode(BB: CB.getParent());
467 BasicBlock *StopAt = CSDTNode ? CSDTNode->getIDom()->getBlock() : nullptr;
468
469 SmallVector<std::pair<BasicBlock *, ConditionsTy>, 2> PredsCS;
470 for (auto *Pred : llvm::reverse(C&: Preds)) {
471 ConditionsTy Conditions;
472 // Record condition on edge BB(CS) <- Pred
473 recordCondition(CB, From: Pred, To: CB.getParent(), Conditions);
474 // Record conditions following Pred's single predecessors.
475 recordConditions(CB, Pred, Conditions, StopAt);
476 PredsCS.push_back(Elt: {Pred, Conditions});
477 }
478
479 if (all_of(Range&: PredsCS, P: [](const std::pair<BasicBlock *, ConditionsTy> &P) {
480 return P.second.empty();
481 }))
482 return {};
483
484 return PredsCS;
485}
486
487static bool tryToSplitCallSite(CallBase &CB, TargetTransformInfo &TTI,
488 DomTreeUpdater &DTU) {
489 // Check if we can split the call site.
490 if (!CB.arg_size() || !canSplitCallSite(CB, TTI))
491 return false;
492
493 auto PredsWithConds = shouldSplitOnPredicatedArgument(CB, DTU);
494 if (PredsWithConds.empty())
495 PredsWithConds = shouldSplitOnPHIPredicatedArgument(CB);
496 if (PredsWithConds.empty())
497 return false;
498
499 splitCallSite(CB, Preds: PredsWithConds, DTU);
500 return true;
501}
502
503static bool doCallSiteSplitting(Function &F, TargetLibraryInfo &TLI,
504 TargetTransformInfo &TTI, DominatorTree &DT) {
505
506 DomTreeUpdater DTU(&DT, DomTreeUpdater::UpdateStrategy::Lazy);
507 bool Changed = false;
508 for (BasicBlock &BB : llvm::make_early_inc_range(Range&: F)) {
509 auto II = BB.getFirstNonPHIOrDbg()->getIterator();
510 auto IE = BB.getTerminator()->getIterator();
511 // Iterate until we reach the terminator instruction. tryToSplitCallSite
512 // can replace BB's terminator in case BB is a successor of itself. In that
513 // case, IE will be invalidated and we also have to check the current
514 // terminator.
515 while (II != IE && &*II != BB.getTerminator()) {
516 CallBase *CB = dyn_cast<CallBase>(Val: &*II++);
517 if (!CB || isa<IntrinsicInst>(Val: CB) || isInstructionTriviallyDead(I: CB, TLI: &TLI))
518 continue;
519
520 Function *Callee = CB->getCalledFunction();
521 if (!Callee || Callee->isDeclaration())
522 continue;
523
524 // Successful musttail call-site splits result in erased CI and erased BB.
525 // Check if such path is possible before attempting the splitting.
526 bool IsMustTail = CB->isMustTailCall();
527
528 Changed |= tryToSplitCallSite(CB&: *CB, TTI, DTU);
529
530 // There're no interesting instructions after this. The call site
531 // itself might have been erased on splitting.
532 if (IsMustTail)
533 break;
534 }
535 }
536 return Changed;
537}
538
539PreservedAnalyses CallSiteSplittingPass::run(Function &F,
540 FunctionAnalysisManager &AM) {
541 auto &TLI = AM.getResult<TargetLibraryAnalysis>(IR&: F);
542 auto &TTI = AM.getResult<TargetIRAnalysis>(IR&: F);
543 auto &DT = AM.getResult<DominatorTreeAnalysis>(IR&: F);
544
545 if (!doCallSiteSplitting(F, TLI, TTI, DT))
546 return PreservedAnalyses::all();
547 PreservedAnalyses PA;
548 PA.preserve<DominatorTreeAnalysis>();
549 return PA;
550}
551