diff --git a/clang/include/clang/CIR/MissingFeatures.h b/clang/include/clang/CIR/MissingFeatures.h index b64e07ff2bfb8..b2b9c0a09a2d7 100644 --- a/clang/include/clang/CIR/MissingFeatures.h +++ b/clang/include/clang/CIR/MissingFeatures.h @@ -153,6 +153,7 @@ struct MissingFeatures { static bool coroEndBuiltinCall() { return false; } static bool coroutineFrame() { return false; } static bool emitBodyAndFallthrough() { return false; } + static bool coroOutsideFrameMD() { return false; } // Various handling of deferred processing in CIRGenModule. static bool cgmRelease() { return false; } diff --git a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp index 930ae55405756..05fb1aedcbf4a 100644 --- a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp @@ -13,6 +13,7 @@ #include "CIRGenFunction.h" #include "mlir/Support/LLVM.h" #include "clang/AST/StmtCXX.h" +#include "clang/AST/StmtVisitor.h" #include "clang/Basic/TargetInfo.h" #include "clang/CIR/Dialect/IR/CIRTypes.h" #include "clang/CIR/MissingFeatures.h" @@ -33,6 +34,65 @@ struct clang::CIRGen::CGCoroData { CIRGenFunction::CGCoroInfo::CGCoroInfo() {} CIRGenFunction::CGCoroInfo::~CGCoroInfo() {} +namespace { +// FIXME: both GetParamRef and ParamReferenceReplacerRAII are good template +// candidates to be shared among LLVM / CIR codegen. + +// Hunts for the parameter reference in the parameter copy/move declaration. +struct GetParamRef : public StmtVisitor { +public: + DeclRefExpr *expr = nullptr; + GetParamRef() {} + void VisitDeclRefExpr(DeclRefExpr *e) { + assert(expr == nullptr && "multilple declref in param move"); + expr = e; + } + void VisitStmt(Stmt *s) { + for (Stmt *c : s->children()) { + if (c) + Visit(c); + } + } +}; + +// This class replaces references to parameters to their copies by changing +// the addresses in CGF.LocalDeclMap and restoring back the original values in +// its destructor. +struct ParamReferenceReplacerRAII { + CIRGenFunction::DeclMapTy savedLocals; + CIRGenFunction::DeclMapTy &localDeclMap; + + ParamReferenceReplacerRAII(CIRGenFunction::DeclMapTy &localDeclMap) + : localDeclMap(localDeclMap) {} + + void addCopy(const DeclStmt *pm) { + // Figure out what param it refers to. + + assert(pm->isSingleDecl()); + const VarDecl *vd = static_cast(pm->getSingleDecl()); + const Expr *initExpr = vd->getInit(); + GetParamRef visitor; + visitor.Visit(const_cast(initExpr)); + assert(visitor.expr); + DeclRefExpr *dreOrig = visitor.expr; + auto *pd = dreOrig->getDecl(); + + auto it = localDeclMap.find(pd); + assert(it != localDeclMap.end() && "parameter is not found"); + savedLocals.insert({pd, it->second}); + + auto copyIt = localDeclMap.find(vd); + assert(copyIt != localDeclMap.end() && "parameter copy is not found"); + it->second = copyIt->getSecond(); + } + + ~ParamReferenceReplacerRAII() { + for (auto &&savedLocal : savedLocals) { + localDeclMap.insert({savedLocal.first, savedLocal.second}); + } + } +}; +} // namespace static void createCoroData(CIRGenFunction &cgf, CIRGenFunction::CGCoroInfo &curCoro, cir::CallOp coroId) { @@ -149,7 +209,47 @@ CIRGenFunction::emitCoroutineBody(const CoroutineBodyStmt &s) { if (s.getReturnStmtOnAllocFailure()) cgm.errorNYI("handle coroutine return alloc failure"); - assert(!cir::MissingFeatures::generateDebugInfo()); - assert(!cir::MissingFeatures::emitBodyAndFallthrough()); + { + assert(!cir::MissingFeatures::generateDebugInfo()); + ParamReferenceReplacerRAII paramReplacer(localDeclMap); + // Create mapping between parameters and copy-params for coroutine + // function. + llvm::ArrayRef paramMoves = s.getParamMoves(); + assert((paramMoves.size() == 0 || (paramMoves.size() == fnArgs.size())) && + "ParamMoves and FnArgs should be the same size for coroutine " + "function"); + // For zipping the arg map into debug info. + assert(!cir::MissingFeatures::generateDebugInfo()); + + // Create parameter copies. We do it before creating a promise, since an + // evolution of coroutine TS may allow promise constructor to observe + // parameter copies. + assert(!cir::MissingFeatures::coroOutsideFrameMD()); + for (auto *pm : paramMoves) { + if (emitStmt(pm, /*useCurrentScope=*/true).failed()) + return mlir::failure(); + paramReplacer.addCopy(cast(pm)); + } + + if (emitStmt(s.getPromiseDeclStmt(), /*useCurrentScope=*/true).failed()) + return mlir::failure(); + // returnValue should be valid as long as the coroutine's return type + // is not void. The assertion could help us to reduce the check later. + assert(returnValue.isValid() == (bool)s.getReturnStmt()); + // Now we have the promise, initialize the GRO. + // We need to emit `get_return_object` first. According to: + // [dcl.fct.def.coroutine]p7 + // The call to get_return_­object is sequenced before the call to + // initial_suspend and is invoked at most once. + // + // So we couldn't emit return value when we emit return statment, + // otherwise the call to get_return_object wouldn't be in front + // of initial_suspend. + if (returnValue.isValid()) + emitAnyExprToMem(s.getReturnValue(), returnValue, + s.getReturnValue()->getType().getQualifiers(), + /*isInit*/ true); + assert(!cir::MissingFeatures::emitBodyAndFallthrough()); + } return mlir::success(); } diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.cpp b/clang/lib/CIR/CodeGen/CIRGenFunction.cpp index 5d5209b9ffb60..cc75acc18c211 100644 --- a/clang/lib/CIR/CodeGen/CIRGenFunction.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenFunction.cpp @@ -632,6 +632,10 @@ cir::FuncOp CIRGenFunction::generateCode(clang::GlobalDecl gd, cir::FuncOp fn, startFunction(gd, retTy, fn, funcType, args, loc, bodyRange.getBegin()); + // Save parameters for coroutine function. + if (body && isa_and_nonnull(body)) + llvm::append_range(fnArgs, funcDecl->parameters()); + if (isa(funcDecl)) { emitDestructorBody(args); } else if (isa(funcDecl)) { diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.h b/clang/lib/CIR/CodeGen/CIRGenFunction.h index f879e580989f7..dab40d37268fa 100644 --- a/clang/lib/CIR/CodeGen/CIRGenFunction.h +++ b/clang/lib/CIR/CodeGen/CIRGenFunction.h @@ -152,6 +152,9 @@ class CIRGenFunction : public CIRGenTypeCache { /// global initializers. mlir::Operation *curFn = nullptr; + /// Save Parameter Decl for coroutine. + llvm::SmallVector fnArgs; + using DeclMapTy = llvm::DenseMap; /// This keeps track of the CIR allocas or globals for local C /// declarations. diff --git a/clang/test/CIR/CodeGen/coro-task.cpp b/clang/test/CIR/CodeGen/coro-task.cpp index 265325f82d7f7..5738c815909ea 100644 --- a/clang/test/CIR/CodeGen/coro-task.cpp +++ b/clang/test/CIR/CodeGen/coro-task.cpp @@ -36,6 +36,12 @@ struct suspend_never { void await_resume() noexcept {} }; +struct string { + int size() const; + string(); + string(char const *s); +}; + } // namespace std namespace folly { @@ -101,7 +107,10 @@ co_invoke_fn co_invoke; }} // namespace folly::coro // CIR-DAG: ![[VoidTask:.*]] = !cir.record" padded {!u8i}> - +// CIR-DAG: ![[IntTask:.*]] = !cir.record" padded {!u8i}> +// CIR-DAG: ![[VoidPromisse:.*]] = !cir.record::promise_type" padded {!u8i}> +// CIR-DAG: ![[IntPromisse:.*]] = !cir.record::promise_type" padded {!u8i}> +// CIR-DAG: ![[StdString:.*]] = !cir.record // CIR: module {{.*}} { // CIR-NEXT: cir.global external @_ZN5folly4coro9co_invokeE = #cir.zero : !rec_folly3A3Acoro3A3Aco_invoke_fn @@ -119,6 +128,7 @@ VoidTask silly_task() { // CIR: cir.func coroutine dso_local @_Z10silly_taskv() -> ![[VoidTask]] // CIR: %[[VoidTaskAddr:.*]] = cir.alloca ![[VoidTask]], {{.*}}, ["__retval"] // CIR: %[[SavedFrameAddr:.*]] = cir.alloca !cir.ptr, !cir.ptr>, ["__coro_frame_addr"] +// CIR: %[[VoidPromisseAddr:.*]] = cir.alloca ![[VoidPromisse]], {{.*}}, ["__promise"] // Get coroutine id with __builtin_coro_id. @@ -138,3 +148,27 @@ VoidTask silly_task() { // CIR: } // CIR: %[[Load0:.*]] = cir.load{{.*}} %[[SavedFrameAddr]] : !cir.ptr>, !cir.ptr // CIR: %[[CoroFrameAddr:.*]] = cir.call @__builtin_coro_begin(%[[CoroId]], %[[Load0]]) + +// Call promise.get_return_object() to retrieve the task object. + +// CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type17get_return_objectEv(%[[VoidPromisseAddr]]) nothrow : {{.*}} -> ![[VoidTask]] +// CIR: cir.store{{.*}} %[[RetObj]], %[[VoidTaskAddr]] : ![[VoidTask]] + +folly::coro::Task byRef(const std::string& s) { + co_return s.size(); +} + +// CIR: cir.func coroutine dso_local @_Z5byRefRKSt6string(%[[ARG:.*]]: !cir.ptr {{.*}}) -> ![[IntTask]] +// CIR: %[[AllocaParam:.*]] = cir.alloca !cir.ptr, {{.*}}, ["s", init, const] +// CIR: %[[IntTaskAddr:.*]] = cir.alloca ![[IntTask]], {{.*}}, ["__retval"] +// CIR: %[[SavedFrameAddr:.*]] = cir.alloca !cir.ptr, !cir.ptr>, ["__coro_frame_addr"] +// CIR: %[[AllocaFnUse:.*]] = cir.alloca !cir.ptr, {{.*}}, ["s", init, const] +// CIR: %[[IntPromisseAddr:.*]] = cir.alloca ![[IntPromisse]], {{.*}}, ["__promise"] +// CIR: cir.store %[[ARG]], %[[AllocaParam]] : !cir.ptr, {{.*}} + +// Call promise.get_return_object() to retrieve the task object. + +// CIR: %[[LOAD:.*]] = cir.load %[[AllocaParam]] : !cir.ptr>, !cir.ptr +// CIR: cir.store {{.*}} %[[LOAD]], %[[AllocaFnUse]] : !cir.ptr, !cir.ptr> +// CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type17get_return_objectEv(%4) nothrow : {{.*}} -> ![[IntTask]] +// CIR: cir.store {{.*}} %[[RetObj]], %[[IntTaskAddr]] : ![[IntTask]]