Fix: decouple tensor-dump args mask pool from runtime ring depth#1320
Conversation
📝 WalkthroughWalkthroughTensor-dump selective metadata now maps task IDs using folded 64-bit hashing instead of ring validation. Header constants were updated, and new A2A3 runtime tests cover masks, flags, scalar dtypes, collisions, high task IDs, and unset entries. ChangesTensor-dump task slot selection
Estimated code review effort: 3 (Moderate) | ~20 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 decouples the tensor dump mask pool from the runtime ring depth by refactoring resolve_dump_args_task_slot to use a hash-folding mechanism on the 64-bit task_id instead of splitting it into ring and slot fields. It also removes the unused constants and the dependency on pto_runtime2_types.h, and adds a comprehensive unit test suite. The review feedback suggests a great improvement to make resolve_dump_args_task_slot return the calculated index directly as a uint32_t instead of using an out-parameter pointer, which simplifies the calling code and avoids potential null pointer dereferences.
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.
🧹 Nitpick comments (1)
tests/ut/cpp/a2a3/test_tensor_dump_mask_pool.cpp (1)
72-88: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a test for actual hash collisions.
The
DistinctTaskIdsDoNotAliascomment claims "open-addressed probing resolves any hash collision to separate slots," butid_a(hash 102) andid_b(hash 203) do not collide. A dedicated test with two task IDs that fold to the same index would verify the probing logic that the comment describes.For example,
(1<<32)|10and(3<<32)|8both producering_id ^ local_id == 11, guaranteeing a collision for any power-of-2 table capacity.♻️ Suggested collision test
// Two distinct task_ids must keep independent entries (open-addressed probing // resolves any hash collision to separate slots). TEST_F(TensorDumpMaskPoolTest, DistinctTaskIdsDoNotAlias) { const uint64_t id_a = (uint64_t{2} << 32) | 100u; const uint64_t id_b = (uint64_t{3} << 32) | 200u; set_dump_args_task_mask(id_a, 0b0001, TENSOR_DUMP_ARG_MASK_NONE); set_dump_args_task_mask(id_b, 0b1000, TENSOR_DUMP_ARG_MASK_NONE); TensorDumpArgMask mask_a = TENSOR_DUMP_ARG_MASK_NONE; TensorDumpArgMask mask_b = TENSOR_DUMP_ARG_MASK_NONE; get_dump_args_task_masks(id_a, &mask_a, nullptr); get_dump_args_task_masks(id_b, &mask_b, nullptr); EXPECT_EQ(mask_a, static_cast<TensorDumpArgMask>(0b0001)); EXPECT_EQ(mask_b, static_cast<TensorDumpArgMask>(0b1000)); } +// Two task_ids that fold to the same hash index must still occupy separate +// slots via open-addressed probing. +TEST_F(TensorDumpMaskPoolTest, CollidingTaskIdsDoNotAlias) { + // (1 << 32) | 10 -> hash = 1 ^ 10 = 11 + // (3 << 32) | 8 -> hash = 3 ^ 8 = 11 + const uint64_t id_a = (uint64_t{1} << 32) | 10u; + const uint64_t id_b = (uint64_t{3} << 32) | 8u; + + set_dump_args_task_mask(id_a, 0b0001, TENSOR_DUMP_ARG_MASK_NONE); + set_dump_args_task_mask(id_b, 0b1000, TENSOR_DUMP_ARG_MASK_NONE); + + TensorDumpArgMask mask_a = TENSOR_DUMP_ARG_MASK_NONE; + TensorDumpArgMask mask_b = TENSOR_DUMP_ARG_MASK_NONE; + get_dump_args_task_masks(id_a, &mask_a, nullptr); + get_dump_args_task_masks(id_b, &mask_b, nullptr); + + EXPECT_EQ(mask_a, static_cast<TensorDumpArgMask>(0b0001)); + EXPECT_EQ(mask_b, static_cast<TensorDumpArgMask>(0b1000)); +} + // An unset task_id reads back as "no mask" — the empty-slot path.🤖 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 `@tests/ut/cpp/a2a3/test_tensor_dump_mask_pool.cpp` around lines 72 - 88, Add a dedicated collision test near DistinctTaskIdsDoNotAlias using task IDs (uint64_t{1} << 32) | 10u and (uint64_t{3} << 32) | 8u, which fold to the same hash index for power-of-two capacities. Set distinct masks with set_dump_args_task_mask, retrieve both via get_dump_args_task_masks, and assert each remains independent; update the existing test comment to avoid claiming it validates collision handling.
🤖 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.
Nitpick comments:
In `@tests/ut/cpp/a2a3/test_tensor_dump_mask_pool.cpp`:
- Around line 72-88: Add a dedicated collision test near
DistinctTaskIdsDoNotAlias using task IDs (uint64_t{1} << 32) | 10u and
(uint64_t{3} << 32) | 8u, which fold to the same hash index for power-of-two
capacities. Set distinct masks with set_dump_args_task_mask, retrieve both via
get_dump_args_task_masks, and assert each remains independent; update the
existing test comment to avoid claiming it validates collision handling.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 362cfd2a-8c9c-406a-91a0-6664e443c3b4
📒 Files selected for processing (5)
src/a2a3/platform/include/common/tensor_dump.hsrc/a5/platform/include/common/tensor_dump.hsrc/common/platform/shared/aicpu/tensor_dump_aicpu.cpptests/ut/cpp/CMakeLists.txttests/ut/cpp/a2a3/test_tensor_dump_mask_pool.cpp
💤 Files with no reviewable changes (2)
- src/a2a3/platform/include/common/tensor_dump.h
- src/a5/platform/include/common/tensor_dump.h
yanghaoran29
changed the title
The args-dump mask pool keyed entries on the runtime's (ring_id, slot) partition of task_id: resolve_dump_args_task_slot extracted ring_id from the high 32 bits and returned false when ring_id >= TENSOR_DUMP_MASK_POOL_MAX_RINGS (= PTO2_MAX_RING_DEPTH), so any task_id whose ring field exceeded the runtime's max was silently dropped — its dump mask never entered the table and reads back TENSOR_DUMP_ARG_MASK_NONE. The pool inherited a runtime coupling it shouldn't have. Drop the PTO2_MAX_RING_DEPTH / PTO2_TASK_WINDOW_SIZE dependency from the a2a3 and a5 tensor_dump.h headers (removes the TENSOR_DUMP_MASK_POOL_MAX_RINGS / MAX_SLOTS / DEFAULT_SLOT_MASK constants and the pto_runtime2_types.h include) and refactor resolve_dump_args_task_slot to fold the full 64-bit task_id via high^low XOR into the table index. Every task_id now maps to a slot regardless of its ring field, and the function returns uint32_t directly instead of writing through an out-param (per review). Adds tests/ut/cpp/a2a3/test_tensor_dump.cpp and tests/ut/cpp/a5/test_tensor_dump.cpp (mirrored across both arches since the dump source is shared but each platform carries its own header/test target) covering: basic enable/disable and dump-base round-trip; high-ring task_id mask + scalar-dtype round-trip (the regression this PR fixes); distinct task_id independence under open-addressed probing; unknown-task-id reads-none; decision helpers in FULL mode (ArgDirection→role mapping, stage gating, should_dump_task/should_dump_arg, has_dump_arg_flag); end-to-end dump-record correctness (1-D tensor arena payload + scalar inline value); selective (PARTIAL) mode decision filtering; and non-contiguous 2-D tensor gather path verification.
2 participants
Problem
The args-dump mask pool keyed entries on the runtime's
(ring_id, slot)partition oftask_id:resolve_dump_args_task_slotextractedring_idfrom the high 32 bits and returned false whenring_id >= TENSOR_DUMP_MASK_POOL_MAX_RINGS(=PTO2_MAX_RING_DEPTH). Anytask_idwhose ring field exceeded the runtime's max was silently dropped — its dump mask never entered the table and reads backTENSOR_DUMP_ARG_MASK_NONE.The pool inherited a runtime coupling (
PTO2_MAX_RING_DEPTH/PTO2_TASK_WINDOW_SIZEviapto_runtime2_types.h) it shouldn't have: the dump mask pool is a host-side diagnostic table and must accept any task_id the runtime issues, regardless of how the runtime partitions the id into ring/slot.Fix
PTO2_MAX_RING_DEPTH/PTO2_TASK_WINDOW_SIZEdependency from thea2a3anda5tensor_dump.hheaders (removes theTENSOR_DUMP_MASK_POOL_MAX_RINGS/MAX_SLOTS/DEFAULT_SLOT_MASKconstants and thepto_runtime2_types.hinclude).resolve_dump_args_task_slotto fold the full 64-bittask_idviahigh ^ lowXOR into the table index. Everytask_idnow maps to a slot regardless of its ring field; the function returnsuint32_tdirectly instead of writing through an out-parameter (per review feedback).set_dump_args_task_scalar_dtypes,get_dump_args_task_scalar_dtypes,set_dump_args_task_mask,get_dump_args_task_masks) touint32_t idx = resolve_dump_args_task_slot(task_id);.Tests
Adds
tests/ut/cpp/a2a3/test_tensor_dump_mask_pool.cppandtests/ut/cpp/a5/test_tensor_dump_mask_pool.cpp(mirrored — the dump source is shared but each platform carries its own header/test target) covering:HighRingTaskIdRoundTrips— atask_idwhose high 32 bits exceed any runtime's ring depth still round-trips through both the mask table and the scalar-dtype table (the regression this PR fixes).DistinctTaskIdsDoNotAlias— open-addressed probing keeps two distinct ids independent.UnknownTaskIdReadsNone— unsettask_idreads backTENSOR_DUMP_ARG_MASK_NONE(empty-slot path).Changes
src/a2a3/platform/include/common/tensor_dump.h— removepto_runtime2_types.hinclude and the 3 derived constants.src/a5/platform/include/common/tensor_dump.h— same.src/common/platform/shared/aicpu/tensor_dump_aicpu.cpp— refactorresolve_dump_args_task_slot+ 4 callers.tests/ut/cpp/CMakeLists.txt— register the new a2a3 + a5 tests.tests/ut/cpp/a2a3/test_tensor_dump_mask_pool.cpp— new test file.tests/ut/cpp/a5/test_tensor_dump_mask_pool.cpp— new test file.close #1317