TVM Pass 总结

tvm.relay.transform文档

ModulePass

将整个程序视作一个单元处理的pass。

LambdaLift、RemoveUnusedFunctions、Inline、EtaExpand、PartialEval、PrintIR、PartitionGraph、ToANormalForm

FunctionPsss

以单个函数为作用域的pass, 每个函数间是相互独立的

DeadCodeElimination、FoldConstant、FuseOps、RewriteAnnotatedOps、ToCPS、ToGraphNormalForm、SimplifyInference、FastMath、InferType、EliminateCommonSubexpr、CombineParallelConv2D、CombineParallelDense、BackwardFoldScaleAxis / ForwardFoldScaleAxis、CanonicalizeOps、AlterOpLayout、 ConvertLayout、 Legalize、 CanonicalizeCast、 MergeComposite、 AnnotateTarget

SequentialPass

FoldScaleAxis

1. tvm.relay.transform.AlterOpLayout（）

   替换运算符的布局或用其他表达式替换基本运算符。此过程可用于计算自定义布局中的卷积或其他常规权重预转换。可分为两个步骤：

• OpAlteration

data[NCHW] ->		conv2d[NCHW]->		broadcast_add[NCHW]->		bias[CHW]->

为运算选择最好的布局并保持输入输出布局不变

data[NCHW] ->		LT[NCHW->NCHW16c] conv2d->LT(NCHW16cNCHW->)		broadcast_add[NCHW]->		bias[CHW]->

• PeepholeOptimizer

Remove the layout_transform ops if possible

seq1 = tvm.transform.Sequential([relay.transform.AlterOpLayout()])
with tvm.transform.PassContext(opt_level=3):
    with tvm.target.create("llvm"):
        mod5 = seq1(mod)
print(mod5)

参考网址
2. tvm.relay.transform.AnnotateTarget( targets )

使用提供的编译器/目标在表达式中注释操作，然后将其用于代码生成。
将表达式中的算子标注编译器或者部署平台，使得算子被wrap为 compiler_begin／subgraph_start 和 compiler_end ／subgraph_end，从而用做之后其他编译器的代码生成。定义作为Wrapper使用。
在 src/relay/pass/annotate_target.cc

# register a dnnl specific target
reg.register("nn.conv2d", "target.dnnl", batch_norm_func)
# register xyzlib specific logic
reg.register("nn.conv2d", "target.xyzlib", batch_norm_func)

mod = relay.transform.AnnotateTarget(target=["dnnl", "xyzlib"])(mod)

https://discuss.tvm.ai/t/rfc-op-based-annotation-for-external-codegen/5709/28

3. tvm.relay.transform.BackwardFoldScaleAxis()

   向后折叠轴缩放为conv2d/dense 的权重。将轴的缩放折叠进卷积或者稠密算子的weights中

   建议在使用前向作为后向目标前调用后向。

   Goal：将轴的缩放（通常由BatchNorm引起）折叠为conv2d的权重。

   e.g.:

   Old:

   %1 = multiply(%x, %scale)
   %2 = conv2d(%1, %w, data_layout="NHWC")
   

   Transformed:

   # scale weight's input channel
   %1 = multiply(%w, expand_dims(%scale, axis=1, num_newaxis=2))
   %2 = conv2d(%x, %1, data_layout="NHWC")
   

4. CanonicalizeCast()

   规范化强制转换表达式以提高Operator fusion的效率

5. CanonicalizeOps()

   将特殊运算符规范化为基本运算符。可以简化随后的分析 e.g. expanding bias_add to expand_dims and broadcast_add.

6. class tvm.relay.transform.ChangeBatch(data, batch_size=16)

   改变batch大小

   data (Dict[relay.Var, int]) – A dictionary of all the params to change. The keys are all params, and the values are which dimension hold the batch.

   batch_size (int) – The batch size to change to

7. CombineParallelConv2D(min_num_branches=3)

   将多个con2d算子合并为一个，会将具有相同输入的卷积合并成一个大的卷积运算。

   **min_num_branches **执行此优化所需的最小并行分支数

