[mlir] Allow for using OpPassManager in pass options

[rriddle]

Summary

This significantly simplifies the boilerplate necessary for passes
to define nested pass pipelines.

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

750

Not related to this change, but I'm curious if collision are expected here and if they should just be ignored?

[silvas]

Thanks River!

mlir/include/mlir/Transforms/Passes.td

84

Should defaultPipelineStr also be an OpPassManager option?

[rriddle]

mlir/include/mlir/Transforms/Passes.td

84

Not right now given that this pipeline is op-agnostic, which we don't support for OpPassManager right now (though I do have it on my list to explore/support).

mlir/lib/Transforms/Inliner.cpp

750

I don't think we should allow collision here. Will send a followup to error out here.

[silvas]

Hey River, I was trying to use this downstream and it isn't working for me

I have code like:

Option<OpPassManager> pipeline{
    *this, "pipeline", llvm::cl::desc("Pipeline to run to a fixed point")};

This gives me errors about OpPassManager not having a default constructor.

[silvas]

You can see my full code here: https://github.com/silvasean/iree/commit/cba1cf90fa8e8cc9b3bdacfd034f7d53c40661ad

[rriddle]

In D122880#3434468, @silvas wrote:

Hey River, I was trying to use this downstream and it isn't working for me

I have code like:

Option<OpPassManager> pipeline{
    *this, "pipeline", llvm::cl::desc("Pipeline to run to a fixed point")};

This gives me errors about OpPassManager not having a default constructor.

Yeah, llvm's command line support doesn't like non-default constructible things (each day I come closer to diverging from it entirely). For this though, I think it should technically work after a default constructor for OpPassManager is added in D123536.

mlir/include/mlir/Pass/PassManager.h

public:
   using const_pass_iterator =
       llvm::pointee_iterator<ArrayRef<std::unique_ptr<Pass>>::const_iterator>;
   const_pass_iterator begin() const;
   const_pass_iterator end() const;
   iterator_range<const_pass_iterator> getPasses() const {
     return {begin(), end()};
   }
 
+  /// Returns true if the pass manager has no passes.
+  bool empty() const { return begin() == end(); }
+
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(StringAttr nestedName);
   OpPassManager &nest(StringRef nestedName);
   template <typename OpT> OpPassManager &nest() {
     return nest(OpT::getOperationName());
   }
 
public:
 
   /// Returns the internal implementation instance.
   detail::OpPassManagerImpl &getImpl();
 
   /// Prints out the passes of the pass manager as the textual representation
   /// of pipelines.
   /// Note: The quality of the string representation depends entirely on the
   /// the correctness of per-pass overrides of Pass::printAsTextualPipeline.
-  void printAsTextualPipeline(raw_ostream &os);
+  void printAsTextualPipeline(raw_ostream &os) const;
 
   /// Raw dump of the pass manager to llvm::errs().
   void dump();
 
   /// Merge the pass statistics of this class into 'other'.
   void mergeStatisticsInto(OpPassManager &other);
 
   /// Register dependent dialects for the current pass manager.

