[AIE2] Partial-word store scheduled into zero-overhead-loop end bundle is dropped on hardware

[FIM43-Redeye]

While building an open-source AIE2 emulator/validator we found what looks like an AIE2 codegen correctness bug in Peano, and wanted to flag it in case it's useful. Happy to help with a reduced test case or a fix if you'd like.

Summary

For simple scalar loops with a sub-word (int8_t/int16_t) output store, Peano's unrolling places the final st.s8/st.s16 in the loop-end (LE) bundle of a zero-overhead (hardware) loop. On real Phoenix (NPU1/AIE2) silicon that store is dropped on every back-edge -- only the final fall-through iteration commits -- so every Nth output element keeps its initial value. xchesscc compiles the same source correctly (it does not park the store at LE).

Reproducer

#include <stdint.h>
extern "C" void k(int8_t* __restrict out) {
    for (int i = 0; i < 128; i++)
        out[i] = i * i;        // int8
}

Built with clang --target=aie2-none-unknown-elf -O2 -c k.cc -o k.o and disassembled with llvm-objdump -d --triple=aie2 k.o, the core loop is a zero-overhead loop whose LE bundle is the st.s8:

00000020 <.LBB0_1>:                       ; loop start (LS)
      20: ... mul r1, r0, r0 ; mov dj0, r0
      30: st.s8 r1, [p0, dj0]
      ...
      7a: mul r2, r1, r1                   ; data for the 4th unrolled element
      7e: nop
00000080 <.L_LEnd0>:                       ; loop end (LE)
      80: nopb ; st.s8 r2, [p0, dj0] ; nops ; nopxm ; nopv   <-- parked at LE

Expected out[i] = (i*i) & 0xff. On hardware, every i % 4 == 3 -- the unrolled element whose store lands in the LE bundle -- reads 0, except the final iteration (which falls through with no back-edge):

i	3	7	11	15	...	127
expected	9	49	121	-31	...	1
NPU	0	0	0	0	...	1

Why we think it's the LE-bundle store

• Only the element whose st.s8 is the LE-bundle store is dropped; stores earlier in the loop body commit fine.
• It is specific to partial-word stores (st.s8/st.s16), which are read-modify-write and commit late (AIE2Schedule.td: II_STHB MemoryCycles<[5,11]>, i.e. E11). A full-word st in the same slot is fine.
• xchesscc does not place the partial-word store in the LE bundle and produces correct output.

This suggests the AIE2 scheduler should avoid placing a partial-word store in the LE bundle of a hardware loop (or account for its E11 commit against the back-edge fetch redirect), similar to how it already keeps the ls/le/lc loop-setup writes a fixed distance from LE.

Environment

• Target: AIE2 / Phoenix (NPU1, Ryzen AI)
• Peano: llvm-aie@0b72475 (clang 21.0.0git)
• Reproduced on real NPU hardware; xchesscc output is correct.

Glad to provide the full objdump, more reproducers, or a reduced test case -- just let us know what's most useful.

[martien-de-jong]

Sorry for the late reply. The case sounds weird. I think you're implying that the loop is unrolled by a factor of four. I would like to see the entire loop, including the loop setup (setting ls, le and lc. Also pass --disassemble-zeroes please, so we see all bundles explicitly.

[FIM43-Redeye]

Unrolled by four, yes. Here's the full dump with --disassemble-zeroes. I added -r as well, since in an unlinked object the movxm ls/le operands print as #0x0 and only the relocations show what they point at.

00000000 <k>:
       0: 55 00 e0 00 00 00                                    movxm  ls, #0x0
                        00000000:  R_AIE_42  .LBB0_1
       6: bb 10 00 38 02 00 00 48 00 04                        mova   r1, #0x20;    movxm le, #0x0
                        00000006:  R_AIE_44  .L_LEnd0
      10: c0 03 00 38 43 b8 02 00 00 00 00 08 00 00 00 00      nopb ; mova r0, #0x0 ; nops ; nopx ; add.nc lc, r1, #0x0 ; nopv

