Conversation
The AICore MMIO register-address tables are device constants: one 8-byte entry per physical sub-core, derived from the driver for a device id and a register kind, and unchanged for as long as the context holds that card. Two of their three owners disagreed with that. `SlotPersistentArgs::regs` held the Ctrl table per pipeline slot, so a depth-two runner derived and uploaded two byte-identical tables and paid the driver query twice. On a5 that query is a `halResMap` loop over every physical core, so the second copy also doubled the live driver-side mappings a worker holds. a2a3's Pmu table had no cache at all: `prepare_execution` queried the driver, allocated and uploaded it on every PMU-enabled run and `cleanup_execution` freed it again, for 600 bytes of addresses that cannot change. hw-native-sys#2222 recorded the reasoning as "the DFX buffers that genuinely vary per run, such as `pmu_reg_addrs`, keep their per-run lifetime", but the Pmu table is built by the same function and the same loop as Ctrl with only a different `addr_type`. It does not vary per run; it is PMU-*gated*. Gating and lifetime are orthogonal, and only the gate belongs to the run. Both tables now live on `DeviceRunnerBase` for the context's lifetime, committed on first use by the arch's `ensure_aicore_reg_table` and released in `finalize_common()` beside the device-wall buffers. Only the driver query stays arch-specific, which is why the entry point is on the subclass: a2a3 maps two MMIO pages selected by an `AicoreRegKind`, a5 one through a different signature. Kernel mode already owns the same table for its own lifetime through `PersistentKernelArgs`, and a context latches one mode permanently, so exactly one of the two owners is ever live. Each table carries an address and a committed flag rather than relying on a non-zero address. `init_aicore_register_addresses` records the address before the host-to-device copy and clears it afterwards only when a failed copy's rollback release also succeeded, so a retained address is owned but unwritten. Release keys on the address, because such a block still has to be freed; reuse keys on the flag, because handing the AICPU handshake an uncommitted table would deadlock the next task on a stream-sync timeout instead of failing the prepare. The retained address is passed back into the driver entry, which is what lets it reuse that block instead of stranding it. One consequence is worth stating: both kinds fold in the live AICore power-gating bitmap, so a committed table is a snapshot taken at first use. That was already true of Ctrl, whose dispatch addresses were never refreshed once a slot committed them. Committing Pmu the same way removes a disagreement rather than creating one — a per-run Pmu query could only ever produce a mask newer than the dispatch table it was measuring. A single shared commit flag needs no lock: `try_reserve_native_run` admits a second reservation only once the first is the active run, so a successor's preparation begins after its predecessor has launched and at most one thread is inside `prepare_execution`. A prepared successor overlaps device execution, not another preparation. `ensure_device_wall_buffer` already relies on this. If submission depth is ever generalised beyond the fixed two slots, every single-guard `ensure_*` on the prepare path has to be revisited together. Verified on a2a3 silicon with a temporary probe, since a passing test cannot distinguish a cached table from a re-derived one. Across three DFX suites, 14 preparations of which 11 enabled PMU produced one Ctrl commit and one Pmu commit, every preparation carrying the same two addresses. The pipeline-slot suite's five preparations spanned both slots and produced one Ctrl commit, both slots reading the same address. Full onboard sweep, the DFX PMU smoke, `vis_isolation`, cpput and pyut all pass; no simulation source is touched, since simulation sizes its register blocks from a run's own worker count and genuinely is per-run. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Warning Review limit reachedNext included review available in 28 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (10)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
ChaoWao
deleted the
refactor/own-the-aicore-register-tables-per-device
branch
September 17, 2026 03:42
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
1 participant
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.
Summary
The AICore MMIO register-address tables are device constants — one 8-byte entry per physical sub-core, derived from the driver for a device id and a register kind, unchanged for as long as the context holds that card. Two of their owners disagreed with that:
SlotPersistentArgs::regsheld the Ctrl table per pipeline slot, so a depth-two runner derived and uploaded two byte-identical tables and paid the driver query twice. On a5 that query is ahalResMaploop over every physical core, so the second copy also doubled the live driver-side mappings a worker holds.prepare_executionqueried the driver, allocated and uploaded it on every PMU-enabled run;cleanup_executionfreed it again — 600 bytes of addresses that cannot change. Refactor: reuse each pipeline slot's device argument blocks across runs #2222 recorded this as "the DFX buffers that genuinely vary per run, such aspmu_reg_addrs", but the Pmu table is built by the same function and loop as Ctrl with only a differentaddr_type. It is PMU-gated, not per-run varying; only the gate belongs to the run.Both tables now live on
DeviceRunnerBasefor the context's lifetime, committed on first use by the arch'sensure_aicore_reg_tableand released infinalize_common()beside the device-wall buffers. Only the driver query stays arch-specific, which is why the entry point is on the subclass. Kernel mode already owns the same table for its own lifetime viaPersistentKernelArgs, and a context latches one mode permanently, so exactly one owner is ever live.Each table carries an address and a committed flag rather than relying on a non-zero address.
init_aicore_register_addressesrecords the address before the H2D copy and clears it afterwards only when a failed copy's rollback release also succeeded, so a retained address is owned but unwritten. Release keys on the address (such a block still has to be freed); reuse keys on the flag (handing the AICPU handshake an uncommitted table would deadlock the next task on a stream-sync timeout rather than fail the prepare). The retained address is passed back into the driver entry so it can reuse that block instead of stranding it.Part of #2254, which proposed this as step 4. Step 1 is #2257.
Testing
A passing test cannot distinguish a cached table from a re-derived one, so reuse was measured with a temporary probe on a2a3 silicon (removed before commit):
--enable-pmu 2pipeline_slotsexamples tests/st -m 'not sdma' --exclude-level 4), zero failures--enable-pmu 2), both arches' path exercised on a2a3tests/st/vis_isolation— the guard forinit_aicore_register_addressesgoing throughacl_to_hal_device_idtests/lint/clang_tidy.py, cpplint,check_retired_names.pyNot verified locally: a5 onboard. This box is a2a3, which is the right silicon for the Pmu half since a5 has no PMU register table; the a5 half is Ctrl-only and relies on CI.
🤖 Generated with Claude Code