[CoroSplit] Handle argument being the frame pointer (PR54523)

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.

Fixes #54523.

Differential Revision: https://reviews.llvm.org/D122375

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -871,11 +871,17 @@ void CoroCloner::create() {
                                   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;
 
@@ -1019,6 +1025,13 @@ void CoroCloner::create() {
   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

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

@@ -0,0 +1,90 @@
+; 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*)