[Coroutines] Optimize the lifespan of temporary co_await object

[lxfind]

Summary

If we ever assign co_await to a temporary variable, such as foo(co_await expr),
we generate AST that looks like this: MaterializedTemporaryExpr(CoawaitExpr(...)).
MaterializedTemporaryExpr would emit an intrinsics that marks the lifetime start of the
temporary storage. However such temporary storage will not be used until co_await is ready
to write the result. Marking the lifetime start way too early causes extra storage to be
put in the coroutine frame instead of the stack.
As you can see from https://godbolt.org/z/zVx_eB, the frame generated for get_big_object2 is 12K, which contains a big_object object unnecessarily.
After this patch, the frame size for get_big_object2 is now only 8K. There are still room for improvements, in particular, GCC has a 4K frame for this function. But that's a separate problem and not addressed in this patch.

The basic idea of this patch is during CoroSplit, look for every local variable in the coroutine created through AllocaInst, identify all the lifetime start/end markers and the use of the variables, and sink the lifetime.start maker to the places as close to the first-ever use as possible.

[junparser]

Rather than doing it here, can we build await_resume call expression with MaterializedTemporaryExpr when expand the coawait expression. That's how gcc does.

[lxfind]

In D82314#2107910, @junparser wrote:

Rather than doing it here, can we build await_resume call expression with MaterializedTemporaryExpr when expand the coawait expression. That's how gcc does.

There doesn't appear to be a way to do that in Clang. It goes from the AST to IR directly, and there needs to be a MaterializedTemporaryExpr to wrap the result of co_await. Could you elaborate on how this might be done in Clang?

[rsmith]

In D82314#2109437, @lxfind wrote:

In D82314#2107910, @junparser wrote:

Rather than doing it here, can we build await_resume call expression with MaterializedTemporaryExpr when expand the coawait expression. That's how gcc does.

There doesn't appear to be a way to do that in Clang. It goes from the AST to IR directly, and there needs to be a MaterializedTemporaryExpr to wrap the result of co_await. Could you elaborate on how this might be done in Clang?

For a call such as:

coro f() { awaitable a; (co_await a).g(); }

we produce an AST like:

