[NPM] Ability to add a pass before a previously registered one

[pratlucas]

Summary

This patch introduces the addPassBefore method to PassManager,
allowing a pass to be registered before another specific one that
may have been registered before it.

This is specially useful for pass plugins, which do not go through the
Pass Registry and lack some of the granularity for controlling when
they should be run.

[aeubanks]

I'd rather not have something like this, especially something that depends on pass names.
What exactly are you trying to do with this?

[pratlucas]

Hi @aeubanks,
As I mentioned, this is mostly targeted on plugin passes. At the moment the NPM only offers the PassBuilder's callbacks level of granularity to control when a plugin pass should be run, which might not be sufficient on some cases. The goal was to provide a more fine grained control over this, allowing plugin implementations to run before other native or plugin passes that migth impact their behavior.
Do you have any suggestions on alternative ways to achieve this? I'm not fond of using pass names here myself and would be more than happy to go for a different implementation.

[aeubanks]

In D102909#2776635, @pratlucas wrote:

Hi @aeubanks,
As I mentioned, this is mostly targeted on plugin passes. At the moment the NPM only offers the PassBuilder's callbacks level of granularity to control when a plugin pass should be run, which might not be sufficient on some cases. The goal was to provide a more fine grained control over this, allowing plugin implementations to run before other native or plugin passes that migth impact their behavior.
Do you have any suggestions on alternative ways to achieve this? I'm not fond of using pass names here myself and would be more than happy to go for a different implementation.

I'd still like to understand exactly what you'd be using this for. As in why does some pass need to run before all other instances of some other pass?

[pratlucas]

I'd still like to understand exactly what you'd be using this for. As in why does some pass need to run before all other instances of some other pass?

The concrete use case I have is based on two separate plugin passes - one of them is responsible for inserting a few inline assembly instructions into functions according to a specific criteria and the second one's job is to make some adjustments to all the inline asm instructions in the module. The former needs to run before the latter to ensure all inline asm instructions are in place before the adjustments take place.
Both of those plugin passes are required to run on the last step of the optimizer pipeline, using the registerOptimizerLastEPCallback extension point, so there's not enough granularity available to ensure those are run in this specific order.

[aeubanks]

In D102909#2808270, @pratlucas wrote:

I'd still like to understand exactly what you'd be using this for. As in why does some pass need to run before all other instances of some other pass?

The concrete use case I have is based on two separate plugin passes - one of them is responsible for inserting a few inline assembly instructions into functions according to a specific criteria and the second one's job is to make some adjustments to all the inline asm instructions in the module. The former needs to run before the latter to ensure all inline asm instructions are in place before the adjustments take place.
Both of those plugin passes are required to run on the last step of the optimizer pipeline, using the registerOptimizerLastEPCallback extension point, so there's not enough granularity available to ensure those are run in this specific order.

If you add both in the same callback passed to registerOptimizerLastEPCallback() you can control the order they're added, e.g.

PB.registerOptimizerLastEPCallback(
    [this](ModulePassManager &PM, PassBuilder::OptimizationLevel Level) {
      FunctionPassManager FPM;
      FPM.addPass(BarPass());
      FPM.addPass(FooPass());
      PM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
    });

[pratlucas]

This solution works well if both passes are always supposed to be run together. If they should also work as standalone passes, it becomes tricky to figure out at what point each of them should be registered.
This could be done using the legacy pass manager on a hacky kind of way by querying for the passes on the PassRegistry, but this is no longer an option in NPM.

[aeubanks]

In D102909#2808270, @pratlucas wrote:

The concrete use case I have is based on two separate plugin passes - one of them is responsible for inserting a few inline assembly instructions into functions according to a specific criteria and the second one's job is to make some adjustments to all the inline asm instructions in the module. The former needs to run before the latter to ensure all inline asm instructions are in place before the adjustments take place.
Both of those plugin passes are required to run on the last step of the optimizer pipeline, using the registerOptimizerLastEPCallback extension point, so there's not enough granularity available to ensure those are run in this specific order.

I'm still not understanding exactly what the criteria you're looking for is. Could you give a specific example of how your example could go wrong using registerOptimizerLastEPCallback? Even if they are to be registered separately via multiple calls to registerOptimizerLastEPCallback, the first registerOptimizerLastEPCallback should happen first.

[pratlucas]

The limitation with the registerOptimizerLastEPCallback ordering is that, when they come from different plugins - i.e. different shared libraries -, it is currently based solely on the order of command line options and can't be defined in the code.

I've managed to untangle the dependency on my concrete scenario, so I'm abandoning this change whilst I think of an alternative for it.

