[Mlir-commits] [mlir] [mlir][async] Lazily create the coroutine destroy-cleanup block (PR #199583)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon May 25 17:20:25 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Matthias Springer (matthias-springer)
<details>
<summary>Changes</summary>
`setupCoroMachinery` previously emitted a `cleanupForDestroy` block unconditionally, alongside the normal `cleanup` block. That block is only ever used as the "destroy" successor of an `async.coro.suspend`, so for coroutines that never suspend (e.g. an `async.func` whose body contains no `async.await`) it ended up unreachable in the lowered CFG.
Make `cleanupForDestroy` mirror the existing `setError` (and `setupSetErrorBlock`) pattern and materialize it lazily via a new `setupCleanupForDestroyBlock` helper, called only from the two places (`outlineExecuteOp` and the `async.await` lowering) that actually wire it up. Store the coroutine id on `CoroMachinery` so the helper can rebuild the block contents without keeping the original `async.coro.id` op around.
Assisted-by: Opus 4.7
---
Full diff: https://github.com/llvm/llvm-project/pull/199583.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp (+33-15)
- (modified) mlir/test/Dialect/Async/async-to-async-runtime.mlir (+33)
``````````diff
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
index 0c5bcfe631c6c..6ed50671fb3b3 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
@@ -90,6 +90,7 @@ struct CoroMachinery {
std::optional<Value> asyncToken; // returned completion token
llvm::SmallVector<Value, 4> returnValues; // returned async values
+ Value coroId; // coroutine id (!async.coro.id value)
Value coroHandle; // coroutine handle (!async.coro.getHandle value)
Block *entry; // coroutine entry block
std::optional<Block *> setError; // set returned values to error state
@@ -115,7 +116,12 @@ struct CoroMachinery {
// If there is resume-specific cleanup logic, it can go into the Cleanup
// block but not the destroy block. Otherwise, it can fail block dominance
// check.
- Block *cleanupForDestroy;
+ //
+ // This block is created lazily by `setupCleanupForDestroyBlock` only when a
+ // suspension point needs a destroy successor, so that functions without any
+ // coroutine suspends (e.g. an `async.func` body with no `await`) don't end
+ // up with dead code.
+ std::optional<Block *> cleanupForDestroy;
Block *suspend; // coroutine suspension block
};
} // namespace
@@ -204,21 +210,16 @@ static CoroMachinery setupCoroMachinery(func::FuncOp func) {
cf::BranchOp::create(builder, originalEntryBlock);
Block *cleanupBlock = func.addBlock();
- Block *cleanupBlockForDestroy = func.addBlock();
Block *suspendBlock = func.addBlock();
// ------------------------------------------------------------------------ //
- // Coroutine cleanup blocks: deallocate coroutine frame, free the memory.
+ // Coroutine cleanup block: deallocate coroutine frame, free the memory.
// ------------------------------------------------------------------------ //
- auto buildCleanupBlock = [&](Block *cb) {
- builder.setInsertionPointToStart(cb);
- CoroFreeOp::create(builder, coroIdOp.getId(), coroHdlOp.getHandle());
-
- // Branch into the suspend block.
- cf::BranchOp::create(builder, suspendBlock);
- };
- buildCleanupBlock(cleanupBlock);
- buildCleanupBlock(cleanupBlockForDestroy);
+ // The matching "destroy" cleanup block is materialized lazily by
+ // `setupCleanupForDestroyBlock` only when a suspend point needs it.
+ builder.setInsertionPointToStart(cleanupBlock);
+ CoroFreeOp::create(builder, coroIdOp.getId(), coroHdlOp.getHandle());
+ cf::BranchOp::create(builder, suspendBlock);
// ------------------------------------------------------------------------ //
// Coroutine suspend block: mark the end of a coroutine and return allocated
@@ -249,11 +250,12 @@ static CoroMachinery setupCoroMachinery(func::FuncOp func) {
machinery.func = func;
machinery.asyncToken = retToken;
machinery.returnValues = retValues;
+ machinery.coroId = coroIdOp.getId();
machinery.coroHandle = coroHdlOp.getHandle();
machinery.entry = entryBlock;
machinery.setError = std::nullopt; // created lazily only if needed
machinery.cleanup = cleanupBlock;
- machinery.cleanupForDestroy = cleanupBlockForDestroy;
+ machinery.cleanupForDestroy = std::nullopt; // created lazily only if needed
machinery.suspend = suspendBlock;
return machinery;
}
@@ -283,6 +285,20 @@ static Block *setupSetErrorBlock(CoroMachinery &coro) {
return *coro.setError;
}
+// Lazily creates the `cleanupForDestroy` block only if a suspension point
+// actually needs a destroy successor. This avoids leaving an unreachable
+// cleanup block behind in coroutines that never suspend.
+static Block *setupCleanupForDestroyBlock(ImplicitLocOpBuilder &builder,
+ CoroMachinery &coro) {
+ if (coro.cleanupForDestroy)
+ return *coro.cleanupForDestroy;
+ OpBuilder::InsertionGuard guard(builder);
+ coro.cleanupForDestroy = builder.createBlock(coro.suspend);
+ CoroFreeOp::create(builder, coro.coroId, coro.coroHandle);
+ cf::BranchOp::create(builder, coro.suspend);
+ return *coro.cleanupForDestroy;
+}
+
//===----------------------------------------------------------------------===//
// async.execute op outlining to the coroutine functions.
//===----------------------------------------------------------------------===//
@@ -373,8 +389,9 @@ outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
RuntimeResumeOp::create(builder, coro.coroHandle);
// Add async.coro.suspend as a suspended block terminator.
+ Block *destroy = setupCleanupForDestroyBlock(builder, coro);
CoroSuspendOp::create(builder, coroSaveOp.getState(), coro.suspend,
- branch.getDest(), coro.cleanupForDestroy);
+ branch.getDest(), destroy);
branch.erase();
}
@@ -614,9 +631,10 @@ class AwaitOpLoweringBase : public OpConversionPattern<AwaitType> {
Block *resume = rewriter.splitBlock(suspended, Block::iterator(op));
// Add async.coro.suspend as a suspended block terminator.
+ Block *destroy = setupCleanupForDestroyBlock(builder, coro);
builder.setInsertionPointToEnd(suspended);
CoroSuspendOp::create(builder, coroSaveOp.getState(), coro.suspend,
- resume, coro.cleanupForDestroy);
+ resume, destroy);
// Split the resume block into error checking and continuation.
Block *continuation = rewriter.splitBlock(resume, Block::iterator(op));
diff --git a/mlir/test/Dialect/Async/async-to-async-runtime.mlir b/mlir/test/Dialect/Async/async-to-async-runtime.mlir
index 36583b2b94a3c..c7734aa044519 100644
--- a/mlir/test/Dialect/Async/async-to-async-runtime.mlir
+++ b/mlir/test/Dialect/Async/async-to-async-runtime.mlir
@@ -498,3 +498,36 @@ async.func @execute_in_async_func(%arg0: f32, %arg1: memref<1xf32>)
// CHECK-SAME: ) -> !async.token
// CHECK: %[[CST:.*]] = arith.constant 0 : index
// CHECK: memref.store %[[VALUE]], %[[MEMREF]][%[[CST]]]
+
+// -----
+
+// An async.func with no suspension points must not leave behind an
+// unreachable `cleanupForDestroy` block. The block is created lazily and is
+// only needed as the destroy successor of an `async.coro.suspend`, so an
+// empty body that never suspends should only have a single cleanup block.
+
+// CHECK-LABEL: @async_func_empty
+async.func @async_func_empty() -> !async.token {
+ return
+}
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin
+// CHECK: cf.br ^[[ORIGIN_ENTRY:.*]]
+
+// CHECK: ^[[ORIGIN_ENTRY]]:
+// CHECK-NEXT: async.runtime.set_available %[[TOKEN]]
+// CHECK-NEXT: cf.br ^[[CLEANUP:.*]]
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK-NEXT: async.coro.free %[[ID]], %[[HDL]]
+// CHECK-NEXT: cf.br ^[[SUSPEND:.*]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK-NEXT: async.coro.end %[[HDL]]
+// CHECK-NEXT: return %[[TOKEN]]
+
+// There must be exactly one async.coro.free op in the lowered function:
+// the destroy-cleanup block (which would contain a second one) should not
+// have been emitted.
+// CHECK-NOT: async.coro.free
``````````
</details>
https://github.com/llvm/llvm-project/pull/199583
More information about the Mlir-commits
mailing list