Simplify interferes with dynamic map ranges

[philip-paul-mueller]

I use dynamic map ranges, to add a offset to a Memlet, which is only known at runtime.
However, I have to perform some operations on this value, before I use it.
For this I use a Tasklet, which is then connected to the Map entry, see picture below.

Before the SDFG is simplified it validates and works as expected.
After simplification it fails to compile.

I found out that simplify will transform the Tasklets into a symbol assignment on an interstate edge.
The problem is that the name used for the dynamic map range symbol, i.e. the connector name on the Map entry, is a mangle version of the data container that actually stores the value.
To make an example, assume that the value of the dynamic offset is stored inside the data container with name dyn_offset_variable and is connected, by a Memlet, to the Map entry and there has the connector name __dyn_offset_symbol.

However, when simplify removes the Tasklet, it will also remove the data container dyn_offset_variable.
Instead it will create a symbol with the name dyn_offset_variable and initialize it with the content (string used as code) of the Tasklet.
The actual problem is that it does not change __dyn_offset_symbol, which is used inside the Map scope, to dyn_offset_variable.
This essentially turns __dyn_offset_symbol into a free symbol which has to be supplied at runtime, never mind that the value was already passed.
Furthermore, DaCe will also complain that the symbol is not inside sdfg.symbols but in sdfg.free_symbols.

My current solution is to simply not mangle the name and always use dyn_offset_variable.
For my current application, JaCe, I think that the Tasklets should remain, to ensure that the Map can be fused with others, which the interstate edge blocks.

The current state is undesirable, since simplify breaks the SDFG.
Thus it should at least change the name inside the scope or what I think is the better solution, to actually keep Tasklets.

slicing_optimized.json
slicing_unoptimized.json

[BenWeber42]

Thanks for the bug report. I think we should separate between the bug and the behavior that would be favorable for you:

1. (the bug) Simplify breaks the SDFG as you described above.
2. (behavior that would be favorable for you) You would prefer that simplify keeps the tasklet and doesn't introduce a separate state with an interstate edge that does the symbol assignments.

I think we should keep the discussion for (2) separate from this issue. I.e., this issue should be about (1).

Regarding (1), do you happen to know which pass introduces the problem? I.e., which pass introduces the symbol assignments on the interstate edge? I believe that pass should also be responsible for making sure that downstream usages of the symbol get adjusted -> If you remove an assignment and introduce it elsewhere but with a different name, then usages of that variable should be renamed accordingly.

[philip-paul-mueller]

It is most likely the Scalar to Symbol Transformation.

[tbennun]

I agree issue (2) should be discussed separately. I would personally also prefer to keep a map with dynamic ranges over a state transition.

However, if we are all in agreement on that, and the issue of scalar-to-symbol promotion is with respect to dynamic map range connectors, then disabling this behavior also fixes the bug in (1). We should add a test, just to be sure.

[philip-paul-mueller]

I think we agree about that.
Here is also a short test example, that can be used ad basis of a test:

import numpy as np
import dace
#================= Reference implementation


def ref_impl(A: np.ndarray, S1: int) -> np.ndarray:
    return A[(S1 + 2):(S1 + 2 + 5), (S1 + 2):(S1 + 2 + 5)]


#================= SDFG implementation


sdfg = dace.SDFG("TEST")

sdfg.add_array("A", shape=(10, 10), dtype=dace.float64, transient=False)
sdfg.add_array("__return", shape=(5, 5), dtype=dace.float64, transient=False)

sdfg.add_scalar("S1", dtype=dace.int32, transient=False)
sdfg.add_scalar("S2", dtype=dace.int32, transient=True)

state = sdfg.add_state(label="TestState", is_start_block=True)

inputs = { "__in":
          dace.Memlet.simple("A", "__i1 + offset, __i2 + offset")
}
outputs = { "__out":
          dace.Memlet.simple("__return", "__i1, __i2")
}

_, mentry, _ = state.add_mapped_tasklet(
    name="Test",
    map_ranges=[("__i1", "0:5"), ("__i2", "0:5")],
    inputs=inputs,
    outputs=outputs,
    code="__out = __in",
    external_edges=True,
)

tskl = dace.nodes.Tasklet(
    "modifier",
    code="__out = __in + 2",
    inputs={"__in": None},
    outputs={"__out": None},
)

S1 = state.add_access("S1")
S2 = state.add_access("S2")


state.add_edge(
    S1,
    None,
    tskl,
    "__in",
    dace.Memlet.simple("S1", "0"),
)
state.add_edge(
    tskl,
    "__out",
    S2,
    None,
    dace.Memlet.simple("S2", "0"),
)
state.add_edge(
    S2,
    None,
    mentry,
    "offset",
    dace.Memlet.simple("S2", "0"),
)
mentry.add_in_connector("offset")

sdfg.validate()

#================= Perform test

A = np.array(np.random.random((10, 10)), dtype=np.float64)
S1 = 2

for i in range(2):
    if i == 1:
        sdfg.simplify()

    res = sdfg(A=A, S1=S1)
    ref = ref_impl(A=A, S1=S1)
    
    assert np.all(res == ref)
    if i == 0:
        print("Passed without simplify.")
    else:
        print("Passed with simplify.")