[Coroutine] Do not CoroElide if there are musttail calls

[lxfind]

Summary

This is to address https://bugs.llvm.org/show_bug.cgi?id=48626.
When there are musttail calls that use parameters aliasing the newly created coroutine frame, the existing implementation will fatal.
We simply cannot perform CoroElide in such cases. In theory a precise analysis can be done to check whether the parameters of the musttail call
actually alias the frame, but it's very hard to do it before the transformation happens. Also in most cases the existence of musttail call is
generated due to symmetric transfers, and in those cases alias analysis won't be able to tell that they don't alias anyway.

[bruno]

Nice correctness improvement!

llvm/lib/Transforms/Coroutines/CoroElide.cpp

259

How about walking the BBs and check by using getTerminatingMustTailCall()?

[lxfind]

llvm/lib/Transforms/Coroutines/CoroElide.cpp

259

That's nice. Thanks!

[lxfind]

address comments

[palebedev]

I confirm this fixes the problem I've originally reported, thank you! Too bad this is yet another CoroElide pessimization...

[ChuanqiXu]

I am a little unclear about this problem. From my point of view, it seems like that there is a Coroutine C elided in a normal function F. Then in the Coroutine Body of C, it would try to switch to itself by symmetric transfer. However, the Coroutine frame of C now is a stack variable. Then the tail call would pass the address of the stack variable whose lifetime has ended, so here is the corruption. Did I understand the situation?

[junparser]

Thanks for fix this. Furthermore, we can add runtime check to see whether this pointer is current coroutine frame, and optimize them separately.

[lxfind]

In D94834#2504128, @ChuanqiXu wrote:

I am a little unclear about this problem. From my point of view, it seems like that there is a Coroutine C elided in a normal function F. Then in the Coroutine Body of C, it would try to switch to itself by symmetric transfer. However, the Coroutine frame of C now is a stack variable. Then the tail call would pass the address of the stack variable whose lifetime has ended, so here is the corruption. Did I understand the situation?

Yes I believe that's an accurate description

[ChuanqiXu]

In D94834#2504992, @lxfind wrote:

In D94834#2504128, @ChuanqiXu wrote:

I am a little unclear about this problem. From my point of view, it seems like that there is a Coroutine C elided in a normal function F. Then in the Coroutine Body of C, it would try to switch to itself by symmetric transfer. However, the Coroutine frame of C now is a stack variable. Then the tail call would pass the address of the stack variable whose lifetime has ended, so here is the corruption. Did I understand the situation?

Yes I believe that's an accurate description

I'm looking into this bug now. Now it is more confusing to me since it seems like that this process shouldn't happen.

The IR generated for the attachment of bug48626 is like:

%19 = tail call noalias nonnull i8* @llvm.coro.begin(token %15, i8* %18), !noalias !1082 ; The handle we are going to elide
// ... in another BB
%35 = call i8* @llvm.coro.subfn.addr(i8* nonnull %19, i8 1) #28 ; destroy %19
// ... in another BB
musttail call fastcc void %51(i8* %retval.sroa.0.0.copyload.i.i.i120) ; dominated by `call i8* @llvm.coro.subfn.addr(i8* nonnull %19, i8 1)`
ret void

It seems like that the argument of the musttail call can't alias %19 in any sense since %19 already been destroyed before the musttail call. Then it makes sense to all situations I can image. Since we only elide the coro handles who wouldn't escape. And we did the escape analysis very conservatively (Not by AA analysis, but by a simple path based escape analysis). So I think it may be better to remove the check in the end of ShouldElide function. It seems like that the memory wouldn't crash due to the argument of the musttail call alias with the coroutine handle we elided. And if any corruption happens, I think is is a bug for the escape analysis.

[lxfind]

In D94834#2519348, @ChuanqiXu wrote:

In D94834#2504992, @lxfind wrote:

In D94834#2504128, @ChuanqiXu wrote:

I am a little unclear about this problem. From my point of view, it seems like that there is a Coroutine C elided in a normal function F. Then in the Coroutine Body of C, it would try to switch to itself by symmetric transfer. However, the Coroutine frame of C now is a stack variable. Then the tail call would pass the address of the stack variable whose lifetime has ended, so here is the corruption. Did I understand the situation?

Yes I believe that's an accurate description

I'm looking into this bug now. Now it is more confusing to me since it seems like that this process shouldn't happen.

The IR generated for the attachment of bug48626 is like:

%19 = tail call noalias nonnull i8* @llvm.coro.begin(token %15, i8* %18), !noalias !1082 ; The handle we are going to elide
// ... in another BB
%35 = call i8* @llvm.coro.subfn.addr(i8* nonnull %19, i8 1) #28 ; destroy %19
// ... in another BB
musttail call fastcc void %51(i8* %retval.sroa.0.0.copyload.i.i.i120) ; dominated by `call i8* @llvm.coro.subfn.addr(i8* nonnull %19, i8 1)`
ret void

