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