[Coroutines][3/6] New pass manager: coro-elide

[modocache]

Summary

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

The third in a series of patches that ports the LLVM coroutines passes
to the new pass manager infrastructure. This patch implements 'coro-elide'.

The new pass manager infrastructure does not implicitly repeat CGSCC
pass pipelines when a function is devirtualized, and so the tests
for the new pass manager that rely on that behavior now explicitly
specify repeat<2>.

[modocache]

As per one of the review comments on D71899, I split out the trivial parts of this patch, such as the 'legacy' pass renaming, and committed them separately. Here's an updated version of this patch without those 'NFC' changes.

[modocache]

Addressed review comments from D71899.

[wenlei]

Comment needs to be updated, otherwise LGTM.

llvm/test/Transforms/Coroutines/coro-elide.ll

8–9

Update this comment now that we rely on CallGraphUpdater instead of devirt for newPM?

[modocache]

Thanks for the review! Yeah, you're right. I think the comment at the top of the file explains both the legacy and new pass manager just fine, so I removed the extra comment I'd added.

llvm/include/llvm/Transforms/Coroutines/CoroElide.h

+//===---- CoroElide.h - Coroutine frame allocation elision ------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// \file
+// This file declares a pass that replaces dynamic allocation of coroutine
+// frames with alloca and replaces calls to llvm.coro.resume and
+// llvm.coro.destroy with direct calls to coroutine sub-functions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_COROUTINES_COROELIDE_H
+#define LLVM_TRANSFORMS_COROUTINES_COROELIDE_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class Function;
+
+struct CoroElidePass : PassInfoMixin<CoroElidePass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_COROUTINES_COROELIDE_H

llvm/lib/Passes/PassBuilder.cpp

 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
 #include "llvm/Transforms/Coroutines/CoroEarly.h"
+#include "llvm/Transforms/Coroutines/CoroElide.h"
 #include "llvm/Transforms/Coroutines/CoroSplit.h"
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/IPO/CalledValuePropagation.h"
 #include "llvm/Transforms/IPO/ConstantMerge.h"
 #include "llvm/Transforms/IPO/CrossDSOCFI.h"
 #include "llvm/Transforms/IPO/DeadArgumentElimination.h"

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("alignment-from-assumptions", AlignmentFromAssumptionsPass())
 FUNCTION_PASS("bdce", BDCEPass())
 FUNCTION_PASS("bounds-checking", BoundsCheckingPass())
 FUNCTION_PASS("break-crit-edges", BreakCriticalEdgesPass())
 FUNCTION_PASS("callsite-splitting", CallSiteSplittingPass())
 FUNCTION_PASS("consthoist", ConstantHoistingPass())
 FUNCTION_PASS("chr", ControlHeightReductionPass())
 FUNCTION_PASS("coro-early", CoroEarlyPass())
+FUNCTION_PASS("coro-elide", CoroElidePass())
 FUNCTION_PASS("correlated-propagation", CorrelatedValuePropagationPass())
 FUNCTION_PASS("dce", DCEPass())
 FUNCTION_PASS("div-rem-pairs", DivRemPairsPass())
 FUNCTION_PASS("dse", DSEPass())
 FUNCTION_PASS("dot-cfg", CFGPrinterPass())
 FUNCTION_PASS("dot-cfg-only", CFGOnlyPrinterPass())
 FUNCTION_PASS("early-cse", EarlyCSEPass(/*UseMemorySSA=*/false))
 FUNCTION_PASS("early-cse-memssa", EarlyCSEPass(/*UseMemorySSA=*/true))

llvm/lib/Transforms/Coroutines/CoroElide.cpp

 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
-// This pass replaces dynamic allocation of coroutine frame with alloca and
-// replaces calls to llvm.coro.resume and llvm.coro.destroy with direct calls
-// to coroutine sub-functions.
-//===----------------------------------------------------------------------===//
 
+#include "llvm/Transforms/Coroutines/CoroElide.h"
 #include "CoroInternal.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
struct Lowerer : coro::LowererBase {
   SmallVector<CoroSubFnInst *, 4> ResumeAddr;
   SmallVector<CoroSubFnInst *, 4> DestroyAddr;
   SmallVector<CoroFreeInst *, 1> CoroFrees;
 
   Lowerer(Module &M) : LowererBase(M) {}
 
   void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA);
   bool shouldElide(Function *F, DominatorTree &DT) const;
+  void collectPostSplitCoroIds(Function *F);
   bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT);
 };
 } // end anonymous namespace
 
 // Go through the list of coro.subfn.addr intrinsics and replace them with the
 // provided constant.
 static void replaceWithConstant(Constant *Value,
                                 SmallVectorImpl<CoroSubFnInst *> &Users) {
bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const {
   }
 
   // If size of the set is the same as total number of coro.begin, that means we
   // found a coro.free or coro.destroy referencing each coro.begin, so we can
   // perform heap elision.
   return ReferencedCoroBegins.size() == CoroBegins.size();
 }
 
