MLIR Operation Pass机制

概况

无论什么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)

注意：对于同一level的op ，会先在第一个op上执行完所有pass后，才会在第二个op上执行pass。

PassManager的嵌套机制解析

PassManager::nest<>

当我们在调用 modulePm.nest<FuncOp>()的时候，实际上是创建了一个新的pass，这个pass的类型是 OpToOpPassAdaptor，这个pass里面包含了一个对应level的Pass Manager，也就是最后我们得到的那个funcPm。

简单看下源码，可以看见这里把创建好的adaptor pass加入到了上层的PassManager，也就是modulePm。

这个adaptor pass的作用节，执行一系列对应level（这里也就是func level）的pass。

OpToOpPassAdaptor

这里看一下adaptor pass的runOnOperation实现。

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

可以看见在 runPipeline,会对当前op执行所有的pass，这里应证了上面提到的这个特性。

PassManager::addNestedPass<>

可以看见这个接口只是把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合并为一个。这是优化的源码，就不做过多解释。

如何启用多线程runPass

一句话来讲就是: 这些op需要带有IsolatedFromAbove trait。

比如你希望FPass1能同时在多个线程，在不同的func op上执行，那么:

你的FPass1的level需要是func level的

func op需要带有IsolatedFromAbove的trait

满足这两个条件，mlir多线程的机制就会自动enable，你如果满足这两个条件，但不希望开启多线程，可以通过MLIRContext::enableMultithreading(false) 关闭。

TODO

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

结语

MLIR的代码读起来真舒服！