Simplify returns an erroneous SDFG

[affifboudaoud]

Describe the bug
I am using the master branch with commit f01b9 (latest). With this version, I try to run a simple program with 3 nested-loops that calculates the result of matrix-matrix-multiplication of the X and Y inputs and puts the result in Z. I try to store all the temporary values of the multiplication in the matrix T. The sum of all the elements of the matrix Z is then put in S.

If you run this program with simplify=False the result of the matrix multiplication is correct. However, if simplify is called on this SDFG, the inputs X and Y are discarded completely and the output is wrong. This can be visualized in the simplified SDFG.

To Reproduce
Here is a test that exposes the issue.

import dace
import numpy as np

size = 32


@dace.program()
def mmm_dace(
    X: dace.float32[size, size],
    Y: dace.float32[size, size],
    Z: dace.float32[size, size],
    S: dace.float32[1],
):
    T: dace.float32[size, size, size] = np.zeros((size,size,size), dtype=dace.float32)
    
    for i in dace.map[0: size]:
        for j in dace.map[0: size]:
            for k in dace.map[0: size]:
                T[i,j,k] = X[i,k] * Y[k,j]
            Z[i,j] = np.sum(T[i,j])
            

    @dace.map(_[0:size, 0:size])
    def summap(i, j):
        s >> S(1, lambda x, y: x + y)[0]
        z << Z[i, j]
        s = z
    
def test_simplify_mmm():
    # Input initialization
    X  = np.random.rand(size, size).astype(np.float32)
    Y  = np.random.rand(size, size).astype(np.float32)
    Z  = np.zeros((size,size), dtype=np.float32)
    S  = np.zeros((1,), dtype=np.float32)
    
    sdfg = mmm_dace.to_sdfg(simplify=False)
    
    sdfg(X=X, Y=Y, Z=Z, S=S)
    
    # Numerically validate the results of the non-simplified SDFG
    assert np.allclose(Z, np.matmul(X, Y))
    assert np.allclose(S, sum(sum(np.matmul(X, Y))))
    
    # Simplify the SDFG
    sdfg.simplify()

    # Input reinitialization just in case
    X  = np.random.rand(size, size).astype(np.float32)
    Y  = np.random.rand(size, size).astype(np.float32)
    Z  = np.zeros((size,size), dtype=np.float32)
    S  = np.zeros((1,), dtype=np.float32)
    
    sdfg(X=X, Y=Y, Z=Z, S=S)

    # Numerically validate the results of the simplified SDFG
    assert np.allclose(Z, np.matmul(X, Y))
    assert np.allclose(S, sum(sum(np.matmul(X, Y))))


if __name__ == "__main__":
    test_simplify_mmm()