[ICP] Handling must tail calls in indirect call promotion

[hoyFB]

Summary

Per the IR convention, a musttail call must precede a ret with an optional bitcast. This was violated by the indirect call promotion optimization which could result an IR like:

; <label>:2192:
  br i1 %2198, label %2199, label %2201, !dbg !226012, !prof !229483

; <label>:2199:                                   ; preds = %2192
  musttail call fastcc void @foo(i8* %2195), !dbg !226012
  br label %2202, !dbg !226012

; <label>:2201:                                   ; preds = %2192
  musttail call fastcc void %2197(i8* %2195), !dbg !226012
  br label %2202, !dbg !226012

; <label>:2202:                                   ; preds = %605, %2201, %2199
  ret void, !dbg !229485

This is being fixed in this change where the return statement goes together with the promoted indirect call. The code generated is like:

; <label>:2192:
  br i1 %2198, label %2199, label %2201, !dbg !226012, !prof !229483

; <label>:2199:                                   ; preds = %2192
  musttail call fastcc void @foo(i8* %2195), !dbg !226012
  ret void, !dbg !229485

; <label>:2201:                                   ; preds = %2192
  musttail call fastcc void %2197(i8* %2195), !dbg !226012
  ret void, !dbg !229485

[davidxl]

Can the code be refactored such that the splitting and cloning part of the code are mostly shared?

[wenlei]

Forgot to mention earlier - this is not unique to AutoFDO, so you may want to change the title.

[hoyFB]

In D79258#2016598, @davidxl wrote:

Can the code be refactored such that the splitting and cloning part of the code are mostly shared?

Yes, I tried that as my first attempt. The change looks like blow and placed closed to the end of the function. Much of the diversion is from cloning the return statement and the bitcast, and also the removal of the merge block. Either way works but I like the current version a bit more. What do you think?

if (OrigInst->isMustTailCall()) {
    auto Next = &(*MergeBlock->begin());
    Value *NewRetVal = NewInst;
    // Move the optional bitcast.
    if (auto *BitCast = dyn_cast_or_null<BitCastInst>(Next)) {
      assert(BitCast->getOperand(0) == OrigInst &&
             "bitcast following musttail call must use the call");
      Next = BitCast->getNextNode();
      auto NewBitCast = BitCast->clone();
      NewBitCast->replaceUsesOfWith(OrigInst, NewInst);
      BitCast->moveBefore(ElseTerm);
      NewBitCast->insertBefore(ThenTerm);
      NewRetVal = NewBitCast;
    }
    // Move return following a musttail call as well.
    ReturnInst *Ret = dyn_cast_or_null<ReturnInst>(Next);
    assert(Ret && "musttail call must precede a ret with an optional bitcast");
    if (!Ret->getReturnValue()) {
      NewRetVal = nullptr;
    }
    auto NewRet = Ret->clone();
    if (Ret->getReturnValue())
      NewRet->replaceUsesOfWith(Ret->getReturnValue(), NewRetVal);
    Ret->moveBefore(ElseTerm);
    NewRet->insertBefore(ThenTerm);
    // Invoke instructions are terminating, so we don't need the terminator
    // instructions that were just created.
    ThenTerm->eraseFromParent();
    ElseTerm->eraseFromParent();
    if (MergeBlock->empty())
      MergeBlock->eraseFromParent();
    return *NewInst;
  }

[hoyFB]

In D79258#2016601, @wenlei wrote:

Forgot to mention earlier - this is not unique to AutoFDO, so you may want to change the title.

Good point. I was thinking [Codegen] ... as the title. Maybe [ICP] ...? Not sure about the convention. @davidxl What do you think? Thanks.

[davidxl]

lgtm

[wenlei]

In D79258#2016607, @hoyFB wrote:

In D79258#2016601, @wenlei wrote:

Forgot to mention earlier - this is not unique to AutoFDO, so you may want to change the title.

Good point. I was thinking [Codegen] ... as the title. Maybe [ICP] ...? Not sure about the convention. @davidxl What do you think? Thanks.

Grepping the history, I found a few titled with [ICP].., so yeah, that's used. [TailCallElim] is also used, though I don't have preference between the two. Looks good otherwise - thanks for the fix!

[smeenai]

Should we merge this to 10.0.1?

llvm/lib/Transforms/Utils/CallPromotionUtils.cpp

 ///     %t0 = invoke i32 %ptr() to label %merge_bb unwind label %unwind_dst
 ///
 ///   merge_bb:
 ///     ; Uses of the original invoke instruction are replaced by uses of the
 ///     ; phi node, and the merge block branches to the normal destination.
 ///     %t2 = phi i32 [ %t0, %else_bb ], [ %t1, %then_bb ]
 ///     br %normal_dst
 ///
