test[next-dace]: Add symbolic LoopBlocking tests

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_loop_blocking.py

@@ -1608,20 +1608,45 @@ def test_loop_blocking_direct_access_node_scalar():
     assert all(np.allclose(ref[name], res[name]) for name in ref)
 
 
-def _make_loop_blocking_sdfg_with_everything() -> tuple[
-    dace.SDFG, dace.SDFGState, dace_nodes.MapEntry, dace_nodes.MapEntry
-]:
+def _make_loop_blocking_sdfg_with_everything(
+    symbolic: bool = False,
+) -> tuple[dace.SDFG, dace.SDFGState, dace_nodes.MapEntry, dace_nodes.MapEntry]:
+    """Creates an SDFG that combines everything `LoopBlocking` has to handle.
+
+    Args:
+        symbolic: Describe the ranges of the outer Map and the shapes of the
+            arrays with symbols instead of constants. The blocking parameter,
+            `__i1`, then runs over `(vertical_end - lev):vertical_end`, i.e.
+            both of its bounds are symbolic expressions and its size, `lev`,
+            is symbolic as well. The range of the inner Map is not affected.
+    """
     sdfg = dace.SDFG(util.unique_name("sdfg_with_inner_semi_independent_map"))
     state = sdfg.add_state(is_start_block=True)
 
-    sdfg.add_array("A", shape=(40, 8), dtype=dace.float64, transient=False)
+    if symbolic:
+        for symbol_name in ["horizontal_start", "horizontal_end", "vertical_end", "lev"]:
+            sdfg.add_symbol(symbol_name, dace.int32)
+        hsize, vsize = "horizontal_end", "vertical_end"
+        hrange, vrange = "horizontal_start:horizontal_end", "(vertical_end - lev):vertical_end"
+    else:
+        hsize, vsize = 40, 8
+        hrange, vrange = "0:40", "0:8"
+
+    # `B` is read indirectly, i.e. its indices are computed from `gt_conn_dummy`
+    #  and `ikoffset`, thus the whole array has to be passed into the Map scope
+    #  and not only the part that is described by `hrange` and `vrange`.
+    b_subset = f"0:{hsize}, 0:{vsize}"
+
+    sdfg.add_array("A", shape=(hsize, vsize), dtype=dace.float64, transient=False)
     for name in "BC":
-        sdfg.add_array(name, shape=(40, 8), dtype=dace.float64, transient=False)
-    sdfg.add_array("inc", shape=(40, 3), dtype=dace.float64, transient=False)
-    sdfg.add_array("inc2", shape=(40, 3), dtype=dace.float64, transient=False)
-    sdfg.add_array("gt_conn_dummy", shape=(40, 2), dtype=dace.int32, transient=False)
-    sdfg.add_array("ikoffset", shape=(40,), dtype=dace.int32, transient=False)
-    sdfg.add_array("S", shape=(40,), dtype=dace.float64, transient=False)
+        sdfg.add_array(name, shape=(hsize, vsize), dtype=dace.float64, transient=False)
+    sdfg.add_array("inc", shape=(hsize, 3), dtype=dace.float64, transient=False)
+    sdfg.add_array("inc2", shape=(hsize, 3), dtype=dace.float64, transient=False)
+    sdfg.add_array("gt_conn_dummy", shape=(hsize, 2), dtype=dace.int32, transient=False)
+    # Note that `ikoffset` is indexed with the blocking parameter, `__i1`, thus
+    #  it follows the vertical dimension and not the horizontal one.
+    sdfg.add_array("ikoffset", shape=(vsize,), dtype=dace.int32, transient=False)
+    sdfg.add_array("S", shape=(hsize,), dtype=dace.float64, transient=False)
     sdfg.add_scalar("t", dtype=dace.float64, transient=True)
     sdfg.add_scalar("tt", dtype=dace.float64, transient=True)
     sdfg.add_scalar("ttt", dtype=dace.float64, transient=True)
