Add loop regions to the frontend's capabilities

dace/frontend/python/newast.py

@@ -1,4 +1,4 @@
-# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved.
+# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved.
 import ast
 from collections import OrderedDict
 import copy
@@ -32,7 +32,7 @@
 from dace.memlet import Memlet
 from dace.properties import LambdaProperty, CodeBlock
 from dace.sdfg import SDFG, SDFGState
-from dace.sdfg.state import ControlFlowBlock, LoopRegion, ControlFlowRegion
+from dace.sdfg.state import BreakBlock, ContinueBlock, ControlFlowBlock, LoopRegion, ControlFlowRegion
 from dace.sdfg.replace import replace_datadesc_names
 from dace.symbolic import pystr_to_symbolic, inequal_symbols
 
@@ -2456,11 +2456,16 @@ def visit_While(self, node: ast.While):
         if test_region is not None:
             iter_end_blocks = set()
             for n in loop_region.nodes():
-                if isinstance(n, LoopRegion.ContinueState):
-                    iter_end_blocks.add(n)
-                    # If it needs to be connected back to the test region, it does no longer need
-                    # to be handled specially and thus is no longer a special continue state.
-                    n.__class__ = SDFGState
+                if isinstance(n, ContinueBlock):
+                    # If it needs to be connected back to the test region, it does no longer need to be handled
+                    # specially and thus is no longer a special continue state. Add an empty state and redirect the
+                    # edges leading into the continue into it.
+                    replacer_state = loop_region.add_state()
+                    iter_end_blocks.add(replacer_state)
+                    for ie in loop_region.in_edges(n):
+                        loop_region.add_edge(ie.src, replacer_state, ie.data)
+                        loop_region.remove_edge(ie)
+                    loop_region.remove_node(n)
             for inner_node in loop_region.nodes():
                 if loop_region.out_degree(inner_node) == 0:
                     iter_end_blocks.add(inner_node)
@@ -2509,7 +2514,7 @@ def _generate_orelse(self, loop_region: LoopRegion, postloop_block: ControlFlowB
         did_break_symbol = 'did_break_' + loop_region.label
         self.sdfg.add_symbol(did_break_symbol, dace.int32)
         for n in loop_region.nodes():
-            if isinstance(n, LoopRegion.BreakState):
+            if isinstance(n, BreakBlock):
                 for iedge in loop_region.in_edges(n):
                     iedge.data.assignments[did_break_symbol] = '1'
         for iedge in self.cfg_target.in_edges(loop_region):
@@ -2528,8 +2533,7 @@ def _generate_orelse(self, loop_region: LoopRegion, postloop_block: ControlFlowB
 
     def visit_Break(self, node: ast.Break):
         if isinstance(self.cfg_target, LoopRegion):
-            break_state = self.cfg_target.add_state('break_%s' % node.lineno, is_break=True)
-            self._on_block_added(break_state)
+            self._on_block_added(self.cfg_target.add_break(f'break_{self.cfg_target.label}_{node.lineno}'))
         else:
             error_msg = "'break' is only supported inside loops "
             if self.nested:
@@ -2539,8 +2543,7 @@ def visit_Break(self, node: ast.Break):
 
     def visit_Continue(self, node: ast.Continue):
         if isinstance(self.cfg_target, LoopRegion):
-            continue_state = self.cfg_target.add_state('continue_%s' % node.lineno, is_continue=True)
-            self._on_block_added(continue_state)
+            self._on_block_added(self.cfg_target.add_continue(f'continue_{self.cfg_target.label}_{node.lineno}'))
         else:
             error_msg = ("'continue' is only supported inside loops ")
             if self.nested:
@@ -4685,6 +4688,10 @@ def visit_Return(self, node: ast.Return):
             ast_name = ast.copy_location(ast.Name(id='__return'), node)
             self._visit_assign(new_node, ast_name, None, is_return=True)
 
+        if not isinstance(self.cfg_target, SDFG):
+            # In a nested control flow region, a return needs to be explicitly marked with a return block.
+            self._on_block_added(self.cfg_target.add_return(f'return_{self.cfg_target.label}_{node.lineno}'))
+
     def visit_With(self, node, is_async=False):
         # "with dace.tasklet" syntax
         if len(node.items) == 1:

dace/frontend/python/preprocessing.py

@@ -935,6 +935,9 @@ def _add_exits(self, until_loop_end: bool, only_one: bool = False) -> List[ast.A
         for stmt in reversed(self.with_statements):
             if until_loop_end and not isinstance(stmt, (ast.With, ast.AsyncWith)):
                 break
+            elif not until_loop_end and isinstance(stmt, (ast.For, ast.While)):
+                break
+
             for mgrname, mgr in reversed(self.context_managers[stmt]):
                 # Call __exit__ (without exception management all three arguments are set to None)
                 exit_call = ast.copy_location(ast.parse(f'{mgrname}.__exit__(None, None, None)').body[0], stmt)

dace/sdfg/sdfg.py

@@ -30,7 +30,7 @@
 from dace.frontend.python import astutils, wrappers
 from dace.sdfg import nodes as nd
 from dace.sdfg.graph import OrderedDiGraph, Edge, SubgraphView
-from dace.sdfg.state import SDFGState, ControlFlowRegion
+from dace.sdfg.state import ControlFlowBlock, SDFGState, ControlFlowRegion
 from dace.sdfg.propagation import propagate_memlets_sdfg
 from dace.distr_types import ProcessGrid, SubArray, RedistrArray
 from dace.dtypes import validate_name
@@ -183,7 +183,7 @@ class InterstateEdge(object):
                            desc="Assignments to perform upon transition (e.g., 'x=x+1; y = 0')")
     condition = CodeProperty(desc="Transition condition", default=CodeBlock("1"))
 
-    def __init__(self, condition: CodeBlock = None, assignments=None):
+    def __init__(self, condition: Optional[Union[CodeBlock, str, ast.AST, list]] = None, assignments=None):
         if condition is None:
             condition = CodeBlock("1")
 
@@ -2686,3 +2686,15 @@ def make_array_memlet(self, array: str):
            :return: a Memlet that fully transfers array
         """
         return dace.Memlet.from_array(array, self.data(array))
+
+    def recheck_using_experimental_blocks(self) -> bool:
+        found_experimental_block = False
+        for node, graph in self.root_sdfg.all_nodes_recursive():
+            if isinstance(graph, ControlFlowRegion) and not isinstance(graph, SDFG):
+                found_experimental_block = True
+                break
+            if isinstance(node, ControlFlowBlock) and not isinstance(node, SDFGState):
+                found_experimental_block = True
+                break
+        self.root_sdfg.using_experimental_blocks = found_experimental_block
+        return found_experimental_block