MLIR笔记（6）

5. 方言与操作

5.1. 方言的概念

在MLIR里，通过Dialect类来抽象方言。具体的每种方言都需要从这个基类派生一个类型，并实现重载自己所需的虚函数。 

MLIR文档里这样描述方言（ MLIR Language Reference - MLIR）：

方言是这样的机制：它融入并扩展MLIR生态系统。它们允许定义新的操作，以及属性与类型。向每个方言给出唯一的名字空间作为定义的每个属性/操作/类型的前缀。例如，Affine方言定义了名字空间affine。

MLIR允许多个方言共存于一个模块中，即使是在主干之外的那些。方言由特定的遍生成与消费。对不同的方言之间以及方言内部的转换，MLIR提供了一个框架。MLIR支持的几个方言：

• Affine dialect
• GPU dialect
• LLVM dialect
• SPIR-V dialect
• Standard dialect
• Vector dialect

 在教程中，还给出了Toy方言的例子。

5.2. 操作的概念

MLIR引入了一个称为操作（operation）的统一概念来描述许多不同的抽象与计算层次。在MLIR系统中，从指令到函数再到模块，一切都塑造为Op。 MLIR没有固定的Op集合，因此允许并鼓励用户自定义扩展Op。编译器遍会保守地对待未知Op，并且MLIR支持通过特征（traits）、特权操作hook和优化接口等方式向遍描述Op语义。MLIR里的操作是完全可扩展的（没有固定的操作列表），并具有应用特定的语义。例如，MLIR支持目标无关操作、仿射（affine）操作，以及目标特定机器操作。

操作的内部表示是简单的：操作由一个唯一字符串标识（如dim、tf.Conv2d、x86.repmovsb、ppc.eieio等），可以返回0或多个结果，接受0或多个操作数，有一个属性字典，有0或多个后继者，以及0或多个封闭的区域。通用打印形式包括所有这些元素，加上一个函数类型来表示结果与操作数的类型。

例子：

// An operation that produces two results.
// The results of %result can be accessed via the <name> `#` <opNo> syntax.
%result:2 = "foo_div"() : () -> (f32, i32)

// Pretty form that defines a unique name for each result.
%foo, %bar = "foo_div"() : () -> (f32, i32)

