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)