[mlir][Pass] Enable printing pass options as part of `-help`.

[rriddle]

Summary

This revision adds support for printing pass options as part of the normal help description. This also moves registered passes and pipelines into different sections of the help.

Example:

Compiler passes to run
  --pass-pipeline                                     -   ...
  Passes:
    --affine-data-copy-generate                       -   ...
    --convert-gpu-to-spirv                            -   ...
      --workgroup-size=<long>                         - ...
    --test-options-pass                               -   ...
      --list=<int>                                    - ...
      --string=<string>                               - ...
      --string-list=<string>                          - ...
  Pass Pipelines:
    --test-options-pass-pipeline                      -   ...
      --list=<int>                                    - ...
      --string=<string>                               - ...
      --string-list=<string>                          - ...

[jpienaar]

Nice, looks much cleaner. Mostly small nits.

mlir/include/mlir/Pass/PassOptions.h

191

Why not make it non-copyable?

mlir/lib/Pass/PassRegistry.cpp

43

So it has numSpaces + 1 empty space before - ? Could llvm::format make this simpler?

199–202

Nit: Hoist out into cmp lambda? I think it would make it easier to read (and may actually be shorter as it reduces whitespace)

550

Any way to make these less magical? E.g., we have 4 here, 5 above, 6 below and 2 for indent (which is matches my editor config so I won't complain 😛 )

[rriddle]

Resolve comments

mlir/include/mlir/Pass/PassOptions.h

191

Done. I originally had this so that users could still call the copy constructor on Pass, but I just overrode that one as well.

mlir/lib/Pass/PassRegistry.cpp

550

I aligned all of the magic. I need to do a followup to the core LLVM CommandLine library at some point to allow for passing in a base indent.

mlir/include/mlir/Pass/Pass.h

public:
 
   /// Returns the main statistics for this pass instance.
   ArrayRef<Statistic *> getStatistics() const { return statistics; }
   MutableArrayRef<Statistic *> getStatistics() { return statistics; }
 
 protected:
   explicit Pass(const PassID *passID, Optional<StringRef> opName = llvm::None)
       : passID(passID), opName(opName) {}
+  Pass(const Pass &other) : Pass(other.passID, other.opName) {}
 
   /// Returns the current pass state.
   detail::PassExecutionState &getPassState() {
     assert(passState && "pass state was never initialized");
     return *passState;
   }
 
   /// Return the MLIR context for the current function being transformed.
private:
   /// The set of statistics held by this pass.
   std::vector<Statistic *> statistics;
 
   /// The pass options registered to this pass instance.
   detail::PassOptions passOptions;
 
   /// Allow access to 'clone' and 'run'.
   friend class OpPassManager;
+
+  /// Allow access to 'passOptions'.
+  friend class PassInfo;
 };
 
 //===----------------------------------------------------------------------===//
 // Pass Model Definitions
 //===----------------------------------------------------------------------===//
 namespace detail {
 /// The opaque CRTP model of a pass. This class provides utilities for derived
 /// pass execution and handles all of the necessary polymorphic API.

mlir/include/mlir/Pass/PassOptions.h

private:
     /// Copy the value from the given option into this one.
     void copyValueFrom(const OptionBase &other) final {
       (*this) = ArrayRef<DataType>(
           (ListOption<DataType> &)(const_cast<OptionBase &>(other)));
     }
   };
 
   PassOptions() = default;
+  /// Delete the copy constructor to avoid copying the internal options map.
+  PassOptions(const PassOptions &) = delete;
+  PassOptions(PassOptions &&) = delete;
 
   /// Copy the option values from 'other' into 'this', where 'other' has the
   /// same options as 'this'.
   void copyOptionValuesFrom(const PassOptions &other);
 
   /// Parse options out as key=value pairs that can then be handed off to the
   /// `llvm::cl` command line passing infrastructure. Everything is space
   /// separated.
   LogicalResult parseFromString(StringRef options);
 
   /// Print the options held by this struct in a form that can be parsed via
   /// 'parseFromString'.
   void print(raw_ostream &os);
 
+  /// Print the help string for the options held by this struct. `descIndent` is
+  /// the indent that the descriptions should be aligned.
+  void printHelp(size_t indent, size_t descIndent) const;
+
+  /// Return the maximum width required when printing the help string.
+  size_t getOptionWidth() const;
+
 private:
   /// A list of all of the opaque options.
   std::vector<OptionBase *> options;
 };
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
 // PassPipelineOptions