|   |   `-CXXMemberCallExpr <col:25, col:40> 'void'
|   |     `-MemberExpr <col:25, col:38> '<bound member function type>' .g 0x55d38948ca98
|   |       `-MaterializeTemporaryExpr <col:25, col:36> 'huge' xvalue
|   |         `-ParenExpr <col:25, col:36> 'huge'
|   |           `-CoawaitExpr <col:26, col:35> 'huge'
|   |             |-DeclRefExpr <col:35> 'awaitable' lvalue Var 0x55d38948d4c8 'a' 'awaitable'
|   |             |-CXXMemberCallExpr <col:35> 'bool'
|   |             | `-MemberExpr <col:35> '<bound member function type>' .await_ready 0x55d38948cee8
|   |             |   `-OpaqueValueExpr <col:35> 'awaitable' lvalue
|   |             |     `-DeclRefExpr <col:35> 'awaitable' lvalue Var 0x55d38948d4c8 'a' 'awaitable'
|   |             |-CXXMemberCallExpr <col:35> 'void'
|   |             | |-MemberExpr <col:35> '<bound member function type>' .await_suspend 0x55d38948d298
|   |             | | `-OpaqueValueExpr <col:35> 'awaitable' lvalue
|   |             | |   `-DeclRefExpr <col:35> 'awaitable' lvalue Var 0x55d38948d4c8 'a' 'awaitable'
|   |             | `-CXXConstructExpr <col:35> 'std::coroutine_handle<>':'std::experimental::coroutines_v1::coroutine_handle<void>' 'void (std::experimental::coroutines_v1::coroutine_handle<void> &&) noexcept'
|   |             |   `-ImplicitCastExpr <col:35> 'std::experimental::coroutines_v1::coroutine_handle<void>' xvalue <DerivedToBase (coroutine_handle)>
|   |             |     `-MaterializeTemporaryExpr <col:35> 'std::experimental::coroutines_v1::coroutine_handle<coro::promise_type>' xvalue
|   |             |       `-CallExpr <col:35> 'std::experimental::coroutines_v1::coroutine_handle<coro::promise_type>'
|   |             |         |-ImplicitCastExpr <col:35> 'std::experimental::coroutines_v1::coroutine_handle<coro::promise_type> (*)(void *) noexcept' <FunctionToPointerDecay>
|   |             |         | `-DeclRefExpr <col:35> 'std::experimental::coroutines_v1::coroutine_handle<coro::promise_type> (void *) noexcept' lvalue CXXMethod 0x55d38948adb0 'from_address' 'std::experimental::coroutines_v1::coroutine_handle<coro::promise_type> (void *) noexcept'
|   |             |         `-CallExpr <col:35> 'void *'
|   |             |           `-ImplicitCastExpr <col:35> 'void *(*)() noexcept' <FunctionToPointerDecay>
|   |             |             `-DeclRefExpr <col:35> 'void *() noexcept' lvalue Function 0x55d38948ef38 '__builtin_coro_frame' 'void *() noexcept'
|   |             `-CXXBindTemporaryExpr <col:35> 'huge' (CXXTemporary 0x55d38948ffb8)
|   |               `-CXXMemberCallExpr <col:35> 'huge'
|   |                 `-MemberExpr <col:35> '<bound member function type>' .await_resume 0x55d38948cff8
|   |                   `-OpaqueValueExpr <col:35> 'awaitable' lvalue
|   |                     `-DeclRefExpr <col:35> 'awaitable' lvalue Var 0x55d38948d4c8 'a' 'awaitable'

See https://godbolt.org/z/hVn2u9.

I think the suggestion is to move the MaterializeTemporaryExpr from being wrapped around the CoawaitExpr to being wrapped around the .await_resume call within it.

Unfortunately, that's not a correct change. The language rules require us to delay materializing the temporary until we see how it is used. For example, in a case such as

g(co_await a);

... no temporary is materialized at all, and the await_resume call instead directly initializes the parameter slot for the function.

---

It seems to me that the same issue (of making large objects unnecessarily live across suspend points) can arise for other similar cases too. For example, consider:

huge(co_await a).g(); // suppose `co_await a` returns something small

Here again we will start the lifetime of the huge temporary before the suspend point, and only initialize it after the resume. Your change won't help here, because it's not the lifetime markers of the value produced by co_await that are the problem.

As a result of the above, I think this is the wrong level at which to perform this optimization. Instead, I think you should consider whether we can move lifetime start markers later (and end markers earlier, for unescaped locals) as part of the coroutine splitting pass.

[lxfind]

@rsmith Thanks. That's a good point. Do you know if there already exists optimization passes in LLVM that attempts to shrink the range of lifetime intrinsics? If so, I am curious why that does not help in this case. Or is it generally unsafe to move the lifetime intrinsics, and we could only do it here with specific context knowledge about coroutines.

[rsmith]

In D82314#2109728, @lxfind wrote:

@rsmith Thanks. That's a good point. Do you know if there already exists optimization passes in LLVM that attempts to shrink the range of lifetime intrinsics? If so, I am curious why that does not help in this case. Or is it generally unsafe to move the lifetime intrinsics, and we could only do it here with specific context knowledge about coroutines.

I don't know for sure, but I would expect someone to have implemented such a pass already. Moving a lifetime start intrinsic later, past instructions that can't possibly reference the object in question, seems like it should always be safe and (presumably) should always be a good thing to do, and similarly for moving lifetime end markers earlier. It could be that such a pass exists but it is run too late in the pass pipeline, so the coroutine split pass doesn't get to take advantage of it.

[junparser]

In D82314#2109893, @rsmith wrote:

In D82314#2109728, @lxfind wrote:

@rsmith Thanks. That's a good point. Do you know if there already exists optimization passes in LLVM that attempts to shrink the range of lifetime intrinsics? If so, I am curious why that does not help in this case. Or is it generally unsafe to move the lifetime intrinsics, and we could only do it here with specific context knowledge about coroutines.

I don't know for sure, but I would expect someone to have implemented such a pass already. Moving a lifetime start intrinsic later, past instructions that can't possibly reference the object in question, seems like it should always be safe and (presumably) should always be a good thing to do, and similarly for moving lifetime end markers earlier. It could be that such a pass exists but it is run too late in the pass pipeline, so the coroutine split pass doesn't get to take advantage of it.

@lxfind, Also lifetime marker of variable are much complex because of the existing of exceptional path(multiple lifetime start & multiple lifetime end) , so it is hard to optimize such cases.

[junparser]

In D82314#2109437, @lxfind wrote:

In D82314#2107910, @junparser wrote:

Rather than doing it here, can we build await_resume call expression with MaterializedTemporaryExpr when expand the coawait expression. That's how gcc does.

There doesn't appear to be a way to do that in Clang. It goes from the AST to IR directly, and there needs to be a MaterializedTemporaryExpr to wrap the result of co_await. Could you elaborate on how this might be done in Clang?

For now, we only wrap coawait expression with MaterializedTemporaryExpr when the kind of result is VK_RValue, We can wrap await_resume call instead in such case when build coawait expression. so in emitSuspendExpression, we can directly emit await_call expression with MaterializedTemporaryExpr.

I think this should work, although i'm not so sure.

[lxfind]

Tackle this problem inside CoroSplit as an optimization. Instead of only handling one particular case, we now look at every local variable in the coroutine, and sink their lifetime start markers when possible. This will bring in more benefits than doing so during IR emit. Confirmed that it works.

[lxfind]

Address test failures

[junparser]

@lxfind, Thank you! And could you please add some testcases?

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1286

It is possible to handle multiple cast instructions as long as they are only used by lifetime marker intrinsic.

[lxfind]

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1286

It's certainly possible. I didn't do it here because a reasonable compiler frontend should never emit multiple cast instructions for the same variable in order to mark lifetime. If there are, they must be used for something else.

[lxfind]

Actually it seems pretty easy to handle the case of multiple BitCastInst, so did it here. Also added a test.

[lxfind]

A few unintended changes

[lxfind]

Remove unused variable

[lewissbaker]

There are still room for improvements, in particular, GCC has a 4K frame for this function.

I think GCC having a smaller coroutine frame is probably because it does not yet implement the allocation-elision optimisation which inlines the nested coroutine frame (which is 4K) into the parent coroutine frame.
I have not looked yet, but I suspect that you'll see within the get_big_object2() code it will perform another heap allocation for the get_big_object() frame.

Until we have more general support for async RVO, I think 8K is probably as good as we're going to be able to get (4K for this coroutine's promise and 4K for child coroutine-frame and its promise).

[lxfind]

In D82314#2117543, @lewissbaker wrote:

There are still room for improvements, in particular, GCC has a 4K frame for this function.

I think GCC having a smaller coroutine frame is probably because it does not yet implement the allocation-elision optimisation which inlines the nested coroutine frame (which is 4K) into the parent coroutine frame.
I have not looked yet, but I suspect that you'll see within the get_big_object2() code it will perform another heap allocation for the get_big_object() frame.

Until we have more general support for async RVO, I think 8K is probably as good as we're going to be able to get (4K for this coroutine's promise and 4K for child coroutine-frame and its promise).

Yes your observation is correct. It's due to the nesting of the child frame.

[junparser]

LGTM， Thank you！

[lxfind]

Rebase before landing

[junparser]

@lxfind This patch causes some mismatch when variable is used in both resume and destroy function. Besides, we should move this patch and the check in buildCoroutineFrame.

[junparser]

In D82314#2124661, @junparser wrote:

@lxfind This patch causes some mismatch when variable is used in both resume and destroy function. Besides, we should move this patch and the check in buildCoroutineFrame.

@lxfind, Would you try to fix this？ If you do not have time, then I'll try do this. Thanks

[lxfind]

In D82314#2124662, @junparser wrote:

In D82314#2124661, @junparser wrote:

@lxfind This patch causes some mismatch when variable is used in both resume and destroy function. Besides, we should move this patch and the check in buildCoroutineFrame.

@lxfind, Would you try to fix this？ If you do not have time, then I'll try do this. Thanks

Could you please help take a look, if you have a local repro? Thanks!

[junparser]

In D82314#2125713, @lxfind wrote:

In D82314#2124662, @junparser wrote:

In D82314#2124661, @junparser wrote:

@lxfind This patch causes some mismatch when variable is used in both resume and destroy function. Besides, we should move this patch and the check in buildCoroutineFrame.

@lxfind, Would you try to fix this？ If you do not have time, then I'll try do this. Thanks

Could you please help take a look, if you have a local repro? Thanks!

of course.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

 #include <iterator>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-split"
 
 namespace {
 
-/// A little helper class for building 
+/// A little helper class for building
 class CoroCloner {
 public:
   enum class Kind {
     /// The shared resume function for a switch lowering.
     SwitchResume,
 
     /// The shared unwind function for a switch lowering.
     SwitchUnwind,

 void CoroCloner::replaceSwiftErrorOps() {
   ::replaceSwiftErrorOps(*NewF, Shape, &VMap);
 }
 
 void CoroCloner::replaceEntryBlock() {
   // In the original function, the AllocaSpillBlock is a block immediately
   // following the allocation of the frame object which defines GEPs for
   // all the allocas that have been moved into the frame, and it ends by
-  // branching to the original beginning of the coroutine.  Make this 
+  // branching to the original beginning of the coroutine.  Make this
   // the entry block of the cloned function.
   auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);
   auto *OldEntry = &NewF->getEntryBlock();
   Entry->setName("entry" + Suffix);
   Entry->moveBefore(OldEntry);
   Entry->getTerminator()->eraseFromParent();
 
   // Clear all predecessors of the new entry block.  There should be
if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
       continue;
     }
     if (++I == N)
       break;
   }
   S.resize(N);
 }
 
+/// For every local variable that has lifetime intrinsics markers, we sink
+/// their lifetime.start marker to the places where the variable is being
+/// used for the first time. Doing so minimizes the lifetime of each variable,
+/// hence minimizing the amount of data we end up putting on the frame.
+static void sinkLifetimeStartMarkers(Function &F) {
+  DominatorTree Dom(F);
+  for (Instruction &I : instructions(F)) {
+    // We look for this particular pattern:
+    //   %tmpX = alloca %.., align ...
+    //   %0 = bitcast %...* %tmpX to i8*
+    //   call void @llvm.lifetime.start.p0i8(i64 ..., i8* nonnull %0) #2
+    if (!isa<AllocaInst>(&I))
+      continue;
+    BitCastInst *CastInst = nullptr;
+    // There can be multiple lifetime start markers for the same variable.
+    SmallPtrSet<IntrinsicInst *, 1> LifetimeStartInsts;
+    // SinkBarriers stores all instructions that use this local variable.
+    // When sinking the lifetime start intrinsics, we can never sink past
+    // these barriers.
+    SmallPtrSet<Instruction *, 4> SinkBarriers;
+    bool Valid = true;
+    auto addSinkBarrier = [&](Instruction *I) {
+      // When adding a new barrier to SinkBarriers, we maintain the case
+      // that no instruction in SinkBarriers dominates another instruction.
+      bool FoundDom = false;
+      SmallPtrSet<Instruction *, 1> ToRemove;
+      for (auto *S : SinkBarriers) {
+        if (Dom.dominates(S, I)) {
+          FoundDom = true;
+          break;
+        } else if (Dom.dominates(I, S)) {
+          ToRemove.insert(S);
+        }
+      }
+      if (!FoundDom) {
+        SinkBarriers.insert(I);
+        for (auto *R : ToRemove) {
+          SinkBarriers.erase(R);
+        }
+      }
+    };
+    for (User *U : I.users()) {
+      if (!isa<BitCastInst>(U))
+        continue;
+      if (CastInst) {
+        // If we have multiple cast instructions for the alloca, don't
+        // deal with it beause it's too complex.
+        Valid = false;
+        break;
+      }
+      CastInst = cast<BitCastInst>(U);
+      for (User *CU : CastInst->users()) {
+        // If we see any user of CastInst that's not lifetime start/end
+        // intrinsics, give up because it's too complex.
+        if (auto *CUI = dyn_cast<IntrinsicInst>(CU)) {
+          if (CUI->getIntrinsicID() == Intrinsic::lifetime_start)
+            LifetimeStartInsts.insert(CUI);
+          else if (CUI->getIntrinsicID() == Intrinsic::lifetime_end)
+            addSinkBarrier(CUI);
+          else
+            Valid = false;
+        } else {
+          Valid = false;
+        }
+      }
+    }
+    if (!Valid || LifetimeStartInsts.empty())
+      continue;
+
+    for (User *U : I.users()) {
+      if (U == CastInst)
+        continue;
+      // Every user of the variable is also a sink barrier.
+      addSinkBarrier(cast<Instruction>(U));
+    }
+
+    // For each sink barrier, we insert a lifetime start marker right
+    // before it.
+    const auto *LifetimeStartInst = *LifetimeStartInsts.begin();
+    for (auto *S : SinkBarriers) {
+      if (auto *IS = dyn_cast<IntrinsicInst>(S)) {
+        if (IS->getIntrinsicID() == Intrinsic::lifetime_end) {
+          // If we have a lifetime end marker in SinkBarriers, meaning it's
+          // not dominated by any other users, we can safely delete it.
+          IS->eraseFromParent();
+          continue;
+        }
+      }
+      LifetimeStartInst->clone()->insertBefore(S);
+    }
+    // All the old markers are no longer necessary.
+    for (auto *S : LifetimeStartInsts) {
+      S->eraseFromParent();
+    }
+    // Put the cast instruction always right after variable declaration
+    // to be safe.
+    CastInst->moveAfter(&I);
+  }
+}
+
 static void splitSwitchCoroutine(Function &F, coro::Shape &Shape,
                                  SmallVectorImpl<Function *> &Clones) {
   assert(Shape.ABI == coro::ABI::Switch);
 
   createResumeEntryBlock(F, Shape);
   auto ResumeClone = createClone(F, ".resume", Shape,
                                  CoroCloner::Kind::SwitchResume);
   auto DestroyClone = createClone(F, ".destroy", Shape,
static coro::Shape splitCoroutine(Function &F,
   // up by uses in unreachable blocks, so remove them as a first pass.
   removeUnreachableBlocks(F);
 
   coro::Shape Shape(F);
   if (!Shape.CoroBegin)
     return Shape;
 
   simplifySuspendPoints(Shape);
+  sinkLifetimeStartMarkers(F);
   buildCoroutineFrame(F, Shape);
   replaceFrameSize(Shape);
 
   // If there are no suspend points, no split required, just remove
   // the allocation and deallocation blocks, they are not needed.
   if (Shape.CoroSuspends.empty()) {
     handleNoSuspendCoroutine(Shape);
   } else {