[Coroutines] [Retcon] Replace CoroBegin with FramePtr after

[ChuanqiXu]

Summary

This tries to fix issue https://github.com/llvm/llvm-project/issues/54523 by repalcing CoroBegin with FramePtr later in splitRetconCoroutine.

[nikic]

Thanks! I like this fix more than D122375, but I'm not familiar with the details of this code.

[rjmccall]

Hmm. Doing this after doing the clones means you're only actually doing the replacement in the ramp function, which I think means it'll explode in the clones if there are uses remaining. On some level, that might be okay, because I don't think the Swift frontend ever actually uses the result of coro.begin in retcon functions except to pass it to coro.end, and retcon lowering of coro.end just ignores the argument. On the other hand, we shouldn't leave something around that'll blow up if the code pattern ever changes.

I think you can just recognize this case (when Shape.FramePtr is an argument) and remove the mapping for that argument from VMap before cloning.

[ChuanqiXu]

In D122383#3413345, @rjmccall wrote:

Hmm. Doing this after doing the clones means you're only actually doing the replacement in the ramp function, which I think means it'll explode in the clones if there are uses remaining. On some level, that might be okay, because I don't think the Swift frontend ever actually uses the result of coro.begin in retcon functions except to pass it to coro.end, and retcon lowering of coro.end just ignores the argument. On the other hand, we shouldn't leave something around that'll blow up if the code pattern ever changes.

I think you can just recognize this case (when Shape.FramePtr is an argument) and remove the mapping for that argument from VMap before cloning.

