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

[Hans Wennborg via llvm-commits]

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

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

>From f6c0a8241074b01835840a3ffb67482e14bc8cdf Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Tue, 8 Apr 2025 14:49:24 +0200
Subject: [PATCH] [Coroutines] Don't assert if coro-early runs more than once

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
---
 .../llvm/Transforms/Coroutines/CoroInstr.h    |  4 +--
 .../Transforms/Coroutines/coro-early-twice.ll | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/Coroutines/coro-early-twice.ll

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)