Wire the AIE trace buffer through the fused full ELF flow - #154
Conversation
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.
andrej
left a comment
andrej
left a comment
There was a problem hiding this comment.
Thanks for submitting this. I see two things that worry me:
- The error about "operators taking the same number of arguments".
- The "pick the IR file with the most write32s" logic.
I think the fundamental problem is that tracing currently does not play nice with aiex.configure. What I read from this PR is that tracing always patches a buffer at a fixed index, but we'd want separate indices for each sub-operator to have its own trace buffer. To me, this PR looks like mostly a workaround around that.
I hate rejecting this work, but could we instead implement a proper fix upstream? For example, make trace aware of aiex.configure and allocate separate trace buffers / different offsets into the same trace buffer for each called runtime sequence. I think at the end of the day that will be fewer lines of code, and it will benefit everyone upstream.
As a solution in the meantime, have you tried the "separate xclbin" dispatch mode for tracing?
| Picks the file with the most ``write32`` ops rather than hardcoding a stage name, | ||
| since those names are aiecc's business and have changed before. |
There was a problem hiding this comment.
That seems like a hack ...
If the trace parser util needs a certain form of IR, we should explicitly request it to be generated using aiecc --get-<whatever it needs>. If it is not exposed as --get-<whatever it needs>, we should enable that in aiecc. Relying on intermediate outputs in <source>.mlir.d seems hacky; the compiler may choose to not generate them and do things all in-memory instead.
| counts = list(trace_slots.values()) | ||
| shared = sorted({n for n in counts if counts.count(n) > 1}) | ||
| if shared: | ||
| raise NotImplementedError( | ||
| "operators taking the same number of arguments would share one trace " | ||
| f"buffer (slots {shared}); trace them in separate dispatches" | ||
| ) |
There was a problem hiding this comment.
Am I understanding correctly that this rejects anything where multiple operators in the same sequence take the same number of arguments?
|
I've sent off an agent to look into a proper upstream fix to help with this, will report back. |
2 participants
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.
trace_sizeinto the sequence, andtrace_sizeandtrace_max_tilesinto stream-dse's codegen, so the design is generated with tracing rather than having it added afterwards.IRON_TRACE_SIZE, and pick how many tiles withIRON_TRACE_NTILES. Off by default, since tracing changes the calling convention and its packet flows need routing a full design may not have spare.This reopens #151, which targeted
update-mlir-aie-latest-nightlyand was closed when that branch merged as #150. It is based ondevelnow, and carries one extra commit for the ELF cache key above.Checked on Strix in a clean venv built exactly as
.github/actions/prereqsdoes, against stream-dse 1.14.0 (what the>=1.13.14floor 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 --checkandreuse lintare clean.