// Invoke a TensorFlow function called tf.scramble with two inputs
// and an attribute "fruit".
%2 = "tf.scramble"(%result#0, %bar) {fruit = "banana"} : (f32, i32) -> f32
5.3. 方言的管理

显然，管理方言最恰当的地方就是MLIRContext，不过MLIRContext只是上下文的接口，真正的实现是MLIRContextImpl。在MLIRContextImpl中有这样一些容器：

DenseMap<StringRef, std::unique_ptr<Dialect>> loadedDialects

DialectRegistry dialectsRegistry

llvm::StringMap<AbstractOperation> registeredOperations

llvm::StringMap<PointerUnion<Dialect *, MLIRContext *>, llvm::BumpPtrAllocator &> identifiers

第一个容器保存已载入的方言对象，第二个容器记录已注册的方言，第三个容器保存已注册的抽象操作，第四个容器记录上下文里可见的标识符。

这里，方言的注册是指MLIRContext知道这个方言的标识符和构造方法，但方言对象并没有构造。方言对象的构造发生在载入时，在载入时刻，不仅构造方言对象，相关的接口也会一并准备好。抽象操作与操作相关，参考操作的管理。

5.3.1. 方言的注册

MLIR提供了一组标准的方言，它们提供了许多有用的功能。为了让程序能方便地使用标准方言，首先，每个程序在main()的入口都要注册标准方言，像这样：

int main(int argc, char **argv) {

  mlir::registerAllDialects();

  …      // 其他初始化

 registerAllDialects()的定义如下：

67  inline void registerAllDialects() {

68    static bool initOnce =

69        ([]() { registerAllDialects(getGlobalDialectRegistry()); }(), true);

70    (void)initOnce;

71  }

69行的getGlobalDialectRegistry()返回一个类型为llvm::ManagedStatic<DialectRegistry>的对象dialectRegistry，它过可视为DialectRegistry的静态对象，这个类通过一个类型为std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>的容器registry来记录标准方言。因此，同一行上的registerAllDialects()重载函数是：

41  inline void registerAllDialects(DialectRegistry &registry) {

42    // clang-format off

43    registry.insert<acc::OpenACCDialect,

44                    AffineDialect,

45                    avx512::AVX512Dialect,

46                    gpu::GPUDialect,

47                    LLVM::LLVMAVX512Dialect,

48                    LLVM::LLVMDialect,

49                    linalg::LinalgDialect,

50                    scf::SCFDialect,

51                    omp::OpenMPDialect,

52                    pdl::PDLDialect,

53                    pdl_interp::PDLInterpDialect,

54                    quant::QuantizationDialect,

55                    spirv::SPIRVDialect,

56                    StandardOpsDialect,

57                    vector::VectorDialect,

58                    NVVM::NVVMDialect,

59                    ROCDL::ROCDLDialect,

60                    SDBMDialect,

61                    shape::ShapeDialect>();

62    // clang-format on

63  }

这个方法列出了MLIR目前实现的标准方言，DialectRegistry通过一系列insert()方法完成注册：

252    template <typename ConcreteDialect, typename OtherDialect,

253                typename... MoreDialects>

254      void insert() {

255        insert<ConcreteDialect>();

256        insert<OtherDialect, MoreDialects...>();

257      }

241    template <typename ConcreteDialect>

242    void insert() {

243      insert(TypeID::get<ConcreteDialect>(),

244             ConcreteDialect::getDialectNamespace(),

245             static_cast<DialectAllocatorFunction>(([](MLIRContext *ctx) {

246               // Just allocate the dialect, the context

247               // takes ownership of it.

248               return ctx->getOrLoadDialect<ConcreteDialect>();

249             })));

250    }

53  void DialectRegistry::insert(TypeID typeID, StringRef name,

54                               DialectAllocatorFunction ctor) {

55    auto inserted = registry.insert(

56        std::make_pair(std::string(name), std::make_pair(typeID, ctor)));

57    if (!inserted.second && inserted.first->second.first != typeID) {

58      llvm::report_fatal_error(

59          "Trying to register different dialects for the same namespace: " +

60          name);

61    }

62  }

 注意，这里只是注册了这些方言的构造方法，并没有把这些方言的Dialect对象构造出来。这是因为Dialect对象的构造需要一个MLIRContext实例，因此要把Dialect对象的构造推迟到MLIRContext对象构造出来后。另外，不是每个程序都需要所有的标准方言，在需要时构造所需的方言才比较合理。所以，DialectRegistry提供了两个函数：loadByName()，loadAll()。前者构造指定名字的标准方言，后者构造所有的标准方言。

5.3.2. 方言的载入

从上面注册的构造方法我们看到，实际执行构造的函数是MLIRContext的getOrLoadDialect()，这也是一系列调用：

69    template <typename T>

70      T *getOrLoadDialect() {

71        return static_cast<T *>(

72          getOrLoadDialect(T::getDialectNamespace(), TypeID::get<T>(), [this]() {

73            std::unique_ptr<T> dialect(new T(this));

74            return dialect;

75          }));

76    }

模板参数T是具体的方言类型，它是Dialect的派生类，Dialect没有定义getDialectNamespace()，派生类必须提供自己的定义。在MLIRContext里这个名字将作为这个方言的身份识别。另外，Dialect及其派生类亦是MLIR类型系统中的组成，它们都有TypeIDMLIRContext::getOrLoadDialect()在2021版本里的定义如下：

511  Dialect *

512  MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,

513                                function_ref<std::unique_ptr<Dialect>()> ctor) {

514    auto &impl = getImpl();

515    // Get the correct insertion position sorted by namespace.

516    std::unique_ptr<Dialect> &dialect = impl.loadedDialects[dialectNamespace];

517  

518    if (!dialect) {

519      LLVM_DEBUG(llvm::dbgs()

520                 << "Load new dialect in Context " << dialectNamespace << "\n");

521  #ifndef NDEBUG

522      if (impl.multiThreadedExecutionContext != 0)

523        llvm::report_fatal_error(

524            "Loading a dialect (" + dialectNamespace +

525            ") while in a multi-threaded execution context (maybe "

526            "the PassManager): this can indicate a "

527            "missing `dependentDialects` in a pass for example.");

528  #endif

529      dialect = ctor();

530      assert(dialect && "dialect ctor failed");

531  

532      // Refresh all the identifiers dialect field, this catches cases where a

533      // dialect may be loaded after identifier prefixed with this dialect name

534      // were already created.

535      llvm::SmallString<32> dialectPrefix(dialectNamespace);

536      dialectPrefix.push_back('.');

537      for (auto &identifierEntry : impl.identifiers)

538        if (identifierEntry.second.is<MLIRContext *>() &&

539            identifierEntry.first().startswith(dialectPrefix))

540          identifierEntry.second = dialect.get();

541  

542      // Actually register the interfaces with delayed registration.

543      impl.dialectsRegistry.registerDelayedInterfaces(dialect.get());

544      return dialect.get();

545    }

546  

547    // Abort if dialect with namespace has already been registered.

548    if (dialect->getTypeID() != dialectID)

549      llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +

550                               "' has already been registered");

551  

552    return dialect.get();

553  }

