[coroutine] add nomerge function attribute to `llvm.coro.save`

[ychen]

Summary

It is illegal to merge two llvm.coro.save calls unless their
llvm.coro.suspend users are also merged. Marks it "nomerge" for
the moment.

This reverts D129025.

Alternative to D129025, which affects other token type users like WinEH.

[ChuanqiXu]

This looks good to me basically. I would love to see an example to test the two coro.save wouldn't be merged.

[ychen]

In D129530#3644074, @ChuanqiXu wrote:

This looks good to me basically. I would love to see an example to test the two coro.save wouldn't be merged.

Yep, hoist-skip-token.ll is testing this.

[ChuanqiXu]

In D129530#3644091, @ychen wrote:

In D129530#3644074, @ChuanqiXu wrote:

This looks good to me basically. I would love to see an example to test the two coro.save wouldn't be merged.

Yep, hoist-skip-token.ll is testing this.

Oh, sorry. I didn't notice that. The current version should be fine.

Personally, I would like to move the test to Transforms/Coroutine and add a new test that the two llvm.coro.suspend are the same but they couldn't be merged. This is not required.

[ychen]

• Address Chuanqi's feedback

[ychen]

In D129530#3644094, @ChuanqiXu wrote:

In D129530#3644091, @ychen wrote:

In D129530#3644074, @ChuanqiXu wrote:

This looks good to me basically. I would love to see an example to test the two coro.save wouldn't be merged.

Yep, hoist-skip-token.ll is testing this.

Oh, sorry. I didn't notice that. The current version should be fine.

Personally, I would like to move the test to Transforms/Coroutine and add a new test that the two llvm.coro.suspend are the same but they couldn't be merged. This is not required.

Sure, sounds good to me. Thanks for the review!

clang/test/CodeGenCoroutines/coro-attributes.cpp

 // RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 -disable-llvm-passes -emit-llvm %s -o - | FileCheck %s
 #include "Inputs/coroutine.h"
 
 using namespace std;
 
 struct coro {
   struct promise_type {
     coro get_return_object();
     suspend_never initial_suspend();
     suspend_never final_suspend() noexcept;
     void return_void();
     static void unhandled_exception();
   };
 };
 
 // CHECK: void @_Z3foov() #[[FOO_ATTR_NUM:[0-9]+]]
+// CHECK: declare token @llvm.coro.save(ptr) #[[SAVE_ATTR_NUM:[0-9]+]]
 // CHECK: attributes #[[FOO_ATTR_NUM]] = { {{.*}} presplitcoroutine
+// CHECK: attributes #[[SAVE_ATTR_NUM]] = { {{.*}}nomerge
 coro foo() {
   co_await suspend_always{};
 }

llvm/docs/Coroutines.rst

