ast_ref added

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
-

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
-

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