def apply(self, sdfg): graph = sdfg.nodes()[self.state_id] if self.expr_index == 0: map_entry = graph.nodes()[self.subgraph[GPUTransformMap._map_entry]] nsdfg_node = helpers.nest_state_subgraph( sdfg, graph, graph.scope_subgraph(map_entry), full_data=self.fullcopy) else: cnode = graph.nodes()[self.subgraph[GPUTransformMap._reduce]] nsdfg_node = helpers.nest_state_subgraph(sdfg, graph, SubgraphView( graph, [cnode]), full_data=self.fullcopy) # Avoiding import loops from dace.transformation.interstate import GPUTransformSDFG transformation = GPUTransformSDFG(0, 0, {}, 0) transformation.register_trans = self.register_trans transformation.sequential_innermaps = self.sequential_innermaps transformation.toplevel_trans = self.toplevel_trans transformation.gpu_id = self.gpu_id transformation.apply(nsdfg_node.sdfg) # Inline back as necessary sdfg.apply_strict_transformations()
def apply(self, graph: SDFGState, sdfg: SDFG): if self.expr_index == 0: map_entry = self.map_entry nsdfg_node = helpers.nest_state_subgraph( sdfg, graph, graph.scope_subgraph(map_entry), full_data=self.fullcopy) else: cnode = self.reduce nsdfg_node = helpers.nest_state_subgraph(sdfg, graph, SubgraphView( graph, [cnode]), full_data=self.fullcopy) # Avoiding import loops from dace.transformation.interstate import GPUTransformSDFG transformation = GPUTransformSDFG(sdfg, 0, -1, {}, 0) transformation.register_trans = self.register_trans transformation.sequential_innermaps = self.sequential_innermaps transformation.toplevel_trans = self.toplevel_trans transformation.apply(nsdfg_node.sdfg, nsdfg_node.sdfg) # Inline back as necessary sdfg.simplify()
def apply(self, _, sdfg: sd.SDFG): # Obtain loop information guard: sd.SDFGState = self.loop_guard body: sd.SDFGState = self.loop_begin # Obtain iteration variable, range, and stride itervar, (start, end, step), _ = find_for_loop(sdfg, guard, body) forward_loop = step > 0 for node in body.nodes(): if isinstance(node, nodes.MapEntry): map_entry = node if isinstance(node, nodes.MapExit): map_exit = node # nest map's content in sdfg map_subgraph = body.scope_subgraph(map_entry, include_entry=False, include_exit=False) nsdfg = helpers.nest_state_subgraph(sdfg, body, map_subgraph, full_data=True) # replicate loop in nested sdfg new_before, new_guard, new_after = nsdfg.sdfg.add_loop( before_state=None, loop_state=nsdfg.sdfg.nodes()[0], loop_end_state=None, after_state=None, loop_var=itervar, initialize_expr=f'{start}', condition_expr=f'{itervar} <= {end}' if forward_loop else f'{itervar} >= {end}', increment_expr=f'{itervar} + {step}' if forward_loop else f'{itervar} - {abs(step)}') # remove outer loop before_guard_edge = nsdfg.sdfg.edges_between(new_before, new_guard)[0] for e in nsdfg.sdfg.out_edges(new_guard): if e.dst is new_after: guard_after_edge = e else: guard_body_edge = e for body_inedge in sdfg.in_edges(body): if body_inedge.src is guard: guard_body_edge.data.assignments.update(body_inedge.data.assignments) sdfg.remove_edge(body_inedge) for body_outedge in sdfg.out_edges(body): sdfg.remove_edge(body_outedge) for guard_inedge in sdfg.in_edges(guard): before_guard_edge.data.assignments.update(guard_inedge.data.assignments) guard_inedge.data.assignments = {} sdfg.add_edge(guard_inedge.src, body, guard_inedge.data) sdfg.remove_edge(guard_inedge) for guard_outedge in sdfg.out_edges(guard): if guard_outedge.dst is body: guard_body_edge.data.assignments.update(guard_outedge.data.assignments) else: guard_after_edge.data.assignments.update(guard_outedge.data.assignments) guard_outedge.data.condition = CodeBlock("1") sdfg.add_edge(body, guard_outedge.dst, guard_outedge.data) sdfg.remove_edge(guard_outedge) sdfg.remove_node(guard) if itervar in nsdfg.symbol_mapping: del nsdfg.symbol_mapping[itervar] if itervar in sdfg.symbols: del sdfg.symbols[itervar] # Add missing data/symbols for s in nsdfg.sdfg.free_symbols: if s in nsdfg.symbol_mapping: continue if s in sdfg.symbols: nsdfg.symbol_mapping[s] = s elif s in sdfg.arrays: desc = sdfg.arrays[s] access = body.add_access(s) conn = nsdfg.sdfg.add_datadesc(s, copy.deepcopy(desc)) nsdfg.sdfg.arrays[s].transient = False nsdfg.add_in_connector(conn) body.add_memlet_path(access, map_entry, nsdfg, memlet=Memlet.from_array(s, desc), dst_conn=conn) else: raise NotImplementedError(f"Free symbol {s} is neither a symbol nor data.") to_delete = set() for s in nsdfg.symbol_mapping: if s not in nsdfg.sdfg.free_symbols: to_delete.add(s) for s in to_delete: del nsdfg.symbol_mapping[s] # propagate scope for correct volumes scope_tree = ScopeTree(map_entry, map_exit) scope_tree.parent = ScopeTree(None, None) # The first execution helps remove apperances of symbols # that are now defined only in the nested SDFG in memlets. propagation.propagate_memlets_scope(sdfg, body, scope_tree) for s in to_delete: if helpers.is_symbol_unused(sdfg, s): sdfg.remove_symbol(s) from dace.transformation.interstate import RefineNestedAccess transformation = RefineNestedAccess() transformation.setup_match(sdfg, 0, sdfg.node_id(body), {RefineNestedAccess.nsdfg: body.node_id(nsdfg)}, 0) transformation.apply(body, sdfg) # Second propagation for refined accesses. propagation.propagate_memlets_scope(sdfg, body, scope_tree)
def apply(self, sdfg): graph = sdfg.nodes()[self.state_id] subgraph = self.subgraph_view(sdfg) map_entries = helpers.get_outermost_scope_maps(sdfg, graph, subgraph) result = StencilTiling.topology(sdfg, graph, map_entries) (children_dict, parent_dict, sink_maps) = result # next up, calculate inferred ranges for each map # for each map entry, this contains a tuple of dicts: # each of those maps from data_name of the array to # inferred outer ranges. An inferred outer range is created # by taking the union of ranges of inner subsets corresponding # to that data and substituting this subset by the min / max of the # parametrized map boundaries # finally, from these outer ranges we can easily calculate # strides and tile sizes required for every map inferred_ranges = defaultdict(dict) # create array of reverse topologically sorted map entries # to iterate over topo_reversed = [] queue = set(sink_maps.copy()) while len(queue) > 0: element = next(e for e in queue if not children_dict[e] - set(topo_reversed)) topo_reversed.append(element) queue.remove(element) for parent in parent_dict[element]: queue.add(parent) # main loop # first get coverage dicts for each map entry # for each map, contains a tuple of two dicts # each of those two maps from data name to outer range coverage = {} for map_entry in map_entries: coverage[map_entry] = StencilTiling.coverage_dicts( sdfg, graph, map_entry, outer_range=True) # we have a mapping from data name to outer range # however we want a mapping from map parameters to outer ranges # for this we need to find out how all array dimensions map to # outer ranges variable_mapping = defaultdict(list) for map_entry in topo_reversed: map = map_entry.map # first find out variable mapping for e in itertools.chain( graph.out_edges(map_entry), graph.in_edges(graph.exit_node(map_entry))): mapping = [] for dim in e.data.subset: syms = set() for d in dim: syms |= symbolic.symlist(d).keys() if len(syms) > 1: raise NotImplementedError( "One incoming or outgoing stencil subset is indexed " "by multiple map parameters. " "This is not supported yet.") try: mapping.append(syms.pop()) except KeyError: # just append None if there is no map symbol in it. # we don't care for now. mapping.append(None) if e.data in variable_mapping: # assert that this is the same everywhere. # else we might run into problems assert variable_mapping[e.data.data] == mapping else: variable_mapping[e.data.data] = mapping # now do mapping data name -> outer range # and from that infer mapping variable -> outer range local_ranges = {dn: None for dn in coverage[map_entry][1].keys()} for data_name, cov in coverage[map_entry][1].items(): local_ranges[data_name] = subsets.union( local_ranges[data_name], cov) # now look at proceeding maps # and union those subsets -> could be larger with stencil indent for child_map in children_dict[map_entry]: if data_name in coverage[child_map][0]: local_ranges[data_name] = subsets.union( local_ranges[data_name], coverage[child_map][0][data_name]) # final assignent: combine local_ranges and variable_mapping # together into inferred_ranges inferred_ranges[map_entry] = {p: None for p in map.params} for data_name, ranges in local_ranges.items(): for param, r in zip(variable_mapping[data_name], ranges): # create new range from this subset and assign rng = subsets.Range((r, )) if param: inferred_ranges[map_entry][param] = subsets.union( inferred_ranges[map_entry][param], rng) # get parameters -- should all be the same params = next(iter(map_entries)).map.params.copy() # define reference range as inferred range of one of the sink maps self.reference_range = inferred_ranges[next(iter(sink_maps))] if self.debug: print("StencilTiling::Reference Range", self.reference_range) # next up, search for the ranges that don't change invariant_dims = [] for idx, p in enumerate(params): different = False if self.reference_range[p] is None: invariant_dims.append(idx) warnings.warn( f"StencilTiling::No Stencil pattern detected for parameter {p}" ) continue for m in map_entries: if inferred_ranges[m][p] != self.reference_range[p]: different = True break if not different: invariant_dims.append(idx) warnings.warn( f"StencilTiling::No Stencil pattern detected for parameter {p}" ) # during stripmining, we will create new outer map entries # for easy access self._outer_entries = set() # with inferred_ranges constructed, we can begin to strip mine for map_entry in map_entries: # Retrieve map entry and exit nodes. map = map_entry.map stripmine_subgraph = { StripMining._map_entry: graph.nodes().index(map_entry) } sdfg_id = sdfg.sdfg_id last_map_entry = None original_schedule = map_entry.schedule self.tile_sizes = [] self.tile_offset_lower = [] self.tile_offset_upper = [] # strip mining each dimension where necessary removed_maps = 0 for dim_idx, param in enumerate(map_entry.map.params): # get current_node tile size if dim_idx >= len(self.strides): tile_stride = symbolic.pystr_to_symbolic(self.strides[-1]) else: tile_stride = symbolic.pystr_to_symbolic( self.strides[dim_idx]) trivial = False if dim_idx in invariant_dims: self.tile_sizes.append(tile_stride) self.tile_offset_lower.append(0) self.tile_offset_upper.append(0) else: target_range_current = inferred_ranges[map_entry][param] reference_range_current = self.reference_range[param] min_diff = symbolic.SymExpr(reference_range_current.min_element()[0] \ - target_range_current.min_element()[0]) max_diff = symbolic.SymExpr(target_range_current.max_element()[0] \ - reference_range_current.max_element()[0]) try: min_diff = symbolic.evaluate(min_diff, {}) max_diff = symbolic.evaluate(max_diff, {}) except TypeError: raise RuntimeError("Symbolic evaluation of map " "ranges failed. Please check " "your parameters and match.") self.tile_sizes.append(tile_stride + max_diff + min_diff) self.tile_offset_lower.append( symbolic.pystr_to_symbolic(str(min_diff))) self.tile_offset_upper.append( symbolic.pystr_to_symbolic(str(max_diff))) # get calculated parameters tile_size = self.tile_sizes[-1] dim_idx -= removed_maps # If map or tile sizes are trivial, skip strip-mining map dimension # special cases: # if tile size is trivial AND we have an invariant dimension, skip if tile_size == map.range.size()[dim_idx] and ( dim_idx + removed_maps) in invariant_dims: continue # trivial map: we just continue if map.range.size()[dim_idx] in [0, 1]: continue if tile_size == 1 and tile_stride == 1 and ( dim_idx + removed_maps) in invariant_dims: trivial = True removed_maps += 1 # indent all map ranges accordingly and then perform # strip mining on these. Offset inner maps accordingly afterwards range_tuple = (map.range[dim_idx][0] + self.tile_offset_lower[-1], map.range[dim_idx][1] - self.tile_offset_upper[-1], map.range[dim_idx][2]) map.range[dim_idx] = range_tuple stripmine = StripMining(sdfg_id, self.state_id, stripmine_subgraph, 0) stripmine.tiling_type = 'ceilrange' stripmine.dim_idx = dim_idx stripmine.new_dim_prefix = self.prefix if not trivial else '' # use tile_stride for both -- we will extend # the inner tiles later stripmine.tile_size = str(tile_stride) stripmine.tile_stride = str(tile_stride) outer_map = stripmine.apply(sdfg) outer_map.schedule = original_schedule # apply to the new map the schedule of the original one map_entry.schedule = self.schedule # if tile stride is 1, we can make a nice simplification by just # taking the overapproximated inner range as inner range # this eliminates the min/max in the range which # enables loop unrolling if tile_stride == 1: map_entry.range[dim_idx] = tuple( symbolic.SymExpr(el._approx_expr) if isinstance( el, symbolic.SymExpr) else el for el in map_entry.range[dim_idx]) # in map_entry: enlarge tiles by upper and lower offset # doing it this way and not via stripmine strides ensures # that the max gets changed as well old_range = map_entry.range[dim_idx] map_entry.range[dim_idx] = ((old_range[0] - self.tile_offset_lower[-1]), (old_range[1] + self.tile_offset_upper[-1]), old_range[2]) # We have to propagate here for correct outer volume and subset sizes _propagate_node(graph, map_entry) _propagate_node(graph, graph.exit_node(map_entry)) # usual tiling pipeline if last_map_entry: new_map_entry = graph.in_edges(map_entry)[0].src mapcollapse_subgraph = { MapCollapse._outer_map_entry: graph.node_id(last_map_entry), MapCollapse._inner_map_entry: graph.node_id(new_map_entry) } mapcollapse = MapCollapse(sdfg_id, self.state_id, mapcollapse_subgraph, 0) mapcollapse.apply(sdfg) last_map_entry = graph.in_edges(map_entry)[0].src # add last instance of map entries to _outer_entries if last_map_entry: self._outer_entries.add(last_map_entry) # Map Unroll Feature: only unroll if conditions are met: # Only unroll if at least one of the inner map ranges is strictly larger than 1 # Only unroll if strides all are one if self.unroll_loops and all(s == 1 for s in self.strides) and any( s not in [0, 1] for s in map_entry.range.size()): l = len(map_entry.params) if l > 1: subgraph = { MapExpansion.map_entry: graph.nodes().index(map_entry) } trafo_expansion = MapExpansion(sdfg.sdfg_id, sdfg.nodes().index(graph), subgraph, 0) trafo_expansion.apply(sdfg) maps = [map_entry] for _ in range(l - 1): map_entry = graph.out_edges(map_entry)[0].dst maps.append(map_entry) for map in reversed(maps): # MapToForLoop subgraph = { MapToForLoop._map_entry: graph.nodes().index(map) } trafo_for_loop = MapToForLoop(sdfg.sdfg_id, sdfg.nodes().index(graph), subgraph, 0) trafo_for_loop.apply(sdfg) nsdfg = trafo_for_loop.nsdfg # LoopUnroll guard = trafo_for_loop.guard end = trafo_for_loop.after_state begin = next(e.dst for e in nsdfg.out_edges(guard) if e.dst != end) subgraph = { DetectLoop._loop_guard: nsdfg.nodes().index(guard), DetectLoop._loop_begin: nsdfg.nodes().index(begin), DetectLoop._exit_state: nsdfg.nodes().index(end) } transformation = LoopUnroll(0, 0, subgraph, 0) transformation.apply(nsdfg) elif self.unroll_loops: warnings.warn( "Did not unroll loops. Either all ranges are equal to " "one or range difference is symbolic.") self._outer_entries = list(self._outer_entries)