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 hoist(parfor_params, loop_body, typemap, wrapped_blocks):
dep_on_param = copy.copy(parfor_params)
hoisted = []
def_once = compute_def_once(loop_body)
(call_table, reverse_call_table) = get_call_table(wrapped_blocks)
for label, block in loop_body.items():
new_block = []
for inst in block.body:
if isinstance(inst, ir.Assign) and inst.target.name in def_once:
if _hoist_internal(inst, dep_on_param, call_table, hoisted,
typemap):
# don't add this instuction to the block since it is hoisted
continue
elif isinstance(inst, parfor.Parfor):
new_init_block = []
if config.DEBUG_ARRAY_OPT == 1:
print("parfor")
inst.dump()
for ib_inst in inst.init_block.body:
if (isinstance(ib_inst, ir.Assign)
and ib_inst.target.name in def_once):
if _hoist_internal(ib_inst, dep_on_param, call_table,
hoisted, typemap):
# don't add this instuction to the block since it is hoisted
continue
new_init_block.append(ib_inst)
inst.init_block.body = new_init_block
new_block.append(inst)
block.body = new_block
return hoisted
def run(self):
""" Finds all calls to StencilFuncs in the IR and converts them to parfor.
"""
from numba.stencil import StencilFunc
# Get all the calls in the function IR.
call_table, _ = get_call_table(self.func_ir.blocks)
stencil_calls = []
stencil_dict = {}
for call_varname, call_list in call_table.items():
if isinstance(call_list[0], StencilFunc):
# Remember all calls to StencilFuncs.
stencil_calls.append(call_varname)
stencil_dict[call_varname] = call_list[0]
if not stencil_calls:
return # return early if no stencil calls found
# find and transform stencil calls
for label, block in self.func_ir.blocks.items():
for i, stmt in reversed(list(enumerate(block.body))):
# Found a call to a StencilFunc.
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'
and stmt.value.func.name in stencil_calls):
kws = dict(stmt.value.kws)
# Create dictionary of input argument number to
# the argument itself.
input_dict = {
i: stmt.value.args[i]
for i in range(len(stmt.value.args))
}
in_args = stmt.value.args
arg_typemap = tuple(self.typemap[i.name] for i in in_args)
for arg_type in arg_typemap:
if isinstance(arg_type, types.BaseTuple):
raise ValueError("Tuple parameters not supported " \
"for stencil kernels in parallel=True mode.")
out_arr = kws.get('out')
# Get the StencilFunc object corresponding to this call.
sf = stencil_dict[stmt.value.func.name]
stencil_blocks, rt, arg_to_arr_dict = get_stencil_blocks(
sf, self.typingctx, arg_typemap, block.scope,
block.loc, input_dict, self.typemap, self.calltypes)
index_offsets = sf.options.get('index_offsets', None)
gen_nodes = self._mk_stencil_parfor(
label, in_args, out_arr, stencil_blocks, index_offsets,
stmt.target, rt, sf, arg_to_arr_dict)
block.body = block.body[:i] + gen_nodes + block.body[i +
1:]
# Found a call to a stencil via numba.stencil().
elif (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'
and guard(find_callname, self.func_ir,
stmt.value) == ('stencil', 'numba')):
# remove dummy stencil() call
stmt.value = ir.Const(0, stmt.loc)
def hoist(parfor_params, loop_body, typemap, wrapped_blocks):
dep_on_param = copy.copy(parfor_params)
hoisted = []
def_once = compute_def_once(loop_body)
(call_table, reverse_call_table) = get_call_table(wrapped_blocks)
for label, block in loop_body.items():
new_block = []
for inst in block.body:
if isinstance(inst, ir.Assign) and inst.target.name in def_once:
if _hoist_internal(inst, dep_on_param, call_table,
hoisted, typemap):
# don't add this instuction to the block since it is hoisted
continue
elif isinstance(inst, parfor.Parfor):
new_init_block = []
if config.DEBUG_ARRAY_OPT == 1:
print("parfor")
inst.dump()
for ib_inst in inst.init_block.body:
if (isinstance(ib_inst, ir.Assign) and
ib_inst.target.name in def_once):
if _hoist_internal(ib_inst, dep_on_param, call_table,
hoisted, typemap):
# don't add this instuction to the block since it is hoisted
continue
new_init_block.append(ib_inst)
inst.init_block.body = new_init_block
new_block.append(inst)
block.body = new_block
return hoisted
def __init__(self, func_ir, typemap, calltypes):
self.func_ir = func_ir
self.typemap = typemap
self.calltypes = calltypes
self._call_table, _ = get_call_table(func_ir.blocks)
self._tuple_table = get_tuple_table(func_ir.blocks)
self._parallel_accesses = set()
self._T_arrs = set()
def run(self):
""" Finds all calls to StencilFuncs in the IR and converts them to parfor.
"""
from numba.stencil import StencilFunc
# Get all the calls in the function IR.
call_table, _ = get_call_table(self.func_ir.blocks)
stencil_calls = []
stencil_dict = {}
for call_varname, call_list in call_table.items():
if isinstance(call_list[0], StencilFunc):
# Remember all calls to StencilFuncs.
stencil_calls.append(call_varname)
stencil_dict[call_varname] = call_list[0]
if not stencil_calls:
return # return early if no stencil calls found
# find and transform stencil calls
for label, block in self.func_ir.blocks.items():
for i, stmt in reversed(list(enumerate(block.body))):
# Found a call to a StencilFunc.
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'
and stmt.value.func.name in stencil_calls):
kws = dict(stmt.value.kws)
# Create dictionary of input argument number to
# the argument itself.
input_dict = {i: stmt.value.args[i] for i in
range(len(stmt.value.args))}
in_args = stmt.value.args
arg_typemap = tuple(self.typemap[i.name] for i in in_args)
for arg_type in arg_typemap:
if isinstance(arg_type, types.BaseTuple):
raise ValueError("Tuple parameters not supported " \
"for stencil kernels in parallel=True mode.")
out_arr = kws.get('out')
# Get the StencilFunc object corresponding to this call.
sf = stencil_dict[stmt.value.func.name]
stencil_ir, rt, arg_to_arr_dict = get_stencil_ir(sf,
self.typingctx, arg_typemap,
block.scope, block.loc, input_dict,
self.typemap, self.calltypes)
index_offsets = sf.options.get('index_offsets', None)
gen_nodes = self._mk_stencil_parfor(label, in_args, out_arr,
stencil_ir, index_offsets, stmt.target, rt, sf,
arg_to_arr_dict)
block.body = block.body[:i] + gen_nodes + block.body[i+1:]
# Found a call to a stencil via numba.stencil().
elif (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'
and guard(find_callname, self.func_ir, stmt.value)
== ('stencil', 'numba')):
# remove dummy stencil() call
stmt.value = ir.Const(0, stmt.loc)
def inline_calls(func_ir):
call_table, _ = ir_utils.get_call_table(func_ir.blocks)
for label, block in func_ir.blocks.items():
for i, stmt in enumerate(block.body):
if isinstance(stmt, ir.Assign):
rhs = stmt.value
if isinstance(rhs, ir.Expr) and rhs.op == 'call':
func = rhs.func.name
if (func in call_table and call_table[func] and isinstance(
call_table[func][0], CPUDispatcher)):
py_func = call_table[func][0].py_func
inline_calls_inner(func_ir, block, stmt, i, py_func)
return # inline_calls_inner will call back recursively
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
