Skip to content

Commit f06bc7b

Browse files
Lancernlanza
Lancern
authored and
lanza
committed
[CIR][CIRGen] Add CIRGen for scalar co_yield expression (#761)
This PR adds CIRGen for scalar `co_yield` expressions.
1 parent 4993f97 commit f06bc7b

File tree

5 files changed

+82
-12
lines changed

5 files changed

+82
-12
lines changed

‎clang/include/clang/CIR/Dialect/IR/CIROps.td

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3156,12 +3156,13 @@ def TryCallOp : CIR_CallOp<"try_call",
31563156

31573157
def AK_Initial : I32EnumAttrCase<"init", 1>;
31583158
def AK_User : I32EnumAttrCase<"user", 2>;
3159-
def AK_Final : I32EnumAttrCase<"final", 3>;
3159+
def AK_Yield : I32EnumAttrCase<"yield", 3>;
3160+
def AK_Final : I32EnumAttrCase<"final", 4>;
31603161

31613162
def AwaitKind : I32EnumAttr<
31623163
"AwaitKind",
31633164
"await kind",
3164-
[AK_Initial, AK_User, AK_Final]> {
3165+
[AK_Initial, AK_User, AK_Yield, AK_Final]> {
31653166
let cppNamespace = "::mlir::cir";
31663167
}
31673168

@@ -3195,9 +3196,10 @@ def AwaitOp : CIR_Op<"await",
31953196
of CIR, e.g. LLVM, should use the `suspend` region to track more
31963197
lower level codegen (e.g. intrinsic emission for coro.save/coro.suspend).
31973198

3198-
There are also 3 flavors of `cir.await` available:
3199+
There are also 4 flavors of `cir.await` available:
31993200
- `init`: compiler generated initial suspend via implicit `co_await`.
32003201
- `user`: also known as normal, representing user written co_await's.
3202+
- `yield`: user written `co_yield` expressions.
32013203
- `final`: compiler generated final suspend via implicit `co_await`.
32023204

32033205
From the C++ snippet we get:

‎clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -492,25 +492,25 @@ buildSuspendExpression(CIRGenFunction &CGF, CGCoroData &Coro,
492492
return awaitRes;
493493
}
494494

495-
RValue CIRGenFunction::buildCoawaitExpr(const CoawaitExpr &E,
496-
AggValueSlot aggSlot,
497-
bool ignoreResult) {
495+
static RValue buildSuspendExpr(CIRGenFunction &CGF,
496+
const CoroutineSuspendExpr &E,
497+
mlir::cir::AwaitKind kind, AggValueSlot aggSlot,
498+
bool ignoreResult) {
498499
RValue rval;
499-
auto scopeLoc = getLoc(E.getSourceRange());
500+
auto scopeLoc = CGF.getLoc(E.getSourceRange());
500501

501502
// Since we model suspend / resume as an inner region, we must store
502503
// resume scalar results in a tmp alloca, and load it after we build the
503504
// suspend expression. An alternative way to do this would be to make
504505
// every region return a value when promise.return_value() is used, but
505506
// it's a bit awkward given that resume is the only region that actually
506507
// returns a value.
507-
mlir::Block *currEntryBlock = currLexScope->getEntryBlock();
508+
mlir::Block *currEntryBlock = CGF.currLexScope->getEntryBlock();
508509
[[maybe_unused]] mlir::Value tmpResumeRValAddr;
509510

510511
// No need to explicitly wrap this into a scope since the AST already uses a
511512
// ExprWithCleanups, which will wrap this into a cir.scope anyways.
512-
rval = buildSuspendExpression(*this, *CurCoro.Data, E,
513-
CurCoro.Data->CurrentAwaitKind, aggSlot,
513+
rval = buildSuspendExpression(CGF, *CGF.CurCoro.Data, E, kind, aggSlot,
514514
ignoreResult, currEntryBlock, tmpResumeRValAddr,
515515
/*forLValue*/ false)
516516
.RV;
@@ -519,7 +519,7 @@ RValue CIRGenFunction::buildCoawaitExpr(const CoawaitExpr &E,
519519
return rval;
520520

521521
if (rval.isScalar()) {
522-
rval = RValue::get(builder.create<mlir::cir::LoadOp>(
522+
rval = RValue::get(CGF.getBuilder().create<mlir::cir::LoadOp>(
523523
scopeLoc, rval.getScalarVal().getType(), tmpResumeRValAddr));
524524
} else if (rval.isAggregate()) {
525525
// This is probably already handled via AggSlot, remove this assertion
@@ -531,6 +531,20 @@ RValue CIRGenFunction::buildCoawaitExpr(const CoawaitExpr &E,
531531
return rval;
532532
}
533533

534+
RValue CIRGenFunction::buildCoawaitExpr(const CoawaitExpr &E,
535+
AggValueSlot aggSlot,
536+
bool ignoreResult) {
537+
return buildSuspendExpr(*this, E, CurCoro.Data->CurrentAwaitKind, aggSlot,
538+
ignoreResult);
539+
}
540+
541+
RValue CIRGenFunction::buildCoyieldExpr(const CoyieldExpr &E,
542+
AggValueSlot aggSlot,
543+
bool ignoreResult) {
544+
return buildSuspendExpr(*this, E, mlir::cir::AwaitKind::yield, aggSlot,
545+
ignoreResult);
546+
}
547+
534548
mlir::LogicalResult CIRGenFunction::buildCoreturnStmt(CoreturnStmt const &S) {
535549
++CurCoro.Data->CoreturnCount;
536550
currLexScope->setCoreturn();

‎clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,9 @@ class ScalarExprEmitter : public StmtVisitor<ScalarExprEmitter, mlir::Value> {
160160
mlir::Value VisitCoawaitExpr(CoawaitExpr *S) {
161161
return CGF.buildCoawaitExpr(*S).getScalarVal();
162162
}
163-
mlir::Value VisitCoyieldExpr(CoyieldExpr *S) { llvm_unreachable("NYI"); }
163+
mlir::Value VisitCoyieldExpr(CoyieldExpr *S) {
164+
return CGF.buildCoyieldExpr(*S).getScalarVal();
165+
}
164166
mlir::Value VisitUnaryCoawait(const UnaryOperator *E) {
165167
llvm_unreachable("NYI");
166168
}

‎clang/lib/CIR/CodeGen/CIRGenFunction.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -900,6 +900,9 @@ class CIRGenFunction : public CIRGenTypeCache {
900900
RValue buildCoawaitExpr(const CoawaitExpr &E,
901901
AggValueSlot aggSlot = AggValueSlot::ignored(),
902902
bool ignoreResult = false);
903+
RValue buildCoyieldExpr(const CoyieldExpr &E,
904+
AggValueSlot aggSlot = AggValueSlot::ignored(),
905+
bool ignoreResult = false);
903906
RValue buildCoroutineIntrinsic(const CallExpr *E, unsigned int IID);
904907
RValue buildCoroutineFrame();
905908

‎clang/test/CIR/CodeGen/coro-task.cpp

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -380,3 +380,52 @@ folly::coro::Task<int> go4() {
380380
// CHECK: }, resume : {
381381
// CHECK: },)
382382
// CHECK: }
383+
384+
folly::coro::Task<void> yield();
385+
folly::coro::Task<void> yield1() {
386+
auto t = yield();
387+
co_yield t;
388+
}
389+
390+
// CHECK: cir.func coroutine @_Z6yield1v() -> !ty_22folly3A3Acoro3A3ATask3Cvoid3E22
391+
392+
// CHECK: cir.await(init, ready : {
393+
// CHECK: }, suspend : {
394+
// CHECK: }, resume : {
395+
// CHECK: },)
396+
397+
// CHECK: cir.scope {
398+
// CHECK-NEXT: %[[#SUSPEND_PTR:]] = cir.alloca !ty_22std3A3Asuspend_always22, !cir.ptr<!ty_22std3A3Asuspend_always22>
399+
// CHECK-NEXT: %[[#AWAITER_PTR:]] = cir.alloca !ty_22folly3A3Acoro3A3ATask3Cvoid3E22, !cir.ptr<!ty_22folly3A3Acoro3A3ATask3Cvoid3E22>
400+
// CHECK-NEXT: %[[#CORO_PTR:]] = cir.alloca !ty_22std3A3Acoroutine_handle3Cvoid3E22, !cir.ptr<!ty_22std3A3Acoroutine_handle3Cvoid3E22>
401+
// CHECK-NEXT: %[[#CORO2_PTR:]] = cir.alloca !ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22, !cir.ptr<!ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22>
402+
// CHECK-NEXT: cir.call @_ZN5folly4coro4TaskIvEC1ERKS2_(%[[#AWAITER_PTR]], %{{.+}}) : (!cir.ptr<!ty_22folly3A3Acoro3A3ATask3Cvoid3E22>, !cir.ptr<!ty_22folly3A3Acoro3A3ATask3Cvoid3E22>) -> ()
403+
// CHECK-NEXT: %[[#AWAITER:]] = cir.load %[[#AWAITER_PTR]] : !cir.ptr<!ty_22folly3A3Acoro3A3ATask3Cvoid3E22>, !ty_22folly3A3Acoro3A3ATask3Cvoid3E22
404+
// CHECK-NEXT: %[[#SUSPEND:]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type11yield_valueES2_(%{{.+}}, %[[#AWAITER]]) : (!cir.ptr<!ty_22folly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type22>, !ty_22folly3A3Acoro3A3ATask3Cvoid3E22) -> !ty_22std3A3Asuspend_always22
405+
// CHECK-NEXT: cir.store %[[#SUSPEND]], %[[#SUSPEND_PTR]] : !ty_22std3A3Asuspend_always22, !cir.ptr<!ty_22std3A3Asuspend_always22>
406+
// CHECK-NEXT: cir.await(yield, ready : {
407+
// CHECK-NEXT: %[[#READY:]] = cir.scope {
408+
// CHECK-NEXT: %[[#A:]] = cir.call @_ZNSt14suspend_always11await_readyEv(%[[#SUSPEND_PTR]]) : (!cir.ptr<!ty_22std3A3Asuspend_always22>) -> !cir.bool
409+
// CHECK-NEXT: cir.yield %[[#A]] : !cir.bool
410+
// CHECK-NEXT: } : !cir.bool
411+
// CHECK-NEXT: cir.condition(%[[#READY]])
412+
// CHECK-NEXT: }, suspend : {
413+
// CHECK-NEXT: %[[#CORO2:]] = cir.call @_ZNSt16coroutine_handleIN5folly4coro4TaskIvE12promise_typeEE12from_addressEPv(%9) : (!cir.ptr<!void>) -> !ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22
414+
// CHECK-NEXT: cir.store %[[#CORO2]], %[[#CORO2_PTR]] : !ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22, !cir.ptr<!ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22>
415+
// CHECK-NEXT: %[[#B:]] = cir.load %[[#CORO2_PTR]] : !cir.ptr<!ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22>, !ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22
416+
// CHECK-NEXT: cir.call @_ZNSt16coroutine_handleIvEC1IN5folly4coro4TaskIvE12promise_typeEEES_IT_E(%[[#CORO_PTR]], %[[#B]]) : (!cir.ptr<!ty_22std3A3Acoroutine_handle3Cvoid3E22>, !ty_22std3A3Acoroutine_handle3Cfolly3A3Acoro3A3ATask3Cvoid3E3A3Apromise_type3E22) -> ()
417+
// CHECK-NEXT: %[[#C:]] = cir.load %[[#CORO_PTR]] : !cir.ptr<!ty_22std3A3Acoroutine_handle3Cvoid3E22>, !ty_22std3A3Acoroutine_handle3Cvoid3E22
418+
// CHECK-NEXT: cir.call @_ZNSt14suspend_always13await_suspendESt16coroutine_handleIvE(%[[#SUSPEND_PTR]], %[[#C]]) : (!cir.ptr<!ty_22std3A3Asuspend_always22>, !ty_22std3A3Acoroutine_handle3Cvoid3E22) -> ()
419+
// CHECK-NEXT: cir.yield
420+
// CHECK-NEXT: }, resume : {
421+
// CHECK-NEXT: cir.call @_ZNSt14suspend_always12await_resumeEv(%[[#SUSPEND_PTR]]) : (!cir.ptr<!ty_22std3A3Asuspend_always22>) -> ()
422+
// CHECK-NEXT: cir.yield
423+
// CHECK-NEXT: },)
424+
// CHECK-NEXT: }
425+
426+
// CHECK: cir.await(final, ready : {
427+
// CHECK: }, suspend : {
428+
// CHECK: }, resume : {
429+
// CHECK: },)
430+
431+
// CHECK: }

0 commit comments

Comments
 (0)