[coroutines] CoroBegin from inner coroutines should be considered for spills

[GorNishanov]

Summary

If a coroutine outer calls another coroutine inner and the inner coroutine body is inlined into the outer, coro.begin from the inner coroutine should be considered for spilling if accessed across suspends.

Prior to this change, coroutine frame building code was not considering any coro.begins for spilling.
With this change, we only ignore coro.begin for the current coroutine, but, any coro.begins that were inlined into the current coroutine are eligible for spills.

Fixes PR34267

lib/Transforms/Coroutines/CoroFrame.cpp

 static bool materializable(Instruction &V) {
   return isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
          isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V);
 }
 
 // Check for structural coroutine intrinsics that should not be spilled into
 // the coroutine frame.
 static bool isCoroutineStructureIntrinsic(Instruction &I) {
-  return isa<CoroIdInst>(&I) || isa<CoroBeginInst>(&I) ||
-         isa<CoroSaveInst>(&I) || isa<CoroSuspendInst>(&I);
+  return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I) ||
+         isa<CoroSuspendInst>(&I);
 }
 
 // For every use of the value that is across suspend point, recreate that value
 // after a suspend point.
 static void rewriteMaterializableInstructions(IRBuilder<> &IRB,
                                               SpillInfo const &Spills) {
   BasicBlock *CurrentBlock = nullptr;
   Instruction *CurrentMaterialization = nullptr;
void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   for (Argument &A : F.args())
     for (User *U : A.users())
       if (Checker.isDefinitionAcrossSuspend(A, U))
         Spills.emplace_back(&A, U);
 
   for (Instruction &I : instructions(F)) {
     // Values returned from coroutine structure intrinsics should not be part
     // of the Coroutine Frame.
-    if (isCoroutineStructureIntrinsic(I))
+    if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
       continue;
     // The Coroutine Promise always included into coroutine frame, no need to
     // check for suspend crossing.
     if (Shape.PromiseAlloca == &I)
       continue;
 
     for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U)) {

test/Transforms/Coroutines/coro-spill-corobegin.ll

+; Check that we can spills coro.begin from an inlined inner coroutine.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+
+%g.Frame = type { void (%g.Frame*)*, void (%g.Frame*)*, i32, i1, i32 }
+
[e-mail] = private constant [3 x void (%g.Frame*)*] [void (%g.Frame*)* @g.dummy, void (%g.Frame*)* @g.dummy, void (%g.Frame*)* @g.dummy]
+
+declare void @g.dummy(%g.Frame*)
+
+define i8* @f() "coroutine.presplit"="1" {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+
+  %innerid = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* bitcast ([3 x void (%g.Frame*)*]* @g.resumers to i8*))
+  %innerhdl = call noalias nonnull i8* @llvm.coro.begin(token %innerid, i8* null)
+  %gframe = bitcast i8* %innerhdl to %g.Frame*
+
+  %tok = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %tok, label %suspend [i8 0, label %resume
+                                i8 1, label %cleanup]
+resume:
+  %gvar.addr = getelementptr inbounds %g.Frame, %g.Frame* %gframe, i32 0, i32 4
+  %gvar = load i32, i32* %gvar.addr
+  call void @print.i32(i32 %gvar)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; See if the i8* for coro.begin was added to f.Frame
+; CHECK-LABEL: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, i1, i1, i8* }
+
+; See if the g's coro.begin was spilled into the frame
+; CHECK-LABEL: @f(
+; CHECK: %innerid = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* bitcast ([3 x void (%g.Frame*)*]* @g.resumers to i8*))
+; CHECK: %innerhdl = call noalias nonnull i8* @llvm.coro.begin(token %innerid, i8* null)
+; CHECK: %[[spilladdr:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
+; CHECK: store i8* %innerhdl, i8** %[[spilladdr]]
+
+; See if the coro.begin was loaded from the frame
+; CHECK-LABEL: @f.resume(
+; CHECK: %[[innerhdlAddr:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %{{.+}}, i32 0, i32 4
+; CHECK: %[[innerhdl:.+]] = load i8*, i8** %[[innerhdlAddr]]
+; CHECK: %[[gframe:.+]] = bitcast i8* %[[innerhdl]] to %g.Frame*
+; CHECK: %[[gvarAddr:.+]] = getelementptr inbounds %g.Frame, %g.Frame* %[[gframe]], i32 0, i32 4
+; CHECK: %[[gvar:.+]] = load i32, i32* %[[gvarAddr]]
+; CHECK: call void @print.i32(i32 %[[gvar]])
+
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare void @print.i32(i32)
+declare void @free(i8*)