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 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_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 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'))