OTFMapFusion: Bugfix for tasklets with None connectors (#1415)

spcl · Dec 2, 2023 · f90cf5b · f90cf5b
1 parent a7b3a75
commit f90cf5b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/dace/transformation/dataflow/otf_map_fusion.py b/dace/transformation/dataflow/otf_map_fusion.py
@@ -297,13 +297,10 @@ def apply(self, graph: SDFGState, sdfg: SDFG):
                             else:
                                 out_connector = edge.src_conn
 
-                            if out_connector not in second_map_entry.out_connectors:
-                                second_map_entry.add_out_connector(out_connector)
+                            if out_connector not in self.second_map_entry.out_connectors:
+                                self.second_map_entry.add_out_connector(out_connector)
                         else:
-                            out_connector = edge.src_conn
-
-                        if out_connector not in self.second_map_entry.out_connectors:
-                            self.second_map_entry.add_out_connector(out_connector)
+                            out_connector = None
 
                         graph.add_edge(second_map_entry, out_connector, node, edge.dst_conn, memlet)
                         graph.remove_edge(edge)

diff --git a/tests/transformations/otf_map_fusion_test.py b/tests/transformations/otf_map_fusion_test.py
@@ -330,6 +330,36 @@ def test_trivial_fusion_nested_sdfg():
     assert (res == res_fused).all()
 
 
+@dace.program
+def trivial_fusion_none_connectors(B: dace.float64[10, 20]):
+    tmp = dace.define_local([10, 20], dtype=B.dtype)
+    for i, j in dace.map[0:10, 0:20]:
+        with dace.tasklet:
+            b >> tmp[i, j]
+            b = 0
+
+    for i, j in dace.map[0:10, 0:20]:
+        with dace.tasklet:
+            a << tmp[i, j]
+            b >> B[i, j]
+            b = a + 2
+
+
+def test_trivial_fusion_none_connectors():
+    sdfg = trivial_fusion_none_connectors.to_sdfg()
+    sdfg.simplify()
+    assert count_maps(sdfg) == 2
+
+    sdfg.apply_transformations(OTFMapFusion)
+    assert count_maps(sdfg) == 1
+
+    B = np.zeros((10, 20))
+    ref = np.zeros((10, 20)) + 2
+
+    sdfg(B=B)
+    assert np.allclose(B, ref)
+
+
 @dace.program
 def undefined_subset(A: dace.float64[10], B: dace.float64[10]):
     tmp = dace.define_local([10], dtype=A.dtype)
@@ -703,6 +733,7 @@ def test_hdiff():
     test_trivial_fusion_permute()
     test_trivial_fusion_not_remove_map()
     test_trivial_fusion_nested_sdfg()
+    test_trivial_fusion_none_connectors()
 
     # Defined subsets
     test_undefined_subset()