[Coroutines][5/6] Add coroutine passes to pipeline

[modocache]

Summary

Depends on https://reviews.llvm.org/D71901.

The fifth in a series of patches that ports the LLVM coroutines passes
to the new pass manager infrastructure.

The first 4 patches allow users to run coroutine passes by invoking, for
example opt -passes=coro-early. However, most of LLVM's tests for
coroutines use an option, opt -enable-coroutines, which adds all 4
coroutine passes to the appropriate legacy pass manager extension points.
This patch does the same, but using the new pass manager: when
coroutine features are enabled and the new pass manager is being used,
this adds the new-pass-manager-compliant coroutine passes to the pass
builder's pipeline.

This allows us to run all coroutine tests using the new pass manager
(besides those that use the coroutine retcon ABI used by the Swift
compiler, which is not yet supported in the new pass manager).

[modocache]

Small cleanup.

[modocache]

Updated to use a new approach: instead of a custom pipeline parser, '-enable-coroutines' now adds coroutine passes to the new pass manager's pass builder. Using this in conjunction with the 'default<O2>' pipeline parser, this patch can now run all coroutine tests in the LLVM test suite -- and they all pass.

[wenlei]

LGTM. Thanks for fixing the coro cleanup issue with ThinLTO here.

llvm/include/llvm/Transforms/Coroutines.h

 //===-- Coroutines.h - Coroutine Transformations ----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 // Declare accessor functions for coroutine lowering passes.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_COROUTINES_H
 #define LLVM_TRANSFORMS_COROUTINES_H
 
 namespace llvm {
 
 class Pass;
+class PassBuilder;
 class PassManagerBuilder;
 
-/// Add all coroutine passes to appropriate extension points.
+/// Register a callback to parse coroutines passes in a pass pipeline.
+void registerCoroutinePipelineParsingCallback(PassBuilder &Builder,
+                                              bool DebugLogging = false);
+
+/// Add all coroutine passes to appropriate extension points using the legacy
+/// pass manager.
 void addCoroutinePassesToExtensionPoints(PassManagerBuilder &Builder);
 
 /// Lower coroutine intrinsics that are not needed by later passes.
 Pass *createCoroEarlyLegacyPass();
 
 /// Split up coroutines into multiple functions driving their state machines.
 Pass *createCoroSplitLegacyPass();
 

llvm/lib/Transforms/Coroutines/Coroutines.cpp

 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/InitializePasses.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Transforms/Coroutines/CoroCleanup.h"
+#include "llvm/Transforms/Coroutines/CoroEarly.h"
+#include "llvm/Transforms/Coroutines/CoroElide.h"
+#include "llvm/Transforms/Coroutines/CoroSplit.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <cassert>
 #include <cstddef>
 #include <utility>
 
 using namespace llvm;
 
 void llvm::initializeCoroutines(PassRegistry &Registry) {
   initializeCoroEarlyLegacyPass(Registry);
   initializeCoroSplitLegacyPass(Registry);
   initializeCoroElideLegacyPass(Registry);
   initializeCoroCleanupLegacyPass(Registry);
 }
 
+void llvm::registerCoroutinePipelineParsingCallback(PassBuilder &Builder,
+                                                    bool DebugLogging) {
+  Builder.registerPipelineParsingCallback(
+      [=](StringRef Name, ModulePassManager &MPM,
+          ArrayRef<PassBuilder::PipelineElement>) {
+        if (Name == "coroutines") {
+          MPM.addPass(createModuleToFunctionPassAdaptor(CoroEarlyPass()));
+
+          CGSCCPassManager CGPM(DebugLogging);
+          CGPM.addPass(CoroSplitPass());
+          CGPM.addPass(createCGSCCToFunctionPassAdaptor(CoroElidePass()));
+          MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+
+          MPM.addPass(createModuleToFunctionPassAdaptor(CoroCleanupPass()));
+          return true;
+        }
+        return false;
+      });
+}
+
 static void addCoroutineOpt0Passes(const PassManagerBuilder &Builder,
                                    legacy::PassManagerBase &PM) {
   PM.add(createCoroSplitLegacyPass());
   PM.add(createCoroElideLegacyPass());
 
   PM.add(createBarrierNoopPass());
   PM.add(createCoroCleanupLegacyPass());
 }

llvm/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
+; RUN: opt < %s -passes='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*)
 

llvm/test/Transforms/Coroutines/restart-trigger.ll

 ; Verifies that the restart trigger that is used by legacy coroutine passes
 ; forces the legacy pass manager to restart IPO pipelines, thereby causing the
 ; same coroutine to be looked at by CoroSplit pass twice.
 ; REQUIRES: asserts
 ; RUN: opt < %s -S -O0 -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -O1 -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck %s
+; RUN: opt < %s -S -passes=coroutines -debug-only=coro-split 2>&1 | FileCheck %s
 
 ; CHECK:      CoroSplit: Processing coroutine 'f' state: 0
 ; CHECK-NEXT: CoroSplit: Processing coroutine 'f' state: 1
 
 define void @f() {
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)

llvm/test/Transforms/Coroutines/smoketest.ll

 ; Test that all coroutine passes run in the correct order at all optimization
 ; levels and -enable-coroutines adds coroutine passes to the pipeline.
 ;
 ; RUN: opt < %s -disable-output -enable-coroutines -debug-pass=Arguments -O0 2>&1 | FileCheck %s
 ; RUN: opt < %s -disable-output -enable-coroutines -debug-pass=Arguments -O1 2>&1 | FileCheck %s
 ; RUN: opt < %s -disable-output -enable-coroutines -debug-pass=Arguments -O2 2>&1 | FileCheck %s
 ; RUN: opt < %s -disable-output -enable-coroutines -debug-pass=Arguments -O3 2>&1 | FileCheck %s
 ; RUN: opt < %s -disable-output -enable-coroutines -debug-pass=Arguments \
 ; RUN:     -coro-early -coro-split -coro-elide -coro-cleanup 2>&1 | FileCheck %s
 ; RUN: opt < %s -disable-output -debug-pass=Arguments 2>&1 \
 ; RUN:     | FileCheck %s -check-prefix=NOCORO
+; RUN: opt < %s -disable-output -passes=coroutines -debug-pass-manager 2>&1 \
+; RUN:     | FileCheck %s -check-prefix=NEWPM
+; RUN: opt < %s -disable-output \
+; RUN:     -passes='coroutines,function(coro-early),cgscc(coro-split),function(coro-elide,coro-cleanup)' \
+; RUN:     -debug-pass-manager 2>&1 | FileCheck %s -check-prefix=NEWPM
 
 ; CHECK: coro-early
 ; CHECK: coro-split
 ; CHECK: coro-elide
 ; CHECK: coro-cleanup
 
 ; NOCORO-NOT: coro-early
 ; NOCORO-NOT: coro-split
 ; NOCORO-NOT: coro-elide
 ; NOCORO-NOT: coro-cleanup
 
+; NEWPM: CoroEarlyPass
+; NEWPM: CoroSplitPass
+; NEWPM: CoroElidePass
+; NEWPM: CoroCleanupPass
+
 define void @foo() {
   ret void
 }

llvm/tools/opt/NewPMDriver.cpp

 #include "llvm/IR/Verifier.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Coroutines.h"
 #include "llvm/Transforms/IPO/ThinLTOBitcodeWriter.h"
 #include "llvm/Transforms/Scalar/LoopPassManager.h"
 #include "llvm/Transforms/Utils/Debugify.h"
 
 using namespace llvm;
 using namespace opt_tool;
 
 static cl::opt<bool>
PB.registerPipelineParsingCallback(
         }
         return false;
       });
 
 #define HANDLE_EXTENSION(Ext)                                                  \
   get##Ext##PluginInfo().RegisterPassBuilderCallbacks(PB);
 #include "llvm/Support/Extension.def"
 
+  // Register a callback that creates coroutines lowering passes as needed.
+  registerCoroutinePipelineParsingCallback(PB, DebugPM);
+
   // Specially handle the alias analysis manager so that we can register
   // a custom pipeline of AA passes with it.
   AAManager AA;
   if (auto Err = PB.parseAAPipeline(AA, AAPipeline)) {
     errs() << Arg0 << ": " << toString(std::move(Err)) << "\n";
     return false;
   }