def test_constructor_type():
mod = tvm.IRModule()
box, constructor = initialize_box_adt(mod)
a = relay.TypeVar("a")
x = relay.Var("x", a)
func = relay.Function([x], constructor(x), box(a), [a])
mod["main"] = func
mod = infer_mod(mod)
func_ty = mod["main"].checked_type
box = mod.get_global_type_var("box")
expected = relay.FuncType([a], box(a), [a])
assert func_ty == expected
def test_adt_match_type_annotations():
mod = tvm.IRModule()
box, constructor = initialize_box_adt(mod)
# the only type annotation is inside the match pattern var
# but that should be enough info
tt = relay.TensorType((2, 2), "float32")
x = relay.Var("x")
mv = relay.Var("mv", tt)
match = relay.Match(
constructor(x),
[
relay.Clause(
relay.PatternConstructor(constructor, [relay.PatternVar(mv)]),
relay.Tuple([]))
],
)
mod["main"] = relay.Function([x], match)
mod = infer_mod(mod)
ft = mod["main"].checked_type
assert ft == relay.FuncType([tt], relay.TupleType([]))
def test_function():
param_names = ['a', 'b', 'c', 'd']
params = tvm.convert(
[relay.Param(relay.Var(n), None) for n in param_names])
ret_type = None
body = None
type_params = tvm.convert([])
fn = relay.Function(params, ret_type, body, type_params)
assert fn.params == params
assert fn.body == body
assert fn.type_params == type_params
assert fn.span == None
str(fn)
def test_function():
param_names = ['a', 'b', 'c', 'd']
params = tvm.convert([relay.Var(n) for n in param_names])
ret_type = relay.TupleType(tvm.convert([]))
body = relay.Tuple(tvm.convert([]))
type_params = tvm.convert([])
fn = relay.Function(params, body, ret_type, type_params)
assert fn.params == params
assert fn.body == body
assert fn.type_params == type_params
assert fn.span == None
str(fn)
check_json_roundtrip(fn)
def test_map():
a = relay.TypeVar("a")
b = relay.TypeVar("b")
lhs = prelude.mod[map].checked_type
rhs = relay.FuncType([relay.FuncType([a], b), rlist(a)], rlist(b), [a, b])
assert lhs == rhs
x = relay.Var("x")
add_one = relay.Function([x], s(x))
res = intrp.evaluate(map(add_one, cons(z(), cons(z(), nil()))))
ones = to_list(res)
assert len(ones) == 2
assert count(ones[0]) == 1 and count(ones[1]) == 1
def test_optional_matching():
x = relay.Var('x')
y = relay.Var('y')
v = relay.Var('v')
condense = relay.Function(
[x, y],
relay.Match(x, [
relay.Clause(relay.PatternConstructor(some, [relay.PatternVar(v)]),
cons(v, y)),
relay.Clause(relay.PatternConstructor(none), y)
]))
res = intrp.evaluate(
foldr(
condense, nil(),
cons(some(build_nat(3)),
cons(none(), cons(some(build_nat(1)), nil())))))
reduced = to_list(res)
assert len(reduced) == 2
assert count(reduced[0]) == 3
assert count(reduced[1]) == 1
def fuse_partitions(pre_mod, mid_mod, post_mod):
"""Combine prefix, middle, and suffix modules into a single module.
The combined module includes an additional `main` that fuses all three
partitions together.
Parameters
----------
pre_mod : tvm.IRModule
Module containing an input quantization function
mid_mod : tvm.IRModule
Module containing core of a quantized inference function
post_mod : tvm.IRModule
Module containing an output dequantization function
Returns
-------
fused_mod : tvm.IRModule
Module containing the input quantization, core quantized inference,
output dequantization, and full quantized inference functions
"""
pre_func = pre_mod['main']
mid_func = mid_mod['main']
post_func = post_mod['main']
# create a module containing the prefix, middle, and suffix partitions
fused_mod = tvm.IRModule(functions={
relay.GlobalVar('quantize_inputs'): pre_func,
relay.GlobalVar('quantized_main'): mid_func,
relay.GlobalVar('dequantize_outputs'): post_func,
})
# construct a `main` that strings together the partitions, such that its
# behaviour is equivalent to `main` in an *unpartitioned* module
scope_builder = relay.ScopeBuilder()
fused_mod_main_params = [relay.Var(param.name_hint) for param in pre_func.params]
quantized_inputs = scope_builder.let('quantized_inputs', relay.Call(
fused_mod.get_global_var('quantize_inputs'),
fused_mod_main_params
))
quantized_outputs = scope_builder.let('quantized_outputs', relay.Call(
fused_mod.get_global_var('quantized_main'),
[relay.TupleGetItem(quantized_inputs, i) for i in range(len(pre_func.ret_type.fields))]
))
dequantized_outputs = scope_builder.let('dequantized_outputs', relay.Call(
fused_mod.get_global_var('dequantize_outputs'),
[quantized_outputs]
))
scope_builder.ret(dequantized_outputs)
fused_mod['main'] = relay.Function(fused_mod_main_params, scope_builder.get())
return fused_mod
def partition_prefix(mod, quantized_dtypes):
"""Extract input quantization expressions from `mod['main']`.
Parameters
----------
mod : tvm.IRModule
Module containing a quantized inference function
quantized_dtypes : Set[str]
Set of data types allowed in quantized operators
Returns
-------
pre_mod : tvm.IRModule
Module containing the input quantization function
mid_mod : tvm.IRModule
Module containing a function with everything except for input quantization
"""
assert len(mod.functions) == 1
func = mod['main']
prefix_cutter = PrefixCutter(func.params, quantized_dtypes)
mid_body = prefix_cutter.visit(func.body)
assert not func.type_params, 'unimplemented'
assert func.attrs is None, 'unimplemented'
mid_func = relay.Function(
relay.analysis.free_vars(mid_body),
mid_body)
mid_mod = tvm.IRModule.from_expr(mid_func)
scope_builder = prefix_cutter.prefix_sb
# make sure we pass through all inputs in the prefix function's return expr
# (even those that don't require quantization)
ret_expr = []
for param in mid_func.params:
if param in prefix_cutter.prefix_binding_map:
# this param required a conversion, so we collected it in the
# prefix cutter pass, and we can use the pass's mapping from mid
# func params to pre func params
ret_expr.append(prefix_cutter.prefix_binding_map[param])
else:
# there was no detected conversion for this argument, so we thread
# it through the prefix function untouched
ret_expr.append(relay.Var(param.name_hint, param.checked_type))
ret_expr = relay.Tuple(ret_expr)
scope_builder.ret(ret_expr)
pre_func_body = scope_builder.get()
pre_func = relay.Function(relay.analysis.free_vars(pre_func_body), pre_func_body)
pre_mod = tvm.IRModule.from_expr(pre_func)
return pre_mod, mid_mod
def test_ref_execution_order():
# we want to have effects execute from left to right
x = relay.Var('x')
y = relay.Var('y')
f = relay.Var('f')
r = relay.Var('r')
expr = relay.Let(f, relay.Function([x, y], x),
# r = 1
relay.Let(r, relay.RefCreate(relay.const(1)),
relay.Tuple([
# should be 1
relay.RefRead(r),
# set r to 2 and read back
seq(relay.RefWrite(r, relay.const(2)),
relay.RefRead(r)),
# set r to 3 and read back
seq(relay.RefWrite(r, relay.const(3)),
relay.RefRead(r)),
# set r to 4 and read as first arg to f
# set r to 5 and read as second arg to f
# f should evaluate to 4
f(
seq(relay.RefWrite(r, relay.const(4)),
relay.RefRead(r)),
seq(relay.RefWrite(r, relay.const(5)),
relay.RefRead(r))),
# read back 5
relay.RefRead(r)
])))
tup_val = run_as_python(expr)
assert_tuple_value(tup_val, 5)
assert_tensor_value(tup_val.fields[0], 1)
assert_tensor_value(tup_val.fields[1], 2)
assert_tensor_value(tup_val.fields[2], 3)
assert_tensor_value(tup_val.fields[3], 4)
assert_tensor_value(tup_val.fields[4], 5)
def test_order():
z = relay.const(3)
y = relay.const(2)
x = relay.const(1)
val = x + y * z
check_eval(val, 7.0)
anf = transform.OptimizeOnExpr(
val, [transform.ToANormalForm(),
transform.InferType()])
a = relay.Var('a', relay.IncompleteType())
b = relay.Var('b', relay.IncompleteType())
c = relay.Var('c', relay.IncompleteType())
d = relay.Var('d', relay.IncompleteType())
e = relay.Var('e', relay.IncompleteType())
expected_output = e
expected_output = relay.Let(e, a + d, expected_output)
expected_output = relay.Let(d, b * c, expected_output)
expected_output = relay.Let(c, z, expected_output)
expected_output = relay.Let(b, y, expected_output)
expected_output = relay.Let(a, x, expected_output)
expected_output = transform.OptimizeOnExpr(expected_output,
transform.InferType())
assert alpha_equal(anf, expected_output)
def test_foldr():
a = relay.TypeVar("a")
b = relay.TypeVar("b")
lhs = prelude.mod[foldr].checked_type
rhs = relay.FuncType([relay.FuncType([a, b], b), b, rlist(a)], b, [a, b])
assert lhs == rhs
x = relay.Var("x")
y = relay.Var("y")
identity = relay.Function([x, y], cons(x, y))
res = intrp.evaluate(
foldr(
identity,
nil(),
cons(
make_nat_expr(prelude, 1),
cons(make_nat_expr(prelude, 2), cons(make_nat_expr(prelude, 3), nil())),
),
)
)
same = to_list(res)
assert len(same) == 3
assert count(same[0]) == 1 and count(same[1]) == 2 and count(same[2]) == 3
def test_function_invalidate():
shape = ()
dtype = "bool"
t = relay.TensorType(shape, dtype)
d = relay.Var("d", t)
r = relay.Var("r")
fetch = relay.Function([], relay.RefRead(r))
fet = relay.Var("fetch")
fet_obscured = relay.Var("fetch_obscured")
u = relay.Var("u")
body = relay.If(d, fet_obscured(), fet_obscured())
body = relay.Let(u, relay.RefWrite(r, relay.const(1)), body)
body = relay.Let(fet_obscured, relay.If(d, fet, fet), body)
body = relay.Let(fet, fetch, body)
body = relay.Let(r, relay.RefCreate(relay.const(0)), body)
f = relay.Function([d], body)
f = infer_type(f)
pe_f = infer_type(partial_evaluate(f))
ex = create_executor()
f_res = ex.evaluate(f)(relay.const(True))
pe_f_res = ex.evaluate(pe_f)(relay.const(True))
np.testing.assert_allclose(f_res.asnumpy(), np.ones_like(f_res.asnumpy()))
np.testing.assert_allclose(pe_f_res.asnumpy(), np.ones_like(pe_f_res.asnumpy()))
def test_wildcard_match_order():
x = relay.Var("x", rlist(nat()))
y = relay.Var("y")
a = relay.Var("a")
return_zero = relay.Function(
[x],
relay.Match(
x,
[
relay.Clause(relay.PatternWildcard(), z()),
relay.Clause(
relay.PatternConstructor(
cons, [relay.PatternVar(y),
relay.PatternVar(a)]), y),
relay.Clause(relay.PatternConstructor(nil), s(z())),
],
),
nat(),
)
res = intrp.evaluate(return_zero(cons(s(z()), nil())))
# wildcard pattern is evaluated first
assert count(res) == 0
def verify_upsampling(dshape,
scale_h,
scale_w,
layout,
method,
align_corners=False):
if layout == "NCHW":
(n, c, h, w) = dshape
x_data = np.random.uniform(size=(n, c, h, w)).astype("float32")
elif layout == "NHWC":
(n, h, w, c) = dshape
x_data = np.random.uniform(size=(n, h, w, c)).astype("float32")
if method == "nearest_neighbor":
ref_res = tvm.topi.testing.upsampling_python(
x_data, (scale_h, scale_w), layout)
else:
ref_res = tvm.topi.testing.bilinear_resize_python(
x_data, (int(round(h * scale_h)), int(round(w * scale_w))),
layout)
x = relay.Var("x", relay.TensorType(dshape, "float32"))
scale_h_var = relay.var("scale_h", relay.TensorType((), "float32"))
scale_w_var = relay.var("scale_h", relay.TensorType((), "float32"))
z = relay.nn.upsampling(x,
scale_h_var,
scale_w_var,
method=method,
layout=layout,
align_corners=align_corners)
zz = run_infer_type(z)
func = relay.Function([x, scale_h_var, scale_w_var], z)
for target, ctx in ctx_list():
if "llvm" not in target: continue
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind,
mod=mod,
ctx=ctx,
target=target)
op_res = intrp.evaluate()(x_data,
np.array(scale_h).astype("float32"),
np.array(scale_w).astype("float32"))
tvm.testing.assert_allclose(op_res.asnumpy(),
ref_res,
rtol=1e-4,
atol=1e-6)
def test_head_cons():
mod = relay.Module()
p = Prelude(mod)
def hd_impl():
a = relay.TypeVar("a")
x = relay.Var("x", p.l(a))
y = relay.Var("y")
z = relay.Var("z")
cons_case = relay.Clause(
relay.PatternConstructor(
p.cons,
[relay.PatternVar(y), relay.PatternVar(z)]), y)
return relay.Function([x], relay.Match(x, [cons_case]), a, [a])
t = relay.TypeVar("t")
x = relay.Var("x", t)
hd = relay.Var("hd")
body = relay.Let(hd, hd_impl(), hd(p.cons(x, p.nil())))
f = relay.Function([x], body, None, [t])
f = infer_type(f, mod=mod)
res = dcpe(f)
assert alpha_equal(res, relay.Function([x], x, t, [t]))
def verify_expand_dims(dshape, dtype, oshape, axis, num_newaxis):
x = relay.Var("x", relay.TensorType(dshape, dtype))
func = relay.Function([x], relay.expand_dims(x, axis, num_newaxis))
for target, dev in tvm.testing.enabled_targets():
if (
dtype == "float16"
and target == "cuda"
and not have_fp16(tvm.cuda(0).compute_version)
):
continue
data = np.random.uniform(size=dshape).astype(dtype)
ref_res = data.reshape(oshape)
op_res = relay.create_executor("graph", device=dev, target=target).evaluate(func)(data)
np.testing.assert_allclose(op_res.numpy(), ref_res, rtol=0.01)
def test_dyn_broadcast_to():
dtype = 'uint8'
rank = 3
shape_type = 'int64'
dyn_shape = relay.Var("shape", relay.ty.TensorType((rank, ), shape_type))
x_shape = (1, )
x = relay.Var("x", relay.ty.TensorType(x_shape, dtype))
z = relay.broadcast_to(x, dyn_shape)
zz = run_infer_type(z)
assert zz.checked_type == relay.ty.TensorType((relay.Any(), ) * rank, dtype)
func = relay.Function([x, dyn_shape], z)
x = np.random.uniform(size=x_shape).astype(dtype)
dyn_shape = (1, ) * rank
ref_res = np.broadcast_to(x, dyn_shape)
for target, ctx in tvm.testing.enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target)
op_res = intrp.evaluate(func)(x, np.array(dyn_shape).astype(shape_type))
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
def verify_zeros_ones(shape, dtype):
for op, ref in [(relay.zeros, np.zeros), (relay.ones, np.ones)]:
rank = len(shape)
dyn_shape = relay.Var("shape",
relay.ty.TensorType((rank, ), "int64"))
y = op(dyn_shape, dtype)
yy = run_infer_type(y)
assert yy.checked_type == relay.ty.TensorType(
(relay.Any(), ) * rank, dtype)
func = relay.Function([dyn_shape], y)
ref_res = ref(shape, dtype)
verify_func(executor_kind, func, [np.array(shape).astype("int64")],
ref_res.astype("int64"))
def __init__(self):
self.shape = tvm.runtime.convert([1, 2, 3])
self.tt = relay.TensorType(self.shape, "float32")
self.int32 = relay.TensorType([], "int32")
self.float32 = relay.TensorType([], "float32")
self.one = relay.const(1.0)
self.two = relay.const(2.0)
self.three = relay.const(3.0)
self.a = relay.Var("a", self.float32)
self.b = relay.Var("b", self.float32)
self.c = relay.Var("c", self.float32)
self.d = relay.Var("d", self.float32)
self.e = relay.Var("e", self.float32)
self.x = relay.Var("x", self.int32)
self.y = relay.Var("y", self.int32)
self.z = relay.Var("z", self.int32)
def test_recursion():
"""
Program:
let f(n: i32, data: f32) -> f32 = {
if (n == 0) {
return data;
} else {
return f(n - 1, log(data));
}
}
f(2, 10000);
"""
f = relay.Var("f")
n = relay.Var("n", e.int32)
data = relay.Var("data", e.float32)
funcbody = relay.If(equal(n, relay.const(0)),
data,
relay.Call(f, [subtract(n, relay.const(1.0)),
log(data)]))
value = relay.Function([n, data], funcbody, e.float32, [])
orig = relay.Let(f, value, relay.Call(f, [relay.const(2.0), relay.const(10000.0)]))
assert alpha_equal(dead_code_elimination(orig), orig)
assert alpha_equal(dead_code_elimination(relay.Let(f, value, e.three)), e.three)
def test_let():
lv = relay.Var('x')
ty = None
arr = tvm.nd.array(10)
value = relay.Constant(arr)
# I would prefer that the order of arguments
# matches syntax let x: t = v in b
let = relay.Let(lv, value, lv, ty)
assert let.var == lv
assert let.value == value
assert let.value_type == ty
assert let.body == lv
assert let.span == None
str(let)
def test_ref():
shape = (10, 10)
dtype = 'float32'
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
r = relay.Var("r")
u = relay.Var("u")
body = relay.RefRead(r)
body = relay.Let(u, relay.RefWrite(r,
relay.RefRead(r) + relay.RefRead(r)),
body)
body = relay.Let(r, relay.RefCreate(x), body)
func = relay.Function([x], body)
func = run_infer_type(func)
back_func = run_infer_type(gradient(func))
assert back_func.checked_type == relay.FuncType(
[t], relay.TupleType([t, relay.TupleType([t])]))
x_nd = rand(dtype, *shape)
ex = create_executor()
forward, (grad_x, ) = ex.evaluate(back_func)(x_nd)
tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy())
tvm.testing.assert_allclose(grad_x.asnumpy(),
2 * np.ones_like(grad_x.asnumpy()))
def test_nested_match_pattern():
mod = tvm.IRModule()
box, box_ctor = init_box_adt(mod)
v = relay.Var("v")
w = relay.Var("w")
match = relay.Let(
v,
box_ctor(box_ctor(relay.const(2))),
relay.Match(
v,
[
relay.Clause(
relay.PatternConstructor(box_ctor, [
relay.PatternConstructor(box_ctor,
[relay.PatternVar(w)])
]),
w,
)
],
),
)
match_val = run_as_python(match, mod)
assert_tensor_value(match_val, 2)
def test_global_function():
m = tvm.IRModule()
shape = (10, 10)
dtype = "float32"
t = relay.TensorType(shape, dtype)
x = relay.Var("x", t)
d = GlobalVar("double")
m[d] = relay.Function([x], x + x)
y = relay.Var("y", t)
q = GlobalVar("q")
m[q] = relay.Function([y], d(d(y)))
g = GlobalVar("grad")
m = tvm.relay.transform.InferType()(m)
m[g] = tvm.relay.transform.gradient(q, m)
m = tvm.relay.transform.InferType()(m)
back_func = m[g]
assert back_func.checked_type == relay.FuncType(
[t], relay.TupleType([t, relay.TupleType([t])]))
ex = create_executor(mod=m)
x = rand(dtype, *shape)
forward, (grad, ) = ex.evaluate(back_func)(x)
tvm.testing.assert_allclose(forward.numpy(), 4 * x.numpy())
tvm.testing.assert_allclose(grad.numpy(), 4 * np.ones_like(x.numpy()))
def test_multiple_constructor_clauses():
mod = tvm.IRModule()
p = Prelude(mod)
_, cons, nil = mod.get_type("List")
v = relay.Var("v")
match = relay.Match(
v,
[
# list of length exactly 1
relay.Clause(
relay.PatternConstructor(cons, [
relay.PatternWildcard(),
relay.PatternConstructor(nil, [])
]),
v,
),
# list of length exactly 2
relay.Clause(
relay.PatternConstructor(
cons,
[
relay.PatternWildcard(),
relay.PatternConstructor(cons, [
relay.PatternWildcard(),
relay.PatternConstructor(nil, [])
]),
],
),
v,
),
# empty list
relay.Clause(relay.PatternConstructor(nil, []), v),
# list of length 2 or more
relay.Clause(
relay.PatternConstructor(
cons,
[
relay.PatternWildcard(),
relay.PatternConstructor(
cons,
[relay.PatternWildcard(),
relay.PatternWildcard()]),
],
),
v,
),
],
)
assert len(unmatched_cases(match, mod)) == 0
def test_ref_write():
# check that the result of a ref write is an empty tuple
v = relay.Var('v')
initial_write = relay.Let(v, relay.RefCreate(relay.Tuple([relay.const(1)])),
relay.RefWrite(v, relay.Tuple([relay.const(2)])))
write_val = run_as_python(initial_write)
assert_tuple_value(write_val, 0)
# now ensure that the value, once written, can be read back
# (we read the value before and after mutation)
w = relay.Var('w')
read_after_write = relay.Let(
v, relay.RefCreate(relay.Tuple([relay.const(1)])),
relay.Let(
w, relay.RefCreate(relay.RefRead(v)),
seq(relay.RefWrite(v, relay.Tuple([relay.const(2)])),
relay.Tuple([relay.RefRead(w), relay.RefRead(v)]))))
read_val = run_as_python(read_after_write)
assert_tuple_value(read_val, 2)
assert_tuple_value(read_val.fields[0], 1)
assert_tuple_value(read_val.fields[1], 1)
assert_tensor_value(read_val.fields[0].fields[0], 1)
assert_tensor_value(read_val.fields[1].fields[0], 2)
def test_let():
lv = relay.Var("x")
ty = None
arr = tvm.nd.array(10)
value = relay.Constant(arr)
# I would prefer that the order of arguments
# matches syntax let x: t = v in b
let = relay.Let(lv, value, lv)
assert let.var == lv
assert let.value == value
assert let.body == lv
assert let.span == None
str(let)
check_json_roundtrip(let)
def verify_upsampling3d(
dshape, scale_d, scale_h, scale_w, layout, method, coord_trans="half_pixel"
):
if layout == "NCDHW":
(n, c, d, h, w) = dshape
x_data = np.random.uniform(size=(n, c, d, h, w)).astype("float32")
elif layout == "NDHWC":
(n, d, h, w, c) = dshape
x_data = np.random.uniform(size=(n, d, h, w, c)).astype("float32")
if method == "nearest_neighbor":
ref_res = tvm.topi.testing.upsampling3d_python(
x_data, (scale_d, scale_h, scale_w), layout
)
else:
ref_res = tvm.topi.testing.trilinear_resize3d_python(
x_data,
(int(round(d * scale_d)), int(round(h * scale_h)), int(round(w * scale_w))),
layout,
)
x = relay.Var("x", relay.TensorType(dshape, "float32"))
scale_d_var = relay.var("scale_d", relay.TensorType((), "float32"))
scale_h_var = relay.var("scale_h", relay.TensorType((), "float32"))
scale_w_var = relay.var("scale_h", relay.TensorType((), "float32"))
z = relay.nn.upsampling3d(
x,
scale_d_var,
scale_h_var,
scale_w_var,
method=method,
layout=layout,
coordinate_transformation_mode=coord_trans,
)
zz = run_infer_type(z)
func = relay.Function([x, scale_d_var, scale_h_var, scale_w_var], z)
for target, dev in enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind, mod=mod, device=dev, target=target)
op_res = intrp.evaluate()(
x_data,
np.array(scale_d).astype("float32"),
np.array(scale_h).astype("float32"),
np.array(scale_w).astype("float32"),
)
tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-4, atol=1e-6)
def partition_suffix(mod, quantized_dtypes):
"""Extract output dequantization expressions from `mod['main']`.
Parameters
----------
mod : tvm.IRModule
Module containing a quantized inference function
quantized_dtypes : Set[str]
Set of data types allowed in quantized operators
Returns
-------
pre_mod : tvm.IRModule
Module containing the input quantization function
mid_mod : tvm.IRModule
Module containing a function with everything except for input quantization
"""
assert len(mod.functions) == 1
func = mod['main']
suffix_cutter = SuffixCutter(quantized_dtypes)
post_body = suffix_cutter.visit(func.body)
assert not func.type_params, 'unimplemented'
assert func.attrs is None, 'unimplemented'
post_func = relay.Function(
relay.analysis.free_vars(post_body),
post_body,
func.ret_type)
post_mod = tvm.IRModule.from_expr(post_func)
mid_body = suffix_cutter.mid_body
if mid_body is None:
# The suffix contains the entire function, meaning there was no
# quantization boundary in the given mod. In this case, we use the
# suffix mod as the middle mod and make the suffix an identity function.
mid_mod = post_mod
post_body = relay.Var('input', mid_mod['main'].ret_type)
post_func = relay.Function(
[post_body],
post_body)
post_mod = tvm.IRModule.from_expr(post_func)
else:
mid_func = relay.Function(
func.params,
mid_body)
mid_mod = tvm.IRModule.from_expr(mid_func)
return mid_mod, post_mod
def verify_upsampling(dshape,
scale_h,
scale_w,
layout,
method,
align_corners=False):
if layout == "NCHW":
(n, c, h, w) = dshape
x_data = np.random.uniform(size=(n, c, h, w)).astype("float32")
elif layout == "NHWC":
(n, h, w, c) = dshape
x_data = np.random.uniform(size=(n, h, w, c)).astype("float32")
ref_res = tvm.topi.testing.resize2d_python(
x_data,
(scale_h, scale_w),
layout,
method[2:] if method[0:2] == "bi" else method,
"align_corners" if align_corners else "asymmetric",
)
x = relay.Var("x", relay.TensorType(dshape, "float32"))
scale_h_var = relay.var("scale_h", relay.TensorType((), "float32"))
scale_w_var = relay.var("scale_h", relay.TensorType((), "float32"))
z = relay.nn.upsampling(x,
scale_h_var,
scale_w_var,
method=method,
layout=layout,
align_corners=align_corners)
zz = run_infer_type(z)
func = relay.Function([x, scale_h_var, scale_w_var], z)
for target, dev in tvm.testing.enabled_targets():
for kind in ["vm", "debug"]:
mod = tvm.ir.IRModule.from_expr(func)
intrp = relay.create_executor(kind,
mod=mod,
device=dev,
target=target)
op_res = intrp.evaluate()(x_data,
np.array(scale_h).astype("float32"),
np.array(scale_w).astype("float32"))
tvm.testing.assert_allclose(op_res.numpy(),
ref_res,
rtol=1e-4,
atol=1e-6)