+void Lowerer::collectPostSplitCoroIds(Function *F) {
+  CoroIds.clear();
+  for (auto &I : instructions(F))
+    if (auto *CII = dyn_cast<CoroIdInst>(&I))
+      if (CII->getInfo().isPostSplit())
+        // If it is the coroutine itself, don't touch it.
+        if (CII->getCoroutine() != CII->getFunction())
+          CoroIds.push_back(CII);
+}
+
 bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
                             DominatorTree &DT) {
   CoroBegins.clear();
   CoroAllocs.clear();
   CoroFrees.clear();
   ResumeAddr.clear();
   DestroyAddr.clear();
 
static bool replaceDevirtTrigger(Function &F) {
   Module &M = *F.getParent();
   Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
   assert(DevirtFn && "coro.devirt.fn not found");
   replaceWithConstant(DevirtFn, DevirtAddr);
 
   return true;
 }
 
-//===----------------------------------------------------------------------===//
-//                              Top Level Driver
-//===----------------------------------------------------------------------===//
+static bool declaresCoroElideIntrinsics(Module &M) {
+  return coro::declaresIntrinsics(M, {"llvm.coro.id"});
+}
+
+PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) {
+  auto &M = *F.getParent();
+  if (!declaresCoroElideIntrinsics(M))
+    return PreservedAnalyses::all();
+
+  Lowerer L(M);
+  L.CoroIds.clear();
+  L.collectPostSplitCoroIds(&F);
+  // If we did not find any coro.id, there is nothing to do.
+  if (L.CoroIds.empty())
+    return PreservedAnalyses::all();
+
+  AAResults &AA = AM.getResult<AAManager>(F);
+  DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
+
+  bool Changed = false;
+  for (auto *CII : L.CoroIds)
+    Changed |= L.processCoroId(CII, AA, DT);
+
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
 
 namespace {
 struct CoroElideLegacy : FunctionPass {
   static char ID;
   CoroElideLegacy() : FunctionPass(ID) {
     initializeCoroElideLegacyPass(*PassRegistry::getPassRegistry());
   }
 
   std::unique_ptr<Lowerer> L;
 
   bool doInitialization(Module &M) override {
-    if (coro::declaresIntrinsics(M, {"llvm.coro.id"}))
+    if (declaresCoroElideIntrinsics(M))
       L = std::make_unique<Lowerer>(M);
     return false;
   }
 
   bool runOnFunction(Function &F) override {
     if (!L)
       return false;
 
     bool Changed = false;
 
     if (F.hasFnAttribute(CORO_PRESPLIT_ATTR))
       Changed = replaceDevirtTrigger(F);
 
     L->CoroIds.clear();
-
-    // Collect all PostSplit coro.ids.
-    for (auto &I : instructions(F))
-      if (auto *CII = dyn_cast<CoroIdInst>(&I))
-        if (CII->getInfo().isPostSplit())
-          // If it is the coroutine itself, don't touch it.
-          if (CII->getCoroutine() != CII->getFunction())
-            L->CoroIds.push_back(CII);
-
+    L->collectPostSplitCoroIds(&F);
     // If we did not find any coro.id, there is nothing to do.
     if (L->CoroIds.empty())
       return Changed;
 
     AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
     DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 
     for (auto *CII : L->CoroIds)

llvm/test/Transforms/Coroutines/coro-elide.ll

 ; Tests that the coro.destroy and coro.resume are devirtualized where possible,
 ; SCC pipeline restarts and inlines the direct calls.
 ; RUN: opt < %s -S -inline -coro-elide -dce | FileCheck %s
+; RUN: opt < %s -S \
+; RUN:   -passes='cgscc(repeat<2>(inline,function(coro-elide,dce)))' \
+; RUN:   | FileCheck %s
 
 declare void @print(i32) nounwind
 
 ; resume part of the coroutine
 define fastcc void @f.resume(i8*) {
   tail call void @print(i32 0)
   ret void
 }
 
 ; destroy part of the coroutine
 define fastcc void @f.destroy(i8*) {

llvm/test/Transforms/Coroutines/coro-heap-elide.ll

 ; Tests that the dynamic allocation and deallocation of the coroutine frame is
 ; elided and any tail calls referencing the coroutine frame has the tail 
 ; call attribute removed.
 ; RUN: opt < %s -S -inline -coro-elide -instsimplify -simplifycfg | FileCheck %s
+; RUN: opt < %s -S \
+; RUN:   -passes='cgscc(inline,function(coro-elide,instsimplify,simplify-cfg))' \
+; RUN:   -aa-pipeline='basic-aa' | FileCheck %s
 
 declare void @print(i32) nounwind
 
 %f.frame = type {i32}
 
 declare void @bar(i8*)
 
 declare fastcc void @f.resume(%f.frame*)