[CIR] Upstream CIR await op #168133
Merged
[CIR] Upstream CIR await op #168133
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Andres-Salamanca requested review from andykaylor, bcardosolopes, lanza and xlauko as code owners 
llvmbot added clang 
Member
| @llvm/pr-subscribers-clang-codegen @llvm/pr-subscribers-clangir Author: None (Andres-Salamanca) ChangesThis PR upstreams Patch is 21.40 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/168133.diff 12 Files Affected:
diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td index 16258513239d9..832f79607a101 100644 --- a/clang/include/clang/CIR/Dialect/IR/CIROps.td +++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td @@ -790,8 +790,8 @@ def CIR_ConditionOp : CIR_Op<"condition", [ //===----------------------------------------------------------------------===// defvar CIR_YieldableScopes = [ - "ArrayCtor", "ArrayDtor", "CaseOp", "DoWhileOp", "ForOp", "GlobalOp", "IfOp", - "ScopeOp", "SwitchOp", "TernaryOp", "WhileOp", "TryOp" + "ArrayCtor", "ArrayDtor", "AwaitOp", "CaseOp", "DoWhileOp", "ForOp", + "GlobalOp", "IfOp", "ScopeOp", "SwitchOp", "TernaryOp", "WhileOp", "TryOp" ]; def CIR_YieldOp : CIR_Op<"yield", [ @@ -2707,6 +2707,102 @@ def CIR_CallOp : CIR_CallOpBase<"call", [NoRegionArguments]> { ]; } +//===----------------------------------------------------------------------===// +// AwaitOp +//===----------------------------------------------------------------------===// + +def CIR_AwaitKind : CIR_I32EnumAttr<"AwaitKind", "await kind", [ + I32EnumAttrCase<"Init", 0, "init">, + I32EnumAttrCase<"User", 1, "user">, + I32EnumAttrCase<"Yield", 2, "yield">, + I32EnumAttrCase<"Final", 3, "final"> +]>; + +def CIR_AwaitOp : CIR_Op<"await",[ + DeclareOpInterfaceMethods<RegionBranchOpInterface>, + RecursivelySpeculatable, NoRegionArguments +]> { + let summary = "Wraps C++ co_await implicit logic"; + let description = [{ + The under the hood effect of using C++ `co_await expr` roughly + translates to: + + ```c++ + // co_await expr; + + auto &&x = CommonExpr(); + if (!x.await_ready()) { + ... + x.await_suspend(...); + ... + } + x.await_resume(); + ``` + + `cir.await` represents this logic by using 3 regions: + - ready: covers veto power from x.await_ready() + - suspend: wraps actual x.await_suspend() logic + - resume: handles x.await_resume() + + Breaking this up in regions allow individual scrutiny of conditions + which might lead to folding some of them out. Lowerings coming out + of CIR, e.g. LLVM, should use the `suspend` region to track more + lower level codegen (e.g. intrinsic emission for coro.save/coro.suspend). + + There are also 4 flavors of `cir.await` available: + - `init`: compiler generated initial suspend via implicit `co_await`. + - `user`: also known as normal, representing user written co_await's. + - `yield`: user written `co_yield` expressions. + - `final`: compiler generated final suspend via implicit `co_await`. + + From the C++ snippet we get: + + ```mlir + cir.scope { + ... // auto &&x = CommonExpr(); + cir.await(user, ready : { + ... // x.await_ready() + }, suspend : { + ... // x.await_suspend() + }, resume : { + ... // x.await_resume() + }) + } + ``` + + Note that resulution of the common expression is assumed to happen + as part of the enclosing await scope. + }]; + + let arguments = (ins CIR_AwaitKind:$kind); + let regions = (region SizedRegion<1>:$ready, + SizedRegion<1>:$suspend, + SizedRegion<1>:$resume); + let assemblyFormat = [{ + `(` $kind `,` + `ready` `:` $ready `,` + `suspend` `:` $suspend `,` + `resume` `:` $resume `,` + `)` + attr-dict + }]; + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<(ins + "cir::AwaitKind":$kind, + CArg<"BuilderCallbackRef", + "nullptr">:$readyBuilder, + CArg<"BuilderCallbackRef", + "nullptr">:$suspendBuilder, + CArg<"BuilderCallbackRef", + "nullptr">:$resumeBuilder + )> + ]; + + let hasVerifier = 1; +} + //===----------------------------------------------------------------------===// // CopyOp //===----------------------------------------------------------------------===// diff --git a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp index 05fb1aedcbf4a..bb55991d9366a 100644 --- a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp @@ -22,6 +22,10 @@ using namespace clang; using namespace clang::CIRGen; struct clang::CIRGen::CGCoroData { + // What is the current await expression kind and how many + // await/yield expressions were encountered so far. + // These are used to generate pretty labels for await expressions in LLVM IR. + cir::AwaitKind currentAwaitKind = cir::AwaitKind::Init; // Stores the __builtin_coro_id emitted in the function so that we can supply // it as the first argument to other builtins. cir::CallOp coroId = nullptr; @@ -249,7 +253,114 @@ CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) { emitAnyExprToMem(s.getReturnValue(), returnValue, s.getReturnValue()->getType().getQualifiers(), /*isInit*/ true); + + assert(!cir::MissingFeatures::ehCleanupScope()); + // FIXME(cir): EHStack.pushCleanup<CallCoroEnd>(EHCleanup); + curCoro.data->currentAwaitKind = cir::AwaitKind::Init; + if (emitStmt(s.getInitSuspendStmt(), /*useCurrentScope=*/true).failed()) + return mlir::failure(); assert(!cir::MissingFeatures::emitBodyAndFallthrough()); } return mlir::success(); } +// Given a suspend expression which roughly looks like: +// +// auto && x = CommonExpr(); +// if (!x.await_ready()) { +// x.await_suspend(...); (*) +// } +// x.await_resume(); +// +// where the result of the entire expression is the result of x.await_resume() +// +// (*) If x.await_suspend return type is bool, it allows to veto a suspend: +// if (x.await_suspend(...)) +// llvm_coro_suspend(); +// +// This is more higher level than LLVM codegen, for that one see llvm's +// docs/Coroutines.rst for more details. +namespace { +struct LValueOrRValue { + LValue lv; + RValue rv; +}; +} // namespace + +static LValueOrRValue +emitSuspendExpression(CIRGenFunction &cgf, CGCoroData &coro, + CoroutineSuspendExpr const &s, cir::AwaitKind kind, + AggValueSlot aggSlot, bool ignoreResult, + mlir::Block *scopeParentBlock, + mlir::Value &tmpResumeRValAddr, bool forLValue) { + mlir::LogicalResult awaitBuild = mlir::success(); + LValueOrRValue awaitRes; + + CIRGenFunction::OpaqueValueMapping binder = + CIRGenFunction::OpaqueValueMapping(cgf, s.getOpaqueValue()); + CIRGenBuilderTy &builder = cgf.getBuilder(); + [[maybe_unused]] cir::AwaitOp awaitOp = cir::AwaitOp::create( + builder, cgf.getLoc(s.getSourceRange()), kind, + /*readyBuilder=*/ + [&](mlir::OpBuilder &b, mlir::Location loc) { + builder.createCondition( + cgf.createDummyValue(loc, cgf.getContext().BoolTy)); + }, + /*suspendBuilder=*/ + [&](mlir::OpBuilder &b, mlir::Location loc) { + cir::YieldOp::create(builder, loc); + }, + /*resumeBuilder=*/ + [&](mlir::OpBuilder &b, mlir::Location loc) { + cir::YieldOp::create(builder, loc); + }); + + assert(awaitBuild.succeeded() && "Should know how to codegen"); + return awaitRes; +} + +static RValue emitSuspendExpr(CIRGenFunction &cgf, + const CoroutineSuspendExpr &e, + cir::AwaitKind kind, AggValueSlot aggSlot, + bool ignoreResult) { + RValue rval; + mlir::Location scopeLoc = cgf.getLoc(e.getSourceRange()); + + // Since we model suspend / resume as an inner region, we must store + // resume scalar results in a tmp alloca, and load it after we build the + // suspend expression. An alternative way to do this would be to make + // every region return a value when promise.return_value() is used, but + // it's a bit awkward given that resume is the only region that actually + // returns a value. + mlir::Block *currEntryBlock = cgf.curLexScope->getEntryBlock(); + [[maybe_unused]] mlir::Value tmpResumeRValAddr; + + // No need to explicitly wrap this into a scope since the AST already uses a + // ExprWithCleanups, which will wrap this into a cir.scope anyways. + rval = emitSuspendExpression(cgf, *cgf.curCoro.data, e, kind, aggSlot, + ignoreResult, currEntryBlock, tmpResumeRValAddr, + /*forLValue*/ false) + .rv; + + if (ignoreResult || rval.isIgnored()) + return rval; + + if (rval.isScalar()) { + rval = RValue::get(cir::LoadOp::create(cgf.getBuilder(), scopeLoc, + rval.getValue().getType(), + tmpResumeRValAddr)); + } else if (rval.isAggregate()) { + // This is probably already handled via AggSlot, remove this assertion + // once we have a testcase and prove all pieces work. + cgf.cgm.errorNYI("emitSuspendExpr Aggregate"); + } else { // complex + cgf.cgm.errorNYI("emitSuspendExpr Complex"); + } + return rval; +} + +RValue CIRGenFunction::emitCoawaitExpr(const CoawaitExpr &e, + AggValueSlot aggSlot, + bool ignoreResult) { + return emitSuspendExpr(*this, e, curCoro.data->currentAwaitKind, aggSlot, + ignoreResult); +} diff --git a/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp b/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp index 4461875fcf678..8925d26bff19d 100644 --- a/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp @@ -154,6 +154,10 @@ class ScalarExprEmitter : public StmtVisitor<ScalarExprEmitter, mlir::Value> { return cgf.emitLoadOfLValue(lv, e->getExprLoc()).getValue(); } + mlir::Value VisitCoawaitExpr(CoawaitExpr *s) { + return cgf.emitCoawaitExpr(*s).getValue(); + } + mlir::Value emitLoadOfLValue(LValue lv, SourceLocation loc) { return cgf.emitLoadOfLValue(lv, loc).getValue(); } diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.h b/clang/lib/CIR/CodeGen/CIRGenFunction.h index b71a28c54dbef..9a4e413fe8f3a 100644 --- a/clang/lib/CIR/CodeGen/CIRGenFunction.h +++ b/clang/lib/CIR/CodeGen/CIRGenFunction.h @@ -1557,6 +1557,9 @@ class CIRGenFunction : public CIRGenTypeCache { void emitForwardingCallToLambda(const CXXMethodDecl *lambdaCallOperator, CallArgList &callArgs); + RValue emitCoawaitExpr(const CoawaitExpr &e, + AggValueSlot aggSlot = AggValueSlot::ignored(), + bool ignoreResult = false); /// Emit the computation of the specified expression of complex type, /// returning the result. mlir::Value emitComplexExpr(const Expr *e); diff --git a/clang/lib/CIR/CodeGen/CIRGenValue.h b/clang/lib/CIR/CodeGen/CIRGenValue.h index ab245a771d72c..e310910457a42 100644 --- a/clang/lib/CIR/CodeGen/CIRGenValue.h +++ b/clang/lib/CIR/CodeGen/CIRGenValue.h @@ -49,6 +49,7 @@ class RValue { bool isScalar() const { return flavor == Scalar; } bool isComplex() const { return flavor == Complex; } bool isAggregate() const { return flavor == Aggregate; } + bool isIgnored() const { return isScalar() && !getValue(); } bool isVolatileQualified() const { return isVolatile; } diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp index 7ba03ce40140c..80b13a4b7fc2e 100644 --- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp +++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp @@ -289,7 +289,10 @@ void cir::ConditionOp::getSuccessorRegions( regions.emplace_back(getOperation(), loopOp->getResults()); } - assert(!cir::MissingFeatures::awaitOp()); + // Parent is an await: condition may branch to resume or suspend regions. + auto await = cast<AwaitOp>(getOperation()->getParentOp()); + regions.emplace_back(&await.getResume(), await.getResume().getArguments()); + regions.emplace_back(&await.getSuspend(), await.getSuspend().getArguments()); } MutableOperandRange @@ -299,8 +302,7 @@ cir::ConditionOp::getMutableSuccessorOperands(RegionSuccessor point) { } LogicalResult cir::ConditionOp::verify() { - assert(!cir::MissingFeatures::awaitOp()); - if (!isa<LoopOpInterface>(getOperation()->getParentOp())) + if (!isa<LoopOpInterface, AwaitOp>(getOperation()->getParentOp())) return emitOpError("condition must be within a conditional region"); return success(); } @@ -1900,6 +1902,19 @@ void cir::FuncOp::print(OpAsmPrinter &p) { mlir::LogicalResult cir::FuncOp::verify() { + if (!isDeclaration() && getCoroutine()) { + bool foundAwait = false; + this->walk([&](Operation *op) { + if (auto await = dyn_cast<AwaitOp>(op)) { + foundAwait = true; + return; + } + }); + if (!foundAwait) + return emitOpError() + << "coroutine body must use at least one cir.await op"; + } + llvm::SmallSet<llvm::StringRef, 16> labels; llvm::SmallSet<llvm::StringRef, 16> gotos; @@ -2116,6 +2131,65 @@ OpFoldResult cir::UnaryOp::fold(FoldAdaptor adaptor) { return {}; } +//===----------------------------------------------------------------------===// +// AwaitOp +//===----------------------------------------------------------------------===// + +void cir::AwaitOp::build(OpBuilder &builder, OperationState &result, + cir::AwaitKind kind, BuilderCallbackRef readyBuilder, + BuilderCallbackRef suspendBuilder, + BuilderCallbackRef resumeBuilder) { + result.addAttribute(getKindAttrName(result.name), + cir::AwaitKindAttr::get(builder.getContext(), kind)); + { + OpBuilder::InsertionGuard guard(builder); + Region *readyRegion = result.addRegion(); + builder.createBlock(readyRegion); + readyBuilder(builder, result.location); + } + + { + OpBuilder::InsertionGuard guard(builder); + Region *suspendRegion = result.addRegion(); + builder.createBlock(suspendRegion); + suspendBuilder(builder, result.location); + } + + { + OpBuilder::InsertionGuard guard(builder); + Region *resumeRegion = result.addRegion(); + builder.createBlock(resumeRegion); + resumeBuilder(builder, result.location); + } +} + +/// Given the region at `index`, or the parent operation if `index` is None, +/// return the successor regions. These are the regions that may be selected +/// during the flow of control. `operands` is a set of optional attributes +/// that correspond to a constant value for each operand, or null if that +/// operand is not a constant. +void cir::AwaitOp::getSuccessorRegions( + mlir::RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) { + // If any index all the underlying regions branch back to the parent + // operation. + if (!point.isParent()) { + regions.push_back( + RegionSuccessor(getOperation(), getOperation()->getResults())); + return; + } + + // FIXME: we want to look at cond region for getting more accurate results + // if the other regions will get a chance to execute. + regions.push_back(RegionSuccessor(&this->getReady())); + regions.push_back(RegionSuccessor(&this->getSuspend())); + regions.push_back(RegionSuccessor(&this->getResume())); +} + +LogicalResult cir::AwaitOp::verify() { + if (!isa<ConditionOp>(this->getReady().back().getTerminator())) + return emitOpError("ready region must end with cir.condition"); + return success(); +} //===----------------------------------------------------------------------===// // CopyOp Definitions diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp index b4afed7019417..6848876c88fea 100644 --- a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp +++ b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp @@ -3778,6 +3778,12 @@ mlir::LogicalResult CIRToLLVMVAArgOpLowering::matchAndRewrite( return mlir::success(); } +mlir::LogicalResult CIRToLLVMAwaitOpLowering::matchAndRewrite( + cir::AwaitOp op, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const { + return mlir::failure(); +} + std::unique_ptr<mlir::Pass> createConvertCIRToLLVMPass() { return std::make_unique<ConvertCIRToLLVMPass>(); } diff --git a/clang/lib/CodeGen/CGValue.h b/clang/lib/CodeGen/CGValue.h index c4ec8d207d2e3..6b381b59e71cd 100644 --- a/clang/lib/CodeGen/CGValue.h +++ b/clang/lib/CodeGen/CGValue.h @@ -64,6 +64,7 @@ class RValue { bool isScalar() const { return Flavor == Scalar; } bool isComplex() const { return Flavor == Complex; } bool isAggregate() const { return Flavor == Aggregate; } + bool isIgnored() const { return isScalar() && !getScalarVal(); } bool isVolatileQualified() const { return IsVolatile; } diff --git a/clang/test/CIR/CodeGen/coro-task.cpp b/clang/test/CIR/CodeGen/coro-task.cpp index 5738c815909ea..01e0786fbda71 100644 --- a/clang/test/CIR/CodeGen/coro-task.cpp +++ b/clang/test/CIR/CodeGen/coro-task.cpp @@ -111,6 +111,8 @@ co_invoke_fn co_invoke; // CIR-DAG: ![[VoidPromisse:.*]] = !cir.record<struct "folly::coro::Task<void>::promise_type" padded {!u8i}> // CIR-DAG: ![[IntPromisse:.*]] = !cir.record<struct "folly::coro::Task<int>::promise_type" padded {!u8i}> // CIR-DAG: ![[StdString:.*]] = !cir.record<struct "std::string" padded {!u8i}> +// CIR-DAG: ![[SuspendAlways:.*]] = !cir.record<struct "std::suspend_always" padded {!u8i}> + // CIR: module {{.*}} { // CIR-NEXT: cir.global external @_ZN5folly4coro9co_invokeE = #cir.zero : !rec_folly3A3Acoro3A3Aco_invoke_fn @@ -153,6 +155,33 @@ VoidTask silly_task() { // CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type17get_return_objectEv(%[[VoidPromisseAddr]]) nothrow : {{.*}} -> ![[VoidTask]] // CIR: cir.store{{.*}} %[[RetObj]], %[[VoidTaskAddr]] : ![[VoidTask]] +// Start a new scope for the actual codegen for co_await, create temporary allocas for +// holding coroutine handle and the suspend_always struct. + +// CIR: cir.scope { +// CIR: %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64} + +// Effectively execute `coawait promise_type::initial_suspend()` by calling initial_suspend() and getting +// the suspend_always struct to use for cir.await. Note that we return by-value since we defer ABI lowering +// to later passes, same is done elsewhere. + +// CIR: %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type15initial_suspendEv(%[[VoidPromisseAddr]]) +// CIR: cir.store{{.*}} %[[Tmp0:.*]], %[[SuspendAlwaysAddr]] + +// +// Here we start mapping co_await to cir.await. +// + +// First regions `ready` has a special cir.yield code to veto suspension. + +// CIR: cir.await(init, ready : { +// CIR: cir.condition({{.*}}) +// CIR: }, suspend : { +// CIR: cir.yield +// CIR: }, resume : { +// CIR: cir.yield +// CIR: },) +// CIR: } folly::coro::Task<int> byRef(const std::string& s) { co_return s.size(); @@ -172,3 +201,13 @@ folly::coro::Task<int> byRef(const std::string& s) { // CIR: cir.store {{.*}} %[[LOAD]], %[[AllocaFnUse]] : !cir.ptr<![[StdString]]>, !cir.ptr<!cir.ptr<![[StdString]]>> // CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type17get_return_objectEv(%4) nothrow : {{.*}} -> ![[IntTask]] // CIR: cir.store {{.*}} %[[RetObj]], %[[IntTaskAddr]] : ![[IntTask]] +// CIR: cir.scope { +// CIR: %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64} +// CIR: %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type15initial_suspendEv(%[[IntPromisseAddr]]) +// CIR: cir.await(init, ready : { +// CIR: cir.condition({{.*}}) +// CIR: }, suspend : { +// CIR: cir.yield +// CIR: }, resume : { +// CIR: cir.yield +// CIR: },) diff --git a/clang/test/CIR/IR/await.cir b/clang/test/CIR/IR/await.cir new file mode 100644 index 0000000000000..c1fb0d6d7c57c --- /dev/null +++ b/clang/test/CIR/IR/await.cir @@ -0,0 +1,21 @@ +// RUN: cir-opt %s --verify-roundtrip | FileCheck %s + +cir.func coroutine @checkPrintParse(%arg0 : !cir.bool) { + cir.await(user, ready : { + cir.condition(%arg0) + }, suspend : { + cir.yield + }, resume : { + cir.yield + },) + cir.return +} + +// CHECK: cir.func coroutine @checkPrintParse +// CHECK: cir.await(user, ready : { +// CH... [truncated] |
Member
| @llvm/pr-subscribers-clang Author: None (Andres-Salamanca) ChangesThis PR upstreams Patch is 21.40 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/168133.diff 12 Files Affected:
diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td index 16258513239d9..832f79607a101 100644 --- a/clang/include/clang/CIR/Dialect/IR/CIROps.td +++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td @@ -790,8 +790,8 @@ def CIR_ConditionOp : CIR_Op<"condition", [ //===----------------------------------------------------------------------===// defvar CIR_YieldableScopes = [ - "ArrayCtor", "ArrayDtor", "CaseOp", "DoWhileOp", "ForOp", "GlobalOp", "IfOp", - "ScopeOp", "SwitchOp", "TernaryOp", "WhileOp", "TryOp" + "ArrayCtor", "ArrayDtor", "AwaitOp", "CaseOp", "DoWhileOp", "ForOp", + "GlobalOp", "IfOp", "ScopeOp", "SwitchOp", "TernaryOp", "WhileOp", "TryOp" ]; def CIR_YieldOp : CIR_Op<"yield", [ @@ -2707,6 +2707,102 @@ def CIR_CallOp : CIR_CallOpBase<"call", [NoRegionArguments]> { ]; } +//===----------------------------------------------------------------------===// +// AwaitOp +//===----------------------------------------------------------------------===// + +def CIR_AwaitKind : CIR_I32EnumAttr<"AwaitKind", "await kind", [ + I32EnumAttrCase<"Init", 0, "init">, + I32EnumAttrCase<"User", 1, "user">, + I32EnumAttrCase<"Yield", 2, "yield">, + I32EnumAttrCase<"Final", 3, "final"> +]>; + +def CIR_AwaitOp : CIR_Op<"await",[ + DeclareOpInterfaceMethods<RegionBranchOpInterface>, + RecursivelySpeculatable, NoRegionArguments +]> { + let summary = "Wraps C++ co_await implicit logic"; + let description = [{ + The under the hood effect of using C++ `co_await expr` roughly + translates to: + + ```c++ + // co_await expr; + + auto &&x = CommonExpr(); + if (!x.await_ready()) { + ... + x.await_suspend(...); + ... + } + x.await_resume(); + ``` + + `cir.await` represents this logic by using 3 regions: + - ready: covers veto power from x.await_ready() + - suspend: wraps actual x.await_suspend() logic + - resume: handles x.await_resume() + + Breaking this up in regions allow individual scrutiny of conditions + which might lead to folding some of them out. Lowerings coming out + of CIR, e.g. LLVM, should use the `suspend` region to track more + lower level codegen (e.g. intrinsic emission for coro.save/coro.suspend). + + There are also 4 flavors of `cir.await` available: + - `init`: compiler generated initial suspend via implicit `co_await`. + - `user`: also known as normal, representing user written co_await's. + - `yield`: user written `co_yield` expressions. + - `final`: compiler generated final suspend via implicit `co_await`. + + From the C++ snippet we get: + + ```mlir + cir.scope { + ... // auto &&x = CommonExpr(); + cir.await(user, ready : { + ... // x.await_ready() + }, suspend : { + ... // x.await_suspend() + }, resume : { + ... // x.await_resume() + }) + } + ``` + + Note that resulution of the common expression is assumed to happen + as part of the enclosing await scope. + }]; + + let arguments = (ins CIR_AwaitKind:$kind); + let regions = (region SizedRegion<1>:$ready, + SizedRegion<1>:$suspend, + SizedRegion<1>:$resume); + let assemblyFormat = [{ + `(` $kind `,` + `ready` `:` $ready `,` + `suspend` `:` $suspend `,` + `resume` `:` $resume `,` + `)` + attr-dict + }]; + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<(ins + "cir::AwaitKind":$kind, + CArg<"BuilderCallbackRef", + "nullptr">:$readyBuilder, + CArg<"BuilderCallbackRef", + "nullptr">:$suspendBuilder, + CArg<"BuilderCallbackRef", + "nullptr">:$resumeBuilder + )> + ]; + + let hasVerifier = 1; +} + //===----------------------------------------------------------------------===// // CopyOp //===----------------------------------------------------------------------===// diff --git a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp index 05fb1aedcbf4a..bb55991d9366a 100644 --- a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp @@ -22,6 +22,10 @@ using namespace clang; using namespace clang::CIRGen; struct clang::CIRGen::CGCoroData { + // What is the current await expression kind and how many + // await/yield expressions were encountered so far. + // These are used to generate pretty labels for await expressions in LLVM IR. + cir::AwaitKind currentAwaitKind = cir::AwaitKind::Init; // Stores the __builtin_coro_id emitted in the function so that we can supply // it as the first argument to other builtins. cir::CallOp coroId = nullptr; @@ -249,7 +253,114 @@ CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) { emitAnyExprToMem(s.getReturnValue(), returnValue, s.getReturnValue()->getType().getQualifiers(), /*isInit*/ true); + + assert(!cir::MissingFeatures::ehCleanupScope()); + // FIXME(cir): EHStack.pushCleanup<CallCoroEnd>(EHCleanup); + curCoro.data->currentAwaitKind = cir::AwaitKind::Init; + if (emitStmt(s.getInitSuspendStmt(), /*useCurrentScope=*/true).failed()) + return mlir::failure(); assert(!cir::MissingFeatures::emitBodyAndFallthrough()); } return mlir::success(); } +// Given a suspend expression which roughly looks like: +// +// auto && x = CommonExpr(); +// if (!x.await_ready()) { +// x.await_suspend(...); (*) +// } +// x.await_resume(); +// +// where the result of the entire expression is the result of x.await_resume() +// +// (*) If x.await_suspend return type is bool, it allows to veto a suspend: +// if (x.await_suspend(...)) +// llvm_coro_suspend(); +// +// This is more higher level than LLVM codegen, for that one see llvm's +// docs/Coroutines.rst for more details. +namespace { +struct LValueOrRValue { + LValue lv; + RValue rv; +}; +} // namespace + +static LValueOrRValue +emitSuspendExpression(CIRGenFunction &cgf, CGCoroData &coro, + CoroutineSuspendExpr const &s, cir::AwaitKind kind, + AggValueSlot aggSlot, bool ignoreResult, + mlir::Block *scopeParentBlock, + mlir::Value &tmpResumeRValAddr, bool forLValue) { + mlir::LogicalResult awaitBuild = mlir::success(); + LValueOrRValue awaitRes; + + CIRGenFunction::OpaqueValueMapping binder = + CIRGenFunction::OpaqueValueMapping(cgf, s.getOpaqueValue()); + CIRGenBuilderTy &builder = cgf.getBuilder(); + [[maybe_unused]] cir::AwaitOp awaitOp = cir::AwaitOp::create( + builder, cgf.getLoc(s.getSourceRange()), kind, + /*readyBuilder=*/ + [&](mlir::OpBuilder &b, mlir::Location loc) { + builder.createCondition( + cgf.createDummyValue(loc, cgf.getContext().BoolTy)); + }, + /*suspendBuilder=*/ + [&](mlir::OpBuilder &b, mlir::Location loc) { + cir::YieldOp::create(builder, loc); + }, + /*resumeBuilder=*/ + [&](mlir::OpBuilder &b, mlir::Location loc) { + cir::YieldOp::create(builder, loc); + }); + + assert(awaitBuild.succeeded() && "Should know how to codegen"); + return awaitRes; +} + +static RValue emitSuspendExpr(CIRGenFunction &cgf, + const CoroutineSuspendExpr &e, + cir::AwaitKind kind, AggValueSlot aggSlot, + bool ignoreResult) { + RValue rval; + mlir::Location scopeLoc = cgf.getLoc(e.getSourceRange()); + + // Since we model suspend / resume as an inner region, we must store + // resume scalar results in a tmp alloca, and load it after we build the + // suspend expression. An alternative way to do this would be to make + // every region return a value when promise.return_value() is used, but + // it's a bit awkward given that resume is the only region that actually + // returns a value. + mlir::Block *currEntryBlock = cgf.curLexScope->getEntryBlock(); + [[maybe_unused]] mlir::Value tmpResumeRValAddr; + + // No need to explicitly wrap this into a scope since the AST already uses a + // ExprWithCleanups, which will wrap this into a cir.scope anyways. + rval = emitSuspendExpression(cgf, *cgf.curCoro.data, e, kind, aggSlot, + ignoreResult, currEntryBlock, tmpResumeRValAddr, + /*forLValue*/ false) + .rv; + + if (ignoreResult || rval.isIgnored()) + return rval; + + if (rval.isScalar()) { + rval = RValue::get(cir::LoadOp::create(cgf.getBuilder(), scopeLoc, + rval.getValue().getType(), + tmpResumeRValAddr)); + } else if (rval.isAggregate()) { + // This is probably already handled via AggSlot, remove this assertion + // once we have a testcase and prove all pieces work. + cgf.cgm.errorNYI("emitSuspendExpr Aggregate"); + } else { // complex + cgf.cgm.errorNYI("emitSuspendExpr Complex"); + } + return rval; +} + +RValue CIRGenFunction::emitCoawaitExpr(const CoawaitExpr &e, + AggValueSlot aggSlot, + bool ignoreResult) { + return emitSuspendExpr(*this, e, curCoro.data->currentAwaitKind, aggSlot, + ignoreResult); +} diff --git a/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp b/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp index 4461875fcf678..8925d26bff19d 100644 --- a/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp @@ -154,6 +154,10 @@ class ScalarExprEmitter : public StmtVisitor<ScalarExprEmitter, mlir::Value> { return cgf.emitLoadOfLValue(lv, e->getExprLoc()).getValue(); } + mlir::Value VisitCoawaitExpr(CoawaitExpr *s) { + return cgf.emitCoawaitExpr(*s).getValue(); + } + mlir::Value emitLoadOfLValue(LValue lv, SourceLocation loc) { return cgf.emitLoadOfLValue(lv, loc).getValue(); } diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.h b/clang/lib/CIR/CodeGen/CIRGenFunction.h index b71a28c54dbef..9a4e413fe8f3a 100644 --- a/clang/lib/CIR/CodeGen/CIRGenFunction.h +++ b/clang/lib/CIR/CodeGen/CIRGenFunction.h @@ -1557,6 +1557,9 @@ class CIRGenFunction : public CIRGenTypeCache { void emitForwardingCallToLambda(const CXXMethodDecl *lambdaCallOperator, CallArgList &callArgs); + RValue emitCoawaitExpr(const CoawaitExpr &e, + AggValueSlot aggSlot = AggValueSlot::ignored(), + bool ignoreResult = false); /// Emit the computation of the specified expression of complex type, /// returning the result. mlir::Value emitComplexExpr(const Expr *e); diff --git a/clang/lib/CIR/CodeGen/CIRGenValue.h b/clang/lib/CIR/CodeGen/CIRGenValue.h index ab245a771d72c..e310910457a42 100644 --- a/clang/lib/CIR/CodeGen/CIRGenValue.h +++ b/clang/lib/CIR/CodeGen/CIRGenValue.h @@ -49,6 +49,7 @@ class RValue { bool isScalar() const { return flavor == Scalar; } bool isComplex() const { return flavor == Complex; } bool isAggregate() const { return flavor == Aggregate; } + bool isIgnored() const { return isScalar() && !getValue(); } bool isVolatileQualified() const { return isVolatile; } diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp index 7ba03ce40140c..80b13a4b7fc2e 100644 --- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp +++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp @@ -289,7 +289,10 @@ void cir::ConditionOp::getSuccessorRegions( regions.emplace_back(getOperation(), loopOp->getResults()); } - assert(!cir::MissingFeatures::awaitOp()); + // Parent is an await: condition may branch to resume or suspend regions. + auto await = cast<AwaitOp>(getOperation()->getParentOp()); + regions.emplace_back(&await.getResume(), await.getResume().getArguments()); + regions.emplace_back(&await.getSuspend(), await.getSuspend().getArguments()); } MutableOperandRange @@ -299,8 +302,7 @@ cir::ConditionOp::getMutableSuccessorOperands(RegionSuccessor point) { } LogicalResult cir::ConditionOp::verify() { - assert(!cir::MissingFeatures::awaitOp()); - if (!isa<LoopOpInterface>(getOperation()->getParentOp())) + if (!isa<LoopOpInterface, AwaitOp>(getOperation()->getParentOp())) return emitOpError("condition must be within a conditional region"); return success(); } @@ -1900,6 +1902,19 @@ void cir::FuncOp::print(OpAsmPrinter &p) { mlir::LogicalResult cir::FuncOp::verify() { + if (!isDeclaration() && getCoroutine()) { + bool foundAwait = false; + this->walk([&](Operation *op) { + if (auto await = dyn_cast<AwaitOp>(op)) { + foundAwait = true; + return; + } + }); + if (!foundAwait) + return emitOpError() + << "coroutine body must use at least one cir.await op"; + } + llvm::SmallSet<llvm::StringRef, 16> labels; llvm::SmallSet<llvm::StringRef, 16> gotos; @@ -2116,6 +2131,65 @@ OpFoldResult cir::UnaryOp::fold(FoldAdaptor adaptor) { return {}; } +//===----------------------------------------------------------------------===// +// AwaitOp +//===----------------------------------------------------------------------===// + +void cir::AwaitOp::build(OpBuilder &builder, OperationState &result, + cir::AwaitKind kind, BuilderCallbackRef readyBuilder, + BuilderCallbackRef suspendBuilder, + BuilderCallbackRef resumeBuilder) { + result.addAttribute(getKindAttrName(result.name), + cir::AwaitKindAttr::get(builder.getContext(), kind)); + { + OpBuilder::InsertionGuard guard(builder); + Region *readyRegion = result.addRegion(); + builder.createBlock(readyRegion); + readyBuilder(builder, result.location); + } + + { + OpBuilder::InsertionGuard guard(builder); + Region *suspendRegion = result.addRegion(); + builder.createBlock(suspendRegion); + suspendBuilder(builder, result.location); + } + + { + OpBuilder::InsertionGuard guard(builder); + Region *resumeRegion = result.addRegion(); + builder.createBlock(resumeRegion); + resumeBuilder(builder, result.location); + } +} + +/// Given the region at `index`, or the parent operation if `index` is None, +/// return the successor regions. These are the regions that may be selected +/// during the flow of control. `operands` is a set of optional attributes +/// that correspond to a constant value for each operand, or null if that +/// operand is not a constant. +void cir::AwaitOp::getSuccessorRegions( + mlir::RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) { + // If any index all the underlying regions branch back to the parent + // operation. + if (!point.isParent()) { + regions.push_back( + RegionSuccessor(getOperation(), getOperation()->getResults())); + return; + } + + // FIXME: we want to look at cond region for getting more accurate results + // if the other regions will get a chance to execute. + regions.push_back(RegionSuccessor(&this->getReady())); + regions.push_back(RegionSuccessor(&this->getSuspend())); + regions.push_back(RegionSuccessor(&this->getResume())); +} + +LogicalResult cir::AwaitOp::verify() { + if (!isa<ConditionOp>(this->getReady().back().getTerminator())) + return emitOpError("ready region must end with cir.condition"); + return success(); +} //===----------------------------------------------------------------------===// // CopyOp Definitions diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp index b4afed7019417..6848876c88fea 100644 --- a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp +++ b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp @@ -3778,6 +3778,12 @@ mlir::LogicalResult CIRToLLVMVAArgOpLowering::matchAndRewrite( return mlir::success(); } +mlir::LogicalResult CIRToLLVMAwaitOpLowering::matchAndRewrite( + cir::AwaitOp op, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const { + return mlir::failure(); +} + std::unique_ptr<mlir::Pass> createConvertCIRToLLVMPass() { return std::make_unique<ConvertCIRToLLVMPass>(); } diff --git a/clang/lib/CodeGen/CGValue.h b/clang/lib/CodeGen/CGValue.h index c4ec8d207d2e3..6b381b59e71cd 100644 --- a/clang/lib/CodeGen/CGValue.h +++ b/clang/lib/CodeGen/CGValue.h @@ -64,6 +64,7 @@ class RValue { bool isScalar() const { return Flavor == Scalar; } bool isComplex() const { return Flavor == Complex; } bool isAggregate() const { return Flavor == Aggregate; } + bool isIgnored() const { return isScalar() && !getScalarVal(); } bool isVolatileQualified() const { return IsVolatile; } diff --git a/clang/test/CIR/CodeGen/coro-task.cpp b/clang/test/CIR/CodeGen/coro-task.cpp index 5738c815909ea..01e0786fbda71 100644 --- a/clang/test/CIR/CodeGen/coro-task.cpp +++ b/clang/test/CIR/CodeGen/coro-task.cpp @@ -111,6 +111,8 @@ co_invoke_fn co_invoke; // CIR-DAG: ![[VoidPromisse:.*]] = !cir.record<struct "folly::coro::Task<void>::promise_type" padded {!u8i}> // CIR-DAG: ![[IntPromisse:.*]] = !cir.record<struct "folly::coro::Task<int>::promise_type" padded {!u8i}> // CIR-DAG: ![[StdString:.*]] = !cir.record<struct "std::string" padded {!u8i}> +// CIR-DAG: ![[SuspendAlways:.*]] = !cir.record<struct "std::suspend_always" padded {!u8i}> + // CIR: module {{.*}} { // CIR-NEXT: cir.global external @_ZN5folly4coro9co_invokeE = #cir.zero : !rec_folly3A3Acoro3A3Aco_invoke_fn @@ -153,6 +155,33 @@ VoidTask silly_task() { // CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type17get_return_objectEv(%[[VoidPromisseAddr]]) nothrow : {{.*}} -> ![[VoidTask]] // CIR: cir.store{{.*}} %[[RetObj]], %[[VoidTaskAddr]] : ![[VoidTask]] +// Start a new scope for the actual codegen for co_await, create temporary allocas for +// holding coroutine handle and the suspend_always struct. + +// CIR: cir.scope { +// CIR: %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64} + +// Effectively execute `coawait promise_type::initial_suspend()` by calling initial_suspend() and getting +// the suspend_always struct to use for cir.await. Note that we return by-value since we defer ABI lowering +// to later passes, same is done elsewhere. + +// CIR: %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type15initial_suspendEv(%[[VoidPromisseAddr]]) +// CIR: cir.store{{.*}} %[[Tmp0:.*]], %[[SuspendAlwaysAddr]] + +// +// Here we start mapping co_await to cir.await. +// + +// First regions `ready` has a special cir.yield code to veto suspension. + +// CIR: cir.await(init, ready : { +// CIR: cir.condition({{.*}}) +// CIR: }, suspend : { +// CIR: cir.yield +// CIR: }, resume : { +// CIR: cir.yield +// CIR: },) +// CIR: } folly::coro::Task<int> byRef(const std::string& s) { co_return s.size(); @@ -172,3 +201,13 @@ folly::coro::Task<int> byRef(const std::string& s) { // CIR: cir.store {{.*}} %[[LOAD]], %[[AllocaFnUse]] : !cir.ptr<![[StdString]]>, !cir.ptr<!cir.ptr<![[StdString]]>> // CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type17get_return_objectEv(%4) nothrow : {{.*}} -> ![[IntTask]] // CIR: cir.store {{.*}} %[[RetObj]], %[[IntTaskAddr]] : ![[IntTask]] +// CIR: cir.scope { +// CIR: %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64} +// CIR: %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type15initial_suspendEv(%[[IntPromisseAddr]]) +// CIR: cir.await(init, ready : { +// CIR: cir.condition({{.*}}) +// CIR: }, suspend : { +// CIR: cir.yield +// CIR: }, resume : { +// CIR: cir.yield +// CIR: },) diff --git a/clang/test/CIR/IR/await.cir b/clang/test/CIR/IR/await.cir new file mode 100644 index 0000000000000..c1fb0d6d7c57c --- /dev/null +++ b/clang/test/CIR/IR/await.cir @@ -0,0 +1,21 @@ +// RUN: cir-opt %s --verify-roundtrip | FileCheck %s + +cir.func coroutine @checkPrintParse(%arg0 : !cir.bool) { + cir.await(user, ready : { + cir.condition(%arg0) + }, suspend : { + cir.yield + }, resume : { + cir.yield + },) + cir.return +} + +// CHECK: cir.func coroutine @checkPrintParse +// CHECK: cir.await(user, ready : { +// CH... [truncated] |
Member
bcardosolopes left a comment
bcardosolopes left a comment There was a problem hiding this comment.
Looks mostly good, couple comments
andykaylor approved these changes Nov 15, 2025
Contributor
andykaylor left a comment
andykaylor left a comment There was a problem hiding this comment.
Looks good to me, just a few nits.
bcardosolopes approved these changes Nov 18, 2025
Andres-Salamanca force-pushed the cir-await branch from 5ddc04f to f73fd50 Compare 
🐧 Linux x64 Test Results
|
aadeshps-mcw pushed a commit to aadeshps-mcw/llvm-project that referenced this pull request Nov 26, 2025
This PR upstreams `cir.await` and adds initial codegen for emitting a skeleton of the ready, suspend, and resume branches. Codegen for these branches is left for a future PR. It also adds a test for the invalid case where a `cir.func` is marked as a coroutine but does not contain a `cir.await` op in its body.
Priyanshu3820 pushed a commit to Priyanshu3820/llvm-project that referenced this pull request Nov 26, 2025
This PR upstreams `cir.await` and adds initial codegen for emitting a skeleton of the ready, suspend, and resume branches. Codegen for these branches is left for a future PR. It also adds a test for the invalid case where a `cir.func` is marked as a coroutine but does not contain a `cir.await` op in its body.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.
This PR upstreams
cir.awaitand adds initial codegen for emitting a skeleton of the ready, suspend, and resume branches. Codegen for these branches is left for a future PR. It also adds a test for the invalid case where acir.funcis marked as a coroutine but does not contain acir.awaitop in its body.