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.core.inline_closurecall._add_definitions(func_ir, block)
work_list.append((_label, block))
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.core.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.core.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):
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 run(self):
typingctx = self.state.typingctx
# save array arg to call
# call_varname -> array
func_ir = self.state.func_ir
blocks = func_ir.blocks
saved_arr_arg = {}
topo_order = find_topo_order(blocks)
replaced = False
for label in topo_order:
block = blocks[label]
new_body = []
for stmt in block.body:
if isinstance(stmt, ir.Assign) and isinstance(
stmt.value, ir.Expr):
lhs = stmt.target.name
rhs = stmt.value
# replace A.func with np.func, and save A in saved_arr_arg
if (rhs.op == "getattr"
and rhs.attr in self.function_name_map
and isinstance(self.typemap[rhs.value.name],
types.npytypes.Array)):
rhs = stmt.value
arr = rhs.value
saved_arr_arg[lhs] = arr
scope = arr.scope
loc = arr.loc
g_dppy_var = ir.Var(scope,
mk_unique_var("$load_global"), loc)
self.typemap[g_dppy_var.name] = types.misc.Module(
numba_dppy)
g_dppy = ir.Global("numba_dppy", numba_dppy, loc)
g_dppy_assign = ir.Assign(g_dppy, g_dppy_var, loc)
dpnp_var = ir.Var(scope, mk_unique_var("$load_attr"),
loc)
self.typemap[dpnp_var.name] = types.misc.Module(
numba_dppy.dpnp)
getattr_dpnp = ir.Expr.getattr(g_dppy_var, "dpnp", loc)
dpnp_assign = ir.Assign(getattr_dpnp, dpnp_var, loc)
rhs.value = dpnp_var
new_body.append(g_dppy_assign)
new_body.append(dpnp_assign)
func_ir._definitions[g_dppy_var.name] = [getattr_dpnp]
func_ir._definitions[dpnp_var.name] = [getattr_dpnp]
# update func var type
func = getattr(numba_dppy.dpnp, rhs.attr)
func_typ = get_dpnp_func_typ(func)
self.typemap.pop(lhs)
self.typemap[lhs] = func_typ
replaced = True
if rhs.op == "call" and rhs.func.name in saved_arr_arg:
# add array as first arg
arr = saved_arr_arg[rhs.func.name]
# update call type signature to include array arg
old_sig = self.calltypes.pop(rhs)
# argsort requires kws for typing so sig.args can't be used
# reusing sig.args since some types become Const in sig
argtyps = old_sig.args[:len(rhs.args)]
kwtyps = {
name: self.typemap[v.name]
for name, v in rhs.kws
}
self.calltypes[rhs] = self.typemap[
rhs.func.name].get_call_type(
typingctx,
[self.typemap[arr.name]] + list(argtyps),
kwtyps,
)
rhs.args = [arr] + rhs.args
new_body.append(stmt)
block.body = new_body
return replaced
def run(self):
"""
This function rewrites the name of NumPy functions that exist in self.function_name_map
e.g np.sum(a) would produce the following:
np.sum() --> numba_dppy.dpnp.sum()
---------------------------------------------------------------------------------------
Numba IR Before Rewrite:
---------------------------------------------------------------------------------------
$2load_global.0 = global(np: <module 'numpy' from 'numpy/__init__.py'>) ['$2load_global.0']
$4load_method.1 = getattr(value=$2load_global.0, attr=sum) ['$2load_global.0', '$4load_method.1']
$8call_method.3 = call $4load_method.1(a, func=$4load_method.1, args=[Var(a, test_rewrite.py:7)],
kws=(), vararg=None) ['$4load_method.1', '$8call_method.3', 'a']
---------------------------------------------------------------------------------------
Numba IR After Rewrite:
---------------------------------------------------------------------------------------
$dppy_replaced_var.0 = global(numba_dppy: <module 'numba_dppy' from 'numba_dppy/__init__.py'>) ['$dppy_replaced_var.0']
$dpnp_var.1 = getattr(value=$dppy_replaced_var.0, attr=dpnp) ['$dpnp_var.1', '$dppy_replaced_var.0']
$4load_method.1 = getattr(value=$dpnp_var.1, attr=sum) ['$4load_method.1', '$dpnp_var.1']
$8call_method.3 = call $4load_method.1(a, func=$4load_method.1, args=[Var(a, test_rewrite.py:7)],
kws=(), vararg=None) ['$4load_method.1', '$8call_method.3', 'a']
---------------------------------------------------------------------------------------
"""
func_ir = self.state.func_ir
blocks = func_ir.blocks
topo_order = find_topo_order(blocks)
replaced = False
for label in topo_order:
block = blocks[label]
saved_arr_arg = {}
new_body = []
for stmt in block.body:
if isinstance(stmt, ir.Assign) and isinstance(
stmt.value, ir.Expr):
lhs = stmt.target.name
rhs = stmt.value
# replace np.FOO with name from self.function_name_map["FOO"]
# e.g. np.sum will be replaced with numba_dppy.dpnp.sum
if (rhs.op == "getattr"
and rhs.attr in self.function_name_map):
module_node = block.find_variable_assignment(
rhs.value.name).value
if (isinstance(module_node, ir.Global)
and module_node.value.__name__
in self.function_name_map[rhs.attr][0]) or (
isinstance(module_node, ir.Expr)
and module_node.attr
in self.function_name_map[rhs.attr][0]):
rhs = stmt.value
rhs.attr = self.function_name_map[rhs.attr][1]
global_module = rhs.value
saved_arr_arg[lhs] = global_module
scope = global_module.scope
loc = global_module.loc
g_dppy_var = ir.Var(scope,
mk_unique_var("$2load_global"),
loc)
# We are trying to rename np.function_name/np.linalg.function_name with
# numba_dppy.dpnp.function_name.
# Hence, we need to have a global variable representing module numba_dppy.
# Next, we add attribute dpnp to global module numba_dppy to
# represent numba_dppy.dpnp.
g_dppy = ir.Global("numba_dppy", numba_dppy, loc)
g_dppy_assign = ir.Assign(g_dppy, g_dppy_var, loc)
dpnp_var = ir.Var(scope,
mk_unique_var("$4load_attr"),
loc)
getattr_dpnp = ir.Expr.getattr(
g_dppy_var, "dpnp", loc)
dpnp_assign = ir.Assign(getattr_dpnp, dpnp_var,
loc)
rhs.value = dpnp_var
new_body.append(g_dppy_assign)
new_body.append(dpnp_assign)
func_ir._definitions[dpnp_var.name] = [
getattr_dpnp
]
func_ir._definitions[g_dppy_var.name] = [g_dppy]
replaced = True
new_body.append(stmt)
block.body = new_body
return replaced
