aie2p: add mm_activation_epilogue, one xclbin with three RTP-selected GEMM epilogues - #3462
aie2p: add mm_activation_epilogue, one xclbin with three RTP-selected GEMM epilogues#3462atassis wants to merge 7 commits into
Conversation
b888d0b to
b47694b
Compare
atassis
changed the title
b47694b to
1f8b34d
Compare
There was a problem hiding this comment.
Pull request overview
Adds a new AIE2P fused GEMM-epilogue kernel plus a complete IRON + host-side programming example that demonstrates selecting among identity/SiLU/GELU at runtime via a single RTP word, keeping all modes in one resident xclbin.
Changes:
- Introduce
aie_kernels/aie2p/mm_activation_epilogue.cc, a runtime-mode-selected f32→f32 epilogue with identity/SiLU/GELU. - Add
programming_examples/ml/mm_activation_epilogueIRON design (@iron.jit) and host testbench demonstrating multi-epoch RTP-driven dispatch. - Add build + lit integration (Makefile/CMake + run lit files) and documentation for the new example.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| programming_examples/ml/mm_activation_epilogue/test.cpp | New C++ host testbench to run the design and verify identity/SiLU/GELU outputs. |
| programming_examples/ml/mm_activation_epilogue/run_strix.lit | Lit runner for on-device JIT execution on NPU2. |
| programming_examples/ml/mm_activation_epilogue/run_makefile.lit | Lit runner covering Makefile build + on-device run flow. |
| programming_examples/ml/mm_activation_epilogue/README.md | Example documentation and usage instructions. |
| programming_examples/ml/mm_activation_epilogue/mm_activation_epilogue.py | IRON design wiring RTP-selected modes across three epochs and verification. |
| programming_examples/ml/mm_activation_epilogue/Makefile | Makefile to build xclbin/insts and the host test executable. |
| programming_examples/ml/mm_activation_epilogue/CMakeLists.txt | CMake build for the host test executable and XRT linking. |
| aie_kernels/aie2p/mm_activation_epilogue.cc | New AIE2P kernel implementing the runtime-selected epilogue math. |
Comments suppressed due to low confidence (1)
programming_examples/ml/mm_activation_epilogue/test.cpp:140
- The memcpy byte count uses sizeof(int) even though instr_v stores uint32_t. This can copy too few/many bytes on platforms where int is not 32-bit. Use sizeof(instr_v[0]) to match the vector element type.
void *bufInstr = bo_instr.map<void *>();
memcpy(bufInstr, instr_v.data(), instr_v.size() * sizeof(int));
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
1f8b34d to
9b7f952
Compare
… GEMM epilogues Fused, runtime-mode-selected epilogue for a matmul's f32 accumulator tile: identity / SiLU (hybrid f32/bf16 precision) / GELU (tanh approximation), all three compiled into ONE kernel body, selected per dispatch by a single RTP word rather than by which xclbin is loaded. SiLU uses a hybrid precision recipe: keep x and the final x*sigmoid multiply in f32, narrow only the sigmoid (bounded to [0,1], so bf16 is accurate there) -- an all-bf16 path rounds the accumulator twice and measurably lost accuracy in my testing, while an all-f32 path hangs on device (f32 elementwise/transcendental ops exceed the per-tile cycle budget on this bf16-native unit). See the kernel's header comment. Adds programming_examples/ml/mm_activation_epilogue (@iron.jit, two cores split a flat input), structurally mirroring ml/scale_shift's RTP-driven multi-phase dispatch (Buffer(use_write_rtp=True) + WorkerRuntimeBarrier), extended from two phases to three and reusing one input/output ObjectFifo pair across all three (the AIE2P compute tile's 2-in/2-out DMA channel budget does not stretch to a separate output per mode, and the real epilogue this mirrors only ever has one input and one output tile anyway).
…e seed Review fixes on the example. devicename: the kernel is aie2p only, but the Makefile never pinned devicename, so makefile-common's default of `npu` was used and the build died on `arch::is(arch::XDNA_2)` evaluating false. run_makefile.lit passes devicename=npu2 explicitly, so CI was green while a plain `make` per the README was broken. Pinned before the include; makefile-common uses `?=`, so a later assignment is a silent no-op. length: the Makefile exposed an overridable `length` fed to the design, but the host hardcoded VOLUME=65536 and `make run` passed no length, so `make length=N` built an xclbin the host buffers did not match. Host now takes -l with the design's own multiple-of-2048 check, and `make run` passes it. seed: srand(time(NULL)) with measured device GELU max abs error 0.0375 against the 0.05 gate is only 1.3x headroom over 65536 samples per run. Fixed at 42. README: said the packaged example had not been run on-device, which was no longer true. Verified on NPU2: `make run` clean at the default, `make run length=8192` clean after rebuilding the xclbin at 8192, -l 3000 and -l 0 rejected.
The host boilerplate that resolves -k to a packaged kernel name dereferences
std::find_if without checking end(), so a -k that matches nothing is
undefined behaviour rather than a diagnostic. init_xrt_load_kernel, which
exists to centralize exactly this, has the same bug; 100 more sites under
programming_examples/ and test/ hand-roll the same six lines.
Add test_utils::get_kernel_name(xclbin, prefix, verbosity), which returns
the matched name or throws listing the kernels that are actually present.
Prefix rather than exact match, because the packaged name is decorated
(-k MLIR_AIE matches "MLIR_AIE:{...}"), so xrt::xclbin::get_kernel() does
not serve. Declared unconditionally and defined under TEST_UTILS_USE_XRT,
matching init_xrt_load_kernel, so the XRT-free build is unaffected.
Ports init_xrt_load_kernel and the mm_activation_epilogue example onto it.
The remaining hand-rolled sites are left alone; they can move over
mechanically once the helper exists.
Wrong -k now prints
No kernel in the xclbin starts with 'NO_SUCH_KERNEL'. Available kernels: MLIR_AIE
instead of dereferencing end(). Device run unchanged: PASS.
5269350 to
87ef84c
Compare
hunhoffe
left a comment
hunhoffe
left a comment
There was a problem hiding this comment.
I'd appreciate a little documentation cleanup before merge, but otherwise looks nice!
| // | ||
| //===----------------------------------------------------------------------===// | ||
| // | ||
| // Fused, runtime-mode-selected GEMM epilogue for aie2p: given the f32 |
There was a problem hiding this comment.
I'm all for good documentation, but this is a little excessive I think. Can you try to make it a bit more concise?
| `aiex.npu.rtp_write`-class dispatch: same xclbin, same `hw_context`, no | ||
| reload. | ||
|
|
||
| In my own project (not this exact packaged example, but the same kernel |
There was a problem hiding this comment.
Same note about documentation; I'm not sure if all this information needs to be included to instruct users how to use these kernels.
|
Ah, did a bit more thinking -- Non-blocking, future-work note: mm_gelu_row/mm_silu_hiprec_row here (aie_kernels/aie2p/mm_activation_epilogue.cc) share almost all of their arithmetic with aie2p/gelu.cc/aie2p/silu.cc — same aie::tanh native-SFU path (unlike the aie2 variants, which use a LUT and genuinely can't share source), just entered from an f32 accumulator instead of a bf16 tile. GELU in particular is close to a literal duplicate once you factor out the narrow-in/widen-out shim; SiLU's hybrid-precision variant here is actually a strict generalization of the existing bf16-only version (feed it bf16 input and x just stays bf16, same as today). This repo already has a clean precedent for exactly this kind of dtype variation: aie_kernels/aie2p/mm.cc is one template<typename T_in, typename T_out, ...> body, with an X-macro (combos(X)) expanding it into dtype-suffixed extern "C" symbols selected at compile time via -D{dtype}_ONLY (wired from python/iron/kernels/linalg.py::mm()). silu.cc/gelu.cc predate that pattern and hardcode bf16 throughout. Not asking for a rework here — just flagging that a natural follow-up would be templating the SiLU/GELU row math on T_in/T_out (bf16 vs f32) the same way mm.cc does, so gelu.cc/silu.cc and this file's row functions become one shared implementation instead of two. |
48% -> 14% comment lines; no other aie2p kernel carries file-header prose.
run_strix.lit already runs the design on device and asserts all three modes; test.cpp re-checked the same modes through a second host API.
|
Done! @andrej, took your #3467 point about the
|
What
Add
mm_activation_epilogue, a fusedfloat32-in/float32-out GEMM epilogue foraie2p: given a matmul's f32 accumulator tile, apply one of{identity, SiLU, GELU}. All three bodies compile into ONE kernel; whichone runs is chosen per dispatch by a single runtime-parameter (RTP) word, not by loading a different xclbin. Plus a new
programming_examples/ml/mm_activation_epilogueexample (@iron.jit, two cores, RTP-driven three-phase dispatch).Why
aie_kernels/aie2phas standalonesilu.cc/gelu.cc(fixed 1024-elementbfloat16in/out, one activation per xclbin) but nothing that (a) reads directly off a matmul's f32 accumulator, or (b) lets more than one activation share a resident program. In a pipeline that dispatches several differently-activated matmuls back to back, one-activation-per-xclbin means a hardware reconfiguration between each. I did not find an existing kernel or example that fuses an activation onto a GEMM's own accumulator tile, or that runtime-selects between activations without a reload (searchedaie_kernels/aie2pandprogramming_examplesat the fetched upstreamtip).
I use this as the epilogue of a resident matmul microkernel in my own project: rather than build one xclbin per activation and reload between FFN layers that use SiLU vs. GELU vs. no activation, one compiled program
serves all three, switched by an RTP write. In my own project's design (not this exact packaged example - see Test) I measured this on NPU2 today: registering ONE
hw_contextand alternating SiLU/GELU dispatchesthrough it stayed numerically correct (rel-L2 0.0083 / 0.0094 against an f32 numpy reference, unchanged after 100 alternating dispatches), the compiled per-core ELF / PDI / CDO binaries were byte-identical between the
two modes except for 4 bytes (the per-core RTP write constant), and the mode switch itself showed no measurable dispatch-time cost beyond ordinary dispatch-to-dispatch noise (alternating-dispatch mean 0.2095 ms vs. a
solo-mode baseline mean of 0.2243 ms - a negative delta, i.e. the "switch" was not distinguishable from no switch at all). A second, independent measurement in the same project - two genuinely separate,
independently-authored kernel bodies (not sibling variants of one activation family) sharing one xclbin via the identical RTP mechanism - found the same result (correct output, byte-identical artifacts except the
RTP payload, statistically zero switch cost), which is why I am confident this generalizes rather than being an artifact of SiLU/GELU specifically.
How
aie_kernels/aie2p/mm_activation_epilogue.cc: three per-row helpers (mm_identity_row,mm_silu_hiprec_row,mm_gelu_row) plus oneextern "C" mm_activation_epilogue_row(c_in, c_out, n, mode)that branches onmode(0/1/2). No bias argument: I fold bias into the producing matmul via K-augmentation in my own pipeline, so by the time this epilogue runs the bias is already summed into the accumulator - documented in the file as acaller-side choice this kernel does not require. SiLU uses a hybrid precision recipe (keep
xand the finalx*sigmoidmultiply in f32, narrow only the sigmoid - bounded to[0,1], sobfloat16is accurate there); the file's header comment records why, including that an all-f32 version of both SiLU and GELU hangs on device (f32 elementwise/transcendental ops exceed the per-tile cycle budget on thisbfloat16-native unit in my testing), which is why GELU staysbfloat16-only rather than getting the same hybrid treatment.programming_examples/ml/mm_activation_epilogue/mm_activation_epilogue.py: two cores split a flatsize-element input, structurally mirroringml/scale_shift's RTP-drivenmulti-phase dispatch (
Buffer(..., use_write_rtp=True)+WorkerRuntimeBarrier, one epoch per host dispatch), extended from two phases to three. Unlikescale_shift, all three phases here reuse ONE input/outputObjectFifopair rather than a separate output per mode - an earlier version gave each mode its own output tile andaieccrejected it (tile requires 1 input/3 output DMA channels, but only 2 input/2 output available), which is itself a useful confirmation of the real epilogue's own design note: one input, one output is the actual DMA budget on this compute tile, and the real 3-mode kernel this example demonstrates only ever needs that one pair regardless of how many modes share it.Test
Compiles clean through the pinned Peano at
--target=aie2p-none-unknown-elf(-Wall -Wextra): only pre-existingaie_apiheader warnings, none from this kernel.mm_activation_epilogue_rowpresent in the object (.text.mm_activation_epilogue_row, 1056 bytes).clang-format(17.0.1) zero diff on both new.cc/.cppfiles.Beyond the object-file compile, I built
mm_activation_epilogue.pyend to end through the fork's IRON Python bindings andaieccagainst this exact commit: real 14361-byte AXLF xclbin (7 sections), a realinsts.bin+insts.elfpair, 2 core ELFs, lowered per-core IR confirmingtarget triple = "aie2p"andaie.device(npu2), and themm_activation_epilogue_rowsymbol present in the linked per-core ELF (2304-byte linked.text).I also transcribed the kernel's exact per-lane arithmetic (the hybrid SiLU path and the
bfloat16GELU path, including manual round-to-nearest-even bit manipulation for every intermediate narrowing the kernel performs)into numpy and checked it against a plain f64 reference over 200000 uniform random samples in
[-4, 4]: SiLU hybrid max abs error 0.0116 (mean 0.0017), GELU max abs error 0.0171 (mean 0.0022), both well inside theexample's own
atol=0.05gate. I separately checked the all-bfloat16SiLU variant (narrowxbefore the sigmoid, the path this kernel deliberately does NOT take) against the same reference: mean error 0.0026,worse than the hybrid's 0.0017 - a real check of the design claim in the header comment, not just an assertion of it.
I have now run this exact packaged example on NPU2. All three modes dispatch clean, no fault: identity is bit-exact against the reference, SiLU comes out at rel-L2 0.00246 and GELU at 0.00341 over a
[-8, 8]sweep, both inside the example's
atol=0.05gate.Worth flagging, because it cuts against the numpy check above: the transcription UNDERPREDICTS the real device error. Restricted to the same
|x| <= 4domain, the device gives SiLU max abs error 0.0196 against thepredicted 0.0116, and GELU 0.0375 against 0.0171 - 1.7x and 2.2x respectively, with mean error about 2.7-2.9x higher. Everything still passes, so this is not a correctness problem, but I would rather state it
than quote the friendlier host numbers: a bit-accurate transcription of the arithmetic is evidently not a substitute for running the thing, and if you are reviewing the accuracy claim, the device figures are the ones to hold me to.
I also checked that the mode select is genuinely doing something, since a kernel that ignored
rtp[0]could pass a lenient tolerance: across the mid-range (|x|in(0.5, 3)), 99.7% of elements have all three modes pairwise separated by more than 1e-3, and the elements that do not separate sit exactly where the functions themselves converge - identity meets SiLU only above x = 6, SiLU meets GELU only below x = -6.The numbers in the Why section above remain from a differently-shaped production design using the identical kernel logic and RTP mechanism; they are corroboration, not a substitute for the run described here.
One question
This example holds three bodies in one program and picks between them with an RTP write, and the switch costs nothing I can measure. I would like to push the same pattern further, and I hit a limit whose cause I cannot determine from outside: past a small number of co-resident programs I run out of locks. Since RTP-selected
bodies are mutually exclusive by construction, only one can ever be issuing, so in principle they could share a lock allocation instead of each reserving its own. Is that something the allocator could reasonably do, or is there a hardware reason the allocation has to be per design? If it is the former and you would take such a change, I am happy to prototype it.
Refs
Closest prior art for the RTP-driven multi-phase dispatch:
ml/scale_shift(Buffer(..., use_write_rtp=True)+WorkerRuntimeBarrier, two phasesinstead of three). Closest prior art for the activation math itself:
silu.cc/gelu.ccin the same directory (standalone, fixed-size, single-mode); this kernel differs by reading directly off an f32 accumulator, writing f32, taking a runtime length, and adding the RTP mode branch plus the SiLU hybrid-precision path.