def test_recursive_func():
mod = tvm.IRModule({})
x = relay.var('x', shape=[], dtype='int32')
fn0 = relay.Function([x], x)
gx = relay.GlobalVar("gx")
mod[gx] = fn0
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = relay.ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype='int32'))
global_call = gx(i)
rec_call = relay.Call(sum_up, [one_less]) + global_call
sb.ret(relay.add(rec_call, i))
func = relay.Function([i],
sb.get(),
ret_type=relay.TensorType([], 'int32'))
func = func.set_attribute("Compiler", tvm.tir.StringImm("a"))
mod[sum_up] = func
iarg = relay.var('i', shape=[], dtype='int32')
mod["main"] = relay.Function([iarg], sum_up(iarg))
call_graph = relay.CallGraph(mod)
assert call_graph.is_recursive(sum_up)
assert call_graph.ref_count(sum_up) == 2
assert call_graph.ref_count(gx) == 1
assert call_graph.ref_count("main") == 0
def get_mod():
mod = tvm.IRModule({})
x = relay.var('x', shape=[], dtype='int32')
fn0 = relay.Function([x], x)
fn0 = fn0.set_attribute("Inline", tvm.tir.IntImm("int32", 1))
gx = relay.GlobalVar("gx")
mod[gx] = fn0
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = relay.ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype="int32"))
global_call = gx(i)
rec_call = relay.Call(sum_up, [one_less]) + global_call
sb.ret(relay.add(rec_call, i))
func = relay.Function([i],
sb.get(),
ret_type=relay.TensorType([], "int32"))
func = func.set_attribute("Inline", tvm.tir.IntImm("int32", 1))
mod[sum_up] = func
iarg = relay.var("i", shape=[], dtype='int32')
mod["main"] = relay.Function([iarg], sum_up(iarg))
return mod
def after():
data = relay.var("data", shape=(1, 32))
eq1 = relay.var("e1", shape=[], dtype="float32")
eq2 = relay.var("e2", shape=[], dtype="float32")
cb_1 = relay.annotation.compiler_begin(eq1, target)
cb_2 = relay.annotation.compiler_begin(eq2, target)
equality_condition = relay.equal(cb_1, cb_2)
ce_1 = relay.annotation.compiler_end(equality_condition, target)
# if condition
cb_3 = relay.annotation.compiler_begin(data, target)
true_branch = relay.tanh(cb_3)
ce_2 = relay.annotation.compiler_end(true_branch, target)
# else condition
cb_4 = relay.annotation.compiler_begin(data, target)
false_branch = relay.sigmoid(cb_4)
ce_3 = relay.annotation.compiler_end(false_branch, target)
if_condition = relay.If(ce_1, ce_2, ce_3)
cb_5 = relay.annotation.compiler_begin(if_condition, target)
erf_out = relay.erf(cb_5)
ce_4 = relay.annotation.compiler_end(erf_out, target)
func = relay.Function([data, eq1, eq2], ce_4)
mod = tvm.IRModule.from_expr(func)
return mod
def test_equal():
i = relay.var("i", shape=[], dtype="int32")
j = relay.var("i", shape=[], dtype="int32")
z = relay.equal(i, j)
func = relay.Function([i, j], z, ret_type=relay.TensorType([], "bool"))
i_data = relay.const(0, "int32")
j_data = relay.const(0, "int32")
def _make_env_find(self, m, rval_t):
ctr = m['ctr']
gv = relay.GlobalVar(f"$_env_find<{ctr.name_hint}>")
env = relay.Var("env", env_type(env_val()))
key = relay.Var("key", relay.ty.scalar_type('int64'))
dft = relay.Var("dft", rval_t)
k = relay.Var("k")
v = relay.Var("v")
r = relay.Var("r")
x = relay.Var("x")
extract_clause = adt.Clause(
adt.PatternConstructor(ctr, [adt.PatternVar(x)]), x)
empty_clause = adt.Clause(adt.PatternConstructor(empty_env, []), dft)
cons_clause = adt.Clause(
adt.PatternConstructor(
cons_env,
[adt.PatternVar(k),
adt.PatternVar(v),
adt.PatternVar(r)]),
relay.If(relay.equal(key, k),
adt.Match(v, [extract_clause], complete=False),
relay.Call(gv, [r, key, dft])))
body = adt.Match(env, [empty_clause, cons_clause])
fn = relay.Function([env, key, dft], body, rval_t)
m['env_find'] = (gv, fn)
return gv, fn
def test_recursion():
"""
Program:
def @f(%n: int32, %data: float32) -> float32 {
if (%n == 0) {
%data
} else {
@f(%n - 1, log(%data))
}
}
"""
sb = relay.ScopeBuilder()
f = relay.GlobalVar("f")
ti32 = relay.scalar_type("int32")
tf32 = relay.scalar_type("float32")
n = relay.var("n", ti32)
data = relay.var("data", tf32)
with sb.if_scope(relay.equal(n, relay.const(0, ti32))):
sb.ret(data)
with sb.else_scope():
sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data)))
mod = tvm.IRModule()
mod[f] = relay.Function([n, data], sb.get())
mod = infer_mod(mod)
assert "@f(%1, %2)" in mod.astext()
assert mod["f"].checked_type == relay.FuncType([ti32, tf32], tf32)
def test_loop():
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 = 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))
result = get_serialized_output(mod, i_data, accum_data)
tvm.testing.assert_allclose(result.asnumpy(), sum(range(1,
loop_bound + 1)))
def test_recursion():
"""
Program:
def @f(%n: int32, %data: float32) -> float32 {
if (%n == 0) {
%data
} else {
@f(%n - 1, log(%data))
}
}
"""
sb = relay.ScopeBuilder()
f = relay.GlobalVar("f")
ti32 = relay.scalar_type("int32")
tf32 = relay.scalar_type("float32")
n = relay.var("n", ti32)
data = relay.var("data", tf32)
with sb.if_scope(relay.equal(n, relay.const(0, ti32))):
sb.ret(data)
with sb.else_scope():
sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data)))
mod = relay.Module()
mod[f] = relay.Function([n, data], sb.get())
assert "@f(%1, %2) /* ty=float32 */" in mod.astext()
assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
def _make_env_update(self, m, rval_t):
ctr = m['ctr']
gv = relay.GlobalVar(f"$_env_update<{ctr.name_hint}>")
env = relay.Var("env", env_type(env_val()))
key = relay.Var("key", relay.ty.scalar_type('int64'))
val = relay.Var("val", rval_t)
k = relay.Var("k")
v = relay.Var("v")
r = relay.Var("r")
empty_clause = adt.Clause(adt.PatternConstructor(empty_env, []),
cons_env(key, ctr(val), env))
cons_clause = adt.Clause(
adt.PatternConstructor(
cons_env,
[adt.PatternVar(k),
adt.PatternVar(v),
adt.PatternVar(r)]),
relay.If(relay.equal(key, k), cons_env(key, ctr(val), env),
cons_env(k, v, relay.Call(gv, [r, key, val]))))
body = adt.Match(env, [empty_clause, cons_clause])
fn = relay.Function([env, key, val], body, env_type(env_val()))
m['env_update'] = (gv, fn)
return gv, fn
def test_recursion():
"""
Program:
let sum_twice(n: i32) -> i32 = {
m = (n * 2)
if (n == 0) {
return m;
} else {
return m + sum(n - 1);
}
}
sum_twice(5);
"""
return # cannot be run as fuse_ops need to recursively visit
mod = relay.Module()
i64 = relay.TensorType((), 'int64')
f = relay.GlobalVar("f")
n = relay.Var("n", i64)
m = n * relay.const(2, 'int64')
funcbody = relay.If(relay.equal(n, relay.const(0, 'int64')), m,
m + f(n - relay.const(1, 'int64')))
value = relay.Function([n], funcbody, i64, [])
mod[f] = value
check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
old_f = mod[f]
mod = transform.ToANormalForm()(mod)
f = mod[f]
check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
def test_loop():
mod = relay.module.Module({})
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
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))
vm = create_vm(mod)
ser = serializer.Serializer(vm)
code, lib = ser.serialize()
deser = deserializer.Deserializer(code, lib)
des_vm = deser.deserialize()
result = veval(des_vm, i_data, accum_data)
tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))
def test_recursion():
"""
Program:
let f(n: i32) -> i32 = {
m = (n * 2)
if (n == 0) {
return m;
} else {
return m + f(n - 1);
}
}
f(5);
"""
mod = tvm.IRModule()
i64 = relay.TensorType((), 'int64')
f = relay.GlobalVar("f")
n = relay.Var("n", i64)
m = n * relay.const(2, 'int64')
funcbody = relay.If(relay.equal(n, relay.const(0, 'int64')), m,
m + f(n - relay.const(1, 'int64')))
value = relay.Function([n], funcbody, i64, [])
mod[f] = value
check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
old_f = mod[f]
mod = transform.ToANormalForm()(mod)
f = mod[f]
check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
def test_loop():
mod = relay.module.Module({})
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
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))
exe = create_exec(mod)
code, lib = exe.save()
des_exec = _vm.Executable.load_exec(code, lib)
des_vm = _vm.VirtualMachine(des_exec)
des_vm.init(tvm.cpu())
result = veval(des_vm, i_data, accum_data)
tvm.testing.assert_allclose(result.asnumpy(), sum(range(1,
loop_bound + 1)))
def test_recursion():
"""
Program:
let f(n: i32) -> i32 = {
m = (n * 2)
if (n == 0) {
return m;
} else {
return m + f(n - 1);
}
}
f(5);
"""
mod = tvm.IRModule()
i64 = relay.TensorType((), "int64")
f = relay.GlobalVar("f")
n = relay.Var("n", i64)
m = n * relay.const(2, "int64")
cond = relay.equal(n, relay.const(0, "int64"))
false_branch = m + f(n - relay.const(1, "int64"))
funcbody = relay.If(cond, m, false_branch)
value = relay.Function([n], funcbody, i64, [])
mod[f] = value
check_eval(f(relay.const(5, "int64")), 30.0, mod=mod)
old_f = mod[f]
mod = transform.ToBasicBlockNormalForm()(mod)
f = mod[f]
check_eval(f(relay.const(5, "int64")), 30.0, mod=mod)
check_basic_block_normal_form(f)
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=mod)
def test_recursion():
"""
Program:
def f(n: i32, data: f32) -> f32 {
if (n == 0) {
return data;
} else {
return f(n - 1, log(data));
}
}
"""
sb = relay.ScopeBuilder()
f = relay.GlobalVar("f")
ti32 = relay.scalar_type("int32")
tf32 = relay.scalar_type("float32")
n = relay.var("n", ti32)
data = relay.var("data", tf32)
with sb.if_scope(relay.equal(n, relay.const(0, ti32))):
sb.ret(data)
with sb.else_scope():
sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data)))
mod = relay.Module()
mod[f] = relay.Function([n, data], sb.get())
assert "%3 = @f(%1, %2)" in mod.astext()
assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
def test_equal():
i = relay.var('i', shape=[], dtype='int32')
j = relay.var('i', shape=[], dtype='int32')
z = relay.equal(i, j)
func = relay.Function([i, j], z, ret_type=relay.TensorType([], 'bool'))
i_data = relay.const(0)
j_data = relay.const(0)
check_eval(func, [i_data, j_data], True)
def test_equal():
i = relay.var('i', shape=[], dtype='int32')
j = relay.var('i', shape=[], dtype='int32')
z = relay.equal(i, j)
func = relay.Function([i, j], z, ret_type=relay.TensorType([], 'bool'))
i_data = relay.const(0, 'int32')
j_data = relay.const(0, 'int32')
check_eval(func, [i_data, j_data], True)
def if_expr(x):
'\n free_var %x: float32\n %0 = equal(%x, 2f);\n if (%0) {\n %1 = add(%x, 1f);\n multiply(%1, 2f)\n } else {\n multiply(%1, 1f)\n }\n '
one = relay.const(1, dtype='float32')
two = relay.const(2, dtype='float32')
v1 = relay.add(x, one)
v2 = relay.equal(x, two)
true_branch = relay.multiply(v1, two)
false_branch = relay.multiply(v1, one)
body = relay.If(v2, true_branch, false_branch)
return body
def expected_if_expr(x):
'\n free_var %x: float32\n let %v1: float32 = add(%x, 1f /* ty=float32 */) /* ty=float32 */;\n %0 = equal(%x, 2f /* ty=float32 */) /* ty=bool */;\n if (%0) {\n multiply(%v1, 2f /* ty=float32 */) /* ty=float32 */\n } else {\n multiply(%v1, 1f /* ty=float32 */) /* ty=float32 */\n }\n '
one = relay.const(1, dtype='float32')
two = relay.const(2, dtype='float32')
v1 = relay.var('v1')
v2 = relay.equal(x, two)
true_branch = relay.multiply(v1, two)
false_branch = relay.multiply(v1, one)
body = relay.If(v2, true_branch, false_branch)
body = relay.Let(v1, relay.add(x, one), body)
return body
def get_func_with_control_flow():
data = relay.var("data", shape=(1, 3, 224, 224))
weight = relay.var("weight", shape=(32, 3, 3, 3))
eq1 = relay.var("e1", shape=[], dtype="float32")
eq2 = relay.var("e2", shape=[], dtype="float32")
eq = relay.equal(eq1, eq2)
true_branch = relay.zeros(shape=(1, 32, 222, 222), dtype="float32")
false_branch = relay.nn.conv2d(data, weight, kernel_size=(3, 3), channels=32)
ife = relay.If(eq, true_branch, false_branch)
out = relay.erf(ife)
return relay.Function([data, weight, eq1, eq2], out)
def before():
data = relay.var("data", shape=(1, 32))
eq1 = relay.var("e1", shape=[], dtype="float32")
eq2 = relay.var("e2", shape=[], dtype="float32")
eq = relay.equal(eq1, eq2)
true_branch = relay.tanh(data)
false_branch = relay.sigmoid(data)
ife = relay.If(eq, true_branch, false_branch)
out = relay.erf(ife)
func = relay.Function([data, eq1, eq2], out)
mod = tvm.IRModule.from_expr(func)
return mod
def test_simple_loop():
mod = tvm.IRModule({})
sum_up = relay.GlobalVar("sum_up")
i = relay.var("i", shape=[], dtype="int32")
sb = ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype="int32"))
rec_call = relay.Call(sum_up, [one_less])
sb.ret(relay.add(rec_call, i))
func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32"))
mod[sum_up] = func
i_data = np.array(10, dtype="int32")
check_eval(sum_up, [i_data], sum(range(1, 11)), mod=mod)
def test_simple_loop():
mod = relay.module.Module({})
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype='int32'))
rec_call = relay.Call(sum_up, [one_less])
sb.ret(relay.add(rec_call, i))
func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], 'int32'))
mod[sum_up] = func
i_data = np.array(10, dtype='int32')
check_eval(sum_up, [i_data], sum(range(1, 11)), mod=mod)
def test_loop():
mod = relay.module.Module({})
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
i_data = np.array(10, dtype='int32')
accum_data = np.array(0, dtype='int32')
check_eval(sum_up, [i_data, accum_data], sum(range(1, 11)), mod=mod)
def get_recursive_count_loop():
mod = tvm.IRModule({})
sum_up = relay.GlobalVar("sum_up")
i = relay.var("i", shape=[], dtype="int32")
sb = relay.ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype="int32"))
rec_call = relay.Call(sum_up, [one_less])
sb.ret(relay.add(rec_call, i))
func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32"))
func = func.with_attr("Inline", tvm.tir.IntImm("int32", 1))
mod[sum_up] = func
iarg = relay.var("i", shape=[], dtype="int32")
mod["main"] = relay.Function([iarg], sum_up(iarg))
return mod, sum_up
def test_loop():
mod = relay.module.Module({})
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))):
sb.ret(accum)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1))
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
i_data = np.array(10, dtype='int32')
accum_data = np.array(0, dtype='int32')
check_eval(sum_up, [i_data, accum_data], sum(range(1, 11)), mod=mod)
def relay_take_grad_inp(c, _nb_indices, _indices, _values):
assert _nb_indices.is_constant(int)
values = c.ref(_values)
r_indices = relay.reshape(c.ref(_indices),
tuple(_indices.abstract.xshape()) + (1, ))
n_rows = _nb_indices.value
n_cols = _values.abstract.xshape()[-1]
outputs = []
indices_dtype = type_to_np_dtype(_indices.abstract.element.xtype())
out_dtype = type_to_np_dtype(_values.abstract.element.xtype())
for i in range(n_rows):
select_entries = relay.equal(r_indices, relay.const(i, indices_dtype))
casted_select = relay.cast(select_entries, out_dtype)
select_dout = relay.multiply(casted_select, values)
reshape_out = relay.reshape(select_dout, (-1, n_cols))
vector = relay.sum(reshape_out, 0)
outputs.append(relay.reshape(vector, (1, n_cols)))
return relay.concatenate(outputs, 0)
def test_count_loop():
mod = relay.module.Module({})
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype='int32'))
rec_call = relay.Call(sum_up, [one_less])
sb.ret(relay.add(rec_call, i))
func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], 'int32'))
mod[sum_up] = func
i_data = np.array(0, dtype='int32')
iarg = relay.var('i', shape=[], dtype='int32')
mod[mod.entry_func] = relay.Function([iarg], sum_up(iarg))
result = veval(mod, i_data)
tvm.testing.assert_allclose(result.asnumpy(), i_data)
def test_count_loop():
mod = relay.module.Module({})
sum_up = relay.GlobalVar('sum_up')
i = relay.var('i', shape=[], dtype='int32')
sb = ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype='int32'))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype='int32'))
rec_call = relay.Call(sum_up, [one_less])
sb.ret(relay.add(rec_call, i))
func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], 'int32'))
mod[sum_up] = func
i_data = np.array(0, dtype='int32')
iarg = relay.var('i', shape=[], dtype='int32')
mod[mod.entry_func] = relay.Function([iarg], sum_up(iarg))
result = veval(mod, i_data)
tvm.testing.assert_allclose(result.asnumpy(), i_data)
def test_recursion():
'\n Program:\n let f(n: i32) -> i32 = {\n m = (n * 2)\n if (n == 0) {\n return m;\n } else {\n return m + f(n - 1);\n }\n }\n f(5);\n '
mod = tvm.IRModule()
i64 = relay.TensorType((), 'int64')
f = relay.GlobalVar('f')
n = relay.Var('n', i64)
m = (n * relay.const(2, 'int64'))
cond = relay.equal(n, relay.const(0, 'int64'))
false_branch = (m + f((n - relay.const(1, 'int64'))))
funcbody = relay.If(cond, m, false_branch)
value = relay.Function([n], funcbody, i64, [])
mod[f] = value
check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
old_f = mod[f]
mod = transform.ToBasicBlockNormalForm()(mod)
f = mod[f]
check_eval(f(relay.const(5, 'int64')), 30.0, mod=mod)
check_basic_block_normal_form(f)
def if_expr(x):
"""
free_var %x: float32
%0 = equal(%x, 2f);
if (%0) {
%1 = add(%x, 1f);
multiply(%1, 2f)
} else {
multiply(%1, 1f)
}
"""
one = relay.const(1, dtype="float32")
two = relay.const(2, dtype="float32")
v1 = relay.add(x, one)
v2 = relay.equal(x, two)
true_branch = relay.multiply(v1, two)
false_branch = relay.multiply(v1, one)
body = relay.If(v2, true_branch, false_branch)
return body
def test_count_loop(target, dev):
mod = tvm.IRModule({})
sum_up = relay.GlobalVar("sum_up")
i = relay.var("i", shape=[], dtype="int32")
sb = ScopeBuilder()
with sb.if_scope(relay.equal(i, relay.const(0, dtype="int32"))):
sb.ret(i)
with sb.else_scope():
one_less = relay.subtract(i, relay.const(1, dtype="int32"))
rec_call = relay.Call(sum_up, [one_less])
sb.ret(relay.add(rec_call, i))
func = relay.Function([i], sb.get(), ret_type=relay.TensorType([], "int32"))
mod[sum_up] = func
i_data = np.array(0, dtype="int32")
iarg = relay.var("i", shape=[], dtype="int32")
mod["main"] = relay.Function([iarg], sum_up(iarg))
result = veval(mod, i_data, device=dev, target=target)
tvm.testing.assert_allclose(result.numpy(), i_data)
check_result(target, dev, [i_data], i_data, mod)
def expected_if_expr(x):
"""
free_var %x: float32
let %v1: float32 = add(%x, 1f /* ty=float32 */) /* ty=float32 */;
%0 = equal(%x, 2f /* ty=float32 */) /* ty=bool */;
if (%0) {
multiply(%v1, 2f /* ty=float32 */) /* ty=float32 */
} else {
multiply(%v1, 1f /* ty=float32 */) /* ty=float32 */
}
"""
one = relay.const(1, dtype="float32")
two = relay.const(2, dtype="float32")
v1 = relay.var("v1")
v2 = relay.equal(x, two)
true_branch = relay.multiply(v1, two)
false_branch = relay.multiply(v1, one)
body = relay.If(v2, true_branch, false_branch)
body = relay.Let(v1, relay.add(x, one), body)
return body
def test_sum_loop():
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
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([i_data, accum_data], sum(range(1, loop_bound + 1)), mod=mod)
def test_sum_loop():
mod = relay.module.Module({})
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
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[mod.entry_func] = relay.Function([iarg, aarg], sum_up(iarg, aarg))
result = veval(mod, i_data, accum_data)
tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))
