Update profiling buffer pool to SPSC#1287
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:
📝 WalkthroughWalkthroughBufferPoolManager and ProfilerAlgorithms are refactored to use per-shard SPSC ring buffers for ready/done/recycled/retired queues, replacing mutex/std::queue mechanisms. Collectors now seed shard-affine recycled lanes via a shard argument to push_recycled. Teardown separates device-allocation release from host-shadow release. Tests and docs are updated accordingly. ChangesShard-local SPSC buffer pool redesign
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 refactors the profiling framework's buffer management system to improve performance and reduce overhead. It replaces mutex-guarded queues with lock-free SPSC (Single Producer Single Consumer) rings for ready, done, and recycled buffer pools. Additionally, it introduces a batched block allocation strategy where a single large device block is allocated, registered once, and carved into fixed-size sub-buffers, significantly reducing HAL registration costs. Range-based host pointer resolution is implemented to support these carved sub-buffers, and the runtime hot path is optimized to be lock-free and allocation-free. All profiling collectors (PMU, L2Swimlane, DepGen, ScopeStats, and TensorDump) have been updated to adopt this new architecture, accompanied by updated documentation and comprehensive unit tests. I have no further feedback to provide as the implementation is robust and well-tested.
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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/common/platform/include/host/buffer_pool_manager.h (1)
633-664: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winFreed allocation base is never removed from
released_allocations_, risking a permanent leak on address reuse.
claim_release_pointer(Line 811-825) insertsrelease_ptrintoreleased_allocations_once, permanently blocking any future claim on that address — butfree_buffer(Line 633-664) removes the entry fromdev_to_host_/block_ranges_and callsops_.free_(release_ptr)without ever erasingrelease_ptrfromreleased_allocations_. If the underlying allocator later hands back the exact same address for a newalloc_and_register_block/register_mappingcall within the same collector lifetime (i.e., before a fullclear_mappings()/release_all_owned()reset),claim_release_pointer/release_pointer_forwill treat that fresh allocation as already released, and any subsequentfree_buffer/release_one_buffercall on it will silently no-op — leaking that device (and possibly host-shadow) allocation for the rest of the run.Erase
release_ptrfromreleased_allocations_insidefree_buffer(undermapping_mutex_) once the release is claimed and device memory is actually freed.Proposed fix
void free_buffer(void *dev_ptr) { if (dev_ptr == nullptr) return; void *release_ptr = nullptr; if (!claim_release_pointer(dev_ptr, &release_ptr)) return; void *host_ptr = nullptr; bool free_host_shadow = false; { std::scoped_lock<std::mutex> lock(mapping_mutex_); auto it = dev_to_host_.find(release_ptr); host_ptr = (it != dev_to_host_.end()) ? it->second : nullptr; if (it != dev_to_host_.end()) { dev_to_host_.erase(it); } block_ranges_.erase(...); free_host_shadow = (host_ptr != nullptr && malloc_shadows_.erase(host_ptr) > 0); + released_allocations_.erase(release_ptr); } if (ops_.free_) { ops_.free_(release_ptr); }Also applies to: 802-825
🤖 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/host/buffer_pool_manager.h` around lines 633 - 664, The free path in free_buffer still leaves the allocation address recorded in released_allocations_, so a later reuse of the same device pointer can be treated as already released and leak the new allocation. Update free_buffer, alongside the existing dev_to_host_ and block_ranges_ cleanup under mapping_mutex_, to also erase release_ptr from released_allocations_ once the release has been claimed and ops_.free_ has run; use claim_release_pointer and free_buffer as the main symbols to locate the change.
🧹 Nitpick comments (1)
src/a2a3/platform/shared/host/pmu_collector.cpp (1)
113-125: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a retire fallback when
push_recycledfails during init, matchingallocate_recycled_batch's own pattern.The return value of
push_recycledis discarded ((void)..., Line 124).BufferPoolManager::allocate_recycled_batchtreats a full recycled ring as recoverable by falling back toretire_unqueued_bufferso the buffer is still teardown-tracked/accounted for; here a full ring just silently drops the surplus buffer from the pool for the remainder of the run (it will still be freed correctly atfinalize()viadev_to_host_, so this isn't a leak — just wasted capacity with no error signal at this call site beyond the manager's internalLOG_ERROR).Proposed fix
} else { // Surplus buffers go into the recycled pool, available to any core. int shard = (num_threads > 0) ? (c % num_threads) : c; - (void)manager_.push_recycled(0, dev_ptr, shard); + if (!manager_.push_recycled(0, dev_ptr, shard)) { + (void)manager_.retire_unqueued_buffer(0, dev_ptr, shard); + } }🤖 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/a2a3/platform/shared/host/pmu_collector.cpp` around lines 113 - 125, The init path in pmu_collector’s buffer distribution silently ignores a failed BufferPoolManager::push_recycled call, unlike BufferPoolManager::allocate_recycled_batch which falls back to retire_unqueued_buffer when the recycled ring is full. Update the surplus-buffer branch in the pmu_collector initialization logic to check the push_recycled return value and, on failure, retire the buffer through the same fallback path so it stays tracked/accounted for instead of being dropped from the pool for the rest of the run.
🤖 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/buffer_pool_manager.h`:
- Around line 153-213: `ProfilerModuleHostQueueCapacity` is using
`Module::kReadyQueueSize` for all ring buffers, which incorrectly ties
recycled/retired pool sizing to the ready queue. Update the capacity trait in
`buffer_pool_manager.h` so `ReadyRing` keeps the ready-queue sizing, but
`DoneRing` and the per-shard `RecycledRing`/retired rings use a capacity derived
from the module’s buffer-pool constants (for example `kSlotCount`,
`kCollectorThreadCount`, or the warm/surplus target) via a separate trait or
specialization. Keep the existing symbols like `SpscRing`,
`ProfilerModuleHostQueueCapacity`, and the ring typedefs intact while splitting
the semantics so warm-up and recycling cannot silently drop buffers.
---
Outside diff comments:
In `@src/common/platform/include/host/buffer_pool_manager.h`:
- Around line 633-664: The free path in free_buffer still leaves the allocation
address recorded in released_allocations_, so a later reuse of the same device
pointer can be treated as already released and leak the new allocation. Update
free_buffer, alongside the existing dev_to_host_ and block_ranges_ cleanup under
mapping_mutex_, to also erase release_ptr from released_allocations_ once the
release has been claimed and ops_.free_ has run; use claim_release_pointer and
free_buffer as the main symbols to locate the change.
---
Nitpick comments:
In `@src/a2a3/platform/shared/host/pmu_collector.cpp`:
- Around line 113-125: The init path in pmu_collector’s buffer distribution
silently ignores a failed BufferPoolManager::push_recycled call, unlike
BufferPoolManager::allocate_recycled_batch which falls back to
retire_unqueued_buffer when the recycled ring is full. Update the surplus-buffer
branch in the pmu_collector initialization logic to check the push_recycled
return value and, on failure, retire the buffer through the same fallback path
so it stays tracked/accounted for instead of being dropped from the pool for the
rest of the run.
🪄 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: d8a2622d-70be-442e-9c2b-cdc9dcdda093
📒 Files selected for processing (21)
docs/dfx/args-dump.mddocs/dfx/dfx-buffer-capacity-audit.mddocs/dfx/l2-swimlane-profiling.mddocs/dfx/pmu-profiling.mddocs/profiling-framework.mdsrc/a2a3/platform/include/host/pmu_collector.hsrc/a2a3/platform/shared/host/pmu_collector.cppsrc/a5/platform/include/host/pmu_collector.hsrc/a5/platform/shared/host/pmu_collector.cppsrc/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/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.cpptests/ut/cpp/CMakeLists.txttests/ut/cpp/common/test_buffer_pool_manager.cpp
zmnobug
force-pushed
the
issue-1251-spsc-buffer-pool
branch
from
a712db9 to
dc445ea
Compare
- Retire undelivered ready-full buffers instead of writing done from drain - Make retired buffers use a mutex-protected exceptional holding pool - Seed single-orchestrator profiler surplus buffers into the owner lane - Add regression coverage and update profiling framework docs
zmnobug
force-pushed
the
issue-1251-spsc-buffer-pool
branch
from
dc445ea to
893b1f1
Compare
Re: issue #1251 子改动 #5(alloc off hot path, on replenish, batched + proactive)已做到的部分:
没做到的部分:
issue #5 明确要求把 watermark 驱动的主动分配放到 replenish 线程上。请按 issue 要求补足这部分实现,或说明不实现的原因。 |
死代码:
|
建议(增强,非阻塞):给
|
播种 shard 与运行期消费 shard 错配(建议按 blockdim/cluster 分组)问题:余量 buffer 的播种 shard 用 叠加本 PR 删掉了 不需要序号严格对齐,只需数量匹配:recycled buffer 在同一
依赖的前提( 改法(L2 的 AicpuTask / AicoreTask 两处播种,PMU 同理): // 播种 shard 必须与运行期 recycled 消费 shard 一致:runtime 按 cluster 对齐分派
// (cluster ci -> AICPU 线程 ci % N,每 cluster = PLATFORM_CORES_PER_BLOCKDIM 个 core)。
// recycled buffer 在 (shard,kind) 内可互换,故只需每条 shard 的“数量”匹配该 lane 的
// core 数,无需 core 序号严格对齐;按 blockdim 分组轮询即可让数量精确吻合(与 AIC/AIV
// 发现顺序无关)。跨 shard 借用(pop_recycled_any)已删,数量不匹配会导致预热期欠货丢样。
int shard = (aicpu_thread_num > 0) ? ((i / PLATFORM_CORES_PER_BLOCKDIM) % aicpu_thread_num) : i;低危(仅预热期、自愈),不阻塞合入;但 AicoreTask/PMU 无 warm_target 兜底,建议一并修掉并加上上述注释。 |
recycled 环容量过配(建议按 blockdim 线程数平分 + 2x 余量,并加超容护栏)问题: using DoneRing = SpscRing<DoneInfo, kPoolQueueCapacity>; // 920
using RecycledRing = SpscRing<void *, kPoolQueueCapacity>; // 920这个
单个 L2 子系统 manager 里 recycled 现约 115 KB,其中绝大部分是过配(SchedPhase 环需 ≤24 却开 920,过配约 38×)。 根因:一个常量(为 done 算的"各 kind 之和")被同时套在了"每格只装单 kind、且只占一条 lane 份额"的 recycled 环上。 当前结论(增量修,动态化留后续):线程/shard 数当前是编译期固定 // recycled 每格只装单 kind、只占一条 lane 份额;按 PLATFORM_MAX_AICPU_THREADS 平分 + 2x 余量。
// (done 仍用“各 kind 之和”,因为它混装所有 kind。)运行期按实际 aicpu_thread_num 动态定容待后续统一改造。
constexpr size_t kRecycledCapacity =
PLATFORM_MAX_CORES * PLATFORM_PROF_BUFFERS_PER_CORE / PLATFORM_MAX_AICPU_THREADS * 2; // AicpuTask=288
using RecycledRing = SpscRing<void *, kRecycledCapacity>;安全前提 + 护栏: if (per_shard_seeded_count > kRecycledCapacity) {
LOG_ERROR("recycled ring under-provisioned: shard=%d kind=%d need=%d cap=%zu "
"(aicpu_thread_num=%d) — switch to runtime-sized rings", ...);
}低危、不阻塞;省内存(L2 recycled 115KB→约 36KB)+ 语义更清晰(done 取和 / recycled 取单 lane 份额),护栏保证未来"核多线程少"配置不会静默丢样。 |
- Add runtime recycled-lane watermark replenishment while keeping drain free_queue publication allocation-free - Split done and recycled host ring capacities for per-kind lane sizing - Align L2 and PMU recycled seeding with cluster-owner shards - Add SPSC concurrency and replenish capacity regression tests - Update profiling documentation to describe the new replenish contract
2 participants
Summary
Testing
CCACHE_DIR=/tmp/simpler-ccache cmake --build tests/ut/cpp/build --target test_buffer_pool_managertests/ut/cpp/build/test_buffer_pool_manager(12/12 passed)CCACHE_DIR=/tmp/simpler-ccache cmake --build tests/ut/cpp/buildctest --test-dir tests/ut/cpp/build -LE requires_hardware --output-on-failure(54/54 passed after rerun outside sandbox; sandboxed run blocked the socket test withOperation not permitted)CCACHE_DIR=/tmp/simpler-ccache .venv/bin/python -m pip install --no-build-isolation -e .git diff --check893b1f16: all checks passed.Benchmark
Baseline vs PR head before review fix
task-submit:task_20260707_203817_92120128835, baseline and current run sequentially on the same NPU 2.upstream/mainat3aa7f26d.a712db9b.3.3754 GB/s, current3.8195 GB/s,+13.16%3.5904 GB/s, current3.8208 GB/s,+6.42%Latest PR head retest
task-submit:task_20260707_231640_262842832235, latest PR head893b1f16plus the same local synthetic WIP, on NPU 3.l2_swimlane_records.jsonasrecords * 32 bytes / 5 ms.3.8208,3.8208,3.8208,1.9456,3.8016 GB/s(run 4 was a retained low outlier).3.3754 GB/s, latest retest3.4419 GB/s,+1.97%3.5904 GB/s, latest retest3.8208 GB/s,+6.42%3.8160 GB/s,+13.05%Notes
close #1251