llvm/include/llvm/IR/PassManager.h

 //===- PassManager.h - Pass management infrastructure -----------*- 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
 ///
public:
   /// happen with nested pass managers of the same type.
   template <typename PassT>
   std::enable_if_t<std::is_same<PassT, PassManager>::value>
   addPass(PassT &&Pass) {
     for (auto &P : Pass.Passes)
       Passes.emplace_back(std::move(P));
   }
 
+  template <typename PassT>
+  std::enable_if_t<!std::is_same<PassT, PassManager>::value>
+  addPassBefore(PassT Pass, StringRef Other) {
+    using PassModelT =
+        detail::PassModel<IRUnitT, PassT, PreservedAnalyses, AnalysisManagerT,
+                          ExtraArgTs...>;
+    auto It = std::find_if(Passes.begin(), Passes.end(),
+      [&](const auto& Pass) { return Pass->name() == Other; });
+    Passes.emplace(It, new PassModelT(std::move(Pass)));
+  }
+
   /// Returns if the pass manager contains any passes.
   bool isEmpty() const { return Passes.empty(); }
 
   static bool isRequired() { return true; }
 
 protected:
   using PassConceptT =
       detail::PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...>;

llvm/unittests/IR/PassManagerTest.cpp

 //===- llvm/unittest/IR/PassManager.cpp - PassManager tests ---------------===//
 //
 // 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/IR/PassManager.h"
struct LambdaPass : public PassInfoMixin<LambdaPass> {
 
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM) {
     return Func(F, AM);
   }
 
   FuncT Func;
 };
 
+struct LambdaModulePass : public PassInfoMixin<LambdaPass> {
+  using FuncT = std::function<PreservedAnalyses(Module &, ModuleAnalysisManager &)>;
+
+  LambdaModulePass(FuncT Func) : Func(std::move(Func)) {}
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
+    return Func(M, AM);
+  }
+
+  FuncT Func;
+};
+
 TEST_F(PassManagerTest, IndirectAnalysisInvalidation) {
   FunctionAnalysisManager FAM;
   ModuleAnalysisManager MAM;
   int FunctionAnalysisRuns = 0, ModuleAnalysisRuns = 0,
       IndirectAnalysisRuns = 0, DoublyIndirectAnalysisRuns = 0;
   FAM.registerPass([&] { return TestFunctionAnalysis(FunctionAnalysisRuns); });
   FAM.registerPass(
       [&] { return TestIndirectFunctionAnalysis(IndirectAnalysisRuns, MAM); });
TEST_F(PassManagerTest, FunctionPassCFGCheckerWrapped) {
   FAM.registerPass([&] { return TargetIRAnalysis(); });
 
   FunctionPassManager InnerFPM;
   InnerFPM.addPass(SimplifyCFGPass());
 
   FPM.addPass(TestSimplifyCFGWrapperPass(InnerFPM));
   FPM.run(*F, FAM);
 }
+
+TEST_F(PassManagerTest, AddPassBefore) {
+  ModuleAnalysisManager MAM;
+  int ModuleAnalysisRuns = 0;
+  PassInstrumentationCallbacks PIC;
+  MAM.registerPass([&] { return TestModuleAnalysis(ModuleAnalysisRuns); });
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+
+  ModulePassManager MPM1;
+
+  int TestModulePassRunCount1 = 0;
+  int CustomModulePassRunCount1 = 0;
+  MPM1.addPass(TestModulePass(TestModulePassRunCount1));
+  MPM1.addPass(LambdaModulePass([&](Module &, ModuleAnalysisManager &){
+    EXPECT_EQ(1, TestModulePassRunCount1);
+    CustomModulePassRunCount1++;
+    return PreservedAnalyses::all();
+  }));
+
+  MPM1.run(*M, MAM);
+  EXPECT_EQ(1, TestModulePassRunCount1);
+  EXPECT_EQ(1, CustomModulePassRunCount1);
+
+  ModulePassManager MPM2;
+
+  int TestModulePassRunCount2 = 0;
+  int CustomModulePassRunCount2 = 0;
+  MPM2.addPass(TestModulePass(TestModulePassRunCount2));
+  MPM2.addPassBefore(LambdaModulePass([&](Module &, ModuleAnalysisManager &){
+    EXPECT_EQ(0, TestModulePassRunCount2);
+    CustomModulePassRunCount2++;
+    return PreservedAnalyses::all();
+  }), "{anonymous}::TestModulePass");
+
+  MPM2.run(*M, MAM);
+  EXPECT_EQ(1, TestModulePassRunCount2);
+  EXPECT_EQ(1, CustomModulePassRunCount2);
+}
+
 }