Add: protect args dump arena reuse with backpressure

[doraemonmj]

Summary

• Protect args-dump payload arenas from wrapping over records that are still in the host collector/writer pipeline.
• Publish a per-AICPU-thread payload watermark before ready-queue handoff, track writer completion per source thread, and acknowledge all threads together only when every written[t] == published[t].
• Reuse the shared DFX freeze/release state machine so arena reuse remains globally coordinated without adding device shared-memory fields or a parked-thread epoch protocol.
• Bound the freeze-release wait and the subsequent payload acknowledgement wait with one existing DFX timeout budget.
• Cover both a2a3 and a5 device/collector paths, including cross-thread count compensation and pop-gate behavior.

Core flow

Device publishes per-thread payload watermarks
                    ↓
Host freezes arena reuse and drains collector/writer work
                    ↓
Every thread satisfies written[t] == published[t]
                    ↓
Host ACKs all thread watermarks and releases the global freeze

Testing

• Full non-hardware C++ suite: 76/76 passed.
• a2a3sim args-dump scene test passed.
• a5sim args-dump scene test passed.
• Ascend910_9392 onboard page-unroll through task-submit, with --dump-args 2, temporary 256 MiB/thread arena, and batch 32:
  • 1 test passed in 61.34 s.
  • Exported 1,024 args and 4.35 GB of payload data.
  • truncated_args=0, dropped_records=0, and dropped_overwrite=0.
• Pre-commit hooks passed, including clang-format, clang-tidy, cpplint, and markdownlint.

Known capacity/throughput limit

With the default 128 MiB/thread arena and the full page-unroll batch 256 workload, the host completed 46 trigger/release pairs with no dropped or overwritten records, but cumulative dump/write waits exceeded the outer 45-second AICPU op-execute timeout. Increasing the arena to 256 MiB/thread and using batch 32 produced a complete passing run. The backpressure protocol protects correctness; very large dump workloads still need enough arena capacity and host write throughput to finish within the outer op timeout.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR adds payload publication and completion watermarks to dump buffers. AICPU arena reuse now waits for host acknowledgement. The host collector refreshes device state on demand and releases backpressure after all payloads are written. Tests cover A2A3, A5, and multi-arena flows.

Changes

Args dump arena backpressure

Layer / File(s)	Summary
Payload watermark contract src/common/platform/include/common/args_dump.h, src/common/platform/include/host/args_dump_collector.h, src/common/platform/shared/host/args_dump_collector.cpp	DumpBufferState stores published and completed payload counts. The collector tracks written payloads and exposes backpressure_release_ready().
AICPU arena capacity and publication src/common/platform/shared/aicpu/args_dump_aicpu.cpp	AICPU counts published payloads, prepares arena capacity, waits for bounded backpressure, and drops records when capacity preparation fails.
Host collection and release src/common/platform/shared/host/args_dump_collector.cpp	The collector refreshes device state on demand, copies arena payloads, and publishes completion counts after all payloads are written.
Protocol documentation and validation docs/dfx/args-dump.md, tests/ut/cpp/a2a3/test_args_dump.cpp, tests/ut/cpp/a5/test_args_dump.cpp, tests/ut/cpp/common/test_args_dump_collector.cpp	Documentation describes freeze-based arena reuse. Tests verify payload preservation, counter updates, and release across multiple arenas.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant AICPU
  participant DumpBufferState
  participant ArgsDumpCollector
  AICPU->>DumpBufferState: publish payload count
  AICPU->>DumpBufferState: wait for completed payload count
  ArgsDumpCollector->>DumpBufferState: publish completed payload count
  DumpBufferState-->>AICPU: allow arena reuse

Loading

Possibly related PRs

• hw-native-sys/simpler#1313: Both modify tensor-dump arena backpressure and payload-preservation logic.
• hw-native-sys/simpler#1342: Both modify the shared DumpBufferState definition.
• hw-native-sys/simpler#1567: Both modify the args-dump payload lifecycle and related tests.

Poem

