[Bug] Dataflow programming pass simulator test but fails at hw_emu test

[RuizeYu05]

Describe the bug

When I try to implement systolic array based flashattention by dataflow programming in Allo. And I've already finished the following computation all in a single systolic array.

1. Score computation (Matrix Multiplication):
   $$S = A B^T$$

2. Row-wise maximum for numerical stability:
   $$m_i = \max_j (S_{ij})$$

3. Scaled exponential (with max-shifting):
   $$P_{ij} = \exp\left((S_{ij} - m_i) \cdot D\right)$$

4. Row-wise sum (Denominator):

$$\mathrm{Sum}_i = \sum_{j} P_{ij}$$

And these computation can both pass simulator and hw_emu tests.

However, after I added the following logic, the systolic array passes simulator tests and fails at hw_emu test. The output matrix O is all zeros.
5. Final Output computation:
$$O = P V$$

Reproduction

My test code is https://github.com/RuizeYu05/FlashAttention-allo/blob/main/systolic/test_systolic_flash_debug.py
(In order to do software simulator test, I mark the allo.exp in above code. The fully implemented code is https://github.com/RuizeYu05/FlashAttention-allo/blob/main/systolic/test_weight_sta_exp.py)

Buggy output

The following page shows software simulator can pass:

The following page shows the output of hw_emu test:

Expected behavior

Since this systolic array can pass software simulator test, it should also pass hw_emu test.

Additional context

No response

[zzzDavid]

A fix for this issue has been submitted in #574.

Root cause: In the generated HLS top() function, kernel calls followed the user's Python definition order. In this design, offchip_storeS (a pure stream consumer) was emitted before the gemm PE kernels that produce the L_S stream data. Vitis HLS #pragma HLS dataflow requires canonical forward-flow ordering (producers before consumers), so the misordering caused hw_emu to produce all-zero outputs for the O matrix. The simulator was unaffected because it uses concurrent threads with blocking FIFOs.

Fix: Sort kernel calls in _build_top() by stream direction — pure producers first, mixed kernels (PEs) in the middle, pure consumers last.

Verified with the exact reproducer from https://github.com/RuizeYu05/FlashAttention-allo/blob/main/systolic/test_systolic_flash_debug.py — both simulator and hw_emu now pass.

[zzzDavid]

Update: I've now also verified the full version with allo.exp (from test_weight_sta_exp.py) against hw_emu with the fix in #574. Both Sum and O matrix are correct — not all zeros.

To be clear, the fix is in the allo framework (allo/dataflow.py), not in the user's design code. Simply reordering kernel definitions in the Python source won't help — the fix sorts kernel calls at the MLIR level in _build_top() based on stream direction analysis.

To test with the fix:

1. Apply [Fix] Sort dataflow kernel calls so producers precede consumers #574 to your allo installation
2. Run your design — no changes needed to your systolic array code

@RuizeYu05 could you try applying the patch from #574 and re-running your hw_emu test?

[zzzDavid]

How to test the fix:

# Clone the fork with the fix and switch to the fix branch
git clone https://github.com/zzzDavid/allo.git
cd allo
git checkout fix/dataflow-kernel-ordering

# Install into your conda env (assuming LLVM is already set up)
pip install -e .

# Run your test — no changes needed to your design code
python test_weight_sta_exp.py

Alternatively, if you want to patch your existing allo installation in-place, the change is a single function in allo/dataflow.py. In _build_top(), right before the # add call functions comment (~line 514), add:

        # Sort kernel functions so that producers come before consumers.
        def _stream_order_key(func_op):
            fname = func_op.attributes["sym_name"].value
            dirs = {d for _, d in stream_info.get(fname, [])}
            if dirs == {"out"}:
                return 0  # pure producer  -> first
            if dirs == {"in"}:
                return 2  # pure consumer  -> last
            return 1  # mixed / no streams -> middle

        sorted_funcs = sorted(funcs, key=_stream_order_key)

Then change the loop from for i, func in enumerate(funcs): to for i, func in enumerate(sorted_funcs): and similarly len(funcs) → len(sorted_funcs).