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

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

前言

因此这篇文章将基于MLIR的ODS文档来讲解ODS中的一些要点，帮助我们更好的了解和上手MLIR。我会把官方文档中需要注意的点拆成一些小的要点。下面文章中提到的TableGen和ODS不做特别区分，ODS中的语法也就是TableGen语法。这里介绍的要点在OneFlow对接MLIR时都或多或少用到了，感兴趣的可以对照着看看OneFlow的这部分源码。​​为什么要使用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支持的类型和表达式可以点这个链接：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]>:$ksize, Confined]>:$strides, TF_AnyStrAttrOf<["SAME", "VALID"]>:$padding, DefaultValuedAttr:$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 :$, ... :$, ...);

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

let results = (outs :$, ...);

还有Operation regions和Operation successors目前我还没用过，暂时不清楚应用场景。Op的特征和约束(Operation traits and constraints)：特征是影响语法或语义的Operation属性。MLIR C++的各种特征在​​mlir::OpTrait​​命名空间中。Operation的特征、接口或者约束涉及多个操作数/属性/结果时，要作为​​Op​​类的第二个模板参数传入。它们都需要继承于​​OpTrait​​类。详见Constraints章节。

4. Operation自动生成的默认构建方法

定义了Operation之后，我们怎么构建呢？ 每一个Operation，都会基于Operation的参数和Operation的返回值自动生成一些​​builers​​。举个例子，给出如下的Operation定义：

def MyOp : ... { let arguments = (ins I32:$i32_operand, F32:$f32_operand, ..., I32Attr:$i32_attr, F32Attr:$f32_attr, ... ); let results = (outs I32:$i32_result, F32:$f32_result, ... );}

下面的​​builders​​被产生：

// All result-types/operands/attributes have one aggregate parameter.// 所有 结果类型/操作数/属性都集合为一个聚合参数。static void build(OpBuilder &odsBuilder, OperationState &odsState, ArrayRef resultTypes, ValueRange operands, ArrayRef attributes);// Each result-type/operand/attribute has a separate parameter. The parameters// for attributes are of mlir::Attribute types.// 每一个 结果类型/操作数/属性 都是一个独立的参数。属性参数为 mlir::Attribute 类型static void build(OpBuilder &odsBuilder, OperationState &odsState, Type i32_result, Type f32_result, ..., Value i32_operand, Value f32_operand, ..., IntegerAttr i32_attr, FloatAttr f32_attr, ...);// Each result-type/operand/attribute has a separate parameter. The parameters// for attributes are raw values unwrapped with mlir::Attribute instances.// (Note that this builder will not always be generated. See the following// explanation for more details.)// 每一个 结果类型/操作数/属性 都是一个独立的参数。// 属性参数是未经 mlir::Attribute 实例包装的原始值。// (注意，该构建器并不总是生成。详见下列解释获得更多细节。)static void build(OpBuilder &odsBuilder, OperationState &odsState, Type i32_result, Type f32_result, ..., Value i32_operand, Value f32_operand, ..., APInt i32_attr, StringRef f32_attr, ...);// Each operand/attribute has a separate parameter but result type is aggregate.// 每一个 操作数/属性 都是一个独立的参数。但是结果全部集合为了一个聚合类型。static void build(OpBuilder &odsBuilder, OperationState &odsState, ArrayRef resultTypes, Value i32_operand, Value f32_operand, ..., IntegerAttr i32_attr, FloatAttr f32_attr, ...);// All operands/attributes have aggregate parameters.// Generated if return type can be inferred.// 这个构建器只有在返回值类型能够被推断出的情况下，才会生成。static void build(OpBuilder &odsBuilder, OperationState &odsState, ValueRange operands, ArrayRef attributes);// (And manually specified builders depending on the specific op.)

上面的代码注释翻译已经解释了这些builder的不同之处。并且可能还存在一些其它的builder，请参考这里的文档进行查看。

5. 自定义builder方法

假设上面生成的C++代码中构造方法没有我们所期待的，这个时候我们就需要自定义builder方法。比如：

def MyOp : Op<"my_op", []> { let arguments = (ins F32Attr:$attr); let builders = [ OpBuilder<(ins "float":$val)> ];}

class MyOp : /*...*/ { /*...*/ static void build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, float val);};

注意，这个builder有两个额外的前置参数。这些参数对于构建Operation很有用。特别地，为了能够通过该方法构建Operation，必须向​​state​​​填充该Operation的属性，操作数，域和返回值类型。​​builder​​​可以用于构建属于Op的任意IR对象，例如类型或嵌套操作。当类型与名字转换为C++代码时，它们应该是有效的C++结构，一个类型(在Op的命名空间中)与一个标识符(例如，​​class​​不是一个有效标识符)。可以在ODS中直接提供builder的实现，使用如下TableGen的代码块：

def MyOp : Op<"my_op", []> { let arguments = (ins F32Attr:$attr); let builders = [ OpBuilder<(ins "float":$val), [{ $_state.addAttribute("attr", $_builder.getF32FloatAttr(val)); }]> ];}