从上面可以看到，构造出来的Dialect对象存放在MLIRContextImpl的loadedDialects容器中（类型DenseMap<StringRef, std::unique_ptr<Dialect>>）。同样，MLIRContext也提供这些函数获取指定方言的Dialect对象：getOrLoadDialect()，loadDialect()等。比如，Toy例子代码有这样的代码片段来构建自己的方言对象：

int dumpMLIR() {

  mlir::MLIRContext context(/*loadAllDialects=*/false);

  // Load our Dialect in this MLIR Context.

  context.getOrLoadDialect<mlir::toy::ToyDialect>();

identifiers 是MLIRContextImpl里有类型为llvm::StringMap<PointerUnion<Dialect *, MLIRContext *>, llvm::BumpPtrAllocator &>的容器。MLIR里的标识符是带有上下文前缀或方言前缀的（以“.”分隔），identifiers容器就是关联标识符与其所在上下文对象或方言对象的，一个操作在创建时首先假设它在一个上下文（参考Identifier::get()），上面537行的for循环检查是否已经创建了具有这个方言名的上下文对象，如果是把它替换为对应的方言对象。

5.3.2.1. 方言接口

接下来，通过DialectRegistry::registerDelayedInterfaces()向MLIRContextImpl注册方言的接口。这里“延迟接口”的意思是只在方言载入（或创建）时才注册接口。

106  void DialectRegistry::registerDelayedInterfaces(Dialect *dialect) const {

107    auto it = interfaces.find(dialect->getTypeID());

108    if (it == interfaces.end())

109      return;

110 

111    // Add an interface if it is not already present.

112    for (const auto &kvp : it->getSecond().dialectInterfaces) {

113      if (dialect->getRegisteredInterface(kvp.first))

114        continue;

115      dialect->addInterface(kvp.second(dialect));

116    }

117 

118    // Add attribute, operation and type interfaces.

119    for (const auto &kvp : it->getSecond().objectInterfaces)

120      kvp.second(dialect->getContext());

121  }

方言可用的接口都保存在DialectRegistry类型为DenseMap<TypeID, DelayedInterfaces>的interfaces容器中，其中DelayedInterfaces是DialectRegistry里这样的一个嵌套定义：

283  struct DelayedInterfaces {

284      /// Dialect interfaces.

285      SmallVector<std::pair<TypeID, DialectInterfaceAllocatorFunction>, 2>

286          dialectInterfaces;

287      /// Attribute/Operation/Type interfaces.

288      SmallVector<std::pair<TypeID, ObjectInterfaceAllocatorFunction>, 2>

289          objectInterfaces;

290    };

在下面我们会看到，方言除了自己的接口，还支持操作/类型/属性的外部模式接口，289行的objectInterfaces是存放这些操作接口的地方。这两个容器用到这两个定义：

30  using DialectInterfaceAllocatorFunction =

31      std::function<std::unique_ptr<DialectInterface>(Dialect *)>;

32  using ObjectInterfaceAllocatorFunction = std::function<void(MLIRContext *)>;

顾名思义，这两个std::function封装的方法用于创建接口对象，在上面115与120行，它们被调用来创建具体的接口对象。Dialect的addInterface()将生成的接口对象保存在registeredInterfaces容器中（类型DenseMap<TypeID, std::unique_ptr<DialectInterface>>）。而120行处，创建的接口对象实际上保存在对应操作/类型/属性抽象对象的容器中，下面会看到。