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