def load_sdfg_from_json(json):
if 'error' in json:
message = ''
if ('message' in json['error']):
message = json['error']['message']
error = {
'error': {
'message': 'Invalid SDFG provided',
'details': message,
}
}
sdfg = None
else:
try:
sdfg = SDFG.from_json(json)
error = None
except Exception as e:
print(traceback.format_exc(), file=sys.stderr)
sys.stderr.flush()
error = {
'error': {
'message': 'Failed to parse the provided SDFG',
'details': get_exception_message(e),
},
}
sdfg = None
return {
'error': error,
'sdfg': sdfg,
}
def apply(self, sdfg):
from dace.transformation.dataflow import TrivialMapElimination
state = sdfg.nodes()[self.state_id]
map_entry = state.nodes()[self.subgraph[MapUnroll._map_entry]]
map_exit = state.exit_node(map_entry)
# Collect all nodes in this weakly connected component
subgraph = sdutil.weakly_connected_component(state, map_entry)
# Save nested SDFGs to JSON, then deserialize them for every copy we
# need to make
nested_sdfgs = {}
for node in subgraph:
if isinstance(node, nodes.NestedSDFG):
nested_sdfgs[node.sdfg] = node.sdfg.to_json()
# Check for local memories that need to be replicated
local_memories = [
name for name in sdutil.local_transients(
sdfg, subgraph, entry_node=map_entry, include_nested=True)
if not isinstance(sdfg.arrays[name], dt.Stream)
and not isinstance(sdfg.arrays[name], dt.View)
]
params = map_entry.map.params
ranges = map_entry.map.range.ranges
constant_ranges = []
for r in ranges:
begin = symbolic.evaluate(r[0], sdfg.constants)
end = symbolic.evaluate(r[1], sdfg.constants)
step = symbolic.evaluate(r[2], sdfg.constants)
end += step # Make non-inclusive
constant_ranges.append(range(begin, end, step))
index_tuples = itertools.product(*constant_ranges)
for t in index_tuples:
suffix = "_" + "_".join(map(str, t))
node_to_unrolled = {}
# Copy all nodes
for node in subgraph:
if isinstance(node, nodes.NestedSDFG):
# Avoid deep-copying the nested SDFG
nsdfg = node.sdfg
# Don't copy the nested SDFG, as we will do this separately
node.sdfg = None
unrolled_node = copy.deepcopy(node)
node.sdfg = nsdfg
# Deserialize into a new SDFG specific to this copy
nsdfg_json = nested_sdfgs[nsdfg]
name = nsdfg_json["attributes"]["name"]
nsdfg_json["attributes"]["name"] += suffix
unrolled_nsdfg = SDFG.from_json(nsdfg_json)
nsdfg_json["attributes"]["name"] = name # Reinstate
# Set all the references
unrolled_nsdfg.parent = state
unrolled_nsdfg.parent_sdfg = sdfg
unrolled_nsdfg.update_sdfg_list([])
unrolled_node.sdfg = unrolled_nsdfg
unrolled_nsdfg.parent_nsdfg_node = unrolled_node
else:
unrolled_node = copy.deepcopy(node)
if node == map_entry:
# Fix the map bounds to only this iteration
unrolled_node.map.range = [(i, i, 1) for i in t]
if (isinstance(node, nodes.AccessNode)
and node.data in local_memories):
# If this is a local memory only used in this subgraph,
# we need to replicate it for each new subgraph
unrolled_name = node.data + suffix
if unrolled_name not in sdfg.arrays:
unrolled_desc = copy.deepcopy(
sdfg.arrays[node.data])
sdfg.add_datadesc(unrolled_name, unrolled_desc)
unrolled_node.data = unrolled_name
state.add_node(unrolled_node)
node_to_unrolled[node] = unrolled_node # Remember mapping
# Copy all edges
for src, src_conn, dst, dst_conn, memlet in subgraph.edges():
src = node_to_unrolled[src]
dst = node_to_unrolled[dst]
memlet = copy.deepcopy(memlet)
if memlet.data in local_memories:
memlet.data = memlet.data + suffix
state.add_edge(src, src_conn, dst, dst_conn, memlet)
# Eliminate the now trivial map
TrivialMapElimination.apply_to(
sdfg,
verify=False,
annotate=False,
save=False,
_map_entry=node_to_unrolled[map_entry])
# Now we can delete the original subgraph. This implicitly also remove
# memlets between nodes
state.remove_nodes_from(subgraph)
# If we added a bunch of new nested SDFGs, reset the internal list
if len(nested_sdfgs) > 0:
sdfg.reset_sdfg_list()
# Remove local memories that were replicated
for mem in local_memories:
sdfg.remove_data(mem)
