Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
getitemname = ir_utils.mk_unique_var("getitem")
getitemvar = ir.Var(scope, getitemname, loc)
getitemcall = ir.Expr.getitem(stmt_index_var,
const_index_vars[dim], loc)
new_body.append(ir.Assign(getitemcall, getitemvar, loc))
# Get the type of this particular part of the index tuple.
one_index_typ = stmt_index_var_typ[dim]
# If the array is indexed with a slice then we
# have to add the index value with a call to
# slice_addition.
if isinstance(one_index_typ, types.misc.SliceType):
sa_var = ir.Var(scope, ir_utils.mk_unique_var("slice_addition"), loc)
sa_func = numba.njit(slice_addition)
sa_func_typ = types.functions.Dispatcher(sa_func)
typemap[sa_var.name] = sa_func_typ
g_sa = ir.Global("slice_addition", sa_func, loc)
new_body.append(ir.Assign(g_sa, sa_var, loc))
slice_addition_call = ir.Expr.call(sa_var, [getitemvar, index_vars[dim]], (), loc)
calltypes[slice_addition_call] = sa_func_typ.get_call_type(self._typingctx, [one_index_typ, types.intp], {})
new_body.append(ir.Assign(slice_addition_call, tmpvar, loc))
else:
acc_call = ir.Expr.binop(operator.add, getitemvar,
index_vars[dim], loc)
new_body.append(ir.Assign(acc_call, tmpvar, loc))
tuple_call = ir.Expr.build_tuple(ind_stencils, loc)
new_body.append(ir.Assign(tuple_call, s_index_var, loc))
new_body.append(ir.Assign(
ir.Expr.getitem(stmt.value.value,s_index_var,loc),
stmt.target,loc))
else:
new_body.append(stmt)
in1 = args[0]
in2 = args[1]
el_typ = self.typemap[lhs.name].dtype
assert self._get_ndims(
in1.name) <= 2 and self._get_ndims(
in2.name) == 1
# loop range correlation is same as first dimention of 1st input
corr = self.array_analysis.array_shape_classes[in1.name][0]
size_var = self.array_analysis.array_size_vars[in1.name][0]
scope = lhs.scope
loc = expr.loc
index_var = ir.Var(scope, mk_unique_var("parfor_index"), lhs.loc)
self.typemap[index_var.name] = types.intp
loopnests = [LoopNest(index_var, 0, size_var, 1, corr)]
init_block = ir.Block(scope, loc)
parfor = Parfor(
loopnests,
init_block,
{},
loc,
self.array_analysis,
index_var)
if self._get_ndims(in1.name) == 2:
# for 2D input, there is an inner loop
# correlation of inner dimension
inner_size_var = self.array_analysis.array_size_vars[in1.name][1]
# loop structure: range block, header block, body
range_label = next_label()
header_label = next_label()
body_label = next_label()
def lower_assign(self, inst):
"""
The returned object must have a new reference
"""
value = inst.value
if isinstance(value, ir.Const):
return self.lower_const(value.value)
elif isinstance(value, ir.Var):
val = self.loadvar(value.name)
self.incref(val)
return val
elif isinstance(value, ir.Expr):
return self.lower_expr(value)
elif isinstance(value, ir.Global):
return self.lower_global(value.name, value.value)
else:
raise NotImplementedError(type(value), value)
# make sure init block isn't removed
init_block_label = min(blocks.keys())
# update loop body blocks
blocks.pop(init_block_label)
parfor.loop_body = blocks
# make sure dummy jump to loop body isn't altered
first_body_label = min(parfor.loop_body.keys())
assert isinstance(parfor.init_block.body[-1], ir.Jump)
# remove dummy jump to loop body
parfor.init_block.body.pop()
# make sure dummy jump back to loop body isn't altered
for block in parfor.loop_body.values():
if (isinstance(block.body[-1], ir.Jump) and
block.body[-1].target == first_body_label):
# remove dummy jump back to loop
block.body.pop()
return
def lower_inst(self, inst):
if isinstance(inst, ir.Assign):
value = self.lower_assign(inst)
self.storevar(value, inst.target.name)
elif isinstance(inst, ir.SetItem):
target = self.loadvar(inst.target.name)
index = self.loadvar(inst.index.name)
value = self.loadvar(inst.value.name)
ok = self.pyapi.object_setitem(target, index, value)
negone = lc.Constant.int_signextend(ok.type, -1)
pred = self.builder.icmp(lc.ICMP_EQ, ok, negone)
with cgutils.if_unlikely(self.builder, pred):
self.return_exception_raised()
elif isinstance(inst, ir.Return):
retval = self.loadvar(inst.value.name)
self.incref(retval)
for label in reversed(order):
for inst in reversed(blocks[label].body):
if isinstance(inst, ir.Assign):
lhs = inst.target.name
rhs = inst.value
if isinstance(rhs, ir.Expr) and rhs.op == 'call':
call_table[rhs.func.name] = []
if isinstance(rhs, ir.Expr) and rhs.op == 'getattr':
if lhs in call_table:
call_table[lhs].append(rhs.attr)
reverse_call_table[rhs.value.name] = lhs
if lhs in reverse_call_table:
call_var = reverse_call_table[lhs]
call_table[call_var].append(rhs.attr)
reverse_call_table[rhs.value.name] = call_var
if isinstance(rhs, ir.Global):
if lhs in call_table:
call_table[lhs].append(rhs.value)
if lhs in reverse_call_table:
call_var = reverse_call_table[lhs]
call_table[call_var].append(rhs.value)
if isinstance(rhs, ir.FreeVar):
if lhs in call_table:
call_table[lhs].append(rhs.value)
if lhs in reverse_call_table:
call_var = reverse_call_table[lhs]
call_table[call_var].append(rhs.value)
if isinstance(rhs, ir.Var):
if lhs in call_table:
call_table[lhs].append(rhs.name)
reverse_call_table[rhs.name] = lhs
if lhs in reverse_call_table:
def filter_distributed_run(filter_node, array_dists, typemap, calltypes, typingctx, targetctx, dist_pass):
# TODO: rebalance if output distributions are 1D instead of 1D_Var
loc = filter_node.loc
bool_arr = filter_node.bool_arr
out = []
for col_name, col_in_var in filter_node.df_in_vars.items():
col_out_var = filter_node.df_out_vars[col_name]
# using getitem like Numba for filtering arrays
# TODO: generate parfor
getitem_call = ir.Expr.getitem(col_in_var, bool_arr, loc)
calltypes[getitem_call] = signature(
typemap[col_out_var.name], # output type
typemap[col_in_var.name], # input type
typemap[bool_arr.name]) # index type
out.append(ir.Assign(getitem_call, col_out_var, loc))
return out
def wrap_loop_body(loop_body):
blocks = loop_body.copy() # shallow copy is enough
first_label = min(blocks.keys())
last_label = max(blocks.keys())
loc = blocks[last_label].loc
blocks[last_label].body.append(ir.Jump(first_label, loc))
return blocks
ast_call = ast.Call(fn_ast_name, ast_args, [], None, None)
return ast_call, env
elif isinstance(expr, ir.Var):
name = expr.name
var = ast.Name(name, ast.Load(),
lineno=expr.loc.line,
col_offset=expr.loc.col if expr.loc.col else 0)
typ = typemap[name]
if isinstance(typ, types.Array):
var = ast.Subscript(
value = var,
slice = ast.Index(value = _mk_tuple([ast.Name(v, ast.Load()) for v in subscripts])),
ctx = ast.Load())
return var, {}
elif isinstance(expr, ir.Const):
return ast.Num(expr.value), {}
raise NotImplementedError(
"Don't know how to translate array expression '%r'" % (expr,))
def _set_getsetitem_index(node, new_ind):
if ((isinstance(node, ir.Expr) and node.op == 'static_getitem') or isinstance(node, ir.StaticSetItem)):
node.index_var = new_ind
node.index = None
return
assert ((isinstance(node, ir.Expr) and node.op == 'getitem') or isinstance(node, ir.SetItem))
node.index = new_ind