0x0. 前言

这篇笔记是阅读Toy Tutorials的第五章之后总结的,这一节主要讲的是将Toy Dialect Lowering的部分Operation Lowering到Affine Dialect,MemRef Dialect和Standard Dialect,而toy.print操作保持不变,所以又被叫作部分Lowering。通过这个Lowering可以将Toy Dialect的Operation更底层的实现逻辑表达出来,以寻求更多的优化机会,得到更好的MLIR表达式。

0x1. Dialect转换

MLIR有众多的Dialect,所以MLIR提供了一个统一的DialectConversion框架来支持这些Dialect互转。这个框架允许将待转换的非法操作转换为合法的操作。为了使用这个框架,我们需要提供两个条件(还有一个可选条件)。

  • 转换目标(Conversation Target)。明确哪些 Dialect 操作是需要合法转换的,不合法的操作需要重写模式(rewrite patterns )来进行合法化。
  • 一组重写模式(Rewrite Pattern)。这是用于将非法操作转换为零个或多个合法操作的一组模式。
  • 类型转换器 (Type Converter)(可选)。如果提供,则用于转换块参数的类型。这一节将不需要此转换。

下面我们来分步介绍具体是如何将目前的MLIR表达式部分Lowering为新的MLIR表达式,并寻求更多的优化机会。

第一步,定义转换目标(Conversion Target)

为了寻求新的优化机会,需要将Toy Dialect中计算密集型操作转换成Affine,MemRef 和 Standard Dialects 的操作组合。代码实现在:mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

void ToyToAffineLoweringPass::runOnFunction() {// 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`, `Arithmetic`, `MemRef`, and `Standard` dialects.target.addLegalDialect<AffineDialect, arith::ArithmeticDialect,memref::MemRefDialect, StandardOpsDialect>();// 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>();target.addDynamicallyLegalOp<toy::PrintOp>([](toy::PrintOp op) {return llvm::none_of(op->getOperandTypes(),[](Type type) { return type.isa<TensorType>(); });});...
}

这个函数实现了将Toy Dialect中的部分计算密集型操作(比如矩阵乘法操作)进行Lowering。我们的目标是要将原始的MLIR表达式Lowering到新的表达式,这里是将原来的计算密集型操作变换为更加靠近底层的操作。首先用mlir::ConversionTarget target(getContext())来定义转换的目标。然后定义在MLIR表达式Lowering的过程中合法的目标,包括特定的Operation或 Dialect,在这个例子中,将MLIR表达式下降到由Affine、MemRef 和 Standard Dialect 三种Dialect下Operation的组合。所以在代码中使用target.addLegalDialect函数将这三个Dialect添加为合法的目标。同时还要定义非法的Dialect,这里定义之前的Toy Dialect为非法的目标,即target.addIllegalDialect<ToyDialect>(),这就暗示如果转换结束MLIR表达式中还存在Toy的操作,则Lowering失败。

一个特殊的点是,在此过程中toy.print在新的目标Dialect中是不支持的,所以这里需要进行保留,不进行Lowering,即target.addLegalOp<PrintOp>()。另外,由于MLIR中单个操作的定义始终优于Dialect的定义,因此上面代码中定义合法,非法以及PrintOp的声明顺序是可随意变的。

第二步,明确转换模式(Conversion Patterns)

在定义了转换目标之后,我们可以定义如何将非法操作转换为合法操作。 与第 3 章介绍的规范化框架类似,DialectConversion 框架也使用 RewritePatterns 来执行转换逻辑。 这些模式可能是之前看到的 RewritePatterns 或特定于转换框架 ConversionPattern 的新型模式。 ConversionPatterns 与传统的 RewritePatterns 不同,因为它们接受一个额外的操作数参数,其中包含已重新映射/替换的操作数。官方文档中给出了Lowering Toy Dialect 中 transpose 操作的例子,代码如下。

