feat(dfx): unified block-on-contention backpressure across DFX subsystems#1313
feat(dfx): unified block-on-contention backpressure across DFX subsystems#1313doraemonmj wants to merge 1 commit into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis change adds opt-in DFX backpressure configuration, extends shared-memory headers and collector initialization, and coordinates device contention with host queue draining, freeze release, collector inflight tracking, and tensor-dump arena epoch barriers. ChangesDFX backpressure configuration
Shared-memory and runtime coordination
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a unified DFX backpressure mechanism (block-on-contention and global-sync freeze) across several DFX subsystems, including L2 swimlane, PMU, tensor dump, dep_gen, and scope_stats. The changes span configuration structures, Python bindings, and host/device-side coordination logic. The reviewer provided critical feedback to address issues in the tight spin loops within profiler_device_engine.h, where reading the system counter on every iteration causes high MMIO overhead and can lead to premature timeouts; gating these reads periodically resolves both issues. Additionally, the reviewer recommended adding read memory barriers (rmb()) in dfx_backpressure_device.h to ensure proper memory visibility on weak-ordering architectures like aarch64, and optimizing host-side checks in profiler_base.h by consolidating multiple DMA reads into a single contiguous operation on non-SVM platforms.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
simpler_setup/scene_test.py (1)
1335-1360: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winStandalone
run_modulepath is missing--enable-dfx-backpressure.The pytest path (
test_run) correctly reads and threadsenable_dfx_backpressure, but the standalonerun_moduleentry point neither registers the--enable-dfx-backpressureargparse argument nor forwards it torun_class_cases. All other DFX flags (--enable-l2-swimlane,--dump-args,--enable-dep-gen,--enable-pmu,--enable-scope-stats,--enable-swimlane-overhead) are present in both paths. Standalone users will silently never get backpressure sincerun_class_casesdefaultsenable_dfx_backpressure=False.🔧 Add the missing argparse argument and forwarding
Add the argparse argument near the other DFX flags (after
--enable-dep-gen):parser.add_argument( "--enable-dep-gen", action="store_true", help="Enable dep_gen capture (SubmitTrace ring, first round only)", ) + parser.add_argument( + "--enable-dfx-backpressure", + action="store_true", + default=False, + help="DFX backpressure: block-on-contention instead of dropping records " + "(currently wired for L2 swimlane). Requires a diagnostic (e.g. --enable-l2-swimlane).", + ) parser.add_argument(Then forward it in the
run_class_casescall insiderun_module:run_class_cases( worker, inst, [case], callable_obj=callable_obj, sub_handles=sub_handles, rounds=args.rounds, skip_golden=args.skip_golden, enable_l2_swimlane=args.enable_l2_swimlane, enable_dump_args=args.dump_args, enable_pmu=args.enable_pmu, enable_dep_gen=args.enable_dep_gen, enable_scope_stats=args.enable_scope_stats, + enable_dfx_backpressure=args.enable_dfx_backpressure, enable_swimlane_overhead=args.enable_swimlane_overhead, )Also applies to: 1549-1563
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simpler_setup/scene_test.py` around lines 1335 - 1360, Standalone run_module lacks the DFX backpressure option and does not pass it through. Add an --enable-dfx-backpressure argparse flag alongside the other DFX options in run_module, then read its value and forward it as enable_dfx_backpressure in the run_class_cases call; keep behavior consistent with test_run.
🧹 Nitpick comments (2)
src/common/platform/include/aicpu/profiler_device_engine.h (1)
53-59: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAdd
spin_relax()to block-mode spin loops.Both block-mode contention paths (
wait_for_ready_queue_spaceandwait_for_free_queue_entry) spin viacontinuewithout a relax hint, unlikepeer_freeze_barrierandwait_for_releaseindfx_backpressure_device.hwhich both callspin_relax(). The tight load loop wastes power on aarch64 and pins a core at 100% on x86_64 sim. Addingdfx_backpressure::spin_relax()beforecontinueis consistent with the other DFX spin paths and is a one-line change per loop.♻️ Proposed fix for both spin loops
In
wait_for_ready_queue_space(line 56-58):if (dfx_backpressure::block_on_contention(header)) { dfx_backpressure::mark_contended(header, &contended_signalled); + dfx_backpressure::spin_relax(); continue; }In
wait_for_free_queue_entry(line 92-94):if (dfx_backpressure::block_on_contention(header)) { dfx_backpressure::mark_contended(header, &contended_signalled); + dfx_backpressure::spin_relax(); continue; }Also applies to: 89-95
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/platform/include/aicpu/profiler_device_engine.h` around lines 53 - 59, Add dfx_backpressure::spin_relax() immediately before continue in both block-mode spin loops, wait_for_ready_queue_space and wait_for_free_queue_entry, matching the existing relaxation pattern used by peer_freeze_barrier and wait_for_release.src/common/platform/include/common/dfx_backpressure_device.h (1)
91-97: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding a
block_on_contentionguard tomark_contendedfor API consistency.
peer_freeze_barrierandwait_for_releaseboth self-guard withblock_on_contention(header), butmark_contendedonly checks*signalledandheader != nullptr. If called in drop mode, it would writecontended = 1to shared memory unnecessarily. Current callers (e.g.,dfx_arena_barrierintensor_dump_aicpu.cpp) gate before calling, so this isn't a live bug — but the inconsistency means the function's contract relies on caller discipline rather than being self-enforcing like its peers.♻️ Optional: add self-guard to `mark_contended`
template <typename Header> inline void mark_contended(Header *header, bool *signalled) { - if (!*signalled && header != nullptr) { + if (!*signalled && block_on_contention(header)) { header->backpressure.contended = 1; *signalled = true; } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/platform/include/common/dfx_backpressure_device.h` around lines 91 - 97, mark_contended should self-enforce the block-on-contention policy like peer_freeze_barrier and wait_for_release. Update mark_contended to require block_on_contention(header) before writing header->backpressure.contended, while preserving the existing signalled and null-header checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/common/platform/include/host/profiler_base.h`:
- Around line 442-443: Make the freeze state accessed by the drain threads and
mgmt_replenish_loop use a host-side atomic or mutex-protected variable instead
of directly sharing header->backpressure.freeze_active; update the device header
from that synchronized host state with the required memory ordering, and have
the drain path consult the synchronized state before calling top_up_free_queue.
In `@src/common/platform/include/host/tensor_dump_collector.h`:
- Around line 222-225: Update the block_on_contention documentation in the
tensor dump collector header to state that arena-lap protection is included
through the epoch-barrier behavior; clarify that only payloads exceeding the
per-thread payload arena remain subject to truncation, removing the statement
that the payload arena is out of scope.
---
Outside diff comments:
In `@simpler_setup/scene_test.py`:
- Around line 1335-1360: Standalone run_module lacks the DFX backpressure option
and does not pass it through. Add an --enable-dfx-backpressure argparse flag
alongside the other DFX options in run_module, then read its value and forward
it as enable_dfx_backpressure in the run_class_cases call; keep behavior
consistent with test_run.
---
Nitpick comments:
In `@src/common/platform/include/aicpu/profiler_device_engine.h`:
- Around line 53-59: Add dfx_backpressure::spin_relax() immediately before
continue in both block-mode spin loops, wait_for_ready_queue_space and
wait_for_free_queue_entry, matching the existing relaxation pattern used by
peer_freeze_barrier and wait_for_release.
In `@src/common/platform/include/common/dfx_backpressure_device.h`:
- Around line 91-97: mark_contended should self-enforce the block-on-contention
policy like peer_freeze_barrier and wait_for_release. Update mark_contended to
require block_on_contention(header) before writing
header->backpressure.contended, while preserving the existing signalled and
null-header checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b20556bc-c604-4553-adf5-1374657a04eb
📒 Files selected for processing (41)
conftest.pypython/bindings/task_interface.cpppython/simpler/worker.pysimpler_setup/scene_test.pysrc/a2a3/platform/include/common/dep_gen.hsrc/a2a3/platform/include/common/l2_swimlane_profiling.hsrc/a2a3/platform/include/common/platform_config.hsrc/a2a3/platform/include/common/pmu_profiling.hsrc/a2a3/platform/include/common/tensor_dump.hsrc/a2a3/platform/include/host/pmu_collector.hsrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/shared/host/pmu_collector.cppsrc/a2a3/platform/sim/host/device_runner.cppsrc/a5/platform/include/common/dep_gen.hsrc/a5/platform/include/common/l2_swimlane_profiling.hsrc/a5/platform/include/common/pmu_profiling.hsrc/a5/platform/include/common/tensor_dump.hsrc/a5/platform/include/host/pmu_collector.hsrc/a5/platform/onboard/host/device_runner.cppsrc/a5/platform/shared/host/pmu_collector.cppsrc/a5/platform/sim/host/device_runner.cppsrc/common/platform/include/aicpu/profiler_device_engine.hsrc/common/platform/include/common/dfx_backpressure_device.hsrc/common/platform/include/common/scope_stats.hsrc/common/platform/include/host/buffer_pool_manager.hsrc/common/platform/include/host/dep_gen_collector.hsrc/common/platform/include/host/l2_swimlane_collector.hsrc/common/platform/include/host/profiler_base.hsrc/common/platform/include/host/scope_stats_collector.hsrc/common/platform/include/host/tensor_dump_collector.hsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/shared/aicpu/l2_swimlane_collector_aicpu.cppsrc/common/platform/shared/aicpu/tensor_dump_aicpu.cppsrc/common/platform/shared/host/dep_gen_collector.cppsrc/common/platform/shared/host/l2_swimlane_collector.cppsrc/common/platform/shared/host/scope_stats_collector.cppsrc/common/platform/shared/host/tensor_dump_collector.cppsrc/common/platform/sim/host/device_runner_base.cppsrc/common/platform/sim/host/device_runner_base.hsrc/common/task_interface/call_config.h
doraemonmj
force-pushed
the
dfxbackpressure
branch
2 times, most recently
from
046005f to
48a8278
Compare
…tems
Add a per-subsystem DfxBackpressureHeader{block_on_contention, freeze_active,
contended} embedded in every DFX DataHeader, plus a host-side global-sync freeze
state machine (update_backpressure_freeze). On real buffer contention a device
lane raises the leader signal; the host opens a freeze so every AICPU lane parks
at its next buffer-switch gate, then releases once all lanes drained and
free_queues refilled — one lane-aligned common-mode gap instead of scattered
per-lane sparseness. Gated behind enable_dfx_backpressure (default off, drops as
before).
tensor_dump additionally guards its payload arena: a pre-write predictive lap
check (dfx_arena_barrier) flushes + parks + resets the epoch before overwriting
not-yet-drained payloads, and the freeze release adds a per-subsystem predicate
(collector_inflight()==0) so the device cannot lap the arena before the
collector has synchronously pulled it into host memory.
1 participant
What
Unify the five DFX subsystems (L2 swimlane, PMU, dep_gen, scope_stats,
tensor_dump) onto a single block-on-contention backpressure mechanism,
gated behind
enable_dfx_backpressure(default off — drops as before).DfxBackpressureHeader{block_on_contention, freeze_active, contended}is embedded in every DFXDataHeader.update_backpressure_freeze):on real buffer contention a device lane raises the leader signal, the host
opens a freeze so every AICPU lane parks at its next buffer-switch gate,
then releases once all lanes drained and free_queues refilled — one
lane-aligned common-mode gap instead of scattered per-lane record loss.
tensor_dumpadditionally guards its payload arena: a predictive lap check(
dfx_arena_barrier) flushes + parks + resets the epoch before overwritingnot-yet-drained payloads; freeze release adds a
collector_inflight()==0predicate so the device cannot lap the arena before the collector has pulled
it into host memory.
Test results
Verified on a2a3 silicon (device isolated via
task-submit).1. flag-OFF has no hot-path overhead — benchmarked merge-base (no
backpressure code) vs this change, both flag-off, same device, 100 rounds,
8 cases (
tensormap_and_ringbuffer), all 8/8 pass:max |ΔEffective| = 2.18%— entirely within run-to-run noise.qwen3_14b_decode/StressBatch16Seq3500: ΔEffective −0.03%.2. flag-ON is safe across all 5 subsystems — ran each subsystem's DFX
smoke with
--enable-dfx-backpressureand all records-per-buffer knobsshrunk to 4 (forces buffer rotation / contention):
dropped == 0)No
HandleTaskTimeout/Task Allocator Deadlock/507018/Timeout(cycles)on any run — the block path is deadlock-free and loses norecords when engaged. Heavy-contention drop→0 evidence (swimlane 4590→0,
PMU 28744→0) was captured during development.
Re-verified the full flag-ON matrix after rebasing onto latest
main(which landed L2-swimlane rotation/wmb changes) — all 5 still pass, no hang.