[CoroSplit] Handle argument being the frame pointer

[nikic]

Summary

If the frame pointer is an argument of the original pointer (which happens with opaque pointers), then we currently first replace the argument with undef, which will prevent later replacement of the old frame pointer with the new one.

Fix this by replacing arguments with some dummy instructions first, and then replacing those with undef later. This gives us a chance to replace the frame pointer before it becomes undef.

This is a big ugly, but I couldn't think of a better solution.

Fixes https://github.com/llvm/llvm-project/issues/54523.

[ChuanqiXu]

I am not familiar with retcon coroutines and opaque pointers. But I am surprised for the sentence If the frame pointer is an argument of the original pointer (which happens with opaque pointers). May I ask where we would make it possible to make the FramePtr to be the argument of the original function with opaque pointers?

[nikic]

In D122375#3404809, @ChuanqiXu wrote:

I am not familiar with retcon coroutines and opaque pointers. But I am surprised for the sentence If the frame pointer is an argument of the original pointer (which happens with opaque pointers). May I ask where we would make it possible to make the FramePtr to be the argument of the original function with opaque pointers?

Without opaque pointers, the frame pointer is a bitcast of the argument. With opaque pointers, the bitcast is not present, and it's directly the argument. So without opaque pointers, the old frame pointer first becomes bitcast undef and then gets replaced with the new frame pointer later. With opaque pointers it becomes undef, which cannot be correctly replaced. This patch tries to artificially add a dummy bitcast to make the replacement work correctly.

[ChuanqiXu]

In D122375#3404813, @nikic wrote:

In D122375#3404809, @ChuanqiXu wrote:

I am not familiar with retcon coroutines and opaque pointers. But I am surprised for the sentence If the frame pointer is an argument of the original pointer (which happens with opaque pointers). May I ask where we would make it possible to make the FramePtr to be the argument of the original function with opaque pointers?

Without opaque pointers, the frame pointer is a bitcast of the argument. With opaque pointers, the bitcast is not present, and it's directly the argument. So without opaque pointers, the old frame pointer first becomes bitcast undef and then gets replaced with the new frame pointer later. With opaque pointers it becomes undef, which cannot be correctly replaced. This patch tries to artificially add a dummy bitcast to make the replacement work correctly.

Oh I guess I understand the problem. The design of retcon ABI allows the user to specify the address of the coroutine frame. So it is possible that the frame is the argument of the original function indeed. Then I found true problem here is that the comment in https://github.com/llvm/llvm-project/blob/be5c3ca7fbaec90fff004af54d3cd5f6c30a9664/llvm/lib/Transforms/Coroutines/CoroSplit.cpp#L874-L876 is invalided. The arguments here shouldn't contain users.

So I believe the problem might be in splitRetconCoroutine. It replaces CoroBegin before cloning. (We treat the value returned by CoroBegin as the frame generally). But it replaces CoroBegin with the argument this time so that the argument still owns the user. So here is the problem. My solution is to replace CoroBegin later after the cloning. I sent https://reviews.llvm.org/D122383. And I add @rjmccall . He is the author of the retcon lowering. Maybe he could give further helps.

[ChuanqiXu]

Given the discussions in D122383, let's pursue this one.

[rjmccall]

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

874–884

Please explain why we need to do this in the comment, so that future maintainers aren't forced to track the commit all the way back to this conversation to understand it. And then the comment on the replacement later can just be something like "all uses of the arguments should have been resolved by this point, so we can safely remove the dummy values".

878

Is it really okay to add an unparented instruction to the value map?

[nikic]

Update comments

[nikic]

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

878

At least I'm not aware of any issues that would cause. Parenting it would be a bit awkward here, as the function is completely empty at this point (doesn't have any basic blocks to insert into).

[rjmccall]

LGTM

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

878

Ah, yes, somehow I missed that this was in clone-specific code. Alright, if it's not causing problems, I guess we can do this.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

 /// some sort.
 void CoroCloner::create() {
   // Create the new function if we don't already have one.
   if (!NewF) {
     NewF = createCloneDeclaration(OrigF, Shape, Suffix,
                                   OrigF.getParent()->end(), ActiveSuspend);
   }
 
-  // Replace all args with undefs. The buildCoroutineFrame algorithm already
-  // rewritten access to the args that occurs after suspend points with loads
-  // and stores to/from the coroutine frame.
-  for (Argument &A : OrigF.args())
-    VMap[&A] = UndefValue::get(A.getType());
+  // Replace all args with dummy instructions. If an argument is the old frame
+  // pointer, the dummy will be replaced by the new frame pointer once it is
+  // computed below. Uses of all other arguments should have already been
+  // rewritten by buildCoroutineFrame() to use loads/stores on the coroutine
+  // frame.
+  SmallVector<Instruction *> DummyArgs;
+  for (Argument &A : OrigF.args()) {
+    DummyArgs.push_back(
+        new BitCastInst(UndefValue::get(A.getType()), A.getType()));
+    VMap[&A] = DummyArgs.back();
+  }
 
   SmallVector<ReturnInst *, 4> Returns;
 
   // Ignore attempts to change certain attributes of the function.
   // TODO: maybe there should be a way to suppress this during cloning?
   auto savedVisibility = NewF->getVisibility();
   auto savedUnnamedAddr = NewF->getUnnamedAddr();
   auto savedDLLStorageClass = NewF->getDLLStorageClass();
void CoroCloner::create() {
 
   // Remap vFrame pointer.
   auto *NewVFrame = Builder.CreateBitCast(
       NewFramePtr, Type::getInt8PtrTy(Builder.getContext()), "vFrame");
   Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);
   if (OldVFrame != NewVFrame)
     OldVFrame->replaceAllUsesWith(NewVFrame);
 
+  // All uses of the arguments should have been resolved by this point,
+  // so we can safely remove the dummy values.
+  for (Instruction *DummyArg : DummyArgs) {
+    DummyArg->replaceAllUsesWith(UndefValue::get(DummyArg->getType()));
+    DummyArg->deleteValue();
+  }
+
   switch (Shape.ABI) {
   case coro::ABI::Switch:
     // Rewrite final suspend handling as it is not done via switch (allows to
     // remove final case from the switch, since it is undefined behavior to
     // resume the coroutine suspended at the final suspend point.
     if (Shape.SwitchLowering.HasFinalSuspend)
       handleFinalSuspend();
     break;

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*)