[llvm] [CoroEarly] Hide promise alloca for later passes (PR #139243)
Weibo He via llvm-commits
llvm-commits at lists.llvm.org
Mon May 12 06:19:55 PDT 2025
https://github.com/NewSigma updated https://github.com/llvm/llvm-project/pull/139243
>From fa9149a8b06665ddd439b4d4b7ca4af3faaef8e6 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Fri, 9 May 2025 18:43:07 +0800
Subject: [PATCH 1/3] [Coro] Hide promise alloca for later passes
---
llvm/docs/Coroutines.rst | 9 ++--
.../llvm/Transforms/Coroutines/CoroShape.h | 11 ++--
llvm/lib/Transforms/Coroutines/CoroEarly.cpp | 53 ++++++++++++++++---
llvm/lib/Transforms/Coroutines/Coroutines.cpp | 14 ++++-
.../AddressSanitizer/skip-coro.ll | 5 +-
llvm/test/Transforms/Coroutines/gh105595.ll | 32 +++++++++++
6 files changed, 106 insertions(+), 18 deletions(-)
create mode 100644 llvm/test/Transforms/Coroutines/gh105595.ll
diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst
index 60e32dc467d27..157b20b7d8d73 100644
--- a/llvm/docs/Coroutines.rst
+++ b/llvm/docs/Coroutines.rst
@@ -2121,10 +2121,11 @@ Coroutine Transformation Passes
===============================
CoroEarly
---------
-The pass CoroEarly lowers coroutine intrinsics that hide the details of the
-structure of the coroutine frame, but, otherwise not needed to be preserved to
-help later coroutine passes. This pass lowers `coro.frame`_, `coro.done`_,
-and `coro.promise`_ intrinsics.
+The CoroEarly pass ensures later middle end passes correctly interpret coroutine
+semantics and lowers coroutine intrinsics that not needed to be preserved to
+help later coroutine passes. This pass lowers `coro.promise`_ that outside the
+coroutine body, `coro.frame`_ and `coro.done`_ intrinsics. It replace uses of
+promise alloca with `coro.promise`_ intrinsic.
.. _CoroSplit:
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
index ea93ced1ce29e..aee918946c70f 100644
--- a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
+++ b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -79,7 +79,8 @@ struct Shape {
// Scan the function and collect the above intrinsics for later processing
void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
- SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
+ SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+ SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
// If for some reason, we were not able to find coro.begin, bailout.
void invalidateCoroutine(Function &F,
SmallVectorImpl<CoroFrameInst *> &CoroFrames);
@@ -87,7 +88,8 @@ struct Shape {
void initABI();
// Remove orphaned and unnecessary intrinsics
void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
- SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
+ SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+ SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
// Field indexes for special fields in the switch lowering.
struct SwitchFieldIndex {
@@ -265,13 +267,14 @@ struct Shape {
explicit Shape(Function &F) {
SmallVector<CoroFrameInst *, 8> CoroFrames;
SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
+ SmallVector<CoroPromiseInst *, 2> CoroPromises;
- analyze(F, CoroFrames, UnusedCoroSaves);
+ analyze(F, CoroFrames, UnusedCoroSaves, CoroPromises);
if (!CoroBegin) {
invalidateCoroutine(F, CoroFrames);
return;
}
- cleanCoroutine(CoroFrames, UnusedCoroSaves);
+ cleanCoroutine(CoroFrames, UnusedCoroSaves, CoroPromises);
}
};
diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
index 5375448d2d2e2..ff1e1c8cc56cf 100644
--- a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -9,6 +9,7 @@
#include "llvm/Transforms/Coroutines/CoroEarly.h"
#include "CoroInternal.h"
#include "llvm/IR/DIBuilder.h"
+#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
@@ -165,6 +166,7 @@ static void setCannotDuplicate(CoroIdInst *CoroId) {
void Lowerer::lowerEarlyIntrinsics(Function &F) {
CoroIdInst *CoroId = nullptr;
+ CoroBeginInst *CoroBegin = nullptr;
SmallVector<CoroFreeInst *, 4> CoroFrees;
bool HasCoroSuspend = false;
for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
@@ -175,6 +177,23 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
switch (CB->getIntrinsicID()) {
default:
continue;
+ case Intrinsic::coro_begin:
+ case Intrinsic::coro_begin_custom_abi: {
+ auto CBI = cast<CoroBeginInst>(&I);
+
+ // Ignore coro id's that aren't pre-split.
+ auto Id = dyn_cast<CoroIdInst>(CBI->getId());
+ if (Id && !Id->getInfo().isPreSplit())
+ break;
+
+ if (CoroBegin)
+ report_fatal_error(
+ "coroutine should have exactly one defining @llvm.coro.begin");
+ CBI->addRetAttr(Attribute::NonNull);
+ CBI->addRetAttr(Attribute::NoAlias);
+ CoroBegin = CBI;
+ break;
+ }
case Intrinsic::coro_free:
CoroFrees.push_back(cast<CoroFreeInst>(&I));
break;
@@ -218,22 +237,44 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
case Intrinsic::coro_destroy:
lowerResumeOrDestroy(*CB, CoroSubFnInst::DestroyIndex);
break;
- case Intrinsic::coro_promise:
- lowerCoroPromise(cast<CoroPromiseInst>(&I));
+ case Intrinsic::coro_promise: {
+ bool OutsideCoro = CoroBegin == nullptr;
+ if (OutsideCoro)
+ lowerCoroPromise(cast<CoroPromiseInst>(&I));
break;
+ }
case Intrinsic::coro_done:
lowerCoroDone(cast<IntrinsicInst>(&I));
break;
}
}
- // Make sure that all CoroFree reference the coro.id intrinsic.
- // Token type is not exposed through coroutine C/C++ builtins to plain C, so
- // we allow specifying none and fixing it up here.
- if (CoroId)
+ if (CoroId) {
+ // Make sure that all CoroFree reference the coro.id intrinsic.
+ // Token type is not exposed through coroutine C/C++ builtins to plain C, so
+ // we allow specifying none and fixing it up here.
for (CoroFreeInst *CF : CoroFrees)
CF->setArgOperand(0, CoroId);
+ if (auto *PA = CoroId->getPromise()) {
+ assert(CoroBegin && "Use Switch-Resumed ABI but missing coro.begin");
+
+ Builder.SetInsertPoint(*CoroBegin->getInsertionPointAfterDef());
+
+ auto *Alignment = Builder.getInt32(PA->getAlign().value());
+ auto *FromPromise = Builder.getInt1(false);
+ SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};
+ auto *PI =
+ Builder.CreateIntrinsic(Builder.getPtrTy(), Intrinsic::coro_promise,
+ Arg, {}, "promise.addr");
+ PA->replaceUsesWithIf(PI, [CoroId](Use &U) {
+ bool IsBitcast = U == U.getUser()->stripPointerCasts();
+ bool IsCoroId = U.getUser() == CoroId;
+ return !IsBitcast && !IsCoroId;
+ });
+ }
+ }
+
// Coroutine suspention could potentially lead to any argument modified
// outside of the function, hence arguments should not have noalias
// attributes.
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 7b59c39283ded..325ad36bc666a 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -192,7 +192,8 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
// Collect "interesting" coroutine intrinsics.
void coro::Shape::analyze(Function &F,
SmallVectorImpl<CoroFrameInst *> &CoroFrames,
- SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
+ SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+ SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
clear();
bool HasFinalSuspend = false;
@@ -286,6 +287,9 @@ void coro::Shape::analyze(Function &F,
}
}
break;
+ case Intrinsic::coro_promise:
+ CoroPromises.push_back(cast<CoroPromiseInst>(II));
+ break;
}
}
}
@@ -477,7 +481,8 @@ void coro::AnyRetconABI::init() {
void coro::Shape::cleanCoroutine(
SmallVectorImpl<CoroFrameInst *> &CoroFrames,
- SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
+ SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+ SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
// The coro.frame intrinsic is always lowered to the result of coro.begin.
for (CoroFrameInst *CF : CoroFrames) {
CF->replaceAllUsesWith(CoroBegin);
@@ -489,6 +494,11 @@ void coro::Shape::cleanCoroutine(
for (CoroSaveInst *CoroSave : UnusedCoroSaves)
CoroSave->eraseFromParent();
UnusedCoroSaves.clear();
+
+ for (auto *PI : CoroPromises) {
+ PI->replaceAllUsesWith(getPromiseAlloca());
+ PI->eraseFromParent();
+ }
}
static void propagateCallAttrsFromCallee(CallInst *Call, Function *Callee) {
diff --git a/llvm/test/Instrumentation/AddressSanitizer/skip-coro.ll b/llvm/test/Instrumentation/AddressSanitizer/skip-coro.ll
index 65b27ee2241dd..ed3af686d5257 100644
--- a/llvm/test/Instrumentation/AddressSanitizer/skip-coro.ll
+++ b/llvm/test/Instrumentation/AddressSanitizer/skip-coro.ll
@@ -15,8 +15,9 @@ define ptr @foo() #0 {
entry:
%__promise = alloca %struct.Promise, align 8
%0 = call token @llvm.coro.id(i32 16, ptr nonnull %__promise, ptr null, ptr null)
- %1 = call ptr @llvm.coro.noop()
- ret ptr %1
+ %1 = call ptr @llvm.coro.begin(token %0, ptr null)
+ %2 = call ptr @llvm.coro.noop()
+ ret ptr %2
}
declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
diff --git a/llvm/test/Transforms/Coroutines/gh105595.ll b/llvm/test/Transforms/Coroutines/gh105595.ll
new file mode 100644
index 0000000000000..cc23419bf7ff4
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/gh105595.ll
@@ -0,0 +1,32 @@
+; Test that store-load operation that crosses suspension point will not be eliminated by DSE
+; Coro result conversion function that attempts to modify promise shall produce this pattern
+; RUN: opt < %s -passes='coro-early,dse' -S | FileCheck %s
+
+define void @fn() presplitcoroutine {
+ %__promise = alloca ptr, align 8
+ %id = call token @llvm.coro.id(i32 16, ptr %__promise, ptr @fn, ptr null)
+ %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+; CHECK: %promise.addr = call ptr @llvm.coro.promise(ptr %hdl, i32 8, i1 false)
+ %save = call token @llvm.coro.save(ptr null)
+ %sp = call i8 @llvm.coro.suspend(token %save, i1 false)
+ %flag = icmp ule i8 %sp, 1
+ br i1 %flag, label %resume, label %suspend
+
+resume:
+; CHECK: call void @use_value(ptr %promise.addr)
+ call void @use_value(ptr %__promise)
+ br label %suspend
+
+suspend:
+ call i1 @llvm.coro.end(ptr null, i1 false, token none)
+; load value when resume
+; CHECK: %null = load ptr, ptr %promise.addr, align 8
+ %null = load ptr, ptr %__promise, align 8
+ call void @use_value(ptr %null)
+; store value when suspend
+; CHECK: store ptr null, ptr %promise.addr, align 8
+ store ptr null, ptr %__promise, align 8
+ ret void
+}
+
+declare void @use_value(ptr)
>From 2908155efeaa065c9c8f90e355ad4820195a7cc3 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Mon, 12 May 2025 11:13:24 +0800
Subject: [PATCH 2/3] Reformat
---
llvm/lib/Transforms/Coroutines/CoroEarly.cpp | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
index ff1e1c8cc56cf..96ec2860f3853 100644
--- a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -9,7 +9,6 @@
#include "llvm/Transforms/Coroutines/CoroEarly.h"
#include "CoroInternal.h"
#include "llvm/IR/DIBuilder.h"
-#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
@@ -264,9 +263,8 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
auto *Alignment = Builder.getInt32(PA->getAlign().value());
auto *FromPromise = Builder.getInt1(false);
SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};
- auto *PI =
- Builder.CreateIntrinsic(Builder.getPtrTy(), Intrinsic::coro_promise,
- Arg, {}, "promise.addr");
+ auto *PI = Builder.CreateIntrinsic(
+ Builder.getPtrTy(), Intrinsic::coro_promise, Arg, {}, "promise.addr");
PA->replaceUsesWithIf(PI, [CoroId](Use &U) {
bool IsBitcast = U == U.getUser()->stripPointerCasts();
bool IsCoroId = U.getUser() == CoroId;
>From 28596aab28c5e53fca955aad61391ecb0abb89bb Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Mon, 12 May 2025 21:16:10 +0800
Subject: [PATCH 3/3] Handle cases that coro.promise may be inlined
---
llvm/lib/Transforms/Coroutines/Coroutines.cpp | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 325ad36bc666a..a7227224eb900 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -495,8 +495,10 @@ void coro::Shape::cleanCoroutine(
CoroSave->eraseFromParent();
UnusedCoroSaves.clear();
+ auto *AI = getPromiseAlloca();
for (auto *PI : CoroPromises) {
- PI->replaceAllUsesWith(getPromiseAlloca());
+ PI->replaceAllUsesWith(PI->isFromPromise() ? cast<Value>(CoroBegin)
+ : cast<Value>(AI));
PI->eraseFromParent();
}
}
More information about the llvm-commits
mailing list