/// Lower the `toy.transpose` operation to an affine loop nest.
struct TransposeOpLowering : public mlir::ConversionPattern {TransposeOpLowering(mlir::MLIRContext *ctx): mlir::ConversionPattern(TransposeOp::getOperationName(), 1, ctx) {}/// Match and rewrite the given `toy.transpose` operation, with the given/// operands that have been remapped from `tensor<...>` to `memref<...>`.mlir::LogicalResultmatchAndRewrite(mlir::Operation *op, ArrayRef<mlir::Value> operands,mlir::ConversionPatternRewriter &rewriter) const final {auto loc = op->getLoc();// Call to a helper function that will lower the current operation to a set// of affine loops. We provide a functor that operates on the remapped// operands, as well as the loop induction variables for the inner most// loop body.lowerOpToLoops(op, operands, rewriter,[loc](mlir::PatternRewriter &rewriter,ArrayRef<mlir::Value> memRefOperands,ArrayRef<mlir::Value> loopIvs) {// Generate an adaptor for the remapped operands of the TransposeOp.// This allows for using the nice named accessors that are generated// by the ODS. This adaptor is automatically provided by the ODS// framework.TransposeOpAdaptor transposeAdaptor(memRefOperands);mlir::Value input = transposeAdaptor.input();// Transpose the elements by generating a load from the reverse// indices.SmallVector<mlir::Value, 2> reverseIvs(llvm::reverse(loopIvs));return rewriter.create<mlir::AffineLoadOp>(loc, input, reverseIvs);});return success();}
};

这段代码就是将tor.transpose操作Lowering到Affine loop nest。这里的lowerOpToLoops函数以及TransposeOpAdaptor的定义是看懂这段代码的关键,感兴趣请查看一下源码。

第三步,在Lowering的模式集合中添加上面定义的转换模式

这段代码在mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp中:

void ToyToAffineLoweringPass::runOnFunction() {...// 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());...

这里提供了一个模式集合patterns来Lowering Toy Dialect的操作,然后在这个集合中添加一系列的Lowering模式,如这里的TransposeOpLowering

第四步,执行真正的Lowering过程

明确了Lowering模式之后就可以执行真正的Lowering过程了。DialectConversion框架提供了几种不同的Lowering模式,这里使用的是部分Lowering,因为这里不会对toy.print操作进行Lowering,因为它不是计算密集形的Operation,只是一个打印操作。下面代码中的mlir::applyPartialConversion(function, target, patterns)表示对当前的MLIR表达式中的Operation应用了Lowering。

void ToyToAffineLoweringPass::runOnFunction() {...// 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.auto function = getFunction();if (mlir::failed(mlir::applyPartialConversion(function, target, patterns)))signalPassFailure();
}

第五步,部分Lowering的注意事项

在Lowering过程中,我们从值类型 TensorType 转换为已分配(类似缓冲区)的类型 MemRefType。但是对于toy.print操作,这里不想Lowering,因为这里主要是处理一些计算密集型算子并寻求优化机会,toy.print只是一个有打印功能的算子。因此这一章节不会Lowering这个算子。但我们知道toy.print操作的定义中只支持打印输入类型为F64Tensor的输入数据,所以现在为了能将其和MemRef Dialect联系,我们需要为其增加一个F64MemRef类型。即修改mlir/examples/toy/Ch5/mlir/Ops.tdtoy.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);...
}

第六步,将上面定义好的部分Lowering功能加到优化pipline里面

这部分代码在mlir/examples/toy/Ch5/toyc.cpp中:

/// Create a pass for lowering operations in the `Affine` and `Std` dialects,
/// for a subset of the Toy IR (e.g. matmul).
std::unique_ptr<Pass> mlir::toy::createLowerToAffinePass() {return std::make_unique<ToyToAffineLoweringPass>();
}if (isLoweringToAffine) {mlir::OpPassManager &optPM = pm.nestmlir::FuncOp();// Partially lower the toy dialect with a few cleanups afterwards.optPM.addPass(mlir::toy::createLowerToAffinePass());optPM.addPass(mlir::createCanonicalizerPass());optPM.addPass(mlir::createCSEPass());...
}

这段代码实现了将ToyToAffineLoweringPass加到优化pipline中,当产生MLIR的命令中使用了-emit=mlir-affine选项,则isLoweringToAffine为真,将执行这个部分Lowering的过程。

引入了部分Lowering之后我们可以观察一下输出的MLIR表达式长什么样子。先看一下原始的MLIR表达式,即mlir/test/Examples/Toy/Ch5/affine-lowering.mlir

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
}

