def can_be_applied(graph: dace.graph.graph.OrderedMultiDiConnectorGraph,
candidate: Dict[dace.graph.nodes.Node, int],
expr_index: int,
sdfg: dace.SDFG,
strict: bool = False):
# A candidate subgraph matches the map-expansion pattern when it includes
# a N-dimensional map, with N greater than one.
map_entry = graph.nodes()[candidate[MapExpansion._map_entry]]
return map_entry.map.get_param_num() > 1
def match_to_str(clc, graph: dace.graph.graph.OrderedMultiDiConnectorGraph,
candidate: Dict[dace.graph.nodes.Node, int]):
node = graph.nodes()[candidate[clc._match_node]]
return str(node)
def merge_maps(
graph: dace.graph.graph.OrderedMultiDiConnectorGraph,
outer_map_entry: dace.graph.nodes.MapEntry,
outer_map_exit: dace.graph.nodes.MapExit,
inner_map_entry: dace.graph.nodes.MapEntry,
inner_map_exit: dace.graph.nodes.MapExit,
param_merge: Callable[[params, params], params] = lambda p1, p2: p1 + p2,
range_merge: Callable[[ranges, ranges], ranges] = lambda r1, r2: type(r1)
(r1.ranges + r2.ranges)
) -> (dace.graph.nodes.MapEntry, dace.graph.nodes.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 = 'merged_' + 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 = dace.graph.nodes.MapEntry(merged_map)
merged_entry.in_connectors = outer_map_entry.in_connectors
merged_entry.out_connectors = outer_map_entry.out_connectors
merged_exit = dace.graph.nodes.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])
# Redirect inner in edges.
inner_in_edges = graph.out_edges(inner_map_entry)
for edge in graph.edges_between(outer_map_entry, inner_map_entry):
if edge.dst_conn is None: # Empty memlets
out_conn = None
else:
out_conn = 'OUT_' + edge.dst_conn[3:]
inner_edge = [e for e in inner_in_edges if e.src_conn == out_conn][0]
graph.remove_edge(edge)
graph.remove_edge(inner_edge)
graph.add_edge(merged_entry, edge.src_conn, inner_edge.dst,
inner_edge.dst_conn, inner_edge.data)
# Redirect inner out edges.
inner_out_edges = graph.in_edges(inner_map_exit)
for edge in graph.edges_between(inner_map_exit, outer_map_exit):
if edge.src_conn is None: # Empty memlets
in_conn = None
else:
in_conn = 'IN_' + edge.src_conn[4:]
inner_edge = [e for e in inner_out_edges if e.dst_conn == in_conn][0]
graph.remove_edge(edge)
graph.remove_edge(inner_edge)
graph.add_edge(inner_edge.src, inner_edge.src_conn, merged_exit,
edge.dst_conn, inner_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 match_to_str(graph: dace.graph.graph.OrderedMultiDiConnectorGraph,
candidate: Dict[dace.graph.nodes.Node, int]):
map_entry = graph.nodes()[candidate[MapExpansion._map_entry]]
return map_entry.map.label + ': ' + str(map_entry.map.params)