现在我们已经熟悉 Toy 语言和它的AST表示,现在让我们看看 MLIR 是怎样帮助编译 Toy 源程序的。
简介:多层中间表示
其他的编译器,像 LLVM,是提供一个固定的预定义类型和指令(通常是底层的像 RISC的指令)。对于一个给定的语言,在发出 LLVM IR之前,执行任何的语言特定的类型检查和分析,或者变换等,是由编译器前端来决定的。例如Clang 将会使用它的 AST 做静态分析 和 变换,例如 通过对 AST 的克隆和重写完成C++ 模版实例化。最后,一个比 C/C++ 还高的层级上且带有构造特性的语言,从它们的 AST 到生成 LLVM IR 可能需要做一个非常重要的下降。因此,多种前端,导致需要重新实现大量的基础设施部件,以便能够支持这些前端的分析和转换。MLIR 通过把 MLIR 设计得具有可扩展性,来应对这个问题。例如,MLIR 中几乎没有预定义的指令(用MLIR的术语,这叫做 operations 操作)和预定义的类型。
与 MLIR 的接口
MLIR 语言参考手册:https://mlir.llvm.org/docs/LangRef/
MLIR 被设计成为一个完全可伸缩的基础设施,这里没有封闭的属性、操作或类型的集合。MLIR 通过 Dialects这个概念来实现其可扩展性。Dialects 为唯一的命名空间下的抽象提供一组机制。
在MLIR中,Operations 是进行抽象和计算的核心单元,这在很多方面都挺像 LLVM 中的 instructions。 Operations 可以具有应用特定的语义,也可以用来表示 LLVM 中所有的核心的 IR 结构:instruction、globals(类似 functions)、modules等等。
下面是 Toy 语言中的 transpose operations 的 MLIR 汇编:
%t_tensor = "toy.transpose"(%tensor) {inplace = true} : (tensor<2x3xf64>) -> tensor<3x2xf64> loc("example/file/path":12:1)
「注:这里只是声明了transpose操作,具体的实现要在 Chapter5中通过 lowering 来下降到其进一步的实现细节」
让我们来解析一下这个 MLIR operations:
. %t_tensor
给这个 operation 的结果所定义的名字(为了避免冲突,它包含一个前缀)。一个 operation 可能定义一个或者多个结果(在 Toy 的上下文中,我们将限制在不超过一个结果的 operations),这些结果都是 SSA 值。这个名字将会在解析时被使用,但是并不是持久的(这个 SSA 值的内存表示不会被追踪)。
. "toy.transpose"
operation 的名字。在其中“.”前面的前缀所定义的命名空间中,这个字符串被期待是独一无二的。这个写法可以读作:toy dialect 中的 transpose operation。
. (%tensor)
0个或多个输入的操作数(或者参数)构成的列表,它们是由其他的 operations 或者对参数 block 的引用所定义的 SSA 值。
. { inplace = true }
0个或多个属性构成的字典,它们是特殊的操作数,永远是常数。这里我们定义了一个boolean 类型的命名为 inplace 的属性,它具有常数值 true。
. (tensor<2x3xf64>) -> tensor<3x2xf64>
这个代表了用函数形式表示的操作类型,圆括号中拼写出了参数的类型,后边跟着返回值的类型。
. loc("example/file/path":12:1)
这是这个操作的定义开头在源码中的位置。
这里展示了一个操作的通常的形式。如上所述,MLIR 中的操作是可扩展的。我们使用了一个最小组的概念来建模 operations,使得 operations 可以被推导和一般性的修改。这些概念如下:
A name for the operation.
A list of SSA operand values.
A list of attributes.
A list of types for result values.
A source location for debugging purposes.
A list of successors blocks (for branches, mostly).
A list of regions (for structural operations like functions).
在 MLIR中,每一个 operation 都有一个强制性的源代码位置与其关联。与 LLVM 不同,debug 使用的 位置信息是元数据,可以背抛弃,在MLIR中,位置信息是一个核心要求,并且又一些API是依赖这些信息的,而且可以修改它们。所以,丢弃位置信息是一个显式的选择,它不会因为失误而发生。
这里提供一个说明:如果一个变换中替换了一个 operation,那么,新的 operation 必须也附带有位置信息。这使得追踪这个 operation 的出处成为可能。
值得注意的 mlir-tool 这个工具——它是用来测试 编译器 passes 的工具——在其输出中,默认并不包含位置信息。而 -mlir-print-debuginfo 这个标志可以指定其输出中包含位置信息。(运行 mlir-opt --help 可以看到更多的选项)
不透明的 API
MLIR 被设计成为可以允许一切 IR 元素,例如,自定义的属性,operations,和类型。同时, IR 元素可以总是规约为上述这些基本概念。如此一来,这就允许 MLIR 去对任意 operation 进行解析,表示和遍历 IR。例如,我们可以把我们的toy 的 operation 从上方移进一个 .mlir 文件,并且通过 mlir-opt 遍历它,这里不需要注册任何与 toy 语言关联的 dialect:
func.func @toy_func(%tensor: tensor<2x3xf64>) -> tensor<3x2xf64> {
%t_tensor = "toy.transpose"(%tensor) { inplace = true } : (tensor<2x3xf64>) -> tensor<3x2xf64>
return %t_tensor : tensor<3x2xf64>
}
在 属性、operations和类型在没有被注册的情况下, MLIR 将会强制一些结构性的约束(例如,dominance 等等),但是在其他方面,它们就是完全不透明的。例如,MLIR 关于一个未注册的 operation 在如下情况中,只知道很少的信息:这个 operation 是否可以作用在一些特别的数据类型上,这个 operation 可以接收多少个操作数,这个 operation 可以产生多少个结果。这种灵活性对于引导性的目的是有用的,但是它通常不建议在成熟的系统中使用。未注册的 operations 在变换和分析时必须保守对待,并且它们在构造和修改上非常困难。
通过精心设计什么是给Toy语言的非法的IR,并且通过观察它不走检测器时的遍历,这个处理的意境可以被观察到:
func.func @main() {
%0 = "toy.print"() : () -> tensor<2x3xf64>
}
这里有多个问题:toy.print operation 不是一个终结操作;它应该接收一个操作数;而且,它不应该返回任何值。在下一节中,我们会使用 MLIR 注册我们 dialect 和 operations,安装进 verifier中,并且添加更好的 APIs 来修改我们的 operations。
定义 Toy 的一个 Dialect
为了与 MLIR 有效地联系,我们 将会定义一个新的 Toy dialect。这个 dialect 将会对 Toy 语言的结构进行建模,同时为高级层次的分析和变换提供一个容易的途径。
/// This is the definition of the Toy dialect. A dialect inherits from
/// mlir::Dialect and registers custom attributes, operations, and types. It can
/// also override virtual methods to change some general behavior, which will be
/// demonstrated in later chapters of the tutorial.
class ToyDialect : public mlir::Dialect {
public:
explicit ToyDialect(mlir::MLIRContext *ctx);
/// Provide a utility accessor to the dialect namespace.
static llvm::StringRef getDialectNamespace() { return "toy"; }
/// An initializer called from the constructor of ToyDialect that is used to
/// register attributes, operations, types, and more within the Toy dialect.
void initialize();
};
上面为 dialect 的 C++ 定义方式,但是 MLIR 也支持通过 tablegen 声明性地定义 dialects。使用 td 声明性的说明是更干净的,因为在定义一个新的 dialect 时,它移除了对大量样板文件的需要。它也使得 dialect 文档的生成变得更容易,它可以直接在 dialect 的旁边做描述。在这种 声明性的格式中,toy 的 dialect 应该如下这样做说明:
// Provide a definition of the 'toy' dialect in the ODS framework so that we
// can define our operations.
def Toy_Dialect : Dialect {
// The namespace of our dialect, this corresponds 1-1 with the string we
// provided in `ToyDialect::getDialectNamespace`.
let name = "toy";
// A short one-line summary of our dialect.
let summary = "A high-level dialect for analyzing and optimizing the "
"Toy language";
// A much longer description of our dialect.
let description = [{
The Toy language is a tensor-based language that allows you to define
functions, perform some math computation, and print results. This dialect
provides a representation of the language that is amenable to analysis and
optimization.
}];
// The C++ namespace that the dialect class definition resides in.
let cppNamespace = "toy";
}
为了看到这会生成什么内容,我们可以运行 mlir-tblgen 命令,带上 gen-dialect-decls 功能,像这样:
${build_root}/bin/mlir-tblgen -gen-dialect-decls \
${mlir_src_root}/examples/toy/Ch2/include/toy/Ops.td \
-I ${mlir_src_root}/include/
在 dialect 被定义之后,它现在可以被加载进一个 MLIRContext之中:
context.loadDialect<ToyDialect>();
默认的话,一个 MLIRContext 只能加载一个 Builtin Dialect,它提供一些核心的 IR 组件,也就是说,其他的 dialects,例如我们自己定义的 Toy dialect,必须被显式地加载才行。
定义 Toy operations
现在我们已经拥有一个 Toy dialect,我们可以开始定义 operations 了。这将允许提供寓意信息,以便系统的其余部分可以连接进去。作为示例,让我们一起浏览一遍 toy.constant operation 的创建。这个 operation 将会在 Toy 语言中表示一个常数值。
%4 = "toy.constant"() {value = dense<1.0> : tensor<2x3xf64>} : () -> tensor<2x3xf64>
这个 operation 接收0个操作数,一个稠密的元素构成的属性,称之为 value,用来表示常熟值,并且返回一个 RankedTensorType 类型的结果。这个操作继承自 CRTP mlir::Op 类,它还有几个可选的 trait 属性来定义它的行为。 Traits 是一个机制,可以用它向 operation 注入额外的行为,例如额外的 访问器,验证和其他行为。关于上边我们讲到的常量 operation,让我们一起看看下边这个可能的定义:
class ConstantOp : public mlir::Op<
/// `mlir::Op` is a CRTP class, meaning that we provide the
/// derived class as a template parameter.
ConstantOp,
/// The ConstantOp takes zero input operands.
mlir::OpTrait::ZeroOperands,
/// The ConstantOp returns a single result.
mlir::OpTrait::OneResult,
/// We also provide a utility `getType` accessor that
/// returns the TensorType of the single result.
mlir::OpTrait::OneTypedResult<TensorType>::Impl> {
public:
/// Inherit the constructors from the base Op class.
using Op::Op;
/// Provide the unique name for this operation. MLIR will use this to register
/// the operation and uniquely identify it throughout the system. The name
/// provided here must be prefixed by the parent dialect namespace followed
/// by a `.`.
static llvm::StringRef getOperationName() { return "toy.constant"; }
/// Return the value of the constant by fetching it from the attribute.
mlir::DenseElementsAttr getValue();
/// Operations may provide additional verification beyond what the attached
/// traits provide. Here we will ensure that the specific invariants of the
/// constant operation are upheld, for example the result type must be
/// of TensorType and matches the type of the constant `value`.
LogicalResult verifyInvariants();
/// Provide an interface to build this operation from a set of input values.
/// This interface is used by the `builder` classes to allow for easily
/// generating instances of this operation:
/// mlir::OpBuilder::create<ConstantOp>(...)
/// This method populates the given `state` that MLIR uses to create
/// operations. This state is a collection of all of the discrete elements
/// that an operation may contain.
/// Build a constant with the given return type and `value` attribute.
static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
mlir::Type result, mlir::DenseElementsAttr value);
/// Build a constant and reuse the type from the given 'value'.
static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
mlir::DenseElementsAttr value);
/// Build a constant by broadcasting the given 'value'.
static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
double value);
};
然后,我们可以在 ToyDialect 的 initializer 函数中注册这个 operation:
void ToyDialect::initialize() {
addOperations<ConstantOp>();
}
Op 与 Operation: 使用 MLIR Operations
现在我们已经定义了一个 operation,我们想要访问并变换它。 在 MLIR 中,这里有两个主要的 class 跟 operations 相关联:Operation 和 Op。Operation 类用于一般意义上建模所有的 operations。它是不透明的,也就是说,它并不描述 特定 operation 的 properties 和 operations 的 types。相反,Operation class 提供通用的 API 给一个operation 实例。另一方面,operation 的每一个特定的类型是由 Op 的派生类来表示的。例如,ConstantOp 的表示一个0输入和一个输出的 operation,它总是被设置成为同样的值。Op的派生类扮演智能指针封装着 Operation*,提供 特定 operation 访问器方法,同时提供operations 的 类型安全的 properties。这意味着,当我们定义我们的Toy operations 时,我们简单地定义一个干净的、语义用途的接口来构建和对接 Operation class。这是为什么我们的 ConstantOp 没有定义class 的字段,这个 operation 的所有数据都存储在所引用的 Operation class 之中。这个设计的一个副作用是我们总是值传递 Op 的派生类,而不是传递引用或指针(值传递是MLIR的特色,同样应用于 attributes和types等)。给定一个一般的 Operation*实例,使用 LLVM 的类型变换设施,我们总是可以得到一个特定的 Op 的实例:
void processConstantOp(mlir::Operation *operation) {
ConstantOp op = llvm::dyn_cast<ConstantOp>(operation);
// This operation is not an instance of `ConstantOp`.
if (!op)
return;
// Get the internal operation instance wrapped by the smart pointer.
mlir::Operation *internalOperation = op.getOperation();
assert(internalOperation == operation &&
"these operation instances are the same");
}
使用 ODS Framework
使用 Operation Definition Specification (ODS) Framework
出了可以特化 mlir::Op 模版,MLIR 还支持使用声明性的方式定义 operations。这是通过 Operation Definition Specificaiton framework 做到的。
也就是把一个 operaiton 通过一个简洁的 TableGen 纪录来特化,在编译的时候,它将被展开成为等价的 mlir::Op C++模版的特化。
使用 ODS framework 是在 MLIR 中定义operations的提倡的方式,因为这很简单,简洁,同时对接 C++ API 的变化表现是稳定的。
让我们一起看看我们 ConstantOp 的等价的 ODS 定义:
使用 ODS时, operations 是通过继承 Op class 来定义的。为了简化我们的 operation 定义,我们将在 Toy dialect中定义一个 operation的基类(因为它们都是 toy dialect 中的 operation,所以这本身是一个共性,故可以存在一个基类。比如它们都会被注册进 同一个 toy dialect)。
// Base class for toy dialect operations. This operation inherits from the base
// `Op` class in OpBase.td, and provides:
// * The parent dialect of the operation.
// * The mnemonic for the operation, or the name without the dialect prefix.
// * A list of traits for the operation.
class Toy_Op<string mnemonic, list<Trait> traits = []> :
Op<Toy_Dialect, mnemonic, traits>;
结合这个初步的定义的代码,我们可以开始定义constant operation。
我们通过继承上述基类 Toy_Op 来定义一个 toy operation。这里我们给 operation 提供了 mnemonic(助记符) 和一个traits 列表。
这个助记符与 Constant::OperationName 这个成员方法提供的名字相匹配,只是需要去掉 dialect 前缀: toy..
与我们的 C++ 定义少了的部分是 ZeroOperands 和 OneResult traits. 这些将会基于我们稍后定义的 arguments 和 results 字段自动推导出来。
def ConstantOp : Toy_Op<"constant"> {
}
到此为止,你可能想知道 TableGen 生成的代码看起来会是什么样子的。带着 -gen-op-decls 或者 -gen-op-defs 动作,简单地运行 mlir-tblgen 命令,具体如下所示:
${build_root}/bin/mlir-tblgen -gen-op-defs ${mlir_src_root}/examples/toy/Ch2/include/toy/Ops.td -I ${mlir_src_root}/include/
依赖于所选择的动作,这将会打印出 ConstantOp class 的声明和它的实现。比较输出的内容和手写的实现,对于开始使用 TableGen 是非常有益的。
定义 Arguments 和 Results
结合刚刚定义的这个 operation 的壳,我们可以给我们的 operation 提供输入和输出。
一个 operation 的输入或者 arguments 可能会是 SSA 的操作数值的 attributes 或者 types。
这个结果对应到一组该 operation 产生的值的类型:
def ConstantOp : Toy_Op<"constant"> {
// The constant operation takes an attribute as the only input.
// `F64ElementsAttr` corresponds to a 64-bit floating-point ElementsAttr.
let arguments = (ins F64ElementsAttr:$value);
// The constant operation returns a single value of TensorType.
// F64Tensor corresponds to a 64-bit floating-point TensorType.
let results = (outs F64Tensor);
}
通过提供一个给 arguments 或 results 的名字,例如 $value, ODS 将会自动生成一个对应的访问器:DenseElementsAttr ConstantOp::value().
添加文档
定义操作的下一个步骤是文档化它。Operations 可以提供 summary 和 description 字段来描述这个 operation 的语义。
这些信息对于这个 dialect 的用户是有益的,甚至可以用来自动生成 Markdown 文档。
def ConstantOp : Toy_Op<"constant"> {
// Provide a summary and description for this operation. This can be used to
// auto-generate documentation of the operations within our dialect.
let summary = "constant operation";
let description = [{
Constant operation turns a literal into an SSA value. The data is attached
to the operation as an attribute. For example:
%0 = "toy.constant"()
{ value = dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64> }
: () -> tensor<2x3xf64>
}];
// The constant operation takes an attribute as the only input.
// `F64ElementsAttr` corresponds to a 64-bit floating-point ElementsAttr.
let arguments = (ins F64ElementsAttr:$value);
// The generic call operation returns a single value of TensorType.
// F64Tensor corresponds to a 64-bit floating-point TensorType.
let results = (outs F64Tensor);
}
验证操作的语义
到此为止,我们已经覆盖了原先的 C++ 的 operation 的定义的大部分内容。下一个需要定义的部分是 verifier。
幸运的是,很想刚才命名的访问器,ODS framework 将会基于我们给定的约束自动生成一大堆必要的验证逻辑。
这意味着我们不需要验证返回类型的结构,甚至输入的 attribute value。在大多数情况下,对于 ODS operations,额外的验证是不需要的。
添加额外的 验证逻辑,一个 operation 可以重载 verifier 字段。这个 verifier 字段允许定义一个 C++ 代码块,作为ConstantOp::verify的一部分来运行。
这个额外的代码块可以假设这个 operation 的所有其他的不变量都已经检查过了:
def ConstantOp : Toy_Op<"constant"> {
// Provide a summary and description for this operation. This can be used to
// auto-generate documentation of the operations within our dialect.
let summary = "constant operation";
let description = [{
Constant operation turns a literal into an SSA value. The data is attached
to the operation as an attribute. For example:
%0 = "toy.constant"()
{ value = dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64> }
: () -> tensor<2x3xf64>
}];
// The constant operation takes an attribute as the only input.
// `F64ElementsAttr` corresponds to a 64-bit floating-point ElementsAttr.
let arguments = (ins F64ElementsAttr:$value);
// The generic call operation returns a single value of TensorType.
// F64Tensor corresponds to a 64-bit floating-point TensorType.
let results = (outs F64Tensor);
// Add additional verification logic to the constant operation. Setting this bit
// to `1` will generate a `::llvm::LogicalResult verify()` declaration on the
// operation class that is called after ODS constructs have been verified, for
// example the types of arguments and results. We implement additional verification
// in the definition of this `verify` method in the C++ source file.
let hasVerifier = 1;
}
附加 build 方法
跟起初 C++ 定义 operation 的例子中相比,最后还缺少的组件是 build methods。
ODS 可以自动地产生一些简单的 build 方法,而且在本例子中,它将会产生我们的第一个 build 方法。其余的情况,我们通过定义 builders 字段来定义。这个字段接收一系列 OpBuilder 对象,这些对象会对应地接受一个字符串,作为C++ 参数列表的一个部分,同时可选的 代码块可以用来指定这个 builder 的内联实现。
def ConstantOp : Toy_Op<"constant"> {
...
// Add custom build methods for the constant operation. These methods populate
// the `state` that MLIR uses to create operations, i.e. these are used when
// using `builder.create<ConstantOp>(...)`.
let builders = [
// Build a constant with a given constant tensor value.
OpBuilder<(ins "DenseElementsAttr":$value), [{
// Call into an autogenerated `build` method.
build(builder, result, value.getType(), value);
}]>,
// Build a constant with a given constant floating-point value. This builder
// creates a declaration for `ConstantOp::build` with the given parameters.
OpBuilder<(ins "double":$value)>
];
}
指定自定义的汇编格式
到此为止,我们可以产生我们的 “Toy IR”。例如如下代码:
# User defined generic function that operates on unknown shaped arguments.
def multiply_transpose(a, b) {
return transpose(a) * transpose(b);
}
def main() {
var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
var b<2, 3> = [1, 2, 3, 4, 5, 6];
var c = multiply_transpose(a, b);
var d = multiply_transpose(b, a);
print(d);
}
可以产生如下的 IR:
module {
"toy.func"() ({
^bb0(%arg0: tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":4:1), %arg1: tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":4:1)):
%0 = "toy.transpose"(%arg0) : (tensor<*xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:10)
%1 = "toy.transpose"(%arg1) : (tensor<*xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:25)
%2 = "toy.mul"(%0, %1) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:25)
"toy.return"(%2) : (tensor<*xf64>) -> () loc("test/Examples/Toy/Ch2/codegen.toy":5:3)
}) {sym_name = "multiply_transpose", type = (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64>} : () -> () loc("test/Examples/Toy/Ch2/codegen.toy":4:1)
"toy.func"() ({
%0 = "toy.constant"() {value = dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":9:17)
%1 = "toy.reshape"(%0) : (tensor<2x3xf64>) -> tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":9:3)
%2 = "toy.constant"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>} : () -> tensor<6xf64> loc("test/Examples/Toy/Ch2/codegen.toy":10:17)
%3 = "toy.reshape"(%2) : (tensor<6xf64>) -> tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":10:3)
%4 = "toy.generic_call"(%1, %3) {callee = @multiply_transpose} : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":11:11)
%5 = "toy.generic_call"(%3, %1) {callee = @multiply_transpose} : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":12:11)
"toy.print"(%5) : (tensor<*xf64>) -> () loc("test/Examples/Toy/Ch2/codegen.toy":13:3)
"toy.return"() : () -> () loc("test/Examples/Toy/Ch2/codegen.toy":8:1)
}) {sym_name = "main", type = () -> ()} : () -> () loc("test/Examples/Toy/Ch2/codegen.toy":8:1)
} loc(unknown)
有一个事情需要注意,我们所有的 Toy operations 都是使用通用的汇编个是打印的。这种格式跟本章开头分析 toy.transpose 的时候显示的一样。MLIR 允许 operations 定义它们独有的汇编格式,或者是声明式的,或者是通过 C++ 命令式的定义。用户自定义的汇编格式允许裁剪调一些通用 IR,使其变得更可读,这通常是删除一些通用格式中需要的一些次要的枝桠。让我们一起浏览一个我们将会简化其 operation 格式的例子。
toy.print
当前的 toy.print 的歌是有点冗长。这里有许多额外的字符是我们想要剥离掉的。让我们先想一下更好的 toy.print 的格式应该是什么样的,然后再看看我们如何实现它。看着基本的 toy.print 的格式,我们得到:
toy.print %5 : tensor<*xf64> loc(...)
这里我们已经剥离掉了不是很有必要的成分,而且变得更加可读了。为了提供一个自定义的汇编格式,一个 operation 可以用 C++ 重写 hasCustomAssemblyFormat 字段,或者重写声明式的 assemblyFormat 字段(tableGen)。我们先看看 C++ 的变体,因为这是声明式的格式内部映射成的样子。
/// Consider a stripped definition of `toy.print` here.
def PrintOp : Toy_Op<"print"> {
let arguments = (ins F64Tensor:$input);
// Divert the printer and parser to `parse` and `print` methods on our operation,
// to be implemented in the .cpp file. More details on these methods is shown below.
let hasCustomAssemblyFormat = 1;
}
一个 C++ 实现的 printer 和 parser如下所示:
/// The 'OpAsmPrinter' class is a stream that will allows for formatting
/// strings, attributes, operands, types, etc.
void PrintOp::print(mlir::OpAsmPrinter &printer) {
printer << "toy.print " << op.input();
printer.printOptionalAttrDict(op.getAttrs());
printer << " : " << op.input().getType();
}
/// The 'OpAsmParser' class provides a collection of methods for parsing
/// various punctuation, as well as attributes, operands, types, etc. Each of
/// these methods returns a `ParseResult`. This class is a wrapper around
/// `LogicalResult` that can be converted to a boolean `true` value on failure,
/// or `false` on success. This allows for easily chaining together a set of
/// parser rules. These rules are used to populate an `mlir::OperationState`
/// similarly to the `build` methods described above.
mlir::ParseResult PrintOp::parse(mlir::OpAsmParser &parser,
mlir::OperationState &result) {
// Parse the input operand, the attribute dictionary, and the type of the
// input.
mlir::OpAsmParser::UnresolvedOperand inputOperand;
mlir::Type inputType;
if (parser.parseOperand(inputOperand) ||
parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() ||
parser.parseType(inputType))
return mlir::failure();
// Resolve the input operand to the type we parsed in.
if (parser.resolveOperand(inputOperand, inputType, result.operands))
return mlir::failure();
return mlir::success();
}
「注,parser 部分比较巧妙,根据实际代码感受true-failure false-success 的效果」
结合 C++ 定义的实现,让我们一起看看这将怎样映射到 声明式的格式。
声明式的格式主要有三个不同的组件构成:
指令(Directives)
一类内置函数,带有可选的一组参数。
字面量(Literals)
一个关键字或者符号,用 ‘’包裹。
变量(Variables)
一个实体,已经通过 operation 自身注册过的,例如,一个参数(属性或者操作数),结果,后继者等。在上述 PrintOp 例子中,一个变量可以是其中的 $input.一个对应 C++ 格式的直接映射是如下这样:
/// Consider a stripped definition of `toy.print` here.
def PrintOp : Toy_Op<"print"> {
let arguments = (ins F64Tensor:$input);
// In the following format we have two directives, `attr-dict` and `type`.
// These correspond to the attribute dictionary and the type of a given
// variable represectively.
let assemblyFormat = "$input attr-dict `:` type($input)";
}
声明格式中还有很多有趣的特性,在实现一个自定义的 C++ 格式之前,务必要先查看理解一下它们。在对我们的 operation 做了一些格式美化之后,我们现在可以得到一个更可读的 toy IR:
module {
toy.func @multiply_transpose(%arg0: tensor<*xf64>, %arg1: tensor<*xf64>) -> tensor<*xf64> {
%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:10)
%1 = toy.transpose(%arg1 : tensor<*xf64>) to tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:25)
%2 = toy.mul %0, %1 : tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:25)
toy.return %2 : tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:3)
} loc("test/Examples/Toy/Ch2/codegen.toy":4:1)
toy.func @main() {
%0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":9:17)
%1 = toy.reshape(%0 : tensor<2x3xf64>) to tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":9:3)
%2 = toy.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64> loc("test/Examples/Toy/Ch2/codegen.toy":10:17)
%3 = toy.reshape(%2 : tensor<6xf64>) to tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":10:3)
%4 = toy.generic_call @multiply_transpose(%1, %3) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":11:11)
%5 = toy.generic_call @multiply_transpose(%3, %1) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":12:11)
toy.print %5 : tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":13:3)
toy.return loc("test/Examples/Toy/Ch2/codegen.toy":8:1)
} loc("test/Examples/Toy/Ch2/codegen.toy":8:1)
} loc(unknown)
上述,我们介绍了在 ODS framework中定义 operations 的几个概念,但是这里还有好几个概念我们没有机会涉及到:regions,variadic 操作数,等等。查看一下 全部的说明可以找到更多细节。
完整的 toy 示例
现在我们可以生成我们的 “Toy IR”。你可以构建 toyc-ch2 然后自己尝试上边的示例:
toyc-ch2 test/Examples/Toy/Ch2/codegen.toy -emit=mlir -mlir-print-debuginfo
我们也可以检查我们的遍历:
toyc-ch2 test/Examples/Toy/Ch2/codegen.toy -emit=mlir -mlir-print-debuginfo 2> codegen.mlir
你也应该在最终的定义文件上使用 mlir-tblgen,并且仔细研究生成的C++代码。
到此为止,MLIR 已经知道了我们的 dialect 和 operations。在下一章中,我们将利用我们的dialect ,为 toy 语言实现一些高级的特定于语言的分析和变换。
注意
「注:根据之前的环境搭建步骤,toyc-ch2 示例中,由 mlir-tblgen 生成于如下文件夹:llvm-project/build_mlir/tools/mlir/examples/toy/Ch2/include/toy/Dialect.cpp.inc
其中 Toy_Dialect 的 td定义:
include "mlir/IR/OpBase.td"
include "mlir/IR/FunctionInterfaces.td"
include "mlir/IR/SymbolInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
// Provide a definition of the 'toy' dialect in the ODS framework so that we
// can define our operations.
def Toy_Dialect : Dialect {
let name = "toy";
let cppNamespace = "::mlir::toy";
let useFoldAPI = kEmitFoldAdaptorFolder;
}
生成指令大体上如此:
inc := -I /home/hipper/ex_mlir/tmp2/llvm-project/mlir/examples/toy/Ch2/include/toy -I/home/hipper/ex_mlir/tmp2/llvm-project/build_mlir/include -I/home/hipper/ex_mlir/tmp2/llvm-project/llvm/include -I/home/hipper/ex_mlir/tmp2/llvm-project/mlir/include -I/home/hipper/ex_mlir/tmp2/llvm-project/build_mlir/tools/mlir/include
$(inc)
input := /home/hipper/ex_mlir/tmp2/llvm-project/mlir/examples/toy/Ch2/include/toy/Ops.td
$(input)
output := tools/mlir/examples/toy/Ch2/include/toy/
$(output)
build_dir := /home/hipper/ex_mlir/tmp2/llvm-project/build_mlir
$(build_dir)
[1/8] $(build_dir)/bin/mlir-tblgen -gen-dialect-decls $(inc) $(input) --write-if-changed -o $(output)/Dialect.h.inc -d $(output)/Dialect.h.inc.d
[2/8] $(build_dir)/bin/mlir-tblgen -gen-op-decls $(inc) $(input) --write-if-changed -o $(output)/Ops.h.inc -d $(output)/Ops.h.inc.d
[3/6] $(build_dir)/bin/mlir-tblgen -gen-dialect-defs $(inc) $(input) --write-if-changed -o $(output)t/Dialect.cpp.inc -d $(output)/Dialect.cpp.inc.d
[4/6] $(build_dir)/bin/mlir-tblgen -gen-op-defs $(inc) $(input) --write-if-changed -o $(output)/Ops.cpp.inc -d $(output)/Ops.cpp.inc.d
生成的 decl 代码如下:(defs 代码在对应的 Dialect.cpp.inc中)
namespace mlir {
namespace toy {
class ToyDialect : public ::mlir::Dialect {
explicit ToyDialect(::mlir::MLIRContext *context);
void initialize();
friend class ::mlir::MLIRContext;
public:
~ToyDialect() override;
static constexpr ::llvm::StringLiteral getDialectNamespace() {
return ::llvm::StringLiteral("toy");
}
};
} // namespace toy
} // namespace mlir
MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::toy::ToyDialect)
」
Chapter1: 生成 toy 语言源程序的 AST
Chapter2: 能够把AST变成SSA的MLIR Dialect IR : toy IR
Chapter3: toy IR 层的优化 opt pass
Chapter4:
Chapter5:
Chapter6:
Chapter7: