TVM圖優化（以Op Fusion為例）

首先給出一個TVM 相關的介紹，這個是Tianqi Chen演講在OSDI18上用的PPT https://files.cnblogs.com/files/jourluohua/Tianqi-Chen-TVM-Stack-Overview.rar

對於圖優化來說，位於整個軟件編譯棧比較高的層次：

首先給出計算圖的定義

Computational graphs: a common way to represent programs in deep learning frameworks

對於圖優化來說，有很多種圖優化手段：Operator Fusion
Constant Parameter Path Pre-Computation
Static Memory Reuse Analysis
Data Layout Transformation
AlterOpLayout
SimplifyInference

這里僅以Operator Fusion做例子介紹

Operator fusion : combine multiple operators together into a single kernel without saving the intermediate results back into global memory

也就說是說算子融合省掉了中間數據的store過程

在TVM中，有三種融合規則：

其中，算子屬於哪一類是算子本身的特性（這個地方不是特別懂，這個屬性有非常多的值），但是能融合的規則只有這三種。

但是這種store是如何減少的，在IR上有明確的體現。

下邊的例子，我會使用tvm.relay來進行介紹，relay是TVM中實現的一種高級IR，可以簡單理解為另一種計算圖表示。其在TVM所處的位置如下圖所示

 

比如，我們假設我們要完成一個y = exp(x+1.0)的計算圖

給出測試代碼(來自於源碼中的test_pass_fuse_ops.py，有改動)：

import tvm
from tvm import relay

def test_fuse_simple():
    """Simple testcase."""
    def before():
        x = relay.var("x", shape=(10, 20))
        y = relay.add(x, relay.const(1, "float32"))
        z = relay.exp(y)
        return relay.Function([x], z)

    def expected():
        x = relay.var("p", shape=(10, 20))
        y = relay.add(x, relay.const(1, "float32"))
        z = relay.exp(y)
        f1 = relay.Function([x], z)
        x = relay.var("x", shape=(10, 20))
        y = relay.Call(f1, [x])
        return relay.Function([x], y)

    z = before()
    z = relay.ir_pass.infer_type(z)
    # print(z.astext())
    zz = relay.ir_pass.fuse_ops(z, opt_level=2)
    print(zz.astext())
    zz = relay.ir_pass.infer_type(zz)
    zz = relay.ir_pass.fuse_ops(zz)
    zz = relay.ir_pass.infer_type(zz)
    after = relay.ir_pass.infer_type(expected())
    # print(after.astext())
    assert relay.ir_pass.alpha_equal(zz, after)

在融合前，其IR(方便用戶看的一種形式，不是真正的IR)

fn (%x: Tensor[(10, 20), float32])
    -> Tensor[(10, 20), float32] {
  %0 = fn(%p0: Tensor[(10, 20), float32],
          %p1: float32)
          -> Tensor[(10, 20), float32] {
    %1 = add(%p0, %p1)
    %1
  }
  %2 = %0(%x, 1f)
  %3 = fn(%p01: Tensor[(10, 20), float32])
          -> Tensor[(10, 20), float32] {
    %4 = exp(%p01)
    %4
  }
  %5 = %3(%2)
  %5
}

 融合后：

fn (%x: Tensor[(10, 20), float32])
    -> Tensor[(10, 20), float32] {
  %0 = fn(%p0: Tensor[(10, 20), float32])
          -> Tensor[(10, 20), float32] {
    %1 = add(%p0, 1f)
    %2 = exp(%1)
    %2
  }
  %3 = %0(%x)
  %3
}

 可以很明顯的發現，省掉了一次數據store過程