8. CombineParallelDense(min_num_branches=3)

   将多个denese operator 符合并为一个

9. ConvertLayout(desired_layouts)

   给定一个dest布局，此过程转换expr，以便将大多数ops输入数据布局更改为dest布局。在理想情况下，只有两个布局转换，一个在开始，一个在结束。

   此pass不是relay.buid的一部分，预期的用途是在framework-to-relay解析器和relay.build模块调用之间调用。这对于只支持/偏好数据布局类型的硬件后端非常有用。

   此过程大多以AlterOpLayout 和 InferCorrectLayout为基础。现在可以为conv2d操作定义新的布局。

   大多数其他算子都使用InferCorrectLayout基础结构来适应其输入布局

   根据运算符对数据布局的敏感性将其分为3类：

   • Layout agnostic -Relu，pow等。这些运算符不受数据布局的影响，无论功能还是性能。

   • Lightly-layout sensitive -pad, concatenate, 像sum这样的reduce操作等。这些运算符具有某些属性，如果我们在其之前进行布局转换，这些属性会在功能上受到影响。但是，就性能而言，差异并不明显。对于这些操作符，仅适应先前的操作符输出数据布局是有益的。

   • Heavily-layout sensitive -卷积，conv2d_transpose等。这些运算符在功能和性能方面都受到数据布局的严重影响。它们还把数据布局作为操作符属性。通常，为这些运算符修改输入数据布局是有益的（如果它不是高效的数据布局），而其余与布局无关的和轻度布局敏感的运算符将适应由这些重布局敏感的输出所控制的布局操作符。

   用法举例：

   # TFlite framework to Relay parser - Default layout is NHWC
   mod, params = relay.frontend.from_tflite(tflite_model,
                                            shape_dict=shape_dict,
                                            dtype_dict=dtype_dict)
   
   # Convert the layout to NCHW
   # RemoveUnunsedFunctions is used to clean up the graph.
   seq = relay.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
                                     relay.transform.ConvertLayout('NCHW')])
   with relay.transform.PassContext(opt_level=3):
       mod = seq(mod)
   
   # Call relay compilation
   with relay.build_config(opt_level=3):
        graph, lib, params = relay.build(mod, target, params=params)
   
   

10. DeadCodeElimination(inline_once=False)

    移除没有使用者的表达（死代码）

    inline_once (Optional*[Bool]*) – Whether to inline binding that occurs only once

    将没有被索引到的 let bindings 去掉，将被使用过一次的 let bindings 作为代码嵌入

11. DenseToSparse(weight_name, weight_shape)

    将nn.dense operation 写为 nn.sparse_dense。用于data_dep_optimization.bsr_dense。参数来自于analysis.sparse_dense.process_params

    • weight_name (Array[String]) – Names of weights which qualified sparse contrains
    • weight_shape (Array[Array[IntImm]]) – Weights shape in BSR format.

12. EliminateCommonSubexpr(fskip=None)

    消除公共子表达式

    int d = (c*b)*12+a+(a+b*c); E=b*c=c*b 为公共子表达式
    int d = E*12+a+(a+E);
    

    fskip (Callable) – 决定是否应跳过表达式的回调函数。

13. EtaExpand(expand_constructor=False, expand_global_var=False)

    为构造函数添加抽象，或者给一个函数添加全局变量。如： square 被转化为 fn (%x: int32) -> int32 { square(x) }

    • expand_constructor (bool) – Whether to expand constructors.
    • expand_global_var (bool) – Whether to expand global variables.

14. FastMath()

    将非线性激活函数替换成近似计算的算子以获得更快的计算速度。一些情景下有损失计算精度的风险。

15. FoldConstant()

    Fold the constant expressions in a Relay program.折叠常量函数，将可以预先计算的数据放在编译器中完成，减少硬件计算

16. FoldScaleAxis()

    将轴的缩放比例折叠为conv2d/dense的权重。此过程将调用向前和向后缩放折叠。