00000020 <.LBB0_1>:
      20: c0 03 00 28 0b 10 7e 10 00 00 00 00 00 00 00 00      nopb ; nopa ; nops ; mul r1, r0, r0 ; mov dj0, r0 ; nopv
      30: bb 00 00 00 00 00 00 40 20 02                        st.s8  r1, [p0, dj0];   nopxm
      3a: 01 00                                                nop
      3c: 01 00                                                nop
      3e: 01 00                                                nop
      40: 01 00                                                nop
      42: 01 00                                                nop
      44: 19 07 02 10                                          add    r1, r0, #0x1
      48: 43 28 4b 10 fe 20 02 00                              mul    r2, r1, r1;      mov dj0, r1
      50: 19 04 11 00                                          st.s8  r2, [p0, dj0]
      54: 19 0b 02 10                                          add    r1, r0, #0x2
      58: 59 96 20 1c                                          mov    dj0, r1
      5c: 01 00                                                nop
      5e: 01 00                                                nop
      60: 99 1f 44 10                                          mul    r2, r1, r1
      64: 01 00                                                nop
      66: 1d 0f 02 80 20 02                                    st.s8  r2, [p0, dj0];   add r1, r0, #0x3
      6c: 43 28 4b 10 9a 00 00 00                              add    r0, r0, #0x4;    mov dj0, r1
      74: 01 00                                                nop
      76: 01 00                                                nop
      78: 01 00                                                nop
      7a: 99 1f 44 10                                          mul    r2, r1, r1
      7e: 01 00                                                nop

00000080 <.L_LEnd0>:
      80: c0 03 00 00 00 00 00 00 00 00 00 80 20 02 00 00      nopb ; st.s8 r2, [p0, dj0] ; nops ; nopxm ; nopv
      90: 01 00                                                nop
      92: 19 18 00 10                                          ret lr

