MLIR Operation Pass机制
🤖

MLIR Operation Pass机制

Tags
MLIR
Pass
PassManager
Created time
Jul 28, 2023 06:29 AM

概况

无论什么pass,本质都是在对IR进行变换。不同的变换关注的层面不一样,有的pass是在站在module level,有的是站在function level。
mlir的pass机制提供了对不同level的pass的管理。正确的编写方式,可以让MLIR自动调度不同level的pass,甚至MLIR还可以多线程调度。

Pass与PassManager的使用简述

所有pass都是由PassManager进行管理。通过接口 addPass 将pass加入PassManager。为了管理不同level的pass,PassManager也分了不同level。
class MPass : public mlir::PassWrapper<MPass, OperationPass<ModuleOp>; mlir::MLIRContext context; mlir::PassManager modulePm(&context); // 默认是moduleOp level modulePm.addPass(std::make_unique<MPass>());
module level的manager无法直接添加其它level的pass。但是pass manager是可以嵌套的。
class FPass : public mlir::PassWrapper<FPass, OperationPass<FuncOp>; auto& funcPm = modulePm.nest<FuncOp>(); funcPm.addPass(std::make_unique<FPass>());
假如有下面这样的IR结构:
module0 func0 func1
有如下的pass:
mlir::PassManager modulePm(&context); modulePm.addPass(std::make_unique<MPass1>()); auto& funcPm = modulePm.nest<FuncOp>(); funcPm.addPass(stD::make_unique<FPass1>()); funcPm.addPass(stD::make_unique<FPass2>()); modulePm.addPass(std::make_unique<MPass2>());
实际在modulePm的内部,pass结构是如下:
modulePm MPass1 funcPm //为什么这里会有funcPm,后面会讲 FPass1 FPass2 MPass2
pass执行流程:
MPass1::runOnOperation(module0) FPass1::runOnOperation(func0) FPass2::runOnOperation(func0) FPass1::runOnOperation(func1) FPass2::runOnOperation(func1) MPass2::runOnOperation(module0)

PassManager的嵌套机制解析

PassManager::nest<>

当我们在调用 modulePm.nest<FuncOp>()的时候,实际上是创建了一个新的pass, 这个pass的类型是 OpToOpPassAdaptor,这个pass里面包含了一个对应level的Pass Manager,也就是最后我们得到的那个funcPm。
notion image
简单看下源码,可以看见这里把创建好的adaptor pass加入到了上层的PassManager,也就是modulePm。
这个adaptor pass的作用节,执行一系列对应level(这里也就是func level)的pass。
 

OpToOpPassAdaptor

这里看一下adaptor pass的runOnOperation实现。
notion image
可以看见这里会遍历当前op的所有block的子op,检查看是否和这个adaptor pass保护的manger是同一个level(对于上面的情况,也就是func level),如果是则执行runPipeline函数。再看一下runPipeline的实现:
notion image

PassManager::addNestedPass<>

notion image
可以看见这个接口只是把nest和addPass这两个接口放在了一起而已。但是值得注意的的是,这个接口没调用一次都会创建一个新的adaptor pass,例如下面的code:
modulePm.addNestedPass<FuncOp>(std::make_unique<FPass1>()); modulePm.addNestedPass<FuncOp>(std::make_unique<FPass2>());
在这里add完之后,modulePm内部的pass结构为:
modulePm funcOpAdaptorPass1 FPass1 funcOpAdaptorPass2 FPass2
这样的Pass结构会影响analysis cache的有效性,不过好在mlir自己已经有优化,会自动将上面的两个adaptor pass合并为一个。这是优化的源码,就不做过多解释。
notion image

如何启用多线程runPass

一句话来讲解释这些op需要带有IsolatedFromAbove trait。
比如你希望FPass1能同时在多个线程,在不同的func op上执行,那么:
  • 你的FPass1的level需要是func level的
满足这两个条件,mlir多线程的机制就会自动enable,你如果满足这两个条件,但不希望开启多线程,可以通过MLIRContext::enableMultithreading(false) 关闭。
 

TODO

Pass和PassManger的具体实现设计、Analysis的机制、Operation::dump的一个性能小坑。Attribute实现,trait实现, dyn_cast…

结语

MLIR的代码读起来真舒服!