def _extend_subgraph_with_access_nodes(state: SDFGState, subgraph: StateSubgraphView) -> StateSubgraphView: """ Expands a subgraph view to include necessary input/output access nodes, using memlet paths. """ sdfg = state.parent result: List[nd.Node] = copy.copy(subgraph.nodes()) queue: Deque[nd.Node] = deque(subgraph.nodes()) # Add all nodes in memlet paths while len(queue) > 0: node = queue.pop() if isinstance(node, nd.AccessNode): if isinstance(node.desc(sdfg), data.View): vnode = sdutil.get_view_node(state, node) result.append(vnode) queue.append(vnode) continue for e in state.in_edges(node): # Special case: IN_* connectors are not traversed further if isinstance(e.dst, (nd.EntryNode, nd.ExitNode)) and e.dst_conn.startswith('IN_'): continue mpath = state.memlet_path(e) new_nodes = [mpe.src for mpe in mpath if mpe.src not in result] result.extend(new_nodes) # Memlet path may end in a code node, continue traversing and expanding graph queue.extend(new_nodes) for e in state.out_edges(node): # Special case: OUT_* connectors are not traversed further if isinstance(e.src, (nd.EntryNode, nd.ExitNode)) and e.src_conn.startswith('OUT_'): continue mpath = state.memlet_path(e) new_nodes = [mpe.dst for mpe in mpath if mpe.dst not in result] result.extend(new_nodes) # Memlet path may end in a code node, continue traversing and expanding graph queue.extend(new_nodes) # Check for mismatch in scopes for node in result: enode = None if isinstance(node, nd.EntryNode) and state.exit_node(node) not in result: enode = state.exit_node(node) if isinstance(node, nd.ExitNode) and state.entry_node(node) not in result: enode = state.entry_node(node) if enode is not None: raise ValueError(f'Cutout cannot expand graph implicitly since "{node}" is in the graph and "{enode}" is ' 'not. Please provide more nodes in the subgraph as necessary.') return StateSubgraphView(state, result)
def _containers_defined_outside(sdfg: SDFG, state: SDFGState, subgraph: StateSubgraphView) -> Set[str]: """ Returns a list of containers set outside the given subgraph. """ # Since we care about containers that are written to, we only need to look at access nodes rather than interstate # edges result: Set[str] = set() for ostate in sdfg.nodes(): for node in ostate.data_nodes(): if ostate is not state or node not in subgraph.nodes(): if ostate.in_degree(node) > 0: result.add(node.data) # Add all new sink nodes of new subgraph for dnode in subgraph.data_nodes(): if subgraph.out_degree(dnode) == 0 and state.out_degree(dnode) > 0: result.add(dnode.data) return result