def copy_function(func, temper=None, module=None): """Copy a Function. `temper` may be given to""" temper = temper or make_temper() f = Function(func.name, list(func.argnames), func.type, temper=temper) valuemap = {} lookup = partial(_lookup, module or func.module, f, valuemap) ### Construct new Blocks for block in func.blocks: new_block = Block(temper(block.name), f) valuemap[block] = new_block f.add_block(new_block) ### Construct new Operations for block in func.blocks: new_block = valuemap[block] for op in block.ops: new_op = Op(op.opcode, op.type, nestedmap(lookup, op.args), result=temper(op.result), parent=new_block) # assert new_op.result != op.result valuemap[op] = new_op new_block.append(new_op) return f
def _build_function(type, args): restype, name, func_args = args[1] argtypes, argnames = func_args[::2], func_args[1::2] type = types.Function(restype, argtypes) f = Function(name, argnames, type) blocks = args[2:-2] # omit the footer and newline for name, ops in blocks: f.add_block(name).extend(ops) return f
def copy_function(func, temper=None, module=None): """Copy a Function. `temper` may be given to""" temper = temper or make_temper() f = Function(func.name, list(func.argnames), func.type, temper=temper) valuemap = {} lookup = partial(_lookup, module or func.module, f, valuemap) ### Construct new Blocks for block in func.blocks: new_block = Block(temper(block.name), f) valuemap[block] = new_block f.add_block(new_block) ### Construct new Operations for block in func.blocks: new_block = valuemap[block] for op in block.ops: if op.opcode == 'phi': # Phi nodes may be circular, or may simply precede some of # their arguments args = [] else: args = nestedmap(lookup, op.args) new_op = Op(op.opcode, op.type, args, result=temper(op.result), parent=new_block) new_op.add_metadata(op.metadata) # assert new_op.result != op.result valuemap[op] = new_op new_block.append(new_op) for old_op in func.ops: if old_op.opcode == 'phi': new_op = valuemap[old_op] new_op.set_args(nestedmap(lookup, old_op.args)) return f
class TestBuilder(unittest.TestCase): def setUp(self): self.f = Function("testfunc", ['a'], types.Function(types.Float32, [types.Int32])) self.b = Builder(self.f) self.b.position_at_end(self.f.add_block('entry')) self.a = self.f.get_arg('a') def test_basic_builder(self): v = self.b.alloca(types.Pointer(types.Float32), []) result = self.b.mul(types.Int32, [self.a, self.a], result='r') c = self.b.convert(types.Float32, [result]) self.b.store(c, v) val = self.b.load(types.Float32, [v]) self.b.ret(val) # print(string(self.f)) self.assertEqual(str(self.f).strip(), basic_expected) def test_splitblock(self): old, new = self.b.splitblock('newblock') with self.b.at_front(old): self.b.add(types.Int32, [self.a, self.a]) with self.b.at_end(new): self.b.div(types.Int32, [self.a, self.a]) # print(string(self.f)) self.assertEqual(split_expected, string(self.f)) def test_loop_builder(self): square = self.b.mul(types.Int32, [self.a, self.a]) c = self.b.convert(types.Float32, [square]) self.b.position_after(square) _, block = self.b.splitblock('start', terminate=True) self.b.position_at_end(block) const = partial(Const, type=types.Int32) cond, body, exit = self.b.gen_loop(const(5), const(10), const(2)) with self.b.at_front(body): self.b.print_(c) with self.b.at_end(exit): self.b.ret(c) # print(string(self.f)) # verify.verify(self.f) # self.assertEqual(loop_expected, string(self.f)) # TestBuilder('test_basic_builder').debug() # TestBuilder('test_splitblock').debug() # TestBuilder('test_loop_builder').debug() # unittest.main()
class PykitIRVisitor(c_ast.NodeVisitor): """ Map pykit IR in the form of polymorphic C to in-memory pykit IR. int function(float x) { int i = 0; /* I am a comment */ while (i < 10) { /*: { "unroll": true } :*/ x = call_external("sqrt", x * x); } return (int) x; } Attributes: """ in_function = False def __init__(self, type_env=None): self.mod = Module() self.type_env = type_env or {} self.func = None self.builder = None self.local_vars = None self.allocas = None self.global_vars = {} self.functions = {} # ______________________________________________________________________ @property def vars(self): if self.in_function: return self.local_vars else: return self.global_vars def enter_func(self): self.in_function = True self.local_vars = {} self.allocas = {} def leave_func(self): self.in_function = False self.mod.add_function(self.func) self.local_vars = None self.allocas = None self.func = None def visit(self, node, type=None): """ Visit a node. :type: Whether we have a type for this opcode, which is an LHS type or a cast. E.g.: (Int) call(...) // cast result = call(...) // assmnt, assuming 'result' is declared result = call(..., call(...)) // second 'call' isn't typed """ self.type = type method = 'visit_' + node.__class__.__name__ visitor = getattr(self, method, self.generic_visit) # if visitor is None: # raise SyntaxError( # "Node %s not supported in %s:%s" % (node, node.coord.file, # node.coord.line)) return visitor(node) def visitif(self, node): if node: return self.visit(node) def visits(self, node): return list(map(self.visit, node)) # ______________________________________________________________________ def alloca(self, varname): if varname not in self.allocas: # Allocate variable with alloca with self.builder.at_front(self.func.blocks[0]): type = self.local_vars[varname] self.allocas[varname] = self.builder.alloca(type, [], varname) return self.allocas[varname] def assign(self, varname, rhs): if not self.in_function: error(rhs, "Assignment only allowed in functions") if varname not in self.allocas: # Allocate variable with alloca with self.builder.at_front(self.func.blocks[0]): type = self.local_vars[varname] self.allocas[varname] = self.builder.alloca(type, [], varname) self.builder.store(self.visit(rhs), self.alloca(varname)) # ______________________________________________________________________ def visit_Decl(self, decl): if decl.name in self.vars: error(decl, "Var '%s' already declared!" % (decl.name, )) type = self.visit(decl.type) self.vars[decl.name] = type if decl.init: self.assign(decl.name, decl.init) return type def visit_TypeDecl(self, decl): return self.visit(decl.type) visit_Typename = visit_TypeDecl def visit_PtrDecl(self, decl): return types.Pointer(self.visit(decl.type.type)) def visit_FuncDecl(self, decl): return types.Function(self.visit(decl.type), self.visits(decl.args.params)) def visit_IdentifierType(self, node): name, = node.names return self.type_env[name] def visit_Typedef(self, node): if node.name in ("Type", "_list"): type = self.type_env[node.name] else: type = self.visit(node.type) if type == types.Type: type = getattr(types, node.name) self.type_env[node.name] = type return type def visit_Template(self, node): left = self.visit(node.left) subtypes = self.visits(node.right) if left is list: return list(subtypes) else: assert issubclass(left, types.Type) subtypes = self.visits(node.right) return left(*subtypes) # ______________________________________________________________________ def visit_FuncDef(self, node): assert not node.param_decls self.enter_func() name = node.decl.name type = self.visit(node.decl.type) argnames = [p.name for p in node.decl.type.args.params] self.func = Function(name, argnames, type) self.func.add_block('entry') self.builder = Builder(self.func) self.builder.position_at_end(self.func.blocks[0]) self.generic_visit(node.body) self.leave_func() # ______________________________________________________________________ def visit_FuncCall(self, node): name = node.name.name if not self.in_typed_context: error( node, "Expected a type for sub-expression " "(add a cast or assignment)") if not hasattr(self.builder, name): error(node, "No opcode %s" % (name, )) self.in_typed_context = False buildop = getattr(self.builder, name) args = self.visits(node.args.exprs) return buildop, args def visit_ID(self, node): if self.in_function: if node.name not in self.local_vars: error(node, "Not a local: %r" % node.name) result = self.alloca(node.name) return self.builder.load(result.type, result) def visit_Cast(self, node): type = self.visit(node.to_type) if isinstance(node.expr, c_ast.FuncCall): self.in_typed_context = True buildop, args = self.visit(node.expr) return buildop(type, args, "temp") else: result = self.visit(node.expr) if result.type == type: return result return self.builder.convert(type, [result], "temp") def visit_Assignment(self, node): if node.op != '=': error(node, "Only assignment with '=' is supported") if not isinstance(node.lvalue, c_ast.ID): error(node, "Canot only assign to a name") self.assign(node.lvalue.name, node.rvalue) def visit_Constant(self, node): type = self.type_env[node.type] const = types.convert(node.value, type) return Const(const) def visit_UnaryOp(self, node): op = defs.unary_defs[node.op] buildop = getattr(self.builder, op) arg = self.visit(node.expr) type = self.type or arg.type return buildop(type, [arg]) def visit_BinaryOp(self, node): op = binary_defs[node.op] buildop = getattr(self.builder, op) left, right = self.visits([node.left, node.right]) if not self.type: assert left.type == right.type, (left, right) return buildop(self.type or left.type, [left, right], "temp") def _loop(self, init, cond, next, body): _, exit_block = self.builder.splitblock("exit") _, body_block = self.builder.splitblock("body") _, cond_block = self.builder.splitblock("cond") self.visitif(init) self.builder.jump(cond_block) with self.builder.at_front(cond_block): cond = self.visit(cond, type=types.Bool) self.builder.cbranch(cond, cond_block, exit_block) with self.builder.at_front(body_block): self.visit(body) self.visitif(next) self.builder.jump(cond_block) self.builder.position_at_end(exit_block) def visit_While(self, node): self._loop(None, node.cond, None, node.stmt) def visit_For(self, node): self._loop(node.init, node.cond, node.next, node.stmt) def visit_Return(self, node): self.builder.ret(self.visit(node.expr))
class PykitIRVisitor(c_ast.NodeVisitor): """ Map pykit IR in the form of polymorphic C to in-memory pykit IR. int function(float x) { int i = 0; /* I am a comment */ while (i < 10) { /*: { "unroll": true } :*/ x = call_external("sqrt", x * x); } return (int) x; } Attributes: """ in_function = False def __init__(self, type_env=None): self.mod = Module() self.type_env = type_env or {} self.func = None self.builder = None self.local_vars = None self.allocas = None self.global_vars = {} self.functions = {} # ______________________________________________________________________ @property def vars(self): if self.in_function: return self.local_vars else: return self.global_vars def enter_func(self): self.in_function = True self.local_vars = {} self.allocas = {} def leave_func(self): self.in_function = False self.mod.add_function(self.func) self.local_vars = None self.allocas = None self.func = None def visit(self, node, type=None): """ Visit a node. :type: Whether we have a type for this opcode, which is an LHS type or a cast. E.g.: (Int) call(...) // cast result = call(...) // assmnt, assuming 'result' is declared result = call(..., call(...)) // second 'call' isn't typed """ self.type = type method = 'visit_' + node.__class__.__name__ visitor = getattr(self, method, self.generic_visit) # if visitor is None: # raise SyntaxError( # "Node %s not supported in %s:%s" % (node, node.coord.file, # node.coord.line)) return visitor(node) def visitif(self, node): if node: return self.visit(node) def visits(self, node): return list(map(self.visit, node)) # ______________________________________________________________________ def alloca(self, varname): if varname not in self.allocas: # Allocate variable with alloca with self.builder.at_front(self.func.blocks[0]): type = self.local_vars[varname] self.allocas[varname] = self.builder.alloca(type, [], varname) return self.allocas[varname] def assign(self, varname, rhs): if not self.in_function: error(rhs, "Assignment only allowed in functions") if varname not in self.allocas: # Allocate variable with alloca with self.builder.at_front(self.func.blocks[0]): type = self.local_vars[varname] self.allocas[varname] = self.builder.alloca(type, [], varname) self.builder.store(self.visit(rhs), self.alloca(varname)) # ______________________________________________________________________ def visit_Decl(self, decl): if decl.name in self.vars: error(decl, "Var '%s' already declared!" % (decl.name,)) type = self.visit(decl.type) self.vars[decl.name] = type if decl.init: self.assign(decl.name, decl.init) return type def visit_TypeDecl(self, decl): return self.visit(decl.type) visit_Typename = visit_TypeDecl def visit_PtrDecl(self, decl): return types.Pointer(self.visit(decl.type.type)) def visit_FuncDecl(self, decl): return types.Function(self.visit(decl.type), self.visits(decl.args.params)) def visit_IdentifierType(self, node): name, = node.names return self.type_env[name] def visit_Typedef(self, node): if node.name in ("Type", "_list"): type = self.type_env[node.name] else: type = self.visit(node.type) if type == types.Type: type = getattr(types, node.name) self.type_env[node.name] = type return type def visit_Template(self, node): left = self.visit(node.left) subtypes = self.visits(node.right) if left is list: return list(subtypes) else: assert issubclass(left, types.Type) subtypes = self.visits(node.right) return left(*subtypes) # ______________________________________________________________________ def visit_FuncDef(self, node): assert not node.param_decls self.enter_func() name = node.decl.name type = self.visit(node.decl.type) argnames = [p.name for p in node.decl.type.args.params] self.func = Function(name, argnames, type) self.func.add_block('entry') self.builder = Builder(self.func) self.builder.position_at_end(self.func.blocks[0]) self.generic_visit(node.body) self.leave_func() # ______________________________________________________________________ def visit_FuncCall(self, node): name = node.name.name if not self.in_typed_context: error(node, "Expected a type for sub-expression " "(add a cast or assignment)") if not hasattr(self.builder, name): error(node, "No opcode %s" % (name,)) self.in_typed_context = False buildop = getattr(self.builder, name) args = self.visits(node.args.exprs) return buildop, args def visit_ID(self, node): if self.in_function: if node.name not in self.local_vars: error(node, "Not a local: %r" % node.name) result = self.alloca(node.name) return self.builder.load(result.type, result) def visit_Cast(self, node): type = self.visit(node.to_type) if isinstance(node.expr, c_ast.FuncCall): self.in_typed_context = True buildop, args = self.visit(node.expr) return buildop(type, args, "temp") else: result = self.visit(node.expr) if result.type == type: return result return self.builder.convert(type, [result], "temp") def visit_Assignment(self, node): if node.op != '=': error(node, "Only assignment with '=' is supported") if not isinstance(node.lvalue, c_ast.ID): error(node, "Canot only assign to a name") self.assign(node.lvalue.name, node.rvalue) def visit_Constant(self, node): type = self.type_env[node.type] const = types.convert(node.value, type) return Const(const) def visit_UnaryOp(self, node): op = defs.unary_defs[node.op] buildop = getattr(self.builder, op) arg = self.visit(node.expr) type = self.type or arg.type return buildop(type, [arg]) def visit_BinaryOp(self, node): op = binary_defs[node.op] buildop = getattr(self.builder, op) left, right = self.visits([node.left, node.right]) if not self.type: assert left.type == right.type, (left, right) return buildop(self.type or left.type, [left, right], "temp") def _loop(self, init, cond, next, body): _, exit_block = self.builder.splitblock("exit") _, body_block = self.builder.splitblock("body") _, cond_block = self.builder.splitblock("cond") self.visitif(init) self.builder.jump(cond_block) with self.builder.at_front(cond_block): cond = self.visit(cond, type=types.Bool) self.builder.cbranch(cond, cond_block, exit_block) with self.builder.at_front(body_block): self.visit(body) self.visitif(next) self.builder.jump(cond_block) self.builder.position_at_end(exit_block) def visit_While(self, node): self._loop(None, node.cond, None, node.stmt) def visit_For(self, node): self._loop(node.init, node.cond, node.next, node.stmt) def visit_Return(self, node): self.builder.ret(self.visit(node.expr))