當(dāng)前位置：首頁 > 人文社科 > 生活经验 >内容正文

生活经验

如何使用TVM Pass Relay

發(fā)布時(shí)間：2023/11/28 生活经验 36 豆豆

生活随笔收集整理的這篇文章主要介紹了如何使用TVM Pass Relay 小編覺得挺不錯(cuò)的,現(xiàn)在分享給大家,幫大家做個(gè)參考.

如何使用TVM Pass Relay
隨著Relay / tir中優(yōu)化遍數(shù)的增加，執(zhí)行并手動(dòng)維護(hù)其依賴關(guān)系變得很棘手。引入了一個(gè)基礎(chǔ)結(jié)構(gòu)來管理優(yōu)化過程，將其應(yīng)用于TVM堆棧中IR的不同層。
Relay / tir程序的優(yōu)化可以以各種粒度應(yīng)用，分別使用tvm.relay.transform.FunctionPass/ tvm.tir.transform.PrimFuncPass和的功能級(jí)別和模塊級(jí)別tvm.transform.ModulePass 。用戶可以依靠在tvm.transform.Sequential relay/ tir程序上應(yīng)用一系列Pass，其中Pass之間的依賴性可以passPass下文解決。
本文主要說明開發(fā)人員如何使用pass infra進(jìn)行特定的優(yōu)化，創(chuàng)建用于Relay程序的優(yōu)化管道。同樣的方法也可以用于tir。
import numpy as np
import tvm
from tvm import te
import tvm.relay as relay
創(chuàng)建一個(gè)示例 relay程序
創(chuàng)建一個(gè)簡單的Relay程序。該程序?qū)⒂糜诒疚闹惺纠母鞣N優(yōu)化。用戶可以編寫一個(gè)tir基本函數(shù)并應(yīng)用tirPass。
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)
為conv2d op注冊(cè)布局更改，在示例中應(yīng)用布局更改通道。alter layout pass如何工作不在本文的討論范圍之內(nèi)。
@relay.op.register_alter_op_layout(“nn.conv2d”, level=101)
def alter_conv2d(attrs, inputs, tinfos, out_type):
data, weight = inputs
new_attrs = dict(attrs)
new_attrs[“data_layout”] = “NCHW16c”
return relay.nn.conv2d(data, weight, **new_attrs)
優(yōu)化程序
現(xiàn)在要優(yōu)化程序。 relay具有許多優(yōu)化功能。將選擇其中一些以應(yīng)用于此示例程序。
有多種優(yōu)化 relay程序的方法。下面將為每個(gè)示例提供示例。
手動(dòng)應(yīng)用優(yōu)化Pass

Let’s first create a relay Module which contains one or multiple Relay

functions for optimization.

f = example()
mod = tvm.IRModule.from_expr(f)

Now we can apply constant folding on the module.

fold_const here is a callback that doesn’t take any parameters.

fold_const = relay.transform.FoldConstant()

Then, we can invoke the pass on the given module. Note that the constant

folding pass works at the function-level. That being said, each function in

the module will be applied with the optimization. Users don’t need to iterate

through individual functions manually to apply this pass.

mod = fold_const(mod)

We can see from the updated program that the constants are folded.

print(mod)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = add(%0, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%3 = add(%1, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%2, %3) / ty=Tensor[(1, 64, 54, 54), float32] /
}
以類似方式應(yīng)用更多優(yōu)化。例如，消除z和z1使用的通用表達(dá)式。
mod = relay.transform.EliminateCommonSubexpr()(mod)
print(mod)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) / ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = add(%0, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%2, %2) / ty=Tensor[(1, 64, 54, 54), float32] */
}
一些優(yōu)化（例如融合）也是參數(shù)化的。例如，選擇級(jí)別0不允許將算子融合在一起。用戶可以傳遞 fuse_opt_level來啟用此功能。
mod = relay.transform.FuseOps(fuse_opt_level=0)(mod)

We can observe that the optimized module contains functions that only have

a signle primitive op.