17. ForwardFoldScaleAxis()

    参考3.BackwardFoldScaleAxis

18. class tvm.relay.transform.FunctionPass

    A pass that works on each tvm.relay.Function in a module. A function pass class should be created through function_pass.

19. FuseOps(fuse_opt_level=- 1)

    根据某些规则，将expr中的operators融合进更大的operator

    融合算子，可以指定融合的优化级别

    fuse_opt_level (int) – The level of fuse optimization. -1 indicates that the level will be inferred from pass context.

    fusion currently doesn’t work well on let bindings. Therefore, we would not be able to fuse operators that were fusable if relay.transform.ToANormalForm() is applied before fusion, as this pass generates let bindings for each expression to canonicalize a Relay program.

20. InferType()

    推断表达式的类型。InferType转换查看运算符的【InferType】属性，确定其输出形状和类型，然后将其传递给下一个运算符InferType属性。获得的结果是一个有显式类别信息的新表达式，以及它的返回类型

​

25. Inline()

    将一个被 inline 标记的全局函数嵌入到 relay 的IR模块中。The global functions that are marked as inline should be always inlined. A cost model will be needed in the future to decide if it is profitable to inline the function

26. LambdaLift()

    将closure提升到全局功能。将局部函数提升为全局函数

    在 src/relay/backend/vm/compiler.cc, src/relay/backend/vm/lambda_lift.cc

27. LazyGradientInit()

    减少梯度张量的内存使用

28. Legalize(legalize_map_attr_name=‘FTVMLegalize’)

    用另一个表达式使表达式合法化。用另一个expr替换一个expr以进行与目标相关的优化。例如，一个expr虽然在表面上等同于另一个expr，但在目标上可以有更好的性能。此pass可用于以目标相关方式使expr合法化。

    legalize_map_attr_name (str) – 与合法化规则函数相对应的Op的attr名称

29. MergeCompilerRegions()

    合并compile（编译）区域

30. MergeComposite(pattern_table)

    将多个模式相匹配的算子合并为一个复合算子。主要用在使用外部代码生成工具时多个算子map到同一个外部算子的情况。定义作为Wrapper使用。
    在 src/relay/pass/annotate_target.cc

    pattern_table (List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Function]**]) A list of (pattern_name, pattern, check) tuples.列表中模式的顺序将决定它们匹配的优先级顺序。“check”是一个检查提取的模式是否匹配的函数。它可以由模式编写器实现，但如果未指定，它将始终返回True。

31. PartialEvaluate()

    计算代码的静态片段。转换可以是Module->Module或Expr->Expr。如果提供了目标表达式，它将直接将输入表达式转换为新表达式。否则，将依赖pass管理器进行改造。

    在编译时评估静态的代码碎片，尽可能多的做常量传播，常量折叠，代码嵌入等优化，以减少运行时开销，获得更多的融合优化。作为代价生成的code量会增加。

32. PartitionGraph()

    将Relay Program 中继程序划分为可在不同后端执行的区域

33. RemoveUnusedFunctions(entry_functions=None)

    移除relay module中未使用的全局relay functions

    entry_functions (list[string]) – The set of entry functions to start from

34. RewriteAnnotatedOps(fallback_device)

    重写带注释的程序，其中注释运算符，例如on_deivce 会标记在哪一种设备上进行调度计算。此过程有助于异构执行，其中可能需要在不同设备上分配不同的运算。

    fallback_device (int) –the fallback device type.及没有注释运算的默认设备

35. SimplifyFCTranspose(target_weight_name)

    Rewrite y = nn.dense(x, transpose(w, [1, 0])) to y = nn.dense(x, wt) This pass is used in data_dep_optimization.simplify_fc_transpose

    weight_name (Array[String]) – Names of weights which qualified y = nn.dense(x, transpose(w, [1, 0])) This parameter is generated by analysis.search_fc_transpose function

36. SimplifyInference()

    简化推理阶段的数据流图。将返回语义上等于输入表达式的简化表达式

