[Coroutines] Run EarlyCSE for changed functions in CoroCleanup (3/3)

[ChuanqiXu]

Summary

After the previous revision (https://reviews.llvm.org/D125292) landed, the optimizations for readnone calls in coroutines would be blocked even if these calls don't cross suspend points.

This patch tries to solve the problem by run EarlyCSE for the functions which is changed in CoroCleanup pass. It wouldn't affect normal functions.

[ChuanqiXu]

Run optimization pass in CoroCleanup conditionally.

[ChuanqiXu]

Prefer https://reviews.llvm.org/D127383

llvm/include/llvm/Transforms/Coroutines/CoroCleanup.h

 
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 class Module;
 
 struct CoroCleanupPass : PassInfoMixin<CoroCleanupPass> {
+  CoroCleanupPass(bool OptimizationEnabled = false) : OptimizationEnabled(OptimizationEnabled) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
   static bool isRequired() { return true; }
+
+private:
+  bool OptimizationEnabled;
 };
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_COROUTINES_COROCLEANUP_H

llvm/lib/Passes/PassBuilder.cpp

 Expected<bool> parseLoopExtractorPassOptions(StringRef Params) {
   return parseSinglePassOption(Params, "single", "LoopExtractor");
 }
 
 Expected<bool> parseLowerMatrixIntrinsicsPassOptions(StringRef Params) {
   return parseSinglePassOption(Params, "minimal", "LowerMatrixIntrinsics");
 }
 
+Expected<bool> parseCoroCleanupPassOptions(StringRef Params) {
+  return parseSinglePassOption(Params, "opt", "CoroCleanup");
+}
+
 Expected<AddressSanitizerOptions> parseASanPassOptions(StringRef Params) {
   AddressSanitizerOptions Result;
   while (!Params.empty()) {
     StringRef ParamName;
     std::tie(ParamName, Params) = Params.split(';');
 
     if (ParamName == "kernel") {
       Result.CompileKernel = true;

llvm/lib/Passes/PassBuilderPipelines.cpp

PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
   if (EnableSyntheticCounts && !PGOOpt)
     MPM.addPass(SyntheticCountsPropagation());
 
   if (EnableModuleInliner)
     MPM.addPass(buildModuleInlinerPipeline(Level, Phase));
   else
     MPM.addPass(buildInlinerPipeline(Level, Phase));
 
-  MPM.addPass(CoroCleanupPass());
+  MPM.addPass(CoroCleanupPass(/*OptimizationEnabled*/ true));
 
   if (EnableMemProfiler && Phase != ThinOrFullLTOPhase::ThinLTOPreLink) {
     MPM.addPass(createModuleToFunctionPassAdaptor(MemProfilerPass()));
     MPM.addPass(ModuleMemProfilerPass());
   }
 
   return MPM;
 }

llvm/lib/Passes/PassRegistry.def

 MODULE_PASS("annotation2metadata", Annotation2MetadataPass())
 MODULE_PASS("openmp-opt", OpenMPOptPass())
 MODULE_PASS("called-value-propagation", CalledValuePropagationPass())
 MODULE_PASS("canonicalize-aliases", CanonicalizeAliasesPass())
 MODULE_PASS("cg-profile", CGProfilePass())
 MODULE_PASS("check-debugify", NewPMCheckDebugifyPass())
 MODULE_PASS("constmerge", ConstantMergePass())
 MODULE_PASS("coro-early", CoroEarlyPass())
-MODULE_PASS("coro-cleanup", CoroCleanupPass())
 MODULE_PASS("cross-dso-cfi", CrossDSOCFIPass())
 MODULE_PASS("deadargelim", DeadArgumentEliminationPass())
 MODULE_PASS("debugify", NewPMDebugifyPass())
 MODULE_PASS("dot-callgraph", CallGraphDOTPrinterPass())
 MODULE_PASS("elim-avail-extern", EliminateAvailableExternallyPass())
 MODULE_PASS("extract-blocks", BlockExtractorPass())
 MODULE_PASS("forceattrs", ForceFunctionAttrsPass())
 MODULE_PASS("function-import", FunctionImportPass())
MODULE_PASS_WITH_PARAMS("hwasan",
                         "kernel;recover")
 MODULE_PASS_WITH_PARAMS("asan-module",
                         "ModuleAddressSanitizerPass",
                         [](AddressSanitizerOptions Opts) {
                           return ModuleAddressSanitizerPass(Opts);
                         },
                         parseASanPassOptions,
                         "kernel")
+MODULE_PASS_WITH_PARAMS("coro-cleanup",
+                        "CoroCleanupPass",
+                        [](bool OptimizationEnabled) {
+                          return CoroCleanupPass(OptimizationEnabled);
+                        },
+                        parseCoroCleanupPassOptions,
+                        "opt")
 #undef MODULE_PASS_WITH_PARAMS
 
 #ifndef CGSCC_ANALYSIS
 #define CGSCC_ANALYSIS(NAME, CREATE_PASS)
 #endif
 CGSCC_ANALYSIS("no-op-cgscc", NoOpCGSCCAnalysis())
 CGSCC_ANALYSIS("fam-proxy", FunctionAnalysisManagerCGSCCProxy())
 CGSCC_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

 //===- CoroCleanup.cpp - Coroutine Cleanup 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Coroutines/CoroCleanup.h"
 #include "CoroInternal.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
+#include "llvm/Transforms/Scalar/EarlyCSE.h"
 #include "llvm/Transforms/Scalar/SimplifyCFG.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-cleanup"
 
 namespace {
 // Created on demand if CoroCleanup pass has work to do.
if (CB->getOperandBundle("coro_readnone")) {
       NewCB->copyMetadata(*CB);
       NewCB->addFnAttr(Attribute::ReadNone);
       NewCB->removeFnAttr(Attribute::ReadOnly);
       NewCB->removeFnAttr(Attribute::InaccessibleMemOnly);
 
       CB->replaceAllUsesWith(NewCB);
       CB->eraseFromParent();
       CB = NewCB;
+
+      Changed = true;
     }
 
     if (auto *II = dyn_cast<IntrinsicInst>(CB)) {
       switch (II->getIntrinsicID()) {
       default:
         continue;
       case Intrinsic::coro_begin:
         II->replaceAllUsesWith(II->getArgOperand(1));
PreservedAnalyses CoroCleanupPass::run(Module &M,
                                        ModuleAnalysisManager &MAM) {
   if (!declaresCoroCleanupIntrinsics(M))
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =
       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
 
   FunctionPassManager FPM;
+  FPM.addPass(EarlyCSEPass());
   FPM.addPass(SimplifyCFGPass());
 
   Lowerer L(M);
   for (auto &F : M)
-    if (L.lower(F))
+    if (L.lower(F) && OptimizationEnabled)
       FPM.run(F, FAM);
 
   return PreservedAnalyses::none();
 }

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

 
 ; CHECK-LABEL: define internal swiftcc void @my_async_functionTQ0_(i8* nocapture readonly swiftasync %0, i8* %1, i8* nocapture readnone %2)
 ; CHECK-O0-LABEL: define internal swiftcc void @my_async_functionTQ0_(i8* swiftasync %0, i8* %1, i8* %2)
 ; CHECK-SAME: !dbg ![[SP2:[0-9]+]] {
 ; CHECK: entryresume.0:
 ; CHECK:   [[CALLER_CONTEXT_ADDR:%.*]] = bitcast i8* %0 to i8**
 ; CHECK:   [[CALLER_CONTEXT:%.*]] = load i8*, i8** [[CALLER_CONTEXT_ADDR]]
 ; CHECK:   [[FRAME_PTR:%.*]] = getelementptr inbounds i8, i8* [[CALLER_CONTEXT]], i64 128
+; CHECK-O0:   [[VECTOR_SPILL_SPILL_ADDR:%.*]] = getelementptr inbounds %my_async_function.Frame, %my_async_function.Frame* {{.*}}, i32 0, i32 1
+; CHECK-O0:   store <4 x double> %vector_spill, <4 x double>* [[VECTOR_SPILL_SPILL_ADDR]]
 ; CHECK-O0:   [[VECTOR_SPILL_ADDR:%.*]] = getelementptr inbounds %my_async_function.Frame, %my_async_function.Frame* {{.*}}, i32 0, i32 1
 ; CHECK-O0:   load <4 x double>, <4 x double>* [[VECTOR_SPILL_ADDR]], align 16
 ; CHECK:   [[CALLEE_CTXT_SPILL_ADDR:%.*]] = getelementptr inbounds i8, i8* [[CALLER_CONTEXT]], i64 160
 ; CHECK:   [[CAST1:%.*]] = bitcast i8* [[CALLEE_CTXT_SPILL_ADDR]] to i8**
 ; CHECK:   [[CALLEE_CTXT_RELOAD:%.*]] = load i8*, i8** [[CAST1]]
 ; CHECK:   [[ACTOR_RELOAD_ADDR:%.*]] = getelementptr inbounds i8, i8* [[CALLER_CONTEXT]], i64 152
 ; CHECK:   [[CAST2:%.*]] = bitcast i8* [[ACTOR_RELOAD_ADDR]] to %async.actor**
 ; CHECK:   [[ACTOR_RELOAD:%.*]] = load %async.actor*, %async.actor** [[CAST2]]

llvm/test/Transforms/Coroutines/coro-debug-frame-variable.ll

-; RUN: opt < %s -passes='default<O0>' -S | FileCheck %s
+; RUN: opt < %s -passes='default<O0>',simplifycfg -S | FileCheck %s
 
 ; Define a function 'f' that resembles the Clang frontend's output for the
 ; following C++ coroutine:
 ;
 ;   void foo() {
 ;     int i = 0;
 ;     ++i;
 ;     int x = {};

llvm/test/Transforms/Coroutines/coro-readnone-06.ll

+; Tests that the readnone function which don't cross suspend points could be optimized expectly.
+; RUN: opt < %s -S -passes='default<O1>' -opaque-pointers | FileCheck %s
+
+define ptr @f() "coroutine.presplit" {
+entry:
+  %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)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [i8 0, label %resume
+                                         i8 1, label %cleanup]
+resume:
+  %i = call i32 @readnone_func() readnone
+  %j = call i32 @readnone_func() readnone
+  %cmp = icmp eq i32 %i, %j
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  br label %cleanup
+
+diff:
+  call void @print_diff()
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: f.resume(
+; CHECK-NEXT: entry
+; CHECK-NEXT:   call i32 @readnone_func()
+; CHECK-NEXT:   call void @print_same()
+; CHECK-NEXT:   call void @free
+; CHECK-NEXT:   ret void
+
+declare i32 @readnone_func() readnone
+
+declare void @print_same()
+declare void @print_diff()
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare void @free(i8*)

llvm/test/Transforms/Coroutines/coro-readnone-07.ll

+; Tests that the calls with "coro_readnone" operand bundles could be optimized
+; correctly by CoroCleanup pass.
+; RUN: opt < %s -S -passes='module(coro-cleanup<opt>)' -opaque-pointers | FileCheck %s
+
+define i1 @f() {
+entry:
+  %i = call i32 @readnone_func() [ "coro_readnone"() ]
+  %j = call i32 @readnone_func() [ "coro_readnone"() ]
+  %cmp = icmp eq i32 %i, %j
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  ret i1 true
+
+diff:
+  call void @print_diff()
+  ret i1 false
+}
+
+; CHECK-LABEL: @f(
+; CHECK-NEXT: entry
+; CHECK-NEXT:   call i32 @readnone_func()
+; CHECK-NEXT:   call void @print_same()
+; CHECK-NEXT:   ret i1 true
+
+declare void @f.resume()
+declare void @f.destroy()
+declare void @f.cleanup()
+
+declare i32 @readnone_func() readnone
+declare void @print_same()
+declare void @print_diff()
+
+declare ptr @llvm.coro.free(token, ptr nocapture readonly)
+declare i32 @llvm.coro.size.i32()
+declare i8 @llvm.coro.suspend(token, i1)
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr writeonly)
+declare i1 @llvm.coro.end(ptr, i1)
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)