def node_path_graph(*args): """ Generates a path graph passing through the input nodes. The function generates a graph using as nodes the input arguments. Subsequently, it creates a path passing through all the nodes, in the same order as they were given in the function input. :param *args: Variable number of nodes or a list of nodes. :return: A directed graph based on the input arguments. @rtype: gr.OrderedDiGraph """ # 1. Create new networkx directed graph. path = gr.OrderedDiGraph() # 2. Place input nodes in a list. if len(args) == 1 and isinstance(args[0], list): # Input is a single list of nodes. input_nodes = args[0] else: # Input is a variable number of nodes. input_nodes = list(args) # 3. Add nodes to the graph. path.add_nodes_from(input_nodes) # 4. Add path edges to the graph. for i in range(len(input_nodes) - 1): path.add_edge(input_nodes[i], input_nodes[i + 1], None) # 5. Return the graph. return path
def expressions(cls): graph = gr.OrderedDiGraph() graph.add_node(cls.transpose_a) graph.add_node(cls.at) graph.add_node(cls.transpose_b) graph.add_node(cls.bt) graph.add_node(cls.a_times_b) graph.add_edge(cls.transpose_a, cls.at, None) graph.add_edge(cls.at, cls.a_times_b, None) graph.add_edge(cls.transpose_b, cls.bt, None) graph.add_edge(cls.bt, cls.a_times_b, None) return [graph]
def _make_dependency_graph(self) -> gr.OrderedDiGraph: """ Makes an ordered dependency graph out of the passes in the pipeline to traverse when applying. """ result = gr.OrderedDiGraph() ptype_to_pass = {type(p): p for p in self.passes} for p in self.passes: if p not in result._nodes: result.add_node(p) for dep in p.depends_on(): # If a type, find it in self.passes if isinstance(dep, type): dep = ptype_to_pass[dep] result.add_edge(dep, p) return result
class CopyToDevice(pattern_matching.Transformation): """ Implements the copy-to-device transformation, which copies a nested SDFG and its dependencies to a given device. The transformation changes all data storage types of a nested SDFG to the given `storage` property, and creates new arrays and copies around the nested SDFG to that storage. """ _nested_sdfg = nodes.NestedSDFG("", graph.OrderedDiGraph(), {}, {}) storage = properties.Property(dtype=dtypes.StorageType, desc="Nested SDFG storage", choices=dtypes.StorageType, from_string=lambda x: dtypes.StorageType[x], default=dtypes.StorageType.Default) @staticmethod def annotates_memlets(): return True @staticmethod def expressions(): return [sdutil.node_path_graph(CopyToDevice._nested_sdfg)] @staticmethod def can_be_applied(graph, candidate, expr_index, sdfg, strict=False): nested_sdfg = graph.nodes()[candidate[CopyToDevice._nested_sdfg]] for edge in graph.all_edges(nested_sdfg): # Stream inputs/outputs not allowed path = graph.memlet_path(edge) if ((isinstance(path[0].src, nodes.AccessNode) and isinstance(sdfg.arrays[path[0].src.data], data.Stream)) or (isinstance(path[-1].dst, nodes.AccessNode) and isinstance(sdfg.arrays[path[-1].dst.data], data.Stream))): return False # WCR outputs with arrays are not allowed if (edge.data.wcr is not None and edge.data.subset.num_elements() != 1): return False return True @staticmethod def match_to_str(graph, candidate): nested_sdfg = graph.nodes()[candidate[CopyToDevice._nested_sdfg]] return nested_sdfg.label def apply(self, sdfg): state = sdfg.nodes()[self.state_id] nested_sdfg = state.nodes()[self.subgraph[CopyToDevice._nested_sdfg]] storage = self.storage created_arrays = set() for _, edge in enumerate(state.in_edges(nested_sdfg)): src, src_conn, dst, dst_conn, memlet = edge dataname = memlet.data if dataname is None: continue memdata = sdfg.arrays[dataname] name = 'device_' + dataname + '_in' if name not in created_arrays: if isinstance(memdata, data.Array): name, _ = sdfg.add_array( 'device_' + dataname + '_in', shape=[ symbolic.overapproximate(r) for r in memlet.bounding_box_size() ], dtype=memdata.dtype, transient=True, storage=storage, find_new_name=True) elif isinstance(memdata, data.Scalar): name, _ = sdfg.add_scalar('device_' + dataname + '_in', dtype=memdata.dtype, transient=True, storage=storage, find_new_name=True) else: raise NotImplementedError created_arrays.add(name) data_node = nodes.AccessNode(name) to_data_mm = dcpy(memlet) from_data_mm = dcpy(memlet) from_data_mm.data = name offset = [] for ind, r in enumerate(memlet.subset): offset.append(r[0]) if isinstance(memlet.subset[ind], tuple): begin = memlet.subset[ind][0] - r[0] end = memlet.subset[ind][1] - r[0] step = memlet.subset[ind][2] from_data_mm.subset[ind] = (begin, end, step) else: from_data_mm.subset[ind] -= r[0] state.remove_edge(edge) state.add_edge(src, src_conn, data_node, None, to_data_mm) state.add_edge(data_node, None, dst, dst_conn, from_data_mm) for _, edge in enumerate(state.out_edges(nested_sdfg)): src, src_conn, dst, dst_conn, memlet = edge dataname = memlet.data if dataname is None: continue memdata = sdfg.arrays[dataname] name = 'device_' + dataname + '_out' if name not in created_arrays: if isinstance(memdata, data.Array): name, _ = sdfg.add_array( name, shape=[ symbolic.overapproximate(r) for r in memlet.bounding_box_size() ], dtype=memdata.dtype, transient=True, storage=storage, find_new_name=True) elif isinstance(memdata, data.Scalar): name, _ = sdfg.add_scalar(name, dtype=memdata.dtype, transient=True, storage=storage) else: raise NotImplementedError created_arrays.add(name) data_node = nodes.AccessNode(name) to_data_mm = dcpy(memlet) from_data_mm = dcpy(memlet) to_data_mm.data = name offset = [] for ind, r in enumerate(memlet.subset): offset.append(r[0]) if isinstance(memlet.subset[ind], tuple): begin = memlet.subset[ind][0] - r[0] end = memlet.subset[ind][1] - r[0] step = memlet.subset[ind][2] to_data_mm.subset[ind] = (begin, end, step) else: to_data_mm.subset[ind] -= r[0] state.remove_edge(edge) state.add_edge(src, src_conn, data_node, None, to_data_mm) state.add_edge(data_node, None, dst, dst_conn, from_data_mm) # Change storage for all data inside nested SDFG to device. change_storage(nested_sdfg.sdfg, storage)