Skip to content

feat(rocr/aie): aie4 agent support - #13740

Open
ypapadop-amd wants to merge 28 commits into
developfrom
users/ypapadop-amd/aie4-support
Open

ypapadop-amd wants to merge 28 commits into
developfrom
users/ypapadop-amd/aie4-support

Conversation

@ypapadop-amd

@ypapadop-amd ypapadop-amd commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

Add aie4 (npu3) support. Unlike aie2/aie2p, which are driven through the kernel driver (KMQ), aie4 is 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

  • AieAqlQueue is renamed AieKmqQueue. The new AieUmqQueue serves aie4 agents, which the driver identifies by device_type == AMDXDNA_DEV_TYPE_UMQ.
  • The UMQ queue buffer (a 128-byte header followed by a 32-slot packet ring) is passed to the driver as the hardware context's umq_bo, and the doorbell is mmaped. aie4 has no device heap.
  • aie4 is cache-coherent with the CPU (XRT's sync_bo is a no-op for this device type), so UMQ submission does no cache flushing.
  • Submission is one packet at a time: publish CERT's write_index, ring the doorbell, then wait for read_index to pass the packet. That's a lock-free check of read_index first, then DRM_IOCTL_AMDXDNA_WAIT_CMD, with a 1 s back-off on EAGAIN (context recovery), as XRT does. There were driver issues submitting more than one packet at time, so chaining will be revisited.
  • Packets are not batched under one doorbell. A full ELF that starts immediately after the previous one finishes produces corrupted output: its array configuration isn't safe to apply over a run that has just ended. XRT's runlist corrupts the same way. Likely the same driver issue.

Packet format

  • hsa_amd_aie_kernel_dispatch_packet_t is now laid over CERT's own host_queue_packet. One layout serves both queue types: on KMQ the device never reads the packet, so its layout is free.
  • Breaking change: field names are unchanged, but completion_signal, kernarg_address and num_kernargs moved, so applications must recompile.
  • On 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.
  • CERT is given no argument list (args_len = 0), as with XRT. A non-zero args_len makes CERT read 3 KiB from args_host_addr and stall if that crosses an unmapped page.

Full-ELF loading and dispatch

  • The ELF loader handles aie4's multi-page control code (.ctrltext/.ctrldata pages) and the 57-bit shim DMA relocation type.
  • Each dispatch patches its kernel arguments into a per-dispatch copy of the control code, including the shim DMA's DDR aperture (0x80000000), as XRT's patch57_aie4 does.
  • If waiting for a packet fails, its control-code copy is kept until the hardware context is destroyed, because CERT may still run it.

Issue Tracking

#13739

Test Plan

rocrtst and updated AIE unit tests.

Test Result

Success.

Submission Checklist

@therock-pr-bot

Copy link
Copy Markdown

✅ All Policy Checks Passed

Check Status Details
📝 PR Description ✅ Pass —
⛔ Forbidden Files ✅ Pass —
🧪 Unit Test ✅ Pass —
🚫 Draft PR 🔜 To Be Enabled —
🚩 Feature Flag 🔜 To Be Enabled —
📊 Code Coverage 🔜 To Be Enabled —

🎉 All policy checks passed!

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Unresolved control-code layout, dispatch safety, and toolchain compatibility issues block approval.

9 open issues
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

Comment thread projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp Outdated
Comment thread projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aie_elf.cpp Outdated
Comment thread projects/rocr-runtime/runtime/hsa-runtime/loader/executable.cpp Outdated
Comment thread projects/rocr-runtime/rocrtst/suites/aie/CMakeLists.txt Outdated
Comment thread projects/rocr-runtime/rocrtst/suites/aie/CMakeLists.txt
Comment thread projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aie_elf.cpp Outdated
Comment thread projects/rocr-runtime/runtime/hsa-runtime/loader/executable.cpp

@cfreeamd cfreeamd left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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.

Comment thread projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd_aie.h
Comment thread projects/rocr-runtime/runtime/hsa-runtime/inc/hsa_ext_amd_aie.h
@ypapadop-amd ypapadop-amd changed the title rocr/aie: aie4 Agent support feat(rocr/aie): aie4 agent support Oct 19, 2026
@ypapadop-amd
ypapadop-amd force-pushed the users/ypapadop-amd/aie4-support branch from 055878e to c4952c8 Compare October 20, 2026 21:21

@cfreeamd cfreeamd left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI-assisted pre-review (for @cfreeamd) — follow-up after commit 055878e

amd_xdna_driver.cpp new UMQ submission path looks well-structured. UMQ_INTER_PACKET_GAP_US is clearly documented with the measured threshold. One question about cross-queue enforcement.

Comment thread projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp Outdated
@ypapadop-amd
ypapadop-amd force-pushed the users/ypapadop-amd/aie4-support branch 2 times, most recently from 15c547f to c248bfd Compare October 21, 2026 14:39
Yiannis Papadopoulos and others added 11 commits October 21, 2026 12:23
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.
ypapadop-amd and others added 16 commits October 21, 2026 12:23
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>
@ypapadop-amd
ypapadop-amd force-pushed the users/ypapadop-amd/aie4-support branch from 2daec39 to 4f5064f Compare October 21, 2026 16:23
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>
@ypapadop-amd

ypapadop-amd commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor Author

XRT vs HSA benchmark results on npu3 (aie4)

Ported rocrtst/suites/aie-benchmarks (branch users/ypapadop-amd/aie-benchmarks) to build and run its vector_scalar_add full-ELF benchmarks against npu3, to validate this PR's UMQ dispatch path against XRT's own aie4 dispatch and compare the two.

Median of 3 runs on real npu3 hardware, dispatch counts 1–32 (µs):

dispatches XRT single-dispatch XRT runlist (batched) HSA full-ELF (batched) HSA + kernarg alloc
1 81.6 37.5 43.1 38.9
2 180.6 57.3 62.7 58.8
4 340.4 98.1 106.4 100.2
8 726.3 193.2 196.2 187.3
16 1573.7 380.0 351.7 349.2
32 3567.3 675.4 667.4 673.2

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 (verify()), not just absence of a crash.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants