def _components(
subgraph: gr.SubgraphView) -> List[Tuple[nodes.Node, nodes.Node]]:
"""
Returns the list of tuples non-array components in this subgraph.
Each element in the list is a 2 tuple of (input node, output node) of
the component.
"""
graph = (subgraph
if isinstance(subgraph, sd.SDFGState) else subgraph.graph)
schildren = subgraph.scope_children()
ns = [(n, graph.exit_node(n)) if isinstance(n, nodes.EntryNode) else
(n, n) for n in schildren[None]
if isinstance(n, (nodes.CodeNode, nodes.EntryNode))]
return ns
def test_simple_program(self):
@dace.program
def multiply(a: dace.float32[N]):
a *= 2
a *= 3
sdfg = multiply.to_sdfg(strict=True)
for state in sdfg.nodes():
if any(isinstance(node, Tasklet) for node in state.nodes()):
break
else:
raise KeyError('State with tasklet not found')
tasklet_nodes = [n for n in state.nodes() if isinstance(n, Tasklet)]
with self.assertRaises(ValueError):
nest_state_subgraph(sdfg, state,
SubgraphView(state, tasklet_nodes))
nest_state_subgraph(sdfg, state, SubgraphView(state,
[tasklet_nodes[0]]))
sdfg.validate()
nest_state_subgraph(sdfg, state, SubgraphView(state,
[tasklet_nodes[1]]))
sdfg.validate()
def test_p1():
N.set(20)
M.set(30)
O.set(50)
P.set(40)
Q.set(42)
R.set(25)
sdfg = subgraph_fusion_parallel.to_sdfg()
sdfg.coarsen_dataflow()
state = sdfg.nodes()[0]
A = np.random.rand(N.get()).astype(np.float64)
B = np.random.rand(M.get()).astype(np.float64)
C = np.random.rand(O.get()).astype(np.float64)
D = np.random.rand(M.get()).astype(np.float64)
E = np.random.rand(N.get()).astype(np.float64)
F = np.random.rand(P.get()).astype(np.float64)
G = np.random.rand(M.get()).astype(np.float64)
H = np.random.rand(P.get()).astype(np.float64)
I = np.random.rand(N.get()).astype(np.float64)
J = np.random.rand(R.get()).astype(np.float64)
X = np.random.rand(N.get()).astype(np.float64)
Y = np.random.rand(M.get()).astype(np.float64)
Z = np.random.rand(P.get()).astype(np.float64)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C, D=D, E=E, F=F, G=G, H=H, I=I, J=J, X=X, Y=Y, Z=Z,\
N=N, M=M, O=O, P=P, R=R,Q=Q)
del csdfg
subgraph = SubgraphView(state, [node for node in state.nodes()])
expansion = MultiExpansion(subgraph)
fusion = SubgraphFusion(subgraph)
me = MultiExpansion(subgraph)
assert me.can_be_applied(sdfg, subgraph)
me.apply(sdfg)
sf = SubgraphFusion(subgraph)
assert sf.can_be_applied(sdfg, subgraph)
sf.apply(sdfg)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C, D=D, E=E, F=F, G=G, H=H, I=I, J=J, X=X, Y=Y, Z=Z,\
N=N, M=M, O=O, P=P, R=R,Q=Q)
print("PASS")
def subgraph_from_maps(sdfg, graph, map_entries, scope_children=None):
"""
Given a list of map entries in a single graph,
return a subgraph view that includes all nodes
inside these maps as well as map entries and exits
as well as adjacent nodes.
"""
if not scope_children:
scope_children = graph.scope_children()
nodes = set()
for map_entry in map_entries:
nodes |= set(scope_children[map_entry])
nodes |= set(e.dst for e in graph.out_edges(graph.exit_node(map_entry)))
nodes |= set(e.src for e in graph.in_edges(map_entry))
nodes.add(map_entry)
return SubgraphView(graph, list(nodes))
def test_p1():
N.set(20)
M.set(30)
O.set(50)
P.set(40)
Q.set(42)
R.set(25)
sdfg = test_program.to_sdfg()
sdfg.apply_strict_transformations()
state = sdfg.nodes()[0]
A = np.random.rand(N.get()).astype(np.float64)
B = np.random.rand(M.get()).astype(np.float64)
C = np.random.rand(O.get()).astype(np.float64)
D = np.random.rand(M.get()).astype(np.float64)
E = np.random.rand(N.get()).astype(np.float64)
F = np.random.rand(P.get()).astype(np.float64)
G = np.random.rand(M.get()).astype(np.float64)
H = np.random.rand(P.get()).astype(np.float64)
I = np.random.rand(N.get()).astype(np.float64)
J = np.random.rand(R.get()).astype(np.float64)
X = np.random.rand(N.get()).astype(np.float64)
Y = np.random.rand(M.get()).astype(np.float64)
Z = np.random.rand(P.get()).astype(np.float64)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C, D=D, E=E, F=F, G=G, H=H, I=I, J=J, X=X, Y=Y, Z=Z,\
N=N, M=M, O=O, P=P, R=R,Q=Q)
subgraph = SubgraphView(state, [node for node in state.nodes()])
expansion = MultiExpansion()
fusion = SubgraphFusion()
assert MultiExpansion.match(sdfg, subgraph)
expansion.apply(sdfg, subgraph)
assert SubgraphFusion.match(sdfg, subgraph)
fusion.apply(sdfg, subgraph)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C, D=D, E=E, F=F, G=G, H=H, I=I, J=J, X=X, Y=Y, Z=Z,\
N=N, M=M, O=O, P=P, R=R,Q=Q)
print("PASS")
def can_be_applied(self, sdfg: SDFG, subgraph: SubgraphView) -> bool:
graph = subgraph.graph
if self.allow_expansion == True:
subgraph_fusion = SubgraphFusion(subgraph)
if subgraph_fusion.can_be_applied(sdfg, subgraph):
# try w/o copy first
return True
expansion = MultiExpansion(subgraph)
expansion.permutation_only = not self.expansion_split
if expansion.can_be_applied(sdfg, subgraph):
# deepcopy
graph_indices = [
i for (i, n) in enumerate(graph.nodes()) if n in subgraph
]
sdfg_copy = copy.deepcopy(sdfg)
graph_copy = sdfg_copy.nodes()[sdfg.nodes().index(graph)]
subgraph_copy = SubgraphView(
graph_copy, [graph_copy.nodes()[i] for i in graph_indices])
expansion.sdfg_id = sdfg_copy.sdfg_id
##sdfg_copy.apply_transformations(MultiExpansion, states=[graph])
#expansion = MultiExpansion(subgraph_copy)
expansion.apply(sdfg_copy)
subgraph_fusion = SubgraphFusion(subgraph_copy)
if subgraph_fusion.can_be_applied(sdfg_copy, subgraph_copy):
return True
stencil_tiling = StencilTiling(subgraph_copy)
if self.allow_tiling and stencil_tiling.can_be_applied(
sdfg_copy, subgraph_copy):
return True
else:
subgraph_fusion = SubgraphFusion(subgraph)
if subgraph_fusion.can_be_applied(sdfg, subgraph):
return True
if self.allow_tiling == True:
stencil_tiling = StencilTiling(subgraph)
if stencil_tiling.can_be_applied(sdfg, subgraph):
return True
return False
def test_offsets_array():
sdfg = dace.SDFG('mapfission_offsets2')
sdfg.add_array('A', [20], dace.float64)
sdfg.add_array('interim', [1], dace.float64, transient=True)
state = sdfg.add_state()
me, mx = state.add_map('outer', dict(i='10:20'))
t1 = state.add_tasklet('addone', {'a'}, {'b'}, 'b = a + 1')
interim = state.add_access('interim')
t2 = state.add_tasklet('addtwo', {'a'}, {'b'}, 'b = a + 2')
aread = state.add_read('A')
awrite = state.add_write('A')
state.add_memlet_path(aread, me, t1, dst_conn='a', memlet=dace.Memlet.simple('A', 'i'))
state.add_edge(t1, 'b', interim, None, dace.Memlet.simple('interim', '0'))
state.add_edge(interim, None, t2, 'a', dace.Memlet.simple('interim', '0'))
state.add_memlet_path(t2, mx, awrite, src_conn='b', memlet=dace.Memlet.simple('A', 'i'))
sdfg.apply_transformations(MapFission)
dace.propagate_memlets_sdfg(sdfg)
sdfg.validate()
# Test
A = np.random.rand(20)
expected = A.copy()
expected[10:] += 3
A_cpy = A.copy()
csdfg = sdfg.compile()
csdfg(A=A_cpy)
del csdfg
print(np.linalg.norm(A_cpy))
print(np.linalg.norm(expected))
assert (np.allclose(A_cpy, expected))
subgraph = SubgraphView(sdfg.nodes()[0], sdfg.nodes()[0].nodes())
sf = SubgraphFusion()
sf.setup_match(subgraph)
assert sf.can_be_applied(sdfg, subgraph)
fusion(sdfg, sdfg.nodes()[0], None)
A_cpy = A.copy()
csdfg = sdfg.compile()
csdfg(A=A_cpy)
assert (np.allclose(A_cpy, expected))
def get_actions(actions, graph, match):
subgraph_node_ids = match.subgraph.values()
subgraph_nodes = [graph.nodes()[nid] for nid in subgraph_node_ids]
for node in subgraph_nodes:
version = 0
while (node, type(match).__name__, match.expr_index,
version) in actions.keys():
version += 1
actions[(node, type(match).__name__, match.expr_index,
version)] = match
subgraph = SubgraphView(graph, subgraph_nodes)
for edge in subgraph.edges():
version = 0
while (edge, type(match).__name__, match.expr_index,
version) in actions.keys():
version += 1
actions[(edge, type(match).__name__, match.expr_index,
version)] = match
return actions
def state_fission(sdfg: SDFG, subgraph: graph.SubgraphView) -> SDFGState:
'''
Given a subgraph, adds a new SDFG state before the state that contains it,
removes the subgraph from the original state, and connects the two states.
:param subgraph: the subgraph to remove.
:return: the newly created SDFG state.
'''
state: SDFGState = subgraph.graph
newstate = sdfg.add_state_before(state)
# Save edges before removing nodes
orig_edges = subgraph.edges()
# Mark boundary access nodes to keep after fission
nodes_to_remove = set(subgraph.nodes())
boundary_nodes = [
n for n in subgraph.nodes()
if len(state.out_edges(n)) > len(subgraph.out_edges(n))
] + [
n for n in subgraph.nodes()
if len(state.in_edges(n)) > len(subgraph.in_edges(n))
]
# Make dictionary of nodes to add to new state
new_nodes = {n: n for n in subgraph.nodes()}
new_nodes.update({b: copy.deepcopy(b) for b in boundary_nodes})
nodes_to_remove -= set(boundary_nodes)
state.remove_nodes_from(nodes_to_remove)
for n in new_nodes.values():
if isinstance(n, nodes.NestedSDFG):
# Set the new parent state
n.sdfg.parent = newstate
newstate.add_nodes_from(new_nodes.values())
for e in orig_edges:
newstate.add_edge(new_nodes[e.src], e.src_conn, new_nodes[e.dst],
e.dst_conn, e.data)
return newstate
def _test_quantitatively(sdfg, graph):
A = np.random.rand(N.get(), M.get(), O.get()).astype(np.float64)
B = np.random.rand(N.get(), M.get(), O.get()).astype(np.float64)
C1 = np.zeros([N.get(), M.get(), O.get()], dtype=np.float64)
C2 = np.zeros([N.get(), M.get(), O.get()], dtype=np.float64)
sdfg.validate()
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C1, N=N, M=M, O=O)
del csdfg
subgraph = SubgraphView(graph, graph.nodes())
sf = SubgraphFusion(subgraph)
assert sf.can_be_applied(sdfg, subgraph)
fusion(sdfg, graph)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C2, N=N, M=M, O=O)
del csdfg
assert np.allclose(C1, C2)
print('PASS')
def _test_quantitatively(sdfg, graph):
A = np.random.rand(N.get()).astype(np.float64)
B = np.random.rand(M.get()).astype(np.float64)
C = np.random.rand(O.get()).astype(np.float64)
out1_base = np.ndarray((N.get(), M.get()), np.float64)
out2_base = np.ndarray((1), np.float64)
out3_base = np.ndarray((N.get(), M.get(), O.get()), np.float64)
out1 = np.ndarray((N.get(), M.get()), np.float64)
out2 = np.ndarray((1), np.float64)
out3 = np.ndarray((N.get(), M.get(), O.get()), np.float64)
csdfg = sdfg.compile()
csdfg(A=A,
B=B,
C=C,
out1=out1_base,
out2=out2_base,
out3=out3_base,
N=N,
M=M,
O=O)
del csdfg
expand_reduce(sdfg, graph)
expand_maps(sdfg, graph)
subgraph = SubgraphView(graph, [node for node in graph.nodes()])
sf = SubgraphFusion()
sf.setup_match(subgraph)
assert sf.can_be_applied(sdfg, subgraph) == True
fusion(sdfg, graph)
sdfg.validate()
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C, out1=out1, out2=out2, out3=out3, N=N, M=M, O=O)
del csdfg
assert np.allclose(out1, out1_base)
assert np.allclose(out2, out2_base)
assert np.allclose(out3, out3_base)
print('PASS')
def test_quantitatively(sdfg):
graph = sdfg.nodes()[0]
A = np.random.rand(N.get()).astype(np.float64)
B = np.random.rand(N.get()).astype(np.float64)
C1 = np.random.rand(N.get()).astype(np.float64)
C2 = np.random.rand(N.get()).astype(np.float64)
D1 = np.random.rand(N.get()).astype(np.float64)
D2 = np.random.rand(N.get()).astype(np.float64)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C1, D=D1, N=N)
subgraph = SubgraphView(graph, [node for node in graph.nodes()])
assert MultiExpansion.can_be_applied(sdfg, subgraph) == True
MultiExpansion(subgraph).apply(sdfg)
assert SubgraphFusion.can_be_applied(sdfg, subgraph) == True
SubgraphFusion(subgraph).apply(sdfg)
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C2, D=D2, N=N)
assert np.allclose(C1, C2)
assert np.allclose(D1, D2)
def test_p1():
sdfg = disjoint_test_1.to_sdfg()
sdfg.simplify()
state = sdfg.nodes()[0]
assert len(sdfg.nodes()) == 1
A = np.random.rand(M.get(), 2).astype(np.float64)
A1 = A.copy()
A2 = A.copy()
csdfg = sdfg.compile()
csdfg(A=A1, N=N, M=M)
del csdfg
subgraph = SubgraphView(state, state.nodes())
sf = SubgraphFusion(subgraph)
assert sf.can_be_applied(sdfg, subgraph)
sf.apply(sdfg)
csdfg = sdfg.compile()
csdfg(A=A2, M=M)
del csdfg
assert np.allclose(A1, A2)
def test_quantitatively(sdfg, graph):
A = np.random.rand(N.get()).astype(np.float64)
B = np.random.rand(M.get()).astype(np.float64)
C = np.random.rand(O.get()).astype(np.float64)
out1_base = np.ndarray((N.get(), M.get()), np.float64)
out2_base = np.ndarray((1), np.float64)
out3_base = np.ndarray((N.get(), M.get(), O.get()), np.float64)
out1 = np.ndarray((N.get(), M.get()), np.float64)
out2 = np.ndarray((1), np.float64)
out3 = np.ndarray((N.get(), M.get(), O.get()), np.float64)
csdfg = sdfg.compile()
csdfg(A=A,
B=B,
C=C,
out1=out1_base,
out2=out2_base,
out3=out3_base,
N=N,
M=M,
O=O)
expand_reduce(sdfg, graph)
expand_maps(sdfg, graph)
sgf = SubgraphFusion()
matcher = sgf.match(sdfg,
SubgraphView(graph, [node for node in graph.nodes()]))
assert matcher == True
fusion(sdfg, graph)
sdfg.validate()
csdfg = sdfg.compile()
csdfg(A=A, B=B, C=C, out1=out1, out2=out2, out3=out3, N=N, M=M, O=O)
assert np.allclose(out1, out1_base)
assert np.allclose(out2, out2_base)
assert np.allclose(out3, out3_base)
print('PASS')
def match(sdfg: SDFG, subgraph: SubgraphView) -> bool:
### get lowest scope maps of subgraph
# grab first node and see whether all nodes are in the same graph
# (or nested sdfgs therein)
graph = subgraph.graph
for node in subgraph.nodes():
if node not in graph.nodes():
return False
# next, get all the maps
maps = helpers.get_highest_scope_maps(sdfg, graph, subgraph)
brng = helpers.common_map_base_ranges(maps)
# if leq than one map found -> fail
if len(maps) <= 1:
return False
# see whether they have common parameters; if not -> fail
if len(brng) == 0:
return False
return True
def test_simple_sdfg_program(self):
sdfg, state, t, me, mx = create_sdfg()
nest_state_subgraph(sdfg, state, SubgraphView(state, state.nodes()))
sdfg.validate()
def apply(self, sdfg):
state: SDFGState = sdfg.nodes()[self.state_id]
nsdfg_node = state.nodes()[self.subgraph[InlineSDFG._nested_sdfg]]
nsdfg: SDFG = nsdfg_node.sdfg
nstate: SDFGState = nsdfg.nodes()[0]
nsdfg_scope_entry = state.entry_node(nsdfg_node)
nsdfg_scope_exit = (state.exit_node(nsdfg_scope_entry)
if nsdfg_scope_entry is not None else None)
#######################################################
# Collect and update top-level SDFG metadata
# Global/init/exit code
for loc, code in nsdfg.global_code.items():
sdfg.append_global_code(code.code, loc)
for loc, code in nsdfg.init_code.items():
sdfg.append_init_code(code.code, loc)
for loc, code in nsdfg.exit_code.items():
sdfg.append_exit_code(code.code, loc)
# Constants
for cstname, cstval in nsdfg.constants.items():
if cstname in sdfg.constants:
if cstval != sdfg.constants[cstname]:
warnings.warn('Constant value mismatch for "%s" while '
'inlining SDFG. Inner = %s != %s = outer' %
(cstname, cstval, sdfg.constants[cstname]))
else:
sdfg.add_constant(cstname, cstval)
# Find original source/destination edges (there is only one edge per
# connector, according to match)
inputs: Dict[str, MultiConnectorEdge] = {}
outputs: Dict[str, MultiConnectorEdge] = {}
input_set: Dict[str, str] = {}
output_set: Dict[str, str] = {}
for e in state.in_edges(nsdfg_node):
inputs[e.dst_conn] = e
input_set[e.data.data] = e.dst_conn
for e in state.out_edges(nsdfg_node):
outputs[e.src_conn] = e
output_set[e.data.data] = e.src_conn
# All transients become transients of the parent (if data already
# exists, find new name)
# Mapping from nested transient name to top-level name
transients: Dict[str, str] = {}
for node in nstate.nodes():
if isinstance(node, nodes.AccessNode):
datadesc = nsdfg.arrays[node.data]
if node.data not in transients and datadesc.transient:
name = sdfg.add_datadesc('%s_%s' %
(nsdfg.label, node.data),
datadesc,
find_new_name=True)
transients[node.data] = name
# All transients of edges between code nodes are also added to parent
for edge in nstate.edges():
if (isinstance(edge.src, nodes.CodeNode)
and isinstance(edge.dst, nodes.CodeNode)):
datadesc = nsdfg.arrays[edge.data.data]
if edge.data.data not in transients and datadesc.transient:
name = sdfg.add_datadesc('%s_%s' %
(nsdfg.label, edge.data.data),
datadesc,
find_new_name=True)
transients[edge.data.data] = name
# Collect nodes to add to top-level graph
new_incoming_edges: Dict[nodes.Node, MultiConnectorEdge] = {}
new_outgoing_edges: Dict[nodes.Node, MultiConnectorEdge] = {}
source_accesses = set()
sink_accesses = set()
for node in nstate.source_nodes():
if (isinstance(node, nodes.AccessNode)
and node.data not in transients):
new_incoming_edges[node] = inputs[node.data]
source_accesses.add(node)
for node in nstate.sink_nodes():
if (isinstance(node, nodes.AccessNode)
and node.data not in transients):
new_outgoing_edges[node] = outputs[node.data]
sink_accesses.add(node)
#######################################################
# Add nested SDFG into top-level SDFG
# Add nested nodes into original state
subgraph = SubgraphView(nstate, [
n for n in nstate.nodes()
if n not in (source_accesses | sink_accesses)
])
state.add_nodes_from(subgraph.nodes())
for edge in subgraph.edges():
state.add_edge(edge.src, edge.src_conn, edge.dst, edge.dst_conn,
edge.data)
#######################################################
# Replace data on inlined SDFG nodes/edges
# Replace symbols using invocation symbol mapping
# Two-step replacement (N -> __dacesym_N --> map[N]) to avoid clashes
for symname, symvalue in nsdfg_node.symbol_mapping.items():
if str(symname) != str(symvalue):
nsdfg.replace(symname, '__dacesym_' + symname)
for symname, symvalue in nsdfg_node.symbol_mapping.items():
if str(symname) != str(symvalue):
nsdfg.replace('__dacesym_' + symname, symvalue)
# Replace data names with their top-level counterparts
repldict = {}
repldict.update(transients)
repldict.update({
k: v.data.data
for k, v in itertools.chain(inputs.items(), outputs.items())
})
for node in subgraph.nodes():
if isinstance(node, nodes.AccessNode) and node.data in repldict:
node.data = repldict[node.data]
for edge in subgraph.edges():
if edge.data.data in repldict:
edge.data.data = repldict[edge.data.data]
#######################################################
# Reconnect inlined SDFG
# If a source/sink node is one of the inputs/outputs, reconnect it,
# replacing memlets in outgoing/incoming paths
modified_edges = set()
modified_edges |= self._modify_memlet_path(new_incoming_edges, nstate,
state, True)
modified_edges |= self._modify_memlet_path(new_outgoing_edges, nstate,
state, False)
# Modify all other internal edges pertaining to input/output nodes
for node in subgraph.nodes():
if isinstance(node, nodes.AccessNode):
if node.data in input_set or node.data in output_set:
if node.data in input_set:
outer_edge = inputs[input_set[node.data]]
else:
outer_edge = outputs[output_set[node.data]]
for edge in state.all_edges(node):
if (edge not in modified_edges
and edge.data.data == node.data):
for e in state.memlet_tree(edge):
if e.data.data == node.data:
e._data = helpers.unsqueeze_memlet(
e.data, outer_edge.data)
# If source/sink node is not connected to a source/destination access
# node, and the nested SDFG is in a scope, connect to scope with empty
# memlets
if nsdfg_scope_entry is not None:
for node in subgraph.nodes():
if state.in_degree(node) == 0:
state.add_edge(nsdfg_scope_entry, None, node, None,
Memlet())
if state.out_degree(node) == 0:
state.add_edge(node, None, nsdfg_scope_exit, None,
Memlet())
# Replace nested SDFG parents with new SDFG
for node in nstate.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent = state
node.sdfg.parent_sdfg = sdfg
node.sdfg.parent_nsdfg_node = node
# Remove all unused external inputs/output memlet paths, as well as
# resulting isolated nodes
removed_in_edges = self._remove_edge_path(state,
inputs,
set(inputs.keys()) -
source_accesses,
reverse=True)
removed_out_edges = self._remove_edge_path(state,
outputs,
set(outputs.keys()) -
sink_accesses,
reverse=False)
# Re-add in/out edges to first/last nodes in subgraph
order = [
x for x in nx.topological_sort(nstate._nx)
if isinstance(x, nodes.AccessNode)
]
for edge in removed_in_edges:
# Find first access node that refers to this edge
node = next(n for n in order if n.data == edge.data.data)
state.add_edge(edge.src, edge.src_conn, node, edge.dst_conn,
edge.data)
for edge in removed_out_edges:
# Find last access node that refers to this edge
node = next(n for n in reversed(order) if n.data == edge.data.data)
state.add_edge(node, edge.src_conn, edge.dst, edge.dst_conn,
edge.data)
#######################################################
# Remove nested SDFG node
state.remove_node(nsdfg_node)
def subgraph_view(self, sdfg: SDFG) -> SubgraphView:
graph = sdfg.sdfg_list[self.sdfg_id]
if self.state_id != -1:
graph = graph.node(self.state_id)
return SubgraphView(graph, [graph.node(idx) for idx in self.subgraph])
def get_transformations(sdfg_json, selected_elements):
# We lazy import DaCe, not to break cyclic imports, but to avoid any large
# delays when booting in daemon mode.
from dace.transformation.optimizer import SDFGOptimizer
from dace.sdfg.graph import SubgraphView
old_meta = utils.disable_save_metadata()
loaded = utils.load_sdfg_from_json(sdfg_json)
if loaded['error'] is not None:
return loaded['error']
sdfg = loaded['sdfg']
optimizer = SDFGOptimizer(sdfg)
matches = optimizer.get_pattern_matches()
transformations = []
docstrings = {}
for transformation in matches:
transformations.append(transformation.to_json())
docstrings[type(transformation).__name__] = transformation.__doc__
selected_states = [
utils.sdfg_find_state_from_element(sdfg, n) for n in selected_elements
if n['type'] == 'state'
]
selected_nodes = [
utils.sdfg_find_node_from_element(sdfg, n) for n in selected_elements
if n['type'] == 'node'
]
selected_sdfg_ids = list(set(elem['sdfgId'] for elem in selected_elements))
selected_sdfg = sdfg
if len(selected_sdfg_ids) > 1:
return {
'transformations': transformations,
'docstrings': docstrings,
'warnings': 'More than one SDFG selected, ignoring subgraph',
}
elif len(selected_sdfg_ids) == 1:
selected_sdfg = sdfg.sdfg_list[selected_sdfg_ids[0]]
subgraph = None
if len(selected_states) > 0:
subgraph = SubgraphView(selected_sdfg, selected_states)
else:
violated = False
state = None
for node in selected_nodes:
if state is None:
state = node.state
elif state != node.state:
violated = True
break
if not violated and state is not None:
subgraph = SubgraphView(state, selected_nodes)
if subgraph is not None:
extensions = SubgraphTransformation.extensions()
for xform in extensions:
xform_data = extensions[xform]
if ('singlestate' in xform_data and xform_data['singlestate']
and len(selected_states) > 0):
continue
xform_obj = xform(subgraph)
if xform_obj.can_be_applied(selected_sdfg, subgraph):
transformations.append(xform_obj.to_json())
docstrings[xform.__name__] = xform_obj.__doc__
utils.restore_save_metadata(old_meta)
return {
'transformations': transformations,
'docstrings': docstrings,
}
def apply(self, sdfg: SDFG):
graph = sdfg.node(self.state_id)
map_exit = graph.node(self.subgraph[AccumulateTransient.map_exit])
outer_map_exit = graph.node(
self.subgraph[AccumulateTransient.outer_map_exit])
# Avoid import loop
from dace.transformation.dataflow.local_storage import OutLocalStorage
array_identity_dict = self.array_identity_dict
# Choose array
array = self.array
if array is not None and len(array) != 0:
array_identity_dict[array] = self.identity
elif ((array is None or len(array) == 0)
and len(array_identity_dict) == 0):
array = next(e.data.data
for e in graph.edges_between(map_exit, outer_map_exit)
if e.data.wcr is not None)
array_identity_dict[array] = self.identity
transients: Dict[str, Any] = {}
for array, identity in array_identity_dict.items():
data_node: nodes.AccessNode = OutLocalStorage.apply_to(
sdfg,
dict(array=array, prefix=self.prefix),
verify=False,
save=False,
node_a=map_exit,
node_b=outer_map_exit)
transients[data_node.data] = identity
if identity is None:
warnings.warn(
'AccumulateTransient did not properly initialize '
'newly-created transient!')
return
sdfg_state: SDFGState = sdfg.node(self.state_id)
map_entry = sdfg_state.entry_node(map_exit)
nested_sdfg: nodes.NestedSDFG = nest_state_subgraph(
sdfg=sdfg,
state=sdfg_state,
subgraph=SubgraphView(
sdfg_state, {map_entry, map_exit}
| sdfg_state.all_nodes_between(map_entry, map_exit)))
nested_sdfg_state: SDFGState = nested_sdfg.sdfg.nodes()[0]
init_state = nested_sdfg.sdfg.add_state_before(nested_sdfg_state)
for data_name, identity in transients.items():
temp_array: Array = sdfg.arrays[data_name]
init_state.add_mapped_tasklet(
name='acctrans_init',
map_ranges={
'_o%d' % i: '0:%s' % symbolic.symstr(d)
for i, d in enumerate(temp_array.shape)
},
inputs={},
code='out = %s' % identity,
outputs={
'out':
dace.Memlet.simple(
data=data_name,
subset_str=','.join([
'_o%d' % i for i, _ in enumerate(temp_array.shape)
]))
},
external_edges=True)
# TODO: use trivial map elimintation here when it will be merged to remove map if it has trivial ranges
return nested_sdfg
def nest_state_subgraph(sdfg: SDFG,
state: SDFGState,
subgraph: SubgraphView,
name: Optional[str] = None,
full_data: bool = False) -> nodes.NestedSDFG:
""" Turns a state subgraph into a nested SDFG. Operates in-place.
:param sdfg: The SDFG containing the state subgraph.
:param state: The state containing the subgraph.
:param subgraph: Subgraph to nest.
:param name: An optional name for the nested SDFG.
:param full_data: If True, nests entire input/output data.
:return: The nested SDFG node.
:raise KeyError: Some or all nodes in the subgraph are not located in
this state, or the state does not belong to the given
SDFG.
:raise ValueError: The subgraph is contained in more than one scope.
"""
if state.parent != sdfg:
raise KeyError('State does not belong to given SDFG')
if subgraph.graph != state:
raise KeyError('Subgraph does not belong to given state')
# Find the top-level scope
scope_tree = state.scope_tree()
scope_dict = state.scope_dict()
scope_dict_children = state.scope_dict(True)
top_scopenode = -1 # Initialized to -1 since "None" already means top-level
for node in subgraph.nodes():
if node not in scope_dict:
raise KeyError('Node not found in state')
# If scope entry/exit, ensure entire scope is in subgraph
if isinstance(node, nodes.EntryNode):
scope_nodes = scope_dict_children[node]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (entry)')
elif isinstance(node, nodes.ExitNode):
entry = state.entry_node(node)
scope_nodes = scope_dict_children[entry] + [entry]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (exit)')
scope_node = scope_dict[node]
if scope_node not in subgraph.nodes():
if top_scopenode != -1 and top_scopenode != scope_node:
raise ValueError(
'Subgraph is contained in more than one scope')
top_scopenode = scope_node
scope = scope_tree[top_scopenode]
###
# Collect inputs and outputs of the nested SDFG
inputs: List[MultiConnectorEdge] = []
outputs: List[MultiConnectorEdge] = []
for node in subgraph.source_nodes():
inputs.extend(state.in_edges(node))
for node in subgraph.sink_nodes():
outputs.extend(state.out_edges(node))
# Collect transients not used outside of subgraph (will be removed of
# top-level graph)
data_in_subgraph = set(n.data for n in subgraph.nodes()
if isinstance(n, nodes.AccessNode))
# Find other occurrences in SDFG
other_nodes = set(
n.data for s in sdfg.nodes() for n in s.nodes()
if isinstance(n, nodes.AccessNode) and n not in subgraph.nodes())
subgraph_transients = set()
for data in data_in_subgraph:
datadesc = sdfg.arrays[data]
if datadesc.transient and data not in other_nodes:
subgraph_transients.add(data)
# All transients of edges between code nodes are also added to nested graph
for edge in subgraph.edges():
if (isinstance(edge.src, nodes.CodeNode)
and isinstance(edge.dst, nodes.CodeNode)):
subgraph_transients.add(edge.data.data)
# Collect data used in access nodes within subgraph (will be referenced in
# full upon nesting)
input_arrays = set()
output_arrays = set()
for node in subgraph.nodes():
if (isinstance(node, nodes.AccessNode)
and node.data not in subgraph_transients):
if state.out_degree(node) > 0:
input_arrays.add(node.data)
if state.in_degree(node) > 0:
output_arrays.add(node.data)
# Create the nested SDFG
nsdfg = SDFG(name or 'nested_' + state.label)
# Transients are added to the nested graph as-is
for name in subgraph_transients:
nsdfg.add_datadesc(name, sdfg.arrays[name])
# Input/output data that are not source/sink nodes are added to the graph
# as non-transients
for name in (input_arrays | output_arrays):
datadesc = copy.deepcopy(sdfg.arrays[name])
datadesc.transient = False
nsdfg.add_datadesc(name, datadesc)
# Connected source/sink nodes outside subgraph become global data
# descriptors in nested SDFG
input_names = []
output_names = []
for edge in inputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = '__in_' + edge.data.data
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
input_names.append(
nsdfg.add_datadesc(name, datadesc, find_new_name=True))
for edge in outputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = '__out_' + edge.data.data
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
output_names.append(
nsdfg.add_datadesc(name, datadesc, find_new_name=True))
###################
# Add scope symbols to the nested SDFG
for v in scope.defined_vars:
if v in sdfg.symbols:
sym = sdfg.symbols[v]
nsdfg.add_symbol(v, sym.dtype)
# Create nested state
nstate = nsdfg.add_state()
# Add subgraph nodes and edges to nested state
nstate.add_nodes_from(subgraph.nodes())
for e in subgraph.edges():
nstate.add_edge(e.src, e.src_conn, e.dst, e.dst_conn, e.data)
# Modify nested SDFG parents in subgraph
for node in subgraph.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent = nstate
node.sdfg.parent_sdfg = nsdfg
# Add access nodes and edges as necessary
edges_to_offset = []
for name, edge in zip(input_names, inputs):
node = nstate.add_read(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge,
nstate.add_edge(node, None, edge.dst,
edge.dst_conn, new_edge)))
for name, edge in zip(output_names, outputs):
node = nstate.add_write(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge,
nstate.add_edge(edge.src, edge.src_conn, node,
None, new_edge)))
# Offset memlet paths inside nested SDFG according to subsets
for original_edge, new_edge in edges_to_offset:
for edge in nstate.memlet_tree(new_edge):
edge.data.data = new_edge.data.data
if not full_data:
edge.data.subset.offset(original_edge.data.subset, True)
# Add nested SDFG node to the input state
nested_sdfg = state.add_nested_sdfg(nsdfg, None,
set(input_names) | input_arrays,
set(output_names) | output_arrays)
# Reconnect memlets to nested SDFG
for name, edge in zip(input_names, inputs):
if full_data:
data = Memlet.from_array(edge.data.data,
sdfg.arrays[edge.data.data])
else:
data = edge.data
state.add_edge(edge.src, edge.src_conn, nested_sdfg, name, data)
for name, edge in zip(output_names, outputs):
if full_data:
data = Memlet.from_array(edge.data.data,
sdfg.arrays[edge.data.data])
else:
data = edge.data
state.add_edge(nested_sdfg, name, edge.dst, edge.dst_conn, data)
# Connect access nodes to internal input/output data as necessary
entry = scope.entry
exit = scope.exit
for name in input_arrays:
node = state.add_read(name)
if entry is not None:
state.add_nedge(entry, node, EmptyMemlet())
state.add_edge(node, None, nested_sdfg, name,
Memlet.from_array(name, sdfg.arrays[name]))
for name in output_arrays:
node = state.add_write(name)
if exit is not None:
state.add_nedge(node, exit, EmptyMemlet())
state.add_edge(nested_sdfg, name, node, None,
Memlet.from_array(name, sdfg.arrays[name]))
# Remove subgraph nodes from graph
state.remove_nodes_from(subgraph.nodes())
# Remove subgraph transients from top-level graph
for transient in subgraph_transients:
del sdfg.arrays[transient]
return nested_sdfg
def can_be_applied(sdfg: SDFG, subgraph: SubgraphView) -> bool:
'''
Fusible if
1. Maps have the same access sets and ranges in order
2. Any nodes in between two maps are AccessNodes only, without WCR
There is at most one AccessNode only on a path between two maps,
no other nodes are allowed
3. The exiting memlets' subsets to an intermediate edge must cover
the respective incoming memlets' subset into the next map
'''
# get graph
graph = subgraph.graph
for node in subgraph.nodes():
if node not in graph.nodes():
return False
# next, get all the maps
map_entries = helpers.get_highest_scope_maps(sdfg, graph, subgraph)
map_exits = [graph.exit_node(map_entry) for map_entry in map_entries]
maps = [map_entry.map for map_entry in map_entries]
# 1. check whether all map ranges and indices are the same
if len(maps) <= 1:
return False
base_map = maps[0]
for map in maps:
if map.get_param_num() != base_map.get_param_num():
return False
if not all(
[p1 == p2 for (p1, p2) in zip(map.params, base_map.params)]):
return False
if not map.range == base_map.range:
return False
# 1.1 check whether all map entries have the same schedule
schedule = map_entries[0].schedule
if not all([entry.schedule == schedule for entry in map_entries]):
return False
# 2. check intermediate feasiblility
# see map_fusion.py for similar checks
# we are being more relaxed here
# 2.1 do some preparation work first:
# calculate all out_nodes and intermediate_nodes
# definition see in apply()
intermediate_nodes = set()
out_nodes = set()
for map_entry, map_exit in zip(map_entries, map_exits):
for edge in graph.out_edges(map_exit):
current_node = edge.dst
if len(graph.out_edges(current_node)) == 0:
out_nodes.add(current_node)
else:
for dst_edge in graph.out_edges(current_node):
if dst_edge.dst in map_entries:
intermediate_nodes.add(current_node)
else:
out_nodes.add(current_node)
# 2.2 topological feasibility:
# For each intermediate and out node: must never reach any map
# entry if it is not connected to map entry immediately
visited = set()
# for memoization purposes
def visit_descendants(graph, node, visited, map_entries):
# if we have already been at this node
if node in visited:
return True
# not necessary to add if there aren't any other in connections
if len(graph.in_edges(node)) > 1:
visited.add(node)
for oedge in graph.out_edges(node):
if not visit_descendants(graph, oedge.dst, visited,
map_entries):
return False
return True
for node in intermediate_nodes | out_nodes:
# these nodes must not lead to a map entry
nodes_to_check = set()
for oedge in graph.out_edges(node):
if oedge.dst not in map_entries:
nodes_to_check.add(oedge.dst)
for forbidden_node in nodes_to_check:
if not visit_descendants(graph, forbidden_node, visited,
map_entries):
return False
# 2.3 memlet feasibility
# For each intermediate node, look at whether inner adjacent
# memlets of the exiting map cover inner adjacent memlets
# of the next entering map.
# We also check for any WCRs on the fly.
for node in intermediate_nodes:
upper_subsets = set()
lower_subsets = set()
# First, determine which dimensions of the memlet ranges
# change with the map, we do not need to care about the other dimensions.
total_dims = len(sdfg.data(node.data).shape)
dims_to_discard = SubgraphFusion.get_invariant_dimensions(
sdfg, graph, map_entries, map_exits, node)
# find upper_subsets
for in_edge in graph.in_edges(node):
# first check for WCRs
if in_edge.data.wcr:
return False
if in_edge.src in map_exits:
edge = graph.memlet_path(in_edge)[-2]
subset_to_add = dcpy(edge.data.subset\
if edge.data.data == node.data\
else edge.data.other_subset)
subset_to_add.pop(dims_to_discard)
upper_subsets.add(subset_to_add)
else:
raise NotImplementedError("Nodes between two maps to be"
"fused with *incoming* edges"
"from outside the maps are not"
"allowed yet.")
# find lower_subsets
for out_edge in graph.out_edges(node):
if out_edge.dst in map_entries:
# cannot use memlet tree here as there could be
# not just one map succedding. Do it manually
for oedge in graph.out_edges(out_edge.dst):
if oedge.src_conn[3:] == out_edge.dst_conn[2:]:
subset_to_add = dcpy(oedge.data.subset \
if edge.data.data == node.data \
else edge.data.other_subset)
subset_to_add.pop(dims_to_discard)
lower_subsets.add(subset_to_add)
upper_iter = iter(upper_subsets)
union_upper = next(upper_iter)
# TODO: add this check at a later point
# We assume that upper_subsets for each data array
# are contiguous
# or do the full check if possible (intersection needed)
'''
# check whether subsets in upper_subsets are adjacent.
# this is a requriement for the current implementation
#try:
# O(n^2*|dims|) but very small amount of subsets anyway
try:
for dim in range(total_dims - len(dims_to_discard)):
ordered_list = [(-1,-1,-1)]
for upper_subset in upper_subsets:
lo = upper_subset[dim][0]
hi = upper_subset[dim][1]
for idx,element in enumerate(ordered_list):
if element[0] <= lo and element[1] >= hi:
break
if element[0] > lo:
ordered_list.insert(idx, (lo,hi))
ordered_list.pop(0)
highest = ordered_list[0][1]
for i in range(len(ordered_list)):
if i < len(ordered_list)-1:
current_range = ordered_list[i]
if current_range[1] > highest:
hightest = current_range[1]
next_range = ordered_list[i+1]
if highest < next_range[0] - 1:
return False
except TypeError:
#return False
'''
# FORNOW: just omit warning if unsure
for lower_subset in lower_subsets:
covers = False
for upper_subset in upper_subsets:
if upper_subset.covers(lower_subset):
covers = True
break
if not covers:
warnings.warn(
f"WARNING: For node {node}, please check assure that"
"incoming memlets cover outgoing ones. Ambiguous check (WIP)."
)
# now take union of upper subsets
for subs in upper_iter:
union_upper = subsets.union(union_upper, subs)
if not union_upper:
# something went wrong using union -- we'd rather abort
return False
# finally check coverage
for lower_subset in lower_subsets:
if not union_upper.covers(lower_subset):
return False
return True
def apply(self, sdfg: SDFG):
subgraph = self.subgraph_view(sdfg)
entry_states_in, entry_states_out = self.get_entry_states(
sdfg, subgraph)
_, exit_states_out = self.get_exit_states(sdfg, subgraph)
entry_state_in = entry_states_in.pop()
entry_state_out = entry_states_out.pop() \
if len(entry_states_out) > 0 else None
exit_state_out = exit_states_out.pop() \
if len(exit_states_out) > 0 else None
launch_state = None
entry_guard_state = None
exit_guard_state = None
# generate entry guard state if needed
if self.include_in_assignment and entry_state_out is not None:
entry_edge = sdfg.edges_between(entry_state_out, entry_state_in)[0]
if len(entry_edge.data.assignments) > 0:
entry_guard_state = sdfg.add_state(
label='{}kernel_entry_guard'.format(
self.kernel_prefix +
'_' if self.kernel_prefix != '' else ''))
sdfg.add_edge(entry_state_out, entry_guard_state,
InterstateEdge(entry_edge.data.condition))
sdfg.add_edge(
entry_guard_state, entry_state_in,
InterstateEdge(None, entry_edge.data.assignments))
sdfg.remove_edge(entry_edge)
# Update SubgraphView
new_node_list = subgraph.nodes()
new_node_list.append(entry_guard_state)
subgraph = SubgraphView(sdfg, new_node_list)
launch_state = sdfg.add_state_before(
entry_guard_state,
label='{}kernel_launch'.format(
self.kernel_prefix +
'_' if self.kernel_prefix != '' else ''))
# generate exit guard state
if exit_state_out is not None:
exit_guard_state = sdfg.add_state_before(
exit_state_out,
label='{}kernel_exit_guard'.format(
self.kernel_prefix +
'_' if self.kernel_prefix != '' else ''))
# Update SubgraphView
new_node_list = subgraph.nodes()
new_node_list.append(exit_guard_state)
subgraph = SubgraphView(sdfg, new_node_list)
if launch_state is None:
launch_state = sdfg.add_state_before(
exit_state_out,
label='{}kernel_launch'.format(
self.kernel_prefix +
'_' if self.kernel_prefix != '' else ''))
# If the launch state doesn't exist at this point then there is no other
# states outside of the kernel, so create a stand alone launch state
if launch_state is None:
assert (entry_state_in is None and exit_state_out is None)
launch_state = sdfg.add_state(label='{}kernel_launch'.format(
self.kernel_prefix + '_' if self.kernel_prefix != '' else ''))
# create sdfg for kernel and fill it with states and edges from
# ssubgraph dfg will be nested at the end
kernel_sdfg = SDFG(
'{}kernel'.format(self.kernel_prefix +
'_' if self.kernel_prefix != '' else ''))
edges = subgraph.edges()
for edge in edges:
kernel_sdfg.add_edge(edge.src, edge.dst, edge.data)
# Setting entry node in nested SDFG if no entry guard was created
if entry_guard_state is None:
kernel_sdfg.start_state = kernel_sdfg.node_id(entry_state_in)
for state in subgraph:
state.parent = kernel_sdfg
# remove the now nested nodes from the outer sdfg and make sure the
# launch state is properly connected to remaining states
sdfg.remove_nodes_from(subgraph.nodes())
if entry_state_out is not None \
and len(sdfg.edges_between(entry_state_out, launch_state)) == 0:
sdfg.add_edge(entry_state_out, launch_state, InterstateEdge())
if exit_state_out is not None \
and len(sdfg.edges_between(launch_state, exit_state_out)) == 0:
sdfg.add_edge(launch_state, exit_state_out, InterstateEdge())
# Handle data for kernel
kernel_data = set(node.data for state in kernel_sdfg
for node in state.nodes()
if isinstance(node, nodes.AccessNode))
# move Streams and Register data into the nested SDFG
# normal data will be added as kernel argument
kernel_args = []
for data in kernel_data:
if (isinstance(sdfg.arrays[data], dace.data.Stream) or
(isinstance(sdfg.arrays[data], dace.data.Array)
and sdfg.arrays[data].storage == StorageType.Register)):
kernel_sdfg.add_datadesc(data, sdfg.arrays[data])
del sdfg.arrays[data]
else:
copy_desc = copy.deepcopy(sdfg.arrays[data])
copy_desc.transient = False
copy_desc.storage = StorageType.Default
kernel_sdfg.add_datadesc(data, copy_desc)
kernel_args.append(data)
# read only data will be passed as input, writeable data will be passed
# as 'output' otherwise kernel cannot write to data
kernel_args_read = set()
kernel_args_write = set()
for data in kernel_args:
data_accesses_read_only = [
node.access == dtypes.AccessType.ReadOnly
for state in kernel_sdfg for node in state
if isinstance(node, nodes.AccessNode) and node.data == data
]
if all(data_accesses_read_only):
kernel_args_read.add(data)
else:
kernel_args_write.add(data)
# Kernel SDFG is complete at this point
if self.validate:
kernel_sdfg.validate()
# Filling launch state with nested SDFG, map and access nodes
map_entry, map_exit = launch_state.add_map(
'{}kernel_launch_map'.format(
self.kernel_prefix + '_' if self.kernel_prefix != '' else ''),
dict(ignore='0'),
schedule=ScheduleType.GPU_Persistent,
)
nested_sdfg = launch_state.add_nested_sdfg(
kernel_sdfg,
sdfg,
kernel_args_read,
kernel_args_write,
)
# Create and connect read only data access nodes
for arg in kernel_args_read:
read_node = launch_state.add_read(arg)
launch_state.add_memlet_path(read_node,
map_entry,
nested_sdfg,
dst_conn=arg,
memlet=Memlet.from_array(
arg, sdfg.arrays[arg]))
# Create and connect writable data access nodes
for arg in kernel_args_write:
write_node = launch_state.add_write(arg)
launch_state.add_memlet_path(nested_sdfg,
map_exit,
write_node,
src_conn=arg,
memlet=Memlet.from_array(
arg, sdfg.arrays[arg]))
# Transformation is done
if self.validate:
sdfg.validate()
def apply(self, graph: SDFGState, sdfg: SDFG):
map_exit = self.map_exit
outer_map_exit = self.outer_map_exit
# Choose array
array = self.array
if array is None or len(array) == 0:
array = next(e.data.data
for e in graph.edges_between(map_exit, outer_map_exit)
if e.data.wcr is not None)
# Avoid import loop
from dace.transformation.dataflow.local_storage import OutLocalStorage
data_node: nodes.AccessNode = OutLocalStorage.apply_to(
sdfg,
dict(array=array),
verify=False,
save=False,
node_a=map_exit,
node_b=outer_map_exit)
if self.identity is None:
warnings.warn('AccumulateTransient did not properly initialize '
'newly-created transient!')
return
sdfg_state: SDFGState = sdfg.node(self.state_id)
map_entry = sdfg_state.entry_node(map_exit)
nested_sdfg: NestedSDFG = nest_state_subgraph(
sdfg=sdfg,
state=sdfg_state,
subgraph=SubgraphView(
sdfg_state, {map_entry, map_exit}
| sdfg_state.all_nodes_between(map_entry, map_exit)))
nested_sdfg_state: SDFGState = nested_sdfg.sdfg.nodes()[0]
init_state = nested_sdfg.sdfg.add_state_before(nested_sdfg_state)
temp_array: Array = sdfg.arrays[data_node.data]
init_state.add_mapped_tasklet(
name='acctrans_init',
map_ranges={
'_o%d' % i: '0:%s' % symstr(d)
for i, d in enumerate(temp_array.shape)
},
inputs={},
code='out = %s' % self.identity,
outputs={
'out':
dace.Memlet.simple(data=data_node.data,
subset_str=','.join([
'_o%d' % i
for i, _ in enumerate(temp_array.shape)
]))
},
external_edges=True)
def calculate_topology(self, subgraph):
''' Calculates topology information of the graph
self._adjacency_list: neighbors dict of outermost scope maps
self._source_maps: outermost scope maps that have in_degree 0 in the subgraph / graph
self._labels: assigns index according to topological ordering (1) + node ID (2) with priorities (1) and (2)
'''
sdfg = self._sdfg
graph = self._graph
self._adjacency_list = {m: set() for m in self._map_entries}
# helper dict needed for a quick build
exit_nodes = {graph.exit_node(me): me for me in self._map_entries}
if subgraph:
proximity_in = set(ie.src for me in self._map_entries
for ie in graph.in_edges(me))
proximity_out = set(ie.dst for me in exit_nodes
for ie in graph.out_edges(me))
extended_subgraph = SubgraphView(
graph,
set(
itertools.chain(subgraph.nodes(), proximity_in,
proximity_out)))
for node in (extended_subgraph.nodes() if subgraph else graph.nodes()):
if isinstance(node, nodes.AccessNode):
adjacent_entries = set()
for e in graph.in_edges(node):
if isinstance(e.src,
nodes.MapExit) and e.src in exit_nodes:
adjacent_entries.add(exit_nodes[e.src])
for e in graph.out_edges(node):
if isinstance(
e.dst,
nodes.MapEntry) and e.dst in self._map_entries:
adjacent_entries.add(e.dst)
# bidirectional mapping
for entry in adjacent_entries:
for other_entry in adjacent_entries:
if entry != other_entry:
self._adjacency_list[entry].add(other_entry)
self._adjacency_list[other_entry].add(entry)
# get DAG children and parents
children_dict = defaultdict(set)
parent_dict = defaultdict(set)
for map_entry in self._map_entries:
map_exit = graph.exit_node(map_entry)
for e in graph.out_edges(map_exit):
if isinstance(e.dst, nodes.AccessNode):
for oe in graph.out_edges(e.dst):
if oe.dst in self._map_entries:
other_entry = oe.dst
children_dict[map_entry].add(other_entry)
parent_dict[other_entry].add(map_entry)
# find out source nodes
self._source_maps = [
me for me in self._map_entries if len(parent_dict[me]) == 0
]
# assign a unique id to each map entry according to topological
# ordering. If on same level, sort according to ID for determinism
self._labels = {} # map -> ID
current_id = 0
while current_id < len(self._map_entries):
# get current ids whose in_degree is 0
candidates = list(me for (me, s) in parent_dict.items()
if len(s) == 0 and me not in self._labels)
candidates.sort(key=lambda me: self._graph.node_id(me))
for c in candidates:
self._labels[c] = current_id
current_id += 1
# remove candidate for each players adjacency list
for c_child in children_dict[c]:
parent_dict[c_child].remove(c)
def greedy_fuse(graph_or_subgraph: GraphViewType,
validate_all: bool,
device: dace.dtypes.DeviceType = dace.dtypes.DeviceType.CPU,
recursive: bool = True,
stencil: bool = False,
stencil_tile=None,
permutations_only: bool = True,
expand_reductions: bool = False) -> None:
'''
Greedily fuses maps of an SDFG or graph, operating in-place.
:param graph_or_subgraph: SDFG, SDFGState or Subgraph
:param validate_all: Validate SDFG or graph at each fusion step
:param device: Device type to specialize for
:param recursive: Fuse recursively within (fused and unfused) scopes
:param stencil: Perform stencil fusion instead of regular fusion
:param stencil_tile: StencilTiling Tile size, default if None
:param permutations_only: Disallow splitting of maps during MultiExpansion stage
:param expand_reductions: Expand all reduce nodes before fusion
'''
debugprint = config.Config.get_bool('debugprint')
if isinstance(graph_or_subgraph, SDFG):
# If we have an SDFG, recurse into graphs
graph_or_subgraph.simplify(validate_all=validate_all)
# MapFusion for trivial cases
graph_or_subgraph.apply_transformations_repeated(
MapFusion, validate_all=validate_all)
# recurse into graphs
for graph in graph_or_subgraph.nodes():
greedy_fuse(graph,
validate_all=validate_all,
device=device,
recursive=recursive,
stencil=stencil,
stencil_tile=stencil_tile,
permutations_only=permutations_only,
expand_reductions=expand_reductions)
else:
# we are in graph or subgraph
sdfg, graph, subgraph = None, None, None
if isinstance(graph_or_subgraph, SDFGState):
sdfg = graph_or_subgraph.parent
sdfg.apply_transformations_repeated(MapFusion,
validate_all=validate_all)
graph = graph_or_subgraph
subgraph = SubgraphView(graph, graph.nodes())
else:
sdfg = graph_or_subgraph.graph.parent
graph = graph_or_subgraph.graph
subgraph = graph_or_subgraph
# create condition function object
fusion_condition = CompositeFusion(SubgraphView(graph, graph.nodes()))
# within SDFGState: greedily enumerate fusible components
# and apply transformation
applied_transformations = 0
reverse = True if stencil else False
if stencil:
# adjust tiling settings
fusion_condition.allow_tiling = True
fusion_condition.schedule_innermaps = dtypes.ScheduleType.Sequential
if device == dtypes.DeviceType.GPU:
fusion_condition.stencil_unroll_loops = True
# tile size
if stencil_tile:
fusion_condition.stencil_strides = stencil_tile
# always only permutate for now with stencil tiles
fusion_condition.expansion_split = False
else:
fusion_condition.allow_tiling = False
# expand reductions
if expand_reductions:
for graph in sdfg.nodes():
for node in graph.nodes():
if isinstance(node,
dace.libraries.standard.nodes.Reduce):
try:
ReduceExpansion.apply_to(sdfg, reduce=node)
except ValueError as e:
pass
# permutation settings
fusion_condition.expansion_split = not permutations_only
condition_function = lambda sdfg, subgraph: fusion_condition.can_be_applied(
sdfg, subgraph)
enumerator = GreedyEnumerator(sdfg,
graph,
subgraph,
condition_function=condition_function)
for map_entries in enumerator:
if len(map_entries) > 1:
current_subgraph = xfsh.subgraph_from_maps(
sdfg, graph, map_entries)
cf = CompositeFusion(current_subgraph)
# transfer settings
cf.allow_tiling = fusion_condition.allow_tiling
cf.schedule_innermaps = fusion_condition.schedule_innermaps
cf.expansion_split = fusion_condition.expansion_split
cf.stencil_strides = fusion_condition.stencil_strides
cf.apply(sdfg)
applied_transformations += 1
if recursive:
global_entry = cf._global_map_entry if len(
map_entries) > 1 else map_entries[0]
greedy_fuse(graph.scope_subgraph(global_entry,
include_entry=False,
include_exit=False),
validate_all=validate_all,
device=device,
recursive=recursive,
stencil=stencil,
stencil_tile=stencil_tile,
permutations_only=permutations_only,
expand_reductions=expand_reductions)
for node in graph_or_subgraph.nodes():
if isinstance(node, nodes.NestedSDFG):
greedy_fuse(node.sdfg,
validate_all=validate_all,
device=device,
stencil=stencil,
stencil_tile=stencil_tile,
recursive=recursive,
permutations_only=permutations_only,
expand_reductions=expand_reductions)
if applied_transformations > 0:
if debugprint:
if stencil:
print(f"Applied {applied_transformations} TileFusion")
else:
print(f"Applied {applied_transformations} SubgraphFusion")
if validate_all:
graph.validate()
def test_simple_sdfg_map(self):
sdfg, state, t, me, mx = create_sdfg()
nest_state_subgraph(sdfg, state, SubgraphView(state, [me, t, mx]))
sdfg.validate()
def nest_state_subgraph(sdfg: SDFG,
state: SDFGState,
subgraph: SubgraphView,
name: Optional[str] = None,
full_data: bool = False) -> nodes.NestedSDFG:
""" Turns a state subgraph into a nested SDFG. Operates in-place.
:param sdfg: The SDFG containing the state subgraph.
:param state: The state containing the subgraph.
:param subgraph: Subgraph to nest.
:param name: An optional name for the nested SDFG.
:param full_data: If True, nests entire input/output data.
:return: The nested SDFG node.
:raise KeyError: Some or all nodes in the subgraph are not located in
this state, or the state does not belong to the given
SDFG.
:raise ValueError: The subgraph is contained in more than one scope.
"""
if state.parent != sdfg:
raise KeyError('State does not belong to given SDFG')
if subgraph is not state and subgraph.graph is not state:
raise KeyError('Subgraph does not belong to given state')
# Find the top-level scope
scope_tree = state.scope_tree()
scope_dict = state.scope_dict()
scope_dict_children = state.scope_children()
top_scopenode = -1 # Initialized to -1 since "None" already means top-level
for node in subgraph.nodes():
if node not in scope_dict:
raise KeyError('Node not found in state')
# If scope entry/exit, ensure entire scope is in subgraph
if isinstance(node, nodes.EntryNode):
scope_nodes = scope_dict_children[node]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (entry)')
elif isinstance(node, nodes.ExitNode):
entry = state.entry_node(node)
scope_nodes = scope_dict_children[entry] + [entry]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (exit)')
scope_node = scope_dict[node]
if scope_node not in subgraph.nodes():
if top_scopenode != -1 and top_scopenode != scope_node:
raise ValueError('Subgraph is contained in more than one scope')
top_scopenode = scope_node
scope = scope_tree[top_scopenode]
###
# Consolidate edges in top scope
utils.consolidate_edges(sdfg, scope)
snodes = subgraph.nodes()
# Collect inputs and outputs of the nested SDFG
inputs: List[MultiConnectorEdge] = []
outputs: List[MultiConnectorEdge] = []
for node in snodes:
for edge in state.in_edges(node):
if edge.src not in snodes:
inputs.append(edge)
for edge in state.out_edges(node):
if edge.dst not in snodes:
outputs.append(edge)
# Collect transients not used outside of subgraph (will be removed of
# top-level graph)
data_in_subgraph = set(n.data for n in subgraph.nodes() if isinstance(n, nodes.AccessNode))
# Find other occurrences in SDFG
other_nodes = set(n.data for s in sdfg.nodes() for n in s.nodes()
if isinstance(n, nodes.AccessNode) and n not in subgraph.nodes())
subgraph_transients = set()
for data in data_in_subgraph:
datadesc = sdfg.arrays[data]
if datadesc.transient and data not in other_nodes:
subgraph_transients.add(data)
# All transients of edges between code nodes are also added to nested graph
for edge in subgraph.edges():
if (isinstance(edge.src, nodes.CodeNode) and isinstance(edge.dst, nodes.CodeNode)):
subgraph_transients.add(edge.data.data)
# Collect data used in access nodes within subgraph (will be referenced in
# full upon nesting)
input_arrays = set()
output_arrays = {}
for node in subgraph.nodes():
if (isinstance(node, nodes.AccessNode) and node.data not in subgraph_transients):
if node.has_reads(state):
input_arrays.add(node.data)
if node.has_writes(state):
output_arrays[node.data] = state.in_edges(node)[0].data.wcr
# Create the nested SDFG
nsdfg = SDFG(name or 'nested_' + state.label)
# Transients are added to the nested graph as-is
for name in subgraph_transients:
nsdfg.add_datadesc(name, sdfg.arrays[name])
# Input/output data that are not source/sink nodes are added to the graph
# as non-transients
for name in (input_arrays | output_arrays.keys()):
datadesc = copy.deepcopy(sdfg.arrays[name])
datadesc.transient = False
nsdfg.add_datadesc(name, datadesc)
# Connected source/sink nodes outside subgraph become global data
# descriptors in nested SDFG
input_names = {}
output_names = {}
global_subsets: Dict[str, Tuple[str, Subset]] = {}
for edge in inputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = edge.data.data
if name not in global_subsets:
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
new_name = nsdfg.add_datadesc(name, datadesc, find_new_name=True)
global_subsets[name] = (new_name, edge.data.subset)
else:
new_name, subset = global_subsets[name]
if not full_data:
new_subset = union(subset, edge.data.subset)
if new_subset is None:
new_subset = Range.from_array(sdfg.arrays[name])
global_subsets[name] = (new_name, new_subset)
nsdfg.arrays[new_name].shape = new_subset.size()
input_names[edge] = new_name
for edge in outputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = edge.data.data
if name not in global_subsets:
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
new_name = nsdfg.add_datadesc(name, datadesc, find_new_name=True)
global_subsets[name] = (new_name, edge.data.subset)
else:
new_name, subset = global_subsets[name]
if not full_data:
new_subset = union(subset, edge.data.subset)
if new_subset is None:
new_subset = Range.from_array(sdfg.arrays[name])
global_subsets[name] = (new_name, new_subset)
nsdfg.arrays[new_name].shape = new_subset.size()
output_names[edge] = new_name
###################
# Add scope symbols to the nested SDFG
defined_vars = set(
symbolic.pystr_to_symbolic(s) for s in (state.symbols_defined_at(top_scopenode).keys()
| sdfg.symbols))
for v in defined_vars:
if v in sdfg.symbols:
sym = sdfg.symbols[v]
nsdfg.add_symbol(v, sym.dtype)
# Add constants to nested SDFG
for cstname, cstval in sdfg.constants.items():
nsdfg.add_constant(cstname, cstval)
# Create nested state
nstate = nsdfg.add_state()
# Add subgraph nodes and edges to nested state
nstate.add_nodes_from(subgraph.nodes())
for e in subgraph.edges():
nstate.add_edge(e.src, e.src_conn, e.dst, e.dst_conn, copy.deepcopy(e.data))
# Modify nested SDFG parents in subgraph
for node in subgraph.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent = nstate
node.sdfg.parent_sdfg = nsdfg
node.sdfg.parent_nsdfg_node = node
# Add access nodes and edges as necessary
edges_to_offset = []
for edge, name in input_names.items():
node = nstate.add_read(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge, nstate.add_edge(node, None, edge.dst, edge.dst_conn, new_edge)))
for edge, name in output_names.items():
node = nstate.add_write(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge, nstate.add_edge(edge.src, edge.src_conn, node, None, new_edge)))
# Offset memlet paths inside nested SDFG according to subsets
for original_edge, new_edge in edges_to_offset:
for edge in nstate.memlet_tree(new_edge):
edge.data.data = new_edge.data.data
if not full_data:
edge.data.subset.offset(global_subsets[original_edge.data.data][1], True)
# Add nested SDFG node to the input state
nested_sdfg = state.add_nested_sdfg(nsdfg, None,
set(input_names.values()) | input_arrays,
set(output_names.values()) | output_arrays.keys())
# Reconnect memlets to nested SDFG
reconnected_in = set()
reconnected_out = set()
empty_input = None
empty_output = None
for edge in inputs:
if edge.data.data is None:
empty_input = edge
continue
name = input_names[edge]
if name in reconnected_in:
continue
if full_data:
data = Memlet.from_array(edge.data.data, sdfg.arrays[edge.data.data])
else:
data = copy.deepcopy(edge.data)
data.subset = global_subsets[edge.data.data][1]
state.add_edge(edge.src, edge.src_conn, nested_sdfg, name, data)
reconnected_in.add(name)
for edge in outputs:
if edge.data.data is None:
empty_output = edge
continue
name = output_names[edge]
if name in reconnected_out:
continue
if full_data:
data = Memlet.from_array(edge.data.data, sdfg.arrays[edge.data.data])
else:
data = copy.deepcopy(edge.data)
data.subset = global_subsets[edge.data.data][1]
data.wcr = edge.data.wcr
state.add_edge(nested_sdfg, name, edge.dst, edge.dst_conn, data)
reconnected_out.add(name)
# Connect access nodes to internal input/output data as necessary
entry = scope.entry
exit = scope.exit
for name in input_arrays:
node = state.add_read(name)
if entry is not None:
state.add_nedge(entry, node, Memlet())
state.add_edge(node, None, nested_sdfg, name, Memlet.from_array(name, sdfg.arrays[name]))
for name, wcr in output_arrays.items():
node = state.add_write(name)
if exit is not None:
state.add_nedge(node, exit, Memlet())
state.add_edge(nested_sdfg, name, node, None, Memlet(data=name, wcr=wcr))
# Graph was not reconnected, but needs to be
if state.in_degree(nested_sdfg) == 0 and empty_input is not None:
state.add_edge(empty_input.src, empty_input.src_conn, nested_sdfg, None, empty_input.data)
if state.out_degree(nested_sdfg) == 0 and empty_output is not None:
state.add_edge(nested_sdfg, None, empty_output.dst, empty_output.dst_conn, empty_output.data)
# Remove subgraph nodes from graph
state.remove_nodes_from(subgraph.nodes())
# Remove subgraph transients from top-level graph
for transient in subgraph_transients:
del sdfg.arrays[transient]
# Remove newly isolated nodes due to memlet consolidation
for edge in inputs:
if state.in_degree(edge.src) + state.out_degree(edge.src) == 0:
state.remove_node(edge.src)
for edge in outputs:
if state.in_degree(edge.dst) + state.out_degree(edge.dst) == 0:
state.remove_node(edge.dst)
return nested_sdfg
def apply(self, sdfg: SDFG):
state: SDFGState = sdfg.nodes()[self.state_id]
nsdfg_node = state.nodes()[self.subgraph[InlineSDFG._nested_sdfg]]
nsdfg: SDFG = nsdfg_node.sdfg
nstate: SDFGState = nsdfg.nodes()[0]
if nsdfg_node.schedule is not dtypes.ScheduleType.Default:
infer_types.set_default_schedule_and_storage_types(
nsdfg, nsdfg_node.schedule)
nsdfg_scope_entry = state.entry_node(nsdfg_node)
nsdfg_scope_exit = (state.exit_node(nsdfg_scope_entry)
if nsdfg_scope_entry is not None else None)
#######################################################
# Collect and update top-level SDFG metadata
# Global/init/exit code
for loc, code in nsdfg.global_code.items():
sdfg.append_global_code(code.code, loc)
for loc, code in nsdfg.init_code.items():
sdfg.append_init_code(code.code, loc)
for loc, code in nsdfg.exit_code.items():
sdfg.append_exit_code(code.code, loc)
# Constants
for cstname, cstval in nsdfg.constants.items():
if cstname in sdfg.constants:
if cstval != sdfg.constants[cstname]:
warnings.warn('Constant value mismatch for "%s" while '
'inlining SDFG. Inner = %s != %s = outer' %
(cstname, cstval, sdfg.constants[cstname]))
else:
sdfg.add_constant(cstname, cstval)
# Find original source/destination edges (there is only one edge per
# connector, according to match)
inputs: Dict[str, MultiConnectorEdge] = {}
outputs: Dict[str, MultiConnectorEdge] = {}
input_set: Dict[str, str] = {}
output_set: Dict[str, str] = {}
for e in state.in_edges(nsdfg_node):
inputs[e.dst_conn] = e
input_set[e.data.data] = e.dst_conn
for e in state.out_edges(nsdfg_node):
outputs[e.src_conn] = e
output_set[e.data.data] = e.src_conn
# Access nodes that need to be reshaped
reshapes: Set(str) = set()
for aname, array in nsdfg.arrays.items():
if array.transient:
continue
edge = None
if aname in inputs:
edge = inputs[aname]
if len(array.shape) > len(edge.data.subset):
reshapes.add(aname)
continue
if aname in outputs:
edge = outputs[aname]
if len(array.shape) > len(edge.data.subset):
reshapes.add(aname)
continue
if edge is not None and not InlineSDFG._check_strides(
array.strides, sdfg.arrays[edge.data.data].strides,
edge.data, nsdfg_node):
reshapes.add(aname)
# Replace symbols using invocation symbol mapping
# Two-step replacement (N -> __dacesym_N --> map[N]) to avoid clashes
for symname, symvalue in nsdfg_node.symbol_mapping.items():
if str(symname) != str(symvalue):
nsdfg.replace(symname, '__dacesym_' + symname)
for symname, symvalue in nsdfg_node.symbol_mapping.items():
if str(symname) != str(symvalue):
nsdfg.replace('__dacesym_' + symname, symvalue)
# All transients become transients of the parent (if data already
# exists, find new name)
# Mapping from nested transient name to top-level name
transients: Dict[str, str] = {}
for node in nstate.nodes():
if isinstance(node, nodes.AccessNode):
datadesc = nsdfg.arrays[node.data]
if node.data not in transients and datadesc.transient:
name = sdfg.add_datadesc('%s_%s' %
(nsdfg.label, node.data),
datadesc,
find_new_name=True)
transients[node.data] = name
# All transients of edges between code nodes are also added to parent
for edge in nstate.edges():
if (isinstance(edge.src, nodes.CodeNode)
and isinstance(edge.dst, nodes.CodeNode)):
if edge.data.data is not None:
datadesc = nsdfg.arrays[edge.data.data]
if edge.data.data not in transients and datadesc.transient:
name = sdfg.add_datadesc('%s_%s' %
(nsdfg.label, edge.data.data),
datadesc,
find_new_name=True)
transients[edge.data.data] = name
# Collect nodes to add to top-level graph
new_incoming_edges: Dict[nodes.Node, MultiConnectorEdge] = {}
new_outgoing_edges: Dict[nodes.Node, MultiConnectorEdge] = {}
source_accesses = set()
sink_accesses = set()
for node in nstate.source_nodes():
if (isinstance(node, nodes.AccessNode)
and node.data not in transients
and node.data not in reshapes):
new_incoming_edges[node] = inputs[node.data]
source_accesses.add(node)
for node in nstate.sink_nodes():
if (isinstance(node, nodes.AccessNode)
and node.data not in transients
and node.data not in reshapes):
new_outgoing_edges[node] = outputs[node.data]
sink_accesses.add(node)
#######################################################
# Replace data on inlined SDFG nodes/edges
# Replace data names with their top-level counterparts
repldict = {}
repldict.update(transients)
repldict.update({
k: v.data.data
for k, v in itertools.chain(inputs.items(), outputs.items())
})
# Add views whenever reshapes are necessary
for dname in reshapes:
desc = nsdfg.arrays[dname]
# To avoid potential confusion, rename protected __return keyword
if dname.startswith('__return'):
newname = f'{nsdfg.name}_ret{dname[8:]}'
else:
newname = dname
newname, _ = sdfg.add_view(newname,
desc.shape,
desc.dtype,
storage=desc.storage,
strides=desc.strides,
offset=desc.offset,
debuginfo=desc.debuginfo,
allow_conflicts=desc.allow_conflicts,
total_size=desc.total_size,
alignment=desc.alignment,
may_alias=desc.may_alias,
find_new_name=True)
repldict[dname] = newname
for node in nstate.nodes():
if isinstance(node, nodes.AccessNode) and node.data in repldict:
node.data = repldict[node.data]
for edge in nstate.edges():
if edge.data.data in repldict:
edge.data.data = repldict[edge.data.data]
# Add extra access nodes for out/in view nodes
for node in nstate.nodes():
if isinstance(node, nodes.AccessNode) and node.data in reshapes:
if nstate.in_degree(node) > 0 and nstate.out_degree(node) > 0:
# Such a node has to be in the output set
edge = outputs[node.data]
# Redirect outgoing edges through access node
out_edges = list(nstate.out_edges(node))
anode = nstate.add_access(edge.data.data)
vnode = nstate.add_access(node.data)
nstate.add_nedge(node, anode, edge.data)
nstate.add_nedge(anode, vnode, edge.data)
for e in out_edges:
nstate.remove_edge(e)
nstate.add_edge(vnode, e.src_conn, e.dst, e.dst_conn,
e.data)
#######################################################
# Add nested SDFG into top-level SDFG
# Add nested nodes into original state
subgraph = SubgraphView(nstate, [
n for n in nstate.nodes()
if n not in (source_accesses | sink_accesses)
])
state.add_nodes_from(subgraph.nodes())
for edge in subgraph.edges():
state.add_edge(edge.src, edge.src_conn, edge.dst, edge.dst_conn,
edge.data)
#######################################################
# Reconnect inlined SDFG
# If a source/sink node is one of the inputs/outputs, reconnect it,
# replacing memlets in outgoing/incoming paths
modified_edges = set()
modified_edges |= self._modify_memlet_path(new_incoming_edges, nstate,
state, True)
modified_edges |= self._modify_memlet_path(new_outgoing_edges, nstate,
state, False)
# Reshape: add connections to viewed data
self._modify_reshape_data(reshapes, repldict, inputs, nstate, state,
True)
self._modify_reshape_data(reshapes, repldict, outputs, nstate, state,
False)
# Modify all other internal edges pertaining to input/output nodes
for node in subgraph.nodes():
if isinstance(node, nodes.AccessNode):
if node.data in input_set or node.data in output_set:
if node.data in input_set:
outer_edge = inputs[input_set[node.data]]
else:
outer_edge = outputs[output_set[node.data]]
for edge in state.all_edges(node):
if (edge not in modified_edges
and edge.data.data == node.data):
for e in state.memlet_tree(edge):
if e.data.data == node.data:
e._data = helpers.unsqueeze_memlet(
e.data, outer_edge.data)
# If source/sink node is not connected to a source/destination access
# node, and the nested SDFG is in a scope, connect to scope with empty
# memlets
if nsdfg_scope_entry is not None:
for node in subgraph.nodes():
if state.in_degree(node) == 0:
state.add_edge(nsdfg_scope_entry, None, node, None,
Memlet())
if state.out_degree(node) == 0:
state.add_edge(node, None, nsdfg_scope_exit, None,
Memlet())
# Replace nested SDFG parents with new SDFG
for node in nstate.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent = state
node.sdfg.parent_sdfg = sdfg
node.sdfg.parent_nsdfg_node = node
# Remove all unused external inputs/output memlet paths, as well as
# resulting isolated nodes
removed_in_edges = self._remove_edge_path(state,
inputs,
set(inputs.keys()) -
source_accesses,
reverse=True)
removed_out_edges = self._remove_edge_path(state,
outputs,
set(outputs.keys()) -
sink_accesses,
reverse=False)
# Re-add in/out edges to first/last nodes in subgraph
order = [
x for x in nx.topological_sort(nstate._nx)
if isinstance(x, nodes.AccessNode)
]
for edge in removed_in_edges:
# Find first access node that refers to this edge
node = next(n for n in order if n.data == edge.data.data)
state.add_edge(edge.src, edge.src_conn, node, edge.dst_conn,
edge.data)
for edge in removed_out_edges:
# Find last access node that refers to this edge
node = next(n for n in reversed(order) if n.data == edge.data.data)
state.add_edge(node, edge.src_conn, edge.dst, edge.dst_conn,
edge.data)
#######################################################
# Remove nested SDFG node
state.remove_node(nsdfg_node)
def can_be_applied(sdfg: SDFG, subgraph: SubgraphView) -> bool:
'''
Fusible if
1. Maps have the same access sets and ranges in order
2. Any nodes in between two maps are AccessNodes only, without WCR
There is at most one AccessNode only on a path between two maps,
no other nodes are allowed
3. The exiting memlets' subsets to an intermediate edge must cover
the respective incoming memlets' subset into the next map.
Also, as a limitation, the union of all exiting memlets'
subsets must be contiguous.
'''
# get graph
graph = subgraph.graph
for node in subgraph.nodes():
if node not in graph.nodes():
return False
# next, get all the maps
map_entries = helpers.get_outermost_scope_maps(sdfg, graph, subgraph)
map_exits = [graph.exit_node(map_entry) for map_entry in map_entries]
maps = [map_entry.map for map_entry in map_entries]
# 1. basic checks:
# 1.1 we need to have at least two maps
if len(maps) <= 1:
return False
'''
# 1.2 Special Case: If we can establish a valid permutation, we can
# skip check 1.3
permutation = self.find_permutation
'''
# 1.3 check whether all maps are the same
base_map = maps[0]
for map in maps:
if map.get_param_num() != base_map.get_param_num():
return False
if not all(
[p1 == p2 for (p1, p2) in zip(map.params, base_map.params)]):
return False
if not map.range == base_map.range:
return False
# 1.3 check whether all map entries have the same schedule
schedule = map_entries[0].schedule
if not all([entry.schedule == schedule for entry in map_entries]):
return False
# 2. check intermediate feasiblility
# see map_fusion.py for similar checks
# with the restrictions below being more relaxed
# 2.1 do some preparation work first:
# calculate all out_nodes and intermediate_nodes
# definition see in apply()
node_config = SubgraphFusion.get_adjacent_nodes(sdfg, graph,
map_entries)
_, intermediate_nodes, out_nodes = node_config
# 2.2 topological feasibility:
if not SubgraphFusion.check_topo_feasibility(
sdfg, graph, map_entries, intermediate_nodes, out_nodes):
return False
# 2.3 memlet feasibility
# For each intermediate node, look at whether inner adjacent
# memlets of the exiting map cover inner adjacent memlets
# of the next entering map.
# We also check for any WCRs on the fly.
for node in intermediate_nodes:
upper_subsets = set()
lower_subsets = set()
# First, determine which dimensions of the memlet ranges
# change with the map, we do not need to care about the other dimensions.
try:
dims_to_discard = SubgraphFusion.get_invariant_dimensions(
sdfg, graph, map_entries, map_exits, node)
except NotImplementedError:
return False
# find upper_subsets
for in_edge in graph.in_edges(node):
in_in_edge = graph.memlet_path(in_edge)[-2]
# first check for WCRs
if in_edge.data.wcr:
# check whether the WCR is actually produced at
# this edge or further up in the memlet path. If not,
# we can still fuse!
subset_params = set(
[str(s) for s in in_in_edge.data.subset.free_symbols])
if any([
p not in subset_params
for p in in_edge.src.map.params
]):
return False
if in_edge.src in map_exits:
subset_to_add = dcpy(in_in_edge.data.subset\
if in_in_edge.data.data == node.data\
else in_in_edge.data.other_subset)
subset_to_add.pop(dims_to_discard)
upper_subsets.add(subset_to_add)
else:
raise NotImplementedError("Nodes between two maps to be"
"fused with *incoming* edges"
"from outside the maps are not"
"allowed yet.")
# find lower_subsets
for out_edge in graph.out_edges(node):
if out_edge.dst in map_entries:
# cannot use memlet tree here as there could be
# not just one map succedding. Do it manually
for oedge in graph.out_edges(out_edge.dst):
if oedge.src_conn[3:] == out_edge.dst_conn[2:]:
subset_to_add = dcpy(oedge.data.subset \
if oedge.data.data == node.data \
else oedge.data.other_subset)
subset_to_add.pop(dims_to_discard)
lower_subsets.add(subset_to_add)
# We assume that upper_subsets are contiguous
# Check for this.
try:
contiguous_upper = find_contiguous_subsets(upper_subsets)
if len(contiguous_upper) > 1:
return False
except TypeError:
warnings.warn(
'Could not determine whether subset is continuous.'
'Exiting Check with False.')
return False
# now take union of upper subsets
upper_iter = iter(upper_subsets)
union_upper = next(upper_iter)
for subs in upper_iter:
union_upper = subsets.union(union_upper, subs)
if not union_upper:
# something went wrong using union -- we'd rather abort
return False
# finally check coverage
# every lower subset must be completely covered by union_upper
for lower_subset in lower_subsets:
if not union_upper.covers(lower_subset):
return False
return True
