diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/__pycache__/helpers.cpython-310.pyc b/__pycache__/helpers.cpython-310.pyc new file mode 100644 index 0000000..7b7d077 Binary files /dev/null and b/__pycache__/helpers.cpython-310.pyc differ diff --git a/batch/ast.py b/batch/ast.py index c02a76d..c358a68 100644 --- a/batch/ast.py +++ b/batch/ast.py @@ -17,16 +17,22 @@ def __init__(self, base): self.base = base self.base.ref_count += 1 - self.batch_size = size[0] + if (len(size) == 2): self.item_type = 'vec' self.dim = size[1] + self.batch_size = size[0] elif len(size) == 3: self.item_type = 'mat' self.dim1 = size[1] self.dim2 = size[2] + self.batch_size = size[0] elif len(size) == 1: self.item_type = 'scal' + self.batch_size = size[0] + elif len(size) == 0: + self.item_type = 'const' + self.batch_size = 0 else: raise TypeError('Batch item type not supported') @@ -69,8 +75,12 @@ def bvm(v1: Batch, v2: Batch): assert v1.dim == v2.dim1 return BatchOp('vec_mul_mat', v1, v2) +def bov(v1: Batch, v2: Batch): + assert v1.item_type == 'vec' and v2.item_type == 'vec' + return BatchOp('vec_outer_vec', v1, v2) + class BatchOp(Batch): - Types = ['scal_mul_vec', 'vec_mul_vec', 'vec_mul_mat'] + list(core.ast.op_mapping.keys()) + Types = ['scal_mul_vec', 'vec_mul_vec', 'vec_mul_mat', 'vec_outer_vec'] + list(core.ast.arith_op.keys()) def __init__(self, op_type, *operators): assert op_type in BatchOp.Types @@ -97,7 +107,7 @@ def __init__(self, op_type, *operators): name = f'{op_type}_' + '_'.join([op.name if hasattr(op, 'name') else '' for op in self.operators]) - if op_type in core.ast.op_mapping: + if op_type in core.ast.arith_op: match op_type: case 'add': res = self.operators[0].base + self.operators[1].base @@ -128,8 +138,12 @@ def __init__(self, op_type, *operators): res = Tensor(name, (bsize, dim), dtype) super().__init__(res) + elif op_type == 'vec_outer_vec': + bsize = self.operators[0].batch_size + res = Tensor(name, (bsize, self.operators[0].dim, self.operators[1].dim ), dtype) + super().__init__(res) + else: # TODO: complete other ops pass self.op_type = op_type - diff --git a/batch/ast2ir.py b/batch/ast2ir.py index 89ca85e..dde619d 100644 --- a/batch/ast2ir.py +++ b/batch/ast2ir.py @@ -1,3 +1,5 @@ +import sys +sys.path.append('/data/backed_up/lihhu/CUKE/cuke') from batch.ast import * from core.ast2ir import * @@ -19,13 +21,16 @@ def gen_ir(node): node.base._gen_ir() node.eval = node.base.eval elif type(node) == BatchOp: - if node.op_type in core.ast.op_mapping: + if node.op_type in core.ast.arith_op: node.operators[0]._gen_ir() node.operators[1]._gen_ir() node.base._gen_ir() + # print(node.operators[1].eval, node.op_type, node.base.compute) node.eval = node.base.eval node.decl = node.base.decl[:] node.compute = node.base.compute[:] + for i in node.compute: + i.ast_ref = node node.base.decl.clear() node.base.compute.clear() @@ -43,9 +48,9 @@ def gen_ir(node): res = bind(node.eval, pre_loop.iterate) inner_loop = Loop(0, node.operators[0].eval.size[1], 1, []) pre_loop.body.append(inner_loop) + pre_loop.ast_ref = node lhs = bind(lhs, inner_loop.iterate) rhs = bind(rhs, inner_loop.iterate) - assign = Assignment(res, Expr(lhs, rhs, '*'), '+') inner_loop.body.append(assign) @@ -63,6 +68,7 @@ def gen_ir(node): res = bind(node.eval, pre_loop.iterate) inner_loop = Loop(0, node.eval.size[1], 1, []) pre_loop.body.append(inner_loop) + pre_loop.ast_ref = node rhs = bind(rhs, inner_loop.iterate) res = bind(res, inner_loop.iterate) @@ -75,6 +81,7 @@ def gen_ir(node): node.operators[1]._gen_ir() size = helpers.get_ir_of_size(node._size()) node.base.eval = node.eval = Ndarray(node.dtype, size) + node.eval.val = 0 node.decl = [Decl(node.eval)] pre_loop = Loop(0, node.eval.size[0], 1, []) node.compute = [pre_loop] @@ -83,6 +90,7 @@ def gen_ir(node): res = bind(node.eval, pre_loop.iterate) loop1 = Loop(0, node.eval.size[1], 1, []) pre_loop.body.append(loop1) + pre_loop.ast_ref = node res = bind(res, loop1.iterate) loop2 = Loop(0, node.operators[0].eval.size[1], 1, []) loop1.body.append(loop2) @@ -92,7 +100,31 @@ def gen_ir(node): assign = Assignment(res, Expr(lhs, rhs, '*'), '+') loop2.body.append(assign) + + elif node.op_type == 'vec_outer_vec': + assert is_bvec(node.operators[0]) and is_bvec(node.operators[1]) + node.operators[0]._gen_ir() + node.operators[1]._gen_ir() + size = helpers.get_ir_of_size(node._size()) + node.base.eval = node.eval = Ndarray(node.dtype, size) + node.decl = [Decl(node.eval)] + pre_loop = Loop(0, node.eval.size[0], 1, []) + node.compute = [pre_loop] + lhs = bind(node.operators[0].eval, pre_loop.iterate) + rhs = bind(node.operators[1].eval, pre_loop.iterate) + res = bind(node.eval, pre_loop.iterate) + loop1 = Loop(0, node.eval.size[1], 1, []) + pre_loop.body.append(loop1) + pre_loop.ast_ref = node + lhs = bind(lhs, loop1.iterate) + res = bind(res, loop1.iterate) + loop2 = Loop(0, node.eval.size[2], 1, []) + loop1.body.append(loop2) + + rhs = bind(rhs, loop2.iterate) + res = bind(res, loop2.iterate) + assign = Assignment(res, Expr(lhs, rhs, '*')) + loop2.body.append(assign) return node - diff --git a/batch/opt/__init__.py b/batch/opt/__init__.py new file mode 100644 index 0000000..5be731f --- /dev/null +++ b/batch/opt/__init__.py @@ -0,0 +1,5 @@ +import batch.opt.fusion_rules +import batch.opt.ir +import batch.opt.node_wise.parallelism +import batch.opt.node_wise.smem +import batch.opt.node_wise.tiling \ No newline at end of file diff --git a/batch/opt/__pycache__/__init__.cpython-310.pyc b/batch/opt/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..bc0ac54 Binary files /dev/null and b/batch/opt/__pycache__/__init__.cpython-310.pyc differ diff --git a/batch/opt/__pycache__/fusion_rules.cpython-310.pyc b/batch/opt/__pycache__/fusion_rules.cpython-310.pyc new file mode 100644 index 0000000..a6cea75 Binary files /dev/null and b/batch/opt/__pycache__/fusion_rules.cpython-310.pyc differ diff --git a/batch/opt/__pycache__/ir.cpython-310.pyc b/batch/opt/__pycache__/ir.cpython-310.pyc new file mode 100644 index 0000000..cef551a Binary files /dev/null and b/batch/opt/__pycache__/ir.cpython-310.pyc differ diff --git a/batch/opt/__pycache__/parallelism.cpython-310.pyc b/batch/opt/__pycache__/parallelism.cpython-310.pyc new file mode 100644 index 0000000..3823759 Binary files /dev/null and b/batch/opt/__pycache__/parallelism.cpython-310.pyc differ diff --git a/batch/opt/__pycache__/smem.cpython-310.pyc b/batch/opt/__pycache__/smem.cpython-310.pyc new file mode 100644 index 0000000..a23ea4c Binary files /dev/null and b/batch/opt/__pycache__/smem.cpython-310.pyc differ diff --git a/batch/opt/__pycache__/tiling.cpython-310.pyc b/batch/opt/__pycache__/tiling.cpython-310.pyc new file mode 100644 index 0000000..528ffce Binary files /dev/null and b/batch/opt/__pycache__/tiling.cpython-310.pyc differ diff --git a/batch/opt/fusion_rules.py b/batch/opt/fusion_rules.py new file mode 100644 index 0000000..90114e3 --- /dev/null +++ b/batch/opt/fusion_rules.py @@ -0,0 +1,748 @@ +from core.ir import * +from batch.ast import * +from batch.ast2ir import * +import codegen + + +def get_same_loop(outloop, fusedloop): + pre_o = outloop + pre_i = fusedloop + oloop = outloop.body[0] + iloop = fusedloop.body[0] + + pre_iter_i = [] + pre_iter_o = [] + + pre_iter_i.append(pre_i.iterate) + pre_iter_o.append(pre_o.iterate) + while (isinstance(oloop, Loop) and isinstance(iloop, Loop)) and oloop.start == iloop.start and oloop.end.__name__ == iloop.end.__name__ and oloop.step == iloop.step: + if len(pre_i.body)>1: + break + pre_o = oloop + pre_i = iloop + pre_iter_i.append(pre_i.iterate) + pre_iter_o.append(pre_o.iterate) + + oloop = oloop.body[-1] + iloop = iloop.body[0] + + return pre_o, pre_i, pre_iter_o, pre_iter_i + +def loop_merge(o_loop, i_loop): + # print(codegen.cpu.to_string(o_loop), codegen.cpu.to_string(i_loop)) + if (isinstance(o_loop, Loop) and isinstance(i_loop, Loop)) and o_loop.start == i_loop.start and o_loop.end.__name__ == i_loop.end.__name__ and o_loop.step == i_loop.step: + + for ii in range(len(i_loop.body)-1): + + change_index(i_loop.body[ii], [o_loop.iterate], [i_loop.iterate]) + o_loop.body.insert(ii, i_loop.body[ii]) + # print('last loop i', count, codegen.cpu.to_string(i_loop.body[-1])) + change_index(i_loop.body[-1], [o_loop.iterate], [i_loop.iterate]) + # print("outer::",codegen.cpu.to_string(o_loop)) + # print("inner::",codegen.cpu.to_string(i_loop)) + if not isinstance(i_loop.body[-1], Loop): + o_loop.body.insert(-1, i_loop.body[-1]) + elif not isinstance(o_loop.body[-1], Loop): + o_loop.body.insert(0, i_loop.body[-1]) + else: + loop_merge(o_loop.body[-1], i_loop.body[-1]) + + +def change_index(iassign, iter_o, iter_i): + if isinstance(iassign, Indexing): + + for idx, item in enumerate(iter_i): + if iassign.idx == item: + iassign.idx = iter_o[idx] + if isinstance(iassign.dobject, Indexing): + change_index(iassign.dobject, iter_o, iter_i) + if isinstance(iassign.idx, Indexing): + change_index(iassign.idx, iter_o, iter_i) + + elif isinstance(iassign, Expr): + # both item.left and item.right + change_index(iassign.left, iter_o, iter_i) + change_index(iassign.right, iter_o, iter_i) + elif isinstance(iassign, Assignment): + # both item.lhs and item.rhs + change_index(iassign.lhs, iter_o, iter_i) + change_index(iassign.rhs, iter_o, iter_i) + elif isinstance(iassign, Loop): + for i in iassign.body: + change_index(i, iter_o, iter_i) + +def change_ref(ir, ast): + ir.astnode = ast + if isinstance(ir, Indexing): + if isinstance(ir.dobject, Indexing): + change_ref(ir.dobject, ast) + + elif isinstance(ir, Expr): + # both item.left and item.right + change_ref(ir.left, ast) + change_ref(ir.right, ast) + elif isinstance(ir, Assignment): + # both item.lhs and item.rhs + change_ref(ir.lhs, ast) + change_ref(ir.rhs, ast) + elif isinstance(ir, Loop): + for i in ir.body: + change_ref(i, ast) + # print('after::', i, i.ast_ref.op_type) + +def swap_arr_to_reg(ir, pre, cur): + if isinstance(ir, Indexing): + temp = ir + while isinstance(temp, Indexing): + temp = temp.dobject + if temp == pre: + return cur + else: + return ir + elif isinstance(ir, Expr): + ir.left = swap_arr_to_reg(ir.left, pre, cur) + ir.right = swap_arr_to_reg(ir.right, pre, cur) + elif isinstance(ir, Assignment): + ir.lhs = swap_arr_to_reg(ir.lhs, pre, cur) + ir.rhs = swap_arr_to_reg(ir.rhs, pre, cur) + elif isinstance(ir, Loop): + for i in range(len(ir.body)): + ir.body[i] = swap_arr_to_reg(ir.body[i], pre, cur) + return ir + +def fuse_elementwise(ast): + + if type(ast) == BatchOp: + if type(ast.operators[0]) == BatchOp: + fuse_elementwise(ast.operators[0]) + if type(ast.operators[1]) == BatchOp: + fuse_elementwise(ast.operators[1]) + else: + return + + if type(ast.operators[1]) == BatchOp and ast.op_type in core.ast.arith_op.keys(): + # fuse operators1 into elementwise + if ast.item_type == 'vec' and ast.operators[1].item_type == ast.item_type: + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[1].compute.clear() + elif ast.item_type == 'scal' and ast.operators[1].item_type == ast.item_type: + # check if type of operator1 is scalar + if ast.operators[1].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[1].compute.clear() + elif ast.item_type not in ['vec', 'scal']: + raise ValueError(f"Tensor type is wrong. Expect ast as 'vec' or 'scal' but found '{ast.item_type}'.") + elif ast.operators[1].item_type not in ['vec', 'scal']: + raise ValueError(f"Tensor type is wrong. Expect operators[1] as 'vec' or 'scal' but found '{ast.operators[1].item_type}'.") + elif ast.operators[1].item_type != ast.item_type: + raise ValueError(f"Tensor shape are not the same. Expect 'vec' or 'scal' but found ast: '{ast.item_type}' and operators[1]: '{ast.operators[1].item_type}'.") + + + if type(ast.operators[0]) == BatchOp and ast.op_type in core.ast.arith_op.keys(): + # fuse operators0 into elementwise + if ast.item_type == 'vec' and ast.operators[0].item_type == ast.item_type: + # check if type of operator1 is vector + if ast.operators[0].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + + iloop.body.clear() + ast.operators[0].compute.clear() + if ast.item_type == 'scal' and ast.operators[0].item_type == ast.item_type: + # check if type of operator1 is scalar + if ast.operators[0].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + + iloop.body.clear() + ast.operators[0].compute.clear() + elif ast.item_type not in ['vec','scal']: + raise ValueError(f"Tensor shape are not the same. Expect ast as 'vec' or 'scal' but found '{ast.item_type}'.") + elif ast.operators[0].item_type not in ['vec','scal']: + raise ValueError(f"Tensor shape are not the same. Expect operators[0] as 'vec' or 'scal' but found '{ast.operators[0].item_type}'.") + elif ast.operators[0].item_type != ast.item_type: + raise ValueError(f"Tensor shape are not the same. Expect 'vec' or 'scal' but found ast: '{ast.item_type}' and operators[0]: '{ast.operators[0].item_type}'.") + +def fuse_bvv(ast): + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + fuse_bvv(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + fuse_bvv(ast.operators[0]) + else: + return + + if type(ast.operators[1]) == BatchOp and ast.op_type == 'vec_mul_vec': + # fuse operators1 into bvv + if ast.item_type == 'scal' and ast.operators[1].item_type == 'vec': + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[1].compute.clear() + elif ast.operators[1].item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect operators[1] as 'vec' but found '{ast.operators[1].item_type}'.") + + if type(ast.operators[0]) == BatchOp and ast.op_type == 'vec_mul_vec': + # fuse operators1 into bvv + if ast.item_type == 'scal' and ast.operators[0].item_type == 'vec': + # check if type of operator0 is vector + if ast.operators[0].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[0].compute.clear() + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect operators[0] as 'vec' but found '{ast.operators[0].item_type}'.") + +def fuse_bsv(ast): + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + fuse_bsv(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + fuse_bsv(ast.operators[0]) + else: + return + + if type(ast) == BatchOp and ast.op_type == 'scal_mul_vec': + # fuse operators into bsv + if ast.item_type == 'vec' and ast.operators[1].item_type == 'vec' and ast.operators[0].item_type == 'scal': + # check if type of operators are scal and vector + if ast.operators[1].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[1].compute.clear() + if ast.operators[0].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[0].compute.clear() + elif ast.operators[0].item_type != 'scal': + raise ValueError(f"Tensor type is wrong. Expect operators[0] as 'scal' but found '{ast.operators[0].item_type}'.") + elif ast.operators[1].item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect operators[1] as 'vec' but found '{ast.operators[0].item_type}'.") + elif ast.item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect ast node as 'vec' but found '{ast.item_type}'.") + +def fuse_bvm(ast): + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + fuse_bvm(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + fuse_bvm(ast.operators[0]) + else: + return + + if type(ast) == BatchOp and ast.op_type == 'vec_mul_mat': + # fuse operators1 into bov + if ast.item_type == 'vec' and ast.operators[0].item_type == 'vec' and ast.operators[1].item_type == 'mat': + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + oloop = outer_loop + iloop = loop + iter_o = [] + iter_i = [] + if (isinstance(oloop, Loop) and isinstance(iloop, Loop)) and oloop.start == iloop.start and oloop.end.__name__ == iloop.end.__name__ and oloop.step == iloop.step: + iter_i.append(iloop.iterate) + iter_o.append(oloop.iterate) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[1].compute.clear() + + if ast.operators[0].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop = outer_loop + iloop = loop + iter_o = [] + iter_i = [] + if (isinstance(oloop, Loop) and isinstance(iloop, Loop)) and oloop.start == iloop.start and oloop.end.__name__ == iloop.end.__name__ and oloop.step == iloop.step: + iter_i.append(iloop.iterate) + iter_o.append(oloop.iterate) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[0].compute.clear() + elif ast.item_type != 'vec': + raise ValueError(f"Tensor shape are not the same. Expect ast node type 'mat' but found '{ast.item_type}'.") + elif ast.operators[1].item_type != 'mat': + raise ValueError(f"Tensor shape are not the same. Expect ast.operators[1] node type 'mat' but found '{ast.operators[1].item_type}'.") + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor shape are not the same. Expect ast.operators[0] node type 'vec' but found '{ast.operators[1].item_type}'.") + +def fuse_bov(ast): + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + fuse_bov(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + fuse_bov(ast.operators[0]) + else: + return + + if type(ast.operators[1]) == BatchOp and ast.op_type == 'vec_outer_vec': + # fuse operators1 into bov + if ast.item_type == 'mat' and ast.operators[1].item_type == 'vec': + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[1].compute.clear() + elif ast.item_type != 'mat': + raise ValueError(f"Tensor shape are not the same. Expect ast node type 'mat' but found '{ast.item_type}'.") + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor shape are not the same. Expect ast.operators[1] node type 'vec' but found '{ast.operators[1].item_type}'.") + + + if type(ast.operators[0]) == BatchOp and ast.op_type == 'vec_outer_vec': + # fuse operators0 into bov + if ast.item_type == 'mat' and ast.operators[0].item_type == 'vec': + # check if type of operator0 is vector + if ast.operators[0].compute and ast.compute: + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + iloop.body.clear() + ast.operators[0].compute.clear() + elif ast.item_type != 'mat': + raise ValueError(f"Tensor shape are not the same. Expect ast node type 'mat' but found '{ast.item_type}'.") + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor shape are not the same. Expect ast.operators[0] node type 'vec' but found '{ast.operators[0].item_type}'.") + + + +def fuse_operators(ast): + + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + # print('op1:', ast.operators[1].op_type) + fuse_operators(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + # print('op0:', ast.operators[0].op_type) + fuse_operators(ast.operators[0]) + else: + return + + # if ast.op_type in core.ast.op_mapping.keys() and type(ast.operators[1]) == Batch and type(ast.operators[0]) == Batch: + # outer_loop = ast.compute[0] + # a = Scalar(ast.eval.dtype) + # pre_arr = ast.eval + # ast.decl.pop(0) + # ast.decl.append(Decl(a)) + # ast.eval = a + # outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + + if type(ast.operators[1]) == BatchOp and ast.op_type in core.ast.arith_op.keys(): + # fuse operators1 into elementwise + if ast.item_type == 'vec' and ast.operators[1].item_type == ast.item_type: + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + for i in ast.operators[1].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + loop_merge(outer_loop, loop) + if ast.operators[1].op_type in core.ast.arith_op.keys() or ast.operators[1].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[1].eval.dtype, val=0) + pre_arr = ast.operators[1].eval + ast.operators[1].decl.pop(0) + ast.operators[1].decl.append(Decl(a)) + ast.operators[1].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[1].compute.clear() + ast.operators[1].valid = False + ast.decl.extend(ast.operators[1].decl) + + elif ast.item_type == 'scal' and ast.operators[1].item_type == ast.item_type: + # check if type of operator1 is scalar + if ast.operators[1].compute and ast.compute: + for i in ast.operators[1].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + loop_merge(outer_loop, loop) + if ast.operators[1].op_type in core.ast.arith_op.keys() or ast.operators[1].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[1].eval.dtype, val=0) + pre_arr = ast.operators[1].eval + ast.operators[1].decl.pop(0) + ast.operators[1].decl.append(Decl(a)) + ast.operators[1].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[1].compute.clear() + ast.operators[1].valid = False + ast.decl.extend(ast.operators[1].decl) + elif ast.item_type not in ['vec', 'scal']: + raise ValueError(f"Tensor type is wrong. Expect ast as 'vec' or 'scal' but found '{ast.item_type}'.") + elif ast.operators[1].item_type not in ['vec', 'scal']: + raise ValueError(f"Tensor type is wrong. Expect operators[1] as 'vec' or 'scal' but found '{ast.operators[1].item_type}'.") + elif ast.operators[1].item_type != ast.item_type: + raise ValueError(f"Tensor shape are not the same. Expect 'vec' or 'scal' but found ast: '{ast.item_type}' and operators[1]: '{ast.operators[1].item_type}'.") + + + if type(ast.operators[0]) == BatchOp and ast.op_type in core.ast.arith_op.keys(): + # fuse operators0 into elementwise + if ast.item_type == 'vec' and ast.operators[0].item_type == ast.item_type: + # check if type of operator0 is vector + if ast.operators[0].compute and ast.compute: + for i in ast.operators[0].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + loop_merge(outer_loop, loop) + if ast.operators[0].op_type in core.ast.arith_op.keys() or ast.operators[0].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[0].eval.dtype, val=0) + pre_arr = ast.operators[0].eval + ast.operators[0].decl.pop(0) + ast.operators[0].decl.append(Decl(a)) + ast.operators[0].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[0].compute.clear() + ast.operators[0].valid = False + ast.decl.extend(ast.operators[0].decl) + + if ast.item_type == 'scal' and ast.operators[0].item_type == ast.item_type: + # check if type of operator1 is scalar + if ast.operators[0].compute and ast.compute: + for i in ast.operators[0].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + loop_merge(outer_loop, loop) + if ast.operators[0].op_type in core.ast.arith_op.keys() or ast.operators[0].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[0].eval.dtype, val=0) + pre_arr = ast.operators[0].eval + ast.operators[0].decl.pop(0) + ast.operators[0].decl.append(Decl(a)) + ast.operators[0].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[0].compute.clear() + ast.operators[0].valid = False + ast.decl.extend(ast.operators[0].decl) + + elif ast.item_type not in ['vec','scal']: + raise ValueError(f"Tensor shape are not the same. Expect ast as 'vec' or 'scal' but found '{ast.item_type}'.") + elif ast.operators[0].item_type not in ['vec','scal']: + raise ValueError(f"Tensor shape are not the same. Expect operators[0] as 'vec' or 'scal' but found '{ast.operators[0].item_type}'.") + elif ast.operators[0].item_type != ast.item_type: + raise ValueError(f"Tensor shape are not the same. Expect 'vec' or 'scal' but found ast: '{ast.item_type}' and operators[0]: '{ast.operators[0].item_type}'.") + + if type(ast.operators[1]) == BatchOp and ast.op_type == 'vec_mul_vec': + # fuse operators1 into bvv + if ast.item_type == 'scal' and ast.operators[1].item_type == 'vec': + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + for i in ast.operators[1].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + loop_merge(outer_loop, loop) + if ast.operators[1].op_type in core.ast.arith_op.keys() or ast.operators[1].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[1].eval.dtype) + pre_arr = ast.operators[1].eval + ast.operators[1].decl.pop(0) + ast.operators[1].decl.append(Decl(a)) + ast.operators[1].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[1].compute.clear() + ast.operators[1].valid = False + ast.decl.extend(ast.operators[1].decl) + + elif ast.operators[1].item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect operators[1] as 'vec' but found '{ast.operators[1].item_type}'.") + + + if type(ast.operators[0]) == BatchOp and ast.op_type == 'vec_mul_vec': + # fuse operators0 into bvv + if ast.item_type == 'scal' and ast.operators[0].item_type == 'vec': + # check if type of operator0 is vector + if ast.operators[0].compute and ast.compute: + for i in ast.operators[0].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + loop_merge(outer_loop, loop) + if ast.operators[0].op_type in core.ast.arith_op.keys() or ast.operators[0].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[0].eval.dtype) + pre_arr = ast.operators[0].eval + ast.operators[0].decl.pop(0) + ast.operators[0].decl.append(Decl(a)) + ast.operators[0].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[0].compute.clear() + ast.operators[0].valid = False + ast.decl.extend(ast.operators[0].decl) + + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect operators[0] as 'vec' but found '{ast.operators[0].item_type}'.") + + if type(ast) == BatchOp and ast.op_type == 'scal_mul_vec': + # fuse operators into bsv + if ast.item_type == 'vec' and ast.operators[1].item_type == 'vec' and ast.operators[0].item_type == 'scal': + # check if type of operators are scal and vector + if ast.operators[1].compute and ast.compute: + for i in ast.operators[1].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + + # oloop, iloop, iter_o, iter_i = get_same_loop(outer_loop, loop) + # for i in iloop.body: + # change_index(i, iter_o, iter_i) + # for i in range(len(iloop.body)): + # oloop.body.insert(i, iloop.body[i]) + # iloop.body.clear() + loop_merge(outer_loop, loop) + if ast.operators[1].op_type in core.ast.arith_op.keys() or ast.operators[1].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[1].eval.dtype) + pre_arr = ast.operators[1].eval + ast.operators[1].decl.pop(0) + ast.operators[1].decl.append(Decl(a)) + ast.operators[1].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[1].compute.clear() + ast.operators[1].eval = None + ast.decl.extend(ast.operators[1].decl) + + if ast.operators[0].compute and ast.compute: + for i in ast.operators[0].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + oloop = outer_loop + iloop = loop + iter_o = [] + iter_i = [] + if (isinstance(oloop, Loop) and isinstance(iloop, Loop)) and oloop.start == iloop.start and oloop.end.__name__ == iloop.end.__name__ and oloop.step == iloop.step: + iter_i.append(iloop.iterate) + iter_o.append(oloop.iterate) + if isinstance(iloop, Loop): + for i in iloop.body: + change_index(i, iter_o, iter_i) + + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + if ast.operators[0].op_type in core.ast.arith_op.keys() or ast.operators[0].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[0].eval.dtype) + pre_arr = ast.operators[0].eval + ast.operators[0].decl.pop(0) + ast.operators[0].decl.append(Decl(a)) + ast.operators[0].eval = a + oloop = swap_arr_to_reg(oloop, pre_arr, a) + # iloop.body.clear() + # ast.operators[0].compute.clear() + ast.operators[0].valid = False + ast.decl.extend(ast.operators[0].decl) + + elif ast.operators[0].item_type != 'scal': + raise ValueError(f"Tensor type is wrong. Expect operators[0] as 'scal' but found '{ast.operators[0].item_type}'.") + elif ast.operators[1].item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect operators[1] as 'vec' but found '{ast.operators[0].item_type}'.") + elif ast.item_type != 'vec': + raise ValueError(f"Tensor type is wrong. Expect ast node as 'vec' but found '{ast.item_type}'.") + + + # if type(ast) == BatchOp and ast.op_type == 'vec_mul_mat': + # # fuse operators1 into bvm + # if ast.item_type == 'vec' and ast.operators[0].item_type == 'vec' and ast.operators[1].item_type == 'mat': + # # check if type of operator1 is vector + # if ast.operators[1].compute and ast.compute: + # for i in ast.operators[1].compute: + # change_ref(i, ast) + # outer_loop = ast.compute[0] + # loop = ast.operators[1].compute[0] + # loop_merge(outer_loop, loop) + # # ast.operators[1].compute.clear() + # ast.operators[1].valid = False + # ast.decl.extend(ast.operators[1].decl) + + # if ast.operators[0].compute and ast.compute: + # print(ast.compute) + # for i in ast.compute: + # print('before:::', codegen.gpu.to_string(i)) + # ast.compute.extend(ast.operators[0].compute) + # for i in ast.compute: + # print('after:::', codegen.gpu.to_string(i)) + # ast.operators[0].valid = False + # for i in ast.operators[0].compute: + # change_ref(i, ast) + # outer_loop = ast.compute[0] + # loop = ast.operators[0].compute[0] + # oloop = outer_loop + # iloop = loop + # iter_o = [] + # iter_i = [] + + # while (isinstance(oloop, Loop) and isinstance(iloop, Loop)) and oloop.start == iloop.start and oloop.end.__name__ == iloop.end.__name__ and oloop.step == iloop.step: + # iter_i.append(iloop.iterate) + # iter_o.append(oloop.iterate) + # if isinstance(oloop.body[-1], Loop) and isinstance(iloop, Loop): + # oloop = oloop.body[-1].body[-1] + # iloop = iloop.body[-1] + # else: + # break + # for i in iloop.body: + # change_index(i, iter_o, iter_i) + # for i in range(len(iloop.body)): + # oloop.body.insert(i, iloop.body[i]) + # # ast.operators[0].compute.clear() + # ast.operators[0].valid = False + # ast.decl.extend(ast.operators[0].decl) + + # elif ast.item_type != 'vec': + # raise ValueError(f"Tensor shape are not the same. Expect ast node type 'mat' but found '{ast.item_type}'.") + # elif ast.operators[1].item_type != 'mat': + # raise ValueError(f"Tensor shape are not the same. Expect ast.operators[1] node type 'mat' but found '{ast.operators[1].item_type}'.") + # elif ast.operators[0].item_type != 'vec': + # raise ValueError(f"Tensor shape are not the same. Expect ast.operators[0] node type 'vec' but found '{ast.operators[1].item_type}'.") + + + if type(ast.operators[1]) == BatchOp and ast.op_type == 'vec_outer_vec': + # fuse operators1 into bov + if ast.item_type == 'mat' and ast.operators[1].item_type == 'vec': + # check if type of operator1 is vector + if ast.operators[1].compute and ast.compute: + for i in ast.operators[1].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[1].compute[0] + oloop = outer_loop + iloop = loop + iter_i = [] + iter_o = [] + while (isinstance(oloop, Loop) and isinstance(iloop, Loop)) and oloop.start == iloop.start and oloop.end.__name__ == iloop.end.__name__ and oloop.step == iloop.step: + iter_i.append(iloop.iterate) + iter_o.append(oloop.iterate) + if isinstance(oloop.body[-1], Loop) and isinstance(iloop, Loop): + oloop = oloop.body[-1].body[-1] + iloop = iloop.body[-1] + else: + break + for i in iloop.body: + change_index(i, iter_o, iter_i) + for i in range(len(iloop.body)): + oloop.body.insert(i, iloop.body[i]) + if ast.operators[1].op_type in core.ast.arith_op.keys() or ast.operators[1].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[1].eval.dtype) + pre_arr = ast.operators[1].eval + ast.operators[1].decl.pop(0) + ast.operators[1].decl.append(Decl(a)) + ast.operators[1].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[1].compute.clear() + ast.operators[1].valid = False + ast.decl.extend(ast.operators[1].decl) + + elif ast.item_type != 'mat': + raise ValueError(f"Tensor shape are not the same. Expect ast node type 'mat' but found '{ast.item_type}'.") + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor shape are not the same. Expect ast.operators[1] node type 'vec' but found '{ast.operators[1].item_type}'.") + + + if type(ast.operators[0]) == BatchOp and ast.op_type == 'vec_outer_vec': + # fuse operators0 into bov + if ast.item_type == 'mat' and ast.operators[0].item_type == 'vec': + # check if type of operator0 is vector + if ast.operators[0].compute and ast.compute: + for i in ast.operators[0].compute: + change_ref(i, ast) + outer_loop = ast.compute[0] + loop = ast.operators[0].compute[0] + + loop_merge(outer_loop, loop) + if ast.operators[0].op_type in core.ast.arith_op.keys() or ast.operators[0].op_type in ["scal_mul_vec", "vec_mul_vec"]: + a = Scalar(ast.operators[0].eval.dtype) + pre_arr = ast.operators[0].eval + ast.operators[0].decl.pop(0) + ast.operators[0].decl.append(Decl(a)) + ast.operators[0].eval = a + outer_loop = swap_arr_to_reg(outer_loop, pre_arr, a) + # ast.operators[0].compute.clear() + ast.operators[0].valid = False + ast.decl.extend(ast.operators[0].decl) + + elif ast.item_type != 'mat': + raise ValueError(f"Tensor shape are not the same. Expect ast node type 'mat' but found '{ast.item_type}'.") + elif ast.operators[0].item_type != 'vec': + raise ValueError(f"Tensor shape are not the same. Expect ast.operators[0] node type 'vec' but found '{ast.operators[0].item_type}'.") \ No newline at end of file diff --git a/batch/opt/ir.py b/batch/opt/ir.py new file mode 100644 index 0000000..900c3b1 --- /dev/null +++ b/batch/opt/ir.py @@ -0,0 +1,98 @@ +from core.ir import * + +class BlockIdy(IR): + def __init__(self): + super().__init__() + +class BlockIdx(IR): + def __init__(self): + super().__init__() + +class BlockDimy(IR): + def __init__(self): + super().__init__() + +class BlockDimx(IR): + def __init__(self): + super().__init__() + +class ThreadIdy(IR): + def __init__(self): + super().__init__() + +class ThreadIdx(IR): + def __init__(self): + super().__init__() + +class SyncThreads(IR): + def __init__(self): + super().__init__() + +class SyncWarps(IR): + def __init__(self): + super().__init__() + +class ShuffleDown(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class ShuffleUp(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class ShuffleXor(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class SaveAtThread(IR): + def __init__(self, src, dst, threadid): + super().__init__() + self.src = src + self.dst = dst + self.threadid = threadid + +class BroadCast(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class Shared(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class Uniq(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class Buffer(IR): + def __init__(self, dobject): + super().__init__() + self.dobject = dobject + +class IF(IR): + def __init__(self, left, condition: Expr, true_var, false_var): + super().__init__() + self.left = left + self.condition = condition + self.true_var = true_var + self.false_var = false_var + +class Pointer(DOject): + def __init__(self, dtype, size): + super().__init__(dtype, size) + self.__name__ = f'ptr{self.dobject_id}' + self.dtype = dtype + self.size = size + + def name(self): + return self.__name__ + +class Access_ptr(): + def __init__(self, dobject:Pointer, idx): + self.dobject = dobject + self.idx = idx \ No newline at end of file diff --git a/batch/opt/node_wise/__pycache__/parallelism.cpython-310.pyc b/batch/opt/node_wise/__pycache__/parallelism.cpython-310.pyc new file mode 100644 index 0000000..3f6d950 Binary files /dev/null and b/batch/opt/node_wise/__pycache__/parallelism.cpython-310.pyc differ diff --git a/batch/opt/node_wise/__pycache__/smem.cpython-310.pyc b/batch/opt/node_wise/__pycache__/smem.cpython-310.pyc new file mode 100644 index 0000000..26ad3f9 Binary files /dev/null and b/batch/opt/node_wise/__pycache__/smem.cpython-310.pyc differ diff --git a/batch/opt/node_wise/__pycache__/tiling.cpython-310.pyc b/batch/opt/node_wise/__pycache__/tiling.cpython-310.pyc new file mode 100644 index 0000000..f6193db Binary files /dev/null and b/batch/opt/node_wise/__pycache__/tiling.cpython-310.pyc differ diff --git a/batch/opt/node_wise/parallelism.py b/batch/opt/node_wise/parallelism.py new file mode 100644 index 0000000..3df97a4 --- /dev/null +++ b/batch/opt/node_wise/parallelism.py @@ -0,0 +1,207 @@ +from core.ir import * +from batch.ast import * +from batch.ast2ir import * +import codegen +from batch.opt.ir import * +# for better optimization on GPU + +def swap_arr_to_reg(ir, pre, cur): + if isinstance(ir, Indexing): + temp = ir + while isinstance(temp, Indexing): + temp = temp.dobject + if temp == pre: + return cur + else: + return ir + elif isinstance(ir, Expr): + ir.left = swap_arr_to_reg(ir.left, pre, cur) + ir.right = swap_arr_to_reg(ir.right, pre, cur) + elif isinstance(ir, Assignment): + ir.lhs = swap_arr_to_reg(ir.lhs, pre, cur) + ir.rhs = swap_arr_to_reg(ir.rhs, pre, cur) + elif isinstance(ir, Loop): + for i in range(len(ir.body)): + ir.body[i] = swap_arr_to_reg(ir.body[i], pre, cur) + return ir + +def find_arr_ind(ir, pre): + if isinstance(ir, Indexing): + temp = ir + while isinstance(temp, Indexing): + temp = temp.dobject + if temp == pre: + return ir + else: + return None + elif isinstance(ir, Expr): + return Expr(find_arr_ind(ir.left, pre), find_arr_ind(ir.right, pre), ir.op) + elif isinstance(ir, Assignment): + return find_arr_ind(ir.lhs, pre) + elif isinstance(ir, Loop): + for i in ir.body: + t = find_arr_ind(i, pre) + if t: + return t + +def if_contain(item, ir): + if isinstance(item, Loop): + flag = False + for i in item.body: + if isinstance(i, Loop): + flag = if_contain(i, ir) + elif i == ir: + flag = True + if flag: + return True + + return False + +def add_thready(ir, arr): + if isinstance(ir, Indexing): + temp = ir + idx_list = [] + while isinstance(temp, Indexing): + idx_list.append(temp.idx) + temp = temp.dobject + if temp == arr: + idx_list = idx_list[::-1] + temp = Indexing(temp, Literal(-1, 'int')) + temp.idx = ThreadIdy() + for i in idx_list: + if isinstance(i, (Scalar, Literal, Indexing)): + temp = Indexing(temp, i) + else: + temp = Indexing(temp, Literal(-1, 'int')) + temp.idx = i + ir = temp + # print('yes', codegen.gpu.to_string(temp)) + elif isinstance(ir, Assignment): + ir.lhs = add_thready(ir.lhs, arr) + ir.rhs = add_thready(ir.rhs, arr) + elif isinstance(ir, Expr): + ir.left = add_thready(ir.left, arr) + ir.right = add_thready(ir.right, arr) + elif isinstance(ir, Loop): + for i in range(len(ir.body)): + ir.body[i] = add_thready(ir.body[i], arr) + return ir + +def add_reduction(ast): + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + add_reduction(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + add_reduction(ast.operators[0]) + else: + return + + # todo: add traverse action to add reduction + if ast.op_type == 'vec_mul_vec': + # this inner_prod node is fused with upper layer + eval = ast.eval + # print(codegen.cpu.to_string(eval), codegen.cpu.to_string(ast.operators[0].eval), codegen.cpu.to_string(ast.operators[1].eval)) + # print(codegen.gpu.to_string(eval), ast.eval, ast.operators[0].eval, ast.operators[1].eval) + # iff eval is scalar, we need to add shfl_sync + if isinstance(ast.eval, Ndarray): + new_compute = [] + for idx, item in enumerate(ast.compute): + new_compute.append(item) + if isinstance(item, Loop): + a = Scalar(ast.eval.dtype) + pre_arr = ast.eval + ast.decl.append(Decl(a)) + t = find_arr_ind(item, pre_arr) + + swap_arr_to_reg(item, pre_arr, a) + new_compute.append(ShuffleDown(a)) + new_compute.append(SaveAtThread(a, t, 0)) + ast.compute = new_compute + elif not ((isinstance(ast.operators[0].eval, Ndarray) or isinstance(ast.operators[1].eval, Ndarray))): + for i in ast.compute: + # search all compute stmts + if isinstance(i, Loop): + for j in i.body: + # search loop body + if isinstance(j, Loop): + # find the stmt of ast node + main_loop = j.astnode.compute + for idx, item in enumerate(main_loop): + + if isinstance(item, Loop) and item == j: + # fused operators + main_loop.insert(idx+1, SyncThreads()) + main_loop.insert(idx+1, BroadCast(eval)) + main_loop.insert(idx+1, ShuffleDown(eval)) + elif isinstance(item, Loop): + # before fuse operators + if if_contain(item, j.body[0]): + main_loop.insert(idx+1, SyncThreads()) + main_loop.insert(idx+1, BroadCast(eval)) + main_loop.insert(idx+1, ShuffleDown(eval)) + if ast.op_type == 'vec_mul_mat': + # print(ast.eval, codegen.gpu.to_string(ast.eval), ast.eval.size) + ast.eval.size.insert(0, Scalar('int', 'C')) + # print(ast.compute[0].astnode.compute) + for i in ast.compute[0].astnode.compute: + # print(codegen.gpu.to_string(i)) + t = add_thready(i, ast.eval) + # print(t, codegen.gpu.to_string(t)) + + + + +def cuda_spec(ast): + if ast.compute and ast.valid: + compute_list = [] + for body in ast.compute: + body_list = [] + if isinstance(body, Loop): + # print(body.iterate, body.iterate.dobject) + ast.decl.append(Decl(body.iterate)) + assign = Assignment(body.iterate, Expr(ThreadIdy(), Expr(BlockDimy(), BlockIdx(), '*'), '+')) + body_list.append(assign) + for item in body.body: + if isinstance(item, Loop) and item.step == 1: + item.start = ThreadIdx() + item.step = BlockDimx() + elif isinstance(item, Loop): + for j in item.body: + if isinstance(j, Loop) and j.step == 1: + j.start = ThreadIdx() + j.step = BlockDimx() + elif isinstance(j, Loop): + for k in j.body: + if isinstance(k, Loop) and k.step == 1: + k.start = ThreadIdx() + k.step = BlockDimx() + + body_list.append(item) + compute_list.extend(body_list) + ast.compute = compute_list + +def add_cuda_spec(ast): + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + add_cuda_spec(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + add_cuda_spec(ast.operators[0]) + else: + return + + cuda_spec(ast) + + +def parallel(ast): + + # print(ast, ast.op_type, ast.compute) + # for i in ast.compute: + # print(i, i.ast_ref, i.ast_ref.compute) + # print(ast.compute) + # for i in ast.compute: + # print(codegen.gpu.to_string(i)) + + add_cuda_spec(ast) + add_reduction(ast) + + \ No newline at end of file diff --git a/batch/opt/node_wise/smem.py b/batch/opt/node_wise/smem.py new file mode 100644 index 0000000..dd2a659 --- /dev/null +++ b/batch/opt/node_wise/smem.py @@ -0,0 +1,411 @@ +from core.ir import * +from batch.ast import * +from batch.ast2ir import * +import codegen +from batch.opt.ir import * + +def find_reuse(ir, smem_list): + if isinstance(ir, Indexing): + # print(find_reuse(ir.dobject, smem_list), find_reuse(ir.idx, smem_list)) + temp = ir.dobject + shape_size = 0 + idx = [ir.idx] + while isinstance(temp, Indexing): + shape_size += 1 + + if isinstance(temp.idx, Indexing): + if temp.idx.dobject.name() == 'r': + data = get_ori_var(ir) + if data in smem_list.keys(): + smem_list[data].append([ir, shape_size, idx]) + else: + smem_list[data] = [[ir, shape_size, idx]] + idx.append(temp.idx) + temp = temp.dobject + + elif isinstance(ir, Assignment): + find_reuse(ir.lhs, smem_list) + find_reuse(ir.rhs, smem_list) + elif isinstance(ir, Expr): + find_reuse(ir.left, smem_list) + find_reuse(ir.right, smem_list) + elif isinstance(ir, Loop): + for i in ir.body: + find_reuse(i, smem_list) + +def get_ori_var(ir): + if isinstance(ir, Indexing): + return get_ori_var(ir.dobject) + elif isinstance(ir, Ndarray): + return ir + + +def if_exist(ir, arr): + if isinstance(ir, Indexing): + if ir == arr: + return True + else: + return False + elif isinstance(ir, Assignment): + return if_exist(ir.lhs, arr) or if_exist(ir.rhs, arr) + elif isinstance(ir, Expr): + return if_exist(ir.left, arr) or if_exist(ir.right, arr) + elif isinstance(ir, Loop): + for i in ir.body: + t = if_exist(i, arr) + if t: + return t + return False + +def change_access(ir, ori, cur): + if isinstance(ir, Indexing): + if ir == ori: + return cur + elif isinstance(ir, Assignment): + ir.lhs = change_access(ir.lhs, ori, cur) + ir.rhs = change_access(ir.rhs, ori, cur) + elif isinstance(ir, Expr): + ir.left = change_access(ir.left, ori, cur) + ir.right = change_access(ir.right, ori, cur) + elif isinstance(ir, Loop): + for i in ir.body: + i = change_access(i, ori, cur) + return ir + +def data_loading(compute_ir, smem_arr, ori_arr): + decl = [] + if ori_arr.keys(): + for i in ori_arr[list(ori_arr.keys())[0]]: + indirect_arr = get_ori_var(i[2][-1]) + decl.append(Decl(Buffer(indirect_arr))) + decl.append(Decl(Uniq(indirect_arr))) + break + t = Scalar('int') + decl.append(Decl(t)) + bufidx = Indexing(Ndarray('int', [Scalar('int', 'batch_size/16'), Scalar('int', 'BlockDim.y')]), Literal(-1, 'int')) + bufidx.dobject = Buffer(indirect_arr) + bufidx.idx = ThreadIdy() + smtm1 = Assignment(t, bufidx) + + main_inequa = Expr(t, Scalar('int', 'C'), '<') + idx_var = Scalar('int') + decl.append(Decl(idx_var)) + idx_false = Indexing(indirect_arr, Literal(-1, 'int')) + idx_false.idx = Expr(t, Scalar('int', 'C'), '-') + idx = IF(idx_var, main_inequa, t, idx_false) + ptr_access = Scalar('int') + decl.append(Decl(ptr_access)) + set_ptr = IF(ptr_access, main_inequa, Scalar('int', 'D'), Scalar('int', 'dim')) + + compute_ir.insert(0, set_ptr) + compute_ir.insert(0, idx) + compute_ir.insert(0, smtm1) + for stmt in compute_ir: + if isinstance(stmt, Loop): + # data loading here + + for i in range(len(smem_arr)): + cur_arr = ori_arr[list(ori_arr.keys())[i]] + if len(smem_arr[i].size) == 2: + temp_stmt = stmt + + for iloop in stmt.body: + if isinstance(iloop, Loop) and isinstance(iloop.body[0], Loop): + if iloop.start == 0 and iloop.end.name() == 'dim' and isinstance(iloop.body[0].start, ThreadIdx) and iloop.body[0].end.name() == 'D': + temp_stmt = iloop + break + for shared_item in cur_arr: + flag = if_exist(temp_stmt, shared_item[0]) + if flag: + store_loop = Loop(Expr(Expr(ThreadIdy(), BlockDimx(), '*'), ThreadIdx(), '+'), Expr(Scalar('float', 'D'), Uniq(shared_item[2][-1].dobject), '*'), Expr(BlockDimx(), BlockDimy(), '*'), []) + + left = smem_arr[i] + left = Indexing(left, Literal(-1, 'int')) + left.idx = Expr(store_loop.iterate, Scalar('int', 'D'), '/') + left = Indexing(left, Literal(-1, 'int')) + left.idx = Expr(store_loop.iterate, Scalar('int', 'D'), '%') + + main_arr = get_ori_var(shared_item[0]) + indirect_arr = get_ori_var(shared_item[2][-1]) + uniq = Ndarray('int', [Scalar('int', 'batch_size/16'), Scalar('int', 'dim')], f'{indirect_arr.__name__}_Uniq') + uniq = Indexing(uniq, Literal(-1, 'int')) + uniq.idx = BlockIdx() + uniq = Indexing(uniq, Literal(-1, 'int')) + uniq.idx = Expr(store_loop.iterate, Scalar('int', 'D'), '/') + main_arr = Indexing(main_arr, Literal(-1, 'int')) + main_arr.idx = uniq + main_arr = Indexing(main_arr, Literal(-1, 'int')) + # main_arr.idx = shared_item[2][0] + main_arr.idx = Expr(temp_stmt.iterate, Expr(store_loop.iterate, Scalar('float', 'D'), '%'), '+') + assign = Assignment(left, main_arr, '') + store_loop.body.append(assign) + + + ptr = Pointer('float *', [Scalar('int', 'C'), ptr_access]) + pmat = IF(ptr, main_inequa, smem_arr[i], get_ori_var(shared_item[0])) + decl.append(Decl(ptr)) + + offset_1 = Scalar('int') + decl.append(Decl(offset_1)) + ofs1 = IF(offset_1, main_inequa, 0, stmt.iterate) + + access_var = Access_ptr(ptr, []) + access_var.idx.append(idx_var) + for kk in range(len(shared_item[2]) - 1): + shared_item[2][kk].left = offset_1 + access_var.idx.append(shared_item[2][kk]) + + stmt = change_access(stmt, shared_item[0], access_var) + + temp_stmt.body.insert(0, ofs1) + temp_stmt.body.insert(0, pmat) + + stmt.body.insert(0, SyncThreads()) + stmt.body.insert(0, store_loop) + elif len(smem_arr[i].size) == 3: + temp_stmt = stmt + # print(codegen.gpu.to_string(temp_stmt)) + for iloop in stmt.body: + if isinstance(iloop, Loop) and isinstance(iloop.body[0], Loop): + if iloop.start == 0 and iloop.end.name() == 'dim' and iloop.body[0].end.name() == 'D': + temp_stmt = iloop + break + # print(codegen.gpu.to_string(temp_stmt)) + for shared_item in cur_arr: + flag = if_exist(temp_stmt, shared_item[0]) + # print(codegen.gpu.to_string(temp_stmt), codegen.gpu.to_string(shared_item[0])) + if flag: + oloop = Loop(Literal(0, 'int'), Literal(2, 'int'), Literal(1, 'int'), []) + store_loop1 = Loop(ThreadIdy(), stmt.step, BlockDimy(), []) + store_loop2 = Loop(ThreadIdx(), stmt.step, BlockDimx(), []) + oloop.body.append(store_loop1) + store_loop1.body.append(store_loop2) + + left = smem_arr[i] + left = Indexing(left, oloop.iterate) + # left.idx = ThreadIdy() + left = Indexing(left, store_loop1.iterate) + left = Indexing(left, store_loop2.iterate) + + main_arr = get_ori_var(shared_item[0]) + indirect_arr = get_ori_var(shared_item[2][-1]) + uniq = Ndarray('int', [Scalar('int', 'batch_size/16'), Scalar('int', 'dim')], f'{indirect_arr.__name__}_Uniq') + uniq = Indexing(uniq, Literal(-1, 'int')) + uniq.idx = BlockIdx() + uniq = Indexing(uniq, oloop.iterate) + # uniq.idx = ThreadIdy() + main_arr = Indexing(main_arr, Literal(-1, 'int')) + main_arr.idx = uniq + main_arr = Indexing(main_arr, Literal(-1, 'int')) + main_arr.idx = Expr(shared_item[2][1].left, store_loop1.iterate, '+') + main_arr = Indexing(main_arr, Literal(-1, 'int')) + main_arr.idx = Expr(shared_item[2][0].left, store_loop2.iterate, '+') + assign = Assignment(left, main_arr, '') + store_loop2.body.append(assign) + + + ptr = Pointer('float *', [Scalar('int', 'C'), ptr_access, ptr_access]) + decl.append(Decl(ptr)) + pmat = IF(ptr, main_inequa, smem_arr[i], get_ori_var(shared_item[0])) + + offset_1 = Scalar('int') + decl.append(Decl(offset_1)) + ofs1 = IF(offset_1, main_inequa, 0, stmt.iterate) + offset_2 = Scalar('int') + decl.append(Decl(offset_2)) + ofs2 = IF(offset_2, main_inequa, 0, temp_stmt.iterate) + + access_var = Access_ptr(ptr, []) + access_var.idx.append(idx_var) + shared_item[2][1].left = offset_2 + access_var.idx.append(shared_item[2][1]) + shared_item[2][0].left = offset_1 + access_var.idx.append(shared_item[2][0]) + + stmt = change_access(stmt, shared_item[0],access_var) + + temp_stmt.body.insert(0, ofs2) + temp_stmt.body.insert(0, ofs1) + temp_stmt.body.insert(0, pmat) + + temp_stmt.body.insert(0, SyncThreads()) + temp_stmt.body.insert(0, oloop) + return decl + +def swap_arr_to_reg(ir, pre, cur): + if isinstance(ir, Indexing): + # todo: add index here + if ir == pre[0]: + temp = Indexing(cur, Literal(-1, 'int')) + # temp.idx = Scalar('int', 'idx') + for i in range(pre[1]): + temp = Indexing(temp, Literal(-1, 'int')) + if isinstance(pre[2][-2-i], Expr): + # pre[2][-2-i].left = f'ofs_{i+1}' + temp.idx = pre[2][-2-i].right + else: + temp.idx = pre[2][-2-i] + return temp + else: + return ir + elif isinstance(ir, Expr): + ir.left = swap_arr_to_reg(ir.left, pre, cur) + ir.right = swap_arr_to_reg(ir.right, pre, cur) + elif isinstance(ir, Assignment): + ir.lhs = swap_arr_to_reg(ir.lhs, pre, cur) + ir.rhs = swap_arr_to_reg(ir.rhs, pre, cur) + elif isinstance(ir, Loop): + for i in range(len(ir.body)): + ir.body[i] = swap_arr_to_reg(ir.body[i], pre, cur) + return ir + +def if_in_ir(ir, arr, smem_dict): + if isinstance(ir, Indexing): + temp = ir + while isinstance(temp, Indexing): + if temp.dobject == arr: + if temp.dobject in smem_dict.keys(): + smem_dict[temp.dobject].append(ir) + else: + smem_dict[temp.dobject] = [ir] + return True + temp = temp.dobject + elif isinstance(ir, Assignment): + return if_in_ir(ir.lhs, arr, smem_dict) or if_in_ir(ir.rhs, arr, smem_dict) + elif isinstance(ir, Expr): + return if_in_ir(ir.left, arr, smem_dict) or if_in_ir(ir.right, arr, smem_dict) + elif isinstance(ir, Loop): + bool_list = [] + for i in ir.body: + t = if_in_ir(i, arr, smem_dict) + bool_list.append(t) + for i in bool_list: + if i: + return True + return False + + + +def add_smem(ast): + + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + add_smem(ast.operators[1]) + if type(ast.operators[0]) == BatchOp: + add_smem(ast.operators[0]) + + if type(ast) == BatchOp and ast.valid: + + var_list = {} + smem_list = [] + for i in ast.compute: + find_reuse(i, var_list) + + # for i in var_list.keys(): + # t = get_ori_var(var_list[i][0][0]) + # print('var_list::', codegen.gpu.to_string(var_list[i][0][0]), var_list[i][0][1], codegen.gpu.to_string(t), t, i) + # for j in var_list[i][0][2]: + # print(codegen.gpu.to_string(j)) + for i in var_list: + size = [Scalar('int', 'C')] + for s in range(var_list[i][0][1]): + size.append(Scalar('int', 'D')) + smema = Ndarray('float', size, f'smem_{i.dobject_id}') + smem_list.append(smema) + ast.decl.append(Decl(Shared(smema))) + + decl = data_loading(ast.compute, smem_list, var_list) + ast.decl.extend(decl) + + if type(ast) == BatchOp and not ast.valid and isinstance(ast.eval, Ndarray): + # change shared memory access + node = ast.compute[0].astnode + ir_dict = {} + smem_list = [] + for i in node.compute: + if if_in_ir(i, ast.eval, ir_dict): + smema = Ndarray('float', [Scalar('int', 'C'), Scalar('int', 'D')], f'smem_{ast.eval.dobject_id}') + smem_list.append(smema) + node.decl.append(Decl(Shared(smema))) + for i in node.decl: + if i.dobject == ast.eval: + node.decl.remove(i) + + for i in range(len(list(ir_dict.keys()))): + key = list(ir_dict.keys())[i] + smem_idx_list = [] + # print(key, codegen.gpu.to_string(key), smem_list[i], codegen.gpu.to_string(smem_list[i])) + for j in ir_dict[key]: + # print('smemlist:::', codegen.gpu.to_string(j)) + idx_list = [] + temp = j + while isinstance(temp, Indexing): + idx_list.append(temp.idx) + temp = temp.dobject + idx_list = idx_list[::-1] + newsmem = smem_list[i] + + for kk in range(len(idx_list)): + newsmem = Indexing(newsmem, Literal(-1, 'int')) + if isinstance(idx_list[kk], (Scalar, ThreadIdy)): + newsmem.idx = idx_list[kk] + elif isinstance(idx_list[kk], Expr): + newsmem.idx = idx_list[kk].right + smem_idx_list.append(newsmem) + for j in node.compute: + for idx in range(len(smem_idx_list)): + j = change_access(j, ir_dict[key][idx], smem_idx_list[idx]) + + if type(ast) == BatchOp and not ast.valid and ast.op_type=='vec_mul_mat': + # print(ast.operators[0], ast.operators[1]) + # print(ast.operators[0].eval, codegen.gpu.to_string(ast.operators[0].eval)) + # print(ast.operators[0].op_type, ast.operators[1].op_type) + # print(ast.compute[0].astnode) + main_loop = ast.compute[0].astnode.compute + smem_dict = {} + for i in main_loop: + # print(codegen.gpu.to_string(i)) + if if_in_ir(i, ast.operators[0].eval, smem_dict): + temp_stmt = i + # find target loop + for iloop in i.body: + if isinstance(iloop, Loop) and isinstance(iloop.body[0], Loop): + if iloop.start == 0 and iloop.end.name() == 'dim' and iloop.body[0].end.name() == 'D': + temp_stmt = iloop + break + # print('eval in this stmt:', temp_stmt, codegen.gpu.to_string(temp_stmt)) + newsmem = ast.operators[0].eval + arr = list(smem_dict.keys())[0] + idx = smem_dict[arr][0] + iter_list = [] + tidx = idx + while isinstance(tidx,Indexing): + iter_list.append(tidx.idx) + tidx = tidx.dobject + # print(arr, codegen.gpu.to_string(arr), idx, codegen.gpu.to_string(idx)) + store_loop = Loop(ThreadIdx(), Scalar('int', 'D'), BlockDimx(), []) + + temp_stmt.body.insert(0, SyncThreads()) + temp_stmt.body.insert(0, store_loop) + smema = Ndarray('float', [Scalar('int', 'C'), Scalar('int', 'D')], f'smem_{ast.operators[0].eval.dobject_id}') + node.decl.append(Decl(Shared(smema))) + + new_access = Indexing(smema, Literal(-1, 'int')) + new_access.idx = ThreadIdy() + new_access = Indexing(new_access, iter_list[0].right) + change_access(temp_stmt, idx, new_access) + + smema = Indexing(smema, Literal(-1, 'int')) + smema.idx = ThreadIdy() + smema = Indexing(smema, store_loop.iterate) + + global_arr = arr + for i in iter_list[::-1]: + global_arr = Indexing(global_arr, Literal(-1, 'int')) + if isinstance(i, Expr): + global_arr.idx = Expr(i.left, i.right, i.op) + else: + global_arr.idx = i + global_arr.idx.right = store_loop.iterate + assign = Assignment(smema, global_arr) + store_loop.body.append(assign) diff --git a/batch/opt/node_wise/tiling.py b/batch/opt/node_wise/tiling.py new file mode 100644 index 0000000..841290d --- /dev/null +++ b/batch/opt/node_wise/tiling.py @@ -0,0 +1,198 @@ +from core.ir import * +from batch.ast import * +from batch.ast2ir import * +import codegen +from batch.opt.ir import * + +def swap_arr_to_reg(ir, pre, cur): + # print(codegen.gpu.to_string(ir), codegen.gpu.to_string(pre), codegen.gpu.to_string(cur)) + if isinstance(ir, Indexing): + temp = ir + while isinstance(temp, Indexing): + # print(codegen.gpu.to_string(temp), codegen.gpu.to_string(temp.idx)) + if temp.idx == pre: + temp.idx = Expr(pre, cur, '+') + temp = temp.dobject + elif isinstance(ir, Expr): + ir.left = swap_arr_to_reg(ir.left, pre, cur) + ir.right = swap_arr_to_reg(ir.right, pre, cur) + elif isinstance(ir, Assignment): + ir.lhs = swap_arr_to_reg(ir.lhs, pre, cur) + ir.rhs = swap_arr_to_reg(ir.rhs, pre, cur) + elif isinstance(ir, Loop): + for i in range(len(ir.body)): + ir.body[i] = swap_arr_to_reg(ir.body[i], pre, cur) + return ir + +def swap_reg_to_arr(ir, pre, cur, iloop): + if isinstance(ir, Scalar): + if ir == pre and iloop.end.name()=='D': + ir = Indexing(cur, iloop.iterate) + elif isinstance(ir, Expr): + ir.left = swap_reg_to_arr(ir.left, pre, cur, iloop) + ir.right = swap_reg_to_arr(ir.right, pre, cur, iloop) + elif isinstance(ir, Assignment): + ir.lhs = swap_reg_to_arr(ir.lhs, pre, cur, iloop) + ir.rhs = swap_reg_to_arr(ir.rhs, pre, cur, iloop) + elif isinstance(ir, Loop): + + for i in range(len(ir.body)): + if ir.end.name() == 'D' and isinstance(ir.body[i], Loop) and ir.body[i].end.name() == 'D': + for j in range(len(ir.body[i].body)): + ir.body[i].body[j] = swap_reg_to_arr(ir.body[i].body[j], pre, cur, ir) + else: + ir.body[i] = swap_reg_to_arr(ir.body[i], pre, cur, ir) + return ir + +def tile_wD(ir): + if isinstance(ir, Loop): + if ir.end.name() == 'dim': + scalar_D = Scalar('int', 'D') + tbody = ir.body + ir.step = scalar_D + new_loop = Loop(0, scalar_D, 1,[]) + for i in range(len(tbody)): + tbody[i] = swap_arr_to_reg(tbody[i], ir.iterate, new_loop.iterate) + # if isinstance(tbody[i], Assignment) and isinstance(tbody[i].rhs, Literal): + # print(codegen.gpu.to_string(tbody[i])) + new_loop.body.extend(tbody) + # iloops.append(new_loop) + ir.body = [new_loop] + return ir + +def recursive_tile(ir): + if isinstance(ir, Loop): + ir = tile_wD(ir) + for i in range(len(ir.body)): + ir.body[i] = recursive_tile(ir.body[i]) + + return ir + +def tile_loops(ir, tile_list): + if isinstance(ir, Loop) and ir.end.name() == 'dim': + scalar_D = Scalar('int', 'D') + ir.step = scalar_D + new_loop = Loop(0, scalar_D, 1,[]) + tile_list.append(ir) + for i in range(len(ir.body)): + ir.body[i] = tile_loops(ir.body[i], tile_list) + return ir + +def swap_loops_tile(ir): + decl = [] + del_decl = [] + replace_pair = [] + if isinstance(ir, Loop): + # print(ir, codegen.gpu.to_string(ir)) + for i in range(len(ir.body)): + # print(ir.body[i], codegen.gpu.to_string(ir.body[i])) + if isinstance(ir.body[i], Loop) and ir.body[i].end.name() != 'dim': + # print(ir.body[i], codegen.gpu.to_string(ir.body[i])) + ir.body[i], temp_decl, temp_del = swap_loops_tile(ir.body[i]) + decl.extend(temp_decl) + del_decl.extend(temp_del) + + elif isinstance(ir.body[i], Loop) and ir.body[i].end.name() == 'dim': + # print(ir.body[i], codegen.gpu.to_string(ir.body[i])) + tile_list = [] + tile_loops(ir.body[i], tile_list) + multi_lv_loop = {} + + + # print(tile_list) + for lidx in range(len(tile_list)): + t = tile_list[lidx] + # print(t, codegen.gpu.to_string(t)) + # add tiled loop here + scalar_D = Scalar('int', 'D') + new_loop = Loop(0, scalar_D, 1,[]) + temp_body = [] + for k in range(len(t.body)): + item = t.body[k] + # print(item, codegen.gpu.to_string(item)) + if isinstance(item, Loop) and item.end.name() == 'dim': + if new_loop.body != []: + # add new_tiled loop and create a new one + temp_body.append(new_loop) + new_loop = Loop(0, scalar_D, 1,[]) + + if item in multi_lv_loop.keys(): + multi_lv_loop[item].append([multi_lv_loop[item][-1][0]+1, t]) + else: + multi_lv_loop[item] = [[1, t]] + temp_body.append(item) + elif t in multi_lv_loop.keys(): + oloop = Loop(0, scalar_D, 1,[]) + for jj in range(len(t.body)): + t.body[jj] = swap_arr_to_reg(t.body[jj], multi_lv_loop[t][0][1].iterate, oloop.iterate) + loop1 = oloop + for i in range(len(multi_lv_loop[t])): + tloop = Loop(0, scalar_D, 1,[]) + loop1.body.append(tloop) + if i+1 < len(multi_lv_loop[t]): + for jj in range(len(t.body)): + t.body[jj] = swap_arr_to_reg(t.body[jj], multi_lv_loop[t][i+1][1].iterate, tloop.iterate) + loop1 = tloop + for jj in range(len(t.body)): + t.body[jj] = swap_arr_to_reg(t.body[jj], t.iterate, tloop.iterate) + tloop.body = t.body + temp_body.append(oloop) + else: + # add tiled loop here + # print(codegen.gpu.to_string(item), item.lhs, item.rhs) + if isinstance(item, Assignment) and isinstance(item.rhs, Literal): + new_lhs = Ndarray(item.lhs.dtype, [scalar_D]) + decl.append(Decl(new_lhs)) + del_decl.append(item.lhs) + replace_pair.append([new_lhs, item.lhs]) + + item = swap_arr_to_reg(item, t.iterate, new_loop.iterate) + new_loop.body.append(item) + if new_loop.body != []: + temp_body.append(new_loop) + tile_list[lidx].body = temp_body + for i in range(len(ir.body)): + if replace_pair: + for jj in replace_pair: + # print(codegen.gpu.to_string(jj[0]), codegen.gpu.to_string(jj[1]), codegen.gpu.to_string(ir.body[i])) + if isinstance(ir.body[i], Loop): + for temp in ir.body[i].body: + temp = swap_reg_to_arr(temp, jj[1], jj[0], ir.body[i]) + else: + ir.body[i] = swap_reg_to_arr(ir.body[i], jj[1], jj[0], ir) + # print(ir.body[i], codegen.gpu.to_string(ir.body[i])) + return ir, decl, del_decl + +def tile_loop(ast, eval_list=[]): + + + if ast.compute and ast.valid: + # if ast.compute: + # print(ast.compute[0], codegen.gpu.to_string(ast.compute[0])) + # print(codegen.gpu.to_string(ast.eval), codegen.gpu.to_string(ast.operators[0].eval), codegen.gpu.to_string(ast.operators[1].eval), ast.operators[0].eval) + # print(ast.eval, codegen.gpu.to_string(ast.eval)) + for i in ast.compute: + # recursive_tile(i) + body, decl, del_decl = swap_loops_tile(i) + for i in range(len(decl)): + # print(decl[i], del_decl[i], codegen.gpu.to_string(decl[i]), codegen.gpu.to_string(del_decl[i])) + eval_list.append([del_decl[i], decl[i].dobject]) + ast.decl.extend(decl) + for dd in ast.decl: + if dd.dobject in del_decl: + ast.decl.remove(dd) + + if type(ast) == BatchOp: + if type(ast.operators[1]) == BatchOp: + tile_loop(ast.operators[1], eval_list) + if type(ast.operators[0]) == BatchOp: + tile_loop(ast.operators[0], eval_list) + else: + return + + if not ast.valid: + # print(ast.op_type, eval_list, ast.eval, ast.decl) + for id, item in enumerate(eval_list): + if item[0] == ast.eval: + ast.eval = item[1] + eval_list.pop(id) \ No newline at end of file diff --git a/batch/opt/sort/__pycache__/mysort.cpython-310.pyc b/batch/opt/sort/__pycache__/mysort.cpython-310.pyc new file mode 100644 index 0000000..02773e9 Binary files /dev/null and b/batch/opt/sort/__pycache__/mysort.cpython-310.pyc differ diff --git a/batch/opt/sort/__pycache__/sort.cpython-310.pyc b/batch/opt/sort/__pycache__/sort.cpython-310.pyc new file mode 100644 index 0000000..e7c4063 Binary files /dev/null and b/batch/opt/sort/__pycache__/sort.cpython-310.pyc differ diff --git a/batch/opt/sort/build_indexing.py b/batch/opt/sort/build_indexing.py new file mode 100644 index 0000000..2132098 --- /dev/null +++ b/batch/opt/sort/build_indexing.py @@ -0,0 +1,31 @@ +import mysort +import torch + +batch_size = 128 + +th = h = torch.randint(0, 9999, (batch_size, )).int().cuda(0) +tr = r = torch.randint(0, 100, (batch_size, )).int().cuda(0) +tt = t = torch.randint(0, 9999, (batch_size, )).int().cuda(0) + +# h = torch.tensor([1, 2, 3, 4]) +# t = torch.tensor([10, 20, 30, 40]) +# r = torch.tensor([3, 1, 4, 2]) + +print(h) +print(t) +print(r) +print('after sorting:::') +sorted_indices = torch.argsort(r) + +sorted_h = h[sorted_indices] +sorted_t = t[sorted_indices] +sorted_r = r[sorted_indices] +print("Sorted h:", sorted_h) +print("Sorted t:", sorted_t) +print(sorted_r) + +runiq = torch.zeros((batch_size//16, 16)).int().cuda(0) +rbuffer = torch.zeros((batch_size//16, 16)).int().cuda(0) +uniq_cnt = torch.zeros((batch_size//16,)).int().cuda(0) + +th, tt, tr = mysort.index_building(th, tt, tr, runiq, rbuffer, uniq_cnt, batch_size, 16, 100) diff --git a/batch/opt/sort/mysort.py b/batch/opt/sort/mysort.py new file mode 100644 index 0000000..0784afa --- /dev/null +++ b/batch/opt/sort/mysort.py @@ -0,0 +1,18 @@ + +import torch +from torch.utils.cpp_extension import load + +sort_func = load(name='sort', sources=['sort.cu']) + +class build_index(torch.autograd.Function): + @staticmethod + def forward(ctx, head, tail, relation, r_uniq, r_buffer, uniq_cnt, n, gs, rel_num): + sort_func.gpu_sort(head, tail, relation, r_uniq, r_buffer, uniq_cnt, n, gs, rel_num) + return head, tail, relation + + @staticmethod + def backward(ctx): + pass + + +index_building = build_index.apply \ No newline at end of file diff --git a/batch/opt/sort/sort.cu b/batch/opt/sort/sort.cu new file mode 100644 index 0000000..6932e6d --- /dev/null +++ b/batch/opt/sort/sort.cu @@ -0,0 +1,121 @@ +#include +#include +#include +#include +#include +#include +using namespace std; + +#define BLOCK_SIZE 256 +#define C 16 +#define T 3 +#define DIV(x, ts) ((x) % (ts) != 0 ? (x) / (ts) + 1 : (x) / (ts)) + +__device__ void swap(int& a, int& b, int& a_idx, int& b_idx) { + int tmp = a; + a = b; + b = tmp; + tmp = a_idx; + a_idx = b_idx; + b_idx = tmp; +} + +__device__ void bitonic_sort(int* arr, int* ord) { + __shared__ int shared_arr[C]; + __shared__ int shared_ord[C]; + + int tid = threadIdx.x; + shared_arr[tid] = arr[tid]; + shared_ord[tid] = tid; + __syncthreads(); + + for (int k = 2; k <= C; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + __syncthreads(); + int ixj = tid ^ j; + if (ixj > tid) { + if ((tid & k) == 0 && shared_arr[tid] > shared_arr[ixj]) + swap(shared_arr[tid], shared_arr[ixj], shared_ord[tid], shared_ord[ixj]); + if ((tid & k) != 0 && shared_arr[tid] < shared_arr[ixj]) + swap(shared_arr[tid], shared_arr[ixj], shared_ord[tid], shared_ord[ixj]); + } + } + } + + __syncthreads(); + arr[tid] = shared_arr[tid]; + ord[shared_ord[tid]] = tid; +} + +__global__ void build_index(torch::PackedTensorAccessor32 indices, torch::PackedTensorAccessor32 r_Uniq, torch::PackedTensorAccessor32 r_Buffer, torch::PackedTensorAccessor32 uniq_cnt, int rel_num){ + __shared__ int idx[C], ord[C], ibuf[C], iuniq[C], pcount[C], ord_uniq[C]; + int tid = threadIdx.x; + idx[tid] = indices[blockIdx.x * C + tid]; + ord[tid] = tid; + ord_uniq[tid] = tid; + pcount[tid] = 0; + __syncthreads(); + + bitonic_sort(idx, ord); + ibuf[tid] = (tid > 0 && idx[tid] > idx[tid-1]) ? 1:0; + __syncthreads(); + + for (int offset = 1; offset < C; offset *= 2) { + __syncthreads(); + if (tid >= offset) { + ibuf[tid] += ibuf[tid - offset]; + } + } + + if (tid == 0) { pcount[ibuf[C-1]+1] = C; } + else if (idx[tid] > idx[tid-1]) { + pcount[ibuf[tid]] = tid; } + iuniq[tid] = rel_num; + __syncthreads(); + + // exceed threshold + if (tid > 0 && pcount[tid]-pcount[tid-1]>T) { + iuniq[tid-1] = idx[pcount[tid]-1]; } + __syncthreads(); + + bitonic_sort(iuniq, ord_uniq); + + int temp = ord_uniq[ibuf[tid]]; + if (iuniq[temp] < rel_num){ + ibuf[tid] = temp; + }else{ + ibuf[tid] = idx[tid] + C; + } + if (iuniq[tid] < rel_num && iuniq[tid+1] == rel_num){ + uniq_cnt[blockIdx.x] = tid+1; + } + r_Buffer[blockIdx.x][tid] = ibuf[ord[tid]]; + r_Uniq[blockIdx.x][tid] = iuniq[tid]; +} + + +void gpu_sort(torch::Tensor head, torch::Tensor tail, torch::Tensor relation, torch::Tensor r_Uniq, torch::Tensor r_Buffer, torch::Tensor uniq_cnt, int batch, int group_size, int rel_num) { + // int batch=4096; + dim3 nblocks(DIV(batch, C)); + dim3 nthreads(32, C); + + torch::Tensor sorted_indices = torch::argsort(relation); + std::cout << "Original Head: " << relation << std::endl; + + head = torch::index_select(head, 0, sorted_indices); + tail = torch::index_select(tail, 0, sorted_indices); + relation = torch::index_select(relation, 0, sorted_indices); + + + std::cout << "Sorted Head: " << relation << std::endl; + + build_index<<< batch/C, C>>>(relation.packed_accessor32(), r_Uniq.packed_accessor32(), r_Buffer.packed_accessor32(), uniq_cnt.packed_accessor32(), rel_num); + + std::cout << "uniq: " << r_Uniq << std::endl; + std::cout << "buffer: " << r_Buffer << std::endl; + std::cout << "uniq_cnt: " << uniq_cnt << std::endl; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("gpu_sort", &gpu_sort, "ss"); +} \ No newline at end of file diff --git a/batch/opt/sort/test b/batch/opt/sort/test new file mode 100755 index 0000000..bd901b9 Binary files /dev/null and b/batch/opt/sort/test differ diff --git a/batch/opt/sort/test.cu b/batch/opt/sort/test.cu new file mode 100644 index 0000000..db208ba --- /dev/null +++ b/batch/opt/sort/test.cu @@ -0,0 +1,187 @@ +#include +#include +#include +#include +using namespace std; + +#define BLOCK_SIZE 256 +#define C 16 +#define T 2 + + +__device__ void swap(long& a, long& b, long& a_idx, long& b_idx) { + int tmp = a; + a = b; + b = tmp; + tmp = a_idx; + a_idx = b_idx; + b_idx = tmp; +} + +__device__ void bitonic_sort(long* arr, long* ord) { + __shared__ long shared_arr[C]; + __shared__ long shared_ord[C]; + + int tid = threadIdx.x; + shared_arr[tid] = arr[tid]; + shared_ord[tid] = tid; + __syncthreads(); + + for (int k = 2; k <= C; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + __syncthreads(); + int ixj = tid ^ j; + if (ixj > tid) { + if ((tid & k) == 0 && shared_arr[tid] > shared_arr[ixj]) + swap(shared_arr[tid], shared_arr[ixj], shared_ord[tid], shared_ord[ixj]); + if ((tid & k) != 0 && shared_arr[tid] < shared_arr[ixj]) + swap(shared_arr[tid], shared_arr[ixj], shared_ord[tid], shared_ord[ixj]); + } + } + } + + __syncthreads(); + arr[tid] = shared_arr[tid]; + ord[shared_ord[tid]] = tid; +} + +__global__ void build_index(long * indices, long * uniq_idx, long * buf_idx, long * uniq_cnt){ + __shared__ long idx[C], ord[C], ibuf[C], iuniq[C], count[C], ord_uniq[C]; + int tid = threadIdx.x; + idx[tid] = indices[blockIdx.x * C + tid]; + ord[tid] = tid; + ord_uniq[tid] = tid; + count[tid] = 0; + __syncthreads(); + + bitonic_sort(idx, ord); + ibuf[tid] = (tid > 0 && idx[tid] > idx[tid-1]) ? 1:0; + __syncthreads(); + + for (int offset = 1; offset < C; offset *= 2) { + __syncthreads(); + if (tid >= offset) { + ibuf[tid] += ibuf[tid - offset]; + } + } + + // buf_idx[blockIdx.x * C + tid] = ibuf[ord[tid]]; + if (tid == 0) { count[ibuf[C-1]+1] = C; } + else if (idx[tid] > idx[tid-1]) { + count[ibuf[tid]] = tid; } + iuniq[tid] = 999; + __syncthreads(); + + // exceed threshold + if (tid > 0 && count[tid]-count[tid-1]>T) { + iuniq[tid-1] = idx[count[tid]-1]; } + __syncthreads(); + + bitonic_sort(iuniq, ord_uniq); + + int temp = ord_uniq[ibuf[tid]]; + // if(threadIdx.x == 0 && blockIdx.x == 0){ + // for(int i=0;i>>(d_data, d_uniq_idx, d_buf_idx, d_uniq_cnt); + + // print result + // for (int i = 0; i < n; i++) { + // printf("%d ", data[i]); + // } + // printf("\n%d\n", n); + + cudaMemcpy(uniq_idx, d_uniq_idx, n*sizeof(long), cudaMemcpyDeviceToHost); + cudaMemcpy(buf_idx, d_buf_idx, n*sizeof(long), cudaMemcpyDeviceToHost); + cudaMemcpy(uniq_cnt, d_uniq_cnt, n/C*sizeof(long), cudaMemcpyDeviceToHost); + printf("unique idx:\n"); + for(int i=0;iab', torch.einsum('ab,ab->a', pemb[r], eemb[h]-eemb[t]), pemb[r]) + # print(y) + + # x = run.gpu.compile_and_run(code, 4096, 512, 0, eemb, h,t, 0, remb, r, pemb) + # print(x) + def transR(): nnodes = Var('nnodes') @@ -60,8 +102,29 @@ def transR(): res = bvm(vh -vt, mr) + vr code = codegen.cpu.print_cpp(res._gen_ir()) - + + ast = res._gen_ir() + opt.fusion_rules.fuse_operators(ast) + # todo decouple operators + opt.node_wise.tiling.tile_loop(ast) + opt.node_wise.parallelism.parallel(ast) + opt.node_wise.smem.add_smem(ast) + # code = codegen.cpu.print_cpp(ast) + code = codegen.gpu.print_cuda(ast) print(code) + # h = torch.randint(0, 9999, (4096, )).cuda(0) + # r = torch.randint(0, 100, (4096, )).cuda(0) + # t = torch.randint(0, 9999, (4096, )).cuda(0) + # eemb = torch.rand((9999, 512)).cuda(0) + # remb = torch.rand((100, 512)).cuda(0) + # pemb = torch.rand((100, 512, 512)).cuda(0) + + # y = torch.einsum('ab,abc->ac', eemb[h] - eemb[t], pemb[r]) + remb[r] + # print(y) + + # x = run.gpu.compile_and_run(code, 4096, 512, 0, eemb, h,t, 0, pemb, r, remb) + # print(x) + def transF(): nnodes = Var('nnodes') @@ -76,12 +139,32 @@ def transF(): vh = Batch(Eemb[h]) vt = Batch(Eemb[t]) vr = Batch(Remb[r]) + + alpha = Const(val=2, dtype='float') + alpha = Batch(alpha) + # alpha = 2 - res = bvv(vh, vt) - bvv(vh - vt, vr) - + res = bvv(vh, vt) - bvv(vh - vt, vr) + code = codegen.cpu.print_cpp(res._gen_ir()) - + ast = res._gen_ir() + opt.fusion_rules.fuse_operators(ast) + opt.node_wise.tiling.tile_loop(ast) + opt.node_wise.parallelism.parallel(ast) + opt.node_wise.smem.add_smem(ast) + code = codegen.gpu.print_cuda(ast) print(code) + # h = torch.randint(0, 9999, (4096, )).cuda(0) + # r = torch.randint(0, 100, (4096, )).cuda(0) + # t = torch.randint(0, 9999, (4096, )).cuda(0) + # eemb = torch.rand((9999, 512)).cuda(0) + # remb = torch.rand((100, 512)).cuda(0) + + # y = torch.einsum('ab,ab->a', eemb[h], eemb[t]) - torch.einsum('ab,ab->a',(eemb[h] - eemb[t]), remb[r]) + # print(y) + + # x = run.gpu.compile_and_run(code, 4096, 512, 0, eemb, h,t, 0, remb, r) + # print(x) def RESCAL(): nnodes = Var('nnodes') @@ -100,14 +183,76 @@ def RESCAL(): res = bvv(bvm(vh, mr), vt) code = codegen.cpu.print_cpp(res._gen_ir()) - + + ast = res._gen_ir() + + opt.fusion_rules.fuse_operators(ast) + opt.node_wise.tiling.tile_loop(ast) + opt.node_wise.parallelism.parallel(ast) + opt.node_wise.smem.add_smem(ast) + # traversal call funcs to opt ir + # code = codegen.cpu.print_cpp(ast) + code = codegen.gpu.print_cuda(ast) print(code) + # h = torch.randint(0, 9999, (4096, )).cuda(0) + # r = torch.randint(0, 100, (4096, )).cuda(0) + # t = torch.randint(0, 9999, (4096, )).cuda(0) + # eemb = torch.rand((9999, 512)).cuda(0) + # remb = torch.rand((100, 512, 512)).cuda(0) + + # y = torch.einsum('ab,ab->a', torch.einsum('ab,abc->ac', eemb[h], remb[r]), eemb[t]) + # print(y, y.shape) + + # x = run.gpu.compile_and_run(code, 4096, 512, 0, eemb, h, 0, remb, r, t) + # print(x) + +def test(): + nnodes = Var('nnodes') + nedges = Var('nedges') + dim = Var('dim') + batch_size = Var('batch_size') + Eemb = Tensor('Eemb', (nnodes, dim)) + Remb = Tensor('Remb', (nedges, dim)) + Proj = Tensor('Proj', (nedges, dim, dim)) + h = Tensor('h', (batch_size, ), dtype='int') + t = Tensor('t', (batch_size, ), dtype='int') + r = Tensor('r', (batch_size, ), dtype='int') + vh = Batch(Eemb[h]) + vt = Batch(Eemb[t]) + vr = Batch(Remb[r]) + vrr = Batch(Remb[r]) + proj_m = Batch(Proj[r]) + proj_h = Batch(Proj[h]) + + # res = vh - vt + vr - vrr + res = bov(vh+vr, vt-vr) + + # code = codegen.cpu.print_cpp(res._gen_ir()) + # print(code) + ast = res._gen_ir() + opt.fusion_rules.fuse_operators(ast) + opt.parallelism.parallel(ast) + + code = codegen.gpu.print_cuda(ast) + print(code) + h = torch.randint(0, 9999, (4096, )).cuda(0) + r = torch.randint(0, 100, (4096, )).cuda(0) + t = torch.randint(0, 9999, (4096, )).cuda(0) + eemb = torch.rand((9999, 512)).cuda(0) + remb = torch.rand((100, 512)).cuda(0) + y = torch.einsum('ab,ac->abc', eemb[h] + remb[r], eemb[t] - remb[r]) + print(y) + + x = run.gpu.compile_and_run(code, 4096, 512, 0, eemb, h, 0, remb, r, t) + print(x) if __name__ == "__main__": - transE() - transH() - transR() - transF() - RESCAL() + # test() # bov success + # transE() # success + # transH() # success + transR() # success + # transF() # success + # RESCAL() # success + \ No newline at end of file diff --git a/codegen/__init__.py b/codegen/__init__.py index 0674dd6..8ef5280 100644 --- a/codegen/__init__.py +++ b/codegen/__init__.py @@ -1 +1,2 @@ -import codegen.cpu \ No newline at end of file +import codegen.cpu +import codegen.gpu \ No newline at end of file diff --git a/codegen/cpu.py b/codegen/cpu.py index f9be9cf..2cfcea6 100644 --- a/codegen/cpu.py +++ b/codegen/cpu.py @@ -1,12 +1,20 @@ from core.ast2ir import * +# from cset.ast2ir import * import helpers import batch +# import cset + def to_string(ir): match ir.__class__.__name__: case 'Expr': - return f"({to_string(ir.left)}" + f" {ir.op} " + f"{to_string(ir.right)})" + if ir.op in arith_op.values(): + return f"({to_string(ir.left)}" + f" {ir.op} " + f"{to_string(ir.right)})" + elif ir.op == 'bigger': + return f"({to_string(ir.left)} > {to_string(ir.right)} ? ({to_string(ir.left)}) : ({to_string(ir.right)}))" + elif ir.op == 'smaller': + return f"({to_string(ir.left)} < {to_string(ir.right)} ? ({to_string(ir.left)}) : ({to_string(ir.right)}))" case 'Assignment': if ir.op is None: return f"{to_string(ir.lhs)} = {to_string(ir.rhs)};\n" @@ -19,59 +27,71 @@ def to_string(ir): code += to_string(e) code += "} \n" return code + case 'FilterLoop': + code = f"for (int {to_string(ir.iterate)} = {to_string(ir.start)}; {to_string(ir.iterate)} < {to_string(ir.end)}; {to_string(ir.iterate)} += {to_string(ir.step)}) {{\n" + for e in ir.body: + if e: + code += to_string(e) + if ir.cond: + code += f"if({to_string(ir.cond)}){{\n" + for e in ir.cond_body: + if e: + code += to_string(e) + code += "} \n" + code += "} \n" + return code + case 'Not': + return f"!{to_string(ir.dobject)}" case 'Scalar' | 'Ndarray' | 'Ref': return ir.name() - case 'Index': - if ir.ind_arr != None: - if type(ir.ind_arr) == Slice: - return f'{to_string(ir.dobject)}[(({to_string(ir.ind_arr.start)})+({to_string(ir.ind_arr.step)})*({to_string(ir.index)}))]' - else: # idx is a Tensor - if ir.index == None: - return f'{to_string(ir.dobject)}[{to_string(ir.ind_arr)}]' - else: - return f'{to_string(ir.dobject)}[{to_string(ir.ind_arr)}[{to_string(ir.index)}]]' + case 'Literal': + return str(ir.val) + case 'Indexing': + if type(ir.dobject) == Slice: + if ir.dobject.step == 1 or (type(ir.dobject.step) == Literal and ir.dobject.step.val == 1): + return f'(({to_string(ir.dobject.start)})+({to_string(ir.idx)}))' + else: + return f'(({to_string(ir.dobject.start)})+({to_string(ir.dobject.step)})*({to_string(ir.idx)}))' else: - return f'{to_string(ir.dobject)}[{to_string(ir.index)}]' + return f'{to_string(ir.dobject)}[{to_string(ir.idx)}]' + case 'Search': + code = f"BinarySearch({to_string(ir.dobject)}, {to_string(ir.start)}, {to_string(ir.end)}, {to_string(ir.item)})" + return code case 'Decl': # variables are passed in as pytorch arguments if type(ir.dobject) == Scalar: if not ir.dobject.is_arg: - # it is a zero or one - if ir.dobject.val != None: - return f"{ir.dobject.dtype} {ir.dobject.name()} = {to_string(ir.dobject.val)};\n" - else: - return f"{ir.dobject.dtype} {ir.dobject.name()};\n" + return f"{ir.dobject.dtype} {ir.dobject.name()};\n" else: return '' elif type(ir.dobject) == Ndarray: code = '' if not ir.dobject.is_arg: - if ir.dobject.val != None: - code = f'torch::Tensor obj_{ir.dobject.name()} = torch::{"ones" if ir.dobject.val == 1 else "zeros"}({{{",".join([to_string(s) for s in ir.dobject.size])}}}, at::k{"Int" if ir.dobject.dtype=="int" else "Float"});\n' - else: - code = f'torch::Tensor obj_{ir.dobject.name()} = torch::empty({{{",".join([to_string(s) for s in ir.dobject.size])}}}, at::k{"Int" if ir.dobject.dtype=="int" else "Float"});\n' - + code = f'torch::Tensor obj_{ir.dobject.name()} = torch::empty({{{",".join([to_string(s) for s in ir.dobject.size])}}}, at::k{"Int" if ir.dobject.dtype=="int" else "Float"});\n' code += f'auto {ir.dobject.name()} = obj_{ir.dobject.name()}.accessor<{ir.dobject.dtype}, {len(ir.dobject.size)}>();\n' return code - # elif type(ir.dobject) == Ref: - # code = f'{ir.dobject.dobject.dtype}* {ir.dobject.name()} = ({ir.dobject.dobject.dtype}*)&{ir.dobject.dobject.addr()}' - # return code + case 'Math': + return f"{ir.type}({to_string(ir.val)})" case _: return str(ir) - - def gen_cpp(ast, ir): def action(node, res): - if type(node) == Var or type(node) == One or type(node) == Zero or type(node) == Ones or type(node) == Zeros or type(node) == Tensor: - res.extend(node.decl) - elif type(node) == TensorOp: - res.extend(node.decl) - res.extend(node.compute) - elif type(node) == batch.ast.BatchOp: - res.extend(node.decl) - res.extend(node.compute) + if node.valid == True: + if type(node) == Var or type(node) == Tensor: + res.extend(node.decl) + elif type(node) == TensorOp: + res.extend(node.decl) + res.extend(node.compute) + elif type(node) == batch.ast.BatchOp: + res.extend(node.decl) + res.extend(node.compute) + elif type(node) == cset.ast.Set: + res.extend(node.decl) + elif type(node) == cset.ast.SetOp: + res.extend(node.decl) + res.extend(node.compute) t = helpers.Traversal(action) ir.extend(t(ast)) diff --git a/codegen/gpu.py b/codegen/gpu.py new file mode 100644 index 0000000..0d70f46 --- /dev/null +++ b/codegen/gpu.py @@ -0,0 +1,179 @@ +from core.ast2ir import * +import helpers +import batch +from batch.opt.ir import * +from codegen.gpu_instructionsets import * + +def to_string(ir): + match ir.__class__.__name__: + case 'Expr': + return f"({to_string(ir.left)}" + f" {ir.op} " + f"{to_string(ir.right)})" + case 'Assignment': + if ir.op is None: + return f"{to_string(ir.lhs)} = {to_string(ir.rhs)};\n" + else: + return f"{to_string(ir.lhs)} {ir.op}= {to_string(ir.rhs)};\n" + case 'Loop': + code = f"for (int {to_string(ir.iterate)} = {to_string(ir.start)}; {to_string(ir.iterate)} < {to_string(ir.end)}; {to_string(ir.iterate)} += {to_string(ir.step)}) {{\n" + # print(ir, ir.body, to_string(ir.body)) + for e in ir.body: + if e: + code += to_string(e) + code += "} \n" + return code + case 'Scalar' | 'Ndarray' | 'Ref': + return ir.name() + case 'Literal': + return str(ir.val) + case 'Indexing': + if type(ir.dobject) == Slice: + return f'(({to_string(ir.dobject.start)})+({to_string(ir.dobject.step)})*({to_string(ir.idx)}))' + # elif type(ir.dobject) == Pointer: + # code = f'{to_string(ir.dobject)}[' + # for i in range(len(ir.dobject.dims)): + # code += f'{to_string(ir.idx)}*{to_string(ir.dobject.dims[i])}' + # if i 0; off >>= 1) {{\n {to_string(ir.dobject)} += __shfl_down_sync(0xffffffff, {to_string(ir.dobject)}, off); \n}}\n' + case 'ShuffleUp': + return f'for (int off = blockDim.x/2; off > 0; off >>= 1) {{\n {to_string(ir.dobject)} += __shfl_up_sync(0xffffffff, {to_string(ir.dobject)}, off); \n}}\n' + case 'ShuffleXor': + return f'for (int off = blockDim.x/2; off > 0; off >>= 1) {{\n {to_string(ir.dobject)} += __shfl_xor_sync(0xffffffff, {to_string(ir.dobject)}, off); \n}}\n' + case 'BroadCast': + return f'{to_string(ir.dobject)} = __shfl_sync(0xffffffff, {to_string(ir.dobject)}, 0);\n' + case 'SaveAtThread': + return f'if (threadIdx.x == {ir.threadid}) {{\n {to_string(Assignment(ir.dst, ir.src))} }}\n' + case 'Uniq' | 'Buffer': + return f'{ir.dobject.__name__}_{ir.__class__.__name__}[{to_string(BlockIdx())}]' + # return f'[{to_string(BlockIdx())}]' + case 'IF': + return f"{to_string(ir.left)} = {to_string(ir.condition)} ? {to_string(ir.true_var)} : {to_string(ir.false_var)};\n" + case 'Pointer': + return f'{ir.name()}' + case 'Access_ptr': + code = f'{to_string(ir.dobject)}[' + for i in range(len(ir.idx)-1): + code += '(' + for i in range(len(ir.idx)): + if i()' if type(args[a]) == Tensor else f'{a}' for a in args]) + # ptrs = ', '.join([f'{args[a].dtype}* {a}' if type(args[a]) == Tensor else f'{args[a].dtype} {a}' for a in args]) + + argsptr = ', '.join([f'obj_{a}.packed_accessor32<{args[a].dtype if args[a].dtype!="int" else "int64_t"}, {len(args[a].ref_size)}, torch::RestrictPtrTraits>()' if type(args[a]) == Tensor else f'{a}' for a in args]) + ptrs = ', '.join([f'torch::PackedTensorAccessor32<{args[a].dtype if args[a].dtype!="int" else "int64_t"}, {len(args[a].ref_size)}, torch::RestrictPtrTraits> {a}' if type(args[a]) == Tensor else f'{args[a].dtype} {a}' for a in args]) + # in cuda kernel: const torch::PackedTensorAccessor32 + # host call cuda: .packed_accessor32() + + code = '' + declare = '' + for d in gpu_ir: + if d: + if type(d) == Decl and type(d.dobject) == Ndarray: + declare += to_string(d) + argsptr += f', obj_{d.dobject.name()}.packed_accessor32<{d.dobject.dtype}, {len(d.dobject.size)}, torch::RestrictPtrTraits>()' + ptrs += f', torch::PackedTensorAccessor32<{d.dobject.dtype}, {len(d.dobject.size)}, torch::RestrictPtrTraits> {d.dobject.name()}' + elif type(d) == Decl and type(d.dobject) in [Buffer, Uniq]: + declare += to_string(d) + argsptr += f', {d.dobject.dobject.__name__}_{d.dobject.__class__.__name__}.packed_accessor32<{d.dobject.dobject.dtype}, 2, torch::RestrictPtrTraits>()' + ptrs += f', torch::PackedTensorAccessor32<{d.dobject.dobject.dtype}, 2, torch::RestrictPtrTraits> {d.dobject.dobject.__name__}_{d.dobject.__class__.__name__}' + else: + code += to_string(d) + # print(declare) + Return = '' + if type(ast.eval) == Scalar: + rtype = ast.dtype + Return = f'return {ast.eval.name()};\n' + elif type(ast.eval) == Ndarray: + rtype = 'torch::Tensor' + Return = f'return obj_{ast.eval.name()};\n' + else: + raise TypeError('wrong output type', ast.eval) + + with open('codegen/gpu_template.cu', 'r') as f: + c_code = f.read() + c_code = c_code.replace('RTYPE', rtype).replace('FNAME', ast.name).replace('ARGS', argscpu).replace('CODE', code).replace('PTR_VARS', argsptr).replace('PTRS', ptrs).replace('DECL', declare).replace('RETURN', Return) + return c_code \ No newline at end of file diff --git a/codegen/gpu_instructionsets.py b/codegen/gpu_instructionsets.py new file mode 100644 index 0000000..e6bd2d3 --- /dev/null +++ b/codegen/gpu_instructionsets.py @@ -0,0 +1,21 @@ +from batch.opt.ir import * +from codegen import * + +def ir2gpu(ir): + match ir.__class__.__name__: + case 'BlockIdx': + return 'blockIdx.x' + case 'BlockIdy': + return 'blockIdx.y' + case 'BlockDimx': + return 'blockDim.x' + case 'BlockDimy': + return 'blockDim.y' + case 'ThreadIdy': + return 'threadIdx.y' + case 'ThreadIdx': + return 'threadIdx.x' + case 'SyncThreads': + return '__syncthreads();\n' + case 'SyncWarps': + return '__syncwarps();\n' \ No newline at end of file diff --git a/codegen/gpu_template.cu b/codegen/gpu_template.cu new file mode 100644 index 0000000..5219403 --- /dev/null +++ b/codegen/gpu_template.cu @@ -0,0 +1,16 @@ +#include + +__global__ void FNAME_kernel(PTRS){ + CODE +} + +RTYPE FNAME(ARGS) +{ + DECL + FNAME_kernel<<< batch_size/16, dim3(32,16) >>>(PTR_VARS); + RETURN +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("run", &FNAME); +} \ No newline at end of file diff --git a/core/ast.py b/core/ast.py index ded7e22..0c19699 100644 --- a/core/ast.py +++ b/core/ast.py @@ -1,18 +1,27 @@ import copy import core +MIN_INT = -2147483648 +MAX_INT = 2147483647 + +arith_op = {'add': '+', 'sub': '-', 'mul': '*', 'floordiv': '/', 'truediv': '/'} +math_op = ['round', 'abs'] +cmp_op = ['bigger', 'smaller'] +func_op = ['index', 'apply', 'reduce', 'aggr', 'einsum', 'setval'] + + def is_int_var(v): - return isinstance(v, Tensor) and v.dtype == 'int' and len(v._size()) == 0 + return isinstance(v, Tensor) and v.dtype == 'int' and len(v.ref_size) == 0 def is_scalar(v): - return isinstance(v, int|float) or (isinstance(v, Tensor) and len(v._size()) == 0) + return isinstance(v, int|float) or (isinstance(v, Tensor) and len(v.ref_size) == 0) def is_1dint_tensor(v): - return isinstance(v, Tensor) and v.dtype == 'int' and len(v._size()) == 1 + return isinstance(v, Tensor) and v.dtype == 'int' and len(v.ref_size) == 1 def eval_const_expr(e): - if type(e) == TensorOp and (e.op_type in op_mapping): + if type(e) == TensorOp and (e.op_type in arith_op): lhs = eval_const_expr(e.operators[0]) if lhs != None: rhs = eval_const_expr(e.operators[1]) @@ -35,7 +44,7 @@ def eval_const_expr(e): else: return None elif type(e) == Const: - return e.eval + return e.val else: return None @@ -54,7 +63,7 @@ def has_same_value(e1, e2): elif type(e1) == TensorOp: if e1.op_type != e2.op_type: return False - elif e1.op_type in op_mapping: + elif e1.op_type in arith_op: return has_same_value(e1.operators[0], e2.operators[0]) and has_same_value(e2.operators[1], e2.operators[1]) else: if len(e1.operators) != len(e2.operators): @@ -77,18 +86,24 @@ def is_same_size(s1, s2): return True -op_mapping = {'add':'+', 'sub':'-', 'mul':'*', 'floordiv':'/', 'truediv':'/'} + +def bigger(x, y): + return TensorOp('bigger', x, y) + +def smaller(x, y): + return TensorOp('smaller', x, y) class ASTNode: nuniq = 0 def __init__(self): self.decl = [] - self.compute = [] self.eval = None self.ref_count = 0 self.id = ASTNode.nuniq ASTNode.nuniq += 1 + self.valid = True + class Tensor(ASTNode): @@ -140,32 +155,37 @@ def apply(self, func, axis=0): if callable(func): from core.ast2ir import gen_ir op = TensorOp('apply', self, func, axis) - gen_ir(op) return op else: raise TypeError('must apply a callable function') def reduce(self, func, init, axis=0): - if callable(func): + if callable(func) and callable(init): from core.ast2ir import gen_ir op = TensorOp('reduce', self, func, init, axis) - gen_ir(op) return op else: raise TypeError('reduce must use a callable function') def sum(self, axis=0): func = lambda x, y: x + y - size = self._size()[:axis] + self._size()[axis+1:] - if (len(size) > 0): - init = Zeros(size, dtype=self.dtype) - else: - init = Const(0, dtype=self.dtype) + init = lambda x: x.setval(0) + return self.reduce(func, init, axis) + + def max(self, axis=0): + func = lambda x, y: bigger(x, y) + init = lambda x: x.setval(MIN_INT) + return self.reduce(func, init, axis) + + def min(self, axis=0): + func = lambda x, y: smaller(x, y) + init = lambda x: x.setval(MAX_INT) return self.reduce(func, init, axis) + def aggr(self, func, init, indices, axis=0, size=None): - if callable(func): - from core.ast2ir import gen_ir + if callable(func) and callable(init): + from core.ast2ir import gen_ir op = TensorOp('aggr', self, func, init, indices, axis, size) gen_ir(op) return op @@ -174,15 +194,24 @@ def aggr(self, func, init, indices, axis=0, size=None): def aggr_sum(self, indices, axis=0, size=None): func = lambda x, y: x + y - s = self._size()[:axis] + self._size()[axis+1:] - if (len(s) > 0): - init = Zeros(s, dtype=self.dtype) - else: - init = Const(0, dtype=self.dtype) + init = lambda x: x.setval(0) + return self.aggr(func, init, indices, axis, size) + + def aggr_max(self, indices, axis=0, size=None): + func = lambda x, y: bigger(x, y) + init = lambda x: x.setval(MIN_INT) + return self.aggr(func, init, indices, axis, size) + + def aggr_min(self, indices, axis=0, size=None): + func = lambda x, y: smaller(x, y) + init = lambda x: x.setval(MAX_INT) return self.aggr(func, init, indices, axis, size) + def setval(self, val): + return TensorOp('setval', self, val) + def _size(self): return self.fix_size + self.ref_size @@ -198,20 +227,16 @@ def size(self): else: return Const(0, dtype='int') - def _gen_ir(self): - return core.ast2ir.gen_ir(self) -class Ones(Tensor): - nones = 0 - def __init__(self, size, dtype='float'): - super.__init__(f'ones_{Ones.nones}', size, dtype, [], False) - Ones.nones += 1 -class Zeros(Tensor): - nzeros = 0 - def __init__(self, size, dtype='float'): - super().__init__(f'zeros_{Zeros.nzeros}', size, dtype, [], False) - Zeros.nzeros += 1 + def round(self): + return TensorOp('round', self) + + def abs(self): + return TensorOp('abs', self) + + def _gen_ir(self): + return core.ast2ir.gen_ir(self) class Var(Tensor): @@ -219,18 +244,6 @@ def __init__(self, name, dtype='int', is_arg=True): super().__init__(name, [], dtype, [], is_arg) -class One(Var): - none = 0 - def __init__(self, dtype='float'): - super().__init__(f'one_{One.none}', dtype, False) - One.none += 1 - -class Zero(Var): - nzero = 0 - def __init__(self, dtype='float'): - super().__init__(f'zero_{Zero.nzero}', dtype, False) - Zero.nzero += 1 - # const is var without name class Const(Var): @@ -238,6 +251,8 @@ class Const(Var): def __init__(self, val, dtype): super().__init__(f'c{Const.nconsts}', dtype) Const.nconsts += 1 + # slice is considered constant because once the slice is created its start, stop, step cannot be reassigned + # however, start, stop, step themselves can be variables if dtype == 'slice': assert type(val.start) == int or is_int_var(val.start) assert type(val.stop) == int or is_int_var(val.stop) @@ -249,15 +264,18 @@ def einsum(exp: str, tensor1, tensor2): return TensorOp('einsum', tensor1, tensor2, exp) class TensorOp(Tensor): - Types = ['index', 'apply', 'reduce', 'aggr', 'einsum'] + list(op_mapping.keys()) + Types = func_op + list(arith_op.keys()) + math_op + cmp_op def __init__(self, op_type, *operators): assert op_type in TensorOp.Types + self.compute = [] + self.output_order = [] # TODO: infer result data type dtype = operators[0].dtype self.operators = [] for opr in operators: + # an index can be referenced multiple times in the ast, we should create duplicate copies so that they can bind with different loop iterates if type(opr) == TensorOp and opr.op_type == 'index' and opr.ref_count >= 1: new_opr = copy.copy(opr) new_opr.ref_count = 1 @@ -270,21 +288,29 @@ def __init__(self, op_type, *operators): if isinstance(opr, ASTNode): opr.ref_count += 1 - # TODO: implement scalar +/-/*/div tensor - if op_type in op_mapping: - ref_size = self.operators[0].fix_size + self.operators[0].ref_size - fix_size = [] + if op_type in arith_op or op_type in cmp_op: + + if type(self.operators[0]) == int: + self.operators[0] = Const(self.operators[0], 'int') + elif type(operators[0]) == float: + self.operators[0] = Const(self.operators[0], 'float') if type(self.operators[1]) == int: self.operators[1] = Const(self.operators[1], 'int') elif type(operators[1]) == float: self.operators[1] = Const(self.operators[1], 'float') - + assert is_same_size(self.operators[0]._size(), self.operators[1]._size()) or len(self.operators[0]._size()) == 0 or len(self.operators[1]._size()) == 0 + if len(self.operators[0]._size()) < len(self.operators[1]._size()): + self.operators[0], self.operators[1] = self.operators[1], self.operators[0] + + ref_size = self.operators[0].fix_size + self.operators[0].ref_size + fix_size = [] + elif op_type == 'einsum': exp = self.operators[2] inputs, output = exp.split('->') input1, input2 = inputs.split(',') - op1_size = self.operators[0].fix_size + self.operators[0].ref_size - op2_size = self.operators[1].fix_size + self.operators[1].ref_size + op1_size = self.operators[0]._size() + op2_size = self.operators[1]._size() ref_size = [] fix_size = [] for i in output: @@ -321,11 +347,18 @@ def __init__(self, op_type, *operators): step = Const(step, 'int') self.operators[1] = Const(slice(start, stop, step), 'slice') - fix_size.append((stop - start)//step) + csize = eval_const_expr((stop - start)//step) + if csize != None: + fix_size.append(csize) + else: + if step.val == 1: + fix_size.append(stop-start) + else: + fix_size.append((stop - start)//step) elif is_int_var(self.operators[1]): self.operators[1] = self.operators[1] elif is_1dint_tensor(self.operators[1]): - fix_size.append(self.operators[1].ref_size[0]) + fix_size.append(self.operators[1]._size()[0]) else: raise TypeError('index must be int, Var of int, or 1d int Tensor') @@ -333,11 +366,6 @@ def __init__(self, op_type, *operators): assert type(self.operators[2]) == int axis = self.operators[2] self.operators[2] = Const(axis, 'int') - # size cannot be determined except for axis dimension, so set -1 - ref_size = [self.operators[0]._size()[axis], -1] - fix_size = [] - # data type also cannot be determined - dtype = None data_size = self.operators[0]._size() item_size = data_size[:axis] + data_size[axis + 1:] @@ -346,7 +374,13 @@ def __init__(self, op_type, *operators): self.operators[0].dtype, [], False) else: item = Var(f'item_of_{self.operators[0].name}', self.operators[0].dtype, False) + + ret = self.operators[1](item) + dtype = ret.dtype + ref_size = [self.operators[0]._size()[axis]] + ret._size() + fix_size = [] self.operators.append(item) + self.operators.append(ret) elif op_type == 'reduce': assert type(self.operators[3]) == int @@ -363,8 +397,10 @@ def __init__(self, op_type, *operators): else: item1 = Var(f'item1_of_{self.operators[0].name}', self.operators[0].dtype, False) item2 = Var(f'item2_of_{self.operators[0].name}', self.operators[0].dtype, False) + self.operators.append(item1) self.operators.append(item2) + self.operators.append(self.operators[1](item1, item2)) elif op_type == 'aggr': assert is_1dint_tensor(self.operators[3]) @@ -372,7 +408,7 @@ def __init__(self, op_type, *operators): axis = self.operators[4] self.operators[4] = Const(axis, 'int') if self.operators[5] == None: - self.operators[5] = self.operators[0]._size()[axis] + self.operators[5] = self.operators[3].ref_size[0] else: assert is_int_var(self.operators[5]) if type(self.operators[5]) == int: @@ -390,6 +426,25 @@ def __init__(self, op_type, *operators): item2 = Var(f'item2_of_{self.operators[0].name}', self.operators[0].dtype, False) self.operators.append(item1) self.operators.append(item2) + self.operators.append(self.operators[1](item1, item2)) + + elif op_type in math_op: + ref_size = self.operators[0].ref_size + fix_size = [] + if op_type == 'round': + dtype = 'int' + elif op_type == 'abs': + dtype = self.operators[0].dtype + + elif op_type == 'setval': + ref_size = self.operators[0].ref_size + fix_size = self.operators[0].fix_size + assert is_scalar(self.operators[1]) + if type(self.operators[1]) == int: + self.operators[1] = Const(self.operators[1], 'int') + elif type(self.operators[1]) == float: + self.operators[1] = Const(self.operators[1], 'float') + name = f'{op_type}_' + '_'.join([op.name if hasattr(op, 'name') else '' for op in self.operators]) @@ -397,7 +452,8 @@ def __init__(self, op_type, *operators): self.op_type = op_type + # call the init function for reduce and aggr + if self.op_type in ('reduce', 'aggr'): + self.operators[2] = self.operators[2](self) - - - + self.input_orders = [None for o in self.operators] \ No newline at end of file diff --git a/core/ast2ir.py b/core/ast2ir.py index 6062bf6..c7eb451 100755 --- a/core/ast2ir.py +++ b/core/ast2ir.py @@ -5,38 +5,61 @@ -def bind(arr: (Ndarray, Index), index, fix_ref=False): - # fix_ref == True means index at current position instead of the first unbind - if type(arr) == Ndarray or index == None or fix_ref: - return Index(arr, index=index) - else: - ref_chain = [arr] - while (type(ref_chain[-1].dobject) != Ndarray): - ref_chain.append(ref_chain[-1].dobject) - for ref in ref_chain[::-1]: - if ref.index == None: - ref.index = index - return arr - return Index(arr, index=index) - - - +def get_first_unbind(index: (Indexing, Ndarray, Slice)): + if type(index) == Indexing: + x = get_first_unbind(index.dobject) + if x != None: + return x + else: + if type(index.idx) == Literal and index.idx.val == -1: + return index + else: + y = get_first_unbind(index.idx) + return y + return None +def bind(index: (Indexing, Ndarray, Slice), idx): + x = get_first_unbind(index) + if x == None: + return Indexing(index, idx) + else: + x.idx = idx + return index + + +def replace_output(ir, old, new): + if type(ir) == list or type(ir) == tuple: + for l in ir: + replace_output(l, old, new) + elif type(ir) == Loop: + replace_output(ir.body, old, new) + elif type(ir) == Assignment: + if ir.lhs == old: + ir.lhs = new + else: + replace_output(ir.lhs, old, new) + elif type(ir) == Indexing: + if ir.dobject == old: + ir.dobject = new + else: + replace_output(ir.dobject, old, new) def gen_ir(node): assert isinstance(node, ASTNode) - if len(node.compute) > 0 or len(node.decl) > 0 or node.eval: + if node.eval or len(node.decl) > 0 or (type(node) == TensorOp and len(node.compute) > 0): return node if type(node) == Const: if node.dtype != 'slice': - node.eval = node.val + assert type(node.val) == int or type(node.val) == float + node.eval = Literal(node.val, node.dtype) else: node.val.start._gen_ir() node.val.stop._gen_ir() node.val.step._gen_ir() node.eval = Slice(node.val.start.eval, node.val.stop.eval, node.val.step.eval) + elif type(node) == Var or (type(node) == Tensor and len(node._size()) == 0): node.eval = Scalar(node.dtype, node.name, node.is_arg) node.decl = [Decl(node.eval)] @@ -47,21 +70,12 @@ def gen_ir(node): node.eval = Ndarray(node.dtype, size, node.name, node.is_arg) node.decl = [Decl(node.eval)] - # here we define two special tensors to simply programming for sum/prod operations - elif (type(node) == Ones or type(node) == Zeros) and len(node._size()) > 0: - size = helpers.get_ir_of_size(node._size()) - node.eval = Ndarray(node.dtype, size, node.name, False, 1 if (type(node) == Ones) else 0) - node.decl = [Decl(node.eval)] - - elif ((type(node) == Ones or type(node) == Zeros) and len(node._size()) == 0) or ((type(node) == One or type(node) == Zero)): - node.eval = Scalar(node.dtype, node.name, False, 1 if (type(node) == Ones or type(node) == One) else 0) - node.decl = [Decl(node.eval)] - elif type(node) == TensorOp: - if node.op_type in op_mapping: + if node.op_type in arith_op or node.op_type in cmp_op: node.operators[0]._gen_ir() node.operators[1]._gen_ir() - # TODO: add support for scalar + tensor + node.input_orders[0] = [] + node.input_orders[1] = [] assert isinstance(node.operators[0], Tensor) and isinstance(node.operators[1], Tensor) if is_same_size(node.operators[0]._size(), node.operators[1]._size()): if len(node._size()) > 0: @@ -81,98 +95,237 @@ def gen_ir(node): rhs = bind(rhs, pre_loop.iterate) res = bind(res, pre_loop.iterate) - op = op_mapping[node.op_type] + if node.op_type in arith_op: + op = arith_op[node.op_type] + else: + op = node.op_type assign = Assignment(res, Expr(lhs, rhs, op)) pre_loop.body.append(assign) else: node.eval = Scalar(node.dtype) node.decl = [Decl(node.eval)] - node.compute = [Assignment(node.eval, Expr(node.operators[0].eval, node.operators[1].eval, op_mapping[node.op_type]))] + if node.op_type in arith_op: + op = arith_op[node.op_type] + else: + op = node.op_type + node.compute = [Assignment(node.eval, Expr(node.operators[0].eval, node.operators[1].eval, op))] + else: + size = helpers.get_ir_of_size(node._size()) + node.eval = Ndarray(node.dtype, size) + node.decl = [Decl(node.eval)] + pre_loop = Loop(0, node.eval.size[0], 1, []) + node.compute = [pre_loop] + lhs = bind(node.operators[0].eval, pre_loop.iterate) + rhs = node.operators[1].eval + res = bind(node.eval, pre_loop.iterate) + for i in range(1, len(node.eval.size)): + loop = Loop(0, node.eval.size[i], 1, []) + pre_loop.body.append(loop) + pre_loop = loop + lhs = bind(lhs, pre_loop.iterate) + res = bind(res, pre_loop.iterate) + + if node.op_type in arith_op: + op = arith_op[node.op_type] + else: + op = node.op_type + assign = Assignment(res, Expr(lhs, rhs, op)) + pre_loop.body.append(assign) + + l = node.compute[0] + for i in range(len(node.eval.size)): + node.output_order.append((i, l)) + node.input_orders[0].append((i, l)) + node.input_orders[1].append((i, l)) + l = l.body[0] + + + + elif node.op_type in math_op: + node.operators[0]._gen_ir() + node.input_orders[0] = [] + if len(node._size()) > 0: + size = helpers.get_ir_of_size(node._size()) + node.eval = Ndarray(node.dtype, size) + node.decl = [Decl(node.eval)] + pre_loop = Loop(0, node.eval.size[0], 1, []) + node.compute = [pre_loop] + val = bind(node.operators[0].eval, pre_loop.iterate) + res = bind(node.eval, pre_loop.iterate) + for i in range(1, len(node.eval.size)): + loop = Loop(0, node.eval.size[i], 1, []) + pre_loop.body.append(loop) + pre_loop = loop + val = bind(val, pre_loop.iterate) + res = bind(res, pre_loop.iterate) + + assign = Assignment(res, Math(val, node.op_type)) + pre_loop.body.append(assign) + + else: + node.eval = Scalar(node.dtype) + node.decl = [Decl(node.eval)] + node.compute = [Assignment(node.eval, Math(node.operators[0].eval, node.op_type))] + + l = node.compute[0] + for i in range(len(node.eval.size)): + node.output_order.append((i, l)) + node.input_orders[0].append((i, l)) + l = l.body[0] + + elif node.op_type == 'setval': + node.operators[0]._gen_ir() + node.operators[1]._gen_ir() + node.input_orders[0] = [] + # node.operators[1] must be a Scalar, so no input_order is needed + + node.eval = node.operators[0].eval + node.decl = node.operators[0].decl[:] + node.operators[0].decl.clear() + val = node.operators[1].eval + + if len(node.ref_size) > 0: + size = helpers.get_ir_of_size(node.ref_size) + pre_loop = Loop(0, size[0], 1, []) + node.compute = [pre_loop] + res = bind(node.eval, pre_loop.iterate) + for i in range(1, len(size)): + loop = Loop(0, size[i], 1, []) + pre_loop.body.append(loop) + pre_loop = loop + res = bind(res, pre_loop.iterate) + + assign = Assignment(res, val) + pre_loop.body.append(assign) + else: + node.compute = [Assignment(node.eval, val)] + + l = node.compute[0] + for i in range(len(node.eval.size)): + node.output_order.append((i, l)) + node.input_orders[0].append((i, l)) + l = l.body[0] + elif node.op_type == 'einsum': node.operators[0]._gen_ir() node.operators[1]._gen_ir() + node.input_orders[0] = [] + node.input_orders[1] = [] + exp = node.operators[2] inputs, output = exp.split('->') input1, input2 = inputs.split(',') all_indices = ''.join(sorted(set(input1 + input2))) all_loops = [] + mapping = {} + for i in range(len(all_indices)): + pos1 = input1.find(all_indices[i]) + pos2 = input2.find(all_indices[i]) + if (pos1 >= 0 and pos2 < 0): + mapping[all_indices[i]] = len(all_loops) + l = Loop(0, node.operators[0].eval.size[pos1], 1, []) + all_loops.append(l) + node.input_orders[0].append((len(node.input_orders[0]), l)) + elif (pos1 < 0 and pos2 >= 0): + mapping[all_indices[i]] = len(all_loops) + l = Loop(0, node.operators[1].eval.size[pos2], 1, []) + all_loops.append(l) + node.input_orders[1].append((len(node.input_orders[1]), l)) + + reduce_begins = len(all_loops) + + for i in range(len(all_indices)): + pos1 = input1.find(all_indices[i]) + pos2 = input2.find(all_indices[i]) + if pos1 >= 0 and pos2 >= 0: + mapping[all_indices[i]] = len(all_loops) + l = Loop(0, node.operators[0].eval.size[pos1], 1, []) + all_loops.append(l) + node.input_orders[0].append((len(node.input_orders[0]), l)) + node.input_orders[1].append((len(node.input_orders[1]), l)) + for i in all_indices: pos1 = input1.find(i) - if pos1 >= 0: - all_loops.append(Loop(0, node.operators[0].eval.size[pos1], 1, [])) - else: - pos2 = input2.find(i) - if pos2 >= 0: - all_loops.append(Loop(0, node.operators[1].eval.size[pos2], 1, [])) - else: - raise IndexError('index not found!') + pos2 = input2.find(i) + if pos1 < 0 and pos2 < 0: + raise IndexError('index not found!') op1 = node.operators[0].eval for i in input1: - idx = all_indices.find(i) - op1 = bind(op1, all_loops[idx].iterate, ) + op1 = bind(op1, all_loops[mapping[i]].iterate) op2 = node.operators[1].eval for i in input2: - idx = all_indices.find(i) - op2 = bind(op2, all_loops[idx].iterate) + op2 = bind(op2, all_loops[mapping[i]].iterate) size = helpers.get_ir_of_size(node._size()) - node.eval = Ndarray(node.dtype, size, val=0) + node.eval = Ndarray(node.dtype, size) node.decl = [Decl(node.eval)] res = node.eval for i in output: - idx = all_indices.find(i) - res = bind(res, all_loops[idx].iterate) + res = bind(res, all_loops[mapping[i]].iterate) body = Assignment(res, Expr(op1, op2, '*'), '+') + init = Assignment(res, 0) + if reduce_begins == 0: + node.compute.append(init) pre_loop = all_loops[0] - node.compute = [pre_loop] + node.compute.append(pre_loop) for i in range(1, len(all_loops)): + if reduce_begins == i: + pre_loop.body.append(init) loop = all_loops[i] pre_loop.body.append(loop) pre_loop = loop pre_loop.body.append(body) + l = node.compute[0] + for i in range(len(node.eval.size)): + node.output_order.append((i, l)) + l = l.body[0] + elif node.op_type == 'index': node.operators[0]._gen_ir() node.operators[1]._gen_ir() - if type(node.operators[1]) == Var or (type(node.operators[1]) == Const and node.operators[1].dtype == 'int'): - node.eval = Index(node.operators[0].eval, index=node.operators[1].eval) - else: # ind_arr can be a slice or Tensor - # TODO: asssert tensor is 1d int? - node.eval = Index(node.operators[0].eval, ind_arr=node.operators[1].eval) + if type(node.operators[1].eval) in (Scalar, Literal, Indexing): + node.eval = Indexing(node.operators[0].eval, node.operators[1].eval) + elif type(node.operators[1].eval) in (Ndarray, Slice): + node.eval = Indexing(node.operators[0].eval, Indexing(node.operators[1].eval, Literal(-1, 'int'))) + else: + raise TypeError('incorrect index type!') elif node.op_type == 'apply': + #TODO: add input_orders for apply, reduce, and aggr - node.operators[0]._gen_ir() - node.operators[2]._gen_ir() + node.operators[0]._gen_ir() # input tensor + node.operators[2]._gen_ir() # axis - axis = node.operators[2].eval + axis = node.operators[2].eval.val outer_loop = Loop(0, node.operators[0].eval.size[axis], 1, []) + + # item is an indexing to the input tensor in axis dimension item = node.operators[3] item.eval = node.operators[0].eval for i in range(axis): - item.eval = bind(item.eval, None) - item.eval = bind(item.eval, outer_loop.iterate, True) + item.eval = Indexing(item.eval, Literal(-1, 'int')) + item.eval = Indexing(item.eval, outer_loop.iterate) - item.decl = [] - ret = node.operators[1](item) - ret._gen_ir() + ret = node.operators[-1] + ret._gen_ir() # generate IR for applied func def action(node, res): - if type(node) == Var or type(node) == One or type(node) == Zero or type(node) == Ones or type(node) == Zeros or type(node) == Tensor: - res.extend(node.decl) - node.decl.clear() - elif type(node) == TensorOp: - res.extend(node.decl) - res.extend(node.compute) - node.decl.clear() - node.compute.clear() + if node.valid == True: + if type(node) == Var or type(node) == Tensor: + res.extend(node.decl) + node.valid = False + elif type(node) == TensorOp: + res.extend(node.decl) + res.extend(node.compute) + node.valid = False t = helpers.Traversal(action) ret_ir = t(ret) @@ -185,48 +338,27 @@ def action(node, res): else: ret_compute.append(ir) - node.operators.append(ret) - node.dtype = ret.dtype - node.ref_size = [node._size()[0]] + ret._size() - node.fix_size = [] outer_loop.body.extend(ret_compute) size = helpers.get_ir_of_size(node._size()) node.eval = Ndarray(ret.eval.dtype, size) node.decl.append(Decl(node.eval)) node.decl.extend(ret_decl) + node.compute = [outer_loop] - # node.eval <= ret.eval res = bind(node.eval, outer_loop.iterate) - if (len(ret.eval.size) > 0): - pre_loop = Loop(0, ret.eval.size[0], 1, []) - outer_loop.body.append(pre_loop) - res = bind(res, pre_loop.iterate) - rhs = bind(ret.eval, pre_loop.iterate) - for i in range(1, len(ret.eval.size)): - loop = Loop(0, ret.eval.size[i], 1, []) - pre_loop.body.append(loop) - pre_loop = loop - res = bind(res, pre_loop.iterate) - rhs = bind(rhs, pre_loop.iterate) - pre_loop.body.append(Assignment(res, rhs)) - else: - assign = Assignment(res, ret.eval) - outer_loop.body.append(assign) + replace_output(node.compute, ret.eval, res) + node.decl = [d for d in node.decl if d.dobject != ret.eval] - scope = outer_loop.body - while len(scope) == 2: - fuse(scope, scope[0], scope[1]) - if type(scope[0]) is not Loop: - break - scope = scope[0].body + # TODO: need test for this + node.output_order = [(0, outer_loop)] + for i in range(len(ret.output_order)): + node.output_order.append((i+1, ret.output_order[i][1])) - node.compute = [outer_loop] elif node.op_type == 'reduce': node.operators[0]._gen_ir() - node.operators[2]._gen_ir() # init node.operators[3]._gen_ir() - axis = node.operators[3].eval + axis = node.operators[3].eval.val size = helpers.get_ir_of_size(node._size()) if len(size) > 0: @@ -235,52 +367,38 @@ def action(node, res): node.eval = Scalar(node.dtype) node.decl.append(Decl(node.eval)) - node.compute = [] - if len(node.eval.size) > 0: - pre_loop = Loop(0, node.eval.size[0], 1, []) - node.compute.append(pre_loop) - res = bind(node.eval, pre_loop.iterate) - rhs = bind(node.operators[2].eval, pre_loop.iterate) - for i in range(1, len(node.eval.size)): - loop = Loop(0, node.eval.size[i], 1, []) - pre_loop.body.append(loop) - pre_loop = loop - res = bind(res, pre_loop.iterate) - rhs = bind(rhs, pre_loop.iterate) - pre_loop.body.append(Assignment(res, rhs)) - else: - assign = Assignment(node.eval, node.operators[2].eval) - node.compute.append(assign) + node.operators[2]._gen_ir() # init + # node.compute.extend(node.operators[2].compute) + # node.decl.extend(node.operators[2].decl) + # node.operators[2].valid = False + # TODO: iterating over the reduction dimension in the outer loop may not give best performance + # TODO: it might be better to make it the innermost loop outer_loop = Loop(0, node.operators[0].eval.size[axis], 1, []) item1 = node.operators[4] item2 = node.operators[5] item1.eval = node.eval item2.eval = node.operators[0].eval - for i in range(axis): # TODO: is this correct? - item2.eval = bind(item2.eval, None) - if type(item2.eval) == Index and type(item2.eval.ind_arr) == Slice: - item2.eval = bind(item2.eval, outer_loop.iterate) - else: - item2.eval = bind(item2.eval, outer_loop.iterate, True) + for i in range(axis): + item2.eval = Indexing(item2.eval, Literal(-1, 'int')) + item2.eval = Indexing(item2.eval, outer_loop.iterate) item2.decl = [] item1.decl = [] - ret = node.operators[1](item1, item2) + ret = node.operators[-1] ret._gen_ir() def action(node, res): - if type(node) == Var or type(node) == One or type(node) == Zero or type(node) == Ones or type(node) == Zeros or type(node) == Tensor: - res.extend(node.decl) - node.decl.clear() - elif type(node) == TensorOp: - res.extend(node.decl) - res.extend(node.compute) - node.decl.clear() - node.compute.clear() - + if node.valid == True: + if type(node) == Var or type(node) == Tensor: + res.extend(node.decl) + node.valid = False + elif type(node) == TensorOp: + res.extend(node.decl) + res.extend(node.compute) + node.valid = False t = helpers.Traversal(action) ret_ir = t(ret) @@ -293,79 +411,56 @@ def action(node, res): else: ret_compute.append(ir) - node.operators.append(ret) outer_loop.body.extend(ret_compute) node.decl.extend(ret_decl) + node.compute.append(outer_loop) - if (len(ret.eval.size) > 0): - pre_loop = Loop(0, ret.eval.size[0], 1, []) - outer_loop.body.append(pre_loop) - res = bind(node.eval, pre_loop.iterate) - rhs = bind(ret.eval, pre_loop.iterate) - for i in range(1, len(ret.eval.size)): - loop = Loop(0, ret.eval.size[i], 1, []) - pre_loop.body.append(loop) - pre_loop = loop - res = bind(res, pre_loop.iterate) - rhs = bind(rhs, pre_loop.iterate) - pre_loop.body.append(Assignment(res, rhs)) - else: - assign = Assignment(node.eval, ret.eval) - outer_loop.body.append(assign) + replace_output(node.compute, ret.eval, node.eval) + node.decl = [d for d in node.decl if d.dobject != ret.eval] + + node.output_order = ret.output_order - node.compute.append(outer_loop) elif node.op_type == 'aggr': node.operators[0]._gen_ir() # input tensor - node.operators[2]._gen_ir() # init node.operators[3]._gen_ir() # indices node.operators[4]._gen_ir() # axis - axis = node.operators[4].eval + axis = node.operators[4].eval.val size = helpers.get_ir_of_size(node._size()) node.eval = Ndarray(node.dtype, size) node.decl.append(Decl(node.eval)) + # this must be called after node.eval is constructed + node.operators[2]._gen_ir() # init - node.compute = [] - # initialize output - pre_loop = Loop(0, node.eval.size[0], 1, []) - node.compute.append(pre_loop) - res = bind(node.eval, pre_loop.iterate) - rhs = node.operators[2].eval - for i in range(1, len(node.eval.size)): - loop = Loop(0, node.eval.size[i], 1, []) - pre_loop.body.append(loop) - pre_loop = loop - res = bind(res, pre_loop.iterate) - rhs = bind(rhs, pre_loop.iterate) - pre_loop.body.append(Assignment(res, rhs)) + # node.compute.extend(node.operators[2].compute) + # node.decl.extend(node.operators[2].decl) + # node.operators[2].valid = False # compute outer_loop = Loop(0, node.operators[0].eval.size[axis], 1, []) item1 = node.operators[6] item2 = node.operators[7] - item1.eval = Index(node.eval, ind_arr=node.operators[3].eval) - item1.eval = bind(item1.eval, outer_loop.iterate) + item1.eval = Indexing(node.eval, Indexing(node.operators[3].eval, outer_loop.iterate)) item2.eval = node.operators[0].eval for i in range(axis): - item2.eval = bind(item2.eval, None) - item2.eval = bind(item2.eval, outer_loop.iterate, True) + item2.eval = Indexing(item2.eval, Literal(-1, 'int')) + item2.eval = Indexing(item2.eval, outer_loop.iterate) item2.decl = [] item1.decl = [] - ret = node.operators[1](item1, item2) + ret = node.operators[-1] ret._gen_ir() def action(node, res): - if type(node) == Var or type(node) == One or type(node) == Zero or type(node) == Ones or type(node) == Zeros or type(node) == Tensor: - res.extend(node.decl) - node.decl.clear() - elif type(node) == TensorOp: - res.extend(node.decl) - res.extend(node.compute) - node.decl.clear() - node.compute.clear() - + if node.valid == True: + if type(node) == Var or type(node) == Tensor: + res.extend(node.decl) + node.valid = False + elif type(node) == TensorOp: + res.extend(node.decl) + res.extend(node.compute) + node.valid = False t = helpers.Traversal(action) ret_ir = t(ret) @@ -378,29 +473,24 @@ def action(node, res): else: ret_compute.append(ir) - node.operators.append(ret) outer_loop.body.extend(ret_compute) node.decl.extend(ret_decl) + node.compute.append(outer_loop) - if (len(ret.eval.size) > 0): - pre_loop = Loop(0, ret.eval.size[0], 1, []) - outer_loop.body.append(pre_loop) - res = bind(Index(node.eval, ind_arr=node.operators[3].eval), outer_loop.iterate) - res = bind(res, pre_loop.iterate) - rhs = bind(ret.eval, pre_loop.iterate) - for i in range(1, len(ret.eval.size)): - loop = Loop(0, ret.eval.size[i], 1, []) - pre_loop.body.append(loop) - pre_loop = loop - res = bind(res, pre_loop.iterate) - rhs = bind(rhs, pre_loop.iterate) - pre_loop.body.append(Assignment(res, rhs)) - else: - res = bind(Index(node.eval, ind_arr=node.operators[3].eval), outer_loop.iterate) - assign = Assignment(res, ret.eval) - outer_loop.body.append(assign) + replace_output(node.compute, ret.eval, item1.eval) + node.decl = [d for d in node.decl if d.dobject != ret.eval] - node.compute.append(outer_loop) - return node + node.output_order = [(0, outer_loop)] + for i in range(len(ret.output_order)): + node.output_order.append((i+1, ret.output_order[i][1])) + # points from IR back to ASTNode + for d in node.decl: + d.astnode = node + + if type(node) == TensorOp: + for s in node.compute: + s.astnode = node + + return node diff --git a/core/ir.py b/core/ir.py index 4963406..387ba4a 100644 --- a/core/ir.py +++ b/core/ir.py @@ -1,17 +1,22 @@ class IR: - pass + def __init__(self): + # astnode tracks the location of this IR in the AST + self.astnode = None class DOject(IR): nobjects = 0 - def __init__(self, dtype: str): + def __init__(self, dtype: str, size: (list, tuple)): + super().__init__() self.dobject_id = DOject.nobjects DOject.nobjects += 1 self.dtype = dtype + self.size = size class Expr(IR): def __init__(self, left, right, op: str): + super().__init__() self.left = left self.right = right self.op = op @@ -20,6 +25,7 @@ def __init__(self, left, right, op: str): class Assignment(IR): def __init__(self, lhs, rhs, op=None): + super().__init__() self.lhs = lhs self.rhs = rhs self.op = op @@ -30,6 +36,7 @@ class Loop(IR): loop_id = 0 def __init__(self, start, end, step, body: list): + super().__init__() self.lid = Loop.loop_id Loop.loop_id += 1 self.start = start @@ -40,34 +47,34 @@ def __init__(self, start, end, step, body: list): class Scalar(DOject): - def __init__(self, dtype: str, name: str = None, is_arg = False, val = None): - super().__init__(dtype) + def __init__(self, dtype: str, name: str = None, is_arg = False): + super().__init__(dtype, []) self.__name__ = name if name else f's{self.dobject_id}' - self.size = [] - self.val = val self.is_arg = is_arg - - def name(self): return self.__name__ - def addr(self): - return self.name() + +class Literal(DOject): + def __init__(self, val: (int, float), dtype: str): + super().__init__(dtype, []) + self.val = val class Slice(IR): def __init__(self, start, stop, step): + super().__init__() self.start = start self.stop = stop self.step = step + self.dtype = 'int' + self.size = [Expr(Expr(self.stop, self.start, '-'), self.step, '/')] class Ndarray(DOject): - def __init__(self, dtype: str, size: tuple, name: str = None, is_arg = False, val = None): - super().__init__(dtype) - self.size = size + def __init__(self, dtype: str, size: tuple, name: str = None, is_arg = False): + super().__init__(dtype, size) self.__name__ = name if name else f'arr{self.dobject_id}' - self.val = val # val is None, 0, or 1 self.is_arg = is_arg def __getitem__(self, item): @@ -76,40 +83,38 @@ def __getitem__(self, item): def name(self): return self.__name__ - def addr(self): - return self.name() +class Math(IR): + def __init__(self, val, type): + self.val = val + self.type = type -class Index(IR): - nindices = 0 - def __init__(self, dobject, index=None, ind_arr=None): +class Indexing(DOject): + def __init__(self, dobject, idx): + assert dobject != None and type(dobject) in (Slice, Ndarray, Indexing) + assert idx != None and type(idx) in (Scalar, Literal, Indexing) self.dobject = dobject - self.index = index - self.ind_arr = ind_arr - self.dtype = self.dobject.dtype - if ind_arr == None: - self.size = dobject.size[1:] - elif type(ind_arr) == Ndarray: - self.size = ind_arr.size + dobject.size[1:] - elif type(ind_arr) == Slice: - s = Expr(Expr(ind_arr.stop, ind_arr.start, '-'), ind_arr.step, '/') - self.size = [s] + dobject.size[1:] - self.index_id = Index.nindices - Index.nindices += 1 - - - def name(self): - return f'ref{self.index_id}_{self.dobject.name()}' - - def addr(self): - if self.ind_arr: - return f'{self.dobject}[{self.ind_arr[0]}]' + self.idx = idx + + if type(self.dobject) in (Ndarray, Slice): + if type(idx) == Literal and idx.val == -1: + # idx is unspecified, which means the Indexing is a range of indice stored in dobject, so the size of Indexing should the same as the dobject + size = dobject.size[:] + self.ref_point = 1 + else: + # idx is a specific Scalar, Literal, or Indexing, in any case, the size of the Indexing operation should be as follows + # ref_point should be the next dimension if the node is further Indexed + size = idx.size + dobject.size[1:] + self.ref_point = len(idx.size) else: - return f'{self.dobject}[0]' - + # dobject is an Indexing + size = dobject.size[:dobject.ref_point] + idx.size + dobject.size[dobject.ref_point+1:] + self.ref_point = dobject.ref_point + len(idx.size) + super().__init__(dobject.dtype, size) class Decl(IR): - def __init__(self, dobject): + def __init__(self, dobject: (Scalar, Ndarray)): + super().__init__() self.dobject = dobject \ No newline at end of file diff --git a/core/test/examples.py b/core/test/examples.py index 2302794..75687e3 100644 --- a/core/test/examples.py +++ b/core/test/examples.py @@ -13,6 +13,7 @@ def func(): ast = func() code = codegen.cpu.print_cpp(ast._gen_ir()) + print(code) A = torch.rand(10, 10) B = torch.rand(10, 10) @@ -78,11 +79,6 @@ def test4(): d = run.cpu.compile_and_run(code, A, i, t) print(d, A[i] + t) -def f5(): - A = Tensor('a', (10, )) - t = Var('t', A.dtype) - - return A[0] + t def test6(): A = Tensor('a', (10, )) @@ -348,20 +344,80 @@ def test17(): d = run.cpu.compile_and_run(code, A, B) print(A[1:10][:, 2:4] + B[1:10][:, 2:4]) + print(d) print(torch.equal(A[1:10][:, 2:4] + B[1:10][:, 2:4], d)) +def test18(): + A = Tensor('A', (100, 20)) + B = Tensor('B', (100, ), dtype='int') + + ast = A[1:10][B[2:4]] + A[1:10][B[1:3]] + print(helpers.get_input_nodes(ast)) + ir = gen_ir(ast) + + code = codegen.cpu.print_cpp(ir) + A = torch.rand(100, 20) + B = torch.randint(0, 20, (100, )).to(torch.int32) + d = run.cpu.compile_and_run(code, A, B) + + print(A[1:10][:, B[2:4]] + A[1:10][:, B[1:3]]) + print(d) + print(torch.equal(A[1:10][:, B[2:4]] + A[1:10][:, B[1:3]], d)) + +def test19(): + nnodes = 100 + nedges = 300 + rowptr = Tensor('rowptr', (nnodes + 1, ), dtype='int') + colidx = Tensor('colidx', (nedges, ), dtype='int') + edge_list = Tensor('edge_list', (10, 2), dtype='int') + ast = colidx[rowptr[edge_list[0][0]]:rowptr[edge_list[0][1]]] + colidx[rowptr[edge_list[0][0]]:rowptr[edge_list[0][1]]] + + print(helpers.get_input_nodes(ast)) + code = codegen.cpu.print_cpp(gen_ir(ast)) + print(code) + + +def test20(): + A = Tensor('A', (100, ), dtype='float') + x = Var('x', dtype='float') + res = A + 10 + code = codegen.cpu.print_cpp(gen_ir(res)) + print(code) + + +def compression(): + input = Tensor('input', (50, 32), dtype='float') + res = (input * 1000).round() + res = res.apply(lambda x:x[0:32]-x[-1:31], axis=0) + res = res.abs().max(axis=1) + code = codegen.cpu.print_cpp(gen_ir(res)) + print(code) + + +def test_math1(): + input = Tensor('input', (50, 32), dtype='float') + res = input[0].abs() + code = codegen.cpu.print_cpp(gen_ir(res)) + print(code) + +def test_math2(): + input = Tensor('input', (50, 32), dtype='float') + input = input.setval(0) + res = input[0].abs() + code = codegen.cpu.print_cpp(gen_ir(res)) + print(code) + + + def apply_test1(): - num_node = 10 num_edges = 20 - max_degree = 20 - rowptr = Tensor('rowptr', (num_node+1,), dtype='int') + length = 50 rowidx = Tensor('rowidx', (num_edges,), dtype='int') colidx = Tensor('colidx', (num_edges,), dtype='int') - edge_idx = Tensor('edge_idx', (num_edges,), dtype='int') + edge_idx = Tensor('edge_idx', (length,), dtype='int') - v0 = rowidx[0] def apply_func(edge_id): v0 = rowidx[edge_id] @@ -370,7 +426,28 @@ def apply_func(edge_id): res = edge_idx.apply(apply_func) code = codegen.cpu.print_cpp(gen_ir(res)) - print(code) + print(helpers.get_input_nodes(res)) + + edge_idx = torch.randint(0, num_edges, (length,)).to(torch.int32) + rowidx = torch.randint(0, 1000, (num_edges,)).to(torch.int32) + colidx = torch.randint(0, 1000, (num_edges,)).to(torch.int32) + + d = run.cpu.compile_and_run(code, edge_idx, rowidx, colidx) + + res = torch.zeros_like(edge_idx) + for i in range(len(edge_idx)): + e = edge_idx[i] + v0 = rowidx[e] + v1 = colidx[e] + 1 + res[i] = v0 + v1 + + print(d) + print(res) + print(torch.equal(d, res)) + + + + def apply_test2(): @@ -396,7 +473,6 @@ def apply_test3(): def apply_func(item): def apply_func2(item2): - print(type(item2)) return C[item2] + B[item2] return item.apply(apply_func2) @@ -419,8 +495,7 @@ def apply_test4(): def apply_func(item): def apply_func2(item2): - print(is_int_var(item2)) - return B[item2] + return B[item2] + 1 return item.apply(apply_func2) @@ -432,11 +507,10 @@ def apply_func2(item2): - def test_aggr1(): A = Tensor('A', (10, 20)) - indices = Tensor('idx', (30, ), dtype='int') - res = A.aggr_sum(indices) + indices = Tensor('idx', (A._size()[0], ), dtype='int') + res = A.aggr_max(indices) code = codegen.cpu.print_cpp(res._gen_ir()) print(code) @@ -656,10 +730,9 @@ def test26(): res = run.cpu.compile_and_run(code, A) print(res, torch.sum(A)) -def test27(): +def reduce_test1(): A = Tensor('a', (10, 20)) - init = Zeros(A[1]._size()) - ast = A.reduce(lambda a,b: a+b, init, axis=1) + ast = A.reduce(lambda a,b: a+b, lambda res: res.setval(0), axis=1) ir = gen_ir(ast) print(helpers.get_input_nodes(ir)) code = codegen.cpu.print_cpp(ir) @@ -669,7 +742,7 @@ def test27(): res = run.cpu.compile_and_run(code, A) print(res - torch.sum(A, dim=1)) -def test28(): +def reduce_test2(): A = Tensor('a', (10, 20, 5)) ast = A.sum(axis=1) ir = gen_ir(ast) @@ -680,6 +753,30 @@ def test28(): res = run.cpu.compile_and_run(code, A) print(res - torch.sum(A, dim=1)) + +def reduce_test3(): + A = Tensor('a', (10, 20, 5)) + ast = A.max(axis=1) + ir = gen_ir(ast) + print(helpers.get_input_nodes(ir)) + code = codegen.cpu.print_cpp(ir) + + A = torch.rand(10, 20, 5) + res = run.cpu.compile_and_run(code, A) + print(res - torch.max(A, dim=1).values) + + +def reduce_test4(): + A = Tensor('a', (10,)) + ast = A.max() + ir = gen_ir(ast) + print(helpers.get_input_nodes(ir)) + code = codegen.cpu.print_cpp(ir) + + A = torch.rand(10, ) + res = run.cpu.compile_and_run(code, A) + print(res - torch.max(A)) + def test29(): A = Tensor('A', (10, )) B = Tensor('B', (10, ), dtype='int') @@ -689,11 +786,72 @@ def test29(): print(code) +def conv1d_v1(): + A = Tensor('a', (100, )) + ast = A[0:97] + A[1:98] + A[2:99] + ir = gen_ir(ast) + print(helpers.get_input_nodes(ir)) + code = codegen.cpu.print_cpp(ir) + print(code) + +def conv1d_v2(width): + A = Tensor('a', (100, )) + res = Tensor('t', A[width:]._size()).setval(0) + for i in range(width): + res = res + A[i:i+97] + ir = gen_ir(res) + print(helpers.get_input_nodes(ir)) + code = codegen.cpu.print_cpp(ir) + print(code) + +def cmp_test(): + A = Tensor('a', (10, 20)) + B = Tensor('b', (10, 20)) + res = bigger(A, B) + + code = codegen.cpu.print_cpp(res._gen_ir()) + print(code) + + if __name__ == "__main__": + # conv1d_v1() + # conv1d_v2(3) # test1() + # test2() + # test3() + # test4() + # test6() + # test7() + # test8() + # test9() + # test10() + # test11() + # test12() + # test13() + # test14() + # test15() + # test16() # test17() + # test18() + # test19() + # test20() + # compression() + # test_math1() + # test_math2() + apply_test1() + apply_test2() + # apply_test3() + # apply_test4() + # reduce_test1() + # reduce_test2() + # reduce_test3() + # reduce_test4() + # test_aggr1() # spmv() # test_einsum1() - test_apply5() - # test27() \ No newline at end of file + # apply_test2() + # test_apply5() + # test27() + + # cmp_test() \ No newline at end of file diff --git a/helpers.py b/helpers.py index 40d9fce..6079b42 100644 --- a/helpers.py +++ b/helpers.py @@ -15,23 +15,19 @@ def _post_traverse(self, node, visited, res): else: visited.add(node) - if type(node) == Var or type(node) == One or type(node) == Zero: + if type(node) == Var: self.action(node, res) elif type(node) == Const: if node.dtype == 'slice': self._post_traverse(node.val.start, visited, res) self._post_traverse(node.val.stop, visited, res) self._post_traverse(node.val.step, visited, res) - elif type(node) == Tensor or type(node) == Ones or type(node) == Zeros: + elif type(node) == Tensor: for s in node.fix_size: self._post_traverse(s, visited, res) for s in node.ref_size: self._post_traverse(s, visited, res) self.action(node, res) - # elif type(node) == Set: - # self._post_traverse(node.storage, visited, res) - # self._post_traverse(node.nelem, visited, res) - # self.action(node, res) elif type(node) == TensorOp: for s in node.fix_size: self._post_traverse(s, visited, res) @@ -46,6 +42,18 @@ def _post_traverse(self, node, visited, res): for c in node.operators: self._post_traverse(c, visited, res) self.action(node, res) + elif type(node) == Set: + self._post_traverse(node.storage, visited, res) + for n in node.nelem: + self._post_traverse(n, visited, res) + self.action(node, res) + elif type(node) == SetOp: + self._post_traverse(node.storage, visited, res) + for n in node.nelem: + self._post_traverse(n, visited, res) + for c in node.operators: + self._post_traverse(c, visited, res) + self.action(node, res) def __call__(self, ast): visited = set() diff --git a/opt/loop.py b/opt/loop.py index 5eaf746..5baf1f8 100644 --- a/opt/loop.py +++ b/opt/loop.py @@ -17,16 +17,15 @@ def rebind_iterate(ir, old, new): rebind_iterate(ir.rhs, old, new) elif type(ir) == Ndarray: rebind_iterate(ir.size, old, new) - elif type(ir) == Ref: + # elif type(ir) == Ref: + # rebind_iterate(ir.dobject, old, new) + elif type(ir) == Indexing: rebind_iterate(ir.dobject, old, new) - elif type(ir) == Index: - rebind_iterate(ir.dobject, old, new) - rebind_iterate(ir.ind_arr, old, new) - if type(ir.index) == Scalar: - if ir.index == old: - ir.index = new + if type(ir.idx) in (Scalar, Literal): + if ir.idx == old: + ir.idx = new else: - rebind_iterate(ir.index, old, new) + rebind_iterate(ir.idx, old, new) elif type(ir) == Slice: rebind_iterate(ir.start, old, new) rebind_iterate(ir.stop, old, new) diff --git a/run/.tmp/cpu_code.cpp b/run/.tmp/cpu_code.cpp index c649bf4..3512403 100644 --- a/run/.tmp/cpu_code.cpp +++ b/run/.tmp/cpu_code.cpp @@ -1,29 +1,22 @@ #include -torch::Tensor reduce_a__zeros_0_c3_item1_of_a_item2_of_a(torch::Tensor obj_a) +torch::Tensor reduce_a___c2_item1_of_a_item2_of_a_add_item1_of_a_item2_of_a(torch::Tensor obj_a) { auto a = obj_a.accessor(); -torch::Tensor obj_zeros_0 = torch::zeros({20}, at::kFloat); -auto zeros_0 = obj_zeros_0.accessor(); -torch::Tensor obj_arr2 = torch::empty({10}, at::kFloat); -auto arr2 = obj_arr2.accessor(); -torch::Tensor obj_arr5 = torch::empty({10}, at::kFloat); -auto arr5 = obj_arr5.accessor(); +torch::Tensor obj_arr4 = torch::empty({10}, at::kFloat); +auto arr4 = obj_arr4.accessor(); for (int _l0 = 0; _l0 < 10; _l0 += 1) { -arr2[_l0] = zeros_0[_l0]; +arr4[_l0] = 0; } for (int _l1 = 0; _l1 < 20; _l1 += 1) { for (int _l2 = 0; _l2 < 10; _l2 += 1) { -arr5[_l2] = (arr2[_l2] + a[_l2][_l1]); +arr4[_l2] = (arr4[_l2] + a[_l2][_l1]); } -for (int _l3 = 0; _l3 < 10; _l3 += 1) { -arr2[_l3] = arr5[_l3]; } -} -return obj_arr2; +return obj_arr4; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("run", &reduce_a__zeros_0_c3_item1_of_a_item2_of_a); + m.def("run", &reduce_a___c2_item1_of_a_item2_of_a_add_item1_of_a_item2_of_a); } \ No newline at end of file diff --git a/run/__init__.py b/run/__init__.py index 2e42cc2..b07f703 100644 --- a/run/__init__.py +++ b/run/__init__.py @@ -1 +1 @@ -from run import cpu \ No newline at end of file +from run import cpu, gpu \ No newline at end of file diff --git a/run/gpu.py b/run/gpu.py new file mode 100644 index 0000000..eeb7e44 --- /dev/null +++ b/run/gpu.py @@ -0,0 +1,8 @@ +from torch.utils.cpp_extension import load + +def compile_and_run(code, *args): + f = open('run/.tmp/cuda_code.cu', 'w') + f.write(code) + f.close() + module = load(name='module', sources=['run/.tmp/cuda_code.cu']) + return module.run(*args) \ No newline at end of file diff --git a/test/examples.py b/test/examples.py index 08dc911..d5718fb 100644 --- a/test/examples.py +++ b/test/examples.py @@ -128,7 +128,17 @@ def f36(): return A.num_elem() + Set(T[40:]).num_elem() + Set(T[1:d]).num_elem() +def compression(): + input = Tensor('input', (1024, )) + input = input.reshape((32, 33)) + x = (input * 1000).round() + y = x.apply(lambda y:y[1:31]-y[0:30], axis=0) + res = y.encoding() + ir = gen_ir(res) + code = codegen.cpu.print_cpp(ir) + code = codegen.cerebra.print_code(ir) + print(code) def conv1d_v1(): @@ -149,5 +159,224 @@ def conv1d_v2(width): code = codegen.cpu.print_cpp(ir) print(code) +import numpy as np +from cset.ast2ir import * +def subgraph_matching_test_code(): + + pattern_size = 4 + + num_node = 10 + num_edges = 20 + + rowptr = Tensor('rowptr', (num_node+1,), dtype='int') + colidx = Tensor('colidx', (num_edges,), dtype='int') + edge_list = Set(Tensor('edge_list', (num_edges, 2), dtype='int')) + + count = Zero(dtype='int') + + class inner_subgraph_matching: + def __init__(self, level, *path): + self.level = level + self.path = list(*path) + + def __call__(self, item): + + if self.level == pattern_size-1: + return count+1 + + if self.level==1: + v0_nb = Set(colidx[rowptr[item[0]]:rowptr[item[0]+1]]) + v1_nb = Set(colidx[rowptr[item[1]]:rowptr[item[1]+1]]) + candidate_set = v0_nb.intersect(v1_nb) + return candidate_set.applyfunc(inner_subgraph_matching(self.level+1, [item[0], item[1]])) + else: + candidate_set = Set(colidx[rowptr[item]:rowptr[item+1]]) + candidate_set = candidate_set.filter(SmallerThan(self.path[-1])) + for v in self.path: + v_nb = Set(colidx[rowptr[v]:rowptr[v+1]]) + candidate_set = candidate_set.intersect(v_nb) + + return candidate_set.applyfunc(inner_subgraph_matching(self.level+1, self.path + [item])) + + res = edge_list.applyfunc(inner_subgraph_matching(1)) + code = codegen.cpu.print_cpp(res._gen_ir()) + print(code) +def triangle_counting(): + np_rowptr = np.fromfile("../MiCo/snap.txt.vertex.bin", dtype=np.int64) + np_colidx = np.fromfile("../MiCo/snap.txt.edge.bin", dtype=np.int32) + + torch_rowptr = torch.from_numpy(np_rowptr, ).to(torch.int32) + torch_colidx =torch.from_numpy(np_colidx) + torch_edge_list = torch.zeros([torch_colidx.shape[0], 2], dtype=torch.int32) + + edge_idx = 0 + for i in range(0, torch_rowptr.shape[0]-1): + for j in range(torch_rowptr[i].item() , torch_rowptr[i+1].item()): + if(torch_colidx[j] + +torch::Tensor sub_add_sub_index_Eemb_h_index_Eemb_t_index_Remb_r_scal_mul_vec_vec_mul_vec_index_Remb_r_sub_index_Eemb_h_index_Eemb_t_index_Remb_r(int batch_size, int dim, int nnodes, torch::Tensor obj_Eemb, torch::Tensor obj_h, torch::Tensor obj_t, int nedges, torch::Tensor obj_Remb, torch::Tensor obj_r) +{ + auto Eemb = obj_Eemb.accessor(); +auto h = obj_h.accessor(); +auto t = obj_t.accessor(); +torch::Tensor obj_arr6 = torch::empty({batch_size,dim}, at::kFloat); +auto arr6 = obj_arr6.accessor(); +for (int _l0 = 0; _l0 < batch_size; _l0 += 1) { +for (int _l1 = 0; _l1 < dim; _l1 += 1) { +arr6[_l0][_l1] = (Eemb[h[_l0]][_l1] - Eemb[t[_l0]][_l1]); +} +} +auto Remb = obj_Remb.accessor(); +auto r = obj_r.accessor(); +torch::Tensor obj_arr12 = torch::empty({batch_size,dim}, at::kFloat); +auto arr12 = obj_arr12.accessor(); +for (int _l2 = 0; _l2 < batch_size; _l2 += 1) { +for (int _l3 = 0; _l3 < dim; _l3 += 1) { +arr12[_l2][_l3] = (arr6[_l2][_l3] + Remb[r[_l2]][_l3]); +} +} +torch::Tensor obj_arr15 = torch::empty({batch_size,dim}, at::kFloat); +auto arr15 = obj_arr15.accessor(); +for (int _l4 = 0; _l4 < batch_size; _l4 += 1) { +for (int _l5 = 0; _l5 < dim; _l5 += 1) { +arr15[_l4][_l5] = (Eemb[h[_l4]][_l5] - Eemb[t[_l4]][_l5]); +} +} +torch::Tensor obj_arr18 = torch::empty({batch_size}, at::kFloat); +auto arr18 = obj_arr18.accessor(); +for (int _l6 = 0; _l6 < batch_size; _l6 += 1) { +for (int _l7 = 0; _l7 < dim; _l7 += 1) { +arr18[_l6] += (Remb[r[_l6]][_l7] * arr15[_l6][_l7]); +} +} +torch::Tensor obj_arr21 = torch::empty({batch_size,dim}, at::kFloat); +auto arr21 = obj_arr21.accessor(); +for (int _l8 = 0; _l8 < batch_size; _l8 += 1) { +for (int _l9 = 0; _l9 < dim; _l9 += 1) { +arr21[_l8][_l9] = (arr18[_l8] * Remb[r[_l8]][_l9]); +} +} +torch::Tensor obj_arr24 = torch::empty({batch_size,dim}, at::kFloat); +auto arr24 = obj_arr24.accessor(); +for (int _l10 = 0; _l10 < batch_size; _l10 += 1) { +for (int _l11 = 0; _l11 < dim; _l11 += 1) { +arr24[_l10][_l11] = (arr12[_l10][_l11] - arr21[_l10][_l11]); +} +} +return obj_arr24; + +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("run", &sub_add_sub_index_Eemb_h_index_Eemb_t_index_Remb_r_scal_mul_vec_vec_mul_vec_index_Remb_r_sub_index_Eemb_h_index_Eemb_t_index_Remb_r); +} +[] +arr6[_l0][_l1] = (Eemb[h[_l0]][_l1] - Eemb[t[_l0]][_l1]); + +[] +arr21[_l8][_l9] = (arr18[_l8] * Remb[r[_l8]][_l9]); + +[, ] +arr6[_l2][_l3] = (Eemb[h[_l2]][_l3] - Eemb[t[_l2]][_l3]); + +arr12[_l2][_l3] = (arr6[_l2][_l3] + Remb[r[_l2]][_l3]); + +#include + +torch::Tensor sub_add_sub_index_Eemb_h_index_Eemb_t_index_Remb_r_scal_mul_vec_vec_mul_vec_index_Remb_r_sub_index_Eemb_h_index_Eemb_t_index_Remb_r(int batch_size, int dim, int nnodes, torch::Tensor obj_Eemb, torch::Tensor obj_h, torch::Tensor obj_t, int nedges, torch::Tensor obj_Remb, torch::Tensor obj_r) +{ + auto Eemb = obj_Eemb.accessor(); +auto h = obj_h.accessor(); +auto t = obj_t.accessor(); +torch::Tensor obj_arr6 = torch::empty({batch_size,dim}, at::kFloat); +auto arr6 = obj_arr6.accessor(); +auto Remb = obj_Remb.accessor(); +auto r = obj_r.accessor(); +torch::Tensor obj_arr12 = torch::empty({batch_size,dim}, at::kFloat); +auto arr12 = obj_arr12.accessor(); +torch::Tensor obj_arr15 = torch::empty({batch_size,dim}, at::kFloat); +auto arr15 = obj_arr15.accessor(); +for (int _l4 = 0; _l4 < batch_size; _l4 += 1) { +for (int _l5 = 0; _l5 < dim; _l5 += 1) { +arr15[_l4][_l5] = (Eemb[h[_l4]][_l5] - Eemb[t[_l4]][_l5]); +} +} +torch::Tensor obj_arr18 = torch::empty({batch_size}, at::kFloat); +auto arr18 = obj_arr18.accessor(); +for (int _l6 = 0; _l6 < batch_size; _l6 += 1) { +for (int _l7 = 0; _l7 < dim; _l7 += 1) { +arr18[_l6] += (Remb[r[_l6]][_l7] * arr15[_l6][_l7]); +} +} +torch::Tensor obj_arr21 = torch::empty({batch_size,dim}, at::kFloat); +auto arr21 = obj_arr21.accessor(); +torch::Tensor obj_arr24 = torch::empty({batch_size,dim}, at::kFloat); +auto arr24 = obj_arr24.accessor(); +for (int _l10 = 0; _l10 < batch_size; _l10 += 1) { +for (int _l11 = 0; _l11 < dim; _l11 += 1) { +arr6[_l2][_l3] = (Eemb[h[_l2]][_l3] - Eemb[t[_l2]][_l3]); +arr12[_l10][_l11] = (arr6[_l10][_l11] + Remb[r[_l10]][_l11]); +arr21[_l10][_l11] = (arr18[_l10] * Remb[r[_l10]][_l11]); +arr24[_l10][_l11] = (arr12[_l10][_l11] - arr21[_l10][_l11]); +} +} +return obj_arr24; + +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("run", &sub_add_sub_index_Eemb_h_index_Eemb_t_index_Remb_r_scal_mul_vec_vec_mul_vec_index_Remb_r_sub_index_Eemb_h_index_Eemb_t_index_Remb_r); +}