def split_critical_edges(func, cfg, phis):
"""
Split critical edges to correctly handle cycles in phis. See 2) above.
"""
b = Builder(func)
for block in cfg.node:
successors = cfg.neighbors(block)
if len(successors) > 1:
# More than one successor, we need to split
# (Alternatively, we could move our copies into the successor block
# if we were the only predecessor, but this seems simpler)
# Split successors with phis
new_succs = {} # old_successor -> new_successor
for succ in successors:
if phis[succ]:
label = func.temp("split_critical")
new_succ = func.new_block(label, after=block)
new_succs[succ] = new_succ
b.position_at_end(new_succ)
b.jump(succ)
# Patch our basic-block terminator to point to new blocks
if new_succs:
terminator = block.terminator
assert terminator.opcode == 'cbranch', terminator
test, truebb, falsebb = terminator.args
terminator.set_args([
test,
new_succs.get(truebb, truebb),
new_succs.get(falsebb, falsebb)
])
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 split_critical_edges(func, cfg, phis):
"""
Split critical edges to correctly handle cycles in phis. See 2) above.
"""
b = Builder(func)
for block in cfg.node:
successors = cfg.neighbors(block)
if len(successors) > 1:
# More than one successor, we need to split
# (Alternatively, we could move our copies into the successor block
# if we were the only predecessor, but this seems simpler)
# Split successors with phis
new_succs = {} # old_successor -> new_successor
for succ in successors:
if phis[succ]:
label = func.temp("split_critical")
new_succ = func.new_block(label, after=block)
new_succs[succ] = new_succ
b.position_at_end(new_succ)
b.jump(succ)
# Patch our basic-block terminator to point to new blocks
if new_succs:
terminator = block.terminator
assert terminator.opcode == 'cbranch', terminator
test, truebb, falsebb = terminator.args
terminator.set_args([test,
new_succs.get(truebb, truebb),
new_succs.get(falsebb, falsebb)])
def move_generator(func, consumer, empty_body):
gen = consumer.generator
gen.unlink()
b = Builder(func)
b.position_at_end(empty_body)
b.emit(gen)
with b.at_end(empty_body):
loop_exit = determine_loop_exit(consumer.loop)
b.jump(loop_exit)
def move_generator(func, consumer, empty_body):
gen = consumer.generator
gen.unlink()
b = Builder(func)
b.position_at_end(empty_body)
b.emit(gen)
with b.at_end(empty_body):
loop_exit = determine_loop_exit(consumer.loop)
b.jump(loop_exit)
def inline(func, call):
"""
Inline the call instruction into func.
:return: { old_op : new_op }
"""
callee = call.args[0]
callblock = call.block
# assert_inlinable(func, call, callee, uses)
builder = Builder(func)
builder.position_before(call)
inline_header, inline_exit = builder.splitblock()
new_callee, valuemap = copy_function(callee, temper=func.temp)
result = rewrite_return(new_callee)
# Fix up arguments
for funcarg, arg in zip(new_callee.args, call.args[1]):
funcarg.replace_uses(arg)
# Copy blocks
new_blocks = list(new_callee.blocks)
after = inline_header
for block in new_blocks:
block.parent = None
func.add_block(block, after=after)
after = block
# Fix up wiring
builder.jump(new_callee.startblock)
with builder.at_end(new_callee.exitblock):
builder.jump(inline_exit)
stretch_exception_block(builder, callblock, new_blocks)
# Fix up final result of call
if result is not None:
# non-void return
result.unlink()
result.result = call.result
call.replace(result)
else:
call.delete()
func.reset_uses()
#verify(func)
return valuemap
def consume_yields(func, consumer, generator_func, valuemap):
b = Builder(func)
copier = lambda x : x
loop = consumer.loop
inlined_values = set(valuemap.values())
for block in func.blocks:
if block in inlined_values:
for op in block.ops:
if op.opcode == 'yield':
# -- Replace 'yield' by the loop body -- #
b.position_after(op)
_, resume = b.splitblock()
# Copy blocks
blocks = [copier(block) for block in loop.blocks]
# Insert blocks
prev = op.block
for block in blocks:
func.add_block(block, after=prev)
prev = block
# Fix wiring
b.jump(blocks[0])
b.position_at_end(blocks[-1])
b.jump(resume)
# We just introduced a bunch of copied blocks
func.reset_uses()
# Update phis with new predecessor
b.replace_predecessor(loop.tail, op.block, loop.head)
b.replace_predecessor(loop.tail, op.block, loop.head)
# Replace next() by value produced by yield
value = op.args[0]
consumer.next.replace_uses(value)
op.delete()
# We don't need these anymore
consumer.next.delete()
def consume_yields(func, consumer, generator_func, valuemap):
b = Builder(func)
copier = lambda x: x
loop = consumer.loop
inlined_values = set(valuemap.values())
for block in func.blocks:
if block in inlined_values:
for op in block.ops:
if op.opcode == 'yield':
# -- Replace 'yield' by the loop body -- #
b.position_after(op)
_, resume = b.splitblock()
# Copy blocks
blocks = [copier(block) for block in loop.blocks]
# Insert blocks
prev = op.block
for block in blocks:
func.add_block(block, after=prev)
prev = block
# Fix wiring
b.jump(blocks[0])
b.position_at_end(blocks[-1])
b.jump(resume)
# We just introduced a bunch of copied blocks
func.reset_uses()
# Update phis with new predecessor
b.replace_predecessor(loop.tail, op.block, loop.head)
b.replace_predecessor(loop.tail, op.block, loop.head)
# Replace next() by value produced by yield
value = op.args[0]
consumer.next.replace_uses(value)
op.delete()
# We don't need these anymore
consumer.next.delete()
def splitblock(block, trailing, name=None, terminate=False, preserve_exc=True):
"""Split the current block, returning (old_block, new_block)"""
from pykit.analysis import cfa
from pykit.ir import Builder
func = block.parent
if block.is_terminated():
successors = cfa.deduce_successors(block)
else:
successors = []
# -------------------------------------------------
# Sanity check
# Allow splitting only after leaders and before terminator
# TODO: error check
# -------------------------------------------------
# Split
blockname = name or func.temp('Block')
newblock = func.new_block(blockname, after=block)
# -------------------------------------------------
# Move ops after the split to new block
for op in trailing:
op.unlink()
newblock.extend(trailing)
if terminate and not block.is_terminated():
# Terminate
b = Builder(func)
b.position_at_end(block)
b.jump(newblock)
# Update phis and preserve exception blocks
patch_phis(block, newblock, successors)
if preserve_exc:
preserve_exceptions(block, newblock)
return block, newblock
def inline(func, call, uses=None):
"""
Inline the call instruction into func.
:param uses: defuse information
"""
callee = call.args[0]
# assert_inlinable(func, call, callee, uses)
builder = Builder(func)
builder.position_before(call)
inline_header, inline_exit = builder.splitblock()
new_callee = copy_function(callee, temper=func.temp)
result = rewrite_return(new_callee)
# Fix up arguments
for funcarg, arg in zip(new_callee.args, call.args[1]):
funcarg.replace_uses(arg)
# Copy blocks
after = inline_header
for block in new_callee.blocks:
block.parent = None
func.add_block(block, after=after)
after = block
# Fix up wiring
builder.jump(new_callee.startblock)
with builder.at_end(new_callee.exitblock):
builder.jump(inline_exit)
# Fix up final result of call
if result is not None:
# non-void return
result.unlink()
result.result = call.result
call.replace(result)
else:
call.delete()
func.reset_uses()
verify(func)
def splitblock(block, trailing, name=None, terminate=False, preserve_exc=True):
"""Split the current block, returning (old_block, new_block)"""
func = block.parent
if block.is_terminated():
successors = deduce_successors(block)
else:
successors = []
# -------------------------------------------------
# Sanity check
# Allow splitting only after leaders and before terminator
# TODO: error check
# -------------------------------------------------
# Split
blockname = name or func.temp('Block')
newblock = func.new_block(blockname, after=block)
# -------------------------------------------------
# Move ops after the split to new block
for op in trailing:
op.unlink()
newblock.extend(trailing)
if terminate and not block.is_terminated():
# Terminate
b = Builder(func)
b.position_at_end(block)
b.jump(newblock)
# Update phis and preserve exception blocks
patch_phis(block, newblock, successors)
if preserve_exc:
preserve_exceptions(block, newblock)
return block, newblock
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 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))
