[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

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

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

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

220

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

Isn't it possible to set isVoid here?

232

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

236

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

246

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

250

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

254–263

Look like we can make a separate function for cloning.

282

if (RI) should be better.

350

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

211–212

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

259

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

311

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.

314

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

327

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

327

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

214

I don't think we need to do that.

264

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.

269

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

275

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

316

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

322

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

339

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

214

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

262

blokcs -> blocks

327

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) {
   // Need 2 predecessors and cannot split an edge from an IndirectBrInst.
   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.
-///
 /// 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.
 ///
 /// From :
 ///
 ///   Header:
 ///     %c = icmp eq i32* %a, null
 ///     br i1 %c %Tail, %TBB
 ///   TBB:
 ///     %c2 = icmp eq i32* %b, null
 ///     br i1 %c %Tail, %End
 ///   Tail:
 ///     %ca = call i1  @callee (i32* %a, i32* %b)
 ///
 ///  to :
 ///
 ///   Header:                          // PredBB1 is Header
 ///     %c = icmp eq i32* %a, null
 ///     br i1 %c %Tail-split1, %TBB
 ///   TBB:                             // PredBB2 is TBB
 ///     %c2 = icmp eq i32* %b, null
 ///     br i1 %c %Tail-split2, %End
 ///   Tail-split1:
 ///     %ca1 = call @callee (i32* null, i32* %b)         // CallInst1
 ///    br %Tail
 ///   Tail-split2:
 ///     %ca2 = call @callee (i32* nonnull %a, i32* null) // CallInst2
 ///    br %Tail
 ///   Tail:
 ///    %p = phi i1 [%ca1, %Tail-split1],[%ca2, %Tail-split2]
 ///
 /// 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) {
+  bool IsMustTail = CS.isMustTailCall();
   Instruction *Instr = CS.getInstruction();
   BasicBlock *TailBB = Instr->getParent();
 
   PHINode *CallPN = nullptr;
+
+  Instruction *StopAt;
+  if (IsMustTail) {
+    // We need to copy everything including the terminator of the Tail block,
+    // since musttail must be followed by optional bitcast and ret.
+    // Insert temporary unreachable instruction to be used as a marker for
+    // `DuplicateInstructionsInSplitBetween`.
+    StopAt = new UnreachableInst(TailBB->getContext(), TailBB);
+  } else {
+    StopAt = &*std::next(Instr->getIterator());
     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]);
+        TailBB, PredBB, StopAt, ValueToValueMaps[i]);
     assert(SplitBlock && "Unexpected new basic block split.");
 
-    Instruction *NewCI =
-        &*std::prev(SplitBlock->getTerminator()->getIterator());
+    Instruction *NewCI;
+
+    if (IsMustTail) {
+      // The pre-last instruction is `unreachable`, do a full lookup to
+      // find the CI's copy.
+      NewCI = dyn_cast<Instruction>(&*ValueToValueMaps[i][Instr]);
+    } else {
+      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();
+  auto I = (&*std::prev(StopAt->getIterator()))->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;
   }
 
+  // At this point the splits have following instructions at the end:
+  //
+  //   %c = call musttail ...
+  //   ret ... %c
+  //   br TBB ; or Tail
+  //
+  // Remove `br`, effectively unlinking the splits from the Tail.
+  // Once this is done - remove the Tail itself.
+  if (IsMustTail) {
+    DEBUG(dbgs() << "clean up musttail splits\n");
+
+    SmallVector<BasicBlock *, 2> Splits(predecessors((TailBB)));
+    assert(Splits.size() == 2 && "Expected exactly 2 splits!");
+
+    for (unsigned i = 0; i < Splits.size(); i++)
+      Splits[i]->getTerminator()->eraseFromParent();
+
+    TailBB->eraseFromParent();
+  }
+
   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();
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: %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
+}