I found another idea. Here I found the second argument of coro.begin is meant to point to the coroutine frame. (https://llvm.org/docs/Coroutines.html#llvm-coro-begin-intrinsic) I found it says that it is ignored by returned-continuation coroutines. But I think it is for the frontend users. We don't invalidate it if we set it in the middle end. And if we could do it, then the codes in CoroCleanup would replace coro.begin with the second argument. Do you think this is available? If yes, I think it might be better than handle the argument case specially.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1777

How about setting the memory argument for CoroBegin here:

Shape.CoroBegin->setOperand(1, RawFramePtr);

Then the CoroBegin would be replaced to RawFramePtr later in https://github.com/llvm/llvm-project/blob/8568d681ae14fb8c7769c80e0196a16934647413/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp#L67-L69

[rjmccall]

Does that actually fix anything? The problem is that the use-list of the frame pointer is being destroyed in the clones if it happens to be an argument of the ramp function. We have to stop that from happening. I don't see a solution to that other than either changing VMap or introducing an intermediate value so that it isn't just a use of an argument, and the former seems much easier.

[ChuanqiXu]

In D122383#3414748, @rjmccall wrote:

Does that actually fix anything? The problem is that the use-list of the frame pointer is being destroyed in the clones if it happens to be an argument of the ramp function. We have to stop that from happening. I don't see a solution to that other than either changing VMap or introducing an intermediate value so that it isn't just a use of an argument, and the former seems much easier.

I think the two strategies (Move the replacement of CoroBegin or set the memory argument for CoroBegin) aim to the solution introducing an intermediate value so that it isn't just a use of an argument, and the intermediate value here is the llvm.coro.begin itself. I think this one if more natural to changing VMap since it might be a noop for other coroutine ABIs.

[rjmccall]

It seems very likely to me that you could contrive a situation under C++ coroutines where the frame pointer is a parameter to the ramp function by the time you run coro splitting, at least under optimization. For example, https://godbolt.org/z/cGnjTbzrY, which is a reasonable implementation if you need to guarantee non-allocation and intend to assert that the size fits in the buffer. This isn't an oddity of retcon lowering, it's a general problem that emerges from the special way that the frame pointer value is assumed to be live throughout the function, which requires it to be protected against this undef'ing of arguments which would otherwise be reasonable.

Thinking about it further, you'll probably have to introduce an intermediate value to do that protection, since the cloner will not be happy if you don't map all the arguments.

[nikic]

In D122383#3415713, @rjmccall wrote:

It seems very likely to me that you could contrive a situation under C++ coroutines where the frame pointer is a parameter to the ramp function by the time you run coro splitting, at least under optimization. For example, https://godbolt.org/z/cGnjTbzrY, which is a reasonable implementation if you need to guarantee non-allocation and intend to assert that the size fits in the buffer. This isn't an oddity of retcon lowering, it's a general problem that emerges from the special way that the frame pointer value is assumed to be live throughout the function, which requires it to be protected against this undef'ing of arguments which would otherwise be reasonable.

Thinking about it further, you'll probably have to introduce an intermediate value to do that protection, since the cloner will not be happy if you don't map all the arguments.

@rjmccall So would D122375 be about what you have in mind here?

[rjmccall]

Yeah. I think you could probably get away with just doing to the FramePtr argument, but it's probably safer to do it unconditionally.

[ChuanqiXu]

In D122383#3415713, @rjmccall wrote:

It seems very likely to me that you could contrive a situation under C++ coroutines where the frame pointer is a parameter to the ramp function by the time you run coro splitting, at least under optimization. For example, https://godbolt.org/z/cGnjTbzrY, which is a reasonable implementation if you need to guarantee non-allocation and intend to assert that the size fits in the buffer. This isn't an oddity of retcon lowering, it's a general problem that emerges from the special way that the frame pointer value is assumed to be live throughout the function, which requires it to be protected against this undef'ing of arguments which would otherwise be reasonable.

Thinking about it further, you'll probably have to introduce an intermediate value to do that protection, since the cloner will not be happy if you don't map all the arguments.

Yeah, your words make sense. Let's abandon this one.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

if (Shape.RetconLowering.IsFrameInlineInStorage) {
     RawFramePtr =
       Builder.CreateBitCast(RawFramePtr, Shape.CoroBegin->getType());
 
     // Stash the allocated frame pointer in the continuation storage.
     auto Dest = Builder.CreateBitCast(Id->getStorage(),
                                       RawFramePtr->getType()->getPointerTo());
     Builder.CreateStore(RawFramePtr, Dest);
   }
 
-  // Map all uses of llvm.coro.begin to the allocated frame pointer.
-  {
-    // Make sure we don't invalidate Shape.FramePtr.
-    TrackingVH<Value> Handle(Shape.FramePtr);
-    Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);
-    Shape.FramePtr = Handle.getValPtr();
-  }
-
   // Create a unique return block.
   BasicBlock *ReturnBB = nullptr;
   SmallVector<PHINode *, 4> ReturnPHIs;
 
   // Create all the functions in order after the main function.
   auto NextF = std::next(F.getIterator());
 
   // Create a continuation function for each of the suspend points.
