WCRToAugAssign: Added memlet propagation

spcl · Nov 3, 2023 · ee6a13d · ee6a13d
1 parent 8140313
commit ee6a13d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/dace/transformation/dataflow/wcr_conversion.py b/dace/transformation/dataflow/wcr_conversion.py
@@ -9,7 +9,7 @@
 from dace import SDFG, SDFGState
 from dace.sdfg.state import StateSubgraphView
 from dace.transformation import helpers
-from dace import propagate_memlet
+from dace.sdfg.propagation import propagate_memlets_state
 
 
 class AugAssignToWCR(transformation.SingleStateTransformation):
@@ -256,6 +256,7 @@ def apply(self, state: SDFGState, sdfg: SDFG):
 
         # Remove input node and connector
         state.remove_memlet_path(inedge)
+        propagate_memlets_state(sdfg, state)
 
         # If outedge leads to non-transient, and this is a nested SDFG,
         # propagate outwards