def _rebalance_arrs(self, array_dists, parfor_dists):
# rebalance an array if it is accessed in a parfor that has output
# arrays or is in a loop
# find sequential loop bodies
cfg = numba.analysis.compute_cfg_from_blocks(self.func_ir.blocks)
loop_bodies = set()
for loop in cfg.loops().values():
loop_bodies |= loop.body
rebalance_arrs = set()
for label, block in self.func_ir.blocks.items():
for inst in block.body:
# TODO: handle hiframes filter etc.
if (isinstance(inst, Parfor)
and parfor_dists[inst.id] == Distribution.OneD_Var):
array_accesses = ir_utils.get_array_accesses(
inst.loop_body)
onedv_arrs = set(
arr for (arr, ind) in array_accesses
if arr in array_dists
and array_dists[arr] == Distribution.OneD_Var)
if (label in loop_bodies
or _arrays_written(onedv_arrs, inst.loop_body)):
rebalance_arrs |= onedv_arrs
if len(rebalance_arrs) != 0:
self._gen_rebalances(rebalance_arrs, self.func_ir.blocks)
return True
return False
def _analyze_parfor(self, parfor, array_dists, parfor_dists):
if parfor.id not in parfor_dists:
parfor_dists[parfor.id] = Distribution.OneD
# analyze init block first to see array definitions
self._analyze_block(parfor.init_block, array_dists, parfor_dists)
out_dist = Distribution.OneD
if self.in_parallel_parfor != -1:
out_dist = Distribution.REP
parfor_arrs = set() # arrays this parfor accesses in parallel
array_accesses = ir_utils.get_array_accesses(parfor.loop_body)
par_index_var = parfor.loop_nests[0].index_variable.name
#stencil_accesses, _ = get_stencil_accesses(parfor, self.typemap)
for (arr, index) in array_accesses:
if index == par_index_var: #or index in stencil_accesses:
parfor_arrs.add(arr)
self._parallel_accesses.add((arr, index))
# multi-dim case
tup_list = guard(find_build_tuple, self.func_ir, index)
if tup_list is not None:
index_tuple = [var.name for var in tup_list]
if index_tuple[0] == par_index_var:
parfor_arrs.add(arr)
self._parallel_accesses.add((arr, index))
if par_index_var in index_tuple[1:]:
out_dist = Distribution.REP
# TODO: check for index dependency
for arr in parfor_arrs:
if arr in array_dists:
out_dist = Distribution(
min(out_dist.value, array_dists[arr].value))
parfor_dists[parfor.id] = out_dist
for arr in parfor_arrs:
if arr in array_dists:
array_dists[arr] = out_dist
# TODO: find prange actually coming from user
# for pattern in parfor.patterns:
# if pattern[0] == 'prange' and not self.in_parallel_parfor:
# parfor_dists[parfor.id] = Distribution.OneD
# run analysis recursively on parfor body
if self.second_pass and out_dist in [
Distribution.OneD, Distribution.OneD_Var
]:
self.in_parallel_parfor = parfor.id
blocks = wrap_parfor_blocks(parfor)
for b in blocks.values():
self._analyze_block(b, array_dists, parfor_dists)
unwrap_parfor_blocks(parfor)
if self.in_parallel_parfor == parfor.id:
self.in_parallel_parfor = -1
return
def _analyze_parfor(self, parfor, array_dists, parfor_dists):
if parfor.id not in parfor_dists:
parfor_dists[parfor.id] = Distribution.OneD
# analyze init block first to see array definitions
self._analyze_block(parfor.init_block, array_dists, parfor_dists)
out_dist = Distribution.OneD
if self.in_parallel_parfor:
out_dist = Distribution.REP
parfor_arrs = set() # arrays this parfor accesses in parallel
array_accesses = ir_utils.get_array_accesses(parfor.loop_body)
par_index_var = parfor.loop_nests[0].index_variable.name
stencil_accesses, _ = get_stencil_accesses(parfor, self.typemap)
for (arr, index) in array_accesses:
if index == par_index_var or index in stencil_accesses:
parfor_arrs.add(arr)
self._parallel_accesses.add((arr, index))
if index in self._tuple_table:
index_tuple = [(var.name if isinstance(var, ir.Var) else var)
for var in self._tuple_table[index]]
if index_tuple[0] == par_index_var:
parfor_arrs.add(arr)
self._parallel_accesses.add((arr, index))
if par_index_var in index_tuple[1:]:
out_dist = Distribution.REP
# TODO: check for index dependency
for arr in parfor_arrs:
if arr in array_dists:
out_dist = Distribution(
min(out_dist.value, array_dists[arr].value))
parfor_dists[parfor.id] = out_dist
for arr in parfor_arrs:
if arr in array_dists:
array_dists[arr] = out_dist
# TODO: find prange actually coming from user
# for pattern in parfor.patterns:
# if pattern[0] == 'prange' and not self.in_parallel_parfor:
# parfor_dists[parfor.id] = Distribution.OneD
# run analysis recursively on parfor body
if self.second_pass and out_dist == Distribution.OneD:
self.in_parallel_parfor = True
blocks = wrap_parfor_blocks(parfor)
for b in blocks.values():
self._analyze_block(b, array_dists, parfor_dists)
unwrap_parfor_blocks(parfor)
self.in_parallel_parfor = False
return
def _analyze_parfor(self, parfor, array_dists, parfor_dists):
if parfor.id not in parfor_dists:
parfor_dists[parfor.id] = Distribution.OneD
# analyze init block first to see array definitions
self._analyze_block(parfor.init_block, array_dists, parfor_dists)
out_dist = Distribution.OneD
parfor_arrs = set() # arrays this parfor accesses in parallel
array_accesses = ir_utils.get_array_accesses(parfor.loop_body)
par_index_var = parfor.loop_nests[0].index_variable.name
stencil_accesses, _ = get_stencil_accesses(parfor.loop_body,
par_index_var)
for (arr, index) in array_accesses:
if index == par_index_var or index in stencil_accesses:
parfor_arrs.add(arr)
self._parallel_accesses.add((arr, index))
if index in self._tuple_table:
index_tuple = [(var.name if isinstance(var, ir.Var) else var)
for var in self._tuple_table[index]]
if index_tuple[0] == par_index_var:
parfor_arrs.add(arr)
self._parallel_accesses.add((arr, index))
if par_index_var in index_tuple[1:]:
out_dist = Distribution.REP
# TODO: check for index dependency
for arr in parfor_arrs:
if arr in array_dists:
out_dist = Distribution(
min(out_dist.value, array_dists[arr].value))
parfor_dists[parfor.id] = out_dist
for arr in parfor_arrs:
if arr in array_dists:
array_dists[arr] = out_dist
# run analysis recursively on parfor body
blocks = wrap_parfor_blocks(parfor)
for b in blocks.values():
self._analyze_block(b, array_dists, parfor_dists)
unwrap_parfor_blocks(parfor)
return