​​$_builder​​​和​​$_state​​​这两个特殊参数等效于​​builder​​​和​​state​​​。​​ins​​​部分中的参数可以被直接使用，比如​​val​​​。builer的c++代码实现会通过替换ODS中的特殊变量来完成，要保证builder ODS实现的其他部分是有效的C++结构。虽然对代码大小没有限制，但我们鼓励只在ODS中内联较短定义的builder，而将定义较长的builder的定义放在C++文件中。最后，如果某些参数需要默认值，可以使用 ​​CArg​​ 定义它们以包装类型和此值，如下所示：

def MyOp : Op<"my_op", []> { let arguments = (ins F32Attr:$attr); let builders = [ OpBuilder<(ins CArg<"float", "0.5f">:$val), [{ $_state.addAttribute("attr", $_builder.getF32FloatAttr(val)); }]> ];}

/// Header file.class MyOp : /*...*/ { /*...*/ static void build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, float val = 0.5f);};/// Source file.MyOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, float val) { state.addAttribute("attr", builder.getF32FloatAttr(val));}

def CallOp : Std_Op<"call", ...> { let arguments = (ins FlatSymbolRefAttr:$callee, Variadic:$args); let results = (outs Variadic); let assemblyFormat = [{ $callee `(` $args `)` attr-dict `:` functional-type($args, results) }];}

主要由三部分组成：

Directives：指令。指令是一种带有可选参数的内置函数。可用的指令有​​attr-dict​​​，​​attr-dict-with-keyword​​​，​​operands​​​，​​ref​​等等。字面值(Literals)。字面值是用``包裹起来的键值或者标点符号。下列是有效的标点符号集合：​​:, ,, =, <, >, (, ), {, }, [, ], ->, ?, +, *​​​ 。​​\n​​标点符号有另起一行的效果。如下：

let assemblyFormat = [{ `{` `\n` ` ` ` ` `this_is_on_a_newline` `\n` `}` attr-dict}];

%results = my.operation { this_is_on_a_newline}

内容为空的字面量可用于删除隐式插入某些字面量元素后的空格。例如​​)​​​或者​​]​​​等等。举个例子，​​]​​​可能出现在输出output的末尾，但它并不是格式中的最后一个元素，在这个例子里可以使用 ​​"]``"​​删除掉后续的空格。

Variables(变量)。变量是注册在Operation上的实体，例如Operation的参数(属性或操作数)，域，结果，后继者，等等。在​​CallOp​​​中，变量代表​​$callee​​​和​​$args​​。属性变量将显示其各自的值类型。除非其值的类型可以构造，在这种情况下，属性变量的值类型可以省略。

7. 自定义指令(Custom Directives) & 可选组(Optional Groups)

在某些情况下，Operations可能具有“可选”信息，例如 属性或一组空的可变参数操作数。 在这些情况下，可以根据此信息的存在将汇编格式的一部分标记为可选。

这两部分比较复杂，我还没用到，所以这里不展开，感兴趣请看官方文档。

8. 类型推断

格式的一项要求是操作数和结果的类型必须始终存在。在某些情况下，可以通过类型约束或其他可用信息来推断变量的类型。 在这些情况下，可以从格式中省略该变量的类型。

Buildable Types。可构建类型。一些类型约束可能只有一种表示，允许它们直接构建； 例如​​I32​​​ 或​​Index​​​类型。 ODS 中的类型可以通过设置​​builderCall​​​ 字段或从​​BuildableType​​ 类继承来将自己标记为可构建。Trait Equality Constraints。特征等价约束。有许多Operations具有在Operations上注册为已知类型相等特征的约束； 例如，​​select​​​ Operation的真、假和结果值通常具有相同的类型。 汇编格式可以检查这些等价的约束以辨别缺失变量的类型。 当前支持的特征有：​​AllTypesMatch​​​、​​TypesMatchWith​​​、​​SameTypeOperands​​​ 和​​SameOperandsAndResultType​​。InferTypeOpInterface。实现​​InferTypeOpInterface​​ 的Operations可以在其汇编格式中省略其结果类型，因为可以从操作数中推断出结果类型。hasCanonicalizer。此布尔字段指示是否已为此Operation定义规范化模式。 如果它是​​1​​​，那么​​::getCanonicalizationPatterns()​​ 应该被定义。hasCanonicalizeMethod。当此布尔字段设置为​​true​​​ 时，表示操作为简单的“matchAndRewrite”样式规范化模式实现了​​canonicalize​​​方法。 如果​​hasCanonicalizer​​​ 为 0，则实现​​::getCanonicalizationPatterns()​​ 的实现来调用此函数。hasFolder。此布尔字段指示是否已为此操作定义了通用折叠规则。 如果它是​​1​​​，那么​​::fold()​​ 应该被定义。

