diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_loop_blocking.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_loop_blocking.py index 0b529919c0..5b8f9e1133 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_loop_blocking.py +++ b/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, + )