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