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 test_assign(self):
a = ir.Assign(self.var_a, self.var_b, self.loc1)
b = ir.Assign(self.var_a, self.var_b, self.loc1)
c = ir.Assign(self.var_a, self.var_b, self.loc2)
d = ir.Assign(self.var_c, self.var_b, self.loc1)
e = ir.Assign(self.var_a, self.var_c, self.loc1)
self.check(a, same=[b, c], different=[d, e])
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 _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 _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 _loop_lift_modify_call_block(liftedloop, block, inputs, outputs, returnto):
"""
Transform calling block from top-level function to call the lifted loop.
"""
scope = block.scope
loc = block.loc
blk = ir.Block(scope=scope, loc=loc)
# load loop
fn = ir.Const(value=liftedloop, loc=loc)
fnvar = scope.make_temp(loc=loc)
blk.append(ir.Assign(target=fnvar, value=fn, loc=loc))
# call loop
args = [scope.get_exact(name) for name in inputs]
callexpr = ir.Expr.call(func=fnvar, args=args, kws=(), loc=loc)
# temp variable for the return value
callres = scope.make_temp(loc=loc)
blk.append(ir.Assign(target=callres, value=callexpr, loc=loc))
# unpack return value
for i, out in enumerate(outputs):
target = scope.get_exact(out)
getitem = ir.Expr.static_getitem(value=callres,
index=i,
index_var=None,
loc=loc)
blk.append(ir.Assign(target=target, value=getitem, loc=loc))
# jump to next block
blk.append(ir.Jump(target=returnto, loc=loc))
return blk
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 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_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 gen_block():
parent = ir.Scope(None, self.loc1)
tmp = ir.Block(parent, self.loc2)
assign1 = ir.Assign(self.var_a, self.var_b, self.loc3)
assign2 = ir.Assign(self.var_a, self.var_c, self.loc3)
assign3 = ir.Assign(self.var_c, self.var_b, self.loc3)
tmp.append(assign1)
tmp.append(assign2)
tmp.append(assign3)
return tmp
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 get_column_read_nodes(c_type, cvar, file_name, i):
loc = cvar.loc
func_text = ('def f(fname):\n col_size = get_column_size_parquet(fname, {})\n'.
format(i))
# generate strings differently
if c_type == string_array_type:
# pass size for easier allocation and distributed analysis
func_text += ' column = read_parquet_str(fname, {}, col_size)\n'.format(
i)
else:
el_type = get_element_type(c_type.dtype)
func_text += ' column = np.empty(col_size, dtype=np.{})\n'.format(
el_type)
func_text += ' status = read_parquet(fname, {}, column, np.int32({}))\n'.format(
i, _type_to_pq_dtype_number[el_type])
loc_vars = {}
exec(func_text, {}, loc_vars)
size_func = loc_vars['f']
_, f_block = compile_to_numba_ir(size_func,
{'get_column_size_parquet': get_column_size_parquet,
'read_parquet': read_parquet,
'read_parquet_str': read_parquet_str, 'np': np,
'StringArray': StringArray}).blocks.popitem()
replace_arg_nodes(f_block, [file_name])
out_nodes = f_block.body[:-3]
for stmt in reversed(out_nodes):
if stmt.target.name.startswith("column"):
assign = ir.Assign(stmt.target, cvar, loc)
break
out_nodes.append(assign)
return out_nodes
def _insert_phi(self):
if self.dfainfo.incomings:
assert len(self.dfainfo.incomings) == 1
incomings = self.cfa.blocks[self.current_block_offset].incoming
phivar = self.dfainfo.incomings[0]
if len(incomings) == 1:
ib = utils.iter_next(iter(incomings))
lingering = self.dfa.infos[ib].stack
assert len(lingering) == 1
iv = lingering[0]
self.store(self.get(iv), phivar)
else:
# Invert the PHI node
for ib in incomings:
lingering = self.dfa.infos[ib].stack
assert len(lingering) == 1
iv = lingering[0]
# Add assignment in incoming block to forward the value
target = self.current_scope.get_or_define('$phi' + phivar,
loc=self.loc)
stmt = ir.Assign(value=self.get(iv),
target=target,
loc=self.loc)
self.blocks[ib].insert_before_terminator(stmt)
self.store(target, phivar)
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 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
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 store(self, value, name):
if self.current_block_offset in self.cfa.backbone:
target = self.current_scope.redefine(name, loc=self.loc)
else:
target = self.current_scope.get_or_define(name, loc=self.loc)
stmt = ir.Assign(value=value, target=target, loc=self.loc)
self.current_block.append(stmt)
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 get_column_read_nodes(c_type, cvar, xe_connect_var, xe_dset_var, i, schema_arr_var):
loc = cvar.loc
func_text = ('def f(xe_connect_var, xe_dset_var, schema_arr):\n')
func_text = (' col_size = get_column_size_xenon(xe_connect_var, xe_dset_var, {})\n'. format(i))
# func_text += ' print(col_size)\n'
# generate strings differently since upfront allocation is not possible
if c_type == string_array_type:
# pass size for easier allocation and distributed analysis
func_text += ' column = read_xenon_str(xe_connect_var, xe_dset_var, {}, col_size, schema_arr)\n'.format(i)
else:
el_type = get_element_type(c_type.dtype)
func_text += ' column = np.empty(col_size, dtype=np.{})\n'.format(el_type)
func_text += ' status = read_xenon_col(xe_connect_var, xe_dset_var, {}, column, schema_arr)\n'.format(i)
loc_vars = {}
exec(func_text, {}, loc_vars)
size_func = loc_vars['f']
_, f_block = compile_to_numba_ir(size_func,
{'get_column_size_xenon': get_column_size_xenon,
'read_xenon_col': read_xenon_col,
'read_xenon_str': read_xenon_str,
'np': np,
}).blocks.popitem()
replace_arg_nodes(f_block, [xe_connect_var, xe_dset_var, schema_arr_var])
out_nodes = f_block.body[:-3]
for stmt in reversed(out_nodes):
if isinstance(stmt, ir.Assign) and stmt.target.name.startswith("column"):
assign = ir.Assign(stmt.target, cvar, loc)
break
out_nodes.append(assign)
return out_nodes
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 _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 _replace_returns(blocks, target, return_label):
"""
Return return statement by assigning directly to target, and a jump.
"""
for label, block in blocks.items():
for i in range(len(block.body)):
stmt = block.body[i]
if isinstance(stmt, ir.Return):
assert(i + 1 == len(block.body))
block.body[i] = ir.Assign(stmt.value, target, stmt.loc)
block.body.append(ir.Jump(return_label, stmt.loc))
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 store(self, value, name, redefine=False):
"""
Store *value* (a Var instance) into the variable named *name*
(a str object).
"""
if redefine or self.current_block_offset in self.cfa.backbone:
target = self.current_scope.redefine(name, loc=self.loc)
else:
target = self.current_scope.get_or_define(name, loc=self.loc)
if isinstance(value, ir.Var):
value = self.assigner.assign(value, target)
stmt = ir.Assign(value=value, target=target, loc=self.loc)
self.current_block.append(stmt)
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 _insert_outgoing_phis(self):
"""
Add assignments to forward requested outgoing values
to subsequent blocks.
"""
for phiname, varname in self.dfainfo.outgoing_phis.items():
target = self.current_scope.get_or_define(phiname, loc=self.loc)
stmt = ir.Assign(value=self.get(varname),
target=target,
loc=self.loc)
if not self.current_block.is_terminated:
self.current_block.append(stmt)
else:
self.current_block.insert_before_terminator(stmt)
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]
