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

你是不是有时候也在想：“机器学习框架是怎么优化计算的？MLIR到底是个啥？”别急，今天咱们就好好聊聊OneFlow和MLIR这对CP是怎么“相爱相杀”的！

首先，OneFlow为什么看上MLIR这位“驸马”呢？ 原来是想通过MLIR这把瑞士军刀，把繁琐的手动写算子的工作给省了。而且，MLIR自带的多重方言（比如Linalg、Tosa）就像瑞士军刀的各个刀片，帮助OneFlow轻松实现加速。比如在图像处理里，Reshape、Conv2D这些操作都能被MLIR优雅地转化成高效的LLVM IR。

那具体怎么“联姻”呢？就是把OneFlow的Job（计算图）和MLIR互转。你可能好奇：“这俩怎么对话的？”秘诀就是注册Pass来来回回转换，比如IRRoundTripBeforeAD和IRRoundTrip，就像翻译官一样，确保双方都能听懂彼此。

运行起来又是怎么回事呢？举个栗子，假如你写了一个前向传播的计算图，OneFlow会先把它转成MLIR，再一步步降到Linalg、Tosa，最后变成LLVM IR执行。这就像是把原始绘画转成草图，再细化成素描，最后上色完成一幅画。整个过程流畅得像一首交响乐！

想试试效果？简单！编译OneFlow时加上MLIR选项，就能看到log文件里长出了优化后的计算图和MLIR表达式。甚至还能用Graphviz看图，Debug起来不要太爽！

总之，OneFlow和MLIR这对CP，把加速计算的事儿搞得明明白白的！想了解更多？快来OneFlow和BBuf的学习仓库看看吧，说不定你也会成为下一位优化大师！😄

撰文 | 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.gitcd oneflow && mkdir build && cd buildcmake-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在整个神经网络的运行流程中均可以作用，这是区别于前向推理框架的重要一点，即训练也可以加速。

在oneflow/api/python/ir.cpp 中有下面两行代码：

REGISTER_JOB_PASS("IRRoundTripBeforeAD", IRRoundTrip);REGISTER_JOB_PASS("IRRoundTrip", IRRoundTrip);

RoundTrip即往返的意思，BeforeAD可以理解为反向之前，kAfterAD 可以理解为反向之后，这里通过将OneFlow Job和MLIR的互转过程注册为OneFlow Job的一个Pass来建立OneFlow计算图和MLIR的联系。在执行OneFlow脚本时，如果想使能MLIR作用于OneFlow计算图，开启ONEFLOW_MLIR_ENABLE_ROUND_TRIP=1环境变量即可。

接下来，要将OneFlow的计算图和MLIR建立联系等价于将OneFlow计算图中的Operation和MLIR中的Operation进行一对一的转换。而MLIR的Operation定义在各级Dialect下，按照MLIR的通用接入原则，我们实现了一个OneFlow Dialect并在OneFlow Dialect上实现了OneFlow Operation到OneFlow Dialect下的Operation的一一映射。

如何定义OneFlow Dialect和Operation这里就不讲了，可以参考MLIR官方文档的Dialects和ODS一节（https://mlir.llvm.org/docs/OpDefinitions/）或者我之前的文章，它们都是基于TableGen规则来完成的。关于MLIR Operation的定义我之前结合OneFlow Dialect的Op定义总结了一个文档（https://github.com/BBuf/tvm_mlir_learn 中） 。

除了Dialect和Operation的定义还有一些其它需要定义的东西，比如OneFlow数据类型到MLIR数据类型映射的定义在oneflow/ir/include/OneFlow/OneFlowEnums.td ，OneFlow Dialect Operation的一些通用前端接口定义在oneflow/ir/include/OneFlow/OneFlowEnums.td。这里我们以Reshape Operation为例子来简单说明一下这个Operation有哪些组成部分：

def OneFlow_ReshapeOp : OneFlow_BaseOp<"reshape", [NoSideEffect, DeclareOpInterfaceMethods]> {  let input = (ins    AnyType:$in  );  let output = (outs    AnyType:$out  );  let attrs = (ins    AnyI64ElementsAttr:$shape  );}

OneFlow_ReshapeOp 这个名字下划线之前的是Dialect的名字，后面是这个Dialect下的Operation的名字。然后这个Operation继承了OneFlow_BaseOp基类，并声明了约束和前端接口，接下来定义了Operation的输入，输出和属性就结束了。

可以发现，OneFlow Dialect Operation的定义和OneFlow User Op是完全一致的，这保证了OneFlow和MLIR互转的合法性。OneFlow Reshape Operation的定义如下：

REGISTER_USER_OP("reshape")    .Input("in")    .Output("out")    .Attr("shape")    ...

OneFlow Job和MLIR的互转实现在oneflow/ir/oneflow-translate，主要做的事情就是遍历Job的OpGraph，对节点和边分别进行处理最后转换成一个MLIR表达式，同时在计算完成后可以基于MLIR表达式重写Job。这里的整体逻辑偏复杂，因为要处理OneFlow Job OpGraph里面各种类型Operation和边的转化，这里不继续深入讲解，因为它也不是我这篇文章要讨论的点，感兴趣的可以直接阅读代码。

3、OneFlow IR如何执行？

在上面Operation定义时是举了一个Reshape的例子，浏览oneflow/ir/include/OneFlow/OneFlowOps.td容易发现这里还定义了一个OneFlow_MlirJitOp，这个自定义的Op就是用来执行MLIR表达式的，它里面实现了CPU和GPU的Kernel（源码在oneflow/ir/oneflow-extension/extension.cpp）用来加载MLIR提供的JIT执行引擎运行最终得到的LLVM IR。那么LLVM IR又是怎么来的呢？这是通过OneFlow MLIR表达式逐级下降之后得来的，具体下降过程如下：

void AddLowerToLinalgMemRefPasses(PassManager& pm) {  pm.addPass(createLowerOneFlowToTosaPass());            // lower-oneflow-to-tosa  pm.addPass(createCSEPass());                           // cse  pm.addNestedPass(tosa::createTosaToLinalg());  // tosa-to-linalg-on-tensors  auto p = createLinalgElementwiseOpFusionPass();  assert(p->initializeOptions("allow-folding-unit-dim-reshapes=true").succeeded());  pm.addNestedPass(std::move(p));                     // linalg-fuse-elementwise-ops  pm.addNestedPass(createLinalgBufferizePass());      // linalg-bufferize  pm.addNestedPass(createTensorBufferizePass());      // tensor-bufferize  pm.addPass(createTensorConstantBufferizePass());            // tensor-constant-bufferize  pm.addPass(createFuncBufferizePass());                      // func-bufferize  pm.addPass(createBufferResultsToOutParamsPass());           // buffer-results-to-out-params  pm.addPass(createCanonicalizerPass());                      // canonicalize  pm.addNestedPass(createFinalizingBufferizePass());  // finalizing