Skip to content

Wire the AIE trace buffer through the fused full ELF flow - #154

Open
asyms wants to merge 5 commits into
amd:develfrom
KULeuven-MICAS:tracing-on-devel
Open

Wire the AIE trace buffer through the fused full ELF flow#154
asyms wants to merge 5 commits into
amd:develfrom
KULeuven-MICAS:tracing-on-devel

Conversation

@asyms

@asyms asyms commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

stream-dse can now emit AIE tracing, but the trace buffer never reaches the device through the fused full ELF flow: lowering appends it to the operator's runtime sequence, and the address gets patched against the dispatched kernel instead, which is the wrapper. This wires the buffer through so a traced design produces real data, and leaves untraced builds exactly as they were.

  • Give the fused runtime sequence a trace buffer per operator, at the index that operator uses. That index is what lowering patches against, so the wrapper has to match it or the trace DMA writes nowhere and the buffer comes back empty.
  • Allocate one buffer per operator, bind them, and sync them back with the outputs.
  • Let the SwiGLU-prefill operator use it: it passes trace_size into the sequence, and trace_size and trace_max_tiles into stream-dse's codegen, so the design is generated with tracing rather than having it added afterwards.
  • Turn tracing on with IRON_TRACE_SIZE, and pick how many tiles with IRON_TRACE_NTILES. Off by default, since tracing changes the calling convention and its packet flows need routing a full design may not have spare.
  • Key both the generated design and the sequence's ELF on whether tracing is on, so traced and untraced builds never share an entry. Without the second one, turning the flag on reuses an untraced ELF and binding the trace buffer fails against a kernel that has no argument for it.
  • Require stream-dse 1.13.14, the first release with the tracing codegen and the tile count argument. This is a floor for every build, not just traced ones.
  • Add unit tests for the slot layout.

This reopens #151, which targeted update-mlir-aie-latest-nightly and was closed when that branch merged as #150. It is based on devel now, and carries one extra commit for the ELF cache key above.

Checked on Strix in a clean venv built exactly as .github/actions/prereqs does, against stream-dse 1.14.0 (what the >=1.13.14 floor resolves to today) and mlir-aie 1.4.2.dev16: pytest -m "not extensive" iron/operators/ iron/tests/ gives 990 passed, a traced k=1 build passes on hardware, and turning tracing on and off again in one build directory works both ways. black --check and reuse lint are clean.

asyms added 5 commits August 26, 2026 09:52
Trace lowering records the buffer's index in the sequence it configures, but that
index is resolved against the kernel the host dispatches, which for a fused build
is the wrapper rather than the operator. The address was patched from an argument
the wrapper did not have, so the trace DMA wrote nowhere and the buffer came back
empty. Give the wrapper the buffer at the same index.
Tracing adds a runtime-sequence argument, so a traced ELF cannot be reused by an
untraced run or the other way round: the two now get different artifact names and
toggling the flag no longer needs the build directory wiped.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant