[PM] Make the the new pass manager support fully generic extra arguments

to run methods, both for transform passes and analysis passes.

This also allows the analysis manager to use a different set of extra
arguments from the pass manager where useful. Consider passes over
analysis produced units of IR like SCCs of the call graph or loops.
Passes of this nature will often want to refer to the analysis result
that was used to compute their IR units (the call graph or LoopInfo).
And for transformations, they may want to communicate special update
information to the outer pass manager. With this change, it becomes
possible to have a run method for a loop pass that looks more like:

  PreservedAnalyses run(Loop &L, AnalysisManager<Loop, LoopInfo> &AM,
                        LoopInfo &LI, LoopUpdateRecord &UR);

And to query the analysis manager like:

    AM.getResult<MyLoopAnalysis>(L, LI);

This makes accessing the known-available analyses convenient and clear,
and it makes passing customized data structures around easy.

My initial use case is going to be in updating the pass manager layers
when the analysis units of IR change. But there are more use cases here
such as having a layer that lets inner passes signal whether certain
additional passes should be run because of particular simplifications
made. Two desires for this have come up in the past: triggering
additional optimization after successfully unrolling loops, and
triggering additional inlining after collapsing indirect calls to direct
calls.

Despite adding this layer of generic extensibility, the *only* change to
existing, simple usage are for places where we forward declare the
AnalysisManager template. We really shouldn't be doing this because of
the fragility exposed here, but currently it makes coping with the
legacy PM code easier.

Differential Revision: http://reviews.llvm.org/D21462

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@279227 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/IR/IRPrintingPasses.h

@@ -30,7 +30,7 @@ class Module;
 class ModulePass;
 class PreservedAnalyses;
 class raw_ostream;
-template <typename IRUnitT> class AnalysisManager;
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
 
 /// \brief Create and return a pass that writes the module to the specified
 /// \c raw_ostream.

include/llvm/IR/PassManager.h

@@ -171,7 +171,7 @@ private:
 };
 
 // Forward declare the analysis manager template.
-template <typename IRUnitT> class AnalysisManager;
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
 
 /// A CRTP mix-in to automatically provide informational APIs needed for
 /// passes.
