[llvm] [Coroutines] Don't assert if coro-early runs more than once (PR #134854)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-coroutines

Author: Hans Wennborg (zmodem)

<details>
<summary>Changes</summary>

The pass may run more than once during ThinLTO for example (see bug). Maybe that means those pass pipelines aren't optimal, but the pass should be resilient against that.

Fixes #<!-- -->134054

---
Full diff: https://github.com/llvm/llvm-project/pull/134854.diff


2 Files Affected:

- (modified) llvm/include/llvm/Transforms/Coroutines/CoroInstr.h (+2-2) 
- (added) llvm/test/Transforms/Coroutines/coro-early-twice.ll (+36) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h
index fbc76219ead86..b8c846596c994 100644
--- a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h
@@ -217,8 +217,8 @@ class CoroIdInst : public AnyCoroIdInst {
     return cast<Function>(getArgOperand(CoroutineArg)->stripPointerCasts());
   }
   void setCoroutineSelf() {
-    assert(isa<ConstantPointerNull>(getArgOperand(CoroutineArg)) &&
-           "Coroutine argument is already assigned");
+    if (!isa<ConstantPointerNull>(getArgOperand(CoroutineArg)))
+      assert(getCoroutine() == getFunction() && "Don't change coroutine.");
     setArgOperand(CoroutineArg, getFunction());
   }
 
diff --git a/llvm/test/Transforms/Coroutines/coro-early-twice.ll b/llvm/test/Transforms/Coroutines/coro-early-twice.ll
new file mode 100644
index 0000000000000..39ec0ccc6fdb8
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-early-twice.ll
@@ -0,0 +1,36 @@
+; RUN: opt < %s -passes='module(coro-early,coro-early)' -S | FileCheck %s
+
+; Check that coro-early can run twice without asserting/crashing.
+
+; CHECK-LABEL: define ptr @f
+; CHECK: call token @llvm.coro.id(i32 0, ptr null, ptr @f, ptr null)
+
+define ptr @f(i32 %n) presplitcoroutine {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc)
+
+  %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sp1, label %suspend [i8 0, label %resume1
+                                  i8 1, label %cleanup]
+resume1:
+  br label %cleanup
+cleanup:
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 0, token none)
+  ret ptr %hdl
+}
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare i32 @llvm.coro.size.i32()
+declare noalias ptr @malloc(i32)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i8  @llvm.coro.suspend(token, i1)
+declare ptr @llvm.coro.free(token, ptr)
+declare void @free(ptr)
+declare i1 @llvm.coro.end(ptr, i1, token)

``````````

</details>


https://github.com/llvm/llvm-project/pull/134854