Run Coverage pass before other *San passes under new pass manager, round 2

[aeubanks]

Summary

This was attempted once before in https://reviews.llvm.org/D79698, but
was reverted due to the coverage pass running in the wrong part of the
pipeline. This commit puts it in the same place as the other sanitizers.

This changes PassBuilder.OptimizerLastEPCallbacks to work on a
ModulePassManager instead of a FunctionPassManager. That is because
SanitizerCoverage cannot (easily) be split into a module pass and a
function pass like some of the other sanitizers since in its current
implementation it conditionally inserts module constructors based on
whether or not it successfully modified functions.

This fixes compiler-rt/test/msan/coverage-levels.cpp under the new pass
manager (last check-msan test).

[leonardchan]

llvm/include/llvm/Passes/PassBuilder.h

597–598

Will need to change the wording on this if this doesn't handle function passes anymore.

llvm/lib/Passes/PassBuilder.cpp

1079

Would it be better to add another SmallVector and register*Calback() for modules in in PassBuilder that we could loop through here instead of changing how these extension points work?

I imagine it would still be meaningful in the future to be able to add function passes at the end of the function optimization pipeline.

[aeubanks]

llvm/include/llvm/Passes/PassBuilder.h

597–598

This doesn't say that registerOptimizerLastEPCallback() handles function passes, it says it handles passes that will run right after all the default function optimization passes. I believe that's still true. The passes were previously added at the end of OptimizePM, now they're added right after OptimizePM, which is the same thing.

llvm/lib/Passes/PassBuilder.cpp

1079

You can still add function passes via createModuleToFunctionPassAdaptor(), which is what I've done here for the existing usages. Given that the number of calls to registerOptimizerLastEPCallback() is small, I don't see a huge benefit to create a both a vector of module passes and function passes that run at the same place. If in the future we end up calling registerOptimizerLastEPCallback() with lots of function passes we can revisit this but for now it seems unnecessary.

This change doesn't change the ordering of any passes aside from sanitizer coverage (I believe), all usages of registerOptimizerLastEPCallback() will still end up putting the passes in the same place in the pipeline.

[leonardchan]

Would be worthwhile to write a small test that asserts Sancov runs before the other sanitizers under the new PM?

Aside from this, LGTM.

llvm/include/llvm/Passes/PassBuilder.h

597–598

Ah I see. You're right. I misread.

llvm/lib/Passes/PassBuilder.cpp

1079

I see. Makes sense.

[aeubanks]

In D80692#2061219, @leonardchan wrote:

Would be worthwhile to write a small test that asserts Sancov runs before the other sanitizers under the new PM?

Aside from this, LGTM.

Asserting that sancov runs before other sanitizers seems like overspecification to me. Really all we want to know is that something like msan + sancov work together, and there are tests for that which we are fixing with this.

[aeubanks]

I foolishly submitted without running check-all, and it turns out this broke a test under check-clang.
Looks like https://reviews.llvm.org/D62888 added tests in sancov-new-pm.c to make sure that sancov + LTO work together, but they won't after this change since the (Thin)LTO pipeline doesn't run things under registerOptimizerLastEPCallback(). So I deleted those tests. Other sanitizers don't work under LTO right now (at least for non-O0) because of the same thing. Does this make sense?

[leonardchan]

In D80692#2061739, @aeubanks wrote:

I foolishly submitted without running check-all, and it turns out this broke a test under check-clang.
Looks like https://reviews.llvm.org/D62888 added tests in sancov-new-pm.c to make sure that sancov + LTO work together, but they won't after this change since the (Thin)LTO pipeline doesn't run things under registerOptimizerLastEPCallback(). So I deleted those tests. Other sanitizers don't work under LTO right now (at least for non-O0) because of the same thing. Does this make sense?

Makes sense. Could you also remove any CHECK lines for RUNs that were removed?

[aeubanks]

Update some CHECKs

