[CallSiteSplitting] properly split musttail calls

[indutny]

Summary

musttail calls can't be naively splitted. The split blocks must
include not only the call instruction itself, but also (optional)
bitcast and return instructions that follow it.

Clone bitcast and ret, place them into the split blocks, and
remove the tail block when done.

[indutny]

A question: what is backport policy of LLVM project? I'd love to see this fix in 5.0.0 .

[indutny]

Move condition to proper function.

[indutny]

Proper implementation!

[indutny]

Fix non-musttail calls.

[junbuml]

Thanks for catching this.

Different cloning using UnreachableInst could be a way to handle musttail. However, I prefer keeping the current way of splitting CS as it as. Instead, we may be able to handle musttail (clone bitcast/retrun and remove TailBB) at the end of splitCallSite() in a separate function.

lib/Transforms/Scalar/CallSiteSplitting.cpp

352

Instead of cloning differently for musttail using UnreachableInst, I think it also possible to keep the current way of splitting as it is, but for musttail, we can handle it (clone bitcast and return and remove TailBB) as a separate function which being called at the end of this function.

423

Here, I think we can call a separate function which handle the musttail case. So we can specifically clone bitcast and return and remove TailBB.

test/Transforms/CallSiteSplitting/musttail.ll

22

Can we add two more test cases :

1. return void .
2. return without bitcast.

[indutny]

Address review comments.

[indutny]

Hello!

Just wanted to say thank you for giving a thorough review! I appreciate it!

[junbuml]

lib/Transforms/Scalar/CallSiteSplitting.cpp

234

You can change "if (PN != nulltr)" into "if (PN)".

Isn't it possible to set isVoid here?

246

