def rewrite_obj_return(func, env): """ Handle returning stack-allocated objects. """ if should_skip(env): return context = env['flypy.typing.context'] restype = env['flypy.typing.restype'] envs = env['flypy.state.envs'] builder = Builder(func) stack_alloc = representation.byref(restype) if stack_alloc: out = func.add_arg(func.temp("out"), opaque_t) context[out] = Pointer[restype] func.type = types.Function(types.Void, func.type.argtypes, False) for arg in func.args: arg.type = opaque_t func.type = types.Function(func.type.restype, (opaque_t,) * len(func.args), False) is_generator = env['flypy.state.generator'] for op in func.ops: if (op.opcode == 'ret' and op.args[0] is not None and stack_alloc and not is_generator): # ret val => # store (load val) out ; ret void [val] = op.args builder.position_before(op) newval = builder.load(val) builder.store(newval, out) op.set_args([None]) # Update context context[newval] = StackVar[context[val]] elif op.opcode == 'call' and op.type != types.Void: # result = call(f, ...) => # alloca result ; call(f, ..., &result) ty = context[op] if conversion.byref(ty): f, args = op.args if not is_flypy_cc(f) or should_skip(envs[f]): continue builder.position_before(op) retval = builder.alloca(opaque_t) builder.position_after(op) op.replace_uses(retval) newargs = args + [retval] op.set_args([f, newargs]) # Update context context[retval] = context[op] context[op] = void
def run(func, env=None, return_block=None): """ Rewrite 'ret' operations into jumps to a return block and assignments to a return variable. """ b = Builder(func) return_block = return_block or func.new_block("pykit.return") # Allocate return variable if not func.type.restype.is_void: with b.at_front(func.startblock): return_var = b.alloca(types.Pointer(func.type.restype), []) b.store(Undef(func.type.restype), return_var) else: return_var = None # Repace 'ret' instructions with jumps and assignments for op in func.ops: if op.opcode == "ret": b.position_after(op) if return_var: b.store(op.args[0], return_var) b.jump(return_block) op.delete() with b.at_end(return_block): if return_var: result = b.load(return_var.type.base, [return_var]) else: result = None b.ret(result)
def run(func, env=None, return_block=None): """ Rewrite 'ret' operations into jumps to a return block and assignments to a return variable. """ b = Builder(func) return_block = return_block or func.new_block("pykit.return") # Allocate return variable if not func.type.restype.is_void: with b.at_front(func.startblock): return_var = b.alloca(types.Pointer(func.type.restype)) b.store(Undef(func.type.restype), return_var) else: return_var = None # Repace 'ret' instructions with jumps and assignments for op in func.ops: if op.opcode == "ret": b.position_after(op) if return_var: b.store(op.args[0], return_var) b.jump(return_block) op.delete() with b.at_end(return_block): if return_var: result = b.load(return_var) else: result = None b.ret(result)
def rewrite_obj_return(func, env): """ Handle returning stack-allocated objects. """ if should_skip(env): return context = env['flypy.typing.context'] restype = env['flypy.typing.restype'] envs = env['flypy.state.envs'] builder = Builder(func) stack_alloc = representation.byref(restype) if stack_alloc: out = func.add_arg(func.temp("out"), opaque_t) context[out] = Pointer[restype] func.type = types.Function(types.Void, func.type.argtypes, False) for arg in func.args: arg.type = opaque_t func.type = types.Function(func.type.restype, (opaque_t, ) * len(func.args), False) is_generator = env['flypy.state.generator'] for op in func.ops: if (op.opcode == 'ret' and op.args[0] is not None and stack_alloc and not is_generator): # ret val => # store (load val) out ; ret void [val] = op.args builder.position_before(op) newval = builder.load(val) builder.store(newval, out) op.set_args([None]) # Update context context[newval] = StackVar[context[val]] elif op.opcode == 'call' and op.type != types.Void: # result = call(f, ...) => # alloca result ; call(f, ..., &result) ty = context[op] if conversion.byref(ty): f, args = op.args if not is_flypy_cc(f) or should_skip(envs[f]): continue builder.position_before(op) retval = builder.alloca(opaque_t) builder.position_after(op) op.replace_uses(retval) newargs = args + [retval] op.set_args([f, newargs]) # Update context context[retval] = context[op] context[op] = void
def lower_fields(func, env): b = Builder(func) opbuilder = OpBuilder() for op in func.ops: if op.opcode not in ("getfield", "setfield"): continue if op.args[0].type.is_pointer: b.position_before(op) # Load the pointer and update the argument p = op.args[0] load = b.load(p) args = [load] + op.args[1:] op.set_args(args) if op.opcode == "setfield": # Write back result b.position_after(op) op.type = load.type b.store(op, p) if op.opcode == "getfield": struct, attr = op.args newop = opbuilder.extractfield(op.type, struct, attr, result=op.result) else: struct, attr, value = op.args newop = opbuilder.insertfield(op.type, struct, attr, value, result=op.result) op.replace(newop)
def lower_fields(func, env=None): b = Builder(func) for op in func.ops: if op.opcode in ("getfield", "setfield") and op.args[0].type.is_pointer: b.position_before(op) p = op.args[0] load = b.load(p.type.base, [p]) args = [load] + op.args[1:] op.set_args(args) if op.opcode == "setfield": b.position_after(op) b.store(op, p)
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 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.new_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)) assert interp.run(self.f, args=[10]) == 100 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]) self.assertEqual(opcodes(self.f), ['add', 'div']) 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) self.assertEqual(interp.run(self.f, args=[10]), 100.0)
def generate_copies(func, phis): """ Emit stores to stack variables in predecessor blocks. """ builder = Builder(func) vars = {} loads = {} # Allocate a stack variable for each phi builder.position_at_beginning(func.startblock) for block in phis: for phi in phis[block]: vars[phi] = builder.alloca(types.Pointer(phi.type)) # Generate loads in blocks containing the phis for block in phis: leaders = list(block.leaders) last_leader = leaders[-1] if leaders else block.head builder.position_after(last_leader) for phi in phis[block]: loads[phi] = builder.load(vars[phi]) # Generate copies (store to stack variables) for block in phis: for phi in phis[block]: preds, args = phi.args var = vars[phi] phi_args = [loads.get(arg, arg) for arg in args] for pred, arg in zip(preds, phi_args): builder.position_before(pred.terminator) builder.store(arg, var) # Replace phis for block in phis: for phi in phis[block]: phi.replace_uses(loads[phi]) phi.delete() return vars, loads
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.startblock): type = types.Pointer(self.local_vars[varname]) result = self.func.temp(varname) self.allocas[varname] = self.builder.alloca(type, [], result) return self.allocas[varname] def assignvar(self, varname, rhs): self.builder.store(rhs, self.alloca(varname)) def assign(self, varname, rhs): if self.in_function: # Local variable type = self.local_vars[varname] self.assignvar(varname, self.visit(rhs, type=type)) else: # Global variable type = self.global_vars[varname] self.mod.add_global(GlobalValue(varname, type=self.type, value=self.visit(rhs, type=type))) # ______________________________________________________________________ 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) elif not self.in_function: extern = decl.storage == 'external' self.mod.add_global(GlobalValue(decl.name, type, external=extern)) 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): if decl.args: params = self.visits(decl.args.params) else: params = [] return types.Function(self.visit(decl.type), 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) if node.decl.type.args: argnames = [p.name or "" for p in node.decl.type.args.params] else: argnames = [] self.func = Function(name, argnames, type) self.func.new_block('entry') self.builder = Builder(self.func) self.builder.position_at_end(self.func.startblock) # Store arguments in stack variables for argname in argnames: self.assignvar(argname, self.func.get_arg(argname)) self.generic_visit(node.body) self.leave_func() # ______________________________________________________________________ def visit_FuncCall(self, node): type = self.type opcode = node.name.name args = self.visits(node.args.exprs) if node.args else [] if opcode == "list": return args elif not type and not ops.is_void(opcode): error(node, "Expected a type for sub-expression " "(add a cast or assignment)") elif not hasattr(self.builder, opcode): if opcode in self.mod.functions: return self.builder.call(type, [self.mod.get_function(opcode), args]) error(node, "No opcode %s" % (opcode,)) buildop = getattr(self.builder, opcode) if ops.is_void(opcode): return buildop(*args) else: return buildop(type or "Unset", args) def visit_ID(self, node): if self.in_function: if node.name in self.local_vars: result = self.alloca(node.name) return self.builder.load(result.type.base, [result]) global_val = (self.mod.get_function(node.name) or self.mod.get_global(node.name)) if not global_val: error(node, "Not a local or global: %r" % node.name) return global_val def visit_Cast(self, node): type = self.visit(node.to_type) if isinstance(node.expr, c_ast.FuncCall): op = self.visit(node.expr, type=type) op.type = type return op else: result = self.visit(node.expr) if result.type == type: return result return self.builder.convert(type, [result]) 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, types.resolve_typedef(type)) if isinstance(const, basestring): const = const[1:-1] # slice away quotes 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]) type = self.type if not type: l, r = map(types.resolve_typedef, [left.type, right.type]) assert l == r, (l, r) if node.op in defs.compare_defs: type = types.Bool return buildop(type or left.type, [left, right]) def visit_If(self, node): cond = self.visit(node.cond) ifpos, elsepos, exit_block = self.builder.ifelse(cond) with ifpos: self.visit(node.iftrue) self.builder.jump(exit_block) with elsepos: if node.iffalse: self.visit(node.iffalse) self.builder.jump(exit_block) self.builder.position_at_end(exit_block) def _loop(self, init, cond, next, body): _, exit_block = self.builder.splitblock(self.func.temp("exit")) _, body_block = self.builder.splitblock(self.func.temp("body")) _, cond_block = self.builder.splitblock(self.func.temp("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, body_block, exit_block) with self.builder.at_front(body_block): self.visit(body) self.visitif(next) bb = self.builder.basic_block if not bb.tail or not ops.is_terminator(bb.tail.opcode): 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): # avoid silly 2to3 rewrite to 'node.__next__' next = getattr(node, 'next') self._loop(node.init, node.cond, next, node.stmt) def visit_Return(self, node): b = self.builder value = self.visit(node.expr) t = self.func.temp b.ret(b.convert(self.func.type.restype, [value]))
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.startblock): type = types.Pointer(self.local_vars[varname]) result = self.func.temp(varname) self.allocas[varname] = self.builder.alloca(type, [], result) return self.allocas[varname] def assignvar(self, varname, rhs): self.builder.store(rhs, self.alloca(varname)) def assign(self, varname, rhs): if self.in_function: # Local variable type = self.local_vars[varname] self.assignvar(varname, self.visit(rhs, type=type)) else: # Global variable type = self.global_vars[varname] self.mod.add_global( GlobalValue(varname, type=self.type, value=self.visit(rhs, type=type))) # ______________________________________________________________________ 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) elif not self.in_function: extern = decl.storage == 'external' self.mod.add_global(GlobalValue(decl.name, type, external=extern)) 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): if decl.args: params = self.visits(decl.args.params) else: params = [] return types.Function(self.visit(decl.type), 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) if node.decl.type.args: argnames = [p.name or "" for p in node.decl.type.args.params] else: argnames = [] self.func = Function(name, argnames, type) self.func.new_block('entry') self.builder = Builder(self.func) self.builder.position_at_end(self.func.startblock) # Store arguments in stack variables for argname in argnames: self.assignvar(argname, self.func.get_arg(argname)) self.generic_visit(node.body) self.leave_func() # ______________________________________________________________________ def visit_FuncCall(self, node): type = self.type opcode = node.name.name args = self.visits(node.args.exprs) if node.args else [] if opcode == "list": return args elif not type and not ops.is_void(opcode): error( node, "Expected a type for sub-expression " "(add a cast or assignment)") elif not hasattr(self.builder, opcode): if opcode in self.mod.functions: return self.builder.call(type, [self.mod.get_function(opcode), args]) error(node, "No opcode %s" % (opcode, )) buildop = getattr(self.builder, opcode) if ops.is_void(opcode): return buildop(*args) else: return buildop(type or "Unset", args) def visit_ID(self, node): if self.in_function: if node.name in self.local_vars: result = self.alloca(node.name) return self.builder.load(result.type.base, [result]) global_val = (self.mod.get_function(node.name) or self.mod.get_global(node.name)) if not global_val: error(node, "Not a local or global: %r" % node.name) return global_val def visit_Cast(self, node): type = self.visit(node.to_type) if isinstance(node.expr, c_ast.FuncCall): op = self.visit(node.expr, type=type) op.type = type return op else: result = self.visit(node.expr) if result.type == type: return result return self.builder.convert(type, [result]) 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, types.resolve_typedef(type)) if isinstance(const, basestring): const = const[1:-1] # slice away quotes 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]) type = self.type if not type: l, r = map(types.resolve_typedef, [left.type, right.type]) assert l == r, (l, r) if node.op in defs.compare_defs: type = types.Bool return buildop(type or left.type, [left, right]) def visit_If(self, node): cond = self.visit(node.cond) ifpos, elsepos, exit_block = self.builder.ifelse(cond) with ifpos: self.visit(node.iftrue) self.builder.jump(exit_block) with elsepos: if node.iffalse: self.visit(node.iffalse) self.builder.jump(exit_block) self.builder.position_at_end(exit_block) def _loop(self, init, cond, next, body): _, exit_block = self.builder.splitblock(self.func.temp("exit")) _, body_block = self.builder.splitblock(self.func.temp("body")) _, cond_block = self.builder.splitblock(self.func.temp("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, body_block, exit_block) with self.builder.at_front(body_block): self.visit(body) self.visitif(next) bb = self.builder.basic_block if not bb.tail or not ops.is_terminator(bb.tail.opcode): 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): # avoid silly 2to3 rewrite to 'node.__next__' next = getattr(node, 'next') self._loop(node.init, node.cond, next, node.stmt) def visit_Return(self, node): b = self.builder value = self.visit(node.expr) t = self.func.temp b.ret(b.convert(self.func.type.restype, [value]))
class TestBuilder(unittest.TestCase): def setUp(self): self.f = Function("testfunc", ['a'], types.Function(types.Float32, [types.Int32], False)) self.b = Builder(self.f) self.b.position_at_end(self.f.new_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(self.a, self.a, result='r') c = self.b.convert(types.Float32, result) self.b.store(c, v) val = self.b.load(v) self.b.ret(val) # print(string(self.f)) assert interp.run(self.f, args=[10]) == 100 def test_splitblock(self): old, new = self.b.splitblock('newblock') with self.b.at_front(old): self.b.add(self.a, self.a) with self.b.at_end(new): self.b.div(self.a, self.a) self.assertEqual(opcodes(self.f), ['add', 'div']) def test_loop_builder(self): square = self.b.mul(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) self.assertEqual(interp.run(self.f, args=[10]), 100.0) def test_splitblock_preserve_phis(self): """ block1: %0 = mul a a jump(newblock) newblock: %1 = phi([block1], [%0]) ret %1 """ square = self.b.mul(self.a, self.a) old, new = self.b.splitblock('newblock', terminate=True) with self.b.at_front(new): phi = self.b.phi(types.Int32, [self.f.startblock], [square]) self.b.ret(phi) # Now split block1 self.b.position_after(square) block1, split = self.b.splitblock(terminate=True) phi, ret = new.ops blocks, values = phi.args self.assertEqual(blocks, [split])
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))