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()))
