def _analyze_op_pair_first(self, scope, equiv_set, expr):
# make dummy lhs since we don't have access to lhs
typ = self.typemap[expr.value.name].first_type
if not isinstance(typ, types.NamedTuple):
return None
lhs = ir.Var(scope, mk_unique_var("tuple_var"), expr.loc)
self.typemap[lhs.name] = typ
rhs = ir.Expr.pair_first(expr.value, expr.loc)
lhs_assign = ir.Assign(rhs, lhs, expr.loc)
#(shape, post) = self._gen_shape_call(equiv_set, lhs, typ.count, )
var = lhs
out = []
size_vars = []
ndims = typ.count
for i in range(ndims):
# get size: Asize0 = A_sh_attr[0]
size_var = ir.Var(var.scope, mk_unique_var(
"{}_size{}".format(var.name, i)), var.loc)
getitem = ir.Expr.static_getitem(lhs, i, None, var.loc)
self.calltypes[getitem] = None
out.append(ir.Assign(getitem, size_var, var.loc))
self._define(equiv_set, size_var, types.intp, getitem)
size_vars.append(size_var)
shape = tuple(size_vars)
return shape, [lhs_assign] + out
def _mk_unique_var(self, prefix):
"""Make unique var name checking self.func_ir._definitions"""
name = mk_unique_var(prefix)
while name in self.func_ir._definitions:
name = mk_unique_var(prefix)
return name
def _handle_f_close_call(self, stmt, lhs_var, rhs):
func_def = guard(get_definition, self.func_ir, rhs.func)
assert func_def is not None
# rare case where function variable is assigned to a new variable
if isinstance(func_def, ir.Var):
rhs.func = func_def
return self._handle_f_close_call(stmt, lhs_var, rhs)
if (isinstance(func_def, ir.Expr) and func_def.op == 'getattr'
and func_def.value.name in self.h5_files
and func_def.attr == 'close'):
f_id = func_def.value
scope = lhs_var.scope
loc = lhs_var.loc
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5close_attr = getattr(g_pio_var, h5close)
h5close_attr_call = ir.Expr.getattr(g_pio_var, "h5close", loc)
attr_var = ir.Var(scope, mk_unique_var("$h5close_attr"), loc)
attr_assign = ir.Assign(h5close_attr_call, attr_var, loc)
# h5close(f_id)
close_call = ir.Expr.call(attr_var, [f_id], (), loc)
close_assign = ir.Assign(close_call, lhs_var, loc)
return [g_pio_assign, attr_assign, close_assign]
return None
def _handle_merge(self, lhs, rhs):
if guard(find_callname, self.func_ir, rhs) == ('merge', 'pandas'):
if len(rhs.args) < 2:
raise ValueError("left and right arguments required for merge")
left_df = rhs.args[0]
right_df = rhs.args[1]
kws = dict(rhs.kws)
if 'on' in kws:
left_on = get_constant(self.func_ir, kws['on'], None)
right_on = left_on
else: # pragma: no cover
if 'left_on' not in kws or 'right_on' not in kws:
raise ValueError("merge 'on' or 'left_on'/'right_on'"
"arguments required")
left_on = get_constant(self.func_ir, kws['left_on'], None)
right_on = get_constant(self.func_ir, kws['right_on'], None)
if left_on is None or right_on is None:
raise ValueError("merge key values should be constant strings")
scope = lhs.scope
loc = lhs.loc
self.df_vars[lhs.name] = {}
# add columns from left to output
for col, _ in self.df_vars[left_df.name].items():
self.df_vars[lhs.name][col] = ir.Var(scope, mk_unique_var(col),
loc)
# add columns from right to output
for col, _ in self.df_vars[right_df.name].items():
self.df_vars[lhs.name][col] = ir.Var(scope, mk_unique_var(col),
loc)
self._update_df_cols()
return [
hiframes_join.Join(lhs.name, left_df.name, right_df.name,
left_on, right_on, self.df_vars, lhs.loc)
]
return None
def convert_seq_to_atleast_3d(self, seq_arg_var):
"""convert array sequence to a sequence of arrays with at least 3d dims
"""
arr_args = self._get_sequence_arrs(seq_arg_var.name)
new_seq = []
for arr in arr_args:
curr_dims = self._get_ndims(arr.name)
if curr_dims < 3:
dummy_var = ir.Var(
arr.scope, mk_unique_var("__dummy_nd"), arr.loc)
self.typemap[dummy_var.name] = self.typemap[arr.name].copy(
ndim=3)
corrs = self.array_shape_classes[arr.name].copy()
if curr_dims == 0:
corrs = [CONST_CLASS] * 3
elif curr_dims == 1: # Numpy adds both prepends and appends a dim
corrs = [CONST_CLASS] + corrs + [CONST_CLASS]
elif curr_dims == 2: # append a dim
corrs = corrs + [CONST_CLASS]
self.array_shape_classes[dummy_var.name] = corrs
new_seq.append(dummy_var)
else:
new_seq.append(arr)
tup_name = mk_unique_var("__dummy_tup")
self.tuple_table[tup_name] = new_seq
return ir.Var(arr_args[0].scope, tup_name, arr_args[0].loc)
def _get_stencil_last_ind(self, dim_size, end_length, gen_nodes, scope,
loc):
last_ind = dim_size
if end_length != 0:
# set last index to size minus stencil size to avoid invalid
# memory access
index_const = ir.Var(scope, mk_unique_var("stencil_const_var"),
loc)
self.typemap[index_const.name] = types.intp
if isinstance(end_length, numbers.Number):
const_assign = ir.Assign(ir.Const(end_length, loc),
index_const, loc)
else:
const_assign = ir.Assign(end_length, index_const, loc)
gen_nodes.append(const_assign)
last_ind = ir.Var(scope, mk_unique_var("last_ind"), loc)
self.typemap[last_ind.name] = types.intp
g_var = ir.Var(scope, mk_unique_var("compute_last_ind_var"), loc)
check_func = numba.njit(_compute_last_ind)
func_typ = types.functions.Dispatcher(check_func)
self.typemap[g_var.name] = func_typ
g_obj = ir.Global("_compute_last_ind", check_func, loc)
g_assign = ir.Assign(g_obj, g_var, loc)
gen_nodes.append(g_assign)
index_call = ir.Expr.call(g_var, [dim_size, index_const], (), loc)
self.calltypes[index_call] = func_typ.get_call_type(
self.typingctx, [types.intp, types.intp], {})
index_assign = ir.Assign(index_call, last_ind, loc)
gen_nodes.append(index_assign)
return last_ind
def _gen_size_call(self, var, i):
out = []
ndims = self._get_ndims(var.name)
# attr call: A_sh_attr = getattr(A, shape)
shape_attr_call = ir.Expr.getattr(var, "shape", var.loc)
attr_var = ir.Var(var.scope,
mk_unique_var(var.name + "_sh_attr" + str(i)),
var.loc)
self.typemap[attr_var.name] = types.containers.UniTuple(
types.intp, ndims)
attr_assign = ir.Assign(shape_attr_call, attr_var, var.loc)
out.append(attr_assign)
# const var for dim: $constA0 = Const(0)
const_node = ir.Const(i, var.loc)
const_var = ir.Var(var.scope,
mk_unique_var("$const" + var.name + str(i)),
var.loc)
self.typemap[const_var.name] = types.intp
const_assign = ir.Assign(const_node, const_var, var.loc)
out.append(const_assign)
# get size: Asize0 = A_sh_attr[0]
size_var = ir.Var(var.scope, mk_unique_var(var.name + "size" + str(i)),
var.loc)
self.typemap[size_var.name] = types.intp
getitem_node = ir.Expr.static_getitem(attr_var, i, const_var, var.loc)
self.calltypes[getitem_node] = None
getitem_assign = ir.Assign(getitem_node, size_var, var.loc)
out.append(getitem_assign)
return out
def inline_array(array_var, expr, stmts, list_vars, dels):
"""Check to see if the given "array_var" is created from a list
of constants, and try to inline the list definition as array
initialization.
Extra statements produced with be appended to "stmts".
"""
callname = guard(find_callname, func_ir, expr)
require(callname and callname[1] == 'numpy' and callname[0] == 'array')
require(expr.args[0].name in list_vars)
ret_type = calltypes[expr].return_type
require(isinstance(ret_type, types.ArrayCompatible) and
ret_type.ndim == 1)
loc = expr.loc
list_var = expr.args[0]
array_typ = typemap[array_var.name]
debug_print("inline array_var = ", array_var, " list_var = ", list_var)
dtype = array_typ.dtype
seq, op = find_build_sequence(func_ir, list_var)
size = len(seq)
size_var = ir.Var(scope, mk_unique_var("size"), loc)
size_tuple_var = ir.Var(scope, mk_unique_var("size_tuple"), loc)
size_typ = types.intp
size_tuple_typ = types.UniTuple(size_typ, 1)
typemap[size_var.name] = size_typ
typemap[size_tuple_var.name] = size_tuple_typ
stmts.append(_new_definition(func_ir, size_var,
ir.Const(size, loc=loc), loc))
stmts.append(_new_definition(func_ir, size_tuple_var,
ir.Expr.build_tuple(items=[size_var], loc=loc), loc))
empty_func = ir.Var(scope, mk_unique_var("empty_func"), loc)
fnty = get_np_ufunc_typ(np.empty)
sig = context.resolve_function_type(fnty, (size_typ,), {})
typemap[empty_func.name] = fnty #
stmts.append(_new_definition(func_ir, empty_func,
ir.Global('empty', np.empty, loc=loc), loc))
empty_call = ir.Expr.call(empty_func, [size_var], {}, loc=loc)
calltypes[empty_call] = typing.signature(array_typ, size_typ)
stmts.append(_new_definition(func_ir, array_var, empty_call, loc))
for i in range(size):
index_var = ir.Var(scope, mk_unique_var("index"), loc)
index_typ = types.intp
typemap[index_var.name] = index_typ
stmts.append(_new_definition(func_ir, index_var,
ir.Const(i, loc), loc))
setitem = ir.SetItem(array_var, index_var, seq[i], loc)
calltypes[setitem] = typing.signature(types.none, array_typ,
index_typ, dtype)
stmts.append(setitem)
stmts.extend(dels)
return True
def gen_empty_like(in_arr, out_arr):
scope = in_arr.scope
loc = in_arr.loc
# g_np_var = Global(numpy)
g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
g_np = ir.Global('np', np, loc)
g_np_assign = ir.Assign(g_np, g_np_var, loc)
# attr call: empty_attr = getattr(g_np_var, empty_like)
empty_attr_call = ir.Expr.getattr(g_np_var, "empty_like", loc)
attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)
attr_assign = ir.Assign(empty_attr_call, attr_var, loc)
# alloc call: out_arr = empty_attr(in_arr)
alloc_call = ir.Expr.call(attr_var, [in_arr], (), loc)
alloc_assign = ir.Assign(alloc_call, out_arr, loc)
return [g_np_assign, attr_assign, alloc_assign]
def gen_stencil_call(in_arr, out_arr, kernel_func, index_offsets, fir_globals,
other_args=None, options=None):
if other_args is None:
other_args = []
if options is None:
options = {}
if index_offsets != [0]:
options['index_offsets'] = index_offsets
scope = in_arr.scope
loc = in_arr.loc
stencil_nodes = []
stencil_nodes += gen_empty_like(in_arr, out_arr)
kernel_var = ir.Var(scope, mk_unique_var("kernel_var"), scope)
if not isinstance(kernel_func, ir.Expr):
kernel_func = ir.Expr.make_function("kernel", kernel_func.__code__,
kernel_func.__closure__, kernel_func.__defaults__, loc)
stencil_nodes.append(ir.Assign(kernel_func, kernel_var, loc))
def f(A, B, f):
numba.stencil(f)(A, out=B)
f_block = compile_to_numba_ir(f, {'numba': numba}).blocks.popitem()[1]
replace_arg_nodes(f_block, [in_arr, out_arr, kernel_var])
stencil_nodes += f_block.body[:-3] # remove none return
setup_call = stencil_nodes[-2].value
stencil_call = stencil_nodes[-1].value
setup_call.kws = list(options.items())
stencil_call.args += other_args
return stencil_nodes
def _gen_col_var(self, out_var, args, col_var):
loc = out_var.loc
scope = out_var.scope
# calculate mean first
mean_var = ir.Var(scope, mk_unique_var("mean_val"), loc)
f_mean_blocks = self._gen_col_mean(mean_var, args, col_var)
f_mean_blocks = add_offset_to_labels(f_mean_blocks, ir_utils._max_label+1)
ir_utils._max_label = max(f_mean_blocks.keys())
m_last_label = find_topo_order(f_mean_blocks)[-1]
remove_none_return_from_block(f_mean_blocks[m_last_label])
def f(A, s, m):
count = 0
for i in numba.parfor.prange(len(A)):
val = A[i]
if not np.isnan(val):
s += (val-m)**2
count += 1
if count <= 1:
s = np.nan
else:
s = s/(count-1)
f_blocks = get_inner_ir(f)
replace_var_names(f_blocks, {'A': col_var.name})
replace_var_names(f_blocks, {'s': out_var.name})
replace_var_names(f_blocks, {'m': mean_var.name})
f_blocks[0].body.insert(0, ir.Assign(ir.Const(0.0, loc), out_var, loc))
# attach first var block to last mean block
f_mean_blocks[m_last_label].body.extend(f_blocks[0].body)
f_blocks.pop(0)
f_blocks = add_offset_to_labels(f_blocks, ir_utils._max_label+1)
# add offset to jump of first f_block since it didn't go through call
f_mean_blocks[m_last_label].body[-1].target += ir_utils._max_label+1
ir_utils._max_label = max(f_blocks.keys())
f_mean_blocks.update(f_blocks)
return f_mean_blocks
def gen_parquet_read(self, file_name, table_types):
import pyarrow.parquet as pq
scope = file_name.scope
loc = file_name.loc
if table_types is None:
fname_def = guard(get_definition, self.func_ir, file_name)
if not isinstance(fname_def, ir.Const) or not isinstance(
fname_def.value, str):
raise ValueError("Parquet schema not available")
file_name_str = fname_def.value
col_names, col_types = parquet_file_schema(file_name_str)
else:
col_names = list(table_types.keys())
col_types = list(table_types.values())
out_nodes = []
col_items = []
for i, cname in enumerate(col_names):
# get column type from schema
c_type = col_types[i]
# create a variable for column and assign type
varname = mk_unique_var(cname)
#self.locals[varname] = c_type
cvar = ir.Var(scope, varname, loc)
col_items.append((cname, cvar))
out_nodes += get_column_read_nodes(c_type, cvar, file_name, i)
return col_items, out_nodes
def replace_var_with_array_in_block(vars, block, typemap, calltypes):
new_block = []
for inst in block.body:
if isinstance(inst, ir.Assign) and inst.target.name in vars:
const_node = ir.Const(0, inst.loc)
const_var = ir.Var(inst.target.scope,
mk_unique_var("$const_ind_0"), inst.loc)
typemap[const_var.name] = types.uintp
const_assign = ir.Assign(const_node, const_var, inst.loc)
new_block.append(const_assign)
setitem_node = ir.SetItem(inst.target, const_var, inst.value,
inst.loc)
calltypes[setitem_node] = signature(
types.none,
types.npytypes.Array(typemap[inst.target.name], 1, "C"),
types.intp, typemap[inst.target.name])
new_block.append(setitem_node)
continue
elif isinstance(inst, parfor.Parfor):
replace_var_with_array_internal(vars, {0: inst.init_block},
typemap, calltypes)
replace_var_with_array_internal(vars, inst.loop_body, typemap,
calltypes)
new_block.append(inst)
return new_block
def _gen_h5close(self, stmt, f_id):
lhs_var = stmt.target
scope = lhs_var.scope
loc = lhs_var.loc
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5close_attr = getattr(g_pio_var, h5close)
h5close_attr_call = ir.Expr.getattr(g_pio_var, "h5close", loc)
attr_var = ir.Var(scope, mk_unique_var("$h5close_attr"), loc)
attr_assign = ir.Assign(h5close_attr_call, attr_var, loc)
# h5close(f_id)
close_call = ir.Expr.call(attr_var, [f_id], (), loc)
close_assign = ir.Assign(close_call, lhs_var, loc)
return [g_pio_assign, attr_assign, close_assign]
def _get_slice_range(self, index_slice, out):
scope = index_slice.scope
loc = index_slice.loc
# start = s.start
start_var = ir.Var(scope, mk_unique_var("$pio_range_start"), loc)
start_attr_call = ir.Expr.getattr(index_slice, "start", loc)
start_assign = ir.Assign(start_attr_call, start_var, loc)
# stop = s.stop
stop_var = ir.Var(scope, mk_unique_var("$pio_range_stop"), loc)
stop_attr_call = ir.Expr.getattr(index_slice, "stop", loc)
stop_assign = ir.Assign(stop_attr_call, stop_var, loc)
# size = stop-start
size_var = ir.Var(scope, mk_unique_var("$pio_range_size"), loc)
size_call = ir.Expr.binop('-', stop_var, start_var, loc)
size_assign = ir.Assign(size_call, size_var, loc)
out += [start_assign, stop_assign, size_assign]
return [start_var], [size_var]
def _gen_rolling_init(self, win_size, func, center):
nodes = []
right_length = 0
scope = win_size.scope
loc = win_size.loc
right_length = ir.Var(scope, mk_unique_var('zero_var'), scope)
nodes.append(ir.Assign(ir.Const(0, loc), right_length, win_size.loc))
def f(w):
return -w + 1
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [win_size])
nodes.extend(f_block.body[:-2]) # remove none return
left_length = nodes[-1].target
if center:
def f(w):
return -(w // 2)
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [win_size])
nodes.extend(f_block.body[:-2]) # remove none return
left_length = nodes[-1].target
def f(w):
return (w // 2)
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [win_size])
nodes.extend(f_block.body[:-2]) # remove none return
right_length = nodes[-1].target
def f(a, b):
return ((a, b), )
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [left_length, right_length])
nodes.extend(f_block.body[:-2]) # remove none return
win_tuple = nodes[-1].target
index_offsets = [right_length]
if func == 'apply':
index_offsets = [left_length]
def f(a):
return (a, )
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, index_offsets)
nodes.extend(f_block.body[:-2]) # remove none return
index_offsets = nodes[-1].target
return index_offsets, win_tuple, nodes
def _gen_h5create_group(self, stmt, f_id):
lhs_var = stmt.target
scope = lhs_var.scope
loc = lhs_var.loc
args = [f_id] + stmt.value.args
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5create_group_attr = getattr(g_pio_var, h5create_group)
h5create_group_attr_call = ir.Expr.getattr(g_pio_var, "h5create_group",
loc)
attr_var = ir.Var(scope, mk_unique_var("$h5create_group_attr"), loc)
attr_assign = ir.Assign(h5create_group_attr_call, attr_var, loc)
# group_id = h5create_group(f_id)
create_group_call = ir.Expr.call(attr_var, args, (), loc)
create_group_assign = ir.Assign(create_group_call, lhs_var, loc)
# add to files since group behavior is same as files for many calls
self.h5_files[lhs_var.name] = "group"
return [g_pio_assign, attr_assign, create_group_assign]
def gen_stencil_call(in_arr, out_arr, code_expr, index_offsets):
scope = in_arr.scope
loc = in_arr.loc
alloc_nodes = gen_empty_like(in_arr, out_arr)
# generate stencil call
# g_numba_var = Global(numba)
g_numba_var = ir.Var(scope, mk_unique_var("$g_numba_var"), loc)
g_dist = ir.Global('numba', numba, loc)
g_numba_assign = ir.Assign(g_dist, g_numba_var, loc)
# attr call: stencil_attr = getattr(g_numba_var, stencil)
stencil_attr_call = ir.Expr.getattr(g_numba_var, "stencil", loc)
stencil_attr_var = ir.Var(scope, mk_unique_var("$stencil_attr"), loc)
stencil_attr_assign = ir.Assign(stencil_attr_call, stencil_attr_var, loc)
# stencil_out = numba.stencil()
stencil_out = ir.Var(scope, mk_unique_var("$stencil_out"), loc)
stencil_call = ir.Expr.call(stencil_attr_var, [in_arr, out_arr], (), loc)
stencil_call.stencil_def = code_expr
stencil_call.index_offsets = index_offsets
stencil_assign = ir.Assign(stencil_call, stencil_out, loc)
return alloc_nodes + [g_numba_assign, stencil_attr_assign, stencil_assign]
def _handle_h5_File_call(self, assign, lhs, rhs):
"""
Handle h5py.File calls like:
f = h5py.File(file_name, mode)
"""
# parallel arg = False for this stage
loc = lhs.loc
scope = lhs.scope
parallel_var = ir.Var(scope, mk_unique_var("$const_parallel"), loc)
parallel_assign = ir.Assign(ir.Const(0, loc), parallel_var, loc)
rhs.args.append(parallel_var)
return [parallel_assign, assign]
def _gen_h5size(self, f_id, dset, ndims, scope, loc, out):
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5size_attr = getattr(g_pio_var, h5size)
h5size_attr_call = ir.Expr.getattr(g_pio_var, "h5size", loc)
attr_var = ir.Var(scope, mk_unique_var("$h5size_attr"), loc)
attr_assign = ir.Assign(h5size_attr_call, attr_var, loc)
out += [g_pio_assign, attr_assign]
size_vars = []
for i in range(ndims):
dim_var = ir.Var(scope, mk_unique_var("$h5_dim_var"), loc)
dim_assign = ir.Assign(ir.Const(np.int32(i), loc), dim_var, loc)
out.append(dim_assign)
size_var = ir.Var(scope, mk_unique_var("$h5_size_var"), loc)
size_vars.append(size_var)
size_call = ir.Expr.call(attr_var, [f_id, dset, dim_var], (), loc)
size_assign = ir.Assign(size_call, size_var, loc)
out.append(size_assign)
return size_vars
def _gen_h5create_dset(self, stmt, f_id):
lhs_var = stmt.target
scope = lhs_var.scope
loc = lhs_var.loc
args = [f_id] + stmt.value.args
# append the dtype arg (e.g. dtype='f8')
assert stmt.value.kws and stmt.value.kws[0][0] == 'dtype'
args.append(stmt.value.kws[0][1])
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5create_dset_attr = getattr(g_pio_var, h5create_dset)
h5create_dset_attr_call = ir.Expr.getattr(g_pio_var, "h5create_dset",
loc)
attr_var = ir.Var(scope, mk_unique_var("$h5create_dset_attr"), loc)
attr_assign = ir.Assign(h5create_dset_attr_call, attr_var, loc)
# dset_id = h5create_dset(f_id)
create_dset_call = ir.Expr.call(attr_var, args, (), loc)
create_dset_assign = ir.Assign(create_dset_call, lhs_var, loc)
self.h5_dsets[lhs_var.name] = (f_id, args[1])
self.h5_dsets_sizes[lhs_var.name] = self.tuple_table[args[2].name]
return [g_pio_assign, attr_assign, create_dset_assign]
def _handle_h5_File_call(self, assign, lhs, rhs):
"""
Handle h5py.File calls like:
f = h5py.File(file_name, mode)
"""
if guard(find_callname, self.func_ir, rhs) == ('File', 'h5py'):
self.h5_files[lhs.name] = rhs.args[0]
# parallel arg = False for this stage
loc = lhs.loc
scope = lhs.scope
parallel_var = ir.Var(scope, mk_unique_var("$const_parallel"), loc)
parallel_assign = ir.Assign(ir.Const(0, loc), parallel_var, loc)
rhs.args.append(parallel_var)
return [parallel_assign, assign]
return None
def _gen_col_std(self, out_var, args, col_var):
loc = out_var.loc
scope = out_var.scope
# calculate var() first
var_var = ir.Var(scope, mk_unique_var("var_val"), loc)
v_nodes = self._gen_col_var(var_var, args, col_var)
def f(a):
a**0.5
s_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(s_block, [var_var])
s_nodes = s_block.body[:-3]
assert len(s_nodes) == 3
s_nodes[-1].target = out_var
return v_nodes + s_nodes
def gen_parquet_read(self, file_name):
import pyarrow.parquet as pq
fname_def = guard(get_definition, self.func_ir, file_name)
if isinstance(fname_def, ir.Const):
assert isinstance(fname_def.value, str)
file_name_str = fname_def.value
col_names, col_types = parquet_file_schema(file_name_str)
scope = file_name.scope
loc = file_name.loc
out_nodes = []
col_items = []
for i, cname in enumerate(col_names):
# get column type from schema
c_type = col_types[i]
# create a variable for column and assign type
varname = mk_unique_var(cname)
self.locals[varname] = c_type
cvar = ir.Var(scope, varname, loc)
col_items.append((cname, cvar))
size_func_text = (
'def f():\n col_size = get_column_size_parquet("{}", {})\n'
.format(file_name_str, i))
size_func_text += ' column = np.empty(col_size, dtype=np.{})\n'.format(
c_type.dtype)
size_func_text += ' status = read_parquet("{}", {}, column)\n'.format(
file_name_str, i)
loc_vars = {}
exec(size_func_text, {}, loc_vars)
size_func = loc_vars['f']
_, f_block = compile_to_numba_ir(
size_func, {
'get_column_size_parquet': get_column_size_parquet,
'read_parquet': read_parquet,
'np': np
}).blocks.popitem()
out_nodes += f_block.body[:-3]
for stmt in out_nodes:
if stmt.target.name.startswith("column"):
assign = ir.Assign(stmt.target, cvar, loc)
break
out_nodes.append(assign)
return col_items, out_nodes
raise ValueError("Parquet schema not available")
def inline_calls_inner(func_ir, block, stmt, i, py_func):
call_expr = stmt.value
scope = block.scope
callee_ir = numba.compiler.run_frontend(py_func)
# relabel callee_ir by adding an offset
max_label = max(func_ir.blocks.keys())
callee_blocks = add_offset_to_labels(callee_ir.blocks, max_label + 1)
callee_ir.blocks = callee_blocks
min_label = min(callee_blocks.keys())
max_label = max(callee_blocks.keys())
# init _max_label global in ir_utils before using next_label()
ir_utils._max_label = max_label
# rename all variables in callee blocks
var_table = get_name_var_table(callee_ir.blocks)
new_var_dict = {}
for name, var in var_table.items():
new_var = scope.define(mk_unique_var(var.name), loc=var.loc)
new_var_dict[name] = new_var
replace_vars(callee_ir.blocks, new_var_dict)
# replace callee arguments
args = list(call_expr.args)
# TODO: replace defaults (add to args)
_replace_args(callee_ir.blocks, args)
# split caller blocks into two
new_block = ir.Block(scope, block.loc)
new_block.body = block.body[i + 1:]
new_label = ir_utils.next_label()
func_ir.blocks[new_label] = new_block
block.body = block.body[:i]
block.body.append(ir.Jump(min_label, stmt.loc))
# replace Return with assignment to LHS
_replace_returns(callee_ir.blocks, stmt.target, new_label)
# insert all new blocks
for label, bl in callee_ir.blocks.items():
func_ir.blocks[label] = bl
# run inline_calls recursively to transform other calls
inline_calls(func_ir)
return
def fix_dependencies(expr, varlist):
"""Double check if all variables in varlist are defined before
expr is used. Try to move constant definition when the check fails.
Bails out by raising GuardException if it can't be moved.
"""
debug_print = _make_debug_print("fix_dependencies")
for label, block in blocks.items():
scope = block.scope
body = block.body
defined = set()
for i in range(len(body)):
inst = body[i]
if isinstance(inst, ir.Assign):
defined.add(inst.target.name)
if inst.value == expr:
new_varlist = []
for var in varlist:
# var must be defined before this inst, or live
# and not later defined.
if (var.name in defined or
(var.name in livemap[label]
and not (var.name in usedefs.defmap[label]))):
debug_print(var.name, " already defined")
new_varlist.append(var)
else:
debug_print(var.name, " not yet defined")
var_def = get_definition(func_ir, var.name)
if isinstance(var_def, ir.Const):
loc = var.loc
new_var = ir.Var(scope,
mk_unique_var("new_var"),
loc)
new_const = ir.Const(var_def.value, loc)
new_vardef = _new_definition(
func_ir, new_var, new_const, loc)
new_body = []
new_body.extend(body[:i])
new_body.append(new_vardef)
new_body.extend(body[i:])
block.body = new_body
new_varlist.append(new_var)
else:
raise GuardException
return new_varlist
# when expr is not found in block
raise GuardException
def gen_parquet_read(self, file_name, lhs):
import pyarrow.parquet as pq
scope = file_name.scope
loc = file_name.loc
table_types = None
# lhs is temporary and will possibly be assigned to user variable
assert lhs.name.startswith('$')
if lhs.name in self.reverse_copies and self.reverse_copies[lhs.name] in self.locals:
table_types = self.locals[self.reverse_copies[lhs.name]]
self.locals.pop(self.reverse_copies[lhs.name])
convert_types = {}
# user-specified type conversion
if lhs.name in self.reverse_copies and (self.reverse_copies[lhs.name]+':convert') in self.locals:
convert_types = self.locals[self.reverse_copies[lhs.name]+':convert']
self.locals.pop(self.reverse_copies[lhs.name]+':convert')
if table_types is None:
fname_def = guard(get_definition, self.func_ir, file_name)
if not isinstance(fname_def, ir.Const) or not isinstance(fname_def.value, str):
raise ValueError("Parquet schema not available")
file_name_str = fname_def.value
col_names, col_types = parquet_file_schema(file_name_str)
else:
col_names = list(table_types.keys())
col_types = list(table_types.values())
out_nodes = []
col_items = []
for i, cname in enumerate(col_names):
# get column type from schema
c_type = col_types[i]
if cname in convert_types:
c_type = convert_types[cname]
# create a variable for column and assign type
varname = mk_unique_var(cname)
#self.locals[varname] = c_type
cvar = ir.Var(scope, varname, loc)
col_items.append((cname, cvar))
out_nodes += get_column_read_nodes(c_type, cvar, file_name, i)
return col_items, out_nodes
def fix_dependencies(expr, varlist):
"""Double check if all variables in varlist are defined before
expr is used. Try to move constant definition when the check fails.
Bails out by raising GuardException if it can't be moved.
"""
debug_print = _make_debug_print("fix_dependencies")
for label, block in blocks.items():
scope = block.scope
body = block.body
defined = set()
for i in range(len(body)):
inst = body[i]
if isinstance(inst, ir.Assign):
defined.add(inst.target.name)
if inst.value == expr:
new_varlist = []
for var in varlist:
# var must be defined before this inst, or live
# and not later defined.
if (var.name in defined or
(var.name in livemap[label] and
not (var.name in usedefs.defmap[label]))):
debug_print(var.name, " already defined")
new_varlist.append(var)
else:
debug_print(var.name, " not yet defined")
var_def = get_definition(func_ir, var.name)
if isinstance(var_def, ir.Const):
loc = var.loc
new_var = ir.Var(scope, mk_unique_var("new_var"), loc)
new_const = ir.Const(var_def.value, loc)
new_vardef = _new_definition(func_ir,
new_var, new_const, loc)
new_body = []
new_body.extend(body[:i])
new_body.append(new_vardef)
new_body.extend(body[i:])
block.body = new_body
new_varlist.append(new_var)
else:
raise GuardException
return new_varlist
# when expr is not found in block
raise GuardException
def replace_return_with_setitem(self, blocks, index_vars, out_name):
"""
Find return statements in the IR and replace them with a SetItem
call of the value "returned" by the kernel into the result array.
Returns the block labels that contained return statements.
"""
ret_blocks = []
for label, block in blocks.items():
scope = block.scope
loc = block.loc
new_body = []
for stmt in block.body:
if isinstance(stmt, ir.Return):
ret_blocks.append(label)
# If 1D array then avoid the tuple construction.
if len(index_vars) == 1:
rvar = ir.Var(scope, out_name, loc)
ivar = ir.Var(scope, index_vars[0], loc)
new_body.append(ir.SetItem(rvar, ivar, stmt.value,
loc))
else:
# Convert the string names of the index variables into
# ir.Var's.
var_index_vars = []
for one_var in index_vars:
index_var = ir.Var(scope, one_var, loc)
var_index_vars += [index_var]
s_index_name = ir_utils.mk_unique_var("stencil_index")
s_index_var = ir.Var(scope, s_index_name, loc)
# Build a tuple from the index ir.Var's.
tuple_call = ir.Expr.build_tuple(var_index_vars, loc)
new_body.append(ir.Assign(tuple_call, s_index_var,
loc))
rvar = ir.Var(scope, out_name, loc)
# Write the return statements original value into
# the array using the tuple index.
si = ir.SetItem(rvar, s_index_var, stmt.value, loc)
new_body.append(si)
else:
new_body.append(stmt)
block.body = new_body
return ret_blocks
def _gen_rebalances(self, rebalance_arrs, blocks):
#
for block in blocks.values():
new_body = []
for inst in block.body:
# TODO: handle hiframes filter etc.
if isinstance(inst, Parfor):
self._gen_rebalances(rebalance_arrs, {0: inst.init_block})
self._gen_rebalances(rebalance_arrs, inst.loop_body)
if isinstance(
inst,
ir.Assign) and inst.target.name in rebalance_arrs:
out_arr = inst.target
self.func_ir._definitions[out_arr.name].remove(inst.value)
# hold inst results in tmp array
tmp_arr = ir.Var(out_arr.scope,
mk_unique_var("rebalance_tmp"),
out_arr.loc)
self.typemap[tmp_arr.name] = self.typemap[out_arr.name]
inst.target = tmp_arr
nodes = [inst]
def f(in_arr): # pragma: no cover
out_a = hpat.distributed_api.rebalance_array(in_arr)
f_block = compile_to_numba_ir(
f, {
'hpat': hpat
}, self.typingctx, (self.typemap[tmp_arr.name], ),
self.typemap, self.calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [tmp_arr])
nodes += f_block.body[:-3] # remove none return
nodes[-1].target = out_arr
# update definitions
dumm_block = ir.Block(out_arr.scope, out_arr.loc)
dumm_block.body = nodes
build_definitions({0: dumm_block},
self.func_ir._definitions)
new_body += nodes
else:
new_body.append(inst)
block.body = new_body
def make_var_name(prefix=None, name=None):
"""
Creates variable name.
If prefix is given it would be extended with postfix to create unique name.
If name is given returning this name.
Either prefix or name could be not None at the same time
"""
assert prefix is None or name is None
var_name = None
if name is not None:
var_name = name
else:
if prefix is None:
prefix = '$_var_'
var_name = mk_unique_var(prefix)
return var_name
def replace_return_with_setitem(self, blocks, index_vars, out_name):
"""
Find return statements in the IR and replace them with a SetItem
call of the value "returned" by the kernel into the result array.
Returns the block labels that contained return statements.
"""
ret_blocks = []
for label, block in blocks.items():
scope = block.scope
loc = block.loc
new_body = []
for stmt in block.body:
if isinstance(stmt, ir.Return):
ret_blocks.append(label)
# If 1D array then avoid the tuple construction.
if len(index_vars) == 1:
rvar = ir.Var(scope, out_name, loc)
ivar = ir.Var(scope, index_vars[0], loc)
new_body.append(ir.SetItem(rvar, ivar, stmt.value, loc))
else:
# Convert the string names of the index variables into
# ir.Var's.
var_index_vars = []
for one_var in index_vars:
index_var = ir.Var(scope, one_var, loc)
var_index_vars += [index_var]
s_index_name = ir_utils.mk_unique_var("stencil_index")
s_index_var = ir.Var(scope, s_index_name, loc)
# Build a tuple from the index ir.Var's.
tuple_call = ir.Expr.build_tuple(var_index_vars, loc)
new_body.append(ir.Assign(tuple_call, s_index_var, loc))
rvar = ir.Var(scope, out_name, loc)
# Write the return statements original value into
# the array using the tuple index.
si = ir.SetItem(rvar, s_index_var, stmt.value, loc)
new_body.append(si)
else:
new_body.append(stmt)
block.body = new_body
return ret_blocks
def inline_closure_call(func_ir, glbls, block, i, callee, typingctx=None,
arg_typs=None, typemap=None, calltypes=None,
work_list=None):
"""Inline the body of `callee` at its callsite (`i`-th instruction of `block`)
`func_ir` is the func_ir object of the caller function and `glbls` is its
global variable environment (func_ir.func_id.func.__globals__).
`block` is the IR block of the callsite and `i` is the index of the
callsite's node. `callee` is either the called function or a
make_function node. `typingctx`, `typemap` and `calltypes` are typing
data structures of the caller, available if we are in a typed pass.
`arg_typs` includes the types of the arguments at the callsite.
"""
scope = block.scope
instr = block.body[i]
call_expr = instr.value
debug_print = _make_debug_print("inline_closure_call")
debug_print("Found closure call: ", instr, " with callee = ", callee)
# support both function object and make_function Expr
callee_code = callee.code if hasattr(callee, 'code') else callee.__code__
callee_defaults = callee.defaults if hasattr(callee, 'defaults') else callee.__defaults__
callee_closure = callee.closure if hasattr(callee, 'closure') else callee.__closure__
# first, get the IR of the callee
callee_ir = get_ir_of_code(glbls, callee_code)
callee_blocks = callee_ir.blocks
# 1. relabel callee_ir by adding an offset
max_label = max(func_ir.blocks.keys())
callee_blocks = add_offset_to_labels(callee_blocks, max_label + 1)
callee_blocks = simplify_CFG(callee_blocks)
callee_ir.blocks = callee_blocks
min_label = min(callee_blocks.keys())
max_label = max(callee_blocks.keys())
# reset globals in ir_utils before we use it
ir_utils._max_label = max_label
debug_print("After relabel")
_debug_dump(callee_ir)
# 2. rename all local variables in callee_ir with new locals created in func_ir
callee_scopes = _get_all_scopes(callee_blocks)
debug_print("callee_scopes = ", callee_scopes)
# one function should only have one local scope
assert(len(callee_scopes) == 1)
callee_scope = callee_scopes[0]
var_dict = {}
for var in callee_scope.localvars._con.values():
if not (var.name in callee_code.co_freevars):
new_var = scope.define(mk_unique_var(var.name), loc=var.loc)
var_dict[var.name] = new_var
debug_print("var_dict = ", var_dict)
replace_vars(callee_blocks, var_dict)
debug_print("After local var rename")
_debug_dump(callee_ir)
# 3. replace formal parameters with actual arguments
args = list(call_expr.args)
if callee_defaults:
debug_print("defaults = ", callee_defaults)
if isinstance(callee_defaults, tuple): # Python 3.5
args = args + list(callee_defaults)
elif isinstance(callee_defaults, ir.Var) or isinstance(callee_defaults, str):
defaults = func_ir.get_definition(callee_defaults)
assert(isinstance(defaults, ir.Const))
loc = defaults.loc
args = args + [ir.Const(value=v, loc=loc)
for v in defaults.value]
else:
raise NotImplementedError(
"Unsupported defaults to make_function: {}".format(defaults))
debug_print("After arguments rename: ")
_debug_dump(callee_ir)
# 4. replace freevar with actual closure var
if callee_closure:
closure = func_ir.get_definition(callee_closure)
debug_print("callee's closure = ", closure)
if isinstance(closure, tuple):
cellget = ctypes.pythonapi.PyCell_Get
cellget.restype = ctypes.py_object
cellget.argtypes = (ctypes.py_object,)
items = tuple(cellget(x) for x in closure)
else:
assert(isinstance(closure, ir.Expr)
and closure.op == 'build_tuple')
items = closure.items
assert(len(callee_code.co_freevars) == len(items))
_replace_freevars(callee_blocks, items)
debug_print("After closure rename")
_debug_dump(callee_ir)
if typingctx:
from numba import compiler
f_typemap, f_return_type, f_calltypes = compiler.type_inference_stage(
typingctx, callee_ir, arg_typs, None)
canonicalize_array_math(callee_ir, f_typemap,
f_calltypes, typingctx)
# remove argument entries like arg.a from typemap
arg_names = [vname for vname in f_typemap if vname.startswith("arg.")]
for a in arg_names:
f_typemap.pop(a)
typemap.update(f_typemap)
calltypes.update(f_calltypes)
_replace_args_with(callee_blocks, args)
# 5. split caller blocks into two
new_blocks = []
new_block = ir.Block(scope, block.loc)
new_block.body = block.body[i + 1:]
new_label = next_label()
func_ir.blocks[new_label] = new_block
new_blocks.append((new_label, new_block))
block.body = block.body[:i]
block.body.append(ir.Jump(min_label, instr.loc))
# 6. replace Return with assignment to LHS
topo_order = find_topo_order(callee_blocks)
_replace_returns(callee_blocks, instr.target, new_label)
# remove the old definition of instr.target too
if (instr.target.name in func_ir._definitions):
func_ir._definitions[instr.target.name] = []
# 7. insert all new blocks, and add back definitions
for label in topo_order:
# block scope must point to parent's
block = callee_blocks[label]
block.scope = scope
_add_definitions(func_ir, block)
func_ir.blocks[label] = block
new_blocks.append((label, block))
debug_print("After merge in")
_debug_dump(func_ir)
if work_list != None:
for block in new_blocks:
work_list.append(block)
return callee_blocks
def _create_gufunc_for_parfor_body(
lowerer,
parfor,
typemap,
typingctx,
targetctx,
flags,
locals,
has_aliases,
index_var_typ):
'''
Takes a parfor and creates a gufunc function for its body.
There are two parts to this function.
1) Code to iterate across the iteration space as defined by the schedule.
2) The parfor body that does the work for a single point in the iteration space.
Part 1 is created as Python text for simplicity with a sentinel assignment to mark the point
in the IR where the parfor body should be added.
This Python text is 'exec'ed into existence and its IR retrieved with run_frontend.
The IR is scanned for the sentinel assignment where that basic block is split and the IR
for the parfor body inserted.
'''
# The parfor body and the main function body share ir.Var nodes.
# We have to do some replacements of Var names in the parfor body to make them
# legal parameter names. If we don't copy then the Vars in the main function also
# would incorrectly change their name.
loop_body = copy.copy(parfor.loop_body)
remove_dels(loop_body)
parfor_dim = len(parfor.loop_nests)
loop_indices = [l.index_variable.name for l in parfor.loop_nests]
# Get all the parfor params.
parfor_params = parfor.params
# Get just the outputs of the parfor.
parfor_outputs = numba.parfor.get_parfor_outputs(parfor, parfor_params)
# Get all parfor reduction vars, and operators.
parfor_redvars, parfor_reddict = numba.parfor.get_parfor_reductions(
parfor, parfor_params, lowerer.fndesc.calltypes)
# Compute just the parfor inputs as a set difference.
parfor_inputs = sorted(
list(
set(parfor_params) -
set(parfor_outputs) -
set(parfor_redvars)))
if config.DEBUG_ARRAY_OPT == 1:
print("parfor_params = ", parfor_params, " ", type(parfor_params))
print("parfor_outputs = ", parfor_outputs, " ", type(parfor_outputs))
print("parfor_inputs = ", parfor_inputs, " ", type(parfor_inputs))
print("parfor_redvars = ", parfor_redvars, " ", type(parfor_redvars))
# Reduction variables are represented as arrays, so they go under
# different names.
parfor_redarrs = []
for var in parfor_redvars:
arr = var + "_arr"
parfor_redarrs.append(arr)
typemap[arr] = types.npytypes.Array(typemap[var], 1, "C")
# Reorder all the params so that inputs go first then outputs.
parfor_params = parfor_inputs + parfor_outputs + parfor_redarrs
if config.DEBUG_ARRAY_OPT == 1:
print("parfor_params = ", parfor_params, " ", type(parfor_params))
print("loop_indices = ", loop_indices, " ", type(loop_indices))
print("loop_body = ", loop_body, " ", type(loop_body))
_print_body(loop_body)
# Some Var are not legal parameter names so create a dict of potentially illegal
# param name to guaranteed legal name.
param_dict = legalize_names(parfor_params + parfor_redvars)
if config.DEBUG_ARRAY_OPT == 1:
print(
"param_dict = ",
sorted(
param_dict.items()),
" ",
type(param_dict))
# Some loop_indices are not legal parameter names so create a dict of potentially illegal
# loop index to guaranteed legal name.
ind_dict = legalize_names(loop_indices)
# Compute a new list of legal loop index names.
legal_loop_indices = [ind_dict[v] for v in loop_indices]
if config.DEBUG_ARRAY_OPT == 1:
print("ind_dict = ", sorted(ind_dict.items()), " ", type(ind_dict))
print(
"legal_loop_indices = ",
legal_loop_indices,
" ",
type(legal_loop_indices))
for pd in parfor_params:
print("pd = ", pd)
print("pd type = ", typemap[pd], " ", type(typemap[pd]))
# Get the types of each parameter.
param_types = [typemap[v] for v in parfor_params]
# if config.DEBUG_ARRAY_OPT==1:
# param_types_dict = { v:typemap[v] for v in parfor_params }
# print("param_types_dict = ", param_types_dict, " ", type(param_types_dict))
# print("param_types = ", param_types, " ", type(param_types))
# Replace illegal parameter names in the loop body with legal ones.
replace_var_names(loop_body, param_dict)
# remember the name before legalizing as the actual arguments
parfor_args = parfor_params
# Change parfor_params to be legal names.
parfor_params = [param_dict[v] for v in parfor_params]
parfor_params_orig = parfor_params
parfor_params = []
ascontig = False
for pindex in range(len(parfor_params_orig)):
if ascontig and pindex < len(parfor_inputs) and isinstance(param_types[pindex], types.npytypes.Array):
parfor_params.append(parfor_params_orig[pindex]+"param")
else:
parfor_params.append(parfor_params_orig[pindex])
# Change parfor body to replace illegal loop index vars with legal ones.
replace_var_names(loop_body, ind_dict)
loop_body_var_table = get_name_var_table(loop_body)
sentinel_name = get_unused_var_name("__sentinel__", loop_body_var_table)
if config.DEBUG_ARRAY_OPT == 1:
print(
"legal parfor_params = ",
parfor_params,
" ",
type(parfor_params))
# Determine the unique names of the scheduling and gufunc functions.
# sched_func_name = "__numba_parfor_sched_%s" % (hex(hash(parfor)).replace("-", "_"))
gufunc_name = "__numba_parfor_gufunc_%s" % (
hex(hash(parfor)).replace("-", "_"))
if config.DEBUG_ARRAY_OPT:
# print("sched_func_name ", type(sched_func_name), " ", sched_func_name)
print("gufunc_name ", type(gufunc_name), " ", gufunc_name)
gufunc_txt = ""
# Create the gufunc function.
gufunc_txt += "def " + gufunc_name + \
"(sched, " + (", ".join(parfor_params)) + "):\n"
for pindex in range(len(parfor_inputs)):
if ascontig and isinstance(param_types[pindex], types.npytypes.Array):
gufunc_txt += (" " + parfor_params_orig[pindex]
+ " = np.ascontiguousarray(" + parfor_params[pindex] + ")\n")
# Add initialization of reduction variables
for arr, var in zip(parfor_redarrs, parfor_redvars):
gufunc_txt += " " + param_dict[var] + \
"=" + param_dict[arr] + "[0]\n"
# For each dimension of the parfor, create a for loop in the generated gufunc function.
# Iterate across the proper values extracted from the schedule.
# The form of the schedule is start_dim0, start_dim1, ..., start_dimN, end_dim0,
# end_dim1, ..., end_dimN
for eachdim in range(parfor_dim):
for indent in range(eachdim + 1):
gufunc_txt += " "
sched_dim = eachdim
gufunc_txt += ("for " +
legal_loop_indices[eachdim] +
" in range(sched[" +
str(sched_dim) +
"], sched[" +
str(sched_dim +
parfor_dim) +
"] + np.uint8(1)):\n")
if config.DEBUG_ARRAY_OPT_RUNTIME:
for indent in range(parfor_dim + 1):
gufunc_txt += " "
gufunc_txt += "print("
for eachdim in range(parfor_dim):
gufunc_txt += "\"" + legal_loop_indices[eachdim] + "\"," + legal_loop_indices[eachdim] + ","
gufunc_txt += ")\n"
# Add the sentinel assignment so that we can find the loop body position
# in the IR.
for indent in range(parfor_dim + 1):
gufunc_txt += " "
gufunc_txt += sentinel_name + " = 0\n"
# Add assignments of reduction variables (for returning the value)
for arr, var in zip(parfor_redarrs, parfor_redvars):
gufunc_txt += " " + param_dict[arr] + \
"[0] = " + param_dict[var] + "\n"
gufunc_txt += " return None\n"
if config.DEBUG_ARRAY_OPT:
print("gufunc_txt = ", type(gufunc_txt), "\n", gufunc_txt)
# Force gufunc outline into existence.
globls = {"np": np}
locls = {}
exec_(gufunc_txt, globls, locls)
gufunc_func = locls[gufunc_name]
if config.DEBUG_ARRAY_OPT:
print("gufunc_func = ", type(gufunc_func), "\n", gufunc_func)
# Get the IR for the gufunc outline.
gufunc_ir = compiler.run_frontend(gufunc_func)
if config.DEBUG_ARRAY_OPT:
print("gufunc_ir dump ", type(gufunc_ir))
gufunc_ir.dump()
print("loop_body dump ", type(loop_body))
_print_body(loop_body)
# rename all variables in gufunc_ir afresh
var_table = get_name_var_table(gufunc_ir.blocks)
new_var_dict = {}
reserved_names = [sentinel_name] + \
list(param_dict.values()) + legal_loop_indices
for name, var in var_table.items():
if not (name in reserved_names):
new_var_dict[name] = mk_unique_var(name)
replace_var_names(gufunc_ir.blocks, new_var_dict)
if config.DEBUG_ARRAY_OPT:
print("gufunc_ir dump after renaming ")
gufunc_ir.dump()
gufunc_param_types = [
numba.types.npytypes.Array(
index_var_typ, 1, "C")] + param_types
if config.DEBUG_ARRAY_OPT:
print(
"gufunc_param_types = ",
type(gufunc_param_types),
"\n",
gufunc_param_types)
gufunc_stub_last_label = max(gufunc_ir.blocks.keys()) + 1
# Add gufunc stub last label to each parfor.loop_body label to prevent
# label conflicts.
loop_body = add_offset_to_labels(loop_body, gufunc_stub_last_label)
# new label for splitting sentinel block
new_label = max(loop_body.keys()) + 1
# If enabled, add a print statement after every assignment.
if config.DEBUG_ARRAY_OPT_RUNTIME:
for label, block in loop_body.items():
new_block = block.copy()
new_block.clear()
loc = block.loc
scope = block.scope
for inst in block.body:
new_block.append(inst)
# Append print after assignment
if isinstance(inst, ir.Assign):
# Only apply to numbers
if typemap[inst.target.name] not in types.number_domain:
continue
# Make constant string
strval = "{} =".format(inst.target.name)
strconsttyp = types.Const(strval)
lhs = ir.Var(scope, mk_unique_var("str_const"), loc)
assign_lhs = ir.Assign(value=ir.Const(value=strval, loc=loc),
target=lhs, loc=loc)
typemap[lhs.name] = strconsttyp
new_block.append(assign_lhs)
# Make print node
print_node = ir.Print(args=[lhs, inst.target], vararg=None, loc=loc)
new_block.append(print_node)
sig = numba.typing.signature(types.none,
typemap[lhs.name],
typemap[inst.target.name])
lowerer.fndesc.calltypes[print_node] = sig
loop_body[label] = new_block
if config.DEBUG_ARRAY_OPT:
print("parfor loop body")
_print_body(loop_body)
wrapped_blocks = wrap_loop_body(loop_body)
hoisted = hoist(parfor_params, loop_body, typemap, wrapped_blocks)
start_block = gufunc_ir.blocks[min(gufunc_ir.blocks.keys())]
start_block.body = start_block.body[:-1] + hoisted + [start_block.body[-1]]
unwrap_loop_body(loop_body)
if config.DEBUG_ARRAY_OPT:
print("After hoisting")
_print_body(loop_body)
# Search all the block in the gufunc outline for the sentinel assignment.
for label, block in gufunc_ir.blocks.items():
for i, inst in enumerate(block.body):
if isinstance(
inst,
ir.Assign) and inst.target.name == sentinel_name:
# We found the sentinel assignment.
loc = inst.loc
scope = block.scope
# split block across __sentinel__
# A new block is allocated for the statements prior to the sentinel
# but the new block maintains the current block label.
prev_block = ir.Block(scope, loc)
prev_block.body = block.body[:i]
# The current block is used for statements after the sentinel.
block.body = block.body[i + 1:]
# But the current block gets a new label.
body_first_label = min(loop_body.keys())
# The previous block jumps to the minimum labelled block of the
# parfor body.
prev_block.append(ir.Jump(body_first_label, loc))
# Add all the parfor loop body blocks to the gufunc function's
# IR.
for (l, b) in loop_body.items():
gufunc_ir.blocks[l] = b
body_last_label = max(loop_body.keys())
gufunc_ir.blocks[new_label] = block
gufunc_ir.blocks[label] = prev_block
# Add a jump from the last parfor body block to the block containing
# statements after the sentinel.
gufunc_ir.blocks[body_last_label].append(
ir.Jump(new_label, loc))
break
else:
continue
break
if config.DEBUG_ARRAY_OPT:
print("gufunc_ir last dump before renaming")
gufunc_ir.dump()
gufunc_ir.blocks = rename_labels(gufunc_ir.blocks)
remove_dels(gufunc_ir.blocks)
if config.DEBUG_ARRAY_OPT:
print("gufunc_ir last dump")
gufunc_ir.dump()
print("flags", flags)
print("typemap", typemap)
old_alias = flags.noalias
if not has_aliases:
if config.DEBUG_ARRAY_OPT:
print("No aliases found so adding noalias flag.")
flags.noalias = True
kernel_func = compiler.compile_ir(
typingctx,
targetctx,
gufunc_ir,
gufunc_param_types,
types.none,
flags,
locals)
flags.noalias = old_alias
kernel_sig = signature(types.none, *gufunc_param_types)
if config.DEBUG_ARRAY_OPT:
print("kernel_sig = ", kernel_sig)
return kernel_func, parfor_args, kernel_sig
def call_parallel_gufunc(lowerer, cres, gu_signature, outer_sig, expr_args,
loop_ranges, redvars, reddict, init_block, index_var_typ):
'''
Adds the call to the gufunc function from the main function.
'''
context = lowerer.context
builder = lowerer.builder
library = lowerer.library
from .parallel import (ParallelGUFuncBuilder, build_gufunc_wrapper,
get_thread_count, _launch_threads, _init)
if config.DEBUG_ARRAY_OPT:
print("make_parallel_loop")
print("args = ", expr_args)
print("outer_sig = ", outer_sig.args, outer_sig.return_type,
outer_sig.recvr, outer_sig.pysig)
print("loop_ranges = ", loop_ranges)
# Build the wrapper for GUFunc
args, return_type = sigutils.normalize_signature(outer_sig)
llvm_func = cres.library.get_function(cres.fndesc.llvm_func_name)
sin, sout = gu_signature
# These are necessary for build_gufunc_wrapper to find external symbols
_launch_threads()
_init()
wrapper_ptr, env, wrapper_name = build_gufunc_wrapper(llvm_func, cres, sin,
sout, {})
cres.library._ensure_finalized()
if config.DEBUG_ARRAY_OPT:
print("parallel function = ", wrapper_name, cres)
# loadvars for loop_ranges
def load_range(v):
if isinstance(v, ir.Var):
return lowerer.loadvar(v.name)
else:
return context.get_constant(types.uintp, v)
num_dim = len(loop_ranges)
for i in range(num_dim):
start, stop, step = loop_ranges[i]
start = load_range(start)
stop = load_range(stop)
assert(step == 1) # We do not support loop steps other than 1
step = load_range(step)
loop_ranges[i] = (start, stop, step)
if config.DEBUG_ARRAY_OPT:
print("call_parallel_gufunc loop_ranges[{}] = ".format(i), start,
stop, step)
cgutils.printf(builder, "loop range[{}]: %d %d (%d)\n".format(i),
start, stop, step)
# Commonly used LLVM types and constants
byte_t = lc.Type.int(8)
byte_ptr_t = lc.Type.pointer(byte_t)
byte_ptr_ptr_t = lc.Type.pointer(byte_ptr_t)
intp_t = context.get_value_type(types.intp)
uintp_t = context.get_value_type(types.uintp)
intp_ptr_t = lc.Type.pointer(intp_t)
uintp_ptr_t = lc.Type.pointer(uintp_t)
zero = context.get_constant(types.uintp, 0)
one = context.get_constant(types.uintp, 1)
one_type = one.type
sizeof_intp = context.get_abi_sizeof(intp_t)
# Prepare sched, first pop it out of expr_args, outer_sig, and gu_signature
sched_name = expr_args.pop(0)
sched_typ = outer_sig.args[0]
sched_sig = sin.pop(0)
if config.DEBUG_ARRAY_OPT:
print("Parfor has potentially negative start", index_var_typ.signed)
if index_var_typ.signed:
sched_type = intp_t
sched_ptr_type = intp_ptr_t
else:
sched_type = uintp_t
sched_ptr_type = uintp_ptr_t
# Call do_scheduling with appropriate arguments
dim_starts = cgutils.alloca_once(
builder, sched_type, size=context.get_constant(
types.uintp, num_dim), name="dims")
dim_stops = cgutils.alloca_once(
builder, sched_type, size=context.get_constant(
types.uintp, num_dim), name="dims")
for i in range(num_dim):
start, stop, step = loop_ranges[i]
if start.type != one_type:
start = builder.sext(start, one_type)
if stop.type != one_type:
stop = builder.sext(stop, one_type)
if step.type != one_type:
step = builder.sext(step, one_type)
# substract 1 because do-scheduling takes inclusive ranges
stop = builder.sub(stop, one)
builder.store(
start, builder.gep(
dim_starts, [
context.get_constant(
types.uintp, i)]))
builder.store(stop, builder.gep(dim_stops,
[context.get_constant(types.uintp, i)]))
sched_size = get_thread_count() * num_dim * 2
sched = cgutils.alloca_once(
builder, sched_type, size=context.get_constant(
types.uintp, sched_size), name="sched")
debug_flag = 1 if config.DEBUG_ARRAY_OPT else 0
scheduling_fnty = lc.Type.function(
intp_ptr_t, [uintp_t, sched_ptr_type, sched_ptr_type, uintp_t, sched_ptr_type, intp_t])
if index_var_typ.signed:
do_scheduling = builder.module.get_or_insert_function(scheduling_fnty,
name="do_scheduling_signed")
else:
do_scheduling = builder.module.get_or_insert_function(scheduling_fnty,
name="do_scheduling_unsigned")
builder.call(
do_scheduling, [
context.get_constant(
types.uintp, num_dim), dim_starts, dim_stops, context.get_constant(
types.uintp, get_thread_count()), sched, context.get_constant(
types.intp, debug_flag)])
# init reduction array allocation here.
nredvars = len(redvars)
ninouts = len(expr_args) - nredvars
redarrs = []
for i in range(nredvars):
redvar_typ = lowerer.fndesc.typemap[redvars[i]]
# we need to use the default initial value instead of existing value in
# redvar if available
init_val = reddict[redvars[i]][0]
if init_val != None:
val = context.get_constant(redvar_typ, init_val)
else:
val = lowerer.loadvar(redvars[i])
typ = context.get_value_type(redvar_typ)
size = get_thread_count()
arr = cgutils.alloca_once(builder, typ,
size=context.get_constant(types.uintp, size))
redarrs.append(arr)
for j in range(size):
dst = builder.gep(arr, [context.get_constant(types.uintp, j)])
builder.store(val, dst)
if config.DEBUG_ARRAY_OPT:
for i in range(get_thread_count()):
cgutils.printf(builder, "sched[" + str(i) + "] = ")
for j in range(num_dim * 2):
cgutils.printf(
builder, "%d ", builder.load(
builder.gep(
sched, [
context.get_constant(
types.intp, i * num_dim * 2 + j)])))
cgutils.printf(builder, "\n")
# Prepare arguments: args, shapes, steps, data
all_args = [lowerer.loadvar(x) for x in expr_args[:ninouts]] + redarrs
num_args = len(all_args)
num_inps = len(sin) + 1
args = cgutils.alloca_once(
builder,
byte_ptr_t,
size=context.get_constant(
types.intp,
1 + num_args),
name="pargs")
array_strides = []
# sched goes first
builder.store(builder.bitcast(sched, byte_ptr_t), args)
array_strides.append(context.get_constant(types.intp, sizeof_intp))
# followed by other arguments
for i in range(num_args):
arg = all_args[i]
aty = outer_sig.args[i + 1] # skip first argument sched
dst = builder.gep(args, [context.get_constant(types.intp, i + 1)])
if i >= ninouts: # reduction variables
builder.store(builder.bitcast(arg, byte_ptr_t), dst)
elif isinstance(aty, types.ArrayCompatible):
ary = context.make_array(aty)(context, builder, arg)
strides = cgutils.unpack_tuple(builder, ary.strides, aty.ndim)
for j in range(len(strides)):
array_strides.append(strides[j])
builder.store(builder.bitcast(ary.data, byte_ptr_t), dst)
else:
if i < num_inps:
# Scalar input, need to store the value in an array of size 1
typ = context.get_data_type(
aty) if aty != types.boolean else lc.Type.int(1)
ptr = cgutils.alloca_once(builder, typ)
builder.store(arg, ptr)
else:
# Scalar output, must allocate
typ = context.get_data_type(
aty) if aty != types.boolean else lc.Type.int(1)
ptr = cgutils.alloca_once(builder, typ)
builder.store(builder.bitcast(ptr, byte_ptr_t), dst)
# Next, we prepare the individual dimension info recorded in gu_signature
sig_dim_dict = {}
occurances = []
occurances = [sched_sig[0]]
sig_dim_dict[sched_sig[0]] = context.get_constant(types.intp, 2 * num_dim)
for var, arg, aty, gu_sig in zip(expr_args[:ninouts], all_args[:ninouts],
outer_sig.args[1:], sin + sout):
if config.DEBUG_ARRAY_OPT:
print("var = ", var, " gu_sig = ", gu_sig)
i = 0
for dim_sym in gu_sig:
if config.DEBUG_ARRAY_OPT:
print("var = ", var, " type = ", aty)
ary = context.make_array(aty)(context, builder, arg)
shapes = cgutils.unpack_tuple(builder, ary.shape, aty.ndim)
sig_dim_dict[dim_sym] = shapes[i]
if not (dim_sym in occurances):
if config.DEBUG_ARRAY_OPT:
print("dim_sym = ", dim_sym, ", i = ", i)
cgutils.printf(builder, dim_sym + " = %d\n", shapes[i])
occurances.append(dim_sym)
i = i + 1
# Prepare shapes, which is a single number (outer loop size), followed by
# the size of individual shape variables.
nshapes = len(sig_dim_dict) + 1
shapes = cgutils.alloca_once(builder, intp_t, size=nshapes, name="pshape")
# For now, outer loop size is the same as number of threads
builder.store(context.get_constant(types.intp, get_thread_count()), shapes)
# Individual shape variables go next
i = 1
for dim_sym in occurances:
if config.DEBUG_ARRAY_OPT:
cgutils.printf(builder, dim_sym + " = %d\n", sig_dim_dict[dim_sym])
builder.store(
sig_dim_dict[dim_sym], builder.gep(
shapes, [
context.get_constant(
types.intp, i)]))
i = i + 1
# Prepare steps for each argument. Note that all steps are counted in
# bytes.
num_steps = num_args + 1 + len(array_strides)
steps = cgutils.alloca_once(
builder, intp_t, size=context.get_constant(
types.intp, num_steps), name="psteps")
# First goes the step size for sched, which is 2 * num_dim
builder.store(context.get_constant(types.intp, 2 * num_dim * sizeof_intp),
steps)
# The steps for all others are 0. (TODO: except reduction results)
for i in range(num_args):
if i >= ninouts: # steps for reduction vars are abi_sizeof(typ)
j = i - ninouts
typ = context.get_value_type(lowerer.fndesc.typemap[redvars[j]])
sizeof = context.get_abi_sizeof(typ)
stepsize = context.get_constant(types.intp, sizeof)
else:
# steps are strides
stepsize = zero
dst = builder.gep(steps, [context.get_constant(types.intp, 1 + i)])
builder.store(stepsize, dst)
for j in range(len(array_strides)):
dst = builder.gep(
steps, [
context.get_constant(
types.intp, 1 + num_args + j)])
builder.store(array_strides[j], dst)
# prepare data
data = builder.inttoptr(zero, byte_ptr_t)
fnty = lc.Type.function(lc.Type.void(), [byte_ptr_ptr_t, intp_ptr_t,
intp_ptr_t, byte_ptr_t])
fn = builder.module.get_or_insert_function(fnty, name=wrapper_name)
if config.DEBUG_ARRAY_OPT:
cgutils.printf(builder, "before calling kernel %p\n", fn)
result = builder.call(fn, [args, shapes, steps, data])
if config.DEBUG_ARRAY_OPT:
cgutils.printf(builder, "after calling kernel %p\n", fn)
scope = init_block.scope
loc = init_block.loc
calltypes = lowerer.fndesc.calltypes
# Accumulate all reduction arrays back to a single value
for i in range(get_thread_count()):
for name, arr in zip(redvars, redarrs):
tmpname = mk_unique_var(name)
src = builder.gep(arr, [context.get_constant(types.intp, i)])
val = builder.load(src)
vty = lowerer.fndesc.typemap[name]
lowerer.fndesc.typemap[tmpname] = vty
lowerer.storevar(val, tmpname)
tmpvar = ir.Var(scope, tmpname, loc)
tmp_assign = ir.Assign(tmpvar, ir.Var(scope, name+"#init", loc), loc)
if name+"#init" not in lowerer.fndesc.typemap:
lowerer.fndesc.typemap[name+"#init"] = vty
lowerer.lower_inst(tmp_assign)
# generate code for combining reduction variable with thread output
for inst in reddict[name][1]:
lowerer.lower_inst(inst)
# TODO: scalar output must be assigned back to corresponding output
# variables
return
def add_indices_to_kernel(self, kernel, index_names, ndim,
neighborhood, standard_indexed):
"""
Transforms the stencil kernel as specified by the user into one
that includes each dimension's index variable as part of the getitem
calls. So, in effect array[-1] becomes array[index0-1].
"""
const_dict = {}
kernel_consts = []
if config.DEBUG_ARRAY_OPT == 1:
print("add_indices_to_kernel", ndim, neighborhood)
ir_utils.dump_blocks(kernel.blocks)
if neighborhood is None:
need_to_calc_kernel = True
else:
need_to_calc_kernel = False
if len(neighborhood) != ndim:
raise ValueError("%d dimensional neighborhood specified for %d " \
"dimensional input array" % (len(neighborhood), ndim))
tuple_table = ir_utils.get_tuple_table(kernel.blocks)
relatively_indexed = set()
for block in kernel.blocks.values():
scope = block.scope
loc = block.loc
new_body = []
for stmt in block.body:
if (isinstance(stmt, ir.Assign) and
isinstance(stmt.value, ir.Const)):
if config.DEBUG_ARRAY_OPT == 1:
print("remembering in const_dict", stmt.target.name,
stmt.value.value)
# Remember consts for use later.
const_dict[stmt.target.name] = stmt.value.value
if ((isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op in ['setitem', 'static_setitem']
and stmt.value.value.name in kernel.arg_names) or
(isinstance(stmt, ir.SetItem)
and stmt.target.name in kernel.arg_names)):
raise ValueError("Assignments to arrays passed to stencil " \
"kernels is not allowed.")
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op in ['getitem', 'static_getitem']
and stmt.value.value.name in kernel.arg_names
and stmt.value.value.name not in standard_indexed):
# We found a getitem from the input array.
if stmt.value.op == 'getitem':
stmt_index_var = stmt.value.index
else:
stmt_index_var = stmt.value.index_var
# allow static_getitem since rewrite passes are applied
#raise ValueError("Unexpected static_getitem in add_indices_to_kernel.")
relatively_indexed.add(stmt.value.value.name)
# Store the index used after looking up the variable in
# the const dictionary.
if need_to_calc_kernel:
assert hasattr(stmt_index_var, 'name')
if stmt_index_var.name in tuple_table:
kernel_consts += [tuple_table[stmt_index_var.name]]
elif stmt_index_var.name in const_dict:
kernel_consts += [const_dict[stmt_index_var.name]]
else:
raise ValueError("Non-constant specified for "
"stencil kernel index.")
if ndim == 1:
# Single dimension always has index variable 'index0'.
# tmpvar will hold the real index and is computed by
# adding the relative offset in stmt.value.index to
# the current absolute location in index0.
index_var = ir.Var(scope, index_names[0], loc)
tmpname = ir_utils.mk_unique_var("stencil_index")
tmpvar = ir.Var(scope, tmpname, loc)
acc_call = ir.Expr.binop('+', stmt_index_var,
index_var, loc)
new_body.append(ir.Assign(acc_call, tmpvar, loc))
new_body.append(ir.Assign(
ir.Expr.getitem(stmt.value.value,tmpvar,loc),
stmt.target,loc))
else:
index_vars = []
sum_results = []
s_index_name = ir_utils.mk_unique_var("stencil_index")
s_index_var = ir.Var(scope, s_index_name, loc)
const_index_vars = []
ind_stencils = []
# Same idea as above but you have to extract
# individual elements out of the tuple indexing
# expression and add the corresponding index variable
# to them and then reconstitute as a tuple that can
# index the array.
for dim in range(ndim):
tmpname = ir_utils.mk_unique_var("const_index")
tmpvar = ir.Var(scope, tmpname, loc)
new_body.append(ir.Assign(ir.Const(dim, loc),
tmpvar, loc))
const_index_vars += [tmpvar]
index_var = ir.Var(scope, index_names[dim], loc)
index_vars += [index_var]
tmpname = ir_utils.mk_unique_var("ind_stencil_index")
tmpvar = ir.Var(scope, tmpname, loc)
ind_stencils += [tmpvar]
getitemname = ir_utils.mk_unique_var("getitem")
getitemvar = ir.Var(scope, getitemname, loc)
getitemcall = ir.Expr.getitem(stmt_index_var,
const_index_vars[dim], loc)
new_body.append(ir.Assign(getitemcall, getitemvar, loc))
acc_call = ir.Expr.binop('+', getitemvar,
index_vars[dim], loc)
new_body.append(ir.Assign(acc_call, tmpvar, loc))
tuple_call = ir.Expr.build_tuple(ind_stencils, loc)
new_body.append(ir.Assign(tuple_call, s_index_var, loc))
new_body.append(ir.Assign(
ir.Expr.getitem(stmt.value.value,s_index_var,loc),
stmt.target,loc))
else:
new_body.append(stmt)
block.body = new_body
if need_to_calc_kernel:
# Find the size of the kernel by finding the maximum absolute value
# index used in the kernel specification.
neighborhood = [[0,0] for _ in range(ndim)]
if len(kernel_consts) == 0:
raise ValueError("Stencil kernel with no accesses to "
"relatively indexed arrays.")
for index in kernel_consts:
if isinstance(index, tuple) or isinstance(index, list):
for i in range(len(index)):
te = index[i]
if isinstance(te, ir.Var) and te.name in const_dict:
te = const_dict[te.name]
if isinstance(te, int):
neighborhood[i][0] = min(neighborhood[i][0], te)
neighborhood[i][1] = max(neighborhood[i][1], te)
else:
raise ValueError(
"Non-constant used as stencil index.")
index_len = len(index)
elif isinstance(index, int):
neighborhood[0][0] = min(neighborhood[0][0], index)
neighborhood[0][1] = max(neighborhood[0][1], index)
index_len = 1
else:
raise ValueError(
"Non-tuple or non-integer used as stencil index.")
if index_len != ndim:
raise ValueError(
"Stencil index does not match array dimensionality.")
return (neighborhood, relatively_indexed)
def _stencil_wrapper(self, result, sigret, return_type, typemap, calltypes, *args):
# Overall approach:
# 1) Construct a string containing a function definition for the stencil function
# that will execute the stencil kernel. This function definition includes a
# unique stencil function name, the parameters to the stencil kernel, loop
# nests across the dimenions of the input array. Those loop nests use the
# computed stencil kernel size so as not to try to compute elements where
# elements outside the bounds of the input array would be needed.
# 2) The but of the loop nest in this new function is a special sentinel
# assignment.
# 3) Get the IR of this new function.
# 4) Split the block containing the sentinel assignment and remove the sentinel
# assignment. Insert the stencil kernel IR into the stencil function IR
# after label and variable renaming of the stencil kernel IR to prevent
# conflicts with the stencil function IR.
# 5) Compile the combined stencil function IR + stencil kernel IR into existence.
# Copy the kernel so that our changes for this callsite
# won't effect other callsites.
(kernel_copy, copy_calltypes) = self.copy_ir_with_calltypes(
self.kernel_ir, calltypes)
# The stencil kernel body becomes the body of a loop, for which args aren't needed.
ir_utils.remove_args(kernel_copy.blocks)
first_arg = kernel_copy.arg_names[0]
in_cps, out_cps = ir_utils.copy_propagate(kernel_copy.blocks, typemap)
name_var_table = ir_utils.get_name_var_table(kernel_copy.blocks)
ir_utils.apply_copy_propagate(
kernel_copy.blocks,
in_cps,
name_var_table,
typemap,
copy_calltypes)
if "out" in name_var_table:
raise ValueError("Cannot use the reserved word 'out' in stencil kernels.")
sentinel_name = ir_utils.get_unused_var_name("__sentinel__", name_var_table)
if config.DEBUG_ARRAY_OPT == 1:
print("name_var_table", name_var_table, sentinel_name)
the_array = args[0]
if config.DEBUG_ARRAY_OPT == 1:
print("_stencil_wrapper", return_type, return_type.dtype,
type(return_type.dtype), args)
ir_utils.dump_blocks(kernel_copy.blocks)
# We generate a Numba function to execute this stencil and here
# create the unique name of this function.
stencil_func_name = "__numba_stencil_%s_%s" % (
hex(id(the_array)).replace("-", "_"),
self.id)
# We will put a loop nest in the generated function for each
# dimension in the input array. Here we create the name for
# the index variable for each dimension. index0, index1, ...
index_vars = []
for i in range(the_array.ndim):
index_var_name = ir_utils.get_unused_var_name("index" + str(i),
name_var_table)
index_vars += [index_var_name]
# Create extra signature for out and neighborhood.
out_name = ir_utils.get_unused_var_name("out", name_var_table)
neighborhood_name = ir_utils.get_unused_var_name("neighborhood",
name_var_table)
sig_extra = ""
if result is not None:
sig_extra += ", {}=None".format(out_name)
if "neighborhood" in dict(self.kws):
sig_extra += ", {}=None".format(neighborhood_name)
# Get a list of the standard indexed array names.
standard_indexed = self.options.get("standard_indexing", [])
if first_arg in standard_indexed:
raise ValueError("The first argument to a stencil kernel must "
"use relative indexing, not standard indexing.")
if len(set(standard_indexed) - set(kernel_copy.arg_names)) != 0:
raise ValueError("Standard indexing requested for an array name "
"not present in the stencil kernel definition.")
# Add index variables to getitems in the IR to transition the accesses
# in the kernel from relative to regular Python indexing. Returns the
# computed size of the stencil kernel and a list of the relatively indexed
# arrays.
kernel_size, relatively_indexed = self.add_indices_to_kernel(
kernel_copy, index_vars, the_array.ndim,
self.neighborhood, standard_indexed)
if self.neighborhood is None:
self.neighborhood = kernel_size
if config.DEBUG_ARRAY_OPT == 1:
print("After add_indices_to_kernel")
ir_utils.dump_blocks(kernel_copy.blocks)
# The return in the stencil kernel becomes a setitem for that
# particular point in the iteration space.
ret_blocks = self.replace_return_with_setitem(kernel_copy.blocks,
index_vars, out_name)
if config.DEBUG_ARRAY_OPT == 1:
print("After replace_return_with_setitem", ret_blocks)
ir_utils.dump_blocks(kernel_copy.blocks)
# Start to form the new function to execute the stencil kernel.
func_text = "def {}({}{}):\n".format(stencil_func_name,
",".join(kernel_copy.arg_names), sig_extra)
# Get loop ranges for each dimension, which could be either int
# or variable. In the latter case we'll use the extra neighborhood
# argument to the function.
ranges = []
for i in range(the_array.ndim):
if isinstance(kernel_size[i][0], int):
lo = kernel_size[i][0]
hi = kernel_size[i][1]
else:
lo = "{}[{}][0]".format(neighborhood_name, i)
hi = "{}[{}][1]".format(neighborhood_name, i)
ranges.append((lo, hi))
# If there are more than one relatively indexed arrays, add a call to
# a function that will raise an error if any of the relatively indexed
# arrays are of different size than the first input array.
if len(relatively_indexed) > 1:
func_text += " raise_if_incompatible_array_sizes(" + first_arg
for other_array in relatively_indexed:
if other_array != first_arg:
func_text += "," + other_array
func_text += ")\n"
# Get the shape of the first input array.
shape_name = ir_utils.get_unused_var_name("full_shape", name_var_table)
func_text += " {} = {}.shape\n".format(shape_name, first_arg)
# If we have to allocate the output array (the out argument was not used)
# then us numpy.full if the user specified a cval stencil decorator option
# or np.zeros if they didn't to allocate the array.
if result is None:
if "cval" in self.options:
cval = self.options["cval"]
if return_type.dtype != typing.typeof.typeof(cval):
raise ValueError(
"cval type does not match stencil return type.")
out_init ="{} = np.full({}, {}, dtype=np.{})\n".format(
out_name, shape_name, cval, return_type.dtype)
else:
out_init ="{} = np.zeros({}, dtype=np.{})\n".format(
out_name, shape_name, return_type.dtype)
func_text += " " + out_init
offset = 1
# Add the loop nests to the new function.
for i in range(the_array.ndim):
for j in range(offset):
func_text += " "
# ranges[i][0] is the minimum index used in the i'th dimension
# but minimum's greater than 0 don't preclude any entry in the array.
# So, take the minimum of 0 and the minimum index found in the kernel
# and this will be a negative number (potentially -0). Then, we do
# unary - on that to get the positive offset in this dimension whose
# use is precluded.
# ranges[i][1] is the maximum of 0 and the observed maximum index
# in this dimension because negative maximums would not cause us to
# preclude any entry in the array from being used.
func_text += ("for {} in range(-min(0,{}),"
"{}[{}]-max(0,{})):\n").format(
index_vars[i],
ranges[i][0],
shape_name,
i,
ranges[i][1])
offset += 1
for j in range(offset):
func_text += " "
# Put a sentinel in the code so we can locate it in the IR. We will
# remove this sentinel assignment and replace it with the IR for the
# stencil kernel body.
func_text += "{} = 0\n".format(sentinel_name)
func_text += " return {}\n".format(out_name)
if config.DEBUG_ARRAY_OPT == 1:
print("new stencil func text")
print(func_text)
# Force the new stencil function into existence.
exec_(func_text) in globals(), locals()
stencil_func = eval(stencil_func_name)
if sigret is not None:
pysig = utils.pysignature(stencil_func)
sigret.pysig = pysig
# Get the IR for the newly created stencil function.
stencil_ir = compiler.run_frontend(stencil_func)
ir_utils.remove_dels(stencil_ir.blocks)
# rename all variables in stencil_ir afresh
var_table = ir_utils.get_name_var_table(stencil_ir.blocks)
new_var_dict = {}
reserved_names = ([sentinel_name, out_name, neighborhood_name,
shape_name] + kernel_copy.arg_names + index_vars)
for name, var in var_table.items():
if not name in reserved_names:
new_var_dict[name] = ir_utils.mk_unique_var(name)
ir_utils.replace_var_names(stencil_ir.blocks, new_var_dict)
stencil_stub_last_label = max(stencil_ir.blocks.keys()) + 1
# Shift lables in the kernel copy so they are guaranteed unique
# and don't conflict with any labels in the stencil_ir.
kernel_copy.blocks = ir_utils.add_offset_to_labels(
kernel_copy.blocks, stencil_stub_last_label)
new_label = max(kernel_copy.blocks.keys()) + 1
# Adjust ret_blocks to account for addition of the offset.
ret_blocks = [x + stencil_stub_last_label for x in ret_blocks]
if config.DEBUG_ARRAY_OPT == 1:
print("ret_blocks w/ offsets", ret_blocks, stencil_stub_last_label)
print("before replace sentinel stencil_ir")
ir_utils.dump_blocks(stencil_ir.blocks)
print("before replace sentinel kernel_copy")
ir_utils.dump_blocks(kernel_copy.blocks)
# Search all the block in the stencil outline for the sentinel.
for label, block in stencil_ir.blocks.items():
for i, inst in enumerate(block.body):
if (isinstance( inst, ir.Assign) and
inst.target.name == sentinel_name):
# We found the sentinel assignment.
loc = inst.loc
scope = block.scope
# split block across __sentinel__
# A new block is allocated for the statements prior to the
# sentinel but the new block maintains the current block
# label.
prev_block = ir.Block(scope, loc)
prev_block.body = block.body[:i]
# The current block is used for statements after sentinel.
block.body = block.body[i + 1:]
# But the current block gets a new label.
body_first_label = min(kernel_copy.blocks.keys())
# The previous block jumps to the minimum labelled block of
# the parfor body.
prev_block.append(ir.Jump(body_first_label, loc))
# Add all the parfor loop body blocks to the gufunc
# function's IR.
for (l, b) in kernel_copy.blocks.items():
stencil_ir.blocks[l] = b
stencil_ir.blocks[new_label] = block
stencil_ir.blocks[label] = prev_block
# Add a jump from all the blocks that previously contained
# a return in the stencil kernel to the block
# containing statements after the sentinel.
for ret_block in ret_blocks:
stencil_ir.blocks[ret_block].append(
ir.Jump(new_label, loc))
break
else:
continue
break
stencil_ir.blocks = ir_utils.rename_labels(stencil_ir.blocks)
ir_utils.remove_dels(stencil_ir.blocks)
assert(isinstance(the_array, types.Type))
array_types = args
new_stencil_param_types = list(array_types)
if config.DEBUG_ARRAY_OPT == 1:
print("new_stencil_param_types", new_stencil_param_types)
ir_utils.dump_blocks(stencil_ir.blocks)
# Compile the combined stencil function with the replaced loop
# body in it.
new_func = compiler.compile_ir(
self._typingctx,
self._targetctx,
stencil_ir,
new_stencil_param_types,
None,
compiler.DEFAULT_FLAGS,
{})
return new_func
def inline_array(array_var, expr, stmts, list_vars, dels):
"""Check to see if the given "array_var" is created from a list
of constants, and try to inline the list definition as array
initialization.
Extra statements produced with be appended to "stmts".
"""
callname = guard(find_callname, func_ir, expr)
require(callname and callname[1] == 'numpy' and callname[0] == 'array')
require(expr.args[0].name in list_vars)
ret_type = calltypes[expr].return_type
require(isinstance(ret_type, types.ArrayCompatible) and
ret_type.ndim == 1)
loc = expr.loc
list_var = expr.args[0]
# Get the type of the array to be created.
array_typ = typemap[array_var.name]
debug_print("inline array_var = ", array_var, " list_var = ", list_var)
# Get the element type of the array to be created.
dtype = array_typ.dtype
# Get the sequence of operations to provide values to the new array.
seq, _ = find_build_sequence(func_ir, list_var)
size = len(seq)
# Create a tuple to pass to empty below to specify the new array size.
size_var = ir.Var(scope, mk_unique_var("size"), loc)
size_tuple_var = ir.Var(scope, mk_unique_var("size_tuple"), loc)
size_typ = types.intp
size_tuple_typ = types.UniTuple(size_typ, 1)
typemap[size_var.name] = size_typ
typemap[size_tuple_var.name] = size_tuple_typ
stmts.append(_new_definition(func_ir, size_var,
ir.Const(size, loc=loc), loc))
stmts.append(_new_definition(func_ir, size_tuple_var,
ir.Expr.build_tuple(items=[size_var], loc=loc), loc))
# The general approach is to create an empty array and then fill
# the elements in one-by-one from their specificiation.
# Get the numpy type to pass to empty.
nptype = types.DType(dtype)
# Create a variable to hold the numpy empty function.
empty_func = ir.Var(scope, mk_unique_var("empty_func"), loc)
fnty = get_np_ufunc_typ(np.empty)
sig = context.resolve_function_type(fnty, (size_typ,), {'dtype':nptype})
typemap[empty_func.name] = fnty
stmts.append(_new_definition(func_ir, empty_func,
ir.Global('empty', np.empty, loc=loc), loc))
# We pass two arguments to empty, first the size tuple and second
# the dtype of the new array. Here, we created typ_var which is
# the dtype argument of the new array. typ_var in turn is created
# by getattr of the dtype string on the numpy module.
# Create var for numpy module.
g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
typemap[g_np_var.name] = types.misc.Module(np)
g_np = ir.Global('np', np, loc)
stmts.append(_new_definition(func_ir, g_np_var, g_np, loc))
# Create var for result of numpy.<dtype>.
typ_var = ir.Var(scope, mk_unique_var("$np_typ_var"), loc)
typemap[typ_var.name] = nptype
dtype_str = str(dtype)
if dtype_str == 'bool':
dtype_str = 'bool_'
# Get dtype attribute of numpy module.
np_typ_getattr = ir.Expr.getattr(g_np_var, dtype_str, loc)
stmts.append(_new_definition(func_ir, typ_var, np_typ_getattr, loc))
# Create the call to numpy.empty passing the size tuple and dtype var.
empty_call = ir.Expr.call(empty_func, [size_var, typ_var], {}, loc=loc)
calltypes[empty_call] = typing.signature(array_typ, size_typ, nptype)
stmts.append(_new_definition(func_ir, array_var, empty_call, loc))
# Fill in the new empty array one-by-one.
for i in range(size):
index_var = ir.Var(scope, mk_unique_var("index"), loc)
index_typ = types.intp
typemap[index_var.name] = index_typ
stmts.append(_new_definition(func_ir, index_var,
ir.Const(i, loc), loc))
setitem = ir.SetItem(array_var, index_var, seq[i], loc)
calltypes[setitem] = typing.signature(types.none, array_typ,
index_typ, dtype)
stmts.append(setitem)
stmts.extend(dels)
return True
def get_stencil_ir(sf, typingctx, args, scope, loc, input_dict, typemap,
calltypes):
"""get typed IR from stencil bytecode
"""
from numba.targets.cpu import CPUContext
from numba.targets.registry import cpu_target
from numba.annotations import type_annotations
from numba.compiler import type_inference_stage
# get untyped IR
stencil_func_ir = sf.kernel_ir.copy()
# copy the IR nodes to avoid changing IR in the StencilFunc object
stencil_blocks = copy.deepcopy(stencil_func_ir.blocks)
stencil_func_ir.blocks = stencil_blocks
name_var_table = ir_utils.get_name_var_table(stencil_func_ir.blocks)
if "out" in name_var_table:
raise ValueError("Cannot use the reserved word 'out' in stencil kernels.")
# get typed IR with a dummy pipeline (similar to test_parfors.py)
targetctx = CPUContext(typingctx)
with cpu_target.nested_context(typingctx, targetctx):
tp = DummyPipeline(typingctx, targetctx, args, stencil_func_ir)
numba.rewrites.rewrite_registry.apply(
'before-inference', tp, tp.func_ir)
tp.typemap, tp.return_type, tp.calltypes = type_inference_stage(
tp.typingctx, tp.func_ir, tp.args, None)
type_annotations.TypeAnnotation(
func_ir=tp.func_ir,
typemap=tp.typemap,
calltypes=tp.calltypes,
lifted=(),
lifted_from=None,
args=tp.args,
return_type=tp.return_type,
html_output=numba.config.HTML)
# make block labels unique
stencil_blocks = ir_utils.add_offset_to_labels(stencil_blocks,
ir_utils.next_label())
min_label = min(stencil_blocks.keys())
max_label = max(stencil_blocks.keys())
ir_utils._max_label = max_label
if config.DEBUG_ARRAY_OPT == 1:
print("Initial stencil_blocks")
ir_utils.dump_blocks(stencil_blocks)
# rename variables,
var_dict = {}
for v, typ in tp.typemap.items():
new_var = ir.Var(scope, mk_unique_var(v), loc)
var_dict[v] = new_var
typemap[new_var.name] = typ # add new var type for overall function
ir_utils.replace_vars(stencil_blocks, var_dict)
if config.DEBUG_ARRAY_OPT == 1:
print("After replace_vars")
ir_utils.dump_blocks(stencil_blocks)
# add call types to overall function
for call, call_typ in tp.calltypes.items():
calltypes[call] = call_typ
arg_to_arr_dict = {}
# replace arg with arr
for block in stencil_blocks.values():
for stmt in block.body:
if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Arg):
if config.DEBUG_ARRAY_OPT == 1:
print("input_dict", input_dict, stmt.value.index,
stmt.value.name, stmt.value.index in input_dict)
arg_to_arr_dict[stmt.value.name] = input_dict[stmt.value.index].name
stmt.value = input_dict[stmt.value.index]
if config.DEBUG_ARRAY_OPT == 1:
print("arg_to_arr_dict", arg_to_arr_dict)
print("After replace arg with arr")
ir_utils.dump_blocks(stencil_blocks)
ir_utils.remove_dels(stencil_blocks)
stencil_func_ir.blocks = stencil_blocks
return stencil_func_ir, sf.get_return_type(args)[0], arg_to_arr_dict
def _add_index_offsets(self, index_list, index_offsets, new_body,
scope, loc):
""" Does the actual work of adding loop index variables to the
relative index constants or variables.
"""
assert len(index_list) == len(index_offsets)
# shortcut if all values are integer
if all([isinstance(v, int) for v in index_list+index_offsets]):
# add offsets in all dimensions
return list(map(add, index_list, index_offsets))
out_nodes = []
index_vars = []
for i in range(len(index_list)):
# new_index = old_index + offset
old_index_var = index_list[i]
if isinstance(old_index_var, int):
old_index_var = ir.Var(scope,
mk_unique_var("old_index_var"), loc)
self.typemap[old_index_var.name] = types.intp
const_assign = ir.Assign(ir.Const(index_list[i], loc),
old_index_var, loc)
out_nodes.append(const_assign)
offset_var = index_offsets[i]
if isinstance(offset_var, int):
offset_var = ir.Var(scope,
mk_unique_var("offset_var"), loc)
self.typemap[offset_var.name] = types.intp
const_assign = ir.Assign(ir.Const(index_offsets[i], loc),
offset_var, loc)
out_nodes.append(const_assign)
if (isinstance(old_index_var, slice)
or isinstance(self.typemap[old_index_var.name],
types.misc.SliceType)):
# only one arg can be slice
assert self.typemap[offset_var.name] == types.intp
index_var = self._add_offset_to_slice(old_index_var, offset_var,
out_nodes, scope, loc)
index_vars.append(index_var)
continue
if (isinstance(offset_var, slice)
or isinstance(self.typemap[offset_var.name],
types.misc.SliceType)):
# only one arg can be slice
assert self.typemap[old_index_var.name] == types.intp
index_var = self._add_offset_to_slice(offset_var, old_index_var,
out_nodes, scope, loc)
index_vars.append(index_var)
continue
index_var = ir.Var(scope,
mk_unique_var("offset_stencil_index"), loc)
self.typemap[index_var.name] = types.intp
index_call = ir.Expr.binop('+', old_index_var,
offset_var, loc)
self.calltypes[index_call] = ir_utils.find_op_typ('+',
[types.intp, types.intp])
index_assign = ir.Assign(index_call, index_var, loc)
out_nodes.append(index_assign)
index_vars.append(index_var)
new_body.extend(out_nodes)
return index_vars
def _replace_stencil_accesses(self, stencil_blocks, parfor_vars, in_args,
index_offsets, stencil_func, arg_to_arr_dict):
""" Convert relative indexing in the stencil kernel to standard indexing
by adding the loop index variables to the corresponding dimensions
of the array index tuples.
"""
in_arr = in_args[0]
in_arg_names = [x.name for x in in_args]
if "standard_indexing" in stencil_func.options:
for x in stencil_func.options["standard_indexing"]:
if x not in arg_to_arr_dict:
raise ValueError("Standard indexing requested for an array " \
"name not present in the stencil kernel definition.")
standard_indexed = [arg_to_arr_dict[x] for x in
stencil_func.options["standard_indexing"]]
else:
standard_indexed = []
if in_arr.name in standard_indexed:
raise ValueError("The first argument to a stencil kernel must use " \
"relative indexing, not standard indexing.")
ndims = self.typemap[in_arr.name].ndim
scope = in_arr.scope
loc = in_arr.loc
# replace access indices, find access lengths in each dimension
need_to_calc_kernel = stencil_func.neighborhood is None
# If we need to infer the kernel size then initialize the minimum and
# maximum seen indices for each dimension to 0. If we already have
# the neighborhood calculated then just convert from neighborhood format
# to the separate start and end lengths format used here.
if need_to_calc_kernel:
start_lengths = ndims*[0]
end_lengths = ndims*[0]
else:
start_lengths = [x[0] for x in stencil_func.neighborhood]
end_lengths = [x[1] for x in stencil_func.neighborhood]
# Get all the tuples defined in the stencil blocks.
tuple_table = ir_utils.get_tuple_table(stencil_blocks)
found_relative_index = False
# For all blocks in the stencil kernel...
for label, block in stencil_blocks.items():
new_body = []
# For all statements in those blocks...
for stmt in block.body:
# Reject assignments to input arrays.
if ((isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op in ['setitem', 'static_setitem']
and stmt.value.value.name in in_arg_names) or
((isinstance(stmt, ir.SetItem) or
isinstance(stmt, ir.StaticSetItem))
and stmt.target.name in in_arg_names)):
raise ValueError("Assignments to arrays passed to stencil kernels is not allowed.")
# We found a getitem for some array. If that array is an input
# array and isn't in the list of standard indexed arrays then
# update min and max seen indices if we are inferring the
# kernel size and create a new tuple where the relative offsets
# are added to loop index vars to get standard indexing.
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op in ['static_getitem', 'getitem']
and stmt.value.value.name in in_arg_names
and stmt.value.value.name not in standard_indexed):
index_list = stmt.value.index
# handle 1D case
if ndims == 1:
index_list = [index_list]
else:
if hasattr(index_list, 'name') and index_list.name in tuple_table:
index_list = tuple_table[index_list.name]
if index_offsets:
index_list = self._add_index_offsets(index_list,
list(index_offsets), new_body, scope, loc)
# update min and max indices
if need_to_calc_kernel:
# all indices should be integer to be able to calculate
# neighborhood automatically
if (isinstance(index_list, ir.Var) or
any([not isinstance(v, int) for v in index_list])):
raise ValueError("Variable stencil index only "
"possible with known neighborhood")
start_lengths = list(map(min, start_lengths,
index_list))
end_lengths = list(map(max, end_lengths, index_list))
found_relative_index = True
# update access indices
index_vars = self._add_index_offsets(parfor_vars,
list(index_list), new_body, scope, loc)
# new access index tuple
if ndims == 1:
ind_var = index_vars[0]
else:
ind_var = ir.Var(scope, mk_unique_var(
"$parfor_index_ind_var"), loc)
self.typemap[ind_var.name] = types.containers.UniTuple(
types.intp, ndims)
tuple_call = ir.Expr.build_tuple(index_vars, loc)
tuple_assign = ir.Assign(tuple_call, ind_var, loc)
new_body.append(tuple_assign)
# getitem return type is scalar if all indices are integer
if all([self.typemap[v.name] == types.intp
for v in index_vars]):
getitem_return_typ = self.typemap[
stmt.value.value.name].dtype
else:
# getitem returns an array
getitem_return_typ = self.typemap[stmt.value.value.name]
# new getitem with the new index var
getitem_call = ir.Expr.getitem(stmt.value.value, ind_var,
loc)
self.calltypes[getitem_call] = signature(
getitem_return_typ,
self.typemap[stmt.value.value.name],
self.typemap[ind_var.name])
stmt.value = getitem_call
new_body.append(stmt)
block.body = new_body
if need_to_calc_kernel and not found_relative_index:
raise ValueError("Stencil kernel with no accesses to " \
"relatively indexed arrays.")
return start_lengths, end_lengths
def _mk_stencil_parfor(self, label, in_args, out_arr, stencil_ir,
index_offsets, target, return_type, stencil_func,
arg_to_arr_dict):
""" Converts a set of stencil kernel blocks to a parfor.
"""
gen_nodes = []
stencil_blocks = stencil_ir.blocks
if config.DEBUG_ARRAY_OPT == 1:
print("_mk_stencil_parfor", label, in_args, out_arr, index_offsets,
return_type, stencil_func, stencil_blocks)
ir_utils.dump_blocks(stencil_blocks)
in_arr = in_args[0]
# run copy propagate to replace in_args copies (e.g. a = A)
in_arr_typ = self.typemap[in_arr.name]
in_cps, out_cps = ir_utils.copy_propagate(stencil_blocks, self.typemap)
name_var_table = ir_utils.get_name_var_table(stencil_blocks)
ir_utils.apply_copy_propagate(
stencil_blocks,
in_cps,
name_var_table,
self.typemap,
self.calltypes)
if config.DEBUG_ARRAY_OPT == 1:
print("stencil_blocks after copy_propagate")
ir_utils.dump_blocks(stencil_blocks)
ir_utils.remove_dead(stencil_blocks, self.func_ir.arg_names, stencil_ir,
self.typemap)
if config.DEBUG_ARRAY_OPT == 1:
print("stencil_blocks after removing dead code")
ir_utils.dump_blocks(stencil_blocks)
# create parfor vars
ndims = self.typemap[in_arr.name].ndim
scope = in_arr.scope
loc = in_arr.loc
parfor_vars = []
for i in range(ndims):
parfor_var = ir.Var(scope, mk_unique_var(
"$parfor_index_var"), loc)
self.typemap[parfor_var.name] = types.intp
parfor_vars.append(parfor_var)
start_lengths, end_lengths = self._replace_stencil_accesses(
stencil_blocks, parfor_vars, in_args, index_offsets, stencil_func,
arg_to_arr_dict)
if config.DEBUG_ARRAY_OPT == 1:
print("stencil_blocks after replace stencil accesses")
ir_utils.dump_blocks(stencil_blocks)
# create parfor loop nests
loopnests = []
equiv_set = self.array_analysis.get_equiv_set(label)
in_arr_dim_sizes = equiv_set.get_shape(in_arr)
assert ndims == len(in_arr_dim_sizes)
for i in range(ndims):
last_ind = self._get_stencil_last_ind(in_arr_dim_sizes[i],
end_lengths[i], gen_nodes, scope, loc)
start_ind = self._get_stencil_start_ind(
start_lengths[i], gen_nodes, scope, loc)
# start from stencil size to avoid invalid array access
loopnests.append(numba.parfor.LoopNest(parfor_vars[i],
start_ind, last_ind, 1))
# We have to guarantee that the exit block has maximum label and that
# there's only one exit block for the parfor body.
# So, all return statements will change to jump to the parfor exit block.
parfor_body_exit_label = max(stencil_blocks.keys()) + 1
stencil_blocks[parfor_body_exit_label] = ir.Block(scope, loc)
exit_value_var = ir.Var(scope, mk_unique_var("$parfor_exit_value"), loc)
self.typemap[exit_value_var.name] = return_type.dtype
# create parfor index var
for_replacing_ret = []
if ndims == 1:
parfor_ind_var = parfor_vars[0]
else:
parfor_ind_var = ir.Var(scope, mk_unique_var(
"$parfor_index_tuple_var"), loc)
self.typemap[parfor_ind_var.name] = types.containers.UniTuple(
types.intp, ndims)
tuple_call = ir.Expr.build_tuple(parfor_vars, loc)
tuple_assign = ir.Assign(tuple_call, parfor_ind_var, loc)
for_replacing_ret.append(tuple_assign)
if config.DEBUG_ARRAY_OPT == 1:
print("stencil_blocks after creating parfor index var")
ir_utils.dump_blocks(stencil_blocks)
# empty init block
init_block = ir.Block(scope, loc)
if out_arr == None:
in_arr_typ = self.typemap[in_arr.name]
shape_name = ir_utils.mk_unique_var("in_arr_shape")
shape_var = ir.Var(scope, shape_name, loc)
shape_getattr = ir.Expr.getattr(in_arr, "shape", loc)
self.typemap[shape_name] = types.containers.UniTuple(types.intp,
in_arr_typ.ndim)
init_block.body.extend([ir.Assign(shape_getattr, shape_var, loc)])
zero_name = ir_utils.mk_unique_var("zero_val")
zero_var = ir.Var(scope, zero_name, loc)
if "cval" in stencil_func.options:
cval = stencil_func.options["cval"]
# TODO: Loosen this restriction to adhere to casting rules.
if return_type.dtype != typing.typeof.typeof(cval):
raise ValueError("cval type does not match stencil return type.")
temp2 = return_type.dtype(cval)
else:
temp2 = return_type.dtype(0)
full_const = ir.Const(temp2, loc)
self.typemap[zero_name] = return_type.dtype
init_block.body.extend([ir.Assign(full_const, zero_var, loc)])
so_name = ir_utils.mk_unique_var("stencil_output")
out_arr = ir.Var(scope, so_name, loc)
self.typemap[out_arr.name] = numba.types.npytypes.Array(
return_type.dtype,
in_arr_typ.ndim,
in_arr_typ.layout)
dtype_g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
self.typemap[dtype_g_np_var.name] = types.misc.Module(np)
dtype_g_np = ir.Global('np', np, loc)
dtype_g_np_assign = ir.Assign(dtype_g_np, dtype_g_np_var, loc)
init_block.body.append(dtype_g_np_assign)
dtype_np_attr_call = ir.Expr.getattr(dtype_g_np_var, return_type.dtype.name, loc)
dtype_attr_var = ir.Var(scope, mk_unique_var("$np_attr_attr"), loc)
self.typemap[dtype_attr_var.name] = types.functions.NumberClass(return_type.dtype)
dtype_attr_assign = ir.Assign(dtype_np_attr_call, dtype_attr_var, loc)
init_block.body.append(dtype_attr_assign)
stmts = ir_utils.gen_np_call("full",
np.full,
out_arr,
[shape_var, zero_var, dtype_attr_var],
self.typingctx,
self.typemap,
self.calltypes)
equiv_set.insert_equiv(out_arr, in_arr_dim_sizes)
init_block.body.extend(stmts)
self.replace_return_with_setitem(stencil_blocks, exit_value_var,
parfor_body_exit_label)
if config.DEBUG_ARRAY_OPT == 1:
print("stencil_blocks after replacing return")
ir_utils.dump_blocks(stencil_blocks)
setitem_call = ir.SetItem(out_arr, parfor_ind_var, exit_value_var, loc)
self.calltypes[setitem_call] = signature(
types.none, self.typemap[out_arr.name],
self.typemap[parfor_ind_var.name],
self.typemap[out_arr.name].dtype
)
stencil_blocks[parfor_body_exit_label].body.extend(for_replacing_ret)
stencil_blocks[parfor_body_exit_label].body.append(setitem_call)
# simplify CFG of parfor body (exit block could be simplified often)
# add dummy return to enable CFG
stencil_blocks[parfor_body_exit_label].body.append(ir.Return(0,
ir.Loc("stencilparfor_dummy", -1)))
stencil_blocks = ir_utils.simplify_CFG(stencil_blocks)
stencil_blocks[max(stencil_blocks.keys())].body.pop()
if config.DEBUG_ARRAY_OPT == 1:
print("stencil_blocks after adding SetItem")
ir_utils.dump_blocks(stencil_blocks)
pattern = ('stencil', [start_lengths, end_lengths])
parfor = numba.parfor.Parfor(loopnests, init_block, stencil_blocks,
loc, parfor_ind_var, equiv_set, pattern, self.flags)
gen_nodes.append(parfor)
gen_nodes.append(ir.Assign(out_arr, target, loc))
return gen_nodes
def _inline_arraycall(func_ir, cfg, visited, loop, enable_prange=False):
"""Look for array(list) call in the exit block of a given loop, and turn list operations into
array operations in the loop if the following conditions are met:
1. The exit block contains an array call on the list;
2. The list variable is no longer live after array call;
3. The list is created in the loop entry block;
4. The loop is created from an range iterator whose length is known prior to the loop;
5. There is only one list_append operation on the list variable in the loop body;
6. The block that contains list_append dominates the loop head, which ensures list
length is the same as loop length;
If any condition check fails, no modification will be made to the incoming IR.
"""
debug_print = _make_debug_print("inline_arraycall")
# There should only be one loop exit
require(len(loop.exits) == 1)
exit_block = next(iter(loop.exits))
list_var, array_call_index, array_kws = _find_arraycall(func_ir, func_ir.blocks[exit_block])
# check if dtype is present in array call
dtype_def = None
dtype_mod_def = None
if 'dtype' in array_kws:
require(isinstance(array_kws['dtype'], ir.Var))
# We require that dtype argument to be a constant of getattr Expr, and we'll
# remember its definition for later use.
dtype_def = get_definition(func_ir, array_kws['dtype'])
require(isinstance(dtype_def, ir.Expr) and dtype_def.op == 'getattr')
dtype_mod_def = get_definition(func_ir, dtype_def.value)
list_var_def = get_definition(func_ir, list_var)
debug_print("list_var = ", list_var, " def = ", list_var_def)
if isinstance(list_var_def, ir.Expr) and list_var_def.op == 'cast':
list_var_def = get_definition(func_ir, list_var_def.value)
# Check if the definition is a build_list
require(isinstance(list_var_def, ir.Expr) and list_var_def.op == 'build_list')
# Look for list_append in "last" block in loop body, which should be a block that is
# a post-dominator of the loop header.
list_append_stmts = []
for label in loop.body:
# We have to consider blocks of this loop, but not sub-loops.
# To achieve this, we require the set of "in_loops" of "label" to be visited loops.
in_visited_loops = [l.header in visited for l in cfg.in_loops(label)]
if not all(in_visited_loops):
continue
block = func_ir.blocks[label]
debug_print("check loop body block ", label)
for stmt in block.find_insts(ir.Assign):
lhs = stmt.target
expr = stmt.value
if isinstance(expr, ir.Expr) and expr.op == 'call':
func_def = get_definition(func_ir, expr.func)
if isinstance(func_def, ir.Expr) and func_def.op == 'getattr' \
and func_def.attr == 'append':
list_def = get_definition(func_ir, func_def.value)
debug_print("list_def = ", list_def, list_def == list_var_def)
if list_def == list_var_def:
# found matching append call
list_append_stmts.append((label, block, stmt))
# Require only one list_append, otherwise we won't know the indices
require(len(list_append_stmts) == 1)
append_block_label, append_block, append_stmt = list_append_stmts[0]
# Check if append_block (besides loop entry) dominates loop header.
# Since CFG doesn't give us this info without loop entry, we approximate
# by checking if the predecessor set of the header block is the same
# as loop_entries plus append_block, which is certainly more restrictive
# than necessary, and can be relaxed if needed.
preds = set(l for l, b in cfg.predecessors(loop.header))
debug_print("preds = ", preds, (loop.entries | set([append_block_label])))
require(preds == (loop.entries | set([append_block_label])))
# Find iterator in loop header
iter_vars = []
iter_first_vars = []
loop_header = func_ir.blocks[loop.header]
for stmt in loop_header.find_insts(ir.Assign):
expr = stmt.value
if isinstance(expr, ir.Expr):
if expr.op == 'iternext':
iter_def = get_definition(func_ir, expr.value)
debug_print("iter_def = ", iter_def)
iter_vars.append(expr.value)
elif expr.op == 'pair_first':
iter_first_vars.append(stmt.target)
# Require only one iterator in loop header
require(len(iter_vars) == 1 and len(iter_first_vars) == 1)
iter_var = iter_vars[0] # variable that holds the iterator object
iter_first_var = iter_first_vars[0] # variable that holds the value out of iterator
# Final requirement: only one loop entry, and we're going to modify it by:
# 1. replacing the list definition with an array definition;
# 2. adding a counter for the array iteration.
require(len(loop.entries) == 1)
loop_entry = func_ir.blocks[next(iter(loop.entries))]
terminator = loop_entry.terminator
scope = loop_entry.scope
loc = loop_entry.loc
stmts = []
removed = []
def is_removed(val, removed):
if isinstance(val, ir.Var):
for x in removed:
if x.name == val.name:
return True
return False
# Skip list construction and skip terminator, add the rest to stmts
for i in range(len(loop_entry.body) - 1):
stmt = loop_entry.body[i]
if isinstance(stmt, ir.Assign) and (stmt.value == list_def or is_removed(stmt.value, removed)):
removed.append(stmt.target)
else:
stmts.append(stmt)
debug_print("removed variables: ", removed)
# Define an index_var to index the array.
# If the range happens to be single step ranges like range(n), or range(m, n),
# then the index_var correlates to iterator index; otherwise we'll have to
# define a new counter.
range_def = guard(_find_iter_range, func_ir, iter_var)
index_var = ir.Var(scope, mk_unique_var("index"), loc)
if range_def and range_def[0] == 0:
# iterator starts with 0, index_var can just be iter_first_var
index_var = iter_first_var
else:
# index_var = -1 # starting the index with -1 since it will incremented in loop header
stmts.append(_new_definition(func_ir, index_var, ir.Const(value=-1, loc=loc), loc))
# Insert statement to get the size of the loop iterator
size_var = ir.Var(scope, mk_unique_var("size"), loc)
if range_def:
start, stop, range_func_def = range_def
if start == 0:
size_val = stop
else:
size_val = ir.Expr.binop(fn='-', lhs=stop, rhs=start, loc=loc)
# we can parallelize this loop if enable_prange = True, by changing
# range function from range, to prange.
if enable_prange and isinstance(range_func_def, ir.Global):
range_func_def.name = 'internal_prange'
range_func_def.value = internal_prange
else:
len_func_var = ir.Var(scope, mk_unique_var("len_func"), loc)
stmts.append(_new_definition(func_ir, len_func_var,
ir.Global('range_iter_len', range_iter_len, loc=loc), loc))
size_val = ir.Expr.call(len_func_var, (iter_var,), (), loc=loc)
stmts.append(_new_definition(func_ir, size_var, size_val, loc))
size_tuple_var = ir.Var(scope, mk_unique_var("size_tuple"), loc)
stmts.append(_new_definition(func_ir, size_tuple_var,
ir.Expr.build_tuple(items=[size_var], loc=loc), loc))
# Insert array allocation
array_var = ir.Var(scope, mk_unique_var("array"), loc)
empty_func = ir.Var(scope, mk_unique_var("empty_func"), loc)
if dtype_def and dtype_mod_def:
# when dtype is present, we'll call emtpy with dtype
dtype_mod_var = ir.Var(scope, mk_unique_var("dtype_mod"), loc)
dtype_var = ir.Var(scope, mk_unique_var("dtype"), loc)
stmts.append(_new_definition(func_ir, dtype_mod_var, dtype_mod_def, loc))
stmts.append(_new_definition(func_ir, dtype_var,
ir.Expr.getattr(dtype_mod_var, dtype_def.attr, loc), loc))
stmts.append(_new_definition(func_ir, empty_func,
ir.Global('empty', np.empty, loc=loc), loc))
array_kws = [('dtype', dtype_var)]
else:
# otherwise we'll call unsafe_empty_inferred
stmts.append(_new_definition(func_ir, empty_func,
ir.Global('unsafe_empty_inferred',
unsafe_empty_inferred, loc=loc), loc))
array_kws = []
# array_var = empty_func(size_tuple_var)
stmts.append(_new_definition(func_ir, array_var,
ir.Expr.call(empty_func, (size_tuple_var,), list(array_kws), loc=loc), loc))
# Add back removed just in case they are used by something else
for var in removed:
stmts.append(_new_definition(func_ir, var, array_var, loc))
# Add back terminator
stmts.append(terminator)
# Modify loop_entry
loop_entry.body = stmts
if range_def:
if range_def[0] != 0:
# when range doesn't start from 0, index_var becomes loop index
# (iter_first_var) minus an offset (range_def[0])
terminator = loop_header.terminator
assert(isinstance(terminator, ir.Branch))
# find the block in the loop body that header jumps to
block_id = terminator.truebr
blk = func_ir.blocks[block_id]
loc = blk.loc
blk.body.insert(0, _new_definition(func_ir, index_var,
ir.Expr.binop(fn='-', lhs=iter_first_var,
rhs=range_def[0], loc=loc),
loc))
else:
# Insert index_var increment to the end of loop header
loc = loop_header.loc
terminator = loop_header.terminator
stmts = loop_header.body[0:-1]
next_index_var = ir.Var(scope, mk_unique_var("next_index"), loc)
one = ir.Var(scope, mk_unique_var("one"), loc)
# one = 1
stmts.append(_new_definition(func_ir, one,
ir.Const(value=1,loc=loc), loc))
# next_index_var = index_var + 1
stmts.append(_new_definition(func_ir, next_index_var,
ir.Expr.binop(fn='+', lhs=index_var, rhs=one, loc=loc), loc))
# index_var = next_index_var
stmts.append(_new_definition(func_ir, index_var, next_index_var, loc))
stmts.append(terminator)
loop_header.body = stmts
# In append_block, change list_append into array assign
for i in range(len(append_block.body)):
if append_block.body[i] == append_stmt:
debug_print("Replace append with SetItem")
append_block.body[i] = ir.SetItem(target=array_var, index=index_var,
value=append_stmt.value.args[0], loc=append_stmt.loc)
# replace array call, by changing "a = array(b)" to "a = b"
stmt = func_ir.blocks[exit_block].body[array_call_index]
# stmt can be either array call or SetItem, we only replace array call
if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr):
stmt.value = array_var
func_ir._definitions[stmt.target.name] = [stmt.value]
return True
