Loop Regions (#1407)

This PR adds `LoopRegion`s to SDFGs. This forms the second core element
of the [plan to make loops first class citizens of
SDFGs](https://github.com/orgs/spcl/projects/10).
`LoopRegion`s are a special class of `ControlFlowRegion`s that represent
different types of loops, meaning the control flow region inside them is
executed a parametric number of times.
A `LoopRegion` _must_ have a conditional expression that determines
whether the region is executed or not. It may additionally have:
- an initialization expression that is run when the region is first
executed (even if the condition to execute its contents does not hold)
- an update expression, executed each time the execution of the region
contents is finished
- a flag indicating that it's inverted, meaning the contents are
executed once before the condition is checked (condition is checked
again after every execution)
- a set of control flow blocks in its contents that, when reached,
finish the execution of the region and do not execute it again even if
the condition still holds (equivalent to `break` in C/Python)
- a set of control flow blocks that end the execution of its contents
and execute the update statement if there is one. After this, the
condition is checked again and the contents may be run again (equivalent
to `continue` in C/Python)

For more general information on control flow regions please refer to
[the
documentation](https://spcldace.readthedocs.io/en/latest/sdfg/ir.html#elements)
or the [PR that introduced
them](#1404).

An example of a tripple loop nest of regular for loops can be seen in
the GEMM program below, showing that a proof-of-concept visualization
for the introduced concepts is already available in the latest release
version of the [VSCode
extension](https://marketplace.visualstudio.com/items?itemName=phschaad.sdfv)
(version 1.6.0 and upwards):


![image](https://github.com/spcl/dace/assets/9193712/9955e3f4-3356-4c52-b715-323330a0e4e4)

As outlined by the [project
plan](https://github.com/orgs/spcl/projects/10), these `LoopRegion`s are
currently _not_ being used by any frontend. They can, however, already
be manually used through the SDFG builder API. According to plan, most
passes and transformations are not currently able to handle the use of
such `LoopRegion`s yet, so their use is still highly experimental. To
allow traditional transformations and passes to work, as well as to be
able to generate codes for SDFGs containing `LoopRegion`s, a
compatibility pass (`InlineLoopRegions`) is available, together with a
utility function `dace.sdfg.utils.inline_loop_blocks` that inlines any
`LoopRegion`s to traditional state machine loops.

In summary, the PR includes:
- [x] Add `LoopRegion`s to SDFGs, a special control flow region to
represent loops
- [x] Pass / Utility function to remove all loop regions from a graph,
turning them back into regular state machine loops (a.k.a., inlining)
- [x] Tests for inlining of loops
- [x] Tests for the functional correctness of different types of loops
(regular for, while, do-while, do-for) and nested loops (gemm - tripple
nested loop)

---------

Co-authored-by: Tal Ben-Nun <[email protected]>

dace/codegen/codegen.py

@@ -6,7 +6,7 @@
 import dace
 from dace import dtypes
 from dace import data
-from dace.sdfg import SDFG
+from dace.sdfg import SDFG, utils as sdutils
 from dace.codegen.targets import framecode
 from dace.codegen.codeobject import CodeObject
 from dace.config import Config
@@ -178,6 +178,9 @@ def generate_code(sdfg, validate=True) -> List[CodeObject]:
                 shutil.move(f"{tmp_dir}/test2.sdfg", "test2.sdfg")
                 raise RuntimeError('SDFG serialization failed - files do not match')
 
+    # Convert any loop constructs with hierarchical loop regions into simple 1-level state machine loops.
+    # TODO (later): Adapt codegen to deal with hierarchical CFGs instead.
+    sdutils.inline_loop_blocks(sdfg)
 
     # Before generating the code, run type inference on the SDFG connectors
     infer_types.infer_connector_types(sdfg)

dace/sdfg/sdfg.py

@@ -2150,6 +2150,7 @@ def compile(self, output_file=None, validate=True) -> 'CompiledSDFG':
 
         # Importing these outside creates an import loop
         from dace.codegen import codegen, compiler
+        from dace.sdfg import utils as sdutils
 
         # Compute build folder path before running codegen
         build_folder = self.build_folder
@@ -2170,6 +2171,10 @@ def compile(self, output_file=None, validate=True) -> 'CompiledSDFG':
             # if the codegen modifies the SDFG (thereby changing its hash)
             sdfg.build_folder = build_folder
 
+            # Convert any loop constructs with hierarchical loop regions into simple 1-level state machine loops.
+            # TODO (later): Adapt codegen to deal with hierarchical CFGs instead.
+            sdutils.inline_loop_blocks(sdfg)
+
             # Rename SDFG to avoid runtime issues with clashing names
             index = 0
             while sdfg.is_loaded():