[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

268 ↗

(On Diff #135804)

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.

358 ↗

(On Diff #135804)

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

21 ↗

(On Diff #135804)

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

224 ↗

(On Diff #135987)

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

Isn't it possible to set isVoid here?

236 ↗

(On Diff #135987)

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

240 ↗

(On Diff #135987)

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

250 ↗

(On Diff #135987)

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

254 ↗

(On Diff #135987)

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

258–267 ↗

(On Diff #135987)

Look like we can make a separate function for cloning.

286 ↗

(On Diff #135987)

if (RI) should be better.

414 ↗

(On Diff #135987)

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 ↗

(On Diff #136101)

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

273 ↗

(On Diff #136101)

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 ↗

(On Diff #136147)

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 ↗

(On Diff #136147)

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

338 ↗

(On Diff #136409)

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

338 ↗

(On Diff #136409)

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

215 ↗

(On Diff #136678)

I don't think we need to do that.

278 ↗

(On Diff #136678)

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.

283 ↗

(On Diff #136678)

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

289 ↗

(On Diff #136678)

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

327 ↗

(On Diff #136678)

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

333 ↗

(On Diff #136678)

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).

350 ↗

(On Diff #136678)

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

215 ↗

(On Diff #136678)

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

303 ↗

(On Diff #136785)

blokcs -> blocks

353 ↗

(On Diff #136785)

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.

llvm/trunk/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 *cloneInstForMustTail(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;
+}
+
+/// Copy mandatory `musttail` return sequence that follows original `CI`, and
+/// link it up to `NewCI` value instead:
+///
+///   * (optional) `bitcast NewCI to ...`
+///   * `ret bitcast or NewCI`
 ///
+/// Insert this sequence right before `SplitBB`'s terminator, which will be
+/// cleaned up later in `splitCallSite` below.
+static void copyMustTailReturn(BasicBlock *SplitBB, Instruction *CI,
+                               Instruction *NewCI) {
+  bool IsVoid = SplitBB->getParent()->getReturnType()->isVoidTy();
+  auto II = std::next(CI->getIterator());
+
+  BitCastInst* BCI = dyn_cast<BitCastInst>(&*II);
+  if (BCI)
+    ++II;
+
+  ReturnInst* RI = dyn_cast<ReturnInst>(&*II);
+  assert(RI && "`musttail` call must be followed by `ret` instruction");
+
+  TerminatorInst *TI = SplitBB->getTerminator();
+  Value *V = NewCI;
+  if (BCI)
+    V = cloneInstForMustTail(BCI, TI, V);
+  cloneInstForMustTail(RI, TI, IsVoid ? nullptr : V);
+
+  // FIXME: remove TI here, `DuplicateInstructionsInSplitBetween` has a bug
+  // that prevents doing this now.
+}
+
 /// 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

 /// 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())
+
+  // `musttail` calls must be followed by optional `bitcast`, and `ret`. The
+  // split blocks will be terminated right after that so there're no users for
+  // this phi in a `TailBB`.
+  if (!IsMustTailCall && 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 (PHINode &PN : TailBB->phis()) {
         }
         ++ArgNo;
       }
     }
     DEBUG(dbgs() << "    " << *NewCI << " in " << SplitBlock->getName()
                  << "\n");
     if (CallPN)
       CallPN->addIncoming(NewCI, SplitBlock);
+
+    // Clone and place bitcast and return instructions before `TI`
+    if (IsMustTailCall)
+      copyMustTailReturn(SplitBlock, Instr, NewCI);
+  }
+
+  NumCallSiteSplit++;
+
+  // FIXME: remove TI in `copyMustTailReturn`
+  if (IsMustTailCall) {
+    // Remove superfluous `br` terminators from the end of the Split blocks
+    for (BasicBlock *SplitBlock : predecessors(TailBB))
+      SplitBlock->getTerminator()->eraseFromParent();
+
+    // Erase the tail block once done with musttail patching
+    TailBB->eraseFromParent();
+    return;
   }
 
   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);
   }
if (!CurrentI->use_empty()) {
       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;
   }
-
-  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())
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;

llvm/trunk/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: %cb2 = bitcast i8* %ca1 to i8*
+;CHECK: ret i8* %cb2
+;CHECK-LABEL: TBB.split
+;CHECK: %ca3 = musttail call i8* @callee(i8* nonnull %a, i8* null)
+;CHECK: %cb4 = bitcast i8* %ca3 to i8*
+;CHECK: ret i8* %cb4
+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
+}