It seems like that the argument of the musttail call can't alias %19 in any sense since %19 already been destroyed before the musttail call. Then it makes sense to all situations I can image. Since we only elide the coro handles who wouldn't escape. And we did the escape analysis very conservatively (Not by AA analysis, but by a simple path based escape analysis). So I think it may be better to remove the check in the end of ShouldElide function. It seems like that the memory wouldn't crash due to the argument of the musttail call alias with the coroutine handle we elided. And if any corruption happens, I think is is a bug for the escape analysis.

It won't crash in this particular case. But as discussed in the bug thread, in theory the symmetric transfer could return the same handle as passed in, and in that case you won't be able to do tail call, right?

[ChuanqiXu]

In D94834#2520137, @lxfind wrote:

In D94834#2519348, @ChuanqiXu wrote:

In D94834#2504992, @lxfind wrote:

In D94834#2504128, @ChuanqiXu wrote:

I am a little unclear about this problem. From my point of view, it seems like that there is a Coroutine C elided in a normal function F. Then in the Coroutine Body of C, it would try to switch to itself by symmetric transfer. However, the Coroutine frame of C now is a stack variable. Then the tail call would pass the address of the stack variable whose lifetime has ended, so here is the corruption. Did I understand the situation?

Yes I believe that's an accurate description

I'm looking into this bug now. Now it is more confusing to me since it seems like that this process shouldn't happen.

The IR generated for the attachment of bug48626 is like:

%19 = tail call noalias nonnull i8* @llvm.coro.begin(token %15, i8* %18), !noalias !1082 ; The handle we are going to elide
// ... in another BB
%35 = call i8* @llvm.coro.subfn.addr(i8* nonnull %19, i8 1) #28 ; destroy %19
// ... in another BB
musttail call fastcc void %51(i8* %retval.sroa.0.0.copyload.i.i.i120) ; dominated by `call i8* @llvm.coro.subfn.addr(i8* nonnull %19, i8 1)`
ret void

It seems like that the argument of the musttail call can't alias %19 in any sense since %19 already been destroyed before the musttail call. Then it makes sense to all situations I can image. Since we only elide the coro handles who wouldn't escape. And we did the escape analysis very conservatively (Not by AA analysis, but by a simple path based escape analysis). So I think it may be better to remove the check in the end of ShouldElide function. It seems like that the memory wouldn't crash due to the argument of the musttail call alias with the coroutine handle we elided. And if any corruption happens, I think is is a bug for the escape analysis.

It won't crash in this particular case. But as discussed in the bug thread, in theory the symmetric transfer could return the same handle as passed in, and in that case you won't be able to do tail call, right?

Yes and my thoughts about this case is that CoroElide pass won't elide the frame. As a general example:

define void @xxx.resume(%FrameTy* %FramePtr) {
entry:
     ; ...
BB:
     %handle = tail call noalias nonnull i8* @llvm.coro.begin(...)
     ; ...
BB1:
     call i8* @llvm.subfn.addr(%handle, 1)
     ; ...
BB.final:
      musttail call fastcc void @resume.call(i8* %val) 
      ret void
}

And I think there are two cases about this example:
(1) The musttail call is dominated by @llvm.subfn.addr(%handle, 1) calls. Then %handle is destroyed at the musttail call, so it makes no sense that the musttail call returns %handle.
(2) The musttail call isn't dominated by @llvm.subfn.addr(%handle, 1) calls. Then the musttail call may returns %handle. And it doesn't matter since %handle won't be elided in this case.

[lxfind]

And I think there are two cases about this example:
(1) The musttail call is dominated by @llvm.subfn.addr(%handle, 1) calls. Then %handle is destroyed at the musttail call, so it makes no sense that the musttail call returns %handle.
(2) The musttail call isn't dominated by @llvm.subfn.addr(%handle, 1) calls. Then the musttail call may returns %handle. And it doesn't matter since %handle won't be elided in this case.

In case (1), why is it that '%handle' must be destroyed at the musttail call?

[ChuanqiXu]

In D94834#2521729, @lxfind wrote:

And I think there are two cases about this example:
(1) The musttail call is dominated by @llvm.subfn.addr(%handle, 1) calls. Then %handle is destroyed at the musttail call, so it makes no sense that the musttail call returns %handle.
(2) The musttail call isn't dominated by @llvm.subfn.addr(%handle, 1) calls. Then the musttail call may returns %handle. And it doesn't matter since %handle won't be elided in this case.

