Homecompile

# 前言

现在开启 MLIR 学习系列。本篇是跟着 Toy 语言学习 MLIR 的第四篇，主要介绍高层级 toy dialect 降级到低层方言的过程，前述内容请参考【MLIR】跟着 Toy 语言学习 MLIR【1】Toy 语言和 Toy Dialect，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换。
学习 MLIR 最好的方式还是照着官方教程来，但是只看 MLIR 官网教程，Ch1 到 Ch6，没有基础的话又让人学着有点吃力。现发现好多细节，不看教程认识不到，所以当前内容跟随官网教程。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程，【编译器】使用 llvm 编译自定义语言【1】构建 AST ，【MLIR】跟着 Toy 语言学习 MLIR【2】pattern 匹配和重写，【MLIR】跟着 Toy 语言学习 MLIR【3】通过接口实现通用转换。
作为初学者，错误在所难免，还望不吝赐教。

# 基本简介

官网 Toy 教程 展示了如何将自定义语言 Toy 借助 MLIR 一步步编译为可执行机器码的过程。Toy 是一种简单的自定义语言，为了简便，其所有数据类型定义为 fp64 类型的 Tensor，支持 +/* 操作和 transpose 等有限的操作。
以下是整个编译降级的过程：
Toy txt -> Toy AST -> Toy Dialect -> Affine Dialect -> llvm Dialect -> llvm IR -> 机器码（通过 JIT 编译）
编译流程图

# 本章内容

前述章节中我们已经将 toy 语言转换成 toy dialect，甚至还添加了 Operation 匹配重写、inline 和 inferShape 接口等，对 toy dialect 进行优化，现在我们迫切地将 toy dialect 转换成实际的代码，最终转换成 llvm，让其投入使用。但本章不会直接将其转换为 llvm，而是将其进行逐步降级 (渐进式降级)，这样我们就可以重用 MLIR 提供的针对中间层级方言的优化技术。
Affine 调方言专为程序中计算密集的部分设计，其功能有限：例如，它不支持表示 Toy 中的 print 内置函数，也不应支持！因此我们可以将 toy dialect 密集计算部分降级到 affine dialect 并进行优化，而 print 内置函数保留原来的 toy dialect。

# 方言转换

MLIR 拥有多种不同的方言，因此有必要建立一个统一的框架来实现它们之间的转换。这时，DialectConversion 框架就派上了用场。DialectConversion 框架提供了一个标准化的、可扩展的 “翻译工具箱”，用于将一个方言（Dialect）的操作降级（Lowering）或转换到另一个方言。
使用这个框架需要提供两个必要内容和一个可选内容。
Conversion Target 转换目标：ConversionTarget 定义了 “什么是合法的”，什么是非法的，然后必须要将非法操作替换为合法操作。
Rewrite Patterns 一组翻译规则 ：RewritePattern 定义了具体的转换逻辑。比如它告诉框架：“当你遇到一个 toy.add 操作时，请按照以下方式将其替换为一组 arith 和 affine 操作。”
Type Converter 跨方言的 “类型翻译” （可选）： 当转换涉及类型变化时（例如从 tensor<...> 到 memref<...>），就需要 TypeConverter 。

# Conversion Target

我们希望将计算密集型的 Toy 操作转换为 Affine、Arith、Func 和 MemRef 方言的组合操作，以便进一步优化。所以首先定义转换目标：

void ToyToAffineLoweringPass::runOnOperation() {
  // The first thing to define is the conversion target. This will define the
  // final target for this lowering.
  mlir::ConversionTarget target(getContext());
  // We define the specific operations, or dialects, that are legal targets for
  // this lowering. In our case, we are lowering to a combination of the
  // `Affine`, `Arith`, `Func`, and `MemRef` dialects.
  target.addLegalDialect<affine::AffineDialect, arith::ArithDialect,
                         func::FuncDialect, memref::MemRefDialect>();  // 将目标方言设置为合法
  // We also define the Toy dialect as Illegal so that the conversion will fail
  // if any of these operations are *not* converted. Given that we actually want
  // a partial lowering, we explicitly mark the Toy operations that don't want
  // to lower, `toy.print`, as *legal*. `toy.print` will still need its operands
  // to be updated though (as we convert from TensorType to MemRefType), so we
  // only treat it as `legal` if its operands are legal.
  target.addIllegalDialect<ToyDialect>();  // 将 toy dialect 设置为非法
  target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op) {  // 将 toy 中的 print 设置为合法
    return llvm::none_of(op->getOperandTypes(), llvm::IsaPred<TensorType>);
  });
  ...
}
我们通过 Pass 来完成方言转换操作， ToyToAffineLoweringPass  是一个 Pass，它重写了 runOnOperation  方法。在前述章节中我们提到过 Pass 需要重写的核心方法就是  runOnOperation 。

在这个方法中指明了哪些方言合法，哪些方言非法，并特别指定了 toy.print 这个操作合法，因为我们没法将其转换到 Affine。


# Conversion Patterns


前述章节方言内的优化（如 transpose (transpose (x)) -> x）的时候，我们使用过  RewritePattern ，它是方言内部的匹配和重写，现在我们需要跨方言的转换  ConversionPattern  。

和 RewritePattern  不同的是， ConversionPattern  接收一个额外的参数（OpAdaptor），该参数包含已重新映射或替换的操作数。

意思就是说，我们将操作（例如 toy.transpose）降级到 Affine 之后，需要将对应的操作数也从 TensorType  降级到  MemRefType  类型，因为 affine 方言的所有内存操作（如 affine.load、affine.store）都要求操作数是 memref 类型。而 OpAdaptor 就是用来提供操作数的类型映射。


/// Lower the `toy.transpose` operation to an affine loop nest.
struct TransposeOpLowering : public OpConversionPattern<toy::TransposeOp> {  // 这次继承的是 OpConversionPattern
  using OpConversionPattern<toy::TransposeOp>::OpConversionPattern;
  LogicalResult  matchAndRewrite(toy::TransposeOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    auto loc = op->getLoc();
    lowerOpToLoops(op, rewriter,
                   [&](OpBuilder &builder, ValueRange loopIvs) {
                     Value input = adaptor.getInput();
                     // Transpose the elements by generating a load from the
                     // reverse indices.
                     SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
                     return affine::AffineLoadOp::create(builder, loc, input,
                                                         reverseIvs);
                   });
    return success();
  }
};
这个模式匹配重写继承的是  OpConversionPattern ，而非之前 transpose(transpose(x)) -> x  时的 OpRewritePattern ，同样是重写 matchAndRewrite  函数，但它多了 适配器  OpAdaptor 。而且我们发现，这个  OpAdaptor  是作为参数传入的，并不需要我们手动实现，实际上， DialectConversion  框架在调用你的  matchAndRewrite  方法时，会自动创建一个  OpAdaptor  实例，并传递当前操作的实际操作数。

接下来就是将这些操作转换匹配全部放在刚才的 Pass 函数里面：


void ToyToAffineLoweringPass::runOnOperation() {
  ...
  // Now that the conversion target has been defined, we just need to provide
  // the set of patterns that will lower the Toy operations.
  mlir::RewritePatternSet patterns(&getContext());
  patterns.add<..., TransposeOpLowering>(&getContext());  // 将 操作转换 Pattern 添加进来
  ...
# Partial Lowering


部分降级。 DialectConversion  框架提供了多种不同的降级模式，这里选择部分降级，因为还有部分（toy.print）不会降级：


void ToyToAffineLoweringPass::runOnOperation() {
  ...
  // With the target and rewrite patterns defined, we can now attempt the
  // conversion. The conversion will signal failure if any of our *illegal*
  // operations were not converted successfully.
  if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, patterns)))
    signalPassFailure();
}
一个 toy dialect 降级的例子，降级之前：


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>
  %2 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
  %3 = toy.mul %2, %2 : tensor<3x2xf64>
  toy.print %3 : tensor<3x2xf64>
  toy.return
}
降级之后：


func.func @main() {
  %cst = arith.constant 1.000000e+00 : f64
  %cst_0 = arith.constant 2.000000e+00 : f64
  %cst_1 = arith.constant 3.000000e+00 : f64
  %cst_2 = arith.constant 4.000000e+00 : f64
  %cst_3 = arith.constant 5.000000e+00 : f64
  %cst_4 = arith.constant 6.000000e+00 : f64
  // Allocating buffers for the inputs and outputs.
  %0 = memref.alloc() : memref<3x2xf64>
  %1 = memref.alloc() : memref<3x2xf64>
  %2 = memref.alloc() : memref<2x3xf64>
  // Initialize the input buffer with the constant values.
  affine.store %cst, %2[0, 0] : memref<2x3xf64>
  affine.store %cst_0, %2[0, 1] : memref<2x3xf64>
  affine.store %cst_1, %2[0, 2] : memref<2x3xf64>
  affine.store %cst_2, %2[1, 0] : memref<2x3xf64>
  affine.store %cst_3, %2[1, 1] : memref<2x3xf64>
  affine.store %cst_4, %2[1, 2] : memref<2x3xf64>
  // Load the transpose value from the input buffer and store it into the
  // next input buffer.
  affine.for %arg0 = 0 to 3 {
    affine.for %arg1 = 0 to 2 {
      %3 = affine.load %2[%arg1, %arg0] : memref<2x3xf64>
      affine.store %3, %1[%arg0, %arg1] : memref<3x2xf64>
    }
  }
  // Multiply and store into the output buffer.
  affine.for %arg0 = 0 to 3 {
    affine.for %arg1 = 0 to 2 {
      %3 = affine.load %1[%arg0, %arg1] : memref<3x2xf64>
      %4 = affine.load %1[%arg0, %arg1] : memref<3x2xf64>
      %5 = arith.mulf %3, %4 : f64
      affine.store %5, %0[%arg0, %arg1] : memref<3x2xf64>
    }
  }
  // Print the value held by the buffer.
  toy.print %0 : memref<3x2xf64>
  memref.dealloc %2 : memref<2x3xf64>
  memref.dealloc %1 : memref<3x2xf64>
  memref.dealloc %0 : memref<3x2xf64>
  return
}
# 部分降级的问题


当一个操作（如 toy.print）没有被完全降级时，如何让它同时支持高层的 tensor 和低层的 memref？

在降级过程中，大部分 toy 操作（如 add、transpose）被转换成了 affine 循环和 memref。但 toy.print 没有被降级，它仍然是一个 toy 方言的操作，此时它的输入出现了矛盾：toy.print 的定义只接受 tensor，但降级后的数据是 memref。

教程提出了三种方案：

1. 从 memref 生成 tensor（隐式复制）


// 降级后的 IR（方案 1）
%memref = ... : memref<2x3xf64>   // 实际数据在 memref 中
%tensor = "toy.materialize"(%memref) : memref<2x3xf64> -> tensor<2x3xf64>  // 隐式复制
toy.print %tensor : tensor<2x3xf64>   // 打印张量
// 优点：
//toy.print 定义不用改。
// 缺点：增加了一次隐式复制，可能隐藏优化机会。
2. 创建 toy.print_memref（新增操作）


// 降级后的 IR（方案 2）
%memref = ... : memref<2x3xf64>
toy.print_memref %memref : memref<2x3xf64>   // 新的打印操作
3. 让 toy.print 同时支持 tensor 和 memref


// 降级后的 IR（方案 3）
%memref = ... : memref<2x3xf64>
toy.print %memref : memref<2x3xf64>   // 直接打印 memref
// 优点：简单，没有隐式复制，也没有重复定义。
// 缺点：toy.print 现在混合了两种抽象层级（tensor 和 memref）。
教程选择了第三种方案，即定义 Print 操作的时候，让其支持两种类型的操作数。


def PrintOp : Toy_Op<"print"> {
  ...
  // The print operation takes an input tensor to print.
  // We also allow a F64MemRef to enable interop during partial lowering.
  let arguments = (ins AnyTypeOf<[F64Tensor, F64MemRef]>:$input);
}
AnyTypeOf<[F64Tensor, F64MemRef]>  表示  PrintOp  的输入可以是  F64Tensor  或  F64MemRef 。


这样，toy.print 在降级前可以接受 tensor，在降级后可以接受 memref，不需要任何额外操作。


# Affine 优化


上述已经将 toy dialect 降级到 Affine dialect，它是正确的，但在效率方面仍有很大提升空间。例如，toy.mul 的降级产生了一些冗余的加载操作。而 MLIR 为已有的 Affine dialect 提供了好用的优化 Pass。我们只需要将 LoopFusion 和 AffineScalarReplacement 通过 passes 添加到流水线后，就能对当前的方言进行优化：


if (isLoweringToAffine) {
    // Partially lower the toy dialect.
    pm.addPass(mlir::toy::createLowerToAffinePass());
    // Add a few cleanups post lowering.
    mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
    optPM.addPass(mlir::createCanonicalizerPass());
    optPM.addPass(mlir::createCSEPass());
    // Add optimizations if enabled.
    if (enableOpt) {
      optPM.addPass(mlir::affine::createLoopFusionPass());  // 添加循环融合 pass
      optPM.addPass(mlir::affine::createAffineScalarReplacementPass());
  // 添加标量替换 pass
    }
  }
LoopFusion  是循环融合 pass，能够将相邻的、有相同循环边界的循环融合为一个循环，减少循环开销，提升数据局部性。

 AffineScalarReplacement  检测临时分配的 memref 是否可以被标量（寄存器）替换，从而减少内存访问。


func.func @main() {
  %cst = arith.constant 1.000000e+00 : f64
  %cst_0 = arith.constant 2.000000e+00 : f64
  %cst_1 = arith.constant 3.000000e+00 : f64
  %cst_2 = arith.constant 4.000000e+00 : f64
  %cst_3 = arith.constant 5.000000e+00 : f64
  %cst_4 = arith.constant 6.000000e+00 : f64
  // Allocating buffers for the inputs and outputs.
  %0 = memref.alloc() : memref<3x2xf64>
  %1 = memref.alloc() : memref<2x3xf64>
  // Initialize the input buffer with the constant values.
  affine.store %cst, %1[0, 0] : memref<2x3xf64>
  affine.store %cst_0, %1[0, 1] : memref<2x3xf64>
  affine.store %cst_1, %1[0, 2] : memref<2x3xf64>
  affine.store %cst_2, %1[1, 0] : memref<2x3xf64>
  affine.store %cst_3, %1[1, 1] : memref<2x3xf64>
  affine.store %cst_4, %1[1, 2] : memref<2x3xf64>
  affine.for %arg0 = 0 to 3 {
    affine.for %arg1 = 0 to 2 {
      // Load the transpose value from the input buffer.
      %2 = affine.load %1[%arg1, %arg0] : memref<2x3xf64>
      // Multiply and store into the output buffer.
      %3 = arith.mulf %2, %2 : f64
      affine.store %3, %0[%arg0, %arg1] : memref<3x2xf64>
    }
  }
  // Print the value held by the buffer.
  toy.print %0 : memref<3x2xf64>
  memref.dealloc %1 : memref<2x3xf64>
  memref.dealloc %0 : memref<3x2xf64>
  return
}
可以通过以下指令，以及是否添加  -opt  来检查优化前后的区别。


toyc-ch5 test/Examples/Toy/Ch5/affine-lowering.mlir -emit=mlir-affine
# 后记


I know someone in the world is waiting for me, although I've no idea of who he is. But I feel happy every day for this. 
                                          -  《返老还童》
本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目 或随便一个项目下提出 issue，并指明哪一篇博客，看到一定及时回复！


Edited on
Give me a cup of [coffee]~(￣▽￣)~*
XianMu WeChat Pay
WeChat Pay
XianMu Alipay
Alipay