clang/lib/CodeGen/BackendUtil.cpp

 #include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/SchedulerRegistry.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ModuleSummaryIndex.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/LTO/LTOBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/BuryPointer.h"
static void addCoroutinePassesAtO0(ModulePassManager &MPM,
 
   MPM.addPass(createModuleToFunctionPassAdaptor(CoroCleanupPass()));
 }
 
 static void addSanitizersAtO0(ModulePassManager &MPM,
                               const Triple &TargetTriple,
                               const LangOptions &LangOpts,
                               const CodeGenOptions &CodeGenOpts) {
+  if (CodeGenOpts.SanitizeCoverageType ||
+      CodeGenOpts.SanitizeCoverageIndirectCalls ||
+      CodeGenOpts.SanitizeCoverageTraceCmp) {
+    auto SancovOpts = getSancovOptsFromCGOpts(CodeGenOpts);
+    MPM.addPass(ModuleSanitizerCoveragePass(
+        SancovOpts, CodeGenOpts.SanitizeCoverageWhitelistFiles,
+        CodeGenOpts.SanitizeCoverageBlacklistFiles));
+  }
+
   auto ASanPass = [&](SanitizerMask Mask, bool CompileKernel) {
     MPM.addPass(RequireAnalysisPass<ASanGlobalsMetadataAnalysis, Module>());
     bool Recover = CodeGenOpts.SanitizeRecover.has(Mask);
     MPM.addPass(createModuleToFunctionPassAdaptor(AddressSanitizerPass(
         CompileKernel, Recover, CodeGenOpts.SanitizeAddressUseAfterScope)));
     bool ModuleUseAfterScope = asanUseGlobalsGC(TargetTriple, CodeGenOpts);
     MPM.addPass(
         ModuleAddressSanitizerPass(CompileKernel, Recover, ModuleUseAfterScope,
if (CodeGenOpts.OptimizationLevel == 0) {
       // Register callbacks to schedule sanitizer passes at the appropriate part of
       // the pipeline.
       // FIXME: either handle asan/the remaining sanitizers or error out
       if (LangOpts.Sanitize.has(SanitizerKind::LocalBounds))
         PB.registerScalarOptimizerLateEPCallback(
             [](FunctionPassManager &FPM, PassBuilder::OptimizationLevel Level) {
               FPM.addPass(BoundsCheckingPass());
             });
+
+      if (CodeGenOpts.SanitizeCoverageType ||
+          CodeGenOpts.SanitizeCoverageIndirectCalls ||
+          CodeGenOpts.SanitizeCoverageTraceCmp) {
+        PB.registerOptimizerLastEPCallback(
+            [this](ModulePassManager &MPM,
+                   PassBuilder::OptimizationLevel Level) {
+              auto SancovOpts = getSancovOptsFromCGOpts(CodeGenOpts);
+              MPM.addPass(ModuleSanitizerCoveragePass(
+                  SancovOpts, CodeGenOpts.SanitizeCoverageWhitelistFiles,
+                  CodeGenOpts.SanitizeCoverageBlacklistFiles));
+            });
+      }
+
       if (LangOpts.Sanitize.has(SanitizerKind::Memory)) {
         int TrackOrigins = CodeGenOpts.SanitizeMemoryTrackOrigins;
         bool Recover = CodeGenOpts.SanitizeRecover.has(SanitizerKind::Memory);
         PB.registerPipelineStartEPCallback(
             [TrackOrigins, Recover](ModulePassManager &MPM) {
               MPM.addPass(MemorySanitizerPass({TrackOrigins, Recover, false}));
             });
         PB.registerOptimizerLastEPCallback(
-            [TrackOrigins, Recover](FunctionPassManager &FPM,
+            [TrackOrigins, Recover](ModulePassManager &MPM,
                                     PassBuilder::OptimizationLevel Level) {
-              FPM.addPass(MemorySanitizerPass({TrackOrigins, Recover, false}));
+              MPM.addPass(createModuleToFunctionPassAdaptor(
+                  MemorySanitizerPass({TrackOrigins, Recover, false})));
             });
       }
       if (LangOpts.Sanitize.has(SanitizerKind::Thread)) {
         PB.registerPipelineStartEPCallback(
             [](ModulePassManager &MPM) { MPM.addPass(ThreadSanitizerPass()); });
         PB.registerOptimizerLastEPCallback(
-            [](FunctionPassManager &FPM, PassBuilder::OptimizationLevel Level) {
-              FPM.addPass(ThreadSanitizerPass());
+            [](ModulePassManager &MPM, PassBuilder::OptimizationLevel Level) {
+              MPM.addPass(
+                  createModuleToFunctionPassAdaptor(ThreadSanitizerPass()));
             });
       }
       if (LangOpts.Sanitize.has(SanitizerKind::Address)) {
         PB.registerPipelineStartEPCallback([&](ModulePassManager &MPM) {
           MPM.addPass(
               RequireAnalysisPass<ASanGlobalsMetadataAnalysis, Module>());
         });
         bool Recover = CodeGenOpts.SanitizeRecover.has(SanitizerKind::Address);
         bool UseAfterScope = CodeGenOpts.SanitizeAddressUseAfterScope;
         PB.registerOptimizerLastEPCallback(
-            [Recover, UseAfterScope](FunctionPassManager &FPM,
+            [Recover, UseAfterScope](ModulePassManager &MPM,
                                      PassBuilder::OptimizationLevel Level) {
-              FPM.addPass(AddressSanitizerPass(
-                  /*CompileKernel=*/false, Recover, UseAfterScope));
+              MPM.addPass(
+                  createModuleToFunctionPassAdaptor(AddressSanitizerPass(
+                      /*CompileKernel=*/false, Recover, UseAfterScope)));
             });
         bool ModuleUseAfterScope = asanUseGlobalsGC(TargetTriple, CodeGenOpts);
         bool UseOdrIndicator = CodeGenOpts.SanitizeAddressUseOdrIndicator;
         PB.registerPipelineStartEPCallback(
             [Recover, ModuleUseAfterScope,
              UseOdrIndicator](ModulePassManager &MPM) {
               MPM.addPass(ModuleAddressSanitizerPass(
                   /*CompileKernel=*/false, Recover, ModuleUseAfterScope,
if (CodeGenOpts.OptimizationLevel == 0) {
         MPM.addPass(CanonicalizeAliasesPass());
         MPM.addPass(NameAnonGlobalPass());
       } else {
         MPM = PB.buildPerModuleDefaultPipeline(Level,
                                                CodeGenOpts.DebugPassManager);
       }
     }
 
-    if (CodeGenOpts.SanitizeCoverageType ||
-        CodeGenOpts.SanitizeCoverageIndirectCalls ||
-        CodeGenOpts.SanitizeCoverageTraceCmp) {
-      auto SancovOpts = getSancovOptsFromCGOpts(CodeGenOpts);
-      MPM.addPass(ModuleSanitizerCoveragePass(
-          SancovOpts, CodeGenOpts.SanitizeCoverageWhitelistFiles,
-          CodeGenOpts.SanitizeCoverageBlacklistFiles));
-    }
-
     if (LangOpts.Sanitize.has(SanitizerKind::HWAddress)) {
       bool Recover = CodeGenOpts.SanitizeRecover.has(SanitizerKind::HWAddress);
       MPM.addPass(HWAddressSanitizerPass(
           /*CompileKernel=*/false, Recover));
     }
     if (LangOpts.Sanitize.has(SanitizerKind::KernelHWAddress)) {
       MPM.addPass(HWAddressSanitizerPass(
           /*CompileKernel=*/true, /*Recover=*/true));

clang/test/CodeGen/sancov-new-pm.c

 // Test that SanitizerCoverage works under the new pass manager.
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=fuzzer %s -fexperimental-new-pass-manager -S -emit-llvm -o - | FileCheck %s --check-prefixes=CHECK,CHECK-O0
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=fuzzer %s -fexperimental-new-pass-manager -O2 -S -emit-llvm -o - | FileCheck %s --check-prefixes=CHECK,CHECK-O2
-// RUN: %clang -target x86_64-linux-gnu -fsanitize=fuzzer %s -fexperimental-new-pass-manager -flto -S -emit-llvm -o - | FileCheck %s --check-prefixes=CHECK,CHECK-O0
-// RUN: %clang -target x86_64-linux-gnu -fsanitize=fuzzer %s -fexperimental-new-pass-manager -flto -O2 -S -emit-llvm -o - | FileCheck %s --check-prefixes=CHECK,CHECK-O2
-// RUN: %clang -target x86_64-linux-gnu -fsanitize=fuzzer %s -fexperimental-new-pass-manager -flto=thin -S -emit-llvm -o - | FileCheck %s --check-prefixes=CHECK,CHECK-O0
-// RUN: %clang -target x86_64-linux-gnu -fsanitize=fuzzer %s -fexperimental-new-pass-manager -flto=thin -O2 -S -emit-llvm -o - | FileCheck %s --check-prefixes=CHECK,CHECK-O2,CHECK-O2-THINLTO
 
 extern void *memcpy(void *, const void *, unsigned long);
 extern int printf(const char *restrict, ...);
 
 int LLVMFuzzerTestOneInput(const unsigned char *data, unsigned long size) {
   unsigned char buf[4];
 
   if (size < 8)
     return 0;
 
   if (data[0] == 'h' && data[1] == 'i' && data[2] == '!') {
     memcpy(buf, data, size);
     printf("test: %.2X\n", buf[0]);
   }
 
   return 0;
 }
 
 // CHECK-DAG: declare void @__sanitizer_cov_pcs_init(i64*, i64*)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_pc_indir(i64)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_cmp1(i8 zeroext, i8 zeroext)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_cmp2(i16 zeroext, i16 zeroext)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_cmp4(i32 zeroext, i32 zeroext)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_cmp8(i64, i64)
-// CHECK-O2-THINLTO-NOT: declare void @__sanitizer_cov_trace_const_cmp1(i8 zeroext, i8 zeroext)
+// CHECK-O2-NOT: declare void @__sanitizer_cov_trace_const_cmp1(i8 zeroext, i8 zeroext)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_const_cmp2(i16 zeroext, i16 zeroext)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_const_cmp4(i32 zeroext, i32 zeroext)
-// CHECK-O2-THINLTO-NOT: declare void @__sanitizer_cov_trace_const_cmp8(i64, i64)
+// CHECK-O2-NOT: declare void @__sanitizer_cov_trace_const_cmp8(i64, i64)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_div4(i32 zeroext)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_div8(i64)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_gep(i64)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_switch(i64, i64*)
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_pc()
 // CHECK-O0-DAG: declare void @__sanitizer_cov_trace_pc_guard(i32*)

llvm/include/llvm/Passes/PassBuilder.h

public:
   /// link-time pipelines).
   void registerPipelineStartEPCallback(
       const std::function<void(ModulePassManager &)> &C) {
     PipelineStartEPCallbacks.push_back(C);
   }
 
   /// Register a callback for a default optimizer pipeline extension point
   ///
   /// This extension point allows adding optimizations at the very end of the
   /// function optimization pipeline. A key difference between this and the
   /// legacy PassManager's OptimizerLast callback is that this extension point
   /// is not triggered at O0. Extensions to the O0 pipeline should append their
   /// passes to the end of the overall pipeline.
   void registerOptimizerLastEPCallback(
-      const std::function<void(FunctionPassManager &, OptimizationLevel)> &C) {
+      const std::function<void(ModulePassManager &, OptimizationLevel)> &C) {
     OptimizerLastEPCallbacks.push_back(C);
   }
 
   /// Register a callback for parsing an AliasAnalysis Name to populate
   /// the given AAManager \p AA
   void registerParseAACallback(
       const std::function<bool(StringRef Name, AAManager &AA)> &C) {
     AAParsingCallbacks.push_back(C);
private:
   SmallVector<std::function<void(LoopPassManager &, OptimizationLevel)>, 2>
       LoopOptimizerEndEPCallbacks;
   SmallVector<std::function<void(FunctionPassManager &, OptimizationLevel)>, 2>
       ScalarOptimizerLateEPCallbacks;
   SmallVector<std::function<void(CGSCCPassManager &, OptimizationLevel)>, 2>
       CGSCCOptimizerLateEPCallbacks;
   SmallVector<std::function<void(FunctionPassManager &, OptimizationLevel)>, 2>
       VectorizerStartEPCallbacks;
-  SmallVector<std::function<void(FunctionPassManager &, OptimizationLevel)>, 2>
+  SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
       OptimizerLastEPCallbacks;
   // Module callbacks
   SmallVector<std::function<void(ModulePassManager &)>, 2>
       PipelineStartEPCallbacks;
   SmallVector<std::function<void(ModuleAnalysisManager &)>, 2>
       ModuleAnalysisRegistrationCallbacks;
   SmallVector<std::function<bool(StringRef, ModulePassManager &,
                                  ArrayRef<PipelineElement>)>,

llvm/lib/Passes/PassBuilder.cpp

ModulePassManager PassBuilder::buildModuleOptimizationPipeline(
   // Optimize PHIs by speculating around them when profitable. Note that this
   // pass needs to be run after any PRE or similar pass as it is essentially
   // inserting redundancies into the program. This even includes SimplifyCFG.
   OptimizePM.addPass(SpeculateAroundPHIsPass());
 
   if (PTO.Coroutines)
     OptimizePM.addPass(CoroCleanupPass());
 
-  for (auto &C : OptimizerLastEPCallbacks)
-    C(OptimizePM, Level);
-
   // Add the core optimizing pipeline.
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(OptimizePM)));
 
+  for (auto &C : OptimizerLastEPCallbacks)
+    C(MPM, Level);
+
   if (PTO.CallGraphProfile)
     MPM.addPass(CGProfilePass());
 
   // Now we need to do some global optimization transforms.
   // FIXME: It would seem like these should come first in the optimization
   // pipeline and maybe be the bottom of the canonicalization pipeline? Weird
   // ordering here.
   MPM.addPass(GlobalDCEPass());

llvm/tools/opt/NewPMDriver.cpp

if (tryParsePipelineText<ModulePassManager>(PB, PipelineStartEPPipeline))
     PB.registerPipelineStartEPCallback(
         [&PB, VerifyEachPass, DebugLogging](ModulePassManager &PM) {
           ExitOnError Err("Unable to parse PipelineStartEP pipeline: ");
           Err(PB.parsePassPipeline(PM, PipelineStartEPPipeline, VerifyEachPass,
                                    DebugLogging));
         });
   if (tryParsePipelineText<FunctionPassManager>(PB, OptimizerLastEPPipeline))
     PB.registerOptimizerLastEPCallback(
-        [&PB, VerifyEachPass, DebugLogging](FunctionPassManager &PM,
+        [&PB, VerifyEachPass, DebugLogging](ModulePassManager &PM,
                                             PassBuilder::OptimizationLevel) {
           ExitOnError Err("Unable to parse OptimizerLastEP pipeline: ");
           Err(PB.parsePassPipeline(PM, OptimizerLastEPPipeline, VerifyEachPass,
                                    DebugLogging));
         });
 }
 
 #define HANDLE_EXTENSION(Ext)                                                  \