def include_new_blocks(blocks,
new_blocks,
label,
new_body,
remove_non_return=True,
work_list=None,
func_ir=None):
inner_blocks = add_offset_to_labels(new_blocks, ir_utils._max_label + 1)
blocks.update(inner_blocks)
ir_utils._max_label = max(blocks.keys())
scope = blocks[label].scope
loc = blocks[label].loc
inner_topo_order = find_topo_order(inner_blocks)
inner_first_label = inner_topo_order[0]
inner_last_label = inner_topo_order[-1]
if remove_non_return:
remove_return_from_block(inner_blocks[inner_last_label])
new_body.append(ir.Jump(inner_first_label, loc))
blocks[label].body = new_body
label = ir_utils.next_label()
blocks[label] = ir.Block(scope, loc)
if remove_non_return:
inner_blocks[inner_last_label].body.append(ir.Jump(label, loc))
# new_body.clear()
if work_list is not None:
topo_order = find_topo_order(inner_blocks)
for _label in topo_order:
block = inner_blocks[_label]
block.scope = scope
numba.inline_closurecall._add_definitions(func_ir, block)
work_list.append((_label, block))
return label
def run(self):
blocks = self.func_ir.blocks
call_table, _ = ir_utils.get_call_table(blocks)
topo_order = find_topo_order(blocks)
for label in topo_order:
new_body = []
for inst in blocks[label].body:
if isinstance(inst, ir.Assign):
out_nodes = self._run_assign(inst, call_table)
if isinstance(out_nodes, list):
new_body.extend(out_nodes)
if isinstance(out_nodes, dict):
label = include_new_blocks(blocks, out_nodes, label,
new_body)
new_body = []
if isinstance(out_nodes, tuple):
gen_blocks, post_nodes = out_nodes
label = include_new_blocks(blocks, gen_blocks, label,
new_body)
new_body = post_nodes
else:
new_body.append(inst)
blocks[label].body = new_body
self.func_ir._definitions = get_definitions(self.func_ir.blocks)
return
def _run_pd_DatetimeIndex(self, assign, lhs, rhs):
"""transform pd.DatetimeIndex() call with string array argument
"""
kws = dict(rhs.kws)
if 'data' in kws:
data = kws['data']
if len(rhs.args) != 0: # pragma: no cover
raise ValueError(
"only data argument suppoted in pd.DatetimeIndex()")
else:
if len(rhs.args) != 1: # pragma: no cover
raise ValueError(
"data argument in pd.DatetimeIndex() expected")
data = rhs.args[0]
def f(str_arr):
numba.parfor.init_prange()
n = len(str_arr)
S = numba.unsafe.ndarray.empty_inferred((n, ))
for i in numba.parfor.internal_prange(n):
S[i] = hpat.pd_timestamp_ext.parse_datetime_str(str_arr[i])
ret = S
f_ir = compile_to_numba_ir(
f, {
'hpat': hpat,
'numba': numba
}, self.typingctx,
(if_series_to_array_type(self.typemap[data.name]), ), self.typemap,
self.calltypes)
topo_order = find_topo_order(f_ir.blocks)
f_ir.blocks[topo_order[-1]].body[-4].target = lhs
replace_arg_nodes(f_ir.blocks[topo_order[0]], [data])
return f_ir.blocks
def run(self):
dprint_func_ir(self.func_ir, "starting hiframes")
topo_order = find_topo_order(self.func_ir.blocks)
for label in topo_order:
self._get_reverse_copies(self.func_ir.blocks[label].body)
new_body = []
for inst in self.func_ir.blocks[label].body:
# df['col'] = arr
if isinstance(
inst,
ir.StaticSetItem) and inst.target.name in self.df_vars:
df_name = inst.target.name
self.df_vars[df_name][inst.index] = inst.value
self._update_df_cols()
elif isinstance(inst, ir.Assign):
out_nodes = self._run_assign(inst)
if isinstance(out_nodes, list):
new_body.extend(out_nodes)
if isinstance(out_nodes, dict):
label = include_new_blocks(self.func_ir.blocks,
out_nodes, label, new_body)
new_body = []
else:
new_body.append(inst)
self.func_ir.blocks[label].body = new_body
self.func_ir._definitions = get_definitions(self.func_ir.blocks)
self.func_ir.df_cols = self.df_cols
#remove_dead(self.func_ir.blocks, self.func_ir.arg_names)
dprint_func_ir(self.func_ir, "after hiframes")
if numba.config.DEBUG_ARRAY_OPT == 1:
print("df_vars: ", self.df_vars)
return
def run(self):
dprint_func_ir(self.func_ir, "starting IO")
topo_order = find_topo_order(self.func_ir.blocks)
for label in topo_order:
new_body = []
# copies are collected before running the pass since
# variables typed in locals are assigned late
self._get_reverse_copies(self.func_ir.blocks[label].body)
for inst in self.func_ir.blocks[label].body:
if isinstance(inst, ir.Assign):
inst_list = self._run_assign(inst)
new_body.extend(inst_list)
elif isinstance(inst, ir.StaticSetItem):
inst_list = self._run_static_setitem(inst)
new_body.extend(inst_list)
else:
new_body.append(inst)
self.func_ir.blocks[label].body = new_body
# iterative remove dead to make sure all extra code (e.g. df vars) is removed
while remove_dead(self.func_ir.blocks, self.func_ir.arg_names,
self.func_ir):
pass
self.func_ir._definitions = get_definitions(self.func_ir.blocks)
dprint_func_ir(self.func_ir, "after IO")
if debug_prints():
print("h5 files: ", self.h5_files)
print("h5 dsets: ", self.h5_dsets)
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 run(self):
blocks = self.func_ir.blocks
array_dists = {}
parfor_dists = {}
topo_order = find_topo_order(blocks)
self._run_analysis(self.func_ir.blocks, topo_order, array_dists, parfor_dists)
self.second_pass = True
self._run_analysis(self.func_ir.blocks, topo_order, array_dists, parfor_dists)
return _dist_analysis_result(array_dists=array_dists, parfor_dists=parfor_dists)
def run(self):
dprint_func_ir(self.func_ir, "starting hiframes")
topo_order = find_topo_order(self.func_ir.blocks)
for label in topo_order:
new_body = []
for inst in self.func_ir.blocks[label].body:
# df['col'] = arr
if isinstance(inst, ir.StaticSetItem) and inst.target.name in self.df_vars:
df_name = inst.target.name
self.df_vars[df_name][inst.index] = inst.value
self._update_df_cols()
elif isinstance(inst, ir.Assign):
out_nodes = self._run_assign(inst)
if isinstance(out_nodes, list):
new_body.extend(out_nodes)
if isinstance(out_nodes, dict):
inner_blocks = add_offset_to_labels(out_nodes, ir_utils._max_label+1)
self.func_ir.blocks.update(inner_blocks)
ir_utils._max_label = max(self.func_ir.blocks.keys())
scope = self.func_ir.blocks[label].scope
loc = self.func_ir.blocks[label].loc
inner_topo_order = find_topo_order(inner_blocks)
inner_first_label = inner_topo_order[0]
inner_last_label = inner_topo_order[-1]
remove_none_return_from_block(inner_blocks[inner_last_label])
new_body.append(ir.Jump(inner_first_label, loc))
self.func_ir.blocks[label].body = new_body
label = ir_utils.next_label()
self.func_ir.blocks[label] = ir.Block(scope, loc)
inner_blocks[inner_last_label].body.append(ir.Jump(label, loc))
new_body = []
else:
new_body.append(inst)
self.func_ir.blocks[label].body = new_body
remove_dead(self.func_ir.blocks, self.func_ir.arg_names)
dprint_func_ir(self.func_ir, "after hiframes")
if config.DEBUG_ARRAY_OPT==1:
print("df_vars: ", self.df_vars)
return
def run(self):
blocks = self.func_ir.blocks
array_dists = {}
parfor_dists = {}
topo_order = find_topo_order(blocks)
save_array_dists = {}
save_parfor_dists = {1: 1} # dummy value
# fixed-point iteration
while array_dists != save_array_dists or parfor_dists != save_parfor_dists:
save_array_dists = copy.copy(array_dists)
save_parfor_dists = copy.copy(parfor_dists)
for label in topo_order:
self._analyze_block(blocks[label], array_dists, parfor_dists)
return _dist_analysis_result(array_dists=array_dists,
parfor_dists=parfor_dists)
def include_new_blocks(blocks, new_blocks, label, new_body):
inner_blocks = add_offset_to_labels(new_blocks, ir_utils._max_label + 1)
blocks.update(inner_blocks)
ir_utils._max_label = max(blocks.keys())
scope = blocks[label].scope
loc = blocks[label].loc
inner_topo_order = find_topo_order(inner_blocks)
inner_first_label = inner_topo_order[0]
inner_last_label = inner_topo_order[-1]
remove_none_return_from_block(inner_blocks[inner_last_label])
new_body.append(ir.Jump(inner_first_label, loc))
blocks[label].body = new_body
label = ir_utils.next_label()
blocks[label] = ir.Block(scope, loc)
inner_blocks[inner_last_label].body.append(ir.Jump(label, loc))
#new_body.clear()
return label
def inline_new_blocks(func_ir, block, i, callee_blocks, work_list=None):
# adopted from inline_closure_call
scope = block.scope
instr = block.body[i]
# 1. relabel callee_ir by adding an offset
callee_blocks = add_offset_to_labels(callee_blocks,
ir_utils._max_label + 1)
callee_blocks = ir_utils.simplify_CFG(callee_blocks)
max_label = max(callee_blocks.keys())
# reset globals in ir_utils before we use it
ir_utils._max_label = max_label
topo_order = find_topo_order(callee_blocks)
# 5. split caller blocks into two
new_blocks = []
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
new_blocks.append((new_label, new_block))
block.body = block.body[:i]
min_label = topo_order[0]
block.body.append(ir.Jump(min_label, instr.loc))
# 6. replace Return with assignment to LHS
numba.inline_closurecall._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
numba.inline_closurecall._add_definitions(func_ir, block)
func_ir.blocks[label] = block
new_blocks.append((label, block))
if work_list is not None:
for block in new_blocks:
work_list.append(block)
return callee_blocks
def run(self):
"""run array shape analysis on the IR and save information in
array_shape_classes, class_sizes, and array_size_vars (see __init__
comments). May generate some array shape calls if necessary.
"""
dprint_func_ir(self.func_ir, "starting array analysis")
if config.DEBUG_ARRAY_OPT == 1:
print("variable types: ", sorted(self.typemap.items()))
print("call types: ", self.calltypes)
topo_order = find_topo_order(self.func_ir.blocks)
for label in topo_order:
self._analyze_block(self.func_ir.blocks[label])
self._merge_equivalent_classes()
self._cleanup_analysis_data()
if config.DEBUG_ARRAY_OPT == 1:
self.dump()
def run(self):
self._init_run()
blocks = self.func_ir.blocks
array_dists = {}
parfor_dists = {}
topo_order = find_topo_order(blocks)
self._run_analysis(self.func_ir.blocks, topo_order, array_dists,
parfor_dists)
self.second_pass = True
self._run_analysis(self.func_ir.blocks, topo_order, array_dists,
parfor_dists)
# rebalance arrays if necessary
if auto_rebalance and Distribution.OneD_Var in array_dists.values():
changed = self._rebalance_arrs(array_dists, parfor_dists)
if changed:
return self.run()
return _dist_analysis_result(array_dists=array_dists,
parfor_dists=parfor_dists)
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)
f_blocks = self._gen_col_var(var_var, args, col_var)
last_label = find_topo_order(f_blocks)[-1]
def f(a):
a ** 0.5
s_blocks = get_inner_ir(f)
replace_var_names(s_blocks, {'a': var_var.name})
remove_none_return_from_block(s_blocks[0])
assert len(s_blocks[0].body) == 2
const_node = s_blocks[0].body[0]
pow_node = s_blocks[0].body[1]
pow_node.target = out_var
f_blocks[last_label].body.insert(-3, const_node)
f_blocks[last_label].body.insert(-3, pow_node)
return f_blocks
def run(self):
dprint_func_ir(self.func_ir, "starting hiframes")
topo_order = find_topo_order(self.func_ir.blocks)
for label in topo_order:
new_body = []
for inst in self.func_ir.blocks[label].body:
# df['col'] = arr
if isinstance(
inst,
ir.StaticSetItem) and inst.target.name in self.df_vars:
df_name = inst.target.name
self.df_vars[df_name][inst.index] = inst.value
self._update_df_cols()
elif isinstance(inst, ir.Assign):
out_nodes = self._run_assign(inst)
if isinstance(out_nodes, list):
new_body.extend(out_nodes)
if isinstance(out_nodes, dict):
label = include_new_blocks(self.func_ir.blocks,
out_nodes, label, new_body)
new_body = []
else:
new_body.append(inst)
self.func_ir.blocks[label].body = new_body
self.func_ir._definitions = _get_definitions(self.func_ir.blocks)
#remove_dead(self.func_ir.blocks, self.func_ir.arg_names)
if config._has_h5py:
io_pass = pio.PIO(self.func_ir, self.locals)
io_pass.run()
remove_dead(self.func_ir.blocks, self.func_ir.arg_names)
DummyFlags = namedtuple('DummyFlags', 'auto_parallel')
inline_pass = InlineClosureCallPass(self.func_ir, DummyFlags(True))
inline_pass.run()
self.typemap, self.return_type, self.calltypes = numba_compiler.type_inference_stage(
self.typingctx, self.func_ir, self.args, None)
self.fix_series_filter(self.func_ir.blocks)
self.func_ir._definitions = _get_definitions(self.func_ir.blocks)
dprint_func_ir(self.func_ir, "after hiframes")
if numba.config.DEBUG_ARRAY_OPT == 1:
print("df_vars: ", self.df_vars)
return
def fix_series_filter(self, blocks):
topo_order = find_topo_order(blocks)
for label in topo_order:
new_body = []
for stmt in blocks[label].body:
# find df['col2'] = df['col1'][arr]
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'getitem'
and stmt.value.value.name in self.df_cols
and stmt.target.name in self.df_cols
and self.is_bool_arr(stmt.value.index.name)):
lhs = stmt.target
in_arr = stmt.value.value
index_var = stmt.value.index
def f(A, B, ind):
for i in numba.parfor.prange(len(A)):
s = 0
if ind[i]:
s = B[i]
else:
s = np.nan
A[i] = s
f_blocks = get_inner_ir(f)
replace_var_names(f_blocks, {'A': lhs.name})
replace_var_names(f_blocks, {'B': in_arr.name})
replace_var_names(f_blocks, {'ind': index_var.name})
alloc_nodes = gen_empty_like(in_arr, lhs)
f_blocks[0].body = alloc_nodes + f_blocks[0].body
label = include_new_blocks(blocks, f_blocks, label,
new_body)
new_body = []
else:
new_body.append(stmt)
blocks[label].body = new_body
return
def run(self):
dprint_func_ir(self.func_ir, "starting IO")
topo_order = find_topo_order(self.func_ir.blocks)
for label in topo_order:
new_body = []
# copies are collected before running the pass since
# variables typed in locals are assigned late
self._get_reverse_copies(self.func_ir.blocks[label].body)
for inst in self.func_ir.blocks[label].body:
if isinstance(inst, ir.Assign):
inst_list = self._run_assign(inst)
new_body.extend(inst_list)
elif isinstance(inst, ir.StaticSetItem):
inst_list = self._run_static_setitem(inst)
new_body.extend(inst_list)
else:
new_body.append(inst)
self.func_ir.blocks[label].body = new_body
#remove_dead(self.func_ir.blocks, self.func_ir.arg_names)
dprint_func_ir(self.func_ir, "after IO")
if config.DEBUG_ARRAY_OPT == 1:
print("h5 files: ", self.h5_files)
print("h5 dsets: ", self.h5_dsets)
def get_inner_ir(func):
# get untyped numba ir
f_ir = numba_compiler.run_frontend(func)
blocks = f_ir.blocks
remove_dels(blocks)
topo_order = find_topo_order(blocks)
first_block = blocks[topo_order[0]]
last_block = blocks[topo_order[-1]]
# remove arg nodes
new_first_body = []
for stmt in first_block.body:
if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Arg):
continue
new_first_body.append(stmt)
first_block.body = new_first_body
# rename all variables to avoid conflict, except args
var_table = get_name_var_table(blocks)
new_var_dict = {}
for name, var in var_table.items():
if not (name in f_ir.arg_names):
new_var_dict[name] = mk_unique_var(name)
replace_var_names(blocks, new_var_dict)
return blocks
def _handle_df_col_calls(self, lhs_name, rhs, assign):
if guard(find_callname, self.func_ir,
rhs) == ('count', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_count_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('fillna', 'hpat.hiframes_api'):
out_arr = rhs.args[0]
in_arr = rhs.args[1]
val = rhs.args[2]
f_blocks = compile_to_numba_ir(
_column_fillna_impl, {
'numba': numba,
'np': np
}, self.typingctx,
(self.typemap[out_arr.name], self.typemap[in_arr.name],
self.typemap[val.name]), self.typemap, self.calltypes).blocks
first_block = min(f_blocks.keys())
replace_arg_nodes(f_blocks[first_block], [out_arr, in_arr, val])
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('column_sum', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_sum_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('mean', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_mean_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('var', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_var_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
return
def _fix_nested_array(func_ir):
"""Look for assignment like: a[..] = b, where both a and b are numpy arrays, and
try to eliminate array b by expanding a with an extra dimension.
"""
blocks = func_ir.blocks
cfg = compute_cfg_from_blocks(blocks)
usedefs = compute_use_defs(blocks)
empty_deadmap = dict([(label, set()) for label in blocks.keys()])
livemap = compute_live_variables(cfg, blocks, usedefs.defmap, empty_deadmap)
def find_array_def(arr):
"""Find numpy array definition such as
arr = numba.unsafe.ndarray.empty_inferred(...).
If it is arr = b[...], find array definition of b recursively.
"""
arr_def = func_ir.get_definition(arr)
_make_debug_print("find_array_def")(arr, arr_def)
if isinstance(arr_def, ir.Expr):
if guard(_find_unsafe_empty_inferred, func_ir, arr_def):
return arr_def
elif arr_def.op == 'getitem':
return find_array_def(arr_def.value)
raise GuardException
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 fix_array_assign(stmt):
"""For assignment like lhs[idx] = rhs, where both lhs and rhs are arrays, do the
following:
1. find the definition of rhs, which has to be a call to numba.unsafe.ndarray.empty_inferred
2. find the source array creation for lhs, insert an extra dimension of size of b.
3. replace the definition of rhs = numba.unsafe.ndarray.empty_inferred(...) with rhs = lhs[idx]
"""
require(isinstance(stmt, ir.SetItem))
require(isinstance(stmt.value, ir.Var))
debug_print = _make_debug_print("fix_array_assign")
debug_print("found SetItem: ", stmt)
lhs = stmt.target
# Find the source array creation of lhs
lhs_def = find_array_def(lhs)
debug_print("found lhs_def: ", lhs_def)
rhs_def = get_definition(func_ir, stmt.value)
debug_print("found rhs_def: ", rhs_def)
require(isinstance(rhs_def, ir.Expr))
if rhs_def.op == 'cast':
rhs_def = get_definition(func_ir, rhs_def.value)
require(isinstance(rhs_def, ir.Expr))
require(_find_unsafe_empty_inferred(func_ir, rhs_def))
# Find the array dimension of rhs
dim_def = get_definition(func_ir, rhs_def.args[0])
require(isinstance(dim_def, ir.Expr) and dim_def.op == 'build_tuple')
debug_print("dim_def = ", dim_def)
extra_dims = [ get_definition(func_ir, x, lhs_only=True) for x in dim_def.items ]
debug_print("extra_dims = ", extra_dims)
# Expand size tuple when creating lhs_def with extra_dims
size_tuple_def = get_definition(func_ir, lhs_def.args[0])
require(isinstance(size_tuple_def, ir.Expr) and size_tuple_def.op == 'build_tuple')
debug_print("size_tuple_def = ", size_tuple_def)
extra_dims = fix_dependencies(size_tuple_def, extra_dims)
size_tuple_def.items += extra_dims
# In-place modify rhs_def to be getitem
rhs_def.op = 'getitem'
rhs_def.value = get_definition(func_ir, lhs, lhs_only=True)
rhs_def.index = stmt.index
del rhs_def._kws['func']
del rhs_def._kws['args']
del rhs_def._kws['vararg']
del rhs_def._kws['kws']
# success
return True
for label in find_topo_order(func_ir.blocks):
block = func_ir.blocks[label]
for stmt in block.body:
if guard(fix_array_assign, stmt):
block.body.remove(stmt)
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 inline_closure_call(self, block, i, callee):
"""Inline the body of `callee` at its callsite (`i`-th instruction of `block`)
"""
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)
func_ir = self.func_ir
# first, get the IR of the callee
callee_ir = self.get_ir_of_code(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_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))
_replace_args_with(callee_blocks, args)
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)
assert(isinstance(closure, ir.Expr)
and closure.op == 'build_tuple')
assert(len(callee.code.co_freevars) == len(closure.items))
debug_print("callee's closure = ", closure)
_replace_freevars(callee_blocks, closure.items)
debug_print("After closure rename")
_debug_dump(callee_ir)
# 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)
return new_blocks
def _fix_nested_array(func_ir):
"""Look for assignment like: a[..] = b, where both a and b are numpy arrays, and
try to eliminate array b by expanding a with an extra dimension.
"""
"""
cfg = compute_cfg_from_blocks(func_ir.blocks)
all_loops = list(cfg.loops().values())
def find_nest_level(label):
level = 0
for loop in all_loops:
if label in loop.body:
level += 1
"""
def find_array_def(arr):
"""Find numpy array definition such as
arr = numba.unsafe.ndarray.empty_inferred(...).
If it is arr = b[...], find array definition of b recursively.
"""
arr_def = func_ir.get_definition(arr)
_make_debug_print("find_array_def")(arr, arr_def)
if isinstance(arr_def, ir.Expr):
if guard(_find_unsafe_empty_inferred, func_ir, arr_def):
return arr_def
elif arr_def.op == 'getitem':
return find_array_def(arr_def.value)
raise GuardException
def fix_array_assign(stmt):
"""For assignment like lhs[idx] = rhs, where both lhs and rhs are arrays, do the
following:
1. find the definition of rhs, which has to be a call to numba.unsafe.ndarray.empty_inferred
2. find the source array creation for lhs, insert an extra dimension of size of b.
3. replace the definition of rhs = numba.unsafe.ndarray.empty_inferred(...) with rhs = lhs[idx]
"""
require(isinstance(stmt, ir.SetItem))
require(isinstance(stmt.value, ir.Var))
debug_print = _make_debug_print("fix_array_assign")
debug_print("found SetItem: ", stmt)
lhs = stmt.target
# Find the source array creation of lhs
lhs_def = find_array_def(lhs)
debug_print("found lhs_def: ", lhs_def)
rhs_def = get_definition(func_ir, stmt.value)
debug_print("found rhs_def: ", rhs_def)
require(isinstance(rhs_def, ir.Expr))
if rhs_def.op == 'cast':
rhs_def = get_definition(func_ir, rhs_def.value)
require(isinstance(rhs_def, ir.Expr))
require(_find_unsafe_empty_inferred(func_ir, rhs_def))
# Find the array dimension of rhs
dim_def = get_definition(func_ir, rhs_def.args[0])
require(isinstance(dim_def, ir.Expr) and dim_def.op == 'build_tuple')
debug_print("dim_def = ", dim_def)
extra_dims = [ get_definition(func_ir, x, lhs_only=True) for x in dim_def.items ]
debug_print("extra_dims = ", extra_dims)
# Expand size tuple when creating lhs_def with extra_dims
size_tuple_def = get_definition(func_ir, lhs_def.args[0])
require(isinstance(size_tuple_def, ir.Expr) and size_tuple_def.op == 'build_tuple')
debug_print("size_tuple_def = ", size_tuple_def)
size_tuple_def.items += extra_dims
# In-place modify rhs_def to be getitem
rhs_def.op = 'getitem'
rhs_def.value = get_definition(func_ir, lhs, lhs_only=True)
rhs_def.index = stmt.index
del rhs_def._kws['func']
del rhs_def._kws['args']
del rhs_def._kws['vararg']
del rhs_def._kws['kws']
# success
return True
for label in find_topo_order(func_ir.blocks):
block = func_ir.blocks[label]
for stmt in block.body:
if guard(fix_array_assign, stmt):
block.body.remove(stmt)
def run(self):
blocks = self.func_ir.blocks
topo_order = find_topo_order(blocks)
for label in topo_order:
new_body = []
for inst in blocks[label].body:
if isinstance(inst, ir.Assign):
out_nodes = self._run_assign(inst)
if isinstance(out_nodes, list):
new_body.extend(out_nodes)
if isinstance(out_nodes, dict):
label = include_new_blocks(blocks, out_nodes, label,
new_body)
new_body = []
if isinstance(out_nodes, tuple):
gen_blocks, post_nodes = out_nodes
label = include_new_blocks(blocks, gen_blocks, label,
new_body)
new_body = post_nodes
else:
new_body.append(inst)
blocks[label].body = new_body
if debug_prints(): # pragma: no cover
print("--- types before Series replacement:", self.typemap)
print("calltypes: ", self.calltypes)
replace_series = {}
for vname, typ in self.typemap.items():
if isinstance(typ, SeriesType):
# print("replacing series type", vname)
new_typ = series_to_array_type(typ)
replace_series[vname] = new_typ
# replace array.call() variable types
if isinstance(typ, types.BoundFunction) and isinstance(
typ.this, SeriesType):
this = series_to_array_type(typ.this)
# TODO: handle string arrays, etc.
assert typ.typing_key.startswith('array.')
attr = typ.typing_key[len('array.'):]
resolver = getattr(ArrayAttribute, 'resolve_' + attr)
# methods are either installed with install_array_method or
# using @bound_function in arraydecl.py
if hasattr(resolver, '__wrapped__'):
resolver = bound_function(typ.typing_key)(
resolver.__wrapped__)
new_typ = resolver(ArrayAttribute(self.typingctx), this)
replace_series[vname] = new_typ
for vname, typ in replace_series.items():
self.typemap.pop(vname)
self.typemap[vname] = typ
replace_calltype = {}
# replace sig of getitem/setitem/... series type with array
for call, sig in self.calltypes.items():
if sig is None:
continue
assert isinstance(sig, Signature)
sig.return_type = if_series_to_array_type(sig.return_type)
sig.args = tuple(map(if_series_to_array_type, sig.args))
# XXX: side effect: force update of call signatures
if isinstance(call, ir.Expr) and call.op == 'call':
# StencilFunc requires kws for typing so sig.args can't be used
# reusing sig.args since some types become Const in sig
argtyps = sig.args[:len(call.args)]
kwtyps = {name: self.typemap[v.name] for name, v in call.kws}
new_sig = self.typemap[call.func.name].get_call_type(
self.typingctx, argtyps, kwtyps)
# calltypes of things like BoundFunction (array.call) need to
# be update for lowering to work
# XXX: new_sig could be None for things like np.int32()
if call in self.calltypes and new_sig is not None:
old_sig = self.calltypes[call]
# fix types with undefined dtypes in empty_inferred, etc.
return_type = _fix_typ_undefs(new_sig.return_type,
old_sig.return_type)
args = tuple(
_fix_typ_undefs(a, b)
for a, b in zip(new_sig.args, old_sig.args))
replace_calltype[call] = Signature(return_type, args,
new_sig.recvr,
new_sig.pysig)
for call, sig in replace_calltype.items():
self.calltypes.pop(call)
self.calltypes[call] = sig
if debug_prints(): # pragma: no cover
print("--- types after Series replacement:", self.typemap)
print("calltypes: ", self.calltypes)
self.func_ir._definitions = get_definitions(self.func_ir.blocks)
return if_series_to_unbox(self.return_type)
def _handle_df_col_calls(self, assign, lhs, rhs, func_name):
if func_name == 'count':
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(
_column_count_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx,
(if_series_to_array_type(self.typemap[in_arr.name]), ),
self.typemap, self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-3].target = assign.target
return f_blocks
if func_name == 'fillna':
out_arr = rhs.args[0]
in_arr = rhs.args[1]
val = rhs.args[2]
f_blocks = compile_to_numba_ir(
_column_fillna_impl, {
'numba': numba,
'np': np
}, self.typingctx,
(if_series_to_array_type(self.typemap[out_arr.name]),
if_series_to_array_type(self.typemap[in_arr.name]),
if_series_to_array_type(self.typemap[val.name])),
self.typemap, self.calltypes).blocks
first_block = min(f_blocks.keys())
replace_arg_nodes(f_blocks[first_block], [out_arr, in_arr, val])
return f_blocks
if func_name == 'column_sum':
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(
_column_sum_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx,
(if_series_to_array_type(self.typemap[in_arr.name]), ),
self.typemap, self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-3].target = assign.target
return f_blocks
if func_name == 'mean':
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(
_column_mean_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx,
(if_series_to_array_type(self.typemap[in_arr.name]), ),
self.typemap, self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-3].target = assign.target
return f_blocks
if func_name == 'var':
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(
_column_var_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx,
(if_series_to_array_type(self.typemap[in_arr.name]), ),
self.typemap, self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-3].target = assign.target
return f_blocks
return [assign]
def fixes_after_typing(self, blocks):
call_table, _ = ir_utils.get_call_table(self.func_ir.blocks)
topo_order = find_topo_order(blocks)
for label in topo_order:
new_body = []
for stmt in blocks[label].body:
# convert str_arr==str into parfor
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'binop'
and stmt.value.fn in ['==', '!=']
and (self.typemap[stmt.value.lhs.name] == string_array_type
or self.typemap[stmt.value.rhs.name] == string_array_type)):
lhs = stmt.value.lhs
rhs = stmt.value.rhs
lhs_access = 'A'
rhs_access = 'B'
len_call = 'A.size'
if self.typemap[lhs.name] == string_array_type:
lhs_access = 'A[i]'
if self.typemap[rhs.name] == string_array_type:
lhs_access = 'B[i]'
len_call = 'B.size'
func_text = 'def f(A, B):\n'
func_text += ' l = {}\n'.format(len_call)
func_text += ' S = np.empty(l, dtype=np.bool_)\n'
func_text += ' for i in numba.parfor.prange(l):\n'
func_text += ' S[i] = {} {} {}\n'.format(lhs_access, stmt.value.fn, rhs_access)
loc_vars = {}
exec(func_text, {}, loc_vars)
f = loc_vars['f']
f_blocks = compile_to_numba_ir(f, {'numba': numba, 'np': np}).blocks
replace_arg_nodes(f_blocks[min(f_blocks.keys())], [lhs, rhs])
label = include_new_blocks(blocks, f_blocks, label, new_body)
new_body = []
# replace == expression with result of parfor (S)
# S is target of last statement in 1st block of f
stmt.value = f_blocks[min(f_blocks.keys())].body[-2].target
# arr = fix_df_array(col) -> arr=col if col is array
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'
and stmt.value.func.name in call_table
and call_table[stmt.value.func.name] ==
['fix_df_array', 'hiframes_api', hpat]
and isinstance(self.typemap[stmt.value.args[0].name],
(types.Array, StringArrayType))):
stmt.value = stmt.value.args[0]
# find df['col2'] = df['col1'][arr]
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op=='getitem'
and stmt.value.value.name in self.df_cols
and stmt.target.name in self.df_cols
and self.is_bool_arr(stmt.value.index.name)):
lhs = stmt.target
in_arr = stmt.value.value
index_var = stmt.value.index
def f(A, B, ind):
for i in numba.parfor.prange(len(A)):
s = 0
if ind[i]:
s = B[i]
else:
s= np.nan
A[i] = s
f_blocks = get_inner_ir(f)
replace_var_names(f_blocks, {'A': lhs.name})
replace_var_names(f_blocks, {'B': in_arr.name})
replace_var_names(f_blocks, {'ind': index_var.name})
alloc_nodes = gen_empty_like(in_arr, lhs)
f_blocks[0].body = alloc_nodes + f_blocks[0].body
label = include_new_blocks(blocks, f_blocks, label, new_body)
new_body = []
else:
new_body.append(stmt)
blocks[label].body = new_body
return
