[mlir][Pass][NFC] Merge OpToOpPassAdaptor and OpToOpPassAdaptorParallel

[rriddle]

Summary

This moves the threading check to runOnOperation. This produces a much cleaner interface for the adaptor pass, and will allow for the ability to enable/disable threading in a much cleaner way in the future.

mlir/lib/Pass/IRPrinting.cpp

private:
   /// being operated on in a pass. This field is only used when the
   /// configuration asked for change detection.
   DenseMap<Pass *, OperationFingerPrint> beforePassFingerPrints;
 };
 } // end anonymous namespace
 
 /// Returns true if the given pass is hidden from IR printing.
 static bool isHiddenPass(Pass *pass) {
-  return isAdaptorPass(pass) || isa<VerifierPass>(pass);
+  return isa<OpToOpPassAdaptor>(pass) || isa<VerifierPass>(pass);
 }
 
 static void printIR(Operation *op, bool printModuleScope, raw_ostream &out,
                     OpPrintingFlags flags) {
   // Check to see if we are printing the top-level module.
   auto module = dyn_cast<ModuleOp>(op);
   if (module && !op->getBlock())
     return module.print(out << "\n", flags);
void IRPrinterInstrumentation::runAfterPass(Pass *pass, Operation *op) {
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
     out << formatv("// *** IR Dump After {0} ***", pass->getName());
     printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags());
     out << "\n\n";
   });
 }
 
 void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass, Operation *op) {
-  if (isAdaptorPass(pass))
+  if (isa<OpToOpPassAdaptor>(pass))
     return;
   if (config->shouldPrintAfterOnlyOnChange())
     beforePassFingerPrints.erase(pass);
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
     out << formatv("// *** IR Dump After {0} Failed ***", pass->getName());
     printIR(op, config->shouldPrintAtModuleScope(), out,
             OpPrintingFlags().printGenericOpForm());

mlir/lib/Pass/Pass.cpp

 void Pass::copyOptionValuesFrom(const Pass *other) {
   passOptions.copyOptionValuesFrom(other->passOptions);
 }
 
 /// Prints out the pass in the textual representation of pipelines. If this is
 /// an adaptor pass, print with the op_name(sub_pass,...) format.
 void Pass::printAsTextualPipeline(raw_ostream &os) {
   // Special case for adaptors to use the 'op_name(sub_passes)' format.
-  if (auto *adaptor = getAdaptorPassBase(this)) {
+  if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(this)) {
     llvm::interleaveComma(adaptor->getPassManagers(), os,
                           [&](OpPassManager &pm) {
                             os << pm.getOpName() << "(";
                             pm.printAsTextualPipeline(os);
                             os << ")";
                           });
     return;
   }

 } // end namespace detail
 } // end namespace mlir
 
 /// Coalesce adjacent AdaptorPasses into one large adaptor. This runs
 /// recursively through the pipeline graph.
 void OpPassManagerImpl::coalesceAdjacentAdaptorPasses() {
   // Bail out early if there are no adaptor passes.
   if (llvm::none_of(passes, [](std::unique_ptr<Pass> &pass) {
-        return isAdaptorPass(pass.get());
+        return isa<OpToOpPassAdaptor>(pass.get());
       }))
     return;
 
   // Walk the pass list and merge adjacent adaptors.
-  OpToOpPassAdaptorBase *lastAdaptor = nullptr;
+  OpToOpPassAdaptor *lastAdaptor = nullptr;
   for (auto it = passes.begin(), e = passes.end(); it != e; ++it) {
     // Check to see if this pass is an adaptor.
-    if (auto *currentAdaptor = getAdaptorPassBase(it->get())) {
+    if (auto *currentAdaptor = dyn_cast<OpToOpPassAdaptor>(it->get())) {
       // If it is the first adaptor in a possible chain, remember it and
       // continue.
       if (!lastAdaptor) {
         lastAdaptor = currentAdaptor;
         continue;
       }
 
       // Otherwise, merge into the existing adaptor and delete the current one.
if (failed(pass->run(op, am)))
       return failure();
   return success();
 }
 
 /// Nest a new operation pass manager for the given operation kind under this
 /// pass manager.
 OpPassManager &OpPassManager::nest(const OperationName &nestedName) {
   OpPassManager nested(nestedName, impl->disableThreads, impl->verifyPasses);
-
-  /// Create an adaptor for this pass. If multi-threading is disabled, then
-  /// create a synchronous adaptor.
-  if (impl->disableThreads || !llvm::llvm_is_multithreaded()) {
-    auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
-    addPass(std::unique_ptr<Pass>(adaptor));
-    return adaptor->getPassManagers().front();
-  }
-
-  auto *adaptor = new OpToOpPassAdaptorParallel(std::move(nested));
-  addPass(std::unique_ptr<Pass>(adaptor));
-  return adaptor->getPassManagers().front();
-}
 OpPassManager &OpPassManager::nest(StringRef nestedName) {
   return nest(OperationName(nestedName, getContext()));
 }
 
 /// Add the given pass to this pass manager. If this pass has a concrete
 /// operation type, it must be the same type as this pass manager.
 void OpPassManager::addPass(std::unique_ptr<Pass> pass) {
   // If this pass runs on a different operation than this pass manager, then

 /// type, or nullptr if one does not exist.
 static OpPassManager *findPassManagerFor(MutableArrayRef<OpPassManager> mgrs,
                                          const OperationName &name) {
   auto it = llvm::find_if(
       mgrs, [&](OpPassManager &mgr) { return mgr.getOpName() == name; });
   return it == mgrs.end() ? nullptr : &*it;
 }
 
-OpToOpPassAdaptorBase::OpToOpPassAdaptorBase(OpPassManager &&mgr) {
+OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr) {
   mgrs.emplace_back(std::move(mgr));
 }
 
 /// Merge the current pass adaptor into given 'rhs'.
-void OpToOpPassAdaptorBase::mergeInto(OpToOpPassAdaptorBase &rhs) {
+void OpToOpPassAdaptor::mergeInto(OpToOpPassAdaptor &rhs) {
   for (auto &pm : mgrs) {
     // If an existing pass manager exists, then merge the given pass manager
     // into it.
     if (auto *existingPM = findPassManagerFor(rhs.mgrs, pm.getOpName())) {
       pm.getImpl().mergeInto(existingPM->getImpl());
     } else {
       // Otherwise, add the given pass manager to the list.
       rhs.mgrs.emplace_back(std::move(pm));
     }
   }
   mgrs.clear();
 
   // After coalescing, sort the pass managers within rhs by name.
   llvm::array_pod_sort(rhs.mgrs.begin(), rhs.mgrs.end(),
                        [](const OpPassManager *lhs, const OpPassManager *rhs) {
                          return lhs->getOpName().getStringRef().compare(
                              rhs->getOpName().getStringRef());
                        });
 }
 
 /// Returns the adaptor pass name.
-std::string OpToOpPassAdaptorBase::getName() {
+std::string OpToOpPassAdaptor::getAdaptorName() {
   std::string name = "Pipeline Collection : [";
   llvm::raw_string_ostream os(name);
   llvm::interleaveComma(getPassManagers(), os, [&](OpPassManager &pm) {
     os << '\'' << pm.getOpName() << '\'';
   });
   os << ']';
   return os.str();
 }
 
-OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr)
-    : OpToOpPassAdaptorBase(std::move(mgr)) {}
-
 /// Run the held pipeline over all nested operations.
 void OpToOpPassAdaptor::runOnOperation() {
+  if (mgrs.front().getImpl().disableThreads || !llvm::llvm_is_multithreaded())
+    runOnOperationImpl();
+  else
+    runOnOperationAsyncImpl();
+}
+
+/// Run this pass adaptor synchronously.
+void OpToOpPassAdaptor::runOnOperationImpl() {
   auto am = getAnalysisManager();
   PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         this};
   auto *instrumentor = am.getPassInstrumentor();
   for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         auto *mgr = findPassManagerFor(mgrs, op.getName());
for (auto &block : region) {
 
         if (failed(result))
           return signalPassFailure();
       }
     }
   }
 }
 
-OpToOpPassAdaptorParallel::OpToOpPassAdaptorParallel(OpPassManager &&mgr)
-    : OpToOpPassAdaptorBase(std::move(mgr)) {}
-
 /// Utility functor that checks if the two ranges of pass managers have a size
 /// mismatch.
 static bool hasSizeMismatch(ArrayRef<OpPassManager> lhs,
                             ArrayRef<OpPassManager> rhs) {
   return lhs.size() != rhs.size() ||
          llvm::any_of(llvm::seq<size_t>(0, lhs.size()),
                       [&](size_t i) { return lhs[i].size() != rhs[i].size(); });
 }
 
-// Run the held pipeline asynchronously across the functions within the module.
-void OpToOpPassAdaptorParallel::runOnOperation() {
+/// Run this pass adaptor synchronously.
+void OpToOpPassAdaptor::runOnOperationAsyncImpl() {
   AnalysisManager am = getAnalysisManager();
 
   // Create the async executors if they haven't been created, or if the main
   // pipeline has changed.
   if (asyncExecutors.empty() || hasSizeMismatch(asyncExecutors.front(), mgrs))
     asyncExecutors.assign(llvm::hardware_concurrency().compute_thread_count(),
                           mgrs);
 
llvm::parallel::for_each(
         }
       });
 
   // Signal a failure if any of the executors failed.
   if (passFailed)
     signalPassFailure();
 }
 
-/// Utility function to convert the given class to the base adaptor it is an
-/// adaptor pass, returns nullptr otherwise.
-OpToOpPassAdaptorBase *mlir::detail::getAdaptorPassBase(Pass *pass) {
-  if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass))
-    return adaptor;
-  if (auto *adaptor = dyn_cast<OpToOpPassAdaptorParallel>(pass))
-    return adaptor;
-  return nullptr;
-}
-
 //===----------------------------------------------------------------------===//
 // PassCrashReproducer
 //===----------------------------------------------------------------------===//
 
 /// Safely run the pass manager over the given module, creating a reproducible
 /// on failure or crash.
 static LogicalResult runWithCrashRecovery(OpPassManager &pm,
                                           ModuleAnalysisManager &am,

mlir/lib/Pass/PassDetail.h

 class VerifierPass : public PassWrapper<VerifierPass, OperationPass<>> {
   void runOnOperation() override;
 };
 
 //===----------------------------------------------------------------------===//
 // OpToOpPassAdaptor
 //===----------------------------------------------------------------------===//
 
-/// A base class for Op-to-Op adaptor passes.
-class OpToOpPassAdaptorBase {
-public:
-  OpToOpPassAdaptorBase(OpPassManager &&mgr);
-  OpToOpPassAdaptorBase(const OpToOpPassAdaptorBase &rhs) = default;
-
-  /// Merge the current pass adaptor into given 'rhs'.
-  void mergeInto(OpToOpPassAdaptorBase &rhs);
-
-  /// Returns the pass managers held by this adaptor.
-  MutableArrayRef<OpPassManager> getPassManagers() { return mgrs; }
-
-  /// Returns the adaptor pass name.
-  std::string getName();
-
-protected:
-  // A set of adaptors to run.
-  SmallVector<OpPassManager, 1> mgrs;
-};
-
-/// An adaptor pass used to run operation passes over nested operations
-/// synchronously on a single thread.
+/// An adaptor pass used to run operation passes over nested operations.
 class OpToOpPassAdaptor
-    : public PassWrapper<OpToOpPassAdaptor, OperationPass<>>,
-      public OpToOpPassAdaptorBase {
+    : public PassWrapper<OpToOpPassAdaptor, OperationPass<>> {
 public:
   OpToOpPassAdaptor(OpPassManager &&mgr);
+  OpToOpPassAdaptor(const OpToOpPassAdaptor &rhs) = default;
 
   /// Run the held pipeline over all operations.
   void runOnOperation() override;
-};
 
-/// An adaptor pass used to run operation passes over nested operations
-/// asynchronously across multiple threads.
-class OpToOpPassAdaptorParallel
-    : public PassWrapper<OpToOpPassAdaptorParallel, OperationPass<>>,
-      public OpToOpPassAdaptorBase {
-public:
-  OpToOpPassAdaptorParallel(OpPassManager &&mgr);
+  /// Merge the current pass adaptor into given 'rhs'.
+  void mergeInto(OpToOpPassAdaptor &rhs);
 
-  /// Run the held pipeline over all operations.
-  void runOnOperation() override;
+  /// Returns the pass managers held by this adaptor.
+  MutableArrayRef<OpPassManager> getPassManagers() { return mgrs; }
 
   /// Return the async pass managers held by this parallel adaptor.
   MutableArrayRef<SmallVector<OpPassManager, 1>> getParallelPassManagers() {
     return asyncExecutors;
   }
 
+  /// Returns the adaptor pass name.
+  std::string getAdaptorName();
+
 private:
-  // A set of executors, cloned from the main executor, that run asynchronously
-  // on different threads.
+  /// Run this pass adaptor synchronously.
+  void runOnOperationImpl();
+
+  /// Run this pass adaptor asynchronously.
+  void runOnOperationAsyncImpl();
+
+  /// A set of adaptors to run.
+  SmallVector<OpPassManager, 1> mgrs;
+
+  /// A set of executors, cloned from the main executor, that run asynchronously
+  /// on different threads. This is used when threading is enabled.
   SmallVector<SmallVector<OpPassManager, 1>, 8> asyncExecutors;
 };
 
-/// Utility function to convert the given class to the base adaptor it is an
-/// adaptor pass, returns nullptr otherwise.
-OpToOpPassAdaptorBase *getAdaptorPassBase(Pass *pass);
-
-/// Utility function to return if a pass refers to an adaptor pass. Adaptor
-/// passes are those that internally execute a pipeline.
-inline bool isAdaptorPass(Pass *pass) {
-  return isa<OpToOpPassAdaptorParallel>(pass) || isa<OpToOpPassAdaptor>(pass);
-}
-
 } // end namespace detail
 } // end namespace mlir
 #endif // MLIR_PASS_PASSDETAIL_H_