In case (1), why is it that '%handle' must be destroyed at the musttail call?

In my mind, @llvm.subfn.addr(%handle, 1) means a call which would destroy %handle. And CoroElide Pass also treat the call as a lifetime end marker for %handle. CoroElide pass only elides coroutine handle who is guarded by @llvm.subfn.addr(%handle, 1) .

[lxfind]

In D94834#2521744, @ChuanqiXu wrote:

In D94834#2521729, @lxfind wrote:

And I think there are two cases about this example:
(1) The musttail call is dominated by @llvm.subfn.addr(%handle, 1) calls. Then %handle is destroyed at the musttail call, so it makes no sense that the musttail call returns %handle.
(2) The musttail call isn't dominated by @llvm.subfn.addr(%handle, 1) calls. Then the musttail call may returns %handle. And it doesn't matter since %handle won't be elided in this case.

In case (1), why is it that '%handle' must be destroyed at the musttail call?

In my mind, @llvm.subfn.addr(%handle, 1) means a call which would destroy %handle. And CoroElide Pass also treat the call as a lifetime end marker for %handle. CoroElide pass only elides coroutine handle who is guarded by @llvm.subfn.addr(%handle, 1) .

I guess that means we could improve this by teaching alias analysis that llvm.subfn.addr kills the handle.

[ChuanqiXu]

In D94834#2521983, @lxfind wrote:

In D94834#2521744, @ChuanqiXu wrote:

In D94834#2521729, @lxfind wrote:

And I think there are two cases about this example:
(1) The musttail call is dominated by @llvm.subfn.addr(%handle, 1) calls. Then %handle is destroyed at the musttail call, so it makes no sense that the musttail call returns %handle.
(2) The musttail call isn't dominated by @llvm.subfn.addr(%handle, 1) calls. Then the musttail call may returns %handle. And it doesn't matter since %handle won't be elided in this case.

In case (1), why is it that '%handle' must be destroyed at the musttail call?

In my mind, @llvm.subfn.addr(%handle, 1) means a call which would destroy %handle. And CoroElide Pass also treat the call as a lifetime end marker for %handle. CoroElide pass only elides coroutine handle who is guarded by @llvm.subfn.addr(%handle, 1) .

I guess that means we could improve this by teaching alias analysis that llvm.subfn.addr kills the handle.

I agree with that is a good solution to use AA if we can teach AA that llvm.subfn.addr(%handle, 1) kills %handle. But for my understanding to AA, it is not easy to do so. Maybe we need to refactor both AA and CoroElide to do that.
Let's go back to this bug. I think we can solve this problem by removing the static check for musttail call in removeTailCallAttribute.

