[mlir] Add support for PassManager options and use them in the Inliner

[rriddle]

Summary

After "any" pass managers were added, we can now default construct op pass
managers. This was the last remaining requirement to properly support
pass manager options. This commit finalizes this support, and updates the
default pass pipeline of the inliner to use a pass manager option instead
of a string.

[mehdi_amini]

mlir/test/Dialect/Affine/inlining.mlir

1

I'm curious about this change: how is it parsed differently if you don't change it?

[rkayaith]

With this change I think we could use cl::opt<OpPassManager> here: https://github.com/llvm/llvm-project/blob/7f0003c19c3a47f484ea9cc929728808b27ace32/mlir/include/mlir/Pass/PassRegistry.h#L262

mlir/lib/Pass/PassRegistry.cpp

413–415

printAsTextualPipeline should already contain the anchor now

flang/include/flang/Tools/CLOptions.inc

void addNestedPassConditionally(
   if (!disabled)
     pm.addNestedPass<OP>(ctor());
 }
 
 } // namespace
 
 namespace fir {
 
-static void defaultFlangInlinerOptPipeline(mlir::OpPassManager &pm) {
+static mlir::OpPassManager buildDefaultFlangInlinerOptPipeline() {
+  mlir::OpPassManager pm;
+
   mlir::GreedyRewriteConfig config;
   config.enableRegionSimplification = false;
   pm.addPass(mlir::createCanonicalizerPass(config));
+
+  return pm;
 }
 
 inline void addCfgConversionPass(mlir::PassManager &pm) {
   addNestedPassConditionally<mlir::func::FuncOp>(
       pm, disableCfgConversion, fir::createFirToCfgPass);
 }
 
 inline void addAVC(mlir::PassManager &pm) {
if (optLevel.isOptimizingForSpeed()) {
     pm.addPass(fir::createAlgebraicSimplificationPass(config));
   }
   pm.addPass(mlir::createCSEPass());
   fir::addMemoryAllocationOpt(pm);
 
   // The default inliner pass adds the canonicalizer pass with the default
   // configuration. Create the inliner pass with tco config.
   llvm::StringMap<mlir::OpPassManager> pipelines;
-  pm.addPass(
-      mlir::createInlinerPass(pipelines, defaultFlangInlinerOptPipeline));
+  pm.addPass(mlir::createInlinerPass(
+      pipelines, buildDefaultFlangInlinerOptPipeline()));
   pm.addPass(fir::createSimplifyRegionLitePass());
   pm.addPass(mlir::createCSEPass());
 
   // convert control flow to CFG form
   fir::addCfgConversionPass(pm);
   pm.addPass(mlir::createConvertSCFToCFPass());
 
   pm.addPass(mlir::createCanonicalizerPass(config));

mlir/include/mlir/Pass/PassManager.h

public:
   /// Construct a new pass manager with the given anchor operation type and
   /// nesting behavior.
   OpPassManager(StringRef name, Nesting nesting = Nesting::Explicit);
   OpPassManager(OperationName name, Nesting nesting = Nesting::Explicit);
   OpPassManager(OpPassManager &&rhs);
   OpPassManager(const OpPassManager &rhs);
   ~OpPassManager();
   OpPassManager &operator=(const OpPassManager &rhs);
+  OpPassManager &operator=(OpPassManager &&rhs);
 
   /// Iterator over the passes in this pass manager.
   using pass_iterator =
       llvm::pointee_iterator<MutableArrayRef<std::unique_ptr<Pass>>::iterator>;
   pass_iterator begin();
   pass_iterator end();
   iterator_range<pass_iterator> getPasses() { return {begin(), end()}; }
 

mlir/include/mlir/Pass/PassOptions.h

 //===----------------------------------------------------------------------===//
 // OpPassManager: Parser
 
 extern template class basic_parser<mlir::OpPassManager>;
 
 template <>
 class parser<mlir::OpPassManager> : public basic_parser<mlir::OpPassManager> {
 public:
-  /// A utility struct used when parsing a pass manager that prevents the need
-  /// for a default constructor on OpPassManager.
-  struct ParsedPassManager {
-    ParsedPassManager();
-    ParsedPassManager(ParsedPassManager &&);
-    ~ParsedPassManager();
-    operator const mlir::OpPassManager &() const {
-      assert(value && "parsed value was invalid");
-      return *value;
-    }
-
-    std::unique_ptr<mlir::OpPassManager> value;
-  };
-  using parser_data_type = ParsedPassManager;
+  using parser_data_type = mlir::OpPassManager;
   using OptVal = OptionValue<mlir::OpPassManager>;
 
   parser(Option &opt) : basic_parser(opt) {}
 
-  bool parse(Option &, StringRef, StringRef arg, ParsedPassManager &value);
+  bool parse(Option &, StringRef, StringRef arg, mlir::OpPassManager &value);
 
   /// Print an instance of the underling option value to the given stream.
   static void print(raw_ostream &os, const mlir::OpPassManager &value);
 
   // Overload in subclass to provide a better default value.
   StringRef getValueName() const override { return "pass-manager"; }
 
-  void printOptionDiff(const Option &opt, mlir::OpPassManager &pm,
+  void printOptionDiff(const Option &opt, const mlir::OpPassManager &pm,
                        const OptVal &defaultValue, size_t globalWidth) const;
 
   // An out-of-line virtual method to provide a 'home' for this class.
   void anchor() override;
 };
 
 } // namespace cl
 } // namespace llvm
 
 #endif // MLIR_PASS_PASSOPTIONS_H_

mlir/include/mlir/Transforms/Passes.h

 // transformation library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_PASSES_H
 #define MLIR_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/LocationSnapshot.h"
 #include "mlir/Transforms/ViewOpGraph.h"
 #include "llvm/Support/Debug.h"
 #include <limits>
 
 namespace mlir {
 
 class GreedyRewriteConfig;

 /// when optimizing callable operations with names matching the key type.
 /// Callable operations with a name not within the provided map will use the
 /// default inliner pipeline during optimization.
 std::unique_ptr<Pass>
 createInlinerPass(llvm::StringMap<OpPassManager> opPipelines);
 /// Creates an instance of the inliner pass, and use the provided pass managers
 /// when optimizing callable operations with names matching the key type.
 /// Callable operations with a name not within the provided map will use the
-/// provided default pipeline builder.
+/// provided default pipeline, which must be anchored on `any`.
 std::unique_ptr<Pass>
 createInlinerPass(llvm::StringMap<OpPassManager> opPipelines,
-                  std::function<void(OpPassManager &)> defaultPipelineBuilder);
+                  OpPassManager defaultPipeline);
 
 /// Creates a pass which performs sparse conditional constant propagation over
 /// nested operations.
 std::unique_ptr<Pass> createSCCPPass();
 
 /// Creates a pass which delete symbol operations that are unreachable. This
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();

mlir/include/mlir/Transforms/Passes.td

let statistics = [
     Statistic<"numDCE", "num-dce'd", "Number of operations DCE'd">
   ];
 }
 
 def Inliner : Pass<"inline"> {
   let summary = "Inline function calls";
   let constructor = "mlir::createInlinerPass()";
   let options = [
-    Option<"defaultPipelineStr", "default-pipeline", "std::string",
-           /*default=*/"", "The default optimizer pipeline used for callables">,
+    Option<"defaultPipeline", "default-pipeline", "OpPassManager",
+           /*default=*/"", "The default optimizer pipeline used for callables (must be anchored on 'any')">,
     ListOption<"opPipelineList", "op-pipelines", "OpPassManager",
                "Callable operation specific optimizer pipelines (in the form "
                "of `dialect.op(pipeline)`)">,
     Option<"maxInliningIterations", "max-iterations", "unsigned",
            /*default=*/"4",
            "Maximum number of iterations when inlining within an SCC">,
   ];
 }

mlir/lib/Pass/Pass.cpp

 OpPassManager::OpPassManager(OperationName name, Nesting nesting)
     : impl(new OpPassManagerImpl(name, nesting)) {}
 OpPassManager::OpPassManager(OpPassManager &&rhs) : impl(std::move(rhs.impl)) {}
 OpPassManager::OpPassManager(const OpPassManager &rhs) { *this = rhs; }
 OpPassManager &OpPassManager::operator=(const OpPassManager &rhs) {
   impl = std::make_unique<OpPassManagerImpl>(*rhs.impl);
   return *this;
 }
+OpPassManager &OpPassManager::operator=(OpPassManager &&rhs) {
+  impl = std::move(rhs.impl);
+  return *this;
+}
 
 OpPassManager::~OpPassManager() = default;
 
 OpPassManager::pass_iterator OpPassManager::begin() {
   return MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.begin();
 }
 OpPassManager::pass_iterator OpPassManager::end() {
   return MutableArrayRef<std::unique_ptr<Pass>>{impl->passes}.end();

mlir/lib/Pass/PassRegistry.cpp

 
 namespace llvm {
 namespace cl {
 template class basic_parser<OpPassManager>;
 } // namespace cl
 } // namespace llvm
 
 bool llvm::cl::parser<OpPassManager>::parse(Option &, StringRef, StringRef arg,
-                                            ParsedPassManager &value) {
+                                            OpPassManager &value) {
+  if (arg.empty())
+    return false;
+
   FailureOr<OpPassManager> pipeline = parsePassPipeline(arg);
   if (failed(pipeline))
     return true;
-  value.value = std::make_unique<OpPassManager>(std::move(*pipeline));
+  value = std::move(*pipeline);
   return false;
 }
 
 void llvm::cl::parser<OpPassManager>::print(raw_ostream &os,
                                             const OpPassManager &value) {
+  os << value.getOpAnchorName() << "(";
   value.printAsTextualPipeline(os);
+  os << ")";
 }
 
 void llvm::cl::parser<OpPassManager>::printOptionDiff(
-    const Option &opt, OpPassManager &pm, const OptVal &defaultValue,
+    const Option &opt, const OpPassManager &pm, const OptVal &defaultValue,
     size_t globalWidth) const {
   printOptionName(opt, globalWidth);
   outs() << "= ";
-  pm.printAsTextualPipeline(outs());
+  print(outs(), pm);
 
   if (defaultValue.hasValue()) {
     outs().indent(2) << " (default: ";
     defaultValue.getValue().printAsTextualPipeline(outs());
     outs() << ")";
   }
   outs() << "\n";
 }
 
 void llvm::cl::parser<OpPassManager>::anchor() {}
 
-llvm::cl::parser<OpPassManager>::ParsedPassManager::ParsedPassManager() =
-    default;
-llvm::cl::parser<OpPassManager>::ParsedPassManager::ParsedPassManager(
-    ParsedPassManager &&) = default;
-llvm::cl::parser<OpPassManager>::ParsedPassManager::~ParsedPassManager() =
-    default;
-
 //===----------------------------------------------------------------------===//
 // TextualPassPipeline Parser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a textual description of a pass pipeline.
 class TextualPipeline {
 public:

mlir/lib/Transforms/Inliner.cpp

 #include "mlir/Transforms/Passes.h.inc"
 } // namespace mlir
 
 #define DEBUG_TYPE "inlining"
 
 using namespace mlir;
 
 /// This function implements the default inliner optimization pipeline.
-static void defaultInlinerOptPipeline(OpPassManager &pm) {
+static OpPassManager buildDefaultInlinerOptPipeline() {
+  OpPassManager pm;
   pm.addPass(createCanonicalizerPass());
+  return pm;
 }
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Tracking
 //===----------------------------------------------------------------------===//
 
 /// Walk all of the used symbol callgraph nodes referenced with the given op.
 static void walkReferencedSymbolNodes(

 // InlinerPass
 //===----------------------------------------------------------------------===//
 
 namespace {
 class InlinerPass : public impl::InlinerBase<InlinerPass> {
 public:
   InlinerPass();
   InlinerPass(const InlinerPass &) = default;
-  InlinerPass(std::function<void(OpPassManager &)> defaultPipeline);
-  InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
+  InlinerPass(OpPassManager &&defaultPM);
+  InlinerPass(OpPassManager &&defaultPM,
               llvm::StringMap<OpPassManager> opPipelines);
   void runOnOperation() override;
 
 private:
   /// Attempt to inline calls within the given scc, and run simplifications,
   /// until a fixed point is reached. This allows for the inlining of newly
   /// devirtualized calls. Returns failure if there was a fatal error during
   /// inlining.
LogicalResult optimizeCallable(CallGraphNode *node,
                                  llvm::StringMap<OpPassManager> &pipelines);
 
   /// Attempt to initialize the options of this pass from the given string.
   /// Derived classes may override this method to hook into the point at which
   /// options are initialized, but should generally always invoke this base
   /// class variant.
   LogicalResult initializeOptions(StringRef options) override;
 
-  /// An optional function that constructs a default optimization pipeline for
-  /// a given operation.
-  std::function<void(OpPassManager &)> defaultPipeline;
   /// A map of operation names to pass pipelines to use when optimizing
   /// callable operations of these types. This provides a specialized pipeline
   /// instead of the default. The vector size is the number of threads used
   /// during optimization.
   SmallVector<llvm::StringMap<OpPassManager>, 8> opPipelines;
 };
 } // namespace
 
-InlinerPass::InlinerPass() : InlinerPass(defaultInlinerOptPipeline) {}
-InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline)
-    : defaultPipeline(std::move(defaultPipeline)) {
+InlinerPass::InlinerPass() : InlinerPass(OpPassManager()) {}
+InlinerPass::InlinerPass(OpPassManager &&defaultPM) {
+  defaultPipeline = std::move(defaultPM);
   opPipelines.push_back({});
-
-  // Initialize the pass options with the provided arguments.
-  if (defaultPipeline) {
-    OpPassManager fakePM("__mlir_fake_pm_op");
-    defaultPipeline(fakePM);
-    llvm::raw_string_ostream strStream(defaultPipelineStr);
-    fakePM.printAsTextualPipeline(strStream);
-  }
 }
 
-InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
+InlinerPass::InlinerPass(OpPassManager &&defaultPM,
                          llvm::StringMap<OpPassManager> opPipelines)
-    : InlinerPass(std::move(defaultPipeline)) {
+    : InlinerPass(std::move(defaultPM)) {
   if (opPipelines.empty())
     return;
 
   // Update the option for the op specific optimization pipelines.
   for (auto &it : opPipelines)
     opPipelineList.addValue(it.second);
   this->opPipelines.emplace_back(std::move(opPipelines));
 }

 }
 
 LogicalResult
 InlinerPass::optimizeCallable(CallGraphNode *node,
                               llvm::StringMap<OpPassManager> &pipelines) {
   Operation *callable = node->getCallableRegion()->getParentOp();
   StringRef opName = callable->getName().getStringRef();
   auto pipelineIt = pipelines.find(opName);
-  if (pipelineIt == pipelines.end()) {
-    // If a pipeline didn't exist, use the default if possible.
-    if (!defaultPipeline)
+
+  if (pipelineIt == pipelines.end()) {
+    if (defaultPipeline.empty())
       return success();
 
-    OpPassManager defaultPM(opName);
-    defaultPipeline(defaultPM);
+    OpPassManager defaultPM(defaultPipeline);
     pipelineIt = pipelines.try_emplace(opName, std::move(defaultPM)).first;
   }
   return runPipeline(pipelineIt->second, callable);
 }
 
 LogicalResult InlinerPass::initializeOptions(StringRef options) {
   if (failed(Pass::initializeOptions(options)))
     return failure();
 
-  // Initialize the default pipeline builder to use the option string.
-  // TODO: Use a generic pass manager for default pipelines, and remove this.
-  if (!defaultPipelineStr.empty()) {
-    std::string defaultPipelineCopy = defaultPipelineStr;
-    defaultPipeline = [=](OpPassManager &pm) {
-      (void)parsePassPipeline(defaultPipelineCopy, pm);
-    };
-  } else if (defaultPipelineStr.getNumOccurrences()) {
-    defaultPipeline = nullptr;
+  // Initialize the default pipeline builder.
+  if (!defaultPipeline.getNumOccurrences()) {
+    defaultPipeline = buildDefaultInlinerOptPipeline();
+  } else if (defaultPipeline.getValue().getOpAnchorName() !=
+             OpPassManager::getAnyOpAnchorName()) {
+    llvm::errs() << "error: Inliner expected default pipeline to be anchored "
+                    "on 'any', got '"
+                 << defaultPipeline.getValue().getOpAnchorName() << "'";
+    return failure();
   }
 
   // Initialize the op specific pass pipelines.
   llvm::StringMap<OpPassManager> pipelines;
   for (OpPassManager pipeline : opPipelineList)
     if (!pipeline.empty())
       pipelines.try_emplace(pipeline.getOpAnchorName(), pipeline);
   opPipelines.assign({std::move(pipelines)});
 
   return success();
 }
 
 std::unique_ptr<Pass> mlir::createInlinerPass() {
   return std::make_unique<InlinerPass>();
 }
 std::unique_ptr<Pass>
 mlir::createInlinerPass(llvm::StringMap<OpPassManager> opPipelines) {
-  return std::make_unique<InlinerPass>(defaultInlinerOptPipeline,
+  return std::make_unique<InlinerPass>(buildDefaultInlinerOptPipeline(),
                                        std::move(opPipelines));
 }
-std::unique_ptr<Pass> mlir::createInlinerPass(
-    llvm::StringMap<OpPassManager> opPipelines,
-    std::function<void(OpPassManager &)> defaultPipelineBuilder) {
-  return std::make_unique<InlinerPass>(std::move(defaultPipelineBuilder),
+std::unique_ptr<Pass>
+mlir::createInlinerPass(llvm::StringMap<OpPassManager> opPipelines,
+                        OpPassManager defaultPipeline) {
+  return std::make_unique<InlinerPass>(std::move(defaultPipeline),
                                        std::move(opPipelines));
 }

mlir/test/Dialect/Affine/inlining.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -inline="default-pipeline=''" | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -inline="default-pipeline=" | FileCheck %s
 
 // Basic test that functions within affine operations are inlined.
 func.func @func_with_affine_ops(%N: index) {
   %c = arith.constant 200 : index
   affine.for %i = 1 to 10 {
     affine.if affine_set<(i)[N] : (i - 2 >= 0, 4 - i >= 0)>(%i)[%c]  {
       %w = affine.apply affine_map<(d0,d1)[s0] -> (d0+d1+s0)> (%i, %i) [%N]
     }

mlir/test/Transforms/inlining.mlir

 // RUN: mlir-opt %s -inline='default-pipeline=''' | FileCheck %s
 // RUN: mlir-opt %s --mlir-disable-threading -inline='default-pipeline=''' | FileCheck %s
 // RUN: mlir-opt %s -inline='default-pipeline=''' -mlir-print-debuginfo -mlir-print-local-scope | FileCheck %s --check-prefix INLINE-LOC
 // RUN: mlir-opt %s -inline | FileCheck %s --check-prefix INLINE_SIMPLIFY
 // RUN: mlir-opt %s -inline='op-pipelines=func.func(canonicalize,cse)' | FileCheck %s --check-prefix INLINE_SIMPLIFY
+// RUN: not mlir-opt %s -inline='default-pipeline=func.func(canonicalize,cse)' 2>&1 | FileCheck %s --check-prefix INLINE_INVALID_DEFAULT
+
+// INLINE_INVALID_DEFAULT: error: Inliner expected default pipeline to be anchored on 'any', got 'func.func'
 
 // Inline a function that takes an argument.
 func.func @func_with_arg(%c : i32) -> i32 {
   %b = arith.addi %c, %c : i32
   return %b : i32
 }
 
 // CHECK-LABEL: func @inline_with_arg