@@ -222,8 +222,11 @@ struct AnalysisInfoMixin : PassInfoMixin<DerivedT> {
 /// that analysis manager to each pass it runs, as well as calling the analysis
 /// manager's invalidation routine with the PreservedAnalyses of each pass it
 /// runs.
-template <typename IRUnitT>
-class PassManager : public PassInfoMixin<PassManager<IRUnitT>> {
+template <typename IRUnitT,
+          typename AnalysisManagerT = AnalysisManager<IRUnitT>,
+          typename... ExtraArgTs>
+class PassManager : public PassInfoMixin<
+                        PassManager<IRUnitT, AnalysisManagerT, ExtraArgTs...>> {
 public:
   /// \brief Construct a pass manager.
   ///
@@ -241,7 +244,8 @@ public:
   }
 
   /// \brief Run all of the passes in this manager over the IR.
-  PreservedAnalyses run(IRUnitT &IR, AnalysisManager<IRUnitT> &AM) {
+  PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
+                        ExtraArgTs... ExtraArgs) {
     PreservedAnalyses PA = PreservedAnalyses::all();
 
     if (DebugLogging)
@@ -252,7 +256,7 @@ public:
         dbgs() << "Running pass: " << Passes[Idx]->name() << " on "
                << IR.getName() << "\n";
 
-      PreservedAnalyses PassPA = Passes[Idx]->run(IR, AM);
+      PreservedAnalyses PassPA = Passes[Idx]->run(IR, AM, ExtraArgs...);
 
       // Update the analysis manager as each pass runs and potentially
       // invalidates analyses. We also update the preserved set of analyses
@@ -278,12 +282,15 @@ public:
   }
 
   template <typename PassT> void addPass(PassT Pass) {
-    typedef detail::PassModel<IRUnitT, PassT> PassModelT;
+    typedef detail::PassModel<IRUnitT, PassT, PreservedAnalyses,
+                              AnalysisManagerT, ExtraArgTs...>
+        PassModelT;
     Passes.emplace_back(new PassModelT(std::move(Pass)));
   }
 
 private:
-  typedef detail::PassConcept<IRUnitT> PassConceptT;
+  typedef detail::PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...>
+      PassConceptT;
 
   PassManager(const PassManager &) = delete;
   PassManager &operator=(const PassManager &) = delete;
@@ -308,10 +315,10 @@ typedef PassManager<Function> FunctionPassManager;
 /// This analysis manager can be used for any IR unit where the address of the
 /// IR unit sufficies as its identity. It manages the cache for a unit of IR via
 /// the address of each unit of IR cached.
-template <typename IRUnitT>
+template <typename IRUnitT, typename... ExtraArgTs>
 class AnalysisManager {
   typedef detail::AnalysisResultConcept<IRUnitT> ResultConceptT;
-  typedef detail::AnalysisPassConcept<IRUnitT> PassConceptT;
+  typedef detail::AnalysisPassConcept<IRUnitT, ExtraArgTs...> PassConceptT;
 
 public:
   // Most public APIs are inherited from the CRTP base class.
@@ -358,11 +365,11 @@ public:
   ///
   /// If there is not a valid cached result in the manager already, this will
   /// re-run the analysis to produce a valid result.
-  template <typename PassT> typename PassT::Result &getResult(IRUnitT &IR) {
+  template <typename PassT>
+  typename PassT::Result &getResult(IRUnitT &IR, ExtraArgTs... ExtraArgs) {
     assert(AnalysisPasses.count(PassT::ID()) &&
            "This analysis pass was not registered prior to being queried");
-
-    ResultConceptT &ResultConcept = getResultImpl(PassT::ID(), IR);
+    ResultConceptT &ResultConcept = getResultImpl(PassT::ID(), IR, ExtraArgs...);
     typedef detail::AnalysisResultModel<IRUnitT, PassT, typename PassT::Result>
         ResultModelT;
     return static_cast<ResultModelT &>(ResultConcept).Result;
@@ -407,7 +414,7 @@ public:
   /// away.
   template <typename PassBuilderT> bool registerPass(PassBuilderT PassBuilder) {
     typedef decltype(PassBuilder()) PassT;
-    typedef detail::AnalysisPassModel<IRUnitT, PassT> PassModelT;
+    typedef detail::AnalysisPassModel<IRUnitT, PassT, ExtraArgTs...> PassModelT;
 
     auto &PassPtr = AnalysisPasses[PassT::ID()];
     if (PassPtr)
@@ -502,7 +509,8 @@ private:
   }
 
   /// \brief Get an analysis result, running the pass if necessary.
-  ResultConceptT &getResultImpl(void *PassID, IRUnitT &IR) {
+  ResultConceptT &getResultImpl(void *PassID, IRUnitT &IR,
+                                ExtraArgTs... ExtraArgs) {
     typename AnalysisResultMapT::iterator RI;
     bool Inserted;
     std::tie(RI, Inserted) = AnalysisResults.insert(std::make_pair(
@@ -515,7 +523,7 @@ private:
       if (DebugLogging)
         dbgs() << "Running analysis: " << P.name() << "\n";
       AnalysisResultListT &ResultList = AnalysisResultLists[&IR];
-      ResultList.emplace_back(PassID, P.run(IR, *this));
+      ResultList.emplace_back(PassID, P.run(IR, *this, ExtraArgs...));
 
       // P.run may have inserted elements into AnalysisResults and invalidated
       // RI.
@@ -607,7 +615,7 @@ typedef AnalysisManager<Function> FunctionAnalysisManager;
 /// Note that the proxy's result is a move-only object and represents ownership
 /// of the validity of the analyses in the \c FunctionAnalysisManager it
 /// provides.
-template <typename AnalysisManagerT, typename IRUnitT>
+template <typename AnalysisManagerT, typename IRUnitT, typename... ExtraArgTs>
 class InnerAnalysisManagerProxy
     : public AnalysisInfoMixin<
           InnerAnalysisManagerProxy<AnalysisManagerT, IRUnitT>> {
@@ -689,7 +697,10 @@ public:
   /// In debug builds, it will also assert that the analysis manager is empty
   /// as no queries should arrive at the function analysis manager prior to
   /// this analysis being requested.
-  Result run(IRUnitT &IR, AnalysisManager<IRUnitT> &) { return Result(*AM); }
+  Result run(IRUnitT &IR, AnalysisManager<IRUnitT, ExtraArgTs...> &,
+             ExtraArgTs...) {
+    return Result(*AM);
+  }
 
 private:
   friend AnalysisInfoMixin<
@@ -699,8 +710,9 @@ private:
   AnalysisManagerT *AM;
 };
 
-template <typename AnalysisManagerT, typename IRUnitT>
-char InnerAnalysisManagerProxy<AnalysisManagerT, IRUnitT>::PassID;
+template <typename AnalysisManagerT, typename IRUnitT, typename... ExtraArgTs>
+char
+    InnerAnalysisManagerProxy<AnalysisManagerT, IRUnitT, ExtraArgTs...>::PassID;
 
 extern template class InnerAnalysisManagerProxy<FunctionAnalysisManager,
                                                 Module>;
@@ -720,7 +732,7 @@ typedef InnerAnalysisManagerProxy<FunctionAnalysisManager, Module>
 /// This proxy *doesn't* manage the invalidation in any way. That is handled by
 /// the recursive return path of each layer of the pass manager and the
 /// returned PreservedAnalysis set.
-template <typename AnalysisManagerT, typename IRUnitT>
+template <typename AnalysisManagerT, typename IRUnitT, typename... ExtraArgTs>
 class OuterAnalysisManagerProxy
     : public AnalysisInfoMixin<
           OuterAnalysisManagerProxy<AnalysisManagerT, IRUnitT>> {
@@ -762,7 +774,10 @@ public:
   /// \brief Run the analysis pass and create our proxy result object.
   /// Nothing to see here, it just forwards the \c AM reference into the
   /// result.
-  Result run(IRUnitT &, AnalysisManager<IRUnitT> &) { return Result(*AM); }
+  Result run(IRUnitT &, AnalysisManager<IRUnitT, ExtraArgTs...> &,
+             ExtraArgTs...) {
+    return Result(*AM);
+  }
 
 private:
   friend AnalysisInfoMixin<
@@ -772,8 +787,9 @@ private:
   const AnalysisManagerT *AM;
 };
 
-template <typename AnalysisManagerT, typename IRUnitT>
-char OuterAnalysisManagerProxy<AnalysisManagerT, IRUnitT>::PassID;
+template <typename AnalysisManagerT, typename IRUnitT, typename... ExtraArgTs>
+char
+    OuterAnalysisManagerProxy<AnalysisManagerT, IRUnitT, ExtraArgTs...>::PassID;
 
 extern template class OuterAnalysisManagerProxy<ModuleAnalysisManager,
                                                 Function>;
@@ -874,15 +890,30 @@ createModuleToFunctionPassAdaptor(FunctionPassT Pass) {
 ///
 /// This is a no-op pass which simply forces a specific analysis pass's result
 /// to be available when it is run.
-template <typename AnalysisT>
-struct RequireAnalysisPass : PassInfoMixin<RequireAnalysisPass<AnalysisT>> {
+template <typename AnalysisT, typename IRUnitT,
+          typename AnalysisManagerT = AnalysisManager<IRUnitT>,
+          typename... ExtraArgTs>
+struct RequireAnalysisPass;
+
+/// A specialization of the RequireAnalysisPass for generic IR unit analysis
+/// managers and pass managers that have no extra arguments.
+///
+/// If there are extra arguments at the pass's run level there may also be
+/// extra arguments to the analysis manager's \c getResult routine. We can't
+/// guess how to effectively map the arguments from one to the other, and so
+/// only the specialization with no extra arguments is provided generically.
+/// Specific patterns of run-method extra arguments and analysis manager extra
+/// arguments will have to be defined as appropriate for those patterns.
+template <typename AnalysisT, typename IRUnitT>
+struct RequireAnalysisPass<AnalysisT, IRUnitT, AnalysisManager<IRUnitT>>
+    : PassInfoMixin<
+          RequireAnalysisPass<AnalysisT, IRUnitT, AnalysisManager<IRUnitT>>> {
   /// \brief Run this pass over some unit of IR.
   ///
   /// This pass can be run over any unit of IR and use any analysis manager
   /// provided they satisfy the basic API requirements. When this pass is
   /// created, these methods can be instantiated to satisfy whatever the
   /// context requires.
-  template <typename IRUnitT>
   PreservedAnalyses run(IRUnitT &Arg, AnalysisManager<IRUnitT> &AM) {
     (void)AM.template getResult<AnalysisT>(Arg);
 
@@ -904,8 +935,8 @@ struct InvalidateAnalysisPass
   /// provided they satisfy the basic API requirements. When this pass is
   /// created, these methods can be instantiated to satisfy whatever the
   /// context requires.
-  template <typename IRUnitT>
-  PreservedAnalyses run(IRUnitT &Arg, AnalysisManager<IRUnitT> &AM) {
+  template <typename IRUnitT, typename AnalysisManagerT, typename... ExtraArgTs>
+  PreservedAnalyses run(IRUnitT &Arg, AnalysisManagerT &AM, ExtraArgTs &&...) {
     // We have to directly invalidate the analysis result as we can't
     // enumerate all other analyses and use the preserved set to control it.
     AM.template invalidate<AnalysisT>(Arg);
@@ -920,8 +951,8 @@ struct InvalidateAnalysisPass
 /// analysis passes to be re-run to produce fresh results if any are needed.
 struct InvalidateAllAnalysesPass : PassInfoMixin<InvalidateAllAnalysesPass> {
   /// \brief Run this pass over some unit of IR.
-  template <typename IRUnitT>
-  PreservedAnalyses run(IRUnitT &, AnalysisManager<IRUnitT> &) {
+  template <typename IRUnitT, typename AnalysisManagerT, typename... ExtraArgTs>
+  PreservedAnalyses run(IRUnitT &, AnalysisManagerT &, ExtraArgTs &&...) {
     return PreservedAnalyses::none();
   }
 };

include/llvm/IR/PassManagerInternal.h

@@ -23,7 +23,7 @@
 
 namespace llvm {
 
-template <typename IRUnitT> class AnalysisManager;
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
 class PreservedAnalyses;
 
 /// \brief Implementation details of the pass manager interfaces.
@@ -31,12 +31,18 @@ namespace detail {
 
 /// \brief Template for the abstract base class used to dispatch
 /// polymorphically over pass objects.
-template <typename IRUnitT> struct PassConcept {
+template <typename IRUnitT, typename AnalysisManagerT, typename... ExtraArgTs>
+struct PassConcept {
   // Boiler plate necessary for the container of derived classes.
   virtual ~PassConcept() {}
 
   /// \brief The polymorphic API which runs the pass over a given IR entity.
-  virtual PreservedAnalyses run(IRUnitT &IR, AnalysisManager<IRUnitT> &AM) = 0;
+  ///
+  /// Note that actual pass object can omit the analysis manager argument if
+  /// desired. Also that the analysis manager may be null if there is no
+  /// analysis manager in the pass pipeline.
+  virtual PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
+                                ExtraArgTs... ExtraArgs) = 0;
 
   /// \brief Polymorphic method to access the name of a pass.
   virtual StringRef name() = 0;
@@ -47,9 +53,9 @@ template <typename IRUnitT> struct PassConcept {
 /// Can be instantiated for any object which provides a \c run method accepting
 /// an \c IRUnitT& and an \c AnalysisManager<IRUnit>&. It requires the pass to
 /// be a copyable object.
-template <typename IRUnitT, typename PassT,
-          typename PreservedAnalysesT = PreservedAnalyses>
-struct PassModel : PassConcept<IRUnitT> {
+template <typename IRUnitT, typename PassT, typename PreservedAnalysesT,
+          typename AnalysisManagerT, typename... ExtraArgTs>
+struct PassModel : PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...> {
   explicit PassModel(PassT Pass) : Pass(std::move(Pass)) {}
   // We have to explicitly define all the special member functions because MSVC
   // refuses to generate them.
@@ -64,8 +70,9 @@ struct PassModel : PassConcept<IRUnitT> {
     return *this;
   }
 
-  PreservedAnalysesT run(IRUnitT &IR, AnalysisManager<IRUnitT> &AM) override {
-    return Pass.run(IR, AM);
+  PreservedAnalysesT run(IRUnitT &IR, AnalysisManagerT &AM,
+                         ExtraArgTs... ExtraArgs) override {
+    return Pass.run(IR, AM, ExtraArgs...);
   }
   StringRef name() override { return PassT::name(); }
   PassT Pass;
@@ -205,14 +212,15 @@ struct AnalysisResultModel<IRUnitT, PassT, ResultT, PreservedAnalysesT, true>
 ///
 /// This concept is parameterized over the IR unit that it can run over and
 /// produce an analysis result.
-template <typename IRUnitT> struct AnalysisPassConcept {
+template <typename IRUnitT, typename... ExtraArgTs> struct AnalysisPassConcept {
   virtual ~AnalysisPassConcept() {}
 
   /// \brief Method to run this analysis over a unit of IR.
   /// \returns A unique_ptr to the analysis result object to be queried by
   /// users.
   virtual std::unique_ptr<AnalysisResultConcept<IRUnitT>>
-  run(IRUnitT &IR, AnalysisManager<IRUnitT> &AM) = 0;
+  run(IRUnitT &IR, AnalysisManager<IRUnitT, ExtraArgTs...> &AM,
+      ExtraArgTs... ExtraArgs) = 0;
 
   /// \brief Polymorphic method to access the name of a pass.
   virtual StringRef name() = 0;
@@ -223,8 +231,8 @@ template <typename IRUnitT> struct AnalysisPassConcept {
 /// Can wrap any type which implements a suitable \c run method. The method
 /// must accept an \c IRUnitT& and an \c AnalysisManager<IRUnitT>& as arguments
 /// and produce an object which can be wrapped in a \c AnalysisResultModel.
-template <typename IRUnitT, typename PassT>
-struct AnalysisPassModel : AnalysisPassConcept<IRUnitT> {
+template <typename IRUnitT, typename PassT, typename... ExtraArgTs>
+struct AnalysisPassModel : AnalysisPassConcept<IRUnitT, ExtraArgTs...> {
   explicit AnalysisPassModel(PassT Pass) : Pass(std::move(Pass)) {}
   // We have to explicitly define all the special member functions because MSVC
   // refuses to generate them.
@@ -247,8 +255,9 @@ struct AnalysisPassModel : AnalysisPassConcept<IRUnitT> {
   ///
   /// The return is wrapped in an \c AnalysisResultModel.
   std::unique_ptr<AnalysisResultConcept<IRUnitT>>
-  run(IRUnitT &IR, AnalysisManager<IRUnitT> &AM) override {
-    return make_unique<ResultModelT>(Pass.run(IR, AM));
+  run(IRUnitT &IR, AnalysisManager<IRUnitT, ExtraArgTs...> &AM,
+      ExtraArgTs... ExtraArgs) override {
+    return make_unique<ResultModelT>(Pass.run(IR, AM, ExtraArgs...));
   }
 
   /// \brief The model delegates to a static \c PassT::name method.

lib/Passes/PassBuilder.cpp

@@ -516,8 +516,9 @@ bool PassBuilder::parseModulePass(ModulePassManager &MPM,
   }
 #define MODULE_ANALYSIS(NAME, CREATE_PASS)                                     \
   if (Name == "require<" NAME ">") {                                           \
-    MPM.addPass(RequireAnalysisPass<                                           \
-                std::remove_reference<decltype(CREATE_PASS)>::type>());        \
+    MPM.addPass(                                                               \
+        RequireAnalysisPass<                                                   \
+            std::remove_reference<decltype(CREATE_PASS)>::type, Module>());    \
     return true;                                                               \
   }                                                                            \
   if (Name == "invalidate<" NAME ">") {                                        \
@@ -578,7 +579,8 @@ bool PassBuilder::parseCGSCCPass(CGSCCPassManager &CGPM,
 #define CGSCC_ANALYSIS(NAME, CREATE_PASS)                                      \
   if (Name == "require<" NAME ">") {                                           \
     CGPM.addPass(RequireAnalysisPass<                                          \
-                 std::remove_reference<decltype(CREATE_PASS)>::type>());       \
+                 std::remove_reference<decltype(CREATE_PASS)>::type,           \
+                 LazyCallGraph::SCC>());                                       \
     return true;                                                               \
   }                                                                            \
   if (Name == "invalidate<" NAME ">") {                                        \
@@ -637,8 +639,9 @@ bool PassBuilder::parseFunctionPass(FunctionPassManager &FPM,
   }
 #define FUNCTION_ANALYSIS(NAME, CREATE_PASS)                                   \
   if (Name == "require<" NAME ">") {                                           \
-    FPM.addPass(RequireAnalysisPass<                                           \
-                std::remove_reference<decltype(CREATE_PASS)>::type>());        \
+    FPM.addPass(                                                               \
+        RequireAnalysisPass<                                                   \
+            std::remove_reference<decltype(CREATE_PASS)>::type, Function>());  \
     return true;                                                               \
   }                                                                            \
   if (Name == "invalidate<" NAME ">") {                                        \
@@ -688,7 +691,7 @@ bool PassBuilder::parseLoopPass(LoopPassManager &LPM, const PipelineElement &E,
 #define LOOP_ANALYSIS(NAME, CREATE_PASS)                                       \
   if (Name == "require<" NAME ">") {                                           \
     LPM.addPass(RequireAnalysisPass<                                           \
-                std::remove_reference<decltype(CREATE_PASS)>::type>());        \
+                std::remove_reference<decltype(CREATE_PASS)>::type, Loop>());  \
     return true;                                                               \
   }                                                                            \
   if (Name == "invalidate<" NAME ">") {                                        \

unittests/IR/PassManagerTest.cpp

@@ -331,4 +331,61 @@ TEST_F(PassManagerTest, Basic) {
 
   EXPECT_EQ(1, ModuleAnalysisRuns);
 }
+
+// A customized pass manager that passes extra arguments through the
+// infrastructure.
+typedef AnalysisManager<Function, int> CustomizedAnalysisManager;
+typedef PassManager<Function, CustomizedAnalysisManager, int, int &>
+    CustomizedPassManager;
+
+class CustomizedAnalysis : public AnalysisInfoMixin<CustomizedAnalysis> {
+public:
+  struct Result {
+    Result(int I) : I(I) {}
+    int I;
+  };
+
+  Result run(Function &F, CustomizedAnalysisManager &AM, int I) {
+    return Result(I);
+  }
+
+private:
+  friend AnalysisInfoMixin<CustomizedAnalysis>;
+  static char PassID;
+};
+
+char CustomizedAnalysis::PassID;
+
+struct CustomizedPass : PassInfoMixin<CustomizedPass> {
+  std::function<void(CustomizedAnalysis::Result &, int &)> Callback;
+
+  template <typename CallbackT>
+  CustomizedPass(CallbackT Callback) : Callback(Callback) {}
+
+  PreservedAnalyses run(Function &F, CustomizedAnalysisManager &AM, int I,
+                        int &O) {
+    Callback(AM.getResult<CustomizedAnalysis>(F, I), O);
+    return PreservedAnalyses::none();
+  }
+};
+
+TEST_F(PassManagerTest, CustomizedPassManagerArgs) {
+  CustomizedAnalysisManager AM;
+  AM.registerPass([&] { return CustomizedAnalysis(); });
+
+  CustomizedPassManager PM;
+
+  // Add an instance of the customized pass that just accumulates the input
+  // after it is round-tripped through the analysis.
+  int Result = 0;
+  PM.addPass(CustomizedPass::CustomizedPass(
+      [](CustomizedAnalysis::Result &R, int &O) { O += R.I; }));
+
+  // Run this over every function with the input of 42.
+  for (Function &F : *M)
+    PM.run(F, AM, 42, Result);
+
+  // And ensure that we accumulated the correct result.
+  EXPECT_EQ(42 * (int)M->size(), Result);
+}
 }