以OneFlow为例探索MLIR的实际开发流程

撰文 | BBuf

原文首发于GiantPandaCV

目录

1、前言

2、OneFlow是如何和MLIR结合的？

3、OneFlow IR如何执行？

         4、总结

---

1、前言

最近在同事shenghang的帮助下做了一点OneFlow IR相关的开发，对MLIR执行部分有一些新的感受，所以尝试分享一下。我之前花了不少时间去理解OneFlow IR的整个架构（可以看我的Toy Tutorials系列），但对OneFloiw IR的JIT的执行这部分一直存疑。最近将OneFlow基于Job（OneFlow的作业函数，不考虑设备的话可以理解为一个计算图）接入MLIR工程实现部分重新进行了梳理，并在shenghang的指导下理解了整个流程。

所以这篇文档我将介绍一下OneFlow和MLIR是如何结合的，如何在OneFlow IR中新增一个图级别的Pass，OneFlow的Operation是如何自动变成MLIR 的Operation的以及为什么OneFlow IR能利用MLIR为计算带来加速等。我对MLIR的了解不算多，2个月前开始接触，有任何错误请大家批评斧正。

本文和 https://github.com/Oneflow-Inc/oneflow & https://github.com/BBuf/tvm_mlir_learn 有关，感兴趣可以star关注一下。

本文提到的Op和Operation是一回事，没有严格区分。

2、OneFlow是如何和MLIR结合的？

在OneFlow中引入MLIR作为OneFlow的IR有诸多优点，不仅可以取代OneFlow中需要通过C++手写的Operation定义减小开发难度，还可以降低Operation定义中一些容器相关的开销。另外我们还可以通过MLIR维护的基础设施（即多重Dialect）来完成对计算图计算的加速。

这里的计算图既可以是Eager的计算图，也可以是Lazy的计算图。由于基于Eager计算图使用MLIR进行加速的工作（即oneflow.jit.xxx）还没有正式开放，我这里仍然以Lazy计算图（Job）为例来讲解OneFlow和MLIR的结合过程。

首先我们需要编译好开启MLIR的OneFlow，编译命令如下：

git clone git@github.com:Oneflow-Inc/oneflow.git
cd oneflow && mkdir build && cd build
cmake-C ../cmake/caches/cn/fast/mlir-cuda-75.cmake -DBUILD_TESTING=ON .. && ninja

然后可以写一个例子进行测试：

os.environ["ONEFLOW_MLIR_ENABLE_ROUND_TRIP"] = '1'
os.environ["ONEFLOW_MLIR_ENABLE_CODEGEN_FUSERS"] = '1'

@flow.unittest.skip_unless_1n1d()
class TestFuseBiasAddGeLUCPUMLIR(oneflow.unittest.TestCase):
    def test_fused_bias_add_gelu_graph(test_case):
        data = np.random.randn(1, 2, 3)
        bias_data = np.random.randn(2)
        x = flow.tensor(data, dtype=flow.float32)
        bias = flow.tensor(bias_data, dtype=flow.float32)
        y_eager = flow.gelu(flow._C.bias_add(x, bias, axis=1))

        class FuseBiasAddGeLUGraph(flow.nn.Graph):
            def __init__(self):
                super().__init__()

            def build(self, x):
                return flow.gelu(flow._C.bias_add(x, bias, axis=1))

        bias_add_gelu = FuseBiasAddGeLUGraph()
        y_lazy = bias_add_gelu(x)
        test_case.assertTrue(np.array_equal(y_eager.numpy(), y_lazy.numpy()))

运行这个例子之后会在当前运行目录下生成一个log文件，里面有一个ir_pass 文件夹记录了经过OneFlow MLIR优化前后的计算图(.prototxt) 以及 MLIR的表达式(*.mlir)，还有一个*.mlir.dot文件可以用graphviz打开来可视化MLIR表达式的计算图。

需要注意的是，如果OneFlow正在执行训练任务，这个log文件夹里不仅包含前向的计算图和MLIR表达式，也会生成后向的计算图和MLIR表达式。所以MLIR在整个神经网络