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")
np = relay.Param(n, e.int32)
data = relay.Var("data")
datap = relay.Param(data, e.float32)
funcbody = relay.If(equal(n, convert(0)), data,
f(subtract(n, convert(1.0)), log(data)))
value = relay.Function([np, datap], e.float32, funcbody, [])
orig = relay.Let(f, funcbody, f(convert(2.0), convert(10000.0)), e.float32)
assert alpha_equal(dead_code_elimination(orig), orig)
assert alpha_equal(
dead_code_elimination(relay.Let(f, funcbody, e.three, e.float32)),
e.three)
def test_let_alpha_equal():
tt1 = relay.TensorType((), "float32")
tt2 = relay.TensorType((), "int8")
v1 = relay.Var("v1")
v1_wtype = relay.Var("v1", tt1)
v2 = relay.Var("v2")
v3 = relay.Var("v3")
let = relay.Let(v1, convert(2), v1)
mapped = relay.Let(v2, convert(2), v2)
assert alpha_equal(let, mapped)
mismatched_var = relay.Let(v2, convert(2), v3)
assert not alpha_equal(let, mismatched_var)
different_value = relay.Let(v2, convert(3), v2)
assert not alpha_equal(let, different_value)
different_body = relay.Let(v2, convert(3), convert(12))
assert not alpha_equal(let, different_body)
# specified types must match
let_with_type = relay.Let(v1_wtype, convert(2), v1_wtype)
same_type = relay.Let(v1_wtype, convert(2), v1_wtype)
assert alpha_equal(let_with_type, same_type)
assert not alpha_equal(let, let_with_type)
v2 = relay.Var("v1", tt2)
different_type = relay.Let(v2, convert(2), v2)
assert not alpha_equal(let_with_type, different_type)
def test_var_alpha_equal():
v1 = relay.Var("v1")
v2 = relay.Var("v2")
# normally only pointer equality
assert alpha_equal(v1, v1)
assert not alpha_equal(v1, v2)
# let node allows for setting the eq_map
l1 = relay.Let(v1, convert(1), v1)
l2 = relay.Let(v2, convert(1), v2)
l3 = relay.Let(v1, convert(1), v2)
assert alpha_equal(l1, l2)
assert not alpha_equal(l1, l3)
def __init__(self):
self.a = relay.Var("a")
self.b = relay.Var("b")
self.c = relay.Var("c")
self.d = relay.Var("d")
self.e = relay.Var("e")
self.x = relay.Var("x")
self.y = relay.Var("y")
self.z = relay.Var("z")
self.shape = tvm.convert([1, 2, 3])
self.tt = relay.TensorType(self.shape, "float32")
self.int32 = relay.TensorType([], "int32")
self.float32 = relay.TensorType([], "float32")
self.one = convert(1.0)
self.two = convert(2.0)
self.three = convert(3.0)
def test_var_alpha_equal():
v1 = relay.Var("v1")
v2 = relay.Var("v2")
# normally only pointer equality
assert alpha_equal(v1, v1)
assert not alpha_equal(v1, v2)
# let node allows for setting the eq_map
l1 = relay.Let(v1, convert(1), v1)
l2 = relay.Let(v2, convert(1), v2)
l3 = relay.Let(v1, convert(1), v2)
assert alpha_equal(l1, l2)
assert not alpha_equal(l1, l3)
# type annotations
tt1 = relay.TensorType([], "int32")
tt2 = relay.TensorType([], "int32")
tt3 = relay.TensorType([], "int64")
v3 = relay.Var("v3", tt1)
v4 = relay.Var("v4", tt2)
v5 = relay.Var("v5", tt3)
l4 = relay.Let(v3, convert(1), v3)
l5 = relay.Let(v4, convert(1), v4)
l6 = relay.Let(v5, convert(1), v5)
# same annotations
assert alpha_equal(l4, l5)
# different annotations
assert not alpha_equal(l4, l6)
# one null annotation
assert not alpha_equal(l1, l4)
def test_recursion():
"""
Program:
def f(n: i32, data: f32) -> f32 {
if (n == 0) {
return f(n - 1, log(data));
} else {
return data;
}
}
f(2, 10000);
"""
b = IRBuilder()
f = b.global_var('f')
n = b.param('n', ty='int32')
data = b.param('data', ty='float32')
with b.decl(f, n, data):
with b.if_scope(equal(n, convert(0))):
b.ret(f(subtract(n, convert(1)), log(data)))
with b.else_scope():
b.ret(data)
b.ret(f(convert(2.0), convert(10000.0)))
assert_decl_has_type(b.env, 'f', func_type(
['int32', 'float32'], 'float32'))
def test_if_alpha_equal():
v1 = relay.Var("v1")
v2 = relay.Var("v2")
if_sample = relay.If(v1, convert(1), relay.Tuple([convert(2), convert(3)]))
same = relay.If(v1, convert(1), relay.Tuple([convert(2), convert(3)]))
assert alpha_equal(if_sample, same)
different_cond = relay.If(v2, convert(1),
relay.Tuple([convert(2), convert(3)]))
assert not alpha_equal(if_sample, different_cond)
different_true = relay.If(v1, convert(2),
relay.Tuple([convert(2), convert(3)]))
assert not alpha_equal(if_sample, different_true)
different_false = relay.If(v1, convert(1), relay.Tuple([]))
assert not alpha_equal(if_sample, different_false)
def test_call_alpha_equal():
v1 = relay.Var("v1")
v2 = relay.Var("v2")
# attrs are compared only by pointer equality
attr1 = tvm.make.node("attrs.TestAttrs", name="attr", padding=(3, 4))
attr2 = tvm.make.node("attrs.TestAttrs", name="attr", padding=(3, 4))
tt1 = relay.TensorType((1, 2, 3), "float32")
tt2 = relay.TensorType((), "int8")
basic_args = [convert(1), convert(2), v2, relay.Tuple([])]
# manually writing out args to ensure that args does not rely on
# pointer equality
call = relay.Call(v1,
[convert(1), convert(2), v2,
relay.Tuple([])], attr1, [tt1])
same = relay.Call(v1, basic_args, attr1, [tt1])
assert alpha_equal(call, same)
different_fn = relay.Call(v2, basic_args, attr1, [tt1])
assert not alpha_equal(call, different_fn)
fewer_args = relay.Call(v1, [convert(1), convert(2), v2], attr1, [tt1])
assert not alpha_equal(call, fewer_args)
reordered_args = relay.Call(
v1,
[convert(2), convert(1), relay.Tuple([]), v2], attr1, [tt1])
assert not alpha_equal(call, reordered_args)
different_args = relay.Call(
v1, [convert(1), convert(2), convert(3)], attr1, [tt1])
assert not alpha_equal(call, different_args)
more_args = relay.Call(
v1,
[convert(1),
convert(2), v2,
relay.Tuple([]),
convert(3),
convert(4)], attr1, [tt1])
assert not alpha_equal(call, more_args)
different_attrs = relay.Call(v1, basic_args, attr2, [tt1])
assert not alpha_equal(call, different_attrs)
no_type_args = relay.Call(v1, basic_args, attr1)
assert not alpha_equal(call, no_type_args)
more_type_args = relay.Call(v1, basic_args, attr1, [tt1, tt2])
assert not alpha_equal(call, more_type_args)
different_type_arg = relay.Call(v1, basic_args, attr1, [tt2])
assert not alpha_equal(call, different_type_arg)
def test_tuple_alpha_equal():
v1 = relay.Var("v1")
v2 = relay.Var("v2")
# unit value is a valid tuple
assert alpha_equal(relay.Tuple([]), relay.Tuple([]))
tup = relay.Tuple([v1, convert(2), convert(3), relay.Tuple([convert(4)])])
same = relay.Tuple([v1, convert(2), convert(3), relay.Tuple([convert(4)])])
assert alpha_equal(tup, same)
# use the eq_map
let_tup = relay.Let(v1, tup, v1)
let_mapped = relay.Let(
v2,
relay.Tuple([v2, convert(2),
convert(3),
relay.Tuple([convert(4)])]), v2)
assert alpha_equal(let_tup, let_mapped)
more_fields = relay.Tuple(
[v1, convert(2),
convert(3), relay.Tuple([convert(4)]), v2])
assert not alpha_equal(tup, more_fields)
fewer_fields = relay.Tuple([v1, convert(2), convert(3)])
assert not alpha_equal(tup, fewer_fields)
different_end = relay.Tuple(
[v1, convert(2), convert(3),
relay.Tuple([convert(5)])])
assert not alpha_equal(tup, different_end)
different_start = relay.Tuple(
[v2, convert(2), convert(3),
relay.Tuple([convert(4)])])
assert not alpha_equal(tup, different_start)
longer_at_end = relay.Tuple(
[v1, convert(2),
convert(3),
relay.Tuple([convert(4), convert(5)])])
assert not alpha_equal(tup, longer_at_end)
def test_constant_alpha_equal():
x = convert(1)
y = convert(2)
assert alpha_equal(x, x)
assert not alpha_equal(x, y)
assert alpha_equal(x, convert(1))
def test_if():
orig = relay.If(convert(True), e.a, e.b)
assert alpha_equal(dead_code_elimination(orig), e.a)
