[Bug][Dataflow] Invalid stream usage in `test_cannon`

[Fangtangtang]

Describe the bug

The program in tests/dataflow/aie/test_cannon.py does not use stream operations correctly. put and get operations for A_pipe and B_pipe are not balanced.

Reproduction

run tests/dataflow/aie/test_cannon.py

Buggy output

The compiled program is functionally correct.

Expected behavior

The program source code should be valid.

Additional context

A validation guard may be required to ensure correct stream operations.

Another related issue: #444

[chhzh123]

This is just a coding style issue?

[Fangtangtang]

no, the program is actually incorrect. There are more A_pipe.put than A_pipe.get.

[Fangtangtang]

when P=2, should be

    @df.region()
    def top(A: Ty[M, K], B: Ty[K, N], C: Ty[M, N]):
        A_init: Stream[Ty[m, k], 2][P, P]
        B_init: Stream[Ty[k, n], 2][P, P]

        A_pipe: Stream[Ty[m, k], 2][P, P]
        B_pipe: Stream[Ty[k, n], 2][P, P]

        C_pipe: Stream[Ty[k, n], 2][P, P]

        @df.kernel(mapping=[P, P], args=[A, B])
        def init(local_A: Ty[M, K] @ LyA, local_B: Ty[K, N] @ LyB):
            pi, pj = df.get_pid()
            A_init[(pi - pj) % P, pj].put(local_A)
            B_init[pi, (pj - pi) % P].put(local_B)

        @df.kernel(mapping=[P, P])
        def cannon0():
            pi, pj = df.get_pid()

            A_out: Ty[m, k] = A_init[pi, pj].get()
            B_out: Ty[k, n] = B_init[pi, pj].get()

            C_pipe[pi, pj].put(allo.matmul(A_out, B_out))

            A_pipe[(pi - 1) % P, pj].put(A_out)
            B_pipe[pi, (pj - 1) % P].put(B_out)

        @df.kernel(mapping=[P, P], args=[C])
        def cannon1(local_C: Ty[M, N] @ LyC):
            pi, pj = df.get_pid()
            A_out: Ty[m, k] = A_pipe[pi, pj].get()
            B_out: Ty[k, n] = B_pipe[pi, pj].get()
            local_C[:, :] = C_pipe[pi, pj].get() + allo.matmul(A_out, B_out)

[chhzh123]

ok, it should be updated

[Fangtangtang]

ok, I'll fix this. I plan to add stream operation validation at the IR level after #510 is completed and the dataflow semantics are encoded into Allo IR.