def test_monomorphic_let(): "Program: let x = 1; return x" b = IRBuilder() x = b.let('x', 1.0, value_type=scalar_type('float64')) b.ret(x) prog, env = b.get() assert_has_type(prog, scalar_type('float64'))
def check_single_op(opfunc): "Program: fn (x : float32) { let t1 = f(x); t1 }" b = IRBuilder() with b.function(('x', 'float32')) as func: x, = func.param_ids() t1 = b.let('t1', opfunc(x)) b.ret(t1) assert_has_type(func.to_func(), func_type(['float32'], 'float32'))
def test_tuple(): ib = IRBuilder() dup = ib.global_var('dup') x = ib.param('x') with ib.decl(dup, x): ib.ret(relay.Tuple([x, x])) # todo: why is this not generalized? fn_ty = func_type([tensor_type()], relay.TupleType([tensor_type(), tensor_type()])) assert_decl_has_type(ib.env, dup, fn_ty)
def test_let(): b = IRBuilder() x = b.let('x', 1) b.ret(x) prog, _ = b.get() assert isinstance(prog, Let) var = prog.var value = prog.value assert var.name_hint == 'x' assert var == prog.body assert isinstance(value, Constant) assert value.data.asnumpy() == np.array(1)
def test_decl(): """Program: def f(x : Tensor[f32, (10, 10)]) { let lx = log(x); return lx; } """ b = IRBuilder() x = b.param('x') with b.decl('f', x): lx = b.let('lx', log(x)) b.ret(lx) _, env = b.get() assert_decl_has_type(env, 'f', func_type(['float32'], 'float32'))
def test_concat(): """ Program: def try_concat2(x: Float(3, 2), y: Float(2, 2)) -> Float(5, 2) { return concat(x, y); } """ ib = IRBuilder() try_concat2 = ib.global_var('try_concat2') x = ib.param('x', ty=tensor_type(3, 2)) y = ib.param('y', ty=tensor_type(2, 2)) with ib.decl(try_concat2, x, y): ib.ret(concat(x, y)) fn_ty = func_type([tensor_type(3, 2), tensor_type(2, 2)], tensor_type(5, 2)) assert_decl_has_type(ib.env, try_concat2, fn_ty)
def test_dual_op(): """Program: fn (x : Tensor[f32, (10, 10)]) { let t1 = log(x); let t2 = add(t1, x); return t1; } """ b = IRBuilder() with b.function(('x', tensor_type(10, 10))) as func: x, = func.param_ids() t1 = b.let('t1', log(x)) t2 = b.let('t2', add(t1, x)) b.ret(t2) assert_has_type(func.to_func(), func_type(['float32'], 'float32'))
def test_add_broadcast_op(): """ Program: fn (x: Tensor[(10, 4), f32], y: Tensor[(5, 10, 1), f32]) -> Tensor[(5, 10, 4), f32] { return x + y; } """ b = IRBuilder() x = b.param('x', tensor_type(10, 4)) y = b.param('y', tensor_type(5, 10, 1)) with b.function(x, y) as func: b.ret(add(x.var, y.var)) b.ret(func) prog, env = b.get() ttype = tensor_type(5, 5, 5) expected_ty = func_type([ttype, ttype], ttype) assert_has_type(func.to_func(), expected_ty)
def check_binary_broadcast_op(opfunc): """ Program: fn (x: Tensor[(10, 4), f32], y: Tensor[(5, 10, 1), f32]) -> Tensor[(5, 10, 4), f32] { return x <op> y; } """ b = IRBuilder() x = b.param('x', tensor_type(10, 4)) y = b.param('y', tensor_type(5, 10, 1)) with b.function(x, y) as func: b.ret(opfunc(x.var, y.var)) b.ret(func) prog, env = b.get() expected_ty = func_type( [tensor_type(10, 4), tensor_type(5, 10, 1)], tensor_type(5, 10, 4)) assert_has_type(func.to_func(), expected_ty)
def check_binary_op(opfunc): """ Program: fn (x, y) { return x <op> y; } """ b = IRBuilder() x = b.param('x', tensor_type(5, 5, 5)) y = b.param('y', tensor_type(5, 5, 5)) with b.function(x, y) as func: b.ret(opfunc(x.var, y.var)) b.ret(func) prog, env = b.get() ttype = tensor_type(5, 5, 5) expected_ty = func_type([ttype, ttype], ttype) assert_has_type(func.to_func(), expected_ty)
def test_recursion(): """ Program: def f(n: i32, data: f32) -> f32 { if (n == 0) { return f(n - 1, log(data)); } else { return data; } } f(2, 10000); """ b = IRBuilder() f = b.global_var('f') n = b.param('n', ty='int32') data = b.param('data', ty='float32') with b.decl(f, n, data): with b.if_scope(equal(n, convert(0))): b.ret(f(subtract(n, convert(1)), log(data))) with b.else_scope(): b.ret(data) b.ret(f(convert(2.0), convert(10000.0))) assert_decl_has_type(b.env, 'f', func_type( ['int32', 'float32'], 'float32'))