mlir/include/mlir/Pass/PassOptions.h

 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include <memory>
 
 namespace mlir {
+class OpPassManager;
+
 namespace detail {
 namespace pass_options {
 /// Parse a string containing a list of comma-delimited elements, invoking the
 /// given parser for each sub-element and passing them to the provided
 /// element-append functor.
 LogicalResult
 parseCommaSeparatedList(llvm::cl::Option &opt, StringRef argName,
                         StringRef optionStr,
public:
 
   /// This class represents a specific pass option, with a provided data type.
   template <typename DataType, typename OptionParser = OptionParser<DataType>>
   class Option
       : public llvm::cl::opt<DataType, /*ExternalStorage=*/false, OptionParser>,
         public OptionBase {
   public:
     template <typename... Args>
-    Option(PassOptions &parent, StringRef arg, Args &&... args)
+    Option(PassOptions &parent, StringRef arg, Args &&...args)
         : llvm::cl::opt<DataType, /*ExternalStorage=*/false, OptionParser>(
               arg, llvm::cl::sub(parent), std::forward<Args>(args)...) {
       assert(!this->isPositional() && !this->isSink() &&
              "sink and positional options are not supported");
       parent.options.push_back(this);
 
       // Set a callback to track if this option has a value.
       this->setCallback([this](const auto &) { this->optHasValue = true; });

 /// to initialize a pass pipeline. See PassPipelineRegistration for usage
 /// details.
 ///
 /// Usage:
 ///
 /// struct MyPipelineOptions : PassPipelineOptions<MyPassOptions> {
 ///   ListOption<int> someListFlag{*this, "flag-name", llvm::cl::desc("...")};
 /// };
-template <typename T> class PassPipelineOptions : public detail::PassOptions {
+template <typename T>
+class PassPipelineOptions : public detail::PassOptions {
 public:
   /// Factory that parses the provided options and returns a unique_ptr to the
   /// struct.
   static std::unique_ptr<T> createFromString(StringRef options) {
     auto result = std::make_unique<T>();
     if (failed(result->parseFromString(options)))
       return nullptr;
     return result;
   }
 };
 
 /// A default empty option struct to be used for passes that do not need to take
 /// any options.
 struct EmptyPipelineOptions : public PassPipelineOptions<EmptyPipelineOptions> {
 };
-
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // MLIR Options
 //===----------------------------------------------------------------------===//
 
 namespace llvm {
 namespace cl {
public:
   parser(Option &opt) : detail::VectorParserBase<std::vector<T>, T>(opt) {}
 };
 template <typename T, unsigned N>
 class parser<SmallVector<T, N>>
     : public detail::VectorParserBase<SmallVector<T, N>, T> {
 public:
   parser(Option &opt) : detail::VectorParserBase<SmallVector<T, N>, T>(opt) {}
 };
-} // end namespace cl
-} // end namespace llvm
 
-#endif // MLIR_PASS_PASSOPTIONS_H_
+//===----------------------------------------------------------------------===//
+// OpPassManager: OptionValue
+
+template <>
+struct OptionValue<mlir::OpPassManager> final : GenericOptionValue {
+  using WrapperType = mlir::OpPassManager;
+
+  OptionValue();
+  OptionValue(const mlir::OpPassManager &value);
+  OptionValue<mlir::OpPassManager> &operator=(const mlir::OpPassManager &rhs);
+  ~OptionValue();
+
+  /// Returns if the current option has a value.
+  bool hasValue() const { return value.get(); }
+
+  /// Returns the current value of the option.
+  mlir::OpPassManager &getValue() const {
+    assert(hasValue() && "invalid option value");
+    return *value;
+  }
+
+  /// Set the value of the option.
+  void setValue(const mlir::OpPassManager &newValue);
+  void setValue(StringRef pipelineStr);
+
+  /// Compare the option with the provided value.
+  bool compare(const mlir::OpPassManager &rhs) const;
+  bool compare(const GenericOptionValue &rhs) const override {
+    const auto &rhsOV =
+        static_cast<const OptionValue<mlir::OpPassManager> &>(rhs);
+    if (!rhsOV.hasValue())
+      return false;
+    return compare(rhsOV.getValue());
+  }
+
+private:
+  void anchor() override;
+
+  /// The underlying pass manager. We use a unique_ptr to avoid the need for the
+  /// full type definition.
+  std::unique_ptr<mlir::OpPassManager> value;
+};
+
+//===----------------------------------------------------------------------===//
+// 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 OptVal = OptionValue<mlir::OpPassManager>;
+
+  parser(Option &opt) : basic_parser(opt) {}
+
+  bool parse(Option &, StringRef, StringRef arg, ParsedPassManager &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,
+                       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.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">,
-    ListOption<"opPipelineStrs", "op-pipelines", "std::string",
+    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

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 = dyn_cast<OpToOpPassAdaptor>(this)) {
-    llvm::interleaveComma(adaptor->getPassManagers(), os,
-                          [&](OpPassManager &pm) {
-                            os << pm.getOpName() << "(";
-                            pm.printAsTextualPipeline(os);
-                            os << ")";
-                          });
+    llvm::interleave(
+        adaptor->getPassManagers(),
+        },
+        [&] { os << ","; });
     return;
   }
   // Otherwise, print the pass argument followed by its options. If the pass
   // doesn't have an argument, print the name of the pass to give some indicator
   // of what pass was run.
   StringRef argument = getArgument();
   if (!argument.empty())
     os << argument;

 /// Return the operation name that this pass manager operates on.
 OperationName OpPassManager::getOpName(MLIRContext &context) const {
   return impl->getOpName(context);
 }
 
 /// Prints out the given passes as the textual representation of a pipeline.
 static void printAsTextualPipeline(ArrayRef<std::unique_ptr<Pass>> passes,
                                    raw_ostream &os) {
-  llvm::interleaveComma(passes, os, [&](const std::unique_ptr<Pass> &pass) {
+  llvm::interleave(
+      passes,
+      [&](const std::unique_ptr<Pass> &pass) {
         pass->printAsTextualPipeline(os);
-  });
+      },
+      [&] { os << ","; });
 }
 
 /// Prints out the passes of the pass manager as the textual representation
 /// of pipelines.
-void OpPassManager::printAsTextualPipeline(raw_ostream &os) {
+void OpPassManager::printAsTextualPipeline(raw_ostream &os) const {
   ::printAsTextualPipeline(impl->passes, os);
 }
 
 void OpPassManager::dump() {
   llvm::errs() << "Pass Manager with " << impl->passes.size() << " passes: ";
   ::printAsTextualPipeline(impl->passes, llvm::errs());
   llvm::errs() << "\n";
 }

mlir/lib/Pass/PassRegistry.cpp

 size_t detail::PassOptions::getOptionWidth() const {
   size_t max = 0;
   for (auto *option : options)
     max = std::max(max, option->getOption()->getOptionWidth());
   return max;
 }
 
 //===----------------------------------------------------------------------===//
+// MLIR Options
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// OpPassManager: OptionValue
+
+llvm::cl::OptionValue<OpPassManager>::OptionValue() = default;
+llvm::cl::OptionValue<OpPassManager>::OptionValue(
+    const mlir::OpPassManager &value) {
+  setValue(value);
+}
+llvm::cl::OptionValue<OpPassManager> &
+llvm::cl::OptionValue<OpPassManager>::operator=(
+    const mlir::OpPassManager &rhs) {
+  setValue(rhs);
+  return *this;
+}
+
+llvm::cl::OptionValue<OpPassManager>::~OptionValue<OpPassManager>() = default;
+
+void llvm::cl::OptionValue<OpPassManager>::setValue(
+    const OpPassManager &newValue) {
+  if (hasValue())
+    *value = newValue;
+  else
+    value = std::make_unique<mlir::OpPassManager>(newValue);
+}
+void llvm::cl::OptionValue<OpPassManager>::setValue(StringRef pipelineStr) {
+  FailureOr<OpPassManager> pipeline = parsePassPipeline(pipelineStr);
+  assert(succeeded(pipeline) && "invalid pass pipeline");
+  setValue(*pipeline);
+}
+
+bool llvm::cl::OptionValue<OpPassManager>::compare(
+    const mlir::OpPassManager &rhs) const {
+  std::string lhsStr, rhsStr;
+  {
+    raw_string_ostream lhsStream(lhsStr);
+    value->printAsTextualPipeline(lhsStream);
+
+    raw_string_ostream rhsStream(rhsStr);
+    rhs.printAsTextualPipeline(rhsStream);
+  }
+
+  // Use the textual format for pipeline comparisons.
+  return lhsStr == rhsStr;
+}
+
+void llvm::cl::OptionValue<OpPassManager>::anchor() {}
+
+//===----------------------------------------------------------------------===//
+// OpPassManager: Parser
+
+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) {
+  FailureOr<OpPassManager> pipeline = parsePassPipeline(arg);
+  if (failed(pipeline))
+    return true;
+  value.value = std::make_unique<OpPassManager>(std::move(*pipeline));
+  return false;
+}
+
+void llvm::cl::parser<OpPassManager>::print(raw_ostream &os,
+                                            const OpPassManager &value) {
+  value.printAsTextualPipeline(os);
+}
+
+void llvm::cl::parser<OpPassManager>::printOptionDiff(
+    const Option &opt, OpPassManager &pm, const OptVal &defaultValue,
+    size_t globalWidth) const {
+  printOptionName(opt, globalWidth);
+  outs() << "= ";
+  pm.printAsTextualPipeline(outs());
+
+  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:
   /// Try to initialize this pipeline with the given pipeline text.

mlir/lib/Transforms/Inliner.cpp

 
 InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
                          llvm::StringMap<OpPassManager> opPipelines)
     : InlinerPass(std::move(defaultPipeline)) {
   if (opPipelines.empty())
     return;
 
   // Update the option for the op specific optimization pipelines.
-  for (auto &it : opPipelines) {
-    std::string pipeline;
-    llvm::raw_string_ostream pipelineOS(pipeline);
-    pipelineOS << it.getKey() << "(";
-    it.second.printAsTextualPipeline(pipelineOS);
-    pipelineOS << ")";
-    opPipelineStrs.addValue(pipeline);
-  }
+  for (auto &it : opPipelines)
+    opPipelineList.addValue(it.second);
   this->opPipelines.emplace_back(std::move(opPipelines));
 }
 
 void InlinerPass::runOnOperation() {
   CallGraph &cg = getAnalysis<CallGraph>();
   auto *context = &getContext();
 
   // The inliner should only be run on operations that define a symbol table,
defaultPipeline = [=](OpPassManager &pm) {
       (void)parsePassPipeline(defaultPipelineCopy, pm);
     };
   } else if (defaultPipelineStr.getNumOccurrences()) {
     defaultPipeline = nullptr;
   }
 
   // Initialize the op specific pass pipelines.
   llvm::StringMap<OpPassManager> pipelines;
-  for (StringRef pipeline : opPipelineStrs) {
-    // Skip empty pipelines.
-    if (pipeline.empty())
-      continue;
-    FailureOr<OpPassManager> pm = parsePassPipeline(pipeline);
-    if (failed(pm))
-      return failure();
-    pipelines.try_emplace(pm->getOpName(), std::move(*pm));
-  }
+  for (OpPassManager pipeline : opPipelineList)
+    if (!pipeline.empty())
+      pipelines.try_emplace(pipeline.getOpName(), pipeline);
   opPipelines.assign({std::move(pipelines)});
 
   return success();
 }
 
 std::unique_ptr<Pass> mlir::createInlinerPass() {
   return std::make_unique<InlinerPass>();
 }

mlir/lib/Transforms/PassDetail.h

 //===- PassDetail.h - Transforms Pass class details -------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TRANSFORMS_PASSDETAIL_H_
 #define TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
 namespace mlir {
 #define GEN_PASS_CLASSES
 #include "mlir/Transforms/Passes.h.inc"
 } // namespace mlir
 
 #endif // TRANSFORMS_PASSDETAIL_H_

mlir/test/Pass/crash-recovery.mlir

 // UNSUPPORTED: asan
 
 // expected-error@below {{Failures have been detected while processing an MLIR pass pipeline}}
 // expected-note@below {{Pipeline failed while executing}}
 module @inner_mod1 {
   module @foo {}
 }
 
-// REPRO: configuration: -pass-pipeline='builtin.module(test-module-pass, test-pass-crash)'
+// REPRO: configuration: -pass-pipeline='builtin.module(test-module-pass,test-pass-crash)'
 
 // REPRO: module @inner_mod1
 // REPRO: module @foo {
 
 // REPRO_LOCAL: configuration: -pass-pipeline='builtin.module(test-pass-crash)'
 
 // REPRO_LOCAL: module @inner_mod1
 // REPRO_LOCAL: module @foo {
 
 // REPRO_LOCAL_DYNAMIC: configuration: -pass-pipeline='builtin.module(test-pass-crash)'
 
 // REPRO_LOCAL_DYNAMIC: module @inner_mod1
 // REPRO_LOCAL_DYNAMIC: module @foo {

mlir/test/Pass/pipeline-options-parsing.mlir

 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(test-module-pass{)' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_1 %s
 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(test-module-pass{test-option=3})' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_2 %s
 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(func.func(test-options-pass{list=3}), test-module-pass{invalid-option=3})' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_3 %s
 // RUN: not mlir-opt %s -pass-pipeline='test-options-pass{list=3 list=notaninteger}' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_4 %s
 // RUN: mlir-opt %s -pass-pipeline='func.func(test-options-pass{list=1,2,3,4 list=5 string=value1 string=value2})'
 // RUN: mlir-opt %s -verify-each=false -pass-pipeline='func.func(test-options-pass{string-list=a list=1,2,3,4 string-list=b,c list=5 string-list=d string=nested_pipeline{arg1=10 arg2=" {} " arg3=true}})' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_1 %s
 // RUN: mlir-opt %s -verify-each=false -test-options-pass-pipeline='list=1 string-list=a,b' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_2 %s
 // RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(func.func(test-options-pass{list=3}), func.func(test-options-pass{list=1,2,3,4}))' -test-dump-pipeline 2>&1 | FileCheck --check-prefix=CHECK_3 %s
 
 // CHECK_ERROR_1: missing closing '}' while processing pass options
 // CHECK_ERROR_2: no such option test-option
 // CHECK_ERROR_3: no such option invalid-option
 // CHECK_ERROR_4: 'notaninteger' value invalid for integer argument
 
 // CHECK_1: test-options-pass{list=1,2,3,4,5 string=nested_pipeline{arg1=10 arg2=" {} " arg3=true} string-list=a,b,c,d}
 // CHECK_2: test-options-pass{list=1 string= string-list=a,b}
-// CHECK_3: builtin.module(func.func(test-options-pass{list=3 string= }), func.func(test-options-pass{list=1,2,3,4 string= }))
+// CHECK_3: builtin.module(func.func(test-options-pass{list=3 string= }),func.func(test-options-pass{list=1,2,3,4 string= }))

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
 
 // Inline a function that takes an argument.
 func @func_with_arg(%c : i32) -> i32 {
   %b = arith.addi %c, %c : i32
   return %b : i32
 }
 
 // CHECK-LABEL: func @inline_with_arg