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)