mlir/include/mlir/Pass/PassRegistry.h

 
 #include "mlir/Pass/PassOptions.h"
 #include <functional>
 
 namespace mlir {
 class OpPassManager;
 class Pass;
 
+namespace detail {
+class PassOptions;
+} // end namespace detail
+
 /// A registry function that adds passes to the given pass manager. This should
 /// also parse options and return success() if parsing succeeded.
 using PassRegistryFunction =
     std::function<LogicalResult(OpPassManager &, StringRef options)>;
 using PassAllocatorFunction = std::function<std::unique_ptr<Pass>()>;
 
 /// A special type used by transformation passes to provide an address that can
 /// act as a unique identifier during pass registration.
public:
 
   /// Returns the command line option that may be passed to 'mlir-opt' that will
   /// cause this pass to run or null if there is no such argument.
   StringRef getPassArgument() const { return arg; }
 
   /// Returns a description for the pass, this never returns null.
   StringRef getPassDescription() const { return description; }
 
+  /// Print the help information for this pass. This includes the argument,
+  /// description, and any pass options. `descIndent` is the indent that the
+  /// descriptions should be aligned.
+  void printHelpStr(size_t indent, size_t descIndent) const;
+
+  /// Return the maximum width required when printing the options of this entry.
+  size_t getOptionWidth() const;
+
 protected:
-  PassRegistryEntry(StringRef arg, StringRef description,
-                    const PassRegistryFunction &builder)
-      : arg(arg), description(description), builder(builder) {}
+  PassRegistryEntry(
+      StringRef arg, StringRef description, const PassRegistryFunction &builder,
+      std::function<void(function_ref<void(const detail::PassOptions &)>)>
+          optHandler)
+      : arg(arg), description(description), builder(builder),
+        optHandler(optHandler) {}
 
 private:
-  // The argument with which to invoke the pass via mlir-opt.
+  /// The argument with which to invoke the pass via mlir-opt.
   StringRef arg;
 
-  // Description of the pass.
+  /// Description of the pass.
   StringRef description;
 
-  // Function to register this entry to a pass manager pipeline.
+  /// Function to register this entry to a pass manager pipeline.
   PassRegistryFunction builder;
+
+  /// Function to invoke a handler for a pass options instance.
+  std::function<void(function_ref<void(const detail::PassOptions &)>)>
+      optHandler;
 };
 
 /// A structure to represent the information of a registered pass pipeline.
 class PassPipelineInfo : public PassRegistryEntry {
 public:
-  PassPipelineInfo(StringRef arg, StringRef description,
-                   const PassRegistryFunction &builder)
-      : PassRegistryEntry(arg, description, builder) {}
+  PassPipelineInfo(
+      StringRef arg, StringRef description, const PassRegistryFunction &builder,
+      std::function<void(function_ref<void(const detail::PassOptions &)>)>
+          optHandler)
+      : PassRegistryEntry(arg, description, builder, optHandler) {}
 };
 
 /// A structure to represent the information for a derived pass class.
 class PassInfo : public PassRegistryEntry {
 public:
   /// PassInfo constructor should not be invoked directly, instead use
   /// PassRegistration or registerPass.
   PassInfo(StringRef arg, StringRef description, const PassID *passID,
            const PassAllocatorFunction &allocator);
 };
 
 //===----------------------------------------------------------------------===//
 // PassRegistration
 //===----------------------------------------------------------------------===//
 
 /// Register a specific dialect pipeline registry function with the system,
 /// typically used through the PassPipelineRegistration template.