assert(RI != nullptr && -> assert(RI &&

250

Although BI is no longer used, ++BI in assert doesn't seem to be a good idea. std::next(BI) should be better.

260

It's not always true because we can have some instructions before the call-site split (e.g., callsite-instructions-before-call.ll).

264

We don't need to do this in the loop.

268–277

Look like we can make a separate function for cloning.

296

if (RI) should be better.

415

The original CallInst must be removed in above while loop (line 394).

[indutny]

Address feedback.

[junbuml]

LGTM.
I want to let Florian also review this as well.

lib/Transforms/Scalar/CallSiteSplitting.cpp

212–224

Since this is only for Return/BitCast for musttail, it will be good to name it more specifically. Something like cloneReturnInstForMustTail ?

273

It would be better if you can clean up DEBUG little bit more. I think DEBUG only for BitCast and Return cloned is good enough.

[indutny]

Rename function, and remove extra debug calls.

[fhahn]

Thanks for the patch. I think it would be simpler (and less code) to integrate musttail handling directly in splitCallSite, rather than patching up afterwards. See my inline comments and please let me know what you think/if I am missing something.

lib/Transforms/Scalar/CallSiteSplitting.cpp

376

For musttail calls, I think we could clone the bitcast & ret instructions here and move them after NewCi. Instead of adding the result to the Call phi node, we just have to updated the cloned bitcast/ret instruction to use NewCI instead of the original one.

379

For musttail calls, we should be able to just delete the BB once we handled all predecessors and we are done.

[indutny]

Hello, @fhahn !

Thank you for taking a look at it. I'll be glad to make the changes. Just one question before I start: it seems that the suggestions of @junbuml and you are to some extent contradictory. At first I was suggested to leave the function as it is, and now the suggestion is to move things back to it. Given this, what do we actually desire it to look like? Let's decide on this before making any additional changes.

Thanks for your time,
Fedor.

[fhahn]

In D43729#1022610, @indutny wrote:

Hello, @fhahn !

Thank you for taking a look at it. I'll be glad to make the changes. Just one question before I start: it seems that the suggestions of @junbuml and you are to some extent contradictory. At first I was suggested to leave the function as it is, and now the suggestion is to move things back to it. Given this, what do we actually desire it to look like? Let's decide on this before making any additional changes.

Ah I missed the full history, but had a look from the start now. I think the initial patch was slightly different to what I outlined. With my suggestion, the only changes to the main loop of splitCallSite would be cloning and moving the bitcast/ret instructions for musttail calls to the newly created predecessor blocks (this should be done in a separate function). And then we can just erase the TailBB, if it does not have any predecessors any longer (which should always the case ATM for musttail calls, as we split all predecessors for now).

If that works as I think, it should be just a couple of lines changed in splitCallSite and we handle musttail directly, rather than after the fact, which IMO makes it easier to see what's going on. I might be missing something though. Jun, what do you think?

[junbuml]

What I didn't want to have initially was creating unreachable to rely on DuplicateInstructionsInSplitBetween to clone BitCast/Ret. As a way to avoid that, I suggested to separate handling musttail from existing cloning. I think what Florian suggested could make fixMustTailCallSite even simpler. If you make fixMustTailCallSite to handle each pred split instead of handling two preds splits, I think you can call the modified fixMustTailCallSite instead of doing CallPN->addIncoming() in splitCallSite();

[indutny]

Sounds like a plan! Will work on it today. Thank you, guys!

[indutny]

Address feedback.

[indutny]

There's a huge FIXME in the code that depends on @fhahn's patch: https://reviews.llvm.org/D43822 . Will remove it as soon as that patch lands. Hopefully it won't deter you from reviewing it ;)

[fhahn]

lib/Transforms/Scalar/CallSiteSplitting.cpp

392

We should only do that for musttail calls, otherwise we crash.

I think we hit a slightly different problem here that won't be fixed by D43822. A different iterator gets invalidated. I'll have a closer look tomorrow.

[indutny]

lib/Transforms/Scalar/CallSiteSplitting.cpp

392

Ooops. I'll take a look at it today, hopefully this won't be in conflict with your work since you plan to start tomorrow.

[indutny]

Proper fixup for pending clone issue.

[fhahn]

Thanks for updating! I left a few more comments, but it looks good overall. I just left a few thoughts about structuring the code to deal with musttail in splitCallSite. Sorry if I was not too clear earlier. Anyways, those are mostly my personal preference. Also, I think doing the deletion after the loop is totally fine and avoids problems.

lib/Transforms/Scalar/CallSiteSplitting.cpp

228

I don't think we need to do that.

348

Personally I would just move the code in this block and the code for the musttail handling in the main loop to a single function, named something like addMusttailReturn. It means we might end up doing a bit of additional work, but it keeps splitCallSite simpler.

353

Not sure if we need those assertions here. If anything is wrong here, then it should have been invalid earlier too.

359

Not sure if we need those assertions here. If anything is wrong here, then it should have been invalid earlier too.

381

I think it would be slightly better to have that code a function addMusttailReturn or something.

387

I don't think D43822 fixes the problem. For now I think it's OK to just remove the terminators after the main loop (and drop the fixme).

404

Please merge that with the block that deletes the terminators. I would drop the fixme and add a comment with a note that we delete it after we split all predecessors to avoid iterator invalidation.

[indutny]

Address feedback.

---

This should do it. Thank you for reviewing this!

lib/Transforms/Scalar/CallSiteSplitting.cpp

228

We need it for fancy names in output IR. Those %ca1, %ca2 come out because of it.

[fhahn]

Great thanks! LGTM. Just 2 small nits that can be addressed before committing, without additional review

lib/Transforms/Scalar/CallSiteSplitting.cpp

346

blokcs -> blocks

392

range based for loop?

[indutny]

Address feedback.

---

Done. Thanks for taking a look! Will land it later today unless there's more to change!

[fhahn]

Great. Please wait a bit with committing so @junbuml has a chance to raise any additional concerns.

[indutny]

Thank you, everyone!

[junbuml]

Sorry for my late response. I was off for past few days. Thanks for working on this.

lib/Transforms/Scalar/CallSiteSplitting.cpp

static bool canSplitCallSite(CallSite CS, TargetTransformInfo &TTI) {
   SmallVector<BasicBlock *, 2> Preds(predecessors(CallSiteBB));
   if (Preds.size() != 2 || isa<IndirectBrInst>(Preds[0]->getTerminator()) ||
       isa<IndirectBrInst>(Preds[1]->getTerminator()))
     return false;
 
   return CallSiteBB->canSplitPredecessors();
 }
 
-/// Return true if the CS is split into its new predecessors.
+static Instruction *cloneReturnInstForMustTail(Instruction *I,
+                                               Instruction *Before, Value *V) {
+  Instruction *Copy = I->clone();
+  Copy->setName(I->getName());
+  Copy->insertBefore(Before);
+  if (V)
+    Copy->setOperand(0, V);
+  return Copy;
+}
+
+/// `musttail` calls must be followed by optional `bitcast`, and `ret`
+/// instructions. `splitCallSite` however produces blocks with copies of the
+/// call followed by `br` instructions, which violates the invariant.
 ///
+/// Restore the invariant in each of the predecessor blocks by moving the
+/// `bitcast` and `ret` from the tail block into the predecessor.
+static void fixMustTailCallSite(BasicBlock *TailBB) {
+  DEBUG(dbgs() << "fix musttail splits:\n");
+
+  bool IsVoid = TailBB->getParent()->getReturnType()->isVoidTy();
+
+  BasicBlock::iterator BI = TailBB->begin();
+  // PHINode is present for non-void return values
+  PHINode *PN = dyn_cast<PHINode>(&*BI);
+  if (PN)
+    ++BI;
+
+  // Optional BitCastInst
+  BitCastInst *BCI = dyn_cast<BitCastInst>(&*BI);
+  if (BCI)
+    ++BI;
+
+  // Mandatory ReturnInst
+  ReturnInst *RI = dyn_cast<ReturnInst>(&*BI);
+  assert(RI && "`musttail` call must be followed by `ret` instruction");
+
+  assert(std::next(BI) == TailBB->end() && "`ret` doesn't end the block");
+
+  SmallVector<BasicBlock *, 2> Splits(predecessors((TailBB)));
+  assert(Splits.size() == 2 && "Expected exactly 2 splits!");
+
+  for (unsigned i = 0; i < Splits.size(); i++) {
+    BasicBlock *Split = Splits[i];
+
+    TerminatorInst *TI = Split->getTerminator();
+
+    CallInst *CI = dyn_cast<CallInst>(&*std::prev(TI->getIterator()));
+    assert(CI && "pred blocks must have the call");
+    assert(IsVoid == CI->getFunctionType()->getReturnType()->isVoidTy() &&
+        "Return type mismatch between musttail call and the caller function");
+    DEBUG(dbgs() << "    feeding CI : " << *CI << " into\n");
+
+    Value *V = CI;
+
+    // Add bitcast
+    if (BCI) {
+      assert(!IsVoid && "Can\'t bitcast void");
+      V = cloneReturnInstForMustTail(BCI, TI, V);
+      DEBUG(dbgs() << "        " << *V << " into\n");
+    }
+
+    // Add return
+    Instruction *NewRI = cloneReturnInstForMustTail(RI, TI,
+                                                    IsVoid ? nullptr : V);
+    DEBUG(dbgs() << "        " << *NewRI << "\n");
+
+    // Remove terminating `br` instruction
+    TI->eraseFromParent();
+
+    if (PN)
+      PN->removeIncomingValue(Split, true);
+  }
+
+  // Remove unused instructions
+  if (RI)
+    RI->eraseFromParent();
+  if (BCI)
+    BCI->eraseFromParent();
+
+  // The block is no longer reachable
+  TailBB->eraseFromParent();
+}
+
 /// For each (predecessor, conditions from predecessors) pair, it will split the
 /// basic block containing the call site, hook it up to the predecessor and
 /// replace the call instruction with new call instructions, which contain
 /// constraints based on the conditions from their predecessors.
 /// For example, in the IR below with an OR condition, the call-site can
 /// be split. In this case, Preds for Tail is [(Header, a == null),
 /// (TBB, a != null, b == null)]. Tail is replaced by 2 split blocks, containing
 /// CallInst1, which has constraints based on the conditions from Head and
 /// CallInst2, which has constraints based on the conditions coming from TBB.
 ///

 /// Note that in case any arguments at the call-site are constrained by its
 /// predecessors, new call-sites with more constrained arguments will be
 /// created in createCallSitesOnPredicatedArgument().
 static void splitCallSite(
     CallSite CS,
     const SmallVectorImpl<std::pair<BasicBlock *, ConditionsTy>> &Preds) {
   Instruction *Instr = CS.getInstruction();
   BasicBlock *TailBB = Instr->getParent();
+  bool IsMustTailCall = CS.isMustTailCall();
 
   PHINode *CallPN = nullptr;
   if (Instr->getNumUses())
     CallPN = PHINode::Create(Instr->getType(), Preds.size(), "phi.call");
 
   DEBUG(dbgs() << "split call-site : " << *Instr << " into \n");
 
   assert(Preds.size() == 2 && "The ValueToValueMaps array has size 2.");
   // ValueToValueMapTy is neither copy nor moveable, so we use a simple array
   // here.
   ValueToValueMapTy ValueToValueMaps[2];
   for (unsigned i = 0; i < Preds.size(); i++) {
     BasicBlock *PredBB = Preds[i].first;
     BasicBlock *SplitBlock = DuplicateInstructionsInSplitBetween(
         TailBB, PredBB, &*std::next(Instr->getIterator()), ValueToValueMaps[i]);
     assert(SplitBlock && "Unexpected new basic block split.");
 
     Instruction *NewCI =
         &*std::prev(SplitBlock->getTerminator()->getIterator());
     CallSite NewCS(NewCI);
     addConditions(NewCS, Preds[i].second);
 
     // Handle PHIs used as arguments in the call-site.
     for (PHINode &PN : TailBB->phis()) {
       unsigned ArgNo = 0;
       for (auto &CI : CS.args()) {
         if (&*CI == &PN) {
           NewCS.setArgument(ArgNo, PN.getIncomingValueForBlock(SplitBlock));
         }
         ++ArgNo;
       }
     }
     DEBUG(dbgs() << "    " << *NewCI << " in " << SplitBlock->getName()
                  << "\n");
     if (CallPN)
       CallPN->addIncoming(NewCI, SplitBlock);
   }
 
   auto *OriginalBegin = &*TailBB->begin();
   // Replace users of the original call with a PHI mering call-sites split.
   if (CallPN) {
     CallPN->insertBefore(OriginalBegin);
     Instr->replaceAllUsesWith(CallPN);
   }
 
   // Remove instructions moved to split blocks from TailBB, from the duplicated
   // call instruction to the beginning of the basic block. If an instruction
   // has any uses, add a new PHI node to combine the values coming from the
   // split blocks. The new PHI nodes are placed before the first original
   // instruction, so we do not end up deleting them. By using reverse-order, we
   // do not introduce unnecessary PHI nodes for def-use chains from the call
   // instruction to the beginning of the block.
   auto I = Instr->getReverseIterator();
   while (I != TailBB->rend()) {
     Instruction *CurrentI = &*I++;
     if (!CurrentI->use_empty()) {
       // If an existing PHI has users after the call, there is no need to create
       // a new one.
       if (isa<PHINode>(CurrentI))
         continue;
       PHINode *NewPN = PHINode::Create(CurrentI->getType(), Preds.size());
       for (auto &Mapping : ValueToValueMaps)
         NewPN->addIncoming(Mapping[CurrentI],
                            cast<Instruction>(Mapping[CurrentI])->getParent());
       NewPN->insertBefore(&*TailBB->begin());
       CurrentI->replaceAllUsesWith(NewPN);
     }
     CurrentI->eraseFromParent();
     // We are done once we handled the first original instruction in TailBB.
     if (CurrentI == OriginalBegin)
       break;
   }
 
+  if (IsMustTailCall)
+    fixMustTailCallSite(TailBB);
+
   NumCallSiteSplit++;
 }
 
 // Return true if the call-site has an argument which is a PHI with only
 // constant incoming values.
 static bool isPredicatedOnPHI(CallSite CS) {
   Instruction *Instr = CS.getInstruction();
   BasicBlock *Parent = Instr->getParent();
   if (Instr != Parent->getFirstNonPHIOrDbg())
     return false;
 
   for (auto &BI : *Parent) {
     if (PHINode *PN = dyn_cast<PHINode>(&BI)) {
       for (auto &I : CS.args())
for (BasicBlock::iterator II = BB.begin(), IE = BB.end(); II != IE;) {
       Instruction *I = &*II++;
       CallSite CS(cast<Value>(I));
       if (!CS || isa<IntrinsicInst>(I) || isInstructionTriviallyDead(I, &TLI))
         continue;
 
       Function *Callee = CS.getCalledFunction();
       if (!Callee || Callee->isDeclaration())
         continue;
+
+      // Successful musttail call-site splits result in erased CI and erased BB.
+      // Check if such path is possible before attempting the splitting.
+      bool IsMustTail = CS.isMustTailCall();
+
       Changed |= tryToSplitCallSite(CS, TTI);
+
+      // There're no interesting instructions after this. The call site
+      // itself might have been erased on splitting.
+      if (IsMustTail)
+        break;
     }
   }
   return Changed;
 }
 
 namespace {
 struct CallSiteSplittingLegacyPass : public FunctionPass {
   static char ID;

test/Transforms/CallSiteSplitting/musttail.ll

+; RUN: opt < %s -callsite-splitting -S | FileCheck %s
+
+;CHECK-LABEL: @caller
+;CHECK-LABEL: Top.split:
+;CHECK: %ca1 = musttail call i8* @callee(i8* null, i8* %b)
+;CHECK: %cb4 = bitcast i8* %ca1 to i8*
+;CHECK: ret i8* %cb4
+;CHECK-LABEL: TBB.split
+;CHECK: %ca2 = musttail call i8* @callee(i8* nonnull %a, i8* null)
+;CHECK: %cb3 = bitcast i8* %ca2 to i8*
+;CHECK: ret i8* %cb3
+define i8* @caller(i8* %a, i8* %b) {
+Top:
+  %c = icmp eq i8* %a, null
+  br i1 %c, label %Tail, label %TBB
+TBB:
+  %c2 = icmp eq i8* %b, null
+  br i1 %c2, label %Tail, label %End
+Tail:
+  %ca = musttail call i8* @callee(i8* %a, i8* %b)
+  %cb = bitcast i8* %ca to i8*
+  ret i8* %cb
+End:
+  ret i8* null
+}
+
+define i8* @callee(i8* %a, i8* %b) noinline {
+  ret i8* %a
+}
+
+;CHECK-LABEL: @no_cast_caller
+;CHECK-LABEL: Top.split:
+;CHECK: %ca1 = musttail call i8* @callee(i8* null, i8* %b)
+;CHECK: ret i8* %ca1
+;CHECK-LABEL: TBB.split
+;CHECK: %ca2 = musttail call i8* @callee(i8* nonnull %a, i8* null)
+;CHECK: ret i8* %ca2
+define i8* @no_cast_caller(i8* %a, i8* %b) {
+Top:
+  %c = icmp eq i8* %a, null
+  br i1 %c, label %Tail, label %TBB
+TBB:
+  %c2 = icmp eq i8* %b, null
+  br i1 %c2, label %Tail, label %End
+Tail:
+  %ca = musttail call i8* @callee(i8* %a, i8* %b)
+  ret i8* %ca
+End:
+  ret i8* null
+}
+
+;CHECK-LABEL: @void_caller
+;CHECK-LABEL: Top.split:
+;CHECK: musttail call void @void_callee(i8* null, i8* %b)
+;CHECK: ret void
+;CHECK-LABEL: TBB.split
+;CHECK: musttail call void @void_callee(i8* nonnull %a, i8* null)
+;CHECK: ret void
+define void @void_caller(i8* %a, i8* %b) {
+Top:
+  %c = icmp eq i8* %a, null
+  br i1 %c, label %Tail, label %TBB
+TBB:
+  %c2 = icmp eq i8* %b, null
+  br i1 %c2, label %Tail, label %End
+Tail:
+  musttail call void @void_callee(i8* %a, i8* %b)
+  ret void
+End:
+  ret void
+}
+
+define void @void_callee(i8* %a, i8* %b) noinline {
+  ret void
+}