表驱动操作定义的目标之一是为每个操作自动生成尽可能多的逻辑和方法。 话虽如此，总会有无法涵盖的长尾案例。 对于这种情况，您可以使用 ​​extraClassDeclaration​​​。 ​​extraClassDeclaration​​ 中的代码将逐字复制到生成的 C++ op 类。

请注意，​​extraClassDeclaration​​ 是一种针对高级用户的长尾案例的机制； 对于尚未实施的广泛适用的情况，改善基础设施是可取的。

10. 生成C++代码

定义文件包含所有的op方法定义，可以通过定义​​GET_OP_CLASSES​​​来包含和启用。 对于每个操作，OpDefinitionsGen 生成一个操作类和一个操作数适配器类(​​operand adaptor class​​​)。 此外，它还包含一个以逗号分隔的所有已定义Operations的列表，可以通过定义 ​​GET_OP_LIST​​ 来包含和启用这些操作。

Class name and namespaces。

对于每个Operation，其生成的C++类名是使用TableGen ​​def​​​为前缀的名字，并删除了Dialect前缀。第一个​​_​​​用作分隔符。例如，对于​​def TF_AddOp​​​，C++类名会是​​AddOp​​​。我们移除了​​TF​​前缀，因为它是多个Operation作用域。其它Dialect也可以定义自己的AddOps。

生成的C++类的namespace将来自Dialect的​​cppNamespace​​​字段。举个例子，如果一个Dialect的​​Namespace​​​是​​A::B​​​，那么该Dialect的Op将被放置在​​namespace A { namespace B { ... } }​​​。 如果Dialect没有指定​​cppNamespace​​，我们就使用方言的名称作为命名空间。

这意味着生成的 C++ 类的名称不一定与Operation 名称中的操作名称完全匹配。 这是为了允许灵活命名以满足编码风格要求。

Operand adaptors

对于每个Operation，MLIR会自动生成一个操作数适配器。这个类解决了访问作为列表值提供的操作数而不使用“魔术“”常量的问题。 操作数适配器引用一个 ​​Value​​​ 数组，并提供与Operation类中名称相同的方法来访问它们。例如，对于二元算术运算，它可以提供 ​​.lhs()​​​ 来访问第一个操作数和 ​​.rhs()​​​ 来访问第二个操作数。 操作数适配器类与Operation类位于同一命名空间中，类的名称由Operation类的名称后面接一个​​Adaptor​​组成。

操作数适配器也可以用于处理Operation的函数模板：

template std::pair zip(BinaryOpTy &&op) { return std::make_pair(op.lhs(), op.rhs());;}void process(AddOp op, ArrayRef newOperands) { zip(op); zip(Adaptor(newOperands)); /*...*/}

在OneFlow中，我们可以看到生成的​​UserOpAdaptor​​代码。里面提供了一系列接口可以访问Operation的操作数以及相关属性。

//===----------------------------------------------------------------------===//// ::mlir::oneflow::UserOp declarations//===----------------------------------------------------------------------===//class UserOpAdaptor {public: UserOpAdaptor(::mlir::ValueRange values, ::mlir::DictionaryAttr attrs, ::mlir::RegionRange regions = {}); UserOpAdaptor(UserOp &op); ::mlir::ValueRange getOperands(); std::pair getODSOperandIndexAndLength(unsigned index); ::mlir::ValueRange getODSOperands(unsigned index); ::mlir::ValueRange data_input(); ::mlir::ValueRange ctrl_inputs(); ::mlir::DictionaryAttr getAttributes(); ::mlir::StringAttr op_name(); ::mlir::BoolAttr trainable(); ::mlir::StringAttr device_tag(); ::mlir::ArrayAttr device_name(); ::mlir::IntegerAttr scope_symbol_id(); ::mlir::ArrayAttr hierarchy(); ::mlir::DenseIntElementsAttr operand_segment_sizes(); ::mlir::DenseIntElementsAttr result_segment_sizes(); ::mlir::StringAttr op_type_name(); ::mlir::ArrayAttr input_lbn_segment_keys(); ::mlir::ArrayAttr input_lbn_segment_sizes(); ::mlir::ArrayAttr output_lbn_segment_keys(); ::mlir::ArrayAttr output_lbn_segment_sizes(); ::mlir::ArrayAttr output_lbns(); ::mlir::LogicalResult verify(::mlir::Location loc);private: ::mlir::ValueRange odsOperands; ::mlir::DictionaryAttr odsAttrs; ::mlir::RegionRange odsRegions;};

11. 约束

约束(Constraint)是表驱动Operation定义中的一个核心概念：Operation验证和图Operation匹配都是基于满足约束。因此，Operation定义和重写规则都直接涉及写入约束。MLIR在​​OpBase.td​​​(​​5 的 32 位整数）多个操作数和结果（例如，第一个结果的形状必须与第一个操作数（可理解为Tensor）相同）操作本身固有的。（例如没有副作用，参考Transpose Op消除那个案例）

我们将它们分别称为单实体约束、多实体约束和特征。

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。