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")