1 //===- PassRegistry.cpp - Pass Registration Utilities ---------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include <utility>
10 
11 #include "mlir/Pass/PassRegistry.h"
12 #include "mlir/Pass/Pass.h"
13 #include "mlir/Pass/PassManager.h"
14 #include "llvm/ADT/DenseMap.h"
15 #include "llvm/Support/Format.h"
16 #include "llvm/Support/ManagedStatic.h"
17 #include "llvm/Support/MemoryBuffer.h"
18 #include "llvm/Support/SourceMgr.h"
19 
20 using namespace mlir;
21 using namespace detail;
22 
23 /// Static mapping of all of the registered passes.
24 static llvm::ManagedStatic<llvm::StringMap<PassInfo>> passRegistry;
25 
26 /// A mapping of the above pass registry entries to the corresponding TypeID
27 /// of the pass that they generate.
28 static llvm::ManagedStatic<llvm::StringMap<TypeID>> passRegistryTypeIDs;
29 
30 /// Static mapping of all of the registered pass pipelines.
31 static llvm::ManagedStatic<llvm::StringMap<PassPipelineInfo>>
32     passPipelineRegistry;
33 
34 /// Utility to create a default registry function from a pass instance.
35 static PassRegistryFunction
36 buildDefaultRegistryFn(const PassAllocatorFunction &allocator) {
37   return [=](OpPassManager &pm, StringRef options,
38              function_ref<LogicalResult(const Twine &)> errorHandler) {
39     std::unique_ptr<Pass> pass = allocator();
40     LogicalResult result = pass->initializeOptions(options);
41     if ((pm.getNesting() == OpPassManager::Nesting::Explicit) &&
42         pass->getOpName() && *pass->getOpName() != pm.getOpName())
43       return errorHandler(llvm::Twine("Can't add pass '") + pass->getName() +
44                           "' restricted to '" + *pass->getOpName() +
45                           "' on a PassManager intended to run on '" +
46                           pm.getOpName() + "', did you intend to nest?");
47     pm.addPass(std::move(pass));
48     return result;
49   };
50 }
51 
52 /// Utility to print the help string for a specific option.
53 static void printOptionHelp(StringRef arg, StringRef desc, size_t indent,
54                             size_t descIndent, bool isTopLevel) {
55   size_t numSpaces = descIndent - indent - 4;
56   llvm::outs().indent(indent)
57       << "--" << llvm::left_justify(arg, numSpaces) << "-   " << desc << '\n';
58 }
59 
60 //===----------------------------------------------------------------------===//
61 // PassRegistry
62 //===----------------------------------------------------------------------===//
63 
64 /// Print the help information for this pass. This includes the argument,
65 /// description, and any pass options. `descIndent` is the indent that the
66 /// descriptions should be aligned.
67 void PassRegistryEntry::printHelpStr(size_t indent, size_t descIndent) const {
68   printOptionHelp(getPassArgument(), getPassDescription(), indent, descIndent,
69                   /*isTopLevel=*/true);
70   // If this entry has options, print the help for those as well.
71   optHandler([=](const PassOptions &options) {
72     options.printHelp(indent, descIndent);
73   });
74 }
75 
76 /// Return the maximum width required when printing the options of this
77 /// entry.
78 size_t PassRegistryEntry::getOptionWidth() const {
79   size_t maxLen = 0;
80   optHandler([&](const PassOptions &options) mutable {
81     maxLen = options.getOptionWidth() + 2;
82   });
83   return maxLen;
84 }
85 
86 //===----------------------------------------------------------------------===//
87 // PassPipelineInfo
88 //===----------------------------------------------------------------------===//
89 
90 void mlir::registerPassPipeline(
91     StringRef arg, StringRef description, const PassRegistryFunction &function,
92     std::function<void(function_ref<void(const PassOptions &)>)> optHandler) {
93   PassPipelineInfo pipelineInfo(arg, description, function,
94                                 std::move(optHandler));
95   bool inserted = passPipelineRegistry->try_emplace(arg, pipelineInfo).second;
96   assert(inserted && "Pass pipeline registered multiple times");
97   (void)inserted;
98 }
99 
100 //===----------------------------------------------------------------------===//
101 // PassInfo
102 //===----------------------------------------------------------------------===//
103 
104 PassInfo::PassInfo(StringRef arg, StringRef description,
105                    const PassAllocatorFunction &allocator)
106     : PassRegistryEntry(
107           arg, description, buildDefaultRegistryFn(allocator),
108           // Use a temporary pass to provide an options instance.
109           [=](function_ref<void(const PassOptions &)> optHandler) {
110             optHandler(allocator()->passOptions);
111           }) {}
112 
113 void mlir::registerPass(const PassAllocatorFunction &function) {
114   std::unique_ptr<Pass> pass = function();
115   StringRef arg = pass->getArgument();
116   if (arg.empty())
117     llvm::report_fatal_error(llvm::Twine("Trying to register '") +
118                              pass->getName() +
119                              "' pass that does not override `getArgument()`");
120   StringRef description = pass->getDescription();
121   PassInfo passInfo(arg, description, function);
122   passRegistry->try_emplace(arg, passInfo);
123 
124   // Verify that the registered pass has the same ID as any registered to this
125   // arg before it.
126   TypeID entryTypeID = pass->getTypeID();
127   auto it = passRegistryTypeIDs->try_emplace(arg, entryTypeID).first;
128   if (it->second != entryTypeID)
129     llvm::report_fatal_error(
130         "pass allocator creates a different pass than previously "
131         "registered for pass " +
132         arg);
133 }
134 
135 /// Returns the pass info for the specified pass argument or null if unknown.
136 const PassInfo *mlir::Pass::lookupPassInfo(StringRef passArg) {
137   auto it = passRegistry->find(passArg);
138   return it == passRegistry->end() ? nullptr : &it->second;
139 }
140 
141 //===----------------------------------------------------------------------===//
142 // PassOptions
143 //===----------------------------------------------------------------------===//
144 
145 /// Out of line virtual function to provide home for the class.
146 void detail::PassOptions::OptionBase::anchor() {}
147 
148 /// Copy the option values from 'other'.
149 void detail::PassOptions::copyOptionValuesFrom(const PassOptions &other) {
150   assert(options.size() == other.options.size());
151   if (options.empty())
152     return;
153   for (auto optionsIt : llvm::zip(options, other.options))
154     std::get<0>(optionsIt)->copyValueFrom(*std::get<1>(optionsIt));
155 }
156 
157 /// Parse in the next argument from the given options string. Returns a tuple
158 /// containing [the key of the option, the value of the option, updated
159 /// `options` string pointing after the parsed option].
160 static std::tuple<StringRef, StringRef, StringRef>
161 parseNextArg(StringRef options) {
162   // Functor used to extract an argument from 'options' and update it to point
163   // after the arg.
164   auto extractArgAndUpdateOptions = [&](size_t argSize) {
165     StringRef str = options.take_front(argSize).trim();
166     options = options.drop_front(argSize).ltrim();
167     return str;
168   };
169   // Try to process the given punctuation, properly escaping any contained
170   // characters.
171   auto tryProcessPunct = [&](size_t &currentPos, char punct) {
172     if (options[currentPos] != punct)
173       return false;
174     size_t nextIt = options.find_first_of(punct, currentPos + 1);
175     if (nextIt != StringRef::npos)
176       currentPos = nextIt;
177     return true;
178   };
179 
180   // Parse the argument name of the option.
181   StringRef argName;
182   for (size_t argEndIt = 0, optionsE = options.size();; ++argEndIt) {
183     // Check for the end of the full option.
184     if (argEndIt == optionsE || options[argEndIt] == ' ') {
185       argName = extractArgAndUpdateOptions(argEndIt);
186       return std::make_tuple(argName, StringRef(), options);
187     }
188 
189     // Check for the end of the name and the start of the value.
190     if (options[argEndIt] == '=') {
191       argName = extractArgAndUpdateOptions(argEndIt);
192       options = options.drop_front();
193       break;
194     }
195   }
196 
197   // Parse the value of the option.
198   for (size_t argEndIt = 0, optionsE = options.size();; ++argEndIt) {
199     // Handle the end of the options string.
200     if (argEndIt == optionsE || options[argEndIt] == ' ') {
201       StringRef value = extractArgAndUpdateOptions(argEndIt);
202       return std::make_tuple(argName, value, options);
203     }
204 
205     // Skip over escaped sequences.
206     char c = options[argEndIt];
207     if (tryProcessPunct(argEndIt, '\'') || tryProcessPunct(argEndIt, '"'))
208       continue;
209     // '{...}' is used to specify options to passes, properly escape it so
210     // that we don't accidentally split any nested options.
211     if (c == '{') {
212       size_t braceCount = 1;
213       for (++argEndIt; argEndIt != optionsE; ++argEndIt) {
214         // Allow nested punctuation.
215         if (tryProcessPunct(argEndIt, '\'') || tryProcessPunct(argEndIt, '"'))
216           continue;
217         if (options[argEndIt] == '{')
218           ++braceCount;
219         else if (options[argEndIt] == '}' && --braceCount == 0)
220           break;
221       }
222       // Account for the increment at the top of the loop.
223       --argEndIt;
224     }
225   }
226   llvm_unreachable("unexpected control flow in pass option parsing");
227 }
228 
229 LogicalResult detail::PassOptions::parseFromString(StringRef options) {
230   // NOTE: `options` is modified in place to always refer to the unprocessed
231   // part of the string.
232   while (!options.empty()) {
233     StringRef key, value;
234     std::tie(key, value, options) = parseNextArg(options);
235     if (key.empty())
236       continue;
237 
238     auto it = OptionsMap.find(key);
239     if (it == OptionsMap.end()) {
240       llvm::errs() << "<Pass-Options-Parser>: no such option " << key << "\n";
241       return failure();
242     }
243     if (llvm::cl::ProvidePositionalOption(it->second, value, 0))
244       return failure();
245   }
246 
247   return success();
248 }
249 
250 /// Print the options held by this struct in a form that can be parsed via
251 /// 'parseFromString'.
252 void detail::PassOptions::print(raw_ostream &os) {
253   // If there are no options, there is nothing left to do.
254   if (OptionsMap.empty())
255     return;
256 
257   // Sort the options to make the ordering deterministic.
258   SmallVector<OptionBase *, 4> orderedOps(options.begin(), options.end());
259   auto compareOptionArgs = [](OptionBase *const *lhs, OptionBase *const *rhs) {
260     return (*lhs)->getArgStr().compare((*rhs)->getArgStr());
261   };
262   llvm::array_pod_sort(orderedOps.begin(), orderedOps.end(), compareOptionArgs);
263 
264   // Interleave the options with ' '.
265   os << '{';
266   llvm::interleave(
267       orderedOps, os, [&](OptionBase *option) { option->print(os); }, " ");
268   os << '}';
269 }
270 
271 /// Print the help string for the options held by this struct. `descIndent` is
272 /// the indent within the stream that the descriptions should be aligned.
273 void detail::PassOptions::printHelp(size_t indent, size_t descIndent) const {
274   // Sort the options to make the ordering deterministic.
275   SmallVector<OptionBase *, 4> orderedOps(options.begin(), options.end());
276   auto compareOptionArgs = [](OptionBase *const *lhs, OptionBase *const *rhs) {
277     return (*lhs)->getArgStr().compare((*rhs)->getArgStr());
278   };
279   llvm::array_pod_sort(orderedOps.begin(), orderedOps.end(), compareOptionArgs);
280   for (OptionBase *option : orderedOps) {
281     // TODO: printOptionInfo assumes a specific indent and will
282     // print options with values with incorrect indentation. We should add
283     // support to llvm::cl::Option for passing in a base indent to use when
284     // printing.
285     llvm::outs().indent(indent);
286     option->getOption()->printOptionInfo(descIndent - indent);
287   }
288 }
289 
290 /// Return the maximum width required when printing the help string.
291 size_t detail::PassOptions::getOptionWidth() const {
292   size_t max = 0;
293   for (auto *option : options)
294     max = std::max(max, option->getOption()->getOptionWidth());
295   return max;
296 }
297 
298 //===----------------------------------------------------------------------===//
299 // TextualPassPipeline Parser
300 //===----------------------------------------------------------------------===//
301 
302 namespace {
303 /// This class represents a textual description of a pass pipeline.
304 class TextualPipeline {
305 public:
306   /// Try to initialize this pipeline with the given pipeline text.
307   /// `errorStream` is the output stream to emit errors to.
308   LogicalResult initialize(StringRef text, raw_ostream &errorStream);
309 
310   /// Add the internal pipeline elements to the provided pass manager.
311   LogicalResult
312   addToPipeline(OpPassManager &pm,
313                 function_ref<LogicalResult(const Twine &)> errorHandler) const;
314 
315 private:
316   /// A functor used to emit errors found during pipeline handling. The first
317   /// parameter corresponds to the raw location within the pipeline string. This
318   /// should always return failure.
319   using ErrorHandlerT = function_ref<LogicalResult(const char *, Twine)>;
320 
321   /// A struct to capture parsed pass pipeline names.
322   ///
323   /// A pipeline is defined as a series of names, each of which may in itself
324   /// recursively contain a nested pipeline. A name is either the name of a pass
325   /// (e.g. "cse") or the name of an operation type (e.g. "builtin.func"). If
326   /// the name is the name of a pass, the InnerPipeline is empty, since passes
327   /// cannot contain inner pipelines.
328   struct PipelineElement {
329     PipelineElement(StringRef name) : name(name) {}
330 
331     StringRef name;
332     StringRef options;
333     const PassRegistryEntry *registryEntry = nullptr;
334     std::vector<PipelineElement> innerPipeline;
335   };
336 
337   /// Parse the given pipeline text into the internal pipeline vector. This
338   /// function only parses the structure of the pipeline, and does not resolve
339   /// its elements.
340   LogicalResult parsePipelineText(StringRef text, ErrorHandlerT errorHandler);
341 
342   /// Resolve the elements of the pipeline, i.e. connect passes and pipelines to
343   /// the corresponding registry entry.
344   LogicalResult
345   resolvePipelineElements(MutableArrayRef<PipelineElement> elements,
346                           ErrorHandlerT errorHandler);
347 
348   /// Resolve a single element of the pipeline.
349   LogicalResult resolvePipelineElement(PipelineElement &element,
350                                        ErrorHandlerT errorHandler);
351 
352   /// Add the given pipeline elements to the provided pass manager.
353   LogicalResult
354   addToPipeline(ArrayRef<PipelineElement> elements, OpPassManager &pm,
355                 function_ref<LogicalResult(const Twine &)> errorHandler) const;
356 
357   std::vector<PipelineElement> pipeline;
358 };
359 
360 } // namespace
361 
362 /// Try to initialize this pipeline with the given pipeline text. An option is
363 /// given to enable accurate error reporting.
364 LogicalResult TextualPipeline::initialize(StringRef text,
365                                           raw_ostream &errorStream) {
366   if (text.empty())
367     return success();
368 
369   // Build a source manager to use for error reporting.
370   llvm::SourceMgr pipelineMgr;
371   pipelineMgr.AddNewSourceBuffer(
372       llvm::MemoryBuffer::getMemBuffer(text, "MLIR Textual PassPipeline Parser",
373                                        /*RequiresNullTerminator=*/false),
374       SMLoc());
375   auto errorHandler = [&](const char *rawLoc, Twine msg) {
376     pipelineMgr.PrintMessage(errorStream, SMLoc::getFromPointer(rawLoc),
377                              llvm::SourceMgr::DK_Error, msg);
378     return failure();
379   };
380 
381   // Parse the provided pipeline string.
382   if (failed(parsePipelineText(text, errorHandler)))
383     return failure();
384   return resolvePipelineElements(pipeline, errorHandler);
385 }
386 
387 /// Add the internal pipeline elements to the provided pass manager.
388 LogicalResult TextualPipeline::addToPipeline(
389     OpPassManager &pm,
390     function_ref<LogicalResult(const Twine &)> errorHandler) const {
391   return addToPipeline(pipeline, pm, errorHandler);
392 }
393 
394 /// Parse the given pipeline text into the internal pipeline vector. This
395 /// function only parses the structure of the pipeline, and does not resolve
396 /// its elements.
397 LogicalResult TextualPipeline::parsePipelineText(StringRef text,
398                                                  ErrorHandlerT errorHandler) {
399   SmallVector<std::vector<PipelineElement> *, 4> pipelineStack = {&pipeline};
400   for (;;) {
401     std::vector<PipelineElement> &pipeline = *pipelineStack.back();
402     size_t pos = text.find_first_of(",(){");
403     pipeline.emplace_back(/*name=*/text.substr(0, pos).trim());
404 
405     // If we have a single terminating name, we're done.
406     if (pos == StringRef::npos)
407       break;
408 
409     text = text.substr(pos);
410     char sep = text[0];
411 
412     // Handle pulling ... from 'pass{...}' out as PipelineElement.options.
413     if (sep == '{') {
414       text = text.substr(1);
415 
416       // Skip over everything until the closing '}' and store as options.
417       size_t close = StringRef::npos;
418       for (unsigned i = 0, e = text.size(), braceCount = 1; i < e; ++i) {
419         if (text[i] == '{') {
420           ++braceCount;
421           continue;
422         }
423         if (text[i] == '}' && --braceCount == 0) {
424           close = i;
425           break;
426         }
427       }
428 
429       // Check to see if a closing options brace was found.
430       if (close == StringRef::npos) {
431         return errorHandler(
432             /*rawLoc=*/text.data() - 1,
433             "missing closing '}' while processing pass options");
434       }
435       pipeline.back().options = text.substr(0, close);
436       text = text.substr(close + 1);
437 
438       // Skip checking for '(' because nested pipelines cannot have options.
439     } else if (sep == '(') {
440       text = text.substr(1);
441 
442       // Push the inner pipeline onto the stack to continue processing.
443       pipelineStack.push_back(&pipeline.back().innerPipeline);
444       continue;
445     }
446 
447     // When handling the close parenthesis, we greedily consume them to avoid
448     // empty strings in the pipeline.
449     while (text.consume_front(")")) {
450       // If we try to pop the outer pipeline we have unbalanced parentheses.
451       if (pipelineStack.size() == 1)
452         return errorHandler(/*rawLoc=*/text.data() - 1,
453                             "encountered extra closing ')' creating unbalanced "
454                             "parentheses while parsing pipeline");
455 
456       pipelineStack.pop_back();
457     }
458 
459     // Check if we've finished parsing.
460     if (text.empty())
461       break;
462 
463     // Otherwise, the end of an inner pipeline always has to be followed by
464     // a comma, and then we can continue.
465     if (!text.consume_front(","))
466       return errorHandler(text.data(), "expected ',' after parsing pipeline");
467   }
468 
469   // Check for unbalanced parentheses.
470   if (pipelineStack.size() > 1)
471     return errorHandler(
472         text.data(),
473         "encountered unbalanced parentheses while parsing pipeline");
474 
475   assert(pipelineStack.back() == &pipeline &&
476          "wrong pipeline at the bottom of the stack");
477   return success();
478 }
479 
480 /// Resolve the elements of the pipeline, i.e. connect passes and pipelines to
481 /// the corresponding registry entry.
482 LogicalResult TextualPipeline::resolvePipelineElements(
483     MutableArrayRef<PipelineElement> elements, ErrorHandlerT errorHandler) {
484   for (auto &elt : elements)
485     if (failed(resolvePipelineElement(elt, errorHandler)))
486       return failure();
487   return success();
488 }
489 
490 /// Resolve a single element of the pipeline.
491 LogicalResult
492 TextualPipeline::resolvePipelineElement(PipelineElement &element,
493                                         ErrorHandlerT errorHandler) {
494   // If the inner pipeline of this element is not empty, this is an operation
495   // pipeline.
496   if (!element.innerPipeline.empty())
497     return resolvePipelineElements(element.innerPipeline, errorHandler);
498   // Otherwise, this must be a pass or pass pipeline.
499   // Check to see if a pipeline was registered with this name.
500   auto pipelineRegistryIt = passPipelineRegistry->find(element.name);
501   if (pipelineRegistryIt != passPipelineRegistry->end()) {
502     element.registryEntry = &pipelineRegistryIt->second;
503     return success();
504   }
505 
506   // If not, then this must be a specific pass name.
507   if ((element.registryEntry = Pass::lookupPassInfo(element.name)))
508     return success();
509 
510   // Emit an error for the unknown pass.
511   auto *rawLoc = element.name.data();
512   return errorHandler(rawLoc, "'" + element.name +
513                                   "' does not refer to a "
514                                   "registered pass or pass pipeline");
515 }
516 
517 /// Add the given pipeline elements to the provided pass manager.
518 LogicalResult TextualPipeline::addToPipeline(
519     ArrayRef<PipelineElement> elements, OpPassManager &pm,
520     function_ref<LogicalResult(const Twine &)> errorHandler) const {
521   for (auto &elt : elements) {
522     if (elt.registryEntry) {
523       if (failed(elt.registryEntry->addToPipeline(pm, elt.options,
524                                                   errorHandler))) {
525         return errorHandler("failed to add `" + elt.name + "` with options `" +
526                             elt.options + "`");
527       }
528     } else if (failed(addToPipeline(elt.innerPipeline, pm.nest(elt.name),
529                                     errorHandler))) {
530       return errorHandler("failed to add `" + elt.name + "` with options `" +
531                           elt.options + "` to inner pipeline");
532     }
533   }
534   return success();
535 }
536 
537 LogicalResult mlir::parsePassPipeline(StringRef pipeline, OpPassManager &pm,
538                                       raw_ostream &errorStream) {
539   TextualPipeline pipelineParser;
540   if (failed(pipelineParser.initialize(pipeline, errorStream)))
541     return failure();
542   auto errorHandler = [&](Twine msg) {
543     errorStream << msg << "\n";
544     return failure();
545   };
546   if (failed(pipelineParser.addToPipeline(pm, errorHandler)))
547     return failure();
548   return success();
549 }
550 
551 FailureOr<OpPassManager> mlir::parsePassPipeline(StringRef pipeline,
552                                                  raw_ostream &errorStream) {
553   // Pipelines are expected to be of the form `<op-name>(<pipeline>)`.
554   size_t pipelineStart = pipeline.find_first_of('(');
555   if (pipelineStart == 0 || pipelineStart == StringRef::npos ||
556       !pipeline.consume_back(")")) {
557     errorStream << "expected pass pipeline to be wrapped with the anchor "
558                    "operation type, e.g. `builtin.module(...)";
559     return failure();
560   }
561 
562   StringRef opName = pipeline.take_front(pipelineStart);
563   OpPassManager pm(opName);
564   if (failed(parsePassPipeline(pipeline.drop_front(1 + pipelineStart), pm)))
565     return failure();
566   return pm;
567 }
568 
569 //===----------------------------------------------------------------------===//
570 // PassNameParser
571 //===----------------------------------------------------------------------===//
572 
573 namespace {
574 /// This struct represents the possible data entries in a parsed pass pipeline
575 /// list.
576 struct PassArgData {
577   PassArgData() = default;
578   PassArgData(const PassRegistryEntry *registryEntry)
579       : registryEntry(registryEntry) {}
580 
581   /// This field is used when the parsed option corresponds to a registered pass
582   /// or pass pipeline.
583   const PassRegistryEntry *registryEntry{nullptr};
584 
585   /// This field is set when instance specific pass options have been provided
586   /// on the command line.
587   StringRef options;
588 
589   /// This field is used when the parsed option corresponds to an explicit
590   /// pipeline.
591   TextualPipeline pipeline;
592 };
593 } // namespace
594 
595 namespace llvm {
596 namespace cl {
597 /// Define a valid OptionValue for the command line pass argument.
598 template <>
599 struct OptionValue<PassArgData> final
600     : OptionValueBase<PassArgData, /*isClass=*/true> {
601   OptionValue(const PassArgData &value) { this->setValue(value); }
602   OptionValue() = default;
603   void anchor() override {}
604 
605   bool hasValue() const { return true; }
606   const PassArgData &getValue() const { return value; }
607   void setValue(const PassArgData &value) { this->value = value; }
608 
609   PassArgData value;
610 };
611 } // namespace cl
612 } // namespace llvm
613 
614 namespace {
615 
616 /// The name for the command line option used for parsing the textual pass
617 /// pipeline.
618 static constexpr StringLiteral passPipelineArg = "pass-pipeline";
619 
620 /// Adds command line option for each registered pass or pass pipeline, as well
621 /// as textual pass pipelines.
622 struct PassNameParser : public llvm::cl::parser<PassArgData> {
623   PassNameParser(llvm::cl::Option &opt) : llvm::cl::parser<PassArgData>(opt) {}
624 
625   void initialize();
626   void printOptionInfo(const llvm::cl::Option &opt,
627                        size_t globalWidth) const override;
628   size_t getOptionWidth(const llvm::cl::Option &opt) const override;
629   bool parse(llvm::cl::Option &opt, StringRef argName, StringRef arg,
630              PassArgData &value);
631 
632   /// If true, this parser only parses entries that correspond to a concrete
633   /// pass registry entry, and does not add a `pass-pipeline` argument, does not
634   /// include the options for pass entries, and does not include pass pipelines
635   /// entries.
636   bool passNamesOnly = false;
637 };
638 } // namespace
639 
640 void PassNameParser::initialize() {
641   llvm::cl::parser<PassArgData>::initialize();
642 
643   /// Add an entry for the textual pass pipeline option.
644   if (!passNamesOnly) {
645     addLiteralOption(passPipelineArg, PassArgData(),
646                      "A textual description of a pass pipeline to run");
647   }
648 
649   /// Add the pass entries.
650   for (const auto &kv : *passRegistry) {
651     addLiteralOption(kv.second.getPassArgument(), &kv.second,
652                      kv.second.getPassDescription());
653   }
654   /// Add the pass pipeline entries.
655   if (!passNamesOnly) {
656     for (const auto &kv : *passPipelineRegistry) {
657       addLiteralOption(kv.second.getPassArgument(), &kv.second,
658                        kv.second.getPassDescription());
659     }
660   }
661 }
662 
663 void PassNameParser::printOptionInfo(const llvm::cl::Option &opt,
664                                      size_t globalWidth) const {
665   // If this parser is just parsing pass names, print a simplified option
666   // string.
667   if (passNamesOnly) {
668     llvm::outs() << "  --" << opt.ArgStr << "=<pass-arg>";
669     opt.printHelpStr(opt.HelpStr, globalWidth, opt.ArgStr.size() + 18);
670     return;
671   }
672 
673   // Print the information for the top-level option.
674   if (opt.hasArgStr()) {
675     llvm::outs() << "  --" << opt.ArgStr;
676     opt.printHelpStr(opt.HelpStr, globalWidth, opt.ArgStr.size() + 7);
677   } else {
678     llvm::outs() << "  " << opt.HelpStr << '\n';
679   }
680 
681   // Print the top-level pipeline argument.
682   printOptionHelp(passPipelineArg,
683                   "A textual description of a pass pipeline to run",
684                   /*indent=*/4, globalWidth, /*isTopLevel=*/!opt.hasArgStr());
685 
686   // Functor used to print the ordered entries of a registration map.
687   auto printOrderedEntries = [&](StringRef header, auto &map) {
688     llvm::SmallVector<PassRegistryEntry *, 32> orderedEntries;
689     for (auto &kv : map)
690       orderedEntries.push_back(&kv.second);
691     llvm::array_pod_sort(
692         orderedEntries.begin(), orderedEntries.end(),
693         [](PassRegistryEntry *const *lhs, PassRegistryEntry *const *rhs) {
694           return (*lhs)->getPassArgument().compare((*rhs)->getPassArgument());
695         });
696 
697     llvm::outs().indent(4) << header << ":\n";
698     for (PassRegistryEntry *entry : orderedEntries)
699       entry->printHelpStr(/*indent=*/6, globalWidth);
700   };
701 
702   // Print the available passes.
703   printOrderedEntries("Passes", *passRegistry);
704 
705   // Print the available pass pipelines.
706   if (!passPipelineRegistry->empty())
707     printOrderedEntries("Pass Pipelines", *passPipelineRegistry);
708 }
709 
710 size_t PassNameParser::getOptionWidth(const llvm::cl::Option &opt) const {
711   size_t maxWidth = llvm::cl::parser<PassArgData>::getOptionWidth(opt) + 2;
712 
713   // Check for any wider pass or pipeline options.
714   for (auto &entry : *passRegistry)
715     maxWidth = std::max(maxWidth, entry.second.getOptionWidth() + 4);
716   for (auto &entry : *passPipelineRegistry)
717     maxWidth = std::max(maxWidth, entry.second.getOptionWidth() + 4);
718   return maxWidth;
719 }
720 
721 bool PassNameParser::parse(llvm::cl::Option &opt, StringRef argName,
722                            StringRef arg, PassArgData &value) {
723   // Handle the pipeline option explicitly.
724   if (argName == passPipelineArg)
725     return failed(value.pipeline.initialize(arg, llvm::errs()));
726 
727   // Otherwise, default to the base for handling.
728   if (llvm::cl::parser<PassArgData>::parse(opt, argName, arg, value))
729     return true;
730   value.options = arg;
731   return false;
732 }
733 
734 //===----------------------------------------------------------------------===//
735 // PassPipelineCLParser
736 //===----------------------------------------------------------------------===//
737 
738 namespace mlir {
739 namespace detail {
740 struct PassPipelineCLParserImpl {
741   PassPipelineCLParserImpl(StringRef arg, StringRef description,
742                            bool passNamesOnly)
743       : passList(arg, llvm::cl::desc(description)) {
744     passList.getParser().passNamesOnly = passNamesOnly;
745     passList.setValueExpectedFlag(llvm::cl::ValueExpected::ValueOptional);
746   }
747 
748   /// Returns true if the given pass registry entry was registered at the
749   /// top-level of the parser, i.e. not within an explicit textual pipeline.
750   bool contains(const PassRegistryEntry *entry) const {
751     return llvm::any_of(passList, [&](const PassArgData &data) {
752       return data.registryEntry == entry;
753     });
754   }
755 
756   /// The set of passes and pass pipelines to run.
757   llvm::cl::list<PassArgData, bool, PassNameParser> passList;
758 };
759 } // namespace detail
760 } // namespace mlir
761 
762 /// Construct a pass pipeline parser with the given command line description.
763 PassPipelineCLParser::PassPipelineCLParser(StringRef arg, StringRef description)
764     : impl(std::make_unique<detail::PassPipelineCLParserImpl>(
765           arg, description, /*passNamesOnly=*/false)) {}
766 PassPipelineCLParser::~PassPipelineCLParser() = default;
767 
768 /// Returns true if this parser contains any valid options to add.
769 bool PassPipelineCLParser::hasAnyOccurrences() const {
770   return impl->passList.getNumOccurrences() != 0;
771 }
772 
773 /// Returns true if the given pass registry entry was registered at the
774 /// top-level of the parser, i.e. not within an explicit textual pipeline.
775 bool PassPipelineCLParser::contains(const PassRegistryEntry *entry) const {
776   return impl->contains(entry);
777 }
778 
779 /// Adds the passes defined by this parser entry to the given pass manager.
780 LogicalResult PassPipelineCLParser::addToPipeline(
781     OpPassManager &pm,
782     function_ref<LogicalResult(const Twine &)> errorHandler) const {
783   for (auto &passIt : impl->passList) {
784     if (passIt.registryEntry) {
785       if (failed(passIt.registryEntry->addToPipeline(pm, passIt.options,
786                                                      errorHandler)))
787         return failure();
788     } else {
789       OpPassManager::Nesting nesting = pm.getNesting();
790       pm.setNesting(OpPassManager::Nesting::Explicit);
791       LogicalResult status = passIt.pipeline.addToPipeline(pm, errorHandler);
792       pm.setNesting(nesting);
793       if (failed(status))
794         return failure();
795     }
796   }
797   return success();
798 }
799 
800 //===----------------------------------------------------------------------===//
801 // PassNameCLParser
802 
803 /// Construct a pass pipeline parser with the given command line description.
804 PassNameCLParser::PassNameCLParser(StringRef arg, StringRef description)
805     : impl(std::make_unique<detail::PassPipelineCLParserImpl>(
806           arg, description, /*passNamesOnly=*/true)) {
807   impl->passList.setMiscFlag(llvm::cl::CommaSeparated);
808 }
809 PassNameCLParser::~PassNameCLParser() = default;
810 
811 /// Returns true if this parser contains any valid options to add.
812 bool PassNameCLParser::hasAnyOccurrences() const {
813   return impl->passList.getNumOccurrences() != 0;
814 }
815 
816 /// Returns true if the given pass registry entry was registered at the
817 /// top-level of the parser, i.e. not within an explicit textual pipeline.
818 bool PassNameCLParser::contains(const PassRegistryEntry *entry) const {
819   return impl->contains(entry);
820 }
821