def run_opt_pass(expr, passes):
passes = passes if isinstance(passes, list) else [passes]
mod = tvm.IRModule.from_expr(expr)
seq = transform.Sequential(passes)
with transform.PassContext(opt_level=3):
mod = seq(mod)
entry = mod["main"]
return entry if isinstance(expr, relay.Function) else entry.body
def test_pass_registration():
passes = [module_pass, function_pass]
opt_level = 2
pass_name = "sequential"
sequential = _transform.Sequential(passes=passes, opt_level=opt_level)
pass_info = sequential.info
assert pass_info.name == pass_name
assert pass_info.opt_level == opt_level
def get_partitoned_mod(mod, params, pattern_table):
# This is required for constant folding
mod["main"] = bind_params_by_name(mod["main"], params)
remove_bn_pass = transform.Sequential([
transform.InferType(),
transform.SimplifyInference(),
transform.FoldConstant(),
transform.FoldScaleAxis(),
])
composite_partition = transform.Sequential([
remove_bn_pass,
transform.MergeComposite(pattern_table),
transform.AnnotateTarget("dnnl"),
transform.PartitionGraph()
])
with relay.build_config(opt_level=3, disabled_pass=["AlterOpLayout"]):
return composite_partition(mod)
def test_only_function_pass():
# Check the subtract function.
passes = [function_pass]
sequential = _transform.Sequential(opt_level=1, passes=passes)
ret_mod = sequential(mod)
_, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
check_func(new_sub, get_ref_sub())
# Check the log function.
log_var, new_log = extract_var_func(ret_mod, v_log.name_hint)
check_func(new_log, get_ref_log())
def test_print_ir():
shape = (1, 2, 3)
tp = relay.TensorType(shape, "float32")
x = relay.var("x", tp)
y = relay.add(x, x)
y = relay.multiply(y, relay.const(2, "float32"))
func = relay.Function([x], y)
seq = _transform.Sequential([
relay.transform.InferType(),
relay.transform.FoldConstant(),
relay.transform.PrintIR(),
relay.transform.DeadCodeElimination()
])
def redirect_output(call):
"""Redirect the C++ logging info."""
import sys
import os
import threading
stderr_fileno = sys.stderr.fileno()
stderr_save = os.dup(stderr_fileno)
stderr_pipe = os.pipe()
os.dup2(stderr_pipe[1], stderr_fileno)
os.close(stderr_pipe[1])
output = ''
def record():
nonlocal output
while True:
data = os.read(stderr_pipe[0], 1024)
if not data:
break
output += data.decode("utf-8")
t = threading.Thread(target=record)
t.start()
call()
os.close(stderr_fileno)
t.join()
os.close(stderr_pipe[0])
os.dup2(stderr_save, stderr_fileno)
os.close(stderr_save)
return output
def run_pass():
mod = relay.Module({"main": func})
with relay.build_config(opt_level=3):
mod = seq(mod)
out = redirect_output(run_pass)
assert "Dumping the module IR" in out
assert "multiply" in out
def dcpe(expr, mod=None, grad=False):
passes = [transform.PartialEvaluate(),
transform.DeadCodeElimination(inline_once=True)]
if grad:
expr = gradient(run_infer_type(expr))
if mod:
assert isinstance(expr, Function)
mod["main"] = expr
seq = transform.Sequential(passes)
mod = seq(mod)
return mod["main"]
return run_opt_pass(expr, passes)
def test_only_module_pass():
passes = [module_pass]
sequential = _transform.Sequential(opt_level=1, passes=passes)
ret_mod = sequential(mod)
# Check the subtract function.
sub_var, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
check_func(new_sub, sub)
# Check the abs function is added.
abs_var, abs_func = get_var_func()
abs_var, new_abs = extract_var_func(ret_mod, abs_var.name_hint)
check_func(new_abs, abs_func)
def destroy_ref(x):
x = run_infer_type(x)
x = to_cps(x)
x = run_infer_type(x)
y = un_cps(x)
y = run_infer_type(y)
x = run_opt_pass(
x,
transform.Sequential([
transform.PartialEvaluate(),
transform.DeadCodeElimination(inline_once=True)
]))
assert Feature.fRefCreate not in detect_feature(x)
def quantize_model(model, params, input_dtype, input_shape, qeval='power2'):
skip_conv_layers = [0]
with relay.quantize.qconfig(store_lowbit_output=False,
skip_conv_layers=skip_conv_layers):
from tvm.relay.quantize.quantize import _bind_params
graph = _bind_params(model['main'], params)
mod = relay.Module.from_expr(graph)
optimize = _transform.Sequential([
_transform.SimplifyInference(),
_transform.FoldConstant(),
_transform.FoldScaleAxis(),
_transform.CanonicalizeOps(),
_transform.FoldConstant()
])
with relay.build_config(opt_level=4):
mod = optimize(mod)
mod = relay.quantize.annotate()(mod)
# find scale
cache_file = '%s_%s_scales.pkl' % (VIDEO_FILE, MODEL_NAME)
if os.path.exists(cache_file):
print("Using cached layer statistics...")
with open(cache_file, 'rb') as f:
scales = pickle.load(f)
else:
print("Compute layer statistics...")
scales = calibrate_on_dataset(mod['main'], params, input_dtype,
input_shape)
with open(cache_file, 'wb') as f:
pickle.dump(scales, f)
if qeval == 'power2':
scales = list(
map(
lambda scale: 2**np.math.ceil(np.math.log(scale, 2))
if scale > 0 else 1.0, scales))
weight_scales = 'power2'
elif qeval == 'max':
weight_scales = 'max'
else:
raise ValueError("Invalid quantiziation eval: " + qeval)
mod['main'] = relay.quantize.calibrate(mod['main'],
weight_scales=weight_scales,
scales=scales)
mod = relay.quantize.realize()(mod)
mod = relay.transform.FoldConstant()(mod)
return mod
def test_checkpoint_alpha_equal_tuple():
xs = [
relay.var("x{}".format(i), relay.TensorType((1, ), "float32"))
for i in range(4)
]
f = relay.Function(
xs,
relay.annotation.checkpoint(
relay.Tuple([relay.add(xs[0], xs[1]),
relay.add(xs[2], xs[3])])))
df = transform.gradient(run_infer_type(f))
# run PE and DCE
with transform.PassContext(opt_level=3):
passes = [
transform.PartialEvaluate(),
transform.DeadCodeElimination(inline_once=True)
]
mod = transform.Sequential(passes)(tvm.IRModule.from_expr(df))
df = mod["main"]
df_parsed = relay.parser.fromtext("""
v0.0.4
fn (%x: Tensor[(1), float32], %y: Tensor[(1), float32],
%z: Tensor[(1), float32], %w: Tensor[(1), float32])
-> ((Tensor[(1), float32], Tensor[(1), float32]),
(Tensor[(1), float32], Tensor[(1), float32],
Tensor[(1), float32], Tensor[(1), float32])) {
let %x1: Tensor[(1), float32] = add(%x, %y) /* ty=Tensor[(1), float32] */;
let %x2: Tensor[(1), float32] = add(%z, %w) /* ty=Tensor[(1), float32] */;
let %x3: Tensor[(1), float32] = zeros_like(%x2) /* ty=Tensor[(1), float32] */;
let %x4: Tensor[(1), float32] = ones_like(%x1) /* ty=Tensor[(1), float32] */;
%0 = (%x1, %x2);
%1 = zeros_like(%x) /* ty=Tensor[(1), float32] */;
%2 = collapse_sum_like(%x4, %x) /* ty=Tensor[(1), float32] */;
%3 = add(%1, %2) /* ty=Tensor[(1), float32] */;
%4 = zeros_like(%y) /* ty=Tensor[(1), float32] */;
%5 = collapse_sum_like(%x4, %y) /* ty=Tensor[(1), float32] */;
%6 = add(%4, %5) /* ty=Tensor[(1), float32] */;
%7 = zeros_like(%z) /* ty=Tensor[(1), float32] */;
%8 = collapse_sum_like(%x3, %z) /* ty=Tensor[(1), float32] */;
%9 = add(%7, %8) /* ty=Tensor[(1), float32] */;
%10 = zeros_like(%w) /* ty=Tensor[(1), float32] */;
%11 = collapse_sum_like(%x3, %w) /* ty=Tensor[(1), float32] */;
%12 = add(%10, %11) /* ty=Tensor[(1), float32] */;
%13 = (%3, %6, %9, %12);
(%0, %13)
}
""")
relay.analysis.assert_alpha_equal(df, df_parsed)
def dcpe(expr, mod=None, grad=False):
passes = [
transform.PartialEvaluate(),
transform.DeadCodeElimination(inline_once=True)
]
if grad:
expr = gradient(expr)
if mod:
assert isinstance(expr, Function)
mod[mod.entry_func] = expr
seq = transform.Sequential(passes)
mod = seq(mod)
return mod[mod.entry_func]
return transform.OptimizeOnExpr(expr, passes)
def test_eta_expand_basic():
x = relay.var('x', 'int32')
orig = relay.Function([x], x)
mod = _module.Module.from_expr(orig)
seq = _transform.Sequential([_transform.EtaExpand()])
with _transform.PassContext(opt_level=3):
mod = seq(mod)
got = mod[mod.entry_func.name_hint]
y = relay.var('y', 'int32')
expected = relay.Function([y], orig(y))
got = relay.ir_pass.infer_type(got, mod)
expected = relay.ir_pass.infer_type(expected, mod)
assert (relay.ir_pass.alpha_equal(got, expected))
def test_eta_expand_basic():
x = relay.var('x', 'int32')
orig = relay.Function([x], x)
mod = _module.Module.from_expr(orig)
seq = _transform.Sequential([_transform.EtaExpand()])
with _transform.PassContext(opt_level=3):
mod = seq(mod)
got = mod["main"]
y = relay.var('y', 'int32')
expected = relay.Function([y], orig(y))
gv = relay.GlobalVar("gv")
mod[gv] = expected
mod = _transform.InferType()(mod)
expected = mod["gv"]
assert (relay.analysis.alpha_equal(got, expected))
def test_fold_batch_norm():
def expected():
data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32"))
weight = relay.const(np.zeros((16, 3, 3, 3)))
bias = relay.const(np.zeros((16, 1, 1)))
conv = relay.nn.conv2d(data=data,
weight=weight,
kernel_size=(3, 3),
channels=16,
padding=(1, 1))
add = relay.add(conv, bias)
return relay.Function(relay.analysis.free_vars(add), add)
remove_bn_pass = transform.Sequential([
relay.transform.InferType(),
relay.transform.SimplifyInference(),
relay.transform.FoldConstant(),
relay.transform.FoldScaleAxis(),
])
data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32"))
weight = relay.var("weight")
bn_gamma = relay.var("bn_gamma")
bn_beta = relay.var("bn_beta")
bn_mmean = relay.var("bn_mean")
bn_mvar = relay.var("bn_var")
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)
def initializer(_, param):
param = np.zeros(param.shape)
mod, params = create_workload(bn_output[0], initializer)
mod["main"] = bind_params_by_name(mod["main"], params)
with relay.build_config(opt_level=3):
mod = remove_bn_pass(mod)
expect = run_infer_type(expected())
assert relay.analysis.graph_equal(mod["main"], expect)
def test_multiple_passes():
# Reset the current module since mod has been polluted by the previous
# function pass.
mod = tvm.IRModule({v_sub: sub, v_log: log})
passes = [module_pass, function_pass]
sequential = _transform.Sequential(opt_level=1, passes=passes)
required = ["mod_transform", "func_transform"]
with relay.build_config(required_pass=required):
ret_mod = sequential(mod)
# Check the abs function is added.
abs_var, abs_func = get_var_func()
abs_var, new_abs = extract_var_func(ret_mod, abs_var.name_hint)
check_func(new_abs, get_ref_abs())
# Check the subtract function is modified correctly.
_, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
check_func(new_sub, get_ref_sub())
# Check the log function is modified correctly.
_, new_log = extract_var_func(ret_mod, v_log.name_hint)
check_func(new_log, get_ref_log())
# Execute the updated subtract function.
x_nd = get_rand(shape, dtype)
y_nd = get_rand(shape, dtype)
ref_res = np.subtract(x_nd.asnumpy() * 2, y_nd.asnumpy() * 2)
for target, ctx in ctx_list():
exe1 = relay.create_executor("graph", ctx=ctx, target=target)
exe2 = relay.create_executor("debug", ctx=ctx, target=target)
res1 = exe1.evaluate(new_sub)(x_nd, y_nd)
tvm.testing.assert_allclose(res1.asnumpy(), ref_res, rtol=1e-5)
res2 = exe2.evaluate(new_sub)(x_nd, y_nd)
tvm.testing.assert_allclose(res2.asnumpy(), ref_res, rtol=1e-5)
# Execute the updated abs function.
x_nd = get_rand((5, 10), dtype)
ref_res = np.abs(x_nd.asnumpy() * 2)
for target, ctx in ctx_list():
exe1 = relay.create_executor("graph", ctx=ctx, target=target)
exe2 = relay.create_executor("debug", ctx=ctx, target=target)
res1 = exe1.evaluate(new_abs)(x_nd)
tvm.testing.assert_allclose(res1.asnumpy(), ref_res, rtol=1e-5)
res2 = exe2.evaluate(new_abs)(x_nd)
tvm.testing.assert_allclose(res2.asnumpy(), ref_res, rtol=1e-5)
def check(shape):
data = relay.var("data", shape=shape, dtype="int8")
conv_weight = relay.var("weight")
bias1 = relay.var("bias1", shape=(16, 1, 1), dtype="int32")
bias2 = relay.var("bias2", shape=(16, 1, 1), dtype="int32")
y = before(data, conv_weight, bias1, bias2)
mod = _module.Module.from_expr(y)
seq = _transform.Sequential([
_transform.InferType(),
_transform.CanonicalizeCast(),
_transform.InferType()
])
with _transform.PassContext(opt_level=3):
mod = seq(mod)
y = mod[mod.entry_func.name_hint]
y_expected = expected(data, conv_weight, bias1, bias2)
y_expected = relay.ir_pass.infer_type(y_expected)
assert relay.ir_pass.alpha_equal(y, y_expected)
def check(shape):
data = relay.var("data", shape=shape, dtype="int8")
conv_weight = relay.var("weight")
bias1 = relay.var("bias1", shape=(16, 1, 1), dtype="int32")
bias2 = relay.var("bias2", shape=(16, 1, 1), dtype="int32")
y = before(data, conv_weight, bias1, bias2)
mod = tvm.IRModule.from_expr(y)
seq = _transform.Sequential([_transform.InferType(), _transform.CanonicalizeCast(),
_transform.InferType()])
with _transform.PassContext(opt_level=3):
mod = seq(mod)
y = mod["main"]
y_expected = expected(data, conv_weight, bias1, bias2)
gv = relay.GlobalVar("expected")
mod[gv] = y_expected
mod = _transform.InferType()(mod)
y_expected = mod["expected"]
assert relay.analysis.alpha_equal(y, y_expected)
def test_after_partial_eval():
"""Test transformation following reverse mode ad and PartialEval"""
mod = tvm.IRModule()
shape = (10, 10)
dtype = 'float32'
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
y = relay.var("y", t)
func = relay.Function([x, y], (x * y) * relay.const(np.ones(shape, dtype)))
func = run_infer_type(func)
back_func = transform.gradient(func)
back_func = run_infer_type(back_func)
mod["main"] = back_func
back_func = mod["main"]
seq = transform.Sequential([
transform.PartialEvaluate(),
transform.LazyGradientInit(),
transform.DeadCodeElimination()
])
mod = seq(mod)
assert mod["main"].checked_type == relay.FuncType(
[t, t], relay.TupleType([t, relay.TupleType([t, t])]))
ex = create_executor(mod=mod)
x = rand(dtype, *shape)
y = rand(dtype, *shape)
(forward), (
grad_x,
grad_y,
) = ex.evaluate(back_func)(x, y)
assert_allclose(forward.asnumpy(), x.asnumpy() * y.asnumpy())
assert_allclose(grad_x.asnumpy(), y.asnumpy())
assert_allclose(grad_y.asnumpy(), x.asnumpy())
def test_sequential_with_scoping():
shape = (1, 2, 3)
c_data = np.array(shape).astype("float32")
tp = relay.TensorType(shape, "float32")
def before():
c = relay.const(c_data)
x = relay.var("x", tp)
y = relay.add(c, c)
y = relay.multiply(y, relay.const(2, "float32"))
y = relay.add(x, y)
z = relay.add(y, c)
z1 = relay.add(y, c)
z2 = relay.add(z, z1)
return relay.Function([x], z2)
def expected():
x = relay.var("x", tp)
c_folded = (c_data + c_data) * 2
y = relay.add(x, relay.const(c_folded))
z = relay.add(y, relay.const(c_data))
z1 = relay.add(z, z)
return relay.Function([x], z1)
seq = _transform.Sequential([
relay.transform.InferType(),
relay.transform.FoldConstant(),
relay.transform.EliminateCommonSubexpr(),
relay.transform.AlterOpLayout()
])
mod = tvm.IRModule({"main": before()})
with relay.build_config(opt_level=3):
with tvm.target.create("llvm"):
mod = seq(mod)
zz = mod["main"]
zexpected = run_infer_type(expected())
assert analysis.alpha_equal(zz, zexpected)
def test_print_debug_callback():
global __TRACE_COUNTER__
shape = (1, 2, 3)
tp = relay.TensorType(shape, "float32")
x = relay.var("x", tp)
y = relay.add(x, x)
y = relay.multiply(y, relay.const(2, "float32"))
func = relay.Function([x], y)
seq = _transform.Sequential([
relay.transform.InferType(),
relay.transform.FoldConstant(),
relay.transform.DeadCodeElimination()
])
assert __TRACE_COUNTER__ == 0
mod = tvm.IRModule({"main": func})
with relay.build_config(opt_level=3, trace=_tracer):
mod = seq(mod)
assert __TRACE_COUNTER__ == 4
def test_print_ir(capfd):
shape = (1, 2, 3)
tp = relay.TensorType(shape, "float32")
x = relay.var("x", tp)
y = relay.add(x, x)
y = relay.multiply(y, relay.const(2, "float32"))
func = relay.Function([x], y)
seq = _transform.Sequential([
relay.transform.InferType(),
relay.transform.FoldConstant(),
relay.transform.PrintIR(),
relay.transform.DeadCodeElimination()
])
mod = tvm.IRModule({"main": func})
with relay.build_config(opt_level=3):
mod = seq(mod)
out = capfd.readouterr().err
assert "Dumping the module IR" in out
assert "multiply" in out
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
op_list = ["nn.batch_norm", "nn.conv2d"]
mod = WhiteListAnnotator(op_list, "test_compiler")(mod)
opt_pass = transform.Sequential([
transform.InferType(),
transform.PartitionGraph(),
transform.SimplifyInference(),
transform.FoldConstant(),
transform.AlterOpLayout(),
])
with relay.build_config(opt_level=3):
mod = opt_pass(mod)
return mod
def test_no_pass():
passes = []
sequential = _transform.Sequential(opt_level=1, passes=passes)
ret_mod = sequential(mod)
mod_func = ret_mod[v_sub]
check_func(sub, mod_func)
"""Workaround for type checker and mutually recursive functions."""
for gv in funcs:
func = funcs[gv]
body = _placeholder_body(func.ret_type)
mod[gv] = relay.Function(func.params, body, func.ret_type)
for gv in funcs:
mod[gv] = funcs[gv]
pass_set = transform.Sequential(
passes=[
transform.SimplifyInference(),
transform.CanonicalizeOps(),
transform.CanonicalizeCast(),
transform.FuseOps(3),
# transform.CombineParallelConv2d(),
transform.AlterOpLayout(),
# transform.RewriteAnnotatedOps(???),
],
opt_level=0)
def optimize(mod):
"""Optimize all the functions in a module.
Modules are the only mutable piece of Relay. We write an
optimization pass over the module which destructively updates each
function while optimizing.
"""
return pass_set(mod)