llvm/lib/Transforms/Coroutines/CoroElide.cpp

 
 // Look for any tail calls referencing the coroutine frame and remove tail
 // attribute from them, since now coroutine frame resides on the stack and tail
 // call implies that the function does not references anything on the stack.
 static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) {
   Function &F = *Frame->getFunction();
   for (Instruction &I : instructions(F))
     if (auto *Call = dyn_cast<CallInst>(&I))
-      if (Call->isTailCall() && operandReferences(Call, Frame, AA)) {
-        // FIXME: If we ever hit this check. Evaluate whether it is more
-        // appropriate to retain musttail and allow the code to compile.
-        if (Call->isMustTailCall())
-          report_fatal_error("Call referring to the coroutine frame cannot be "
-                             "marked as musttail");
+      if (Call->isTailCall() && operandReferences(Call, Frame, AA))
         Call->setTailCall(false);
-      }
-}
 
 // Given a resume function @f.resume(%f.frame* %frame), returns the size
 // and expected alignment of %f.frame type.
 static std::pair<uint64_t, Align> getFrameLayout(Function *Resume) {
   // Prefer to pull information from the function attributes.
   auto Size = Resume->getParamDereferenceableBytes(0);
   auto Align = Resume->getParamAlign(0);
 
for (auto &It : DestroyAddr) {
     if (!ReferencedCoroBegins.count(It.first) &&
         !hasEscapePath(It.first, Terminators))
       ReferencedCoroBegins.insert(It.first);
   }
 
   // If size of the set is the same as total number of coro.begin, that means we
   // found a coro.free or coro.destroy referencing each coro.begin, so we can
   // perform heap elision.
-  return ReferencedCoroBegins.size() == CoroBegins.size();
+  if (ReferencedCoroBegins.size() != CoroBegins.size())
+    return false;
+
+  // If any call in the function is a musttail call, it usually won't work
+  // because we cannot drop the tailcall attribute, and a tail call will reuse
+  // the entire stack where we are going to put the new frame. In theory a more
+  // precise analysis can be done to check whether the new frame aliases with
+  // the call, however it's challenging to do so before the elision actually
+  // happened.
+  for (BasicBlock &BB : *F)
+    if (BB.getTerminatingMustTailCall())
+      return false;
+
+  return true;
 }
 
 void Lowerer::collectPostSplitCoroIds(Function *F) {
   CoroIds.clear();
   CoroSuspendSwitches.clear();
   for (auto &I : instructions(F)) {
     if (auto *CII = dyn_cast<CoroIdInst>(&I))
       if (CII->getInfo().isPostSplit())

llvm/test/Transforms/Coroutines/coro-elide-musttail.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; Only run with new pass manager since old pass manager's alias analysis isn't
+; powerful enough to tell that the tailcall's arguments don't alias the frame.
+;
+; RUN: opt < %s -passes='coro-elide' -S | FileCheck %s
+
+%"bar.Frame" = type { void (%"bar.Frame"*)*, void (%"bar.Frame"*)*, %"struct.coroutine<false, int>::promise_type", i1 }
+%"struct.coroutine<false, int>::promise_type" = type { i32 }
+%"foo.Frame" = type { void (%"foo.Frame"*)*, void (%"foo.Frame"*)*, %"struct.coroutine<true, int>::promise_type", i1, %"struct.coroutine<true, int>::promise_type::final_awaitable" }
+%"struct.coroutine<true, int>::promise_type" = type { i32 }
+%"struct.coroutine<true, int>::promise_type::final_awaitable" = type { i8 }
+@"bar.resumers" = private constant [3 x void (%"bar.Frame"*)*] [void (%"bar.Frame"*)* @"bar.resume", void (%"bar.Frame"*)* undef, void (%"bar.Frame"*)* undef]
+
+declare dso_local void @"bar"() align 2
+declare dso_local fastcc void @"bar.resume"(%"bar.Frame"*) align 2
+
+define internal fastcc void @foo.resume_musttail(%"foo.Frame"* %FramePtr) {
+; CHECK-LABEL: @foo.resume_musttail(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call token @llvm.coro.id
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i1 @llvm.coro.alloc(token [[TMP0]])
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call i8* @llvm.coro.begin(token [[TMP0]], i8* null)
+; CHECK-NEXT:    [[CALL34:%.*]] = call i8* undef()
+; CHECK-NEXT:    musttail call fastcc void undef(i8* [[CALL34]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = tail call token @llvm.coro.id(i32 16, i8* null, i8* bitcast (void ()* @"bar" to i8*), i8* bitcast ([3 x void (%"bar.Frame"*)*]* @"bar.resumers" to i8*))
+  %1 = tail call i1 @llvm.coro.alloc(token %0)
+  %2 = tail call i8* @llvm.coro.begin(token %0, i8* null)
+  call i8* @llvm.coro.subfn.addr(i8* %2, i8 1)
+  %call34 = call i8* undef()
+  musttail call fastcc void undef(i8* %call34)
+  ret void
+}
+
+define internal fastcc void @foo.resume_no_musttail(%"foo.Frame"* %FramePtr) {
+; CHECK-LABEL: @foo.resume_no_musttail(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca [24 x i8], align 8
+; CHECK-NEXT:    [[VFRAME:%.*]] = bitcast [24 x i8]* [[TMP0]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = call token @llvm.coro.id
+; CHECK-NEXT:    [[CALL34:%.*]] = call i8* undef()
+; CHECK-NEXT:    call fastcc void undef(i8* [[CALL34]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = tail call token @llvm.coro.id(i32 16, i8* null, i8* bitcast (void ()* @"bar" to i8*), i8* bitcast ([3 x void (%"bar.Frame"*)*]* @"bar.resumers" to i8*))
+  %1 = tail call i1 @llvm.coro.alloc(token %0)
+  %2 = tail call i8* @llvm.coro.begin(token %0, i8* null)
+  call i8* @llvm.coro.subfn.addr(i8* %2, i8 1)
+  %call34 = call i8* undef()
+  tail call fastcc void undef(i8* %call34)
+  ret void
+}
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #0
+
+; Function Attrs: argmemonly nounwind readonly
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
+
+; Function Attrs: nounwind
+declare i1 @llvm.coro.alloc(token) #2
+
+; Function Attrs: nounwind
+declare i8* @llvm.coro.begin(token, i8* writeonly) #2
+
+; Function Attrs: argmemonly nounwind readonly
+declare i8* @llvm.coro.subfn.addr(i8* nocapture readonly, i8) #1
+
+attributes #0 = { argmemonly nofree nosync nounwind willreturn }
+attributes #1 = { argmemonly nounwind readonly }
+attributes #2 = { nounwind }