feat(a5): add URMA deferred completion backend for tensormap_and_ringbuffer - #1392
Merged
jvjhfhg merged 6 commits intoJul 20, 2026
Merged
Conversation
📝 WalkthroughWalkthroughThe change adds URMA communication workspace management, asynchronous TGET/TPUT submission, event-handle polling, backend-cookie propagation through completion queues, and architecture-specific host build configuration. ChangesURMA completion integration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant URMARequest
participant CompletionKernel
participant CompletionMailbox
participant AsyncWait
participant URMAScheduler
URMARequest->>CompletionKernel: submit TGET or TPUT
CompletionKernel->>CompletionMailbox: register event handle and workspace cookie
CompletionMailbox->>AsyncWait: drain completion condition
AsyncWait->>URMAScheduler: poll event handle
URMAScheduler-->>AsyncWait: PENDING, READY, or FAILED
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 introduces support for URMA (Ultra Remote Memory Access) workspaces and completions in the host runtime and tensormap/ringbuffer runtime. It adds URMA workspace management, a new URMA completion kernel and scheduler, and integrates a 'backend_cookie' field across completion mailbox messages and deferred entries to support URMA event handles. Additionally, the CMake configuration is updated to support newer CANN packages and architecture-specific paths. A bug was identified in the updated CMake library-finding logic where finding only 'libhccl' but not 'libhcomm' erroneously triggers a fatal error instead of falling back correctly.
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: 7
🤖 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/a5/platform/onboard/host/CMakeLists.txt`:
- Around line 187-192: Update the Ascend HCCL discovery conditional around
HCOMM_LIB to add a branch for HCCL_LIB when the preferred libhccl target is
unavailable. Set HCCL_LINK_TARGETS from HCCL_LIB and emit the appropriate
fallback warning, while retaining the existing HCOMM_LIB fallback and fatal
error only when neither library is found.
In `@src/a5/platform/onboard/host/comm_hccl.cpp`:
- Around line 1076-1085: Update the error paths in the domain allocation flow,
including the aclrtMalloc and aclrtMemcpy failure branches near
release_own_vmm_window, to call release_domain_peer_windows(*out) before
returning -1. Ensure peer-imported physical handles and virtual mappings are
released while preserving the existing cleanup operations.
- Around line 990-991: Update the domain workspace initialization near domain
context setup to use the explicitly configured sdma_workspace when SDMA is
enabled, rather than inheriting h->host_ctx.workSpace; ensure non-dense domains
receive no base URMA workspace while preserving the corresponding workspace size
behavior.
In
`@src/a5/runtime/tensormap_and_ringbuffer/runtime/backend/urma/urma_completion_kernel.h`:
- Around line 202-212: Update submit_chunked_urma_request to validate that
desc.src and desc.dst have equal logical element counts before calculating
total_bytes or submitting any transfer. On mismatch, defer
PTO2_ERROR_ASYNC_COMPLETION_INVALID through ctx and return false; preserve the
existing contiguous 1D validation and transfer flow for matching lengths.
- Around line 191-193: Update submit_urma_request_once() to propagate a false
result from register_urma_async_event() instead of unconditionally returning
true. On registration failure, drain the event before returning false so the DMA
is not left untracked and the chunked submission loop stops; preserve
defer_flush() and the successful return path for successful registration.
In
`@src/a5/runtime/tensormap_and_ringbuffer/runtime/backend/urma/urma_completion_scheduler.h`:
- Around line 192-201: Update the CQE processing loop so a failed completion is
consumed before returning failure: advance next_tail past the failed CQE and
call update_tail_info with the new tail before returning
CompletionPollState::FAILED. Preserve the existing failure status and
successful-CQE handling.
- Around line 141-159: Update the cache invalidation in the completion-context
setup around UrmaCqCtx and UrmaWqCtx so it covers every cache line spanned by
each context, not only the line containing its first byte. Use the actual
sizeof(UrmaCqCtx) and sizeof(UrmaWqCtx) ranges when calling
cache_invalidate_range, while preserving the existing cache_line alignment and
pointer calculations.
🪄 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: b6541b26-b476-410a-a953-92cab3fbcab4
📒 Files selected for processing (9)
src/a5/platform/onboard/host/CMakeLists.txtsrc/a5/platform/onboard/host/comm_hccl.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox_types.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/backend/urma/urma_completion_kernel.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/backend/urma/urma_completion_scheduler.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
Comment on lines
+202
to
+212
| inline __aicore__ bool submit_chunked_urma_request(AsyncCtx &ctx, UrmaRequestDescriptor<DstTensor, SrcTensor> desc) { | ||
| using RawDType = typename DstTensor::RawDType; | ||
| static_assert(std::is_same_v<RawDType, typename SrcTensor::RawDType>, "URMA transfer requires matching dtypes"); | ||
|
|
||
| if (!is_flat_contiguous_1d(desc.dst) || !is_flat_contiguous_1d(desc.src)) { | ||
| pto2::detail::defer_error(ctx, PTO2_ERROR_ASYNC_COMPLETION_INVALID); | ||
| return false; | ||
| } | ||
|
|
||
| const uint64_t elem_count = tensor_element_count(desc.dst); | ||
| const uint64_t total_bytes = elem_count * sizeof(RawDType); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Reject mismatched source and destination lengths.
Only desc.dst determines the transfer and slice length. A shorter source can therefore produce slices beyond its logical extent, while a longer source is silently truncated.
Proposed fix
- const uint64_t elem_count = tensor_element_count(desc.dst);
+ const uint64_t elem_count = tensor_element_count(desc.dst);
+ if (elem_count != tensor_element_count(desc.src)) {
+ pto2::detail::defer_error(ctx, PTO2_ERROR_ASYNC_COMPLETION_INVALID);
+ return false;
+ }
const uint64_t total_bytes = elem_count * sizeof(RawDType);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| inline __aicore__ bool submit_chunked_urma_request(AsyncCtx &ctx, UrmaRequestDescriptor<DstTensor, SrcTensor> desc) { | |
| using RawDType = typename DstTensor::RawDType; | |
| static_assert(std::is_same_v<RawDType, typename SrcTensor::RawDType>, "URMA transfer requires matching dtypes"); | |
| if (!is_flat_contiguous_1d(desc.dst) || !is_flat_contiguous_1d(desc.src)) { | |
| pto2::detail::defer_error(ctx, PTO2_ERROR_ASYNC_COMPLETION_INVALID); | |
| return false; | |
| } | |
| const uint64_t elem_count = tensor_element_count(desc.dst); | |
| const uint64_t total_bytes = elem_count * sizeof(RawDType); | |
| inline __aicore__ bool submit_chunked_urma_request(AsyncCtx &ctx, UrmaRequestDescriptor<DstTensor, SrcTensor> desc) { | |
| using RawDType = typename DstTensor::RawDType; | |
| static_assert(std::is_same_v<RawDType, typename SrcTensor::RawDType>, "URMA transfer requires matching dtypes"); | |
| if (!is_flat_contiguous_1d(desc.dst) || !is_flat_contiguous_1d(desc.src)) { | |
| pto2::detail::defer_error(ctx, PTO2_ERROR_ASYNC_COMPLETION_INVALID); | |
| return false; | |
| } | |
| const uint64_t elem_count = tensor_element_count(desc.dst); | |
| if (elem_count != tensor_element_count(desc.src)) { | |
| pto2::detail::defer_error(ctx, PTO2_ERROR_ASYNC_COMPLETION_INVALID); | |
| return false; | |
| } | |
| const uint64_t total_bytes = elem_count * sizeof(RawDType); |
🤖 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/a5/runtime/tensormap_and_ringbuffer/runtime/backend/urma/urma_completion_kernel.h`
around lines 202 - 212, Update submit_chunked_urma_request to validate that
desc.src and desc.dst have equal logical element counts before calculating
total_bytes or submitting any transfer. On mismatch, defer
PTO2_ERROR_ASYNC_COMPLETION_INVALID through ctx and return false; preserve the
existing contiguous 1D validation and transfer flow for matching lengths.
Comment on lines
+192
to
+201
| const uint8_t substatus = static_cast<uint8_t>((dw0 >> 16) & 0xFFu); | ||
| const uint8_t status = static_cast<uint8_t>((dw0 >> 24) & 0xFFu); | ||
| if (status != 0 || substatus != 0) { | ||
| return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID}; | ||
| } | ||
| next_tail++; | ||
| } | ||
|
|
||
| if (next_tail != cur_tail) { | ||
| update_tail_info(cq_ctx, wq_ctx, next_tail); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Consume failed CQEs before returning failure.
The early return leaves cur_tail on the same failed CQE, causing every later poll on that CQ to encounter it again and preventing queue retirement.
Proposed fix
const uint8_t substatus = static_cast<uint8_t>((dw0 >> 16) & 0xFFu);
const uint8_t status = static_cast<uint8_t>((dw0 >> 24) & 0xFFu);
if (status != 0 || substatus != 0) {
+ ++next_tail;
+ update_tail_info(cq_ctx, wq_ctx, next_tail);
return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID};
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| const uint8_t substatus = static_cast<uint8_t>((dw0 >> 16) & 0xFFu); | |
| const uint8_t status = static_cast<uint8_t>((dw0 >> 24) & 0xFFu); | |
| if (status != 0 || substatus != 0) { | |
| return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID}; | |
| } | |
| next_tail++; | |
| } | |
| if (next_tail != cur_tail) { | |
| update_tail_info(cq_ctx, wq_ctx, next_tail); | |
| const uint8_t substatus = static_cast<uint8_t>((dw0 >> 16) & 0xFFu); | |
| const uint8_t status = static_cast<uint8_t>((dw0 >> 24) & 0xFFu); | |
| if (status != 0 || substatus != 0) { | |
| +next_tail; | |
| update_tail_info(cq_ctx, wq_ctx, next_tail); | |
| return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID}; | |
| } |
🤖 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/a5/runtime/tensormap_and_ringbuffer/runtime/backend/urma/urma_completion_scheduler.h`
around lines 192 - 201, Update the CQE processing loop so a failed completion is
consumed before returning failure: advance next_tail past the failed CQE and
call update_tail_info with the new tail before returning
CompletionPollState::FAILED. Preserve the existing failure status and
successful-CQE handling.
Collaborator
|
Requesting changes: the host-side URMA workspace needs the same opt-in macro gating as the SDMA overlay — the a5 CI environment is not ready for it (see #1315). The reason for gating is environmental, not code quality: as documented in #1315, the current st-onboard-a5 CI CANN drops do not expose working async primitives (CANN 9.1.T500:
Asks:
|
10 tasks
jvjhfhg
changed the title
jvjhfhg
approved these changes
Jul 20, 2026
jvjhfhg
pushed a commit
that referenced
this pull request
Jul 22, 2026
Add a minimal two-rank smoke test under examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo, mirroring the SDMA async completion demo but driving the URMA deferred completion backend from #1392: - Each rank stages its input inside the HCCL/URMA communication domain (256 B reserved head + input window). - The producer issues a deferred UrmaTget from the peer rank's input window into local out, resolving the remote address through pto2::urma_backend::peer_mr_ptr on the URMA workspace. - The consumer depends on out and computes result = out + 1; host-side validation checks both against the peer input, exercising URMA completion polling and deferred-release dependency handling. Isolation per #1315: the demo is examples-only, self-skips unless SIMPLER_ENABLE_PTO_URMA_WORKSPACE=ON (env-var gate at collection, RuntimeError in run(), workSpace==0 self-skip in the kernel), and stays out of CI until the URMA overlay is enabled for the a5 environment.
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
2 participants
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
Testing
git diff --check