From c3bbbd1e7084b630c4cb9797e798991f5941f0d7 Mon Sep 17 00:00:00 2001 From: majin0824 Date: Wed, 26 Aug 2026 20:46:46 -0700 Subject: [PATCH] Feature: enable A5 SDMA and URMA together Provision SDMA and URMA workspaces together in the default A5 runtime, with transport-specific CommContext fields and a communicator-lifetime URMA arena for derived domain slices. Map reordered domain ranks back to communicator ranks, return arena slices only after backend teardown, and keep successful or partially initialized URMA state alive until HCCL destruction. Provision both transports best-effort: a workspace that fails to initialize leaves its CommContext pair zero for the kernels to self-skip, and a domain window is never denied over a missing transport. URMA agrees that outcome across ranks through the base_urma_ready rendezvous, which now carries each rank's result, so no rank issues an RDMA against a peer holding no registration. Size the arena from Worker(comm_arena_size=...), defaulting to the backend's 200 MiB, since every live domain is a slice of it. Drop the per-domain Path-D allocator and its announce helpers, which the arena path leaves unreachable. Reject a non-zero window_offset on the backends whose windows are owned by the allocation, and make their derived-window range checks overflow-safe like a5's. Cover dual-workspace contexts, non-zero derived offsets, rank remapping, arena allocation and reuse, arena sizing, and HCCL-to-URMA teardown ordering. --- docs/capability-survey.md | 61 +- docs/comm-domain.md | 58 +- docs/getting-started.md | 4 +- docs/tensormap-and-ringbuffer-a2a3-vs-a5.md | 22 +- docs/user/reference/python-api.md | 1 + .../a5/tensormap_and_ringbuffer/README.md | 2 +- .../sdma_async_completion_demo/README.md | 8 +- .../kernels/aiv/kernel_sdma_tget_async.cpp | 4 +- .../test_sdma_async_completion_demo.py | 21 - .../urma_deferred_completion_demo/README.md | 76 +-- .../kernels/aiv/kernel_urma_tget_async.cpp | 12 +- .../test_urma_deferred_completion_demo.py | 29 +- python/bindings/task_interface.cpp | 9 +- python/simpler/task_interface.py | 13 +- python/simpler/worker.py | 172 ++++- simpler_setup/runtime_builder.py | 9 +- simpler_setup/runtime_compiler.py | 11 +- src/a2a3/platform/onboard/host/comm_hccl.cpp | 24 +- src/a5/platform/onboard/host/CMakeLists.txt | 20 +- src/a5/platform/onboard/host/comm_hccl.cpp | 615 +++++------------- src/common/platform_comm/comm.h | 40 +- src/common/platform_comm/comm_context.h | 44 +- src/common/platform_comm/comm_sim.cpp | 28 +- src/common/worker/chip_worker.cpp | 4 +- src/common/worker/chip_worker.h | 4 +- tests/ut/py/test_runtime_builder.py | 20 +- .../py/test_worker/test_dynamic_alloc_hw.py | 27 +- tests/ut/py/test_worker/test_host_worker.py | 196 +++++- tests/ut/py/test_worker/test_platform_comm.py | 54 +- 29 files changed, 871 insertions(+), 717 deletions(-) diff --git a/docs/capability-survey.md b/docs/capability-survey.md index 6c13f2fef0..b45f02e8d9 100644 --- a/docs/capability-survey.md +++ b/docs/capability-survey.md @@ -220,8 +220,8 @@ values yield `SIMPLER_ERROR_ASYNC_COMPLETION_INVALID`. | Engine | a2a3 | a5 | Status | | ------ | ---- | -- | ------ | | COUNTER (default) | registered | registered | **Shipped** — `tests/st/worker/comm_domain/async_notify` runs onboard on both architectures; `tests/st/worker/comm_domain/deferred_notify` runs in sim on both and onboard on a2a3, through the `st-onboard-*` / `st-sim-*` jobs in `ci.yml`. Routed by `CASES[*]["platforms"]`, no `skipif` | -| SDMA | build macro forced ON; runtime opt-in | `option(... OFF)` | a2a3 **Shipped** (the "SDMA pytest (a2a3)" step in `ci.yml`); a5 not built | -| URMA | absent | full implementation | **Gated** — see below | +| SDMA | build macro forced ON; runtime opt-in | built and provisioned with communication contexts | a2a3 **Shipped** (the "SDMA pytest (a2a3)" step in `ci.yml`); a5 demo runs in the ordinary A5 sweep | +| URMA | absent | built and provisioned with communication contexts, best-effort | A5-only; exercised by `urma_deferred_completion_demo` without a build or environment gate. A rank set where registration fails runs SDMA-only, and that demo fails its golden rather than skipping — so a green run is evidence URMA moved data | | ROCE, CCU | enum only | enum only | **Name only** | **a2a3 SDMA is opt-in at runtime**, not "always on": the provider is always @@ -235,23 +235,33 @@ without, traced to a single 300,000 ms remote TRS event timeout ([investigations/2026-07-a2a3-sdma-fault-teardown.md](investigations/2026-07-a2a3-sdma-fault-teardown.md), issue #1425). -**a5 URMA is real code that cannot execute.** The scheduler walks `UrmaCqCtx` +**a5 SDMA and URMA coexist in one communication context.** The scheduler walks `UrmaCqCtx` CQEs checking the owner bit, advances the tail and rings the doorbell (`src/a5/.../backend/urma/urma_completion_scheduler.h:133-215`); the kernel -submits `TGET_ASYNC`/`TPUT_ASYNC` with 256 MB chunking. Both -sit behind `PTO_URMA_SUPPORTED`, which is **defined nowhere in this repo and -nowhere in the installed CANN pto headers**, so the `#else` branch returns -`SIMPLER_ERROR_ASYNC_COMPLETION_INVALID` immediately. The host overlay macro is -fully wired, so turning it on does not help — the device path stays unreachable. -a5's SDMA and URMA overlays are mutually exclusive by CMake `FATAL_ERROR` -because `CommContext` exposes a single `workSpace` pair -(`src/a5/platform/onboard/host/CMakeLists.txt:49-53`). - -**HCCL is bootstrap, not data movement.** The complete set of functions called -is `HcclGetRootInfo`, `HcclCommInitRootInfo`, `HcclBarrier`, `HcclCommDestroy`. -There is no `HcclAllReduce` / `AllGather` / `Send` / `Recv` anywhere; every -shipped collective is a hand-written AIV kernel that computes a peer pointer -from the symmetric window +submits `TGET_ASYNC`/`TPUT_ASYNC` with 256 MB chunking. The +pinned PTO-ISA defines `PTO_URMA_SUPPORTED` for DAV_3510. The host provisions +the process-global SDMA workspace and communicator-scoped URMA workspace before +uploading `CommContext`; the leading `sdmaWorkSpace` pair carries SDMA, +and the appended `urmaWorkSpace` pair carries URMA. A derived context also +carries `urmaWindowOffset`, translating its domain-local window offsets back +to the registered base MR. URMA metadata is indexed by communicator rank, and +derived contexts map domain-local ranks onto those communicator ranks. Dynamic +domains are slices of the registered arena, so sequential and concurrent +domains do not create extra HCCL memory registrations or channels. Engine-specific kernels run +from the same default build without an environment selector. Both workspaces +are best-effort: a transport that fails to initialize leaves its `CommContext` +pair zero and its kernels self-skip, and the domain is handed out regardless. +URMA additionally agrees that outcome across ranks through the +`base_urma_ready` rendezvous, so no rank issues an RDMA against a peer that +holds no registration. + +**HCCL is control-plane setup, not collective data movement.** Communicator +lifecycle uses `HcclGetRootInfo`, `HcclCommInitRootInfo`, `HcclBarrier`, and +`HcclCommDestroy`; A5 URMA setup additionally uses `HcclCommMemReg`, +`HcclRankGraphGetLinks`, and `HcclChannelAcquire` to register memory and create +transport channels. There is no `HcclAllReduce` / `AllGather` / `Send` / `Recv` +in the collective data path; every shipped collective is a hand-written AIV +kernel that computes a peer pointer from the symmetric window (`tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ring_kernel.cpp:59-61`). **Fabric is undocumented.** a2a3 prefers `alloc_windows_via_fabric()` with @@ -261,10 +271,10 @@ only when Fabric is unsupported occurrences of "fabric". No `.md` in the repo describes Fabric as a memory-sharing mechanism. -One ABI fact constrains portability: the AICore→AICPU completion struct -diverged — `DeferredCompletionEntry` is 24 bytes on a2a3 and 32 on a5, the extra -8 being the `backend_cookie` URMA's poll needs. A URMA-on-a2a3 port must widen a -struct that is currently frozen by `static_assert`. +The AICore→AICPU completion ABI is currently aligned: +`DeferredCompletionEntry` is 32 bytes on both a2a3 and a5, and both carry +`backend_cookie`. A2/A3 does not currently register or implement the URMA +completion backend despite sharing this ABI field. ## Open questions @@ -273,8 +283,9 @@ Unresolved after this survey, in rough order of how much they block: 1. **What are `ASYNC_ENGINE_ROCE` and `ASYNC_ENGINE_CCU` for?** No design doc, investigation entry, or code comment says whether they are reserved slots or leftovers from a dropped design. -2. **Has a5 URMA ever run on silicon?** No CI run, test artifact, or - investigation attests to it. +2. **How broad is a5 URMA coverage?** The two-rank deferred-completion demo has + passed bring-up on A5 and is selected by the ordinary A5 sweep after this + change, but there is no retained multi-topology or long-running artifact. 3. **Which CANN mitigation closed issue #822, and is Path B usable on CANN 9.0.0?** The doc says "CANN-side mitigation landed" without naming it, and nobody re-ran the repro. @@ -282,9 +293,7 @@ Unresolved after this survey, in rough order of how much they block: one-process-per-`(arch, runtime)` ChipWorker model rests on it, but it is asserted only from CANN source paths that are not vendored here and no in-repo probe detects it. -5. **Where would `PTO_URMA_SUPPORTED` ever be defined?** Not in this repo, not - in the installed CANN pto headers. -6. **Which platforms lack Fabric support?** Stated only as a code comment, never +5. **Which platforms lack Fabric support?** Stated only as a code comment, never enumerated, and a5's divergence to the V1 handle route is undocumented. ## Documentation drift found while compiling this survey diff --git a/docs/comm-domain.md b/docs/comm-domain.md index b7db6f67c4..ec56c37391 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -213,21 +213,59 @@ symmetric window is realized: | Aspect | Sim | HCCL (onboard) | | ------ | --- | -------------- | -| Window memory | POSIX shm + `ftruncate`, mmap'd per rank | a2a3: Fabric V2 handle exchange (`ACL_MEM_SHARE_HANDLE_TYPE_FABRIC`), falling back to VMM + shareable-handle IPC where Fabric is unsupported. a5: VMM shareable handles only. Cross-card P2P via `aclrtDeviceEnablePeerAccess` on both | +| Window memory | POSIX shm + `ftruncate`, mmap'd per rank | a2a3: Fabric V2 handle exchange (`ACL_MEM_SHARE_HANDLE_TYPE_FABRIC`), falling back to VMM + shareable-handle IPC where Fabric is unsupported. a5: slices of one per-rank VMM arena registered for the communicator lifetime, 200 MiB unless `Worker(comm_arena_size=…)` says otherwise. Cross-card P2P via `aclrtDeviceEnablePeerAccess` on both | | Subset barrier | shm-header atomic, `allocation_id`-scoped | file barriers, `allocation_id`-scoped | -| Window init | window zeroed before the subset barrier (`memset`) | window zeroed before the handle is announced (`aclrtMemset`) | -| Async-DMA workspace | n/a | a2a3: opt-in per Worker (`enable_sdma`); a5: SDMA by default, URMA as an opt-in alternative | +| Window init | window zeroed before the subset barrier (`memset`) | a2a3: window zeroed before the handle is announced. a5: the base arena is registered and mapped during communicator init; every newly assigned or reused slice is zeroed at domain allocation before that domain is returned to the caller | +| Async-DMA workspace | n/a | a2a3: opt-in per Worker (`enable_sdma`); a5: communicator-scoped SDMA and URMA workspaces are inherited by derived contexts, whose rank map supports arbitrary subsets/reorderings. Either may be absent — see below | The window is zero-initialized on both backends so scratch/signal protocols see a known starting state (matching the historical static-path contract). -The wipe happens before the window becomes reachable by any peer — before the -shareable handle is announced on HCCL, before the `ready_count` barrier on sim. -A peer that clears the subset barrier can return, launch its kernel and store a -barrier signal into this rank's window immediately; a wipe issued after that -point can erase a signal the owner has not yet waited on, and the owner then -waits on it forever. The rank skew that opens that window grows with host load, -so the resulting hang shows up only under a loaded box. +### A5: arena size and transport availability + +The A5 arena is allocated once per communicator and every domain is a slice of +it, so **`comm_arena_size` is the ceiling on all live domains combined**, not a +per-domain limit. A request that no free run can satisfy raises `MemoryError` +rather than falling back to a private allocation. The default is 200 MiB; pass +`Worker(comm_arena_size=)` to change it. The backend rounds the request +up to VMM granularity and reports the granted size back, which is what the +slice allocator uses. + +Neither async transport is required for a domain. SDMA and URMA are each +provisioned best-effort at `comm_alloc_windows`, and a transport that fails to +initialize simply leaves its `CommContext` pair (`sdmaWorkSpace` / +`urmaWorkSpace`) zero; the kernels guard on that and self-skip. A communicator +that came up SDMA-only still hands out windows normally. + +A self-skipping kernel moves no data, so this degrades availability of a +transport, never correctness of a workload: anything that depended on the +transfer sees unwritten output and fails its own verification. The a5 +`sdma_async_completion_demo` and `urma_deferred_completion_demo` both fail +their goldens on a skipped transfer rather than passing vacuously, which is +what makes a green run of either one evidence that its engine ran. + +URMA's outcome is agreed across ranks rather than decided per rank: the +`base_urma_ready` rendezvous carries each rank's result and every rank takes +the AND. A rank whose own registration succeeded while a peer's failed +therefore publishes a zero `urmaWorkSpace` too, because issuing an RDMA against +a peer that holds no registration is worse than not issuing one. That rank +still keeps its workspace manager alive, since `Finalize` must run after +`HcclCommDestroy` either way. + +On sim and a2a3, the wipe happens before the window becomes reachable by any +peer — before the `ready_count` barrier on sim and before the shareable handle +is announced on a2a3 HCCL. A5 is different: its base arena is already +registered and peer-mapped when each rank clears a slice, both on the slice's +first assignment and after reuse. The base arena itself is not published as a +domain. Safety comes from the parent waiting for every member's allocation RPC +(including that clear) to complete before it publishes the +`CommDomainHandle`; no task can access the slice during either wipe. + +Once a domain is published, a peer can launch its kernel and store a barrier +signal into this rank's window immediately. A wipe issued after that point can +erase a signal the owner has not yet waited on, and the owner then waits on it +forever. The rank skew that opens that window grows with host load, so the +resulting hang shows up only under a loaded box. On a2a3, async-DMA resources are a Worker-level opt-in, not a communication-domain property. Construct the Worker with `enable_sdma=True` and diff --git a/docs/getting-started.md b/docs/getting-started.md index 88d3b6305e..fc6f333462 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -55,8 +55,8 @@ To use a different PTO-ISA revision, update `pto_isa.pin` to the desired diff and applies the same revision to install-time runtime builds and run-time kernel compilation. -For platforms that embed PTO-ISA headers into onboard host runtimes (a2a3 -always; a5 when an async workspace overlay is ON), builds record the actual +For platforms that embed PTO-ISA headers into onboard host runtimes (a2a3 and +a5), builds record the actual PTO-ISA git HEAD used for each runtime in `build/lib/pto_isa_build.json`. This JSON is artifact provenance, not a second configuration source. Lookup of those runtimes **requires** the metadata file: if it is missing, or if it diff --git a/docs/tensormap-and-ringbuffer-a2a3-vs-a5.md b/docs/tensormap-and-ringbuffer-a2a3-vs-a5.md index 53012416c9..67fa0d4937 100644 --- a/docs/tensormap-and-ringbuffer-a2a3-vs-a5.md +++ b/docs/tensormap-and-ringbuffer-a2a3-vs-a5.md @@ -109,7 +109,7 @@ The functional differences group into the following themes: | Cache coherence | Hardware coherence model | Yes | Retain the required invalidate/flush operations on A2/A3; do not copy unnecessary maintenance operations to A5 | | PMU collection | Hardware PMU and platform collection protocol | Yes | Retain the different counter counts, readers, and FIN submission paths | | System counter and DMB | Hardware timing and register layout | Yes | Use the constants for each platform | -| URMA completion | A5-specific implementation and product capability gate | Yes, for now | Retain the A5 path; do not claim that URMA is available in the default build | +| URMA completion | A5-specific implementation and product capability | Yes, for now | Retain the A5 path; the default onboard build provisions it together with SDMA | | Next-block prefetch | A2/A3-only performance optimization | No | Retain on A2/A3; validate on A5 before considering a port | | Scheduler progress publication | AICPU topology and measured publication cost | No | Retain A5's 16-task batching; keep per-advance publication on A2/A3, where the portable implementation showed no significant benefit | | Fatal teardown | Software reliability strategy | No | Retain the current implementations; decide whether to converge after measuring the worst-case A5 teardown time | @@ -184,18 +184,17 @@ platform's `platform/include/common/platform_config.h`. | `runtime/scheduler/scheduler_completion.cpp` | After FIN, A2/A3 invokes the AICPU MMIO reader for eight counters; A5 commits the ten-counter slot written by the AICore | | `platform/shared/aicpu/pmu_collector_aicpu.cpp` | Implements the A2/A3 direct MMIO read and the A5 staging-slot consumption paths | -### Optional A5-Specific URMA Backend +### A5-Specific URMA Backend A5 contains the source path for issuing URMA completion requests, creating deferred entries, forwarding FIN, and polling/retiring CQ entries. A2/A3 currently registers only the COUNTER and SDMA completion backends. -The repository does not currently define `PTO_URMA_SUPPORTED`. A5 therefore -compiles the shared ABI, mailbox, CQ polling/retirement, and related paths, but -the kernel path that successfully issues URMA PTO instructions is unreachable. -The current state is "implemented but disabled by default." It neither means -that the default A5 build supports URMA nor proves that the A2/A3 hardware does -not support URMA. +The pinned PTO-ISA defines `PTO_URMA_SUPPORTED` for DAV_3510. A5 compiles the +request path together with the shared ABI, mailbox, and CQ polling/retirement +paths. Its host communication context provisions SDMA and URMA workspaces +together, so the default build supports both engines without a selector. This +does not imply that A2/A3 hardware supports URMA. Both platforms already share `CompletionToken::backend_cookie`, `ASYNC_ENGINE_URMA`, the 32-byte `DeferredCompletionEntry`, end-to-end cookie @@ -341,7 +340,6 @@ comparison. Examples and tests unique to either platform primarily reflect chip-feature validation and test-porting progress; they cannot be used to infer whether the runtime supports a shared algorithm. -For example, A5 has `urma_deferred_completion_demo`, but this does not mean -that the current build defines `PTO_URMA_SUPPORTED`. Likewise, the absence of a -workload on one platform does not automatically mean that the corresponding -runtime capability is unavailable. +For example, A5's `urma_deferred_completion_demo` exercises an A5-only +capability. The absence of a corresponding workload on another platform does +not automatically establish a hardware limitation there. diff --git a/docs/user/reference/python-api.md b/docs/user/reference/python-api.md index b162eda7ad..5618319de1 100644 --- a/docs/user/reference/python-api.md +++ b/docs/user/reference/python-api.md @@ -39,6 +39,7 @@ into `**config` and validated later. The recognized keys: | `device_ids` | L3+ | one chip child process per entry | | `num_sub_workers` | L3+ | host-side Python callables to fork | | `enable_sdma` | a2a3 | provisions the SDMA workspace; defaults to `False` | +| `comm_arena_size` | a5 | bytes of the communicator-lifetime arena every communication domain is sliced from; `0` (default) takes the backend's 200 MiB. This is the ceiling on all live domains combined | | `heap_ring_size` | all | heap ring sizing | | `remote_heap_ring_size`, `remote_session_timeout_s` | L4 | remote-session sizing and timeout | diff --git a/examples/a5/tensormap_and_ringbuffer/README.md b/examples/a5/tensormap_and_ringbuffer/README.md index 0de941d811..cfa12b23ea 100644 --- a/examples/a5/tensormap_and_ringbuffer/README.md +++ b/examples/a5/tensormap_and_ringbuffer/README.md @@ -38,7 +38,7 @@ direct `Worker` communication-domain walkthrough from construction through | Example | Mechanism | | ------- | --------- | | [`sdma_async_completion_demo/`](sdma_async_completion_demo/) | `TGET_ASYNC` from a peer's window slot over SDMA, completion registered via `defer_pto_async_event`. Enabled by default on a5 onboard. | -| [`urma_deferred_completion_demo/`](urma_deferred_completion_demo/) | The same protocol over **URMA** — `kernel_consumer.cpp` is byte-identical to the SDMA demo's, so the transport is the only variable. The two overlays are **mutually exclusive in one build**, so comparing them means rebuilding. | +| [`urma_deferred_completion_demo/`](urma_deferred_completion_demo/) | The same protocol over **URMA** — `kernel_consumer.cpp` is byte-identical to the SDMA demo's, so the transport is the only variable. It runs from the same default build as the SDMA demo. | The cross-architecture notification-counter and deferred-notify watchdogs live under [`tests/st/worker/comm_domain/`](../../../tests/st/worker/comm_domain/). diff --git a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/README.md b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/README.md index b5d1beb46c..e7f2f1bc4f 100644 --- a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/README.md +++ b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/README.md @@ -20,16 +20,12 @@ Every rank's window is laid out identically, so an offset is rank-independent. ## Requirements -The a5 host runtime includes the async-SDMA workspace by default: +The A5 host runtime includes both async-SDMA and URMA workspaces by default: | Gate | Effect | | ---- | ------ | | `CASES[*]["platforms"] = ["a5"]` | deselected on any other `--platform` | | `CASES[*]["config"]["device_count"] = 2` | needs two dies | -| `@pytest.mark.skipif(_urma_workspace_enabled())` | skipped when `SIMPLER_ENABLE_PTO_URMA_WORKSPACE` selects the URMA backend | - -URMA replaces SDMA in a URMA build. Rebuild without -`SIMPLER_ENABLE_PTO_URMA_WORKSPACE` before running this demo. ```bash pytest examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo \ @@ -40,5 +36,5 @@ Wrap the hardware run in `task-submit` on a shared box. ## Compare with -- [`../urma_deferred_completion_demo/`](../urma_deferred_completion_demo/) — the same protocol over URMA. `kernel_consumer.cpp` is byte-identical; only the transfer kernel, its completion header, and the build flag differ. **The two overlays are mutually exclusive in one build**, so comparing them means rebuilding — that README has the detail. +- [`../urma_deferred_completion_demo/`](../urma_deferred_completion_demo/) — the same protocol over URMA. `kernel_consumer.cpp` is byte-identical; only the transfer kernel, its completion header, and its workspace field differ. Both demos run from the same build without environment changes. - [`examples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/`](../../../a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/) — the a2a3 port of this demo, which needs no overlay flag. diff --git a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/kernels/aiv/kernel_sdma_tget_async.cpp b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/kernels/aiv/kernel_sdma_tget_async.cpp index c3d82b59ff..8704156206 100644 --- a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/kernels/aiv/kernel_sdma_tget_async.cpp +++ b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/kernels/aiv/kernel_sdma_tget_async.cpp @@ -45,7 +45,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in int rank = static_cast(comm_ctx->rankId); int nranks = static_cast(comm_ctx->rankNum); // A null workspace indicates that host-side SDMA provisioning failed. - if (nranks != 2 || comm_ctx->workSpace == 0) { + if (nranks != 2 || comm_ctx->sdmaWorkSpace == 0) { pipe_barrier(PIPE_ALL); return; } @@ -67,6 +67,6 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in AsyncCtx async_ctx = get_async_ctx(args); send_request_entry( async_ctx, - SdmaTget(local_global, remote_global, scratch_tile, reinterpret_cast<__gm__ uint8_t *>(comm_ctx->workSpace)) + SdmaTget(local_global, remote_global, scratch_tile, reinterpret_cast<__gm__ uint8_t *>(comm_ctx->sdmaWorkSpace)) ); } diff --git a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py index 1c07a71725..1b36afaf38 100644 --- a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py +++ b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py @@ -15,8 +15,6 @@ ``result = out + 1``. """ -import os - import pytest import torch from simpler.task_interface import ArgDirection as D @@ -28,24 +26,9 @@ N = 128 * 128 NRANKS = 2 DTYPE_NBYTES = 4 -_URMA_WORKSPACE_ENV = "SIMPLER_ENABLE_PTO_URMA_WORKSPACE" -_WORKSPACE_TRUTHY = {"1", "ON", "TRUE", "YES"} - - -def _urma_workspace_enabled(): - return os.environ.get(_URMA_WORKSPACE_ENV, "").upper() in _WORKSPACE_TRUTHY - - -def _require_sdma_workspace(): - if _urma_workspace_enabled(): - raise RuntimeError( - "sdma_async_completion_demo requires the default SDMA backend; " - f"unset {_URMA_WORKSPACE_ENV} and rebuild simpler." - ) def sdma_async_completion_orch_fn(orch, callables, task_args, config): - _require_sdma_workspace() input_nbytes = N * DTYPE_NBYTES with orch.allocate_domain( name="default", @@ -71,10 +54,6 @@ def sdma_async_completion_orch_fn(orch, callables, task_args, config): @pytest.mark.sdma(worker_workspace=False) -@pytest.mark.skipif( - _urma_workspace_enabled(), - reason="SDMA demo requires the default SDMA backend; unset SIMPLER_ENABLE_PTO_URMA_WORKSPACE and rebuild simpler.", -) @scene_test(level=3, runtime="tensormap_and_ringbuffer") class TestSdmaAsyncCompletionDemo(SceneTestCase): CALLABLE = { diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/README.md b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/README.md index c4ceef269b..b380fa14a6 100644 --- a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/README.md +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/README.md @@ -24,65 +24,59 @@ directories. Only the transfer kernel and the orchestration change: | ---- | ---- | ---- | | Transfer kernel | `kernel_urma_tget_async.cpp` | `kernel_sdma_tget_async.cpp` | | Completion header | `backend/urma/urma_completion_kernel.h` | `backend/sdma/sdma_completion_kernel.h` | -| Build selection | `SIMPLER_ENABLE_PTO_URMA_WORKSPACE=ON` | default | +| Workspace field | `urmaWorkSpace` | `sdmaWorkSpace` | Read them side by side and the transport is the only variable — which is exactly what you want when deciding which one a workload should use. -## SDMA is the default; URMA is the opt-in alternative - -This is the part to know before trying to run either. - -```cmake -option(SIMPLER_ENABLE_PTO_URMA_WORKSPACE "..." OFF) -if(SIMPLER_ENABLE_PTO_URMA_WORKSPACE) - set(SIMPLER_ENABLE_PTO_SDMA_WORKSPACE OFF) -else() - set(SIMPLER_ENABLE_PTO_SDMA_WORKSPACE ON) -endif() -``` - -— `src/a5/platform/onboard/host/CMakeLists.txt`. **One host runtime build can -carry one overlay, not both**, because `CommContext` has a single -`workSpace` / `workSpaceSize` pair to hand the kernel. So the two demos cannot -pass in the same build; comparing them means rebuilding. - -The kernel checks for this at run time too: `workSpace == 0` means the overlay -was not built in. - -## Consequently it is doubly gated, and never runs in CI +## SDMA and URMA are both available + +The A5 host runtime provisions both workspaces when it creates a communication +domain. `CommContext` keeps the SDMA workspace in the leading `sdmaWorkSpace` +pair and the URMA workspace in the appended `urmaWorkSpace` pair, so both +engines are usable from the same build and domain. + +Provisioning is best-effort on both: a transport that fails to initialize +leaves its pair zero and the transfer kernel takes its `urmaWorkSpace == 0` +self-skip. This demo then **fails**, and loudly — `compute_golden` expects +`out_rank == in_{1-rank}`, so a skipped TGET leaves `out` all zeros against a +non-zero golden. A green run of this test is therefore real evidence that URMA +moved data. When it fails, check the host log for +`URMA workspace unavailable on at least one rank` to tell an absent transport +apart from a wrong one. + +URMA metadata follows communicator-rank order, while each derived context +carries an explicit domain-rank-to-communicator-rank map. Subsets and reordered +domains therefore use the same communicator-scoped registration safely. + +The A5 worker registers one per-rank arena for the communicator (200 MiB by +default, `Worker(comm_arena_size=…)` to change it, and the backend reports back +the size it granted) and carves ordinary dynamic domains from it. The first 256 +bytes are reserved, so this demo's normal `allocate_domain` path performs a +real URMA TGET through a non-zero derived offset. Its domain worker order is +also reversed to `[1, 0]`, covering domain-rank-to-communicator-rank remapping, +peer MR selection, and `UrmaTget` in one transfer. Releasing a domain frees its +context and returns its slice; the HCCL registration and channels live until +communicator teardown. + +The test has only the ordinary scene constraints: | Gate | Effect | | ---- | ------ | | `CASES[*]["platforms"] = ["a5"]` | deselected on any other `--platform` | | `CASES[*]["config"]["device_count"] = 2` | needs two dies | -| `@pytest.mark.skipif(not _urma_workspace_enabled())` | skipped unless `SIMPLER_ENABLE_PTO_URMA_WORKSPACE` is one of `1` / `ON` / `TRUE` / `YES` in the environment | -| `urma_deferred_completion_orch_fn()` raises | re-checks the env var immediately before allocating the communication domain | - -Since the URMA CMake option defaults `OFF`, a stock build skips this test even on a5 -hardware. **A green CI run says nothing about URMA.** Treat it as a manual -bring-up check, not as coverage. ## Run ```bash -# 1. rebuild the host runtime with the overlay on -SIMPLER_ENABLE_PTO_URMA_WORKSPACE=ON pip install --no-build-isolation -e . - -# 2. run it, with the same variable visible to pytest (the skipif reads the env) -SIMPLER_ENABLE_PTO_URMA_WORKSPACE=ON \ - pytest examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo \ +pytest examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo \ --platform a5 --device 0-1 ``` -The variable is needed twice for different reasons: `runtime_builder.py` -forwards it to CMake so the overlay is compiled in, and the test reads it from -the environment to decide whether to skip. Setting only one of the two gives -you either a skipped test or a `workSpace == 0` failure. - Wrap the hardware run in `task-submit` on a shared box. ## See also [`../sdma_async_completion_demo/`](../sdma_async_completion_demo/) — the -default SDMA variant of the same protocol. +SDMA variant of the same protocol. It can run immediately before this test +without rebuilding or changing the environment. diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp index 9580862b16..d558fdefd7 100644 --- a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/kernels/aiv/kernel_urma_tget_async.cpp @@ -49,10 +49,8 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); __gm__ CommContext *comm_ctx = reinterpret_cast<__gm__ CommContext *>(args[2]); - // A null workspace means the host runtime was not built with the URMA - // backend; self-skip rather than dereferencing it. if (comm_ctx == nullptr || comm_ctx->rankNum != 2 || comm_ctx->rankId >= comm_ctx->rankNum || - comm_ctx->workSpace == 0 || comm_ctx->windowsIn[comm_ctx->rankId] == 0) { + comm_ctx->urmaWorkSpace == 0 || comm_ctx->windowsIn[comm_ctx->rankId] == 0) { pipe_barrier(PIPE_ALL); return; } @@ -60,9 +58,11 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ float *local_input = tensor_data(input_tensor); __gm__ float *local_out = tensor_data(out_tensor); uint32_t peer_rank = 1u - comm_ctx->rankId; + uint32_t urma_peer_rank = comm_ctx->urmaRankMap[peer_rank]; uint64_t input_offset = reinterpret_cast(local_input) - comm_ctx->windowsIn[comm_ctx->rankId]; + uint64_t registered_input_offset = comm_ctx->urmaWindowOffset + input_offset; __gm__ float *remote_input = pto2::urma_backend::peer_mr_ptr( - reinterpret_cast<__gm__ uint8_t *>(comm_ctx->workSpace), peer_rank, input_offset + reinterpret_cast<__gm__ uint8_t *>(comm_ctx->urmaWorkSpace), urma_peer_rank, registered_input_offset ); using FlatShape = Shape<1, 1, 1, 1, kElems>; @@ -75,6 +75,8 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in AsyncCtx async_ctx = get_async_ctx(args); (void)send_request_entry( async_ctx, - UrmaTget(local_global, remote_global, reinterpret_cast<__gm__ uint8_t *>(comm_ctx->workSpace), peer_rank) + UrmaTget( + local_global, remote_global, reinterpret_cast<__gm__ uint8_t *>(comm_ctx->urmaWorkSpace), urma_peer_rank + ) ); } diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py index cd31b08a72..fbc1a9eedc 100644 --- a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py @@ -14,9 +14,6 @@ event. The consumer depends on that output and writes ``result = out + 1``. """ -import os - -import pytest import torch from simpler.task_interface import ArgDirection as D from simpler.task_interface import CommBufferSpec, DataType, TaskArgs, TensorArgType @@ -28,28 +25,15 @@ NRANKS = 2 DTYPE_NBYTES = 4 URMA_DATA_OFFSET_NBYTES = 64 * DTYPE_NBYTES -_URMA_WORKSPACE_ENV = "SIMPLER_ENABLE_PTO_URMA_WORKSPACE" -_WORKSPACE_TRUTHY = {"1", "ON", "TRUE", "YES"} - - -def _urma_workspace_enabled(): - return os.environ.get(_URMA_WORKSPACE_ENV, "").upper() in _WORKSPACE_TRUTHY - - -def _require_urma_workspace_enabled(): - if not _urma_workspace_enabled(): - raise RuntimeError( - "urma_deferred_completion_demo requires host runtime built with " - f"{_URMA_WORKSPACE_ENV}=ON; set it before rebuilding simpler." - ) def urma_deferred_completion_orch_fn(orch, callables, task_args, config): - _require_urma_workspace_enabled() input_nbytes = N * DTYPE_NBYTES with orch.allocate_domain( name="urma_deferred_completion", - workers=list(range(NRANKS)), + # Reverse the domain order so the real TGET covers domain-rank to + # communicator-rank remapping rather than only checking its metadata. + workers=list(reversed(range(NRANKS))), window_size=max(URMA_DATA_OFFSET_NBYTES + input_nbytes, 4 * 1024 * 1024), buffers=[ CommBufferSpec( @@ -66,6 +50,8 @@ def urma_deferred_completion_orch_fn(orch, callables, task_args, config): orch.copy_to(handle[rank].buffers["input_window"], _rehosted_buffer_for(task_args, source)) for rank in range(NRANKS): domain = handle[rank] + if domain.window_offset == 0: + raise AssertionError("URMA regression requires a non-zero derived arena offset") args = TaskArgs() args.add_tensor(domain.buffers["input_window"].tensor((N,), DataType.FLOAT32), TensorArgType.INPUT) args.add_tensor( @@ -78,11 +64,6 @@ def urma_deferred_completion_orch_fn(orch, callables, task_args, config): orch.submit_next_level(callables.urma_deferred_completion, args, config, worker=rank) -@pytest.mark.skipif( - not _urma_workspace_enabled(), - reason="URMA workspace overlay not enabled (set SIMPLER_ENABLE_PTO_URMA_WORKSPACE=ON to run). " - "See docs/a5-sdma-overlay.md (#1315).", -) @scene_test(level=3, runtime="tensormap_and_ringbuffer") class TestUrmaDeferredCompletionDemo(SceneTestCase): CALLABLE = { diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 0da92d4996..fafc10e56d 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -3174,12 +3174,13 @@ NB_MODULE(_task_interface, m) { .def( "comm_alloc_domain_windows", [](ChipWorker &self, uint64_t comm_handle, uint64_t allocation_id, const std::vector &rank_ids, - uint32_t domain_rank, size_t window_size, uint64_t commit_flag_address) { + uint32_t domain_rank, size_t window_offset, size_t window_size, uint64_t commit_flag_address) { if (commit_flag_address == 0 || commit_flag_address % alignof(uint64_t) != 0) { throw std::invalid_argument("comm_alloc_domain_windows: commit flag address is invalid"); } - auto [device_ctx, local_window_base] = - self.comm_alloc_domain_windows(comm_handle, allocation_id, rank_ids, domain_rank, window_size); + auto [device_ctx, local_window_base] = self.comm_alloc_domain_windows( + comm_handle, allocation_id, rank_ids, domain_rank, window_offset, window_size + ); __atomic_store_n( reinterpret_cast(static_cast(commit_flag_address)), uint64_t{1}, __ATOMIC_RELEASE @@ -3187,7 +3188,7 @@ NB_MODULE(_task_interface, m) { return nb::make_tuple(device_ctx, local_window_base); }, nb::arg("comm_handle"), nb::arg("allocation_id"), nb::arg("rank_ids"), nb::arg("domain_rank"), - nb::arg("window_size"), nb::arg("commit_flag_address"), + nb::arg("window_offset"), nb::arg("window_size"), nb::arg("commit_flag_address"), "Collectively allocate a fresh per-rank pool for a subset; returns " "(device_ctx, local_window_base) for this rank and publishes the commit flag before result conversion." ) diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 64b41df177..1fc48323cf 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -912,17 +912,21 @@ def _validate_remote_sidecar_access(args: TaskArgs, remote_sidecar: _RemoteTaskA class _CommContextStruct(ctypes.Structure): _fields_ = [ - ("workSpace", ctypes.c_uint64), - ("workSpaceSize", ctypes.c_uint64), + ("sdmaWorkSpace", ctypes.c_uint64), + ("sdmaWorkSpaceSize", ctypes.c_uint64), ("rankId", ctypes.c_uint32), ("rankNum", ctypes.c_uint32), ("winSize", ctypes.c_uint64), ("windowsIn", ctypes.c_uint64 * COMM_MAX_RANK_NUM), ("windowsOut", ctypes.c_uint64 * COMM_MAX_RANK_NUM), + ("urmaWorkSpace", ctypes.c_uint64), + ("urmaWorkSpaceSize", ctypes.c_uint64), + ("urmaWindowOffset", ctypes.c_uint64), + ("urmaRankMap", ctypes.c_uint32 * COMM_MAX_RANK_NUM), ] -assert ctypes.sizeof(_CommContextStruct) == 1056 +assert ctypes.sizeof(_CommContextStruct) == 1336 def scalar_to_uint64(value) -> int: @@ -982,6 +986,9 @@ class ChipDomainContext: # Each named window slice as a device ``VMM_WINDOW`` Buffer owned by this chip. Name a task # arg with ``buffers[name].tensor(shapes, dtype)`` and dispatch it only to this chip (``domain_rank``). buffers: dict[str, Buffer] + # Offset inside a backend-owned persistent arena. Zero for backends that + # allocate one independent window per domain. + window_offset: int = 0 class CommDomainHandle: diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 785bb74cc5..f28faa0981 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -629,9 +629,12 @@ def _require_copy_span(extent: int, offset: int, nbytes: int, *, side: str, api: _CTRL_OP_NAMES[_CTRL_GLOBAL_DOMAIN_NODE] = "global_domain" _CTRL_OP_NAMES[_CTRL_DELEGATED_REGION] = "delegated_region" -# Layout of the CTRL_COMM_INIT request shm. -_COMM_INIT_HEADER = struct.Struct(" None: req_buf = req_shm.buf assert req_buf is not None try: - (allocation_id, rank_count, domain_rank, window_size, buffer_count) = _DOMAIN_REQ_HEADER.unpack_from(req_buf, 0) + (allocation_id, rank_count, domain_rank, window_offset, window_size, buffer_count) = ( + _DOMAIN_REQ_HEADER.unpack_from(req_buf, 0) + ) # Layout: header | buffer_nbytes[buffer_count] (u64) | rank_ids[rank_count] (u32) nbytes_offset = _DOMAIN_REQ_HEADER.size nbytes_struct = struct.Struct(f"<{buffer_count}Q") if buffer_count else struct.Struct("") @@ -2323,6 +2341,7 @@ def _handle_ctrl_alloc_domain(cw: ChipWorker, buf: memoryview) -> None: int(allocation_id), rank_ids, int(domain_rank), + int(window_offset), int(window_size), _buffer_field_addr(reply_buf, _OFF_DOMAIN_REPLY_COMMITTED), ) @@ -2347,7 +2366,7 @@ def _handle_ctrl_alloc_domain(cw: ChipWorker, buf: memoryview) -> None: reply_shm.close() -def _handle_ctrl_comm_init(cw: ChipWorker, buf: memoryview) -> None: +def _handle_ctrl_comm_init(cw: ChipWorker, buf: memoryview, chip_platform: str) -> None: """CTRL_COMM_INIT handler — drives `cw.comm_init` on the chip child. Idempotent: ``ChipWorker.comm_init`` itself caches the handle and returns @@ -2359,20 +2378,30 @@ def _handle_ctrl_comm_init(cw: ChipWorker, buf: memoryview) -> None: req_buf = req_shm.buf assert req_buf is not None try: - (rank, nranks) = _COMM_INIT_HEADER.unpack_from(req_buf, 0) + rank, nranks, requested_arena_size = _COMM_INIT_HEADER.unpack_from(req_buf, 0) # rootinfo_path is the rest of the shm, NUL-terminated. raw = bytes(req_buf[_COMM_INIT_HEADER.size :]) nul = raw.find(b"\x00") rootinfo_path = raw[: nul if nul >= 0 else len(raw)].decode("utf-8", "replace") + + handle = cw.comm_init(int(rank), int(nranks), rootinfo_path) + if handle == 0: + raise RuntimeError("comm_init returned 0 handle for hidden base communicator") + cw._comm_base_handle_cached = int(handle) + if chip_platform == "a5": + if not getattr(cw, "_comm_base_windows_ready", False): + device_ctx = cw.comm_alloc_windows(int(handle), int(requested_arena_size)) + if device_ctx == 0: + raise RuntimeError("comm_alloc_windows returned 0 for the A5 persistent communication arena") + cw._comm_base_windows_ready = True + arena_size = int(cw.comm_get_window_size(int(handle))) + if arena_size <= 0: + raise RuntimeError(f"comm_get_window_size returned invalid A5 arena size {arena_size}") + _COMM_INIT_HEADER.pack_into(req_buf, 0, int(rank), int(nranks), arena_size) finally: req_buf.release() req_shm.close() - handle = cw.comm_init(int(rank), int(nranks), rootinfo_path) - if handle == 0: - raise RuntimeError("comm_init returned 0 handle for hidden base communicator") - cw._comm_base_handle_cached = int(handle) - @dataclass class _L2GlobalDomain: @@ -2753,7 +2782,7 @@ def _handle_ctrl_release_domain(cw: ChipWorker, buf: memoryview) -> None: req_buf = req_shm.buf assert req_buf is not None try: - (allocation_id, rank_count, domain_rank, _ws, _bc) = _DOMAIN_REQ_HEADER.unpack_from(req_buf, 0) + (allocation_id, rank_count, domain_rank, _offset, _ws, _bc) = _DOMAIN_REQ_HEADER.unpack_from(req_buf, 0) finally: req_buf.release() req_shm.close() @@ -2996,7 +3025,7 @@ def handle_control( # noqa: PLR0912, PLR0915 -- one branch per control sub-comm elif sub_cmd == _CTRL_RELEASE_DOMAIN: _handle_ctrl_release_domain(cw, buf) elif sub_cmd == _CTRL_COMM_INIT: - _handle_ctrl_comm_init(cw, buf) + _handle_ctrl_comm_init(cw, buf, chip_platform) elif sub_cmd == _CTRL_DELEGATED_REGION: _handle_ctrl_delegated_region_terminal(buf, provider_transaction_table, provider_region_store) elif sub_cmd == _CTRL_COMMITTED_DEVICE_MEMORY: @@ -4443,7 +4472,7 @@ class Worker: add_worker() before init(). """ - def __init__( + def __init__( # noqa: PLR0915 -- lifecycle and communication state initialization stays centralized self, level: int, **config, @@ -4665,6 +4694,10 @@ def __init__( # keep ``Worker.init()`` cheap — it only forks chip children and # starts the C++ scheduler; no comm work happens there. self._comm_base_ready: bool = False + self._comm_arena_mu = threading.Lock() + self._comm_arena_size: int = 0 + self._comm_arena_free: list[tuple[int, int]] = [] + self._comm_arena_allocations: dict[int, tuple[int, int]] = {} self._endpoint_registry: EndpointRegistry | None = None self._endpoint_registry_epoch: int = 0 @@ -8673,7 +8706,12 @@ def _ensure_comm_base(self) -> None: device_ids = self._config.get("device_ids", []) rootinfo_path = self._comm_plan_rootinfo_path() - # Layout: header (rank, nranks) + NUL-terminated rootinfo_path bytes. + # Layout: header (rank, nranks, arena size) followed by NUL-terminated + # rootinfo_path bytes. The A5 child overwrites the arena-size field + # with what the backend actually allocated. + requested_arena_size = int(self._config.get("comm_arena_size", 0) or 0) + if requested_arena_size < 0: + raise ValueError(f"comm_arena_size must be non-negative, got {requested_arena_size}") path_bytes = rootinfo_path.encode("utf-8") + b"\x00" req_size = _COMM_INIT_HEADER.size + len(path_bytes) @@ -8684,7 +8722,7 @@ def initialize(request_owner: _SharedMemoryOwner) -> None: request_shms[chip_idx] = req req_buf = req.buf assert req_buf is not None - _COMM_INIT_HEADER.pack_into(req_buf, 0, int(chip_idx), int(len(device_ids))) + _COMM_INIT_HEADER.pack_into(req_buf, 0, int(chip_idx), int(len(device_ids)), requested_arena_size) req_buf[_COMM_INIT_HEADER.size : _COMM_INIT_HEADER.size + len(path_bytes)] = path_bytes dw = self._worker @@ -8705,6 +8743,16 @@ def dispatch(chip_idx: int) -> None: f"first error chip={first[0]}: {first[1]}" ) + if str(self._config.get("platform", "")) == "a5": + arena_sizes: set[int] = set() + for chip_idx in range(len(device_ids)): + req_buf = request_shms[chip_idx].buf + assert req_buf is not None + arena_sizes.add(int(_COMM_INIT_HEADER.unpack_from(req_buf, 0)[2])) + if len(arena_sizes) != 1: + raise RuntimeError(f"A5 backend reported inconsistent communication arena sizes: {arena_sizes}") + self._initialize_comm_arena(arena_sizes.pop()) + _run_with_owned_shared_memory( len(device_ids), initialize, @@ -8712,6 +8760,63 @@ def dispatch(chip_idx: int) -> None: after_success=lambda: setattr(self, "_comm_base_ready", True), ) + def _initialize_comm_arena(self, arena_size: int) -> None: + """Initialize the A5 slice allocator from the backend's actual window size.""" + arena_size = int(arena_size) + if arena_size <= _A5_COMM_ARENA_RESERVED: + raise RuntimeError( + f"A5 backend communication arena is too small: {arena_size} bytes " + f"(must exceed {_A5_COMM_ARENA_RESERVED})" + ) + with self._comm_arena_mu: + if self._comm_arena_allocations: + raise RuntimeError("cannot reinitialize the A5 communication arena while domain slices are live") + self._comm_arena_size = arena_size + self._comm_arena_free = [(_A5_COMM_ARENA_RESERVED, arena_size - _A5_COMM_ARENA_RESERVED)] + + def _reserve_comm_arena(self, allocation_id: int, window_size: int) -> int: + """Reserve one shared A5 arena offset; other backends keep offset zero.""" + if str(self._config.get("platform", "")) != "a5": + return 0 + alignment = _A5_COMM_ARENA_ALIGNMENT + reserved_size = ((int(window_size) + alignment - 1) // alignment) * alignment + with self._comm_arena_mu: + if self._comm_arena_size == 0: + raise RuntimeError("A5 communication arena is not initialized") + for index, (offset, extent) in enumerate(self._comm_arena_free): + if extent < reserved_size: + continue + self._comm_arena_allocations[allocation_id] = (offset, reserved_size) + if extent == reserved_size: + del self._comm_arena_free[index] + else: + self._comm_arena_free[index] = (offset + reserved_size, extent - reserved_size) + return offset + raise MemoryError( + f"A5 communication arena exhausted: requested {window_size} bytes " + f"({self._comm_arena_size - _A5_COMM_ARENA_RESERVED} bytes usable per rank, shared by every live " + f"domain). Raise it with Worker(comm_arena_size=)." + ) + + def _release_comm_arena(self, allocation_id: int) -> None: + """Return a successfully released A5 domain slice and coalesce holes.""" + if str(self._config.get("platform", "")) != "a5": + return + with self._comm_arena_mu: + allocation = self._comm_arena_allocations.pop(allocation_id, None) + if allocation is None: + return + self._comm_arena_free.append(allocation) + self._comm_arena_free.sort() + merged: list[tuple[int, int]] = [] + for offset, extent in self._comm_arena_free: + if merged and merged[-1][0] + merged[-1][1] == offset: + prior_offset, prior_extent = merged[-1] + merged[-1] = (prior_offset, prior_extent + extent) + else: + merged.append((offset, extent)) + self._comm_arena_free = merged + def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm staging + dispatch + reply unpack; splitting obscures the fail-fast ordering self, *, @@ -8735,6 +8840,7 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm with self._alloc_id_lock: allocation_id = self._next_alloc_id self._next_alloc_id += 1 + window_offset = self._reserve_comm_arena(allocation_id, window_size) # Stage per-chip request shms (domain_rank differs per chip) and a # per-chip reply shm. We let the chip child write back its own slot. @@ -8767,6 +8873,7 @@ def allocate(staged_shms: _SharedMemoryOwner) -> CommDomainHandle: int(allocation_id), int(len(workers)), int(worker_to_rank[chip_idx]), # domain_rank + int(window_offset), int(window_size), int(buffer_count), ) @@ -8846,6 +8953,7 @@ def allocate(staged_shms: _SharedMemoryOwner) -> CommDomainHandle: ) for i, b in enumerate(buffers) }, + window_offset=int(window_offset), ) handle.contexts = contexts finally: @@ -8881,12 +8989,22 @@ def publish_provenance() -> None: for buf in ctx.buffers.values(): self._record_device_alloc(buf, domain_allocation_id=allocation_id) - published_handle = _run_with_owned_shared_memory( - len(workers) * 2, - allocate, - name_prefix="shm-domain-alloc-lifecycle-", - after_success=publish_provenance, - ) + try: + published_handle = _run_with_owned_shared_memory( + len(workers) * 2, + allocate, + name_prefix="shm-domain-alloc-lifecycle-", + after_success=publish_provenance, + ) + except BaseException: # noqa: BLE001 + # A partial backend commit remains owned by its live handle and is + # reclaimed by the run/close sweep. Only a pre-commit failure may + # return the arena slice immediately. + if handle is None or ( + self._live_domains.get(name) is not handle and resources.live_domains.get(name) is not handle + ): + self._release_comm_arena(allocation_id) + raise assert handle is not None and published_handle is handle return handle @@ -9095,6 +9213,7 @@ def release(request_owner: _SharedMemoryOwner) -> None: int(handle.allocation_id), int(handle._domain_size), # noqa: SLF001 -- backend release identity belongs to the handle int(handle._domain_ranks[chip_idx]), # noqa: SLF001 -- preserve the allocation-time rank + 0, # window_offset — ignored on release 0, # window_size — ignored on release 0, # buffer_count — ignored on release ) @@ -9105,6 +9224,7 @@ def release(request_owner: _SharedMemoryOwner) -> None: op="release", allocation_id=handle.allocation_id, ) + self._release_comm_arena(handle.allocation_id) def retire_live_handle() -> None: if self._live_domains.get(handle.name) is handle: diff --git a/simpler_setup/runtime_builder.py b/simpler_setup/runtime_builder.py index 9602a17cb4..005dd73275 100644 --- a/simpler_setup/runtime_builder.py +++ b/simpler_setup/runtime_builder.py @@ -24,6 +24,7 @@ logger = logging.getLogger(__name__) _GIT_COMMIT_FILE = ".git_commit" +_RETIRED_A5_URMA_WORKSPACE_ENV = "SIMPLER_ENABLE_PTO_URMA_WORKSPACE" _DEFAULT_PROFILING_CONFIG = { "SIMPLER_DFX": "1", "SIMPLER_ORCH_PROFILING": "0", @@ -164,6 +165,11 @@ def __init__(self, platform: str = "a2a3"): """ self.platform = platform self._arch, self._variant = parse_platform(platform) + if self._arch == "a5" and self._variant == "onboard" and _RETIRED_A5_URMA_WORKSPACE_ENV in os.environ: + logger.warning( + "%s is retired; A5 onboard always builds SDMA and URMA workspaces", + _RETIRED_A5_URMA_WORKSPACE_ENV, + ) runtime_root = PROJECT_ROOT self.runtime_root = runtime_root @@ -391,9 +397,6 @@ def _compile_target(target: str) -> Path: if target == "host": if build_pto_isa_commit: defines["SIMPLER_PTO_ISA_BUILD_COMMIT"] = build_pto_isa_commit - for opt_in_define in ("SIMPLER_ENABLE_PTO_URMA_WORKSPACE",): - if os.environ.get(opt_in_define, "").upper() in {"1", "ON", "TRUE", "YES"}: - defines[opt_in_define] = "ON" cmake_defines = defines or None # compile() adds a {target}/ subdirectory inside build_dir cache_dir = self._CACHE_DIR / arch / variant / name diff --git a/simpler_setup/runtime_compiler.py b/simpler_setup/runtime_compiler.py index 24505a36d2..02b3604e54 100644 --- a/simpler_setup/runtime_compiler.py +++ b/simpler_setup/runtime_compiler.py @@ -180,8 +180,8 @@ def _init_a2a3(self): """Initialize toolchains for real a2a3 hardware.""" env_manager.ensure("ASCEND_HOME_PATH") # a2a3 onboard host_runtime hard-depends on pto-isa headers + CANN-9.0 - # aclnn syms (cf. src/a2a3/platform/onboard/host/CMakeLists.txt - # SIMPLER_ENABLE_PTO_SDMA_WORKSPACE marker). Resolve the pinned managed + # aclnn syms (cf. src/a2a3/platform/onboard/host/CMakeLists.txt). + # Resolve the pinned managed # checkout once; RuntimeBuilder passes it to CMake as -DPTO_ISA_ROOT= # (#1403 — do not smuggle via os.environ). from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: PLC0415 @@ -219,9 +219,10 @@ def _init_a2a3sim(self): def _init_a5(self): """Initialize toolchains for real a5 hardware.""" env_manager.ensure("ASCEND_HOME_PATH") - # a5 onboard always embeds the PTO-ISA SDMA workspace. Use the pinned - # managed checkout so build metadata and load-time validation match - # the headers compiled into host_runtime (#1351, #1403). + # a5 onboard always embeds the PTO-ISA SDMA and URMA workspaces. Use + # the pinned managed checkout so build metadata and load-time + # validation match the headers compiled into host_runtime (#1351, + # #1403). from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: PLC0415 self.pto_isa_root = ensure_pto_isa_root(verbose=True) diff --git a/src/a2a3/platform/onboard/host/comm_hccl.cpp b/src/a2a3/platform/onboard/host/comm_hccl.cpp index ec85e275c5..2ba84260e6 100644 --- a/src/a2a3/platform/onboard/host/comm_hccl.cpp +++ b/src/a2a3/platform/onboard/host/comm_hccl.cpp @@ -373,11 +373,11 @@ static aclError release_domain_windows(DomainAllocation *alloc) { static aclError release_base_windows(CommHandle h) { aclError first_error = release_vmm_windows(&h->base_peer_windows); record_first_error(release_vmm_window(&h->base_local_window), &first_error); - const uint64_t workspace = h->host_ctx.workSpace; - const uint64_t workspace_size = h->host_ctx.workSpaceSize; + const uint64_t workspace = h->host_ctx.sdmaWorkSpace; + const uint64_t workspace_size = h->host_ctx.sdmaWorkSpaceSize; h->host_ctx = CommContext{}; - h->host_ctx.workSpace = workspace; - h->host_ctx.workSpaceSize = workspace_size; + h->host_ctx.sdmaWorkSpace = workspace; + h->host_ctx.sdmaWorkSpaceSize = workspace_size; return first_error; } @@ -1367,7 +1367,8 @@ extern "C" int comm_derive_context( ); return -1; } - if (window_offset + window_size > static_cast(h->host_ctx.winSize)) { + if (window_offset > static_cast(h->host_ctx.winSize) || + window_size > static_cast(h->host_ctx.winSize) - window_offset) { LOG_ERROR( "[comm rank %d] comm_derive_context: window range [%zu, %zu) exceeds base window size %llu", h->rank, window_offset, window_offset + window_size, static_cast(h->host_ctx.winSize) @@ -1390,6 +1391,7 @@ extern "C" int comm_derive_context( } ctx.windowsIn[i] = h->host_ctx.windowsIn[base_rank] + window_offset; ctx.windowsOut[i] = h->host_ctx.windowsOut[base_rank] + window_offset; + ctx.urmaRankMap[i] = base_rank; } void *newDevMem = nullptr; @@ -1433,9 +1435,19 @@ extern "C" int comm_barrier(CommHandle h) { extern "C" int comm_alloc_domain_windows( CommHandle h, uint64_t allocation_id, const uint32_t *rank_ids, size_t rank_count, uint32_t domain_rank, - size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out + size_t window_offset, size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out ) try { if (!h || !rank_ids || !device_ctx_out || !local_window_base_out) return -1; + // Every allocation owns its window here, so there is no arena to offset + // into. Reject a non-zero request instead of returning a window that + // silently starts somewhere else than the caller asked for. + if (window_offset != 0) { + LOG_ERROR( + "[comm rank %d] alloc_domain: window_offset=%zu is not supported by the a2a3 backend", h->rank, + window_offset + ); + return -1; + } if (rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || window_size == 0) { LOG_ERROR( "[comm rank %d] alloc_domain: bad args (rank_count=%zu domain_rank=%u window_size=%zu)", h->rank, diff --git a/src/a5/platform/onboard/host/CMakeLists.txt b/src/a5/platform/onboard/host/CMakeLists.txt index 978afcfe4b..44e11c50cf 100644 --- a/src/a5/platform/onboard/host/CMakeLists.txt +++ b/src/a5/platform/onboard/host/CMakeLists.txt @@ -36,17 +36,6 @@ else() endif() set(ASCEND_ARCH_HOME "${ASCEND_HOME_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux") -option(SIMPLER_ENABLE_PTO_URMA_WORKSPACE "Enable a5 PTO URMA workspace overlay" OFF) -# SDMA is the default a5 async-workspace backend. URMA remains an explicit -# alternative because CommContext exposes one workSpace/workSpaceSize pair. -# SIMPLER_ENABLE_PTO_SDMA_WORKSPACE is an internal compile marker, not a -# user-configurable option. -if(SIMPLER_ENABLE_PTO_URMA_WORKSPACE) - set(SIMPLER_ENABLE_PTO_SDMA_WORKSPACE OFF) -else() - set(SIMPLER_ENABLE_PTO_SDMA_WORKSPACE ON) -endif() - # Pin-resolved checkout path is passed by RuntimeBuilder as -DPTO_ISA_ROOT= # (#1403). Do not read $ENV{PTO_ISA_ROOT} — ambient exports are not the pin. if(NOT PTO_ISA_ROOT) @@ -130,15 +119,8 @@ target_compile_options(host_runtime # src/common/platform/shared/host/platform_compile_info.cpp. target_compile_definitions(host_runtime PRIVATE SIMPLER_PLATFORM_NAME="a5") -if(SIMPLER_ENABLE_PTO_SDMA_WORKSPACE) - target_compile_definitions(host_runtime PRIVATE SIMPLER_ENABLE_PTO_SDMA_WORKSPACE=1) -endif() -if(SIMPLER_ENABLE_PTO_URMA_WORKSPACE) - target_compile_definitions(host_runtime PRIVATE SIMPLER_ENABLE_PTO_URMA_WORKSPACE=1) -endif() - # Bake the resolved pto-isa commit into the compile command because the -# selected async-workspace backend embeds pto-isa headers into this .so. A +# async-workspace backends embed pto-isa headers into this .so. A # pto-isa pin bump leaves runtime HEAD # — and often header mtimes — untouched, so without this define ccache # (compiler_check=mtime) can serve a stale object. Mirror of the a2a3 block diff --git a/src/a5/platform/onboard/host/comm_hccl.cpp b/src/a5/platform/onboard/host/comm_hccl.cpp index d8cefd20e0..3dea505414 100644 --- a/src/a5/platform/onboard/host/comm_hccl.cpp +++ b/src/a5/platform/onboard/host/comm_hccl.cpp @@ -46,12 +46,10 @@ #include "acl/acl.h" #include "hccl/hccl_comm.h" #include "hccl/hccl_types.h" -#ifdef SIMPLER_ENABLE_PTO_SDMA_WORKSPACE #include "pto/comm/async/sdma/sdma_workspace_manager.hpp" -#endif -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE #include "pto/comm/async/urma/urma_workspace_manager.hpp" -#endif + +#include // Thin wrappers around the HCCL public APIs we use. Kept as a translation // layer in case we need to swap (e.g., InitConfig variant) later. @@ -83,9 +81,12 @@ struct DomainAllocation { // can cycle repeatedly within one comm handle before any device reset, so // these are released explicitly at domain teardown rather than left to reset. std::vector> peer_windows; -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE - std::unique_ptr urma_workspace; -#endif + // A5 dynamic domains are slices of the communicator-owned base arena. + // Such slices own only their device context, never the VMM mapping or + // the communicator-scoped URMA registration. + bool arena_slice = false; + size_t window_offset = 0; + size_t window_size = 0; CommContext *device_ctx = nullptr; // aclrtMalloc'd CommContext mirror }; @@ -102,14 +103,11 @@ struct CommHandle_ { CommContext host_ctx{}; CommContext *device_ctx = nullptr; bool owns_device_ctx = false; + bool base_windows_attempted = false; std::vector derived_contexts; std::unordered_map> domain_allocations; -#ifdef SIMPLER_ENABLE_PTO_SDMA_WORKSPACE std::unique_ptr sdma_workspace; -#endif -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE std::unique_ptr urma_workspace; -#endif }; // ============================================================================ @@ -223,6 +221,56 @@ static bool file_barrier( return true; } +// Collective AND over one boolean per rank, carried in the barrier marker's +// payload. The marker is published by rename so a reader never observes a +// half-written payload. The return value reports the rendezvous itself; +// *agreed receives the AND, which is meaningful only when it returns true. +static bool file_barrier_agree( + const std::string &rootinfo_path, int rank, int nranks, const std::string &tag, uint64_t run_token, + bool local_value, bool *agreed, int timeout_sec = 120 +) { + const std::string my_marker = barrier_marker_path(rootinfo_path, run_token, tag, rank); + const std::string tmp = my_marker + ".tmp." + std::to_string(getpid()); + { + std::ofstream f(tmp, std::ios::trunc); + f << (local_value ? '1' : '0'); + if (!f.good()) { + std::remove(tmp.c_str()); + LOG_ERROR("[comm rank %d] file_barrier_agree('%s'): marker write failed", rank, tag.c_str()); + return false; + } + } + if (std::rename(tmp.c_str(), my_marker.c_str()) != 0) { + std::remove(tmp.c_str()); + LOG_ERROR("[comm rank %d] file_barrier_agree('%s'): marker publish failed", rank, tag.c_str()); + return false; + } + + bool all_true = true; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_sec); + for (int r = 0; r < nranks; ++r) { + const std::string marker = barrier_marker_path(rootinfo_path, run_token, tag, r); + while (true) { + std::ifstream f(marker); + char value = '\0'; + if (f.good() && f.get(value)) { + all_true = all_true && value == '1'; + break; + } + if (std::chrono::steady_clock::now() >= deadline) { + LOG_ERROR( + "[comm rank %d] file_barrier_agree('%s') timed out after %ds waiting for rank %d", rank, + tag.c_str(), timeout_sec, r + ); + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + } + *agreed = all_true; + return true; +} + // Release one VMM window — either an own-rank allocation or a peer import; the // teardown sequence is identical. `va` must be the mapped address returned by // aclrtMapMem and `handle` the backing physical handle (from aclrtMallocPhysical @@ -658,70 +706,12 @@ static int alloc_windows_via_ipc(CommHandle h, uint64_t win_size) { // ============================================================================ // Per-domain dynamic allocation (for orch.allocate_domain). // -// Same Path-D VMM dance as alloc_windows_via_ipc, but on a fresh per-allocation -// local buffer. Every barrier filename and announce filename is scoped by -// allocation_id so concurrent allocations from different orch.allocate_domain -// calls do not collide. Participation is by subset (domain_rank within -// rank_count), so non-members of the subset are not involved. +// A domain is a slice of the communicator-lifetime base arena, so it carries no +// per-allocation VMM mapping and no IPC handshake: the memory is mapped and +// peer-visible from comm_alloc_windows onwards. The only cross-rank step a +// domain still owns is the release barrier. // ============================================================================ -// Announce file path scoped by allocation_id so two concurrent allocations -// from different orch calls do not collide. Same dir + cleanup-friendly -// prefix as the base-comm IPC announce. -static std::string -domain_announce_path(const std::string &rootinfo, uint64_t allocation_id, uint32_t domain_rank, uint64_t run_token) { - return handshake_dir(rootinfo) + "/barrier_" + handshake_prefix(rootinfo) + "_alloc_" + - std::to_string(allocation_id) + "_ipc_announce_" + run_token_hex(run_token) + "_" + - std::to_string(domain_rank) + ".ready"; -} - -static bool domain_write_announce( - const std::string &rootinfo, uint64_t allocation_id, uint32_t domain_rank, uint64_t run_token, int32_t pid, - int32_t device_id, uint64_t shareable_handle -) { - IpcAnnounceFile a{}; - a.magic = kIpcAnnounceMagic; - a.pid = pid; - a.rank = domain_rank; - a.device_id = device_id; - a.shareable_handle = shareable_handle; - std::string p = domain_announce_path(rootinfo, allocation_id, domain_rank, run_token); - std::string tmp = p + ".tmp." + std::to_string(getpid()); - { - std::ofstream f(tmp, std::ios::binary | std::ios::trunc); - f.write(reinterpret_cast(&a), sizeof(a)); - if (!f.good()) { - std::remove(tmp.c_str()); - return false; - } - } - if (std::rename(tmp.c_str(), p.c_str()) != 0) { - std::remove(tmp.c_str()); - return false; - } - return true; -} - -static bool domain_read_announce( - const std::string &rootinfo, uint64_t allocation_id, uint32_t peer_domain_rank, uint64_t run_token, - IpcAnnounceFile *out, int timeout_sec = 60 -) { - std::string p = domain_announce_path(rootinfo, allocation_id, peer_domain_rank, run_token); - for (int i = 0; i < timeout_sec * 10; ++i) { - std::ifstream f(p, std::ios::binary); - if (f.good()) { - IpcAnnounceFile a{}; - f.read(reinterpret_cast(&a), sizeof(a)); - if (f.good() && a.magic == kIpcAnnounceMagic && a.rank == peer_domain_rank) { - *out = a; - return true; - } - } - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - return false; -} - // Tag helper for allocation-scoped file barriers. Tag is fed straight into // `file_barrier`, which already namespaces the marker filename by // rootinfo prefix + run_token + rank, so adding allocation_id to `tag` is @@ -736,12 +726,11 @@ static std::string domain_barrier_tag(uint64_t allocation_id, const char *phase) // first call allocates. Requires CANN to expose working // aclnnShmemSdmaStarsQuery primitives. static void ensure_sdma_workspace(CommHandle h) { -#ifdef SIMPLER_ENABLE_PTO_SDMA_WORKSPACE if (h->sdma_workspace) return; h->sdma_workspace = std::make_unique(); if (h->sdma_workspace->Init()) { - h->host_ctx.workSpace = reinterpret_cast(h->sdma_workspace->GetWorkspaceAddr()); - h->host_ctx.workSpaceSize = 16 * 1024; + h->host_ctx.sdmaWorkSpace = reinterpret_cast(h->sdma_workspace->GetWorkspaceAddr()); + h->host_ctx.sdmaWorkSpaceSize = 16 * 1024; } else { // SDMA workspace initialization failed - this may occur due to: // 1. Missing ACL symbols in libopapi.so (CANN version compatibility) @@ -750,9 +739,6 @@ static void ensure_sdma_workspace(CommHandle h) { // The system gracefully degrades to non-SDMA mode when this occurs. h->sdma_workspace.reset(); } -#else - (void)h; -#endif } // Callable-declared workspace injection is not available on a5 yet. Its URMA @@ -776,7 +762,6 @@ extern "C" int dma_workspace_provision(uint32_t required_mask, uint64_t *addr_ou extern "C" void dma_workspace_release(void *handle) { (void)handle; } -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE static uint64_t urma_workspace_bytes(uint32_t rank_count) { using namespace pto::comm::urma; constexpr uint32_t qp_num = 1; @@ -785,14 +770,6 @@ static uint64_t urma_workspace_bytes(uint32_t rank_count) { (2ULL * sizeof(UrmaWQCtx) * qp_num + 2ULL * sizeof(UrmaCqCtx) * qp_num + sizeof(UrmaMemInfo) * qp_num); } -static bool rank_ids_are_dense_prefix(const uint32_t *rank_ids, size_t rank_count) { - if (rank_ids == nullptr) return false; - for (size_t i = 0; i < rank_count; ++i) { - if (rank_ids[i] != static_cast(i)) return false; - } - return true; -} - static bool init_urma_workspace( CommHandle h, uint32_t rank_id, uint32_t rank_count, void *symmetric_addr, uint64_t symmetric_size, std::unique_ptr &workspace @@ -803,21 +780,24 @@ static bool init_urma_workspace( return false; } - auto manager = std::make_unique(); - if (!manager->Init(h->hccl_comm, rank_id, rank_count, symmetric_addr, symmetric_size)) { + // Transfer ownership before Init: it may register HCCL memory or acquire + // channels before a later allocation throws. The handle must retain that + // partial state until HcclCommDestroy runs. + workspace = std::make_unique(); + const bool initialized = workspace->Init(h->hccl_comm, rank_id, rank_count, symmetric_addr, symmetric_size); + if (!initialized) { LOG_WARN( "[comm rank %d] URMA workspace init failed (rank_id=%u rank_count=%u size=%llu)", h->rank, rank_id, rank_count, static_cast(symmetric_size) ); return false; } - workspace = std::move(manager); return true; } static bool ensure_base_urma_workspace(CommHandle h) { if (h == nullptr) return false; - if (h->urma_workspace) return h->host_ctx.workSpace != 0 && h->host_ctx.workSpaceSize != 0; + if (h->urma_workspace) return h->host_ctx.urmaWorkSpace != 0 && h->host_ctx.urmaWorkSpaceSize != 0; void *local_buf = reinterpret_cast(static_cast(h->host_ctx.windowsIn[h->rank])); if (!init_urma_workspace( h, static_cast(h->rank), static_cast(h->nranks), local_buf, h->host_ctx.winSize, @@ -825,337 +805,27 @@ static bool ensure_base_urma_workspace(CommHandle h) { )) { return false; } - h->host_ctx.workSpace = reinterpret_cast(h->urma_workspace->GetWorkspaceAddr()); - h->host_ctx.workSpaceSize = urma_workspace_bytes(static_cast(h->nranks)); - return h->host_ctx.workSpace != 0 && h->host_ctx.workSpaceSize != 0; + h->host_ctx.urmaWorkSpace = reinterpret_cast(h->urma_workspace->GetWorkspaceAddr()); + h->host_ctx.urmaWorkSpaceSize = urma_workspace_bytes(static_cast(h->nranks)); + return h->host_ctx.urmaWorkSpace != 0 && h->host_ctx.urmaWorkSpaceSize != 0; } -#endif - -static void reset_domain_urma_workspace(DomainAllocation &alloc) { -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE - alloc.urma_workspace.reset(); -#else - (void)alloc; -#endif -} - -static void reset_base_urma_workspace(CommHandle h) { -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE - h->urma_workspace.reset(); -#else - (void)h; -#endif -} - -// Performs the per-allocation Path-D dance for one subset rank. rank_ids -// must list participating BASE-COMM rank ids in domain rank order; this -// rank's domain_rank must match its base rank for the same invariant -// alloc_windows_via_ipc relies on (rank_ids[domain_rank] == h->rank). -// -// Failure paths tear down the own VMM window if it was mapped, plus every peer -// import already recorded on `out` (release_domain_peer_windows). On success -// the peer imports live on `out->peer_windows` for comm_release_domain_windows. -static int domain_alloc_via_ipc( - CommHandle h, uint64_t allocation_id, const uint32_t *rank_ids, size_t rank_count, uint32_t domain_rank, - uint64_t win_size, DomainAllocation *out -) { - const std::string &rootinfo = h->rootinfo_path; - const uint64_t run_token = h->run_token; - const int subset_n = static_cast(rank_count); - const int my_dr = static_cast(domain_rank); - - int32_t myDevice = -1; - if (aclrtGetDevice(&myDevice) != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: aclrtGetDevice failed", h->rank); - return -1; - } - - // VMM own-window allocation; see alloc_windows_via_ipc for step-by-step - // rationale. - aclrtPhysicalMemProp prop{}; - prop.handleType = ACL_MEM_HANDLE_TYPE_NONE; - prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED; - prop.memAttr = ACL_HBM_MEM_NORMAL; - prop.location.id = static_cast(myDevice); - prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE; - size_t granularity = 0; - aclError aret = aclrtMemGetAllocationGranularity(&prop, ACL_RT_MEM_ALLOC_GRANULARITY_MINIMUM, &granularity); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: GetAllocationGranularity -> %d", h->rank, static_cast(aret)); - return -1; - } - const uint64_t aligned_size = - granularity == 0 ? win_size : ((win_size + granularity - 1) / granularity) * granularity; - - aclrtDrvMemHandle handle = nullptr; - aret = aclrtMallocPhysical(&handle, aligned_size, &prop, 0); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: MallocPhysical -> %d", h->rank, static_cast(aret)); - return -1; - } - void *localBuf = nullptr; - aret = aclrtReserveMemAddress(&localBuf, aligned_size, 0, nullptr, 0); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: ReserveMemAddress -> %d", h->rank, static_cast(aret)); - aclrtFreePhysical(handle); - return -1; - } - aret = aclrtMapMem(localBuf, aligned_size, 0, handle, 0); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: MapMem -> %d", h->rank, static_cast(aret)); - aclrtReleaseMemAddress(localBuf); - aclrtFreePhysical(handle); - return -1; - } - // Driver-visible id, as in alloc_windows_via_ipc above; the peer mappings reuse this descriptor. - aclrtMemAccessDesc accessDesc{}; - accessDesc.flags = ACL_RT_MEM_ACCESS_FLAGS_READWRITE; - accessDesc.location.type = ACL_MEM_LOCATION_TYPE_DEVICE; - accessDesc.location.id = static_cast(pto::acl_to_hal_device_id(myDevice)); - aret = aclrtMemSetAccess(localBuf, aligned_size, &accessDesc, 1); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: MemSetAccess -> %d", h->rank, static_cast(aret)); - release_own_vmm_window(localBuf, handle); - return -1; - } - uint64_t shareableHandle = 0; - aret = aclrtMemExportToShareableHandle( - handle, ACL_MEM_HANDLE_TYPE_NONE, ACL_RT_VMM_EXPORT_FLAG_DISABLE_PID_VALIDATION, &shareableHandle - ); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: ExportToShareableHandle -> %d", h->rank, static_cast(aret)); - release_own_vmm_window(localBuf, handle); - return -1; - } - // Wipe before the handle is published. Publication is the point from which - // a peer may import this window and store into it — a barrier signal lands - // there as soon as that peer's kernel runs — so a wipe issued any later can - // erase a signal the owner has not yet waited on. Kernels take the zeroed - // tail as the initial value of their barrier-signal slots. The full - // granularity-aligned mapped range is zeroed to match ctx.winSize. - aret = aclrtMemset(localBuf, aligned_size, 0, aligned_size); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: aclrtMemset -> %d", h->rank, static_cast(aret)); - release_own_vmm_window(localBuf, handle); - return -1; - } - - const int32_t myPid = static_cast(getpid()); - if (!domain_write_announce(rootinfo, allocation_id, domain_rank, run_token, myPid, myDevice, shareableHandle)) { - LOG_ERROR("[comm rank %d] alloc_domain: write_announce failed", h->rank); - release_own_vmm_window(localBuf, handle); - return -1; - } - std::vector peers(subset_n); - for (int p = 0; p < subset_n; ++p) { - if (p == my_dr) { - peers[p].magic = kIpcAnnounceMagic; - peers[p].pid = myPid; - peers[p].rank = domain_rank; - peers[p].device_id = myDevice; - peers[p].shareable_handle = shareableHandle; - continue; - } - if (!domain_read_announce(rootinfo, allocation_id, static_cast(p), run_token, &peers[p])) { - LOG_ERROR("[comm rank %d] alloc_domain: read_announce(peer_dr=%d) timed out", h->rank, p); - release_own_vmm_window(localBuf, handle); - return -1; - } - } - - // Enable cross-card P2P for every domain peer, then a best-effort - // confirmation poll. The orch-only allocate_domain model has no base - // comm_alloc_windows to own the P2P route, so each allocation must - // (idempotently) ensure it. aclrtDeviceEnablePeerAccess is process-global - // and per device-pair; once any allocation opens a pair, later ones simply - // observe it. The enable is the operative call (resolves the peer physical - // id via the HCCL adapter and opens the HCCS route). - for (int p = 0; p < subset_n; ++p) { - if (p == my_dr) continue; - aclError r = aclrtDeviceEnablePeerAccess(peers[p].device_id, 0); - if (r != ACL_SUCCESS) { - LOG_WARN( - "[comm rank %d] alloc_domain: EnablePeerAccess(peer_dev=%d) -> %d", h->rank, peers[p].device_id, - static_cast(r) - ); - } - } - // See the device-remap note in alloc_windows_via_ipc: under - // ASCEND_VISIBLE_DEVICES remapping aclrtDevicePeerAccessStatus cannot - // resolve a peer that this fork'd single-device process never set, so it - // reports status=0 even when P2P is up. Confirm quickly where reliable, - // else fall through with a warning; the file_barrier below synchronizes - // every rank's enable and a dead link surfaces as a kernel-side hang. - for (int p = 0; p < subset_n; ++p) { - if (p == my_dr) continue; - const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); - while (true) { - int32_t status = 0; - aclError r = aclrtDevicePeerAccessStatus(myDevice, peers[p].device_id, &status); - if (r != ACL_SUCCESS) { - LOG_ERROR( - "[comm rank %d] alloc_domain: PeerAccessStatus(local_dev=%d peer_dev=%d) -> %d", h->rank, myDevice, - peers[p].device_id, static_cast(r) - ); - release_own_vmm_window(localBuf, handle); - return -1; - } - if (status == 1) break; - if (std::chrono::steady_clock::now() >= deadline) { - LOG_WARN( - "[comm rank %d] alloc_domain: P2P status unconfirmed peer_dr=%d peer_dev=%d status=%d " - "(proceeding after best-effort enable attempt, see device-remap note)", - h->rank, p, peers[p].device_id, status - ); - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } - } - - // With DISABLE_PID_VALIDATION the import can proceed once peers have - // published their shareable handles (read above) and P2P is up. - if (!file_barrier(rootinfo, my_dr, subset_n, domain_barrier_tag(allocation_id, "p2p_ready"), run_token)) { - release_own_vmm_window(localBuf, handle); - return -1; - } - - out->rank = my_dr; - out->nranks = subset_n; - out->local_buf = localBuf; - out->own_handle = handle; - // Build a host-side CommContext for the subset and upload it as device_ctx. - // PTO-ISA async SDMA ops (SdmaTget) read the scratch workspace off - // CommContext::workSpace. The dynamic-domain path does not go through - // comm_alloc_windows, so provision the workspace here; without it a - // freshly zero-initialized per-domain ctx would leave workSpace == 0 and - // those kernels early-return on the workSpace guard. - ensure_sdma_workspace(h); - - uint64_t domain_workspace_addr = 0; - uint64_t domain_workspace_size = 0; -#ifdef SIMPLER_ENABLE_PTO_SDMA_WORKSPACE - if (h->sdma_workspace) { - domain_workspace_addr = reinterpret_cast(h->sdma_workspace->GetWorkspaceAddr()); - domain_workspace_size = 16 * 1024; - } -#endif -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE - if (rank_ids_are_dense_prefix(rank_ids, rank_count)) { - if (!init_urma_workspace( - h, domain_rank, static_cast(rank_count), localBuf, aligned_size, out->urma_workspace - )) { - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - domain_workspace_addr = reinterpret_cast(out->urma_workspace->GetWorkspaceAddr()); - domain_workspace_size = urma_workspace_bytes(static_cast(rank_count)); - } else { - LOG_WARN("[comm rank %d] alloc_domain: URMA workspace disabled for non-dense rank mapping", h->rank); - } -#endif - - CommContext ctx{}; - ctx.rankId = domain_rank; - ctx.rankNum = static_cast(subset_n); - ctx.winSize = aligned_size; - ctx.workSpace = domain_workspace_addr; - ctx.workSpaceSize = domain_workspace_size; - ctx.windowsIn[my_dr] = reinterpret_cast(localBuf); - // Import each peer's shareable handle onto our device; see the symmetry - // note in alloc_windows_via_ipc (one win_size, shared chip granularity). - for (int p = 0; p < subset_n; ++p) { - if (p == my_dr) continue; - aclrtDrvMemHandle peerHandle = nullptr; - aret = aclrtMemImportFromShareableHandle(peers[p].shareable_handle, myDevice, &peerHandle); - if (aret != ACL_SUCCESS) { - LOG_ERROR( - "[comm rank %d] alloc_domain: ImportFromShareableHandle(peer_dr=%d) -> %d", h->rank, p, - static_cast(aret) - ); - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - void *peerVa = nullptr; - aret = aclrtReserveMemAddress(&peerVa, aligned_size, 0, nullptr, 0); - if (aret != ACL_SUCCESS) { - LOG_ERROR( - "[comm rank %d] alloc_domain: peer ReserveMemAddress(peer_dr=%d) -> %d", h->rank, p, - static_cast(aret) - ); - aclrtFreePhysical(peerHandle); - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - aret = aclrtMapMem(peerVa, aligned_size, 0, peerHandle, 0); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: peer MapMem(peer_dr=%d) -> %d", h->rank, p, static_cast(aret)); - aclrtReleaseMemAddress(peerVa); - aclrtFreePhysical(peerHandle); - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - aret = aclrtMemSetAccess(peerVa, aligned_size, &accessDesc, 1); - if (aret != ACL_SUCCESS) { - LOG_ERROR( - "[comm rank %d] alloc_domain: peer MemSetAccess(peer_dr=%d) -> %d", h->rank, p, static_cast(aret) - ); - aclrtUnmapMem(peerVa); - aclrtReleaseMemAddress(peerVa); - aclrtFreePhysical(peerHandle); - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - out->peer_windows.emplace_back(peerVa, peerHandle); - ctx.windowsIn[p] = reinterpret_cast(peerVa); - } - - void *newDevMem = nullptr; - aret = aclrtMalloc(&newDevMem, sizeof(CommContext), ACL_MEM_MALLOC_HUGE_FIRST); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: ctx aclrtMalloc -> %d", h->rank, static_cast(aret)); - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - aret = aclrtMemcpy(newDevMem, sizeof(CommContext), &ctx, sizeof(CommContext), ACL_MEMCPY_HOST_TO_DEVICE); - if (aret != ACL_SUCCESS) { - LOG_ERROR("[comm rank %d] alloc_domain: ctx Memcpy H2D -> %d", h->rank, static_cast(aret)); - aclrtFree(newDevMem); - reset_domain_urma_workspace(*out); - release_domain_peer_windows(*out); - release_own_vmm_window(localBuf, handle); - return -1; - } - out->device_ctx = reinterpret_cast(newDevMem); - return 0; -} +static void reset_base_urma_workspace(CommHandle h) { h->urma_workspace.reset(); } } // namespace extern "C" int comm_alloc_windows(CommHandle h, size_t win_size, uint64_t *device_ctx_out) try { if (!h || !device_ctx_out) return -1; - // Idempotency guard: comm_alloc_windows is not re-entrant. The localBuf - // allocated by alloc_windows_via_ipc is owned by the handle's windowsIn[] - // entries and is only reclaimed at aclrtResetDevice; calling this twice - // would leak a full per-rank pool. device_ctx is set on first success. - if (h->device_ctx != nullptr) { - LOG_ERROR("[comm rank %d] comm_alloc_windows: already allocated on this handle", h->rank); + // Idempotency guard: even a failed URMA setup may already have registered + // the still-live base VMM window. Refuse a same-handle retry so the fixed + // HCCL tag can never collide with a partial first attempt; comm_destroy + // tears down the communicator before a fresh handle may try again. + if (h->base_windows_attempted) { + LOG_ERROR("[comm rank %d] comm_alloc_windows: allocation was already attempted on this handle", h->rank); return -1; } + h->base_windows_attempted = true; // Path D: DIY symmetric pool on stable ACL VMM shareable handles + // EnablePeerAccess. Replaced the prior HcclAllocComResourceByTiling @@ -1165,13 +835,35 @@ extern "C" int comm_alloc_windows(CommHandle h, size_t win_size, uint64_t *devic const uint64_t effective_win_size = win_size != 0 ? static_cast(win_size) : kDefaultIpcWinSize; if (alloc_windows_via_ipc(h, effective_win_size) != 0) return -1; - // Optional PTO-ISA async SDMA workspace pre-allocation (overlays the comm - // backend's output; comm-side flow does not care about workSpace). + // Both PTO-ISA async transport workspaces are materialized before the + // CommContext is uploaded, and both degrade the same way: a workspace that + // fails to initialise leaves its CommContext pair zero, and the kernels + // take their `sdmaWorkSpace == 0` / `urmaWorkSpace == 0` self-skip. A + // window is never denied over a missing transport. ensure_sdma_workspace(h); -#ifdef SIMPLER_ENABLE_PTO_URMA_WORKSPACE - if (!ensure_base_urma_workspace(h)) return -1; - if (!file_barrier(h->rootinfo_path, h->rank, h->nranks, "base_urma_ready", h->run_token)) return -1; -#endif + for (int rank = 0; rank < h->nranks; ++rank) { + h->host_ctx.urmaRankMap[rank] = static_cast(rank); + } + // URMA's outcome is agreed across ranks rather than decided locally: a rank + // that kept a live registration would issue an RDMA against a peer that has + // none. A rank whose Init succeeded still keeps its manager — Finalize runs + // after HcclCommDestroy either way — and only publishes a zero workspace. + const bool local_urma_ready = ensure_base_urma_workspace(h); + bool urma_ready = false; + if (!file_barrier_agree( + h->rootinfo_path, h->rank, h->nranks, "base_urma_ready", h->run_token, local_urma_ready, &urma_ready + )) { + return -1; + } + if (!urma_ready) { + LOG_WARN( + "[comm rank %d] comm_alloc_windows: URMA workspace unavailable on at least one rank; " + "this communicator is SDMA-only", + h->rank + ); + h->host_ctx.urmaWorkSpace = 0; + h->host_ctx.urmaWorkSpaceSize = 0; + } void *newDevMem = nullptr; aclError aRet = aclrtMalloc(&newDevMem, sizeof(CommContext), ACL_MEM_MALLOC_HUGE_FIRST); @@ -1221,7 +913,8 @@ extern "C" int comm_derive_context( ); return -1; } - if (window_offset + window_size > static_cast(h->host_ctx.winSize)) { + if (window_offset > static_cast(h->host_ctx.winSize) || + window_size > static_cast(h->host_ctx.winSize) - window_offset) { LOG_ERROR( "[comm rank %d] comm_derive_context: window range [%zu, %zu) exceeds base window size %llu", h->rank, window_offset, window_offset + window_size, static_cast(h->host_ctx.winSize) @@ -1230,8 +923,11 @@ extern "C" int comm_derive_context( } CommContext ctx{}; - ctx.workSpace = h->host_ctx.workSpace; - ctx.workSpaceSize = h->host_ctx.workSpaceSize; + ctx.sdmaWorkSpace = h->host_ctx.sdmaWorkSpace; + ctx.sdmaWorkSpaceSize = h->host_ctx.sdmaWorkSpaceSize; + ctx.urmaWorkSpace = h->host_ctx.urmaWorkSpace; + ctx.urmaWorkSpaceSize = h->host_ctx.urmaWorkSpaceSize; + ctx.urmaWindowOffset = window_offset; ctx.rankId = domain_rank; ctx.rankNum = static_cast(rank_count); ctx.winSize = window_size; @@ -1246,6 +942,7 @@ extern "C" int comm_derive_context( } ctx.windowsIn[i] = h->host_ctx.windowsIn[base_rank] + window_offset; ctx.windowsOut[i] = h->host_ctx.windowsOut[base_rank] + window_offset; + ctx.urmaRankMap[i] = base_rank; } void *newDevMem = nullptr; @@ -1289,7 +986,7 @@ extern "C" int comm_barrier(CommHandle h) { extern "C" int comm_alloc_domain_windows( CommHandle h, uint64_t allocation_id, const uint32_t *rank_ids, size_t rank_count, uint32_t domain_rank, - size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out + size_t window_offset, size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out ) try { if (!h || !rank_ids || !device_ctx_out || !local_window_base_out) return -1; if (rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || window_size == 0) { @@ -1313,18 +1010,66 @@ extern "C" int comm_alloc_domain_windows( ); return -1; } - // The base communicator only needs comm_init to have run (rootinfo_path - // + run_token are set, used to scope barrier filenames). We do NOT - // require comm_alloc_windows on the base in the orch-only model — the - // dynamic alloc path does its own per-allocation aclrtMalloc + IPC dance. - if (h->rootinfo_path.empty() || h->hccl_comm == nullptr) { - LOG_ERROR("[comm rank %d] alloc_domain: base communicator not initialised", h->rank); + // The arena, not any async transport, is what a domain slice needs. A + // communicator that came up SDMA-only still hands out windows. + if (h->device_ctx == nullptr || h->host_ctx.winSize == 0 || h->host_ctx.windowsIn[h->rank] == 0) { + LOG_ERROR("[comm rank %d] alloc_domain: persistent base arena is not initialised", h->rank); + return -1; + } + if (window_offset > static_cast(h->host_ctx.winSize) || + window_size > static_cast(h->host_ctx.winSize) - window_offset) { + LOG_ERROR( + "[comm rank %d] alloc_domain: arena range [%zu, %zu) exceeds base window size %llu", h->rank, window_offset, + window_offset + window_size, static_cast(h->host_ctx.winSize) + ); + return -1; + } + for (const auto &entry : h->domain_allocations) { + const auto &live = entry.second; + if (!live->arena_slice) continue; + const size_t live_end = live->window_offset + live->window_size; + const size_t requested_end = window_offset + window_size; + if (window_offset < live_end && live->window_offset < requested_end) { + LOG_ERROR( + "[comm rank %d] alloc_domain: arena range [%zu, %zu) overlaps live allocation_id=%llu " + "range [%zu, %zu)", + h->rank, window_offset, requested_end, static_cast(entry.first), + live->window_offset, live_end + ); + return -1; + } + } + + void *local_slice = reinterpret_cast( + static_cast(h->host_ctx.windowsIn[h->rank] + static_cast(window_offset)) + ); + aclError clear_ret = aclrtMemset(local_slice, window_size, 0, window_size); + if (clear_ret != ACL_SUCCESS) { + LOG_ERROR( + "[comm rank %d] alloc_domain: arena-slice aclrtMemset failed: %d", h->rank, static_cast(clear_ret) + ); return -1; } auto alloc = std::make_unique(); - int rc = domain_alloc_via_ipc(h, allocation_id, rank_ids, rank_count, domain_rank, window_size, alloc.get()); + uint64_t derived_ctx = 0; + int rc = comm_derive_context(h, rank_ids, rank_count, domain_rank, window_offset, window_size, &derived_ctx); if (rc != 0) return rc; + auto *derived = reinterpret_cast(derived_ctx); + auto derived_it = std::find(h->derived_contexts.begin(), h->derived_contexts.end(), derived); + if (derived_it == h->derived_contexts.end()) { + aclrtFree(derived); + LOG_ERROR("[comm rank %d] alloc_domain: derived context ownership was not recorded", h->rank); + return -1; + } + h->derived_contexts.erase(derived_it); + alloc->rank = static_cast(domain_rank); + alloc->nranks = static_cast(rank_count); + alloc->local_buf = local_slice; + alloc->arena_slice = true; + alloc->window_offset = window_offset; + alloc->window_size = window_size; + alloc->device_ctx = derived; *device_ctx_out = reinterpret_cast(alloc->device_ctx); *local_window_base_out = reinterpret_cast(alloc->local_buf); @@ -1374,12 +1119,11 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co aclError aret = aclrtFree(alloc->device_ctx); if (aret != ACL_SUCCESS && rc == 0) rc = -1; } - reset_domain_urma_workspace(*alloc); // local_buf and every peer import are VMM-mapped VAs, not aclrtMalloc // pointers: unmap + release the VA reservation, then free the physical // handle. release_domain_peer_windows(*alloc); - if (alloc->local_buf) { + if (alloc->local_buf && !alloc->arena_slice) { release_own_vmm_window(alloc->local_buf, alloc->own_handle); alloc->local_buf = nullptr; alloc->own_handle = nullptr; @@ -1441,12 +1185,10 @@ extern "C" int comm_destroy(CommHandle h) try { for (auto &kv : h->domain_allocations) { auto &alloc = kv.second; if (alloc->device_ctx) aclrtFree(alloc->device_ctx); - reset_domain_urma_workspace(*alloc); release_domain_peer_windows(*alloc); - if (alloc->local_buf) release_own_vmm_window(alloc->local_buf, alloc->own_handle); + if (alloc->local_buf && !alloc->arena_slice) release_own_vmm_window(alloc->local_buf, alloc->own_handle); } h->domain_allocations.clear(); - reset_base_urma_workspace(h); if (h->hccl_comm) { HcclResult hret = hccl_comm_destroy(h->hccl_comm); if (hret != HCCL_SUCCESS) { @@ -1454,6 +1196,9 @@ extern "C" int comm_destroy(CommHandle h) try { if (rc == 0) rc = -1; } } + // UrmaWorkspaceManager owns memory registered with the communicator and + // channel-derived state, so its Finalize must run after HCCL teardown. + reset_base_urma_workspace(h); // NOTE: we do NOT destroy h->stream — it is caller-owned. // We also do NOT call aclrtResetDevice / aclFinalize here. Device/ACL diff --git a/src/common/platform_comm/comm.h b/src/common/platform_comm/comm.h index b13be6c75b..66cd0b257a 100644 --- a/src/common/platform_comm/comm.h +++ b/src/common/platform_comm/comm.h @@ -191,6 +191,9 @@ int comm_get_window_size(CommHandle h, size_t *size_out); * device_ctx_out points to a backend-owned device CommContext that remains * valid until comm_destroy(base). * + * A5's URMA workspace is indexed in communicator-rank order; CommContext's + * rank map translates arbitrary domain subsets/reorderings to that index. + * * @param h Allocated base communicator handle. * @param rank_ids Base-communicator rank ids in domain rank order. * @param rank_count Number of domain ranks. @@ -208,22 +211,31 @@ int comm_derive_context( /** * Allocate a fresh per-rank symmetric window pool for a subset of ranks. * - * Unlike comm_alloc_windows() which allocates the single base pool once at - * bootstrap, this allocates an additional pool for a dynamically-derived - * domain (a subset of the base communicator). Multiple concurrent + * Backends may allocate an additional pool or derive a slice from the base + * communicator's persistent pool. Multiple concurrent * allocations are disambiguated by `allocation_id`, which is mixed into * every internal handshake / barrier filename so a second allocation * does not collide with the first. * * This is a collective operation across the subset only: every * participating chip must call this with matching arguments; non-members - * of the subset do NOT call it. Internal file barriers synchronise the - * subset, so the parent (orchestrator) only needs to dispatch and wait - * for completion — it does not need to broker the cross-rank handshake. + * of the subset do NOT call it. + * + * How the subset synchronises differs by backend, and only the release side + * is uniform. A2/A3 and sim rendezvous inside this call, so the caller needs + * no barrier of its own. A5 does not: an arena slice is already mapped and + * peer-visible, so the call only clears the slice and builds a context, and + * what keeps a peer from writing into a slice its owner has not yet cleared is + * the caller joining every member's call before it publishes the domain. That + * holds because one Worker process owns every A5 rank of a communicator. A + * deployment that spread those ranks across processes would need a barrier + * here, and would also need `window_offset` agreed across those processes — + * the offsets come from a free list private to the calling Worker. * - * On HCCL this performs aclrtMalloc + the same Path-D IPC pattern as - * comm_alloc_windows but on a fresh per-allocation buffer. On sim it - * shm_opens a fresh POSIX shm scoped by `allocation_id`. + * A2/A3 HCCL performs the same Path-D IPC pattern as comm_alloc_windows on a + * fresh buffer. A5 derives a slice from its communicator-lifetime arena so + * its URMA memory registration and channels cannot outlive their memory. On + * sim this shm_opens a fresh POSIX shm scoped by `allocation_id`. * * Resources allocated here remain owned by the base handle; either an * explicit comm_release_domain_windows() or final comm_destroy(base) @@ -238,6 +250,14 @@ int comm_derive_context( * domain order (length rank_count). * @param rank_count Number of subset members. * @param domain_rank This caller's dense rank in the subset. + * @param window_offset Offset in a communicator-owned persistent + * arena, for backends that derive domain views. + * Every rank of the subset must pass the same + * value, and the range must not overlap another + * live allocation. Backends whose windows are + * owned by the allocation have no arena to + * offset into and **reject** a non-zero value + * rather than ignoring it. * @param window_size Bytes per rank. Backend must allocate * exactly this size; no auto-rounding. * @param device_ctx_out Receives a device pointer to a new @@ -249,7 +269,7 @@ int comm_derive_context( */ int comm_alloc_domain_windows( CommHandle h, uint64_t allocation_id, const uint32_t *rank_ids, size_t rank_count, uint32_t domain_rank, - size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out + size_t window_offset, size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out ); /** diff --git a/src/common/platform_comm/comm_context.h b/src/common/platform_comm/comm_context.h index d9aad3be1e..43b1e04f4b 100644 --- a/src/common/platform_comm/comm_context.h +++ b/src/common/platform_comm/comm_context.h @@ -24,26 +24,39 @@ * layout is owned end-to-end by simpler. * - comm_sim.cpp: same shape, filled with malloc'd host pointers. * - * The layout is shared with pto-isa's parallel HcclDeviceContext - * declaration and must stay byte-equivalent with it. + * The leading layout through windowsOut is shared with pto-isa's parallel + * HcclDeviceContext declaration. Simpler-owned transport fields are appended + * after that compatible prefix. */ #pragma once #include #include +#include static constexpr uint32_t COMM_MAX_RANK_NUM = 64; struct CommContext { - uint64_t workSpace; - uint64_t workSpaceSize; + uint64_t sdmaWorkSpace; + uint64_t sdmaWorkSpaceSize; uint32_t rankId; uint32_t rankNum; uint64_t winSize; uint64_t windowsIn[COMM_MAX_RANK_NUM]; uint64_t windowsOut[COMM_MAX_RANK_NUM]; + + uint64_t urmaWorkSpace; + uint64_t urmaWorkSpaceSize; + // Byte displacement of this context's windowsIn[] view from the symmetric + // memory registered in urmaWorkSpace. Zero for the base context; non-zero + // for a derived arena slice. + uint64_t urmaWindowOffset; + // Map a domain-local rank to the rank used by the communicator-scoped + // URMA workspace. Base contexts contain the identity map; derived + // contexts may select/reorder any communicator ranks. + uint32_t urmaRankMap[COMM_MAX_RANK_NUM]; }; // The struct itself lives in this repo, so on the surface these asserts look @@ -52,11 +65,12 @@ struct CommContext { // this header at the same time: // // 1. The pto-isa repo carries a parallel declaration (HcclDeviceContext) -// that must be byte-equivalent to this struct -- pto-isa kernels read -// windowsIn[]/winSize/rankId via that mirror. Any insert/reorder here -// that is not matched in pto-isa silently shifts the device-side field -// offsets and corrupts MTE2 reads. The locks below pin our side; the -// pto-isa side should add its own mirror asserts. +// that must be prefix-compatible with this struct -- pto-isa kernels read +// windowsIn[]/winSize/rankId via that mirror. Field names may differ, but +// any insert/reorder before the simpler-owned tail that is not matched in +// pto-isa silently shifts +// the device-side field offsets and corrupts MTE2 reads. The locks below +// pin our side; pto-isa should add its own mirror asserts. // // 2. Device kernels (AICore / AICPU) compiled with CCEC may apply slightly // different alignment rules than host gcc. A host-side sizeof/offset @@ -65,11 +79,17 @@ struct CommContext { // Treat the numbers below as a tripwire: changing them is a deliberate act // that forces the editor to coordinate the matching change on the pto-isa // side, not a routine "oh I just added a field" edit. -static_assert(sizeof(CommContext) == 1056, "CommContext size shifted"); -static_assert(offsetof(CommContext, workSpace) == 0, "CommContext layout drift"); -static_assert(offsetof(CommContext, workSpaceSize) == 8, "CommContext layout drift"); +static_assert(std::is_trivially_copyable_v, "CommContext must remain trivially copyable"); +static_assert(std::is_standard_layout_v, "CommContext must remain standard layout"); +static_assert(sizeof(CommContext) == 1336, "CommContext size shifted"); +static_assert(offsetof(CommContext, sdmaWorkSpace) == 0, "CommContext layout drift"); +static_assert(offsetof(CommContext, sdmaWorkSpaceSize) == 8, "CommContext layout drift"); static_assert(offsetof(CommContext, rankId) == 16, "CommContext layout drift"); static_assert(offsetof(CommContext, rankNum) == 20, "CommContext layout drift"); static_assert(offsetof(CommContext, winSize) == 24, "CommContext layout drift"); static_assert(offsetof(CommContext, windowsIn) == 32, "CommContext layout drift"); static_assert(offsetof(CommContext, windowsOut) == 544, "CommContext layout drift"); +static_assert(offsetof(CommContext, urmaWorkSpace) == 1056, "CommContext layout drift"); +static_assert(offsetof(CommContext, urmaWorkSpaceSize) == 1064, "CommContext layout drift"); +static_assert(offsetof(CommContext, urmaWindowOffset) == 1072, "CommContext layout drift"); +static_assert(offsetof(CommContext, urmaRankMap) == 1080, "CommContext layout drift"); diff --git a/src/common/platform_comm/comm_sim.cpp b/src/common/platform_comm/comm_sim.cpp index 50befcec0b..dc693aa6bc 100644 --- a/src/common/platform_comm/comm_sim.cpp +++ b/src/common/platform_comm/comm_sim.cpp @@ -380,8 +380,8 @@ extern "C" int comm_alloc_windows(CommHandle h, size_t win_size, uint64_t *devic // cross-rank address-agreement assert is specifically an HCCL-GVA // invariant and is not expected to hold (nor intended to run) under sim. auto &ctx = h->host_ctx; - ctx.workSpace = 0; - ctx.workSpaceSize = 0; + ctx.sdmaWorkSpace = 0; + ctx.sdmaWorkSpaceSize = 0; ctx.rankId = static_cast(h->rank); ctx.rankNum = static_cast(h->nranks); ctx.winSize = win_size; @@ -446,7 +446,8 @@ extern "C" int comm_derive_context( ); return -1; } - if (window_offset + window_size > static_cast(h->host_ctx.winSize)) { + if (window_offset > static_cast(h->host_ctx.winSize) || + window_size > static_cast(h->host_ctx.winSize) - window_offset) { std::fprintf( stderr, "[comm_sim rank %d] comm_derive_context: window range [%zu, %zu) exceeds base window size %llu\n", h->rank, window_offset, window_offset + window_size, static_cast(h->host_ctx.winSize) @@ -456,8 +457,8 @@ extern "C" int comm_derive_context( auto *ctx = new (std::nothrow) CommContext{}; if (ctx == nullptr) return -1; - ctx->workSpace = h->host_ctx.workSpace; - ctx->workSpaceSize = h->host_ctx.workSpaceSize; + ctx->sdmaWorkSpace = h->host_ctx.sdmaWorkSpace; + ctx->sdmaWorkSpaceSize = h->host_ctx.sdmaWorkSpaceSize; ctx->rankId = domain_rank; ctx->rankNum = static_cast(rank_count); ctx->winSize = window_size; @@ -473,6 +474,7 @@ extern "C" int comm_derive_context( } ctx->windowsIn[i] = h->host_ctx.windowsIn[base_rank] + window_offset; ctx->windowsOut[i] = h->host_ctx.windowsOut[base_rank] + window_offset; + ctx->urmaRankMap[i] = base_rank; } h->derived_contexts.push_back(ctx); @@ -529,9 +531,19 @@ extern "C" int comm_barrier(CommHandle h) { extern "C" int comm_alloc_domain_windows( CommHandle h, uint64_t allocation_id, const uint32_t *rank_ids, size_t rank_count, uint32_t domain_rank, - size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out + size_t window_offset, size_t window_size, uint64_t *device_ctx_out, uint64_t *local_window_base_out ) try { if (h == nullptr || rank_ids == nullptr || device_ctx_out == nullptr || local_window_base_out == nullptr) return -1; + // Every allocation owns its shm here, so there is no arena to offset into. + // Reject a non-zero request instead of returning a window that silently + // starts somewhere else than the caller asked for. + if (window_offset != 0) { + std::fprintf( + stderr, "[comm_sim rank %d] alloc_domain: window_offset=%zu is not supported by the sim backend\n", h->rank, + window_offset + ); + return -1; + } if (rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || window_size == 0) { std::fprintf( stderr, "[comm_sim rank %d] alloc_domain: bad args (rank_count=%zu domain_rank=%u window_size=%zu)\n", @@ -634,8 +646,8 @@ extern "C" int comm_alloc_domain_windows( auto *win_base = static_cast(base) + HEADER_SIZE; alloc->host_ctx = std::make_unique(); auto &ctx = *alloc->host_ctx; - ctx.workSpace = 0; - ctx.workSpaceSize = 0; + ctx.sdmaWorkSpace = 0; + ctx.sdmaWorkSpaceSize = 0; ctx.rankId = domain_rank; ctx.rankNum = static_cast(rank_count); ctx.winSize = window_size; diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 20d66ba1bd..10e2d0ae47 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -1199,7 +1199,7 @@ uint64_t ChipWorker::comm_derive_context( std::pair ChipWorker::comm_alloc_domain_windows( uint64_t comm_handle, uint64_t allocation_id, const std::vector &rank_ids, uint32_t domain_rank, - size_t window_size + size_t window_offset, size_t window_size ) { if (comm_alloc_domain_windows_fn_ == nullptr) { throw std::runtime_error("comm_alloc_domain_windows is not supported by this runtime"); @@ -1217,7 +1217,7 @@ std::pair ChipWorker::comm_alloc_domain_windows( uint64_t local_window_base = 0; int rc = comm_alloc_domain_windows_fn_( reinterpret_cast(comm_handle), allocation_id, rank_ids.data(), rank_ids.size(), domain_rank, - window_size, &device_ctx, &local_window_base + window_offset, window_size, &device_ctx, &local_window_base ); if (rc != 0) { throw std::runtime_error("comm_alloc_domain_windows failed with code " + std::to_string(rc)); diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 0725a787ab..e9da0f66bb 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -210,7 +210,7 @@ class ChipWorker { /// participating ranks call this; non-members of the subset must not. std::pair comm_alloc_domain_windows( uint64_t comm_handle, uint64_t allocation_id, const std::vector &rank_ids, uint32_t domain_rank, - size_t window_size + size_t window_offset, size_t window_size ); /// Pair to `comm_alloc_domain_windows`: collectively free the per-rank /// pool and the device CommContext, then drop the allocation record. @@ -285,7 +285,7 @@ class ChipWorker { using CommGetWindowSizeFn = int (*)(void *, size_t *); using CommDeriveContextFn = int (*)(void *, const uint32_t *, size_t, uint32_t, size_t, size_t, uint64_t *); using CommAllocDomainWindowsFn = - int (*)(void *, uint64_t, const uint32_t *, size_t, uint32_t, size_t, uint64_t *, uint64_t *); + int (*)(void *, uint64_t, const uint32_t *, size_t, uint32_t, size_t, size_t, uint64_t *, uint64_t *); using CommReleaseDomainWindowsFn = int (*)(void *, uint64_t, size_t, uint32_t); using CommGlobalDomainPrepareFn = int (*)(uint64_t, uint32_t, uint32_t, size_t, uint32_t, CommGlobalDomainDescriptor *, uint64_t *); diff --git a/tests/ut/py/test_runtime_builder.py b/tests/ut/py/test_runtime_builder.py index 8d73e280a7..36eca19bbf 100644 --- a/tests/ut/py/test_runtime_builder.py +++ b/tests/ut/py/test_runtime_builder.py @@ -587,13 +587,14 @@ def test_a5_default_build_writes_pto_isa_metadata(self, MockCompiler, tmp_path, assert calls == [(RuntimeBuilder._LIB_DIR, "/tmp/pto-isa", ["a5/onboard/test_rt"])] @patch("simpler_setup.runtime_builder.RuntimeCompiler") - def test_a5_default_host_build_passes_pto_isa_cmake_define(self, MockCompiler, tmp_path, monkeypatch): - """a5 host ccache key includes the pinned PTO-ISA commit.""" + def test_a5_host_build_ignores_retired_urma_workspace_env(self, MockCompiler, tmp_path, monkeypatch, caplog): + """a5 always builds both communication workspaces without an opt-in define.""" from simpler_setup import pto_isa # noqa: PLC0415 from simpler_setup.runtime_builder import RuntimeBuilder # noqa: PLC0415 pin = "b" * 40 self._make_runtime(tmp_path, "a5") + monkeypatch.setenv("SIMPLER_ENABLE_PTO_URMA_WORKSPACE", "ON") mock_instance = MockCompiler.get_instance.return_value mock_instance.compile.side_effect = lambda target, *a, **kw: (Path(kw["output_dir"]) / f"lib{target}.so") monkeypatch.setattr(pto_isa, "read_pto_isa_pin", lambda: pin) @@ -606,7 +607,9 @@ def test_a5_default_host_build_passes_pto_isa_cmake_define(self, MockCompiler, t host_call = next(call for call in mock_instance.compile.call_args_list if call.args[0] == "host") assert host_call.kwargs["cmake_defines"]["SIMPLER_PTO_ISA_BUILD_COMMIT"] == pin assert "SIMPLER_ENABLE_PTO_SDMA_WORKSPACE" not in host_call.kwargs["cmake_defines"] + assert "SIMPLER_ENABLE_PTO_URMA_WORKSPACE" not in host_call.kwargs["cmake_defines"] assert host_call.kwargs["cmake_defines"]["PTO_ISA_ROOT"] == "/tmp/pto-isa" + assert "A5 onboard always builds SDMA and URMA workspaces" in caplog.text @patch("simpler_setup.runtime_builder.RuntimeCompiler") def test_sim_direct_build_does_not_write_pto_isa_metadata(self, MockCompiler, tmp_path, monkeypatch): @@ -717,11 +720,10 @@ def test_a2a3_onboard_folds_in_pto_isa_commit(self, monkeypatch): assert builder._build_cache_stamp() == "runtime_sha:pto-isa=isa_sha" def test_a5_default_folds_in_pto_isa_commit(self, monkeypatch): - """a5 default SDMA workspace folds the pto-isa pin into the cache stamp.""" + """a5 default DMA workspaces fold the pto-isa pin into the cache stamp.""" import simpler_setup.runtime_builder as rb_module # noqa: PLC0415 from simpler_setup import pto_isa # noqa: PLC0415 - monkeypatch.delenv("SIMPLER_ENABLE_PTO_URMA_WORKSPACE", raising=False) monkeypatch.setattr(rb_module, "_get_git_head", lambda _root: "runtime_sha") monkeypatch.setattr(pto_isa, "read_pto_isa_pin", lambda: "isa_sha") @@ -774,21 +776,11 @@ def test_non_a2a3_onboard_returns_empty(self, monkeypatch): def test_a5_default_reads_pin(self, monkeypatch): from simpler_setup import pto_isa # noqa: PLC0415 - monkeypatch.delenv("SIMPLER_ENABLE_PTO_URMA_WORKSPACE", raising=False) monkeypatch.setattr(pto_isa, "read_pto_isa_pin", lambda: "isa_sha") builder = self._make_builder("a5") assert builder._resolve_build_pto_isa_commit() == "isa_sha" - def test_a5_urma_overlay_on_reads_pin(self, monkeypatch): - """URMA overlay also embeds pto-isa, so it reads the pin too (#1392).""" - from simpler_setup import pto_isa # noqa: PLC0415 - - monkeypatch.setenv("SIMPLER_ENABLE_PTO_URMA_WORKSPACE", "ON") - monkeypatch.setattr(pto_isa, "read_pto_isa_pin", lambda: "isa_sha") - builder = self._make_builder("a5") - assert builder._resolve_build_pto_isa_commit() == "isa_sha" - def test_a2a3_onboard_reads_pin(self, monkeypatch): from simpler_setup import pto_isa # noqa: PLC0415 diff --git a/tests/ut/py/test_worker/test_dynamic_alloc_hw.py b/tests/ut/py/test_worker/test_dynamic_alloc_hw.py index 63728ab1b0..cebdb14610 100644 --- a/tests/ut/py/test_worker/test_dynamic_alloc_hw.py +++ b/tests/ut/py/test_worker/test_dynamic_alloc_hw.py @@ -42,8 +42,9 @@ def test_two_rank_allocate_release_round_trip(st_platform, st_device_ids): """End-to-end 2-rank hardware alloc + release round trip. - A3 performs two sequential allocations on the same base communicator to - check Fabric mapping release. A5 keeps the existing single allocation. + Performs two sequential allocations on the same base communicator. The + second reverses communicator ranks, covering the A5 URMA rank remap while + also checking that an arena slice can be released and reused. """ from simpler.task_interface import CallConfig, CommBufferSpec from simpler.worker import Worker @@ -58,11 +59,14 @@ def test_two_rank_allocate_release_round_trip(st_platform, st_device_ids): captures: list[dict[str, object]] = [] + worker_orders = [tuple(range(nranks)), tuple(reversed(range(nranks)))] + def orch_fn(orch, _args, _cfg): captured: dict[str, object] = {} + order = worker_orders[len(captures)] with orch.allocate_domain( name="tp", - workers=list(range(nranks)), + workers=list(order), window_size=4 * 1024 * 1024, buffers=[ CommBufferSpec(name="scratch", dtype="float32", count=16, nbytes=64), @@ -77,6 +81,7 @@ def orch_fn(orch, _args, _cfg): "domain_size": tp[chip_idx].domain_size, "device_ctx": int(tp[chip_idx].device_ctx), "local_window_base": int(tp[chip_idx].local_window_base), + "window_offset": int(tp[chip_idx].window_offset), "buffer_bases": {name: h.base for name, h in tp[chip_idx].buffers.items()}, } for chip_idx in tp.workers @@ -93,7 +98,7 @@ def orch_fn(orch, _args, _cfg): ) try: worker.init() - repetitions = 2 if st_platform == "a2a3" else 1 + repetitions = len(worker_orders) for _ in range(repetitions): worker.run(orch_fn, args=None, config=CallConfig()) finally: @@ -102,18 +107,22 @@ def orch_fn(orch, _args, _cfg): assert len(captures) == repetitions if repetitions == 2: assert captures[0]["alloc_id"] != captures[1]["alloc_id"] - for captured in captures: + for capture_index, captured in enumerate(captures): assert captured["released_after_with"] is True - assert captured["workers"] == tuple(range(nranks)) + expected_order = worker_orders[capture_index] + assert captured["workers"] == expected_order contexts: dict[int, dict[str, object]] = captured["contexts"] # type: ignore[assignment] - # Dense domain ranks follow worker order. - assert contexts[0]["domain_rank"] == 0 - assert contexts[1]["domain_rank"] == 1 + # Dense domain ranks follow the caller's worker order, including a + # non-prefix/reordered A5 domain. + for domain_rank, chip_idx in enumerate(expected_order): + assert contexts[chip_idx]["domain_rank"] == domain_rank for chip_idx in range(nranks): ctx = contexts[chip_idx] assert ctx["device_ctx"] != 0, f"chip {chip_idx}: device_ctx is 0" assert ctx["local_window_base"] != 0, f"chip {chip_idx}: local_window_base is 0" + if st_platform == "a5": + assert ctx["window_offset"] != 0, f"chip {chip_idx}: A5 domain did not use a derived arena slice" # Buffers are carved sequentially from the local pool. ptrs = ctx["buffer_bases"] assert isinstance(ptrs, dict) diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 4e3d3aa19c..a2c6503b5b 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -5860,6 +5860,192 @@ def _worker(): worker._worker = cast(Any, object()) return worker + def test_a5_comm_arena_reuses_and_coalesces_released_slices(self): + worker = self._worker() + worker._config = {"platform": "a5"} + arena_size = 4096 + worker._initialize_comm_arena(arena_size) + + first = worker._reserve_comm_arena(10, 300) + second = worker._reserve_comm_arena(11, 300) + assert first == worker_mod._A5_COMM_ARENA_RESERVED + assert second == first + 512 + + worker._release_comm_arena(10) + assert worker._reserve_comm_arena(12, 256) == first + worker._release_comm_arena(11) + worker._release_comm_arena(12) + assert worker._comm_arena_free == [ + ( + worker_mod._A5_COMM_ARENA_RESERVED, + arena_size - worker_mod._A5_COMM_ARENA_RESERVED, + ) + ] + + def test_a5_comm_base_uses_the_backend_reported_arena_size(self, monkeypatch): + worker = self._worker() + worker._config = {"platform": "a5", "device_ids": [0, 1]} + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + arena_size = 8192 + + def control_comm_init(chip_idx, request_name): + request = SharedMemory(name=request_name) + request_buf = request.buf + assert request_buf is not None + try: + rank, nranks, requested_size = worker_mod._COMM_INIT_HEADER.unpack_from(request_buf, 0) + assert (rank, nranks, requested_size) == (chip_idx, 2, 0) + worker_mod._COMM_INIT_HEADER.pack_into(request_buf, 0, rank, nranks, arena_size) + finally: + request_buf.release() + request.close() + + worker._worker = cast(Any, SimpleNamespace(control_comm_init=control_comm_init)) + + worker._ensure_comm_base() + + assert worker._comm_base_ready + assert worker._comm_arena_size == arena_size + assert worker._comm_arena_free == [ + ( + worker_mod._A5_COMM_ARENA_RESERVED, + arena_size - worker_mod._A5_COMM_ARENA_RESERVED, + ) + ] + + def test_a5_comm_init_child_reports_comm_get_window_size(self): + rootinfo_path = "/tmp/comm-rootinfo" + path_bytes = rootinfo_path.encode() + b"\x00" + request = SharedMemory(create=True, size=worker_mod._COMM_INIT_HEADER.size + len(path_bytes)) + request_buf = request.buf + assert request_buf is not None + worker_mod._COMM_INIT_HEADER.pack_into(request_buf, 0, 1, 2, 0) + path_start = worker_mod._COMM_INIT_HEADER.size + request_buf[path_start : path_start + len(path_bytes)] = path_bytes + mailbox = memoryview(bytearray(worker_mod._OFF_ARGS + worker_mod._CTRL_SHM_NAME_BYTES)) + encoded_name = request.name.encode() + mailbox[worker_mod._OFF_ARGS : worker_mod._OFF_ARGS + len(encoded_name)] = encoded_name + calls = [] + chip_worker = SimpleNamespace( + comm_init=lambda rank, nranks, path: calls.append(("init", rank, nranks, path)) or 17, + comm_alloc_windows=lambda handle, size: calls.append(("alloc", handle, size)) or 23, + comm_get_window_size=lambda handle: calls.append(("size", handle)) or 8192, + ) + + try: + worker_mod._handle_ctrl_comm_init(cast(Any, chip_worker), mailbox, "a5") + + assert worker_mod._COMM_INIT_HEADER.unpack_from(request_buf, 0) == (1, 2, 8192) + assert calls == [("init", 1, 2, rootinfo_path), ("alloc", 17, 0), ("size", 17)] + assert chip_worker._comm_base_handle_cached == 17 + assert chip_worker._comm_base_windows_ready is True + finally: + mailbox.release() + request_buf.release() + request.close() + request.unlink() + + def test_a5_comm_base_forwards_the_configured_arena_size(self, monkeypatch): + worker = self._worker() + requested = 1 << 20 + worker._config = {"platform": "a5", "device_ids": [0, 1], "comm_arena_size": requested} + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + granted = requested + 4096 # the backend rounds up to VMM granularity + seen: list[int] = [] + + def control_comm_init(chip_idx, request_name): + request = SharedMemory(name=request_name) + request_buf = request.buf + assert request_buf is not None + try: + rank, nranks, requested_size = worker_mod._COMM_INIT_HEADER.unpack_from(request_buf, 0) + seen.append(int(requested_size)) + worker_mod._COMM_INIT_HEADER.pack_into(request_buf, 0, rank, nranks, granted) + finally: + request_buf.release() + request.close() + + worker._worker = cast(Any, SimpleNamespace(control_comm_init=control_comm_init)) + + worker._ensure_comm_base() + + assert seen == [requested, requested] + # The allocator sizes itself from what the backend granted, not from + # what was asked for. + assert worker._comm_arena_size == granted + + def test_a5_comm_base_rejects_a_negative_arena_size(self, monkeypatch): + worker = self._worker() + worker._config = {"platform": "a5", "device_ids": [0], "comm_arena_size": -1} + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + worker._worker = cast(Any, SimpleNamespace(control_comm_init=lambda *_a: None)) + + with pytest.raises(ValueError, match="comm_arena_size must be non-negative"): + worker._ensure_comm_base() + + def test_a5_comm_init_child_forwards_the_requested_arena_size(self): + rootinfo_path = "/tmp/comm-rootinfo" + path_bytes = rootinfo_path.encode() + b"\x00" + request = SharedMemory(create=True, size=worker_mod._COMM_INIT_HEADER.size + len(path_bytes)) + request_buf = request.buf + assert request_buf is not None + worker_mod._COMM_INIT_HEADER.pack_into(request_buf, 0, 1, 2, 4096) + path_start = worker_mod._COMM_INIT_HEADER.size + request_buf[path_start : path_start + len(path_bytes)] = path_bytes + mailbox = memoryview(bytearray(worker_mod._OFF_ARGS + worker_mod._CTRL_SHM_NAME_BYTES)) + encoded_name = request.name.encode() + mailbox[worker_mod._OFF_ARGS : worker_mod._OFF_ARGS + len(encoded_name)] = encoded_name + calls = [] + chip_worker = SimpleNamespace( + comm_init=lambda rank, nranks, path: calls.append(("init", rank, nranks, path)) or 17, + comm_alloc_windows=lambda handle, size: calls.append(("alloc", handle, size)) or 23, + comm_get_window_size=lambda handle: calls.append(("size", handle)) or 8192, + ) + + try: + worker_mod._handle_ctrl_comm_init(cast(Any, chip_worker), mailbox, "a5") + + assert ("alloc", 17, 4096) in calls + # The reply carries what the backend granted, not the request. + assert worker_mod._COMM_INIT_HEADER.unpack_from(request_buf, 0) == (1, 2, 8192) + finally: + mailbox.release() + request_buf.release() + request.close() + request.unlink() + + def test_a5_arena_exhaustion_names_the_size_knob(self): + worker = self._worker() + worker._config = {"platform": "a5"} + worker._initialize_comm_arena(4096) + + with pytest.raises(MemoryError, match=r"comm_arena_size"): + worker._reserve_comm_arena(1, 1 << 20) + + def test_non_a5_domain_does_not_consume_the_a5_arena(self): + worker = self._worker() + worker._config = {"platform": "a2a3"} + + assert worker._reserve_comm_arena(10, 4096) == 0 + assert worker._comm_arena_allocations == {} + + def test_a5_precommit_failure_returns_its_arena_slice(self, monkeypatch): + worker = self._worker() + worker._config = {"platform": "a5", "device_ids": [0]} + worker._initialize_comm_arena(8192) + worker._building_run_resources = worker_mod._RunResources() + monkeypatch.setattr(worker, "_ensure_comm_base", lambda: None) + + def fail_staging(*_args, **_kwargs): + raise RuntimeError("staging failed") + + monkeypatch.setattr(worker_mod, "_run_with_owned_shared_memory", fail_staging) + + with pytest.raises(RuntimeError, match="staging failed"): + worker._allocate_domain(name="d", workers=(0,), window_size=4096, buffers=[]) + + assert worker._comm_arena_allocations == {} + def test_a_partial_domain_allocation_keeps_its_original_release_ranks(self, monkeypatch): """Two chips of three committed a window and no handle exists. @@ -5935,8 +6121,8 @@ def test_the_child_publishes_its_commit_before_anything_that_can_fail(self): try: req_buf = cast(Any, request.buf) worker_mod._DOMAIN_REQ_HEADER.pack_into( - req_buf, 0, 7, 1, 0, 64, 1 - ) # allocation_id, rank_count, domain_rank, window_size, buffer_count + req_buf, 0, 7, 1, 0, 0, 64, 1 + ) # allocation_id, rank_count, domain_rank, window_offset, window_size, buffer_count # One buffer larger than the window: the carve raises after the # collective has already committed. struct.pack_into(" int: + """Validate A5 URMA view translation and communicator-rank remapping.""" + derived_offset = 256 + derived_size = 1024 + derived_ctx_ptr = worker.comm_derive_context(comm, list(range(nranks)), rank, derived_offset, derived_size) + derived_ctx = _CommContext() + worker.copy_from(ctypes.addressof(derived_ctx), derived_ctx_ptr, ctypes.sizeof(derived_ctx)) + for peer in range(nranks): + expected = int(host_ctx.windowsIn[peer]) + derived_offset + if derived_ctx.windowsIn[peer] != expected: + raise AssertionError( + f"derived windowsIn[{peer}]=0x{derived_ctx.windowsIn[peer]:x}, expected 0x{expected:x}" + ) + if platform != "a5": + return int(derived_ctx.urmaWindowOffset) + if derived_ctx.urmaWorkSpace != host_ctx.urmaWorkSpace: + raise AssertionError("derived context did not inherit the base URMA workspace") + if derived_ctx.urmaWindowOffset != derived_offset: + raise AssertionError(f"urmaWindowOffset={derived_ctx.urmaWindowOffset}, expected {derived_offset}") + reversed_ranks = list(reversed(range(nranks))) + reversed_ctx_ptr = worker.comm_derive_context(comm, reversed_ranks, nranks - 1 - rank, derived_offset, derived_size) + reversed_ctx = _CommContext() + worker.copy_from(ctypes.addressof(reversed_ctx), reversed_ctx_ptr, ctypes.sizeof(reversed_ctx)) + if list(reversed_ctx.urmaRankMap[:nranks]) != reversed_ranks: + raise AssertionError(f"urmaRankMap={list(reversed_ctx.urmaRankMap[:nranks])}, expected {reversed_ranks}") + return int(derived_ctx.urmaWindowOffset) def _rank_entry( @@ -75,6 +108,7 @@ def _rank_entry( nranks: int, device_id: int, bins, + platform: str, rootinfo_path: str, result_queue: mp.Queue, # type: ignore[type-arg] ) -> None: @@ -136,6 +170,13 @@ def _rank_entry( result["local_base"] = int(local_base) result["rank_id"] = int(host_ctx.rankId) result["rank_num"] = int(host_ctx.rankNum) + result["sdma_workspace"] = int(host_ctx.sdmaWorkSpace) + result["urma_workspace"] = int(host_ctx.urmaWorkSpace) + + # A derived context shifts its visible windows while URMA remains + # registered against the base symmetric window. The explicit offset + # lets the kernel translate a domain-local offset back into that MR. + result["derived_urma_window_offset"] = _check_derived_context(worker, comm, host_ctx, rank, nranks, platform) # Barrier. The C++ HCCL UT observed CANN error 507018 here on some # builds; that bug is tracked independently. Surface the failure to @@ -189,6 +230,7 @@ def test_two_rank_comm_lifecycle(st_platform, st_device_ids): nranks, int(st_device_ids[rank]), bins, + st_platform, rootinfo_path, result_queue, ), @@ -218,6 +260,10 @@ def test_two_rank_comm_lifecycle(st_platform, st_device_ids): r = results_by_rank[rank] if not r.get("ok"): pytest.fail(f"rank {rank} failed at stage {r.get('stage')!r}:\n{r.get('error', '(no traceback)')}") + if st_platform == "a5": + assert r["sdma_workspace"] != 0, f"rank {rank} has no SDMA workspace" + assert r["urma_workspace"] != 0, f"rank {rank} has no URMA workspace" + assert r["derived_urma_window_offset"] == 256 # Each rank's own-slot invariant (windowsIn[rank] == local_base) is # asserted inside _rank_entry; all peer slots are already checked to be