37. ToANormalForm()

    Turn Graph Normal Form expression into A Normal Form Expression. 将一个数据流图转化为行政范式 (Administrative Normal Form, A-Normal Form, ANF)。将一个表达从隐式共享的图的形式转化为显式共享， 也就是 ANF。

    根表达式的作用域是全局作用域。任何非根表达式的作用域都是其作用域中最不常见的祖先。值按每个作用域中的post-DFS顺序排序。

38. ToCPS(expr, mod=None)

    将表达式转换为连续传递样式continuation passing style(CPS)

    CPS 意思是每一个函数将不会直接返回结果，而是传递一个另外的函数，作为参数，然后将结果传到下一个续集。这样每一个函数调用时将会多一个参数，表示其余的计算。每一个中间计算都会被传入一个续集。每个中间计算都会被传入一个续集。

39. ToGraphNormalForm()

    去除所有的 let binding, 并将所有的变量转化为直接的指针索引。返回的表达叫做 graph normal form。

40. build_config(opt_level=2, required_pass=None, disabled_pass=None, trace=None)

    通过设置配置变量配置生成行为。在TVM v0.7中被弃用，直接使用tvm.transform.PassContext.

    opt_level (int, optional) – Optimization level.

    OPT_PASS_LEVEL = {
        "SimplifyInference": 0,
        "OpFusion": 1,`FuseOps`
        "FoldConstant": 2,
        "FoldScaleAxis": 3,
        "AlterOpLayout": 3,
        "CanonicalizeOps": 3,
        "CanonicalizeCast": 3,
        "EliminateCommonSubexpr": 3,
        "CombineParallelConv2D": 4,
        "CombineParallelDense": 4,
        "FastMath": 4
    }
    
    only the passes that have optimization level less or equal to 3 will be executed by default under tvm.transform.Sequential.
    seq=...
    with tvm.transform.PassContext(opt_level=3):
        mod2 = seq(mod)
    print(mod2)
    ##disable EliminateCommonSubexpr
    with tvm.transform.PassContext(opt_level=3, disabled_pass=["EliminateCommonSubexpr"]):
        mod3 = seq(mod)
    print(mod3)
    

    required_pass (set of str*,* optional) – Optimization passes that are required regardless of optimization level.

    disabled_pass (set of str*,* optional) – Optimization passes to be disabled during optimization.

    trace (Callable[**[IRModule, PassInfo, bool]**, None]) – A tracing function for debugging or introspection.

41. function_pass(pass_func=None, opt_level=None, name=None, required=None)

    提供pass-func时，此函数返回回调。否则，它将使用给定的优化函数返回已创建的函数过程。

    Example…

42. gradient(expr, mod=None, mode=‘higher_order’)

    转换输入函数，返回一个计算原始结果的函数，与输入的梯度配对

    • expr (tvm.relay.Expr) – The input expression, which is a Function or a GlobalVar.
    • mod (Optional*[tvm.IRModule]*) –
    • mode (Optional*[String]*) – The mode of the automatic differentiation algorithm. 自动微分算法 ‘first_order’ only works on first order code, but will not produce reference nor closure. ‘higher_order’ works on all code using reference and closure.

    Returns expr – The transformed expression.

43. to_cps(func, mod=None)

    转化为CPS表达式。Every intermediate compute will be passed to a continuation.

    • func (tvm.relay.Function) – The input function.
    • mod (Optional*[tvm.IRModule]*) – The global module.

    Returns result – The output function.

44. un_cps(func)

    将cps函数转换为不带continuation参数的函数。

    注意，这不会提供与之前的cps完全相同的接口：如果输入/输出是高阶的，它们仍然是cps形式。

    func (tvm.relay.Function) – The input function

    Returns result – The output function

参考链接：
https://www.zhihu.com/search?type=content&q=Relay%20Pass%20in%20TVM
https://github.com/apache/incubator-tvm/pull/2020
https://blog.csdn.net/weixin_42164269/article/details/104291698
https://tvm.apache.org/docs/dev/convert_layout.html