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 initialize(self):
"""Initialize pykit untypes structures"""
# Setup Function
sig = types.Function(types.Opaque, [types.Opaque] * len(self.argnames),
False)
self.dst = Function(func_name(self.func), self.argnames, sig)
# Setup Builder
self.builder = Builder(self.dst)
# Setup Blocks
for offset in self.bytecode.labels:
name = blockname(self.func, offset)
block = self.dst.new_block(name)
self.blocks[offset] = block
self.stacks[block] = []
# Setup Variables
self.builder.position_at_beginning(self.dst.startblock)
for varname in self.varnames:
stackvar = self.builder.alloca(types.Pointer(types.Opaque),
result=self.dst.temp(varname))
self.allocas[varname] = stackvar
# Initialize function arguments
if varname in self.argnames:
self.builder.store(self.dst.get_arg(varname), stackvar)
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 test_exc_rewrite(self):
func = Function("foo", [], types.Function(types.Void, (), False))
entry = func.new_block("entry")
catch_block = func.new_block("catch")
b = Builder(func)
with b.at_front(entry):
b.exc_setup([catch_block])
b.exc_throw(Const(StopIteration, types.Exception))
with b.at_front(catch_block):
b.exc_catch([Const(Exception, types.Exception)])
local_exceptions.run(func, {})
self.assertNotIn('exc_throw', opcodes(func))
def lower_into_runtime(self, op, insert_decl=False):
"""
Lower op into a runtime call.
:param decl: Whether to insert an external declaration if not present
"""
_verify_args(op)
if insert_decl and not self.mod.get_global(op.opcode):
signature = types.Function(op.type, [arg.type for arg in op.args])
self.mod.add_global(
GlobalValue(op.opcode, signature, external=True))
call = self.builder.gen_call_external(op.opcode, op.args, op.result)
op.replace(call)
def from_expr(graph, expr_context, env):
"""
Map a Blaze expression graph to blaze AIR
Parameters
----------
graph: blaze.expr.Op
Expression graph
expr_context: ExprContext
Context of the expression
ctx: ExecutionContext
"""
inputs = expr_context.params
# -------------------------------------------------
# Types
argtypes = [operand.dshape for operand in inputs]
signature = types.Function(graph.dshape, argtypes, varargs=False)
# -------------------------------------------------
# Setup function
name = "expr"
argnames = ["e%d" % i for i in range(len(inputs))]
f = Function(name, argnames, signature)
builder = Builder(f)
builder.position_at_beginning(f.new_block('entry'))
# -------------------------------------------------
# Generate function
valuemap = dict(
(expr, f.get_arg("e%d" % i)) for i, expr in enumerate(inputs))
_from_expr(graph, f, builder, valuemap)
retval = valuemap[graph]
builder.ret(retval)
# Update environment with runtime arguments
runtime_args = [expr_context.terms[input] for input in inputs]
env['runtime.args'] = dict(zip(f.args, runtime_args))
return f
def __init__(self, function):
self.theano = function
argnames = ["arg%d" % i for i in range(len(function.inputs))]
argtypes = [map_type(i.type) for i in function.inputs]
restype = type_from_outputs(function.outputs)
signature = types.Function(restype, argtypes)
# Setup up pykit function
self.func = ir.Function("theano_func", argnames, signature)
self.builder = ir.Builder(self.func)
self.builder.position_at_end(self.func.new_block('entry'))
# Theano Variable -> PyKit Operation
self.values = {}
self.update_dict = {} # Z_new[...] = Z_old[...]
def from_ctypes_type(ctypes_type, memo=None):
"""
Convert a ctypes type to a pykit type.
Supported are structs, unit types (int/float)
"""
if memo is None:
memo = {}
if hashable(ctypes_type) and ctypes_type in memo:
return memo[ctypes_type]
if hashable(ctypes_type) and ctypes_type in ctypes_map:
result = ctypes_map[ctypes_type]
elif ctypes_type is ctypes.c_void_p:
result = types.Pointer(types.Void)
elif is_ctypes_array_type(ctypes_type):
result = types.Array(from_ctypes_type(ctypes_type._type_, memo),
ctypes_type._length_)
elif is_ctypes_pointer_type(ctypes_type):
result = types.Pointer(from_ctypes_type(ctypes_type._type_, memo))
elif is_ctypes_struct_type(ctypes_type):
# pre-order caching for recursive data structures
f_names = []
f_types = []
result = types.Struct(f_names, f_types)
memo[ctypes_type] = result
fields = [(name, from_ctypes_type(field_type, memo))
for name, field_type in ctypes_type._fields_]
fieldnames, fieldtypes = zip(*fields) or (('dummy', ), (types.Int8, ))
f_names.extend(fieldnames)
f_types.extend(fieldtypes)
elif is_ctypes_function_type(ctypes_type):
c_restype = from_ctypes_type(ctypes_type._restype_, memo)
c_argtypes = [
from_ctypes_type(argty, memo) for argty in ctypes_type._argtypes_
]
result = types.Function(c_restype, c_argtypes, False)
else:
raise NotImplementedError(ctypes_type)
memo[ctypes_type] = result
return result
def impl(py_func, argtypes):
# TODO: do this better
from flypy.compiler import representation_type
ll_argtypes = [representation_type(x) for x in argtypes]
argnames = list(string.ascii_letters[:len(argtypes)])
# Determine return type
if restype_func:
result_type = restype_func(argtypes)
else:
result_type = restype or ll_argtypes[0]
type = ptypes.Function(result_type, tuple(ll_argtypes), False)
func = ir.Function(name, argnames, type)
func.new_block("entry")
b = ir.Builder(func)
b.position_at_beginning(func.startblock)
implementation(b, argtypes, *func.args)
return func
def rewrite_externs(func, env):
"""Rewrite external symbol references
Base on flypy.compiler.lower.constants.rewrite_constants
"""
if env['flypy.state.opaque']:
return
target = env['flypy.target']
# For each operation
for op in func.ops:
# Only for call operation
if op.opcode == 'call':
# For each constant
constants = ir.collect_constants(op)
new_constants = []
for c in constants:
extern = c.const
if extern_support.is_extern_symbol(extern):
# Make a declare-only function
argtypes = extern.type.argtypes
restype = extern.type.restype
if target == "cpu":
# Install external symbol for CPU target
extern.install()
functype = ptypes.Function(lltype(restype),
[lltype(t) for t in argtypes],
extern.type.varargs)
# Note: Global value should really be part inserted into
# a module. But there are no module support at this point.
replacment = ir.GlobalValue(extern.name, functype,
external=True)
else:
# No change
replacment = c
# Add replacement
new_constants.append(replacment)
# Replace
ir.substitute_args(op, constants, new_constants)
def lltyping(func, env):
"""Annotate the function with the low-level representation types"""
if not env['flypy.state.opaque']:
context = env['flypy.typing.context']
resolve = partial(resolve_type, context)
for arg in func.args:
resolve(arg)
for op in func.ops:
if op.opcode == 'exc_catch':
continue
op.replace(resolve(op))
op.set_args(nestedmap(resolve, op.args))
restype = env['flypy.typing.restype']
if conversion.byref(restype):
ll_restype = ptypes.Void
else:
ll_restype = compiler.representation_type(restype)
func.type = ptypes.Function(ll_restype,
[arg.type for arg in func.args], False)
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 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 resolve_restype(func, env):
"""Figure out the return type and update the context and environment"""
context = env['flypy.typing.context']
restype = env['flypy.typing.restype']
signature = env['flypy.typing.signature']
typeset = context['return']
inferred_restype = signature.restype
if restype is None:
restype = inferred_restype
elif inferred_restype != restype:
try:
[restype] = unify([(inferred_restype, restype)])
except UnificationError, e:
raise TypeError("Annotated result type %s does not match inferred "
"type %s for function %r: %s" %
(restype, inferred_restype, func.name, e))
if isinstance(restype, set):
raise TypeError("Undetermined return type for function %s" %
(func.name, ))
env['flypy.typing.restype'] = restype
if restype == void or env['flypy.state.generator']:
_, argtypes, varargs = func.type
func.type = types.Function(types.Void, argtypes, varargs)
def visit_FuncDecl(self, decl):
return types.Function(self.visit(decl.type),
self.visits(decl.args.params))
def add_arg(self, argname, argtype):
self.argnames.append(argname)
argtypes = tuple(self.type.argtypes) + (argtype, )
self.type = types.Function(self.type.restype, argtypes,
self.type.varargs)
return self.get_arg(argname)