1. 简介
MLIR是Multi-layer IR的缩写,它是基于LLVM IR发展的一个中间语言形式,是Clang/LLVM的发明者Chris Lattner在加盟谷歌后又一个重要的发明。MLIR是较之LLVM IR更为灵活的深度学习编译器架构。
其他编译器,像LLVM(参考Kaleidoscope tutorial),提供一组固定的预定义的类型以及(通常低级/类RISC)指令。在发布LLVM IR之前,由特定语言的前端来执行所有语言特定的类型检查、分析或转换。例如,Clang的AST不仅用来执行静态分析,还用于转换,比如使用AST克隆与重写的C++模板具现。最后,具有比C/C++更高级结构的语言可能要求从它们的AST经过重要的(non-trivial)降级来产生LLVM IR。
因此,多个前端最终重新实现基础架构重要的部分,以支持这些分析与转换。而MLIR通过可扩展性的设计来应对这些情况。因此,只有少数几个预定义指令(MLIR术语里的操作,operation)以及类型。
LLVM本身有一整套相当复杂的命令行选项解析机制,通过llvm::cl::AddLiteralOption(),可以向LLVM的命令行解析器注册新的命令行选项,让它为我们提供一整套选项处理功能。这正是MLIR命令行选项着手的地方!注意,它是与LLVM本身的命令行选项解析分开的,LLVM虽然也是调用这个方法,但LLVM组织的方式是不一样的。
为了解析与MLIR相关的ODS定义,首先MLIR提供了一个GenRegistration定义,这个结构体提供的唯一方法就是构造函数:
31 mlir::GenRegistration::GenRegistration(StringRef arg, StringRef description,
32 GenFunction function) {
33 generatorRegistry->emplace_back(arg, description, function);
34 }
这里generatorRegistry是一个静态变量(由ManagedStatic类提供封装):
29 static llvm::ManagedStatic<std::vector<GenInfo>> generatorRegistry;
其中,GenInfo用于封装各种代码生成器,它只有3个域:arg(选项名,字符串类型),description(选项描述,字符串类型),generator(代码生成器,一个可执行的对象)。
因此在OpDefinitionsGen.cpp里,我们可以看到这样的GenRegistration对象声明:
2255 static mlir::GenRegistration
2256 genOpDecls("gen-op-decls", "Generate op declarations",
2257 [](const RecordKeeper &records, raw_ostream &os) {
2258 return emitOpDecls(records, os);
2259 });
2260
2261 static mlir::GenRegistration genOpDefs("gen-op-defs", "Generate op definitions",
2262 [](const RecordKeeper &records,
2263 raw_ostream &os) {
2264 return emitOpDefs(records, os);
2265 });
这样,程序在初始化时会向generatorRegistry添加这些GenInfo对象。那么generatorRegistry怎么样被调动起来的呢?我们看一下mlri-tblgen.cpp,这是TableGen语言代码生成器的源代码所在:
75 int main(int argc, char **argv) {
76 llvm::InitLLVM y(argc, argv);
77 llvm::cl::opt<const mlir::GenInfo *, false, mlir::GenNameParser> generator(
78 "", llvm::cl::desc("Generator to run"));
79 cl::ParseCommandLineOptions(argc, argv);
80 ::generator = generator.getValue();
81
82 return TableGenMain(argv[0], &MlirTableGenMain);
83 }
LLVM有一套极其复杂的命令行选项处理机制,这里我们只能简要说一下。
首先,76行的llvm::InitLLVM类型的局部变量y是初始化LLVM的必要模块,与命令行解析的关系不大。77行的generator就是命令行选项解析机制的一部分,它的类型是cl::opt,我们看一下这个类型定义开头的几行(CommandLine.h):
1404 template <class DataType, bool ExternalStorage = false,
1405 class ParserClass = parser<DataType>>
1406 class opt : public Option,
1407 public opt_storage<DataType, ExternalStorage,
1408 std::is_class<DataType>::value> {
1409 ParserClass Parser;
它的构造函数是这样定义的:
1482 template <class... Mods>
1483 explicit opt(const Mods &... Ms)
1484 : Option(Optional, NotHidden), Parser(*this) {
1485 apply(this, Ms...);
1486 done();
1487 }
基类Option描述了选项属性,而opt_storage则保存了命令行上出现选项的具体信息(它是一个模板类,这里的特化以所服务的信息类型为基类,在这个上下文里就是GenInfo的派生类),这里因为在声明opt_storage基类时把ExternalStorage指定为false,因此generator的opt_storage部分将用于保存命令行上出现的选项所对应的GenInfo实例。
同时,构造函数里指定opt使用的Parser是mlir::GenNameParser:
23 struct GenNameParser : public llvm::cl::parser<const GenInfo *>
它的构造函数是这样的:
36 GenNameParser::GenNameParser(llvm::cl::Option &opt)
37 : llvm::cl::parser<const GenInfo *>(opt) {
38 for (const auto &kv : *generatorRegistry) {
39 addLiteralOption(kv.getGenArgument(), &kv, kv.getGenDescription());
40 }
41 }
在37用作基类的cl::parser是这样定义的:
795 template <class DataType> class parser : public generic_parser_base
在37行调用了它的构造函数:
807 parser(Option &O) : generic_parser_base(O) {}
其基类generic_parser_base的构造函数是这样的:
710 generic_parser_base(Option &O) : Owner(O) {}
Owner是generic_parser_base里Option&类型的成员,在这个上下文里,它绑定到了generator变量的Option部分。这样GenNameParser和generator就关联起来了。在GenNameParser构造函数39行调用的addLiteralOption()(这是类parser的方法)处理注册在generatorRegistry里的GenInfo对象(作为参数V):
842 template <class DT>
843 void addLiteralOption(StringRef Name, const DT &V, StringRef HelpStr) {
844 assert(findOption(Name) == Values.size() && "Option already exists!");
845 OptionInfo X(Name, static_cast<DataType>(V), HelpStr);
846 Values.push_back(X);
847 AddLiteralOption(Owner, Name);
848 }
846行的Values是类parser里的一个SmallVector<OptionInfo, 8>类型的容器,类型OptionInfo也是GenInfo的一个封装类。847行的AddLiteralOption()是cl名字空间里的全局函数:
436 void cl::AddLiteralOption(Option &O, StringRef Name) {
437 GlobalParser->addLiteralOption(O, Name);
438 }
这里的GlobalParser也是一个“静态”变量:
434 static ManagedStatic<CommandLineParser> GlobalParser
它的addLiteralOption()方法的定义是:
198 void addLiteralOption(Option &Opt, StringRef Name) {
199 if (Opt.Subs.empty())
200 addLiteralOption(Opt, &*TopLevelSubCommand, Name);
201 else {
202 for (auto SC : Opt.Subs)
203 addLiteralOption(Opt, SC, Name);
204 }
205 }
Option支持选项组的概念,在使用选项组时它的Subs容器不为空。我们这里不使用选项组,因此generator的Option部分在200行通过一个重载的方法添加到TopLevelSubCommand对象(代表第一级选项)的OptionsMap容器中:
178 void addLiteralOption(Option &Opt, SubCommand *SC, StringRef Name) {
179 if (Opt.hasArgStr())
180 return;
181 if (!SC->OptionsMap.insert(std::make_pair(Name, &Opt)).second) {
182 errs() << ProgramName << ": CommandLine Error: Option '" << Name
183 << "' registered more than once!\n";
184 report_fatal_error("inconsistency in registered CommandLine options");
185 }
186
187 // If we're adding this to all sub-commands, add it to the ones that have
188 // already been registered.
189 if (SC == &*AllSubCommands) {
190 for (auto *Sub : RegisteredSubCommands) {
191 if (SC == Sub)
192 continue;
193 addLiteralOption(Opt, Sub, Name);
194 }
195 }
196 }
在181行,选项名与generator在容器OptionsMap(类型StringMap<Option *>)里关联起来(因此,所有这些MLIR选项都由generator这个变量提供处理支持)。
回到opt的构造函数。接下来,对调用参数调用apply(),这是一个变长参数模板函数,它会根据每个参数分别调用applicator::opt构造函数:
1291 template <class Opt, class Mod, class... Mods>
1292 void apply(Opt *O, const Mod &M, const Mods &... Ms) {
1293 applicator<Mod>::opt(M, *O);
1294 apply(O, Ms...);
1295 }
1296
1297 template <class Opt, class Mod> void apply(Opt *O, const Mod &M) {
1298 applicator<Mod>::opt(M, *O);
1299 }
这里,我们给出的字符串是""(main()的78行),因此1298行调用的下面这个特化版本:
1258 template <> struct applicator<StringRef > {
1259 template <class Opt> static void opt(StringRef Str, Opt &O) {
1260 O.setArgStr(Str);
1261 }
1262 };
因此,1260行的setArgStr()会把Option的ArgStr设置为空字符串,因为这个Option表示的是命令本身,在generic_parser_base::printOptionInfo()里会使用ArgStr相应显示帮助信息。第二个参数是这个命令的描述,它调用这个applicator:
1243 template <class Mod> struct applicator {
1244 template <class Opt> static void opt(const Mod &M, Opt &O) { M.apply(O); }
1245 };
它调用llvm::cl::desc的apply()设置帮助字符串(help string):
403 void apply(Option &O) const { O.setDescription(Desc); }
到这里,命令行选项解析准备工作就完成了。注意,这一切都在main()被调用前完成的。现在,万事俱备,只欠东风,直到main()登场。
Main()的调用参数就包含了命令行选项,在main()的79行调用cl::ParseCommandLineOptions()来解析命令行参数:
1290 bool cl::ParseCommandLineOptions(int argc, const char *const *argv,
1291 StringRef Overview, raw_ostream *Errs,
1292 const char *EnvVar,
1293 bool LongOptionsUseDoubleDash) {
…
1311 // Parse all options.
1312 return GlobalParser->ParseCommandLineOptions(NewArgc, &NewArgv[0], Overview,
1313 Errs, LongOptionsUseDoubleDash);
1314 }
跳过对环境变量的处理,处理的主体是1312行GlobalParser的ParseCommandLineOptions()。这个函数比较大且复杂,我们不细看代码。大致上,这个函数在循环体中依次比对命令行选项与OptionsMap里保存的选项,如果发现匹配就调用generator的addOccurrence()方法。这个方法进而调用handleOccurrence():
1411 bool handleOccurrence(unsigned pos, StringRef ArgName,
1412 StringRef Arg) override {
1413 typename ParserClass::parser_data_type Val =
1414 typename ParserClass::parser_data_type();
1415 if (Parser.parse(*this, ArgName, Arg, Val))
1416 return true; // Parse error!
1417 this->setValue(Val);
1418 this->setPosition(pos);
1419 Callback(Val); // 这里我们没有注册回调,因此是空函数
1420 return false;
1421 }
1415行调用GenNameParser的parse()(实际上就是基类parser的方法):
824 bool parse(Option &O, StringRef ArgName, StringRef Arg, DataType &V) {
825 StringRef ArgVal;
826 if (Owner.hasArgStr())
827 ArgVal = Arg;
828 else
829 ArgVal = ArgName;
830
831 for (size_t i = 0, e = Values.size(); i != e; ++i)
832 if (Values[i].Name == ArgVal) {
833 V = Values[i].V.getValue();
834 return false;
835 }
836
837 return O.error("Cannot find option named '" + ArgVal + "'!");
838 }
前面addLiteralOption()向Values容器添加了与选项相关的OptionInfo实例(实例封装了对应的GenInfo对象),这里在这个容器里查找名字相匹配的OptionInfo实例,并获取对应的GenInfo对象。在1417行这个对象被parser的setValue()保存在指定的成员里,setValue()的定义如下:
1355 template <class DataType> // 有几个特化版本,当前上下文是这个
1356 class opt_storage<DataType, false, true> : public DataType {
1357 public:
1358 OptionValue<DataType> Default;
1359
1360 template <class T> void setValue(const T &V, bool initial = false) {
1361 DataType::operator=(V);
1362 if (initial)
1363 Default = V;
1364 }
1365
1366 DataType &getValue() { return *this; }
1367 const DataType &getValue() const { return *this; }
1368
1369 const OptionValue<DataType> &getDefault() const { return Default; }
1370 };
在main()的80行通过getValue()获取这个对象,保存在全局变量generator里。随后在MlirTableGenMain()里调用它的Invoke()方法:
67 static bool MlirTableGenMain(raw_ostream &os, RecordKeeper &records) {
68 if (!generator) {
69 os << records;
70 return false;
71 }
72 return generator->invoke(records, os);
73 }
接着Invoke()调用构造GenRegistration对象时传入的可执行体,比如下面标绿的部分。
2255 static mlir::GenRegistration
2256 genOpDecls("gen-op-decls", "Generate op declarations",
2257 [](const RecordKeeper &records, raw_ostream &os) {
2258 return emitOpDecls(records, os);
2259 });
上述过程对所有的命令行选项依次进行,待这一切完成后,main()继续往下完成自己的使命。
-
- 方言间转换的例子
另一个生猛的例子是用于方言间转换的命令行解析的TranslateFromMLIRRegistration,它是MLIR方言转换框架的一部分。为了提供更大灵活性(这些是使用命令行选项进行更复杂处理所需的),MLIR提供了以下的命令行选项注册框架。TranslateFromMLIRRegistration是为MLIR到LLVM之间的方言转换服务的,它只有一个构造函数。类似的转换都需要提供自己的注册方法。
95 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
96 StringRef name, const TranslateFromMLIRFunction &function,
97 std::function<void(DialectRegistry &)> dialectRegistration) {
98 registerTranslation(name, [function, dialectRegistration](
99 llvm::SourceMgr &sourceMgr, raw_ostream &output,
100 MLIRContext *context) {
101 DialectRegistry registry;
102 dialectRegistration(registry);
103 context->appendDialectRegistry(registry);
104 auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
105 if (!module || failed(verify(*module)))
106 return failure();
107 return function(module.get(), output);
108 });
109 }
注意上面不同颜色标注的代码片段,它们对应下面代码中标注了同样颜色的片段。
98行的registerTranslation()完成类似的注册:
39 static void registerTranslation(StringRef name,
40 const TranslateFunction &function) {
41 auto &translationRegistry = getTranslationRegistry();
42 if (translationRegistry.find(name) != translationRegistry.end())
43 llvm::report_fatal_error(
44 "Attempting to overwrite an existing <file-to-file> function");
45 assert(function &&
46 "Attempting to register an empty translate <file-to-file> function");
47 translationRegistry[name] = function;
48 }
41行的getTranslationRegistry()封装了一个静态变量:
33 static llvm::StringMap<TranslateFunction> &getTranslationRegistry() {
34 static llvm::StringMap<TranslateFunction> translationRegistry;
35 return translationRegistry;
36 }
上面的构造函数完成所谓选项名与处理方法的注册。具体的,在convertToLLVMIR.cpp里进行了这样的声明,令人赞赏的是,这几行代码就完成了MLIR到LLVM IR的转换(当然,里面有复杂的处理与调用关系,但至少表面上看起来简单、干净):
22 namespace mlir {
23 void registerToLLVMIRTranslation() {
24 TranslateFromMLIRRegistration registration(
25 "mlir-to-llvmir",
26 [](ModuleOp module, raw_ostream &output) {
27 llvm::LLVMContext llvmContext;
28 auto llvmModule = translateModuleToLLVMIR(module, llvmContext);
29 if (!llvmModule)
30 return failure();
31
32 llvmModule->print(output, nullptr);
33 return success();
34 },
35 [](DialectRegistry ®istry) {
36 registerAllToLLVMIRTranslations(registry);
37 });
38 }
39 } // namespace mlir
注意25行,这个转换是由mlir-translate工具执行的,25行就是给到这个工具的命令行选项,即命令“mlir-translate -mlir-to-llvmir”将完成mlir到llvm IR的转换。与之配合,在mlirTranslateMain()里需要这个代码片段:
157 // Add flags for all the registered translations.
158 llvm::cl::opt<const TranslateFunction *, false, TranslationParser>
159 translationRequested("", llvm::cl::desc("Translation to perform"),
160 llvm::cl::Required);
161 registerAsmPrinterCLOptions();
162 registerMLIRContextCLOptions();
163 llvm::cl::ParseCommandLineOptions(argc, argv, toolName);
在158行的llvm::cl::opt类型对象translationRequested包含了一个TranslationParser类型的成员。显然TranslationParser也必须是llvm::cl::parser的派生类(这样可以利用它的parse()以及相关的方法),因此TranslationParser只需要实现自己的构造函数与printOptionInfo()方法:
115 TranslationParser::TranslationParser(llvm::cl::Option &opt)
116 : llvm::cl::parser<const TranslateFunction *>(opt) {
117 for (const auto &kv : getTranslationRegistry())
118 addLiteralOption(kv.first(), &kv.second, kv.first());
119 }
同样需要在118行通过addLiteralOption()向GlobalParser(CommandLineParser对象)注册这些选项,告诉它,这些选项由TranslationParser提供处理方法。到这里与方言转换相关的处理就完成。后续选项的解析与处理就是公共的,在命令行上发现相关选项后,将调用上面标色的可执行体。