TVM/Relay 的 PartitionGraph()(mod) 函数讨论整理
TVM/Relay 的图形分区功能。以下简单示例，错误信息。
PartitionGraph() 函数指定图形是用带有 AnnotateTarget([“target”]) 函数的目标注释的。编写了以下示例，以便能够将“add”运算符划分为一个单独的功能函数（使用relay模式语言，或遍历 AST，将add划分为一个单独的relay函数），试图了解 PartitionGraph() 如何在简单情况下工作。
这是代码：
graph_type =1
def _register_external_op_helper(op_name, supported=True):
@tvm.ir.register_op_attr(op_name, "target.special")
def _func_wrapper(attrs, args):
return supported
return _func_wrapper
_register_external_op_helper("add")
_register_external_op_helper("subtract")
if graph_type == 1:
# this is test case for graph type 1
print("Graph type 1")
# graph 1: true branch
x1 = relay.var('x', shape=(10, 1))
y1 = relay.var('y', shape=(10, 1))
# graph 2: false branch
x2 = relay.var('x', shape=(10, 1))
y2 = relay.var('y', shape=(10, 1))
f1 = relay.op.add(x1, y1)
f2 = relay.op.multiply(x2, y2)
cond = relay.var('c')
result = relay.If(cond, true_branch=f1, false_branch=f2)
f = relay.Function([], result)
mod = tvm.IRModule({"main": f})
mod = relay.transform.AnnotateTarget(["special"])(mod) # ==> It GIVES ERROR here
mod = relay.transform.PartitionGraph()(mod) #
这是错误信息。
Graph type 1
Traceback (most recent call last):
File "C:\Program Files\JetBrains\PyCharm 2020.1.2\plugins\python\helpers\pydev\pydevd.py", line 1438, in _exec
pydev_imports.execfile(file, globals, locals) # execute the script
File "C:\Program Files\JetBrains\PyCharm 2020.1.2\plugins\python\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "C:/repos/tvm23/tvm/graph_opt/subgraph/PartitionGraphTry.py", line 48, in <module>
mod = relay.transform.AnnotateTarget(["special"])(mod) # Output: Figure 2
File "C:\repos\tvm23\tvm\python\tvm\ir\transform.py", line 127, in __call__
return _ffi_transform_api.RunPass(self, mod)
File "C:\repos\tvm23\tvm\python\tvm\_ffi\_ctypes\packed_func.py", line 237, in __call__
raise get_last_ffi_error()
tvm._ffi.base.TVMError: Traceback (most recent call last):
File "C:\repos\tvm23\tvm\src\ir\module.cc", line 192
TVMError: Check failed: fv.size() == 0 (5 vs. 0) : There are free variables: [Var(c, ty=TensorType([], bool)), Var(x, ty=TensorType([10, 1], float32)), Var(y, ty=TensorType([10, 1], float32)), Var(x, ty=TensorType([10, 1], float32)), Var(y, ty=TensorType([10, 1], float32))] in function: #[version = "0.0.5"]
fn () -> Tensor[(10, 1), float32] {
free_var %c: bool;
if (%c) {
free_var %x: Tensor[(10, 1), float32];
free_var %y: Tensor[(10, 1), float32];
add(%x, %y) /* ty=Tensor[(10, 1), float32] */
} else {
free_var %x1: Tensor[(10, 1), float32];
free_var %y1: Tensor[(10, 1), float32];
multiply(%x1, %y1) /* ty=Tensor[(10, 1), float32] */
}
}
可能的错误原因
1） the if/else handling in this pass might not be correct.
2） apache/incubator-tvm/blob/main/tests/python/relay/test_pass_annotate_target.py
f = relay.Function([x], out)
mod = tvm.IRModule.from_expr(f)
return mod
mod = transform.AnnotateTarget("A")(before())
mod = transform.AnnotateTarget("B")(mod)
expected = transform.AnnotateTarget(["A", "B"])(before())
assert tvm.ir.structural_equal(expected, mod)
def test_if_else():
target = "test_if_else"
@tvm.ir.register_op_attr("equal", "target." + target)
def relu(attrs, args): # pylint: disable=unused-variable
return True
@tvm.ir.register_op_attr("tanh", "target." + target)
def tanh(attrs, args): # pylint: disable=unused-variable
return True
3） Isn’t it simply a problem of free variables? I suggest replacing
f = relay.Function([], result)
with
f = relay.Function(relay.analysis.free_vars(result), result)
4） 现在调通了。
想确认
1) relay中 PartitionGraph() 函数的功能，
2) ParitioGraph() 是否可用于特定用例。
这是对 PartitionGraph() 函数如何工作的理解：
这是主要问题：
- 注释是按运算符类型完成的，例如“add”，而不是运算符实例。例如，如果true 和 false 分支中有两个“add”运算符，并且想将 true 和 false 分支分开，PartitionGraph() 可以帮助吗？可以覆盖 ExprMutator 类中的 visit_if() 函数，实现刚刚描述的内容，正在为更复杂的问题寻找更高级的解决方案。
ParitioGraph() 似乎有限，基于附加到每个运算符种类的注释进行分区。
理想情况下，想要一个解决方案，执行以下操作：
- 将基于用户提供的表达式注释的relay IRModule 划分为单独的 Relay IR 函数（或 IRModule）
在 mod = tvm.IRModule({“main”: f}) 之后
print(mod)
def @main(%c, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32], %x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) {
if (%c) {
add(%x, %y)
} else {
multiply(%x1, %y1)
}
}
注释后： mod = relay.transform.AnnotateTarget([“special”])(mod)
print(mod)
def @main(%c: bool, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32], %x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) -> Tensor[(10, 1), float32] {
%0 = annotation.compiler_begin(%c, meta[relay.attrs.CompilerAttrs][0]) /* ty=bool */;
%9 = if (%0) {
%1 = annotation.compiler_begin(%x, meta[relay.attrs.CompilerAttrs][1]) /* ty=Tensor[(10, 1), float32] */;
%2 = annotation.compiler_begin(%y, meta[relay.attrs.CompilerAttrs][2]) /* ty=Tensor[(10, 1), float32] */;
%3 = add(%1, %2) /* ty=Tensor[(10, 1), float32] */;
%4 = annotation.compiler_end(%3, meta[relay.attrs.CompilerAttrs][3]) /* ty=Tensor[(10, 1), float32] */;
annotation.compiler_begin(%4, meta[relay.attrs.CompilerAttrs][4]) /* ty=Tensor[(10, 1), float32] */
} else {
%5 = annotation.compiler_begin(%x1, meta[relay.attrs.CompilerAttrs][5]) /* ty=Tensor[(10, 1), float32] */;
%6 = annotation.compiler_begin(%y1, meta[relay.attrs.CompilerAttrs][6]) /* ty=Tensor[(10, 1), float32] */;
%7 = multiply(%5, %6) /* ty=Tensor[(10, 1), float32] */;
%8 = annotation.compiler_end(%7, meta[relay.attrs.CompilerAttrs][7]) /* ty=Tensor[(10, 1), float32] */;
annotation.compiler_begin(%8, meta[relay.attrs.CompilerAttrs][8]) /* ty=Tensor[(10, 1), float32] */
};
annotation.compiler_end(%9, meta[relay.attrs.CompilerAttrs][9]) /* ty=Tensor[(10, 1), float32] */
}
在 mod = relay.transform.PartitionGraph()(mod) 之后
def @special_0(%special_0_i0: Tensor[(10, 1), float32], %special_0_i1: Tensor[(10, 1), float32], global_symbol="special_0", Primitive=1, Compiler="special", Inline=1) -> Tensor[(10, 1), float32] {
add(%special_0_i0, %special_0_i1) /* ty=Tensor[(10, 1), float32] */
}
def @main(%c: bool, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32], %x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) -> Tensor[(10, 1), float32] {
if (%c) {
@special_0(%x, %y) /* ty=Tensor[(10, 1), float32] */
} else {
multiply(%x1, %y1) /* ty=Tensor[(10, 1), float32] */
}
}
5） 这正是 PartitionGraph 所做的。
这是因为只调用AnnotateTarget-> PartitionGraph。还有另一个过程叫做MergeCompilerRegion删除不必要的注释，所以应该通过AnnotateTarget-> MergeCompilerRegion-> PartitionGraph。
示例的预期结果应该是：
def @special_0(%special_0_i0: Tensor[(10, 1), float32], %special_0_i1: Tensor[(10, 1), float32], global_symbol="special_0", Primitive=1, Compiler="special", Inline=1) -> Tensor[(10, 1), float32] {
add(%special_0_i0, %special_0_i1) /* ty=Tensor[(10, 1), float32] */
}
def @special_1(%special_0_i0: Tensor[(10, 1), float32], %special_0_i1: Tensor[(10, 1), float32], global_symbol="special_0", Primitive=1, Compiler="special", Inline=1) -> Tensor[(10, 1), float32] {
multiply(%special_0_i0, %special_0_i1) /* ty=Tensor[(10, 1), float32] */
}
def @main(%c: bool, %x: Tensor[(10, 1), float32], %y: Tensor[(10, 1), float32], %x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32]) -> Tensor[(10, 1), float32] {
if (%c) {
@special_0(%x, %y) /* ty=Tensor[(10, 1), float32] */
} else {
@special_1(%x1, %y1) /* ty=Tensor[(10, 1), float32] */
}
}
如果不是，可能有一些问题/错误需要修复。
6） 曾尝试使用 MergeCompilerRegion，但有以下代码的错误。
下面的代码有效（注释掉了 MergeCompilerRegion），产生带有 UNMERGED @special _ 定义的输出。理想情况下，希望为 true 分支中的表达式设置一个分区，为 false 分支获取另一个分区。
def _register_external_op_helper(op_name, supported=True):
@tvm.ir.register_op_attr(op_name, "target.special")
def _func_wrapper(attrs, args):
return supported
return _func_wrapper
_register_external_op_helper("multiply")
_register_external_op_helper("add")
_register_external_op_helper("subtract")
if graph_type == 1:
# this is test case for graph type 1
print("Graph type 1")
# graph 1: true branch
x1 = relay.var('x1', shape=(10, 1))
y1 = relay.var('y1', shape=(10, 1))
f1 = relay.op.multiply(x1, y1)
x3 = relay.var('x3', shape=(10, 1))
y3 = relay.var('y3', shape=(10, 1))
f3 = relay.op.multiply(x3, y3)
true_branch = relay.op.add(f1, f3)
# graph 2: false branch
x2 = relay.var('x2', shape=(10, 1))
y2 = relay.var('y2', shape=(10, 1))
f2 = relay.op.add(x2, y2)
x4 = relay.var('x4', shape=(10, 1))
y4 = relay.var('y4', shape=(10, 1))
f4 = relay.op.add(x4, y4)
false_branch = relay.op.add(f2, f4)
cond = relay.var('c')
result = relay.If(cond, true_branch=true_branch, false_branch=false_branch)
# f = relay.Function([], result)
f = relay.Function(relay.analysis.free_vars(result), result)
mod = tvm.IRModule({"main": f})
mod = relay.transform.AnnotateTarget(["special"])(mod) # Output: Figure 2
#mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.PartitionGraph()(mod) # Output: Figure 4
这是注销注释 MergeCompilerRegions 函数时得到的错误
Graph type 1
Traceback (most recent call last):
File "C:/repos/tvm23/tvm/graph_opt/subgraph/PartitionGraphTry.py", line 62, in <module>
mod = relay.transform.MergeCompilerRegions()(mod)
File "C:\repos\tvm23\tvm\python\tvm\ir\transform.py", line 127, in __call__
return _ffi_transform_api.RunPass(self, mod)
File "C:\repos\tvm23\tvm\python\tvm\_ffi\_ctypes\packed_func.py", line 237, in __call__
raise get_last_ffi_error()
tvm._ffi.base.TVMError: TVMError: Cannot find the corresponding region for end annotation:
#[version = "0.0.5"]
free_var %c: bool;
%0 = annotation.compiler_begin(%c, meta[relay.attrs.CompilerAttrs][0]) /* ty=bool */;
%25 = if (%0) {
free_var %x1: Tensor[(10, 1), float32];
%1 = annotation.compiler_begin(%x1, meta[relay.attrs.CompilerAttrs][1]) /* ty=Tensor[(10, 1), float32] */;
free_var %y1: Tensor[(10, 1), float32];
%2 = annotation.compiler_begin(%y1, meta[relay.attrs.CompilerAttrs][2]) /* ty=Tensor[(10, 1), float32] */;
%3 = multiply(%1, %2) /* ty=Tensor[(10, 1), float32] */;
%4 = annotation.compiler_end(%3, meta[relay.attrs.CompilerAttrs][3]) /* ty=Tensor[(10, 1), float32] */;
%5 = annotation.compiler_begin(%4, meta[relay.attrs.CompilerAttrs][4]) /* ty=Tensor[(10, 1), float32] */;
free_var %x3: Tensor[(10, 1), float32];
%6 = annotation.compiler_begin(%x3, meta[relay.attrs.CompilerAttrs][5]) /* ty=Tensor[(10, 1), float32] */;
free_var %y3: Tensor[(10, 1), float32];
%7 = annotation.compiler_begin(%y3, meta[relay.attrs.CompilerAttrs][6]) /* ty=Tensor[(10, 1), float32] */;
%8 = multiply(%6, %7) /* ty=Tensor[(10, 1), float32] */;
%9 = annotation.compiler_end(%8, meta[relay.attrs.CompilerAttrs][7]) /* ty=Tensor[(10, 1), float32] */;
%10 = annotation.compiler_begin(%9, meta[relay.attrs.CompilerAttrs][8]) /* ty=Tensor[(10, 1), float32] */;
%11 = add(%5, %10) /* ty=Tensor[(10, 1), float32] */;
%12 = annotation.compiler_end(%11, meta[relay.attrs.CompilerAttrs][9]) /* ty=Tensor[(10, 1), float32] */;
annotation.compiler_begin(%12, meta[relay.attrs.CompilerAttrs][10]) /* ty=Tensor[(10, 1), float32] */
} else {
free_var %x2: Tensor[(10, 1), float32];
%13 = annotation.compiler_begin(%x2, meta[relay.attrs.CompilerAttrs][11]) /* ty=Tensor[(10, 1), float32] */;
free_var %y2: Tensor[(10, 1), float32];
%14 = annotation.compiler_begin(%y2, meta[relay.attrs.CompilerAttrs][12]) /* ty=Tensor[(10, 1), float32] */;
%15 = add(%13, %14) /* ty=Tensor[(10, 1), float32] */;
%16 = annotation.compiler_end(%15, meta[relay.attrs.CompilerAttrs][13]) /* ty=Tensor[(10, 1), float32] */;
%17 = annotation.compiler_begin(%16, meta[relay.attrs.CompilerAttrs][14]) /* ty=Tensor[(10, 1), float32] */;
free_var %x4: Tensor[(10, 1), float32];
%18 = annotation.compiler_begin(%x4, meta[relay.attrs.CompilerAttrs][15]) /* ty=Tensor[(10, 1), float32] */;
free_var %y4: Tensor[(10, 1), float32];
%19 = annotation.compiler_begin(%y4, meta[relay.attrs.CompilerAttrs][16]) /* ty=Tensor[(10, 1), float32] */;
%20 = add(%18, %19) /* ty=Tensor[(10, 1), float32] */;
%21 = annotation.compiler_end(%20, meta[relay.attrs.CompilerAttrs][17]) /* ty=Tensor[(10, 1), float32] */;
%22 = annotation.compiler_begin(%21, meta[relay.attrs.CompilerAttrs][18]) /* ty=Tensor[(10, 1), float32] */;
%23 = add(%17, %22) /* ty=Tensor[(10, 1), float32] */;
%24 = annotation.compiler_end(%23, meta[relay.attrs.CompilerAttrs][19]) /* ty=Tensor[(10, 1), float32] */;
annotation.compiler_begin(%24, meta[relay.attrs.CompilerAttrs][20]) /* ty=Tensor[(10, 1), float32] */
};
annotation.compiler_end(%25, meta[relay.attrs.CompilerAttrs][21]) /* ty=Tensor[(10, 1), float32] */
/* For debugging purposes the metadata section has been omitted.
* If you would like to see the full metadata section you can set the
* option to `True` when invoking `astext`.
*/
Process finished with exit code 1
7） 删除了if 语句，现在它可以工作了。
这是否意味着有些 MergeCompilerRegions 还不完全支持 if。
这是有效的代码。
# this is test case for graph type 1
print("Graph type 1")
# graph 1: true branch
x1 = relay.var('x1', shape=(10, 1))
y1 = relay.var('y1', shape=(10, 1))
f1 = relay.op.multiply(x1, y1)
x3 = relay.var('x3', shape=(10, 1))
y3 = relay.var('y3', shape=(10, 1))
f3 = relay.op.multiply(x3, y3)
true_branch = relay.op.add(f1, f3)
# graph 2: false branch
x2 = relay.var('x2', shape=(10, 1))
y2 = relay.var('y2', shape=(10, 1))
f2 = relay.op.add(x2, y2)
x4 = relay.var('x4', shape=(10, 1))
y4 = relay.var('y4', shape=(10, 1))
f4 = relay.op.add(x4, y4)
false_branch = relay.op.add(f2, f4)
cond = relay.var('c')
#result = relay.If(cond, true_branch=true_branch, false_branch=false_branch)
result = true_branch
#f = relay.Function([], result)
f = relay.Function(relay.analysis.free_vars(result), result)
mod = tvm.IRModule({"main": f})
mod = relay.transform.AnnotateTarget(["special"])(mod) # Output: Figure 2
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.PartitionGraph()(mod) # Output: Figure 4
这是不起作用的代码。
# this is test case for graph type 1
print("Graph type 1")
# graph 1: true branch
x1 = relay.var('x1', shape=(10, 1))
y1 = relay.var('y1', shape=(10, 1))
f1 = relay.op.multiply(x1, y1)
x3 = relay.var('x3', shape=(10, 1))
y3 = relay.var('y3', shape=(10, 1))
f3 = relay.op.multiply(x3, y3)
true_branch = relay.op.add(f1, f3)
# graph 2: false branch
x2 = relay.var('x2', shape=(10, 1))
y2 = relay.var('y2', shape=(10, 1))
f2 = relay.op.add(x2, y2)
x4 = relay.var('x4', shape=(10, 1))
y4 = relay.var('y4', shape=(10, 1))
f4 = relay.op.add(x4, y4)
false_branch = relay.op.add(f2, f4)
cond = relay.var('c')
result = relay.If(cond, true_branch=true_branch, false_branch=false_branch)
#result = true_branch
#f = relay.Function([], result)
f = relay.Function(relay.analysis.free_vars(result), result)
mod = tvm.IRModule({"main": f})
mod = relay.transform.AnnotateTarget(["special"])(mod) # Output: Figure 2
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.PartitionGraph()(mod) # Output: Figure 4
8） 正在 if 节点上工作，这应该已经修复了。
是否尝试过使用最新提交的主分支？
这是使用的脚本：
import tvm
from tvm import relay
def _register_external_op_helper(op_name, supported=True):
@tvm.ir.register_op_attr(op_name, "target.special")
def _func_wrapper(expr):
return supported
return _func_wrapper
_register_external_op_helper("add")
_register_external_op_helper("subtract")
# graph 1: true branch
x1 = relay.var('x1', shape=(10, 1))
y1 = relay.var('y1', shape=(10, 1))
f1 = relay.op.multiply(x1, y1)
x3 = relay.var('x3', shape=(10, 1))
y3 = relay.var('y3', shape=(10, 1))
f3 = relay.op.multiply(x3, y3)
true_branch = relay.op.add(f1, f3)
# graph 2: false branch
x2 = relay.var('x2', shape=(10, 1))
y2 = relay.var('y2', shape=(10, 1))
f2 = relay.op.add(x2, y2)
x4 = relay.var('x4', shape=(10, 1))
y4 = relay.var('y4', shape=(10, 1))
f4 = relay.op.add(x4, y4)
false_branch = relay.op.add(f2, f4)
cond = relay.var('c')
result = relay.If(cond, true_branch=true_branch, false_branch=false_branch)
f = relay.Function(relay.analysis.free_vars(result), result)
mod = tvm.IRModule({"main": f})
mod = relay.transform.AnnotateTarget(["special"])(mod)
mod = relay.transform.MergeCompilerRegions()(mod)
mod = relay.transform.PartitionGraph()(mod)
print(mod)
这是输出，看起来不错。
def @main(%c: bool, %x1: Tensor[(10, 1), float32], %y1: Tensor[(10, 1), float32], %x3: Tensor[(10, 1), float32], %y3: Tensor[(10, 1), float32], %x2: Tensor[(10, 1), float32], %y2: Tensor[(10, 1), float32], %x4: Tensor[(10, 1), float32], %y4: Tensor[(10, 1), float32]) -> Tensor[(10, 1), float32] {
if (%c) {
%0 = multiply(%x1, %y1) /* ty=Tensor[(10, 1), float32] */;
%1 = multiply(%x3, %y3) /* ty=Tensor[(10, 1), float32] */;
@special_0(%0, %1) /* ty=Tensor[(10, 1), float32] */
} else {
@special_2(%x2, %y2, %x4, %y4) /* ty=Tensor[(10, 1), float32] */
}
}
def @special_0(%special_0_i0: Tensor[(10, 1), float32], %special_0_i1: Tensor[(10, 1), float32], global_symbol="special_0", Primitive=1, Compiler="special", Inline=1) -> Tensor[(10, 1), float32] {
add(%special_0_i0, %special_0_i1) /* ty=Tensor[(10, 1), float32] */
}
def @special_2(%special_2_i0: Tensor[(10, 1), float32], %special_2_i1: Tensor[(10, 1), float32], %special_2_i2: Tensor[(10, 1), float32], %special_2_i3: Tensor[(10, 1), float32], global_symbol="special_2", Primitive=1, Compiler="special", Inline=1) -> Tensor[(10, 1), float32] {
%2 = add(%special_2_i0, %special_2_i1) /* ty=Tensor[(10, 1), float32] */;
%3 = add(%special_2_i2, %special_2_i3) /* ty=Tensor[(10, 1), float32] */;
add(%2, %3) /* ty=Tensor[(10, 1), float32] */
}
参考链接：
https://discuss.tvm.apache.org/t/understanding-tvm-relays-partitiongraph-mod-function/8290/10