-void registerPassPipeline(StringRef arg, StringRef description,
-                          const PassRegistryFunction &function);
+void registerPassPipeline(
+    StringRef arg, StringRef description, const PassRegistryFunction &function,
+    std::function<void(function_ref<void(const detail::PassOptions &)>)>
+        optHandler);
 
 /// Register a specific dialect pass allocator function with the system,
 /// typically used through the PassRegistration template.
 void registerPass(StringRef arg, StringRef description, const PassID *passID,
                   const PassAllocatorFunction &function);
 
 /// PassRegistration provides a global initializer that registers a Pass
 /// allocation routine for a concrete pass instance. The third argument is
 /// optional and provides a callback to construct a pass that does not have
 /// a default constructor.
 ///
 /// Usage:
 ///
 ///   /// At namespace scope.
 ///   static PassRegistration<MyPass> reg("my-pass", "My Pass Description.");
 ///
 template <typename ConcretePass> struct PassRegistration {
-
   PassRegistration(StringRef arg, StringRef description,
                    const PassAllocatorFunction &constructor) {
     registerPass(arg, description, PassID::getID<ConcretePass>(), constructor);
   }
 
   PassRegistration(StringRef arg, StringRef description)
       : PassRegistration(arg, description,
                          [] { return std::make_unique<ConcretePass>(); }) {}

 ///
 ///   static PassPipelineRegistration Unused("unused", "Unused pass",
 ///                                          pipelineBuilder);
 template <typename Options = EmptyPipelineOptions>
 struct PassPipelineRegistration {
   PassPipelineRegistration(
       StringRef arg, StringRef description,
       std::function<void(OpPassManager &, const Options &options)> builder) {
-    registerPassPipeline(arg, description,
-                         [builder](OpPassManager &pm, StringRef optionsStr) {
-                           Options options;
-                           if (failed(options.parseFromString(optionsStr)))
-                             return failure();
-                           builder(pm, options);
-                           return success();
-                         });
+    registerPassPipeline(
+        arg, description,
+        },
+        [](function_ref<void(const detail::PassOptions &)> optHandler) {
+          optHandler(Options());
   }
 };
 
 /// Convenience specialization of PassPipelineRegistration for EmptyPassOptions
 /// that does not pass an empty options struct to the pass builder function.
 template <> struct PassPipelineRegistration<EmptyPipelineOptions> {
   PassPipelineRegistration(StringRef arg, StringRef description,
                            std::function<void(OpPassManager &)> builder) {
-    registerPassPipeline(arg, description,
-                         [builder](OpPassManager &pm, StringRef optionsStr) {
-                           if (!optionsStr.empty())
-                             return failure();
-                           builder(pm);
-                           return success();
-                         });
+    registerPassPipeline(
+        arg, description,
+        },
+        [](function_ref<void(const detail::PassOptions &)>) {});
   }
 };
 
 /// This function parses the textual representation of a pass pipeline, and adds
 /// the result to 'pm' on success. This function returns failure if the given
 /// pipeline was invalid. 'errorStream' is the output stream used to emit errors
 /// found during parsing.
 LogicalResult parsePassPipeline(StringRef pipeline, OpPassManager &pm,

mlir/lib/Pass/PassRegistry.cpp

buildDefaultRegistryFn(const PassAllocatorFunction &allocator) {
   return [=](OpPassManager &pm, StringRef options) {
     std::unique_ptr<Pass> pass = allocator();
     LogicalResult result = pass->initializeOptions(options);
     pm.addPass(std::move(pass));
     return result;
   };
 }
 
+/// Utility to print the help string for a specific option.
+void printOptionHelp(StringRef arg, StringRef desc, size_t indent,
+                     size_t descIndent, bool isTopLevel) {
+  size_t numSpaces = descIndent - indent - 4;
+  llvm::outs().indent(indent)
+      << "--" << llvm::left_justify(arg, numSpaces) << "-   " << desc << '\n';
+}
+
+//===----------------------------------------------------------------------===//
+// PassRegistry
+//===----------------------------------------------------------------------===//
+
+/// Print the help information for this pass. This includes the argument,
+/// description, and any pass options. `descIndent` is the indent that the
+/// descriptions should be aligned.
+void PassRegistryEntry::printHelpStr(size_t indent, size_t descIndent) const {
+  printOptionHelp(getPassArgument(), getPassDescription(), indent, descIndent,
+                  /*isTopLevel=*/true);
+  // If this entry has options, print the help for those as well.
+  optHandler([=](const PassOptions &options) {
+    options.printHelp(indent, descIndent);
+  });
+}
+
+/// Return the maximum width required when printing the options of this
+/// entry.
+size_t PassRegistryEntry::getOptionWidth() const {
+  size_t maxLen = 0;
+  optHandler([&](const PassOptions &options) mutable {
+    maxLen = options.getOptionWidth() + 2;
+  });
+  return maxLen;
+}
+
 //===----------------------------------------------------------------------===//
 // PassPipelineInfo
 //===----------------------------------------------------------------------===//
 
-void mlir::registerPassPipeline(StringRef arg, StringRef description,
-                                const PassRegistryFunction &function) {
-  PassPipelineInfo pipelineInfo(arg, description, function);
+void mlir::registerPassPipeline(
+    StringRef arg, StringRef description, const PassRegistryFunction &function,
+    std::function<void(function_ref<void(const PassOptions &)>)> optHandler) {
+  PassPipelineInfo pipelineInfo(arg, description, function, optHandler);
   bool inserted = passPipelineRegistry->try_emplace(arg, pipelineInfo).second;
   assert(inserted && "Pass pipeline registered multiple times");
   (void)inserted;
 }
 
 //===----------------------------------------------------------------------===//
 // PassInfo
 //===----------------------------------------------------------------------===//
 
 PassInfo::PassInfo(StringRef arg, StringRef description, const PassID *passID,
                    const PassAllocatorFunction &allocator)
-    : PassRegistryEntry(arg, description, buildDefaultRegistryFn(allocator)) {}
+    : PassRegistryEntry(
+          arg, description, buildDefaultRegistryFn(allocator),
+          // Use a temporary pass to provide an options instance.
+          [=](function_ref<void(const PassOptions &)> optHandler) {
+            optHandler(allocator()->passOptions);
+          }) {}
 
 void mlir::registerPass(StringRef arg, StringRef description,
                         const PassID *passID,
                         const PassAllocatorFunction &function) {
   PassInfo passInfo(arg, description, passID, function);
   bool inserted = passRegistry->try_emplace(passID, passInfo).second;
   assert(inserted && "Pass registered multiple times");
   (void)inserted;

 /// Print the options held by this struct in a form that can be parsed via
 /// 'parseFromString'.
 void detail::PassOptions::print(raw_ostream &os) {
   // If there are no options, there is nothing left to do.
   if (OptionsMap.empty())
     return;
 
   // Sort the options to make the ordering deterministic.
-  SmallVector<OptionBase *, 4> orderedOptions(options.begin(), options.end());
-  llvm::array_pod_sort(orderedOptions.begin(), orderedOptions.end(),
-                       [](OptionBase *const *lhs, OptionBase *const *rhs) {
-                         return (*lhs)->getArgStr().compare(
-                             (*rhs)->getArgStr());
-                       });
+  SmallVector<OptionBase *, 4> orderedOps(options.begin(), options.end());
+  auto compareOptionArgs = [](OptionBase *const *lhs, OptionBase *const *rhs) {
+    return (*lhs)->getArgStr().compare((*rhs)->getArgStr());
+  };
+  llvm::array_pod_sort(orderedOps.begin(), orderedOps.end(), compareOptionArgs);
 
   // Interleave the options with ' '.
   os << '{';
   interleave(
-      orderedOptions, os, [&](OptionBase *option) { option->print(os); }, " ");
+      orderedOps, os, [&](OptionBase *option) { option->print(os); }, " ");
   os << '}';
 }
 
+/// Print the help string for the options held by this struct. `descIndent` is
+/// the indent within the stream that the descriptions should be aligned.
+void detail::PassOptions::printHelp(size_t indent, size_t descIndent) const {
+  // Sort the options to make the ordering deterministic.
+  SmallVector<OptionBase *, 4> orderedOps(options.begin(), options.end());
+  auto compareOptionArgs = [](OptionBase *const *lhs, OptionBase *const *rhs) {
+    return (*lhs)->getArgStr().compare((*rhs)->getArgStr());
+  };
+  llvm::array_pod_sort(orderedOps.begin(), orderedOps.end(), compareOptionArgs);
+  for (OptionBase *option : orderedOps) {
+    // TODO(riverriddle) printOptionInfo assumes a specific indent and will
+    // print options with values with incorrect indentation. We should add
+    // support to llvm::cl::Option for passing in a base indent to use when
+    // printing.
+    llvm::outs().indent(indent);
+    option->getOption()->printOptionInfo(descIndent - indent);
+  }
+}
+
+/// Return the maximum width required when printing the help string.
+size_t detail::PassOptions::getOptionWidth() const {
+  size_t max = 0;
+  for (auto *option : options)
+    max = std::max(max, option->getOption()->getOptionWidth());
+  return max;
+}
+
 //===----------------------------------------------------------------------===//
 // TextualPassPipeline Parser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a textual description of a pass pipeline.
 class TextualPipeline {
 public:

 /// Adds command line option for each registered pass or pass pipeline, as well
 /// as textual pass pipelines.
 struct PassNameParser : public llvm::cl::parser<PassArgData> {
   PassNameParser(llvm::cl::Option &opt) : llvm::cl::parser<PassArgData>(opt) {}
 
   void initialize();
   void printOptionInfo(const llvm::cl::Option &opt,
                        size_t globalWidth) const override;
+  size_t getOptionWidth(const llvm::cl::Option &opt) const override;
   bool parse(llvm::cl::Option &opt, StringRef argName, StringRef arg,
              PassArgData &value);
 };
 } // namespace
 
 void PassNameParser::initialize() {
   llvm::cl::parser<PassArgData>::initialize();
 
   /// Add an entry for the textual pass pipeline option.
   addLiteralOption(passPipelineArg, PassArgData(),
                    "A textual description of a pass pipeline to run");
 
   /// Add the pass entries.
   for (const auto &kv : *passRegistry) {
     addLiteralOption(kv.second.getPassArgument(), &kv.second,
                      kv.second.getPassDescription());
   }
   /// Add the pass pipeline entries.
   for (const auto &kv : *passPipelineRegistry) {
     addLiteralOption(kv.second.getPassArgument(), &kv.second,
                      kv.second.getPassDescription());
   }
 }
 
-void PassNameParser::printOptionInfo(const llvm::cl::Option &O,
-                                     size_t GlobalWidth) const {
-  PassNameParser *TP = const_cast<PassNameParser *>(this);
-  llvm::array_pod_sort(TP->Values.begin(), TP->Values.end(),
-                       [](const PassNameParser::OptionInfo *VT1,
-                          const PassNameParser::OptionInfo *VT2) {
-                         return VT1->Name.compare(VT2->Name);
-                       });
-  llvm::cl::parser<PassArgData>::printOptionInfo(O, GlobalWidth);
+void PassNameParser::printOptionInfo(const llvm::cl::Option &opt,
+                                     size_t globalWidth) const {
+  // Print the information for the top-level option.
+  if (opt.hasArgStr()) {
+    llvm::outs() << "  --" << opt.ArgStr;
+    opt.printHelpStr(opt.HelpStr, globalWidth, opt.ArgStr.size() + 7);
+  } else {
+    llvm::outs() << "  " << opt.HelpStr << '\n';
+  }
+
+  // Print the top-level pipeline argument.
+  printOptionHelp(passPipelineArg,
+                  "A textual description of a pass pipeline to run",
+                  /*indent=*/4, globalWidth, /*isTopLevel=*/!opt.hasArgStr());
+
+  // Functor used to print the ordered entries of a registration map.
+  auto printOrderedEntries = [&](StringRef header, auto &map) {
+    llvm::SmallVector<PassRegistryEntry *, 32> orderedEntries;
+    for (auto &kv : map)
+      orderedEntries.push_back(&kv.second);
+    llvm::array_pod_sort(
+        orderedEntries.begin(), orderedEntries.end(),
+        [](PassRegistryEntry *const *lhs, PassRegistryEntry *const *rhs) {
+          return (*lhs)->getPassArgument().compare((*rhs)->getPassArgument());
+
+    llvm::outs().indent(4) << header << ":\n";
+    for (PassRegistryEntry *entry : orderedEntries)
+      entry->printHelpStr(/*indent=*/6, globalWidth);
+  };
+
+  // Print the available passes.
+  printOrderedEntries("Passes", *passRegistry);
+
+  // Print the available pass pipelines.
+  if (!passPipelineRegistry->empty())
+    printOrderedEntries("Pass Pipelines", *passPipelineRegistry);
+}
+
+size_t PassNameParser::getOptionWidth(const llvm::cl::Option &opt) const {
+  size_t maxWidth = llvm::cl::parser<PassArgData>::getOptionWidth(opt) + 2;
+
+  // Check for any wider pass or pipeline options.
+  for (auto &entry : *passRegistry)
+    maxWidth = std::max(maxWidth, entry.second.getOptionWidth() + 4);
+  for (auto &entry : *passPipelineRegistry)
+    maxWidth = std::max(maxWidth, entry.second.getOptionWidth() + 4);
+  return maxWidth;
 }
 
 bool PassNameParser::parse(llvm::cl::Option &opt, StringRef argName,
                            StringRef arg, PassArgData &value) {
   // Handle the pipeline option explicitly.
   if (argName == passPipelineArg)
     return failed(value.pipeline.initialize(arg, llvm::errs()));