A rabbit guards the frozen ground,
While payloads wait in queues around.
Counts rise, then hosts reply,
Safe arenas turn nearby.
“Hop!” says the hare, “the dump stays whole!”

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the primary change: protecting args-dump arena reuse with backpressure.
Description check	✅ Passed	The description directly explains the backpressure design, implementation flow, affected paths, testing, and known capacity limits.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

src/common/platform/shared/host/args_dump_collector.cpp (1)

279-295: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift

Consider copying only the newly-written arena span instead of the full arena on every collection.

Once write_offset exceeds ai.size, bytes_to_copy is clamped to the full arena size on every call, so each buffer collection re-copies the entire per-thread arena (up to 128 MiB by default) rather than only the bytes written since the last refresh. With PLATFORM_DUMP_BUFFERS_PER_THREAD = 8, this can multiply the actual host-device copy volume roughly 8x per arena lap. Track the last-copied offset per thread and copy only the delta (handling the wrap case) to reduce I/O on the a5 rtMemcpy/memcpy transport.

Please confirm whether the actual copy volume introduced by this per-collection full-arena refresh was measured against the prior full-shm-mirror-per-tick baseline it replaces.

🤖 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/shared/host/args_dump_collector.cpp` around lines 279 -
295, Update the arena refresh logic around arenas_ and the write_offset
calculation to retain a last-copied offset for each thread and copy only the
newly written span, including correct wrap-around handling when the write cursor
laps ai.size. Preserve the existing bounds checks and device-to-host copy
behavior, and update the tracked offset after each successful refresh. Measure
and report the resulting copy volume against the prior full-shm-mirror-per-tick
baseline.

🤖 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/shared/host/args_dump_collector.cpp`:
- Around line 286-289: Check the return value of profiling_copy_from_device in
the dump-buffer state flow before using state->arena_write_offset; on failure,
follow the existing error-handling behavior used by equivalent calls in
backpressure_release_ready() and stop processing rather than copying bytes from
a stale offset.
- Around line 615-650: Update ArgsDumpCollector::backpressure_release_ready() to
acknowledge each dump thread independently rather than gating all progress on
written_payload_count_ equaling the summed published_payload_count. Remove or
bypass the global-count check, and advance each thread’s completed_payload_count
when that thread’s published data has been written/read, preserving the existing
per-thread device copy and failure handling.

In `@tests/ut/cpp/a2a3/test_args_dump.cpp`:
- Around line 235-249: The host-simulation lambda copies arena, so its race
check reads a stale snapshot instead of the shared buffer. In
tests/ut/cpp/a2a3/test_args_dump.cpp lines 235-249, update the std::thread host
capture to capture arena by reference; apply the identical change in
tests/ut/cpp/a5/test_args_dump.cpp lines 235-249 so payload_preserved_before_ack
validates the live arena contents.

---

Nitpick comments:
In `@src/common/platform/shared/host/args_dump_collector.cpp`:
- Around line 279-295: Update the arena refresh logic around arenas_ and the
write_offset calculation to retain a last-copied offset for each thread and copy
only the newly written span, including correct wrap-around handling when the
write cursor laps ai.size. Preserve the existing bounds checks and
device-to-host copy behavior, and update the tracked offset after each
successful refresh. Measure and report the resulting copy volume against the
prior full-shm-mirror-per-tick baseline.

🪄 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 Plus

Run ID: 9ac06d8a-bd04-4051-96d6-c535362218ab

📥 Commits

Reviewing files that changed from the base of the PR and between b5261a7 and f7eb998.

📒 Files selected for processing (8)

• docs/dfx/args-dump.md
• src/common/platform/include/common/args_dump.h
• src/common/platform/include/host/args_dump_collector.h
• src/common/platform/shared/aicpu/args_dump_aicpu.cpp
• src/common/platform/shared/host/args_dump_collector.cpp
• tests/ut/cpp/a2a3/test_args_dump.cpp
• tests/ut/cpp/a5/test_args_dump.cpp
• tests/ut/cpp/common/test_args_dump_collector.cpp