Repository navigation
Switch HSA/ROCR dispatch to hsaco/kernel-object loading - #158
Merged
Merged
Conversation
Replace the raw PDI+instruction-address dispatch path with ROCr's code-object-reader/executable/freeze flow: HsaRuntime now loads an hsaco file into a frozen HSA executable and dispatches by the kernel object HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT resolves it to, instead of filling in raw PDI/instruction-stream device addresses. triton_npu_hsa_prepare(hsaco_path, kernel_name) replaces the old (pdi_path, insts_path) signature; triton_npu_hsa_patch_insts is removed along with the dev-pool PDI/insts loading machinery it depended on, since there is no longer a per-dispatch instruction stream to patch. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Update hsa_launcher.py to generate C++ code that calls triton_npu_hsa_prepare with hsaco_path and kernel_name parameters, replacing the previous pdi and insts parameters. This matches the ABI change made in Task 1 to the shared HSA runtime. Changes: - Update py_set_paths function to accept hsaco_path and kernel_name - Update comment describing g_program handle - Update PyMethodDef doc string for set_paths - Update module docstring to reflect new API - Update _generate_hsa_launcher docstring to explain kernel_name handling Verified with Python string template substitution check. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
driver.py now lets the HSA runtime auto-detect its output format the same way XRT does (ELF on npu2, xclbin on npu1) instead of forcing a raw "pdi" format, packs aiecc's artifacts into an hsaco via aie.compiler.hsaco.pack, caches it alongside the other compiled artifacts, and passes it to the generated launcher's set_paths(hsaco_path, kernel_name). Also fixes a latent bug in the cache-hit path: the "hsaco missing, force a rebuild" guard only cleared cached_artifacts, not the separately-gated cache_path the rebuild branch actually checks. Verified the new _pack_hsaco() helper directly against real previously compiled PDI+insts and full-ELF artifacts (both produce a structurally valid, correctly-tagged hsaco section per aie.compiler.hsaco.dump). Could not get a full live Triton-JIT compile through this path to green in this session's ad hoc venv -- aircc itself fails the same generic MLIR legalization error for every kernel tried, reproduced by invoking it directly with no Triton/driver.py code involved, so this is an environment issue for Task 4 (real hardware) to resolve, not a defect in this change. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The final whole-branch review caught a real gap the plan missed: this hsaco/kernel-object switch silently broke examples/llm_q4nx/hsa_decode.py, the HSA decode path for llama-3.2-1b/3b, qwen3-4b, and gemma3-4b. It called triton_npu_hsa_prepare(pdi_path, insts_path) and patched 248 L-dependent instruction words per token in the runtime-owned buffer the dispatch packet pointed at -- a buffer that no longer exists once instructions live inside a frozen ROCr executable. Porting this to a per-context-length packed-executable cache (mirroring mlir-aie's own DispatchTime[T] LRU) was considered and explicitly not done: in a normal single-sequence decode the context length increments every token with no repeats, so an LRU buys nothing and every token would pay a full ROCr executable load instead of the old in-place memcpy -- an unmeasured, likely-large regression on a path benchmarked at ~51 tok/s. Instead: HsaProgram now raises immediately and clearly, before any weight loading; the four registry entries that used to advertise HSA decode support are flipped off (so the harness's existing up-front guard catches this before a multi-minute weight load, not after); the guard's error message covers both rejection reasons instead of asserting the wrong one; and the llama32_1b README is updated to say so. Also folds in the final review's other findings: AMD_TRITON_NPU_OUTPUT_FORMAT =pdi under the HSA runtime now gets a clear rejection (previously it would silently reach the xclbin-shaped packing branch and fail inside xclbinutil with a confusing error), the aie.compiler.hsaco import failure now names what it needs, and a few stale comments/docstrings are corrected (_pack_hsaco's Raises section, NPULauncher's PDI-era comment, get_npu_cache_dir's missing mention of kernel.hsaco). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
The core ROCr dispatch seam lacks hardware validation, and the new packing verification has unresolved issues.
4 open findings
What changed in this PR
Migrates HSA dispatch from raw PDI/instruction addresses to hsaco-loaded ROCr kernel objects while disabling incompatible LLM decode paths.
Changes:
- Adds hsaco packing, caching, loading, and kernel-object dispatch.
- Updates generated launcher ABI and output-format handling.
- Disables unsupported per-token instruction-patching decode paths.
| File | Description |
|---|---|
scripts/test_hsa_hsaco_pack.py |
Adds hsaco packing verification script. |
examples/llm_q4nx/registry.py |
Disables HSA for affected models. |
examples/llm_q4nx/hsa_decode.py |
Replaces legacy patching with an unsupported-path error. |
examples/llm_q4nx/harness.py |
Updates early HSA rejection messaging. |
examples/llama32_1b_q4nx/README.md |
Documents HSA decode removal. |
amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h |
Updates the public HSA preparation ABI. |
amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp |
Implements hsaco executable loading and kernel-object dispatch. |
amd_triton_npu/backend/hsa_launcher.py |
Passes hsaco paths and kernel names. |
amd_triton_npu/backend/driver.py |
Packs and caches hsaco artifacts. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
ypapadop-amd
marked this pull request as draft
October 9, 2026 17:13
Resolves conflicts between this branch's hsaco/kernel-object HSA dispatch switch and main's new HSA-runtime features (single-op NPUChain, mlir-air's fused prefill, Gemma4-E2B decode), all added after this branch's fork point and all built against the old raw PDI+insts ABI this branch replaces. Conflicts: - HsaRuntime.cpp: main added PDI deduplication (shared_ptr<PoolBuffer> across programs prepared from the same file) to the old prepare()/triton_npu_hsa_ program struct. Superseded, not merged: the hsaco model has no separate PDI pool allocation to deduplicate (PDI bytes are embedded in the packed hsaco and loaded by ROCr itself), so this branch's LoadedExecutable-based struct wins outright. Removed the now-orphaned pdis_ member and its two call sites that had merged in cleanly (no conflict marker) alongside it. - harness.py, registry.py: kept this branch's more detailed HSA-rejection message (names both possible reasons) while keeping main's new supports_xrt_zero_copy=True fields and comment style. Found via this merge (git's auto-merge produced valid-but-broken code, since the two branches touched disjoint regions of the same underlying problem): - registry.py: Gemma4-E2B's new supports_hsa=True (main's "Run Gemma4-E2B on the HSA runtime") routes through the exact same HsaProgram/_programs() path the other four now-disabled models do (via engine="ple"'s _make_hsa_ple_decoder_class) -- flipped off for the same reason. - multilaunch.py: HsaChainRunner (single-op NPUChain-on-HSA) calls triton_npu_hsa_prepare with raw (pdi_path, insts_path) and does no hsaco packing at all; separately, this branch's _get_output_format no longer returns "pdi" for the hsa runtime, so it was also about to get the wrong artifact shape entirely. Ported properly (not disabled): unlike the per-token decode case, a chain compiles once and is dispatched statically, exactly the pattern this branch's Task 3 already solved for the main compile path -- MultiLaunchBuilder.compile() now packs and caches a kernel.hsaco the same way compile_module does, for link_profile="hsa". - hsa_xrt.py (new file, mlir-air's fused-prefill-engine shim on HSA): architecturally the same per-call instruction-stream-patching problem as the decode path (bo._prepared() calls the removed triton_npu_hsa_prepare/ patch_insts). Disabled the same way (fail fast, clear error), not ported: porting would need the same new from-memory C ABI entry point the decode port would have, and while prefill's repeating op shapes might suit an LRU better than decode's ever-increasing context length, that is an unmeasured hypothesis, not something to build speculatively. Also, while in here: moved the HSA-decode-unsupported check to the very start of HsaFusedDecoder/HsaPleDecoder's __init__, before the superclass's own (expensive) weight-loading constructor runs -- the earlier version's claim that HsaProgram's own raise happened "before any weight loading" was only true via harness.py's registry-gated guard, not for a caller that constructs a decoder class directly. Removed the two dangling references to a design doc that was never going to be part of this repo (per an explicit earlier decision to keep process docs untracked) in favor of pointing at the in-tree docstrings that carry the same explanation. Rewrote scripts/test_hsa_hsaco_pack.py to compile through XRT (with npu_config.target set, bypassing device detection) and call _pack_hsaco directly, rather than through compile_module's HSA link_profile, which unconditionally builds the C++ HSA runtime against a real ROCR install even in compile_only mode -- exactly the dependency this test was meant to avoid needing. Verified: HsaRuntime.cpp still builds cleanly against the target ROCr (rocm-systems @ users/ypapadop-amd/aie-hsaco) with the same symbol check as before (triton_npu_hsa_prepare present, triton_npu_hsa_patch_insts absent). All touched Python files pass py_compile. The rewritten packing test is reasoned through but unexecuted in this session (same toolchain limitation documented on this branch's earlier commits) -- flagging honestly rather than claiming a run that didn't happen. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
# Conflicts: # amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp # examples/llama32_1b_q4nx/README.md # examples/llm_q4nx/hsa_decode.py
ypapadop-amd
marked this pull request as ready for review
October 21, 2026 16:19
…repare_bytes From a code review of the hsaco/kernel-object merge: - compile_module forced a rebuild when kernel.hsaco alone was missing from a cache hit, but not when the format-specific artifacts (aie.elf/aie.xclbin) were themselves missing while the compiled __npu_dispatch.so survived -- that left cache_hsaco_path/hsaco_kernel_name unassigned, raising UnboundLocalError instead of rebuilding. Reproduced on hardware (deleted a cached kernel's aie.elf, kept its .so) and confirmed fixed. - _scratchpad_relocs never checked the scratchpad's declared size was a multiple of 4; HsaElfProgram then viewed it as whole uint32 words via integer division, silently truncating a misbuilt one instead of failing where the build is actually broken. Verified against the real decode ELF (passes) and a corrupted copy (raises). Also, while in there: - Consolidated the hsaco-kernel-name-derivation + pack + cache sequence, duplicated across compile_module's and MultiLaunchBuilder.compile's cache-hit and rebuild branches, into _hsaco_kernel_name/_cached_hsaco/ _pack_and_cache_hsaco. - Added triton_npu_hsa_prepare_bytes (loads an hsaco already in memory, uncached -- there's no stable key for a buffer, and each one HsaElfProgram builds is single-use anyway) so hsa_decode.py no longer writes its hsaco to a temp file solely to read it back through the path-based ABI. Verified on hardware: matmul HSA compile from cold and warm caches, the cache-corruption regression above, and the Qwen2.5-3B HSA decode (scratchpad ELF, now loaded via prepare_bytes) all still produce correct results. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…spatch # Conflicts: # examples/llm_q4nx/decode_build.py
1 task done
erwei-xilinx
added a commit
to ROCm/rocm-systems
that referenced
this pull request
Oct 25, 2026
… variables (#14004) ## Motivation A full ELF can relocate the address of a buffer it does not carry. The common case is the control scratchpad (`scratch-pad-ctrl`), which holds the runtime parameters that a design reads in its preamble. MLIR-AIR's LLM decode uses it to pass the context length to one ELF for every token. The AIE hsaco loader refused any such relocation ("unsupported relocation symbol"). So a design with scratchpad parameters could not be loaded. Applications worked around this: they allocated the scratchpad, wrote its device address into the ELF, pointed the relocation at symbol 0 so that the loader skipped it, and repacked the hsaco for each allocation (for example amd/Triton-XDNA#158). ## Technical Details A relocation whose symbol is neither a PDI nor an argument index is now an **external variable**. It is resolved the same way as a GPU code object's undefined symbols: the application allocates the buffer and defines it before the load. ```cpp hsa_amd_memory_pool_allocate(dev_pool, size, 0, &scratchpad); hsa_executable_agent_global_variable_define(exe, aie_agent, "scratch-pad-ctrl", scratchpad); hsa_executable_load_agent_code_object(exe, aie_agent, reader, nullptr, nullptr); ``` - **Parse** (`amd_aie_elf.cpp`): the parser records each external symbol's patch sites and `st_size` in `Kernel::external_symbols`. Only `Address64` is accepted, as for PDIs. A site must hold 8 bytes inside the control code. All the symbols that name one variable must have the same `st_size`. - **Load** (`executable.cpp`): the loader looks up each name in `agent_symbols_`, the table that `hsa_executable_agent_global_variable_define` fills. If the variable is not defined, the load fails with `HSA_STATUS_ERROR_VARIABLE_UNDEFINED`. If its `st_size` bytes are not inside one allocation that the AIE agent can reach, the load fails with `HSA_STATUS_ERROR_INVALID_ALLOCATION`. Otherwise the loader writes the variable's agent address plus the addend into the control code, next to the PDI addresses, and records its BO handle. - **Dispatch** (`amd_xdna_driver.cpp`): `BuildFullElfCommand` flushes each external variable from the CPU cache and adds its BO to the command. So the host can change a parameter between dispatches with an ordinary store, and the application does not flush. The packet still lists only the kernel's arguments. - **Docs**: `docs/aie-hsaco-format.md` sections 7.4, 7.5, 10 and 11, and a new section 9.3. No change to the hsaco format, the packer, or the packet ABI. Two points for reviewers: 1. The command processor copies the scratchpad itself. On aie2p it does so only from the device-SVM pool (dev heap). A scratchpad in the data pool loads, but the dispatch fails. This PR documents the requirement but does not enforce it, because enforcement needs an allocation-to-region lookup that `Runtime` does not have. Should the loader refuse the load? 2. The new test design follows MLIR-AIR's sequence: load the empty device, then the design, then `sync_parameters()`, run, and load the empty device again. Without the final empty load, aie2p firmware gives each dispatch the scratchpad value of the dispatch before it. XRT gives the same result for one ELF run through a new `xrt::run` for each dispatch, so this comes from the firmware, not from this change. It is noted in the design file. ## Issue Tracking Follow-up to #12846 (AIE hsaco support). Independent of #13908. ## Test Plan - `elf_parse_test`: 3 new tests on a real scratchpad artifact (`vector_scalar_add_param.py`, a new IRON design that is built only as a full ELF): - The scratchpad is parsed as an external symbol, with the correct sites and size. - Any scheme other than `Address64` is refused. - An out-of-range site is refused. - `test_dispatch`: new `ExternalVariableTest` cases on hardware: - The load refuses an undefined scratchpad, a scratchpad in host memory, and a scratchpad that runs past the end of its allocation or of a VMM mapping. - One loaded kernel is dispatched 5 times, and the host changes the parameter with a plain store before each dispatch. Each output is checked against that dispatch's value. - The full AIE suite, as a regression check. - `TryLoad`/`LoadAndFreeze` take an optional list of variables to define before the load. `add_aie_kernel`'s full-ELF half is split out as `add_aie_full_elf` so that a full-ELF-only design can use it. ## Test Result On an aie2p (Krackan) NPU with `hsa-runtime64` built from this branch: - `elf_parse_test`: 16/16 (13 before this PR). - `aie_code_test`: 12/12. - `test_dispatch`: 56/56, run several times. - `test_memory`: 40/40. - `test_error_callback`: 2/2. - `ExternalVariableTest*` and the existing `FullElf*` tests: 20 consecutive runs each, with no failures. During development, one run of `ExternalVariableTest*` hung, right after the artifacts were rebuilt. It did not happen again in the 30 runs after it. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/rocm-systems/blob/develop/CONTRIBUTING.md. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
ROCR's AIE loader resolves a full ELF's control scratchpad against a variable defined with hsa_executable_agent_global_variable_define, and flushes it on every dispatch (ROCm/rocm-systems#14004). So the decode allocates the scratchpad, loads the cached hsaco with it defined, and writes L with plain stores. This drops the ELF relocation rewriting in driver.py, the per-decoder repack and prepare_bytes, the host-side flush, and the pointer_info probe it needed. Pin rocm-sdk-core 10.3.0a20261027, the first nightly with both ROCm/rocm-systems#14004 and #13908. On an older ROCR, repeated full-ELF dispatches give wrong results and the decode fails to load. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
3 of 4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Switches the HSA runtime from dispatching AIE kernels by raw PDI/instruction addresses to ROCR's hsaco model: aircc's artifact (xclbin on npu1, full ELF on npu2) is packed into an hsaco with mlir-aie's
aie-hsaco, loaded through the HSA code-object-reader/executable flow, and dispatched by kernel object. Current ROCR has removed raw-address dispatch, so this is required to keep the HSA path working.Changes
HsaRuntime.{h,cpp}: load an hsaco and dispatch by kernel object.triton_npu_hsa_prepare(hsaco_path, kernel_name)replaces(pdi_path, insts_path);triton_npu_hsa_prepare_with_variablesalso defines agent variables before loading. The in-tree full-ELF loader (AieFullElf.h) andtriton_npu_hsa_patch_instsare removed; ROCR does that work now.driver.py,hsa_launcher.py,multilaunch.py: HSA uses the same output format as XRT, packs and caches the hsaco next to the other artifacts, and generated launchers pass(hsaco_path, kernel_name).examples/llm_q4nx/hsa_decode.py): the context length goes through the full ELF's control scratchpad. The host allocates it and defines it asscratch-pad-ctrlwhen loading; ROCR flushes it per dispatch. The patched-instruction-stream decode is removed.utils/rocm-requirements.txt:rocm-sdk-core==10.3.0a20261027, the first nightly with feat(rocr/aie): resolve a full ELF's external symbols against defined variables ROCm/rocm-systems#14004 (scratchpad as external variable) and #13908 (PDI placement fix). Older nightlies give wrong results on repeated full-ELF dispatches or fail to load the decode.Not supported on HSA after this PR
ENGINE=auto RUNTIME=hsa):hsa_xrt.bo._preparedraises, since it relied on patching instruction streams in place.supports_hsa=False): its PLE engine has no full-ELF build.Test plan
On Krackan with the pinned wheel:
AMD_TRITON_NPU_RUNTIME=hsa:axpy,matmul_bf16_m64_n64_k64,test_softmaxpass.examples/llm_q4nx/test_decode_artifact.pypasses.scripts/test_hsa_hsaco_pack.pyfails in aircc ("failed to legalize func.func"), independent of this PR's runtime changes.🤖 Generated with Claude Code