def apply(self, graph: SDFGState, sdfg: SDFG): import dace.libraries.blas as blas transpose_a = self.transpose_a _at = self.at transpose_b = self.transpose_b _bt = self.bt a_times_b = self.a_times_b for src, src_conn, _, _, memlet in graph.in_edges(transpose_a): graph.add_edge(src, src_conn, a_times_b, '_b', memlet) graph.remove_node(transpose_a) for src, src_conn, _, _, memlet in graph.in_edges(transpose_b): graph.add_edge(src, src_conn, a_times_b, '_a', memlet) graph.remove_node(transpose_b) graph.remove_node(_at) graph.remove_node(_bt) for _, _, dst, dst_conn, memlet in graph.out_edges(a_times_b): subset = dcpy(memlet.subset) subset.squeeze() size = subset.size() shape = [size[1], size[0]] break tmp_name, tmp_arr = sdfg.add_temp_transient(shape, a_times_b.dtype) tmp_acc = graph.add_access(tmp_name) transpose_c = blas.Transpose('_Transpose_', a_times_b.dtype) for edge in graph.out_edges(a_times_b): _, _, dst, dst_conn, memlet = edge graph.remove_edge(edge) graph.add_edge(transpose_c, '_out', dst, dst_conn, memlet) graph.add_edge(a_times_b, '_c', tmp_acc, None, dace.Memlet.from_array(tmp_name, tmp_arr)) graph.add_edge(tmp_acc, None, transpose_c, '_inp', dace.Memlet.from_array(tmp_name, tmp_arr))
def consolidate_edges_scope( state: SDFGState, scope_node: Union[nd.EntryNode, nd.ExitNode]) -> int: """ Union scope-entering memlets relating to the same data node in a scope. This effectively reduces the number of connectors and allows more transformations to be performed, at the cost of losing the individual per-tasklet memlets. :param state: The SDFG state in which the scope to consolidate resides. :param scope_node: The scope node whose edges will be consolidated. :return: Number of edges removed. """ if scope_node is None: return 0 data_to_conn = {} consolidated = 0 if isinstance(scope_node, nd.EntryNode): outer_edges = state.in_edges inner_edges = state.out_edges remove_outer_connector = scope_node.remove_in_connector remove_inner_connector = scope_node.remove_out_connector prefix, oprefix = 'IN_', 'OUT_' else: outer_edges = state.out_edges inner_edges = state.in_edges remove_outer_connector = scope_node.remove_out_connector remove_inner_connector = scope_node.remove_in_connector prefix, oprefix = 'OUT_', 'IN_' edges_by_connector = collections.defaultdict(list) connectors_to_remove = set() for e in inner_edges(scope_node): edges_by_connector[e.src_conn].append(e) if e.data.data not in data_to_conn: data_to_conn[e.data.data] = e.src_conn elif data_to_conn[e.data.data] != e.src_conn: # Need to consolidate connectors_to_remove.add(e.src_conn) for conn in connectors_to_remove: e = edges_by_connector[conn][0] # Outer side of the scope - remove edge and union subsets target_conn = prefix + data_to_conn[e.data.data][len(oprefix):] conn_to_remove = prefix + conn[len(oprefix):] remove_outer_connector(conn_to_remove) out_edge = next(ed for ed in outer_edges(scope_node) if ed.dst_conn == target_conn) edge_to_remove = next(ed for ed in outer_edges(scope_node) if ed.dst_conn == conn_to_remove) out_edge.data.subset = sbs.union(out_edge.data.subset, edge_to_remove.data.subset) state.remove_edge(edge_to_remove) consolidated += 1 # Inner side of the scope - remove and reconnect remove_inner_connector(e.src_conn) for e in edges_by_connector[conn]: e._src_conn = data_to_conn[e.data.data] return consolidated
def apply(self, graph: SDFGState, sdfg: SDFG): tasklet = self.tasklet map_exit = self.map_exit outer_map_exit = self.outer_map_exit memlet = None edge = None for e in graph.out_edges(map_exit): memlet = e.data # TODO: What if there's more than one? if e.dst == outer_map_exit and isinstance(sdfg.arrays[memlet.data], data.Stream): edge = e break tasklet_memlet = None for e in graph.out_edges(tasklet): tasklet_memlet = e.data if tasklet_memlet.data == memlet.data: break bbox = map_exit.map.range.bounding_box_size() bbox_approx = [symbolic.overapproximate(dim) for dim in bbox] dataname = memlet.data # Create the new node: Temporary stream and an access node newname, _ = sdfg.add_stream('trans_' + dataname, sdfg.arrays[memlet.data].dtype, bbox_approx[0], storage=sdfg.arrays[memlet.data].storage, transient=True, find_new_name=True) snode = graph.add_access(newname) to_stream_mm = copy.deepcopy(memlet) to_stream_mm.data = snode.data tasklet_memlet.data = snode.data if self.with_buffer: newname_arr, _ = sdfg.add_transient('strans_' + dataname, [bbox_approx[0]], sdfg.arrays[memlet.data].dtype, find_new_name=True) anode = graph.add_access(newname_arr) to_array_mm = copy.deepcopy(memlet) to_array_mm.data = anode.data graph.add_edge(snode, None, anode, None, to_array_mm) else: anode = snode # Reconnect, assuming one edge to the stream graph.remove_edge(edge) graph.add_edge(map_exit, edge.src_conn, snode, None, to_stream_mm) graph.add_edge(anode, None, outer_map_exit, edge.dst_conn, memlet) return
def remove_edge_and_dangling_path(state: SDFGState, edge: MultiConnectorEdge): """ Removes an edge and all of its parent edges in a memlet path, cleaning dangling connectors and isolated nodes resulting from the removal. :param state: The state in which the edge exists. :param edge: The edge to remove. """ mtree = state.memlet_tree(edge) inwards = (isinstance(edge.src, nd.EntryNode) or isinstance(edge.dst, nd.EntryNode)) # Traverse tree upwards, removing edges and connectors as necessary curedge = mtree while curedge is not None: e = curedge.edge state.remove_edge(e) if inwards: neighbors = [ neighbor for neighbor in state.out_edges(e.src) if e.src_conn == neighbor.src_conn ] else: neighbors = [ neighbor for neighbor in state.in_edges(e.dst) if e.dst_conn == neighbor.dst_conn ] if len(neighbors) > 0: # There are still edges connected, leave as-is break # Remove connector and matching outer connector if inwards: if e.src_conn: e.src.remove_out_connector(e.src_conn) e.src.remove_in_connector('IN' + e.src_conn[3:]) else: if e.dst_conn: e.dst.remove_in_connector(e.dst_conn) e.src.remove_out_connector('OUT' + e.dst_conn[2:]) # Continue traversing upwards curedge = curedge.parent else: # Check if an isolated node have been created at the root and remove root_edge = mtree.root().edge root_node: nd.Node = root_edge.src if inwards else root_edge.dst if state.degree(root_node) == 0: state.remove_node(root_node)
def apply(self, graph: SDFGState, sdfg: SDFG): in_array = self.in_array med_array = self.med_array out_array = self.out_array # Modify all edges that point to in_array to point to out_array for in_edge in graph.in_edges(in_array): # Make all memlets that write to in_array write to out_array instead tree = graph.memlet_tree(in_edge) for te in tree: if te.data.data == in_array.data: te.data.data = out_array.data # Redirect edge to in_array graph.remove_edge(in_edge) graph.add_edge(in_edge.src, in_edge.src_conn, out_array, None, in_edge.data) graph.remove_node(med_array) graph.remove_node(in_array)
def merge_maps( graph: SDFGState, outer_map_entry: nd.MapEntry, outer_map_exit: nd.MapExit, inner_map_entry: nd.MapEntry, inner_map_exit: nd.MapExit, param_merge: Callable[[ParamsType, ParamsType], ParamsType] = lambda p1, p2: p1 + p2, range_merge: Callable[[RangesType, RangesType], RangesType] = lambda r1, r2: type(r1) (r1.ranges + r2.ranges) ) -> (nd.MapEntry, nd.MapExit): """ Merges two maps (their entries and exits). It is assumed that the operation is valid. """ outer_map = outer_map_entry.map inner_map = inner_map_entry.map # Create merged map by inheriting attributes from outer map and using # the merge functions for parameters and ranges. merged_map = copy.deepcopy(outer_map) merged_map.label = outer_map.label merged_map.params = param_merge(outer_map.params, inner_map.params) merged_map.range = range_merge(outer_map.range, inner_map.range) merged_entry = nd.MapEntry(merged_map) merged_entry.in_connectors = outer_map_entry.in_connectors merged_entry.out_connectors = outer_map_entry.out_connectors merged_exit = nd.MapExit(merged_map) merged_exit.in_connectors = outer_map_exit.in_connectors merged_exit.out_connectors = outer_map_exit.out_connectors graph.add_nodes_from([merged_entry, merged_exit]) # Handle the case of dynamic map inputs in the inner map inner_dynamic_map_inputs = dynamic_map_inputs(graph, inner_map_entry) for edge in inner_dynamic_map_inputs: remove_conn = (len( list(graph.out_edges_by_connector(edge.src, edge.src_conn))) == 1) conn_to_remove = edge.src_conn[4:] if remove_conn: merged_entry.remove_in_connector('IN_' + conn_to_remove) merged_entry.remove_out_connector('OUT_' + conn_to_remove) merged_entry.add_in_connector( edge.dst_conn, inner_map_entry.in_connectors[edge.dst_conn]) outer_edge = next( graph.in_edges_by_connector(outer_map_entry, 'IN_' + conn_to_remove)) graph.add_edge(outer_edge.src, outer_edge.src_conn, merged_entry, edge.dst_conn, outer_edge.data) if remove_conn: graph.remove_edge(outer_edge) # Redirect inner in edges. for edge in graph.out_edges(inner_map_entry): if edge.src_conn is None: # Empty memlets graph.add_edge(merged_entry, edge.src_conn, edge.dst, edge.dst_conn, edge.data) continue # Get memlet path and edge path = graph.memlet_path(edge) ind = path.index(edge) # Add an edge directly from the previous source connector to the # destination graph.add_edge(merged_entry, path[ind - 1].src_conn, edge.dst, edge.dst_conn, edge.data) # Redirect inner out edges. for edge in graph.in_edges(inner_map_exit): if edge.dst_conn is None: # Empty memlets graph.add_edge(edge.src, edge.src_conn, merged_exit, edge.dst_conn, edge.data) continue # Get memlet path and edge path = graph.memlet_path(edge) ind = path.index(edge) # Add an edge directly from the source to the next destination # connector graph.add_edge(edge.src, edge.src_conn, merged_exit, path[ind + 1].dst_conn, edge.data) # Redirect outer edges. change_edge_dest(graph, outer_map_entry, merged_entry) change_edge_src(graph, outer_map_exit, merged_exit) # Clean-up graph.remove_nodes_from( [outer_map_entry, outer_map_exit, inner_map_entry, inner_map_exit]) return merged_entry, merged_exit
def apply(self, graph: SDFGState, sdfg: SDFG): node_a = self.node_a node_b = self.node_b prefix = self.prefix # Determine direction of new memlet scope_dict = graph.scope_dict() propagate_forward = sd.scope_contains_scope(scope_dict, node_a, node_b) array = self.array if array is None or len(array) == 0: array = next(e.data.data for e in graph.edges_between(node_a, node_b) if e.data.data is not None and e.data.wcr is None) original_edge = None invariant_memlet = None for edge in graph.edges_between(node_a, node_b): if array == edge.data.data: original_edge = edge invariant_memlet = edge.data break if invariant_memlet is None: for edge in graph.edges_between(node_a, node_b): original_edge = edge invariant_memlet = edge.data warnings.warn('Array %s not found! Using array %s instead.' % (array, invariant_memlet.data)) array = invariant_memlet.data break if invariant_memlet is None: raise NameError('Array %s not found!' % array) if self.create_array: # Add transient array new_data, _ = sdfg.add_transient( name=prefix + invariant_memlet.data, shape=[ symbolic.overapproximate(r).simplify() for r in invariant_memlet.bounding_box_size() ], dtype=sdfg.arrays[invariant_memlet.data].dtype, find_new_name=True) else: new_data = prefix + invariant_memlet.data data_node = nodes.AccessNode(new_data) # Store as fields so that other transformations can use them self._local_name = new_data self._data_node = data_node to_data_mm = copy.deepcopy(invariant_memlet) from_data_mm = copy.deepcopy(invariant_memlet) offset = subsets.Indices([r[0] for r in invariant_memlet.subset]) # Reconnect, assuming one edge to the access node graph.remove_edge(original_edge) if propagate_forward: graph.add_edge(node_a, original_edge.src_conn, data_node, None, to_data_mm) new_edge = graph.add_edge(data_node, None, node_b, original_edge.dst_conn, from_data_mm) else: new_edge = graph.add_edge(node_a, original_edge.src_conn, data_node, None, to_data_mm) graph.add_edge(data_node, None, node_b, original_edge.dst_conn, from_data_mm) # Offset all edges in the memlet tree (including the new edge) for edge in graph.memlet_tree(new_edge): edge.data.subset.offset(offset, True) edge.data.data = new_data return data_node
def apply(self, graph: SDFGState, sdfg: SDFG): """ This method applies the mapfusion transformation. Other than the removal of the second map entry node (SME), and the first map exit (FME) node, it has the following side effects: 1. Any transient adjacent to both FME and SME with degree = 2 will be removed. The tasklets that use/produce it shall be connected directly with a scalar/new transient (if the dataflow is more than a single scalar) 2. If this transient is adjacent to FME and SME and has other uses, it will be adjacent to the new map exit post fusion. Tasklet-> Tasklet edges will ALSO be added as mentioned above. 3. If an access node is adjacent to FME but not SME, it will be adjacent to new map exit post fusion. 4. If an access node is adjacent to SME but not FME, it will be adjacent to the new map entry node post fusion. """ first_exit = self.first_map_exit first_entry = graph.entry_node(first_exit) second_entry = self.second_map_entry second_exit = graph.exit_node(second_entry) intermediate_nodes = set() for _, _, dst, _, _ in graph.out_edges(first_exit): intermediate_nodes.add(dst) assert isinstance(dst, nodes.AccessNode) # Check if an access node refers to non transient memory, or transient # is used at another location (cannot erase) do_not_erase = set() for node in intermediate_nodes: if sdfg.arrays[node.data].transient is False: do_not_erase.add(node) else: for edge in graph.in_edges(node): if edge.src != first_exit: do_not_erase.add(node) break else: for edge in graph.out_edges(node): if edge.dst != second_entry: do_not_erase.add(node) break # Find permutation between first and second scopes perm = self.find_permutation(first_entry.map, second_entry.map) params_dict = {} for index, param in enumerate(first_entry.map.params): params_dict[param] = second_entry.map.params[perm[index]] # Replaces (in memlets and tasklet) the second scope map # indices with the permuted first map indices. # This works in two passes to avoid problems when e.g., exchanging two # parameters (instead of replacing (j,i) and (i,j) to (j,j) and then # i,i). second_scope = graph.scope_subgraph(second_entry) for firstp, secondp in params_dict.items(): if firstp != secondp: replace(second_scope, secondp, '__' + secondp + '_fused') for firstp, secondp in params_dict.items(): if firstp != secondp: replace(second_scope, '__' + secondp + '_fused', firstp) # Isolate First exit node ############################ edges_to_remove = set() nodes_to_remove = set() for edge in graph.in_edges(first_exit): tree = graph.memlet_tree(edge) access_node = tree.root().edge.dst if access_node not in do_not_erase: out_edges = [ e for e in graph.out_edges(access_node) if e.dst == second_entry ] # In this transformation, there can only be one edge to the # second map assert len(out_edges) == 1 # Get source connector to the second map connector = out_edges[0].dst_conn[3:] new_dsts = [] # Look at the second map entry out-edges to get the new # destinations for e in graph.out_edges(second_entry): if e.src_conn[4:] == connector: new_dsts.append(e) if not new_dsts: # Access node is not used in the second map nodes_to_remove.add(access_node) continue # Add a transient scalar/array self.fuse_nodes(sdfg, graph, edge, new_dsts[0].dst, new_dsts[0].dst_conn, new_dsts[1:]) edges_to_remove.add(edge) # Remove transient node between the two maps nodes_to_remove.add(access_node) else: # The case where intermediate array node cannot be removed # Node will become an output of the second map exit out_e = tree.parent.edge conn = second_exit.next_connector() graph.add_edge( second_exit, 'OUT_' + conn, out_e.dst, out_e.dst_conn, dcpy(out_e.data), ) second_exit.add_out_connector('OUT_' + conn) graph.add_edge(edge.src, edge.src_conn, second_exit, 'IN_' + conn, dcpy(edge.data)) second_exit.add_in_connector('IN_' + conn) edges_to_remove.add(out_e) edges_to_remove.add(edge) # If the second map needs this node, link the connector # that generated this to the place where it is needed, with a # temp transient/scalar for memlet to be generated for out_e in graph.out_edges(second_entry): second_memlet_path = graph.memlet_path(out_e) source_node = second_memlet_path[0].src if source_node == access_node: self.fuse_nodes(sdfg, graph, edge, out_e.dst, out_e.dst_conn) ### # First scope exit is isolated and can now be safely removed for e in edges_to_remove: graph.remove_edge(e) graph.remove_nodes_from(nodes_to_remove) graph.remove_node(first_exit) # Isolate second_entry node ########################### for edge in graph.in_edges(second_entry): tree = graph.memlet_tree(edge) access_node = tree.root().edge.src if access_node in intermediate_nodes: # Already handled above, can be safely removed graph.remove_edge(edge) continue # This is an external input to the second map which will now go # through the first map. conn = first_entry.next_connector() graph.add_edge(edge.src, edge.src_conn, first_entry, 'IN_' + conn, dcpy(edge.data)) first_entry.add_in_connector('IN_' + conn) graph.remove_edge(edge) for out_enode in tree.children: out_e = out_enode.edge graph.add_edge( first_entry, 'OUT_' + conn, out_e.dst, out_e.dst_conn, dcpy(out_e.data), ) graph.remove_edge(out_e) first_entry.add_out_connector('OUT_' + conn) ### # Second node is isolated and can now be safely removed graph.remove_node(second_entry) # Fix scope exit to point to the right map second_exit.map = first_entry.map