print(mod)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%0 = fn (%p0: Tensor[(1, 64, 56, 56), float32], %p1: Tensor[(64, 64, 3, 3), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] /
};
%1 = %0(%x, %weight) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = fn (%p01: Tensor[(1, 64, 54, 54), float32], %p11: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
add(%p01, %p11) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%3 = %2(%1, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%4 = fn (%p02: Tensor[(1, 64, 54, 54), float32], %p12: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
add(%p02, %p12) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%5 = %4(%3, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%6 = fn (%p03: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
add(%p03, %p03) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%6(%5) / ty=Tensor[(1, 64, 54, 54), float32] */
}
使用順序來應(yīng)用Pass序列
應(yīng)用Pass實(shí)際上是乏味的，需要用戶更好地了解之間的依賴性。例如，融合目前不適用于let綁定。如果relay.transform.ToANormalForm()在融合之前應(yīng)用算子，無法融合在一起，此過程為每個(gè)表達(dá)式生成let綁定，以規(guī)范化Relay程序。
Relaytvm.transform.Sequentialpass指定每個(gè)遍歷，將打包為整體來減輕開發(fā)人員顯式處理這些問題的負(fù)擔(dān)。例如，現(xiàn)在可以使用以下順序樣式應(yīng)用。tvm.transform.Sequential與torch.nn.sequential 和mxnet.gluon.block類似。例如，torch.nn.sequential用于包含一系列PyTorch模塊，這些模塊將被添加，以構(gòu)建網(wǎng)絡(luò)，著重于網(wǎng)絡(luò)層。取而代之的是tvm.transform.Sequential，下面的過程中的基礎(chǔ)工作于優(yōu)化過程。

Now let’s execute some passes through :py:class:`tvm.transform.Sequential`

f = example()
mod = tvm.IRModule.from_expr(f)

Glob the interested passes.

seq = tvm.transform.Sequential(
[
relay.transform.FoldConstant(),
relay.transform.EliminateCommonSubexpr(),
relay.transform.FuseOps(fuse_opt_level=2),
]
)
mod1 = seq(mod)
print(mod1)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%4 = fn (%p0: Tensor[(1, 64, 56, 56), float32], %p1: Tensor[(64, 64, 3, 3), float32], %p2: Tensor[(1, 64, 54, 54), float32], %p3: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = add(%0, %p2) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, %p3) / ty=Tensor[(1, 64, 54, 54), float32] /;
%3 = add(%1, %p3) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%2, %3) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%4(%x, %weight, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /
}
從轉(zhuǎn)換后的Relay程序中，可以看到仍然有兩個(gè)相同的加法運(yùn)算。這是EliminateCommonSubexpr 未實(shí)際執(zhí)行。只有優(yōu)化級(jí)別小于或等于2的過程才被執(zhí)行 tvm.transform.Sequential。下面的pass提供了一個(gè)配置界面，供用戶自定義要執(zhí)行的優(yōu)化級(jí)別。
with tvm.transform.PassContext(opt_level=3):
mod2 = seq(mod)
print(mod2)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%3 = fn (%p0: Tensor[(1, 64, 56, 56), float32], %p1: Tensor[(64, 64, 3, 3), float32], %p2: Tensor[(1, 64, 54, 54), float32], %p3: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) / ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = add(%0, %p2) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, %p3) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%2, %2) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%3(%x, %weight, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /
}
可以看到僅保留了兩個(gè)相同的加法之一。
用戶可以使用disabled_pass配置有選擇地禁用某些pass，這類似于通用編譯器（例如Clang和GCC）使用的-fno-xxx選項(xiàng)。例如，可以禁用EliminateCommonSubexpr，如下所示。打印的模塊將再次顯示兩個(gè)相同的加法運(yùn)算。
with tvm.transform.PassContext(opt_level=3, disabled_pass=[“EliminateCommonSubexpr”]):
mod3 = seq(mod)
print(mod3)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%4 = fn (%p0: Tensor[(1, 64, 56, 56), float32], %p1: Tensor[(64, 64, 3, 3), float32], %p2: Tensor[(1, 64, 54, 54), float32], %p3: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) / ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = add(%0, %p2) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, %p3) / ty=Tensor[(1, 64, 54, 54), float32] /;
%3 = add(%1, %p3) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%2, %3) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%4(%x, %weight, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] */
}
應(yīng)用的Pass與目標(biāo)無關(guān)。下文的Pass還提供了具有目標(biāo)意識(shí)的方法。例如，布局變更階段屬于這種類別。
with tvm.transform.PassContext(opt_level=3):
mod4 = seq(mod)
print(mod4)

