def expand_reduce(sdfg: dace.SDFG, graph: dace.SDFGState, subgraph: Union[SubgraphView, List[SubgraphView]] = None, **kwargs): subgraph = graph if not subgraph else subgraph if not isinstance(subgraph, list): subgraph = [subgraph] for sg in subgraph: reduce_nodes = [] for node in sg.nodes(): if isinstance(node, stdlib.Reduce): rexp = ReduceExpansion(sdfg, sdfg.sdfg_id, sdfg.node_id(graph), {ReduceExpansion.reduce: graph.node_id(node)}, 0) if not rexp.can_be_applied(graph, 0, sdfg): print(f"WARNING: Cannot expand reduce node {node}:" "can_be_applied() failed.") continue reduce_nodes.append(node) trafo_reduce = ReduceExpansion(sdfg, sdfg.sdfg_id, sdfg.node_id(graph), {}, 0) for (property, val) in kwargs.items(): setattr(trafo_reduce, property, val) for reduce_node in reduce_nodes: trafo_reduce.expand(sdfg, graph, reduce_node) if isinstance(sg, SubgraphView): sg.nodes().remove(reduce_node) sg.nodes().append(trafo_reduce._reduce) sg.nodes().append(trafo_reduce._outer_entry)
def test_p1(in_transient, out_transient): sdfg = reduction_test_1.to_sdfg() sdfg.apply_strict_transformations() state = sdfg.nodes()[0] for node in state.nodes(): if isinstance(node, dace.libraries.standard.nodes.Reduce): reduce_node = node assert ReduceExpansion.can_be_applied(state, \ {ReduceExpansion._reduce: state.nodes().index(reduce_node)}, \ 0, \ sdfg) == True A = np.random.rand(M.get(), N.get()).astype(np.float64) B = np.random.rand(M.get(), N.get()).astype(np.float64) C1 = np.zeros([N.get()], dtype=np.float64) C2 = np.zeros([N.get()], dtype=np.float64) csdfg = sdfg.compile() csdfg(A=A, B=B, C=C1, N=N, M=M) del csdfg expand_reduce(sdfg, state, create_in_transient=in_transient, create_out_transient=out_transient) csdfg = sdfg.compile() csdfg(A=A, B=B, C=C2, N=N, M=M) del csdfg assert np.linalg.norm(C1) > 0.01 assert np.allclose(C1, C2)
def test_p1(): sdfg = test_program.to_sdfg() sdfg.apply_strict_transformations() state = sdfg.nodes()[0] for node in state.nodes(): if isinstance(node, dace.libraries.standard.nodes.Reduce): reduce_node = node assert ReduceExpansion.can_be_applied(state, \ {ReduceExpansion._reduce: state.nodes().index(reduce_node)}, \ 0, \ sdfg) == True A = np.random.rand(M.get(), N.get()).astype(np.float64) B = np.random.rand(M.get(), N.get()).astype(np.float64) C1 = np.zeros([N.get()], dtype=np.float64) C2 = np.zeros([N.get()], dtype=np.float64) csdfg = sdfg.compile() csdfg(A=A, B=B, C=C1, N=N, M=M) expand_reduce(sdfg, state) csdfg = sdfg.compile() csdfg(A=A, B=B, C=C2, N=N, M=M) assert np.allclose(C1, C2) print(np.linalg.norm(C1)) print(np.linalg.norm(C2)) print("PASS")