From dfa80d9f56f24fd783e1efcb886f912eaea5283d Mon Sep 17 00:00:00 2001 From: Lukas Truemper Date: Sun, 3 Sep 2023 11:28:09 +0200 Subject: [PATCH] MapFission: WCR in NestedSDFG fix --- dace/transformation/dataflow/map_fission.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dace/transformation/dataflow/map_fission.py b/dace/transformation/dataflow/map_fission.py index a50605812b..15a16099a9 100644 --- a/dace/transformation/dataflow/map_fission.py +++ b/dace/transformation/dataflow/map_fission.py @@ -264,7 +264,10 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): edge_to_outer[edge] = path[eindex - 1] else: # Nested SDFGs use the internal map edges of the node - outer_edge = next(e for e in graph.in_edges(nsdfg_node) if e.dst_conn == edge.src.data) + outer_edge = [e for e in graph.in_edges(nsdfg_node) if e.dst_conn == edge.src.data] + if not outer_edge: + outer_edge = [e for e in graph.out_edges(nsdfg_node) if e.src_conn == edge.src.data and e.data.wcr is not None] + outer_edge = outer_edge[0] edge_to_outer[edge] = outer_edge for edge in external_edges_exit: