Teach the inliner how to preserve musttail invariants

[chandlerc]

lib/IR/Instructions.cpp

327 ↗

(On Diff #8816)

Please commit this separately with a separate test? Maybe a unittest?

lib/Transforms/Utils/InlineFunction.cpp

483 ↗

(On Diff #8816)

This is called in a context that needs it to not go truly linear... And it doesn't. I almost want to call it 'isImmediatelyPrecededByMustTailCall', but the bitcast thing makes that a lie. Maybe just some comments?

687–690 ↗

(On Diff #8816)

I assume the after cases should have "h"? Actually, shouldn't these be "f -> g -> f" => "f -> f"?

692–693 ↗

(On Diff #8816)

I would sink this comment to the loop body.

744–745 ↗

(On Diff #8816)

Do all of the passes using lifetime markers know this?

It might be useful to separate this into its own patch with its own test case.

795 ↗

(On Diff #8816)

You say "should" but then "assert" below. I think the should is correct though...

804 ↗

(On Diff #8816)

So, I don't think you can *always* assert. What about a musttail call to a function that just happens to return the result of a non-musttail call? Above, you set the flag if any call is musttail, so it seems you should just continue if this call site isn't a musttail.

Similarly, I don't think you can assert that there will be a callinst here.

[rnk]

lib/IR/Instructions.cpp

327 ↗

(On Diff #8816)

Done.

lib/Transforms/Utils/InlineFunction.cpp

483 ↗

(On Diff #8816)

Comments sound good to me.

687–690 ↗

(On Diff #8816)

Yeah. Let's go with f -> g -> f. That's the example described in the paragraph above. The reader can extrapolate a larger cyclical call graph.

692–693 ↗

(On Diff #8816)

done

744–745 ↗

(On Diff #8816)

Sure. I'll split this into a followup patch to highlight the test case that fails without it. It fixes the byval and dynalloca test cases.

795 ↗

(On Diff #8816)

The assert was correct. I reworked this to be simpler.

804 ↗

(On Diff #8816)

We are looping over the second partition here, i.e. returns that are preceded by musttail. So, all of these things have to be true, or there are bugs in isPrecededByMustTailCall.

It sounds like this was confusing, so I'll restructure it without partition. That was silly.

[rnk]

• review
• split out lifetime.end change and test cases

[chandlerc]

Ok, this makes more sense now. Thahnks for the clue-bat.

I want to ruminate on the bit about not needing stackrestore. I think I believe you, but I don't fully understand why yet and I would like to. But below are some easy comments to just clean up code.

lib/Transforms/Utils/InlineFunction.cpp

697 ↗

(On Diff #9130)

nit: invert and use continue to reduce indentation

794 ↗

(On Diff #9130)

nit: as above, invert and continue to reduce indentation.

795–800 ↗

(On Diff #9130)

Rather than duplicating isPrecededByMustTailCall, make that function getPrecedingMustTailCallInst or some such, and have it return null when there is no call? That should make this loop simplify quite a bit.

[rnk]

• reduce indentation
• return the preceding musttail call

lib/Transforms/Utils/InlineFunction.cpp

697 ↗

(On Diff #9130)

done

794 ↗

(On Diff #9130)

done

795–800 ↗

(On Diff #9130)

Sure. It helps a bit, but I have to get the bitcast instr so I can erase it.

[chandlerc]

Sorry for continually getting distracted here. I think this is essentially fine. Feel free to submit whenever, and we can figure out the comments and the confusing parts as we go. I don't think there is anything wrong here.

lib/Transforms/Utils/InlineFunction.cpp

482–483 ↗

(On Diff #9136)

Stale comment.

694–705 ↗

(On Diff #9136)

Would it make sense to sink this comment down to the TailCallKind bit below?

780–784 ↗

(On Diff #9136)

It's really surprising to me that this applies to *all* dynamic allocas rather than just to inalloca arguments... Maybe I'm just missing the point, but I'd appreciate a comment clarifying why this makes sense.

Is there a test case that covers this? I can't tell if some of the CHECK-NEXT ones actually cover this case.

[chandlerc]

(and actually mork it as LGTM in phab...)

[rnk]

Thanks, committing. I'll update the follow-on, http://reviews.llvm.org/D3630, in a sec.

lib/Transforms/Utils/InlineFunction.cpp

482–483 ↗

(On Diff #9136)

fixed

694–705 ↗

(On Diff #9136)

Seems reasonable.

780–784 ↗

(On Diff #9136)

I'm not really following.

Here's what could happen when inlining this fragment:

%a = alloca i8, i32 %n
call void @use_a(i8* %a)
musttail call void @bar()
ret void

Without this change, after inlining, we get this:

%sp = call i8* @llvm.stacksave()
%a = alloca i8, i32 %n
call void @use_a(i8* %a)
musttail call void @bar()
call void @llvm.stackrestore(i8*)
ret void

First, this fails musttail verification, because now the musttail call isn't in tail position.

The stackrestore is completely redundant with the ret. Unlike other rets, this ret is going to remain in the inlined function. It isn't removed later and merged back into the normal control flow.

Would it make more sense to reframe the comment to talk about the fact that this ret isn't going a way and stackrestore before ret is pointless? The way I think about it is that we just did a musttail call which already did the work of clearing the stack, so there's no reason to restore the stack.

Anyway, this should have been split into the @lifetime.end change, because it's covered by those tests. I'll go do that.

[rnk]

Closed by commit rL208910 (authored by @rnk).

llvm/trunk/lib/Transforms/Utils/InlineFunction.cpp

 
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/CallSite.h"
+#include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
for (BasicBlock::iterator BI = FI->begin(), BE = FI->end();
           DVI->setOperand(2, createInlinedVariable(DVI->getVariable(), 
                                                    InlinedAt, Ctx));
         }
       }
     }
   }
 }
 
+/// Returns a musttail call instruction if one immediately precedes the given
+/// return instruction with an optional bitcast instruction between them.
+static CallInst *getPrecedingMustTailCall(ReturnInst *RI) {
+  Instruction *Prev = RI->getPrevNode();
+  if (!Prev)
+    return nullptr;
+
+  if (Value *RV = RI->getReturnValue()) {
+    if (RV != Prev)
+      return nullptr;
+
+    // Look through the optional bitcast.
+    if (auto *BI = dyn_cast<BitCastInst>(Prev)) {
+      RV = BI->getOperand(0);
+      Prev = BI->getPrevNode();
+      if (!Prev || RV != Prev)
+        return nullptr;
+    }
+  }
+
+  if (auto *CI = dyn_cast<CallInst>(Prev)) {
+    if (CI->isMustTailCall())
+      return CI;
+  }
+  return nullptr;
+}
+
 /// InlineFunction - This function inlines the called function into the basic
 /// block of the caller.  This returns false if it is not possible to inline
 /// this call.  The program is still in a well defined state if this occurs
 /// though.
 ///
 /// Note that this only does one level of inlining.  For example, if the
 /// instruction 'call B' is inlined, and 'B' calls 'C', then the call to 'C' now
 /// exists in the instruction stream.  Similarly this will inline a recursive
bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
   
   const Function *CalledFunc = CS.getCalledFunction();
   if (!CalledFunc ||              // Can't inline external function or indirect
       CalledFunc->isDeclaration() || // call, or call to a vararg function!
       CalledFunc->getFunctionType()->isVarArg()) return false;
 
   // If the call to the callee is not a tail call, we must clear the 'tail'
   // flags on any calls that we inline.
-  bool MustClearTailCallFlags =
-    !(isa<CallInst>(TheCall) && cast<CallInst>(TheCall)->isTailCall());
+  CallInst::TailCallKind CallSiteTailKind = CallInst::TCK_None;
+  if (CallInst *CI = dyn_cast<CallInst>(TheCall))
+    CallSiteTailKind = CI->getTailCallKind();
+  bool MustClearTailCallFlags = false;
 
   // If the call to the callee cannot throw, set the 'nounwind' flag on any
   // calls that we inline.
   bool MarkNoUnwind = CS.doesNotThrow();
 
   BasicBlock *OrigBB = TheCall->getParent();
   Function *Caller = OrigBB->getParent();
 
for (BasicBlock::iterator I = FirstNewBlock->begin(),
       // that the instructions aren't removed from the symbol table, then
       // reinserted.
       Caller->getEntryBlock().getInstList().splice(InsertPoint,
                                                    FirstNewBlock->getInstList(),
                                                    AI, I);
     }
   }
 
+  bool InlinedMustTailCalls = false;
+  if (InlinedFunctionInfo.ContainsCalls) {
+    for (Function::iterator BB = FirstNewBlock, E = Caller->end(); BB != E;
+         ++BB) {
+      for (Instruction &I : *BB) {
+        CallInst *CI = dyn_cast<CallInst>(&I);
+        if (!CI)
+          continue;
+
+        // We need to reduce the strength of any inlined tail calls.  For
+        // musttail, we have to avoid introducing potential unbounded stack
+        // growth.  For example, if functions 'f' and 'g' are mutually recursive
+        // with musttail, we can inline 'g' into 'f' so long as we preserve
+        // musttail on the cloned call to 'f'.  If either the inlined call site
+        // or the cloned call site is *not* musttail, the program already has
+        // one frame of stack growth, so it's safe to remove musttail.  Here is
+        // a table of example transformations:
+        //
+        //    f -> musttail g -> musttail f  ==>  f -> musttail f
+        //    f -> musttail g ->     tail f  ==>  f ->     tail f
+        //    f ->          g -> musttail f  ==>  f ->          f
+        //    f ->          g ->     tail f  ==>  f ->          f
+        CallInst::TailCallKind ChildTCK = CI->getTailCallKind();
+        ChildTCK = std::min(CallSiteTailKind, ChildTCK);
+        CI->setTailCallKind(ChildTCK);
+        InlinedMustTailCalls |= CI->isMustTailCall();
+
+        // Calls inlined through a 'nounwind' call site should be marked
+        // 'nounwind'.
+        if (MarkNoUnwind)
+          CI->setDoesNotThrow();
+      }
+    }
+  }
+
   // Leave lifetime markers for the static alloca's, scoping them to the
   // function we just inlined.
   if (InsertLifetime && !IFI.StaticAllocas.empty()) {
     IRBuilder<> builder(FirstNewBlock->begin());
     for (unsigned ai = 0, ae = IFI.StaticAllocas.size(); ai != ae; ++ai) {
       AllocaInst *AI = IFI.StaticAllocas[ai];
 
       // If the alloca is already scoped to something smaller than the whole
for (unsigned ai = 0, ae = IFI.StaticAllocas.size(); ai != ae; ++ai) {
               UINT64_MAX / AllocaArraySize >= AllocaTypeSize) {
             AllocaSize = ConstantInt::get(Type::getInt64Ty(AI->getContext()),
                                           AllocaArraySize * AllocaTypeSize);
           }
         }
       }
 
       builder.CreateLifetimeStart(AI, AllocaSize);
-      for (unsigned ri = 0, re = Returns.size(); ri != re; ++ri) {
-        IRBuilder<> builder(Returns[ri]);
-        builder.CreateLifetimeEnd(AI, AllocaSize);
-      }
+      for (ReturnInst *RI : Returns)
+        IRBuilder<>(RI).CreateLifetimeEnd(AI, AllocaSize);
     }
   }
 
   // If the inlined code contained dynamic alloca instructions, wrap the inlined
   // code with llvm.stacksave/llvm.stackrestore intrinsics.
   if (InlinedFunctionInfo.ContainsDynamicAllocas) {
     Module *M = Caller->getParent();
     // Get the two intrinsics we care about.
     Function *StackSave = Intrinsic::getDeclaration(M, Intrinsic::stacksave);
     Function *StackRestore=Intrinsic::getDeclaration(M,Intrinsic::stackrestore);
 
     // Insert the llvm.stacksave.
     CallInst *SavedPtr = IRBuilder<>(FirstNewBlock, FirstNewBlock->begin())
       .CreateCall(StackSave, "savedstack");
 
     // Insert a call to llvm.stackrestore before any return instructions in the
     // inlined function.
-    for (unsigned i = 0, e = Returns.size(); i != e; ++i) {
-      IRBuilder<>(Returns[i]).CreateCall(StackRestore, SavedPtr);
-    }
-  }
-
-  // If we are inlining tail call instruction through a call site that isn't
-  // marked 'tail', we must remove the tail marker for any calls in the inlined
-  // code.  Also, calls inlined through a 'nounwind' call site should be marked
-  // 'nounwind'.
-  if (InlinedFunctionInfo.ContainsCalls &&
-      (MustClearTailCallFlags || MarkNoUnwind)) {
-    for (Function::iterator BB = FirstNewBlock, E = Caller->end();
-         BB != E; ++BB)
-      for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
-        if (CallInst *CI = dyn_cast<CallInst>(I)) {
-          if (MustClearTailCallFlags)
-            CI->setTailCall(false);
-          if (MarkNoUnwind)
-            CI->setDoesNotThrow();
-        }
+    for (ReturnInst *RI : Returns)
+      IRBuilder<>(RI).CreateCall(StackRestore, SavedPtr);
   }
 
   // If we are inlining for an invoke instruction, we must make sure to rewrite
   // any call instructions into invoke instructions.
   if (InvokeInst *II = dyn_cast<InvokeInst>(TheCall))
     HandleInlinedInvoke(II, FirstNewBlock, InlinedFunctionInfo);
 
+  // Handle any inlined musttail call sites.  In order for a new call site to be
+  // musttail, the source of the clone and the inlined call site must have been
+  // musttail.  Therefore it's safe to return without merging control into the
+  // phi below.
+  if (InlinedMustTailCalls) {
+    // Check if we need to bitcast the result of any musttail calls.
+    Type *NewRetTy = Caller->getReturnType();
+    bool NeedBitCast = !TheCall->use_empty() && TheCall->getType() != NewRetTy;
+
+    // Handle the returns preceded by musttail calls separately.
+    SmallVector<ReturnInst *, 8> NormalReturns;
+    for (ReturnInst *RI : Returns) {
+      CallInst *ReturnedMustTail = getPrecedingMustTailCall(RI);
+      if (!ReturnedMustTail) {
+        NormalReturns.push_back(RI);
+        continue;
+      }
+      if (!NeedBitCast)
+        continue;
+
+      // Delete the old return and any preceding bitcast.
+      BasicBlock *CurBB = RI->getParent();
+      auto *OldCast = dyn_cast_or_null<BitCastInst>(RI->getReturnValue());
+      RI->eraseFromParent();
+      if (OldCast)
+        OldCast->eraseFromParent();
+
+      // Insert a new bitcast and return with the right type.
+      IRBuilder<> Builder(CurBB);
+      Builder.CreateRet(Builder.CreateBitCast(ReturnedMustTail, NewRetTy));
+    }
+
+    // Leave behind the normal returns so we can merge control flow.
+    std::swap(Returns, NormalReturns);
+  }
+
   // If we cloned in _exactly one_ basic block, and if that block ends in a
   // return instruction, we splice the body of the inlined callee directly into
   // the calling basic block.
   if (Returns.size() == 1 && std::distance(FirstNewBlock, Caller->end()) == 1) {
     // Move all of the instructions right before the call.
     OrigBB->getInstList().splice(TheCall, FirstNewBlock->getInstList(),
                                  FirstNewBlock->begin(), FirstNewBlock->end());
     // Remove the cloned basic block.
if (Returns.size() > 1) {
     // No returns, but something is using the return value of the call.  Just
     // nuke the result.
     TheCall->replaceAllUsesWith(UndefValue::get(TheCall->getType()));
   }
 
   // Since we are now done with the Call/Invoke, we can delete it.
   TheCall->eraseFromParent();
 
+  // If we inlined any musttail calls and the original return is now
+  // unreachable, delete it.  It can only contain a bitcast and ret.
+  if (InlinedMustTailCalls && pred_begin(AfterCallBB) == pred_end(AfterCallBB))
+    AfterCallBB->eraseFromParent();
+
   // We should always be able to fold the entry block of the function into the
   // single predecessor of the block...
   assert(cast<BranchInst>(Br)->isUnconditional() && "splitBasicBlock broken!");
   BasicBlock *CalleeEntry = cast<BranchInst>(Br)->getSuccessor(0);
 
   // Splice the code entry block into calling block, right before the
   // unconditional branch.
   CalleeEntry->replaceAllUsesWith(OrigBB);  // Update PHI nodes

llvm/trunk/test/Transforms/Inline/inline-tail.ll

-; RUN: opt < %s -inline -S | not grep tail
+; RUN: opt < %s -inline -S | FileCheck %s
 
-declare void @bar(i32*)
+; We have to apply the less restrictive TailCallKind of the call site being
+; inlined and any call sites cloned into the caller.
 
-define internal void @foo(i32* %P) {
-        tail call void @bar( i32* %P )
-        ret void
+; No tail marker after inlining, since test_capture_c captures an alloca.
+; CHECK: define void @test_capture_a(
+; CHECK-NOT: tail
+; CHECK: call void @test_capture_c(
+
+declare void @test_capture_c(i32*)
+define internal void @test_capture_b(i32* %P) {
+  tail call void @test_capture_c(i32* %P)
+  ret void
+}
+define void @test_capture_a() {
+  %A = alloca i32     ; captured by test_capture_b
+  call void @test_capture_b(i32* %A)
+  ret void
+}
+
+; No musttail marker after inlining, since the prototypes don't match.
+; CHECK: define void @test_proto_mismatch_a(
+; CHECK-NOT: musttail
+; CHECK: call void @test_proto_mismatch_c(
+
+declare void @test_proto_mismatch_c(i32*)
+define internal void @test_proto_mismatch_b(i32* %p) {
+  musttail call void @test_proto_mismatch_c(i32* %p)
+  ret void
+}
+define void @test_proto_mismatch_a() {
+  call void @test_proto_mismatch_b(i32* null)
+  ret void
+}
+
+; After inlining through a musttail call site, we need to keep musttail markers
+; to prevent unbounded stack growth.
+; CHECK: define void @test_musttail_basic_a(
+; CHECK: musttail call void @test_musttail_basic_c(
+
+declare void @test_musttail_basic_c(i32* %p)
+define internal void @test_musttail_basic_b(i32* %p) {
+  musttail call void @test_musttail_basic_c(i32* %p)
+  ret void
+}
+define void @test_musttail_basic_a(i32* %p) {
+  musttail call void @test_musttail_basic_b(i32* %p)
 }
 
-define void @caller() {
-        %A = alloca i32         ; <i32*> [#uses=1]
-        call void @foo( i32* %A )
-        ret void
+; We can't merge the returns.
+; CHECK: define void @test_multiret_a(
+; CHECK: musttail call void @test_multiret_c(
+; CHECK-NEXT: ret void
+; CHECK: musttail call void @test_multiret_d(
+; CHECK-NEXT: ret void
+
+declare void @test_multiret_c(i1 zeroext %b)
+declare void @test_multiret_d(i1 zeroext %b)
+define internal void @test_multiret_b(i1 zeroext %b) {
+  br i1 %b, label %c, label %d
+c:
+  musttail call void @test_multiret_c(i1 zeroext %b)
+  ret void
+d:
+  musttail call void @test_multiret_d(i1 zeroext %b)
+  ret void
+}
+define void @test_multiret_a(i1 zeroext %b) {
+  musttail call void @test_multiret_b(i1 zeroext %b)
 }
 
+; We have to avoid bitcast chains.
+; CHECK: define i32* @test_retptr_a(
+; CHECK: musttail call i8* @test_retptr_c(
+; CHECK-NEXT: bitcast i8* {{.*}} to i32*
+; CHECK-NEXT: ret i32*
+
+declare i8* @test_retptr_c()
+define internal i16* @test_retptr_b() {
+  %rv = musttail call i8* @test_retptr_c()
+  %v = bitcast i8* %rv to i16*
+  ret i16* %v
+}
+define i32* @test_retptr_a() {
+  %rv = musttail call i16* @test_retptr_b()
+  %v = bitcast i16* %rv to i32*
+  ret i32* %v
+}
+
+; Combine the last two cases: multiple returns with pointer bitcasts.
+; CHECK: define i32* @test_multiptrret_a(
+; CHECK: musttail call i8* @test_multiptrret_c(
+; CHECK-NEXT: bitcast i8* {{.*}} to i32*
+; CHECK-NEXT: ret i32*
+; CHECK: musttail call i8* @test_multiptrret_d(
+; CHECK-NEXT: bitcast i8* {{.*}} to i32*
+; CHECK-NEXT: ret i32*
+
+declare i8* @test_multiptrret_c(i1 zeroext %b)
+declare i8* @test_multiptrret_d(i1 zeroext %b)
+define internal i16* @test_multiptrret_b(i1 zeroext %b) {
+  br i1 %b, label %c, label %d
+c:
+  %c_rv = musttail call i8* @test_multiptrret_c(i1 zeroext %b)
+  %c_v = bitcast i8* %c_rv to i16*
+  ret i16* %c_v
+d:
+  %d_rv = musttail call i8* @test_multiptrret_d(i1 zeroext %b)
+  %d_v = bitcast i8* %d_rv to i16*
+  ret i16* %d_v
+}
+define i32* @test_multiptrret_a(i1 zeroext %b) {
+  %rv = musttail call i16* @test_multiptrret_b(i1 zeroext %b)
+  %v = bitcast i16* %rv to i32*
+  ret i32* %v
+}
+
+; Inline a musttail call site which contains a normal return and a musttail call.
+; CHECK: define i32 @test_mixedret_a(
+; CHECK: br i1 %b
+; CHECK: musttail call i32 @test_mixedret_c(
+; CHECK-NEXT: ret i32
+; CHECK: call i32 @test_mixedret_d(i1 zeroext %b)
+; CHECK: add i32 1,
+; CHECK-NOT: br
+; CHECK: ret i32
+
+declare i32 @test_mixedret_c(i1 zeroext %b)
+declare i32 @test_mixedret_d(i1 zeroext %b)
+define internal i32 @test_mixedret_b(i1 zeroext %b) {
+  br i1 %b, label %c, label %d
+c:
+  %c_rv = musttail call i32 @test_mixedret_c(i1 zeroext %b)
+  ret i32 %c_rv
+d:
+  %d_rv = call i32 @test_mixedret_d(i1 zeroext %b)
+  %d_rv1 = add i32 1, %d_rv
+  ret i32 %d_rv1
+}
+define i32 @test_mixedret_a(i1 zeroext %b) {
+  %rv = musttail call i32 @test_mixedret_b(i1 zeroext %b)
+  ret i32 %rv
+}