def test_tensorrt_dynamic_batch_conv():
if skip_codegen_test():
return
batches_to_test = [1, 1, 0, 2, 3, 0, 1, 3, 2]
x_shape = (relay.Any(), 32, 8, 8)
x_data = np.ones([max(batches_to_test)] +
list(x_shape)[1:]).astype("float32")
k_shape = (16, 32, 3, 3)
params = {"kernel": np.random.uniform(-1, 1, k_shape).astype("float32")}
result_arr = [{} for _ in range(len(batches_to_test))]
for use_trt in [True, False]:
x = relay.var("x", shape=x_shape, dtype="float32")
kernel = relay.var("kernel", shape=k_shape, dtype="float32")
out = relay.nn.conv2d(x,
kernel,
channels=16,
kernel_size=(3, 3),
groups=1)
f = relay.Function([x, kernel], out)
mod = tvm.IRModule()
mod["main"] = f
if use_trt:
mod, _ = tensorrt.partition_for_tensorrt(mod, params)
if not skip_runtime_test():
with relay.build_config(opt_level=3):
relay_exec = relay.create_executor("vm",
mod=mod,
ctx=tvm.cpu(0),
target="llvm")
for i, batch_size in enumerate(batches_to_test):
result_arr[i][use_trt] = relay_exec.evaluate()(
x_data[:batch_size, ...], **params)
if not skip_runtime_test():
for i in range(len(batches_to_test)):
assert_result_dict_holds(result_arr[i])
def test_relay_strided_slice_legalize(ifm_shape, begin, end):
ifm = relay.var("ifm", shape=ifm_shape, dtype="int8")
strided_slice = relay.op.strided_slice(ifm, begin, end)
func = relay.Function([ifm], strided_slice)
mod = tvm.IRModule()
mod["main"] = func
mod = relay.transform.InferType()(mod)
strided_slice_pattern_table = [
(
ethosu.StridedSliceParams.composite_name,
ethosu.strided_slice_pattern(),
lambda pat: ethosu.StridedSliceParams(pat).is_valid(),
),
]
mod = partition_ethosu_by_table(mod, strided_slice_pattern_table)
mod["tvmgen_default_ethos_u_main_0"] = dataflow_pattern.rewrite(
legalize.StridedSliceRewriter(), mod["tvmgen_default_ethos_u_main_0"]
)
mod["tvmgen_default_ethos_u_main_0"] = dataflow_pattern.rewrite(
legalize.NoOpRewriter(), mod["tvmgen_default_ethos_u_main_0"]
)
mod = relay.transform.InferType()(mod)
ext_func = mod["tvmgen_default_ethos_u_main_0"]
identity = ext_func.body
assert identity.op.name == "contrib.ethosu.identity"
# check that the strided_slice is still there
strided_slice = identity.args[0]
assert strided_slice.op.name == "strided_slice"
# check that identity's output shape matches strided slice's output shape
slice_shape = [a - b for a, b in zip(end, begin)]
assert list(identity.checked_type.shape) == slice_shape
def test_function_taking_adt_ref_tuple():
mod = tvm.IRModule()
prelude = relay.prelude.Prelude(mod)
intrp = create_executor("debug", mod)
_, cons, nil = prelude.mod.get_type("List")
nil_value = ConstructorValue(nil.tag, [], nil)
cons_value = ConstructorValue(
cons.tag,
[nd.array(np.random.rand(1, 10).astype("float32")), nil_value],
cons,
)
ref_value = RefValue(nd.array(np.random.rand(1, 10).astype("float32")))
tuple_value = container.tuple_object(
[nd.array(np.random.rand(1, 10).astype("float32")) for _ in range(10)]
)
id_func = intrp.evaluate(prelude.id)
res_nil = id_func(nil_value)
assert res_nil.tag == nil_value.tag
assert len(res_nil.fields) == 0
res_cons = id_func(cons_value)
assert res_cons.tag == cons_value.tag
assert len(res_cons.fields) == len(cons_value.fields)
tvm.testing.assert_allclose(res_cons.fields[0].numpy(), cons_value.fields[0].numpy())
assert isinstance(res_cons.fields[1], ConstructorValue)
assert res_cons.fields[1].tag == nil.tag
assert len(res_cons.fields[1].fields) == 0
res_ref = id_func(ref_value)
tvm.testing.assert_allclose(res_ref.value.numpy(), ref_value.value.numpy())
res_tuple = id_func(tuple_value)
for i in range(10):
tvm.testing.assert_allclose(res_tuple[i].numpy(), tuple_value[i].numpy())
def test_extern_dnnl():
if not tvm.get_global_func("relay.ext.dnnl", True):
print("skip because DNNL codegen is not available")
return
dtype = "float32"
ishape = (1, 32, 14, 14)
w1shape = (32, 1, 3, 3)
data0 = relay.var("data0", shape=(ishape), dtype=dtype)
weight0 = relay.var("weight0", shape=(w1shape), dtype=dtype)
data1 = relay.var("data0", shape=(ishape), dtype=dtype)
weight1 = relay.var("weight0", shape=(w1shape), dtype=dtype)
weight2 = relay.var("weight1", shape=(w1shape), dtype=dtype)
depthwise_conv2d_1 = relay.nn.conv2d(
data1, weight1, kernel_size=(3, 3), padding=(1, 1), groups=32
)
depthwise_conv2d_2 = relay.nn.conv2d(
depthwise_conv2d_1, weight2, kernel_size=(3, 3), padding=(1, 1), groups=32
)
out = relay.add(depthwise_conv2d_1, depthwise_conv2d_2)
f = relay.Function([data1, weight1, weight2], out)
ref_mod = tvm.IRModule()
ref_mod["main"] = f
f = set_external_func_attr(f, "dnnl", "dnnl_0")
call = relay.Call(f, [data0, weight0, weight0])
mod = tvm.IRModule.from_expr(call)
i_data = np.random.uniform(0, 1, ishape).astype(dtype)
w_data = np.random.uniform(0, 1, w1shape).astype(dtype)
ref_ex = relay.create_executor("graph", mod=ref_mod, ctx=tvm.cpu())
ref_res = ref_ex.evaluate()(i_data, w_data, w_data)
check_result(
mod, {"data0": i_data, "weight0": w_data}, (1, 32, 14, 14), ref_res.asnumpy(), tol=1e-5
)
def test_tensorrt_simple():
if skip_codegen_test():
return
dtype = "float32"
xshape = (1, 3, 2, 2)
yshape = (1, 3, 1, 1)
zshape = (1, 1, 1, 1)
x = relay.var("x", shape=(xshape), dtype=dtype)
y = relay.var("y", shape=(yshape), dtype=dtype)
z = relay.var("z", shape=(zshape), dtype=dtype)
w = z * (x + y)
out = relay.nn.relu(w)
f = relay.Function([x, y, z], out)
x_data = np.random.uniform(-1, 1, xshape).astype(dtype)
y_data = np.random.uniform(-1, 1, yshape).astype(dtype)
z_data = np.random.uniform(-1, 1, zshape).astype(dtype)
result_dict = dict()
for mode in ["vm", "graph"]:
for use_trt in [True, False]:
mod = tvm.IRModule()
mod["main"] = f
result_key = mode + ("_trt" if use_trt else "")
if use_trt:
mod, config = tensorrt.partition_for_tensorrt(mod)
with tvm.transform.PassContext(
opt_level=3, config={"relay.ext.tensorrt.options": config}
):
relay_exec = relay.create_executor(mode, mod=mod, ctx=tvm.gpu(0), target="cuda")
else:
with tvm.transform.PassContext(opt_level=3):
relay_exec = relay.create_executor(mode, mod=mod, ctx=tvm.gpu(0), target="cuda")
if not skip_runtime_test():
result_dict[result_key] = relay_exec.evaluate()(x_data, y_data, z_data)
if not skip_runtime_test():
assert_result_dict_holds(result_dict)
def test_too_specific_match():
mod = tvm.IRModule()
p = Prelude(mod)
_, cons, nil = mod.get_type("List")
v = relay.Var("v")
match = relay.Match(
v,
[
relay.Clause(
relay.PatternConstructor(
cons,
[
relay.PatternWildcard(),
relay.PatternConstructor(
cons, [relay.PatternWildcard(), relay.PatternWildcard()]
),
],
),
v,
)
],
)
unmatched = unmatched_cases(match, mod)
# will not match nil or a list of length 1
nil_found = False
single_length_found = False
assert len(unmatched) == 2
for case in unmatched:
assert isinstance(case, relay.PatternConstructor)
if case.constructor == nil:
nil_found = True
if case.constructor == cons:
assert isinstance(case.patterns[1], relay.PatternConstructor)
assert case.patterns[1].constructor == nil
single_length_found = True
def partition(dpu_target):
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
params = {}
params["weight"] = np.random.rand(16, 3, 3, 3).astype("float32")
params["bn_gamma"] = np.random.rand(16).astype("float32")
params["bn_beta"] = np.random.rand(16).astype("float32")
params["bn_mean"] = np.random.rand(16).astype("float32")
params["bn_var"] = np.random.rand(16).astype("float32")
mod = annotation(mod, params, dpu_target)
opt_pass = tvm.transform.Sequential([
transform.MergeCompilerRegions(),
transform.PartitionGraph(),
])
with tvm.transform.PassContext(opt_level=3):
mod = opt_pass(mod)
return mod
def test_all_primary_operands_tensor_constants():
dtype = "int8"
shape = (1, 3, 3, 32)
x0 = generate_variable("x0", shape, dtype)
x1 = generate_variable("x1", shape, dtype)
z1 = make_binary_op(
relay.qnn.op.add,
x0,
x1,
input_0_scale=0.0128,
input_0_zero_point=32,
input_1_scale=0.256,
input_1_zero_point=-64,
)
lf = relay.Function([x0, x1], z1, relay.TensorType(shape, dtype))
lf = set_composite_func_attr(lf, "cmsis-nn.qnn_add")
rng = np.random.default_rng(12345)
y0 = relay.const(rng.integers(-128, high=127, size=shape, dtype=dtype))
y1 = relay.const(rng.integers(-128, high=127, size=shape, dtype=dtype))
c0 = relay.Call(lf, [y0, y1])
ef = relay.Function([], c0, relay.TensorType(shape, dtype))
ev = relay.GlobalVar("external_function")
ef = set_external_func_attr(ef, "cmsis-nn", ev.name_hint)
c = relay.Call(ev, [])
mf = relay.Function([], c, relay.TensorType(shape, dtype))
mv = relay.GlobalVar("main")
mod = tvm.IRModule()
mod[ev] = ef
mod[mv] = mf
mod = relay.transform.InferType()(mod)
mod = ScalarToTensorConstants()(mod)
new_mod = relay.transform.InferType()(mod)
assert tvm.ir.structural_equal(mod[ev].body, new_mod[ev].body)
def run(dtype):
mod = tvm.IRModule()
p = Prelude(mod)
# tensor array
v1 = relay.var("v1")
v2 = relay.var("v2")
v3 = relay.var("v2")
tensor_array = p.get_var("tensor_array", dtype)
tensor_array1 = tensor_array(relay.const(3))
write_func = p.get_var("tensor_array_write", dtype)
split_func = p.get_var("tensor_array_split", dtype)
tensor2 = p.get_var("tensor2", dtype)
tensor_array1 = write_func(tensor_array1, relay.const(0), tensor2(v1))
tensor_array1 = write_func(tensor_array1, relay.const(1), tensor2(v2))
tensor_array1 = write_func(tensor_array1, relay.const(2), tensor2(v3))
# value tensor
value = relay.var("value")
# lengths tensor
ta_len = relay.var("length")
# create the scatter function
tensor_array_split = split_func(tensor_array1, tensor2(value), ta_len)
mod["main"] = relay.Function([v1, v2, v3, value, ta_len], tensor_array_split)
# initialize and check
v1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
v2_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
v3_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype)
value_data = np.random.uniform(low=0.0, high=8.0, size=(4, 3)).astype(dtype)
length_data = np.array([2, 2], dtype="int32")
expected = np.concatenate([value_data, v3_data])
expected = np.split(expected, indices_or_sections=[2, 4])
check_tensor_array(
mod, expected, *(v1_data, v2_data, v3_data, value_data, length_data), dtype=dtype
)
def test_before_partial_eval():
"""Test transformation before 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)
func = run_infer_type(func)
back_func = transform.gradient(func)
back_func = run_infer_type(back_func)
mod["main"] = back_func
seq = tvm.transform.Sequential([
transform.LazyGradientInit(),
transform.PartialEvaluate(),
transform.DeadCodeElimination()
])
mod = seq(mod)
back_func = mod["main"]
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_constant_propagation():
ones = np.ones(shape=(8, 8), dtype="float32")
def expected():
mod = tvm.IRModule()
x = relay.const(ones)
y = relay.var("y", shape=(8, 8))
x0 = relay.const(ones)
y0 = relay.var("y0", shape=(8, 8))
add = x0 + y0
# Function that uses C compiler
func = relay.Function([y0], add)
func = set_func_attr(func, "ccompiler", "ccompiler_0")
glb_0 = relay.GlobalVar("ccompiler_0")
mod[glb_0] = func
add_call = relay.Call(glb_0, [y])
log = relay.log(add_call)
main = relay.Function([y], log)
mod["main"] = main
return mod
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 = WhiteListAnnotator(["add"], "ccompiler")(mod)
mod = transform.PartitionGraph()(mod)
expected_mod = expected()
assert tvm.ir.structural_equal(mod, expected_mod, map_free_vars=True)
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))
def expected():
mod = tvm.IRModule()
# function 0
f0_i0 = relay.var(target + "_0_i0", shape=(10, 10))
f0_i1 = relay.var(target + "_0_i1")
f0_i2 = relay.var(target + "_0_i2")
f0_i3 = relay.var(target + "_0_i3")
f0_i4 = relay.var(target + "_0_i4")
f0_n0 = relay.nn.batch_norm(f0_i0, f0_i1, f0_i2, f0_i3, f0_i4)
f0_n1 = f0_n0[1]
f0_n2 = relay.nn.relu(f0_n0[0])
f0_o0 = relay.Tuple([f0_n2, f0_n1])
func0 = relay.Function([f0_i0, f0_i1, f0_i2, f0_i3, f0_i4], f0_o0)
func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Compiler", target)
func0 = func0.with_attr("global_symbol", target + "_0")
gv0 = relay.GlobalVar(target + "_0")
mod[gv0] = func0
# body
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")
function_out = gv0(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)
get_out0 = relay.TupleGetItem(function_out, 0)
get_out1 = relay.TupleGetItem(function_out, 1)
out_2 = relay.tanh(get_out1)
out_3 = relay.log(get_out1)
out = relay.Tuple([get_out0, out_2, out_3])
func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar],
out)
mod["main"] = func
return mod
def test_recursive_concat():
"""
fn @concat_loop(%i: int32, %st: (any, 1)) -> (any, 1) {
if (%i < 10) {
let %i = reshape(cast(i, "float32"), newshape=(1, ))
let %new_st = concatenate((st, i), axis=0)
concat_loop(%i + 1, )
} else {
st
}
}
"""
# Initial Values.
i = relay.var('i', shape=(), dtype='int32')
st = relay.var('st', shape=(relay.Any(), 1), dtype='int32')
def _cond(i, st):
return relay.op.min(relay.op.less(i, int32(10)))
def _body(i, st):
i_vec = relay.op.reshape(i, (1,1))
ret = relay.op.concatenate([st, i_vec], axis=0)
return i + int32(1), ret
loop = while_loop(_cond, [i, st], _body)
start = relay.var('start', shape=(), dtype='int32')
body = loop(start, relay.op.reshape(relay.const(0), newshape=(1, 1)))
func = relay.Function([start], relay.TupleGetItem(body, 1))
mod = tvm.IRModule()
mod["main"] = func
data = np.array(0.0, dtype='int32')
ref = np.array([0] + list(range(10))).reshape((11, 1)).astype("int32")
# TODO(@jroesch): After LambdaLift pass, TypeInfer pass will fail
# so currently we cannot run this test case on VM
for kind in ["debug"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(data)
np.testing.assert_allclose(result.asnumpy(), ref)
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_constant_shape_with_external_codegen():
mod = tvm.IRModule()
shape = (relay.Any(), 25)
dtype = "float32"
# external function
x = relay.var("x", shape=shape, dtype=dtype)
weight = relay.const(np.random.rand(5, 25).astype("float32"), dtype="float32")
out = relay.nn.dense(x, weight)
f1 = relay.Function([x], out)
f1 = f1.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
f1 = f1.with_attr("Inline", tvm.tir.IntImm("int32", 1))
f1 = f1.with_attr("Compiler", "a")
glb_f1 = relay.GlobalVar("f1")
mod[glb_f1] = f1
mod = relay.transform.InferType()(mod)
# Main function
x = relay.var("x", shape=shape, dtype=dtype)
mod["main"] = relay.Function([x], glb_f1(x))
comp = relay.vm.VMCompiler()
opt_mod, _ = comp.optimize(mod, target="llvm")
assert "shape_func" in opt_mod.astext(False)
def verify_any_conv2d(
data_shape,
kernel_shape,
strides,
padding,
dilation,
static_data_shape,
ref_out_shape,
):
mod = tvm.IRModule()
dtype = "float32"
data = relay.var("data", shape=data_shape, dtype=dtype)
kernel = relay.var("kernel", shape=kernel_shape, dtype=dtype)
y = relay.nn.conv2d(data,
kernel,
strides,
padding,
dilation,
kernel_size=kernel_shape[2:4])
mod["main"] = relay.Function([data, kernel], y)
data_np = np.random.uniform(size=static_data_shape).astype(dtype)
kernel_np = np.random.uniform(size=kernel_shape).astype(dtype)
check_result([data_np, kernel_np], mod, ref_out_shape, assert_shape=True)
def verify_any_reshape(x_shape,
newshape,
x_np_shape,
out_shape,
variable_newshape=False):
x = relay.var("x", shape=x_shape, dtype="float32")
relu_x = relay.nn.relu(x)
data = np.random.uniform(size=x_np_shape).astype("float32")
params = [x]
args = [data]
if variable_newshape:
newshape_var = relay.var("newshape",
shape=(len(newshape), ),
dtype="int64")
params.append(newshape_var)
args.append(np.array(newshape, dtype="int64"))
newshape = newshape_var
y = relay.reshape(relu_x, newshape=newshape)
mod = tvm.IRModule()
mod["main"] = relay.Function(params, y)
check_result(args, mod, data, flatten=True)
def test_multiple_type_param_defn():
glob_typ_var = relay.GlobalTypeVar("Either")
typ_var_a = relay.TypeVar("A")
typ_var_b = relay.TypeVar("B")
prog = relay.TypeData(
glob_typ_var,
[typ_var_a, typ_var_b],
[
relay.Constructor("Left", [typ_var_a], glob_typ_var),
relay.Constructor("Right", [typ_var_b], glob_typ_var),
],
)
mod = tvm.IRModule()
mod[glob_typ_var] = prog
assert_parse_module_as(
"""
type Either[A, B] {
Left(A),
Right(B),
}
""",
mod,
)
def test_list_update():
expected = list(range(10))
mod = tvm.IRModule()
p = Prelude(mod)
nil = p.nil
cons = p.cons
update = p.update
l = nil()
# create zero initialized list
for i in range(len(expected)):
l = cons(relay.const(0), l)
# set value
for i, v in enumerate(expected):
l = update(l, relay.const(i), relay.const(v))
f = relay.Function([], l)
mod["main"] = f
result = veval(mod)
tvm.testing.assert_allclose(vmobj_to_list(result), np.array(expected))
def test_sum_loop(target, dev):
mod = tvm.IRModule({})
sum_up = relay.GlobalVar("sum_up")
i = relay.var("i", shape=[], dtype="int32")
accum = relay.var("accum", shape=[], dtype="int32")
sb = ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, "int32"))):
sb.ret(accum)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, "int32"))
new_accum = relay.add(accum, i)
sb.ret(relay.Call(sum_up, [one_less, new_accum]))
func = relay.Function([i, accum], sb.get())
mod[sum_up] = func
mod = relay.transform.InferType()(mod)
loop_bound = 0
i_data = np.array(loop_bound, dtype="int32")
accum_data = np.array(0, dtype="int32")
iarg = relay.var("i", shape=[], dtype="int32")
aarg = relay.var("accum", shape=[], dtype="int32")
mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg))
check_result(target, dev, [i_data, accum_data],
sum(range(1, loop_bound + 1)), mod)
def test_list_update():
expected = list(range(10))
mod = tvm.IRModule()
p = Prelude(mod)
_, cons, nil = mod.get_type("List")
update = mod.get_global_var("update")
l = nil()
# create zero initialized list
for i in range(len(expected)):
l = cons(relay.const(0), l)
# set value
for i, v in enumerate(expected):
l = update(l, relay.const(i), relay.const(v))
f = relay.Function([], l)
mod["main"] = f
for tgt, dev in tvm.testing.enabled_targets():
result = veval(mod, device=dev, target=tgt)
tvm.testing.assert_allclose(vmobj_to_list(result), np.array(expected))
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_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 = tvm.transform.Sequential([
relay.transform.InferType(),
relay.transform.FoldConstant(),
tvm.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 "PrintIR" in out
assert "multiply" in out
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", "ccompiler_0")
glb_0 = relay.GlobalVar("ccompiler_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 test_loop_free_var(target, dev):
x = relay.var("x", shape=(), dtype="int32")
i = relay.var("i", shape=(), dtype="int32")
s = relay.var("s", shape=(), dtype="int32")
def cond(i, _):
return i < relay.const(10, dtype="int32")
def body_no_free_var(i, acc):
incr = relay.const(1, "int32")
return i + incr, acc + i
def body_with_free_var(i, acc):
incr = relay.const(1, "int32")
return i + incr, acc + x
for args, body, expected in zip([[], [1]], [body_no_free_var, body_with_free_var], [45, 10]):
loop = while_loop(cond, [i, s], body)
tup = loop(relay.const(0, dtype="int32"), relay.zeros(shape=(), dtype="int32"))
ret = relay.TupleGetItem(tup, 1)
mod = tvm.IRModule()
mod["main"] = relay.Function(relay.analysis.free_vars(ret), ret)
check_result(target, dev, args, expected, mod)
def get_mod():
x1 = relay.var("x1", shape=None)
x2 = relay.var("x2", shape=None)
z1 = x1 + x2
lf = relay.Function([x1, x2], z1, relay.TensorType((), "float32"))
lf = set_composite_func_attr(lf, "cmsis-nn.qnn_add")
y0 = relay.expr.const(5, "float32")
y1 = relay.expr.const(3, "float32")
c0 = relay.Call(lf, [y0, y1])
ef = relay.Function([], c0, relay.TensorType((), "float32"))
ev = relay.GlobalVar("external_function")
ef = set_external_func_attr(ef, "foo", ev.name_hint)
c = relay.Call(ev, [])
mf = relay.Function([], c, relay.TensorType((), "float32"))
mv = relay.GlobalVar("main")
mod = tvm.IRModule()
mod[ev] = ef
mod[mv] = mf
mod = relay.transform.InferType()(mod)
return mod
def test_prim_func_pass():
@tvm.tir.transform.prim_func_pass(opt_level=1)
class TestReplaceFunc:
"""Simple test function to replace one argument to another."""
def __init__(self, new_func):
self.new_func = new_func
def transform_function(self, func, mod, ctx):
return self.new_func
x = te.var('x')
y = te.var('y')
b = tvm.tir.decl_buffer((x, ), "float32")
stmt = tvm.tir.LetStmt(x, 10, tvm.tir.Evaluate(x + 1))
func = tvm.tir.PrimFunc([x, y, b], stmt)
new_func = tvm.tir.PrimFunc([x, y, b], tvm.tir.Evaluate(0))
mod = tvm.IRModule({"main": func})
mod = TestReplaceFunc(new_func)(mod)
assert tvm.ir.structural_equal(mod["main"].body, new_func.body)
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_mixed_input_type():
mod = tvm.IRModule()
dtype = "float32"
static_data_shape = (9, 4)
data_shape = (relay.Any(), 4)
tensor_type = relay.TensorType(data_shape, dtype)
tuple_type = relay.TupleType([tensor_type, tensor_type])
data0 = relay.var("d0",
type_annotation=relay.TupleType(
[tuple_type, tensor_type]))
data1 = relay.var("d1", shape=(relay.Any(), 4), dtype=dtype)
data_tuple = relay.expr.TupleWrapper(data0, 2)
nested_data_tuple = relay.expr.TupleWrapper(data_tuple[0], 2)
y = nested_data_tuple[1] * data_tuple[1] + data1
mod["main"] = relay.Function([data0, data1], y)
data_np0 = np.random.uniform(size=static_data_shape).astype(dtype)
data_np1 = np.random.uniform(size=static_data_shape).astype(dtype)
ref_out_shape = (9, 4)
for kind in ["vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()([[data_np0, data_np0], data_np0], data_np1)
assert result.asnumpy().shape == ref_out_shape, \
"Shape mismatch: expect %s but got %s." % (str(ref_out_shape), str(result.asnumpy().shape))
def make_partitioned_function(relay_op):
ifm0 = relay.analysis.free_vars(relay_op)
ifm_shape = ifm0[0].type_annotation.shape
ifm = relay.var("ifm", shape=ifm_shape, dtype="int8")
glb_ethosu = relay.GlobalVar("tvmgen_default_ethosu_main_0")
func = (relay.Function(ifm0, relay_op).with_attr("Inline", 1).with_attr(
"Compiler",
"ethos-u").with_attr("global_symbol",
"tvmgen_default_ethosu_main_0").with_attr(
"Primitive", 1))
mod = tvm.IRModule()
mod[glb_ethosu] = func
mod = relay.transform.InferType()(mod)
call = relay.Call(glb_ethosu, [ifm])
mod["main"] = relay.Function([ifm], call)
mod = relay.transform.InferType()(mod)
return mod
