[Coroutines] Part 15a: Lower coro.subfn.addr in CoroCleanup

[GorNishanov]

Summary

Not all coro.subfn.addr intrinsics can be eliminated in CoroElide through devirtualization. Those that remain need to be lowered in CoroCleanup.

[GorNishanov]

Ping

[majnemer]

LGTM

lib/Transforms/Coroutines/CoroCleanup.cpp

 //===- CoroCleanup.cpp - Coroutine Cleanup Pass ---------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass lowers all remaining coroutine intrinsics.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Scalar.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-cleanup"
 
 namespace {
 // Created on demand if CoroCleanup pass has work to do.
 struct Lowerer : coro::LowererBase {
-  Lowerer(Module &M) : LowererBase(M) {}
+  IRBuilder<> Builder;
+  Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
   bool lowerRemainingCoroIntrinsics(Function &F);
 };
 }
 
 static void simplifyCFG(Function &F) {
   llvm::legacy::FunctionPassManager FPM(F.getParent());
   FPM.add(createCFGSimplificationPass());
 
   FPM.doInitialization();
   FPM.run(F);
   FPM.doFinalization();
 }
 
+static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
+  Builder.SetInsertPoint(SubFn);
+  Value *FrameRaw = SubFn->getFrame();
+  int Index = SubFn->getIndex();
+
+  auto *FrameTy = StructType::get(
+      SubFn->getContext(), {Builder.getInt8PtrTy(), Builder.getInt8PtrTy()});
+  PointerType *FramePtrTy = FrameTy->getPointerTo();
+
+  Builder.SetInsertPoint(SubFn);
+  auto *FramePtr = Builder.CreateBitCast(FrameRaw, FramePtrTy);
+  auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
+  auto *Load = Builder.CreateLoad(Gep);
+
+  SubFn->replaceAllUsesWith(Load);
+}
+
 bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
   bool Changed = false;
 
   for (auto IB = inst_begin(F), E = inst_end(F); IB != E;) {
     Instruction &I = *IB++;
     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
       switch (II->getIntrinsicID()) {
       default:
         continue;
       case Intrinsic::coro_begin:
         II->replaceAllUsesWith(II->getArgOperand(1));
         break;
       case Intrinsic::coro_free:
         II->replaceAllUsesWith(II->getArgOperand(1));
         break;
       case Intrinsic::coro_alloc:
         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
         break;
       case Intrinsic::coro_id:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
+      case Intrinsic::coro_subfn_addr:
+        lowerSubFn(Builder, cast<CoroSubFnInst>(II));
+        break;
       }
       II->eraseFromParent();
       Changed = true;
     }
   }
 
   if (Changed) {
     // After replacement were made we can cleanup the function body a little.
struct CoroCleanup : FunctionPass {
   CoroCleanup() : FunctionPass(ID) {}
 
   std::unique_ptr<Lowerer> L;
 
   // This pass has work to do only if we find intrinsics we are going to lower
   // in the module.
   bool doInitialization(Module &M) override {
     if (coro::declaresIntrinsics(M, {"llvm.coro.alloc", "llvm.coro.begin",
-                                     "llvm.coro.free", "llvm.coro.id"}))
+                                     "llvm.coro.subfn.addr", "llvm.coro.free",
+                                     "llvm.coro.id"}))
       L = llvm::make_unique<Lowerer>(M);
     return false;
   }
 
   bool runOnFunction(Function &F) override {
     if (L)
       return L->lowerRemainingCoroIntrinsics(F);
     return false;

lib/Transforms/Coroutines/Coroutines.cpp

 #ifndef NDEBUG
 static bool isCoroutineIntrinsicName(StringRef Name) {
   // NOTE: Must be sorted!
   static const char *const CoroIntrinsics[] = {
       "llvm.coro.alloc",   "llvm.coro.begin",   "llvm.coro.destroy",
       "llvm.coro.done",    "llvm.coro.end",     "llvm.coro.frame",
       "llvm.coro.free",    "llvm.coro.id",      "llvm.coro.param",
       "llvm.coro.promise", "llvm.coro.resume",  "llvm.coro.save",
-      "llvm.coro.size",    "llvm.coro.suspend",
+      "llvm.coro.size",    "llvm.coro.subfn.addr", "llvm.coro.suspend",
   };
   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
 }
 #endif
 
 // Verifies if a module has named values listed. Also, in debug mode verifies
 // that names are intrinsic names.
 bool coro::declaresIntrinsics(Module &M,

test/Transforms/Coroutines/coro-cleanup.ll

+; Make sure that all library helper coro intrinsics are lowered.
+; RUN: opt < %s -O0 -enable-coroutines -S | FileCheck %s
+
+; CHECK-LABEL: @uses_library_support_coro_intrinsics(
+; CHECK-NOT:     @llvm.coro
+; CHECK:         ret void
+define void @uses_library_support_coro_intrinsics(i8* %hdl) {
+entry:
+  call void @llvm.coro.resume(i8* %hdl)
+  call void @llvm.coro.destroy(i8* %hdl)
+  call i1 @llvm.coro.done(i8* %hdl)
+  ret void
+}
+
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)
+declare i1 @llvm.coro.done(i8*)
+