Fix loop detection

spcl · Jan 9, 2025 · 508f41b · 508f41b
1 parent 35aebfb
commit 508f41b
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/dace/transformation/interstate/loop_detection.py b/dace/transformation/interstate/loop_detection.py
@@ -17,20 +17,20 @@ class DetectLoop(transformation.PatternTransformation):
     """ Detects a for-loop construct from an SDFG. """
 
     # Always available
-    loop_begin = transformation.PatternNode(sd.SDFGState)
-    exit_state = transformation.PatternNode(sd.SDFGState)
+    loop_begin = transformation.PatternNode(ControlFlowBlock)
+    exit_state = transformation.PatternNode(ControlFlowBlock)
 
     # Available for natural loops
-    loop_guard = transformation.PatternNode(sd.SDFGState)
+    loop_guard = transformation.PatternNode(ControlFlowBlock)
 
     # Available for rotated loops
-    loop_latch = transformation.PatternNode(sd.SDFGState)
+    loop_latch = transformation.PatternNode(ControlFlowBlock)
 
     # Available for rotated and self loops
-    entry_state = transformation.PatternNode(sd.SDFGState)
+    entry_state = transformation.PatternNode(ControlFlowBlock)
 
     # Available for explicit-latch rotated loops
-    loop_break = transformation.PatternNode(sd.SDFGState)
+    loop_break = transformation.PatternNode(ControlFlowBlock)
 
     @classmethod
     def expressions(cls):
@@ -260,6 +260,10 @@ def detect_rotated_loop(self, graph: ControlFlowRegion, multistate_loop: bool,
         if latch_outedges[0].data.condition_sympy() != (sp.Not(latch_outedges[1].data.condition_sympy())):
             return None
 
+        # Make sure the backedge (i.e, one of the condition edges) goes from the latch to the beginning state.
+        if latch_outedges[0].dst is not self.loop_begin and latch_outedges[1].dst is not self.loop_begin:
+            return None
+
         # All nodes inside loop must be dominated by loop start
         dominators = nx.dominance.immediate_dominators(graph.nx, graph.start_block)
         if begin is ltest: