def __init__(self, *, name="f", **kwargs): name = f"{name}_{Network.cnt}" Network.cnt += 1 if len(Network.stack) is not 0: mod = Network.stack[-1].mod p = Network.stack[-1].p else: mod = Module() p = Prelude(mod) add_nat_definitions(p) self.mod = mod self.p = p self.inputs = [] self.weights = OrderedSet() self.sub_network = OrderedSet() self.f = relay.GlobalVar(name) self.recurse = relay.Var("recurse") self.use_recurse = False self.ret_type = None body = self.build(**kwargs) assert isinstance(body, relay.Expr) if self.use_recurse: inputs = [copy_var(v) for v in self.inputs] body = relay.Let( self.recurse, relay.Function(inputs, self.call_from_outside(*inputs)), body) self.mod[self.f] = relay.Function(self.inputs + self.all_weights(), body, self.ret_type)
def test_nat_3(): mod = Module() p = Prelude(mod) add_nat_definitions(p) cfunc = compile(Function([], p.s(p.s(p.s(p.z())))), mod) output = cfunc() assert nat_to_int(output) == 3
def test_double(): mod = Module() p = Prelude(mod) add_nat_definitions(p) orig = p.double(make_nat_expr(p, 3)) res = dcpe(orig, mod=mod) assert alpha_equal(res, make_nat_expr(p, 6))
def test_42(): mod = Module() func = Function([], relay.const(42)) cfunc = compile(func, mod) output = cfunc() np.testing.assert_allclose(output.asnumpy(), np.array(42.0, dtype='float32'))
def test_add_convert(): mod = Module() p = Prelude(mod) add_nat_definitions(p) cfunc = compile(p.add, mod) output = cfunc(int_to_nat(p, 12), int_to_nat(p, 34)) assert nat_to_int(output) == 46
def test_identity(): mod = Module() x = var('x', shape=()) func = Function([x], x) cfunc = compile(func, mod) a = tvm.nd.array(np.array(1.0, dtype='float32')) output = cfunc(a) np.testing.assert_allclose(output.asnumpy(), a.asnumpy())
def test_int_mult_3(): mod = Module() x = var('x', dtype='int32', shape=()) func = Function([x], x * relay.const(3)) cfunc = compile(func, mod) a = tvm.nd.array(np.array(4, dtype='int32')) output = cfunc(a) np.testing.assert_allclose(output.asnumpy(), np.array(12, dtype='int32'))
def test_nat_add(): mod = Module() p = Prelude(mod) cfunc = compile( Function([], p.add(p.s(p.s(p.s(p.z()))), p.s(p.s(p.s(p.s(p.z())))))), mod) output = cfunc() assert nat_to_int(output) == 7
def test_constructor_tag_round_trip(): mod1 = Module() p1 = Prelude(mod1) add_nat_definitions(p1) mod2 = Module() p2 = Prelude(mod2) add_nat_definitions(p2) # ensure hashes match across modules ctors1 = constructor_list(p1) ctors2 = constructor_list(p2) for i in range(len(ctors1)): tag = ctors1[i].tag ctor = mod2.get_constructor(tag) assert ctor == ctors2[i] assert ctor.name_hint == ctors1[i].name_hint
def test_add_42(): mod = Module() x = var('x', shape=()) func = Function([x], x + relay.const(42.0)) cfunc = compile(func, mod) a = tvm.nd.array(np.array(42.0, dtype='float32')) output = cfunc(a) np.testing.assert_allclose(output.asnumpy(), np.array(84.0, dtype='float32'))
def test_map(): mod = Module() p = Prelude(mod) f = Var("f") orig = p.map(f, p.cons(const(1), p.cons(const(2), p.cons(const(3), p.nil())))) expected = p.cons(f(const(1)), p.cons(f(const(2)), p.cons(f(const(3)), p.nil()))) assert alpha_equal(dcpe(orig, mod=mod), expected)
def test_loop(): mod = Module() t = TypeVar("t") x = Var("x", t) loop = GlobalVar("loop") mod[loop] = Function([x], loop(x), t, [t]) res = dcpe(loop(const(1)), mod=mod) expected = Call(loop, [const(1)], None, [None]) assert alpha_equal(res, expected)
def test_head_cons(): mod = Module() p = Prelude(mod) hd = p.hd t = TypeVar("t") x = Var("x", t) body = hd(p.cons(x, p.nil())) f = Function([x], body, None, [t]) res = dcpe(f, mod) assert alpha_equal(res, Function([x], x, t, [t]))
def test_double(): mod = Module() x = var('x', shape=()) double = GlobalVar('double') mod[double] = Function([x], x + x) x = var('x', shape=()) cfunc = compile(Function([x], double(double(x))), mod) a = tvm.nd.array(np.array(1.5, dtype='float32')) output = cfunc(a) np.testing.assert_allclose(output.asnumpy(), np.array(6.0, dtype='float32'))
def test_tuple(): mod = Module() cfunc = compile( Function([], relay.TupleGetItem( relay.Tuple([ relay.const(3, dtype='int32'), relay.const(4.0, dtype='float32') ]), 1)), mod) np.testing.assert_allclose(cfunc().asnumpy(), np.array(4.0, dtype='float32'))
def test_global_match_nat_id(): mod = Module() p = Prelude(mod) add_nat_definitions(p) nat = p.nat() x = Var("x", nat) z_case = Clause(PatternConstructor(p.z, []), p.z()) s_case = Clause(PatternConstructor(p.s, [PatternVar(x)]), p.s(x)) orig = Match(make_nat_expr(p, 3), [z_case, s_case]) res = dcpe(orig, mod=mod) assert alpha_equal(res, make_nat_expr(p, 3))
def test_add(): mod = Module() x = var('x', shape=()) y = var('y', shape=()) z = x + y func = Function([x, y], z) cfunc = compile(func, mod) a = tvm.nd.array(np.array(1.0, dtype='float32')) b = tvm.nd.array(np.array(1.0, dtype='float32')) c = tvm.nd.array(np.array(2.0, dtype='float32')) output = cfunc(a, b) np.testing.assert_allclose(output.asnumpy(), c.asnumpy())
def test_compose(): mod = Module() p = Prelude(mod) 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_recur_sum_global(): mod = Module() x = var('x', dtype='int32', shape=()) sum = GlobalVar('sum') c = relay.const(0) mod[sum] = Function([x], relay.If(op.less(x, c), c, x + sum(x - relay.const(1))), relay.TensorType(dtype='int32', shape=())) cfunc = compile(Function([], sum(relay.const(10))), mod) output = cfunc() np.testing.assert_allclose(output.asnumpy(), np.array(55, dtype='int32'))
def test_recur_sum_local(): mod = Module() x = var('x', dtype='int32', shape=()) t = relay.TensorType(dtype='int32', shape=()) sum = relay.Var('sum', type_annotation=relay.FuncType([t], t)) c = relay.const(0) func = Function([x], relay.If(op.less(x, c), c, x + sum(x - relay.const(1))), t) body = relay.Let(sum, func, sum(relay.const(10))) cfunc = compile(Function([], body), mod) output = cfunc() np.testing.assert_allclose(output.asnumpy(), np.array(55, dtype='int32'))
def test_swap_loop(): mod = Module() p = Prelude(mod) add_nat_definitions(p) nat = p.nat() x = Var("x", nat) y = Var("y", nat) loop = GlobalVar("loop") mod[loop] = Function([x, y], loop(y, x), nat) prog = loop(make_nat_expr(p, 1), make_nat_expr(p, 2)) res = dcpe(prog, mod=mod) assert alpha_equal(prog, res)
def test_nat_id(): mod = Module() p = Prelude(mod) add_nat_definitions(p) nat = p.nat() x = Var("x", nat) y = Var("y", nat) nat_id = GlobalVar("nat_id") mod[nat_id] = Function([x], x) orig = nat_id(make_nat_expr(p, 3)) res = dcpe(orig, mod=mod) assert alpha_equal(res, make_nat_expr(p, 3))
def test_loop(): mod = Module() t = TypeVar("t") x = Var("x", t) loop = GlobalVar("loop") mod[loop] = Function([x], loop(x), t, [t]) expected = Call(loop, [const(1)]) mod["main"] = Function([], expected) expected = mod["main"].body call = Function([], loop(const(1))) res = dcpe(call, mod=mod) assert alpha_equal(res.body, expected)
def test_mult_op(): mod = Module() x = var('x', shape=()) y = var('y', shape=()) z = x + y zz = op.exp(z) func = Function([x, y], zz) cfunc = compile(func, mod) a = tvm.nd.array(np.array(1.0, dtype='float32')) b = tvm.nd.array(np.array(1.0, dtype='float32')) output = cfunc(a, b) np.testing.assert_allclose(output.asnumpy(), np.exp(a.asnumpy() + b.asnumpy()))
def test_match_nat_id(): mod = Module() p = Prelude(mod) add_nat_definitions(p) nat = p.nat() x = Var("x", nat) y = Var("y", nat) nat_id = GlobalVar("nat_id") z_case = Clause(PatternConstructor(p.z, []), p.z()) s_case = Clause(PatternConstructor(p.s, [PatternVar(y)]), p.s(y)) mod[nat_id] = Function([x], Match(x, [z_case, s_case])) orig = nat_id(make_nat_expr(p, 3)) res = dcpe(orig, mod=mod) assert alpha_equal(res, make_nat_expr(p, 3))
def test_map(): mod = Module() p = Prelude(mod) f = GlobalVar("f") t = TypeVar("t") a = Var("a", t) mod[f] = Function([a], a, t, [t]) orig = p.map(f, p.cons(const(1), p.cons(const(2), p.cons(const(3), p.nil())))) expected = p.cons((const(1)), p.cons((const(2)), p.cons((const(3)), p.nil()))) expected = Function([], expected) mod["main"] = expected expected = mod["main"] orig = Function([], orig) res = dcpe(orig, mod=mod) assert alpha_equal(res.body, expected.body)
def test_abs(): mod = Module() x = var('x', shape=()) func = Function([x], relay.If(op.less(x, relay.const(0.0)), relay.const(-1.0) * x, x)) cfunc = compile(func, mod) a = tvm.nd.array(np.array(12.0, dtype='float32')) output = cfunc(a) np.testing.assert_allclose(output.asnumpy(), np.array(12.0, dtype='float32')) a = tvm.nd.array(np.array(-34.0, dtype='float32')) output = cfunc(a) np.testing.assert_allclose(output.asnumpy(), np.array(34.0, dtype='float32'))
def double_example(): # Declare a Relay module. mod = Module() # Implement the double function. x = var('x', shape=()) double = GlobalVar('double') mod[double] = Function([x], x + x) # Generate a function which calls double twice. x = var('x', shape=()) f = Function([x], double(double(x))) # Compile the function. cfunc = compile(f, mod) a = tvm.nd.array(np.array(1.5, dtype='float32')) print(cfunc(a).asnumpy())
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 = Module() 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_local_local_rec_outer_scope(): mod = Module() x = var('x', dtype='int32', shape=()) t = relay.TensorType(dtype='int32', shape=()) sum = relay.Var('sum', type_annotation=relay.FuncType([t], t)) c = relay.const(0) # we define a locally recursive function inside another function's scope # and have that function return the closure of the locally recursive function inner_func = Function([x], relay.If(op.less(x, c), c, x + sum(x - relay.const(1))), t) outer_func_body = relay.Let(sum, inner_func, sum) outer_func = Function([], outer_func_body) f = relay.Var('f') body = relay.Let(f, outer_func(), f(relay.const(10))) cfunc = compile(Function([], body), mod) output = cfunc() np.testing.assert_allclose(output.asnumpy(), np.array(55, dtype='int32'))