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 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 eval_expr(expr):
if hasattr(expr, 'wrapper'):
expr = expr.wrapper
assert isinstance(expr, Expr), "Not an expression-- %s : %s" % \
(expr, type(expr))
def expr_Const():
return expr.value
def _strides(numpy_array):
strides_bytes = numpy_array.strides
itemsize = numpy_array.dtype.itemsize
return tuple(stride / itemsize for stride in strides_bytes)
def expr_Strides():
value = eval_expr(expr.array)
return _strides(value)
def expr_Attribute():
value = eval_expr(expr.value)
if expr.name == 'offset':
if value.base is None:
return 0
else:
offset_bytes = value.ctypes.data - value.base.ctypes.data
return offset_bytes / value.dtype.itemsize
elif expr.name == 'data':
return np.ravel(value)
elif expr.name == 'strides':
return _strides(value)
elif expr.name == 'step':
step = getattr(value, 'step', 1)
if step is None: step = 1
return step
elif isinstance(value, tuple):
if expr.name.startswith('elt'):
field = int(expr.name[3:])
else:
field = int(expr.name)
return value[field]
else:
assert hasattr(value, expr.name), "Missing attribute %s from value %s" % (expr.name, value)
return getattr(value, expr.name)
def expr_Alloc():
count = eval_expr(expr.count)
arr = np.empty(shape = (count,), dtype = expr.elt_type.dtype)
return arr
def expr_AllocArray():
shape = eval_expr(expr.shape)
assert isinstance(shape, tuple), "Expected tuple, got %s" % (shape,)
assert isinstance(expr.elt_type, ScalarT), \
"Expected scalar element type for AllocArray, got %s" % (expr.elt_type,)
dtype = expr.elt_type.dtype
return np.ndarray(shape = shape, dtype = dtype)
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
bytes_per_elt = dtype.itemsize
byte_offset = offset * bytes_per_elt
byte_strides = tuple(si * bytes_per_elt for si in strides)
if False:
print expr
print "data", data
print "shape(data)", np.shape(data)
print "shape", shape
print "strides", strides
print "offset", offset
print "itemsize", bytes_per_elt
print "byte strides", byte_strides
print "byte offset", byte_offset
if isinstance(data, np.ndarray):
data = data.data
return np.ndarray(shape = shape,
offset = byte_offset,
buffer = data,
strides = byte_strides,
dtype = np.dtype(dtype))
def expr_Array():
elt_values = map(eval_expr, expr.elts)
return np.array(elt_values)
def expr_ConstArray():
shape = eval_expr(expr.shape)
value = eval_expr(expr.value)
return np.ones(shape, dtype = expr.value.type.dtype) * value
def expr_ConstArrayLike():
array = eval_expr(expr.array)
value = eval_expr(expr.value)
return np.ones_like(array, dtype = expr.value.type.elt_type.dtype) * value
def expr_TypeValue():
t = expr.type_value
assert isinstance(t, ScalarT), \
"Parakeet only supports scalar types as values, not %s" % expr
return t.dtype
def expr_Shape():
return np.shape(eval_expr(expr.array))
def expr_Reshape():
array = eval_expr(expr.array)
shape = eval_expr(expr.shape)
return array.reshape(shape)
def expr_Index():
array = eval_expr(expr.value)
index = eval_expr(expr.index)
return array[index]
def expr_PrimCall():
arg_values = eval_args(expr.args)
result = expr.prim.fn (*arg_values)
return result
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, (UntypedFn, TypedFn)):
fundef = expr.fn
else:
assert isinstance(expr.fn, str)
fundef = UntypedFn.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_Select():
cond = eval_expr(expr.cond)
trueval = eval_expr(expr.true_value)
falseval = eval_expr(expr.false_value)
return trueval if cond else falseval
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, 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 normalize_axes(axis, args):
if isinstance(axis, Expr):
axis = eval_expr(axis)
if isinstance(axis, tuple):
axes = axis
else:
axes = (axis,) * len(args)
assert len(axes) == len(args)
def expr_Map():
fn = eval_expr(expr.fn)
args = [eval_expr(arg) for arg in expr.args]
axes = normalize_axes(expr.axis, args)
largest_rank = 0
largest_arg = None
iter_space = None
for arg, axis in zip(args, axes):
rank = rankof(arg)
if rank > largest_rank and (axis is None or axis < rank):
largest_rank = rank
largest_arg = arg
if axis is None:
iter_space = largest_arg.shape
else:
iter_space = [largest_arg.shape[axis]]
if largest_rank == 0:
return eval_fn(fn, args)
results = []
for idx in np.ndindex(iter_space):
elt_args = []
for arg, axis in zip(args, axes):
r = rankof(arg)
if r == 0:
elt_args.append(arg)
if axis is None:
curr_indices = idx[-(largest_rank - r):]
elt_args.append(arg[tuple(curr_indices)])
elif r > axis:
indices = [slice(None) if j != axis else idx[0] for j in xrange(rankof(arg)) ]
elt_args.append(arg[tuple(indices)])
else:
elt_args.append(arg)
results.append(eval_fn(fn, elt_args))
return np.array(results)
def expr_Reduce():
fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
init = eval_expr(expr.init) if expr.init else None
args = [eval_expr(arg) for arg in expr.args]
if isinstance(expr.axis, (int, long)):
axis = expr.axis
else:
axis = eval_expr(expr.axis)
if axis is None:
args = [np.ravel(arg) for arg in args]
axis = 0
largest_rank = 0
largest_arg = None
for arg in args:
rank = rankof(arg)
if rank > largest_rank:
largest_rank = rank
largest_arg = arg
if largest_rank == 0:
acc = eval_fn(fn, args)
if init is None: return acc
else: return eval_fn(combine, [init, acc])
niters = largest_arg.shape[axis]
if init is not None: acc = init
else: acc = None
for i in xrange(niters):
elt_args = []
for arg in args:
r = rankof(arg)
if r > 0:
indices = [slice(None) if j != axis else i for j in xrange(rankof(arg)) ]
elt_args.append(arg[tuple(indices)])
else:
elt_args.append(arg)
elt_result = eval_fn(fn, elt_args)
if acc is not None:
acc = eval_fn(combine, [acc, elt_result])
else:
acc = elt_result
return acc
def expr_Scan():
assert False, "Scan not implemented"
def expr_IndexMap():
fn = eval_expr(expr.fn)
shape = eval_expr(expr.shape)
dtype = expr.type.elt_type.dtype
result = np.empty(shape, dtype = dtype)
for idx in np.ndindex(shape):
result[idx] = eval_fn(fn, (idx,))
return result
def expr_IndexReduce():
fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
shape = eval_expr(expr.shape)
if not isinstance(shape, (list, tuple) ):
shape = [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
fn_name = "expr_" + expr.__class__.__name__
dispatch_fn = locals()[fn_name]
result = dispatch_fn()
# 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 eval_expr(expr):
if hasattr(expr, 'wrapper'):
expr = expr.wrapper
assert isinstance(expr, Expr), "Not an expression-- %s : %s" % \
(expr, type(expr))
def expr_Const():
return expr.value
def _strides(numpy_array):
strides_bytes = numpy_array.strides
itemsize = numpy_array.dtype.itemsize
return tuple(stride / itemsize for stride in strides_bytes)
def expr_Strides():
value = eval_expr(expr.array)
return _strides(value)
def expr_Attribute():
value = eval_expr(expr.value)
if expr.name == 'offset':
if value.base is None:
return 0
else:
offset_bytes = value.ctypes.data - value.base.ctypes.data
return offset_bytes / value.dtype.itemsize
elif expr.name == 'data':
return np.ravel(value)
elif expr.name == 'strides':
return _strides(value)
elif expr.name == 'step':
step = getattr(value, 'step', 1)
if step is None: step = 1
return step
elif isinstance(value, tuple):
if expr.name.startswith('elt'):
field = int(expr.name[3:])
else:
field = int(expr.name)
return value[field]
else:
assert hasattr(
value, expr.name), "Missing attribute %s from value %s" % (
expr.name, value)
return getattr(value, expr.name)
def expr_Alloc():
count = eval_expr(expr.count)
arr = np.empty(shape=(count, ), dtype=expr.elt_type.dtype)
return arr
def expr_AllocArray():
shape = eval_expr(expr.shape)
assert isinstance(shape,
tuple), "Expected tuple, got %s" % (shape, )
assert isinstance(expr.elt_type, ScalarT), \
"Expected scalar element type for AllocArray, got %s" % (expr.elt_type,)
dtype = expr.elt_type.dtype
return np.ndarray(shape=shape, dtype=dtype)
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
bytes_per_elt = dtype.itemsize
byte_offset = offset * bytes_per_elt
byte_strides = tuple(si * bytes_per_elt for si in strides)
if False:
print expr
print "data", data
print "shape(data)", np.shape(data)
print "shape", shape
print "strides", strides
print "offset", offset
print "itemsize", bytes_per_elt
print "byte strides", byte_strides
print "byte offset", byte_offset
if isinstance(data, np.ndarray):
data = data.data
return np.ndarray(shape=shape,
offset=byte_offset,
buffer=data,
strides=byte_strides,
dtype=np.dtype(dtype))
def expr_Array():
elt_values = map(eval_expr, expr.elts)
return np.array(elt_values)
def expr_ConstArray():
shape = eval_expr(expr.shape)
value = eval_expr(expr.value)
return np.ones(shape, dtype=expr.value.type.dtype) * value
def expr_ConstArrayLike():
array = eval_expr(expr.array)
value = eval_expr(expr.value)
return np.ones_like(array,
dtype=expr.value.type.elt_type.dtype) * value
def expr_TypeValue():
t = expr.type_value
assert isinstance(t, ScalarT), \
"Parakeet only supports scalar types as values, not %s" % expr
return t.dtype
def expr_Shape():
return np.shape(eval_expr(expr.array))
def expr_Reshape():
array = eval_expr(expr.array)
shape = eval_expr(expr.shape)
return array.reshape(shape)
def expr_Index():
array = eval_expr(expr.value)
index = eval_expr(expr.index)
return array[index]
def expr_PrimCall():
arg_values = eval_args(expr.args)
result = expr.prim.fn(*arg_values)
return result
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, (UntypedFn, TypedFn)):
fundef = expr.fn
else:
assert isinstance(expr.fn, str)
fundef = UntypedFn.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_Select():
cond = eval_expr(expr.cond)
trueval = eval_expr(expr.true_value)
falseval = eval_expr(expr.false_value)
return trueval if cond else falseval
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, 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 normalize_axes(axis, args):
if isinstance(axis, Expr):
axis = eval_expr(axis)
if isinstance(axis, tuple):
axes = axis
else:
axes = (axis, ) * len(args)
assert len(axes) == len(args)
def expr_Map():
fn = eval_expr(expr.fn)
args = [eval_expr(arg) for arg in expr.args]
axes = normalize_axes(expr.axis, args)
largest_rank = 0
largest_arg = None
iter_space = None
for arg, axis in zip(args, axes):
rank = rankof(arg)
if rank > largest_rank and (axis is None or axis < rank):
largest_rank = rank
largest_arg = arg
if axis is None:
iter_space = largest_arg.shape
else:
iter_space = [largest_arg.shape[axis]]
if largest_rank == 0:
return eval_fn(fn, args)
results = []
for idx in np.ndindex(iter_space):
elt_args = []
for arg, axis in zip(args, axes):
r = rankof(arg)
if r == 0:
elt_args.append(arg)
if axis is None:
curr_indices = idx[-(largest_rank - r):]
elt_args.append(arg[tuple(curr_indices)])
elif r > axis:
indices = [
slice(None) if j != axis else idx[0]
for j in xrange(rankof(arg))
]
elt_args.append(arg[tuple(indices)])
else:
elt_args.append(arg)
results.append(eval_fn(fn, elt_args))
return np.array(results)
def expr_Reduce():
fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
init = eval_expr(expr.init) if expr.init else None
args = [eval_expr(arg) for arg in expr.args]
if isinstance(expr.axis, (int, long)):
axis = expr.axis
else:
axis = eval_expr(expr.axis)
if axis is None:
args = [np.ravel(arg) for arg in args]
axis = 0
largest_rank = 0
largest_arg = None
for arg in args:
rank = rankof(arg)
if rank > largest_rank:
largest_rank = rank
largest_arg = arg
if largest_rank == 0:
acc = eval_fn(fn, args)
if init is None: return acc
else: return eval_fn(combine, [init, acc])
niters = largest_arg.shape[axis]
if init is not None: acc = init
else: acc = None
for i in xrange(niters):
elt_args = []
for arg in args:
r = rankof(arg)
if r > 0:
indices = [
slice(None) if j != axis else i
for j in xrange(rankof(arg))
]
elt_args.append(arg[tuple(indices)])
else:
elt_args.append(arg)
elt_result = eval_fn(fn, elt_args)
if acc is not None:
acc = eval_fn(combine, [acc, elt_result])
else:
acc = elt_result
return acc
def expr_Scan():
assert False, "Scan not implemented"
def expr_IndexMap():
fn = eval_expr(expr.fn)
shape = eval_expr(expr.shape)
dtype = expr.type.elt_type.dtype
result = np.empty(shape, dtype=dtype)
for idx in np.ndindex(shape):
result[idx] = eval_fn(fn, (idx, ))
return result
def expr_IndexReduce():
fn = eval_expr(expr.fn)
combine = eval_expr(expr.combine)
shape = eval_expr(expr.shape)
if not isinstance(shape, (list, tuple)):
shape = [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
fn_name = "expr_" + expr.__class__.__name__
dispatch_fn = locals()[fn_name]
result = dispatch_fn()
# 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