seq1 = tvm.transform.Sequential([relay.transform.AlterOpLayout()])
with tvm.transform.PassContext(opt_level=3):
with tvm.target.Target(“l(fā)lvm”):
mod5 = seq1(mod)
print(mod5)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%3 = fn (%p0: Tensor[(1, 64, 56, 56), float32], %p1: Tensor[(64, 64, 3, 3), float32], %p2: Tensor[(1, 64, 54, 54), float32], %p3: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = add(%0, %p2) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, %p3) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%2, %2) / ty=Tensor[(1, 64, 54, 54), float32] /
};
%3(%x, %weight, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /, meta[relay.Constant][1] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /
}
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%0 = layout_transform(%x, src_layout=“NCHW”, dst_layout=“NCHW16c”) / ty=Tensor[(1, 4, 56, 56, 16), float32] /;
%1 = nn.conv2d(%0, %weight, padding=[0, 0, 0, 0], data_layout=“NCHW16c”) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
%2 = add(meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%3 = multiply(%2, 2f / ty=float32 /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%4 = layout_transform(%3, src_layout=“NCHW”, dst_layout=“NCHW16c”) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
%5 = add(%1, %4) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
%6 = layout_transform(meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /, src_layout=“NCHW”, dst_layout=“NCHW16c”) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
%7 = add(%5, %6) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
%8 = add(%5, %6) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
%9 = add(%7, %8) / ty=Tensor[(1, 4, 54, 54, 16), float32] /;
layout_transform(%9, src_layout=“NCHW16c”, dst_layout=“NCHW”) / ty=Tensor[(1, 64, 54, 54), float32] */
}
使用Python Decorator實(shí)施Pass
下一個(gè)示例說明了如何使用Python裝飾器pass傳遞基礎(chǔ)流程來編排定制的優(yōu)化管道。極大地簡化了Pass的實(shí)施。例如，用戶可以簡單地定義一個(gè)修飾的類，進(jìn)行功能級(jí)別的優(yōu)化，如以下示例所示。transform_function包裝一個(gè)類，以用c的倍數(shù)替換所有常量。調(diào)用自定義過程時(shí)，將訪問給定模塊中的每個(gè)函數(shù)，并且將替換函數(shù)中的每個(gè)常量。
@relay.transform.function_pass(opt_level=1)
class CustomPipeline:
“”“Simple test function to replace one argument to another.”""

def __init__(self, multiplier):self.multiplier = multiplier# This function can define a pass.
def transform_function(self, func, mod, ctx):obj = selfclass ReplaceConstant(tvm.relay.ExprMutator):def visit_constant(self, c):return relay.multiply(obj.multiplier, c)return ReplaceConstant().visit(func)

f = example()
mod = tvm.IRModule.from_expr(f)
custom_pass = CustomPipeline(multiplier=relay.const(3, “float32”))
assert custom_pass.info.name == “CustomPipeline”
mod3 = custom_pass(mod)
print(mod3)
輸出：
def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] {
%0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] /;
%1 = multiply(3f / ty=float32 /, meta[relay.Constant][0] / ty=Tensor[(1, 64, 54, 54), float32] /) / ty=Tensor[(1, 64, 54, 54), float32] /;
%2 = add(%1, %1) / ty=Tensor[(1, 64, 54, 54), float32] /;
%3 = multiply(3f / ty=float32 /, 2f / ty=float32 /) / ty=float32 /;
%4 = multiply(%2, %3) / ty=Tensor[(1, 64, 54, 54), float32] /;
%5 = add(%0, %4) / ty=Tensor[(1, 64, 54, 54), float32] /;
%6 = add(%5, %1) / ty=Tensor[(1, 64, 54, 54), float32] /;
%7 = add(%5, %1) / ty=Tensor[(1, 64, 54, 54), float32] /;
add(%6, %7) / ty=Tensor[(1, 64, 54, 54), float32] */
}
調(diào)試Pass
TVM為用戶提供了一個(gè)插用式的調(diào)試通道，在pass特殊通道（PrintIR）來轉(zhuǎn)儲(chǔ)整個(gè)模塊的IR之后，將IR打印出來。順序傳遞示例的略微修改版本，類似于以下內(nèi)容，以啟用IR轉(zhuǎn)儲(chǔ)以進(jìn)行FoldConstant優(yōu)化。
f = example()
mod = tvm.IRModule.from_expr(f)
seq = tvm.transform.Sequential(
[
relay.transform.FoldConstant(),
tvm.transform.PrintIR(),
relay.transform.EliminateCommonSubexpr(),
relay.transform.FuseOps(),
relay.transform.AlterOpLayout(),
]
)

By inserting the `PrintIR` pass after `FoldConstant`, the pass infra will

dump out the module IR when `FoldConstant` is done. Users can plug in this

pass after any pass they want to debug for viewing the optimization effect.

There is a more flexible debugging mechanism also exposed by the build configuration

object. One can pass a tracing function which can be used to execute arbitrary code

before and/or after each pass. A tracing function will receive a :py::class:`tvm.IRModule`,

a :py:class:`tvm.transform.PassInfo` object,

and a boolean indicating whether you are executing before, or after a pass.

An example is below.

def print_ir(mod, info, is_before):
“”“Print the name of the pass, the IR, only before passes execute.”""
if is_before:
print(“Running pass: {}”, info)
print(mod)

with tvm.transform.PassContext(opt_level=3, trace=print_ir):
with tvm.target.Target(“l(fā)lvm”):
# Perform the optimizations.
mod = seq(mod)
print(mod)

print(“done”)
輸出：
Running pass: {} The meta data of the pass: pass name: FoldConstantopt_level: 2required passes: [
]

def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) {
%0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]);
%1 = add(meta[relay.Constant][0], meta[relay.Constant][0]);
%2 = multiply(%1, 2f);
%3 = add(%0, %2);
%4 = add(%3, meta[relay.Constant][0]);
%5 = add(%3, meta[relay.Constant][0]);
add(%4, %5)
}