+/// An indirect musttail call is processed slightly differently in that:
+/// 1. No merge block needed for the orginal and the cloned callsite, since
+///    either one ends the flow. No phi node is needed either.
+/// 2. The return statement following the original call site is duplicated too
+///    and placed immediately after the cloned call site per the IR convention.
+///
+/// For example, the musttail call instruction below:
+///
+///   orig_bb:
+///     %t0 = musttail call i32 %ptr()
+///     ...
+///
+/// Is replaced by the following:
+///
+///   cond_bb:
+///     %cond = icmp eq i32 ()* %ptr, @func
+///     br i1 %cond, %then_bb, %orig_bb
+///
+///   then_bb:
+///     ; The clone of the original call instruction is placed in the "then"
+///     ; block. It is not yet promoted.
+///     %t1 = musttail call i32 %ptr()
+///     ret %t1
+///
+///   orig_bb:
+///     ; The original call instruction stays in its original block.
+///     %t0 = musttail call i32 %ptr()
+///     ret %t0
 static CallBase &versionCallSite(CallBase &CB, Value *Callee,
                                  MDNode *BranchWeights) {
 
   IRBuilder<> Builder(&CB);
   CallBase *OrigInst = &CB;
   BasicBlock *OrigBlock = OrigInst->getParent();
 
   // Create the compare. The called value and callee must have the same type to
   // be compared.
   if (CB.getCalledOperand()->getType() != Callee->getType())
     Callee = Builder.CreateBitCast(Callee, CB.getCalledOperand()->getType());
   auto *Cond = Builder.CreateICmpEQ(CB.getCalledOperand(), Callee);
 
+  if (OrigInst->isMustTailCall()) {
+    // Create an if-then structure. The original instruction stays in its block,
+    // and a clone of the original instruction is placed in the "then" block.
+    Instruction *ThenTerm =
+        SplitBlockAndInsertIfThen(Cond, &CB, false, BranchWeights);
+    BasicBlock *ThenBlock = ThenTerm->getParent();
+    ThenBlock->setName("if.true.direct_targ");
+    CallBase *NewInst = cast<CallBase>(OrigInst->clone());
+    NewInst->insertBefore(ThenTerm);
+
+    // Place a clone of the optional bitcast after the new call site.
+    Value *NewRetVal = NewInst;
+    auto Next = OrigInst->getNextNode();
+    if (auto *BitCast = dyn_cast_or_null<BitCastInst>(Next)) {
+      assert(BitCast->getOperand(0) == OrigInst &&
+             "bitcast following musttail call must use the call");
+      auto NewBitCast = BitCast->clone();
+      NewBitCast->replaceUsesOfWith(OrigInst, NewInst);
+      NewBitCast->insertBefore(ThenTerm);
+      NewRetVal = NewBitCast;
+      Next = BitCast->getNextNode();
+    }
+
+    // Place a clone of the return instruction after the new call site.
+    ReturnInst *Ret = dyn_cast_or_null<ReturnInst>(Next);
+    assert(Ret && "musttail call must precede a ret with an optional bitcast");
+    auto NewRet = Ret->clone();
+    if (Ret->getReturnValue())
+      NewRet->replaceUsesOfWith(Ret->getReturnValue(), NewRetVal);
+    NewRet->insertBefore(ThenTerm);
+
+    // A return instructions is terminating, so we don't need the terminator
+    // instruction just created.
+    ThenTerm->eraseFromParent();
+
+    return *NewInst;
+  }
+
   // Create an if-then-else structure. The original instruction is moved into
   // the "else" block, and a clone of the original instruction is placed in the
   // "then" block.
   Instruction *ThenTerm = nullptr;
   Instruction *ElseTerm = nullptr;
   SplitBlockAndInsertIfThenElse(Cond, &CB, &ThenTerm, &ElseTerm, BranchWeights);
   BasicBlock *ThenBlock = ThenTerm->getParent();
   BasicBlock *ElseBlock = ElseTerm->getParent();

llvm/test/Transforms/PGOProfile/indirect_call_promotion_musttail.ll

+; RUN: opt < %s -pgo-icall-prom -S | FileCheck %s --check-prefix=ICALL-PROM
+; RUN: opt < %s -passes=pgo-icall-prom -S | FileCheck %s --check-prefix=ICALL-PROM
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@foo = common global i32* ()* null, align 8
+
+declare i32* @func1()
+
+declare i32* @func2()
+
+declare i32* @func3()
+
+declare i32* @func4()
+
+define i32* @bar() {
+entry:
+  %tmp = load i32* ()*, i32* ()** @foo, align 8
+; ICALL-PROM:   [[CMP1:%[0-9]+]] = icmp eq i32* ()* %tmp, @func4
+; ICALL-PROM:   br i1 [[CMP1]], label %if.true.direct_targ, label %[[L1:[0-9]+]], !prof [[BRANCH_WEIGHT1:![0-9]+]]
+; ICALL-PROM: if.true.direct_targ:
+; ICALL-PROM:   [[DIRCALL_RET1:%[0-9]+]] = musttail call i32* @func4()
+; ICALL-PROM:   ret i32* [[DIRCALL_RET1]]
+; ICALL-PROM: [[L1]]:
+; ICALL-PROM:   [[CMP2:%[0-9]+]] = icmp eq i32* ()* %tmp, @func2
+; ICALL-PROM:   br i1 [[CMP2]], label %if.true.direct_targ1, label %[[L2:[0-9]+]], !prof [[BRANCH_WEIGHT2:![0-9]+]]
+; ICALL-PROM: if.true.direct_targ1:
+; ICALL-PROM:   [[DIRCALL_RET2:%[0-9]+]] = musttail call i32* @func2()
+; ICALL-PROM:   ret i32* [[DIRCALL_RET2]]
+; ICALL-PROM: [[L2]]:
+; ICALL-PROM:   [[CMP3:%[0-9]+]] = icmp eq i32* ()* %tmp, @func3
+; ICALL-PROM:   br i1 [[CMP3]], label %if.true.direct_targ2, label %[[L3:[0-9]+]], !prof [[BRANCH_WEIGHT3:![0-9]+]]
+; ICALL-PROM: if.true.direct_targ2:
+; ICALL-PROM:   [[DIRCALL_RET3:%[0-9]+]] = musttail call i32* @func3()
+; ICALL-PROM:   ret i32* [[DIRCALL_RET3]]
+; ICALL-PROM: [[L3]]:
+; ICALL-PROM:   %call = musttail call i32* %tmp()
+; ICALL-PROM:   ret i32* %call
+  %call = musttail call i32* %tmp(), !prof !1
+  ret i32* %call
+}
+
+define i64* @bar2() {
+entry:
+  %tmp = load i32* ()*, i32* ()** @foo, align 8
+; ICALL-PROM:   [[CMP1:%[0-9]+]] = icmp eq i32* ()* %tmp, @func4
+; ICALL-PROM:   br i1 [[CMP1]], label %if.true.direct_targ, label %[[L4:[0-9]+]], !prof [[BRANCH_WEIGHT4:![0-9]+]]
+; ICALL-PROM: if.true.direct_targ:
+; ICALL-PROM:   [[DIRCALL_RET1:%[0-9]+]] = musttail call i32* @func4()
+; ICALL-PROM:   [[DIRCALL_RET2:%[0-9]+]] = bitcast i32* [[DIRCALL_RET1]] to i64*
+; ICALL-PROM:   ret i64* [[DIRCALL_RET2]]
+; ICALL-PROM: [[L4]]:
+; ICALL-PROM:   %call = musttail call i32* %tmp()
+; ICALL-PROM:   %rv = bitcast i32* %call to i64*
+; ICALL-PROM:   ret i64* %rv
+  %call = musttail call i32* %tmp(), !prof !2
+  %rv = bitcast i32* %call to i64*
+  ret i64* %rv
+}
+
+!1 = !{!"VP", i32 0, i64 1600, i64 7651369219802541373, i64 1030, i64 -4377547752858689819, i64 410, i64 -6929281286627296573, i64 150, i64 -2545542355363006406, i64 10}
+!2 = !{!"VP", i32 0, i64 100, i64 7651369219802541373, i64 100}
+
+; ICALL-PROM: [[BRANCH_WEIGHT1]] = !{!"branch_weights", i32 1030, i32 570}
+; ICALL-PROM: [[BRANCH_WEIGHT2]] = !{!"branch_weights", i32 410, i32 160}
+; ICALL-PROM: [[BRANCH_WEIGHT3]] = !{!"branch_weights", i32 150, i32 10}
+; ICALL-PROM: [[BRANCH_WEIGHT4]] = !{!"branch_weights", i32 100, i32 0}