def apply(self, graph: SDFGState, sdfg: SDFG):
first_map_entry = graph.entry_node(self.first_map_exit)
intermediate_dnodes = set()
for _, _, node, _, _ in graph.out_edges(self.first_map_exit):
if not isinstance(node, nds.AccessNode):
continue
intermediate_dnodes.add(node)
self._update_in_connectors(graph, intermediate_dnodes)
self._replicate_first_map(sdfg, graph, first_map_entry,
intermediate_dnodes)
graph.remove_nodes_from(
graph.all_nodes_between(first_map_entry, self.first_map_exit)
| {first_map_entry, self.first_map_exit})
for node in graph.nodes():
if not isinstance(node, nds.AccessNode):
continue
if graph.in_degree(node) == 0 and graph.out_degree(node) == 0:
graph.remove_node(node)
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):
"""
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
