本笔记由学习MLIR Tutorials总结而成,欢迎批评指正。
Chapter1: Toy语言和AST
MLIR提供了一种Toy语言来说明MLIR的定义和执行的流程。Toy语言是一种基于张量的语言,我们可以使用它来定义函数,执行一些数学计算以及输出结果。下面要介绍的例子中限制Tensor的维度是<=2的,并且Toy语言中唯一的数据类型是64位浮点类型,对应C语言中的"double"。另外Values
是不可以重写的,即每个操作都会返回一个新分配的值,并自动管理释放。直接看下面这个例子:
def main() {
# Define a variable `a` with shape <2, 3>, initialized with the literal value.
# The shape is inferred from the supplied literal.
var a = [[1, 2, 3], [4, 5, 6]];
# b is identical to a, the literal tensor is implicitly reshaped: defining new
# variables is the way to reshape tensors (element count must match).
var b<2, 3> = [1, 2, 3, 4, 5, 6];
# transpose() and print() are the only builtin, the following will transpose
# a and b and perform an element-wise multiplication before printing the result.
print(transpose(a) * transpose(b));
}
类型检查是通过类型推断静态执行的。Toy语言只需在必要时指定Tensor形状的类型声明。下面定义了一个multiply_transpose
函数,注意这个函数里面参数a
和b
的形状我们预先都是不知道的,只有调用这个函数时我们才知道,可以关注一下下面例子中的shape变化。
# User defined generic function that operates on unknown shaped arguments.
def multiply_transpose(a, b) {
return transpose(a) * transpose(b);
}
def main() {
# Define a variable `a` with shape <2, 3>, initialized with the literal value.
var a = [[1, 2, 3], [4, 5, 6]];
var b<2, 3> = [1, 2, 3, 4, 5, 6];
# This call will specialize `multiply_transpose` with <2, 3> for both
# arguments and deduce a return type of <3, 2> in initialization of `c`.
var c = multiply_transpose(a, b);
# A second call to `multiply_transpose` with <2, 3> for both arguments will
# reuse the previously specialized and inferred version and return <3, 2>.
var d = multiply_transpose(b, a);
# A new call with <3, 2> (instead of <2, 3>) for both dimensions will
# trigger another specialization of `multiply_transpose`.
var e = multiply_transpose(b, c);
# Finally, calling into `multiply_transpose` with incompatible shape will
# trigger a shape inference error.
var f = multiply_transpose(transpose(a), c);
}
然后我们可以使用下面的命令来产生这个Toy语言程序的AST:
cd llvm-project/build/bin
./toyc-ch1 ../../mlir/test/Examples/Toy/Ch1/ast.toy --emit=ast
前提是要构建好llvm-project工程,构建过程按照https://mlir.llvm.org/getting_started/ 这里的方法操作即可,这里再列一下完整过程:
$ git clone https://github.com/llvm/llvm-project.git
$ mkdir llvm-project/build
$ cd llvm-project/build
$ cmake -G "Unix Makefiles" ../llvm \
-DLLVM_ENABLE_PROJECTS=mlir \
-DLLVM_BUILD_EXAMPLES=ON \
-DLLVM_TARGETS_TO_BUILD="host" \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_ENABLE_ASSERTIONS=ON
$ cmake --build . --target check-mlir
上面Toy程序产生的AST长下面这样:
Module:
Function
Proto 'multiply_transpose' @../../mlir/test/Examples/Toy/Ch1/ast.toy:4:1
Params: [a, b]
Block {
Return
BinOp: * @../../mlir/test/Examples/Toy/Ch1/ast.to