然后执行./toyc-ch5 ../../mlir/test/Examples/Toy/Ch5/affine-lowering.mlir -emit=mlir-affine之后就可以查看应用了本节的部分Lowering之后的MLIR表达式了。如下所示:

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 : f64affine.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
}

对其简单解读一下,首先这里有六个f64数据类型的常量,使用%cst开头的变量来表示。然后为输入和输出分配缓冲区,如%0 = memref.alloc() : memref<3x2xf64>这里就是申请了一个类型为memref<3x2xf64>的一段缓冲区。然后将之前声明的六个数据依次存入上述分配的缓冲区中,affine.store指明操作是保存数据操作。之后,第一个循环,将加载的输入数据(数据加载操作affine.load),保存到另一个数据容器中,最终实现转置操作。接着,第二个循环,加载之前定义在两个数据容器中的数据,相乘并存放到输出的数据容器中。最终使用toy.print打印结果,并释放缓冲区。

在Affine Dialect中寻求优化机会

使用了Affine Dialect之后,我们可以将Operation更底层的逻辑展示出来,将代码中的冗余更轻易的暴露出来。这里可以优化的地方为两个循环嵌套的循环边界相同,可以进行循环融合。若在同一个循环中进行处理,减少循环的次数,同时减少多余的数据容器的分配,也减少数据加载的耗时,必然会提高程序的运行效率。

这里可以类比TVM的scheduler中各种循环相关的源语。

这里我们要讲上述提到的循环融合以及减少多余的数据容器的分配和加载的优化加入到pipline中,具体在mlir/examples/toy/Ch5/toyc.cpp中的如下代码:

if (isLoweringToAffine) {mlir::OpPassManager &optPM = pm.nestmlir::FuncOp();// Partially lower the toy dialect with a few cleanups afterwards.optPM.addPass(mlir::toy::createLowerToAffinePass());optPM.addPass(mlir::createCanonicalizerPass());optPM.addPass(mlir::createCSEPass());// Add optimizations if enabled.if (enableOpt) {optPM.addPass(mlir::createLoopFusionPass());optPM.addPass(mlir::createMemRefDataFlowOptPass());}
}

注意到最后一个if条件,添加了createLoopFusionPasscreateMemRefDataFlowOptPass这两种MLIR自带的Pass分别完成了相同循环边界融合优化以及对于MemRef的数据流优化功能。我们比较关心多了这两个优化之后上面的MLIR表达式会变成什么样子,具体来说我们只需要在上面那个生成部分Lowering的MLIR表达式的命令中额外加一个-opt选项就可以生成加入了这两个优化Pass的新的MLIR表达式了。命令如下:./toyc-ch5 ../../mlir/test/Examples/Toy/Ch5/affine-lowering.mlir -emit=mlir-affine -opt 。生成的MLIR表达式为:

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 : f64affine.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
}

我们可以看到,这里删除了多余的数据缓冲区的分配,两个循环嵌套被融合在一起,另外去除了一些不必要的数据加载操作。

所以我们通过部分Lowering确实寻求到了更多的优化机会,使得MLIR表达式的运行效率更高。

参考

  • https://mlir.llvm.org/docs/Tutorials/Toy/Ch-5/
  • https://zhuanlan.zhihu.com/p/362749628