mlir/lib/Pass/PassStatistics.cpp

static void printPassEntry(raw_ostream &os, unsigned indent, StringRef pass,
   }
 }
 
 /// Print the statistics results in a list form, where each pass is sorted by
 /// name.
 static void printResultsAsList(raw_ostream &os, OpPassManager &pm) {
   llvm::StringMap<std::vector<Statistic>> mergedStats;
   std::function<void(Pass *)> addStats = [&](Pass *pass) {
-    auto *adaptor = getAdaptorPassBase(pass);
+    auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass);
 
     // If this is not an adaptor, add the stats to the list if there are any.
     if (!adaptor) {
       auto statistics = pass->getStatistics();
       if (statistics.empty())
         return;
 
       auto &passEntry = mergedStats[pass->getName()];
for (auto &statData : passAndStatistics)
     printPassEntry(os, /*indent=*/2, statData.first, statData.second);
 }
 
 /// Print the results in pipeline mode that mirrors the internal pass manager
 /// structure.
 static void printResultsAsPipeline(raw_ostream &os, OpPassManager &pm) {
   std::function<void(unsigned, Pass *)> printPass = [&](unsigned indent,
                                                         Pass *pass) {
-    // Handle the case of an adaptor pass.
-    if (auto *adaptor = getAdaptorPassBase(pass)) {
+    if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass)) {
       // If this adaptor has more than one internal pipeline, print an entry for
       // it.
       auto mgrs = adaptor->getPassManagers();
       if (mgrs.size() > 1) {
-        printPassEntry(os, indent, adaptor->getName());
+        printPassEntry(os, indent, adaptor->getAdaptorName());
         indent += 2;
       }
 
       // Print each of the children passes.
       for (OpPassManager &mgr : mgrs) {
         auto name = ("'" + mgr.getOpName().getStringRef() + "' Pipeline").str();
         printPassEntry(os, indent, name);
         for (Pass &pass : mgr.getPasses())

 /// Merge the pass statistics of this class into 'other'.
 void OpPassManager::mergeStatisticsInto(OpPassManager &other) {
   auto passes = getPasses(), otherPasses = other.getPasses();
 
   for (auto passPair : llvm::zip(passes, otherPasses)) {
     Pass &pass = std::get<0>(passPair), &otherPass = std::get<1>(passPair);
 
     // If this is an adaptor, then recursively merge the pass managers.
-    if (auto *adaptorPass = getAdaptorPassBase(&pass)) {
-      auto *otherAdaptorPass = getAdaptorPassBase(&otherPass);
+    if (auto *adaptorPass = dyn_cast<OpToOpPassAdaptor>(&pass)) {
+      auto *otherAdaptorPass = cast<OpToOpPassAdaptor>(&otherPass);
       for (auto mgrs : llvm::zip(adaptorPass->getPassManagers(),
                                  otherAdaptorPass->getPassManagers()))
         std::get<0>(mgrs).mergeStatisticsInto(std::get<1>(mgrs));
       continue;
     }
     // Otherwise, merge the statistics for the current pass.
     assert(pass.statistics.size() == otherPass.statistics.size());
     for (unsigned i = 0, e = pass.statistics.size(); i != e; ++i) {
       assert(pass.statistics[i]->getName() ==
              StringRef(otherPass.statistics[i]->getName()));
       *otherPass.statistics[i] += *pass.statistics[i];
       *pass.statistics[i] = 0;
     }
   }
 }
 
 /// Prepare the statistics of passes within the given pass manager for
 /// consumption(e.g. dumping).
 static void prepareStatistics(OpPassManager &pm) {
   for (Pass &pass : pm.getPasses()) {
-    OpToOpPassAdaptorBase *adaptor = getAdaptorPassBase(&pass);
+    OpToOpPassAdaptor *adaptor = dyn_cast<OpToOpPassAdaptor>(&pass);
     if (!adaptor)
       continue;
     MutableArrayRef<OpPassManager> nestedPms = adaptor->getPassManagers();
 
-    // If this is a parallel adaptor, merge the statistics from the async
-    // pass managers into the main nested pass managers.
-    if (auto *parallelAdaptor = dyn_cast<OpToOpPassAdaptorParallel>(&pass)) {
-      for (auto &asyncPM : parallelAdaptor->getParallelPassManagers()) {
-        for (unsigned i = 0, e = asyncPM.size(); i != e; ++i)
-          asyncPM[i].mergeStatisticsInto(nestedPms[i]);
-      }
-    }
+    // Merge the statistics from the async pass managers into the main nested
+    // pass managers.
+    for (auto &asyncPM : adaptor->getParallelPassManagers()) {
 
     // Prepare the statistics of each of the nested passes.
     for (OpPassManager &nestedPM : nestedPms)
       prepareStatistics(nestedPM);
   }
 }
 
 /// Dump the statistics of the passes within this pass manager.

mlir/lib/Pass/PassTiming.cpp

assert(parentTimer->children.size() == 1 &&
          "expected a single pipeline timer");
   pipelinesToMerge[parentInfo].push_back(
       std::move(*parentTimer->children.begin()));
   rootTimers.erase(tid);
 }
 
 /// Start a new timer for the given pass.
 void PassTiming::startPassTimer(Pass *pass) {
-  auto kind = isAdaptorPass(pass) ? TimerKind::PipelineCollection
+  auto kind = isa<OpToOpPassAdaptor>(pass) ? TimerKind::PipelineCollection
                                            : TimerKind::PassOrAnalysis;
   Timer *timer = getTimer(pass, kind, [pass]() -> std::string {
-    if (auto *adaptor = getAdaptorPassBase(pass))
-      return adaptor->getName();
+    if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass))
+      return adaptor->getAdaptorName();
     return std::string(pass->getName());
   });
 
   // We don't actually want to time the adaptor passes, they gather their total
   // from their held passes.
-  if (!isAdaptorPass(pass))
+  if (!isa<OpToOpPassAdaptor>(pass))
     timer->start();
 }
 
 /// Start a new timer for the given analysis.
 void PassTiming::startAnalysisTimer(StringRef name, TypeID id) {
   Timer *timer = getTimer(id.getAsOpaquePointer(), TimerKind::PassOrAnalysis,
                           [name] { return "(A) " + name.str(); });
   timer->start();
 }
 
 /// Stop a pass timer.
 void PassTiming::runAfterPass(Pass *pass, Operation *) {
   Timer *timer = popLastActiveTimer();
 
-  // If this is an OpToOpPassAdaptorParallel, then we need to merge in the
-  // timing data for the pipelines running on other threads.
-  if (isa<OpToOpPassAdaptorParallel>(pass)) {
+  // If this is a pass adaptor, then we need to merge in the timing data for the
+  // pipelines running on other threads.
+  if (isa<OpToOpPassAdaptor>(pass)) {
     auto toMerge = pipelinesToMerge.find({llvm::get_threadid(), pass});
     if (toMerge != pipelinesToMerge.end()) {
       for (auto &it : toMerge->second)
         timer->mergeChild(std::move(it));
       pipelinesToMerge.erase(toMerge);
     }
     return;
   }
 
-  // Adaptor passes aren't timed directly, so we don't need to stop their
-  // timers.
-  if (!isAdaptorPass(pass))
-    timer->stop();
 }
 
 /// Stop a timer.
 void PassTiming::runAfterAnalysis(StringRef, TypeID, Operation *) {
   popLastActiveTimer()->stop();
 }
 
 /// Utility to print the timer heading information.