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_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_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 __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_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_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_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_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_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_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_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_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_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_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 test_ref():
mod = 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_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'))
def test_get_valid_counts():
# Based on test_get_valid_counts in tvm's test_op_level5.
# Tests the case of a packed func returning a Relay tuple.
# Only checks the shapes of the output because the reference implementation
# is long and inconvenient.
shape = (1, 2500, 6)
score_threshold = 0
id_index = 0
score_index = 1
np_data = np.random.uniform(low=-2, high=2, size=shape).astype("float32")
mod = Module()
cfunc = compile(
relay.Function([],
relay.vision.get_valid_counts(relay.const(np_data),
score_threshold, id_index,
score_index).astuple()),
mod)
relay_out = cfunc()
out1 = relay_out[0].asnumpy()
out2 = relay_out[1].asnumpy()
assert out1.shape == (shape[0], )
assert out2.shape == shape
def __init__(self, do_aot, use_gpu, *args):
assert isinstance(do_aot, bool)
assert isinstance(use_gpu, bool)
self.mod = Module()
self.prelude = Prelude(self.mod)
self.use_gpu = use_gpu
self.context = tvm.gpu(0) if use_gpu else tvm.cpu(0)
self.target = tvm.target.cuda() if use_gpu else tvm.target.create('llvm')
self.executor = create_executor(mod=self.mod, ctx=self.context, target=self.target)
self.parameters = []
self.forward_var = relay.GlobalVar('forward_var')
# Set up forward pass.
inputs, body, ret_type = self.compute(*args)
self.inputs = inputs
forward_compute = relay.Function(inputs + list([p[0] for p in self.parameters]), body, ret_type)
self.mod[self.forward_var] = forward_compute
self.mod['main'] = self.mod[self.forward_var]
if do_aot:
self.forward = aot.compile(self.forward_var, self.mod, ctx=self.context, tgt=self.target)
else:
self.forward = self.executor.evaluate(self.forward_var)
self.args = [None] * len(inputs) + list([p[1] for p in self.parameters])