【从零开始学深度学习编译器】十四,MLIR Toy Tutorials学习笔记之部分Lowering相关推荐

  1. 【从零开始学深度学习编译器】十二,MLIR Toy Tutorials学习笔记一

    本笔记由学习MLIR Tutorials总结而成,欢迎批评指正. Chapter1: Toy语言和AST MLIR提供了一种Toy语言来说明MLIR的定义和执行的流程.Toy语言是一种基于张量的语言, ...

  2. 【从零开始学深度学习编译器】十五,MLIR Toy Tutorials学习笔记之Lowering到LLVM IR

    0x0. 前言 在上一节中,我们将Toy Dialect的部分Operation Lowering到Affine Dialect,MemRef Dialect和Standard Dialect,而to ...

  3. 从零开始学USB(二十四、USB接口HID类设备(六)_物理描述符)

    物理描述符也称作实体描述符.被用来描述设备的行为特性. 物理描述符是一种数据结构,它提供有关激活控件或控件的人体特定部分的信息. 例如,物理描述符可能指示右手拇指用于激活按钮5.应用程序可以使用此信息 ...

  4. 从零开始学ios开发(十四):Navigation Controllers and Table Views(上)

    这一篇我们将学习一个新的控件Navigation Controller,很多时候Navigation Controller是和Table View紧密结合在一起的,因此在学习Navigation Co ...

  5. 从零开始学USB(二十四、最简单的USB驱动开始)

    18年底,花了2个月啃学习USB2.0协议.现在终于开始要派上用场了,开始和代码打交道了. 当然,现在只是从简单的开始usb设备驱动开始,后面再慢慢的分析usb总线,usb的hub之类. 先给出代码, ...

  6. 【从零开始学深度学习编译器】十三,如何在MLIR里面写Pass?

    [GiantPandaCV导语]这篇文章是学习了比较久然后按照自己的理解步骤重新总结了下来,主要是MLIR Toy Tutorials第3,4篇文章的内容.这里主要讲解了如何在MLIR中自定义Pass ...

  7. 从零开始编写深度学习库(四)Eigen::Tensor学习使用及代码重构

    从零开始编写深度学习库(四)Eigen::Tensor学习使用及代码重构 博客:http://blog.csdn.net/hjimce 微博:黄锦池-hjimce   qq:1393852684 一. ...

  8. 花书+吴恩达深度学习(十四)卷积神经网络 CNN 之经典案例(LetNet-5, AlexNet, VGG-16, ResNet, Inception Network)

    目录 0. 前言 1. LeNet-5 2. AlexNet 3. VGG-16 4. ResNet 残差网络 5. Inception Network 如果这篇文章对你有一点小小的帮助,请给个关注, ...

  9. 华为鸿蒙系统HarmonyOS学习之十四:方舟编译器

    华为鸿蒙系统HarmonyOS学习之十四:方舟编译器 方舟编译器是华为自研的一个支持多种编程语言,多种芯片平台的联合编译编程平台,而经过方舟编译器编译适配后的APP,运行效率会大大提高,拥有更为流畅的 ...

最新文章

  1. (0001) iOS 开发之收集第三方资源篇
  2. Spring Framework 5.3.6、5.2.14 发布
  3. 一位技术人员成长的烦恼及我的分析(转)
  4. mmz-asio4delphi死链接的解决办法
  5. Tensroflow随笔-测试集
  6. 20145308刘昊阳 《Java程序设计》实验五报告
  7. 使用.Net6中的System.Text.Json遇到几个常见问题及解决方案
  8. JS_Java_Oracle_MySQL截取字符串的函数
  9. mybatis 配置 mysql连接池_spring 5.x 系列第5篇 —— 整合 mybatis + druid 连接池 (xml配置方式)...
  10. Oracle 数据怎么实时同步到 Kafka | 亲测干货分享建议收藏
  11. html5,用或不用,它都在那里
  12. oracle数据库导入gson包
  13. css3动画由浅入深总结
  14. c语言设计一个酒店管理系统,C语言酒店管理系统设计.docx
  15. 面向接口编程与面向实现编程
  16. 自己编写的android汉字转拼音类(超全字库)
  17. 恩智浦智能车竞赛电磁直立组国一之速墨二队2018回顾与总结
  18. 错误1919配置ODBC数据源时发生错误Visio database samples,ODBC错误6:部件在注册
  19. SCI 投稿全过程信件模板一览
  20. 阿里云道哥:我曾经是饱受网络攻击摧残的创业者,现在我要帮他们顶住

热门文章

  1. Django应用与分布式路由
  2. 网站内容快速填充的技巧和方法
  3. c语言str系列函数
  4. python计算excel平均值_python计算excel平均值和标准差
  5. fastjson不序列化空属性,解决办法
  6. NIUSHOP wap端底部导航设置
  7. App Extension
  8. 给深度学习入门者的Python快速教程 - numpy和Matplotlib篇
  9. [Eclipse]GEF入门系列(二、GEF概述)
  10. 当元宇宙时代来临,才真正让这些新技术跳出了互联网的牵绊