def specialize(untyped, args, kwargs = {}, optimize = True):
"""
Translate, specialize and begin to optimize the given function for the types
of the supplies arguments.
Return the untyped and typed representations, along with all the
arguments in a linear order.
"""
if not isinstance(untyped, UntypedFn):
import ast_conversion
untyped = ast_conversion.translate_function_value(untyped)
arg_values, arg_types = prepare_args(untyped, args, kwargs)
# convert the awkward mix of positional, named, and starargs
# into a positional sequence of arguments
linear_args = untyped.args.linearize_without_defaults(arg_values)
# propagate types through function representation and all
# other functions it calls
typed_fn = type_inference.specialize(untyped, arg_types)
if optimize:
from .. transforms.pipeline import normalize
# apply high level optimizations
typed_fn = normalize.apply(typed_fn)
return typed_fn, linear_args
def specialize(untyped, args, kwargs = {}, optimize = False):
"""
Translate, specialize and begin to optimize the given function for the types
of the supplies arguments.
Return the untyped and typed representations, along with all the
arguments in a linear order.
"""
if not isinstance(untyped, UntypedFn):
untyped = ast_conversion.translate_function_value(untyped)
arg_values, arg_types = prepare_args(untyped, args, kwargs)
# convert the awkward mix of positional, named, and starargs
# into a positional sequence of arguments
linear_args = untyped.args.linearize_without_defaults(arg_values)
# propagate types through function representation and all
# other functions it calls
typed_fn = type_inference.specialize(untyped, arg_types)
if optimize:
from .. transforms.pipeline import normalize
# apply high level optimizations
typed_fn = normalize.apply(typed_fn)
return typed_fn, linear_args
def run_python_fn(fn, args, kwargs = None, backend = None): """ Given a python function, run it in Parakeet on the supplied args """ # translate from the Python AST to Parakeet's untyped format untyped = ast_conversion.translate_function_value(fn) return run_untyped_fn(untyped, args, kwargs, backend)
def run_python_fn(fn, args, kwargs = None, backend = None): """ Given a python function, run it in Parakeet on the supplied args """ # translate from the Python AST to Parakeet's untyped format untyped = ast_conversion.translate_function_value(fn) return run_untyped_fn(untyped, args, kwargs, backend)
def run_python_fn(python_fn, args, kwds): untyped = ast_conversion.translate_function_value(python_fn) # should eventually roll this up into something cleaner, since # top-level functions are really acting like closures over their # global dependencies global_args = [python_fn.func_globals[n] for n in untyped.nonlocals] all_positional = global_args + list(args) actuals = args.FormalArgs(all_positional, kwds) return eval_fn(untyped, actuals)
def prepare_args(fn, args, kwargs):
"""
Fetch the function's nonlocals and return an ActualArgs object of both the arg
values and their types
"""
#assert not isinstance(fn, TypedFn), "[prepare_args] Only works for untyped functions"
if not isinstance(fn, UntypedFn):
fn = ast_conversion.translate_function_value(fn)
nonlocals = tuple(fn.python_nonlocals())
arg_values = ActualArgs(nonlocals + tuple(args), kwargs)
arg_types = arg_values.transform(_typeof)
return arg_values, arg_types
def prepare_args(fn, args, kwargs):
"""
Fetch the function's nonlocals and return an ActualArgs object of both the arg
values and their types
"""
#assert not isinstance(fn, TypedFn), "[prepare_args] Only works for untyped functions"
if not isinstance(fn, UntypedFn):
fn = ast_conversion.translate_function_value(fn)
nonlocals = tuple(fn.python_nonlocals())
arg_values = ActualArgs(nonlocals + tuple(args), kwargs)
arg_types = arg_values.transform(_typeof)
return arg_values, arg_types
def __call__(self, *args, **kwargs):
if '_backend' in kwargs:
backend_name = kwargs['_backend']
del kwargs['_backend']
else:
backend_name = None
if self.untyped is None:
import ast_conversion
self.untyped = ast_conversion.translate_function_value(self.fn)
typed_fn, linear_args = specialize(self.untyped, args, kwargs)
return run_typed_fn(typed_fn, linear_args, backend_name)
def get_fundef(fn):
"""
Get the function definition in case I want to pass in the name of an untyped
function or an untranslated python fn.
"""
import ast_conversion
if isinstance(fn, str):
assert fn in syntax.Fn.registry, "Function not found: %s" % fn
return syntax.Fn.registry[fn]
elif not isinstance(fn, syntax.Fn):
return ast_conversion.translate_function_value(fn)
else:
return fn
def prepare_args(fn, args, kwargs):
"""
Fetch the function's nonlocals and return an ActualArgs object of both the arg
values and their types
"""
if isinstance(fn, syntax.Fn):
untyped = fn
else:
# translate from the Python AST to Parakeet's untyped format
untyped = ast_conversion.translate_function_value(fn)
nonlocals = list(untyped.python_nonlocals())
arg_values = ActualArgs(nonlocals + list(args), kwargs)
# get types of all inputs
arg_types = arg_values.transform(type_conv.typeof)
return untyped, arg_values, arg_types
def from_python(self, python_fn):
import ast_conversion
untyped_fundef = ast_conversion.translate_function_value(python_fn)
closure_args = untyped_fundef.python_nonlocals()
closure_arg_types = map(type_conv.typeof, closure_args)
closure_t = make_closure_type(untyped_fundef, closure_arg_types)
closure_id = id_of_closure_type(closure_t)
def field_value(closure_arg):
obj = type_conv.from_python(closure_arg)
parakeet_type = type_conv.typeof(closure_arg)
if isinstance(parakeet_type, StructT):
return ctypes.pointer(obj)
else:
return obj
converted_args = [field_value(closure_arg) for closure_arg in closure_args]
return closure_t.ctypes_repr(closure_id, *converted_args)
def prepare_args(fn, args, kwargs):
"""
Fetch the function's nonlocals and return an ActualArgs object of both the arg
values and their types
"""
assert not isinstance(fn, TypedFn), "[prepare_args] Only works for untyped functions"
if not isinstance(fn, UntypedFn):
import ast_conversion
fn = ast_conversion.translate_function_value(fn)
nonlocals = list(fn.python_nonlocals())
arg_values = ActualArgs(nonlocals + list(args), kwargs)
# get types of all inputs
def _typeof(arg):
if hasattr(arg, 'type') and isinstance(arg.type, Type):
return arg.type
else:
return type_conv.typeof(arg)
arg_types = arg_values.transform(_typeof)
return arg_values, arg_types
closure_set = closure_type.ClosureSet(fn)
else:
assert isinstance(fn, closure_type.ClosureSet), \
"Invoke expected closure, but got %s" % (fn,)
closure_set = fn
result_type = Unknown
for closure_t in closure_set.closures:
typed_fundef = specialize(closure_t, arg_types)
result_type = result_type.combine(typed_fundef.return_type)
_invoke_type_cache[key] = result_type
return result_type
def identity(x):
return x
untyped_identity_function = ast_conversion.translate_function_value(identity)
class Annotator(Transform):
def __init__(self, tenv, var_map):
Transform.__init__(self)
self.type_env = tenv
self.var_map = var_map
def transform_expr(self, expr):
if not isinstance(expr, syntax.Expr):
expr = ast_conversion.value_to_syntax(expr)
result = Transform.transform_expr(self, expr)
def typeof_fn(f):
import ast_conversion
untyped_fn = ast_conversion.translate_function_value(f)
closure_args = untyped_fn.python_nonlocals()
closure_arg_types = map(typeof, closure_args)
return make_closure_type(untyped_fn, closure_arg_types)
def eval_expr(expr):
# print ">>", expr
if hasattr(expr, 'wrapper'):
expr = expr.wrapper
assert isinstance(expr, syntax.Expr), "Not an expression-- %s : %s" % \
(expr, type(expr))
def expr_Const():
return expr.value
def expr_Attribute():
value = eval_expr(expr.value)
if expr.name == 'offset':
if value.base is None:
return 0
else:
return value.ctypes.data - value.base.ctypes.data
else:
return getattr(value, expr.name)
def expr_ArrayView():
data = eval_expr(expr.data)
shape = eval_expr(expr.shape)
strides = eval_expr(expr.strides)
offset = eval_expr(expr.offset)
dtype = expr.type.elt_type.dtype
return np.ndarray(shape = shape,
offset = offset,
buffer = data,
strides = strides,
dtype = np.dtype(dtype))
def expr_Array():
elt_values = map(eval_expr, expr.elts)
return np.array(elt_values)
def expr_Index():
array = eval_expr(expr.value)
index = eval_expr(expr.index)
return array[index]
def expr_PrimCall():
return expr.prim.fn (*eval_args(expr.args))
def expr_Slice():
return slice(eval_expr(expr.start), eval_expr(expr.stop),
eval_expr(expr.step))
def expr_Var():
return env[expr.name]
def expr_Call():
fn = eval_expr(expr.fn)
arg_values = eval_args(expr.args)
return eval_fn(fn, arg_values)
def expr_Closure():
if isinstance(expr.fn, (syntax.Fn, syntax.TypedFn)):
fundef = expr.fn
else:
assert isinstance(expr.fn, str)
fundef = syntax.Fn.registry[expr.fn]
closure_arg_vals = map(eval_expr, expr.args)
return ClosureVal(fundef, closure_arg_vals)
def expr_Fn():
return ClosureVal(expr, [])
def expr_TypedFn():
return ClosureVal(expr, [])
def expr_Cast():
x = eval_expr(expr.value)
t = expr.type
assert isinstance(t, ScalarT)
# use numpy's conversion function
return t.dtype.type(x)
def expr_Struct():
assert expr.type, "Expected type on %s!" % expr
assert isinstance(expr.type, StructT), \
"Expected %s : %s to be a struct" % (expr, expr.type)
elts = map(eval_expr, expr.args)
return expr.type.ctypes_repr(elts)
def expr_Tuple():
return tuple(map(eval_expr, expr.elts))
def expr_TupleProj():
return eval_expr(expr.tuple)[expr.index]
def expr_ClosureElt():
assert isinstance(expr.closure, syntax.Expr), \
"Invalid closure expression-- %s : %s" % \
(expr.closure, type(expr.closure))
clos = eval_expr(expr.closure)
return clos.fixed_args[expr.index]
def expr_Range():
return np.arange(eval_expr(expr.start), eval_expr(expr.stop), eval_expr(expr.step))
def expr_Len():
return len(eval_expr(expr.value))
def expr_IndexMap():
fn = eval_expr(expr.fn)
shape = eval_expr(expr.shape)
ranges = [xrange(n) for n in shape]
def wrap_idx(idx):
if len(idx) == 1:
idx = idx[0]
return eval_fn(fn, (idx,))
elts = [wrap_idx(idx) for idx in itertools.product(*ranges)]
return np.array(elts).reshape((shape))
def expr_IndexReduce():
fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
shape = eval_expr(expr.shape)
ranges = [xrange(n) for n in shape]
acc = eval_if_expr(expr.init)
for idx in itertools.product(*ranges):
if len(idx) == 1:
idx = idx[0]
elt = eval_fn(fn, (idx,))
if acc is None:
acc = elt
else:
elt = eval_fn(combine, (acc, elt))
return elt
def expr_Map():
fn = eval_expr(expr.fn)
args = eval_args(expr.args)
axis = eval_if_expr(expr.axis)
return adverb_evaluator.eval_map(fn, args, axis)
def expr_AllPairs():
fn = eval_expr(expr.fn)
x,y = eval_args(expr.args)
axis = eval_if_expr(expr.axis)
return adverb_evaluator.eval_allpairs(fn, x, y, axis)
def expr_Reduce():
map_fn = eval_expr(expr.fn)
combine_fn = eval_expr(expr.combine)
args = eval_args(expr.args)
init = eval_expr(expr.init) if expr.init else None
axis = eval_if_expr(expr.axis)
if axis is None:
args = [np.ravel(x) for x in args]
axis = 0
return adverb_evaluator.eval_reduce(map_fn, combine_fn, init, args, axis)
def expr_Scan():
map_fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
emit = eval_expr(expr.emit)
args = eval_args(expr.args)
init = eval_expr(expr.init)
axis = eval_if_expr(expr.axis)
return adverb_evaluator.eval_scan(map_fn, combine, emit, init, args, axis)
result = dispatch(expr, 'expr')
# we don't support python function's inside parakeet,
# they have to be translated into Parakeet functions
if isinstance(result, types.FunctionType):
fundef = ast_conversion.translate_function_value(result)
return ClosureVal(fundef, fundef.python_nonlocals())
else:
return result
def typeof_fn(f): import ast_conversion untyped_fn = ast_conversion.translate_function_value(f) closure_args = untyped_fn.python_nonlocals() closure_arg_types = map(type_conv.typeof, closure_args) return make_closure_type(untyped_fn, closure_arg_types)
def eval_expr(expr):
if hasattr(expr, 'wrapper'):
expr = expr.wrapper
assert isinstance(expr, syntax.Expr), "Not an expression-- %s : %s" % \
(expr, type(expr))
def expr_Const():
return expr.value
def expr_Attribute():
value = eval_expr(expr.value)
return getattr(value, expr.name)
def expr_Array():
elt_values = map(eval_expr, expr.elts)
return np.array(elt_values)
def expr_Index():
array = eval_expr(expr.value)
index = eval_expr(expr.index)
return array[index]
def expr_PrimCall():
return expr.prim.fn (*eval_args(expr.args))
def expr_Slice():
return slice(eval_expr(expr.start), eval_expr(expr.stop),
eval_expr(expr.step))
def expr_Var():
return env[expr.name]
def expr_Call():
fn = eval_expr(expr.fn)
arg_values = eval_args(expr.args)
return eval_fn(fn, arg_values)
def expr_Closure():
if isinstance(expr.fn, (syntax.Fn, syntax.TypedFn)):
fundef = expr.fn
else:
assert isinstance(expr.fn, str)
fundef = syntax.Fn.registry[expr.fn]
closure_arg_vals = map(eval_expr, expr.args)
return ClosureVal(fundef, closure_arg_vals)
def expr_Fn():
return ClosureVal(expr, [])
def expr_TypedFn():
return ClosureVal(expr, [])
def expr_Cast():
x = eval_expr(expr.value)
t = expr.type
assert isinstance(t, ScalarT)
# use numpy's conversion function
return t.dtype.type(x)
def expr_Struct():
assert expr.type, "Expected type on %s!" % expr
assert isinstance(expr.type, StructT), \
"Expected %s : %s to be a struct" % (expr, expr.type)
elts = map(eval_expr, expr.args)
return expr.type.ctypes_repr(elts)
def expr_Tuple():
return tuple(map(eval_expr, expr.elts))
def expr_TupleProj():
return eval_expr(expr.tuple)[expr.index]
def expr_ClosureElt():
assert isinstance(expr.closure, syntax.Expr), \
"Invalid closure expression-- %s : %s" % \
(expr.closure, type(expr.closure))
clos = eval_expr(expr.closure)
return clos.fixed_args[expr.index]
def expr_Map():
fn = eval_expr(expr.fn)
args = eval_args(expr.args)
axis = syntax_helpers.unwrap_constant(expr.axis)
return adverb_evaluator.eval_map(fn, args, axis)
def expr_AllPairs():
fn = eval_expr(expr.fn)
x,y = eval_args(expr.args)
axis = syntax_helpers.unwrap_constant(expr.axis)
return adverb_evaluator.eval_allpairs(fn, x, y, axis)
def expr_Reduce():
map_fn = eval_expr(expr.fn)
combine_fn = eval_expr(expr.combine)
args = eval_args(expr.args)
init = eval_expr(expr.init) if expr.init else None
axis = syntax_helpers.unwrap_constant(expr.axis)
return adverb_evaluator.eval_reduce(map_fn, combine_fn, init, args, axis)
def expr_Scan():
map_fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
emit = eval_expr(expr.emit)
args = eval_args(expr.args)
init = eval_expr(expr.init)
axis = syntax_helpers.unwrap_constant(expr.axis)
return adverb_evaluator.eval_scan(map_fn, combine, emit, init, args, axis)
result = dispatch(expr, 'expr')
# we don't support python function's inside parakeet,
# they have to be translated into Parakeet functions
if isinstance(result, types.FunctionType):
fundef = ast_conversion.translate_function_value(result)
return ClosureVal(fundef, fundef.python_nonlocals())
else:
return result
