def _gen_rebalances(self, rebalance_arrs, blocks): # for block in blocks.values(): new_body = [] for inst in block.body: # TODO: handle hiframes filter etc. if isinstance(inst, Parfor): self._gen_rebalances(rebalance_arrs, {0: inst.init_block}) self._gen_rebalances(rebalance_arrs, inst.loop_body) if isinstance( inst, ir.Assign) and inst.target.name in rebalance_arrs: out_arr = inst.target self.func_ir._definitions[out_arr.name].remove(inst.value) # hold inst results in tmp array tmp_arr = ir.Var(out_arr.scope, mk_unique_var("rebalance_tmp"), out_arr.loc) self.typemap[tmp_arr.name] = self.typemap[out_arr.name] inst.target = tmp_arr nodes = [inst] def f(in_arr): # pragma: no cover out_a = hpat.distributed_api.rebalance_array(in_arr) f_block = compile_to_numba_ir( f, { 'hpat': hpat }, self.typingctx, (self.typemap[tmp_arr.name], ), self.typemap, self.calltypes).blocks.popitem()[1] replace_arg_nodes(f_block, [tmp_arr]) nodes += f_block.body[:-3] # remove none return nodes[-1].target = out_arr # update definitions dumm_block = ir.Block(out_arr.scope, out_arr.loc) dumm_block.body = nodes build_definitions({0: dumm_block}, self.func_ir._definitions) new_body += nodes else: new_body.append(inst) block.body = new_body
def _replace_stencil_accesses(self, stencil_ir, parfor_vars, in_args, index_offsets, stencil_func, arg_to_arr_dict): """ Convert relative indexing in the stencil kernel to standard indexing by adding the loop index variables to the corresponding dimensions of the array index tuples. """ stencil_blocks = stencil_ir.blocks in_arr = in_args[0] in_arg_names = [x.name for x in in_args] if "standard_indexing" in stencil_func.options: for x in stencil_func.options["standard_indexing"]: if x not in arg_to_arr_dict: raise ValueError("Standard indexing requested for an array " \ "name not present in the stencil kernel definition.") standard_indexed = [ arg_to_arr_dict[x] for x in stencil_func.options["standard_indexing"] ] else: standard_indexed = [] if in_arr.name in standard_indexed: raise ValueError("The first argument to a stencil kernel must use " \ "relative indexing, not standard indexing.") ndims = self.typemap[in_arr.name].ndim scope = in_arr.scope loc = in_arr.loc # replace access indices, find access lengths in each dimension need_to_calc_kernel = stencil_func.neighborhood is None # If we need to infer the kernel size then initialize the minimum and # maximum seen indices for each dimension to 0. If we already have # the neighborhood calculated then just convert from neighborhood format # to the separate start and end lengths format used here. if need_to_calc_kernel: start_lengths = ndims * [0] end_lengths = ndims * [0] else: start_lengths = [x[0] for x in stencil_func.neighborhood] end_lengths = [x[1] for x in stencil_func.neighborhood] # Get all the tuples defined in the stencil blocks. tuple_table = ir_utils.get_tuple_table(stencil_blocks) found_relative_index = False # For all blocks in the stencil kernel... for label, block in stencil_blocks.items(): new_body = [] # For all statements in those blocks... for stmt in block.body: # Reject assignments to input arrays. if ((isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op in ['setitem', 'static_setitem'] and stmt.value.value.name in in_arg_names) or ((isinstance(stmt, ir.SetItem) or isinstance(stmt, ir.StaticSetItem)) and stmt.target.name in in_arg_names)): raise ValueError( "Assignments to arrays passed to stencil kernels is not allowed." ) # We found a getitem for some array. If that array is an input # array and isn't in the list of standard indexed arrays then # update min and max seen indices if we are inferring the # kernel size and create a new tuple where the relative offsets # are added to loop index vars to get standard indexing. if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op in ['static_getitem', 'getitem'] and stmt.value.value.name in in_arg_names and stmt.value.value.name not in standard_indexed): index_list = stmt.value.index # handle 1D case if ndims == 1: index_list = [index_list] else: if hasattr(index_list, 'name') and index_list.name in tuple_table: index_list = tuple_table[index_list.name] # indices can be inferred as constant in simple expressions # like -c where c is constant # handled here since this is a common stencil index pattern stencil_ir._definitions = ir_utils.build_definitions( stencil_blocks) index_list = [ _get_const_index_expr(stencil_ir, self.func_ir, v) for v in index_list ] if index_offsets: index_list = self._add_index_offsets( index_list, list(index_offsets), new_body, scope, loc) # update min and max indices if need_to_calc_kernel: # all indices should be integer to be able to calculate # neighborhood automatically if (isinstance(index_list, ir.Var) or any( [not isinstance(v, int) for v in index_list])): raise ValueError( "Variable stencil index only " "possible with known neighborhood") start_lengths = list( map(min, start_lengths, index_list)) end_lengths = list(map(max, end_lengths, index_list)) found_relative_index = True # update access indices index_vars = self._add_index_offsets( parfor_vars, list(index_list), new_body, scope, loc) # new access index tuple if ndims == 1: ind_var = index_vars[0] else: ind_var = ir.Var( scope, mk_unique_var("$parfor_index_ind_var"), loc) self.typemap[ind_var.name] = types.containers.UniTuple( types.intp, ndims) tuple_call = ir.Expr.build_tuple(index_vars, loc) tuple_assign = ir.Assign(tuple_call, ind_var, loc) new_body.append(tuple_assign) # getitem return type is scalar if all indices are integer if all([ self.typemap[v.name] == types.intp for v in index_vars ]): getitem_return_typ = self.typemap[ stmt.value.value.name].dtype else: # getitem returns an array getitem_return_typ = self.typemap[ stmt.value.value.name] # new getitem with the new index var getitem_call = ir.Expr.getitem(stmt.value.value, ind_var, loc) self.calltypes[getitem_call] = signature( getitem_return_typ, self.typemap[stmt.value.value.name], self.typemap[ind_var.name]) stmt.value = getitem_call new_body.append(stmt) block.body = new_body if need_to_calc_kernel and not found_relative_index: raise ValueError("Stencil kernel with no accesses to " \ "relatively indexed arrays.") return start_lengths, end_lengths
def _init_run(self): self.func_ir._definitions = build_definitions(self.func_ir.blocks) self._parallel_accesses = set() self._T_arrs = set() self.second_pass = False self.in_parallel_parfor = -1
def _replace_stencil_accesses(self, stencil_ir, parfor_vars, in_args, index_offsets, stencil_func, arg_to_arr_dict): """ Convert relative indexing in the stencil kernel to standard indexing by adding the loop index variables to the corresponding dimensions of the array index tuples. """ stencil_blocks = stencil_ir.blocks in_arr = in_args[0] in_arg_names = [x.name for x in in_args] if "standard_indexing" in stencil_func.options: for x in stencil_func.options["standard_indexing"]: if x not in arg_to_arr_dict: raise ValueError("Standard indexing requested for an array " \ "name not present in the stencil kernel definition.") standard_indexed = [arg_to_arr_dict[x] for x in stencil_func.options["standard_indexing"]] else: standard_indexed = [] if in_arr.name in standard_indexed: raise ValueError("The first argument to a stencil kernel must use " \ "relative indexing, not standard indexing.") ndims = self.typemap[in_arr.name].ndim scope = in_arr.scope loc = in_arr.loc # replace access indices, find access lengths in each dimension need_to_calc_kernel = stencil_func.neighborhood is None # If we need to infer the kernel size then initialize the minimum and # maximum seen indices for each dimension to 0. If we already have # the neighborhood calculated then just convert from neighborhood format # to the separate start and end lengths format used here. if need_to_calc_kernel: start_lengths = ndims*[0] end_lengths = ndims*[0] else: start_lengths = [x[0] for x in stencil_func.neighborhood] end_lengths = [x[1] for x in stencil_func.neighborhood] # Get all the tuples defined in the stencil blocks. tuple_table = ir_utils.get_tuple_table(stencil_blocks) found_relative_index = False # For all blocks in the stencil kernel... for label, block in stencil_blocks.items(): new_body = [] # For all statements in those blocks... for stmt in block.body: # Reject assignments to input arrays. if ((isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op in ['setitem', 'static_setitem'] and stmt.value.value.name in in_arg_names) or ((isinstance(stmt, ir.SetItem) or isinstance(stmt, ir.StaticSetItem)) and stmt.target.name in in_arg_names)): raise ValueError("Assignments to arrays passed to stencil kernels is not allowed.") # We found a getitem for some array. If that array is an input # array and isn't in the list of standard indexed arrays then # update min and max seen indices if we are inferring the # kernel size and create a new tuple where the relative offsets # are added to loop index vars to get standard indexing. if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op in ['static_getitem', 'getitem'] and stmt.value.value.name in in_arg_names and stmt.value.value.name not in standard_indexed): index_list = stmt.value.index # handle 1D case if ndims == 1: index_list = [index_list] else: if hasattr(index_list, 'name') and index_list.name in tuple_table: index_list = tuple_table[index_list.name] # indices can be inferred as constant in simple expressions # like -c where c is constant # handled here since this is a common stencil index pattern stencil_ir._definitions = ir_utils.build_definitions(stencil_blocks) index_list = [_get_const_index_expr( stencil_ir, self.func_ir, v) for v in index_list] if index_offsets: index_list = self._add_index_offsets(index_list, list(index_offsets), new_body, scope, loc) # update min and max indices if need_to_calc_kernel: # all indices should be integer to be able to calculate # neighborhood automatically if (isinstance(index_list, ir.Var) or any([not isinstance(v, int) for v in index_list])): raise ValueError("Variable stencil index only " "possible with known neighborhood") start_lengths = list(map(min, start_lengths, index_list)) end_lengths = list(map(max, end_lengths, index_list)) found_relative_index = True # update access indices index_vars = self._add_index_offsets(parfor_vars, list(index_list), new_body, scope, loc) # new access index tuple if ndims == 1: ind_var = index_vars[0] else: ind_var = ir.Var(scope, mk_unique_var( "$parfor_index_ind_var"), loc) self.typemap[ind_var.name] = types.containers.UniTuple( types.intp, ndims) tuple_call = ir.Expr.build_tuple(index_vars, loc) tuple_assign = ir.Assign(tuple_call, ind_var, loc) new_body.append(tuple_assign) # getitem return type is scalar if all indices are integer if all([self.typemap[v.name] == types.intp for v in index_vars]): getitem_return_typ = self.typemap[ stmt.value.value.name].dtype else: # getitem returns an array getitem_return_typ = self.typemap[stmt.value.value.name] # new getitem with the new index var getitem_call = ir.Expr.getitem(stmt.value.value, ind_var, loc) self.calltypes[getitem_call] = signature( getitem_return_typ, self.typemap[stmt.value.value.name], self.typemap[ind_var.name]) stmt.value = getitem_call new_body.append(stmt) block.body = new_body if need_to_calc_kernel and not found_relative_index: raise ValueError("Stencil kernel with no accesses to " \ "relatively indexed arrays.") return start_lengths, end_lengths