pass的作用
pass是编译器中用来对程序进行优化或分析的步骤,一个程序通常要经过多个pass,每个pass也可能执行多遍。关于TVM中的pass机制,请参考TVM的官方文档:
https://tvm.apache.org/docs/arch/pass_infra.htmltvm.apache.org/docs/arch/pass_infra.html
本文粗略统计了Relay高层次pass,主要想搞明白在高层次计算图优化中,编译器都做了哪些事。
关于pass执行前后程序的变化,可以参考relay的测试文件夹,文中示例均来自tvm/tests/python/relay/ ,由于精力有限,我只选取个别进行了输入输出的对比。
pass统计
DeadCodeElimination
名称 |
中文名 |
作用 |
备注 |
DeadCodeElimination |
死节点消除 |
删除未被使用到的表达式 |
实现在tvm/src/relay/transforms/dead_code.cc |
实现流程应该是用UsageVisitor遍历每个节点,计数每个节点被用到的次数,当其为0时,在EliminatorMutator中将其去除。(这里没有详细探究实现细节,仅凭类名和函数名做合理猜测,可能与真实实现有出入,若想了解其具体实现,请参阅源代码,下面的pass同)
LazyGradientInit
名称 |
中文名 |
作用 |
备注 |
LazyGradientInit |
梯度延迟初始化 |
减小梯度张量的内存开销,对ones, ones_like, zeros, zeros_like 这些算子只有使用到时再实例化 |
实现在tvm/src/relay/transforms/lazy_gradient_init.cc |
FoldConstant
名称 |
中文名 |
作用 |
备注 |
FoldConstant(还有个FoldConstantExpr,针对常量表达式) |
常量折叠 |
折叠relay中的常量。 |
实现tvm/src/relay/transforms/fold_constant.cc |
这里代码中注释专门强调与QNN(TVM中导入已量化模型的框架)有关,QNN有些静态子图不能盲目折叠。
SplitArgs
名称 |
中文名 |
作用 |
备注 |
SplitArgs |
分割参数 |
将有大量参数的函数分割成小块 |
实现tvm/src/relay/transforms/split_args.cc |
根据对应python文件的注释,该pass在常量折叠时使用
FuseOps
名称 |
中文名 |
作用 |
备注 |
FuseOps(其逆操作为DefuseOps) |
算子融合 |
按照规则将算子融合为复合算子 |
实现tvm/src/relay/transforms/fuse_ops.cc |
在TVM中每个算子都有个类型,论文中是injective(1对1映射、如add、exp)、reduction(多对少的映射,如sum、max、min)、complex-out-fusable(逐元素复用映射到输出,比如conv2d)、opaque(不能融合,比如sort),融合规则如下
![img](/2023/01/07/relay-pass/v2-a9a599dbfc1cc3bd908219894c7f21db_720w.png)
不晓得现在融合方式有没有变(代码太复杂了,短时间没有理清(T_T)
名称 |
中文名 |
作用 |
备注 |
ToBasicBlockNormalForm |
将表达式转换为基本块范式 |
|
实现tvm/src/relay/transforms/to_basic_block_normal_form.cc |
名称 |
中文名 |
作用 |
备注 |
ToANormalForm |
转换为A范式 |
将图范式转换为A范式的形式 |
tvm/src/relay/transforms/to_a_normal_form.cc |
google了一下A -Norm,是函数式编译器中程序的中间表示
为什么这么做可参阅这里:
https://matt.might.net/articles/a-normalization/
ToCPS
名称 |
中文名 |
作用 |
备注 |
ToCPS |
转换为CPS(continuation passing style) |
将表达式转换为连续传递风格 |
tvm/src/relay/transforms/to_cps.cc |
关于CPS,戳这里https://en.wikipedia.org/wiki/Continuation-passing_style
名称 |
中文名 |
作用 |
备注 |
ToGraphNormalForm |
转换为图范式 |
将A范式转换为图范式 |
tvm/src/relay/transforms/to_graph_normal_form.cc |
PartialEval
名称 |
中文名 |
作用 |
备注 |
PartialEval |
部分求值 |
在编译时做尽可能多的计算,减少运行时的开销 |
tvm/src/relay/transforms/partial_eval.cc |
SimplifyInference
名称 |
中文名 |
作用 |
备注 |
SimplifyInference |
简化推理 |
对各种归一化层做了重写 |
tvm/src/relay/transforms/simplify_inference.cc |
对batch_norm、dropout、instance_norm、layer_norm、group_norm、l2_norm做了重写
Example:tvm/tests/python/relay/test_pass_simplify_inference.py
# 原始输入 def @main(%x: Tensor[(10, 10), float32], %gamma: Tensor[(10), float32], %beta: Tensor[(10), float32], %moving_mean: Tensor[(10), float32], %moving_var: Tensor[(10), float32]) { %0 = add(%x, 1f); %1 = nn.batch_norm(%0, %gamma, %beta, %moving_mean, %moving_var, epsilon=0.01f); %2 = %1.0; %3 = nn.dropout(%2); %3.0 }
#执行该pass def @main(%x: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, %gamma: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %beta: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %moving_mean: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %moving_var: Tensor[(10), float32] /* ty=Tensor[(10), float32] */) -> Tensor[(10, 10), float32] { %0 = add(%moving_var, 0.01f /* ty=float32 */) /* ty=Tensor[(10), float32] */; %1 = sqrt(%0) /* ty=Tensor[(10), float32] */; %2 = divide(1f /* ty=float32 */, %1) /* ty=Tensor[(10), float32] */; %3 = add(%x, 1f /* ty=float32 */) /* ty=Tensor[(10, 10), float32] */; %4 = multiply(%2, %gamma) /* ty=Tensor[(10), float32] */; %5 = negative(%moving_mean) /* ty=Tensor[(10), float32] */; %6 = multiply(%5, %4) /* ty=Tensor[(10), float32] */; %7 = multiply(%3, %4) /* ty=Tensor[(10, 10), float32] */; %8 = add(%6, %beta) /* ty=Tensor[(10), float32] */; add(%7, %8) /* ty=Tensor[(10, 10), float32] */ }
|
FastMath
名称 |
中文名 |
作用 |
备注 |
FastMath |
快速数值计算 |
将非线性激活函数用更快的近似方法代替 |
tvm/src/relay/transforms/fast_math.cc |
包括exp(自然指数)、erf(误差函数)、tanh、softmax的fast版本
DynamicToStatic
名称 |
中文名 |
作用 |
备注 |
DynamicToStatic |
动态转静态 |
若动态算子的输入是静态,将其转换为静态算子,重新执行类型推理和常量折叠 |
tvm/src/relay/transforms/dynamic_to_static.cc |
包含算子:reshape、squeeze、tile、topk、broadcast_to、zeros、ones、one_hot、image.resize2d、full、nn.upsampling、nn.upsampling3d、nn.pad、strided_slice、sparse_to_dense
InferType
名称 |
中文名 |
作用 |
备注 |
InferType(InferTypeLocal类似) |
推断类型 |
填入明确的类型信息 |
tvm/src/relay/transforms/type_infer.cc |
EliminateCommonSubexpr
名称 |
中文名 |
作用 |
备注 |
EliminateCommonSubexpr |
消除公共子表达式 |
若多个表达式是相同(属性相同且输入相同)的值且在多处使用,创建一个公共的变量替换这些表达式 |
tvm/src/relay/transforms/eliminate_common_subexpr.cc |
CombineParallelConv2D
名称 |
中文名 |
作用 |
备注 |
CombineParallelConv2D(CombineParallelBatchMatmul类似) |
合并并行卷积 |
将共享相同输入节点和相同参数的卷积(除了输出通道的数量可以不同)替换为单一的卷积。新的2d卷积的权重是原始权重的concat。conv2d之后的加法广播运算也会尽可能合并。这可以防止在有多个卷积分支的网络中启动多个内核,如Inception。 |
tvm/src/relay/transforms/combine_parallel_conv2d.cc |
CombineParallelDense
名称 |
中文名 |
作用 |
备注 |
CombineParallelDense |
合并并行Dense层 |
取代了共享相同输入节点、相同形状、没有定义 “ units “的dense操作,只用一个batch MM就可以。 这可以防止在有多个dense分支的网络中启动多个内核,例如BERT |
tvm/src/relay/transforms/combine_parallel_dense.cc |
FoldScaleAxis
名称 |
中文名 |
作用 |
备注 |
FoldScaleAxis(ForwardFoldScaleAxis、BackwardFoldScaleAxis类似) |
折叠在轴上的缩放 |
当某个轴上带有放缩时,将其融合进dense或conv2d的权重 |
tvm/src/relay/transforms/fold_scale_axis.cc |
# tvm/tests/python/relay/test_pass_fold_scale_axis.py
#最初的输入 def @main(%x: Tensor[(2, 4), float32] /* ty=Tensor[(2, 4), float32] */, %weight: Tensor[(3, 4), float32] /* ty=Tensor[(3, 4), float32] */, %in_bias: Tensor[(4), float32] /* ty=Tensor[(4), float32] */) -> Tensor[(2, 3), float32] { %0 = multiply(%x, meta[relay.Constant][0] /* ty=Tensor[(4), float32] */) /* ty=Tensor[(2, 4), float32] */; %1 = nn.relu(%0) /* ty=Tensor[(2, 4), float32] */; %2 = add(%1, %in_bias) /* ty=Tensor[(2, 4), float32] */; nn.dense(%2, %weight, units=None) /* ty=Tensor[(2, 3), float32] */ }
#结果 def @main(%x: Tensor[(2, 4), float32] /* ty=Tensor[(2, 4), float32] */, %weight: Tensor[(3, 4), float32] /* ty=Tensor[(3, 4), float32] */, %in_bias: Tensor[(4), float32] /* ty=Tensor[(4), float32] */) -> Tensor[(2, 3), float32] { %0 = nn.relu(%x); %1 = divide(%in_bias, meta[relay.Constant][0] /* ty=Tensor[(4), float32] */); %2 = add(%0, %1); %3 = multiply(%weight, meta[relay.Constant][0] /* ty=Tensor[(4), float32] */); nn.dense(%2, %3, units=None) }
|
口头验证一下结果:
输入的计算流程相当于:假设f是relu函数
AX -> f(AX) -> f(AX)+B ->W(f(AX)+B) = Wf(AX)+WB
执行该pass后:
f(X) -> B/A -> f(X)+B/A -> WA -> WA(f(X)+B/A) =WA(fX)+WB 由于Relu是留下正数,所以上下两个的计算结果是相等的
通过该pass,将在运行时对输入的计算转换为对权重的计算,然后这部分可以在编译时就对权重进行修改,从而节约了运行时开销。BatchNorm分解后也可以用该pass融合进权重。
CanonicalizeOps
名称 |
中文名 |
作用 |
备注 |
CanonicalizeOps |
规范化算子 |
将算子转换为其简化版本 |
tvm/src/relay/transforms/canonicalize_ops.cc |
在实现中仅涉及bias_add_op,将其转换为expand_dims+boradcast_add
AlterOpLayout
名称 |
中文名 |
作用 |
备注 |
AlterOpLayout(ConvertLayout类似) |
替换算子布局 |
用于计算自定义布局的卷积或其他通用权重的预变换。 |
tvm/src/relay/transforms/alter_op_layout.cc |
AutoSchedulerLayoutRewrite
名称 |
中文名 |
作用 |
备注 |
AutoSchedulerLayoutRewrite(MetaScheduleLayoutRewrite类似) |
AutoScheduler布局重写 |
根据AutoScheduler(TVM的无模板自动调优器)生成的tile结构进行对应权重的存储布局重写 |
tvm/src/relay/transforms/auto_scheduler_layout_rewrite.cc |
包括Conv2D、Conv2DWinograd、Conv3D、Matmul、Dense、BatchMatmul
Legalize
名称 |
中文名 |
作用 |
备注 |
Legalize |
合法化 |
用另一个表达式替换一个表达式,以实现目标平台相关的优化 |
tvm/src/relay/transforms/legalize.cc |
该pass主要用于QNN,关于该pass的作用可查阅QNN的论文。举例几个QNN合法化的算子:
路径在tvm/src/relay/qnn/op/
qnn.dense、qnn.batch_matmul、qnn.conv2d等
CanonicalizeCast
名称 |
中文名 |
作用 |
备注 |
CanonicalizeCast |
规范化cast |
规范化cast表达式,使得算子融合更有效 |
tvm/src/relay/transforms/canonicalize_cast.cc |
EtaExpand
名称 |
中文名 |
作用 |
备注 |
EtaExpand |
|
为构造函数添加抽象,或者给一个函数添加全局变量 |
tvm/src/relay/transforms/eta_expand.cc |
PartitionGraph
名称 |
中文名 |
作用 |
备注 |
PartitionGraph |
划分图 |
根据插入的标注节点(即compiler_begin和compiler_end),将一个输入函数分割成多个函数。这些节点被用作边界,将Relay函数划分为多个区域,可以卸载到不同的加速器后端。 每一个被分割的函数,也就是区域,都将被视为外部函数,并且它们将使用所提供的编译器进行代码生成。 |
tvm/src/relay/transforms/partition_graph.cc |
该pass的作用可见专栏中的BYOC
Inline
名称 |
中文名 |
作用 |
备注 |
Inline |
内联 |
将全局函数内联到Relay IR Module |
tvm/src/relay/transforms/inline.cc |
RemoveUnusedFunctions
名称 |
中文名 |
作用 |
备注 |
RemoveUnusedFunctions |
删除未用到的函数 |
|
tvm/src/relay/backend/vm/removed_unused_funcs.cc |
SimplifyExpr
名称 |
中文名 |
作用 |
备注 |
SimplifyExpr |
简化表达式 |
简化Relay表达式 |
tvm/src/relay/transforms/simplify_expr.cc |
主要包含以下简化:
ConcretizeZerosLikeRewrite ConcretizeOnesLikeRewrite ConcretizeFullLikeRewrite ConcretizeReshapeLikeRewrite ConcretizeCollapseSumLikeRewrite ConcretizeBroadcastToLikeRewrite ConcretizeCastLikeRewrite SimplifyRSqrt EliminateIdentityRewrite SimplifyReshape SimplifyTranspose SimplifySameCast SimplifyConsecutiveCast FullElementwise SwitchAddMultiply SimplifyAdjacentMultiplyOrAdd SimplifyDQArgMax SimplifyDQArgMin SimplifyDQArgSort SimplifyClipAndConsecutiveCast SimplifyCastClip
|
RelayToTIRTargetHook
名称 |
中文名 |
作用 |
备注 |
RelayToTIRTargetHook |
Relay到TIR目标的hook |
处理已在IRModule内的函数上注册的目标hook |
tvm/src/relay/transforms/target_hooks.cc |
ManifestAlloc
名称 |
中文名 |
作用 |
备注 |
ManifestAlloc(ManifestLifetimes类似) |
表现内存分配 |
|
tvm/src/relay/transforms/memory_alloc.cc |
PlanDevices
|
中文名 |
作用 |
备注 |
PlanDevices |
计划设备 |
确定每个relay子表达式运行和存储结果的虚拟设备(VirtualDevice) |
tvm/src/relay/transforms/device_planner.cc |
FlattenAtrousConv
|
中文名 |
作用 |
备注 |
FlattenAtrousConv |
展平空洞卷积 |
这种转换将空洞卷积扁平化,它对应于操作序列”space_to_batch_nd”->”conv2d””batch_to_space_nd”,并将它们转换成带有修改过的 “dilation “和重新计算的 “padding “参数的卷积子图。 |
tvm/src/relay/transforms/flatten_atrous_conv.cc |
AnnotateUsedMemory
|
中文名 |
作用 |
备注 |
AnnotateUsedMemory |
标注使用的存储 |
通过分析每个函数调用的输入/输出张量的有效性并计算这些张量所需的总内存量,标注每个函数调用点的最小所需内存。不支持动态类型。 |
tvm/src/relay/backend/annotate_used_memory.cc |
示例:
Before: def @main(%input: Tensor[(1, 2, 2, 4), int8]) -> Tensor[(1, 2, 2, 4), int8] { let %x_0 = fn (%x: Tensor[(1, 2, 2, 4), int8], Primitive=1) -> Tensor[(1, 2, 2, 4), int8] { nn.max_pool2d(%x, pool_size=[1, 1], padding=[0, 0, 0, 0]) }; let %x_1 = %x_0(%input); %x_1 }
After: def @main(%input: Tensor[(1, 2, 2, 4), int8], io_used_memory=32) -> Tensor[(1, 2, 2, 4), int8] { let %x_0: fn (%x: Tensor[(1, 2, 2, 4), int8], Primitive=1, used_memory=[32]) -> Tensor[(1, 2, 2, 4), int8] { nn.max_pool2d(%x, pool_size=[1, 1], padding=[0, 0, 0, 0]) }; let %x_1: Tensor[(1, 2, 2, 4), int8] = %x_0(%input); %x_1 }
|
CapturePostDfsIndexInSpans
|
中文名 |
作用 |
备注 |
CapturePostDfsIndexInSpans |
|
|
tvm/src/relay/transforms/capture_postdfsindex_in_spans.cc |
AnnotateMemoryScope
|
中文名 |
作用 |
备注 |
AnnotateMemoryScope |
标注内存范围 |
|
tvm/src/relay/transforms/annotate_texture_storage.cc |
RemoveStandaloneReshapes
|
中文名 |
作用 |
备注 |
RemoveStandaloneReshapes |
删除独立reshape |
|
tvm/src/relay/transforms/remove_standalone_reshapes.cc |
Pass简要分类
粗糙的分类,很可能不正确,仅供参考
主要作用 |
包含pass |
化简,去除运行时冗余,加快计算 |
DeadCodeElimination、LazyGradientInit、FoldConstant、SplitArgs、PartialEval、SimplifyInference、FastMath、DynamicToStatic、EliminateCommonSubexpr、FoldScaleAxis、CanonicalizeOps、CanonicalizeCast、RemoveUnusedFunctions、SimplifyExpr、RemoveStandaloneReshapes |
程序语言形式相关 |
ToBasicBlockNormalForm、ToANormalForm、ToCPS、ToGraphNormalForm、InferType、EtaExpand、Inline |
编译器相关的,设备、存储 |
AlterOpLayout、AutoSchedulerLayoutRewrite、PartitionGraph、RelayToTIRTargetHook、ManifestAlloc、AnnotateUsedMemory、AnnotateMemoryScope |
算子内核生成、改进的 |
FuseOps、CombineParallelConv2D、CombineParallelDense、FlattenAtrousConv |
执行检查 |
Legalize |
有些pass我也没搞懂啥作用,上述表格对应位置就空下了,在此抛砖引玉。