def test_ref_write():
# check that the result of a ref write is an empty tuple
v = relay.Var('v')
initial_write = relay.Let(v,
relay.RefCreate(relay.Tuple([relay.const(1)])),
relay.RefWrite(v, relay.Tuple([relay.const(2)])))
write_val = run_as_python(initial_write)
assert_adt_len(write_val, 0)
# now ensure that the value, once written, can be read back
# (we read the value before and after mutation)
w = relay.Var('w')
read_after_write = relay.Let(
v, relay.RefCreate(relay.Tuple([relay.const(1)])),
relay.Let(
w, relay.RefCreate(relay.RefRead(v)),
seq(relay.RefWrite(v, relay.Tuple([relay.const(2)])),
relay.Tuple([relay.RefRead(w),
relay.RefRead(v)]))))
read_val = run_as_python(read_after_write)
assert_adt_len(read_val, 2)
assert_adt_len(read_val[0], 1)
assert_adt_len(read_val[1], 1)
assert_tensor_value(read_val[0][0], 1)
assert_tensor_value(read_val[1][0], 2)
def test_if_ref():
shape = ()
dtype = "bool"
t = relay.TensorType(shape, dtype)
d = relay.Var("d", t)
r = relay.Var("r")
update = relay.Function([],
relay.RefWrite(r,
relay.RefRead(r) +
relay.RefRead(r)))
u = relay.Var("u")
body = relay.If(d, u(), u())
eff = relay.Var("eff")
body = relay.Let(eff, body, relay.RefRead(r))
f = relay.Function([d],
relay.Let(r, relay.RefCreate(relay.const(1)),
relay.Let(u, update, body)))
f = infer_type(f)
pe_f = infer_type(partial_evaluate(f))
ex = create_executor()
f_res = ex.evaluate(f)(relay.const(True))
pe_f_res = ex.evaluate(pe_f)(relay.const(True))
np.testing.assert_allclose(f_res.asnumpy(),
2 * np.ones_like(f_res.asnumpy()))
np.testing.assert_allclose(pe_f_res.asnumpy(),
2 * np.ones_like(pe_f_res.asnumpy()))
def test_match_effect_exactly_once():
mod = tvm.IRModule()
p = Prelude(mod)
_, cons, nil = p.mod.get_type("List")
# 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, cons(relay.Tuple([]), relay.RefRead(r))),
relay.RefRead(r))
match = relay.Let(
r,
relay.RefCreate(nil()),
relay.Match(
data,
[
relay.Clause(relay.PatternConstructor(nil, []),
relay.const(0)),
relay.Clause(
relay.PatternConstructor(cons, [
relay.PatternWildcard(),
relay.PatternConstructor(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_ref():
d = relay.Var("d")
r = relay.Var("r")
x = relay.Var("x")
body = relay.RefRead(r)
body = relay.Let(x, relay.RefWrite(r, relay.RefRead(r) * relay.RefRead(r)), body)
body = relay.Let(r, relay.RefCreate(d), body)
square = relay.Function([d], body)
assert alpha_equal(dcpe(square), relay.Function([d], d * d))
def test_ref():
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
check_eval(body, 3)
check_eval(to_a_normal_form(body), 3)
def test_ref():
mod = tvm.IRModule()
three_with_ref = relay.GlobalVar("three_with_ref")
i = relay.Var("i")
iv = relay.Var("iv")
u = relay.Var("u")
uv = relay.Var("uv")
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
mod[three_with_ref] = relay.Function([], body)
def test_ref():
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
check_eval(body, 3)
opt_body = run_opt_pass(body, transform.ToANormalForm())
check_eval(opt_body, 3)
def test_ref():
i = relay.Var("i")
iv = relay.Var("iv")
u = relay.Var("u")
uv = relay.Var("uv")
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
check_eval(body, 3)
opt_body = run_opt_pass(body, transform.ToBasicBlockNormalForm())
check_eval(opt_body, 3)
check_basic_block_normal_form(opt_body)
def test_ref():
mod = relay.Module()
three_with_ref = relay.GlobalVar('three_with_ref')
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2)), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1)), body)
mod[three_with_ref] = relay.Function([], body)
check_eval(three_with_ref, [], 3, mod=mod)
def test_ref():
mod = relay.Module()
three_with_ref = relay.GlobalVar('three_with_ref')
i = relay.Var('i')
iv = relay.Var('iv')
u = relay.Var('u')
uv = relay.Var('uv')
body = relay.add(iv, uv)
body = relay.Let(uv, relay.RefRead(i), body)
body = relay.Let(u, relay.RefWrite(i, relay.const(2, dtype='int32')), body)
body = relay.Let(iv, relay.RefRead(i), body)
body = relay.Let(i, relay.RefCreate(relay.const(1, dtype='int32')), body)
mod[three_with_ref] = relay.Function([], body)
cfunc = compile(three_with_ref, mod)
output = cfunc()
np.testing.assert_allclose(output.asnumpy(), np.array(3, dtype='int32'))
def test_ref():
shape = (10, 10)
dtype = 'float32'
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
r = relay.Var("r")
u = relay.Var("u")
body = relay.RefRead(r)
body = relay.Let(u, relay.RefWrite(r, relay.RefRead(r) + relay.RefRead(r)), body)
body = relay.Let(r, relay.RefCreate(x), body)
func = relay.Function([x], body)
back_func = relay.ir_pass.infer_type(gradient(func))
assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])]))
x_nd = rand(dtype, *shape)
ex = create_executor()
forward, (grad_x,) = ex.evaluate(back_func)(x_nd)
tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy())
tvm.testing.assert_allclose(grad_x.asnumpy(), 2 * np.ones_like(grad_x.asnumpy()))
def test_ref():
x = relay.var("x", "float32")
y = relay.var("y", "float32")
r = relay.RefCreate(x)
st = relay.scalar_type("float32")
assert relay.ir_pass.infer_type(r).checked_type == relay.RefType(st)
g = relay.RefRead(r)
assert relay.ir_pass.infer_type(g).checked_type == st
w = relay.RefWrite(r, y)
assert relay.ir_pass.infer_type(w).checked_type == relay.TupleType([])
def test_if():
# we will have effects in the blocks to ensure only the intended one is executed
true_cond = relay.const(True)
false_cond = relay.const(False)
v = relay.Var('v')
true_branch = seq(relay.RefWrite(v, relay.const(1)), relay.RefRead(v))
false_branch = seq(relay.RefWrite(v, relay.const(2)), relay.RefRead(v))
true_expr = relay.Let(v, relay.RefCreate(relay.const(0)),
relay.If(true_cond, true_branch, false_branch))
false_expr = relay.Let(v, relay.RefCreate(relay.const(0)),
relay.If(false_cond, true_branch, false_branch))
true_val = run_as_python(true_expr)
assert_tensor_value(true_val, 1)
false_val = run_as_python(false_expr)
assert_tensor_value(false_val, 2)
def test_split():
"""Test that the result is well formed."""
x = relay.var("x", shape=(6, 9))
y = relay.split(x, 3).astuple()
a = relay.TupleGetItem(y, 0)
b = relay.TupleGetItem(y, 1)
c = relay.TupleGetItem(y, 2)
mod = relay.module.Module()
mod["main"] = relay.Function([x], a + relay.RefRead(relay.RefCreate(b)) + c)
mod = transform.FuseOps()(mod)
def test_cps_pe():
def destroy_ref(x):
x = run_infer_type(x)
x = to_cps(x)
x = run_infer_type(x)
y = un_cps(x)
y = run_infer_type(y)
# TODO(mbs): Revisit once DCE can eliminate dead writes.
x = run_opt_pass(
x,
tvm.transform.Sequential(
[
transform.PartialEvaluate(),
transform.InferType(),
transform.DeadCodeElimination(inline_once=True, ignore_impurity=True),
]
),
)
assert Feature.fRefCreate not in detect_feature(x)
unit = relay.Function([], relay.const(0.0, dtype="float32"))
f_ref = relay.Var("f_ref")
one = relay.const(1.0, dtype="float32")
two = relay.const(2.0, dtype="float32")
cond = relay.var(shape=(), dtype="uint1", name_hint="cond")
true_branch = relay.RefWrite(f_ref, relay.Function([], one))
false_branch = relay.RefWrite(f_ref, relay.Function([], two))
if_expr = relay.If(cond, true_branch, false_branch)
stmt = relay.Let(
f_ref,
relay.RefCreate(unit),
relay.Let(relay.Var("x"), if_expr, relay.Call(relay.RefRead(f_ref), [])),
)
F = relay.Function([cond], stmt)
destroy_ref(F)
G = relay.Function([cond], relay.If(cond, one, two))
G = run_infer_type(G)
G = relay.transform.gradient(G)
destroy_ref(G)
x = relay.var("x", shape=(1, 16))
y = relay.var("y", shape=(1, 16))
z = relay.var("z", shape=(1, 16))
cond = relay.var("cond", shape=(), dtype="uint1")
H = relay.If(cond, x, y)
H = relay.add(H, z)
H = relay.Function([cond, x, y, z], H)
H = run_infer_type(H)
H = relay.transform.gradient(H)
destroy_ref(H)
def test_ref():
t = relay.TensorType([], "float32")
d = relay.Var("d", t)
r = relay.Var("r", relay.RefType(t))
x = relay.Var("x")
body = relay.RefRead(r)
body = Let(x, RefWrite(r, RefRead(r) * RefRead(r)), body)
body = Let(r, RefCreate(d), body)
square = Function([d], body)
expected = run_opt_pass(Function([d], d * d), transform.InferType())
assert tvm.ir.structural_equal(dcpe(square), expected)
def test_ref():
t = relay.TensorType([], "float32")
d = relay.Var("d", t)
r = relay.Var("r", relay.RefType(t))
x = relay.Var("x")
body = relay.RefRead(r)
body = Let(x, RefWrite(r, RefRead(r) * RefRead(r)), body)
body = Let(r, RefCreate(d), body)
square = Function([d], body)
expected = transform.OptimizeOnExpr(Function([d], d * d),
transform.InferType())
assert alpha_equal(dcpe(square), expected)
def test_ref():
t = relay.TensorType([], "float32")
d = relay.Var("d", t)
r = relay.Var("r", relay.RefType(t))
x = relay.Var("x")
body = relay.RefRead(r)
body = Let(x, RefWrite(r, RefRead(r) * RefRead(r)), body)
body = Let(r, RefCreate(d), body)
square = Function([d], body)
expected = run_opt_pass(Function([d], d * d), transform.InferType())
# TODO(mbs): Revisit once DCE eliminates dead writes.
actual = dcpe(square, ignore_impurity=True)
assert tvm.ir.structural_equal(actual, expected)
def test_cps_pe():
def destroy_ref(x):
x = run_infer_type(x)
x = to_cps(x)
x = run_infer_type(x)
y = un_cps(x)
y = run_infer_type(y)
x = run_opt_pass(
x,
transform.Sequential([
transform.PartialEvaluate(),
transform.DeadCodeElimination(inline_once=True)
]))
assert Feature.fRefCreate not in detect_feature(x)
unit = relay.Function([], relay.const(0., dtype='float32'))
f_ref = relay.Var("f_ref")
one = relay.const(1., dtype='float32')
two = relay.const(2., dtype='float32')
cond = relay.var(shape=(), dtype='uint1', name_hint='cond')
true_branch = relay.RefWrite(f_ref, relay.Function([], one))
false_branch = relay.RefWrite(f_ref, relay.Function([], two))
if_expr = relay.If(cond, true_branch, false_branch)
stmt = relay.Let(
f_ref, relay.RefCreate(unit),
relay.Let(relay.Var("x"), if_expr, relay.Call(relay.RefRead(f_ref),
[])))
F = relay.Function([cond], stmt)
destroy_ref(F)
G = relay.Function([cond], relay.If(cond, one, two))
G = run_infer_type(G)
G = relay.transform.gradient(G)
destroy_ref(G)
x = relay.var("x", shape=(1, 16))
y = relay.var("y", shape=(1, 16))
z = relay.var("z", shape=(1, 16))
cond = relay.var("cond", shape=(), dtype='uint1')
H = relay.If(cond, x, y)
H = relay.add(H, z)
H = relay.Function([cond, x, y, z], H)
H = run_infer_type(H)
H = relay.transform.gradient(H)
destroy_ref(H)
def test_arbitrary_let_nesting():
# something that is tricky to do in Python but comes naturally in Relay
mod = relay.Module()
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 test_function_invalidate():
shape = ()
dtype = "bool"
t = relay.TensorType(shape, dtype)
d = relay.Var("d", t)
r = relay.Var("r")
fetch = relay.Function([], relay.RefRead(r))
fet = relay.Var("fetch")
fet_obscured = relay.Var("fetch_obscured")
u = relay.Var("u")
body = relay.If(d, fet_obscured(), fet_obscured())
body = relay.Let(u, relay.RefWrite(r, relay.const(1)), body)
body = relay.Let(fet_obscured, relay.If(d, fet, fet), body)
body = relay.Let(fet, fetch, body)
body = relay.Let(r, relay.RefCreate(relay.const(0)), body)
f = relay.Function([d], body)
f = infer_type(f)
pe_f = infer_type(partial_evaluate(f))
ex = create_executor()
f_res = ex.evaluate(f)(relay.const(True))
pe_f_res = ex.evaluate(pe_f)(relay.const(True))
np.testing.assert_allclose(f_res.asnumpy(), np.ones_like(f_res.asnumpy()))
np.testing.assert_allclose(pe_f_res.asnumpy(), np.ones_like(pe_f_res.asnumpy()))
def test_ref_execution_order():
# we want to have effects execute from left to right
x = relay.Var("x")
y = relay.Var("y")
f = relay.Var("f")
r = relay.Var("r")
expr = relay.Let(
f,
relay.Function([x, y], x),
# r = 1
relay.Let(
r,
relay.RefCreate(relay.const(1)),
relay.Tuple([
# should be 1
relay.RefRead(r),
# set r to 2 and read back
seq(relay.RefWrite(r, relay.const(2)), relay.RefRead(r)),
# set r to 3 and read back
seq(relay.RefWrite(r, relay.const(3)), relay.RefRead(r)),
# set r to 4 and read as first arg to f
# set r to 5 and read as second arg to f
# f should evaluate to 4
f(
seq(relay.RefWrite(r, relay.const(4)), relay.RefRead(r)),
seq(relay.RefWrite(r, relay.const(5)), relay.RefRead(r)),
),
# read back 5
relay.RefRead(r),
]),
),
)
tup_val = run_as_python(expr)
assert_adt_len(tup_val, 5)
assert_tensor_value(tup_val[0], 1)
assert_tensor_value(tup_val[1], 2)
assert_tensor_value(tup_val[2], 3)
assert_tensor_value(tup_val[3], 4)
assert_tensor_value(tup_val[4], 5)
def relay_universe_getitem(c, u, h):
return relay.RefRead(c.ref(h))
def test_ref_read():
v = relay.Var('v')
assign = relay.Let(v, relay.RefCreate(relay.Tuple([])), relay.RefRead(v))
read_val = run_as_python(assign)
assert_adt_len(read_val, 0)