[MLIR] Add custom option parser type to PassOption

[ezhulenev]

Summary

Pass custom DataType parser to pass options, similar to llvm:🆑:opt class.

[rriddle]

Thanks Eugene, looks great.

[jpienaar]

Thanks!

mlir/include/mlir/Pass/Pass.h

public:
   /// this is a generic OperationPass.
   Optional<StringRef> getOpName() const { return opName; }
 
   //===--------------------------------------------------------------------===//
   // Options
   //===--------------------------------------------------------------------===//
 
   /// This class represents a specific pass option, with a provided data type.
-  template <typename DataType>
-  struct Option : public detail::PassOptions::Option<DataType> {
+  template <typename DataType,
+            typename OptionParser = detail::PassOptions::OptionParser<DataType>>
+  struct Option : public detail::PassOptions::Option<DataType, OptionParser> {
     template <typename... Args>
     Option(Pass &parent, StringRef arg, Args &&... args)
-        : detail::PassOptions::Option<DataType>(parent.passOptions, arg,
-                                                std::forward<Args>(args)...) {}
-    using detail::PassOptions::Option<DataType>::operator=;
+        : detail::PassOptions::Option<DataType, OptionParser>(
+              parent.passOptions, arg, std::forward<Args>(args)...) {}
+    using detail::PassOptions::Option<DataType, OptionParser>::operator=;
   };
   /// This class represents a specific pass option that contains a list of
   /// values of the provided data type.
-  template <typename DataType>
-  struct ListOption : public detail::PassOptions::ListOption<DataType> {
+  template <typename DataType,
+            typename OptionParser = detail::PassOptions::OptionParser<DataType>>
+  struct ListOption
+      : public detail::PassOptions::ListOption<DataType, OptionParser> {
     template <typename... Args>
     ListOption(Pass &parent, StringRef arg, Args &&... args)
-        : detail::PassOptions::ListOption<DataType>(
+        : detail::PassOptions::ListOption<DataType, OptionParser>(
               parent.passOptions, arg, std::forward<Args>(args)...) {}
-    using detail::PassOptions::ListOption<DataType>::operator=;
+    using detail::PassOptions::ListOption<DataType, OptionParser>::operator=;
   };
 
   /// Attempt to initialize the options of this pass from the given string.
   LogicalResult initializeOptions(StringRef options);
 
   /// 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 printAsTextualPipeline(raw_ostream &os);

mlir/include/mlir/Pass/PassOptions.h

struct GenericOptionParser : public llvm::cl::parser<DataType> {
     Optional<StringRef> findArgStrForValue(const DataType &value) {
       for (auto &it : this->Values)
         if (it.V.compare(value))
           return it.Name;
       return llvm::None;
     }
   };
 
-  /// The specific parser to use depending on llvm::cl parser used. This is only
-  /// necessary because we need to provide additional methods for certain data
-  /// type parsers.
-  /// TODO(riverriddle) We should upstream the methods in GenericOptionParser to
-  /// avoid the need to do this.
-  template <typename DataType>
-  using OptionParser =
-      std::conditional_t<std::is_base_of<llvm::cl::generic_parser_base,
-                                         llvm::cl::parser<DataType>>::value,
-                         GenericOptionParser<DataType>,
-                         llvm::cl::parser<DataType>>;
-
   /// Utility methods for printing option values.
   template <typename DataT>
   static void printValue(raw_ostream &os, GenericOptionParser<DataT> &parser,
                          const DataT &value) {
     if (Optional<StringRef> argStr = parser.findArgStrForValue(value))
       os << argStr;
     else
       llvm_unreachable("unknown data value for option");
   }
   template <typename DataT, typename ParserT>
   static void printValue(raw_ostream &os, ParserT &parser, const DataT &value) {
     os << value;
   }
   template <typename ParserT>
   static void printValue(raw_ostream &os, ParserT &parser, const bool &value) {
     os << (value ? StringRef("true") : StringRef("false"));
   }
 
 public:
-  /// This class represents a specific pass option, with a provided data type.
+  /// The specific parser to use depending on llvm::cl parser used. This is only
+  /// necessary because we need to provide additional methods for certain data
+  /// type parsers.
+  /// TODO(riverriddle) We should upstream the methods in GenericOptionParser to
+  /// avoid the need to do this.
   template <typename DataType>
-  class Option : public llvm::cl::opt<DataType, /*ExternalStorage=*/false,
-                                      OptionParser<DataType>>,
-                 public OptionBase {
+  using OptionParser =
+      std::conditional_t<std::is_base_of<llvm::cl::generic_parser_base,
+                                         llvm::cl::parser<DataType>>::value,
+                         GenericOptionParser<DataType>,
+                         llvm::cl::parser<DataType>>;
+
+  /// 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:
     template <typename... Args>
     Option(PassOptions &parent, StringRef arg, Args &&... args)
-        : llvm::cl::opt<DataType, /*ExternalStorage=*/false,
-                        OptionParser<DataType>>(arg, llvm::cl::sub(parent),
-                                                std::forward<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);
     }
     using llvm::cl::opt<DataType, /*ExternalStorage=*/false,
-                        OptionParser<DataType>>::operator=;
+                        OptionParser>::operator=;
     ~Option() override = default;
 
   private:
     /// Return the main option instance.
     const llvm::cl::Option *getOption() const final { return this; }
 
     /// Print the name and value of this option to the given stream.
     void print(raw_ostream &os) final {
       os << this->ArgStr << '=';
       printValue(os, this->getParser(), this->getValue());
     }
 
     /// Copy the value from the given option into this one.
     void copyValueFrom(const OptionBase &other) final {
-      this->setValue(static_cast<const Option<DataType> &>(other).getValue());
+      this->setValue(static_cast<const Option<DataType, OptionParser> &>(other)
+                         .getValue());
     }
   };
 
   /// This class represents a specific pass option that contains a list of
   /// values of the provided data type.
-  template <typename DataType>
-  class ListOption : public llvm::cl::list<DataType, /*StorageClass=*/bool,
-                                           OptionParser<DataType>>,
-                     public OptionBase {
+  template <typename DataType, typename OptionParser = OptionParser<DataType>>
+  class ListOption
+      : public llvm::cl::list<DataType, /*StorageClass=*/bool, OptionParser>,
   public:
     template <typename... Args>
     ListOption(PassOptions &parent, StringRef arg, Args &&... args)
-        : llvm::cl::list<DataType, /*StorageClass=*/bool,
-                         OptionParser<DataType>>(arg, llvm::cl::sub(parent),
-                                                 std::forward<Args>(args)...) {
+        : llvm::cl::list<DataType, /*StorageClass=*/bool, 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);
     }
     ~ListOption() override = default;
 
     /// Allow assigning from an ArrayRef.
-    ListOption<DataType> &operator=(ArrayRef<DataType> values) {
+    ListOption<DataType, OptionParser> &operator=(ArrayRef<DataType> values) {
       (*this)->assign(values.begin(), values.end());
       return *this;
     }
 
     std::vector<DataType> *operator->() { return &*this; }
 
   private:
     /// Return the main option instance.
     const llvm::cl::Option *getOption() const final { return this; }
 
     /// Print the name and value of this option to the given stream.
     void print(raw_ostream &os) final {
       os << this->ArgStr << '=';
       auto printElementFn = [&](const DataType &value) {
         printValue(os, this->getParser(), value);
       };
       interleave(*this, os, printElementFn, ",");
     }
 
     /// 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)));
+          (ListOption<DataType, OptionParser> &)(const_cast<OptionBase &>(
+              other)));
     }
   };
 
   PassOptions() = default;
   /// Delete the copy constructor to avoid copying the internal options map.
   PassOptions(const PassOptions &) = delete;
   PassOptions(PassOptions &&) = delete;