MLIR Pass机制 | 实在之路

type

status

date

slug

summary

概况

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

module level的manager无法直接添加其它level的pass。但是pass manager是可以嵌套的。

假如有下面这样的IR结构:

有如下的pass：

实际在modulePm的内部，pass结构是如下：

pass执行流程：

注意：对于同一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：

在这里add完之后，modulePm内部的pass结构为：

这样的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的代码读起来真舒服！

🤖MLIR Pass机制

概况