AugAssingToWCR: permissive option for free maps

spcl · Sep 2, 2023 · 9c73a2f · 9c73a2f
1 parent 7c9e8e2
commit 9c73a2f
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/dace/transformation/dataflow/wcr_conversion.py b/dace/transformation/dataflow/wcr_conversion.py
@@ -134,9 +134,10 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False):
 
                 # If in map, only match if the subset is independent of any
                 # map indices (otherwise no conflict)
-                if (expr_index == 1 and len(outedge.data.subset.free_symbols
-                                            & set(me.map.params)) == len(me.map.params)):
-                    continue
+                if expr_index == 1:
+                    if not permissive and len(outedge.data.subset.free_symbols & set(me.map.params)) == len(
+                            me.map.params):
+                        continue
 
                 return True
         else:

diff --git a/tests/transformations/wcr_conversion_test.py b/tests/transformations/wcr_conversion_test.py
@@ -222,3 +222,26 @@ def sdfg_aug_assign_dependent_map(A: dace.float64[32], B: dace.float64[32]):
 
     applied = sdfg.apply_transformations_repeated(AugAssignToWCR)
     assert applied == 1
+
+
+def test_free_map_permissive():
+
+    @dace.program
+    def sdfg_free_map_permissive(A: dace.float64[32], B: dace.float64[32]):
+        for i in dace.map[0:32]:
+            with dace.tasklet(language=dace.Language.CPP):
+                a << A[i]
+                k << B[i]
+                b >> A[i]
+                """
+                b = k * a;
+                """
+
+    sdfg = sdfg_free_map_permissive.to_sdfg()
+    sdfg.simplify()
+
+    applied = sdfg.apply_transformations_repeated(AugAssignToWCR, permissive=False)
+    assert applied == 0
+
+    applied = sdfg.apply_transformations_repeated(AugAssignToWCR, permissive=True)
+    assert applied == 1