10b57cec5SDimitry Andric //===- CoroCleanup.cpp - Coroutine Cleanup Pass ---------------------------===//
20b57cec5SDimitry Andric //
30b57cec5SDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
40b57cec5SDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
50b57cec5SDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
60b57cec5SDimitry Andric //
70b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
80b57cec5SDimitry Andric
95ffd83dbSDimitry Andric #include "llvm/Transforms/Coroutines/CoroCleanup.h"
100b57cec5SDimitry Andric #include "CoroInternal.h"
110b57cec5SDimitry Andric #include "llvm/IR/IRBuilder.h"
120b57cec5SDimitry Andric #include "llvm/IR/InstIterator.h"
1381ad6265SDimitry Andric #include "llvm/IR/PassManager.h"
1481ad6265SDimitry Andric #include "llvm/IR/Function.h"
1581ad6265SDimitry Andric #include "llvm/Transforms/Scalar/SimplifyCFG.h"
160b57cec5SDimitry Andric
170b57cec5SDimitry Andric using namespace llvm;
180b57cec5SDimitry Andric
190b57cec5SDimitry Andric #define DEBUG_TYPE "coro-cleanup"
200b57cec5SDimitry Andric
210b57cec5SDimitry Andric namespace {
220b57cec5SDimitry Andric // Created on demand if CoroCleanup pass has work to do.
230b57cec5SDimitry Andric struct Lowerer : coro::LowererBase {
240b57cec5SDimitry Andric IRBuilder<> Builder;
Lowerer__anon60aa7eab0111::Lowerer250b57cec5SDimitry Andric Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
2681ad6265SDimitry Andric bool lower(Function &F);
270b57cec5SDimitry Andric };
280b57cec5SDimitry Andric }
290b57cec5SDimitry Andric
lowerSubFn(IRBuilder<> & Builder,CoroSubFnInst * SubFn)300b57cec5SDimitry Andric static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
310b57cec5SDimitry Andric Builder.SetInsertPoint(SubFn);
325f757f3fSDimitry Andric Value *FramePtr = SubFn->getFrame();
330b57cec5SDimitry Andric int Index = SubFn->getIndex();
340b57cec5SDimitry Andric
355f757f3fSDimitry Andric auto *FrameTy = StructType::get(SubFn->getContext(),
365f757f3fSDimitry Andric {Builder.getPtrTy(), Builder.getPtrTy()});
370b57cec5SDimitry Andric
380b57cec5SDimitry Andric Builder.SetInsertPoint(SubFn);
390b57cec5SDimitry Andric auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
400b57cec5SDimitry Andric auto *Load = Builder.CreateLoad(FrameTy->getElementType(Index), Gep);
410b57cec5SDimitry Andric
420b57cec5SDimitry Andric SubFn->replaceAllUsesWith(Load);
430b57cec5SDimitry Andric }
440b57cec5SDimitry Andric
lower(Function & F)4581ad6265SDimitry Andric bool Lowerer::lower(Function &F) {
4681ad6265SDimitry Andric bool IsPrivateAndUnprocessed = F.isPresplitCoroutine() && F.hasLocalLinkage();
470b57cec5SDimitry Andric bool Changed = false;
480b57cec5SDimitry Andric
49349cc55cSDimitry Andric for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
500b57cec5SDimitry Andric if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
510b57cec5SDimitry Andric switch (II->getIntrinsicID()) {
520b57cec5SDimitry Andric default:
530b57cec5SDimitry Andric continue;
540b57cec5SDimitry Andric case Intrinsic::coro_begin:
550b57cec5SDimitry Andric II->replaceAllUsesWith(II->getArgOperand(1));
560b57cec5SDimitry Andric break;
570b57cec5SDimitry Andric case Intrinsic::coro_free:
580b57cec5SDimitry Andric II->replaceAllUsesWith(II->getArgOperand(1));
590b57cec5SDimitry Andric break;
600b57cec5SDimitry Andric case Intrinsic::coro_alloc:
610b57cec5SDimitry Andric II->replaceAllUsesWith(ConstantInt::getTrue(Context));
620b57cec5SDimitry Andric break;
63349cc55cSDimitry Andric case Intrinsic::coro_async_resume:
64349cc55cSDimitry Andric II->replaceAllUsesWith(
65349cc55cSDimitry Andric ConstantPointerNull::get(cast<PointerType>(I.getType())));
66349cc55cSDimitry Andric break;
670b57cec5SDimitry Andric case Intrinsic::coro_id:
688bcb0991SDimitry Andric case Intrinsic::coro_id_retcon:
698bcb0991SDimitry Andric case Intrinsic::coro_id_retcon_once:
70e8d8bef9SDimitry Andric case Intrinsic::coro_id_async:
710b57cec5SDimitry Andric II->replaceAllUsesWith(ConstantTokenNone::get(Context));
720b57cec5SDimitry Andric break;
730b57cec5SDimitry Andric case Intrinsic::coro_subfn_addr:
740b57cec5SDimitry Andric lowerSubFn(Builder, cast<CoroSubFnInst>(II));
750b57cec5SDimitry Andric break;
76349cc55cSDimitry Andric case Intrinsic::coro_end:
77349cc55cSDimitry Andric case Intrinsic::coro_suspend_retcon:
78349cc55cSDimitry Andric if (IsPrivateAndUnprocessed) {
79349cc55cSDimitry Andric II->replaceAllUsesWith(UndefValue::get(II->getType()));
80349cc55cSDimitry Andric } else
81349cc55cSDimitry Andric continue;
82349cc55cSDimitry Andric break;
83fe6060f1SDimitry Andric case Intrinsic::coro_async_size_replace:
84fe6060f1SDimitry Andric auto *Target = cast<ConstantStruct>(
85fe6060f1SDimitry Andric cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
86fe6060f1SDimitry Andric ->getInitializer());
87fe6060f1SDimitry Andric auto *Source = cast<ConstantStruct>(
88fe6060f1SDimitry Andric cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
89fe6060f1SDimitry Andric ->getInitializer());
90fe6060f1SDimitry Andric auto *TargetSize = Target->getOperand(1);
91fe6060f1SDimitry Andric auto *SourceSize = Source->getOperand(1);
92fe6060f1SDimitry Andric if (TargetSize->isElementWiseEqual(SourceSize)) {
93fe6060f1SDimitry Andric break;
94fe6060f1SDimitry Andric }
95fe6060f1SDimitry Andric auto *TargetRelativeFunOffset = Target->getOperand(0);
96fe6060f1SDimitry Andric auto *NewFuncPtrStruct = ConstantStruct::get(
97fe6060f1SDimitry Andric Target->getType(), TargetRelativeFunOffset, SourceSize);
98fe6060f1SDimitry Andric Target->replaceAllUsesWith(NewFuncPtrStruct);
99fe6060f1SDimitry Andric break;
1000b57cec5SDimitry Andric }
1010b57cec5SDimitry Andric II->eraseFromParent();
1020b57cec5SDimitry Andric Changed = true;
1030b57cec5SDimitry Andric }
1040b57cec5SDimitry Andric }
1050b57cec5SDimitry Andric
1060b57cec5SDimitry Andric return Changed;
1070b57cec5SDimitry Andric }
1080b57cec5SDimitry Andric
declaresCoroCleanupIntrinsics(const Module & M)1095ffd83dbSDimitry Andric static bool declaresCoroCleanupIntrinsics(const Module &M) {
110fe6060f1SDimitry Andric return coro::declaresIntrinsics(
111fe6060f1SDimitry Andric M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
112fe6060f1SDimitry Andric "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
113fcaf7f86SDimitry Andric "llvm.coro.id.async", "llvm.coro.id.retcon.once",
114fcaf7f86SDimitry Andric "llvm.coro.async.size.replace", "llvm.coro.async.resume"});
1155ffd83dbSDimitry Andric }
1165ffd83dbSDimitry Andric
run(Module & M,ModuleAnalysisManager & MAM)11781ad6265SDimitry Andric PreservedAnalyses CoroCleanupPass::run(Module &M,
11881ad6265SDimitry Andric ModuleAnalysisManager &MAM) {
11981ad6265SDimitry Andric if (!declaresCoroCleanupIntrinsics(M))
1205ffd83dbSDimitry Andric return PreservedAnalyses::all();
1215ffd83dbSDimitry Andric
12281ad6265SDimitry Andric FunctionAnalysisManager &FAM =
12381ad6265SDimitry Andric MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
12481ad6265SDimitry Andric
12581ad6265SDimitry Andric FunctionPassManager FPM;
12681ad6265SDimitry Andric FPM.addPass(SimplifyCFGPass());
12781ad6265SDimitry Andric
12806c3fb27SDimitry Andric PreservedAnalyses FuncPA;
12906c3fb27SDimitry Andric FuncPA.preserveSet<CFGAnalyses>();
13006c3fb27SDimitry Andric
13181ad6265SDimitry Andric Lowerer L(M);
13206c3fb27SDimitry Andric for (auto &F : M) {
13306c3fb27SDimitry Andric if (L.lower(F)) {
13406c3fb27SDimitry Andric FAM.invalidate(F, FuncPA);
13581ad6265SDimitry Andric FPM.run(F, FAM);
13606c3fb27SDimitry Andric }
13706c3fb27SDimitry Andric }
13881ad6265SDimitry Andric
1395ffd83dbSDimitry Andric return PreservedAnalyses::none();
1405ffd83dbSDimitry Andric }
141