【从零开始学深度学习编译器】十六，MLIR ODS要点总结上篇

前言

在【从零开始学深度学习编译器】十二，MLIR Toy Tutorials学习笔记一 中提到MLIR是通过Dialect来统一各种不同级别的IR，即负责定义各种Operation（算子）。然后对Dialect和Operation的定义又是通过TabelGen规范构造的，通过TableGen驱动MLIR的Operation定义也被称作ODS（ Operation Definition Specification) 。我们目前只是简单认识了Toy Tutorials的Dialect和Operation是如何通过ODS定义的，但对ODS本身的语法以及一些限制都没有太多了解，这就导致在看一些相关工程的Operation定义时时常陷入迷惑，不知道某个字段是什么含义，或者说自定义Op的时候的应当如何声明操作数和Attr（举个例子，要将卷积的groups参数设置为可选的属性，应该怎么做）。

因此这篇文章将基于MLIR的ODS文档来讲解ODS中的一些要点，帮助我们更好的了解和上手MLIR。我会把官方文档中需要注意的点拆成一些小的要点。下面文章中提到的TableGen和ODS不做特别区分，ODS中的语法也就是TableGen语法。这里介绍的要点在OneFlow对接MLIR时都或多或少用到了，感兴趣的可以对照着看看OneFlow的这部分源码。https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/ir/include/OneFlow/OneFlowOps.td。

1. 为什么要使用ODS来定义Operation

在MLIR中要定义Operation支持用C++直接定义以及基于ODS框架定义两种方法。使用C++直接定义要求我们继承基类Op的一些构造方法并重写，对于每一个Op都要写一段C++代码。可以想到这样做整个系统的Op定义部分会非常冗余，产生大量可重复代码并且可读性也会比较差。如果基于ODS来定义Operation，我们只需要将Op定义按照ODS的规范统一写到一个td文件中，然后使用MLIR提供的代码生成工具自动生成Operation的C++定义，这种完全auto codegen的方式非常优雅的实现了Operation定义并且需要用户操心的东西（也就是ODS的语法规范）更加直观。

ODS是MLIR定义Operation的不二选择，因此我们有必要学习ODS的语法规范。

2. TableGen语法

一个TableGen文件（以.td结尾）包含以下一些语法：

• TableGen class 类似于C++的class，可以作为模板或者基类去派生子类。
• TableGen def 类似于C++的对象。以用一个TableGen class的特化来声明，例如，def MyDef: MyClass<...>;，也可以单独使用def MyDef;。它不能用作模板，也不能作为基类去派生子类。
• TableGen dag 是一种专门用于有向无环图元素的类型。一个dag类型带有一个操作符和零个或者多个参数。语法形如(operator arg0, arg1, argN.)，其中operator可以是任意的TableGen def。参数可以是任何东西，包括dag本身。我们可以将名称附加到操作符和参数上，如(MyOp:$op_name MyArg:$arg_name)。

想了解更多TableGen支持的类型和表达式可以点这个链接：https://llvm.org/docs/TableGen/ProgRef.html。

3. Operation定义

MLIR定义了几个公共的结构用于帮助定义Operation，并通过TableGen backend : OpDefinitionsGen提供它们的语义。这些公共结构在文件OpBase.td中定义。主要包括：

• Op类：这是定义Operation时使用的主要结构。在特化该类时，通过下述结构的帮助，指定与Operation有关的所有事实。
• Dialect类：归属于同一个逻辑组的Operation会被放置在同一个Dialect下。Dialect包含了方言等级信息。
• OpTrait类及其子类：它们用于指定Operation的特殊属性和约束，包括Operation是否具有副作用、Op的输出是否与输入具有相同的形状等。
• ins/outs标记：这是OpDefinitionsGen后端内置的两个特殊标记，分别引导操作数(operands)/属性(attributes)、结果(results)的定义。
• TypeConstraint类及其子类：它们用于指定对操作数(operands)或结果(results)的约束。一个值得注意的子类是Type，它代表通用C++类型的约束。
• AttrConstraint类及其子类：它们用于指定对属性(attributes)的约束。一个值得注意的子类是Attr，它代表值为通用类型的属性的约束。

一个Operation是通过特化Op类定义的，特化后的Op类包含它需要的所有字段的具体内容。举个例子，tf.AvgPool定义如下：

def TF_AvgPoolOp : TF_Op<"AvgPool", [NoSideEffect]> {
   
  let summary = "Performs average pooling on the input.";

  let description = [{
   
Each entry in `output` is the mean of the corresponding size `ksize`
window in `value`.
  }];

  let arguments = (ins
    TF_FpTensor:$value,

    Confined<I64ArrayAttr, [ArrayMinCount<4>]>:$ksize,
    Confined<I64ArrayAttr, [ArrayMinCount<4>]>:$strides,
    TF_AnyStrAttrOf<["SAME", "VALID"]>:$padding,
    DefaultValuedAttr<TF_ConvertDataFormatAttr, "NHWC">:$data_format
  );

  let results = (outs
    TF_FpTensor:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

下面描述一下定义一个Operation所需的所有字段。有关支持的字段的完整列表，请参阅Op类的定义(就是OpBase.td)。

• Operation name : 就是Operation的名字，比如TensorFlow Dialect中的tf.Add。
• Operation documentation : Operation的文档描述，包含summary和description两种，大家看下就懂，不多说。
• Operation arguments ： Operation的参数，一个Operation有两种参数一种是operands即操作数，一种是attributes属性参数。其中属性参数又分为Natural attributes和Derived attributes两种，前者为自然属性必须指定比如卷积的输出通道数，后者为派生属性比如输出Tensor的形状。

操作数和属性都在dag类型的arguments中被指定，以ins引导：

let arguments = (ins
  <type-constraint>:$<operand-name>,
  ...
  <attr-constraint>:$<attr-name>,
  ...
);

这里<type-constraint>是一个来自TypeConstraint类层次的TableGen def。与此类似的，<attr-constraint>是一个来自AttrConstraint类层次的TableGen def。在Constraints章节有更多详细内容。

• 可变操作数。定义一个可变操作数，需要用Variadic<...>把TypeConstraint包起来。通常，Operation是没有可变操作数或者只有一个可变操作数。对于后一种情况，可以通过静态可变操作数的定义很容易的推导出动态可变操作数。但是，如果一个Operation有多个可变长度操作数(可选的或可变长度的)，在没有来自该操作的进一步信息的情况下，就不可能将动态操作数归因于相应的静态可变长度操作数定义。因此，需要用SameVariadicOperandSize或AttrSizedOperandSegments特征来表明所有的可变长度操作数都有与之对应的动态值。
• 可选操作数。定义一个可选操作数，需要用Optional<...>把TypeConstraint包起来。解释和可变操作数一样。
• 可选属性。定义一个可选属性，需要使用OptionalAttr<...> 把AttrConstraint包起来。
• 带默认值的可选属性。使用DefaultValuedAttr<..., "...">把AttrConstraint包起来。DefaultValuedAttr的第二个参数应该是包含C++默认值的字符串。举个例子，一个单精度浮点默认值需要被指定为“0.5f”，一个整型数组的默认值需要被指定为