Repository navigation
feat(rocr/aie): aie4 agent support - #13740
ypapadop-amd wants to merge 28 commits into
Conversation
✅ All Policy Checks Passed
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved control-code layout, dispatch safety, and toolchain compatibility issues block approval.
9 open issues
Ignores CERT completion errors and retires failed packets · New Applies AIE4 page bias too late for ctrltext relocations · New Accepts relocation schemes incompatible with selected architecture · New Unmaps doorbell resources while submission remains blocked · New Misorders and misaligns AIE4 firmware pages · New Unbounded npu1 match misclassifies npu10 and npu11 · New Creates test artifact for unsupported AIE4 packer · New Rejects valid 57-bit relocations near target buffer end · New Drops multi-column entry points during UMQ dispatch · New
What changed in this PR
Adds aie4/npu3 support to ROCR through user-mode queues, extending the existing aie2/aie2p kernel-mode dispatch implementation.
Changes:
- Adds UMQ device discovery, queue management, and synchronous dispatch.
- Updates the packet ABI and full-ELF loader for AIE4 control code.
- Extends kernel generation, build configuration, and AIE tests.
Paths below are relative to projects/rocr-runtime/.
| File | Description |
|---|---|
runtime/hsa-runtime/loader/executable.hpp |
Tracks system-allocated control-code buffers. |
runtime/hsa-runtime/loader/executable.cpp |
Loads and releases AIE4 control code. |
runtime/hsa-runtime/inc/hsa_ext_amd_aie.h |
Defines the shared KMQ/UMQ packet layout. |
runtime/hsa-runtime/core/runtime/amd_aie_umq_queue.cpp |
Implements UMQ lifecycle and submission. |
runtime/hsa-runtime/core/runtime/amd_aie_kmq_queue.cpp |
Renames the existing KMQ implementation. |
runtime/hsa-runtime/core/runtime/amd_aie_elf.cpp |
Parses AIE4 pages and relocations. |
runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp |
Selects queue type and available memory pools. |
runtime/hsa-runtime/core/inc/amd_xdna_driver.h |
Declares UMQ driver interfaces. |
runtime/hsa-runtime/core/inc/amd_aie_umq_queue.h |
Declares the UMQ queue class. |
runtime/hsa-runtime/core/inc/amd_aie_section.h |
Extends kernel descriptors for AIE4. |
runtime/hsa-runtime/core/inc/amd_aie_kmq_queue.h |
Renames the KMQ queue class. |
runtime/hsa-runtime/core/inc/amd_aie_elf.h |
Models control-code pages and patch sites. |
runtime/hsa-runtime/core/inc/amd_aie_agent.h |
Makes queue-size limits device-dependent. |
runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp |
Implements AIE4 discovery, patching, and dispatch. |
runtime/hsa-runtime/CMakeLists.txt |
Registers KMQ and UMQ sources. |
rocrtst/suites/aie/vector_scalar_mul.py |
Adds npu3 kernel generation. |
rocrtst/suites/aie/vector_scalar_add.py |
Adds npu3 kernel generation. |
rocrtst/suites/aie/test_pointer_info.cc |
Separates and adapts pointer-info tests. |
rocrtst/suites/aie/test_memory.cc |
Adapts memory tests to heapless devices. |
rocrtst/suites/aie/test_dispatch.cc |
Extends full-ELF dispatch coverage. |
rocrtst/suites/aie/elf_parse_test.cc |
Adapts parser tests for AIE4. |
rocrtst/suites/aie/CMakeLists.txt |
Adds npu3 tooling, artifacts, and test targets. |
rocrtst/suites/aie/aie_test_env.h |
Expands test-helper documentation. |
rocrtst/suites/aie/aie_code_test.cc |
Tests AIE4 full-ELF code objects. |
Review effort: Balanced
cfreeamd
left a comment
There was a problem hiding this comment.
🤖 AI-assisted pre-review (for @cfreeamd)
The UMQ queue implementation (amd_aie_umq_queue.cpp/.h, amd_xdna_driver.cpp) and the aie4 ELF loading path in executable.cpp are the load-bearing changes. DestroyUserModeQueue properly munmaps the doorbell, destroys the hwctx, and frees orphaned ctrl-copy allocations. The static asserts confirming the host_queue_packet/hsa_amd_aie_kernel_dispatch_packet_t overlay are correct and intentional for the UMQ zero-copy ring path.
[Required] PR title rocr/aie: aie4 Agent support does not follow the required type(component): format — ROCm squash-merges so this becomes the permanent commit title. Please use e.g. feat(rocr): aie4 agent support.
055878e to
c4952c8
Compare
15c547f to
c248bfd
Compare
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…through UMQ aie4's full ELF bundles each kernel's control code across several ctrltext/ctrldata page pairs instead of aie2p's single ctrltext buffer, and CERT does not patch argument addresses into it itself (confirmed on real hardware: it faults instead). Extend amd_aie_elf::Parse() to model a kernel's control code as a list of pages, rebasing ctrldata-targeting relocations by comparing their raw offset against the page's ctrltext size rather than by which section their symbol is bound to (every relocation symbol is bound to ctrltext, even when it targets ctrldata). At executable load, concatenate an aie4 kernel's pages into one flat, pristine control-code buffer placed through the agent's system allocator (aie4 has no device heap to carve a CODE_AGENT segment from), and move the aie2p single-page shape assumption KMQ's BuildFullElfCommand relies on into a load-time check, since the shape cannot change between dispatches of the same arch. Dispatch through UMQ by having AieUmqQueue::SubmitPackets hand each packet to the new XdnaDriver::TranslateUmqPackets, which copies the kernel's pristine control code and patches this dispatch's argument addresses into the copy with the new PatchShimDma57Aie4 (aie4's 57-bit equivalent of PatchShimDma48) -- the same per-dispatch-copy pattern KMQ's BuildFullElfCommand already uses for aie2p. Validates the packet's kernel object and kernarg count before allocating anything, matching KMQ's ValidateBatchPacket/BuildFullElfCommand checks and error codes.
Build test_dispatch and elf_parse_test for npu3 (AIE4) instead of skipping them outright: elf_parse_test's PDI-specific cases self-skip (aie4's model has no separate PDI construct), and test_dispatch's PDI + instruction sequence and multi-arch cases self-skip the same way every other kernel-toolchain-dependent test already does, since npu3's toolchain flow builds only full-ELF artifacts. test_dispatch's full-ELF tests run against aie4 for real: make_aie_packet picks the UMQ opcode for an aie4 agent, FullElfDispatchTest accepts aie4 alongside aie2p, and the ~9 tests needing a non-allocatable NPU-local dev pool (aie4 has none) fall back to the allocatable data pool instead, where the semantics allow it. The handful of tests that need a real blocking dispatch-and-wait are skipped on aie4 with a clear reason: CERT advances its own read index but nothing yet bridges that back to the HSA completion signal. vector_scalar_add.py skips its design's npu_load_pdi call on npu3: aie4's full ELF has no PDI to load that way, and including it anyway stalls dispatch on real hardware. Default Peano discovery to the standard llvm-aie wheel layout for every arch; npu3 alone falls back to mlir-aie-xrlabs's experimental 'lnx64.o/tools/peano' layout when the standard one is not found.
libhsakmt only registers a CPU-pool allocation as a real KFD buffer object when at least one discrete-GPU KFD node exists (hsakmt_is_dgpu); with none, it reports the allocation as HSA_EXT_POINTER_TYPE_RESERVED_ADDR instead of HSA_EXT_POINTER_TYPE_HSA, which hsa_amd_pointer_info passes through unchanged. That is a libhsakmt classification gap on GPU-less systems, unrelated to AIE, so skip the test there rather than trying to work around it.
…orruption Chaining multiple packets under one UMQ doorbell ring let CERT execute them with less separation than a real hardware barrier, corrupting the second-and-later packets' output on real hardware (confirmed: the mismatch location moved between otherwise-identical runs, pointing at an execution race rather than a fixed addressing bug). Submission is now fully sequential -- translate, ring the doorbell, and block for completion one packet at a time -- which fixes SameKernelDifferentArgsInOneBatch and ElfFullQueueDispatch. While here, restructure AieUmqQueue::SubmitPackets to mirror AieKmqQueue::SubmitPackets's shape: the new XdnaDriver::SubmitUmqCmdChain owns the whole device-facing submission loop and reports num_completed the same way KMQ's chain-splitting does, so the queue layer is symmetric between the two backends. WaitUmqCmd and RetireCompletedPackets, now only called from within amd_xdna_driver.cpp, became internal (a static free function next to KMQ's WaitCommand, and a static free function respectively) instead of public XdnaDriver API. FullElfInterleaveTest.ElfInterleavedKernels remains skipped: its "+384" additive corruption when alternating two differently-compiled kernels is a different bug than the chain race above. A KMQ-mirroring fix (rebuilding the hardware context on kernel-identity change) was tried and found to trade the silent corruption for a worse, unexplained deterministic hang; it was reverted, and the skip reason documents why. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… sizing comment AieUmqQueue's reaper thread was removed when UMQ submission became synchronous; the comment explaining why a UMQ queue is sized to XdnaDriver::UmqMaxCmds() still pointed at it. Point at SubmitUmqCmdChain, which now confirms CERT's retirement instead. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…dKernels on aie4 On npu3, aiecc's full-ELF control code already configures the array in its first job. vector_scalar_mul.py additionally emitted npu_load_pdi, so the mul ELF applied the same configuration twice, and every run wrote back stretches of its output shifted by 128 elements (the +384 seen at scale 3). vector_scalar_add.py already skipped load_pdi on npu3, which is why only the mul kernel was affected. This was not an interleaving or ROCr bug: the mul ELF was corrupt even when dispatched alone on a fresh queue. An XRT host program (one hw_context per ELF) reproduces the identical corruption with the old ELF and passes with the fixed one. Also corrects vector_scalar_add.py's comment, which claimed aie4 has no PDI and that its full ELF carries no PDI section. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
VMemMapNPU, VMemDoubleMap and VMemReleaseBeforeUnmap were switched to the NPU-local dev-heap pool for aie4, which made them fail on aie2p: hsa_amd_vmem_handle_create exports the BO as a dma-buf, and amdxdna's dev-heap BOs have no .export (DRM_IOCTL_PRIME_HANDLE_TO_FD returns EOPNOTSUPP). On aie4 there is no dev heap at all, so the tests were skipped there instead. Invert the pool filter so these tests skip the dev-heap pool and pick a system-backed one, which every AIE agent has, so they run on every architecture instead of failing on aie2p and skipping on aie4 (verified on aie2p/npu4). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
VMemMapNPU, VMemDoubleMap and VMemReleaseBeforeUnmap reserved their address range with flags 0, which reserves it through hsaKmtAllocMemoryAlign. On a machine with no GPU (the npu3/aie4 box) that call fails and the reserve returns HSA_STATUS_ERROR_OUT_OF_RESOURCES, so the tests failed there once they stopped skipping. Pass HSA_AMD_VMEM_ADDRESS_NO_REGISTER, as the file's other NPU VMem tests already do, so the reservation is a plain host one. Verified on aie4/npu3: test_memory 30 passed, 9 skipped, 0 failed (3 runs). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- ELF reader: aie4 relocation offsets count from the page payload, 16 bytes past the control-code header, across ctrltext then ctrldata, as aiebu does. Argument sites must use the arch's own patch scheme, and the bounds check uses the scheme's width (two dwords for the 57-bit one). - aie4 pages are ordered by page number and placed at number * 8192, matching aiebu's elf_reader_gen2plus; oversized pages, duplicate pages and multi-column control code are rejected. - UMQ submission checks CERT's return code after the wait and fails the packet without retiring it unless it is COMPLETED. - Test CMake: bound the npu1 device match so npu10/npu11 are not taken for npu1, and build multiarch_elf.hsaco only if the packer accepts the other arch. - Add parser tests for the aie4 offset, bound, scheme, multi-column and page-order rules. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
CERT retiring a packet doesn't mean the array has settled: a full ELF's control code reconfigures memtile DMA channels that a run just ended may still be draining, corrupting output if the next packet starts too soon. Measured on aie4 with vsadd: <=500us corrupts every time, >=800us is clean. Add a 2ms margin, paid unconditionally after every packet since the next one may come from a different queue sharing the same device. Also store UmqMetadata's owning agent instead of a bound deallocator. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… the kernel instead The gap was standing in for a missing compiler feature, not a hardware limit: a full-ELF kernel's control code never reset its own memtile DMA channels before reusing them, so CERT retiring a packet was not enough evidence the array was safe to reconfigure for the next one. Confirmed on real aie4 hardware that aiex.dma_channel_reset_for on every objectFifo removes the requirement entirely -- not just reduces it -- down to no gap at all between back-to-back dispatches. - vector_scalar_add.py, vector_scalar_mul.py: reset every objectFifo's DMA channel at the start of the runtime sequence on npu3, documented as required now that the runtime no longer covers for a kernel that skips it. - amd_xdna_driver.cpp: drop UMQ_INTER_PACKET_GAP_US and the sleep after every packet. - CreateUserModeQueue: report a likely cause (kernel-mode submission) when hwctx creation fails with EINVAL, checked via xdna_ioctl's own HSA_STATUS_ERROR_INVALID_ARGUMENT instead of re-reading errno. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…cket Per-packet submission was standing in for the removed inter-packet gap, not for a separate CERT chain-race bug: with the kernel now resetting its own DMA channels, a whole batch of UMQ packets can be translated and readied before ringing the doorbell once, instead of once per packet. Confirmed on real aie4 hardware across every multi-packet pattern in the suite (same-kernel batches, full-ring chains, two independent queues, interleaved kernels), including with the submission loop forced to batch unconditionally. This changes SubmitUmqPackets' failure contract for a translate failure: since nothing is rung until every packet in the batch has translated, a bad packet partway through now aborts the whole batch (num_completed stays 0) rather than leaving the packets before it already run. A wait failure (CERT already notified) is unaffected. Added ElfMalformedPacketMidBatchAbortsWholeBatch to cover the new case: a malformed packet in the middle of three leaves all three outputs untouched, including the ones before it a per-packet design would already have run. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ate failure Replaces the manual loop over already-translated packets with a MAKE_NAMED_SCOPE_GUARD, matching how CreateUserModeQueue already frees its hwctx on a failed setup. Dismissed once the doorbell is rung, after which the wait loop owns freeing each copy explicitly as before. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… already waited on ctrl_copies_guard was dismissed right after ringing the doorbell, before the wait loop below it. Since every packet in a batch is handed to CERT at once (not one at a time), a wait failure partway through the batch left every later packet's control-code copy unreachable and unfreed -- a real leak. Keep the guard armed through the whole wait loop instead, nulling each entry once it is freed there so the guard can never free one twice. Along the way: - TranslateUmqPacket now hands back its ctrl_copy through an out-parameter, instead of the caller reconstructing it from the packet's post-translate kernel_object fields. - write_index is stored once per batch, with its final value, instead of once per packet: nothing reads it before the single doorbell ring below. - UMQ releases its own completion signal directly instead of going through RetireCompletedPackets, which was only ever called here with num_pkts=1 and flush_arguments=false. RetireCompletedPackets' flush_arguments parameter is dropped too, since both remaining (KMQ) callers always passed true. Verified on real npu3 hardware: full suite (107 tests) passed 5/5 runs after each change, plus 5 extra runs targeted at the batch-sensitive tests (ElfMalformedPacketMidBatchAbortsWholeBatch, ElfMultiDispatch, ElfMultiDispatchWrapAround, SameKernelDifferentArgsInOneBatch, ElfFullQueueDispatch, ElfInterleavedKernels). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2daec39 to
4f5064f
Compare
Removes repetition, duplication, and filler words from comments added over the course of this branch, across the driver, a couple of queue/ELF spots, a test, and the npu3 kernel/build files. No code logic touched -- comment text only. Notably fixes one comment that had drifted stale (RetireCompletedPackets' doc still claimed SubmitUmqPackets calls it, no longer true) and reconciles two ctrl_copies_guard comments that had ended up contradicting each other. Verified on real npu3 hardware: clean rebuild, full suite (107 tests) passed 5/5 runs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
XRT vs HSA benchmark results on npu3 (aie4)Ported Median of 3 runs on real npu3 hardware, dispatch counts 1–32 (µs):
The fair comparison is XRT's runlist path vs HSA, since both batch the whole submission behind one wait (XRT's plain single-dispatch path pays a separate wait per dispatch and is shown only for reference). HSA tracks XRT within single-digit µs or noise at every dispatch count, and is marginally faster at 16/32. All benchmark runs verified their own output ( |


Motivation
Add
aie4(npu3) support. Unlike aie2/aie2p, which are driven through the kernel driver (KMQ),aie4is dispatched through an user-mode queue (UMQ). The queue's ring is handed to the device's command processor (CERT), which reads packets in place, and the host submits work by writing a doorbell, much like a GPU hardware queue.Technical Details
Queues and driver
AieAqlQueueis renamedAieKmqQueue. The newAieUmqQueueservesaie4agents, which the driver identifies bydevice_type == AMDXDNA_DEV_TYPE_UMQ.aie4has no device heap.aie4is cache-coherent with the CPU (XRT's sync_bo is a no-op for this device type), so UMQ submission does no cache flushing.Packet format
hsa_amd_aie_kernel_dispatch_packet_tis now laid over CERT's ownhost_queue_packet. One layout serves both queue types: on KMQ the device never reads the packet, so its layout is free.aie4, the application owns header (including barrier and fence bits), opcode and count, and zeroes the reserved fields. The runtime writes only CERT's completion address (reserved2) and replaces kernel_object with the control-code address.Full-ELF loading and dispatch
Issue Tracking
#13739
Test Plan
rocrtst and updated AIE unit tests.
Test Result
Success.
Submission Checklist