ai-content-maker/.venv/Lib/site-packages/numba/stencils/stencilparfor.py

950 lines
43 KiB
Python
Raw Normal View History

2024-05-03 04:18:51 +03:00
#
# Copyright (c) 2017 Intel Corporation
# SPDX-License-Identifier: BSD-2-Clause
#
import numbers
import copy
import types as pytypes
from operator import add
import operator
import numpy as np
import numba.parfors.parfor
from numba.core import types, ir, rewrites, config, ir_utils
from numba.core.typing.templates import infer_global, AbstractTemplate
from numba.core.typing import signature
from numba.core import utils, typing
from numba.core.ir_utils import (get_call_table, mk_unique_var,
compile_to_numba_ir, replace_arg_nodes, guard,
find_callname, require, find_const, GuardException)
from numba.core.errors import NumbaValueError
from numba.core.utils import OPERATORS_TO_BUILTINS
from numba.np import numpy_support
def _compute_last_ind(dim_size, index_const):
if index_const > 0:
return dim_size - index_const
else:
return dim_size
class StencilPass(object):
def __init__(self, func_ir, typemap, calltypes, array_analysis, typingctx,
targetctx, flags):
self.func_ir = func_ir
self.typemap = typemap
self.calltypes = calltypes
self.array_analysis = array_analysis
self.typingctx = typingctx
self.targetctx = targetctx
self.flags = flags
def run(self):
""" Finds all calls to StencilFuncs in the IR and converts them to parfor.
"""
from numba.stencils.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():
for one_call in call_list:
if isinstance(one_call, StencilFunc):
# Remember all calls to StencilFuncs.
stencil_calls.append(call_varname)
stencil_dict[call_varname] = one_call
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 replace_return_with_setitem(self, blocks, exit_value_var,
parfor_body_exit_label):
"""
Find return statements in the IR and replace them with a SetItem
call of the value "returned" by the kernel into the result array.
Returns the block labels that contained return statements.
"""
for label, block in blocks.items():
scope = block.scope
loc = block.loc
new_body = []
for stmt in block.body:
if isinstance(stmt, ir.Return):
# previous stmt should have been a cast
prev_stmt = new_body.pop()
assert (isinstance(prev_stmt, ir.Assign)
and isinstance(prev_stmt.value, ir.Expr)
and prev_stmt.value.op == 'cast')
new_body.append(ir.Assign(prev_stmt.value.value, exit_value_var, loc))
new_body.append(ir.Jump(parfor_body_exit_label, loc))
else:
new_body.append(stmt)
block.body = new_body
def _mk_stencil_parfor(self, label, in_args, out_arr, stencil_ir,
index_offsets, target, return_type, stencil_func,
arg_to_arr_dict):
""" Converts a set of stencil kernel blocks to a parfor.
"""
gen_nodes = []
stencil_blocks = stencil_ir.blocks
if config.DEBUG_ARRAY_OPT >= 1:
print("_mk_stencil_parfor", label, in_args, out_arr, index_offsets,
return_type, stencil_func, stencil_blocks)
ir_utils.dump_blocks(stencil_blocks)
in_arr = in_args[0]
# run copy propagate to replace in_args copies (e.g. a = A)
in_arr_typ = self.typemap[in_arr.name]
in_cps, out_cps = ir_utils.copy_propagate(stencil_blocks, self.typemap)
name_var_table = ir_utils.get_name_var_table(stencil_blocks)
ir_utils.apply_copy_propagate(
stencil_blocks,
in_cps,
name_var_table,
self.typemap,
self.calltypes)
if config.DEBUG_ARRAY_OPT >= 1:
print("stencil_blocks after copy_propagate")
ir_utils.dump_blocks(stencil_blocks)
ir_utils.remove_dead(stencil_blocks, self.func_ir.arg_names, stencil_ir,
self.typemap)
if config.DEBUG_ARRAY_OPT >= 1:
print("stencil_blocks after removing dead code")
ir_utils.dump_blocks(stencil_blocks)
# create parfor vars
ndims = self.typemap[in_arr.name].ndim
scope = in_arr.scope
loc = in_arr.loc
parfor_vars = []
for i in range(ndims):
parfor_var = ir.Var(scope, mk_unique_var(
"$parfor_index_var"), loc)
self.typemap[parfor_var.name] = types.intp
parfor_vars.append(parfor_var)
start_lengths, end_lengths = self._replace_stencil_accesses(
stencil_ir, parfor_vars, in_args, index_offsets, stencil_func,
arg_to_arr_dict)
if config.DEBUG_ARRAY_OPT >= 1:
print("stencil_blocks after replace stencil accesses")
print("start_lengths:", start_lengths)
print("end_lengths:", end_lengths)
ir_utils.dump_blocks(stencil_blocks)
# create parfor loop nests
loopnests = []
equiv_set = self.array_analysis.get_equiv_set(label)
in_arr_dim_sizes = equiv_set.get_shape(in_arr)
assert ndims == len(in_arr_dim_sizes)
start_inds = []
last_inds = []
for i in range(ndims):
last_ind = self._get_stencil_last_ind(in_arr_dim_sizes[i],
end_lengths[i], gen_nodes, scope, loc)
start_ind = self._get_stencil_start_ind(
start_lengths[i], gen_nodes, scope, loc)
start_inds.append(start_ind)
last_inds.append(last_ind)
# start from stencil size to avoid invalid array access
loopnests.append(numba.parfors.parfor.LoopNest(parfor_vars[i],
start_ind, last_ind, 1))
# We have to guarantee that the exit block has maximum label and that
# there's only one exit block for the parfor body.
# So, all return statements will change to jump to the parfor exit block.
parfor_body_exit_label = max(stencil_blocks.keys()) + 1
stencil_blocks[parfor_body_exit_label] = ir.Block(scope, loc)
exit_value_var = ir.Var(scope, mk_unique_var("$parfor_exit_value"), loc)
self.typemap[exit_value_var.name] = return_type.dtype
# create parfor index var
for_replacing_ret = []
if ndims == 1:
parfor_ind_var = parfor_vars[0]
else:
parfor_ind_var = ir.Var(scope, mk_unique_var(
"$parfor_index_tuple_var"), loc)
self.typemap[parfor_ind_var.name] = types.containers.UniTuple(
types.intp, ndims)
tuple_call = ir.Expr.build_tuple(parfor_vars, loc)
tuple_assign = ir.Assign(tuple_call, parfor_ind_var, loc)
for_replacing_ret.append(tuple_assign)
if config.DEBUG_ARRAY_OPT >= 1:
print("stencil_blocks after creating parfor index var")
ir_utils.dump_blocks(stencil_blocks)
# empty init block
init_block = ir.Block(scope, loc)
if out_arr is None:
in_arr_typ = self.typemap[in_arr.name]
shape_name = ir_utils.mk_unique_var("in_arr_shape")
shape_var = ir.Var(scope, shape_name, loc)
shape_getattr = ir.Expr.getattr(in_arr, "shape", loc)
self.typemap[shape_name] = types.containers.UniTuple(types.intp,
in_arr_typ.ndim)
init_block.body.extend([ir.Assign(shape_getattr, shape_var, loc)])
zero_name = ir_utils.mk_unique_var("zero_val")
zero_var = ir.Var(scope, zero_name, loc)
if "cval" in stencil_func.options:
cval = stencil_func.options["cval"]
# TODO: Loosen this restriction to adhere to casting rules.
cval_ty = typing.typeof.typeof(cval)
if not self.typingctx.can_convert(cval_ty, return_type.dtype):
raise ValueError("cval type does not match stencil return type.")
temp2 = return_type.dtype(cval)
else:
temp2 = return_type.dtype(0)
full_const = ir.Const(temp2, loc)
self.typemap[zero_name] = return_type.dtype
init_block.body.extend([ir.Assign(full_const, zero_var, loc)])
so_name = ir_utils.mk_unique_var("stencil_output")
out_arr = ir.Var(scope, so_name, loc)
self.typemap[out_arr.name] = numba.core.types.npytypes.Array(
return_type.dtype,
in_arr_typ.ndim,
in_arr_typ.layout)
dtype_g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
self.typemap[dtype_g_np_var.name] = types.misc.Module(np)
dtype_g_np = ir.Global('np', np, loc)
dtype_g_np_assign = ir.Assign(dtype_g_np, dtype_g_np_var, loc)
init_block.body.append(dtype_g_np_assign)
return_type_name = numpy_support.as_dtype(
return_type.dtype).type.__name__
if return_type_name == 'bool':
return_type_name = 'bool_'
dtype_np_attr_call = ir.Expr.getattr(dtype_g_np_var, return_type_name, loc)
dtype_attr_var = ir.Var(scope, mk_unique_var("$np_attr_attr"), loc)
self.typemap[dtype_attr_var.name] = types.functions.NumberClass(return_type.dtype)
dtype_attr_assign = ir.Assign(dtype_np_attr_call, dtype_attr_var, loc)
init_block.body.append(dtype_attr_assign)
stmts = ir_utils.gen_np_call("empty",
np.empty,
out_arr,
[shape_var, dtype_attr_var],
self.typingctx,
self.typemap,
self.calltypes)
# ------------------
# Generate the code to fill just the border with zero_var.
# Generate a none var to use in slicing.
none_var = ir.Var(scope, mk_unique_var("$none_var"), loc)
none_assign = ir.Assign(ir.Const(None, loc), none_var, loc)
stmts.append(none_assign)
self.typemap[none_var.name] = types.none
# Generate a zero var to use in slicing.
zero_index_var = ir.Var(scope, mk_unique_var("$zero_index_var"), loc)
zero_index_assign = ir.Assign(ir.Const(0, loc), zero_index_var, loc)
stmts.append(zero_index_assign)
self.typemap[zero_index_var.name] = types.intp
# Generate generic ":" slice.
# ---- Generate var to hold slice func var.
slice_func_var = ir.Var(scope, mk_unique_var("$slice_func_var"), loc)
slice_fn_ty = self.typingctx.resolve_value_type(slice)
self.typemap[slice_func_var.name] = slice_fn_ty
slice_g = ir.Global('slice', slice, loc)
slice_assign = ir.Assign(slice_g, slice_func_var, loc)
stmts.append(slice_assign)
# ---- Generate call to slice func.
sig = self.typingctx.resolve_function_type(slice_fn_ty,
(types.none,) * 2,
{})
slice_callexpr = ir.Expr.call(func=slice_func_var,
args=(none_var, none_var),
kws=(),
loc=loc)
self.calltypes[slice_callexpr] = sig
# ---- Generate slice var
slice_var = ir.Var(scope, mk_unique_var("$slice"), loc)
self.typemap[slice_var.name] = types.slice2_type
slice_assign = ir.Assign(slice_callexpr, slice_var, loc)
stmts.append(slice_assign)
def handle_border(slice_fn_ty,
dim,
scope,
loc,
slice_func_var,
stmts,
border_inds,
border_tuple_items,
other_arg,
other_first):
# Handle the border for start or end of the index range.
# ---- Generate call to slice func.
sig = self.typingctx.resolve_function_type(
slice_fn_ty,
(types.intp,) * 2,
{})
si = border_inds[dim]
assert(isinstance(si, (int, ir.Var)))
si_var = ir.Var(scope, mk_unique_var("$border_ind"), loc)
self.typemap[si_var.name] = types.intp
if isinstance(si, int):
si_assign = ir.Assign(ir.Const(si, loc), si_var, loc)
else:
si_assign = ir.Assign(si, si_var, loc)
stmts.append(si_assign)
slice_callexpr = ir.Expr.call(
func=slice_func_var,
args=(other_arg, si_var) if other_first else (si_var, other_arg),
kws=(),
loc=loc)
self.calltypes[slice_callexpr] = sig
# ---- Generate slice var
border_slice_var = ir.Var(scope, mk_unique_var("$slice"), loc)
self.typemap[border_slice_var.name] = types.slice2_type
slice_assign = ir.Assign(slice_callexpr, border_slice_var, loc)
stmts.append(slice_assign)
border_tuple_items[dim] = border_slice_var
border_ind_var = ir.Var(scope, mk_unique_var(
"$border_index_tuple_var"), loc)
self.typemap[border_ind_var.name] = types.containers.UniTuple(
types.slice2_type, ndims)
tuple_call = ir.Expr.build_tuple(border_tuple_items, loc)
tuple_assign = ir.Assign(tuple_call, border_ind_var, loc)
stmts.append(tuple_assign)
setitem_call = ir.SetItem(out_arr, border_ind_var, zero_var, loc)
self.calltypes[setitem_call] = signature(
types.none, self.typemap[out_arr.name],
self.typemap[border_ind_var.name],
self.typemap[out_arr.name].dtype
)
stmts.append(setitem_call)
# For each dimension, add setitem to set border values.
for dim in range(in_arr_typ.ndim):
# First, fill all entries with ":".
start_tuple_items = [slice_var] * in_arr_typ.ndim
last_tuple_items = [slice_var] * in_arr_typ.ndim
handle_border(slice_fn_ty,
dim,
scope,
loc,
slice_func_var,
stmts,
start_inds,
start_tuple_items,
zero_index_var,
True)
handle_border(slice_fn_ty,
dim,
scope,
loc,
slice_func_var,
stmts,
last_inds,
last_tuple_items,
in_arr_dim_sizes[dim],
False)
# ------------------
equiv_set.insert_equiv(out_arr, in_arr_dim_sizes)
init_block.body.extend(stmts)
else: # out is present
if "cval" in stencil_func.options: # do out[:] = cval
cval = stencil_func.options["cval"]
# TODO: Loosen this restriction to adhere to casting rules.
cval_ty = typing.typeof.typeof(cval)
if not self.typingctx.can_convert(cval_ty, return_type.dtype):
msg = "cval type does not match stencil return type."
raise NumbaValueError(msg)
# get slice ref
slice_var = ir.Var(scope, mk_unique_var("$py_g_var"), loc)
slice_fn_ty = self.typingctx.resolve_value_type(slice)
self.typemap[slice_var.name] = slice_fn_ty
slice_g = ir.Global('slice', slice, loc)
slice_assigned = ir.Assign(slice_g, slice_var, loc)
init_block.body.append(slice_assigned)
sig = self.typingctx.resolve_function_type(slice_fn_ty,
(types.none,) * 2,
{})
callexpr = ir.Expr.call(func=slice_var, args=(), kws=(),
loc=loc)
self.calltypes[callexpr] = sig
slice_inst_var = ir.Var(scope, mk_unique_var("$slice_inst"),
loc)
self.typemap[slice_inst_var.name] = types.slice2_type
slice_assign = ir.Assign(callexpr, slice_inst_var, loc)
init_block.body.append(slice_assign)
# get const val for cval
cval_const_val = ir.Const(return_type.dtype(cval), loc)
cval_const_var = ir.Var(scope, mk_unique_var("$cval_const"),
loc)
self.typemap[cval_const_var.name] = return_type.dtype
cval_const_assign = ir.Assign(cval_const_val,
cval_const_var, loc)
init_block.body.append(cval_const_assign)
# do setitem on `out` array
setitemexpr = ir.StaticSetItem(out_arr, slice(None, None),
slice_inst_var, cval_const_var,
loc)
init_block.body.append(setitemexpr)
sig = signature(types.none, self.typemap[out_arr.name],
self.typemap[slice_inst_var.name],
self.typemap[out_arr.name].dtype)
self.calltypes[setitemexpr] = sig
self.replace_return_with_setitem(stencil_blocks, exit_value_var,
parfor_body_exit_label)
if config.DEBUG_ARRAY_OPT >= 1:
print("stencil_blocks after replacing return")
ir_utils.dump_blocks(stencil_blocks)
setitem_call = ir.SetItem(out_arr, parfor_ind_var, exit_value_var, loc)
self.calltypes[setitem_call] = signature(
types.none, self.typemap[out_arr.name],
self.typemap[parfor_ind_var.name],
self.typemap[out_arr.name].dtype
)
stencil_blocks[parfor_body_exit_label].body.extend(for_replacing_ret)
stencil_blocks[parfor_body_exit_label].body.append(setitem_call)
# simplify CFG of parfor body (exit block could be simplified often)
# add dummy return to enable CFG
dummy_loc = ir.Loc("stencilparfor_dummy", -1)
ret_const_var = ir.Var(scope, mk_unique_var("$cval_const"), dummy_loc)
cval_const_assign = ir.Assign(ir.Const(0, loc=dummy_loc), ret_const_var, dummy_loc)
stencil_blocks[parfor_body_exit_label].body.append(cval_const_assign)
stencil_blocks[parfor_body_exit_label].body.append(
ir.Return(ret_const_var, dummy_loc),
)
stencil_blocks = ir_utils.simplify_CFG(stencil_blocks)
stencil_blocks[max(stencil_blocks.keys())].body.pop()
if config.DEBUG_ARRAY_OPT >= 1:
print("stencil_blocks after adding SetItem")
ir_utils.dump_blocks(stencil_blocks)
pattern = ('stencil', [start_lengths, end_lengths])
parfor = numba.parfors.parfor.Parfor(loopnests, init_block, stencil_blocks,
loc, parfor_ind_var, equiv_set, pattern, self.flags)
gen_nodes.append(parfor)
gen_nodes.append(ir.Assign(out_arr, target, loc))
return gen_nodes
def _get_stencil_last_ind(self, dim_size, end_length, gen_nodes, scope,
loc):
last_ind = dim_size
if end_length != 0:
# set last index to size minus stencil size to avoid invalid
# memory access
index_const = ir.Var(scope, mk_unique_var("stencil_const_var"),
loc)
self.typemap[index_const.name] = types.intp
if isinstance(end_length, numbers.Number):
const_assign = ir.Assign(ir.Const(end_length, loc),
index_const, loc)
else:
const_assign = ir.Assign(end_length, index_const, loc)
gen_nodes.append(const_assign)
last_ind = ir.Var(scope, mk_unique_var("last_ind"), loc)
self.typemap[last_ind.name] = types.intp
g_var = ir.Var(scope, mk_unique_var("compute_last_ind_var"), loc)
check_func = numba.njit(_compute_last_ind)
func_typ = types.functions.Dispatcher(check_func)
self.typemap[g_var.name] = func_typ
g_obj = ir.Global("_compute_last_ind", check_func, loc)
g_assign = ir.Assign(g_obj, g_var, loc)
gen_nodes.append(g_assign)
index_call = ir.Expr.call(g_var, [dim_size, index_const], (), loc)
self.calltypes[index_call] = func_typ.get_call_type(
self.typingctx, [types.intp, types.intp], {})
index_assign = ir.Assign(index_call, last_ind, loc)
gen_nodes.append(index_assign)
return last_ind
def _get_stencil_start_ind(self, start_length, gen_nodes, scope, loc):
if isinstance(start_length, int):
return abs(min(start_length, 0))
def get_start_ind(s_length):
return abs(min(s_length, 0))
f_ir = compile_to_numba_ir(get_start_ind, {}, self.typingctx,
self.targetctx, (types.intp,), self.typemap,
self.calltypes)
assert len(f_ir.blocks) == 1
block = f_ir.blocks.popitem()[1]
replace_arg_nodes(block, [start_length])
gen_nodes += block.body[:-2]
ret_var = block.body[-2].value.value
return ret_var
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 _add_index_offsets(self, index_list, index_offsets, new_body,
scope, loc):
""" Does the actual work of adding loop index variables to the
relative index constants or variables.
"""
assert len(index_list) == len(index_offsets)
# shortcut if all values are integer
if all([isinstance(v, int) for v in index_list+index_offsets]):
# add offsets in all dimensions
return list(map(add, index_list, index_offsets))
out_nodes = []
index_vars = []
for i in range(len(index_list)):
# new_index = old_index + offset
old_index_var = index_list[i]
if isinstance(old_index_var, int):
old_index_var = ir.Var(scope,
mk_unique_var("old_index_var"), loc)
self.typemap[old_index_var.name] = types.intp
const_assign = ir.Assign(ir.Const(index_list[i], loc),
old_index_var, loc)
out_nodes.append(const_assign)
offset_var = index_offsets[i]
if isinstance(offset_var, int):
offset_var = ir.Var(scope,
mk_unique_var("offset_var"), loc)
self.typemap[offset_var.name] = types.intp
const_assign = ir.Assign(ir.Const(index_offsets[i], loc),
offset_var, loc)
out_nodes.append(const_assign)
if (isinstance(old_index_var, slice)
or isinstance(self.typemap[old_index_var.name],
types.misc.SliceType)):
# only one arg can be slice
assert self.typemap[offset_var.name] == types.intp
index_var = self._add_offset_to_slice(old_index_var, offset_var,
out_nodes, scope, loc)
index_vars.append(index_var)
continue
if (isinstance(offset_var, slice)
or isinstance(self.typemap[offset_var.name],
types.misc.SliceType)):
# only one arg can be slice
assert self.typemap[old_index_var.name] == types.intp
index_var = self._add_offset_to_slice(offset_var, old_index_var,
out_nodes, scope, loc)
index_vars.append(index_var)
continue
index_var = ir.Var(scope,
mk_unique_var("offset_stencil_index"), loc)
self.typemap[index_var.name] = types.intp
index_call = ir.Expr.binop(operator.add, old_index_var,
offset_var, loc)
self.calltypes[index_call] = self.typingctx.resolve_function_type(
operator.add, (types.intp, types.intp), {})
index_assign = ir.Assign(index_call, index_var, loc)
out_nodes.append(index_assign)
index_vars.append(index_var)
new_body.extend(out_nodes)
return index_vars
def _add_offset_to_slice(self, slice_var, offset_var, out_nodes, scope,
loc):
if isinstance(slice_var, slice):
f_text = """def f(offset):
return slice({} + offset, {} + offset)
""".format(slice_var.start, slice_var.stop)
loc = {}
exec(f_text, {}, loc)
f = loc['f']
args = [offset_var]
arg_typs = (types.intp,)
else:
def f(old_slice, offset):
return slice(old_slice.start + offset, old_slice.stop + offset)
args = [slice_var, offset_var]
slice_type = self.typemap[slice_var.name]
arg_typs = (slice_type, types.intp,)
_globals = self.func_ir.func_id.func.__globals__
f_ir = compile_to_numba_ir(f, _globals, self.typingctx, self.targetctx,
arg_typs, self.typemap, self.calltypes)
_, block = f_ir.blocks.popitem()
replace_arg_nodes(block, args)
new_index = block.body[-2].value.value
out_nodes.extend(block.body[:-2]) # ignore return nodes
return new_index
def get_stencil_ir(sf, typingctx, args, scope, loc, input_dict, typemap,
calltypes):
"""get typed IR from stencil bytecode
"""
from numba.core.cpu import CPUContext
from numba.core.registry import cpu_target
from numba.core.annotations import type_annotations
from numba.core.typed_passes import type_inference_stage
# get untyped IR
stencil_func_ir = sf.kernel_ir.copy()
# copy the IR nodes to avoid changing IR in the StencilFunc object
stencil_blocks = copy.deepcopy(stencil_func_ir.blocks)
stencil_func_ir.blocks = stencil_blocks
name_var_table = ir_utils.get_name_var_table(stencil_func_ir.blocks)
if "out" in name_var_table:
raise ValueError("Cannot use the reserved word 'out' in stencil kernels.")
# get typed IR with a dummy pipeline (similar to test_parfors.py)
from numba.core.registry import cpu_target
targetctx = cpu_target.target_context
tp = DummyPipeline(typingctx, targetctx, args, stencil_func_ir)
rewrites.rewrite_registry.apply('before-inference', tp.state)
tp.state.typemap, tp.state.return_type, tp.state.calltypes, _ = type_inference_stage(
tp.state.typingctx, tp.state.targetctx, tp.state.func_ir,
tp.state.args, None)
type_annotations.TypeAnnotation(
func_ir=tp.state.func_ir,
typemap=tp.state.typemap,
calltypes=tp.state.calltypes,
lifted=(),
lifted_from=None,
args=tp.state.args,
return_type=tp.state.return_type,
html_output=config.HTML)
# make block labels unique
stencil_blocks = ir_utils.add_offset_to_labels(stencil_blocks,
ir_utils.next_label())
min_label = min(stencil_blocks.keys())
max_label = max(stencil_blocks.keys())
ir_utils._the_max_label.update(max_label)
if config.DEBUG_ARRAY_OPT >= 1:
print("Initial stencil_blocks")
ir_utils.dump_blocks(stencil_blocks)
# rename variables,
var_dict = {}
for v, typ in tp.state.typemap.items():
new_var = ir.Var(scope, mk_unique_var(v), loc)
var_dict[v] = new_var
typemap[new_var.name] = typ # add new var type for overall function
ir_utils.replace_vars(stencil_blocks, var_dict)
if config.DEBUG_ARRAY_OPT >= 1:
print("After replace_vars")
ir_utils.dump_blocks(stencil_blocks)
# add call types to overall function
for call, call_typ in tp.state.calltypes.items():
calltypes[call] = call_typ
arg_to_arr_dict = {}
# replace arg with arr
for block in stencil_blocks.values():
for stmt in block.body:
if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Arg):
if config.DEBUG_ARRAY_OPT >= 1:
print("input_dict", input_dict, stmt.value.index,
stmt.value.name, stmt.value.index in input_dict)
arg_to_arr_dict[stmt.value.name] = input_dict[stmt.value.index].name
stmt.value = input_dict[stmt.value.index]
if config.DEBUG_ARRAY_OPT >= 1:
print("arg_to_arr_dict", arg_to_arr_dict)
print("After replace arg with arr")
ir_utils.dump_blocks(stencil_blocks)
ir_utils.remove_dels(stencil_blocks)
stencil_func_ir.blocks = stencil_blocks
return stencil_func_ir, sf.get_return_type(args)[0], arg_to_arr_dict
class DummyPipeline(object):
def __init__(self, typingctx, targetctx, args, f_ir):
from numba.core.compiler import StateDict
self.state = StateDict()
self.state.typingctx = typingctx
self.state.targetctx = targetctx
self.state.args = args
self.state.func_ir = f_ir
self.state.typemap = None
self.state.return_type = None
self.state.calltypes = None
def _get_const_index_expr(stencil_ir, func_ir, index_var):
"""
infer index_var as constant if it is of a expression form like c-1 where c
is a constant in the outer function.
index_var is assumed to be inside stencil kernel
"""
const_val = guard(
_get_const_index_expr_inner, stencil_ir, func_ir, index_var)
if const_val is not None:
return const_val
return index_var
def _get_const_index_expr_inner(stencil_ir, func_ir, index_var):
"""inner constant inference function that calls constant, unary and binary
cases.
"""
require(isinstance(index_var, ir.Var))
# case where the index is a const itself in outer function
var_const = guard(_get_const_two_irs, stencil_ir, func_ir, index_var)
if var_const is not None:
return var_const
# get index definition
index_def = ir_utils.get_definition(stencil_ir, index_var)
# match inner_var = unary(index_var)
var_const = guard(
_get_const_unary_expr, stencil_ir, func_ir, index_def)
if var_const is not None:
return var_const
# match inner_var = arg1 + arg2
var_const = guard(
_get_const_binary_expr, stencil_ir, func_ir, index_def)
if var_const is not None:
return var_const
raise GuardException
def _get_const_two_irs(ir1, ir2, var):
"""get constant in either of two IRs if available
otherwise, throw GuardException
"""
var_const = guard(find_const, ir1, var)
if var_const is not None:
return var_const
var_const = guard(find_const, ir2, var)
if var_const is not None:
return var_const
raise GuardException
def _get_const_unary_expr(stencil_ir, func_ir, index_def):
"""evaluate constant unary expr if possible
otherwise, raise GuardException
"""
require(isinstance(index_def, ir.Expr) and index_def.op == 'unary')
inner_var = index_def.value
# return -c as constant
const_val = _get_const_index_expr_inner(stencil_ir, func_ir, inner_var)
op = OPERATORS_TO_BUILTINS[index_def.fn]
return eval("{}{}".format(op, const_val))
def _get_const_binary_expr(stencil_ir, func_ir, index_def):
"""evaluate constant binary expr if possible
otherwise, raise GuardException
"""
require(isinstance(index_def, ir.Expr) and index_def.op == 'binop')
arg1 = _get_const_index_expr_inner(stencil_ir, func_ir, index_def.lhs)
arg2 = _get_const_index_expr_inner(stencil_ir, func_ir, index_def.rhs)
op = OPERATORS_TO_BUILTINS[index_def.fn]
return eval("{}{}{}".format(arg1, op, arg2))