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