static void splitRetconCoroutine(Function &F, coro::Shape &Shape,
 
   assert(Clones.size() == Shape.CoroSuspends.size());
   for (size_t i = 0, e = Shape.CoroSuspends.size(); i != e; ++i) {
     auto Suspend = Shape.CoroSuspends[i];
     auto Clone = Clones[i];
 
     CoroCloner(F, "resume." + Twine(i), Shape, Clone, Suspend).create();
   }
+
+  // Map all uses of llvm.coro.begin to the allocated frame pointer.
+  {
+    // Make sure we don't invalidate Shape.FramePtr.
+    TrackingVH<Value> Handle(Shape.FramePtr);
+    Shape.CoroBegin->replaceAllUsesWith(RawFramePtr);
+    Shape.FramePtr = Handle.getValPtr();
+  }
 }
 
 namespace {
   class PrettyStackTraceFunction : public PrettyStackTraceEntry {
     Function &F;
   public:
     PrettyStackTraceFunction(Function &F) : F(F) {}
     void print(raw_ostream &OS) const override {

llvm/test/Transforms/Coroutines/coro-retcon-alloca-opaque-ptr.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+
+; RUN: opt < %s -enable-coroutines -passes='default<O2>' -opaque-pointers=1 -S | FileCheck %s
+
+target datalayout = "p:64:64:64"
+
+declare {i8*, i8*, i32} @prototype_f(i8*, i1)
+define {i8*, i8*, i32} @f(i8* %buffer, i32 %n) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  coro.return:
+; CHECK-NEXT:    [[N_VAL_SPILL_ADDR:%.*]] = getelementptr inbounds [[F_FRAME:%.*]], ptr [[BUFFER:%.*]], i64 0, i32 1
+; CHECK-NEXT:    store i32 [[N:%.*]], ptr [[N_VAL_SPILL_ADDR]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call ptr @allocate(i32 [[N]])
+; CHECK-NEXT:    store ptr [[TMP0]], ptr [[BUFFER]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue { ptr, ptr, i32 } { ptr @f.resume.0, ptr undef, i32 undef }, ptr [[TMP0]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { ptr, ptr, i32 } [[TMP1]], i32 [[N]], 2
+; CHECK-NEXT:    ret { ptr, ptr, i32 } [[TMP2]]
+;
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, i8* %buffer, i8* bitcast ({i8*, i8*, i32} (i8*, i1)* @prototype_f to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  br label %loop
+
+loop:
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  %alloca = call token @llvm.coro.alloca.alloc.i32(i32 %n.val, i32 8)
+  %ptr = call i8* @llvm.coro.alloca.get(token %alloca)
+  %unwind = call i1 (...) @llvm.coro.suspend.retcon.i1(i8* %ptr, i32 %n.val)
+  call void @llvm.coro.alloca.free(token %alloca)
+  br i1 %unwind, label %cleanup, label %resume
+
+resume:
+  %inc = add i32 %n.val, 1
+  br label %loop
+
+cleanup:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
+
+
+declare {i8*, i32} @prototype_g(i8*, i1)
+define {i8*, i32} @g(i8* %buffer, i32 %n) {
+; CHECK-LABEL: @g(
+; CHECK-NEXT:  coro.return:
+; CHECK-NEXT:    store i32 [[N:%.*]], ptr [[BUFFER:%.*]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = zext i32 [[N]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca i8, i64 [[TMP0]], align 8
+; CHECK-NEXT:    tail call void @use(ptr nonnull [[TMP1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { ptr, i32 } { ptr @g.resume.0, i32 undef }, i32 [[N]], 1
+; CHECK-NEXT:    ret { ptr, i32 } [[TMP2]]
+;
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, i8* %buffer, i8* bitcast ({i8*, i32} (i8*, i1)* @prototype_g to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  br label %loop
+
+loop:
+  %n.val = phi i32 [ %n, %entry ], [ %inc, %resume ]
+  %alloca = call token @llvm.coro.alloca.alloc.i32(i32 %n.val, i32 8)
+  %ptr = call i8* @llvm.coro.alloca.get(token %alloca)
+  call void @use(i8* %ptr)
+  call void @llvm.coro.alloca.free(token %alloca)
+  %unwind = call i1 (...) @llvm.coro.suspend.retcon.i1(i32 %n.val)
+  br i1 %unwind, label %cleanup, label %resume
+
+resume:
+  %inc = add i32 %n.val, 1
+  br label %loop
+
+cleanup:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
+
+declare token @llvm.coro.id.retcon(i32, i32, i8*, i8*, i8*, i8*)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.suspend.retcon.i1(...)
+declare void @llvm.coro.suspend.retcon.isVoid(...)
+declare i1 @llvm.coro.end(i8*, i1)
+declare i8* @llvm.coro.prepare.retcon(i8*)
+declare token @llvm.coro.alloca.alloc.i32(i32, i32)
+declare i8* @llvm.coro.alloca.get(token)
+declare void @llvm.coro.alloca.free(token)
+
+declare noalias i8* @allocate(i32 %size)
+declare void @deallocate(i8* %ptr)
+
+declare void @print(i32)
+declare void @use(i8*)

llvm/test/Transforms/Coroutines/coro-retcon-resume-values.ll

 ; CHECK-NEXT:    [[INPUT_RELOAD_ADDR_I:%.*]] = getelementptr inbounds [[F_FRAME]], %f.Frame* [[FRAMEPTR_I1]], i64 0, i32 2
 ; CHECK-NEXT:    [[INPUT_RELOAD_I:%.*]] = load i32, i32* [[INPUT_RELOAD_ADDR_I]], align 4, !noalias !3
 ; CHECK-NEXT:    [[N_VAL3_RELOAD_ADDR_I:%.*]] = getelementptr inbounds [[F_FRAME]], %f.Frame* [[FRAMEPTR_I1]], i64 0, i32 1
 ; CHECK-NEXT:    [[N_VAL3_RELOAD_I:%.*]] = load i32, i32* [[N_VAL3_RELOAD_ADDR_I]], align 4, !noalias !3
 ; CHECK-NEXT:    [[SUM8_I:%.*]] = add i32 [[N_VAL3_RELOAD_I]], [[INPUT_RELOAD_I]]
 ; CHECK-NEXT:    store i32 [[SUM8_I]], i32* [[N_VAL3_RELOAD_ADDR_I]], align 4, !noalias !3
 ; CHECK-NEXT:    store i32 4, i32* [[INPUT_RELOAD_ADDR_I]], align 4, !noalias !3
 ; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata [[META6:![0-9]+]])
-; CHECK-NEXT:    [[FRAMEPTR_I4:%.*]] = load %f.Frame*, %f.Frame** [[TMP2]], align 8, !alias.scope !6
-; CHECK-NEXT:    [[INPUT_RELOAD_ADDR13_I:%.*]] = getelementptr inbounds [[F_FRAME]], %f.Frame* [[FRAMEPTR_I4]], i64 0, i32 2
+; CHECK-NEXT:    [[FRAMEPTR_I2:%.*]] = load %f.Frame*, %f.Frame** [[TMP2]], align 8, !alias.scope !6
+; CHECK-NEXT:    [[VFRAME_I:%.*]] = bitcast %f.Frame* [[FRAMEPTR_I2]] to i8*
+; CHECK-NEXT:    [[INPUT_RELOAD_ADDR13_I:%.*]] = getelementptr inbounds [[F_FRAME]], %f.Frame* [[FRAMEPTR_I2]], i64 0, i32 2
 ; CHECK-NEXT:    [[INPUT_RELOAD14_I:%.*]] = load i32, i32* [[INPUT_RELOAD_ADDR13_I]], align 4, !noalias !6
-; CHECK-NEXT:    [[N_VAL3_RELOAD_ADDR11_I:%.*]] = getelementptr inbounds [[F_FRAME]], %f.Frame* [[FRAMEPTR_I4]], i64 0, i32 1
+; CHECK-NEXT:    [[N_VAL3_RELOAD_ADDR11_I:%.*]] = getelementptr inbounds [[F_FRAME]], %f.Frame* [[FRAMEPTR_I2]], i64 0, i32 1
 ; CHECK-NEXT:    [[N_VAL3_RELOAD12_I:%.*]] = load i32, i32* [[N_VAL3_RELOAD_ADDR11_I]], align 4, !noalias !6
 ; CHECK-NEXT:    [[SUM7_I:%.*]] = add i32 [[N_VAL3_RELOAD12_I]], [[INPUT_RELOAD14_I]]
 ; CHECK-NEXT:    call void @print(i32 [[SUM7_I]]), !noalias !6
-; CHECK-NEXT:    [[TMP3:%.*]] = bitcast %f.Frame* [[FRAMEPTR_I4]] to i8*
-; CHECK-NEXT:    call void @deallocate(i8* [[TMP3]]), !noalias !6
+; CHECK-NEXT:    call void @deallocate(i8* [[VFRAME_I]]), !noalias !6
 ; CHECK-NEXT:    ret i32 0
 ;
 entry:
   %0 = alloca [8 x i8], align 4
   %buffer = bitcast [8 x i8]* %0 to i8*
   %prepare = call i8* @llvm.coro.prepare.retcon(i8* bitcast (i8* (i8*, i32)* @f to i8*))
   %f = bitcast i8* %prepare to i8* (i8*, i32)*
   %cont0 = call i8* %f(i8* %buffer, i32 1)