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 _replace_stencil_accesses(self, stencil_ir, 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.
"""
stencil_blocks = stencil_ir.blocks
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]
# indices can be inferred as constant in simple expressions
# like -c where c is constant
# handled here since this is a common stencil index pattern
stencil_ir._definitions = ir_utils.build_definitions(
stencil_blocks)
index_list = [
_get_const_index_expr(stencil_ir, self.func_ir, v)
for v in index_list
]
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 _init_run(self):
self.func_ir._definitions = build_definitions(self.func_ir.blocks)
self._parallel_accesses = set()
self._T_arrs = set()
self.second_pass = False
self.in_parallel_parfor = -1
def _replace_stencil_accesses(self, stencil_ir, 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.
"""
stencil_blocks = stencil_ir.blocks
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]
# indices can be inferred as constant in simple expressions
# like -c where c is constant
# handled here since this is a common stencil index pattern
stencil_ir._definitions = ir_utils.build_definitions(stencil_blocks)
index_list = [_get_const_index_expr(
stencil_ir, self.func_ir, v) for v in index_list]
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