@@ -1648,9 +1673,9 @@ def _make_loop_blocking_sdfg_with_everything() -> tuple[
     #  enough. The only reason for doing it is, that the Memlets inside and outside
     #  the inner Map scope can refer to different data, and the outside Memlet
     #  is dependent.
-    sdfg.add_array("T", shape=(8,), dtype=dace.float64, transient=True)
+    sdfg.add_array("T", shape=(vsize,), dtype=dace.float64, transient=True)
 
-    me, mx = state.add_map("main_comp", ndrange={"__i0": "0:40", "__i1": "0:8"})
+    me, mx = state.add_map("main_comp", ndrange={"__i0": hrange, "__i1": vrange})
 
     # The inner computation on its own is not useful.
     ime, imx = state.add_map("inner_comp", ndrange={"__inner": "0:3"})
@@ -1711,15 +1736,15 @@ def _make_loop_blocking_sdfg_with_everything() -> tuple[
         code="__out = __in1 + __in2 + __in3 + __in4 + __in5",
     )
 
-    state.add_edge(A, None, me, "IN_A", dace.Memlet("A[0:40, 0:8]"))
+    state.add_edge(A, None, me, "IN_A", dace.Memlet(f"A[{hrange}, {vrange}]"))
     state.add_edge(me, "OUT_A", ime, "IN_A", dace.Memlet("A[__i0, __i1]"))
     me.add_scope_connectors("A")
 
-    state.add_edge(inc, None, me, "IN_inc", dace.Memlet("inc[0:40, 0:3]"))
+    state.add_edge(inc, None, me, "IN_inc", dace.Memlet(f"inc[{hrange}, 0:3]"))
     state.add_edge(me, "OUT_inc", ime, "IN_inc", dace.Memlet("inc[__i0, 0:3]"))
     me.add_scope_connectors("inc")
 
-    state.add_edge(inc2, None, me, "IN_inc2", dace.Memlet("inc2[0:40, 0:3]"))
+    state.add_edge(inc2, None, me, "IN_inc2", dace.Memlet(f"inc2[{hrange}, 0:3]"))
     state.add_edge(me, "OUT_inc2", ime, "IN_inc2", dace.Memlet("inc2[__i0, 0:3]"))
     me.add_scope_connectors("inc2")
 
@@ -1742,11 +1767,11 @@ def _make_loop_blocking_sdfg_with_everything() -> tuple[
 
     state.add_edge(itlet, "__out", t, None, dace.Memlet("t[0]"))
 
-    state.add_edge(B, None, me, "IN_B", dace.Memlet("B[0:40, 0:8]"))
+    state.add_edge(B, None, me, "IN_B", dace.Memlet(f"B[{b_subset}]"))
     state.add_edge(
-        gt_conn_dummy, None, me, "IN_gt_conn_dummy", dace.Memlet("gt_conn_dummy[0:40, 0:2]")
+        gt_conn_dummy, None, me, "IN_gt_conn_dummy", dace.Memlet(f"gt_conn_dummy[{hrange}, 0:2]")
     )
-    state.add_edge(me, "OUT_B", indirectaccesstlet, "__in", dace.Memlet("B[0:40, 0:8]"))
+    state.add_edge(me, "OUT_B", indirectaccesstlet, "__in", dace.Memlet(f"B[{b_subset}]"))
     state.add_edge(
         me,
         "OUT_gt_conn_dummy",
@@ -1772,11 +1797,11 @@ def _make_loop_blocking_sdfg_with_everything() -> tuple[
     state.add_edge(tttt, None, dtlet, "__in4", dace.Memlet("tttt[0]"))
 
     state.add_edge(dtlet, "__out", mx, "IN_C", dace.Memlet("C[__i0, __i1]"))
-    state.add_edge(mx, "OUT_C", C, None, dace.Memlet("C[0:40, 0:8]"))
+    state.add_edge(mx, "OUT_C", C, None, dace.Memlet(f"C[{hrange}, {vrange}]"))
     mx.add_scope_connectors("C")
 
-    state.add_edge(ikoffset, None, me, "IN_ikoffset", dace.Memlet("ikoffset[0:40]"))
-    state.add_edge(me, "OUT_B", indirectaccesstletkoffset, "__in", dace.Memlet("B[0:40, 0:8]"))
+    state.add_edge(ikoffset, None, me, "IN_ikoffset", dace.Memlet(f"ikoffset[{vrange}]"))
+    state.add_edge(me, "OUT_B", indirectaccesstletkoffset, "__in", dace.Memlet(f"B[{b_subset}]"))
     state.add_edge(
         me,
         "OUT_gt_conn_dummy",
@@ -1791,7 +1816,7 @@ def _make_loop_blocking_sdfg_with_everything() -> tuple[
         "__koffset",
         dace.Memlet("ikoffset[__i1]"),
     )
-    state.add_edge(me, "OUT_B", indirectaccesstletkoffset_B, "__in", dace.Memlet("B[0:40, 0:8]"))
+    state.add_edge(me, "OUT_B", indirectaccesstletkoffset_B, "__in", dace.Memlet(f"B[{b_subset}]"))
     state.add_edge(
         me,
         "OUT_gt_conn_dummy",
@@ -1819,24 +1844,31 @@ def _make_loop_blocking_sdfg_with_everything() -> tuple[
     return sdfg, state, me, ime
 
 
-@pytest.mark.parametrize(
-    "require_independent_nodes,promote_independent_memlets,independent_node_threshold",
-    [
-        (True, True, 0),
-        (True, False, 0),
-        (False, True, 0),
-        (False, False, 0),
-        (True, True, 7),
-        (False, True, 7),
-    ],
-)
-def test_loop_blocking_sdfg_with_everything(
+_LOOP_BLOCKING_EVERYTHING_PARAMS = [
+    (True, True, 0),
+    (True, False, 0),
+    (False, True, 0),
+    (False, False, 0),
+    (True, True, 7),
+    (False, True, 7),
+]
+
+
+def _check_loop_blocking_sdfg_with_everything(
+    sdfg: dace.SDFG,
+    state: dace.SDFGState,
+    me: dace_nodes.MapEntry,
+    ime: dace_nodes.MapEntry,
     require_independent_nodes: bool,
     promote_independent_memlets: bool,
     independent_node_threshold: int,
-):
-    sdfg, state, me, ime = _make_loop_blocking_sdfg_with_everything()
+) -> None:
+    """Applies `LoopBlocking` to `_make_loop_blocking_sdfg_with_everything()` and checks it.
 
+    The checks are shared between the constant and the symbolic version of the
+    SDFG, because the outcome of the transformation must not depend on whether
+    the ranges of the Map are constants or symbolic expressions.
+    """
     scope_dict_before = state.scope_dict()
     assert scope_dict_before[ime] is me
 
@@ -1884,3 +1916,34 @@ def test_loop_blocking_sdfg_with_everything(
     new_scope_of_inner_map = state.scope_dict()[ime]
     assert isinstance(new_scope_of_inner_map, dace_nodes.MapEntry)
     assert new_scope_of_inner_map is not (me if count == 1 else None)
+
+
+@pytest.mark.parametrize(
+    "require_independent_nodes,promote_independent_memlets,independent_node_threshold",
+    _LOOP_BLOCKING_EVERYTHING_PARAMS,
+)
+@pytest.mark.parametrize("symbolic", [False, True], ids=["literal_ranges", "symbolic_ranges"])
+def test_loop_blocking_sdfg_with_everything(
+    require_independent_nodes: bool,
+    promote_independent_memlets: bool,
+    independent_node_threshold: int,
+    symbolic: bool,
+):
+    """Applies `LoopBlocking` to `_make_loop_blocking_sdfg_with_everything()` and checks it.
+
+    If `symbolic` is `True` then the blocking parameter of the outer Map, `__i1`,
+    runs over `(vertical_end - lev):vertical_end`, i.e. both bounds are symbolic
+    expressions and the size of the range, `lev`, is symbolic as well. The outcome
+    of the transformation must not depend on whether the ranges of the Map are
+    constants or symbolic expressions.
+    """
+    sdfg, state, me, ime = _make_loop_blocking_sdfg_with_everything(symbolic=symbolic)
+    _check_loop_blocking_sdfg_with_everything(
+        sdfg,
+        state,
+        me,
+        ime,
+        require_independent_nodes=require_independent_nodes,
+        promote_independent_memlets=promote_independent_memlets,
+        independent_node_threshold=independent_node_threshold,
+    )