[AIE2P][PostPipeliner] II=18 vs chess II=11 on bf16/BFP16 matmul-accumulate loop — false loop-carried anti-deps from pipelining-unaware register allocation

[erwei-xilinx]

Summary

A bf16 matmul-accumulate loop (the QK step of a flash-attention kernel: 4 accumulators, 2 transposed-B operands per iteration) is software-pipelined to II=11 by chess but only II=18 by Peano's post-RA AIEPostPipeliner on AIE2P — ~64% worse steady-state, and ~5% slower end-to-end on the full attention kernel on hardware (Strix/NPU2).

The binding constraint is a false loop-carried anti-dependency created by pipelining-unaware register allocation: register allocation runs before the post-RA pipeliner and reuses a register for a loaded operand that aliases a live MAC operand, even though free registers are available.

Environment

• Target: aie2p-none-unknown-elf (Strix / AIE2P)
• Toolchains compared: xchesscc (aie2p) vs Peano clang++
• Reproduced on Peano 742b6c9b19d52b9c851ee30198e1466906bd19aa and aie-public bf0a7459
• Flags: -O2 -std=c++20 -DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16

Minimal reproducer (mm2.cc)

#include <aie_api/aie.hpp>
using bf16 = bfloat16;
#ifndef __chess__
#define chess_prepare_for_pipelining
#define chess_loop_range(...)
#endif
static constexpr unsigned COLA = 16;

extern "C" void qk_accum(bf16 *__restrict pQ1, bf16 *__restrict pQ2,
                         bf16 *__restrict pK1, bf16 *__restrict pK2,
                         bf16 *__restrict pSout) {
    using MMUL = aie::mmul<8, 8, 8, bf16, bf16, accauto>;
    auto Q0 = aie::load_v<MMUL::size_A>(pQ1); pQ1 += MMUL::size_A;
    auto Q1 = aie::load_v<MMUL::size_A>(pQ2); pQ2 += MMUL::size_A;
    auto K0 = aie::transpose(aie::load_v<MMUL::size_B>(pK1), 8, 8); pK1 += MMUL::size_B*2;
    auto K1 = aie::transpose(aie::load_v<MMUL::size_B>(pK2), 8, 8); pK2 += MMUL::size_B*2;
    MMUL C00(aie::zeros<bf16,MMUL::size_C>()), C01(aie::zeros<bf16,MMUL::size_C>()),
         C10(aie::zeros<bf16,MMUL::size_C>()), C11(aie::zeros<bf16,MMUL::size_C>());
    C00.mac(Q0,K0); C01.mac(Q0,K1); C10.mac(Q1,K0); C11.mac(Q1,K1);
    for (unsigned i = 1; i < COLA; ++i)
    chess_prepare_for_pipelining chess_loop_range(COLA,) {
        Q0 = aie::load_v<MMUL::size_A>(pQ1); pQ1 += MMUL::size_A;
        Q1 = aie::load_v<MMUL::size_A>(pQ2); pQ2 += MMUL::size_A;
        K0 = aie::transpose(aie::load_v<MMUL::size_B>(pK1), 8, 8); pK1 += MMUL::size_B*2;
        K1 = aie::transpose(aie::load_v<MMUL::size_B>(pK2), 8, 8); pK2 += MMUL::size_B*2;
        C00.mac(Q0,K0); C01.mac(Q0,K1); C10.mac(Q1,K0); C11.mac(Q1,K1);
    }
    aie::store_v(pSout,                   C00.to_vector<bf16>());
    aie::store_v(pSout +   MMUL::size_C,  C01.to_vector<bf16>());
    aie::store_v(pSout + 2*MMUL::size_C,  C10.to_vector<bf16>());
    aie::store_v(pSout + 3*MMUL::size_C,  C11.to_vector<bf16>());
}

# Peano  -> steady-state loop II = 18 bundles
clang++ -O2 -std=c++20 --target=aie2p-none-unknown-elf -DNDEBUG \
  -DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16 -I $AIEOPT_DIR/include -c mm2.cc -o p.o

# chess  -> steady-state loop II = 11 bundles
xchesscc_wrapper aie2p -DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16 -c mm2.cc -o c.o

A 1-accumulator control loop pipelines better on Peano than chess (II=2 vs 3), so this is specific to register-heavy loops, not a blanket deficiency.

Root cause (-mllvm -debug-only=postpipeliner)

PostPipeliner: ResMII=8
Backedge 29 -> 5  circuit=15  kind=Anti  reg=$x3  lat=-6
   SRC[29]: $dm3 = VMAC_f ... $ex3, $ex5 ...      (accumulator MAC reads ex3)
   DST[5]:  $x3  = VLDB ... [p3, 256]             (next-iteration K load writes x3)
RecMII=15   =>  II=18

ResMII=8, but the post-pipeliner is bound by RecMII=15. The binding recurrence is a false loop-carried anti-dependency: the BFP16 conversion operand $ex3 aliases the K-load destination $x3, so the next iteration's load is serialized behind this iteration's MAC. The allocator reused $x3 even though x0/x1/x2 are free in the loop body. chess co-schedules and allocates, so it never forms this aliasing and reaches II=11.

Why this resists the obvious workarounds (all verified)

• Pruning the anti-dependency is not sufficient. Hand-renaming $ex3→$ex0 (a free register) in the pre-postmisched MIR drops RecMII 15 -> 3, but II stays 18: scheduleOtherIterations then fails for every II in 8..17. The loop is register-footprint-bound (~9 of 12 X-registers live per iteration), so the post-RA pipeliner cannot overlap iterations without renaming / modulo variable expansion, which it does not perform. A naive second rename ($ex6→$ex1) actually raised RecMII to 22, i.e. ad-hoc renaming without liveness awareness creates new conflicts — this needs pipeline-aware allocation, not dependence pruning.

• A feasible pre-RA schedule exists but is discarded. The pre-RA MachinePipeliner (SMS) finds a valid II=8 schedule (rec=1, no recurrence problem on virtual registers), but ZeroOverheadLoop::shouldUseSchedule -> preferPostPipeliner() rejects it and defers to the post-pipeliner, which then lands at II=18.

Ask

Make register allocation pipeline-aware for these loops (or add post-pipeline register renaming / modulo variable expansion to AIEPostPipeliner) so register-heavy modulo loops reach chess-comparable II. Relatedly, reconsider preferPostPipeliner() discarding a feasible pre-RA II=8 SMS schedule in favor of a post-RA schedule that is markedly worse.

I can attach full -debug-only=postpipeliner/pipeliner logs and a self-contained build/compare script on request.

[F-Stuckmann]

We are working on improving postpipeliner performance on this gemm variant.

One note on the provided source code: typically manually peeling the first iteration only increases Program Memory and overall decreases Performance because a) the prologue to the postpipeliner gets larger than it should be (PM impact) and b) the steady state pipeline gets the highest performance for an iteration compared to running an iteration with no overlaps (your Performance penalty are 25 cycles for the peeled first iteration vs 7-8 of the postpipeliner). For the Performance considerations of peeled vs non-peeled please check out the code diff of %bb.0 between reference.txt and no-peel.txt.
Nonetheless you are reporting an existing postpipeliner shortcoming and we are activly investigating it.

Additionally, In your case peeling the first iteration has an additional penalty that pointer updates cannot be properly combined and thus we end up with unnecessary ptradds in the loop which limit the final performance (for your report they are irrelevant since the Accumulator registers are the issue, but they limit the improvements that are possible and are probably also performance penalties in your reference as well). See attached image for the unnecessary ptradds.

Is there a reason you peeled the first iteration out of the gemm loop?