[Fix][Codegen][AIE] Fix Problematic Optimization on Stream Operations

[Fangtangtang]

Description

This PR partially fixes #601

Problems

Previous optimization on stream operation lowering is too aggressive and didn't consider inter-compute-tile fifo reuse. The leads to the bug mentioned in #601

Proposed Solutions

(What changes have you made)

Examples

import os
import allo
import allo.dataflow as df
from allo.backend.aie import is_available
from allo.ir.types import int16, Stream
import numpy as np

Ty = int16
M, N, K = 64, 16, 16

def make_atb_top(rho):
    assert M % rho == 0
    Ma = M // rho

    @df.region()
    def top(A: Ty[M, K], B: Ty[K, N], C: Ty[M, N]):
        pipe_a: Stream[Ty[Ma, K], 2][rho]
        pipe_b: Stream[Ty[K, N], 2][rho]
        pipe_c: Stream[Ty[Ma, N], 2][rho]

        @df.kernel(mapping=[1], args=[A])
        def load_a(local_A: Ty[M, K]):
            with allo.meta_for(rho) as i:
                pipe_a[i].put(local_A[i * Ma : (i + 1) * Ma, :])

        @df.kernel(mapping=[1], args=[B])
        def load_b(local_B: Ty[K, N]):
            with allo.meta_for(rho) as i:
                pipe_b[i].put(local_B)

        @df.kernel(mapping=[rho])
        def compute():
            pk = df.get_pid()
            local_A: Ty[Ma, K] = pipe_a[pk].get()
            local_B: Ty[K, N] = pipe_b[pk].get()
            pipe_c[pk].put(allo.matmul(local_A, local_B))

        @df.kernel(mapping=[1], args=[C])
        def store_c(local_C: Ty[M, N]):
            with allo.meta_for(rho) as i:
                local_C[i * Ma : (i + 1) * Ma, :] = pipe_c[i].get()

    return top


def run_atb(rho):
    top = make_atb_top(rho)
    mapping_primitives = None
    if rho > 1:
        mapping_primitives = [("bundle", [f"compute_{i}" for i in range(rho)])]

    A = np.random.randint(0, 64, (M, K)).astype(np.int16)
    B = np.random.randint(0, 64, (K, N)).astype(np.int16)
    C = np.zeros((M, N)).astype(np.int16)

    if is_available():
        os.environ["FORCE_UNROLL_INDEX"] = "1"
        mod = df.build(top, target="aie", mapping_primitives=mapping_primitives)
        mod(A, B, C)
        del os.environ["FORCE_UNROLL_INDEX"]
        np.testing.assert_allclose(C, A @ B, atol=1e-5)
        print(f"rho={rho} PASSED!")
    else:
        print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.")

run_atb(2)

Checklist

Please make sure to review and check all of these items:

• PR's title starts with a category (e.g. [Bugfix], [IR], [Builder], etc)
• All changes have test coverage (It would be good to provide ~2 different test cases to test the robustness of your code)
• Pass the formatting check locally
• Code is well-documented

[Copilot]

Pull request overview

This PR targets an AIE backend miscompile related to overly aggressive stream-operation lowering, particularly when FIFOs are reused across compute tiles (as described in issue #601). It refactors parts of the stream op handling in the AIE MLIR codegen and adds a focused regression test for multi-subtile (rho > 1) ATB dataflow.

Changes:

• Adjust AIE stream op lowering logic in mlir_codegen.py to avoid incorrect optimization in inter-compute-tile FIFO reuse scenarios.
• Add a new AIE dataflow regression test covering rho in {1,2,4} with bundling and forced unroll behavior.
• Minor test refactor/rename in existing AIE matrix test to keep naming consistent.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
allo/backend/aie/mlir_codegen.py	Refines stream op classification and FIFO selection logic during lowering to mitigate incorrect optimizations with FIFO reuse.
tests/dataflow/aie/test_matrix.py	Renames a region function (top1 → top) in an existing ATB GEMM test.
tests/dataflow/aie/test_mapping_atb.py	Adds a regression test for ATB mapping/bundling across multiple rho values to catch the #601 symptom.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

Suppressed comments (5)

allo/backend/aie/mlir_codegen.py:449

• inter_ct_fifo processing assumes every collected argument has at least one use (uses_[0]) and that the first use is a recognized stream op. If an argument is unused (or its first use is not memref.store/load/copy), this will raise IndexError or propagate is_put=None into the fifo tuple selection. Add guards and skip empty/unrecognized arguments, and skip the fifo group entirely if no stream uses remain.

                for arg in args:
                    uses_ = list(arg.uses)
                    is_put, is_tensor = check_stream_op_type(arg, uses_[0].owner)
                    uses.extend(uses_)

tests/dataflow/aie/test_mapping_atb.py:71

• This PR targets incorrect behavior for rho > 1 and the original report includes a failing case at rho = 8. Adding 8 here helps ensure the regression is actually covered for the worst-case subtile factor.

@pytest.mark.parametrize("rho", [1, 2, 4])

tests/dataflow/aie/test_mapping_atb.py:64

• If df.build(...) or mod(...) raises, FORCE_UNROLL_INDEX will remain set for the rest of the pytest session and can affect unrelated tests. Use try/finally and restore any prior value instead of unconditionally deleting.

        os.environ["FORCE_UNROLL_INDEX"] = "1"
        mod = df.build(top, target="aie", mapping_primitives=mapping_primitives)
        mod(A, B, C)
        del os.environ["FORCE_UNROLL_INDEX"]

tests/dataflow/aie/test_mapping_atb.py:111

• If df.build(...) or mod(...) raises, FORCE_UNROLL_INDEX will remain set for the rest of the pytest session and can affect unrelated tests. Use try/finally and restore any prior value instead of unconditionally deleting.

        os.environ["FORCE_UNROLL_INDEX"] = "1"
        mod = df.build(top, target="aie", mapping_primitives=mapping_primitives)
        mod(A, B, C)
        del os.environ["FORCE_UNROLL_INDEX"]

tests/dataflow/aie/test_mapping_atb.py:17

• This PR targets incorrect behavior for rho > 1 and the original report includes a failing case at rho = 8. Adding 8 here helps ensure the regression is actually covered for the worst-case subtile factor.

This issue also appears on line 71 of the same file.

@pytest.mark.parametrize("rho", [1, 2, 4])