计算图分区#
import numpy as np
import tvm
from tvm.relay.backend import te_compiler
from tvm.relay.backend.runtime import Runtime
import tvm.relay.testing
from tvm import relay
from tvm import runtime as tvm_runtime
from tvm.relay import transform
from tvm.contrib import utils
from tvm.relay.expr_functor import ExprMutator
from tvm.relay.op.annotation import compiler_begin, compiler_end
from tvm.relay.op.contrib.register import get_pattern_table
from tvm.relay.build_module import bind_params_by_name
Show code cell content
# Leverage the pass manager to write a simple allowed list based annotator
@transform.function_pass(opt_level=0)
class AllowedListAnnotator:
def __init__(self, op_list, compiler):
assert isinstance(op_list, (list, tuple, set))
self.op_list = op_list
self.compiler = compiler
def transform_function(self, func, mod, dev):
annotator = self
class Annotator(tvm.relay.ExprMutator):
def visit_call(self, call):
op_name = call.op.name
if op_name in annotator.op_list:
new_args = []
for arg in call.args:
ann = compiler_begin(super().visit(arg), annotator.compiler)
new_args.append(ann)
new_call = relay.Call(call.op, new_args, call.attrs, call.type_args)
return compiler_end(new_call, annotator.compiler)
else:
return super().visit_call(call)
return Annotator().visit(func)
class WholeGraphAnnotator(ExprMutator):
"""
An annotator that creates a compiler for an entire graph.
"""
def __init__(self, compiler):
super(WholeGraphAnnotator, self).__init__()
self.compiler = compiler
self.last_call = True
def visit_call(self, call):
curr_last = self.last_call
self.last_call = False
params = []
for arg in call.args:
param = super().visit(arg)
if isinstance(param, relay.expr.Var):
param = compiler_begin(param, self.compiler)
params.append(param)
new_call = relay.Call(call.op, params, call.attrs)
if curr_last:
new_call = compiler_end(new_call, self.compiler)
return new_call
class MobileNetAnnotator(ExprMutator):
"""
Annotate mobilenet until global_avg_pool.
"""
def __init__(self, compiler):
super(MobileNetAnnotator, self).__init__()
self.compiler = compiler
self.compiler_open = False
def visit_call(self, call):
if call.op.name == "nn.global_avg_pool2d":
self.compiler_open = True
compiler_open = self.compiler_open
params = []
for arg in call.args:
param = super().visit(arg)
if call.op.name == "nn.global_avg_pool2d":
param = compiler_end(param, self.compiler)
if compiler_open and isinstance(param, relay.expr.Var):
param = compiler_begin(param, self.compiler)
params.append(param)
new_call = relay.Call(call.op, params, call.attrs)
return new_call
Show code cell content
import os
def check_result(
mod,
map_inputs,
out_shape,
result,
tol=1e-5,
target="llvm",
device=tvm.cpu(),
params=None,
runtime=Runtime("cpp"),
):
def update_lib(lib):
test_dir = os.path.dirname(os.path.realpath(os.path.expanduser("__file__")))
source_dir = os.path.join(test_dir, "..", "..", "..")
contrib_path = os.path.join(source_dir, "src", "runtime", "contrib")
kwargs = {}
kwargs["options"] = ["-O2", "-std=c++17", "-I" + contrib_path]
tmp_path = utils.tempdir()
lib_name = "lib.so"
lib_path = tmp_path.relpath(lib_name)
lib.export_library(lib_path, fcompile=False, **kwargs)
lib = tvm_runtime.load_module(lib_path)
return lib
def check_vm_result():
te_compiler.get().clear()
with tvm.transform.PassContext(opt_level=3):
exe = relay.vm.compile(mod, target=target, params=params)
code, lib = exe.save()
lib = update_lib(lib)
exe = tvm_runtime.vm.Executable.load_exec(code, lib)
vm = tvm_runtime.vm.VirtualMachine(exe, device)
outs = vm.run(**map_inputs)
outs = outs if isinstance(outs, tvm_runtime.container.ADT) else [outs]
results = result if isinstance(result, list) else [result]
for out, ref in zip(outs, results):
tvm.testing.assert_allclose(out.numpy(), ref, rtol=tol, atol=tol)
def check_graph_executor_result():
te_compiler.get().clear()
with tvm.transform.PassContext(opt_level=3):
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
lib = update_lib(lib)
rt_mod = tvm.contrib.graph_executor.create(json, lib, device)
for name, data in map_inputs.items():
rt_mod.set_input(name, data)
rt_mod.set_input(**param)
rt_mod.run()
out_shapes = out_shape if isinstance(out_shape, list) else [out_shape]
results = result if isinstance(result, list) else [result]
for idx, shape in enumerate(out_shapes):
out = tvm.nd.empty(shape, device=device)
out = rt_mod.get_output(idx, out)
tvm.testing.assert_allclose(out.numpy(), results[idx], rtol=tol, atol=tol)
check_vm_result()
check_graph_executor_result()
测试外部 C 编译器处理单一算子#
@transform.function_pass(opt_level=0)
class MyAnnotator:
def transform_function(self, func, mod, dev):
class Annotator(tvm.relay.ExprMutator):
def visit_call(self, call):
new_args = []
for arg in call.args:
ann = compiler_begin(self.visit(arg), "ccompiler")
new_args.append(ann)
new_call = relay.Call(call.op, new_args)
return compiler_end(new_call, "ccompiler")
return Annotator().visit(func)
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
z = x + y
f = relay.Function([x, y], z)
x_data = np.random.rand(8, 8).astype("float32")
y_data = np.random.rand(8, 8).astype("float32")
mod = tvm.IRModule()
mod["main"] = f
mod = MyAnnotator()(mod)
mod = transform.PartitionGraph()(mod)
check_result(mod, {"x": x_data, "y": y_data}, (8, 8), x_data + y_data)
[17:36:37] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module
/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the new recommended usage.
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
测试外部 C 编译器处理默认算子#
def set_func_attr(func, compile_name, symbol_name):
func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func = func.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func = func.with_attr("Compiler", compile_name)
func = func.with_attr("global_symbol", symbol_name)
return func
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
add = x + y
log = relay.log(add)
exp = relay.exp(add)
concat = relay.concatenate([log, exp], axis=0)
f = relay.Function([x, y], concat)
mod = tvm.IRModule()
mod["main"] = f
mod = AllowedListAnnotator(["add", "subtract", "multiply"], "ccompiler")(mod)
mod = transform.PartitionGraph()(mod)
fused_mod = transform.FuseOps(2)(mod)
x_data = np.random.rand(8, 8).astype("float32")
y_data = np.random.rand(8, 8).astype("float32")
np_add = x_data + y_data
res = np.concatenate([np.log(np_add), np.exp(np_add)])
check_result(mod, {"x": x_data, "y": y_data}, (16, 8), res)
/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the new recommended usage.
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
测试外部编译器处理经过清理的算子#
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
add = x + y
log = relay.log(add)
exp = relay.exp(add)
concat = relay.concatenate([log, exp], axis=0)
f = relay.Function([x, y], concat)
mod = tvm.IRModule()
mod["main"] = f
mod = AllowedListAnnotator(["add", "subtract", "multiply"], "unsanitary-name++")(mod)
mod = transform.PartitionGraph()(mod)
fused_mod = transform.FuseOps(2)(mod)
fused_mod.show()
def @main(%x: Tensor[(8, 8), float32] /* ty=Tensor[(8, 8), float32] */, %y: Tensor[(8, 8), float32] /* ty=Tensor[(8, 8), float32] */) -> Tensor[(16, 8), float32] {
%3 = @tvmgen_default_unsanitary_name___main_0(%x, %y) /* ty=Tensor[(8, 8), float32] */;
%4 = fn (%p0: Tensor[(8, 8), float32] /* ty=Tensor[(8, 8), float32] */, Primitive=1) -> Tensor[(16, 8), float32] {
%0 = log(%p0) /* ty=Tensor[(8, 8), float32] */;
%1 = exp(%p0) /* ty=Tensor[(8, 8), float32] */;
%2 = (%0, %1) /* ty=(Tensor[(8, 8), float32], Tensor[(8, 8), float32]) */;
concatenate(%2) /* ty=Tensor[(16, 8), float32] */
} /* ty=fn (Tensor[(8, 8), float32]) -> Tensor[(16, 8), float32] */;
%4(%3) /* ty=Tensor[(16, 8), float32] */
}
def @tvmgen_default_unsanitary_name___main_0(%unsanitary-name++_0_i0: Tensor[(8, 8), float32] /* ty=Tensor[(8, 8), float32] */, %unsanitary-name++_0_i1: Tensor[(8, 8), float32] /* ty=Tensor[(8, 8), float32] */, Compiler="unsanitary-name++", Primitive=1, Inline=1, global_symbol="tvmgen_default_unsanitary_name___main_0") -> Tensor[(8, 8), float32] {
add(%unsanitary-name++_0_i0, %unsanitary-name++_0_i1) /* ty=Tensor[(8, 8), float32] */
}
测试外部 C 编译器处理多个函数#
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
add = x + y
log = relay.log(add)
exp = relay.exp(add)
concat = relay.concatenate([log, exp], axis=0)
f = relay.Function([x, y], concat)
mod = tvm.IRModule()
mod["main"] = f
# define second function
a = relay.var("a", shape=(16, 16))
b = relay.var("b", shape=(16, 16))
add = a + b
log = relay.log(add)
exp = relay.exp(add)
concat = relay.concatenate([log, exp], axis=0)
f2 = relay.Function([a, b], concat)
mod["subfunction"] = f2
mod = AllowedListAnnotator(["add", "subtract", "multiply"], "ccompiler")(mod)
mod = transform.PartitionGraph()(mod)
fused_mod = transform.FuseOps(2)(mod)
x_data = np.random.rand(8, 8).astype("float32")
y_data = np.random.rand(8, 8).astype("float32")
np_add = x_data + y_data
res = np.concatenate([np.log(np_add), np.exp(np_add)])
check_result(mod, {"x": x_data, "y": y_data}, (16, 8), res)
/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the new recommended usage.
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
测试外部 C 编译器#
x = relay.var("x", shape=(2, 2))
y = relay.var("y", shape=(2, 2))
z = x + x
p = y * y
f = relay.Function([x, y], p - z)
x_data = np.random.rand(2, 2).astype("float32")
y_data = np.random.rand(2, 2).astype("float32")
mod = tvm.IRModule()
mod["main"] = f
mod = AllowedListAnnotator(["add", "subtract", "multiply"], "ccompiler")(mod)
mod = transform.PartitionGraph()(mod)
check_result(mod, {"x": x_data, "y": y_data}, (2, 2), (y_data * y_data) - (x_data + x_data))
[17:36:40] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module
/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the new recommended usage.
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
测试函数提升#
Show code cell content
def partition():
data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32"))
weight = relay.var("weight", relay.TensorType((16, 3, 3, 3), "float32"))
bn_gamma = relay.var("bn_gamma", relay.TensorType((16,), "float32"))
bn_beta = relay.var("bn_beta", relay.TensorType((16,), "float32"))
bn_mmean = relay.var("bn_mean", relay.TensorType((16,), "float32"))
bn_mvar = relay.var("bn_var", relay.TensorType((16,), "float32"))
conv = relay.nn.conv2d(
data=data, weight=weight, kernel_size=(3, 3), channels=16, padding=(1, 1)
)
bn_output = relay.nn.batch_norm(conv, bn_gamma, bn_beta, bn_mmean, bn_mvar)
func = relay.Function(
[data, weight, bn_gamma, bn_beta, bn_mmean, bn_mvar], bn_output.astuple()
)
mod = tvm.IRModule()
mod["main"] = func
mod = relay.transform.InferType()(mod)
op_list = ["nn.batch_norm", "nn.conv2d"]
mod = AllowedListAnnotator(op_list, "test_compiler")(mod)
opt_pass = tvm.transform.Sequential(
[
transform.InferType(),
transform.PartitionGraph(),
transform.SimplifyInference(),
transform.FoldConstant(),
transform.AlterOpLayout(),
]
)
with tvm.transform.PassContext(opt_level=3):
mod = opt_pass(mod)
return mod
partitioned = partition()
partitioned.show()
def @main(%data: Tensor[(1, 3, 224, 224), float32] /* ty=Tensor[(1, 3, 224, 224), float32] */, %weight: Tensor[(16, 3, 3, 3), float32] /* ty=Tensor[(16, 3, 3, 3), float32] */, %bn_gamma: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_beta: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_mean: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_var: Tensor[(16), float32] /* ty=Tensor[(16), float32] */) -> (Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) {
%0 = @tvmgen_default_test_compiler_main_0(%data, %weight) /* ty=Tensor[(1, 16, 224, 224), float32] */;
@tvmgen_default_test_compiler_main_2(%0, %bn_gamma, %bn_beta, %bn_mean, %bn_var) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */
}
def @tvmgen_default_test_compiler_main_0(%test_compiler_0_i0: Tensor[(1, 3, 224, 224), float32] /* ty=Tensor[(1, 3, 224, 224), float32] */, %test_compiler_0_i1: Tensor[(16, 3, 3, 3), float32] /* ty=Tensor[(16, 3, 3, 3), float32] */, Compiler="test_compiler", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_compiler_main_0") -> Tensor[(1, 16, 224, 224), float32] {
nn.conv2d(%test_compiler_0_i0, %test_compiler_0_i1, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 224, 224), float32] */
}
def @tvmgen_default_test_compiler_main_2(%test_compiler_2_i0: Tensor[(1, 16, 224, 224), float32] /* ty=Tensor[(1, 16, 224, 224), float32] */, %test_compiler_2_i1: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_compiler_2_i2: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_compiler_2_i3: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_compiler_2_i4: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, Compiler="test_compiler", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_compiler_main_2") -> (Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) {
nn.batch_norm(%test_compiler_2_i0, %test_compiler_2_i1, %test_compiler_2_i2, %test_compiler_2_i3, %test_compiler_2_i4) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */
}
测试函数提升内联#
def partition():
data = relay.var("data", relay.TensorType((1, 16, 224, 224), "float32"))
bn_gamma = relay.var("bn_gamma", relay.TensorType((16,), "float32"))
bn_beta = relay.var("bn_beta", relay.TensorType((16,), "float32"))
bn_mmean = relay.var("bn_mean", relay.TensorType((16,), "float32"))
bn_mvar = relay.var("bn_var", relay.TensorType((16,), "float32"))
bn_output = relay.nn.batch_norm(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)
func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar], bn_output.astuple())
mod = tvm.IRModule()
mod["main"] = func
op_list = ["nn.batch_norm", "nn.conv2d"]
mod = AllowedListAnnotator(op_list, "test_compiler")(mod)
opt_pass = tvm.transform.Sequential(
[
transform.InferType(),
transform.PartitionGraph(),
transform.SimplifyInference(),
transform.FoldConstant(),
transform.AlterOpLayout(),
transform.Inline(),
]
)
with tvm.transform.PassContext(opt_level=3):
mod = opt_pass(mod)
return mod
partitioned = partition()
partitioned.show()
def @main(%data: Tensor[(1, 16, 224, 224), float32] /* ty=Tensor[(1, 16, 224, 224), float32] */, %bn_gamma: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_beta: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_mean: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_var: Tensor[(16), float32] /* ty=Tensor[(16), float32] */) -> (Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) {
%0 = fn (%test_compiler_0_i0: Tensor[(1, 16, 224, 224), float32] /* ty=Tensor[(1, 16, 224, 224), float32] */, %test_compiler_0_i1: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_compiler_0_i2: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_compiler_0_i3: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_compiler_0_i4: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, Compiler="test_compiler", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_compiler_main_0") -> (Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) {
nn.batch_norm(%test_compiler_0_i0, %test_compiler_0_i1, %test_compiler_0_i2, %test_compiler_0_i3, %test_compiler_0_i4) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */
};
%0(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var)
}
测试常量传播#
ones = np.ones(shape=(8, 8), dtype="float32")
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
add = x + y
log = relay.log(add)
f = relay.Function([x, y], log)
f = bind_params_by_name(f, {"x": tvm.nd.array(ones)})
mod = tvm.IRModule()
mod["main"] = f
mod = AllowedListAnnotator(["add"], "ccompiler")(mod)
mod = transform.PartitionGraph()(mod)
mod = relay.transform.InferType()(mod)
y_data = np.random.rand(8, 8).astype("float32")
np_add = ones + y_data
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the new recommended usage.
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
测试多输出#
Show code cell content
def create_graph():
data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32"))
weight = relay.var("weight", relay.TensorType((16, 3, 3, 3), "float32"))
bn_gamma = relay.var("bn_gamma", relay.TensorType((16,), "float32"))
bn_beta = relay.var("bn_beta", relay.TensorType((16,), "float32"))
bn_mean = relay.var("bn_mean", relay.TensorType((16,), "float32"))
bn_var = relay.var("bn_var", relay.TensorType((16,), "float32"))
data_cb = compiler_begin(data, "test_target")
weight_cb = compiler_begin(weight, "test_target")
bn_gamma_cb = compiler_begin(bn_gamma, "test_target")
bn_beta_cb = compiler_begin(bn_beta, "test_target")
bn_mean_cb = compiler_begin(bn_mean, "test_target")
bn_var_cb = compiler_begin(bn_var, "test_target")
conv_o = relay.nn.conv2d(
data=data_cb, weight=weight_cb, kernel_size=(3, 3), channels=16, padding=(1, 1)
)
bn_o = relay.nn.batch_norm(conv_o, bn_gamma_cb, bn_beta_cb, bn_mean_cb, bn_var_cb)
relu_o = relay.nn.relu(bn_o[0])
relu_o_ce = compiler_end(relu_o, "test_target")
bn_omean = bn_o[1]
rebn_omean_ce = compiler_end(bn_omean, "test_target")
bn_ovar = bn_o[2]
bn_ovar_ce = compiler_end(bn_ovar, "test_target")
dummy_mean_abs = relay.abs(rebn_omean_ce)
dummy_ovar_abs = relay.abs(bn_ovar_ce)
dummy_tuple = relay.Tuple((relu_o_ce, dummy_mean_abs, dummy_ovar_abs))
func = relay.Function([data, weight, bn_gamma, bn_beta, bn_mean, bn_var], dummy_tuple)
return func
mod = tvm.IRModule()
mod["main"] = create_graph()
partitioned = transform.PartitionGraph()(mod)
partitioned.show()
def @main(%data: Tensor[(1, 3, 224, 224), float32] /* ty=Tensor[(1, 3, 224, 224), float32] */, %weight: Tensor[(16, 3, 3, 3), float32] /* ty=Tensor[(16, 3, 3, 3), float32] */, %bn_gamma: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_beta: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_mean: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %bn_var: Tensor[(16), float32] /* ty=Tensor[(16), float32] */) -> (Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) {
%0 = @tvmgen_default_test_target_main_0(%data, %weight, %bn_gamma, %bn_beta, %bn_mean, %bn_var) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */;
%1 = %0.1 /* ty=Tensor[(16), float32] */;
%2 = %0.2 /* ty=Tensor[(16), float32] */;
%3 = %0.0 /* ty=Tensor[(1, 16, 224, 224), float32] */;
%4 = abs(%1) /* ty=Tensor[(16), float32] */;
%5 = abs(%2) /* ty=Tensor[(16), float32] */;
(%3, %4, %5) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */
}
def @tvmgen_default_test_target_main_0(%test_target_0_i0: Tensor[(1, 3, 224, 224), float32] /* ty=Tensor[(1, 3, 224, 224), float32] */, %test_target_0_i1: Tensor[(16, 3, 3, 3), float32] /* ty=Tensor[(16, 3, 3, 3), float32] */, %test_target_0_i2: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_target_0_i3: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_target_0_i4: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, %test_target_0_i5: Tensor[(16), float32] /* ty=Tensor[(16), float32] */, Compiler="test_target", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_target_main_0") -> (Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) {
%6 = nn.conv2d(%test_target_0_i0, %test_target_0_i1, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 224, 224), float32] */;
%7 = nn.batch_norm(%6, %test_target_0_i2, %test_target_0_i3, %test_target_0_i4, %test_target_0_i5) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */;
%8 = %7.0 /* ty=Tensor[(1, 16, 224, 224), float32] */;
%9 = nn.relu(%8) /* ty=Tensor[(1, 16, 224, 224), float32] */;
%10 = %7.1 /* ty=Tensor[(16), float32] */;
%11 = %7.2 /* ty=Tensor[(16), float32] */;
(%9, %10, %11) /* ty=(Tensor[(1, 16, 224, 224), float32], Tensor[(16), float32], Tensor[(16), float32]) */
}
测试混合单输出和多输出#
def create_graph():
data = relay.var("data", shape=(10, 10))
cb_1 = compiler_begin(data, "test_target")
O_1 = relay.abs(cb_1)
ce_2 = compiler_end(O_1, "test_target")
O_2 = relay.nn.relu(O_1)
ce_3 = compiler_end(O_2, "test_target")
X = relay.tanh(ce_2)
cb_3 = compiler_begin(ce_3, "test_target")
cb_4 = compiler_begin(X, "test_target")
O_3 = relay.add(cb_3, cb_4)
ce_4 = compiler_end(O_3, "test_target")
func = relay.Function([data], ce_4)
return func
mod = tvm.IRModule()
mod["main"] = create_graph()
mod = transform.InferType()(mod)
partitioned = transform.PartitionGraph()(mod)
partitioned.show()
def @main(%data: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */) -> Tensor[(10, 10), float32] {
%0 = @tvmgen_default_test_target_main_0(%data) /* ty=(Tensor[(10, 10), float32], Tensor[(10, 10), float32]) */;
%1 = %0.1 /* ty=Tensor[(10, 10), float32] */;
%2 = %0.0 /* ty=Tensor[(10, 10), float32] */;
%3 = tanh(%1) /* ty=Tensor[(10, 10), float32] */;
@tvmgen_default_test_target_main_1(%2, %3) /* ty=Tensor[(10, 10), float32] */
}
def @tvmgen_default_test_target_main_0(%test_target_0_i0: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, Compiler="test_target", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_target_main_0") -> (Tensor[(10, 10), float32], Tensor[(10, 10), float32]) {
%4 = abs(%test_target_0_i0) /* ty=Tensor[(10, 10), float32] */;
%5 = nn.relu(%4) /* ty=Tensor[(10, 10), float32] */;
(%5, %4) /* ty=(Tensor[(10, 10), float32], Tensor[(10, 10), float32]) */
}
def @tvmgen_default_test_target_main_1(%test_target_1_i0: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, %test_target_1_i1: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, Compiler="test_target", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_target_main_1") -> Tensor[(10, 10), float32] {
add(%test_target_1_i0, %test_target_1_i1) /* ty=Tensor[(10, 10), float32] */
}
测试一个输出被多次使用的情况#
def expected_same_output_region():
mod = tvm.IRModule()
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
z = relay.var("z", shape=(8, 8))
x0 = relay.var("x0", shape=(8, 8))
y0 = relay.var("y0", shape=(8, 8))
log = relay.log(x0)
sub = x0 - y0
mul = log * sub
# The partitioned graph contains log, subtract, and multiply
func = relay.Function([x0, y0], mul)
func = set_func_attr(func, "ccompiler", "tvmgen_default_ccompiler_main_0")
glb_0 = relay.GlobalVar("tvmgen_default_ccompiler_main_0")
mod[glb_0] = func
mod = transform.InferType()(mod)
add = x + y
call = relay.Call(glb_0, [add, z])
main = relay.Function([x, y, z], call)
mod["main"] = main
mod = transform.InferType()(mod)
return mod
def expected_different_output_region():
mod = tvm.IRModule()
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
z = relay.var("z", shape=(8, 8))
# The partitioned graph contains log
i0 = relay.var("i0", shape=(8, 8))
log = relay.log(i0)
func = relay.Function([i0], log)
func = set_func_attr(func, "ccompiler", "tvmgen_default_ccompiler_main_0")
glb_0 = relay.GlobalVar("tvmgen_default_ccompiler_main_0")
mod[glb_0] = func
mod = transform.InferType()(mod)
# The partitioned graph contains subtract
x0 = relay.var("x0", shape=(8, 8))
y0 = relay.var("y0", shape=(8, 8))
sub = x0 - y0
func = relay.Function([x0, y0], sub)
func = set_func_attr(func, "ccompiler", "tvmgen_default_ccompiler_main_1")
glb_1 = relay.GlobalVar("tvmgen_default_ccompiler_main_1")
mod[glb_1] = func
mod = transform.InferType()(mod)
add = x + y
call_log = relay.Call(glb_0, [add])
call_sub = relay.Call(glb_1, [add, z])
main = relay.Function([x, y, z], call_log * call_sub)
mod["main"] = main
mod = transform.InferType()(mod)
return mod
def get_mod():
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
z = relay.var("z", shape=(8, 8))
add = x + y
sub = add - z
log = relay.log(add)
sub1 = log * sub
f = relay.Function([x, y, z], sub1)
mod = tvm.IRModule()
mod["main"] = f
return mod
def test_same_output_region():
mod = get_mod()
mod = AllowedListAnnotator(["subtract", "log", "multiply"], "ccompiler")(mod)
mod = transform.MergeCompilerRegions()(mod)
mod = transform.PartitionGraph()(mod)
expected_mod = expected_same_output_region()
tvm.ir.assert_structural_equal(mod, expected_mod, map_free_vars=True)
def test_different_output_region():
mod = get_mod()
mod = AllowedListAnnotator(["subtract", "log"], "ccompiler")(mod)
mod = transform.MergeCompilerRegions()(mod)
mod = transform.PartitionGraph()(mod)
expected_mod = expected_different_output_region()
tvm.ir.assert_structural_equal(mod, expected_mod, map_free_vars=True)
test_same_output_region()
test_different_output_region()
测试重复输出的情况#
target = "test_duplicate_outputs"
@tvm.ir.register_op_attr("abs", "target." + target)
def abs(expr): # pylint: disable=unused-variable
return True
def create_graph():
data = relay.var("data", shape=(10, 10))
x = relay.abs(data)
out_1 = relay.nn.relu(x)
out_2 = relay.tanh(x)
out_3 = relay.log(x)
out = relay.Tuple([out_1, out_2, out_3])
func = relay.Function([data], out)
return func
mod = tvm.IRModule()
mod["main"] = create_graph()
seq = tvm.transform.Sequential(
[
transform.AnnotateTarget(target),
transform.MergeCompilerRegions(),
transform.PartitionGraph(),
]
)
partitioned = seq(mod)
partitioned.show()
def @main(%data: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */) -> (Tensor[(10, 10), float32], Tensor[(10, 10), float32], Tensor[(10, 10), float32]) {
%0 = @tvmgen_default_test_duplicate_outputs_main_0(%data) /* ty=Tensor[(10, 10), float32] */;
%1 = nn.relu(%0) /* ty=Tensor[(10, 10), float32] */;
%2 = tanh(%0) /* ty=Tensor[(10, 10), float32] */;
%3 = log(%0) /* ty=Tensor[(10, 10), float32] */;
(%1, %2, %3) /* ty=(Tensor[(10, 10), float32], Tensor[(10, 10), float32], Tensor[(10, 10), float32]) */
}
def @tvmgen_default_test_duplicate_outputs_main_0(%test_duplicate_outputs_0_i0: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, Compiler="test_duplicate_outputs", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_duplicate_outputs_main_0") -> Tensor[(10, 10), float32] {
abs(%test_duplicate_outputs_0_i0) /* ty=Tensor[(10, 10), float32] */
}
测试重复合并和 TupleGetItem 的情况#
target = "test_duplicate_merge_and_tuplegetitem"
@tvm.ir.register_op_attr("nn.batch_norm", "target." + target)
def batch_norm(expr): # pylint: disable=unused-variable
return True
@tvm.ir.register_op_attr("nn.relu", "target." + target)
def relu(expr): # pylint: disable=unused-variable
return True
def create_graph():
data = relay.var("data", shape=(10, 10))
bn_gamma = relay.var("bn_gamma")
bn_beta = relay.var("bn_beta")
bn_mmean = relay.var("bn_mean")
bn_mvar = relay.var("bn_var")
x = relay.nn.batch_norm(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)
out_1 = relay.nn.relu(x[0])
bn_out_1 = x[1]
out_2 = relay.tanh(bn_out_1)
out_3 = relay.log(bn_out_1)
out = relay.Tuple([out_1, out_2, out_3])
func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar], out)
return func
mod = tvm.IRModule()
mod["main"] = create_graph()
mod = transform.InferType()(mod)
seq = tvm.transform.Sequential(
[
transform.AnnotateTarget(target),
transform.MergeCompilerRegions(),
transform.PartitionGraph(),
]
)
partitioned = seq(mod)
partitioned.show()
def @main(%data: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, %bn_gamma: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %bn_beta: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %bn_mean: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %bn_var: Tensor[(10), float32] /* ty=Tensor[(10), float32] */) -> (Tensor[(10, 10), float32], Tensor[(10), float32], Tensor[(10), float32]) {
%0 = @tvmgen_default_test_duplicate_merge_and_tuplegetitem_main_0(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var) /* ty=(Tensor[(10, 10), float32], Tensor[(10), float32]) */;
%1 = %0.1 /* ty=Tensor[(10), float32] */;
%2 = %0.0 /* ty=Tensor[(10, 10), float32] */;
%3 = tanh(%1) /* ty=Tensor[(10), float32] */;
%4 = log(%1) /* ty=Tensor[(10), float32] */;
(%2, %3, %4) /* ty=(Tensor[(10, 10), float32], Tensor[(10), float32], Tensor[(10), float32]) */
}
def @tvmgen_default_test_duplicate_merge_and_tuplegetitem_main_0(%test_duplicate_merge_and_tuplegetitem_0_i0: Tensor[(10, 10), float32] /* ty=Tensor[(10, 10), float32] */, %test_duplicate_merge_and_tuplegetitem_0_i1: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %test_duplicate_merge_and_tuplegetitem_0_i2: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %test_duplicate_merge_and_tuplegetitem_0_i3: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, %test_duplicate_merge_and_tuplegetitem_0_i4: Tensor[(10), float32] /* ty=Tensor[(10), float32] */, Compiler="test_duplicate_merge_and_tuplegetitem", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_duplicate_merge_and_tuplegetitem_main_0") -> (Tensor[(10, 10), float32], Tensor[(10), float32]) {
%5 = nn.batch_norm(%test_duplicate_merge_and_tuplegetitem_0_i0, %test_duplicate_merge_and_tuplegetitem_0_i1, %test_duplicate_merge_and_tuplegetitem_0_i2, %test_duplicate_merge_and_tuplegetitem_0_i3, %test_duplicate_merge_and_tuplegetitem_0_i4) /* ty=(Tensor[(10, 10), float32], Tensor[(10), float32], Tensor[(10), float32]) */;
%6 = %5.0 /* ty=Tensor[(10, 10), float32] */;
%7 = nn.relu(%6) /* ty=Tensor[(10, 10), float32] */;
%8 = %5.1 /* ty=Tensor[(10), float32] */;
(%7, %8) /* ty=(Tensor[(10, 10), float32], Tensor[(10), float32]) */
}
测试常量元组#
@tvm.ir.register_op_attr("qnn.concatenate", "target.const_tuples")
def add(expr): # pylint: disable=unused-variable
return True
def create_graph():
a = relay.var("a", shape=(10, 10), dtype="uint8")
b = relay.var("b", shape=(10, 10), dtype="uint8")
a1 = relay.abs(a)
zeroi = relay.const(1, "int32")
zerof = relay.const(0, "float32")
con = relay.qnn.op.concatenate(
(a1, b),
input_scales=(zerof, zerof),
input_zero_points=(zeroi, zeroi),
output_scale=zerof,
output_zero_point=zeroi,
axis=1,
)
f = relay.Function([a, b], con)
mod = tvm.IRModule.from_expr(f)
mod = transform.InferType()(mod)
return mod
seq = tvm.transform.Sequential(
[
transform.AnnotateTarget("const_tuples"),
transform.InferType(),
transform.MergeCompilerRegions(),
transform.PartitionGraph(),
]
)
partitioned = seq(create_graph())
concat = partitioned["tvmgen_default_const_tuples_main_0"].body
assert type(concat.args[1]) == relay.Tuple
assert type(concat.args[2]) == relay.Tuple
assert type(concat.args[3]) == relay.Constant
assert type(concat.args[4]) == relay.Constant
测试扁平化元组输出#
target = "test_flatten_tuple_output"
@tvm.ir.register_op_attr("split", "target." + target)
def split(expr): # pylint: disable=unused-variable
return True
@tvm.ir.register_op_attr("abs", "target." + target)
def abs(expr): # pylint: disable=unused-variable
return True
def create_graph():
a = relay.var("a", shape=(10, 10), dtype="uint8")
a_split = relay.split(a, 2)
a_split_0 = relay.TupleGetItem(a_split.astuple(), 0)
a_split_0_abs = relay.abs(a_split_0)
a_con = relay.concatenate(a_split, 0)
a_split_0_relu = relay.nn.relu(a_split_0_abs)
out = relay.Tuple((a_con, a_split_0_relu))
f = relay.Function([a], out)
mod = tvm.IRModule.from_expr(f)
mod = transform.InferType()(mod)
return mod
seq = tvm.transform.Sequential(
[
transform.AnnotateTarget(target),
transform.MergeCompilerRegions(),
transform.PartitionGraph(),
]
)
partitioned = seq(create_graph())
partitioned = transform.InferType()(partitioned)
partitioned.show()
def @main(%a: Tensor[(10, 10), uint8] /* ty=Tensor[(10, 10), uint8] */) -> (Tensor[(10, 10), uint8], Tensor[(5, 10), uint8]) {
%0 = @tvmgen_default_test_flatten_tuple_output_main_0(%a) /* ty=(Tensor[(5, 10), uint8], Tensor[(5, 10), uint8], Tensor[(5, 10), uint8]) */;
%1 = %0.0 /* ty=Tensor[(5, 10), uint8] */;
%2 = %0.1 /* ty=Tensor[(5, 10), uint8] */;
%3 = (%1, %2) /* ty=(Tensor[(5, 10), uint8], Tensor[(5, 10), uint8]) */;
%4 = %0.2 /* ty=Tensor[(5, 10), uint8] */;
%5 = concatenate(%3) /* ty=Tensor[(10, 10), uint8] */;
%6 = nn.relu(%4) /* ty=Tensor[(5, 10), uint8] */;
(%5, %6) /* ty=(Tensor[(10, 10), uint8], Tensor[(5, 10), uint8]) */
}
def @tvmgen_default_test_flatten_tuple_output_main_0(%test_flatten_tuple_output_0_i0: Tensor[(10, 10), uint8] /* ty=Tensor[(10, 10), uint8] */, Compiler="test_flatten_tuple_output", Primitive=1, Inline=1, global_symbol="tvmgen_default_test_flatten_tuple_output_main_0") -> (Tensor[(5, 10), uint8], Tensor[(5, 10), uint8], Tensor[(5, 10), uint8]) {
%7 = split(%test_flatten_tuple_output_0_i0, indices_or_sections=meta[runtime.BoxInt][0]) /* ty=(Tensor[(5, 10), uint8], Tensor[(5, 10), uint8]) */;
%8 = %7.0 /* ty=Tensor[(5, 10), uint8] */;
%9 = %7.0 /* ty=Tensor[(5, 10), uint8] */;
%10 = %7.1 /* ty=Tensor[(5, 10), uint8] */;
%11 = abs(%8) /* ty=Tensor[(5, 10), uint8] */;
(%9, %10, %11) /* ty=(Tensor[(5, 10), uint8], Tensor[(5, 10), uint8], Tensor[(5, 10), uint8]) */
}
测试元组输出的执行#
测试具有元组输出的子图的 C 代码生成和运行时。
a = relay.var("a", shape=(10, 10), dtype="float32")
b = relay.var("b", shape=(10, 10), dtype="float32")
ba = relay.annotation.compiler_begin(a, "ccompiler")
bb = relay.annotation.compiler_begin(b, "ccompiler")
add = relay.add(ba, bb)
sub = relay.subtract(ba, bb)
out = relay.Tuple((add, sub))
eout = relay.annotation.compiler_end(out, "ccompiler")
func = relay.Function([a, b], eout)
mod = tvm.IRModule()
mod["main"] = func
mod = transform.InferType()(mod)
mod = transform.PartitionGraph()(mod)
a_data = np.random.rand(10, 10).astype("float32")
b_data = np.random.rand(10, 10).astype("float32")
check_result(
mod,
{"a": a_data, "b": b_data},
[(10, 10), (10, 10)],
[(a_data + b_data), (a_data - b_data)],
)
[17:40:17] /media/pc/data/lxw/ai/tvm/src/relay/backend/vm/compiler.cc:1199: All lowered functions have been build by BYOC -- generating an empty TVM module
/tmp/ipykernel_2738903/3034817802.py:45: DeprecationWarning: legacy graph executor behavior of producing json / lib / params will be removed in the next release. Please see documents of tvm.contrib.graph_executor.GraphModule for the new recommended usage.
json, lib, param = relay.build(mod, target=target, params=params, runtime=runtime)
测试外部优化#
def Optimize(mod):
return relay.transform.FoldConstant()(mod)
tvm.register_func("relay.ext.test_target.optimize", Optimize)
x = relay.var("x", shape=(2, 2))
y0 = relay.var("y0", shape=(2, 2))
y1 = relay.var("y1", shape=(2, 2))
yy0 = relay.annotation.compiler_begin(y0, "test_target")
yy1 = relay.annotation.compiler_begin(y1, "test_target")
z = yy0 + yy1
end = relay.annotation.compiler_end(z, "test_target")
f = relay.Function([x, y0, y1], end * x)
c = np.ones(shape=(2, 2), dtype="float32")
f = bind_params_by_name(f, {"y0": tvm.nd.array(c), "y1": tvm.nd.array(c)})
mod = tvm.IRModule()
mod["main"] = f
mod = transform.InferType()(mod)
mod = transform.PartitionGraph()(mod)
try:
t0 = mod["tvmgen_default_test_target_main_0"]
except:
raise KeyError("test_target_main_0 not found")
assert isinstance(t0.body, relay.Constant)
expected = np.empty([2, 2])
expected.fill(2)
tvm.testing.assert_allclose(t0.body.data.numpy(), expected, rtol=1e-5, atol=1e-5)
测试保留类型导入#
测试确保在 BYOC 管道中保留类型定义和导入。
from tvm.relay.prelude import Prelude, StaticTensorArrayOps
def run(dtype, shape):
mod = tvm.IRModule()
p = Prelude(mod)
static_tensor_array_ops = StaticTensorArrayOps(p, dtype, shape)
static_tensor_array_ops.register()
tensor_array = p.get_global_var_static("tensor_array", dtype, shape)
tensor = p.get_tensor_ctor_static("tensor_constructor", dtype, shape)
write = p.get_global_var_static("tensor_array_write", dtype, shape)
gather = p.get_global_var_static("tensor_array_gather", dtype, shape)
v = relay.var("v")
indice = relay.var("indice")
init_tensor_array = tensor_array(relay.const(3))
tensor_array1 = write(init_tensor_array, relay.const(0), tensor(v))
tensor_array2 = write(tensor_array1, relay.const(1), tensor(v))
tensor_array3 = write(tensor_array2, relay.const(2), tensor(v))
out = gather(tensor_array3, indice)
mod["main"] = relay.Function([v, indice], out)
mod = transform.RemoveUnusedFunctions()(mod)
mod = transform.PartitionGraph()(mod)
run("float32", [2, 3])
测试不绑定常量#
def get_net(prefix, data, out_channel):
weight = relay.var(prefix + "weight")
bn_gamma = relay.var(prefix + "bn_gamma")
bn_beta = relay.var(prefix + "bn_beta")
bn_mmean = relay.var(prefix + "bn_mean")
bn_mvar = relay.var(prefix + "bn_var")
layer = relay.nn.conv2d(
data=data, weight=weight, kernel_size=(3, 3), channels=out_channel, padding=(1, 1)
)
bn_output = relay.nn.batch_norm(layer, bn_gamma, bn_beta, bn_mmean, bn_mvar)
out = relay.nn.relu(bn_output[0])
return relay.Function(relay.analysis.free_vars(out), out)
def get_partitoned_mod(mod, params, pattern_table, bind_constants):
mod["main"] = bind_params_by_name(mod["main"], params)
remove_bn_pass = tvm.transform.Sequential(
[
transform.InferType(),
transform.SimplifyInference(),
transform.FoldConstant(),
transform.FoldScaleAxis(),
]
)
composite_partition = tvm.transform.Sequential(
[
remove_bn_pass,
transform.MergeComposite(pattern_table),
transform.AnnotateTarget("dnnl"),
transform.PartitionGraph(bind_constants=bind_constants),
]
)
with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
return composite_partition(mod)
data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32"))
net = get_net("block_", data, 8)
mod, params = tvm.relay.testing.create_workload(net)
mod = get_partitoned_mod(mod, params, get_pattern_table("dnnl"), bind_constants=True)
len(mod["main"].body.args) == 1
mod = get_partitoned_mod(mod, params, get_pattern_table("dnnl"), bind_constants=False)
len(mod["main"].body.args) == 3
False