TVM Pass Infra#

参考：

Relay/tir 程序的优化可以应用在不同的粒度上，即函数级 tvm.relay.transform.FunctionPass/tvm.tir.transform.PrimFuncPass 和模块级 tvm.transform.ModulePass。或者用户可以依赖于 tvm.transform.Sequential 在 Relay/tir 程序上应用 pass 序列，其中 pass 之间的依赖性可以由 pass infra 解析。

函数级 Pass#

当提供 pass_func 时，function_pass() 函数返回回调函数。否则，它将使用给定的优化函数返回创建的函数级 pass。

参数：

pass_func：变换函数或变换类。
opt_level：优化级别。
name：pass 的名称。名称可以为空。
required：依赖的 pass 列表。

直接看例子。

创建函数 func：

import numpy as np
import tvm
from tvm import relay
from tvm.ir import IRModule


a = relay.var("a", shape=(10, 20))
b = relay.var("b", shape=(10, 20))
c = a + b
d = c * c
func = relay.Function([a, b], d)
input_mod = IRModule.from_expr(func)
input_mod

#[version = "0.0.5"]
def @main(%a: Tensor[(10, 20), float32], %b: Tensor[(10, 20), float32]) {
  %0 = add(%a, %b);
  multiply(%0, %0)
}

创建函数级的 pass：

@relay.transform.function_pass(opt_level=1)
class TestReplaceFunc:
    def __init__(self, new_func):
        self.new_func = new_func

    def transform_function(self, func, mod, ctx):
        # 为了演示，将 func 转换为 new_func
        return self.new_func

此 pass 仅将 func 变换为 new_func。

创建 new_func：

x = relay.var("x", shape=(10, 20))
new_func = relay.Function([x], relay.log(x))

fpass 是特殊的 pass，它将每个函数替换为 new_func：

fpass = TestReplaceFunc(new_func)

现在，input_mod 中的每个函数都被 new_func 替换：

res_mod = fpass(input_mod)
res_mod

#[version = "0.0.5"]
def @main(%x: Tensor[(10, 20), float32] /* ty=Tensor[(10, 20), float32] */) -> Tensor[(10, 20), float32] {
  log(%x) /* ty=Tensor[(10, 20), float32] */
}

也可以通过装饰用户定义的 transform 函数来创建函数级 pass：

@relay.transform.function_pass(opt_level=2)
def transform(func, mod, ctx):
    # 自定义变换
    x = relay.var("x", shape=(10, 20))
    new_func = relay.Function([x], relay.log(x))
    return new_func

function_pass = transform
assert isinstance(function_pass, relay.transform.FunctionPass)
assert function_pass.info.opt_level == 2

给定模块 input_mod，优化可以如下调用：

updated_mod = function_pass(input_mod)
updated_mod

#[version = "0.0.5"]
def @main(%x: Tensor[(10, 20), float32] /* ty=Tensor[(10, 20), float32] */) -> Tensor[(10, 20), float32] {
  log(%x) /* ty=Tensor[(10, 20), float32] */
}

模块级 Pass#

模块级 Pass tvm.transform.module_pass() 与 function_pass() 的定义和使用很相似。也分为类模式和函数模式两种。

类模式：

@tvm.transform.module_pass(opt_level=2)
class CustomPipeline:
    def __init__(self, enable_fold):
        self.enable_fold = enable_fold
        self.cse = relay.transform.EliminateCommonSubexpr()
        self.const_fold = relay.transform.FoldConstant()

    def transform_module(self, mod, ctx):
        mod = self.cse(mod)
        if self.enable_fold:
            mod = self.const_fold(mod)
        return mod

# 创建定制的 pipeline 实例
pipeline = CustomPipeline(enable_fold=False)
assert isinstance(pipeline, tvm.transform.ModulePass)

def example():
    shape = (1, 64, 54, 54)
    c_data = np.empty(shape).astype("float32")
    c = relay.const(c_data)
    weight = relay.var("weight", shape=(64, 64, 3, 3))
    x = relay.var("x", relay.TensorType((1, 64, 56, 56), "float32"))
    conv = relay.nn.conv2d(x, weight)
    y = relay.add(c, c)
    y = relay.multiply(y, relay.const(2, "float32"))
    y = relay.add(conv, y)
    z = relay.add(y, c)
    z1 = relay.add(y, c)
    z2 = relay.add(z, z1)
    return relay.Function([x, weight], z2)

m = IRModule.from_expr(example())

运行 pipeline：

output_module = pipeline(m)

函数模式：

@tvm.transform.module_pass(opt_level=2)
def transform(mod, ctx):
    x = relay.var("x", shape=(2,), dtype="float32")
    func = relay.Function([x], relay.abs(x))
    new_mod = IRModule()
    new_mod['var'] = func
    new_mod.update(mod)
    return new_mod

module_pass = transform
assert isinstance(module_pass, tvm.transform.ModulePass)
assert module_pass.info.opt_level == 2

# 给定模块 `m`，优化可以如下调用：
updated_mod = module_pass(m)
# 现在，函数 `abs` 应该被添加到模块 `m` 中。

updated_mod["var"]

fn (%x: Tensor[(2), float32]) {
  abs(%x)
}

TVM Pass Infra

目录

TVM Pass Infra#

函数级 Pass#

模块级 Pass#