def test_ref():
d = relay.Var("d")
r = relay.Var("r")
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)
assert alpha_equal(dcpe(square), Function([d], d * d))
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_triangle_number():
t = relay.TensorType([], "int32")
x = Var("x", t)
f_var = Var("f")
f = Function([x], If(op.equal(x, const(0)), const(0), x + f_var(x - const(1))))
orig = run_infer_type(Let(f_var, f, f_var(const(10))))
assert_alpha_equal(dcpe(orig), const(55))
def test_if_ref():
shape = ()
dtype = "bool"
t = TensorType(shape, dtype)
d = Var("d", t)
r = Var("r")
update = Function([], RefWrite(r, RefRead(r) + RefRead(r)))
u = Var("u")
body = If(d, u(), u())
eff = Var("eff")
body = Let(eff, body, RefRead(r))
f = Function([d], Let(r, RefCreate(const(1)), Let(u, update, body)))
pe_f = tipe(f)
f_res = create_executor().evaluate(f)(const(True))
pe_f_res = create_executor().evaluate(pe_f)(const(True))
np.testing.assert_allclose(f_res.numpy(), 2 * np.ones_like(f_res.numpy()))
np.testing.assert_allclose(pe_f_res.numpy(), 2 * np.ones_like(pe_f_res.numpy()))
def test_function_invalidate():
shape = ()
dtype = "bool"
t = TensorType(shape, dtype)
d = Var("d", t)
r = Var("r")
fetch = Function([], RefRead(r))
fet = Var("fetch")
fet_obscured = Var("fetch_obscured")
u = Var("u")
body = If(d, fet_obscured(), fet_obscured())
body = Let(u, RefWrite(r, const(1)), body)
body = Let(fet_obscured, If(d, fet, fet), body)
body = Let(fet, fetch, body)
body = Let(r, RefCreate(const(0)), body)
f = Function([d], body)
pe_f = tipe(f)
f_res = create_executor().evaluate(f)(const(True))
pe_f_res = create_executor().evaluate(pe_f)(const(True))
np.testing.assert_allclose(f_res.numpy(), np.ones_like(f_res.numpy()))
np.testing.assert_allclose(pe_f_res.numpy(), np.ones_like(pe_f_res.numpy()))
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]))