ls = .LBB0_1 = 0x20, le = .L_LEnd0 = 0x80, lc = 32 (mova r1, #32 then add.nc lc, r1, #0). 32 x 4 = 128.

Two numbers matter. The body is le - ls = 0x60 = 96 bytes, six 16-byte fetch packets. The setup-to-LE distance is 0x80 - 0x10 = 112 bytes, which is exactly the 7-bundle minimum from AIE2InstrInfo.cpp:1412.

So the setup distance holds, but without any slack, which turns out to matter: when the setup bundle sits immediately before the header, the smallest body the rule permits is 112 - 16 = 96 bytes. Tight loops land in the affected size by construction. That's also why hand-written body-size sweeps didn't reproduce it on our end. Peano parks a register op at LE in most kernels we tried instead of a store.

For a control, same source with #pragma clang loop unroll(disable), nothing else changed. (The on-hardware harness wants a two-arg kernel, so out sits in p1 rather than p0 from here on; the loop is otherwise bundle-for-bundle the same.) lc = 128, ls = 0x20, le = 0x80 -- the same 96-byte body and st.s8 -- but the store moves to 0x30 and the LE bundle is empty:

      30: nopb ; st.s8 r1, [p1, dj0] ; nops ; nopxm ; nopv
      ...
00000080 <.L_LEnd0>:
      80: nopb ; nopa ; nops ; nopxm ; nopv

On hardware today (Phoenix NPU1, out zero-initialised):

• Unrolled: 31 of 128 elements come back zero, exactly i % 4 == 3 for i = 3..123. i = 127 falls through with no back-edge and commits. Identical over four runs.
• Control: exact match, no errors.

So it isn't the body size since both are 96 bytes, and it isn't st.s8 on its own since both use it. The one thing that changes is whether the partial-word store sits in the LE bundle. xchesscc lands on the same shape as the control: no unroll, 96-byte body, store at LS, LE bundle all nops.

The read-modify-write width doesn't seem to be what separates them either, if that's worth ruling out. In the unrolled version all four st.s8s of an iteration hit the same aligned 32-bit word (bytes 4k..4k+3); three commit and only the one at LE drops. And the control writes those same four bytes across four iterations, so it does the same-word read-modify-write across three back-edges, and comes back exact.

One correction to my first post, I said the store is dropped on every back-edge but that's stronger than what we have. Across ~3,500 store-at-LE kernels fuzzed against silicon, bodies larger than 96 bytes commit (4/4) and 96-byte bodies flush in 107 of 109 cases. The two exceptions are structurally indistinguishable from the flushes: same body, same LE bundle, same producer. Producer latency, body size, and LE-bundle composition each looked like the discriminator and each failed on a larger sample. It's a cycle-exact pipeline effect, and I don't know if any rule stated from the disassembly alone will capture it, 96-byte threshold included.

Whatever the exact rule turns out to be, every case we've seen goes away if the partial-word store isn't in the LE bundle, which is where the control and the Chess output both end up. That's basically what I wanted to bring to your attention.

I went looking for an existing constraint and couldn't find one. verifyDistanceToLoopEnd in AIEMachineAlignment.cpp seems to only enforce the ls/le/lc setup distance, and I didn't see anything else restricting what lands in the last bundle. Am I missing somewhere it's already handled?

Happy to send the xclbins, the control, or a reduced lit test, whichever is useful.

Peano 0b72475 (clang 21.0.0git), Phoenix NPU1, XRT 2.23.0, NPU firmware 1.5.5.391. xchesscc from aietools 2025.2.

[martien-de-jong]

Thank you for the thorough analysis. My penny's worth:

• There's some ambiguity in the spec as to the definition of that distance. For instance, is loop end the value of the le register, or the first bundle after the loop.
• Your case only does the store in the last bundle. No pointer update or other induction stuff.
• The length of the loop is dominated by the partial store safety distance. Regular stores will have totally different bundles and quantization, moving the margin away from the exact 112.
• The loop seems to execute correctly. The observable is that the side-effect of the last bundle is missing. This may be because of a resource conflict on the store hardware or of missing the last bundle altogether. My bet is on the latter.
• the lc setup sits in an aligned bundle. The 112 byte margin is an entire number of bundles, related to fetch enough bundles to bridge the branch overhead. I don't know the details of what gets buffered, but it may lose/overwrite stages of the decoded instruction.
• Different view: partial store has one of the deepest pipelines, in particular larger than 7, which is the number of full bundles sitting in 112 bytes. It is a compromised instruction to be avoided as much as possible, and in practice only used to do once-only set up of kernel parameters. Perhaps there is an undocumented constraint connected to this.

I'll be on holiday next week. For followup I would ask you to ping @khallouh or @katerynamuts, who have been finding out timing details.
Perhaps they know of similar findings in other AIE versions. They are busy though. If you want to experiment yourself, you could try to make PseudoLoopEnd conflict with the part-word store resource that protects the RMW quirk.

[mludevid]

I just had a quick look at this and I'm surprised by the missing cycle distance between the last part word store in an iteration and the first part word store in the next iteration. As far as I know we guarantee an 8 cycle distance with a resource conflict, why is that not being honored on the backedge?

In the unrolled version all four st.s8s of an iteration hit the same aligned 32-bit word (bytes 4k..4k+3);

Who guarantees that though? As far as I can tell the incoming pointer does not have to be 32-bit aligned. I'm not sure what pointer your test is passing in though, it might very well be 32-bit aligned, in which case this may not be the cause for this specific problem. But without alignment guarantees on that int8 pointer I would expect the scheduler to have to keep the 8 cycle distance between all part word stores to make sure we don't loose any data.

[FIM43-Redeye]

Good catch, thanks much- int8* carries no alignment guarantee, yup. Unclear phrasing on my part. In this particular test the objectfifo output buffer happens to land 128-byte aligned (0x72000), so the four st.s8s of an iteration do all hit one 32-bit word, BUT that's the allocator instead of a promise from source.

Worth also noting that this doesn't look like unaligned-straddle data loss, since the pointer's cleanly aligned. The drop happens specifically with a perfectly aligned buffer.

On your actual question: I went cycle-counting and it lines up with your surprise. Within an iteration, the part-word stores are ~8 cycles apart (you can see the nop padding between them in the dump). But across the back edge, the LE-bundle store and the next iteration's first store land only ~2 cycles apart. So the 8-cycle distance is honored in the linear schedule and collapses on the wrap -- the scheduler seems to treat the LE store as the last part-word store before fall-through and pads nothing after it.

Whether that ~2-cycle collision is itself the cause, or whether it's more about the fetch redirect (the >96-byte bodies committing points somewhere in that direction), is past what I can tell from the disassembly -- leaving that for @khallouh / @katerynamuts. In the meantime I'm going to try martien's suggestion on my end, making PseudoLoopEnd conflict with the part-word-store resource, and will report back what I find on the same hardware setup.

[FIM43-Redeye]

Following up on the scheduler side, since I had the pieces to dig in. The short version is that the loop-aware scheduler does see a back-edge hazard here, but a partial-word store landing in the loop-end bundle isn't part of what it treats as illegal, so it converges with the store there. Details and a question below, plus a couple of fix directions if it's useful.

What the scheduler does today

Compiling the reproducer with -debug-only=loop-aware, the interblock scheduler flags the loop-end st.s8 as a back-edge resource conflict:

Conflicting MI at Bottom-up cycle=0: ST_S8_ag_idx renamable $r2 ... (store (s8) into %ir.arrayidx.3)

arrayidx.3 is the i+3 element, the store that ends up in the .L_LEnd0 bundle. (The first store, arrayidx, is also flagged from the other side, but it's the .3 one that lands at LE.)

It then iterates, bumping the per-instruction latency margin, and converges properly rather than by hitting a cap: the resource margin stays 0, it runs well under MaxExpensiveIterations (35), and there's no software-pipelining fallback. The final schedule it considers legal still has the store in the loop-end bundle.

So this looks like the scheduler resolving the resource hazard it does model, the 7-cycle PART_WORD_STORE occupation (AIE2Schedule.td:138, InstrStage<7, [PART_WORD_STORE]>), while "a partial-word store occupying the ZOL end bundle" just isn't a condition in its legality model.

The part I can't reconcile from the outside

The PART_WORD_STORE resource is meant to keep partial-word stores ~7 cycles apart (matching the E5-load/E11-store RMW window). But in the emitted loop the LE-bundle store sits ~2 bundles from the next iteration's first store across the back-edge, which is inside that window. So either the overlay isn't enforcing the 7-cycle spacing across the wrap the way it does within the block, or the LE-bundle case falls outside what it checks. Which is it? I can't tell from the disassembly and it's probably where a fix belongs.

Why the store sits at LE (and why it's easy to move)

The store is a sink -- nothing in the loop consumes it -- so the bottom-up scheduler places it last, serialized behind the other stores by the same PART_WORD_STORE resource. Bumping its latency margin doesn't help, because a sink has no in-block successor to push against. But there's plenty of NOP slack in the body, so it's not pinned: a single trailing bundle would relocate it out of LE. That's the same shape xchesscc and a Peano #pragma unroll(disable) build both already produce (store at LS, all-NOP LE bundle).

Fix direction ideas

I haven't implemented any of these, but I wanted to toss out a few ideas:

• Post-schedule fixup (my instinct). After a ZOL body is scheduled, if a partial-word store is in the loop-end bundle, insert a NOP bundle after it. Smallest blast radius -- only fires on the exact bug, grows the body by one bundle, and larger bodies commit correctly on hardware anyway. .L_LEnd0 rides the last real bundle (AIEBaseAsmPrinter.cpp:~86), so the NOP becomes the loop-end bundle and the store moves out; the added bundle only increases the setup-to-LE distance, which is a minimum (verifyDistanceToLoopEnd), so nothing trips.
• Scheduling constraint. Teach the hazard recognizer to forbid a partial-word store in the last bundle of a ZOL body, so it's prevented up front rather than fixed after.
• Extend the interblock resource model. Make the back-edge overlay treat the LE-bundle partial-word store as the conflict it is, so the existing convergence machinery resolves it. Most integrated, but it's in the fixpoint core and likely the most delicate.

Happy to prototype whichever direction you'd prefer, or leave it with you if it's something you'd rather handle inside the scheduler yourselves. Either way I'm glad the earlier detail was useful.