MapFission: WCR in NestedSDFG fix

spcl · Sep 3, 2023 · 1ea713f · 1ea713f
1 parent c5ca99a
commit 1ea713f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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: