MLIR笔记（1）

笔记（1）"/>

MLIR笔记（1）

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）以及类型。

1. 命令行选项的解析
   1. 原理

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()继续往下完成自己的使命。

1. 1. 方言间转换的例子

另一个生猛的例子是用于方言间转换的命令行解析的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 &registry) {

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提供处理方法。到这里与方言转换相关的处理就完成。后续选项的解析与处理就是公共的，在命令行上发现相关选项后，将调用上面标色的可执行体。