::
 
   declare token @llvm.coro.save(i8* <handle>)
 
 Overview:
 """""""""
 
 The '``llvm.coro.save``' marks the point where a coroutine need to update its
 state to prepare for resumption to be considered suspended (and thus eligible
-for resumption).
+for resumption). It is illegal to merge two '``llvm.coro.save``' calls unless their
+'``llvm.coro.suspend``' users are also merged. So '``llvm.coro.save``' is currently
+tagged with the `no_merge` function attribute.
 
 Arguments:
 """"""""""
 
 The first argument points to a coroutine handle of the enclosing coroutine.
 
 Semantics:
 """"""""""

llvm/include/llvm/IR/Intrinsics.td

 def int_coro_end_async
     : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_i1_ty, llvm_vararg_ty], []>;
 
 def int_coro_frame : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_noop : Intrinsic<[llvm_ptr_ty], [], [IntrNoMem]>;
 def int_coro_size : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
 def int_coro_align : Intrinsic<[llvm_anyint_ty], [], [IntrNoMem]>;
 
-def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], []>;
+def int_coro_save : Intrinsic<[llvm_token_ty], [llvm_ptr_ty], [IntrNoMerge]>;
 def int_coro_suspend : Intrinsic<[llvm_i8_ty], [llvm_token_ty, llvm_i1_ty], []>;
 def int_coro_suspend_retcon : Intrinsic<[llvm_any_ty], [llvm_vararg_ty], []>;
 def int_coro_prepare_retcon : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
                                         [IntrNoMem]>;
 def int_coro_alloca_alloc : Intrinsic<[llvm_token_ty],
                                       [llvm_anyint_ty, llvm_i32_ty], []>;
 def int_coro_alloca_get : Intrinsic<[llvm_ptr_ty], [llvm_token_ty], []>;
 def int_coro_alloca_free : Intrinsic<[], [llvm_token_ty], []>;

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

bool SimplifyCFGOpt::HoistThenElseCodeToIf(BranchInst *BI,
   }
 
   do {
     // If we are hoisting the terminator instruction, don't move one (making a
     // broken BB), instead clone it, and remove BI.
     if (I1->isTerminator())
       goto HoistTerminator;
 
-    // Hoisting token-returning instructions would obscure the origin.
-    if (I1->getType()->isTokenTy())
-      return Changed;
-
     // If we're going to hoist a call, make sure that the two instructions we're
     // commoning/hoisting are both marked with musttail, or neither of them is
     // marked as such. Otherwise, we might end up in a situation where we hoist
     // from a block where the terminator is a `ret` to a block where the terminator
     // is a `br`, and `musttail` calls expect to be followed by a return.
     auto *C1 = dyn_cast<CallInst>(I1);
     auto *C2 = dyn_cast<CallInst>(I2);
     if (C1 && C2)

llvm/test/Transforms/Coroutines/coro-save-nomerge.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes='simplifycfg<hoist-common-insts>' -S | FileCheck %s
+
+declare token @llvm.coro.save(ptr) #0
+declare i8 @llvm.coro.suspend(token, i1)
+
+define void @final_nonfinal_suspend(i32 %x) {
+; CHECK-LABEL: @final_nonfinal_suspend(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[X:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[AWAIT_SUSPEND:%.*]], label [[FINAL_SUSPEND:%.*]]
+; CHECK:       await.suspend:
+; CHECK-NEXT:    [[TMP0:%.*]] = call token @llvm.coro.save(ptr null)
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.coro.suspend(token [[TMP0]], i1 false)
+; CHECK-NEXT:    br label [[CORO_RET:%.*]]
+; CHECK:       final.suspend:
+; CHECK-NEXT:    [[TMP2:%.*]] = call token @llvm.coro.save(ptr null)
+; CHECK-NEXT:    [[TMP3:%.*]] = call i8 @llvm.coro.suspend(token [[TMP2]], i1 true)
+; CHECK-NEXT:    br label [[CORO_RET]]
+; CHECK:       coro.ret:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp = icmp slt i32 %x, 0
+  br i1 %cmp, label %await.suspend, label %final.suspend
+
+await.suspend:
+  %0 = call token @llvm.coro.save(ptr null)
+  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
+  br label %coro.ret
+
+final.suspend:
+  %2 = call token @llvm.coro.save(ptr null)
+  %3 = call i8 @llvm.coro.suspend(token %2, i1 true)
+  br label %coro.ret
+
+coro.ret:
+  ret void
+}
+
+define void @both_nonfinal_suspend(i32 %x) {
+; CHECK-LABEL: @both_nonfinal_suspend(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[X:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[AWAIT_SUSPEND:%.*]], label [[FINAL_SUSPEND:%.*]]
+; CHECK:       await.suspend:
+; CHECK-NEXT:    [[TMP0:%.*]] = call token @llvm.coro.save(ptr null)
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.coro.suspend(token [[TMP0]], i1 false)
+; CHECK-NEXT:    br label [[CORO_RET:%.*]]
+; CHECK:       final.suspend:
+; CHECK-NEXT:    [[TMP2:%.*]] = call token @llvm.coro.save(ptr null)
+; CHECK-NEXT:    [[TMP3:%.*]] = call i8 @llvm.coro.suspend(token [[TMP2]], i1 false)
+; CHECK-NEXT:    br label [[CORO_RET]]
+; CHECK:       coro.ret:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp = icmp slt i32 %x, 0
+  br i1 %cmp, label %await.suspend, label %final.suspend
+
+await.suspend:
+  %0 = call token @llvm.coro.save(ptr null)
+  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
+  br label %coro.ret
+
+final.suspend:
+  %2 = call token @llvm.coro.save(ptr null)
+  %3 = call i8 @llvm.coro.suspend(token %2, i1 false)
+  br label %coro.ret
+
+coro.ret:
+  ret void
+}
+
+attributes #0 = { nomerge }

llvm/test/Transforms/SimplifyCFG/hoist-skip-token.ll

-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -passes='simplifycfg<hoist-common-insts>' -S | FileCheck %s
-
-declare token @llvm.coro.save(ptr)
-declare i8 @llvm.coro.suspend(token, i1)
-
-define void @f(i32 %x) {
-; CHECK-LABEL: @f(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[X:%.*]], 0
-; CHECK-NEXT:    br i1 [[CMP]], label [[AWAIT_SUSPEND:%.*]], label [[FINAL_SUSPEND:%.*]]
-; CHECK:       await.suspend:
-; CHECK-NEXT:    [[TMP0:%.*]] = call token @llvm.coro.save(ptr null)
-; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.coro.suspend(token [[TMP0]], i1 false)
-; CHECK-NEXT:    br label [[CORO_RET:%.*]]
-; CHECK:       final.suspend:
-; CHECK-NEXT:    [[TMP2:%.*]] = call token @llvm.coro.save(ptr null)
-; CHECK-NEXT:    [[TMP3:%.*]] = call i8 @llvm.coro.suspend(token [[TMP2]], i1 true)
-; CHECK-NEXT:    br label [[CORO_RET]]
-; CHECK:       coro.ret:
-; CHECK-NEXT:    ret void
-;
-entry:
-  %cmp = icmp slt i32 %x, 0
-  br i1 %cmp, label %await.suspend, label %final.suspend
-
-await.suspend:
-  %0 = call token @llvm.coro.save(ptr null)
-  %1 = call i8 @llvm.coro.suspend(token %0, i1 false)
-  br label %coro.ret
-
-final.suspend:
-  %2 = call token @llvm.coro.save(ptr null)
-  %3 = call i8 @llvm.coro.suspend(token %2, i1 true)
-  br label %coro.ret
-
-coro.ret:
-  ret void
-}