[llvm] [Coroutines] Only rematerialize when operands are free after the suspend (PR #209195)
Christian Ulmann via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 13 22:52:39 PDT 2026
https://github.com/Dinistro updated https://github.com/llvm/llvm-project/pull/209195
>From a2b5bc5d76b89fb2d254d8ad8e8a7382f7a0cb8f Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Fri, 10 Jul 2026 21:42:50 +0200
Subject: [PATCH 1/2] [Coroutines] Only rematerialize when operands are free
after the suspend
CoroSplit rematerializes materializable values across suspend points based
only on their opcode, with no profitability check. When the cone bottoms out
in non-materializable leaves that are otherwise dead across the suspend, this
forces those leaves into the coroutine frame instead of spilling the single
value.
Rematerialize a value only when all of its operands are free after the suspend:
a constant, an argument, a value that already crosses a suspend independently,
or a materializable value whose own operands are all free. Otherwise spill the
value itself.
This also fixes SuspendCrossingInfo::isDefinitionAcrossSuspend(Value&), which
fell through to an llvm_unreachable when no user crossed a suspend instead of
returning false (latent, as the overload had no in-tree callers). The guard
uses this overload to test whether an operand already crosses a suspend on its
own.
coro-retcon-resume-values.ll shows the effect end to end: the retcon frame now
fits the inline buffer, dropping a heap allocation.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Coroutines/SuspendCrossingInfo.h | 5 +-
.../Coroutines/MaterializationUtils.cpp | 99 ++++++++++++-
.../coro-materialize-check-operands.ll | 133 ++++++++++++++++++
.../Coroutines/coro-retcon-resume-values.ll | 17 +--
4 files changed, 234 insertions(+), 20 deletions(-)
create mode 100644 llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll
diff --git a/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h b/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h
index 4ed3266c46fe8..f9ded818479e4 100644
--- a/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h
+++ b/llvm/include/llvm/Transforms/Coroutines/SuspendCrossingInfo.h
@@ -187,10 +187,13 @@ class SuspendCrossingInfo {
for (User *U : Arg->users())
if (isDefinitionAcrossSuspend(*Arg, U))
return true;
- } else if (auto *Inst = dyn_cast<Instruction>(&V)) {
+ return false;
+ }
+ if (auto *Inst = dyn_cast<Instruction>(&V)) {
for (User *U : Inst->users())
if (isDefinitionAcrossSuspend(*Inst, U))
return true;
+ return false;
}
llvm_unreachable(
diff --git a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp
index 1827c702d2c77..bb180d4409522 100644
--- a/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp
+++ b/llvm/lib/Transforms/Coroutines/MaterializationUtils.cpp
@@ -11,7 +11,11 @@
#include "llvm/Transforms/Coroutines/MaterializationUtils.h"
#include "CoroInternal.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instruction.h"
@@ -28,6 +32,82 @@ using namespace coro;
// "coro-frame", which results in leaner debug spew.
#define DEBUG_TYPE "coro-suspend-crossing"
+// Returns true if \p Root is available in the resume function without
+// introducing a new spill. This holds when a value is:
+// 1. A constant: Genuinely free, as it is materialized as an immediate in the
+// resume function.
+// 2. An argument: Treated as free heuristically. An argument crossing the
+// suspend is actually spilled, but treating it as free preserves beneficial
+// shared-operand rematerialization
+// 3. A value that already crosses a suspend point for an independent use, i.e.,
+// it is spilled regardless
+// 4. A materializable value all of whose operands are themselves free.
+static bool isFreeAfterSuspend(
+ Value *Root, const std::function<bool(Instruction &)> &Materializable,
+ const SuspendCrossingInfo &Checker, SmallDenseMap<Value *, bool> &Memo) {
+ SmallVector<Value *> Stack;
+ // Materializable nodes whose operands have been scheduled but not yet folded.
+ // Distinguishes a node being visited a second time (operands ready) from the
+ // first visit, and lets an operand reached through a cycle resolve to false.
+ SmallPtrSet<Value *, 16> Opened;
+ Stack.push_back(Root);
+ while (!Stack.empty()) {
+ Value *V = Stack.back();
+ if (Memo.contains(V)) {
+ Stack.pop_back();
+ continue;
+ }
+
+ // Leaves that resolve without inspecting operands. Non-instructions
+ // (constants, arguments) are free; a value that already crosses a suspend
+ // independently is spilled regardless, so referencing its slot is free.
+ auto *I = dyn_cast<Instruction>(V);
+ if (!I || Checker.isDefinitionAcrossSuspend(*I)) {
+ Memo[V] = true;
+ Stack.pop_back();
+ continue;
+ }
+ if (!Materializable(*I)) {
+ Memo[V] = false;
+ Stack.pop_back();
+ continue;
+ }
+
+ // Materializable: free iff all operands are free. On the first visit,
+ // schedule the unresolved operands above V and revisit V once they are
+ // folded.
+ if (Opened.insert(V).second) {
+ for (Use &U : I->operands())
+ if (!Memo.contains(U.get()))
+ Stack.push_back(U.get());
+ continue;
+ }
+
+ // Second visit: operands are resolved, except any reached through a cycle
+ // (only possible via non-materializable PHIs, which never open). An
+ // unresolved operand is treated as not free, conservatively.
+ Stack.pop_back();
+ Memo[V] = all_of(I->operands(), [&](Use &U) {
+ auto It = Memo.find(U.get());
+ return It != Memo.end() && It->second;
+ });
+ }
+ return Memo.lookup(Root);
+}
+
+// Returns true if every operand of \p I is free after the suspend. This is the
+// profitability condition for rematerializing \p I: recomputing it then forces
+// no new value into the coroutine frame. Note this deliberately does NOT
+// consult isDefinitionAcrossSuspend(I) on I itself -- every remat candidate
+// crosses a suspend by definition, so that would make the check a no-op.
+static bool allOperandsFreeAfterSuspend(
+ Instruction &I, const std::function<bool(Instruction &)> &Materializable,
+ const SuspendCrossingInfo &Checker, SmallDenseMap<Value *, bool> &Memo) {
+ return all_of(I.operands(), [&](Use &U) {
+ return isFreeAfterSuspend(U.get(), Materializable, Checker, Memo);
+ });
+}
+
namespace {
// RematGraph is used to construct a DAG for rematerializable instructions
@@ -52,10 +132,16 @@ struct RematGraph {
RematNodeMap Remats;
const std::function<bool(Instruction &)> &MaterializableCallback;
SuspendCrossingInfo &Checker;
+ // Shared across all RematGraphs and the candidate scan in
+ // doRematerializations: isFreeAfterSuspend is context-free and the IR is not
+ // mutated until rematerialization, so results can be cached across uses.
+ SmallDenseMap<Value *, bool> &FreeMemo;
RematGraph(const std::function<bool(Instruction &)> &MaterializableCallback,
- Instruction *I, SuspendCrossingInfo &Checker)
- : MaterializableCallback(MaterializableCallback), Checker(Checker) {
+ Instruction *I, SuspendCrossingInfo &Checker,
+ SmallDenseMap<Value *, bool> &FreeMemo)
+ : MaterializableCallback(MaterializableCallback), Checker(Checker),
+ FreeMemo(FreeMemo) {
std::unique_ptr<RematNode> FirstNode = std::make_unique<RematNode>(I);
EntryNode = FirstNode.get();
std::deque<std::unique_ptr<RematNode>> WorkList;
@@ -80,7 +166,9 @@ struct RematGraph {
for (auto &Def : N->Node->operands()) {
Instruction *D = dyn_cast<Instruction>(Def.get());
if (!D || !MaterializableCallback(*D) ||
- !Checker.isDefinitionAcrossSuspend(*D, FirstUse))
+ !Checker.isDefinitionAcrossSuspend(*D, FirstUse) ||
+ !allOperandsFreeAfterSuspend(*D, MaterializableCallback, Checker,
+ FreeMemo))
continue;
if (auto It = Remats.find(D); It != Remats.end()) {
@@ -316,9 +404,12 @@ void coro::doRematerializations(
// See if there are materializable instructions across suspend points
// We record these as the starting point to also identify materializable
// defs of uses in these operations
+ SmallDenseMap<Value *, bool> FreeMemo;
for (Instruction &I : instructions(F)) {
if (!IsMaterializable(I))
continue;
+ if (!allOperandsFreeAfterSuspend(I, IsMaterializable, Checker, FreeMemo))
+ continue;
for (User *U : I.users())
if (Checker.isDefinitionAcrossSuspend(I, U))
Spills[&I].push_back(cast<Instruction>(U));
@@ -350,7 +441,7 @@ void coro::doRematerializations(
// Constructor creates the whole RematGraph for the given Use
auto RematUPtr =
- std::make_unique<RematGraph>(IsMaterializable, U, Checker);
+ std::make_unique<RematGraph>(IsMaterializable, U, Checker, FreeMemo);
LLVM_DEBUG(dbgs() << "***** Next remat group *****\n";
ReversePostOrderTraversal<RematGraph *> RPOT(RematUPtr.get());
diff --git a/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll b/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll
new file mode 100644
index 0000000000000..7d6e7dd79c420
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll
@@ -0,0 +1,133 @@
+; RUN: opt %s -passes='cgscc(coro-split)' -S | FileCheck %s
+
+; Verify -coro-remat-check-operands only rematerializes a value across a suspend
+; when doing so introduces no new spill.
+
+target datalayout = "e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+; %v's operands are two call results that are otherwise dead across the suspend.
+; Rematerializing %v would force both calls' results into the frame, so spill %v.
+
+; CHECK-LABEL: @f_nonfree(
+; CHECK: %[[V:.*]] = add i32 %[[C1:.*]], %[[C2:.*]]
+; CHECK: store i32 %[[V]]
+; CHECK-NOT: store i32 %[[C1]]
+; CHECK-NOT: store i32 %[[C2]]
+define ptr @f_nonfree() presplitcoroutine {
+entry:
+ %id = call token @llvm.coro.id(i32 0, ptr null, ptr @f_nonfree, 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)
+ %c1 = call i32 @opaque()
+ %c2 = call i32 @opaque()
+ %v = add i32 %c1, %c2
+ %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+ switch i8 %sp1, label %suspend [i8 0, label %resume
+ i8 1, label %cleanup]
+resume:
+ call void @print(i32 %v)
+ br label %cleanup
+cleanup:
+ %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+ call void @free(ptr %mem)
+ br label %suspend
+suspend:
+ call void @llvm.coro.end(ptr %hdl, i1 0, token none)
+ ret ptr %hdl
+}
+
+; %inc's only non-constant operand is an argument -> free -> still rematerialized
+; with the flag on: %n is spilled, %inc is not.
+
+; CHECK-LABEL: @f_free(
+; CHECK-SAME: %[[N:[[:alnum:]]+]]
+; CHECK: store i32 %[[N]]
+; CHECK-NOT: store i32
+; CHECK: switch
+define ptr @f_free(i32 %n) presplitcoroutine {
+entry:
+ %id = call token @llvm.coro.id(i32 0, ptr null, ptr @f_free, 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)
+ %inc = add i32 %n, 1
+ %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+ switch i8 %sp1, label %suspend [i8 0, label %resume
+ i8 1, label %cleanup]
+resume:
+ call void @print(i32 %inc)
+ br label %cleanup
+cleanup:
+ %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+ call void @free(ptr %mem)
+ br label %suspend
+suspend:
+ call void @llvm.coro.end(ptr %hdl, i1 0, token none)
+ ret ptr %hdl
+}
+
+; %v is materializable and still rematerialized: its operand %m crosses the
+; suspend on its own (used directly in resume), so %m is "free as an operand,
+; i.e., its frame slot is referenced at no new cost. But %m must NOT itself be
+; rematerialized, because %m's leaf %leaf is a call result that is otherwise
+; dead across the suspend.
+
+; CHECK-LABEL: @f_nested(
+; CHECK: %[[M:.*]] = add i32 %[[LEAF:.*]], 1
+; CHECK-NOT: store i32 %[[LEAF]]
+; CHECK: store i32 %[[M]]
+; CHECK-NOT: store i32 %[[LEAF]]
+define ptr @f_nested() presplitcoroutine {
+entry:
+ %id = call token @llvm.coro.id(i32 0, ptr null, ptr @f_nested, 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)
+ %leaf = call i32 @opaque()
+ %m = add i32 %leaf, 1
+ %v = add i32 %m, 2
+ %sp1 = call i8 @llvm.coro.suspend(token none, i1 false)
+ switch i8 %sp1, label %suspend [i8 0, label %resume
+ i8 1, label %cleanup]
+resume:
+ call void @print(i32 %v)
+ call void @print(i32 %m)
+ br label %cleanup
+cleanup:
+ %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+ call void @free(ptr %mem)
+ br label %suspend
+suspend:
+ call void @llvm.coro.end(ptr %hdl, i1 0, token none)
+ ret ptr %hdl
+}
+
+; coro-split emits all ramp functions before any of the resume/destroy/cleanup
+; thunks, so these checks (which must scan forward) come after all defines above.
+
+; CHECK-LABEL: @f_nonfree.resume(
+; CHECK: %[[V_RELOAD:.*]] = load i32
+; CHECK-NOT: add i32
+; CHECK: call void @print(i32 %[[V_RELOAD]])
+
+; CHECK-LABEL: @f_free.resume(
+; CHECK: %[[N_RELOAD:.*]] = load i32
+; CHECK: add i32 %[[N_RELOAD]], 1
+
+; CHECK-LABEL: @f_nested.resume(
+; CHECK: %[[M_RELOAD:.*]] = load i32
+; CHECK-NOT: add i32 %{{.*}}, 1
+; CHECK: add i32 %[[M_RELOAD]], 2
+; CHECK-NOT: add i32 %{{.*}}, 1
+
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8 @llvm.coro.suspend(token, i1)
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare ptr @llvm.coro.begin(token, ptr)
+declare void @llvm.coro.end(ptr, i1, token)
+declare noalias ptr @malloc(i32)
+declare i32 @opaque()
+declare void @print(i32)
+declare void @free(ptr)
diff --git a/llvm/test/Transforms/Coroutines/coro-retcon-resume-values.ll b/llvm/test/Transforms/Coroutines/coro-retcon-resume-values.ll
index 2f04453d69c4b..b03e0c85af293 100644
--- a/llvm/test/Transforms/Coroutines/coro-retcon-resume-values.ll
+++ b/llvm/test/Transforms/Coroutines/coro-retcon-resume-values.ll
@@ -4,9 +4,7 @@
define ptr @f(ptr %buffer, i32 %n) {
; CHECK-LABEL: @f(
; CHECK-NEXT: coro.return:
-; CHECK-NEXT: [[TMP0:%.*]] = tail call ptr @allocate(i32 12)
-; CHECK-NEXT: store ptr [[TMP0]], ptr [[BUFFER:%.*]], align 8
-; CHECK-NEXT: store i32 [[N:%.*]], ptr [[TMP0]], align 4
+; CHECK-NEXT: store i32 [[N:%.*]], ptr [[TMP0:%.*]], align 4
; CHECK-NEXT: ret ptr @f.resume.0
;
entry:
@@ -36,18 +34,7 @@ cleanup:
define i32 @main() {
; CHECK-LABEL: @main(
; CHECK-NEXT: entry:
-; CHECK-NEXT: [[TMP0:%.*]] = tail call ptr @allocate(i32 12)
-; CHECK-NEXT: store i32 1, ptr [[TMP0]], align 4
-; CHECK-NEXT: [[N_VAL3_SPILL_ADDR_I:%.*]] = getelementptr inbounds nuw i8, ptr [[TMP0]], i64 4
-; CHECK-NEXT: store i32 1, ptr [[N_VAL3_SPILL_ADDR_I]], align 4, !noalias [[META0:![0-9]+]]
-; CHECK-NEXT: [[INPUT_SPILL_ADDR_I:%.*]] = getelementptr inbounds nuw i8, ptr [[TMP0]], i64 8
-; CHECK-NEXT: store i32 2, ptr [[INPUT_SPILL_ADDR_I]], align 4, !noalias [[META0]]
-; CHECK-NEXT: [[INPUT_RELOAD_ADDR13_I:%.*]] = getelementptr inbounds nuw i8, ptr [[TMP0]], i64 8
-; CHECK-NEXT: [[N_VAL3_RELOAD_ADDR11_I:%.*]] = getelementptr inbounds nuw i8, ptr [[TMP0]], i64 4
-; CHECK-NEXT: store i32 3, ptr [[N_VAL3_RELOAD_ADDR11_I]], align 4, !noalias [[META3:![0-9]+]]
-; CHECK-NEXT: store i32 4, ptr [[INPUT_RELOAD_ADDR13_I]], align 4, !noalias [[META3]]
-; CHECK-NEXT: tail call void @print(i32 7), !noalias [[META6:![0-9]+]]
-; CHECK-NEXT: tail call void @deallocate(ptr nonnull [[TMP0]]), !noalias [[META6]]
+; CHECK-NEXT: tail call void @print(i32 7), !noalias [[META0:![0-9]+]], !inline_history [[META3:![0-9]+]]
; CHECK-NEXT: ret i32 0
;
entry:
>From 5dd58c9c5dc0911b942d706567211ecae82a12a7 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Tue, 14 Jul 2026 07:52:26 +0200
Subject: [PATCH 2/2] test cleanup
---
.../Coroutines/coro-materialize-check-operands.ll | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll b/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll
index 7d6e7dd79c420..50f934a0a6049 100644
--- a/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll
+++ b/llvm/test/Transforms/Coroutines/coro-materialize-check-operands.ll
@@ -1,9 +1,7 @@
; RUN: opt %s -passes='cgscc(coro-split)' -S | FileCheck %s
-; Verify -coro-remat-check-operands only rematerializes a value across a suspend
-; when doing so introduces no new spill.
-
-target datalayout = "e-m:e-p:64:64-i64:64-f80:128-n8:16:32:64-S128"
+; Verify that values are only rematerializes across a suspend when doing so
+; introduces no new spill.
; %v's operands are two call results that are otherwise dead across the suspend.
; Rematerializing %v would force both calls' results into the frame, so spill %v.
More information about the llvm-commits
mailing list