def test_var(self):
a = ir.Var(None, 'foo', self.loc1)
b = ir.Var(None, 'foo', self.loc1)
c = ir.Var(None, 'foo', self.loc2)
d = ir.Var(ir.Scope(None, ir.unknown_loc), 'foo', self.loc1)
e = ir.Var(None, 'bar', self.loc1)
self.check(a, same=[b, c, d], different=[e])
def _analyze_op_pair_first(self, scope, equiv_set, expr):
# make dummy lhs since we don't have access to lhs
typ = self.typemap[expr.value.name].first_type
if not isinstance(typ, types.NamedTuple):
return None
lhs = ir.Var(scope, mk_unique_var("tuple_var"), expr.loc)
self.typemap[lhs.name] = typ
rhs = ir.Expr.pair_first(expr.value, expr.loc)
lhs_assign = ir.Assign(rhs, lhs, expr.loc)
#(shape, post) = self._gen_shape_call(equiv_set, lhs, typ.count, )
var = lhs
out = []
size_vars = []
ndims = typ.count
for i in range(ndims):
# get size: Asize0 = A_sh_attr[0]
size_var = ir.Var(var.scope, mk_unique_var(
"{}_size{}".format(var.name, i)), var.loc)
getitem = ir.Expr.static_getitem(lhs, i, None, var.loc)
self.calltypes[getitem] = None
out.append(ir.Assign(getitem, size_var, var.loc))
self._define(equiv_set, size_var, types.intp, getitem)
size_vars.append(size_var)
shape = tuple(size_vars)
return shape, [lhs_assign] + out
def _handle_f_close_call(self, stmt, lhs_var, rhs):
func_def = guard(get_definition, self.func_ir, rhs.func)
assert func_def is not None
# rare case where function variable is assigned to a new variable
if isinstance(func_def, ir.Var):
rhs.func = func_def
return self._handle_f_close_call(stmt, lhs_var, rhs)
if (isinstance(func_def, ir.Expr) and func_def.op == 'getattr'
and func_def.value.name in self.h5_files
and func_def.attr == 'close'):
f_id = func_def.value
scope = lhs_var.scope
loc = lhs_var.loc
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5close_attr = getattr(g_pio_var, h5close)
h5close_attr_call = ir.Expr.getattr(g_pio_var, "h5close", loc)
attr_var = ir.Var(scope, mk_unique_var("$h5close_attr"), loc)
attr_assign = ir.Assign(h5close_attr_call, attr_var, loc)
# h5close(f_id)
close_call = ir.Expr.call(attr_var, [f_id], (), loc)
close_assign = ir.Assign(close_call, lhs_var, loc)
return [g_pio_assign, attr_assign, close_assign]
return None
def _handle_merge(self, lhs, rhs):
if guard(find_callname, self.func_ir, rhs) == ('merge', 'pandas'):
if len(rhs.args) < 2:
raise ValueError("left and right arguments required for merge")
left_df = rhs.args[0]
right_df = rhs.args[1]
kws = dict(rhs.kws)
if 'on' in kws:
left_on = get_constant(self.func_ir, kws['on'], None)
right_on = left_on
else: # pragma: no cover
if 'left_on' not in kws or 'right_on' not in kws:
raise ValueError("merge 'on' or 'left_on'/'right_on'"
"arguments required")
left_on = get_constant(self.func_ir, kws['left_on'], None)
right_on = get_constant(self.func_ir, kws['right_on'], None)
if left_on is None or right_on is None:
raise ValueError("merge key values should be constant strings")
scope = lhs.scope
loc = lhs.loc
self.df_vars[lhs.name] = {}
# add columns from left to output
for col, _ in self.df_vars[left_df.name].items():
self.df_vars[lhs.name][col] = ir.Var(scope, mk_unique_var(col),
loc)
# add columns from right to output
for col, _ in self.df_vars[right_df.name].items():
self.df_vars[lhs.name][col] = ir.Var(scope, mk_unique_var(col),
loc)
self._update_df_cols()
return [
hiframes_join.Join(lhs.name, left_df.name, right_df.name,
left_on, right_on, self.df_vars, lhs.loc)
]
return None
def mk_loop_header(typemap, phi_var, calltypes, scope, loc):
"""make a block that is a loop header updating iteration variables.
target labels in branch need to be set.
"""
# iternext_var = iternext(phi_var)
iternext_var = ir.Var(scope, mk_unique_var("$iternext_var"), loc)
typemap[iternext_var.name] = types.containers.Pair(
types.intp, types.boolean)
iternext_call = ir.Expr.iternext(phi_var, loc)
calltypes[iternext_call] = signature(
types.containers.Pair(
types.intp,
types.boolean),
types.range_iter64_type)
iternext_assign = ir.Assign(iternext_call, iternext_var, loc)
# pair_first_var = pair_first(iternext_var)
pair_first_var = ir.Var(scope, mk_unique_var("$pair_first_var"), loc)
typemap[pair_first_var.name] = types.intp
pair_first_call = ir.Expr.pair_first(iternext_var, loc)
pair_first_assign = ir.Assign(pair_first_call, pair_first_var, loc)
# pair_second_var = pair_second(iternext_var)
pair_second_var = ir.Var(scope, mk_unique_var("$pair_second_var"), loc)
typemap[pair_second_var.name] = types.boolean
pair_second_call = ir.Expr.pair_second(iternext_var, loc)
pair_second_assign = ir.Assign(pair_second_call, pair_second_var, loc)
# phi_b_var = pair_first_var
phi_b_var = ir.Var(scope, mk_unique_var("$phi"), loc)
typemap[phi_b_var.name] = types.intp
phi_b_assign = ir.Assign(pair_first_var, phi_b_var, loc)
# branch pair_second_var body_block out_block
branch = ir.Branch(pair_second_var, -1, -1, loc)
header_block = ir.Block(scope, loc)
header_block.body = [iternext_assign, pair_first_assign,
pair_second_assign, phi_b_assign, branch]
return header_block
def mk_range_block(typemap, start, stop, step, calltypes, scope, loc):
"""make a block that initializes loop range and iteration variables.
target label in jump needs to be set.
"""
# g_range_var = Global(range)
g_range_var = ir.Var(scope, mk_unique_var("$range_g_var"), loc)
typemap[g_range_var.name] = get_global_func_typ(range)
g_range = ir.Global('range', range, loc)
g_range_assign = ir.Assign(g_range, g_range_var, loc)
arg_nodes, args = _mk_range_args(typemap, start, stop, step, scope, loc)
# range_call_var = call g_range_var(start, stop, step)
range_call = ir.Expr.call(g_range_var, args, (), loc)
calltypes[range_call] = typemap[g_range_var.name].get_call_type(
typing.Context(), [types.intp] * len(args), {})
#signature(types.range_state64_type, types.intp)
range_call_var = ir.Var(scope, mk_unique_var("$range_c_var"), loc)
typemap[range_call_var.name] = types.iterators.RangeType(types.intp)
range_call_assign = ir.Assign(range_call, range_call_var, loc)
# iter_var = getiter(range_call_var)
iter_call = ir.Expr.getiter(range_call_var, loc)
calltypes[iter_call] = signature(types.range_iter64_type,
types.range_state64_type)
iter_var = ir.Var(scope, mk_unique_var("$iter_var"), loc)
typemap[iter_var.name] = types.iterators.RangeIteratorType(types.intp)
iter_call_assign = ir.Assign(iter_call, iter_var, loc)
# $phi = iter_var
phi_var = ir.Var(scope, mk_unique_var("$phi"), loc)
typemap[phi_var.name] = types.iterators.RangeIteratorType(types.intp)
phi_assign = ir.Assign(iter_var, phi_var, loc)
# jump to header
jump_header = ir.Jump(-1, loc)
range_block = ir.Block(scope, loc)
range_block.body = arg_nodes + [g_range_assign, range_call_assign,
iter_call_assign, phi_assign, jump_header]
return range_block
def convert_seq_to_atleast_3d(self, seq_arg_var):
"""convert array sequence to a sequence of arrays with at least 3d dims
"""
arr_args = self._get_sequence_arrs(seq_arg_var.name)
new_seq = []
for arr in arr_args:
curr_dims = self._get_ndims(arr.name)
if curr_dims < 3:
dummy_var = ir.Var(
arr.scope, mk_unique_var("__dummy_nd"), arr.loc)
self.typemap[dummy_var.name] = self.typemap[arr.name].copy(
ndim=3)
corrs = self.array_shape_classes[arr.name].copy()
if curr_dims == 0:
corrs = [CONST_CLASS] * 3
elif curr_dims == 1: # Numpy adds both prepends and appends a dim
corrs = [CONST_CLASS] + corrs + [CONST_CLASS]
elif curr_dims == 2: # append a dim
corrs = corrs + [CONST_CLASS]
self.array_shape_classes[dummy_var.name] = corrs
new_seq.append(dummy_var)
else:
new_seq.append(arr)
tup_name = mk_unique_var("__dummy_tup")
self.tuple_table[tup_name] = new_seq
return ir.Var(arr_args[0].scope, tup_name, arr_args[0].loc)
def _gen_size_call(self, var, i):
out = []
ndims = self._get_ndims(var.name)
# attr call: A_sh_attr = getattr(A, shape)
shape_attr_call = ir.Expr.getattr(var, "shape", var.loc)
attr_var = ir.Var(var.scope,
mk_unique_var(var.name + "_sh_attr" + str(i)),
var.loc)
self.typemap[attr_var.name] = types.containers.UniTuple(
types.intp, ndims)
attr_assign = ir.Assign(shape_attr_call, attr_var, var.loc)
out.append(attr_assign)
# const var for dim: $constA0 = Const(0)
const_node = ir.Const(i, var.loc)
const_var = ir.Var(var.scope,
mk_unique_var("$const" + var.name + str(i)),
var.loc)
self.typemap[const_var.name] = types.intp
const_assign = ir.Assign(const_node, const_var, var.loc)
out.append(const_assign)
# get size: Asize0 = A_sh_attr[0]
size_var = ir.Var(var.scope, mk_unique_var(var.name + "size" + str(i)),
var.loc)
self.typemap[size_var.name] = types.intp
getitem_node = ir.Expr.static_getitem(attr_var, i, const_var, var.loc)
self.calltypes[getitem_node] = None
getitem_assign = ir.Assign(getitem_node, size_var, var.loc)
out.append(getitem_assign)
return out
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 mk_alloc(typemap, calltypes, lhs, size_var, dtype, scope, loc):
"""generate an array allocation with np.empty() and return list of nodes.
size_var can be an int variable or tuple of int variables.
"""
out = []
ndims = 1
size_typ = types.intp
if isinstance(size_var, tuple):
if len(size_var) == 1:
size_var = size_var[0]
size_var = convert_size_to_var(size_var, typemap, scope, loc, out)
else:
# tuple_var = build_tuple([size_var...])
ndims = len(size_var)
tuple_var = ir.Var(scope, mk_unique_var("$tuple_var"), loc)
if typemap:
typemap[tuple_var.name] = types.containers.UniTuple(
types.intp, ndims)
# constant sizes need to be assigned to vars
new_sizes = [
convert_size_to_var(s, typemap, scope, loc, out)
for s in size_var
]
tuple_call = ir.Expr.build_tuple(new_sizes, loc)
tuple_assign = ir.Assign(tuple_call, tuple_var, loc)
out.append(tuple_assign)
size_var = tuple_var
size_typ = types.containers.UniTuple(types.intp, ndims)
# g_np_var = Global(numpy)
g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
if typemap:
typemap[g_np_var.name] = types.misc.Module(numpy)
g_np = ir.Global('np', numpy, loc)
g_np_assign = ir.Assign(g_np, g_np_var, loc)
# attr call: empty_attr = getattr(g_np_var, empty)
empty_attr_call = ir.Expr.getattr(g_np_var, "empty", loc)
attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)
if typemap:
typemap[attr_var.name] = get_np_ufunc_typ(numpy.empty)
attr_assign = ir.Assign(empty_attr_call, attr_var, loc)
# alloc call: lhs = empty_attr(size_var, typ_var)
typ_var = ir.Var(scope, mk_unique_var("$np_typ_var"), loc)
if typemap:
typemap[typ_var.name] = types.functions.NumberClass(dtype)
# assuming str(dtype) returns valid np dtype string
np_typ_getattr = ir.Expr.getattr(g_np_var, str(dtype), loc)
typ_var_assign = ir.Assign(np_typ_getattr, typ_var, loc)
alloc_call = ir.Expr.call(attr_var, [size_var, typ_var], (), loc)
if calltypes:
calltypes[alloc_call] = typemap[attr_var.name].get_call_type(
typing.Context(),
[size_typ, types.functions.NumberClass(dtype)], {})
# signature(
# types.npytypes.Array(dtype, ndims, 'C'), size_typ,
# types.functions.NumberClass(dtype))
alloc_assign = ir.Assign(alloc_call, lhs, loc)
out.extend([g_np_assign, attr_assign, typ_var_assign, alloc_assign])
return out
def inline_array(array_var, expr, stmts, list_vars, dels):
"""Check to see if the given "array_var" is created from a list
of constants, and try to inline the list definition as array
initialization.
Extra statements produced with be appended to "stmts".
"""
callname = guard(find_callname, func_ir, expr)
require(callname and callname[1] == 'numpy' and callname[0] == 'array')
require(expr.args[0].name in list_vars)
ret_type = calltypes[expr].return_type
require(isinstance(ret_type, types.ArrayCompatible) and
ret_type.ndim == 1)
loc = expr.loc
list_var = expr.args[0]
array_typ = typemap[array_var.name]
debug_print("inline array_var = ", array_var, " list_var = ", list_var)
dtype = array_typ.dtype
seq, op = find_build_sequence(func_ir, list_var)
size = len(seq)
size_var = ir.Var(scope, mk_unique_var("size"), loc)
size_tuple_var = ir.Var(scope, mk_unique_var("size_tuple"), loc)
size_typ = types.intp
size_tuple_typ = types.UniTuple(size_typ, 1)
typemap[size_var.name] = size_typ
typemap[size_tuple_var.name] = size_tuple_typ
stmts.append(_new_definition(func_ir, size_var,
ir.Const(size, loc=loc), loc))
stmts.append(_new_definition(func_ir, size_tuple_var,
ir.Expr.build_tuple(items=[size_var], loc=loc), loc))
empty_func = ir.Var(scope, mk_unique_var("empty_func"), loc)
fnty = get_np_ufunc_typ(np.empty)
sig = context.resolve_function_type(fnty, (size_typ,), {})
typemap[empty_func.name] = fnty #
stmts.append(_new_definition(func_ir, empty_func,
ir.Global('empty', np.empty, loc=loc), loc))
empty_call = ir.Expr.call(empty_func, [size_var], {}, loc=loc)
calltypes[empty_call] = typing.signature(array_typ, size_typ)
stmts.append(_new_definition(func_ir, array_var, empty_call, loc))
for i in range(size):
index_var = ir.Var(scope, mk_unique_var("index"), loc)
index_typ = types.intp
typemap[index_var.name] = index_typ
stmts.append(_new_definition(func_ir, index_var,
ir.Const(i, loc), loc))
setitem = ir.SetItem(array_var, index_var, seq[i], loc)
calltypes[setitem] = typing.signature(types.none, array_typ,
index_typ, dtype)
stmts.append(setitem)
stmts.extend(dels)
return True
def gen_empty_like(in_arr, out_arr):
scope = in_arr.scope
loc = in_arr.loc
# g_np_var = Global(numpy)
g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
g_np = ir.Global('np', np, loc)
g_np_assign = ir.Assign(g_np, g_np_var, loc)
# attr call: empty_attr = getattr(g_np_var, empty_like)
empty_attr_call = ir.Expr.getattr(g_np_var, "empty_like", loc)
attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)
attr_assign = ir.Assign(empty_attr_call, attr_var, loc)
# alloc call: out_arr = empty_attr(in_arr)
alloc_call = ir.Expr.call(attr_var, [in_arr], (), loc)
alloc_assign = ir.Assign(alloc_call, out_arr, loc)
return [g_np_assign, attr_assign, alloc_assign]
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 test_inline_update_target_def(self):
def test_impl(a):
if a == 1:
b = 2
else:
b = 3
return b
func_ir = compiler.run_frontend(test_impl)
blocks = list(func_ir.blocks.values())
for block in blocks:
for i, stmt in enumerate(block.body):
# match b = 2 and replace with lambda: 2
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Var)
and guard(find_const, func_ir, stmt.value) == 2):
# replace expr with a dummy call
func_ir._definitions[stmt.target.name].remove(stmt.value)
stmt.value = ir.Expr.call(
ir.Var(block.scope, "myvar", loc=stmt.loc), (), (),
stmt.loc)
func_ir._definitions[stmt.target.name].append(stmt.value)
#func = g.py_func#
inline_closure_call(func_ir, {}, block, i, lambda: 2)
break
self.assertEqual(len(func_ir._definitions['b']), 2)
def gen_parquet_read(self, file_name, table_types):
import pyarrow.parquet as pq
scope = file_name.scope
loc = file_name.loc
if table_types is None:
fname_def = guard(get_definition, self.func_ir, file_name)
if not isinstance(fname_def, ir.Const) or not isinstance(
fname_def.value, str):
raise ValueError("Parquet schema not available")
file_name_str = fname_def.value
col_names, col_types = parquet_file_schema(file_name_str)
else:
col_names = list(table_types.keys())
col_types = list(table_types.values())
out_nodes = []
col_items = []
for i, cname in enumerate(col_names):
# get column type from schema
c_type = col_types[i]
# create a variable for column and assign type
varname = mk_unique_var(cname)
#self.locals[varname] = c_type
cvar = ir.Var(scope, varname, loc)
col_items.append((cname, cvar))
out_nodes += get_column_read_nodes(c_type, cvar, file_name, i)
return col_items, out_nodes
def _gen_h5close(self, stmt, f_id):
lhs_var = stmt.target
scope = lhs_var.scope
loc = lhs_var.loc
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5close_attr = getattr(g_pio_var, h5close)
h5close_attr_call = ir.Expr.getattr(g_pio_var, "h5close", loc)
attr_var = ir.Var(scope, mk_unique_var("$h5close_attr"), loc)
attr_assign = ir.Assign(h5close_attr_call, attr_var, loc)
# h5close(f_id)
close_call = ir.Expr.call(attr_var, [f_id], (), loc)
close_assign = ir.Assign(close_call, lhs_var, loc)
return [g_pio_assign, attr_assign, close_assign]
def gen_stencil_call(in_arr, out_arr, kernel_func, index_offsets, fir_globals,
other_args=None, options=None):
if other_args is None:
other_args = []
if options is None:
options = {}
if index_offsets != [0]:
options['index_offsets'] = index_offsets
scope = in_arr.scope
loc = in_arr.loc
stencil_nodes = []
stencil_nodes += gen_empty_like(in_arr, out_arr)
kernel_var = ir.Var(scope, mk_unique_var("kernel_var"), scope)
if not isinstance(kernel_func, ir.Expr):
kernel_func = ir.Expr.make_function("kernel", kernel_func.__code__,
kernel_func.__closure__, kernel_func.__defaults__, loc)
stencil_nodes.append(ir.Assign(kernel_func, kernel_var, loc))
def f(A, B, f):
numba.stencil(f)(A, out=B)
f_block = compile_to_numba_ir(f, {'numba': numba}).blocks.popitem()[1]
replace_arg_nodes(f_block, [in_arr, out_arr, kernel_var])
stencil_nodes += f_block.body[:-3] # remove none return
setup_call = stencil_nodes[-2].value
stencil_call = stencil_nodes[-1].value
setup_call.kws = list(options.items())
stencil_call.args += other_args
return stencil_nodes
class CheckEquality(unittest.TestCase):
var_a = ir.Var(None, 'a', ir.unknown_loc)
var_b = ir.Var(None, 'b', ir.unknown_loc)
var_c = ir.Var(None, 'c', ir.unknown_loc)
var_d = ir.Var(None, 'd', ir.unknown_loc)
var_e = ir.Var(None, 'e', ir.unknown_loc)
loc1 = ir.Loc('mock', 1, 0)
loc2 = ir.Loc('mock', 2, 0)
loc3 = ir.Loc('mock', 3, 0)
def check(self, base, same=[], different=[]):
for s in same:
self.assertTrue(base == s)
for d in different:
self.assertTrue(base != d)
def _get_slice_range(self, index_slice, out):
scope = index_slice.scope
loc = index_slice.loc
# start = s.start
start_var = ir.Var(scope, mk_unique_var("$pio_range_start"), loc)
start_attr_call = ir.Expr.getattr(index_slice, "start", loc)
start_assign = ir.Assign(start_attr_call, start_var, loc)
# stop = s.stop
stop_var = ir.Var(scope, mk_unique_var("$pio_range_stop"), loc)
stop_attr_call = ir.Expr.getattr(index_slice, "stop", loc)
stop_assign = ir.Assign(stop_attr_call, stop_var, loc)
# size = stop-start
size_var = ir.Var(scope, mk_unique_var("$pio_range_size"), loc)
size_call = ir.Expr.binop('-', stop_var, start_var, loc)
size_assign = ir.Assign(size_call, size_var, loc)
out += [start_assign, stop_assign, size_assign]
return [start_var], [size_var]
def replace_return_with_setitem(self, blocks, index_vars, out_name):
"""
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.
"""
ret_blocks = []
for label, block in blocks.items():
scope = block.scope
loc = block.loc
new_body = []
for stmt in block.body:
if isinstance(stmt, ir.Return):
ret_blocks.append(label)
# If 1D array then avoid the tuple construction.
if len(index_vars) == 1:
rvar = ir.Var(scope, out_name, loc)
ivar = ir.Var(scope, index_vars[0], loc)
new_body.append(ir.SetItem(rvar, ivar, stmt.value,
loc))
else:
# Convert the string names of the index variables into
# ir.Var's.
var_index_vars = []
for one_var in index_vars:
index_var = ir.Var(scope, one_var, loc)
var_index_vars += [index_var]
s_index_name = ir_utils.mk_unique_var("stencil_index")
s_index_var = ir.Var(scope, s_index_name, loc)
# Build a tuple from the index ir.Var's.
tuple_call = ir.Expr.build_tuple(var_index_vars, loc)
new_body.append(ir.Assign(tuple_call, s_index_var,
loc))
rvar = ir.Var(scope, out_name, loc)
# Write the return statements original value into
# the array using the tuple index.
si = ir.SetItem(rvar, s_index_var, stmt.value, loc)
new_body.append(si)
else:
new_body.append(stmt)
block.body = new_body
return ret_blocks
def _gen_rolling_init(self, win_size, func, center):
nodes = []
right_length = 0
scope = win_size.scope
loc = win_size.loc
right_length = ir.Var(scope, mk_unique_var('zero_var'), scope)
nodes.append(ir.Assign(ir.Const(0, loc), right_length, win_size.loc))
def f(w):
return -w + 1
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [win_size])
nodes.extend(f_block.body[:-2]) # remove none return
left_length = nodes[-1].target
if center:
def f(w):
return -(w // 2)
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [win_size])
nodes.extend(f_block.body[:-2]) # remove none return
left_length = nodes[-1].target
def f(w):
return (w // 2)
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [win_size])
nodes.extend(f_block.body[:-2]) # remove none return
right_length = nodes[-1].target
def f(a, b):
return ((a, b), )
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, [left_length, right_length])
nodes.extend(f_block.body[:-2]) # remove none return
win_tuple = nodes[-1].target
index_offsets = [right_length]
if func == 'apply':
index_offsets = [left_length]
def f(a):
return (a, )
f_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(f_block, index_offsets)
nodes.extend(f_block.body[:-2]) # remove none return
index_offsets = nodes[-1].target
return index_offsets, win_tuple, nodes
def convert_size_to_var(size_var, typemap, scope, loc, nodes):
if isinstance(size_var, int):
new_size = ir.Var(scope, mk_unique_var("$alloc_size"), loc)
if typemap:
typemap[new_size.name] = types.intp
size_assign = ir.Assign(ir.Const(size_var, loc), new_size, loc)
nodes.append(size_assign)
return new_size
assert isinstance(size_var, ir.Var)
return size_var
def _gen_h5create_group(self, stmt, f_id):
lhs_var = stmt.target
scope = lhs_var.scope
loc = lhs_var.loc
args = [f_id] + stmt.value.args
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5create_group_attr = getattr(g_pio_var, h5create_group)
h5create_group_attr_call = ir.Expr.getattr(g_pio_var, "h5create_group",
loc)
attr_var = ir.Var(scope, mk_unique_var("$h5create_group_attr"), loc)
attr_assign = ir.Assign(h5create_group_attr_call, attr_var, loc)
# group_id = h5create_group(f_id)
create_group_call = ir.Expr.call(attr_var, args, (), loc)
create_group_assign = ir.Assign(create_group_call, lhs_var, loc)
# add to files since group behavior is same as files for many calls
self.h5_files[lhs_var.name] = "group"
return [g_pio_assign, attr_assign, create_group_assign]
def gen_stencil_call(in_arr, out_arr, code_expr, index_offsets):
scope = in_arr.scope
loc = in_arr.loc
alloc_nodes = gen_empty_like(in_arr, out_arr)
# generate stencil call
# g_numba_var = Global(numba)
g_numba_var = ir.Var(scope, mk_unique_var("$g_numba_var"), loc)
g_dist = ir.Global('numba', numba, loc)
g_numba_assign = ir.Assign(g_dist, g_numba_var, loc)
# attr call: stencil_attr = getattr(g_numba_var, stencil)
stencil_attr_call = ir.Expr.getattr(g_numba_var, "stencil", loc)
stencil_attr_var = ir.Var(scope, mk_unique_var("$stencil_attr"), loc)
stencil_attr_assign = ir.Assign(stencil_attr_call, stencil_attr_var, loc)
# stencil_out = numba.stencil()
stencil_out = ir.Var(scope, mk_unique_var("$stencil_out"), loc)
stencil_call = ir.Expr.call(stencil_attr_var, [in_arr, out_arr], (), loc)
stencil_call.stencil_def = code_expr
stencil_call.index_offsets = index_offsets
stencil_assign = ir.Assign(stencil_call, stencil_out, loc)
return alloc_nodes + [g_numba_assign, stencil_attr_assign, stencil_assign]
def _handle_h5_File_call(self, assign, lhs, rhs):
"""
Handle h5py.File calls like:
f = h5py.File(file_name, mode)
"""
# parallel arg = False for this stage
loc = lhs.loc
scope = lhs.scope
parallel_var = ir.Var(scope, mk_unique_var("$const_parallel"), loc)
parallel_assign = ir.Assign(ir.Const(0, loc), parallel_var, loc)
rhs.args.append(parallel_var)
return [parallel_assign, assign]
def _gen_h5size(self, f_id, dset, ndims, scope, loc, out):
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5size_attr = getattr(g_pio_var, h5size)
h5size_attr_call = ir.Expr.getattr(g_pio_var, "h5size", loc)
attr_var = ir.Var(scope, mk_unique_var("$h5size_attr"), loc)
attr_assign = ir.Assign(h5size_attr_call, attr_var, loc)
out += [g_pio_assign, attr_assign]
size_vars = []
for i in range(ndims):
dim_var = ir.Var(scope, mk_unique_var("$h5_dim_var"), loc)
dim_assign = ir.Assign(ir.Const(np.int32(i), loc), dim_var, loc)
out.append(dim_assign)
size_var = ir.Var(scope, mk_unique_var("$h5_size_var"), loc)
size_vars.append(size_var)
size_call = ir.Expr.call(attr_var, [f_id, dset, dim_var], (), loc)
size_assign = ir.Assign(size_call, size_var, loc)
out.append(size_assign)
return size_vars
def _gen_h5create_dset(self, stmt, f_id):
lhs_var = stmt.target
scope = lhs_var.scope
loc = lhs_var.loc
args = [f_id] + stmt.value.args
# append the dtype arg (e.g. dtype='f8')
assert stmt.value.kws and stmt.value.kws[0][0] == 'dtype'
args.append(stmt.value.kws[0][1])
# g_pio_var = Global(hpat.pio_api)
g_pio_var = ir.Var(scope, mk_unique_var("$pio_g_var"), loc)
g_pio = ir.Global('pio_api', hpat.pio_api, loc)
g_pio_assign = ir.Assign(g_pio, g_pio_var, loc)
# attr call: h5create_dset_attr = getattr(g_pio_var, h5create_dset)
h5create_dset_attr_call = ir.Expr.getattr(g_pio_var, "h5create_dset",
loc)
attr_var = ir.Var(scope, mk_unique_var("$h5create_dset_attr"), loc)
attr_assign = ir.Assign(h5create_dset_attr_call, attr_var, loc)
# dset_id = h5create_dset(f_id)
create_dset_call = ir.Expr.call(attr_var, args, (), loc)
create_dset_assign = ir.Assign(create_dset_call, lhs_var, loc)
self.h5_dsets[lhs_var.name] = (f_id, args[1])
self.h5_dsets_sizes[lhs_var.name] = self.tuple_table[args[2].name]
return [g_pio_assign, attr_assign, create_dset_assign]
def _mk_range_args(typemap, start, stop, step, scope, loc):
nodes = []
if isinstance(stop, ir.Var):
g_stop_var = stop
else:
assert isinstance(stop, int)
g_stop_var = ir.Var(scope, mk_unique_var("$range_stop"), loc)
if typemap:
typemap[g_stop_var.name] = types.intp
stop_assign = ir.Assign(ir.Const(stop, loc), g_stop_var, loc)
nodes.append(stop_assign)
if start == 0 and step == 1:
return nodes, [g_stop_var]
if isinstance(start, ir.Var):
g_start_var = start
else:
assert isinstance(start, int)
g_start_var = ir.Var(scope, mk_unique_var("$range_start"), loc)
if typemap:
typemap[g_start_var.name] = types.intp
start_assign = ir.Assign(ir.Const(start, loc), g_start_var)
nodes.append(start_assign)
if step == 1:
return nodes, [g_start_var, g_stop_var]
if isinstance(step, ir.Var):
g_step_var = step
else:
assert isinstance(step, int)
g_step_var = ir.Var(scope, mk_unique_var("$range_step"), loc)
if typemap:
typemap[g_step_var.name] = types.intp
step_assign = ir.Assign(ir.Const(step, loc), g_step_var)
nodes.append(step_assign)
return nodes, [g_start_var, g_stop_var, g_step_var]
def _handle_h5_File_call(self, assign, lhs, rhs):
"""
Handle h5py.File calls like:
f = h5py.File(file_name, mode)
"""
if guard(find_callname, self.func_ir, rhs) == ('File', 'h5py'):
self.h5_files[lhs.name] = rhs.args[0]
# parallel arg = False for this stage
loc = lhs.loc
scope = lhs.scope
parallel_var = ir.Var(scope, mk_unique_var("$const_parallel"), loc)
parallel_assign = ir.Assign(ir.Const(0, loc), parallel_var, loc)
rhs.args.append(parallel_var)
return [parallel_assign, assign]
return None
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)
v_nodes = self._gen_col_var(var_var, args, col_var)
def f(a):
a**0.5
s_block = compile_to_numba_ir(f, {}).blocks.popitem()[1]
replace_arg_nodes(s_block, [var_var])
s_nodes = s_block.body[:-3]
assert len(s_nodes) == 3
s_nodes[-1].target = out_var
return v_nodes + s_nodes
