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 test_match_effect_exactly_once(): mod = tvm.IRModule() p = Prelude(mod) # the list should be of length 1! # Unless we mistakenly execute the data clause more than once r = relay.Var('r') data = seq(relay.RefWrite(r, p.cons(relay.Tuple([]), relay.RefRead(r))), relay.RefRead(r)) match = relay.Let( r, relay.RefCreate(p.nil()), relay.Match(data, [ relay.Clause(relay.PatternConstructor(p.nil, []), relay.const(0)), relay.Clause( relay.PatternConstructor(p.cons, [ relay.PatternWildcard(), relay.PatternConstructor(p.nil, []) ]), relay.const(1)), relay.Clause(relay.PatternWildcard(), relay.const(2)) ])) match_val = run_as_python(match, mod) assert_tensor_value(match_val, 1)
def test_add(): mod = relay.Module() p = Prelude(mod) nat = p.nat add = p.add s = p.s z = p.z ctx = tvm.context("llvm", 0) intrp = create_executor(mod=mod, ctx=ctx, target="llvm") assert mod[add].checked_type == relay.FuncType([nat(), nat()], nat()) assert count(intrp.evaluate(add(s(z()), s(z())))) == 2 assert count(intrp.evaluate(to_a_normal_form(add(s(z()), s(z())), mod))) == 2 assert "let" in mod[add].astext()
def test_list_tl_empty_list(): mod = relay.Module() p = Prelude(mod) nil = p.nil l = p.l tl = p.tl f = relay.Function([], tl(nil())) mod["main"] = f result = veval(mod) print(result)
def test_list_sum(): mod = tvm.IRModule() p = Prelude(mod) nil = p.nil cons = p.cons sum = p.sum l = cons(relay.const(1), cons(relay.const(2), cons(relay.const(3), nil()))) f = relay.Function([], sum(l)) mod["main"] = f for tgt, ctx in tvm.testing.enabled_targets(): result = veval(mod, ctx=ctx, target=tgt) tvm.testing.assert_allclose(result.asnumpy(), 6)
def test_list_tl_empty_list(): mod = tvm.IRModule() p = Prelude(mod) nil = p.nil l = p.l tl = p.tl f = relay.Function([], tl(nil())) mod["main"] = f for tgt, ctx in tvm.testing.enabled_targets(): result = veval(mod, ctx=ctx, target=tgt)
def test_pow(): mod = tvm.IRModule() p = Prelude(mod) add_nat_definitions(p) shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) double = relay.Function([x], x + x) i = relay.var("i", t) func = relay.Function([i], p.nat_iterate(double, make_nat_expr(p, 3))(i)) mod["main"] = func mod["main"] = gradient(mod["main"], mod=mod) m = transform.InferType()(mod) back_func = m["main"] assert back_func.checked_type == relay.FuncType( [t], relay.TupleType([t, relay.TupleType([t])])) i_nd = rand(dtype, *shape) ex = create_executor(mod=mod) forward, (grad_i, ) = ex.evaluate(back_func)(i_nd) tvm.testing.assert_allclose(forward.asnumpy(), 8 * i_nd.asnumpy()) tvm.testing.assert_allclose(grad_i.asnumpy(), 8 * np.ones_like(grad_i.asnumpy()))
def test_missing_in_the_middle(): 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, ), # empty list relay.Clause(relay.PatternConstructor(nil, []), v), # list of length 3 or more relay.Clause( relay.PatternConstructor( cons, [ relay.PatternWildcard(), relay.PatternConstructor( cons, [ relay.PatternWildcard(), relay.PatternConstructor( cons, [relay.PatternWildcard(), relay.PatternWildcard()] ), ], ), ], ), v, ), ], ) # fails to match a list of length exactly two unmatched = unmatched_cases(match, mod) assert len(unmatched) == 1 assert isinstance(unmatched[0], relay.PatternConstructor) assert unmatched[0].constructor == cons assert isinstance(unmatched[0].patterns[1], relay.PatternConstructor) assert unmatched[0].patterns[1].constructor == cons assert isinstance(unmatched[0].patterns[1].patterns[1], relay.PatternConstructor) assert unmatched[0].patterns[1].patterns[1].constructor == nil
def test_compose(): mod = Module() p = Prelude(mod) add_nat_definitions(p) x = relay.Var('x') inc = GlobalVar('inc') mod[inc] = Function([x], p.s(x)) x = relay.Var('x') func = GlobalVar('func') f = Function([x], relay.Call(p.compose(inc, p.double), [x])) mod[func] = f cfunc = compile(func, mod) assert nat_to_int(cfunc(p.s(p.s(p.z())))) == 5
def test_nat_add(): mod = relay.Module() p = Prelude(mod) add_nat_definitions(p) nat = p.nat add = p.add s = p.s z = p.z ctx = tvm.context("llvm", 0) intrp = create_executor(mod=mod, ctx=ctx, target="llvm") assert mod[add].checked_type == relay.FuncType([nat(), nat()], nat()) assert count(p, intrp.evaluate(add(s(z()), s(z())))) == 2 assert count(p, intrp.evaluate(to_a_normal_form(add(s(z()), s(z())), mod))) == 2 assert Feature.fLet in detect_feature(mod[add])
def test_match_nat_id(): mod = tvm.IRModule() p = Prelude(mod) p.mod.import_from_std("nat.rly") nat, z, s = p.mod.get_type("nat") x = Var("x", nat()) y = Var("y", nat()) nat_id = GlobalVar("nat_id") z_case = Clause(PatternConstructor(z, []), z()) s_case = Clause(PatternConstructor(s, [PatternVar(y)]), s(y)) mod[nat_id] = Function([x], Match(x, [z_case, s_case])) orig = nat_id(make_nat_expr(p, 3)) orig = Function([], orig) res = dcpe(orig, mod=mod) assert tvm.ir.structural_equal(res.body, make_nat_expr(p, 3))
def test_head_cons(): mod = Module() p = Prelude(mod) def hd_impl(): a = TypeVar("a") x = Var("x", p.l(a)) y = Var("y") z = Var("z") cons_case = Clause( PatternConstructor(p.cons, [PatternVar(y), PatternVar(z)]), y) y = Var("y") z = Var("z") return Function([x], Match(x, [cons_case]), a, [a]) t = TypeVar("t") x = Var("x", t) hd = Var("hd") body = Let(hd, hd_impl(), hd(p.cons(x, p.nil()))) f = Function([x], body, None, [t]) f = infer_type(f, mod=mod) res = dcpe(f) assert alpha_equal(res, Function([x], x, t, [t]))
def test_arbitrary_let_nesting(): # something that is tricky to do in Python but comes naturally in Relay mod = tvm.IRModule() p = Prelude(mod) x = relay.Var("x") r = relay.Var("r") y = relay.Var("y") z = relay.Var("z") expr = relay.Tuple([ relay.Let(x, relay.Tuple([relay.const(1), relay.const(2)]), relay.TupleGetItem(x, 1)), relay.Let( r, relay.RefCreate(relay.const(1)), seq(relay.RefWrite(r, relay.const(3)), relay.RefRead(r)), ), relay.Let(y, p.id(relay.Let(z, relay.const(4), z)), y), ]) tup_val = run_as_python(expr, mod) assert_adt_len(tup_val, 3) assert_tensor_value(tup_val[0], 2) assert_tensor_value(tup_val[1], 3) assert_tensor_value(tup_val[2], 4)
def convert_list_to_vmobj(py_lst): def wrap_nd_array(arr): return tvm.nd.array(arr, ctx=tvm.cpu(0)) mod = tvm.IRModule() prelude = Prelude(mod) adt_lst = ADT(prelude.nil.tag, []) for elem in reversed(py_lst): if isinstance(elem, np.ndarray): vmobj = wrap_nd_array(elem) elif isinstance(elem, tuple): vmobj = tuple_object([wrap_nd_array(e) for e in elem]) elif isinstance(elem, list): vmobj = convert_list_to_vmobj(elem) adt_lst = ADT(prelude.cons.tag, [vmobj, adt_lst]) return adt_lst
def test_list_map(target, dev): mod = tvm.IRModule() p = Prelude(mod) x = relay.var("x", "int32") add_one_func = relay.Function([x], relay.const(1) + x) _, cons, nil = mod.get_type("List") map = mod.get_global_var("map") l = cons(relay.const(2), cons(relay.const(1), nil())) f = relay.Function([], map(add_one_func, l)) mod["main"] = f result = veval(mod, device=dev, target=target) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([3, 2]))
def test_list_foldr(target, dev): mod = tvm.IRModule() p = Prelude(mod) _, cons, nil = mod.get_type("List") foldr = mod.get_global_var("foldr") x = relay.var("x") y = relay.var("y") identity_func = relay.Function([x, y], cons(x, y)) l = cons(relay.const(1), cons(relay.const(2), cons(relay.const(3), nil()))) f = relay.Function([], foldr(identity_func, nil(), l)) mod["main"] = f result = veval(mod, device=dev, target=target) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([1, 2, 3]))
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_list_foldl(): mod = tvm.IRModule() p = Prelude(mod) _, cons, nil = mod.get_type("List") foldl = mod.get_global_var("foldl") x = relay.var("x") y = relay.var("y") rev_dup_func = relay.Function([y, x], cons(x, cons(x, y))) l = cons(relay.const(1), cons(relay.const(2), cons(relay.const(3), nil()))) f = relay.Function([], foldl(rev_dup_func, nil(), l)) mod["main"] = f for tgt, ctx in tvm.testing.enabled_targets(): result = veval(mod, ctx=ctx, target=tgt) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([3, 3, 2, 2, 1, 1]))
def test_nat_add(): mod = tvm.IRModule() p = Prelude(mod) p.mod.import_from_std("nat.rly") nat, z, s = p.mod.get_type("nat") add = p.mod.get_global_var("nat_add") dev = tvm.device("llvm", 0) intrp = create_executor(mod=mod, device=dev, target="llvm") assert mod[add].checked_type == relay.FuncType([nat(), nat()], nat()) assert count(p, intrp.evaluate(add(s(z()), s(z())))) == 2 expr = add(s(z()), s(z())) f = relay.GlobalVar("f") mod[f] = relay.Function([], expr) mod = transform.ToANormalForm()(mod) expr = mod["f"] assert count(p, intrp.evaluate(expr.body)) == 2 assert Feature.fLet in detect_feature(mod[add])
def test_list_foldl(): mod = tvm.IRModule() p = Prelude(mod) nil = p.nil cons = p.cons foldl = p.foldl x = relay.var("x") y = relay.var("y") rev_dup_func = relay.Function([y, x], cons(x, cons(x, y))) l = cons(relay.const(1), cons(relay.const(2), cons(relay.const(3), nil()))) f = relay.Function([], foldl(rev_dup_func, nil(), l)) mod["main"] = f result = veval(mod) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([3, 3, 2, 2, 1, 1]))
def test_list_map(): mod = relay.Module() p = Prelude(mod) x = relay.var('x', 'int32') add_one_func = relay.Function([x], relay.const(1) + x) nil = p.nil cons = p.cons map = p.map l = cons(relay.const(2), cons(relay.const(1), nil())) f = relay.Function([], map(add_one_func, l)) mod["main"] = f result = veval(mod) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([3, 2]))
def test_prelude(): p = Prelude() feats = detect_feature(p.mod) assert feats == set([ Feature.fVar, Feature.fGlobalVar, Feature.fConstant, Feature.fTuple, Feature.fTupleGetItem, Feature.fFunction, Feature.fOp, Feature.fCall, Feature.fLet, Feature.fIf, Feature.fConstructor, Feature.fMatch, ])
def test_list_tl(target, dev): mod = tvm.IRModule() p = Prelude(mod) l, cons, nil = mod.get_type("List") tl = mod.get_global_var("tl") one2 = cons(relay.const(1), nil()) one3 = cons(relay.const(2), one2) one4 = cons(relay.const(3), one3) f = relay.Function([], tl(one4)) mod["main"] = f result = veval(mod, device=dev, target=target) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([2, 1]))
def test_list_hd(target, dev): mod = tvm.IRModule() p = Prelude(mod) l, cons, nil = mod.get_type("List") hd = mod.get_global_var("hd") one2 = cons(relay.const(1), nil()) one3 = cons(relay.const(2), one2) one4 = cons(relay.const(3), one3) three = hd(one4) f = relay.Function([], three) mod["main"] = f result = veval(mod, device=dev, target=target) tvm.testing.assert_allclose(result.numpy(), 3)
def test_list_foldr(): mod = relay.Module() p = Prelude(mod) nil = p.nil cons = p.cons foldr = p.foldr x = relay.var("x") y = relay.var("y") identity_func = relay.Function([x, y], cons(x, y)) l = cons(relay.const(1), cons(relay.const(2), cons(relay.const(3), nil()))) f = relay.Function([], foldr(identity_func, nil(), l)) mod["main"] = f result = veval(mod) tvm.testing.assert_allclose(vmobj_to_list(result), np.array([1, 2, 3]))
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_global_var("tensor_array", dtype) tensor_array1 = tensor_array(relay.const(3)) write_func = p.get_global_var("tensor_array_write", dtype) scatter_func = p.get_global_var("tensor_array_scatter", dtype) tensor2 = p.get_tensor_ctor("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)) # indices array index = relay.var("index") # values array value_0 = relay.var("value_0") value_1 = relay.var("value_1") values_array = tensor_array(relay.const(2)) values_array = write_func(values_array, relay.const(0), tensor2(value_0)) values_array = write_func(values_array, relay.const(1), tensor2(value_1)) # create the scatter function tensor_array_scatter = scatter_func(tensor_array1, index, values_array) mod["main"] = relay.Function([v1, v2, v3, index, value_0, value_1], tensor_array_scatter) # 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) index_data = np.array([0, 1], dtype="int32") val1_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype) val2_data = np.random.uniform(low=0.0, high=8.0, size=(2, 3)).astype(dtype) expected = [val1_data, val2_data, v3_data] check_tensor_array( mod, expected, *(v1_data, v2_data, v3_data, index_data, val1_data, val2_data), dtype=dtype, )
def test_map(): mod = tvm.IRModule() p = Prelude(mod) rlist, cons, nil = p.mod.get_type("List") rmap = p.mod.get_global_var("map") f = GlobalVar("f") t = TypeVar("t") a = Var("a", t) mod[f] = Function([a], a, t, [t]) orig = rmap(f, cons(const(1), cons(const(2), cons(const(3), nil())))) expected = cons((const(1)), cons((const(2)), cons((const(3)), nil()))) expected = Function([], expected) mod["main"] = expected mod = transform.InferType()(mod) expected = mod["main"] orig = Function([], orig) res = dcpe(orig, mod=mod) assert tvm.ir.structural_equal(res.body, expected.body)
def test_constructor_tag_differences(): # ensure that if we have the type data for a given ADT, the tags # for the constructors of the *same ADT* are simple offsets from # each other mod = tvm.IRModule() p = Prelude(mod) add_nat_definitions(p) adts = adt_list(p) for adt in adts: data = mod[adt] for i in range(len(data.constructors) - 1): ctor1 = data.constructors[i] ctor2 = data.constructors[i + 1] assert ctor2.tag - ctor1.tag == 1 # make sure there is something present at the MSB assert ctor1.tag - i != 0 assert ctor2.tag - (i + 1) != 0
def test_list_hd(): mod = tvm.IRModule() p = Prelude(mod) l, cons, nil = mod.get_type("List") hd = mod.get_global_var("hd") one2 = cons(relay.const(1), nil()) one3 = cons(relay.const(2), one2) one4 = cons(relay.const(3), one3) three = hd(one4) f = relay.Function([], three) mod["main"] = f for tgt, ctx in tvm.testing.enabled_targets(): result = veval(mod, ctx=ctx, target=tgt) tvm.testing.assert_allclose(result.asnumpy(), 3)
def test_list_nth(target, dev): expected = list(range(10)) for i in range(len(expected)): mod = tvm.IRModule() p = Prelude(mod) _, cons, nil = mod.get_type("List") nth = mod.get_global_var("nth") l = nil() for i in reversed(expected): l = cons(relay.const(i), l) f = relay.Function([], nth(l, relay.const(i))) mod["main"] = f result = veval(mod, device=dev, target=target) tvm.testing.assert_allclose(result.numpy(), expected[i])
def test_list_nth(): expected = list(range(10)) for i in range(len(expected)): mod = relay.Module() p = Prelude(mod) nil = p.nil cons = p.cons nth = p.nth l = nil() for i in reversed(expected): l = cons(relay.const(i), l) f = relay.Function([], nth(l, relay.const(i))) mod["main"] = f result = veval(mod) tvm.testing.assert_allclose(result.asnumpy(), expected[i])
def test_pow(): mod = relay.Module() p = Prelude(mod) shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) double = relay.Function([x], x + x) i = relay.var("i", t) func = relay.Function([i], relay.Call(p.iterate(double, p.s(p.s(p.s(p.z())))), [i])) back_func = relay.ir_pass.infer_type(gradient(func, mod=mod), mod=mod) assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])])) i_nd = rand(dtype, *shape) ex = create_executor(mod=mod) forward, (grad_i,) = ex.evaluate(back_func)(i_nd) tvm.testing.assert_allclose(forward.asnumpy(), 8 * i_nd.asnumpy()) tvm.testing.assert_allclose(grad_i.asnumpy(), 8 * np.ones_like(grad_i.asnumpy()))