From 277a9e26cf844955e9c3072848678581c9e2af6e Mon Sep 17 00:00:00 2001 From: Kur0x Date: Wed, 12 Aug 2026 17:21:10 +0800 Subject: [PATCH] perf(moonep): replace reduce grad with owner pull --- .../2026-08-11-reduce-grad-owner-pull.md | 577 +++++++ ...026-08-11-reduce-grad-owner-pull-design.md | 543 ++++++ .../moonep_torch/tilexr_moonep/abi.py | 47 +- .../moonep_torch/tilexr_moonep/compat.py | 33 +- .../moonep_torch/tilexr_moonep/runtime.py | 181 +- .../moonep_torch/tilexr_moonep/torch_api.py | 165 +- src/comm/comm_wrap.cpp | 36 + src/comm/tilexr_comm.cpp | 37 + src/comm/tilexr_comm.h | 3 + src/comm/udma/tilexr_udma_context.cpp | 269 ++- src/comm/udma/tilexr_udma_context.h | 10 + src/comm/udma/tilexr_udma_transport.cpp | 406 +++-- src/comm/udma/tilexr_udma_transport.h | 13 + src/include/tilexr_api.h | 9 + src/include/tilexr_moonep.h | 70 +- src/include/tilexr_udma.h | 179 +- src/include/tilexr_udma_reg.h | 177 ++ .../host/prefetch_weight_layout.cpp | 6 +- src/moonep/reduce_grad/CMakeLists.txt | 4 +- .../reduce_grad/common/reduce_grad_common.h | 174 +- .../reduce_grad/host/reduce_grad_host.cpp | 462 +++-- .../reduce_grad/host/reduce_grad_host.h | 41 +- .../reduce_grad/host/reduce_grad_launch.cpp | 79 +- .../reduce_grad/host/reduce_grad_layout.cpp | 259 ++- .../reduce_grad/host/reduce_grad_layout.h | 6 +- .../tilexr_moonep_reduce_grad_kernel.cpp | 1484 +++++++++++------ tests/moonep/CMakeLists.txt | 5 +- tests/moonep/demo/run_a5.sh | 4 +- tests/moonep/demo/tilexr_moonep_flow_demo.cpp | 123 +- tests/moonep/python/fakes.py | 99 +- tests/moonep/python/test_ffi_unittest.py | 83 +- .../python/test_reduce_grad_benchmark.py | 82 + .../python/test_tilexr_correctness_adapter.py | 17 +- tests/moonep/python/test_upstream_compat.py | 15 +- tests/moonep/python/unittest_smoke.py | 35 +- .../unit/test_tilexr_moonep_abi_layout.cpp | 38 +- .../moonep/unit/test_tilexr_moonep_c_header.c | 8 + .../test_tilexr_moonep_kernel_sources.cpp | 63 +- ...est_tilexr_moonep_prefetch_weight_host.cpp | 12 + .../test_tilexr_moonep_reduce_grad_host.cpp | 587 +++---- .../test_tilexr_moonep_reduce_grad_launch.cpp | 169 +- .../test_tilexr_moonep_reduce_grad_layout.cpp | 247 +-- .../unit/test_tilexr_moonep_sources.cpp | 12 +- tests/udma/CMakeLists.txt | 87 + tests/udma/build.sh | 5 + tests/udma/demo/README.md | 26 + .../demo/run_tilexr_udma_profile_probe_mpi.sh | 181 ++ tests/udma/demo/tilexr_udma_profile_probe.cpp | 767 +++++++++ .../demo/tilexr_udma_profile_probe_kernel.cpp | 165 ++ .../unit/test_tilexr_udma_demo_sources.cpp | 35 + .../udma/unit/test_tilexr_udma_device_api.cpp | 94 ++ tests/udma/unit/test_tilexr_udma_registry.cpp | 99 ++ .../unit/test_tilexr_udma_source_guard.cpp | 79 +- tools/moonep/reduce_grad_benchmark.py | 801 +++++++++ tools/moonep/tilexr_backend.py | 20 - 55 files changed, 7411 insertions(+), 1817 deletions(-) create mode 100644 docs/plans/2026-08-11-reduce-grad-owner-pull.md create mode 100644 docs/specs/2026-08-11-reduce-grad-owner-pull-design.md create mode 100644 tests/moonep/python/test_reduce_grad_benchmark.py create mode 100644 tests/udma/demo/run_tilexr_udma_profile_probe_mpi.sh create mode 100644 tests/udma/demo/tilexr_udma_profile_probe.cpp create mode 100644 tests/udma/demo/tilexr_udma_profile_probe_kernel.cpp create mode 100644 tools/moonep/reduce_grad_benchmark.py diff --git a/docs/plans/2026-08-11-reduce-grad-owner-pull.md b/docs/plans/2026-08-11-reduce-grad-owner-pull.md new file mode 100644 index 00000000..f7e2db15 --- /dev/null +++ b/docs/plans/2026-08-11-reduce-grad-owner-pull.md @@ -0,0 +1,577 @@ +# ReduceGrad Owner-Pull Implementation Plan + +Date: 2026-08-11 + +## Classification + +Durable plan. Retain this document with the implementation because it records +cross-module UDMA ownership, compatibility boundaries, hardware gates, and the +performance evidence required for future maintenance. + +Execution status on 2026-08-12: Tasks 1 through 11 produced a correct four-card +owner-pull implementation and reproducible benchmark/probe tooling. The first +explicit 3-QP control measured `2764.82/2843.38 us`, but latest `main` uses a +fixed 32-QP shared domain. Separating that transport domain from three active +lanes and mapping them to physical QPs `{0, 1, 16}` reduced P50 to +`1536-1562 us` in three `20 x 50` runs, versus retained pinned-native P50 +`2724-2737 us`. This clears the numerical 3% gate against the retained +baseline. The final TileXR and retained native artifacts used different +`torch_npu` builds, so a strict same-runtime native rebuild and interleaved A/B +is the immediate next validation step. Sender-PUT is conditional future work +and is not implemented in this patch. + +## Goal and Scope + +Implement the approved design in +`docs/specs/2026-08-11-reduce-grad-owner-pull-design.md`: + +- replace TileXR MoonEP ReduceGrad with persistent-profile, large-chunk UDMA + owner-pull on Ascend950; +- fuse gate, up, and down in one kernel and one cross-rank barrier; +- remove hot-path synchronization, reduce-buffer tail copies, registration + switching, sender/receiver roles, UDMA push, per-chunk acknowledgement GETs, + and the old peer transport from ReduceGrad; +- preserve upstream Python behavior and ordered FP32 accumulation; +- prove steady-state P50 and P99 are reproducibly lower than native Ascend + MoonEP commit `a49538a45e5c5bdc82aa6ae02548f99e72ec67eb`. + +The initial workspace is `D:/1-项目/TileXR/TileXR`. Remote validation uses +both `root@141.61.53.106` and `root@141.61.53.110` under the same unique path: + +`/tmp/TileXR-reducegrad-20260811-44dc37b` + +## Non-Goals + +- Do not change CANN 9.1 or C++14 compatibility. +- Do not compile or link anything from `reference/` into TileXR. +- Do not redesign Dispatch, PrefetchWeight, Combine, EP, collectives, SDMA, or + ordinary peer-memory transports except for the minimum profile API needed to + let their existing registration behavior coexist with ReduceGrad. +- Do not add a 910B, peer-window, UDMA-push, or host collective ReduceGrad + fallback. +- Do not claim UDMA data-plane validation from local unit tests, simulator, or + 910B hardware. +- Do not include one-time allocation, MR registration/import, or binary + registration in steady-state latency. + +## Authoritative References + +- Approved design: + `docs/specs/2026-08-11-reduce-grad-owner-pull-design.md` +- Repository validation guidance: `docs/BUILD_VERIFICATION.md` +- Existing registered-memory design: + `docs/specs/2026-08-06-udma-registered-memory-multi-qp-design.md` +- UDMA Host lifecycle: `src/comm/udma/tilexr_udma_context.cpp` and + `src/comm/udma/tilexr_udma_transport.cpp` +- Device WQE path: `src/include/tilexr_udma.h` and + `src/include/tilexr_udma_types.h` +- Public communicator ABI: `src/include/tilexr_api.h` +- MoonEP public ABI: `src/include/tilexr_moonep.h` +- Current code to replace: `src/moonep/reduce_grad/` +- Torch integration: `integrations/moonep_torch/tilexr_moonep/` +- Native baseline kernel: + `reference/ascend-moonep-dev/kernels/moonep_grad_reduce.cpp` +- Native baseline benchmark: + `reference/ascend-moonep-dev/benchmarks/bench_grad_reduce.py` +- Installed target evidence: + `/usr/local/Ascend/cann-9.1.T560/include/pto/npu/comm/async/urma/` + +## Global Constraints + +- Preserve ordinary `TileXRUDMARegister` behavior for existing consumers. +- Profile registration is collective, forbidden in `InitThread`, and must have + explicit quiescent-lifetime rules. +- UDMA WQEs are built entirely in UB, copied to SQ with MTE3, and followed by a + precise MTE3-to-scalar dependency before an `st_dev` doorbell. +- One AIV leader owns each QP. No two blocks update the same SQ/CQ head or tail. +- The new kernel applies contributions in ascending `source_rank * B + slot` + order. Prefetch may finish out of order; FP32 additions may not. +- Shared synchronization flags use `TileXRCommNextMagic`; they are never reset. +- Keep all workspace arithmetic in checked 64-bit Host code and keep each WQE + transfer length within the 32-bit SGE field. +- Preserve unrelated user files and changes in the working tree. + +## Task 1: Add Persistent UDMA Profile Contracts and Host Unit Tests + +**Objective and role:** Define the smallest reusable contract that can keep a +ReduceGrad staging MR and three source MRs registered simultaneously while +binding per-QP local and remote region roles. This is the foundation for every +later task. + +**Background and prerequisites:** The current runtime has one active +`RegistrationState`, one device registry, and a handle fixed to zero. Existing +`UDMAWQCtx.localTokenId` and per-peer/per-QP `UDMAMemInfo` already permit +different local and remote registrations per QP. + +**Modification scope:** + +- `src/include/tilexr_api.h` +- `src/include/tilexr_udma_reg.h` +- `src/comm/tilexr_comm.h` +- `src/comm/tilexr_comm.cpp` +- `src/comm/comm_wrap.cpp` +- `src/comm/udma/tilexr_udma_context.h` +- `src/comm/udma/tilexr_udma_context.cpp` +- `src/comm/udma/tilexr_udma_transport.h` +- `src/comm/udma/tilexr_udma_transport.cpp` +- focused files under `tests/udma/unit/` + +**Work:** + +1. Add bounded profile descriptors: local regions, region count, and one + `(local_region, remote_region)` binding per QP. +2. Add collective register/unregister/query APIs returning nonzero persistent + handles and explicit device `UDMAInfo`/registry views. +3. Generalize registration state so legacy active registration and persistent + profiles can coexist and be cleaned independently. +4. Register every local MR on each needed HCCP context, exchange every remote + MR descriptor/key, import remote regions, and build per-QP device metadata + from the declared bindings. +5. Make partial failure rollback deterministic and preserve cleanup-pending + state when HCCP cleanup itself fails. +6. Keep legacy handle-zero behavior and `CommArgs` publication unchanged. + +**Constraints and non-goals:** Do not make profile selection a mutable global +operation. Do not add dynamic allocation in device code. Do not widen unrelated +registry consumers to multi-region semantics. + +**Acceptance and verification:** + +- Unit tests cover invalid region counts, invalid QP bindings, overflow, + partial registration/import failure, independent cleanup, legacy coexistence, + and exact per-QP token/remote metadata selection. +- Existing UDMA registration and source-guard tests remain green. +- Run: + + `cmake -S . -B build -DTILEXR_BUILD_TESTS=ON` + + `cmake --build build -j` + + `ctest --test-dir build -R 'udma' --output-on-failure` + +**Artifacts and downstream interfaces:** A profile handle plus immutable Host +and device views consumed by Tasks 2, 4, and 6. + +## Task 2: Add Deferred UDMA READ Batching + +**Objective and role:** Let one QP leader post all contributor GET WQEs, ring +one doorbell, and wait one ordered completion frontier. + +**Background and prerequisites:** Depends on Task 1's profile device view. The +existing PUT path has deferred publication helpers, while GET always rings its +doorbell immediately. + +**Modification scope:** + +- `src/include/tilexr_udma.h` +- `tests/udma/unit/test_tilexr_udma_device_api.cpp` +- `tests/udma/unit/test_tilexr_udma_source_guard.cpp` + +**Work:** + +1. Add a profile-aware registered-range and remote-address accessor. +2. Add deferred `UDMAGetNbiOnQp` that publishes a READ WQE without ringing. +3. Return or expose the submitted completion frontier needed for one bounded + CQ poll after a batch. +4. Reuse the existing doorbell and quiet primitives without scalar SQ stores. +5. Add exact validation for zero length, null pointers, 32-bit length, QP + ownership, and registered local/remote ranges. + +**Constraints and non-goals:** Do not change legacy immediate GET semantics. +Do not allow UB addresses as local SGE addresses. + +**Acceptance and verification:** Unit tests prove N deferred GETs advance the SQ +without ringing, one flush rings the final head, and one quiet consumes the +expected CQ frontier. Source guards verify MTE3 publication precedes `st_dev`. + +**Artifacts and downstream interfaces:** Device functions used only by a +single owner block per QP in Task 5. + +## Task 3: Build the Multi-MR Hardware Proof + +**Objective and role:** Prove the approved architecture on CANN 9.1.T560 and +Ascend950 before replacing the operator. + +**Background and prerequisites:** Depends on Tasks 1 and 2. This is the design's +mandatory stop gate. + +**Modification scope:** + +- a focused hardware probe under `tests/udma/demo/` or the existing UDMA demo + location; +- minimal CMake/script wiring for that probe; +- no ReduceGrad production code yet. + +**Work:** + +1. Allocate four independent 2 MiB-aligned device regions per rank: staging + plus three patterned sources. +2. Register one persistent profile with at least three QPs, binding staging as + every local region and one distinct remote source region per QP. +3. On two ranks, issue 48 KiB, 256 KiB, 1 MiB, 2 MiB, 4 MiB, 8 MiB, and 16 MiB + deferred READ batches through every binding. +4. Verify byte-exact destination data and independent profile/legacy-region + lifetime. +5. Record transfer and combined staging-consumption timing in machine-readable + output. + +**Remote execution:** Create the common directory on both hosts, synchronize +the same source snapshot, source `scripts/common_env.sh`, and build with +CANN `/usr/local/Ascend/cann-9.1.T560`. Do not exclude cards solely because +`npu-smi` reports `Alarm`. + +**Acceptance and verification:** All three remote region bindings transfer +correct data through one shared QP set, a legacy registration remains usable, +and unregistering either handle does not invalidate the other. + +**Stop condition:** If HCCP rejects simultaneous MRs, the target cannot bind +the selected local/remote tokens, or shared QP state is not independently safe, +stop implementation and return to design review. Do not start Task 4 and do not +restore or optimize the old ReduceGrad. + +**Artifacts and downstream interfaces:** Raw JSON/CSV size sweep, target logs, +and the checked-in chunk-size hypothesis used by Task 4. + +## Task 4: Replace ReduceGrad Layout, Host Contract, and Launch ABI + +**Objective and role:** Remove the old transport model from Host code and +describe only the owner-pull workspace/profile/kernel contract. + +**Background and prerequisites:** Depends on a passing Task 3. The public +upstream Python signature stays stable, but the native internal ABI must carry +the three reduce-buffer sources and persistent profile view. + +**Modification scope:** + +- `src/include/tilexr_moonep.h` +- `src/moonep/reduce_grad/common/reduce_grad_common.h` +- `src/moonep/reduce_grad/host/reduce_grad_layout.h` +- `src/moonep/reduce_grad/host/reduce_grad_layout.cpp` +- `src/moonep/reduce_grad/host/reduce_grad_host.h` +- `src/moonep/reduce_grad/host/reduce_grad_host.cpp` +- `src/moonep/reduce_grad/host/reduce_grad_launch.cpp` +- ReduceGrad Host/layout/ABI unit tests under `tests/moonep/unit/` + +**Work:** + +1. Delete peer/UDMA transport selection, sender control counts, outbound and + inbound offsets, ACK state, and the 1 MiB threshold from ReduceGrad layout. +2. Define checked workspace layout for per-QP lane state and two banks of + `rank_size * chunk_bytes` payload staging. +3. Require A5 UDMA and at least one QP per projection for multi-rank launch. +4. Allocate QPs to projections deterministically, initially proportional to + projection bytes with stable tie-breaking. +5. Add the three local reduce-buffer slice descriptors and persistent profile + handle/device pointers to preparation and launch context. +6. Keep one-time preparation separate from hot launch. Validate pointer identity + and reject stale preparation without synchronizing or registering. +7. Launch only the replacement pure AIV binary through + `rtKernelLaunchWithFlagV2`. + +**Constraints and non-goals:** Existing C symbols may delegate to the new +engine for compatibility, but no entry can launch or contain the old protocol. +Do not add topology-only transport selection. + +**Acceptance and verification:** Unit tests cover the minimum-rank rejection and workspace arithmetic at 4/8/16 +ranks, QP counts, all projection sizes, alignment, pointer mismatch, unsupported +capability, profile mismatch, and launch argument layout. Source tests assert +old protocol identifiers and acknowledgement paths are absent. + +**Artifacts and downstream interfaces:** One immutable launch context and +kernel argument block for Task 5; Torch-facing preparation inputs for Task 6. + +## Task 5: Replace the Ascend C ReduceGrad Kernel + +**Objective and role:** Implement fused gate/up/down owner-pull with batched +large UDMA reads, helper-group accumulation, one global barrier, and parallel +clear. + +**Background and prerequisites:** Depends on Tasks 2 and 4. The current +`tilexr_moonep_reduce_grad_kernel.cpp` is replaced rather than incrementally +retained. + +**Modification scope:** + +- `src/moonep/reduce_grad/kernels/tilexr_moonep_reduce_grad_kernel.cpp` +- minimal common headers owned by ReduceGrad; +- kernel source/compile guards in `tests/moonep/unit/`. + +**Work:** + +1. Implement deterministic projection-to-lane work indexing over + `(local_expert, row_chunk)`. +2. Let one leader block own each QP and batch remote contributors in flattened + plan order into ping/pong GM banks. +3. Support contributor waves when a plan contains more contributors than one + rank-sized bank can represent. +4. Pipeline next-bank UDMA with current-bank helper computation using + magic/epoch-tagged lane flags and bounded cache maintenance. +5. Partition UB subtiles across helper blocks. Load output once, add all + contributions in order, and store once per subtile and wave. +6. Use accumulator plus ping/pong input UB buffers and precise MTE2/vector/MTE3 + events. Remove broad barriers from inner copy/add loops. +7. Locally synchronize, execute a one-leader-per-rank TileXR barrier, locally + release workers, and clear three local live slot sets in parallel. +8. Publish the first device failure to status and ensure every block still + reaches required local barriers to avoid deadlock. + +**Constraints and non-goals:** No sender role, UDMA PUT, per-chunk ACK, 8-byte +poll GET, peer packed-record path, or scalar/direct-GM SQ/doorbell writes. No +early return before a collective participant boundary. + +**Acceptance and verification:** + +- Bisheng compiles the kernel for Ascend950 with CANN 9.1.T560. +- Source guards prove only leaders touch QPs and the deleted protocol is absent. +- Hardware correctness passes sparse, mixed, heavy, full, empty-local, tail, + repeated-magic, and multi-wave plans at 4 ranks. +- Device status is zero and no CQ error, timeout, vector exception, hang, or + stale flag appears. + +**Artifacts and downstream interfaces:** Embedded replacement AIV binary with +the Task 4 argument contract. + +## Task 6: Remove Torch Hot-Path Copies and Registration Switching + +**Objective and role:** Ensure the measured Python `reduce_grad` interval +contains only stream-ordered preparation already enqueued by the caller and the +new kernel launch. + +**Background and prerequisites:** Depends on Tasks 1 and 4. Current code copies +`reduce_buffer[rank]` into full-gradient tail rows and synchronizes before +re-confirming active registration on every call. + +**Modification scope:** + +- `integrations/moonep_torch/tilexr_moonep/abi.py` +- `integrations/moonep_torch/tilexr_moonep/runtime.py` +- `integrations/moonep_torch/tilexr_moonep/torch_api.py` +- compatibility adapter only where required by the unchanged public signature; +- focused Python tests under `tests/moonep/python/`. + +**Work:** + +1. Pass the three reduce buffers directly through the native ABI. +2. Key preparation by allocation identity, local slice pointer, size, shape, + plan dimensions, QP mapping, and chunk size. +3. Register the persistent ReduceGrad profile once during preparation and retain + all source/staging owners until completion or explicit re-prepare/close. +4. Remove the full-gradient tail copy and post-launch legacy zeroing from the + native path. +5. Remove device synchronize and active-region registration calls from every + prepared hot launch. +6. Require quiescence only when pointer identity changes, a profile is destroyed, + or the runtime closes. +7. Preserve asynchronous event behavior, in-flight ownership, idempotent close, + and device-status propagation. + +**Constraints and non-goals:** Do not hide first-use registration inside timed +benchmark iterations. Do not weaken tensor/device/contiguity validation. + +**Acceptance and verification:** Python unit tests prove one preparation causes +one profile registration, repeated launches cause none, changed pointers require +re-prepare, hot launch does not synchronize, and all ownership/error paths +release resources exactly once. + +**Artifacts and downstream interfaces:** Stable upstream `Buffer.reduce_grad` +behavior with an explicitly preparable steady-state path. + +## Task 7: Add Correctness and Native-Baseline Benchmark Harnesses + +**Objective and role:** Provide identical, auditable correctness and timing +boundaries for TileXR and the pinned native baseline. + +**Background and prerequisites:** Can begin after Task 4's ABI is stable and is +completed after Tasks 5 and 6. + +**Modification scope:** + +- focused additions under `tools/moonep/`; +- focused tests under `tests/moonep/python/`; +- `scripts/README.md` or MoonEP validation documentation for commands; +- no modifications under `reference/ascend-moonep-dev` beyond ignored build + products. + +**Work:** + +1. Add an isolated ReduceGrad runner that can invoke native or TileXR with the + same generated plan and tensors. +2. Prepare and register before warmup; use device events around only the three + native launches or one fused TileXR launch. +3. Gather every iteration's latency and compute the cross-rank maximum before + P50/P99 aggregation. +4. Emit raw JSON including commit IDs, CANN/driver/device metadata, dimensions, + plan density, warmup, iterations, chunk, QPs, block layout, and one-time + preparation latency. +5. Add exact mutation checks for owned rows, live local slots, unused slots, + non-local slots, and ordered FP32 results. +6. Cover native dedicated 3584x3072 cases and the 7168x2048 E=384/B=48 primary + case. + +**Constraints and non-goals:** No host wall-clock comparisons for the pass gate. +No rank-zero-only timing. No different plan, initialization, or synchronization +between implementations. + +**Acceptance and verification:** Unit-test generation/statistics locally, then +run correctness at 4 ranks. Raw results must make timing boundaries and +excluded one-time work explicit. + +**Artifacts and downstream interfaces:** Reproducible benchmark command and raw +result schema consumed by Task 9. + +## Task 8: Run Local Build, Unit, ABI, and Source Verification + +**Objective and role:** Catch host, layout, lifecycle, source-rule, and Python +regressions before consuming hardware time. + +**Background and prerequisites:** Depends on Tasks 1, 2, 4, 5, 6, and 7. + +**Modification scope:** Tests and fixes within the files already owned by prior +tasks. No unrelated refactor. + +**Verification:** + +1. Initialize submodules and source `scripts/common_env.sh` where the local + environment supports CANN. +2. Configure with MoonEP and tests enabled. +3. Build and install. +4. Run all UDMA and MoonEP unit/source/ABI tests. +5. Run focused Python unit tests for FFI, Torch preparation, compatibility, and + benchmark statistics. +6. Run `git diff --check` and inspect the complete scoped diff. + +If the Windows workspace lacks CANN, run host-independent Python/source tests +locally and perform the full build on both target hosts. Report that boundary +without treating it as hardware proof. + +**Acceptance and verification:** All runnable tests pass; any test requiring +unavailable local CANN is listed and then passed remotely in Task 9. + +**Artifacts and downstream interfaces:** A single synchronized source snapshot +ready for remote validation. + +## Task 9: Build and Validate on Both Ascend950 Hosts + +**Objective and role:** Establish target-matched correctness before tuning. + +**Background and prerequisites:** Depends on Task 8. Use the exact same source +snapshot and remote path on both machines. + +**Remote work:** + +1. Resolve and verify `/tmp/TileXR-reducegrad-20260811-44dc37b` on each host, + then create it and synchronize the workspace while excluding `.git`, local + build output, ignored native build output, and unrelated untracked archives. +2. Initialize submodules or synchronize their required contents. +3. Source `scripts/common_env.sh`; verify CANN 9.1.T560, Bisheng, driver, and NPU + visibility. +4. Configure, build, install, and run focused UDMA/MoonEP tests on both hosts. +5. Run the hardware profile probe. +6. Run isolated correctness at four ranks on physical devices 0-3 of host 106. + Retain the physical device mapping with the results. + +**Constraints and non-goals:** Preserve both hosts' unrelated `/tmp` contents. +Do not treat `Alarm` alone as an unhealthy card. Do not claim 8/16-rank proof +from the required 4-rank run. + +**Acceptance and verification:** Both hosts build the same source; all focused +tests pass; 4-rank exact correctness passes; repeated runs do not hang or +consume stale flags. + +**Artifacts and downstream interfaces:** Build logs, environment metadata, +correctness JSON, and a stable candidate for performance tuning. + +## Task 10: Tune, Profile, and Enforce the Faster-Than-Native Gate + +**Objective and role:** Select measured defaults and prove the user-visible +performance requirement rather than merely improving the old implementation. + +**Background and prerequisites:** Depends on Task 9 correctness. + +**Work:** + +1. Build native MoonEP at the pinned baseline commit in a separate ignored + remote comparison directory. +2. Sweep 48 KiB through 16 MiB transfer sizes, then tune chunk size, QP count, + projection allocation, helper count, UB tile, bank overlap, batch depth, + barrier, and clear partitioning one group at a time. +3. Re-run correctness after every kernel-affecting tuning change. +4. For each primary 4-rank case, run 20 warmups and at least 50 + device-event iterations, repeated three times. +5. Use `msprof` on representative native and TileXR runs. Verify MiB-scale + batched reads, next-bank/current-compute overlap, no ACK GETs, and no hidden + registration/synchronize inside the measured stage. +6. Check in only defaults justified by the retained raw evidence. + +**Pass gate:** + +- TileXR cross-rank-max P50 and P99 are lower than native in every repeated + primary 4-rank run. +- Median advantage is at least 3%. +- All correctness checks remain green. +- One-time preparation is separately reported. + +**Failure handling:** Continue targeted tuning while evidence identifies an +in-scope bottleneck. If the correct new architecture cannot meet the pass gate, +report the measured limitation and do not label the task complete. Do not +restore the deleted implementation. + +**Artifacts and downstream interfaces:** Raw JSON/CSV, profiler outputs, +comparison summary, and final checked-in tuning constants. + +## Task 11: Final Documentation and Completion Verification + +**Objective and role:** Align maintained documentation with the replacement and +verify the exact final tree after the last tuning edit. + +**Modification scope:** + +- `docs/moonep/DISPATCH_COMBINE.md` +- `docs/BUILD_VERIFICATION.md` +- relevant script documentation; +- approved spec and this plan only for factual implementation notes. + +**Work:** + +1. Remove descriptions of the old ReduceGrad publish/ack protocol. +2. Document preparation lifetime, UDMA-only multi-rank requirement, workspace + sizing, and benchmark commands. +3. Record validation scope as Ascend950/CANN 9.1.T560 on the actual 4-rank + runs. +4. Run the complete focused local and remote verification again after the final + edit, including `git diff --check` and status review. +5. Use `superpowers-neo-verification-before-completion` before describing the + work as complete or faster than native. + +**Acceptance and verification:** Documentation matches behavior, no old +ReduceGrad source path remains, all final tests pass, and the retained evidence +satisfies the faster-than-native gate. + +## Dependency Order + +`Task 1 -> Task 2 -> Task 3 (hardware stop gate) -> Task 4 -> Task 5 -> Task 6` + +Task 7 may start after Task 4 and finishes after Tasks 5 and 6. Task 8 depends +on all implementation and harness work. Tasks 9, 10, and 11 are sequential. + +## Handoff Checklist + +- Approved design and this plan are retained with the implementation. +- Old ReduceGrad transport code and source guards are removed, not disabled. +- Persistent profiles are independently owned and cleaned. +- Hot launches perform no synchronize, MR registration, or legacy tail copy. +- Remote directories on both hosts contain the same source snapshot. +- Correctness evidence covers 4 ranks on physical devices 0-3. +- Performance evidence compares the pinned native commit with identical timing + boundaries and contains cross-rank-max P50/P99. +- No completion claim is made unless the reproducible performance gate passes. + +## Retained Tuning Conclusions + +Do not restart owner-pull tuning from the original variable sweep. The +implementation outcome, exact measurements, GET calibration hashes, rejected +chunk/bank/QP/route/helper/contributor-ready experiments, and sender-PUT stop +gate are recorded in +`docs/specs/2026-08-11-reduce-grad-owner-pull-design.md` under +"Implementation Outcome and Experiment Ledger". Treat that section as the +starting evidence for the next phase. diff --git a/docs/specs/2026-08-11-reduce-grad-owner-pull-design.md b/docs/specs/2026-08-11-reduce-grad-owner-pull-design.md new file mode 100644 index 00000000..36e7791b --- /dev/null +++ b/docs/specs/2026-08-11-reduce-grad-owner-pull-design.md @@ -0,0 +1,543 @@ +# Performance-First ReduceGrad Owner-Pull Design + +Date: 2026-08-11 + +## Status + +Approved by the user on 2026-08-11. The owner-pull implementation and exact +ordered-FP32 validation are complete. On latest `main`, the fixed shared-QP +domain reduced TileXR P50 to `1536-1562 us` in three `20 x 50` runs, versus +the retained pinned-native P50 of `2724-2737 us`. This clears the numerical +3% gate against the retained baseline, but the native artifact used +`torch_npu 2.7.1.post5.dev20260730` while the final TileXR runs used +`2.7.1.post4.dev20260417`; a strict same-runtime native rerun remains open. +Sender-PUT is therefore a conditional next stage, not part of this patch. + +## Goal + +Replace the current TileXR MoonEP ReduceGrad implementation with an Ascend950 +implementation whose steady-state latency is reproducibly lower than native +Ascend MoonEP at commit `a49538a45e5c5bdc82aa6ae02548f99e72ec67eb`. + +The replacement is performance-first and A5/Ascend950-specific. The existing +sender/receiver ReduceGrad kernel, peer-window transport, UDMA push protocol, +per-chunk acknowledgements, and transport fallback are removed. They are not +kept as a runtime fallback. + +The upstream Python `Buffer.reduce_grad` behavior remains compatible: + +- gate, up, and down FP32 gradients are reduced in place; +- only experts owned by the local rank are updated; +- contributions are accumulated in ascending `(source_rank, slot)` order; +- only live slots in the local rank's reduce buffers are cleared; +- unused slots and non-local storage remain unchanged; +- asynchronous stream ordering remains valid. + +## Baseline and Measured Workload + +The comparison source is the ignored checkout +`reference/ascend-moonep-dev` at commit +`a49538a45e5c5bdc82aa6ae02548f99e72ec67eb`. + +Native MoonEP's kernel in `kernels/moonep_grad_reduce.cpp` uses: + +- a fixed 12,288-element FP32 tile, or 48 KiB per remote GET; +- `aclshmemx_mte_get_nbi` directly from remote symmetric memory to UB; +- one 48 KiB accumulator and 48 KiB ping/pong input tiles; +- one output load and store per UB tile; +- source-rank/slot ordered accumulation; +- 32 AIV workers by default; +- a global vector barrier before clearing local live slots. + +The complete native `Buffer.reduce_grad` launches this kernel three times, +once for gate, up, and down. TileXR can fuse those projections and pay one +cross-rank completion barrier. + +The relevant matrix sizes are large even though native MoonEP issues small +individual GETs: + +| Shape | One FP32 expert gradient | Native GETs per gradient | +| --- | ---: | ---: | +| 3584 x 3072 | 42 MiB | 896 | +| 7168 x 2048 | 56 MiB | about 1,195 | + +For three contributors, one 56 MiB expert causes 168 MiB of remote reads per +projection. This is the large-transfer regime in which UDMA must be measured +against UB-memory rather than rejected based on the 48 KiB native tile size. + +## Confirmed Target Constraints + +The target hosts use CANN 9.1.T560 and Ascend950DT. The installed header +`include/pto/npu/comm/async/urma/urma_async_intrin.hpp` declares the public +URMA GET path with GM source and destination pointers: + +```cpp +uint64_t __urma_get_async( + __gm__ uint8_t *dst, + __gm__ uint8_t *src, + uint64_t transferSize, + const UrmaExecContext &execCtx); +``` + +The WQE SGE contains a 32-bit length, a local MR token, and a GM virtual +address. Therefore the design does not assume that UDMA can write directly to +UB. UDMA reads land in registered, non-cacheable GM staging and the AIV vector +pipeline consumes that staging through MTE2. + +The current TileXR UDMA registration model exposes one active region. Its +per-QP device state already contains the two pieces needed for a better model: + +- `UDMAWQCtx.localTokenId` selects the local SGE memory region; +- the per-peer, per-QP `UDMAMemInfo` selects the remote target memory region. + +The new design uses these per-QP fields to bind a local staging MR to one of +three remote gradient MRs without switching registration on the hot path. + +## Why the Current Implementation Is Replaced + +The current kernel performs the following sequence for every chunk: + +1. Copy an unregistered gradient chunk into registered outbound GM. +2. UDMA PUT the chunk into the owner's inbound GM. +3. Quiet the QP. +4. Accumulate the inbound chunk. +5. Poll an 8-byte acknowledgement with UDMA GET plus quiet. +6. Reuse one of two stages and clear the source. + +It also assigns receiver work by local expert only, serializing projections and +chunks for an expert, and uses broad pipeline barriers in dense copy/add/clear +loops. + +There is a separate host-side problem: the Python hot path performs a device +synchronize before re-confirming ReduceGrad registration. Dispatch, +PrefetchWeight, and ReduceGrad use different active memory regions, so a full +flow can repeatedly switch UDMA registration profiles. Stage event timing then +includes preceding work and registration coordination. Both the kernel +protocol and this hot-path profile switching are removed. + +## Chosen Architecture + +### 1. Persistent UDMA Profiles + +Extend the UDMA runtime with persistent, handle-addressed profiles. A profile +owns one or more local memory registrations, their imported remote regions, a +device registry, and a device `UDMAInfo` image. Multiple profiles may coexist +and share the same hardware QPs. A kernel receives the profile's device +pointers explicitly; selecting a profile does not mutate `CommArgs`, register +memory, synchronize a device, or exchange host metadata. + +The existing single-region registration API remains available to unrelated +operators. ReduceGrad does not use its replace-active behavior. + +The ReduceGrad profile contains four region roles on every rank: + +| Region | Local purpose | Remote purpose | +| --- | --- | --- | +| staging | UDMA READ destination and lane state | none | +| gate source | local live gate reduce-buffer slice | remote gate source | +| up source | local live up reduce-buffer slice | remote up source | +| down source | local live down reduce-buffer slice | remote down source | + +Each QP binding chooses the staging region as its local SGE token and exactly +one projection region as its remote token. With eight QPs, the initial mapping +is 3/3/2 lanes across gate/up/down; the measured mapping may be changed when +projection byte sizes differ. + +Registration happens during `prepare_reduce_grad`, before warmup and timed +iterations. The preparation signature includes tensor addresses, local slice +addresses, byte sizes, rank/plan dimensions, chunk size, and QP mapping. A hot +call with mismatched pointers or sizes fails validation. It does not +implicitly synchronize or re-register. + +The internal native ReduceGrad ABI is advanced to carry the three reduce +buffer source pointers and the persistent profile handle/device view. The +public upstream-shaped Python signature is unchanged. Existing exported C ABI +symbols may adapt to the new engine, but no exported entry may launch the old +kernel or old protocol. + +### 2. Large-Chunk Owner Pull + +The owner pulls contributions. No sender AIV, sender notification, completion +GET, or sender acknowledgement exists. + +For each projection, local expert, and MiB-scale row chunk, the owner scans the +small `[R, B]` plan in ascending flattened order and identifies contributors. +It posts deferred UDMA READ WQEs for remote contributors into that lane's +registered staging area, rings the QP doorbell once for the batch, and polls +one ordered completion frontier. + +The initial chunk candidates are 1, 2, 4, and 8 MiB. A target-hardware sweep +selects the checked-in default. The choice must include transfer setup, CQ +polling, GM staging consumption, and vector accumulation; UDMA link bandwidth +alone is insufficient. + +A work item normally has no more than one slot per source rank. To preserve the +general contract, a plan with more contributors than one staging wave can hold +is processed in consecutive waves. Waves and contributions retain flattened +`source_rank * B + slot` order. This is a new-path slow case, not a fallback to +the deleted implementation. + +### 3. QP Lanes and AIV Helper Groups + +One leader AIV exclusively owns each QP, eliminating concurrent SQ head/tail +updates. The remaining AIVs are divided into helper groups, one group per QP +lane. With eight QPs and 64 AIVs, the initial shape is eight leaders and seven +helpers per lane. + +Each lane has two GM staging banks. A leader performs this pipeline: + +1. Batch all remote reads for work item N into bank ping. +2. Publish ping readiness to its helper group after CQ completion. +3. Batch work item N+1 into bank pong while helpers compute N. +4. Wait for helper completion before reusing ping. +5. Alternate banks until the projection lane's work list is empty. + +Each bank reserves `rank_size * chunk_bytes` for one normal contribution wave. +At 16 ranks, eight lanes, two banks, and a 2 MiB chunk, payload staging is 512 +MiB. Workspace sizing is checked for overflow and reported explicitly. + +Helpers partition the staged chunk into UB tiles. For each UB tile a helper: + +1. Loads the owned output tile once into FP32 UB. +2. Loads each staged or local contribution in ordered sequence. +3. Accumulates with FP32 vector add. +4. Stores the owned output tile once. + +The UB layout keeps an accumulator plus ping/pong input tiles. MTE2-to-vector +and vector-to-MTE dependencies use precise events. `PIPE_ALL` is reserved for +boundaries whose target headers require it, not used as the default dependency. + +### 4. Completion and Clear + +Every lane first completes all UDMA reads and all helper work. A local AIV +barrier then ensures no worker can still read a source slot. + +One rank leader publishes one magic-tagged TileXR completion flag and waits for +the corresponding flag from every rank. A rank with non-zero device status +publishes the adjacent failure step; peers still satisfy the same wait, then +detect the non-success value and set their local status. Thus every rank skips +source clearing if any rank fails, without a second barrier or a new mailbox. +This replaces per-AIV global participation. The magic comes from +`TileXRCommNextMagic`; shared flag storage is never reset. + +After the global barrier, all AIVs partition the three local reduce-buffer +slices by projection, live slot, and tile and clear them in parallel. Unused +local slots and every non-local slice are untouched. + +## API and Ownership Rules + +- The ReduceGrad V2 ABI is an unreleased development interface in this branch. + The owner-pull replacement changes it in place; no compatibility wrapper for + the deleted ReduceGrad protocol is retained. +- Multi-rank ReduceGrad requires A5/Ascend950 UDMA and at least one QP per + projection. Missing capability returns `NOT_SUPPORT`; there is no old + ReduceGrad fallback. +- ReduceGrad requires at least four ranks. Rank counts from one through three + return `NOT_SUPPORT` during preparation and have no Kernel path. Eight and + larger rank counts remain supported when the hardware provides usable peer + windows and at least one UDMA QP per projection. +- Registered source slices and staging allocations must remain alive and + unchanged until the last ReduceGrad event completes. +- Re-preparation requires a quiescent communicator and is outside steady-state + timing. +- Profile registration remains unsupported in `InitThread` mode. +- UDMA WQEs are assembled in UB and published through MTE3. Doorbells use + `st_dev` only after WQE publication completes. +- Runtime RPATH and pure-AICore binary registration/launch rules remain + unchanged. + +## Implementation Scope + +Expected affected areas are: + +- `src/comm/udma`: persistent profile registration, cleanup, imports, and + per-QP local/remote region binding; +- `src/include`: profile descriptors and device accessors while preserving + C++14 and existing single-region consumers; +- `src/moonep/reduce_grad`: replacement layout, host validation, kernel args, + kernel, binary launch, and removal of old transport code; +- `integrations/moonep_torch/tilexr_moonep`: preparation cache keyed by pointer + identity, direct use of reduce buffers, and removal of hot-path synchronize, + tail copy, and registration switching; +- `tools/moonep` and `tests/moonep`: isolated baseline harness, correctness + checks, timing hygiene, and source guards proving the old protocol is absent; +- documentation describing the new UDMA-only multi-rank contract. + +`reference/` remains comparison-only and is never compiled or linked into an +active target. + +## Validation Strategy + +### Correctness + +Validate four ranks on physical devices 0-3 of the target host. Include empty, sparse, +mixed, heavy, and full plans; aligned and tail chunks; and all three +projections. Check: + +- bitwise equality with the source-rank/slot ordered FP32 reference where the + same operation order is used; +- only locally owned expert rows change; +- all and only local live reduce slots become zero; +- unused and non-local slots remain byte-identical; +- repeated magic rounds do not consume stale flags; +- asynchronous launch and explicit synchronization report device failures. + +### Performance + +Build native MoonEP from the pinned baseline commit and TileXR from the same +test directory on both hosts. Use identical dimensions, plan tensors, rank +count, stream synchronization, warmup, iterations, and device-event timing +boundaries. + +Primary cases include: + +- MoonEP's dedicated 3584 x 3072 suite; +- E=384, H=7168, H'=2048, K=8, S=8192, B=48; +- 4 ranks on physical devices 0-3 of one host; +- 20 warmup iterations and at least 50 measured iterations. + +Report the cross-rank maximum per iteration, then P50 and P99. Run each primary +case at least three times. One-time allocation, MR registration/import, and +kernel binary registration are reported separately and excluded from +steady-state latency. + +Completion requires all of the following: + +1. TileXR P50 and P99 are both lower than native MoonEP for the primary 4-rank + cases in every repeated run. +2. The median advantage is at least 3%, so measurement noise is not presented + as a win. +3. No correctness condition above fails. +4. Device-event timelines show no host/device synchronize or registration in + the timed ReduceGrad path. +5. `msprof` confirms batched MiB-scale UDMA reads, overlap between the next + staging bank and current vector work, and no per-chunk acknowledgement GET. + +### Tuning Sequence + +Tune one variable group at a time and retain raw JSON/CSV results: + +1. UDMA versus native UB-memory transfer sweep from 48 KiB through 16 MiB. +2. Chunk size with one QP, including staging and accumulation. +3. QP count and projection-to-QP allocation. +4. Helper count per lane and UB tile size. +5. Ping/pong overlap and WQE batch depth. +6. End barrier and local clear partitioning. + +## Stop Conditions and Risks + +The first implementation milestone is a target-hardware proof that several +simultaneously registered MRs can supply a staging local token and +projection-specific remote tokens on shared QPs. If CANN/HCCP rejects that +profile, or if independent profiles cannot safely share the QPs, implementation +returns to design review. It must not silently substitute the old protocol. + +Other material risks are: + +- MR/import resource limits with four regions per rank; +- a UDMA/GM staging crossover larger than the real contribution chunks; +- staging workspace pressure at 16 ranks; +- helper-group flag/cache overhead erasing pipeline overlap; +- a low-QP environment that cannot allocate one lane per projection; +- an apparent win caused by different timing boundaries or by excluding work + that native MoonEP includes. + +No performance claim is made until the 4-rank target-hardware evidence +satisfies the stated gate. The physical-card mapping and excluded faulty +devices are retained with the raw results. + +## Implementation Outcome and Experiment Ledger + +### Implemented Owner-Pull Optimizations + +The delivered implementation replaces the previous ReduceGrad protocol with: + +- persistent, handle-addressed multi-MR UDMA profiles, avoiding hot-path + registration switching and device synchronization; +- batched deferred owner GETs with one doorbell and completion frontier per + participating source/QP; +- one fused gate/up/down launch, deterministic projection-to-QP allocation, + ping/pong registered GM staging, and leader/helper AIV groups; +- separation of the fixed 32-QP shared transport profile from three active + ReduceGrad lanes, mapped to physical QPs `{0, 1, 16}` to preserve the + measured six-port/six-port/two-port route mix without allocating workspace + for inactive QPs; +- an 8 MiB production default selected by the retained chunk sweep, while + preserving the explicit chunk-size override for targeted validation; +- ordered FP32 accumulation in flattened `(source_rank, slot)` order; +- one magic-tagged cross-rank completion barrier followed by parallel clearing + of local live slots, with the same flag carrying a failure step so any rank + failure prevents clearing on every rank; +- retained peer-memory `SyncCollectives`; the ReduceGrad data plane itself has + no legacy sender/receiver, PUT/ACK, or peer-packed-record fallback. + +Cold failure paths were tightened without changing the owner-pull data loop: a +failed kernel launch drains the already-enqueued asynchronous status memset +once; a data, profile, or workspace configuration failure skips lane work but +still reaches the existing collective barrier, whose adjacent failure step +prevents every rank from clearing source data; and ReduceGrad device statuses +`1..5` poison the Python context. These are correctness boundaries, not +performance optimizations, and add no stream synchronization to a successful +launch or another cross-rank barrier. + +These changes passed exact ordered-FP32 correctness on four physical cards. +The pre-shared-domain control runtime is: + +`/tmp/TileXR-reducegrad-20260811-44dc37b/install-control-2bank-runtime-20260812` + +For four ranks, one expert per rank, one slot, mixed plan, FP32 +`3584 x 3072`, and an 8 MiB transfer setting, owner-pull measured +P50/P99 `2764.82/2843.38 us`. Native MoonEP measured approximately +`2720-2737 us`; the required 3% gate was approximately `<2655 us`. +That explicit 3-QP owner-pull control was close to native but did not pass the +acceptance gate. + +Latest `main` initializes MoonEP with a fixed 32-QP shared domain. The first +delivery build rejected that domain during workspace query with +`TileXR ret=-6` because transport QPs and active lanes were incorrectly the +same count. The final integration keeps all 32 QPs in the persistent profile, +uses only logical lanes `{0,1,2}`, and maps them to physical QPs `{0,1,16}`. +Inactive QPs have valid harmless bindings and are never scheduled. This keeps +the 4-rank, 8 MiB workspace at `203423744` bytes rather than scaling it to 32 +lanes. + +Final artifacts are on `141.61.53.106` under: + +`/tmp/TileXR-reducegrad-delivery-20260812-8d22775/artifacts` + +| Run | Warmup x measured | P50 | P99 | Result | +| --- | ---: | ---: | ---: | --- | +| `final-perf-shared32-run1` | `20 x 50` | `1562.40 us` | `1662.54 us` | Passed | +| `final-perf-shared32-run2` | `20 x 50` | `1550.89 us` | `1632.76 us` | Passed | +| `final-perf-shared32-run3` | `20 x 50` | `1536.22 us` | `1595.37 us` | Passed | +| pinned native `perf-20x50` | `20 x 50` | `2723.68 us` | `2755.00 us` | Passed | +| pinned native interleaved A | `20 x 50` | `2736.30 us` | `2768.92 us` | Passed | +| pinned native interleaved B | `20 x 50` | `2737.48 us` | `2800.86 us` | Passed | + +The worst final TileXR P50 is `42.6%` below the best retained native P50; the +worst TileXR P99 is `39.7%` below the best retained native P99. The final +`final-exact-4card-shared32-guarded` run passed exact ordered-FP32 validation +on all four ranks with P50/P99 `1565.46/1567.53 us` over its three smoke +samples. Those three samples are correctness evidence, not a replacement for +the three formal performance runs. + +After the final review fix, the default-parameter regression +`final-post-review-default8m-exact` again passed exact ordered-FP32 validation +on four physical cards, including live-source clearing. It intentionally +omitted `--chunk-bytes`; the reported layout selected `8388608` bytes and the +same `203423744`-byte workspace. Its three smoke samples measured P50/P99 +`1614.31/2004.99 us`; these samples are correctness evidence, not a formal +performance run. The summary SHA256 is +`019f7c4a41d86544a1cdefe008dee7fa4c7935272c0b4d95d062d51d244762bc`. +The final ReduceGrad shared-library SHA256 is +`33764d7e8149c4c7f23a1f58916e8484e3b6b74906e58893d98c092651fa82c0` +on both validation hosts. + +The native comparison source remains pinned at +`a49538a45e5c5bdc82aa6ae02548f99e72ec67eb`. A same-runtime rerun was attempted +on the final host. The old extension required the system `GLIBCXX_3.4.30`, +PyTorch and torch-npu library search paths, and then still failed under +`torch_npu 2.7.1.post4.dev20260417` with duplicate NPU backend registration. +No reference source or server environment was modified to bypass that +incompatibility. Therefore the large measured advantage is actionable, but a +strict same-runtime A/B remains a release-quality follow-up. + +### UDMA GET Calibration + +The retained probe is under `/tmp/TileXR-get-probe-20260812-44dc37b` with +artifacts in `artifacts/get-calibration-run1`. + +- Probe SHA256: + `37aead24de783550dbe61859f5a345651030b4d6bb5a307324b0d894d752be7a`. +- `libtile-comm.so` SHA256: + `2615b4103dfdbe3c047c7c40da6bb525afa1e260e09a5334ec6826db4b51fdb4`. +- High-bandwidth-route QP GET reached approximately `273 GB/s`. +- Low-route QP GET reached approximately `95.6 GB/s`. +- GET plus staging consumption reached only approximately `28-35 GB/s`. +- Byte-exact validation passed. + +This isolates the remaining owner-pull bottleneck to staging consumption and +owner-side accumulation/synchronization rather than raw UDMA GET bandwidth. + +### Performance Experiments Not To Repeat + +The following owner-pull tuning directions were tested and rejected because +they produced no repeatable improvement or regressed the control: + +| Experiment | Observation | Decision | +| --- | --- | --- | +| Chunk-size sweep | 8 MiB control: `2764.82/2843.38 us`. 9, 10, 10.5, 11, 12, 14, and 16 MiB gave P50 `2801`, `2765`, `2867`, `2804`, `2790`, `2861`, and `2890 us`; 2 MiB was `2860 us`. | Keep 8 MiB. Do not repeat a blind chunk sweep. Artifacts: `control-chunk-*`, `reducegrad-async-two-bank-chunk*`, and `reducegrad-accum-pipeline-v1/perf-screen-chunk16m`. | +| Three or twelve staging banks | Three banks reached `2762.27/2792.89 us`, statistically flat against two banks; twelve banks regressed to `2861.70/3034.28 us` and consumed more workspace/state. | Keep two banks. Artifacts: `reducegrad-async-three-bank-v1` and `reducegrad-async-twelve-bank-v1`. | +| Six or eight active QPs | Six-QP 7 MiB: `2872.42/2914.68 us`; six-QP 8 MiB: `2900.12/3391.89 us`; eight-QP 4 MiB: `2947.77/3049.88 us`. Four-QP `6/3/1` was `2772.12/2866.72 us`. | Keep three active lanes. Do not add QPs without a new data-flow hypothesis. | +| All-six-port or all-two-port routing | All-six-port three-QP was `2762.69/2879.98 us`, flat/noisier than control; all-two-port was `2787.76/2922.89 us`. | Link route alone is not the old owner-pull bottleneck. Keep `{0,1,16}` in shared mode. | +| QP rotation/interleaving | Rotating QPs per work item produced no stable gain in the local sweep and was excluded from production. | Keep stable single-owner QP assignment. Revisit only with a different data-flow model. | +| 32 AIV and helper-shape tuning | 32 AIV at 8 MiB regressed to `2806.93/2929.62 us`; 48 AIV output-prefetch was `2784.40/2901.04 us`. | Keep 64 AIV. Do not retune helpers in isolation. | +| Accumulator/data-copy micro-pipelines | Accumulator pipeline `2809.00/2900.88 us`; aligned DataCopy `2810.22/2954.08 us`; bank skew `2785.19/2960.88 us`. | Rejected. The extra pipeline/cache bookkeeping did not improve end to end. | +| Output/input prefetch variants | Output-prefetch variants ranged from P50 `2767.53` to `2790.40 us`; input+output prefetch was `2796.69/2931.12 us`. | Keep the checked-in two-bank consume loop; do not layer more prefetch without profiling a new bottleneck. | +| Inline contributor cache | Helper contributor cache was `2752.98/2852.92 us`, not repeatably better; fully inline contributors regressed to `2898.88/3064.55 us`. | Rejected and reverted. | +| Immediate doorbell, WQE batch, polling backoff | Immediate doorbell `2783.25/3008.42 us`; alternate WQE batch `2910.37/3009.26 us`; polling backoff `2819.27/2986.26 us`. | Retain deferred batch publication and current completion polling. | +| Leader overlap or single coordinator | Leader overlap `2866.05/2967.99 us`; one coordinator `3301.58/3359.10 us`. | Retain one leader per lane with distributed helpers. | +| Early helper release | Produced vector-core exception `507035` from out-of-bounds internal-buffer access; the reverted control was still `2825.93/3262.03 us`. | Unsafe and rejected. Do not retry without a redesigned ownership/lifetime proof. | +| Contributor-ready overlap | Exact ordered FP32 passed, but P50/P99 regressed to `2763.22/2903.76 us` and helper wait rose to `295-323 us`. | Rejected and reverted. Artifacts: `/tmp/TileXR-reducegrad-source-ready-20260812`. | + +The rejected contributor-ready artifacts remain under +`/tmp/TileXR-reducegrad-source-ready-20260812`. Local QP-rotation and chunk +sweep helper scripts were experimental only and are intentionally excluded +from this patch. + +### Debugging History Not To Repeat + +These were correctness investigations, not viable optimization candidates: + +- The first sparse/mixed implementations showed staged data was correct while + helper-visible metadata or done tokens were stale. The final protocol uses + cache-line-separated magic-tagged tokens and explicit cache maintenance. +- Increasing `waitIterations` did not fix missing helper completion; it only + prolonged the timeout. Do not treat this class of failure as a timeout-tuning + problem. +- MTE barriers around the bank item did not substitute for the required GM + visibility/cache-line protocol. The `correctness-8r-empty-*` and + `reducegrad-4r-debug-v10` through `v22` artifacts are diagnostic history. +- Exact correctness initially stopped before launch at the default 512 MiB + inspection cap: the tested tensors require `660602880` bytes. Use at least + `805306368` for this `3584 x 3072`, four-rank case. +- Latest-main initially failed before launch with `ret=-6`. The cause was not + workspace overflow: shared initialization exposes 32 transport QPs while + the old owner-pull layout allowed at most eight lanes. Keep transport-QP + count separate from active-lane count and require the `UDMA_SHARED_QP` + capability before applying the fixed `{0,1,16}` physical mapping. + +## Recommended Next Stage: Sender-PUT Fan-In + +This section is a conditional design recommendation and stop gate only. It is +not part of the implemented owner-pull patch. Because shared-QP owner-pull now +measures below the `<1.9 ms` sender-PUT exploration gate, sender-PUT should not +replace it on the current workload without first demonstrating a material, +repeatable advantage under a strict same-runtime A/B. + +Use three non-owner ranks as senders and one rank as owner. Each sender PUTs +gate, up, and down directly into the owner's registered staging. Use routes +`6,6,6,2` and distribute payload approximately 3:1 between the aggregate +six-port paths and the two-port path. Allocate whole-row, non-reused staging so +senders need no per-chunk acknowledgement. Each source publishes one +magic-tagged ready indication only after all three projections are visible. +The owner then accumulates strictly in `(source_rank, slot)` order and retains +the existing peer-memory `SyncCollectives` completion path. + +For the validated `3584 x 3072` FP32 case, one projection row is 42 MiB. One +sender transfers 126 MiB for three projections, and the owner receives +`3 projections x 3 contributors x 42 MiB = 378 MiB` per ReduceGrad round. +This is firmly a large-transfer regime, so UDMA is preferred over UB-memory for +the data plane even though owner-side consumption must still be optimized. + +Before production ReduceGrad is redesigned, build an isolated four-physical- +card exploratory probe with three senders and one owner. It must pass exact +ordered-FP32 validation and achieve P50 `<1.9 ms` for the 378 MiB fan-in. Stop +and return to design review if either condition fails; do not integrate a +sender-PUT production path based only on raw link-bandwidth results. + +The immediate next validation step is instead to rebuild pinned native MoonEP +for the exact final `torch_npu` environment, then repeat three interleaved +`20 x 50` runs on the same physical cards. Only pursue sender-PUT if that strict +comparison invalidates the shared-QP owner-pull gain or a broader workload +matrix exposes a new bottleneck. diff --git a/integrations/moonep_torch/tilexr_moonep/abi.py b/integrations/moonep_torch/tilexr_moonep/abi.py index 0c2ac639..9c2f33a3 100644 --- a/integrations/moonep_torch/tilexr_moonep/abi.py +++ b/integrations/moonep_torch/tilexr_moonep/abi.py @@ -15,7 +15,6 @@ TILEXR_MOONEP_FLAG_ZERO_COPY = 1 << 2 TILEXR_MOONEP_FLAG_COMBINE_PUBLISH_ONLY = 1 << 3 TILEXR_MOONEP_FLAG_COMBINE_CONSUME_ONLY = 1 << 4 -TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES = 1 << 20 class TileXRMoonEPDType(IntEnum): @@ -33,12 +32,6 @@ class TileXRMoonEPStage(IntFlag): REDUCE_GRAD = 1 << 4 -class TileXRMoonEPReduceGradTransport(IntEnum): - NONE = 0 - PEER = 1 - UDMA = 2 - - class TileXRMoonEPTensorV1(ctypes.Structure): _fields_ = [ ("structSize", ctypes.c_uint32), @@ -169,16 +162,29 @@ class TileXRMoonEPReduceGradWorkspaceInfoV2(ctypes.Structure): ("workspaceBytes", ctypes.c_uint64), ("workspaceAlignment", ctypes.c_uint64), ("udmaChunkBytes", ctypes.c_uint64), - ("peerWindowBytes", ctypes.c_uint64), - ("peerHalfBytes", ctypes.c_uint64), - ("peerSlotStrideBytes", ctypes.c_uint64), + ("laneStateBytes", ctypes.c_uint64), + ("laneStateStrideBytes", ctypes.c_uint64), + ("bankStrideBytes", ctypes.c_uint64), + ("laneStrideBytes", ctypes.c_uint64), ("rowBytes", ctypes.c_uint64 * 3), - ("transports", ctypes.c_uint32 * 3), + ("chunkCounts", ctypes.c_uint64 * 3), + ("projectionQpCounts", ctypes.c_uint32 * 3), + ("qpCount", ctypes.c_uint32), ("blockDim", ctypes.c_uint32), + ("reserved", ctypes.c_uint32), ] -class TileXRMoonEPReduceGradArgsV2(ctypes.Structure): +class TileXRMoonEPReduceGradSourceSliceV2(ctypes.Structure): + _fields_ = [ + ("data", ctypes.c_void_p), + ("bytes", ctypes.c_uint64), + ("registrationBase", ctypes.c_void_p), + ("registrationBytes", ctypes.c_uint64), + ] + + +class TileXRMoonEPReduceGradPrepareArgsV2(ctypes.Structure): _fields_ = [ ("structSize", ctypes.c_uint32), ("abiVersion", ctypes.c_uint32), @@ -187,11 +193,26 @@ class TileXRMoonEPReduceGradArgsV2(ctypes.Structure): ("gate", ctypes.POINTER(TileXRMoonEPTensorV1)), ("up", ctypes.POINTER(TileXRMoonEPTensorV1)), ("down", ctypes.POINTER(TileXRMoonEPTensorV1)), + ("sources", TileXRMoonEPReduceGradSourceSliceV2 * 3), ("workspace", ctypes.c_void_p), ("workspaceBytes", ctypes.c_uint64), + ("requestedUdmaChunkBytes", ctypes.c_uint64), + ("flags", ctypes.c_uint64), + ] + + +class TileXRMoonEPReduceGradArgsV2(ctypes.Structure): + _fields_ = [ + ("structSize", ctypes.c_uint32), + ("abiVersion", ctypes.c_uint32), + ("prepared", ctypes.c_void_p), + ("plan", ctypes.POINTER(TileXRMoonEPPlanV1)), + ("gate", ctypes.POINTER(TileXRMoonEPTensorV1)), + ("up", ctypes.POINTER(TileXRMoonEPTensorV1)), + ("down", ctypes.POINTER(TileXRMoonEPTensorV1)), + ("sources", TileXRMoonEPReduceGradSourceSliceV2 * 3), ("status", ctypes.POINTER(TileXRMoonEPTensorV1)), ("waitIterations", ctypes.c_uint64), - ("requestedUdmaChunkBytes", ctypes.c_uint64), ("flags", ctypes.c_uint64), ] diff --git a/integrations/moonep_torch/tilexr_moonep/compat.py b/integrations/moonep_torch/tilexr_moonep/compat.py index 2e089075..5fd9a751 100644 --- a/integrations/moonep_torch/tilexr_moonep/compat.py +++ b/integrations/moonep_torch/tilexr_moonep/compat.py @@ -464,28 +464,9 @@ def reduce_grad( if not isinstance(plan, MoonEPCommPlan): raise AssertionError("Buffer.reduce_grad: plan is required") native_plan = plan._require_native() - full_gradients = (full_gate_grad, full_up_grad, full_down_grad) - reduce_buffers = ( - gate_reduce_buffer, - up_reduce_buffer, - down_reduce_buffer, - ) - full_slot_snapshots = tuple( - value.narrow(0, self.E, self.B).clone() for value in full_gradients - ) - local_reduce_buffers = tuple( - value[self._context.planner_group_rank] for value in reduce_buffers - ) - local_reduce_snapshots = tuple( - value.clone() for value in local_reduce_buffers - ) - valid_slots = native_plan.experts_to_copy[ - self._context.planner_group_rank - ] >= 0 - - self._native_buffer.reduce_grad( + event = self._native_buffer.reduce_grad( plan=native_plan, - async_finish=False, + async_finish=bool(async_finish), full_gate_grad=full_gate_grad, full_up_grad=full_up_grad, full_down_grad=full_down_grad, @@ -493,17 +474,9 @@ def reduce_grad( up_reduce_buffer=up_reduce_buffer, down_reduce_buffer=down_reduce_buffer, ) - for full_gradient, snapshot in zip(full_gradients, full_slot_snapshots): - full_gradient.narrow(0, self.E, self.B).copy_(snapshot) - for local_buffer, snapshot in zip( - local_reduce_buffers, local_reduce_snapshots - ): - local_buffer.copy_(snapshot) - mask_shape = (self.B,) + (1,) * (len(local_buffer.shape) - 1) - local_buffer.masked_fill_(valid_slots.reshape(mask_shape), 0.0) if async_finish: return _CompletionEvent( - self._record_event(), + event, self._native_buffer.synchronize, ) self._native_buffer.synchronize() diff --git a/integrations/moonep_torch/tilexr_moonep/runtime.py b/integrations/moonep_torch/tilexr_moonep/runtime.py index 6b1ac304..4e3f421d 100644 --- a/integrations/moonep_torch/tilexr_moonep/runtime.py +++ b/integrations/moonep_torch/tilexr_moonep/runtime.py @@ -18,7 +18,8 @@ TileXRMoonEPPlanningArgsV1, TileXRMoonEPPrefetchWeightArgsV1, TileXRMoonEPReduceGradArgsV2, - TileXRMoonEPReduceGradTransport, + TileXRMoonEPReduceGradPrepareArgsV2, + TileXRMoonEPReduceGradSourceSliceV2, TileXRMoonEPReduceGradWorkspaceInfoV2, TileXRMoonEPReduceGradWorkspaceQueryV2, TileXRMoonEPStage, @@ -111,28 +112,33 @@ class ReduceGradWorkspaceInfo: workspace_bytes: int workspace_alignment: int udma_chunk_bytes: int - peer_window_bytes: int - peer_half_bytes: int - peer_slot_stride_bytes: int + lane_state_bytes: int + lane_state_stride_bytes: int + bank_stride_bytes: int + lane_stride_bytes: int row_bytes: tuple[int, int, int] - transports: tuple[str, str, str] + chunk_counts: tuple[int, int, int] + projection_qp_counts: tuple[int, int, int] + qp_count: int block_dim: int - @property - def uses_udma(self) -> bool: - return "udma" in self.transports - def as_dict(self) -> dict[str, object]: return { "workspace_bytes": self.workspace_bytes, "workspace_alignment": self.workspace_alignment, "udma_chunk_bytes": self.udma_chunk_bytes, - "peer_window_bytes": self.peer_window_bytes, - "peer_half_bytes": self.peer_half_bytes, - "peer_slot_stride_bytes": self.peer_slot_stride_bytes, + "lane_state_bytes": self.lane_state_bytes, + "lane_state_stride_bytes": self.lane_state_stride_bytes, + "bank_stride_bytes": self.bank_stride_bytes, + "lane_stride_bytes": self.lane_stride_bytes, "row_bytes": dict(zip(("gate", "up", "down"), self.row_bytes)), - "transports": dict(zip(("gate", "up", "down"), self.transports)), + "chunk_counts": dict(zip(("gate", "up", "down"), self.chunk_counts)), + "projection_qp_counts": dict( + zip(("gate", "up", "down"), self.projection_qp_counts) + ), + "qp_count": self.qp_count, "block_dim": self.block_dim, + "transport": "udma", "registration_in_timed_path": False, } @@ -201,10 +207,6 @@ def __init__( self._closed = False self._comm = ctypes.c_void_p() self._udma_qp_count = 0 - self._udma_handle: ctypes.c_uint32 | None = None - self._udma_workspace = None - self._udma_workspace_ptr = 0 - self._udma_workspace_bytes = 0 self._active_udma_owner: str | None = None self._active_udma_pointer = 0 self._active_udma_bytes = 0 @@ -374,6 +376,15 @@ def _configure_symbols(self) -> None: ctypes.POINTER(TileXRMoonEPReduceGradWorkspaceInfoV2), ] self._moonep_lib.TileXRMoonEpReduceGradGetWorkspaceSizeV2.restype = ctypes.c_int + self._moonep_lib.TileXRMoonEpReduceGradPrepareV2.argtypes = [ + ctypes.POINTER(TileXRMoonEPReduceGradPrepareArgsV2), + ctypes.POINTER(ctypes.c_void_p), + ] + self._moonep_lib.TileXRMoonEpReduceGradPrepareV2.restype = ctypes.c_int + self._moonep_lib.TileXRMoonEpReduceGradDestroyPreparedV2.argtypes = [ + ctypes.c_void_p, + ] + self._moonep_lib.TileXRMoonEpReduceGradDestroyPreparedV2.restype = ctypes.c_int self._moonep_lib.TileXRMoonEpReduceGradV2.argtypes = [ ctypes.POINTER(TileXRMoonEPReduceGradArgsV2), ctypes.c_void_p, @@ -836,98 +847,111 @@ def reduce_grad_workspace_info( ctypes.byref(query), ctypes.byref(info) ) self._check("TileXRMoonEpReduceGradGetWorkspaceSizeV2", ret) - transport_names = { - int(TileXRMoonEPReduceGradTransport.NONE): "none", - int(TileXRMoonEPReduceGradTransport.PEER): "peer", - int(TileXRMoonEPReduceGradTransport.UDMA): "udma", - } - try: - transports = tuple(transport_names[int(value)] for value in info.transports) - except KeyError as exc: - raise TileXRMoonEPError( - "TileXRMoonEpReduceGradGetWorkspaceSizeV2", -1, - f"unknown transport code {int(exc.args[0])}", - ) from exc return ReduceGradWorkspaceInfo( workspace_bytes=int(info.workspaceBytes), workspace_alignment=int(info.workspaceAlignment), udma_chunk_bytes=int(info.udmaChunkBytes), - peer_window_bytes=int(info.peerWindowBytes), - peer_half_bytes=int(info.peerHalfBytes), - peer_slot_stride_bytes=int(info.peerSlotStrideBytes), + lane_state_bytes=int(info.laneStateBytes), + lane_state_stride_bytes=int(info.laneStateStrideBytes), + bank_stride_bytes=int(info.bankStrideBytes), + lane_stride_bytes=int(info.laneStrideBytes), row_bytes=tuple(int(value) for value in info.rowBytes), - transports=transports, + chunk_counts=tuple(int(value) for value in info.chunkCounts), + projection_qp_counts=tuple( + int(value) for value in info.projectionQpCounts + ), + qp_count=int(info.qpCount), block_dim=int(info.blockDim), ) - def register_reduce_grad_workspace( - self, workspace, required_bytes: int, *, owner_token: object | None = None - ) -> None: - with self._reduce_grad_lock: - self._require_reduce_grad_workspace_owner( - owner_token, "register_reduce_grad_workspace" - ) - pointer = int(workspace.data_ptr()) - available = tensor_nbytes(workspace) - if required_bytes <= 0 or available < required_bytes: - raise ValueError( - f"ReduceGrad workspace has {available} bytes, requires {required_bytes}" - ) - if self._udma_handle is not None and ( - pointer != self._udma_workspace_ptr or - available != self._udma_workspace_bytes - ): - raise RuntimeError("a different TileXR UDMA workspace is already registered") - handle = self._activate_udma_region( - pointer, available, "reduce_grad", f"workspace_bytes={available}" + @staticmethod + def _reduce_grad_source_slices( + sources, registrations + ) -> tuple[TileXRMoonEPReduceGradSourceSliceV2, ...]: + slices = [] + for source, registration in zip(sources, registrations): + value = TileXRMoonEPReduceGradSourceSliceV2() + value.data = tensor_ptr(source) + value.bytes = tensor_nbytes(source) + value.registrationBase = tensor_ptr(registration) + value.registrationBytes = tensor_nbytes(registration) + slices.append(value) + return tuple(slices) + + def prepare_reduce_grad( + self, + context, + plan, + gradients, + sources, + source_registrations, + workspace, + *, + requested_udma_chunk_bytes: int = 0, + ) -> int: + plan_v1 = self._plan_v1(context, plan) + descriptors = [ + make_tensor_v1(getattr(gradients, name)) for name in ("gate", "up", "down") + ] + source_slices = self._reduce_grad_source_slices(sources, source_registrations) + args = initialize_struct(TileXRMoonEPReduceGradPrepareArgsV2()) + args.comm = void_p(self.comm_ptr) + args.plan = ctypes.pointer(plan_v1) + args.gate = ctypes.pointer(descriptors[0]) + args.up = ctypes.pointer(descriptors[1]) + args.down = ctypes.pointer(descriptors[2]) + for index, source in enumerate(source_slices): + args.sources[index] = source + args.workspace = tensor_ptr(workspace) + args.workspaceBytes = tensor_nbytes(workspace) + args.requestedUdmaChunkBytes = int(requested_udma_chunk_bytes) + args.flags = TILEXR_MOONEP_FLAG_NONE + prepared = ctypes.c_void_p() + ret = self._moonep_lib.TileXRMoonEpReduceGradPrepareV2( + ctypes.byref(args), ctypes.byref(prepared) + ) + self._check("TileXRMoonEpReduceGradPrepareV2", ret) + if not prepared.value: + raise TileXRMoonEPError( + "TileXRMoonEpReduceGradPrepareV2", -1, "successful prepare returned null" ) - self._udma_handle = None if handle is None else ctypes.c_uint32(handle) - self._udma_workspace = workspace - self._udma_workspace_ptr = pointer - self._udma_workspace_bytes = available + return int(prepared.value) - def unregister_reduce_grad_workspace( - self, *, owner_token: object | None = None - ) -> None: - with self._reduce_grad_lock: - self._require_reduce_grad_workspace_owner( - owner_token, "unregister_reduce_grad_workspace" - ) - if self._udma_handle is None: - return - self._deactivate_udma_region("reduce_grad") - self._udma_handle = None - self._udma_workspace = None - self._udma_workspace_ptr = 0 - self._udma_workspace_bytes = 0 + def destroy_reduce_grad(self, prepared: int | None) -> None: + if prepared is None: + return + ret = self._moonep_lib.TileXRMoonEpReduceGradDestroyPreparedV2( + void_p(prepared) + ) + self._check("TileXRMoonEpReduceGradDestroyPreparedV2", ret) def reduce_grad( self, context, plan, gradients, - workspace, + sources, + source_registrations, + prepared: int, stream_ptr: int, wait_iterations: int, - *, - requested_udma_chunk_bytes: int = 0, ) -> None: plan_v1 = self._plan_v1(context, plan) descriptors = [ make_tensor_v1(getattr(gradients, name)) for name in ("gate", "up", "down") ] + source_slices = self._reduce_grad_source_slices(sources, source_registrations) status = make_tensor_v1(plan.reduce_grad_status) args = initialize_struct(TileXRMoonEPReduceGradArgsV2()) - args.comm = void_p(self.comm_ptr) + args.prepared = void_p(prepared) args.plan = ctypes.pointer(plan_v1) args.gate = ctypes.pointer(descriptors[0]) args.up = ctypes.pointer(descriptors[1]) args.down = ctypes.pointer(descriptors[2]) - args.workspace = void_p(None if workspace is None else int(workspace.data_ptr())) - args.workspaceBytes = 0 if workspace is None else tensor_nbytes(workspace) + for index, source in enumerate(source_slices): + args.sources[index] = source args.status = ctypes.pointer(status) args.waitIterations = int(wait_iterations) - args.requestedUdmaChunkBytes = int(requested_udma_chunk_bytes) args.flags = TILEXR_MOONEP_FLAG_NONE ret = self._moonep_lib.TileXRMoonEpReduceGradV2( ctypes.byref(args), void_p(stream_ptr) @@ -944,7 +968,6 @@ def close(self) -> None: if self._closed: return if self._comm.value: - self.unregister_reduce_grad_workspace() self._deactivate_udma_region() ret = self._comm_lib.TileXRCommDestroy(self._comm) self._check("TileXRCommDestroy", ret, f"rank={self.rank}") diff --git a/integrations/moonep_torch/tilexr_moonep/torch_api.py b/integrations/moonep_torch/tilexr_moonep/torch_api.py index f422330e..cae1942e 100644 --- a/integrations/moonep_torch/tilexr_moonep/torch_api.py +++ b/integrations/moonep_torch/tilexr_moonep/torch_api.py @@ -491,6 +491,8 @@ def __init__( self._reduce_grad_inflight = False self._reduce_grad_workspace = None self._reduce_grad_workspace_allocation = None + self._reduce_grad_prepared: int | None = None + self._reduce_grad_prepared_refs: tuple[Any, ...] = () self._reduce_grad_signature = None self._reduce_grad_info = None self._quiesced = True @@ -563,7 +565,7 @@ def _zeros(self, shape: tuple[int, ...], dtype): def _aligned_workspace(self, size_bytes: int, alignment: int): if size_bytes <= 0 or alignment <= 0: raise ValueError("ReduceGrad workspace size and alignment must be positive") - allocation = self._zeros((size_bytes + alignment - 1,), self._torch.uint8) + allocation = self._empty((size_bytes + alignment - 1,), self._torch.uint8) offset = (-int(allocation.data_ptr())) % alignment workspace = allocation.narrow(0, offset, size_bytes) if int(workspace.data_ptr()) % alignment != 0: @@ -984,11 +986,35 @@ def prepare_reduce_grad( self._require_open() self._validate_plan(plan) self._validate_projection_buffers(gradients, reduce=True) + sources, source_registrations = self._reduce_grad_source_regions(gradients) shape_signature = ( - tuple(_shape(getattr(gradients, name)) for name in ("gate", "up", "down")), + tuple( + ( + int(getattr(gradients, name).data_ptr()), + _shape(getattr(gradients, name)), + ) + for name in ("gate", "up", "down") + ), + tuple( + (int(source.data_ptr()), int(source.numel()) * int(source.element_size())) + for source in sources + ), + tuple( + ( + int(registration.data_ptr()), + int(registration.numel()) * int(registration.element_size()), + ) + for registration in source_registrations + ), + int(plan.experts_to_copy.data_ptr()), + int(plan.n), + int(plan.topk), self.requested_udma_chunk_bytes, ) - if self._reduce_grad_signature == shape_signature: + if ( + self._reduce_grad_signature == shape_signature + and self._reduce_grad_prepared is not None + ): return self._reduce_grad_info info = self.runtime.reduce_grad_workspace_info( @@ -997,34 +1023,68 @@ def prepare_reduce_grad( gradients, requested_udma_chunk_bytes=self.requested_udma_chunk_bytes, ) - if self._reduce_grad_signature is not None: - if not self._quiesced: - self.synchronize() - self.runtime.unregister_reduce_grad_workspace(owner_token=_owner_token) - self._reduce_grad_workspace = None - self._reduce_grad_workspace_allocation = None - self._reduce_grad_signature = None - self._reduce_grad_info = None - - if info.uses_udma: - if not self._quiesced: - self.synchronize() - workspace, allocation = self._aligned_workspace( - info.workspace_bytes, info.workspace_alignment - ) - self.quiesce() - self.runtime.register_reduce_grad_workspace( - workspace, info.workspace_bytes, owner_token=_owner_token - ) - self._reduce_grad_workspace = workspace - self._reduce_grad_workspace_allocation = allocation - else: - self._reduce_grad_workspace = None - self._reduce_grad_workspace_allocation = None + if not self._quiesced: + self.synchronize() + if self._reduce_grad_prepared is not None: + self.runtime.destroy_reduce_grad(self._reduce_grad_prepared) + self._reduce_grad_prepared = None + self._reduce_grad_prepared_refs = () + self._reduce_grad_workspace = None + self._reduce_grad_workspace_allocation = None + self._reduce_grad_signature = None + self._reduce_grad_info = None + + workspace, allocation = self._aligned_workspace( + info.workspace_bytes, info.workspace_alignment + ) + prepared = self.runtime.prepare_reduce_grad( + self.context, + plan, + gradients, + sources, + source_registrations, + workspace, + requested_udma_chunk_bytes=self.requested_udma_chunk_bytes, + ) + self._reduce_grad_workspace = workspace + self._reduce_grad_workspace_allocation = allocation + self._reduce_grad_prepared = prepared + self._reduce_grad_prepared_refs = ( + plan, + gradients, + *sources, + *source_registrations, + workspace, + allocation, + ) self._reduce_grad_signature = shape_signature self._reduce_grad_info = info return info + def _reduce_grad_source_regions( + self, gradients: ProjectionBuffers + ) -> tuple[tuple[Any, Any, Any], tuple[Any, Any, Any]]: + names = ("gate_reduce", "up_reduce", "down_reduce") + reduce_buffers = tuple(getattr(gradients, name) for name in names) + if all(value is None for value in reduce_buffers): + begin = self.context.expert_count + count = self.context.prefetch_slots + registrations = tuple( + getattr(gradients, name) for name in ("gate", "up", "down") + ) + sources = tuple( + getattr(gradients, name).narrow(0, begin, count) + for name in ("gate", "up", "down") + ) + return sources, registrations + sources = tuple( + value[self.context.planner_group_rank] for value in reduce_buffers + ) + return sources, reduce_buffers + + def _reduce_grad_sources(self, gradients: ProjectionBuffers) -> tuple[Any, Any, Any]: + return self._reduce_grad_source_regions(gradients)[0] + @property def reduce_grad_info(self): return self._reduce_grad_info @@ -1051,16 +1111,6 @@ def reduce_grad( down_reduce_buffer, ) self._validate_projection_buffers(gradients, reduce=True) - legacy_names = ("gate_reduce", "up_reduce", "down_reduce") - legacy_buffers = [getattr(gradients, name) for name in legacy_names] - if all(value is not None for value in legacy_buffers): - begin = self.context.expert_count - end = begin + self.context.prefetch_slots - for full_name, reduce_name in zip(("gate", "up", "down"), legacy_names): - getattr(gradients, full_name)[begin:end].copy_( - getattr(gradients, reduce_name)[self.context.planner_group_rank] - ) - self.runtime._acquire_reduce_grad(self._reduce_grad_owner_token) self._reduce_grad_token_held = True retained = False @@ -1068,14 +1118,14 @@ def reduce_grad( self.prepare_reduce_grad( plan, gradients, _owner_token=self._reduce_grad_owner_token ) - if self._reduce_grad_info is not None and self._reduce_grad_info.uses_udma: - self._synchronize_device() - self.runtime.register_reduce_grad_workspace( - self._reduce_grad_workspace, - self._reduce_grad_info.workspace_bytes, - owner_token=self._reduce_grad_owner_token, - ) - self._retain(plan, gradients, self._reduce_grad_workspace) + sources, source_registrations = self._reduce_grad_source_regions(gradients) + self._retain( + plan, + gradients, + sources, + source_registrations, + self._reduce_grad_workspace, + ) retained = True if all(existing is not plan for existing in self._pending_reduce_plans): self._pending_reduce_plans.append(plan) @@ -1083,15 +1133,13 @@ def reduce_grad( self.context, plan, gradients, - self._reduce_grad_workspace, + sources, + source_registrations, + self._reduce_grad_prepared, self._stream_ptr(), self.wait_iterations, - requested_udma_chunk_bytes=self.requested_udma_chunk_bytes, ) self._reduce_grad_inflight = True - if all(value is not None for value in legacy_buffers): - for value in legacy_buffers: - value[self.context.planner_group_rank].zero_() return self._record_event() if async_finish else None except Exception: if retained: @@ -1214,8 +1262,11 @@ def check_pending_status(self) -> None: expected, ) ) - for plan in self._pending_reduce_plans: - statuses.append((plan.epoch, int(plan.reduce_grad_status.item()), 0)) + reduce_statuses = [ + (plan.epoch, int(plan.reduce_grad_status.item()), 0) + for plan in self._pending_reduce_plans + ] + statuses.extend(reduce_statuses) self._pending_refs.clear() self._pending_plans.clear() self._pending_reduce_plans.clear() @@ -1228,7 +1279,9 @@ def check_pending_status(self) -> None: if status != expected ] if failed: - if any(status in (2005, 2006, 2007) for _, status, _ in failed): + if any(status in (2005, 2006, 2007) for _, status, _ in failed) or any( + status in (1, 2, 3, 4, 5) for _, status, _ in reduce_statuses + ): self.context.mark_poisoned() details = ", ".join( f"epoch {epoch}: actual {status}, expected {expected}" @@ -1256,7 +1309,9 @@ def close(self) -> None: sync_error = exc try: try: - self.runtime.unregister_reduce_grad_workspace() + self.runtime.destroy_reduce_grad(self._reduce_grad_prepared) + self._reduce_grad_prepared = None + self._reduce_grad_prepared_refs = () finally: try: if self._registered_projections is not None: @@ -1278,6 +1333,8 @@ def close(self) -> None: self._registered_projections = None self._reduce_grad_workspace = None self._reduce_grad_workspace_allocation = None + self._reduce_grad_prepared = None + self._reduce_grad_prepared_refs = () if sync_error is not None: raise sync_error diff --git a/src/comm/comm_wrap.cpp b/src/comm/comm_wrap.cpp index c6be0d32..efbad453 100644 --- a/src/comm/comm_wrap.cpp +++ b/src/comm/comm_wrap.cpp @@ -191,6 +191,42 @@ int TileXRUDMAUnregister(TileXRCommPtr comm, TileXRUDMAMemHandle handle) return c->UnregisterUDMAMemory(handle); } +int TileXRUDMAProfileRegister(TileXRCommPtr comm, + const TileXR::TileXRUDMAProfileDesc *desc, TileXRUDMAProfileHandle *handle) +{ + if (comm == nullptr || desc == nullptr || handle == nullptr) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister invalid input"; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + *handle = 0; + auto* c = static_cast(comm); + return c->RegisterUDMAProfile(*desc, handle); +} + +int TileXRUDMAProfileUnregister(TileXRCommPtr comm, TileXRUDMAProfileHandle handle) +{ + if (comm == nullptr || handle == 0) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileUnregister invalid input"; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + auto* c = static_cast(comm); + return c->UnregisterUDMAProfile(handle); +} + +int TileXRUDMAProfileQuery(TileXRCommPtr comm, TileXRUDMAProfileHandle handle, + TileXR::TileXRUDMAProfileView *view) +{ + if (view != nullptr) { + *view = TileXR::TileXRUDMAProfileView {}; + } + if (comm == nullptr || handle == 0 || view == nullptr) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileQuery invalid input"; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + auto* c = static_cast(comm); + return c->QueryUDMAProfile(handle, view); +} + int TileXRUDMAGetQpCount(TileXRCommPtr comm, uint32_t *qpCount) { if (qpCount != nullptr) { diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index c80338ce..f7418ffc 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -385,6 +385,43 @@ int TileXRComm::UnregisterUDMAMemory(TileXRUDMAMemHandle handle) return udmaContext_->UnregisterMemory(handle); } +int TileXRComm::RegisterUDMAProfile( + const TileXRUDMAProfileDesc &desc, TileXRUDMAProfileHandle *handle) +{ + if (!inited_) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister requires initialized communicator"; + return TILEXR_ERROR_NOT_INITIALIZED; + } + if (handle == nullptr) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (udmaContext_ == nullptr || !udmaContext_->IsAvailable()) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister called while UDMA is unavailable"; + return TILEXR_ERROR_NOT_SUPPORT; + } + return udmaContext_->RegisterProfile(desc, handle); +} + +int TileXRComm::UnregisterUDMAProfile(TileXRUDMAProfileHandle handle) +{ + if (handle == 0 || udmaContext_ == nullptr) { + return TILEXR_ERROR_NOT_FOUND; + } + return udmaContext_->UnregisterProfile(handle); +} + +int TileXRComm::QueryUDMAProfile( + TileXRUDMAProfileHandle handle, TileXRUDMAProfileView *view) const +{ + if (view == nullptr || handle == 0) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (udmaContext_ == nullptr || !udmaContext_->IsAvailable()) { + return TILEXR_ERROR_NOT_SUPPORT; + } + return udmaContext_->QueryProfile(handle, view); +} + int TileXRComm::GetUDMAQpCount(uint32_t *qpCount) const { if (qpCount == nullptr) { diff --git a/src/comm/tilexr_comm.h b/src/comm/tilexr_comm.h index 3e49d8ef..c56e0abf 100644 --- a/src/comm/tilexr_comm.h +++ b/src/comm/tilexr_comm.h @@ -47,6 +47,9 @@ class TileXRComm { int64_t NextMagic(); int RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMemHandle *handle); int UnregisterUDMAMemory(TileXRUDMAMemHandle handle); + int RegisterUDMAProfile(const TileXRUDMAProfileDesc &desc, TileXRUDMAProfileHandle *handle); + int UnregisterUDMAProfile(TileXRUDMAProfileHandle handle); + int QueryUDMAProfile(TileXRUDMAProfileHandle handle, TileXRUDMAProfileView *view) const; int GetUDMAQpCount(uint32_t *qpCount) const; GM_ADDR GetUDMARegistryPtr() const; const TileXRUDMARegistry* GetUDMARegistryHost() const; diff --git a/src/comm/udma/tilexr_udma_context.cpp b/src/comm/udma/tilexr_udma_context.cpp index b6a028df..12cca20f 100644 --- a/src/comm/udma/tilexr_udma_context.cpp +++ b/src/comm/udma/tilexr_udma_context.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -20,6 +21,12 @@ namespace TileXR { +struct TileXRUDMAContext::ProfileRecord { + TileXRUDMAProfileRegistry registry {}; + GM_ADDR registryDev = nullptr; + bool cleanupPending = false; +}; + TileXRUDMAContext::TileXRUDMAContext() = default; TileXRUDMAContext::~TileXRUDMAContext() @@ -94,15 +101,19 @@ void TileXRUDMAContext::Shutdown() TILEXR_LOG(ERROR) << "TileXR UDMA shutdown failed to clear comm args: " << publishRet; } + const int profileCleanupRet = CleanupAllProfiles(); int memoryCleanupRet = TILEXR_SUCCESS; if (transport_ != nullptr) { memoryCleanupRet = transport_->CleanupAllMemory(); } const int registryCleanupRet = CleanupAllRegistries(); - if (memoryCleanupRet != TILEXR_SUCCESS || registryCleanupRet != TILEXR_SUCCESS || + if (profileCleanupRet != TILEXR_SUCCESS || memoryCleanupRet != TILEXR_SUCCESS || + registryCleanupRet != TILEXR_SUCCESS || (transport_ != nullptr && transport_->HasMemoryCleanupPending()) || - udmaRegistryDev_ != nullptr || !retiredRegistryDevs_.empty()) { + (transport_ != nullptr && transport_->HasProfileCleanupPending()) || + udmaRegistryDev_ != nullptr || !retiredRegistryDevs_.empty() || !profiles_.empty()) { TILEXR_LOG(ERROR) << "TileXR UDMA shutdown retains memory cleanup state" + << ", profile ret " << profileCleanupRet << ", transport ret " << memoryCleanupRet << ", registry ret " << registryCleanupRet << ", registry ptr " << reinterpret_cast(udmaRegistryDev_) @@ -115,6 +126,8 @@ void TileXRUDMAContext::Shutdown() } retiredRegistryDevs_.clear(); + profiles_.clear(); + nextProfileHandle_ = 1; udmaRegistryDev_ = nullptr; registry_ = TileXRUDMARegistry {}; registeredPtr_ = nullptr; @@ -354,6 +367,216 @@ int TileXRUDMAContext::UnregisterMemory(TileXRUDMAMemHandle handle) return TILEXR_SUCCESS; } +int TileXRUDMAContext::RegisterProfile( + const TileXRUDMAProfileDesc& desc, TileXRUDMAProfileHandle* handle) +{ + if (handle == nullptr) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + *handle = 0; + if (!IsAvailable() || transport_ == nullptr) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister called while UDMA is unavailable"; + return TILEXR_ERROR_NOT_SUPPORT; + } + if (options_.threadMode) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister is not supported in InitThread mode"; + return TILEXR_ERROR_NOT_SUPPORT; + } + if (options_.exchange == nullptr) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister requires live socket exchange"; + return TILEXR_ERROR_INTERNAL; + } + if (transport_->HasProfileCleanupPending()) { + TILEXR_LOG(ERROR) << "TileXRUDMAProfileRegister requires pending profile cleanup first"; + return TILEXR_ERROR_INTERNAL; + } + + const uint32_t qpCount = transport_->GetQpCount(); + int localStatus = UDMAProfileDescValid(&desc, qpCount) + ? TILEXR_SUCCESS : TILEXR_ERROR_PARA_CHECK_FAIL; + int ret = AgreeStatus(localStatus); + if (ret != TILEXR_SUCCESS) { + return ret; + } + + std::vector allDescs(options_.rankSize); + ret = options_.exchange->AllGather(&desc, 1, allDescs.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } + localStatus = TILEXR_SUCCESS; + for (int rank = 0; rank < options_.rankSize; ++rank) { + if (!UDMAProfileDescValid(&allDescs[rank], qpCount) || + !UDMAProfileContractsEqual(allDescs[0], allDescs[rank])) { + TILEXR_LOG(ERROR) << "TileXR UDMA profile contract mismatch at rank " << rank; + localStatus = TILEXR_ERROR_PARA_CHECK_FAIL; + break; + } + } + ret = AgreeStatus(localStatus); + if (ret != TILEXR_SUCCESS) { + return ret; + } + + const TileXRUDMAProfileHandle candidateHandle = NextProfileHandle(); + localStatus = candidateHandle == 0 ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; + ret = AgreeStatus(localStatus); + if (ret != TILEXR_SUCCESS) { + return ret; + } + std::vector allHandles(options_.rankSize); + ret = options_.exchange->AllGather(&candidateHandle, 1, allHandles.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } + for (int rank = 0; rank < options_.rankSize; ++rank) { + if (allHandles[rank] != candidateHandle) { + localStatus = TILEXR_ERROR_INTERNAL; + break; + } + } + ret = AgreeStatus(localStatus); + if (ret != TILEXR_SUCCESS) { + return ret; + } + + ret = transport_->PrepareProfile(desc); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "TileXR UDMA profile memory registration failed: " << ret; + return ret; + } + + std::unique_ptr candidate(new (std::nothrow) ProfileRecord()); + localStatus = candidate == nullptr ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; + if (candidate != nullptr) { + candidate->registry.rankSize = static_cast(options_.rankSize); + candidate->registry.regionCount = desc.regionCount; + candidate->registry.qpCount = qpCount; + for (uint32_t qp = 0; qp < qpCount; ++qp) { + candidate->registry.qpBindings[qp] = desc.qpBindings[qp]; + } + for (int rank = 0; rank < options_.rankSize; ++rank) { + for (uint32_t region = 0; region < desc.regionCount; ++region) { + const size_t index = static_cast(rank) * + TILEXR_UDMA_PROFILE_MAX_REGIONS + region; + candidate->registry.regions[index] = allDescs[rank].regions[region]; + candidate->registry.regions[index].registrationBase = nullptr; + candidate->registry.regions[index].registrationBytes = 0; + } + } + const aclError allocRet = aclrtMalloc(reinterpret_cast(&candidate->registryDev), + sizeof(candidate->registry), ACL_MEM_MALLOC_HUGE_FIRST); + if (allocRet != ACL_SUCCESS) { + TILEXR_LOG(ERROR) << "aclrtMalloc UDMA profile registry failed: " << allocRet; + localStatus = TILEXR_ERROR_INTERNAL; + } + } + if (localStatus == TILEXR_SUCCESS) { + const aclError copyRet = aclrtMemcpy(candidate->registryDev, sizeof(candidate->registry), + &candidate->registry, sizeof(candidate->registry), ACL_MEMCPY_HOST_TO_DEVICE); + if (copyRet != ACL_SUCCESS) { + TILEXR_LOG(ERROR) << "aclrtMemcpy UDMA profile registry failed: " << copyRet; + localStatus = TILEXR_ERROR_INTERNAL; + } + } + + int agreedStatus = AgreeStatus(localStatus); + if (agreedStatus != TILEXR_SUCCESS) { + int cleanupRet = transport_->AbortPreparedProfile(); + if (candidate != nullptr) { + const int registryRet = FreeDeviceRegistry(candidate->registryDev); + if (registryRet != TILEXR_SUCCESS) { + RetainRegistry(candidate->registryDev); + } + if (cleanupRet == TILEXR_SUCCESS && registryRet != TILEXR_SUCCESS) { + cleanupRet = registryRet; + } + } + const int agreedCleanupStatus = AgreeStatus(cleanupRet); + if (agreedCleanupStatus != TILEXR_SUCCESS || transport_->HasProfileCleanupPending()) { + TILEXR_LOG(ERROR) << "TileXR UDMA profile candidate cleanup remains pending"; + } + return agreedStatus; + } + + localStatus = transport_->GetPreparedProfileInfoDev() == nullptr + ? TILEXR_ERROR_NOT_INITIALIZED + : transport_->CommitPreparedProfile(candidateHandle); + agreedStatus = AgreeStatus(localStatus); + if (agreedStatus != TILEXR_SUCCESS) { + int cleanupRet = localStatus == TILEXR_SUCCESS + ? transport_->CleanupProfile(candidateHandle) + : transport_->AbortPreparedProfile(); + const int registryRet = FreeDeviceRegistry(candidate->registryDev); + if (registryRet != TILEXR_SUCCESS) { + RetainRegistry(candidate->registryDev); + } + if (cleanupRet == TILEXR_SUCCESS && registryRet != TILEXR_SUCCESS) { + cleanupRet = registryRet; + } + (void)AgreeStatus(cleanupRet); + return agreedStatus; + } + + profiles_.emplace(candidateHandle, std::move(candidate)); + nextProfileHandle_ = candidateHandle == std::numeric_limits::max() + ? 1U : candidateHandle + 1U; + *handle = candidateHandle; + return TILEXR_SUCCESS; +} + +int TileXRUDMAContext::UnregisterProfile(TileXRUDMAProfileHandle handle) +{ + const auto it = profiles_.find(handle); + if (handle == 0 || it == profiles_.end() || transport_ == nullptr) { + return TILEXR_ERROR_NOT_FOUND; + } + ProfileRecord& record = *it->second; + const bool retrying = record.cleanupPending; + record.cleanupPending = true; + int transportRet = transport_->CleanupProfile(handle); + if (retrying && transportRet == TILEXR_ERROR_NOT_FOUND) { + transportRet = TILEXR_SUCCESS; + } + const int registryRet = FreeDeviceRegistry(record.registryDev); + int ret = transportRet; + if (ret == TILEXR_SUCCESS && registryRet != TILEXR_SUCCESS) { + ret = registryRet; + } + if (ret == TILEXR_SUCCESS) { + profiles_.erase(it); + } + return ret; +} + +int TileXRUDMAContext::QueryProfile( + TileXRUDMAProfileHandle handle, TileXRUDMAProfileView* view) const +{ + if (view == nullptr) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + *view = TileXRUDMAProfileView {}; + const auto it = profiles_.find(handle); + if (handle == 0 || it == profiles_.end() || it->second->cleanupPending || + transport_ == nullptr) { + return TILEXR_ERROR_NOT_FOUND; + } + const ProfileRecord& record = *it->second; + const GM_ADDR infoDev = transport_->GetProfileInfoDev(handle); + if (infoDev == nullptr || record.registryDev == nullptr || + !UDMAProfileRegistryValid(&record.registry, options_.rankSize, + record.registry.regionCount, transport_->GetQpCount())) { + return TILEXR_ERROR_NOT_INITIALIZED; + } + view->rankSize = record.registry.rankSize; + view->regionCount = record.registry.regionCount; + view->qpCount = transport_->GetQpCount(); + view->infoDev = infoDev; + view->registryDev = record.registryDev; + view->registryHost = &record.registry; + return TILEXR_SUCCESS; +} + GM_ADDR TileXRUDMAContext::GetRegistryDev() const { return lifecycle_ == Lifecycle::MemoryReady ? udmaRegistryDev_ : nullptr; @@ -489,6 +712,48 @@ int TileXRUDMAContext::CleanupAllRegistries() return firstError; } +int TileXRUDMAContext::CleanupAllProfiles() +{ + int firstError = TILEXR_SUCCESS; + for (auto it = profiles_.begin(); it != profiles_.end();) { + ProfileRecord& record = *it->second; + const bool retrying = record.cleanupPending; + record.cleanupPending = true; + int transportRet = transport_ == nullptr + ? TILEXR_ERROR_NOT_FOUND : transport_->CleanupProfile(it->first); + if (retrying && transportRet == TILEXR_ERROR_NOT_FOUND) { + transportRet = TILEXR_SUCCESS; + } + const int registryRet = FreeDeviceRegistry(record.registryDev); + int ret = transportRet; + if (ret == TILEXR_SUCCESS && registryRet != TILEXR_SUCCESS) { + ret = registryRet; + } + if (firstError == TILEXR_SUCCESS && ret != TILEXR_SUCCESS) { + firstError = ret; + } + if (ret == TILEXR_SUCCESS) { + it = profiles_.erase(it); + } else { + ++it; + } + } + return firstError; +} + +TileXRUDMAProfileHandle TileXRUDMAContext::NextProfileHandle() const +{ + TileXRUDMAProfileHandle candidate = nextProfileHandle_ == 0 ? 1U : nextProfileHandle_; + for (size_t checked = 0; checked <= profiles_.size(); ++checked) { + if (profiles_.count(candidate) == 0) { + return candidate; + } + candidate = candidate == std::numeric_limits::max() + ? 1U : candidate + 1U; + } + return 0; +} + void TileXRUDMAContext::RetainRegistry(GM_ADDR& registryDev) { if (registryDev == nullptr) { diff --git a/src/comm/udma/tilexr_udma_context.h b/src/comm/udma/tilexr_udma_context.h index 55ac0319..41eb8431 100644 --- a/src/comm/udma/tilexr_udma_context.h +++ b/src/comm/udma/tilexr_udma_context.h @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -58,12 +59,17 @@ class TileXRUDMAContext { int RegisterMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMemHandle* handle); int UnregisterMemory(TileXRUDMAMemHandle handle); + int RegisterProfile(const TileXRUDMAProfileDesc& desc, TileXRUDMAProfileHandle* handle); + int UnregisterProfile(TileXRUDMAProfileHandle handle); + int QueryProfile(TileXRUDMAProfileHandle handle, TileXRUDMAProfileView* view) const; GM_ADDR GetRegistryDev() const; const TileXRUDMARegistry* GetRegistryHost() const; uint32_t GetQpCount() const; private: + struct ProfileRecord; + enum class Lifecycle { Unavailable, TransportReady, @@ -77,8 +83,10 @@ class TileXRUDMAContext { int FreeDeviceRegistry(GM_ADDR& registryDev) const; int CleanupRetiredRegistries(); int CleanupAllRegistries(); + int CleanupAllProfiles(); void RetainRegistry(GM_ADDR& registryDev); void EnterCleanupPending(const char* reason); + TileXRUDMAProfileHandle NextProfileHandle() const; TileXRUDMAContextOptions options_ {}; Lifecycle lifecycle_ = Lifecycle::Unavailable; @@ -88,6 +96,8 @@ class TileXRUDMAContext { GM_ADDR registeredPtr_ = nullptr; size_t registeredBytes_ = 0; TileXRUDMARegistry registry_ {}; + std::map> profiles_; + TileXRUDMAProfileHandle nextProfileHandle_ = 1; std::unique_ptr transport_; }; diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 4633aa6c..3fd8700b 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -128,15 +128,21 @@ struct TileXRUDMATransport::SharedQpState { UDMACQCtx localCq {}; }; -struct TileXRUDMATransport::RegistrationState { +struct TileXRUDMATransport::RegisteredRegionState { GM_ADDR localPtr = nullptr; size_t bytes = 0; std::map localRegistrations; std::map, void*> remoteMemHandles; std::map localMemInfoByEid; +}; + +struct TileXRUDMATransport::RegistrationState { + std::vector regions; + std::vector qpBindings; std::vector memoryImage; GM_ADDR infoDev = nullptr; uint32_t infoSize = 0; + bool cleanupPending = false; }; TileXRUDMATransport::TileXRUDMATransport() = default; @@ -1362,7 +1368,8 @@ int TileXRUDMATransport::BuildRegistrationUDMAInfo(RegistrationState& registrati { const size_t entryCount = static_cast(options_.rankSize) * qpCount_; if (eidCount_ == 0 || eidTableDev_ == nullptr || - registration.memoryImage.size() != entryCount) { + registration.memoryImage.size() != entryCount || + registration.qpBindings.size() != qpCount_) { return TILEXR_ERROR_INTERNAL; } @@ -1401,8 +1408,13 @@ int TileXRUDMATransport::BuildRegistrationUDMAInfo(RegistrationState& registrati } else if (peer != options_.rank) { localEid = peerLocalEid_.at(peer); } - const auto registrationIt = registration.localRegistrations.find(localEid); - if (registrationIt == registration.localRegistrations.end()) { + const uint32_t localRegion = registration.qpBindings[qpIdx].localRegion; + if (localRegion >= registration.regions.size()) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + const auto& region = registration.regions[localRegion]; + const auto registrationIt = region.localRegistrations.find(localEid); + if (registrationIt == region.localRegistrations.end()) { return TILEXR_ERROR_NOT_INITIALIZED; } sq[index].localTokenId = registrationIt->second.tokenId; @@ -1438,58 +1450,72 @@ int TileXRUDMATransport::BuildRegistrationUDMAInfo(RegistrationState& registrati int TileXRUDMATransport::RegisterMemoryOnContexts(RegistrationState& registration) { - for (const auto& ctxEntry : ctxHandleByEid_) { - const uint32_t eidIndex = ctxEntry.first; - void* tokenHandle = tokenHandleByEid_[eidIndex]; - MrRegInfoT mrInfo {}; - mrInfo.in.mem.addr = reinterpret_cast(registration.localPtr); - mrInfo.in.mem.size = registration.bytes; - mrInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - mrInfo.in.ub.tokenIdHandle = tokenHandle; - mrInfo.in.ub.flags.bs.cacheable = 0; - mrInfo.in.ub.flags.bs.access = MEM_SEG_ACCESS_DEFAULT; - mrInfo.in.ub.flags.bs.nonPin = options_.nonPinRegistration ? 1 : 0; - mrInfo.in.ub.flags.bs.userIova = 0; - mrInfo.in.ub.flags.bs.tokenIdValid = 1; - mrInfo.in.ub.flags.bs.tokenPolicy = MEM_SEG_TOKEN_PLAIN_TEXT; - void* lmemHandle = nullptr; - int ret = loader_.RaCtxLmemRegister(ctxEntry.second, &mrInfo, &lmemHandle); - if (ret != 0 || lmemHandle == nullptr) { - constexpr uintptr_t twoMiB = UINT64_C(2) << 20; - TILEXR_LOG(ERROR) << "RaCtxLmemRegister failed for eid " << eidIndex - << ", ret " << ret - << ", bytes " << registration.bytes - << ", ptr " << reinterpret_cast(registration.localPtr) - << ", handle " << lmemHandle - << ", ptr modulo 2 MiB " - << (reinterpret_cast(registration.localPtr) % twoMiB); - return TILEXR_ERROR_INTERNAL; - } + for (size_t regionIndex = 0; regionIndex < registration.regions.size(); ++regionIndex) { + auto& region = registration.regions[regionIndex]; + for (const auto& ctxEntry : ctxHandleByEid_) { + const uint32_t eidIndex = ctxEntry.first; + void* tokenHandle = tokenHandleByEid_[eidIndex]; + MrRegInfoT mrInfo {}; + mrInfo.in.mem.addr = reinterpret_cast(region.localPtr); + mrInfo.in.mem.size = region.bytes; + mrInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + mrInfo.in.ub.tokenIdHandle = tokenHandle; + mrInfo.in.ub.flags.bs.cacheable = 0; + mrInfo.in.ub.flags.bs.access = MEM_SEG_ACCESS_DEFAULT; + mrInfo.in.ub.flags.bs.nonPin = options_.nonPinRegistration ? 1 : 0; + mrInfo.in.ub.flags.bs.userIova = 0; + mrInfo.in.ub.flags.bs.tokenIdValid = 1; + mrInfo.in.ub.flags.bs.tokenPolicy = MEM_SEG_TOKEN_PLAIN_TEXT; + void* lmemHandle = nullptr; + int ret = loader_.RaCtxLmemRegister(ctxEntry.second, &mrInfo, &lmemHandle); + if (ret != 0 || lmemHandle == nullptr) { + constexpr uintptr_t twoMiB = UINT64_C(2) << 20; + TILEXR_LOG(ERROR) << "RaCtxLmemRegister failed for region " << regionIndex + << ", eid " << eidIndex + << ", ret " << ret + << ", bytes " << region.bytes + << ", ptr " << reinterpret_cast(region.localPtr) + << ", handle " << lmemHandle + << ", ptr modulo 2 MiB " + << (reinterpret_cast(region.localPtr) % twoMiB); + return TILEXR_ERROR_INTERNAL; + } - RegMemResultInfo result {}; - result.address = reinterpret_cast(registration.localPtr); - result.size = registration.bytes; - result.lmemHandle = lmemHandle; - result.key = mrInfo.out.key; - result.tokenId = mrInfo.out.ub.tokenId; - result.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - result.targetSegHandle = mrInfo.out.ub.targetSegHandle; - result.tokenIdHandle = tokenHandle; - result.cacheable = 0; - result.access = MEM_SEG_ACCESS_DEFAULT; - registration.localRegistrations[eidIndex] = result; - registration.localMemInfoByEid[eidIndex] = BuildMemInfo(result); - } - return registration.localRegistrations.empty() ? TILEXR_ERROR_NOT_FOUND : TILEXR_SUCCESS; + RegMemResultInfo result {}; + result.address = reinterpret_cast(region.localPtr); + result.size = region.bytes; + result.lmemHandle = lmemHandle; + result.key = mrInfo.out.key; + result.tokenId = mrInfo.out.ub.tokenId; + result.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + result.targetSegHandle = mrInfo.out.ub.targetSegHandle; + result.tokenIdHandle = tokenHandle; + result.cacheable = 0; + result.access = MEM_SEG_ACCESS_DEFAULT; + region.localRegistrations[eidIndex] = result; + region.localMemInfoByEid[eidIndex] = BuildMemInfo(result); + } + } + return registration.regions.empty() ? TILEXR_ERROR_NOT_FOUND : TILEXR_SUCCESS; } int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration) { - if (registration.localRegistrations.empty()) { + if (registration.regions.empty() || registration.qpBindings.size() != qpCount_) { return TILEXR_ERROR_NOT_FOUND; } - const auto& localByEid = registration.localRegistrations; - uint32_t localCount = static_cast(localByEid.size()); + size_t localCountValue = 0; + for (const auto& region : registration.regions) { + if (region.localRegistrations.size() > + std::numeric_limits::max() - localCountValue) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + localCountValue += region.localRegistrations.size(); + } + if (localCountValue == 0 || localCountValue > std::numeric_limits::max()) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + const uint32_t localCount = static_cast(localCountValue); std::vector allCounts(options_.rankSize); int ret = options_.exchange->AllGather(&localCount, 1, allCounts.data()); if (ret != TILEXR_SUCCESS) { @@ -1501,6 +1527,7 @@ int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration } struct ExchangedMrInfo { + uint32_t regionIndex; uint32_t eidIndex; uint32_t valid; RegMemResultInfo mr; @@ -1508,11 +1535,14 @@ int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration std::vector local(maxCount); uint32_t idx = 0; - for (const auto& entry : localByEid) { - local[idx].eidIndex = entry.first; - local[idx].valid = 1; - local[idx].mr = entry.second; - ++idx; + for (uint32_t regionIndex = 0; regionIndex < registration.regions.size(); ++regionIndex) { + for (const auto& entry : registration.regions[regionIndex].localRegistrations) { + local[idx].regionIndex = regionIndex; + local[idx].eidIndex = entry.first; + local[idx].valid = 1; + local[idx].mr = entry.second; + ++idx; + } } if (static_cast(options_.rankSize) > std::numeric_limits::max() / maxCount) { @@ -1530,8 +1560,13 @@ int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration for (uint32_t qpIdx = 0; qpIdx < qpCount_; ++qpIdx) { const size_t selfIndex = RouteIndex(options_.rank, qpIdx); const uint32_t localEid = explicitConfig_ ? localRouteByPeerQp_[selfIndex] : fallbackEid; - const auto localMemIt = registration.localMemInfoByEid.find(localEid); - if (localMemIt == registration.localMemInfoByEid.end()) { + const uint32_t remoteRegion = registration.qpBindings[qpIdx].remoteRegion; + if (remoteRegion >= registration.regions.size()) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + const auto& region = registration.regions[remoteRegion]; + const auto localMemIt = region.localMemInfoByEid.find(localEid); + if (localMemIt == region.localMemInfoByEid.end()) { return TILEXR_ERROR_INTERNAL; } registration.memoryImage[selfIndex] = localMemIt->second; @@ -1549,10 +1584,15 @@ int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration ? localRouteByPeerQp_[imageIndex] : peerLocalEid_[peer]; const uint32_t remoteEid = explicitConfig_ ? remoteRouteByPeerQp_[imageIndex] : peerRemoteEid_[peer]; + const uint32_t remoteRegion = registration.qpBindings[qpIdx].remoteRegion; + if (remoteRegion >= registration.regions.size()) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } const ExchangedMrInfo* remote = nullptr; for (uint32_t i = 0; i < allCounts[peer]; ++i) { const auto& candidate = all[static_cast(peer) * maxCount + i]; - if (candidate.valid != 0 && candidate.eidIndex == remoteEid) { + if (candidate.valid != 0 && candidate.regionIndex == remoteRegion && + candidate.eidIndex == remoteEid) { remote = &candidate; break; } @@ -1565,7 +1605,8 @@ int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration return TILEXR_ERROR_INTERNAL; } const auto importKey = std::make_tuple(peer, localEid, remoteEid); - if (registration.remoteMemHandles.count(importKey) == 0) { + auto& region = registration.regions[remoteRegion]; + if (region.remoteMemHandles.count(importKey) == 0) { MrImportInfoT importInfo {}; importInfo.in.key = remote->mr.key; importInfo.in.ub.tokenValue = remote->mr.tokenValue; @@ -1580,7 +1621,7 @@ int TileXRUDMATransport::ExchangeAndImportMemory(RegistrationState& registration << ", ret " << ret << ", handle " << remoteHandle; return TILEXR_ERROR_INTERNAL; } - registration.remoteMemHandles[importKey] = remoteHandle; + region.remoteMemHandles[importKey] = remoteHandle; } uint32_t tpn = 0; @@ -1634,45 +1675,66 @@ int TileXRUDMATransport::AgreeRegistrationStatus(int localStatus) const return TILEXR_SUCCESS; } -int TileXRUDMATransport::PrepareMemory(GM_ADDR localPtr, size_t bytes) +int TileXRUDMATransport::PrepareRegistration(const TileXRUDMAProfileDesc& desc, + std::unique_ptr& registration) { - if (!available_ || localPtr == nullptr || bytes == 0) { - return TILEXR_ERROR_NOT_FOUND; - } - if (preparedRegistration_ != nullptr || !retiredRegistrations_.empty()) { - TILEXR_LOG(ERROR) << "TileXR UDMA cannot prepare memory while cleanup is pending"; - return TILEXR_ERROR_INTERNAL; + if (!available_ || !UDMAProfileDescValid(&desc, qpCount_) || registration != nullptr) { + return TILEXR_ERROR_PARA_CHECK_FAIL; } - preparedRegistration_.reset(new (std::nothrow) RegistrationState()); - int localStatus = preparedRegistration_ == nullptr ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; - if (preparedRegistration_ != nullptr) { - preparedRegistration_->localPtr = localPtr; - preparedRegistration_->bytes = bytes; - localStatus = RegisterMemoryOnContexts(*preparedRegistration_); + registration.reset(new (std::nothrow) RegistrationState()); + int localStatus = registration == nullptr ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; + if (registration != nullptr) { + registration->regions.resize(desc.regionCount); + for (uint32_t region = 0; region < desc.regionCount; ++region) { + registration->regions[region].localPtr = + UDMAProfileRegistrationBase(desc.regions[region]); + registration->regions[region].bytes = static_cast( + UDMAProfileRegistrationBytes(desc.regions[region])); + } + registration->qpBindings.assign(desc.qpBindings, + desc.qpBindings + desc.qpBindingCount); + localStatus = RegisterMemoryOnContexts(*registration); } int agreedStatus = AgreeRegistrationStatus(localStatus); if (agreedStatus != TILEXR_SUCCESS) { - const int cleanupRet = AbortPreparedMemory(); + const int cleanupRet = CleanupRegistrationPtr(registration); return cleanupRet == TILEXR_SUCCESS ? agreedStatus : cleanupRet; } - localStatus = ExchangeAndImportMemory(*preparedRegistration_); + localStatus = ExchangeAndImportMemory(*registration); agreedStatus = AgreeRegistrationStatus(localStatus); if (agreedStatus != TILEXR_SUCCESS) { - const int cleanupRet = AbortPreparedMemory(); + const int cleanupRet = CleanupRegistrationPtr(registration); return cleanupRet == TILEXR_SUCCESS ? agreedStatus : cleanupRet; } - localStatus = BuildRegistrationUDMAInfo(*preparedRegistration_); + localStatus = BuildRegistrationUDMAInfo(*registration); agreedStatus = AgreeRegistrationStatus(localStatus); if (agreedStatus != TILEXR_SUCCESS) { - const int cleanupRet = AbortPreparedMemory(); + const int cleanupRet = CleanupRegistrationPtr(registration); return cleanupRet == TILEXR_SUCCESS ? agreedStatus : cleanupRet; } return TILEXR_SUCCESS; } +int TileXRUDMATransport::PrepareMemory(GM_ADDR localPtr, size_t bytes) +{ + if (!available_ || localPtr == nullptr || bytes == 0) { + return TILEXR_ERROR_NOT_FOUND; + } + if (preparedRegistration_ != nullptr || !retiredRegistrations_.empty()) { + TILEXR_LOG(ERROR) << "TileXR UDMA cannot prepare memory while cleanup is pending"; + return TILEXR_ERROR_INTERNAL; + } + TileXRUDMAProfileDesc desc {}; + desc.regionCount = 1; + desc.qpBindingCount = qpCount_; + desc.regions[0].base = localPtr; + desc.regions[0].bytes = bytes; + return PrepareRegistration(desc, preparedRegistration_); +} + GM_ADDR TileXRUDMATransport::GetPreparedUDMAInfoDev() const { return preparedRegistration_ == nullptr ? nullptr : preparedRegistration_->infoDev; @@ -1696,6 +1758,35 @@ int TileXRUDMATransport::AbortPreparedMemory() return CleanupRegistrationPtr(preparedRegistration_); } +int TileXRUDMATransport::PrepareProfile(const TileXRUDMAProfileDesc& desc) +{ + if (preparedProfile_ != nullptr) { + TILEXR_LOG(ERROR) << "TileXR UDMA cannot prepare a profile while profile cleanup is pending"; + return TILEXR_ERROR_INTERNAL; + } + return PrepareRegistration(desc, preparedProfile_); +} + +GM_ADDR TileXRUDMATransport::GetPreparedProfileInfoDev() const +{ + return preparedProfile_ == nullptr ? nullptr : preparedProfile_->infoDev; +} + +int TileXRUDMATransport::CommitPreparedProfile(TileXRUDMAProfileHandle handle) +{ + if (handle == 0 || preparedProfile_ == nullptr || preparedProfile_->infoDev == nullptr || + profiles_.count(handle) != 0) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + profiles_.emplace(handle, std::move(preparedProfile_)); + return TILEXR_SUCCESS; +} + +int TileXRUDMATransport::AbortPreparedProfile() +{ + return CleanupRegistrationPtr(preparedProfile_); +} + int TileXRUDMATransport::CleanupLocalRegistrations(std::map& byEid) { int firstError = TILEXR_SUCCESS; @@ -1731,33 +1822,38 @@ int TileXRUDMATransport::CleanupLocalRegistrations(std::mapsecond == nullptr) { - it = registration.remoteMemHandles.erase(it); - continue; - } - const int peer = std::get<0>(it->first); - const uint32_t localEid = std::get<1>(it->first); - const auto ctxIt = ctxHandleByEid_.find(localEid); - if (ctxIt == ctxHandleByEid_.end() || ctxIt->second == nullptr) { - TILEXR_LOG(ERROR) << "Cannot unimport UDMA remote memory for peer " << peer - << " without local eid context " << localEid; - if (firstError == TILEXR_SUCCESS) { - firstError = TILEXR_ERROR_INTERNAL; + for (size_t regionIndex = 0; regionIndex < registration.regions.size(); ++regionIndex) { + auto& handles = registration.regions[regionIndex].remoteMemHandles; + for (auto it = handles.begin(); it != handles.end();) { + if (it->second == nullptr) { + it = handles.erase(it); + continue; } - ++it; - continue; - } - const int ret = loader_.RaCtxRmemUnimport(ctxIt->second, it->second); - if (ret != 0) { - TILEXR_LOG(ERROR) << "RaCtxRmemUnimport failed for peer " << peer << ", ret " << ret; - if (firstError == TILEXR_SUCCESS) { - firstError = TILEXR_ERROR_INTERNAL; + const int peer = std::get<0>(it->first); + const uint32_t localEid = std::get<1>(it->first); + const auto ctxIt = ctxHandleByEid_.find(localEid); + if (ctxIt == ctxHandleByEid_.end() || ctxIt->second == nullptr) { + TILEXR_LOG(ERROR) << "Cannot unimport UDMA remote memory for region " + << regionIndex << ", peer " << peer + << " without local eid context " << localEid; + if (firstError == TILEXR_SUCCESS) { + firstError = TILEXR_ERROR_INTERNAL; + } + ++it; + continue; } - ++it; - continue; + const int ret = loader_.RaCtxRmemUnimport(ctxIt->second, it->second); + if (ret != 0) { + TILEXR_LOG(ERROR) << "RaCtxRmemUnimport failed for region " << regionIndex + << ", peer " << peer << ", ret " << ret; + if (firstError == TILEXR_SUCCESS) { + firstError = TILEXR_ERROR_INTERNAL; + } + ++it; + continue; + } + it = handles.erase(it); } - it = registration.remoteMemHandles.erase(it); } return firstError; } @@ -1779,18 +1875,30 @@ int TileXRUDMATransport::FreeDeviceInfo(GM_ADDR& infoDev) const int TileXRUDMATransport::CleanupRegistration(RegistrationState& registration) { + registration.cleanupPending = true; int firstError = CleanupRemoteImports(registration); - const int localRet = CleanupLocalRegistrations(registration.localRegistrations); - if (firstError == TILEXR_SUCCESS && localRet != TILEXR_SUCCESS) { - firstError = localRet; + for (auto& region : registration.regions) { + const int localRet = CleanupLocalRegistrations(region.localRegistrations); + if (firstError == TILEXR_SUCCESS && localRet != TILEXR_SUCCESS) { + firstError = localRet; + } } const int infoRet = FreeDeviceInfo(registration.infoDev); if (firstError == TILEXR_SUCCESS && infoRet != TILEXR_SUCCESS) { firstError = infoRet; } - if (registration.remoteMemHandles.empty() && registration.localRegistrations.empty()) { - registration.localMemInfoByEid.clear(); + const bool resourcesClean = std::all_of(registration.regions.begin(), registration.regions.end(), + [](const RegisteredRegionState& region) { + return region.remoteMemHandles.empty() && region.localRegistrations.empty(); + }); + if (resourcesClean && registration.infoDev == nullptr) { + for (auto& region : registration.regions) { + region.localMemInfoByEid.clear(); + } + registration.regions.clear(); + registration.qpBindings.clear(); registration.memoryImage.clear(); + registration.cleanupPending = false; } return firstError; } @@ -1801,8 +1909,8 @@ int TileXRUDMATransport::CleanupRegistrationPtr(std::unique_ptrremoteMemHandles.empty() && - registration->localRegistrations.empty() && registration->infoDev == nullptr) { + if (ret == TILEXR_SUCCESS && registration->regions.empty() && + registration->infoDev == nullptr) { registration.reset(); } return ret; @@ -1816,8 +1924,7 @@ int TileXRUDMATransport::CleanupRetiredMemory() if (firstError == TILEXR_SUCCESS && ret != TILEXR_SUCCESS) { firstError = ret; } - if (ret == TILEXR_SUCCESS && (*it)->remoteMemHandles.empty() && - (*it)->localRegistrations.empty() && (*it)->infoDev == nullptr) { + if (ret == TILEXR_SUCCESS && (*it)->regions.empty() && (*it)->infoDev == nullptr) { it = retiredRegistrations_.erase(it); } else { ++it; @@ -1841,6 +1948,56 @@ int TileXRUDMATransport::CleanupAllMemory() return firstError; } +int TileXRUDMATransport::CleanupProfile(TileXRUDMAProfileHandle handle) +{ + const auto it = profiles_.find(handle); + if (handle == 0 || it == profiles_.end()) { + return TILEXR_ERROR_NOT_FOUND; + } + const int ret = CleanupRegistration(*it->second); + if (ret == TILEXR_SUCCESS && it->second->regions.empty() && + it->second->infoDev == nullptr) { + profiles_.erase(it); + } + return ret; +} + +int TileXRUDMATransport::CleanupAllProfiles() +{ + int firstError = CleanupRegistrationPtr(preparedProfile_); + for (auto it = profiles_.begin(); it != profiles_.end();) { + const int ret = CleanupRegistration(*it->second); + if (firstError == TILEXR_SUCCESS && ret != TILEXR_SUCCESS) { + firstError = ret; + } + if (ret == TILEXR_SUCCESS && it->second->regions.empty() && + it->second->infoDev == nullptr) { + it = profiles_.erase(it); + } else { + ++it; + } + } + return firstError; +} + +GM_ADDR TileXRUDMATransport::GetProfileInfoDev(TileXRUDMAProfileHandle handle) const +{ + const auto it = profiles_.find(handle); + return it == profiles_.end() || it->second->cleanupPending ? nullptr : it->second->infoDev; +} + +bool TileXRUDMATransport::HasProfileCleanupPending() const +{ + if (preparedProfile_ != nullptr && preparedProfile_->cleanupPending) { + return true; + } + return std::any_of(profiles_.begin(), profiles_.end(), + [](const std::pair>& entry) { + return entry.second != nullptr && entry.second->cleanupPending; + }); +} + int TileXRUDMATransport::UnregisterMemory(GM_ADDR localPtr) { if (localPtr == nullptr) { @@ -2070,6 +2227,11 @@ void TileXRUDMATransport::Shutdown() << ", ptr " << reinterpret_cast(GetRegisteredMemoryPtr()) << ", bytes " << GetRegisteredMemoryBytes(); } + const int profileCleanupRet = CleanupAllProfiles(); + if (profileCleanupRet != TILEXR_SUCCESS || HasProfileCleanupPending()) { + TILEXR_LOG(ERROR) << "TileXR UDMA shutdown retains persistent profile resources: " + << profileCleanupRet; + } int queueCleanupRet = CleanupQueues(); if (queueCleanupRet != TILEXR_SUCCESS) { TILEXR_LOG(ERROR) << "TileXR UDMA queue cleanup will be retried: " << queueCleanupRet; @@ -2121,15 +2283,15 @@ GM_ADDR TileXRUDMATransport::GetBaseUDMAInfoDev() const GM_ADDR TileXRUDMATransport::GetRegisteredMemoryPtr() const { - if (activeRegistration_ != nullptr) { - return activeRegistration_->localPtr; + if (activeRegistration_ != nullptr && !activeRegistration_->regions.empty()) { + return activeRegistration_->regions[0].localPtr; } - if (preparedRegistration_ != nullptr) { - return preparedRegistration_->localPtr; + if (preparedRegistration_ != nullptr && !preparedRegistration_->regions.empty()) { + return preparedRegistration_->regions[0].localPtr; } for (const auto& registration : retiredRegistrations_) { - if (registration != nullptr) { - return registration->localPtr; + if (registration != nullptr && !registration->regions.empty()) { + return registration->regions[0].localPtr; } } return nullptr; @@ -2137,15 +2299,15 @@ GM_ADDR TileXRUDMATransport::GetRegisteredMemoryPtr() const size_t TileXRUDMATransport::GetRegisteredMemoryBytes() const { - if (activeRegistration_ != nullptr) { - return activeRegistration_->bytes; + if (activeRegistration_ != nullptr && !activeRegistration_->regions.empty()) { + return activeRegistration_->regions[0].bytes; } - if (preparedRegistration_ != nullptr) { - return preparedRegistration_->bytes; + if (preparedRegistration_ != nullptr && !preparedRegistration_->regions.empty()) { + return preparedRegistration_->regions[0].bytes; } for (const auto& registration : retiredRegistrations_) { - if (registration != nullptr) { - return registration->bytes; + if (registration != nullptr && !registration->regions.empty()) { + return registration->regions[0].bytes; } } return 0; diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index efecd812..8a50e682 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -51,6 +51,14 @@ class TileXRUDMATransport { int CleanupRetiredMemory(); int CleanupAllMemory(); int UnregisterMemory(GM_ADDR localPtr); + int PrepareProfile(const TileXRUDMAProfileDesc& desc); + GM_ADDR GetPreparedProfileInfoDev() const; + int CommitPreparedProfile(TileXRUDMAProfileHandle handle); + int AbortPreparedProfile(); + int CleanupProfile(TileXRUDMAProfileHandle handle); + int CleanupAllProfiles(); + GM_ADDR GetProfileInfoDev(TileXRUDMAProfileHandle handle) const; + bool HasProfileCleanupPending() const; void Shutdown(); bool IsAvailable() const; @@ -66,6 +74,7 @@ class TileXRUDMATransport { struct PerEidState; struct PerPeerQpState; struct SharedQpState; + struct RegisteredRegionState; struct RegistrationState; int AgreeInitStatus(int localStatus) const; @@ -90,6 +99,8 @@ class TileXRUDMATransport { int BuildRegistrationUDMAInfo(RegistrationState& registration); int RegisterMemoryOnContexts(RegistrationState& registration); int ExchangeAndImportMemory(RegistrationState& registration); + int PrepareRegistration(const TileXRUDMAProfileDesc& desc, + std::unique_ptr& registration); int AgreeRegistrationStatus(int localStatus) const; int CleanupLocalRegistrations(std::map& byEid); int CleanupRemoteImports(RegistrationState& registration); @@ -130,6 +141,8 @@ class TileXRUDMATransport { std::unique_ptr activeRegistration_; std::unique_ptr preparedRegistration_; std::vector> retiredRegistrations_; + std::unique_ptr preparedProfile_; + std::map> profiles_; GM_ADDR udmaInfoDev_ = nullptr; GM_ADDR baseUDMAInfoDev_ = nullptr; GM_ADDR eidTableDev_ = nullptr; diff --git a/src/include/tilexr_api.h b/src/include/tilexr_api.h index 30a9a68c..8c875b67 100644 --- a/src/include/tilexr_api.h +++ b/src/include/tilexr_api.h @@ -20,6 +20,7 @@ extern "C" { typedef void *TileXRCommPtr; typedef uint32_t TileXRUDMAMemHandle; +typedef uint32_t TileXRUDMAProfileHandle; #define TILEXRUNIQUE_ID_BYTES 128 typedef struct { char internal[TILEXRUNIQUE_ID_BYTES]; } TileXRUniqueId; @@ -46,6 +47,14 @@ int TileXRUDMARegister(TileXRCommPtr comm, GM_ADDR localPtr, size_t bytes, TileX int TileXRUDMAUnregister(TileXRCommPtr comm, TileXRUDMAMemHandle handle); +int TileXRUDMAProfileRegister(TileXRCommPtr comm, const TileXR::TileXRUDMAProfileDesc *desc, + TileXRUDMAProfileHandle *handle); + +int TileXRUDMAProfileUnregister(TileXRCommPtr comm, TileXRUDMAProfileHandle handle); + +int TileXRUDMAProfileQuery(TileXRCommPtr comm, TileXRUDMAProfileHandle handle, + TileXR::TileXRUDMAProfileView *view); + int TileXRUDMAGetQpCount(TileXRCommPtr comm, uint32_t *qpCount); int TileXRGetUDMARegistryDev(TileXRCommPtr comm, GM_ADDR ®istryPtr); diff --git a/src/include/tilexr_moonep.h b/src/include/tilexr_moonep.h index f89630b8..56c982bf 100644 --- a/src/include/tilexr_moonep.h +++ b/src/include/tilexr_moonep.h @@ -16,7 +16,8 @@ typedef void *TileXRCommPtr; #define TILEXR_MOONEP_ABI_VERSION_V2 UINT32_C(2) #define TILEXR_MOONEP_MAX_TENSOR_RANK UINT32_C(4) #define TILEXR_MOONEP_FLAG_NONE UINT64_C(0) -#define TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES UINT64_C(1048576) +#define TILEXR_MOONEP_REDUCE_GRAD_DEFAULT_CHUNK_BYTES UINT64_C(8388608) +#define TILEXR_MOONEP_REDUCE_GRAD_WORKSPACE_ALIGNMENT UINT64_C(2097152) #define TILEXR_MOONEP_FLAG_BUILD_DEDUP (UINT64_C(1) << 0) #define TILEXR_MOONEP_FLAG_SKIP_INTER_RANK_SYNC (UINT64_C(1) << 1) #define TILEXR_MOONEP_FLAG_ZERO_COPY (UINT64_C(1) << 2) @@ -45,16 +46,10 @@ typedef enum TileXRMoonEpStage { TILEXR_MOONEP_STAGE_REDUCE_GRAD = 1u << 4 } TileXRMoonEpStage; -typedef enum TileXRMoonEpReduceGradTransport { - TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_NONE = 0, - TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER = 1, - TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA = 2 -} TileXRMoonEpReduceGradTransport; - typedef enum TileXRMoonEpReduceGradDeviceStatus { TILEXR_MOONEP_REDUCE_GRAD_DEVICE_SUCCESS = 0, TILEXR_MOONEP_REDUCE_GRAD_DEVICE_INVALID_STATE = 1, - TILEXR_MOONEP_REDUCE_GRAD_DEVICE_PEER_TIMEOUT = 2, + TILEXR_MOONEP_REDUCE_GRAD_DEVICE_UDMA_TIMEOUT = 2, TILEXR_MOONEP_REDUCE_GRAD_DEVICE_UDMA_CQ_ERROR = 3 } TileXRMoonEpReduceGradDeviceStatus; @@ -173,15 +168,28 @@ typedef struct TileXRMoonEpReduceGradWorkspaceInfoV2 { uint64_t workspaceBytes; uint64_t workspaceAlignment; uint64_t udmaChunkBytes; - uint64_t peerWindowBytes; - uint64_t peerHalfBytes; - uint64_t peerSlotStrideBytes; + uint64_t laneStateBytes; + uint64_t laneStateStrideBytes; + uint64_t bankStrideBytes; + uint64_t laneStrideBytes; uint64_t rowBytes[3]; - uint32_t transports[3]; + uint64_t chunkCounts[3]; + uint32_t projectionQpCounts[3]; + uint32_t qpCount; uint32_t blockDim; + uint32_t reserved; } TileXRMoonEpReduceGradWorkspaceInfoV2; -typedef struct TileXRMoonEpReduceGradArgsV2 { +typedef struct TileXRMoonEpReduceGradSourceSliceV2 { + void *data; + uint64_t bytes; + void *registrationBase; + uint64_t registrationBytes; +} TileXRMoonEpReduceGradSourceSliceV2; + +typedef void *TileXRMoonEpReduceGradPreparedV2; + +typedef struct TileXRMoonEpReduceGradPrepareArgsV2 { uint32_t structSize; uint32_t abiVersion; TileXRCommPtr comm; @@ -189,11 +197,24 @@ typedef struct TileXRMoonEpReduceGradArgsV2 { TileXRMoonEpTensorV1 *gate; TileXRMoonEpTensorV1 *up; TileXRMoonEpTensorV1 *down; + TileXRMoonEpReduceGradSourceSliceV2 sources[3]; void *workspace; uint64_t workspaceBytes; + uint64_t requestedUdmaChunkBytes; + uint64_t flags; +} TileXRMoonEpReduceGradPrepareArgsV2; + +typedef struct TileXRMoonEpReduceGradArgsV2 { + uint32_t structSize; + uint32_t abiVersion; + TileXRMoonEpReduceGradPreparedV2 prepared; + const TileXRMoonEpPlanV1 *plan; + TileXRMoonEpTensorV1 *gate; + TileXRMoonEpTensorV1 *up; + TileXRMoonEpTensorV1 *down; + TileXRMoonEpReduceGradSourceSliceV2 sources[3]; TileXRMoonEpTensorV1 *status; uint64_t waitIterations; - uint64_t requestedUdmaChunkBytes; uint64_t flags; } TileXRMoonEpReduceGradArgsV2; @@ -233,12 +254,21 @@ int TileXRMoonEpReduceGradGetWorkspaceSizeV2( TileXRMoonEpReduceGradWorkspaceInfoV2 *info); /* - * Enqueues asynchronously and is non-reentrant per communicator. Successful - * calls on one stream may be queued in order. Before switching streams, - * unregistering workspace, destroying the communicator, or releasing device - * buffers referenced by args, the caller must synchronize the last used stream. - * If the status reset is queued but kernel launch fails, this function attempts - * to drain that stream; a drain failure is reported as a runtime error. + * Collective for multi-rank communicators. Registers one persistent UDMA + * profile and returns an immutable prepared handle. All referenced allocations + * must remain alive until the handle is destroyed after stream quiescence. + */ +int TileXRMoonEpReduceGradPrepareV2( + const TileXRMoonEpReduceGradPrepareArgsV2 *args, + TileXRMoonEpReduceGradPreparedV2 *prepared); + +int TileXRMoonEpReduceGradDestroyPreparedV2( + TileXRMoonEpReduceGradPreparedV2 prepared); + +/* + * Enqueues asynchronously from an immutable prepared handle. The hot path does + * not synchronize or register memory. The caller must quiesce before destroying + * the handle or releasing any referenced allocation. */ int TileXRMoonEpReduceGradV2(const TileXRMoonEpReduceGradArgsV2 *args, aclrtStream stream); diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index ae470f02..f585e93b 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -104,6 +104,65 @@ __aicore__ inline __gm__ TileXRUDMARegistry* GetUDMARegistry(const __gm__ CommAr return reinterpret_cast<__gm__ TileXRUDMARegistry*>(args->udmaRegistryPtr); } +__aicore__ inline bool UDMAProfileRegistryValid( + const __gm__ CommArgs* args, const __gm__ UDMAInfo* udmaInfo, + const __gm__ TileXRUDMAProfileRegistry* registry) +{ + if (args == nullptr || (args->extraFlag & ExtraFlag::UDMA) == 0U || + udmaInfo == nullptr || registry == nullptr || + registry->magic != TILEXR_UDMA_PROFILE_MAGIC || + registry->version != TILEXR_UDMA_PROFILE_VERSION || + args->rankSize <= 0 || args->rankSize > TILEXR_MAX_RANK_SIZE || + args->rank < 0 || args->rank >= args->rankSize || + registry->rankSize != static_cast(args->rankSize) || + registry->regionCount == 0 || + registry->regionCount > TILEXR_UDMA_PROFILE_MAX_REGIONS || + registry->qpCount == 0 || registry->qpCount > TILEXR_UDMA_DEVICE_MAX_QP_COUNT || + udmaInfo->qpNum != registry->qpCount || udmaInfo->sqPtr == 0U || + udmaInfo->memPtr == 0U) { + return false; + } + return true; +} + +__aicore__ inline const __gm__ TileXRUDMAProfileRegionDesc* UDMAProfileGetRegion( + const __gm__ TileXRUDMAProfileRegistry* registry, int rank, uint32_t region) +{ + if (registry == nullptr || rank < 0 || static_cast(rank) >= registry->rankSize || + region >= registry->regionCount || registry->rankSize > TILEXR_MAX_RANK_SIZE || + registry->regionCount > TILEXR_UDMA_PROFILE_MAX_REGIONS) { + return nullptr; + } + const uint64_t index = static_cast(rank) * + TILEXR_UDMA_PROFILE_MAX_REGIONS + region; + return ®istry->regions[index]; +} + +__aicore__ inline bool UDMAProfileRegisteredRangeValid( + const __gm__ TileXRUDMAProfileRegistry* registry, int rank, uint32_t region, + uint64_t byteOffset, uint64_t byteCount) +{ + const __gm__ TileXRUDMAProfileRegionDesc* desc = + UDMAProfileGetRegion(registry, rank, region); + if (desc == nullptr || desc->base == nullptr || desc->bytes == 0 || + byteOffset > desc->bytes) { + return false; + } + return byteCount <= desc->bytes - byteOffset; +} + +__aicore__ inline __gm__ uint8_t* UDMAProfileRegisteredAddr( + const __gm__ TileXRUDMAProfileRegistry* registry, int rank, uint32_t region, + uint64_t byteOffset) +{ + const __gm__ TileXRUDMAProfileRegionDesc* desc = + UDMAProfileGetRegion(registry, rank, region); + if (desc == nullptr || desc->base == nullptr || byteOffset > desc->bytes) { + return nullptr; + } + return reinterpret_cast<__gm__ uint8_t*>(desc->base + byteOffset); +} + __aicore__ inline bool UDMARegisteredRangeValid( const __gm__ TileXRUDMARegistry* registry, int targetRank, uint64_t byteOffset, uint64_t byteCount) { @@ -534,15 +593,20 @@ __attribute__((always_inline)) inline __aicore__ uint32_t UDMAWrite( localAddr, pe, qpIdx, messageLen, nullptr, sqeFlag); } -__aicore__ inline uint32_t UDMARead( +__attribute__((always_inline)) inline __aicore__ uint32_t UDMARead( const __gm__ CommArgs* args, const AscendC::LocalTensor& wqeScratch, __gm__ uint8_t* localAddr, __gm__ uint8_t* remoteAddr, - uint32_t pe, uint32_t qpIdx, uint64_t messageLen) + uint32_t pe, uint32_t qpIdx, uint64_t messageLen, bool ringDoorbell = true) { if (TILEXR_UDMA_ARCH_SUPPORTED && UDMAQueueOperationValid( args, static_cast(pe), qpIdx)) { - return UDMAPostSend(GetUDMAInfo(args), wqeScratch, remoteAddr, - localAddr, pe, qpIdx, messageLen, nullptr, + if (ringDoorbell) { + return UDMAPostSend(GetUDMAInfo(args), wqeScratch, + remoteAddr, localAddr, pe, qpIdx, messageLen, nullptr, + TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION); + } + return UDMAPostSend(GetUDMAInfo(args), wqeScratch, + remoteAddr, localAddr, pe, qpIdx, messageLen, nullptr, TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION); } return TILEXR_UDMA_STATUS_INVALID; @@ -682,6 +746,113 @@ __aicore__ inline void UDMAGetRegisteredNbi( UDMAGetNbi(args, wqeScratch, sourceRank, localDst, byteOffset, byteCount); } +__attribute__((always_inline)) inline __aicore__ bool UDMAProfileGetOperationValid( + const __gm__ CommArgs* args, const __gm__ UDMAInfo* udmaInfo, + const __gm__ TileXRUDMAProfileRegistry* registry, + int sourceRank, uint32_t qpIdx, uint32_t localRegion, uint64_t localByteOffset, + uint32_t remoteRegion, uint64_t remoteByteOffset, uint64_t byteCount) +{ + if (!TILEXR_UDMA_ARCH_SUPPORTED || + !UDMAProfileRegistryValid(args, udmaInfo, registry) || + sourceRank < 0 || sourceRank >= args->rankSize || sourceRank == args->rank || + qpIdx >= registry->qpCount || byteCount == 0U || byteCount > 0xFFFFFFFFULL) { + return false; + } + const __gm__ TileXRUDMAProfileQpBinding* binding = ®istry->qpBindings[qpIdx]; + if (binding->localRegion != localRegion || binding->remoteRegion != remoteRegion) { + return false; + } + return UDMAProfileRegisteredRangeValid(registry, args->rank, localRegion, + localByteOffset, byteCount) && + UDMAProfileRegisteredRangeValid(registry, sourceRank, remoteRegion, + remoteByteOffset, byteCount); +} + +__attribute__((always_inline)) inline __aicore__ uint32_t UDMAProfileGetNbiOnQpDeferred( + const __gm__ CommArgs* args, __gm__ UDMAInfo* udmaInfo, + const __gm__ TileXRUDMAProfileRegistry* registry, + const AscendC::LocalTensor& wqeScratch, + int sourceRank, uint32_t qpIdx, uint32_t localRegion, uint64_t localByteOffset, + uint32_t remoteRegion, uint64_t remoteByteOffset, uint64_t byteCount) +{ + if (!UDMAProfileGetOperationValid(args, udmaInfo, registry, sourceRank, qpIdx, + localRegion, localByteOffset, remoteRegion, remoteByteOffset, byteCount)) { + return TILEXR_UDMA_STATUS_INVALID; + } + __gm__ uint8_t* localAddr = UDMAProfileRegisteredAddr( + registry, args->rank, localRegion, localByteOffset); + __gm__ uint8_t* remoteAddr = UDMAProfileRegisteredAddr( + registry, sourceRank, remoteRegion, remoteByteOffset); + return UDMAPostSend(udmaInfo, wqeScratch, + remoteAddr, localAddr, static_cast(sourceRank), qpIdx, byteCount, + nullptr, TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION); +} + +__aicore__ inline uint32_t UDMAProfileCompletionFrontier( + const __gm__ CommArgs* args, __gm__ UDMAInfo* udmaInfo, + const __gm__ TileXRUDMAProfileRegistry* registry, int sourceRank, uint32_t qpIdx) +{ + if (!UDMAProfileRegistryValid(args, udmaInfo, registry) || sourceRank < 0 || + sourceRank >= args->rankSize || sourceRank == args->rank || qpIdx >= registry->qpCount) { + return 0U; + } + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx( + udmaInfo, static_cast(sourceRank), qpIdx); + if (qpCtxEntry == nullptr || qpCtxEntry->wqeCntAddr == 0U) { + return 0U; + } + return ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); +} + +__aicore__ inline uint32_t UDMAProfileFlushQpDoorbell( + const __gm__ CommArgs* args, __gm__ UDMAInfo* udmaInfo, + const __gm__ TileXRUDMAProfileRegistry* registry, int sourceRank, uint32_t qpIdx) +{ + if (!UDMAProfileRegistryValid(args, udmaInfo, registry) || sourceRank < 0 || + sourceRank >= args->rankSize || sourceRank == args->rank || qpIdx >= registry->qpCount) { + return TILEXR_UDMA_STATUS_INVALID; + } + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx( + udmaInfo, static_cast(sourceRank), qpIdx); + if (qpCtxEntry == nullptr || qpCtxEntry->headAddr == 0U || + qpCtxEntry->tailAddr == 0U || qpCtxEntry->dbAddr == 0U || + qpCtxEntry->depth != TILEXR_UDMA_SQ_BB_COUNT) { + return TILEXR_UDMA_STATUS_INVALID; + } + const uint32_t head = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->headAddr), 0); + const uint32_t tail = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->tailAddr), 0); + if (head - tail > qpCtxEntry->depth) { + return TILEXR_UDMA_STATUS_INVALID; + } + UDMARingDoorbell(head, qpCtxEntry); + return TILEXR_UDMA_STATUS_SUCCESS; +} + +__aicore__ inline uint32_t UDMAProfileQuietStatusOnQpUntil( + const __gm__ CommArgs* args, __gm__ UDMAInfo* udmaInfo, + const __gm__ TileXRUDMAProfileRegistry* registry, + int sourceRank, uint32_t qpIdx, uint32_t completionFrontier) +{ + if (!UDMAProfileRegistryValid(args, udmaInfo, registry) || sourceRank < 0 || + sourceRank >= args->rankSize || sourceRank == args->rank || qpIdx >= registry->qpCount || + udmaInfo->scqPtr == 0U) { + return TILEXR_UDMA_STATUS_INVALID; + } + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx( + udmaInfo, static_cast(sourceRank), qpIdx); + if (qpCtxEntry == nullptr || qpCtxEntry->wqeCntAddr == 0U || + qpCtxEntry->depth != TILEXR_UDMA_SQ_BB_COUNT) { + return TILEXR_UDMA_STATUS_INVALID; + } + const uint32_t submitted = ld_dev( + reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + if (submitted - completionFrontier > qpCtxEntry->depth) { + return TILEXR_UDMA_STATUS_INVALID; + } + return UDMAPollCQ(udmaInfo, static_cast(sourceRank), qpIdx, + completionFrontier); +} + template __aicore__ inline uint32_t UDMAPutSignalNbiOnQp( const __gm__ CommArgs* args, const AscendC::LocalTensor& wqeScratch, diff --git a/src/include/tilexr_udma_reg.h b/src/include/tilexr_udma_reg.h index ccdca377..6b2d081c 100644 --- a/src/include/tilexr_udma_reg.h +++ b/src/include/tilexr_udma_reg.h @@ -8,6 +8,7 @@ #include #include +#include #include "comm_args.h" @@ -16,6 +17,10 @@ namespace TileXR { constexpr uint32_t TILEXR_UDMA_REGISTRY_MAGIC = 0x54585255U; // TXRU constexpr uint32_t TILEXR_UDMA_REGISTRY_VERSION = 1U; constexpr uint32_t TILEXR_UDMA_MAX_REGIONS = 1U; +constexpr uint32_t TILEXR_UDMA_PROFILE_MAGIC = 0x54585052U; // TXPR +constexpr uint32_t TILEXR_UDMA_PROFILE_VERSION = 1U; +constexpr uint32_t TILEXR_UDMA_PROFILE_MAX_REGIONS = 8U; +constexpr uint32_t TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS = 32U; struct TileXRUDMARegionDesc { GM_ADDR base = nullptr; @@ -30,6 +35,178 @@ struct TileXRUDMARegistry { TileXRUDMARegionDesc regions[TILEXR_MAX_RANK_SIZE] = {}; }; +struct TileXRUDMAProfileRegionDesc { + GM_ADDR base = nullptr; + uint64_t bytes = 0; + GM_ADDR registrationBase = nullptr; + uint64_t registrationBytes = 0; +}; + +struct TileXRUDMAProfileQpBinding { + uint32_t localRegion = 0; + uint32_t remoteRegion = 0; +}; + +struct TileXRUDMAProfileDesc { + uint32_t version = TILEXR_UDMA_PROFILE_VERSION; + uint32_t regionCount = 0; + uint32_t qpBindingCount = 0; + uint32_t reserved = 0; + TileXRUDMAProfileRegionDesc regions[TILEXR_UDMA_PROFILE_MAX_REGIONS] = {}; + TileXRUDMAProfileQpBinding qpBindings[TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS] = {}; +}; + +struct TileXRUDMAProfileRegistry { + uint32_t magic = TILEXR_UDMA_PROFILE_MAGIC; + uint32_t version = TILEXR_UDMA_PROFILE_VERSION; + uint32_t rankSize = 0; + uint32_t regionCount = 0; + uint32_t qpCount = 0; + uint32_t reserved = 0; + TileXRUDMAProfileQpBinding qpBindings[TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS] = {}; + TileXRUDMAProfileRegionDesc + regions[TILEXR_MAX_RANK_SIZE * TILEXR_UDMA_PROFILE_MAX_REGIONS] = {}; +}; + +struct TileXRUDMAProfileView { + uint32_t version = TILEXR_UDMA_PROFILE_VERSION; + uint32_t rankSize = 0; + uint32_t regionCount = 0; + uint32_t qpCount = 0; + GM_ADDR infoDev = nullptr; + GM_ADDR registryDev = nullptr; + const TileXRUDMAProfileRegistry* registryHost = nullptr; +}; + +inline bool UDMAProfileRegionValid(const TileXRUDMAProfileRegionDesc& region) +{ + if (region.base == nullptr || region.bytes == 0) { + return false; + } + const uintptr_t base = reinterpret_cast(region.base); + if (region.bytes > static_cast(std::numeric_limits::max() - base)) { + return false; + } + if (region.registrationBase == nullptr && region.registrationBytes == 0) { + return true; + } + if (region.registrationBase == nullptr || region.registrationBytes == 0) { + return false; + } + const uintptr_t registrationBase = + reinterpret_cast(region.registrationBase); + if (region.registrationBytes > static_cast( + std::numeric_limits::max() - registrationBase) || + base < registrationBase) { + return false; + } + const uint64_t offset = static_cast(base - registrationBase); + return offset <= region.registrationBytes && + region.bytes <= region.registrationBytes - offset; +} + +inline GM_ADDR UDMAProfileRegistrationBase(const TileXRUDMAProfileRegionDesc& region) +{ + return region.registrationBase == nullptr ? region.base : region.registrationBase; +} + +inline uint64_t UDMAProfileRegistrationBytes(const TileXRUDMAProfileRegionDesc& region) +{ + return region.registrationBase == nullptr ? region.bytes : region.registrationBytes; +} + +inline bool UDMAProfileDescValid(const TileXRUDMAProfileDesc* desc, uint32_t expectedQpCount) +{ + if (desc == nullptr || desc->version != TILEXR_UDMA_PROFILE_VERSION || + desc->regionCount == 0 || desc->regionCount > TILEXR_UDMA_PROFILE_MAX_REGIONS || + expectedQpCount == 0 || expectedQpCount > TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS || + desc->qpBindingCount != expectedQpCount) { + return false; + } + for (uint32_t region = 0; region < desc->regionCount; ++region) { + if (!UDMAProfileRegionValid(desc->regions[region])) { + return false; + } + } + for (uint32_t qp = 0; qp < desc->qpBindingCount; ++qp) { + if (desc->qpBindings[qp].localRegion >= desc->regionCount || + desc->qpBindings[qp].remoteRegion >= desc->regionCount) { + return false; + } + } + return true; +} + +inline bool UDMAProfileContractsEqual( + const TileXRUDMAProfileDesc& lhs, const TileXRUDMAProfileDesc& rhs) +{ + if (lhs.version != rhs.version || lhs.regionCount != rhs.regionCount || + lhs.qpBindingCount != rhs.qpBindingCount) { + return false; + } + for (uint32_t qp = 0; qp < lhs.qpBindingCount; ++qp) { + if (lhs.qpBindings[qp].localRegion != rhs.qpBindings[qp].localRegion || + lhs.qpBindings[qp].remoteRegion != rhs.qpBindings[qp].remoteRegion) { + return false; + } + } + return true; +} + +inline bool UDMAProfileRegistryValid(const TileXRUDMAProfileRegistry* registry, + int expectedRankSize, uint32_t expectedRegionCount, uint32_t expectedQpCount = 0) +{ + if (registry == nullptr || registry->magic != TILEXR_UDMA_PROFILE_MAGIC || + registry->version != TILEXR_UDMA_PROFILE_VERSION || expectedRankSize <= 0 || + expectedRankSize > TILEXR_MAX_RANK_SIZE || + registry->rankSize != static_cast(expectedRankSize) || + expectedRegionCount == 0 || expectedRegionCount > TILEXR_UDMA_PROFILE_MAX_REGIONS || + registry->regionCount != expectedRegionCount || registry->qpCount == 0 || + registry->qpCount > TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS || + (expectedQpCount != 0 && registry->qpCount != expectedQpCount)) { + return false; + } + for (uint32_t qp = 0; qp < registry->qpCount; ++qp) { + if (registry->qpBindings[qp].localRegion >= registry->regionCount || + registry->qpBindings[qp].remoteRegion >= registry->regionCount) { + return false; + } + } + return true; +} + +inline const TileXRUDMAProfileRegionDesc* UDMAProfileRegion( + const TileXRUDMAProfileRegistry* registry, int rank, uint32_t region) +{ + if (registry == nullptr || rank < 0 || static_cast(rank) >= registry->rankSize || + region >= registry->regionCount || registry->rankSize > TILEXR_MAX_RANK_SIZE || + registry->regionCount > TILEXR_UDMA_PROFILE_MAX_REGIONS) { + return nullptr; + } + const size_t index = static_cast(rank) * TILEXR_UDMA_PROFILE_MAX_REGIONS + region; + return ®istry->regions[index]; +} + +inline bool UDMAProfileRegionContains(const TileXRUDMAProfileRegistry* registry, + int rank, uint32_t region, uint64_t byteOffset, uint64_t byteCount) +{ + const auto* desc = UDMAProfileRegion(registry, rank, region); + if (desc == nullptr || !UDMAProfileRegionValid(*desc) || byteOffset > desc->bytes) { + return false; + } + return byteCount <= desc->bytes - byteOffset; +} + +inline GM_ADDR UDMAProfileRemoteAddr(const TileXRUDMAProfileRegistry* registry, + int rank, uint32_t region, uint64_t byteOffset) +{ + const auto* desc = UDMAProfileRegion(registry, rank, region); + if (desc == nullptr || byteOffset > desc->bytes) { + return nullptr; + } + return desc->base + byteOffset; +} + inline bool UDMARegistryValid(const TileXRUDMARegistry *registry, int expectedRankSize) { return registry != nullptr && diff --git a/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp b/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp index b7ec8cf2..542d64e5 100644 --- a/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp +++ b/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp @@ -147,8 +147,10 @@ int TileXRMoonEpBuildPrefetchWeightLayout( } const bool hasOverride = blockDimOverride != nullptr && blockDimOverride[0] != '\0'; - uint32_t workers = qpNum < kPrefetchWeightMaxWorkers ? - qpNum : kPrefetchWeightMaxWorkers; + uint32_t workers = kPrefetchWeightMaxWorkers; + while (workers > qpNum) { + workers >>= 1; + } if (!ParseWorkerOverride(blockDimOverride, &workers) || workers > qpNum || (hasOverride && workers > static_cast(args.plan->b))) { return TILEXR_MOONEP_ERROR_INVALID_ARGUMENT; diff --git a/src/moonep/reduce_grad/CMakeLists.txt b/src/moonep/reduce_grad/CMakeLists.txt index 8c88bac7..7a9647f2 100644 --- a/src/moonep/reduce_grad/CMakeLists.txt +++ b/src/moonep/reduce_grad/CMakeLists.txt @@ -59,9 +59,11 @@ tilexr_add_moonep_kernel(tilexr_moonep_reduce_grad_kernel INCLUDES ${TILEXR_MOONEP_REDUCE_GRAD_KERNEL_INCLUDES} DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/common/reduce_grad_common.h" - "${CMAKE_SOURCE_DIR}/src/moonep/common/moonep_peer_window.h" "${CMAKE_SOURCE_DIR}/src/include/comm_args.h" "${CMAKE_SOURCE_DIR}/src/include/tilexr_sync.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_udma.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_udma_reg.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_udma_types.h" ) add_library(tilexr-moonep-reduce-grad SHARED diff --git a/src/moonep/reduce_grad/common/reduce_grad_common.h b/src/moonep/reduce_grad/common/reduce_grad_common.h index 068eb3f6..95dca817 100644 --- a/src/moonep/reduce_grad/common/reduce_grad_common.h +++ b/src/moonep/reduce_grad/common/reduce_grad_common.h @@ -3,40 +3,72 @@ #include +#include "tilexr_udma_reg.h" #include "tilexr_udma_types.h" namespace TileXRMoonEp { -constexpr uint64_t kReduceGradUdmaThresholdBytes = UINT64_C(1) << 20; -constexpr uint64_t kReduceGradDefaultUdmaChunkBytes = UINT64_C(4) << 20; -constexpr uint64_t kReduceGradStateWindowBytes = UINT64_C(1) << 20; -constexpr uint64_t kReduceGradDataAsFlagRecordBytes = 512; -constexpr uint64_t kReduceGradDataAsFlagPayloadBytes = 480; +constexpr uint64_t kReduceGradDefaultChunkBytes = UINT64_C(8) << 20; constexpr uint64_t kReduceGradUdmaAlignment = 512; -// Keep the caller-owned MR compatible with older Ascend950 registration granularity. -constexpr uint64_t kReduceGradUdmaWorkspaceAlignment = UINT64_C(2) << 20; -constexpr uint64_t kReduceGradUdmaStageCount = 2; -constexpr uint64_t kReduceGradUdmaSignalStageStride = +constexpr uint64_t kReduceGradWorkspaceAlignment = UINT64_C(2) << 20; +constexpr uint64_t kReduceGradBankCount = 2; +constexpr uint64_t kReduceGradLaneStateStrideBytes = UINT64_C(8) << 10; +constexpr uint64_t kReduceGradLaneFlagStrideBytes = TileXR::TILEXR_UDMA_CACHE_LINE_SIZE; +constexpr uint64_t kReduceGradBankReadyOffset = 0; +constexpr uint64_t kReduceGradDoneTokenBytes = TileXR::TILEXR_UDMA_CACHE_LINE_SIZE; -constexpr uint64_t kReduceGradUdmaPeerStateBytes = - 3 * kReduceGradUdmaStageCount * kReduceGradUdmaSignalStageStride; -constexpr uint64_t kReduceGradUdmaReadyOffset = 0; -constexpr uint64_t kReduceGradUdmaCompletionOffset = - kReduceGradUdmaStageCount * kReduceGradUdmaSignalStageStride; -constexpr uint64_t kReduceGradUdmaPollScratchOffset = - 2 * kReduceGradUdmaStageCount * kReduceGradUdmaSignalStageStride; -static_assert(kReduceGradUdmaPollScratchOffset + - kReduceGradUdmaStageCount * kReduceGradUdmaSignalStageStride <= - kReduceGradUdmaPeerStateBytes, "ReduceGrad UDMA peer state exceeds its stride"); -constexpr uint64_t kReduceGradKernelTileBytes = UINT64_C(30) * 1024; +constexpr uint64_t kReduceGradDoneTokenCount = 32; +constexpr uint64_t kReduceGradDoneBankStrideBytes = + kReduceGradDoneTokenCount * kReduceGradDoneTokenBytes; +constexpr uint64_t kReduceGradBankDoneOffset = + kReduceGradBankCount * kReduceGradLaneFlagStrideBytes; +constexpr uint64_t kReduceGradBankItemOffset = + kReduceGradBankDoneOffset + + kReduceGradBankCount * kReduceGradDoneBankStrideBytes; +constexpr uint64_t kReduceGradLaneErrorOffset = + kReduceGradBankItemOffset + + kReduceGradBankCount * kReduceGradLaneFlagStrideBytes; +static_assert(kReduceGradLaneErrorOffset + kReduceGradLaneFlagStrideBytes <= + kReduceGradLaneStateStrideBytes, "ReduceGrad lane state exceeds its stride"); + +constexpr uint64_t kReduceGradKernelTileBytes = UINT64_C(40) * 1024; constexpr uint32_t kReduceGradProjectionCount = 3; +constexpr uint32_t kReduceGradProfileRegionCount = 4; +constexpr uint32_t kReduceGradStagingRegion = 0; +constexpr int64_t kReduceGradMinRankCount = 4; constexpr uint32_t kReduceGradMaxUdmaQpCount = 8; +constexpr uint32_t kReduceGradMaxTransportQpCount = + TileXR::TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS; +constexpr uint32_t kReduceGradMinMultiRankQpCount = kReduceGradProjectionCount; constexpr int64_t kReduceGradMaxAivBlockCount = 64; -constexpr uint32_t kReduceGradTransportPeer = 1; -constexpr uint32_t kReduceGradTransportUdma = 2; constexpr int32_t kReduceGradDeviceInvalidState = 1; -constexpr int32_t kReduceGradDevicePeerTimeout = 2; +constexpr int32_t kReduceGradDeviceUdmaTimeout = 2; constexpr int32_t kReduceGradDeviceUdmaCqError = 3; +constexpr int32_t kReduceGradDeviceLeaderTimeout = 4; +constexpr int32_t kReduceGradDeviceHelperTimeout = 5; +constexpr int32_t kReduceGradBarrierStep = 0x5247; +constexpr int32_t kReduceGradBarrierFailureStep = kReduceGradBarrierStep + 1; + +enum ReduceGradBankItemKind : uint32_t { + kReduceGradBankWork = 0, + kReduceGradBankTerminal = 1, +}; + +struct alignas(TileXR::TILEXR_UDMA_CACHE_LINE_SIZE) ReduceGradBankItem { + uint64_t token = 0; + uint64_t chunkIndex = 0; + uint64_t chunkBytes = 0; + uint32_t kind = kReduceGradBankTerminal; + uint32_t projection = 0; + uint32_t localExpert = 0; + uint32_t waveStart = 0; + uint32_t contributorCount = 0; + uint32_t remoteContributorCount = 0; + uint64_t reserved[2] = {}; +}; + +static_assert(sizeof(ReduceGradBankItem) == kReduceGradLaneFlagStrideBytes, + "ReduceGrad bank item must occupy exactly one cache line"); enum ReduceGradProjection : uint32_t { kReduceGradGate = 0, @@ -51,63 +83,59 @@ struct ReduceGradLayout { int64_t expertsPerRank = 0; int64_t prefetchSlots = 0; int64_t blockDim = 0; - int64_t controlBlockCount = 0; uint64_t rowElements[kReduceGradProjectionCount] = {}; uint64_t rowBytes[kReduceGradProjectionCount] = {}; - uint32_t transports[kReduceGradProjectionCount] = {}; - uint32_t udmaQpCount = 0; - - uint64_t peerWindowBytes = 0; - uint64_t peerRecordBaseOffset = 0; - uint64_t peerHalfBytes = 0; - uint64_t peerSlotStrideBytes = 0; - uint64_t peerChunkPayloadBytes = 0; - uint64_t peerChunkCounts[kReduceGradProjectionCount] = {}; - - uint64_t udmaStateOffset = 0; - uint64_t udmaOutboundOffset = 0; - uint64_t udmaInboundOffset = 0; - uint64_t udmaChunkBytes = 0; - uint64_t udmaChunkCounts[kReduceGradProjectionCount] = {}; + uint64_t chunkCounts[kReduceGradProjectionCount] = {}; + uint32_t projectionQpBase[kReduceGradProjectionCount] = {}; + uint32_t projectionQpCounts[kReduceGradProjectionCount] = {}; + uint32_t qpProjection[kReduceGradMaxUdmaQpCount] = {}; + uint32_t lanePhysicalQps[kReduceGradMaxUdmaQpCount] = {}; + uint32_t transportQpCount = 0; + uint32_t qpCount = 0; + uint32_t laneCount = 0; + + uint64_t laneStateBytes = 0; + uint64_t stagingOffset = 0; + uint64_t bankStrideBytes = 0; + uint64_t laneStrideBytes = 0; + uint64_t chunkBytes = 0; uint64_t workspaceBytes = 0; }; struct ReduceGradKernelArgs { - uint8_t *commArgs; - uint8_t *expertsToCopy; - uint8_t *gate; - uint8_t *up; - uint8_t *down; - uint8_t *workspace; - uint8_t *status; - int64_t rank; - int64_t rankSize; - int64_t expertCount; - int64_t expertsPerRank; - int64_t prefetchSlots; - int64_t controlBlockCount; - uint64_t gateRowElements; - uint64_t upRowElements; - uint64_t downRowElements; - uint64_t gateRowBytes; - uint64_t upRowBytes; - uint64_t downRowBytes; - uint32_t gateTransport; - uint32_t upTransport; - uint32_t downTransport; - uint32_t udmaQpCount; - uint64_t peerRecordBaseOffset; - uint64_t peerHalfBytes; - uint64_t peerSlotStrideBytes; - uint64_t peerChunkPayloadBytes; - uint64_t udmaStateOffset; - uint64_t udmaOutboundOffset; - uint64_t udmaInboundOffset; - uint64_t udmaChunkBytes; - uint64_t workspaceBytes; - uint64_t waitIterations; - int64_t magic; + uint8_t *commArgs = nullptr; + uint8_t *profileInfo = nullptr; + uint8_t *profileRegistry = nullptr; + uint8_t *expertsToCopy = nullptr; + uint8_t *gradients[kReduceGradProjectionCount] = {}; + uint8_t *sources[kReduceGradProjectionCount] = {}; + uint8_t *workspace = nullptr; + uint8_t *status = nullptr; + + int64_t rank = 0; + int64_t rankSize = 0; + int64_t expertCount = 0; + int64_t expertsPerRank = 0; + int64_t prefetchSlots = 0; + uint64_t rowElements[kReduceGradProjectionCount] = {}; + uint64_t rowBytes[kReduceGradProjectionCount] = {}; + uint64_t chunkCounts[kReduceGradProjectionCount] = {}; + uint32_t projectionQpBase[kReduceGradProjectionCount] = {}; + uint32_t projectionQpCounts[kReduceGradProjectionCount] = {}; + uint32_t lanePhysicalQps[kReduceGradMaxUdmaQpCount] = {}; + uint32_t transportQpCount = 0; + uint32_t qpCount = 0; + uint32_t laneCount = 0; + + uint64_t laneStateBytes = 0; + uint64_t stagingOffset = 0; + uint64_t bankStrideBytes = 0; + uint64_t laneStrideBytes = 0; + uint64_t chunkBytes = 0; + uint64_t workspaceBytes = 0; + uint64_t waitIterations = 0; + int64_t magic = 0; }; } // namespace TileXRMoonEp diff --git a/src/moonep/reduce_grad/host/reduce_grad_host.cpp b/src/moonep/reduce_grad/host/reduce_grad_host.cpp index dc11f173..929623a2 100644 --- a/src/moonep/reduce_grad/host/reduce_grad_host.cpp +++ b/src/moonep/reduce_grad/host/reduce_grad_host.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "acl/acl_rt.h" #include "comm_args.h" @@ -32,6 +33,44 @@ bool CheckedAdd(uint64_t lhs, uint64_t rhs, uint64_t *out) return true; } +bool AddressRangeValid(const void *base, uint64_t bytes) +{ + if (base == nullptr || bytes == 0 || + bytes > static_cast(std::numeric_limits::max())) { + return false; + } + const uintptr_t address = reinterpret_cast(base); + return static_cast(bytes) <= std::numeric_limits::max() - address; +} + +bool AddressRangeContains(const void *outerBase, uint64_t outerBytes, + const void *innerBase, uint64_t innerBytes) +{ + if (!AddressRangeValid(outerBase, outerBytes) || + !AddressRangeValid(innerBase, innerBytes)) { + return false; + } + const uintptr_t outer = reinterpret_cast(outerBase); + const uintptr_t inner = reinterpret_cast(innerBase); + if (inner < outer) { + return false; + } + const uint64_t offset = static_cast(inner - outer); + return offset <= outerBytes && innerBytes <= outerBytes - offset; +} + +void ResolveSourceRegistration(const TileXRMoonEpReduceGradSourceSliceV2 &source, + void **base, uint64_t *bytes) +{ + if (source.registrationBase == nullptr && source.registrationBytes == 0) { + *base = source.data; + *bytes = source.bytes; + return; + } + *base = source.registrationBase; + *bytes = source.registrationBytes; +} + bool ValidatePlan(const TileXRMoonEpPlanV1 *plan, const TileXR::CommArgs &commArgs) { if (plan == nullptr || plan->structSize < sizeof(*plan) || @@ -39,7 +78,10 @@ bool ValidatePlan(const TileXRMoonEpPlanV1 *plan, const TileXR::CommArgs &commAr plan->k <= 0 || plan->n % plan->k != 0 || plan->e <= 0 || plan->b <= 0 || plan->r <= 0 || plan->r > TileXR::TILEXR_MAX_RANK_SIZE || plan->r != commArgs.rankSize || plan->e % plan->r != 0 || - plan->b > plan->e / plan->r || plan->expertsToCopy == nullptr) { + plan->e > std::numeric_limits::max() || + plan->b > std::numeric_limits::max() || + plan->b > std::numeric_limits::max() / plan->r || + plan->expertsToCopy == nullptr) { return false; } uint64_t capacity = 0; @@ -86,11 +128,41 @@ bool ValidateStatus(const TileXRMoonEpTensorV1 *status) status->shape[2] == 0 && status->shape[3] == 0; } +bool TensorIdentityEqual(const TileXRMoonEpTensorV1 &expected, + const TileXRMoonEpTensorV1 *actual) +{ + if (actual == nullptr || actual->structSize < sizeof(*actual) || + actual->abiVersion != expected.abiVersion || actual->data != expected.data || + actual->elementCount != expected.elementCount || actual->dtype != expected.dtype || + actual->rank != expected.rank) { + return false; + } + for (uint32_t dim = 0; dim < TILEXR_MOONEP_MAX_TENSOR_RANK; ++dim) { + if (actual->shape[dim] != expected.shape[dim]) { + return false; + } + } + return true; +} + +bool SourceIdentityEqual(const TileXRMoonEpReduceGradSourceSliceV2 &expected, + const TileXRMoonEpReduceGradSourceSliceV2 &actual) +{ + return actual.data == expected.data && actual.bytes == expected.bytes && + actual.registrationBase == expected.registrationBase && + actual.registrationBytes == expected.registrationBytes; +} + bool IsA5(const TileXR::CommArgs &commArgs) { return (commArgs.extraFlag & TileXR::ExtraFlag::TOPO_910A5) != 0; } +bool UsesSharedQps(const TileXR::CommArgs &commArgs) +{ + return (commArgs.extraFlag & TileXR::ExtraFlag::UDMA_SHARED_QP) != 0; +} + bool LocalityValid(const TileXR::CommArgs &commArgs) { return commArgs.rankSize > 0 && commArgs.rankSize <= TileXR::TILEXR_MAX_RANK_SIZE && @@ -100,80 +172,75 @@ bool LocalityValid(const TileXR::CommArgs &commArgs) commArgs.localRank >= 0 && commArgs.localRank < commArgs.localRankSize; } -bool UsesTransport(const ReduceGradLayout &layout, uint32_t transport) +bool ProfileViewMatches(const ReduceGradPreparedContext &context, + const TileXR::TileXRUDMAProfileView &view) { - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { - if (layout.transports[q] == transport) { - return true; - } - } - return false; -} - -bool PeerWindowsReady(const TileXR::CommArgs &commArgs) -{ - for (int32_t rank = 0; rank < commArgs.rankSize; ++rank) { - if (commArgs.peerMems[rank] == nullptr) { - return false; - } - } - return true; + return view.version == TileXR::TILEXR_UDMA_PROFILE_VERSION && + view.rankSize == static_cast(context.layout.rankSize) && + view.regionCount == kReduceGradProfileRegionCount && + view.qpCount == context.layout.transportQpCount && + view.infoDev == context.profileView.infoDev && + view.registryDev == context.profileView.registryDev && + view.registryHost == context.profileView.registryHost && + TileXR::UDMAProfileRegistryValid(view.registryHost, context.layout.rankSize, + kReduceGradProfileRegionCount, context.layout.transportQpCount); } -int ValidateRegisteredWorkspace(const ReduceGradParams ¶ms, - const TileXR::CommArgs &commArgs, const ReduceGradLayout &layout, - const TileXR::TileXRUDMARegistry **registryOut) +bool PreparedRegistryMatches(const ReduceGradPreparedContext &context) { - if (registryOut == nullptr) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + const auto *registry = context.profileView.registryHost; + if (!TileXR::UDMAProfileRegistryValid(registry, context.layout.rankSize, + kReduceGradProfileRegionCount, context.layout.transportQpCount)) { + return false; } - *registryOut = nullptr; - if (!UsesTransport(layout, TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA)) { - return params.workspace == nullptr && params.workspaceBytes == 0 ? - TileXR::TILEXR_SUCCESS : TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + const int rank = static_cast(context.layout.rank); + const auto *staging = TileXR::UDMAProfileRegion(registry, rank, kReduceGradStagingRegion); + if (staging == nullptr || staging->base != static_cast(context.workspace) || + staging->bytes != context.workspaceBytes) { + return false; } - if (params.workspace == nullptr || params.workspaceBytes < layout.workspaceBytes || - (commArgs.extraFlag & TileXR::ExtraFlag::UDMA) == 0 || - commArgs.udmaInfoPtr == nullptr || commArgs.udmaRegistryPtr == nullptr) { - return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + const auto *source = TileXR::UDMAProfileRegion(registry, rank, projection + 1); + if (source == nullptr || source->base != static_cast( + context.sources[projection].data) || + source->bytes != context.sources[projection].bytes) { + return false; + } } - - const TileXR::TileXRUDMARegistry *registry = nullptr; - const int ret = TileXRGetUDMARegistryHost(params.comm, ®istry); - if (ret != TileXR::TILEXR_SUCCESS) { - return ret; + uint32_t qpRegions[kReduceGradMaxTransportQpCount] = {}; + for (uint32_t qp = 0; qp < context.layout.transportQpCount; ++qp) { + qpRegions[qp] = kReduceGradGate + 1U; } - if (!TileXR::UDMARegistryValid(registry, commArgs.rankSize) || - registry->regions[commArgs.rank].base != static_cast(params.workspace) || - registry->regions[commArgs.rank].bytes != params.workspaceBytes) { - return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + for (uint32_t lane = 0; lane < context.layout.laneCount; ++lane) { + qpRegions[context.layout.lanePhysicalQps[lane]] = + context.layout.qpProjection[lane] + 1U; } - const size_t registeredBytes = registry->regions[0].bytes; - for (int32_t rank = 0; rank < commArgs.rankSize; ++rank) { - const auto ®ion = registry->regions[rank]; - if (reinterpret_cast(region.base) % kReduceGradUdmaWorkspaceAlignment != 0 || - region.bytes != registeredBytes || - !TileXR::UDMARegionContains(registry, rank, 0, layout.workspaceBytes)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + for (uint32_t qp = 0; qp < context.layout.transportQpCount; ++qp) { + if (registry->qpBindings[qp].localRegion != kReduceGradStagingRegion || + registry->qpBindings[qp].remoteRegion != qpRegions[qp]) { + return false; } } - *registryOut = registry; - return TileXR::TILEXR_SUCCESS; + return true; } void FillWorkspaceInfo(const ReduceGradLayout &layout, TileXRMoonEpReduceGradWorkspaceInfoV2 *info) { info->workspaceBytes = layout.workspaceBytes; - info->workspaceAlignment = kReduceGradUdmaWorkspaceAlignment; - info->udmaChunkBytes = layout.udmaChunkBytes; - info->peerWindowBytes = layout.peerWindowBytes; - info->peerHalfBytes = layout.peerHalfBytes; - info->peerSlotStrideBytes = layout.peerSlotStrideBytes; + info->workspaceAlignment = kReduceGradWorkspaceAlignment; + info->udmaChunkBytes = layout.chunkBytes; + info->laneStateBytes = layout.laneStateBytes; + info->laneStateStrideBytes = kReduceGradLaneStateStrideBytes; + info->bankStrideBytes = layout.bankStrideBytes; + info->laneStrideBytes = layout.laneStrideBytes; + info->qpCount = layout.qpCount; info->blockDim = static_cast(layout.blockDim); - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { - info->rowBytes[q] = layout.rowBytes[q]; - info->transports[q] = layout.transports[q]; + info->reserved = 0; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + info->rowBytes[projection] = layout.rowBytes[projection]; + info->chunkCounts[projection] = layout.chunkCounts[projection]; + info->projectionQpCounts[projection] = layout.projectionQpCounts[projection]; } } @@ -182,7 +249,7 @@ void FillWorkspaceInfo(const ReduceGradLayout &layout, int TileXRMoonEpPrepareReduceGradLayout(TileXRCommPtr comm, const TileXRMoonEpPlanV1 *plan, const TileXRMoonEpTensorV1 *const gradients[kReduceGradProjectionCount], - uint64_t requestedUdmaChunkBytes, ReduceGradLayout *layout) + uint64_t requestedChunkBytes, ReduceGradLayout *layout) { if (comm == nullptr || gradients == nullptr || layout == nullptr) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; @@ -204,76 +271,210 @@ int TileXRMoonEpPrepareReduceGradLayout(TileXRCommPtr comm, if (!LocalityValid(*commArgs)) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } + if (commArgs->rankSize < kReduceGradMinRankCount) { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; + } uint64_t rowElements[kReduceGradProjectionCount] = {}; uint64_t rowCountValue = 0; if (!CheckedAdd(static_cast(plan->e), static_cast(plan->b), - &rowCountValue) || - rowCountValue > static_cast(INT64_MAX)) { + &rowCountValue) || rowCountValue > static_cast(INT64_MAX)) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } const int64_t rowCount = static_cast(rowCountValue); - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { - if (!ValidateGradient(gradients[q], rowCount, &rowElements[q])) { + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + if (!ValidateGradient(gradients[projection], rowCount, &rowElements[projection])) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } } - ret = TileXRMoonEpBuildReduceGradLayout(commArgs->rank, commArgs->rankSize, - plan->e, plan->b, rowElements, TileXRMoonEpReduceGradPeerWindowBytes(), - requestedUdmaChunkBytes, layout); - if (ret != TileXR::TILEXR_SUCCESS || - !UsesTransport(*layout, TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA)) { - return ret; - } - uint32_t qpCount = 0; - ret = TileXRUDMAGetQpCount(comm, &qpCount); + uint32_t transportQpCount = 0; + if ((commArgs->extraFlag & TileXR::ExtraFlag::UDMA) == 0) { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; + } + ret = TileXRUDMAGetQpCount(comm, &transportQpCount); if (ret != TileXR::TILEXR_SUCCESS) { return ret; } - if (qpCount == 0 || qpCount > kReduceGradMaxUdmaQpCount) { - return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + if (transportQpCount < kReduceGradMinMultiRankQpCount || + transportQpCount > kReduceGradMaxTransportQpCount || + (transportQpCount > kReduceGradMaxUdmaQpCount && + (transportQpCount != kReduceGradMaxTransportQpCount || + !UsesSharedQps(*commArgs)))) { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; } - layout->udmaQpCount = qpCount; - return TileXR::TILEXR_SUCCESS; + return TileXRMoonEpBuildReduceGradLayout(commArgs->rank, commArgs->rankSize, + plan->e, plan->b, rowElements, transportQpCount, requestedChunkBytes, layout); } -int TileXRMoonEpPrepareReduceGradLaunchContext(const ReduceGradParams ¶ms, - ReduceGradLaunchContext *context) +int TileXRMoonEpCreateReduceGradPreparedContext( + const ReduceGradPrepareParams ¶ms, ReduceGradPreparedContext **context) { - if (context == nullptr || params.stream == nullptr || params.waitIterations == 0 || - !ValidateStatus(params.status)) { + if (context == nullptr) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - *context = ReduceGradLaunchContext {}; + *context = nullptr; const TileXRMoonEpTensorV1 *gradients[kReduceGradProjectionCount] = { params.gradients[kReduceGradGate], params.gradients[kReduceGradUp], params.gradients[kReduceGradDown]}; - int ret = TileXRMoonEpPrepareReduceGradLayout(params.comm, params.plan, gradients, - params.requestedUdmaChunkBytes, &context->layout); + ReduceGradLayout layout {}; + int ret = TileXRMoonEpPrepareReduceGradLayout(params.comm, params.plan, + gradients, params.requestedChunkBytes, &layout); if (ret != TileXR::TILEXR_SUCCESS) { return ret; } - ret = TileXRGetCommArgsHost(params.comm, context->hostArgs); - if (ret != TileXR::TILEXR_SUCCESS || context->hostArgs == nullptr) { - *context = ReduceGradLaunchContext {}; - return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + + if (!AddressRangeValid(params.workspace, params.workspaceBytes) || + params.workspaceBytes < layout.workspaceBytes || + reinterpret_cast(params.workspace) % + kReduceGradWorkspaceAlignment != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - if (UsesTransport(context->layout, TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER) && - !PeerWindowsReady(*context->hostArgs)) { - *context = ReduceGradLaunchContext {}; - return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + uint64_t requiredBytes = 0; + void *registrationBase = nullptr; + uint64_t registrationBytes = 0; + ResolveSourceRegistration(params.sources[projection], + ®istrationBase, ®istrationBytes); + if (!CheckedMul(static_cast(layout.prefetchSlots), + layout.rowBytes[projection], &requiredBytes) || + params.sources[projection].bytes != requiredBytes || + !AddressRangeValid(params.sources[projection].data, + params.sources[projection].bytes) || + !AddressRangeContains(registrationBase, registrationBytes, + params.sources[projection].data, params.sources[projection].bytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + + ReduceGradPreparedContext *next = new (std::nothrow) ReduceGradPreparedContext(); + if (next == nullptr) { + return TileXR::TILEXR_ERROR_INTERNAL; + } + next->comm = params.comm; + next->layout = layout; + next->planN = params.plan->n; + next->planK = params.plan->k; + next->expertsToCopy = params.plan->expertsToCopy; + next->workspace = params.workspace; + next->workspaceBytes = params.workspaceBytes; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + next->gradients[projection] = *params.gradients[projection]; + next->sources[projection] = params.sources[projection]; + } + + ret = TileXRGetCommArgsHost(params.comm, next->hostArgs); + if (ret == TileXR::TILEXR_SUCCESS && next->hostArgs == nullptr) { + ret = TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + if (ret == TileXR::TILEXR_SUCCESS) { + ret = TileXRGetCommArgsDev(params.comm, next->devArgs); + if (ret == TileXR::TILEXR_SUCCESS && next->devArgs == nullptr) { + ret = TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + } + if (ret != TileXR::TILEXR_SUCCESS) { + delete next; + return ret; + } + + TileXR::TileXRUDMAProfileDesc desc {}; + desc.regionCount = kReduceGradProfileRegionCount; + desc.qpBindingCount = layout.transportQpCount; + desc.regions[kReduceGradStagingRegion].base = static_cast(params.workspace); + desc.regions[kReduceGradStagingRegion].bytes = params.workspaceBytes; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + desc.regions[projection + 1].base = static_cast( + params.sources[projection].data); + desc.regions[projection + 1].bytes = params.sources[projection].bytes; + void *registrationBase = nullptr; + uint64_t registrationBytes = 0; + ResolveSourceRegistration(params.sources[projection], + ®istrationBase, ®istrationBytes); + desc.regions[projection + 1].registrationBase = + static_cast(registrationBase); + desc.regions[projection + 1].registrationBytes = registrationBytes; + } + for (uint32_t qp = 0; qp < layout.transportQpCount; ++qp) { + desc.qpBindings[qp].localRegion = kReduceGradStagingRegion; + desc.qpBindings[qp].remoteRegion = kReduceGradGate + 1U; + } + for (uint32_t lane = 0; lane < layout.laneCount; ++lane) { + desc.qpBindings[layout.lanePhysicalQps[lane]].remoteRegion = + layout.qpProjection[lane] + 1U; + } + ret = TileXRUDMAProfileRegister(params.comm, &desc, &next->profileHandle); + if (ret == TileXR::TILEXR_SUCCESS) { + ret = TileXRUDMAProfileQuery(params.comm, next->profileHandle, + &next->profileView); + } + if (ret == TileXR::TILEXR_SUCCESS && + (!ProfileViewMatches(*next, next->profileView) || + !PreparedRegistryMatches(*next))) { + ret = TileXR::TILEXR_ERROR_NOT_INITIALIZED; } - ret = ValidateRegisteredWorkspace(params, *context->hostArgs, context->layout, - &context->registry); if (ret != TileXR::TILEXR_SUCCESS) { - *context = ReduceGradLaunchContext {}; + if (next->profileHandle != 0) { + const int cleanupRet = TileXRUDMAProfileUnregister( + params.comm, next->profileHandle); + if (cleanupRet != TileXR::TILEXR_SUCCESS) { + ret = cleanupRet; + } + } + delete next; return ret; } - ret = TileXRGetCommArgsDev(params.comm, context->devArgs); - if (ret != TileXR::TILEXR_SUCCESS || context->devArgs == nullptr) { - *context = ReduceGradLaunchContext {}; - return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + + *context = next; + return TileXR::TILEXR_SUCCESS; +} + +int TileXRMoonEpDestroyReduceGradPreparedContext(ReduceGradPreparedContext *context) +{ + if (context == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (context->profileHandle != 0) { + const int ret = TileXRUDMAProfileUnregister(context->comm, context->profileHandle); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + } + delete context; + return TileXR::TILEXR_SUCCESS; +} + +int TileXRMoonEpValidateReduceGradLaunch(const ReduceGradLaunchParams ¶ms, + const ReduceGradPreparedContext &context) +{ + if (params.stream == nullptr || params.waitIterations == 0 || + !ValidateStatus(params.status) || params.plan == nullptr || + params.plan->structSize < sizeof(*params.plan) || + params.plan->abiVersion != TILEXR_MOONEP_ABI_VERSION_V1 || + params.plan->n != context.planN || params.plan->k != context.planK || + params.plan->r != context.layout.rankSize || + params.plan->e != context.layout.expertCount || + params.plan->b != context.layout.prefetchSlots || + params.plan->expertsToCopy != context.expertsToCopy) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + if (!TensorIdentityEqual(context.gradients[projection], + params.gradients[projection]) || + !SourceIdentityEqual(context.sources[projection], params.sources[projection])) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + if (context.layout.rankSize > 1) { + TileXR::TileXRUDMAProfileView current {}; + const int ret = TileXRUDMAProfileQuery( + context.comm, context.profileHandle, ¤t); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + if (!ProfileViewMatches(context, current) || !PreparedRegistryMatches(context)) { + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } } return TileXR::TILEXR_SUCCESS; } @@ -302,32 +503,73 @@ extern "C" int TileXRMoonEpReduceGradGetWorkspaceSizeV2( return TileXR::TILEXR_SUCCESS; } +extern "C" int TileXRMoonEpReduceGradPrepareV2( + const TileXRMoonEpReduceGradPrepareArgsV2 *args, + TileXRMoonEpReduceGradPreparedV2 *prepared) +{ + if (prepared != nullptr) { + *prepared = nullptr; + } + if (args == nullptr || prepared == nullptr || args->structSize < sizeof(*args) || + args->abiVersion != TILEXR_MOONEP_ABI_VERSION_V2 || args->comm == nullptr || + args->plan == nullptr || args->gate == nullptr || args->up == nullptr || + args->down == nullptr || args->flags != TILEXR_MOONEP_FLAG_NONE) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + TileXRMoonEp::ReduceGradPrepareParams params {}; + params.comm = args->comm; + params.plan = args->plan; + params.gradients[TileXRMoonEp::kReduceGradGate] = args->gate; + params.gradients[TileXRMoonEp::kReduceGradUp] = args->up; + params.gradients[TileXRMoonEp::kReduceGradDown] = args->down; + params.workspace = args->workspace; + params.workspaceBytes = args->workspaceBytes; + params.requestedChunkBytes = args->requestedUdmaChunkBytes; + for (uint32_t projection = 0; + projection < TileXRMoonEp::kReduceGradProjectionCount; ++projection) { + params.sources[projection] = args->sources[projection]; + } + TileXRMoonEp::ReduceGradPreparedContext *context = nullptr; + const int ret = TileXRMoonEp::TileXRMoonEpCreateReduceGradPreparedContext( + params, &context); + if (ret == TileXR::TILEXR_SUCCESS) { + *prepared = static_cast(context); + } + return ret; +} + +extern "C" int TileXRMoonEpReduceGradDestroyPreparedV2( + TileXRMoonEpReduceGradPreparedV2 prepared) +{ + return TileXRMoonEp::TileXRMoonEpDestroyReduceGradPreparedContext( + static_cast(prepared)); +} + extern "C" int TileXRMoonEpReduceGradV2(const TileXRMoonEpReduceGradArgsV2 *args, aclrtStream stream) { if (args == nullptr || args->structSize < sizeof(*args) || - args->abiVersion != TILEXR_MOONEP_ABI_VERSION_V2 || args->comm == nullptr || + args->abiVersion != TILEXR_MOONEP_ABI_VERSION_V2 || args->prepared == nullptr || args->plan == nullptr || args->gate == nullptr || args->up == nullptr || args->down == nullptr || args->status == nullptr || args->waitIterations == 0 || args->flags != TILEXR_MOONEP_FLAG_NONE || stream == nullptr) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - TileXRMoonEp::ReduceGradParams params {}; - params.comm = args->comm; + TileXRMoonEp::ReduceGradLaunchParams params {}; params.plan = args->plan; params.gradients[TileXRMoonEp::kReduceGradGate] = args->gate; params.gradients[TileXRMoonEp::kReduceGradUp] = args->up; params.gradients[TileXRMoonEp::kReduceGradDown] = args->down; - params.workspace = args->workspace; - params.workspaceBytes = args->workspaceBytes; params.status = args->status; params.waitIterations = args->waitIterations; - params.requestedUdmaChunkBytes = args->requestedUdmaChunkBytes; params.stream = stream; - - TileXRMoonEp::ReduceGradLaunchContext context {}; - int ret = TileXRMoonEp::TileXRMoonEpPrepareReduceGradLaunchContext(params, &context); + for (uint32_t projection = 0; + projection < TileXRMoonEp::kReduceGradProjectionCount; ++projection) { + params.sources[projection] = args->sources[projection]; + } + auto *context = static_cast(args->prepared); + int ret = TileXRMoonEp::TileXRMoonEpValidateReduceGradLaunch(params, *context); if (ret != TileXR::TILEXR_SUCCESS) { return ret; } @@ -338,9 +580,11 @@ extern "C" int TileXRMoonEpReduceGradV2(const TileXRMoonEpReduceGradArgsV2 *args if (memsetRet != ACL_SUCCESS) { return TileXR::TILEXR_ERROR_MKIRT; } - ret = TileXRMoonEp::TileXRMoonEpLaunchReduceGradKernel(params, context); - if (ret != TileXR::TILEXR_SUCCESS && aclrtSynchronizeStream(stream) != ACL_SUCCESS) { + const int launchRet = TileXRMoonEp::TileXRMoonEpLaunchReduceGradKernel( + params, *context); + if (launchRet != TileXR::TILEXR_SUCCESS && + aclrtSynchronizeStream(stream) != ACL_SUCCESS) { return TileXR::TILEXR_ERROR_MKIRT; } - return ret; + return launchRet; } diff --git a/src/moonep/reduce_grad/host/reduce_grad_host.h b/src/moonep/reduce_grad/host/reduce_grad_host.h index 3301b1df..1d898903 100644 --- a/src/moonep/reduce_grad/host/reduce_grad_host.h +++ b/src/moonep/reduce_grad/host/reduce_grad_host.h @@ -11,35 +11,58 @@ namespace TileXRMoonEp { -struct ReduceGradParams { +struct ReduceGradPrepareParams { TileXRCommPtr comm = nullptr; const TileXRMoonEpPlanV1 *plan = nullptr; TileXRMoonEpTensorV1 *gradients[kReduceGradProjectionCount] = {}; + TileXRMoonEpReduceGradSourceSliceV2 sources[kReduceGradProjectionCount] = {}; void *workspace = nullptr; uint64_t workspaceBytes = 0; + uint64_t requestedChunkBytes = 0; +}; + +struct ReduceGradLaunchParams { + const TileXRMoonEpPlanV1 *plan = nullptr; + TileXRMoonEpTensorV1 *gradients[kReduceGradProjectionCount] = {}; + TileXRMoonEpReduceGradSourceSliceV2 sources[kReduceGradProjectionCount] = {}; TileXRMoonEpTensorV1 *status = nullptr; uint64_t waitIterations = 0; - uint64_t requestedUdmaChunkBytes = 0; aclrtStream stream = nullptr; }; -struct ReduceGradLaunchContext { +struct ReduceGradPreparedContext { + TileXRCommPtr comm = nullptr; TileXR::CommArgs *hostArgs = nullptr; GM_ADDR devArgs = nullptr; - const TileXR::TileXRUDMARegistry *registry = nullptr; + TileXRUDMAProfileHandle profileHandle = 0; + TileXR::TileXRUDMAProfileView profileView {}; ReduceGradLayout layout {}; + + int64_t planN = 0; + int64_t planK = 0; + void *expertsToCopy = nullptr; + TileXRMoonEpTensorV1 gradients[kReduceGradProjectionCount] = {}; + TileXRMoonEpReduceGradSourceSliceV2 sources[kReduceGradProjectionCount] = {}; + void *workspace = nullptr; + uint64_t workspaceBytes = 0; }; int TileXRMoonEpPrepareReduceGradLayout(TileXRCommPtr comm, const TileXRMoonEpPlanV1 *plan, const TileXRMoonEpTensorV1 *const gradients[kReduceGradProjectionCount], - uint64_t requestedUdmaChunkBytes, ReduceGradLayout *layout); + uint64_t requestedChunkBytes, ReduceGradLayout *layout); + +int TileXRMoonEpCreateReduceGradPreparedContext( + const ReduceGradPrepareParams ¶ms, ReduceGradPreparedContext **context); + +int TileXRMoonEpDestroyReduceGradPreparedContext( + ReduceGradPreparedContext *context); -int TileXRMoonEpPrepareReduceGradLaunchContext(const ReduceGradParams ¶ms, - ReduceGradLaunchContext *context); +int TileXRMoonEpValidateReduceGradLaunch(const ReduceGradLaunchParams ¶ms, + const ReduceGradPreparedContext &context); -int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradParams ¶ms, - const ReduceGradLaunchContext &context); +int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradLaunchParams ¶ms, + const ReduceGradPreparedContext &context); } // namespace TileXRMoonEp diff --git a/src/moonep/reduce_grad/host/reduce_grad_launch.cpp b/src/moonep/reduce_grad/host/reduce_grad_launch.cpp index b7c280af..849b6911 100644 --- a/src/moonep/reduce_grad/host/reduce_grad_launch.cpp +++ b/src/moonep/reduce_grad/host/reduce_grad_launch.cpp @@ -63,8 +63,8 @@ int EnsureReduceGradKernelRegistered() } // namespace -int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradParams ¶ms, - const ReduceGradLaunchContext &context) +int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradLaunchParams ¶ms, + const ReduceGradPreparedContext &context) { const int registerRet = EnsureReduceGradKernelRegistered(); if (registerRet != TileXR::TILEXR_SUCCESS) { @@ -72,49 +72,50 @@ int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradParams ¶ms, } int64_t magic = 0; - const int ret = TileXRCommNextMagic(params.comm, &magic); + const int ret = TileXRCommNextMagic(context.comm, &magic); if (ret != TileXR::TILEXR_SUCCESS) { return ret; } const ReduceGradLayout &layout = context.layout; static_assert(sizeof(ReduceGradKernelArgs) <= std::numeric_limits::max(), "ReduceGrad kernel argument block exceeds Runtime V2 argsSize"); - ReduceGradKernelArgs args { - context.devArgs, - reinterpret_cast(params.plan->expertsToCopy), - reinterpret_cast(params.gradients[kReduceGradGate]->data), - reinterpret_cast(params.gradients[kReduceGradUp]->data), - reinterpret_cast(params.gradients[kReduceGradDown]->data), - reinterpret_cast(params.workspace), - reinterpret_cast(params.status->data), - layout.rank, - layout.rankSize, - layout.expertCount, - layout.expertsPerRank, - layout.prefetchSlots, - layout.controlBlockCount, - layout.rowElements[kReduceGradGate], - layout.rowElements[kReduceGradUp], - layout.rowElements[kReduceGradDown], - layout.rowBytes[kReduceGradGate], - layout.rowBytes[kReduceGradUp], - layout.rowBytes[kReduceGradDown], - layout.transports[kReduceGradGate], - layout.transports[kReduceGradUp], - layout.transports[kReduceGradDown], - layout.udmaQpCount, - layout.peerRecordBaseOffset, - layout.peerHalfBytes, - layout.peerSlotStrideBytes, - layout.peerChunkPayloadBytes, - layout.udmaStateOffset, - layout.udmaOutboundOffset, - layout.udmaInboundOffset, - layout.udmaChunkBytes, - layout.workspaceBytes, - params.waitIterations, - magic, - }; + ReduceGradKernelArgs args {}; + args.commArgs = context.devArgs; + args.profileInfo = context.profileView.infoDev; + args.profileRegistry = context.profileView.registryDev; + args.expertsToCopy = reinterpret_cast(context.expertsToCopy); + args.workspace = reinterpret_cast(context.workspace); + args.status = reinterpret_cast(params.status->data); + args.rank = layout.rank; + args.rankSize = layout.rankSize; + args.expertCount = layout.expertCount; + args.expertsPerRank = layout.expertsPerRank; + args.prefetchSlots = layout.prefetchSlots; + args.transportQpCount = layout.transportQpCount; + args.qpCount = layout.qpCount; + args.laneCount = layout.laneCount; + args.laneStateBytes = layout.laneStateBytes; + args.stagingOffset = layout.stagingOffset; + args.bankStrideBytes = layout.bankStrideBytes; + args.laneStrideBytes = layout.laneStrideBytes; + args.chunkBytes = layout.chunkBytes; + args.workspaceBytes = layout.workspaceBytes; + args.waitIterations = params.waitIterations; + args.magic = magic; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + args.gradients[projection] = reinterpret_cast( + context.gradients[projection].data); + args.sources[projection] = reinterpret_cast( + context.sources[projection].data); + args.rowElements[projection] = layout.rowElements[projection]; + args.rowBytes[projection] = layout.rowBytes[projection]; + args.chunkCounts[projection] = layout.chunkCounts[projection]; + args.projectionQpBase[projection] = layout.projectionQpBase[projection]; + args.projectionQpCounts[projection] = layout.projectionQpCounts[projection]; + } + for (uint32_t lane = 0; lane < layout.laneCount; ++lane) { + args.lanePhysicalQps[lane] = layout.lanePhysicalQps[lane]; + } rtArgsEx_t argsInfo {}; argsInfo.args = &args; diff --git a/src/moonep/reduce_grad/host/reduce_grad_layout.cpp b/src/moonep/reduce_grad/host/reduce_grad_layout.cpp index 9ebd3529..7a99b1d4 100644 --- a/src/moonep/reduce_grad/host/reduce_grad_layout.cpp +++ b/src/moonep/reduce_grad/host/reduce_grad_layout.cpp @@ -1,12 +1,9 @@ #include "reduce_grad_layout.h" #include -#include #include #include -#include "comm_args.h" -#include "tilexr_moonep.h" #include "tilexr_types.h" namespace TileXRMoonEp { @@ -30,11 +27,6 @@ bool CheckedMul(uint64_t lhs, uint64_t rhs, uint64_t *out) return true; } -uint64_t AlignDown(uint64_t value, uint64_t alignment) -{ - return alignment == 0 ? 0 : value / alignment * alignment; -} - bool AlignUp(uint64_t value, uint64_t alignment, uint64_t *out) { if (out == nullptr || alignment == 0) { @@ -54,10 +46,73 @@ bool DivideRoundUp(uint64_t value, uint64_t divisor, uint64_t *out) return true; } -bool ResolveBlockDim(int64_t rankSize, bool, int64_t *blockDim, - int64_t *controlBlockCount) +bool ScoreGreater(uint64_t lhsBytes, uint32_t lhsCount, + uint64_t rhsBytes, uint32_t rhsCount) +{ + const uint64_t lhsQuotient = lhsBytes / lhsCount; + const uint64_t rhsQuotient = rhsBytes / rhsCount; + if (lhsQuotient != rhsQuotient) { + return lhsQuotient > rhsQuotient; + } + const uint64_t lhsRemainder = lhsBytes % lhsCount; + const uint64_t rhsRemainder = rhsBytes % rhsCount; + return lhsRemainder * rhsCount > rhsRemainder * lhsCount; +} + +bool AllocateProjectionQps(ReduceGradLayout *layout) +{ + if (layout == nullptr || layout->qpCount < kReduceGradMinMultiRankQpCount || + layout->qpCount > kReduceGradMaxUdmaQpCount) { + return false; + } + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + layout->projectionQpCounts[projection] = 1; + } + for (uint32_t assigned = kReduceGradProjectionCount; + assigned < layout->qpCount; ++assigned) { + uint32_t selected = 0; + for (uint32_t projection = 1; projection < kReduceGradProjectionCount; ++projection) { + if (ScoreGreater(layout->rowBytes[projection], + layout->projectionQpCounts[projection], + layout->rowBytes[selected], layout->projectionQpCounts[selected])) { + selected = projection; + } + } + ++layout->projectionQpCounts[selected]; + } + + uint32_t cursor = 0; + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + layout->projectionQpBase[projection] = cursor; + for (uint32_t lane = 0; lane < layout->projectionQpCounts[projection]; ++lane) { + if (cursor >= layout->qpCount || cursor >= kReduceGradMaxUdmaQpCount) { + return false; + } + layout->qpProjection[cursor++] = projection; + } + } + return cursor == layout->qpCount; +} + +bool MapPhysicalQps(ReduceGradLayout *layout) +{ + if (layout == nullptr || layout->laneCount == 0 || + layout->laneCount > kReduceGradMaxUdmaQpCount || + layout->transportQpCount < layout->laneCount) { + return false; + } + for (uint32_t lane = 0; lane < layout->laneCount; ++lane) { + layout->lanePhysicalQps[lane] = lane; + } + if (layout->transportQpCount == kReduceGradMaxTransportQpCount) { + layout->lanePhysicalQps[kReduceGradDown] = 16U; + } + return true; +} + +bool ResolveBlockDim(uint32_t laneCount, int64_t *blockDim) { - if (blockDim == nullptr || controlBlockCount == nullptr) { + if (blockDim == nullptr || laneCount == 0) { return false; } const char *value = std::getenv("TILEXR_MOONEP_REDUCE_GRAD_BLOCK_DIM"); @@ -71,43 +126,33 @@ bool ResolveBlockDim(int64_t rankSize, bool, int64_t *blockDim, } selected = static_cast(parsed); } - if (selected <= 0 || selected > kReduceGradMaxAivBlockCount) { + const uint64_t minimum = 2U * laneCount; + if (selected <= 0 || selected > kReduceGradMaxAivBlockCount || + static_cast(selected) < minimum) { return false; } - - int64_t controls = 0; - if (rankSize > 1) { - if (selected < 2) { - return false; - } - controls = rankSize - 1; - if (controls >= selected) { - controls = selected - 1; - } - } *blockDim = selected; - *controlBlockCount = controls; return true; } } // namespace -uint64_t TileXRMoonEpReduceGradPeerWindowBytes() -{ - return TileXR::IPC_BUFF_MAX_SIZE > 0 ? - static_cast(TileXR::IPC_BUFF_MAX_SIZE) : 0; -} - int TileXRMoonEpBuildReduceGradLayout(int64_t rank, int64_t rankSize, int64_t expertCount, int64_t prefetchSlots, const uint64_t rowElements[kReduceGradProjectionCount], - uint64_t peerWindowBytes, uint64_t requestedUdmaChunkBytes, + uint32_t transportQpCount, uint64_t requestedChunkBytes, ReduceGradLayout *out) { - if (out == nullptr || rowElements == nullptr || rankSize <= 0 || + if (out == nullptr || rowElements == nullptr || rankSize < kReduceGradMinRankCount || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || rank < 0 || rank >= rankSize || expertCount <= 0 || expertCount % rankSize != 0 || prefetchSlots <= 0 || - prefetchSlots > expertCount / rankSize) { + expertCount > std::numeric_limits::max() || + prefetchSlots > std::numeric_limits::max() || + prefetchSlots > std::numeric_limits::max() / rankSize || + transportQpCount < kReduceGradMinMultiRankQpCount || + transportQpCount > kReduceGradMaxTransportQpCount || + (transportQpCount > kReduceGradMaxUdmaQpCount && + transportQpCount != kReduceGradMaxTransportQpCount)) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } @@ -117,119 +162,62 @@ int TileXRMoonEpBuildReduceGradLayout(int64_t rank, int64_t rankSize, next.expertCount = expertCount; next.expertsPerRank = expertCount / rankSize; next.prefetchSlots = prefetchSlots; - next.peerWindowBytes = peerWindowBytes; + next.transportQpCount = transportQpCount; + next.qpCount = transportQpCount == kReduceGradMaxTransportQpCount ? + kReduceGradProjectionCount : transportQpCount; + next.laneCount = next.qpCount; - bool usesPeer = false; - bool usesUdma = false; - uint64_t maxUdmaRowBytes = 0; - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { uint64_t bytes = 0; - if (rowElements[q] == 0 || !CheckedMul(rowElements[q], sizeof(float), &bytes)) { + if (rowElements[projection] == 0 || + !CheckedMul(rowElements[projection], sizeof(float), &bytes)) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - next.rowElements[q] = rowElements[q]; - next.rowBytes[q] = bytes; - if (rankSize == 1 || bytes <= kReduceGradUdmaThresholdBytes) { - next.transports[q] = TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER; - usesPeer = true; - } else { - next.transports[q] = TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA; - usesUdma = true; - if (bytes > maxUdmaRowBytes) { - maxUdmaRowBytes = bytes; - } - } + next.rowElements[projection] = rowElements[projection]; + next.rowBytes[projection] = bytes; } - if (!ResolveBlockDim(rankSize, usesUdma, &next.blockDim, &next.controlBlockCount)) { + if (!ResolveBlockDim(next.laneCount, &next.blockDim)) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - - if (usesPeer) { - if (peerWindowBytes <= kReduceGradStateWindowBytes) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - next.peerRecordBaseOffset = kReduceGradStateWindowBytes; - next.peerHalfBytes = AlignDown( - (peerWindowBytes - kReduceGradStateWindowBytes) / 2, - kReduceGradDataAsFlagRecordBytes); - uint64_t incomingSlots = 0; - if (!CheckedMul(static_cast(rankSize), - static_cast(prefetchSlots), &incomingSlots)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - next.peerSlotStrideBytes = AlignDown(next.peerHalfBytes / incomingSlots, - kReduceGradDataAsFlagRecordBytes); - if (next.peerSlotStrideBytes < kReduceGradDataAsFlagRecordBytes || - !CheckedMul(next.peerSlotStrideBytes / kReduceGradDataAsFlagRecordBytes, - kReduceGradDataAsFlagPayloadBytes, &next.peerChunkPayloadBytes) || - next.peerChunkPayloadBytes == 0) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { - if (next.transports[q] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - !DivideRoundUp(next.rowBytes[q], next.peerChunkPayloadBytes, - &next.peerChunkCounts[q])) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - } + if (!AllocateProjectionQps(&next)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (!MapPhysicalQps(&next)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - if (usesUdma) { - uint64_t desiredChunk = requestedUdmaChunkBytes == 0 ? - kReduceGradDefaultUdmaChunkBytes : requestedUdmaChunkBytes; - if (desiredChunk < kReduceGradUdmaThresholdBytes || desiredChunk > UINT32_MAX) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - if (desiredChunk > maxUdmaRowBytes) { - desiredChunk = maxUdmaRowBytes; - } - if (!AlignUp(desiredChunk, kReduceGradUdmaAlignment, &next.udmaChunkBytes) || - next.udmaChunkBytes > UINT32_MAX) { + const uint64_t desiredChunk = requestedChunkBytes == 0 ? + kReduceGradDefaultChunkBytes : requestedChunkBytes; + if (desiredChunk == 0 || desiredChunk > UINT32_MAX || + !AlignUp(desiredChunk, kReduceGradUdmaAlignment, &next.chunkBytes) || + next.chunkBytes > UINT32_MAX) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { + if (!DivideRoundUp(next.rowBytes[projection], next.chunkBytes, + &next.chunkCounts[projection])) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } + } - uint64_t stateBytes = 0; - uint64_t stageCount = 0; - uint64_t payloadBytes = 0; - uint64_t cursor = 0; - if (!CheckedMul(static_cast(rankSize), kReduceGradUdmaPeerStateBytes, - &stateBytes) || - !AlignUp(stateBytes, kReduceGradUdmaAlignment, &cursor) || - !CheckedMul(static_cast(rankSize), 2, &stageCount) || - !CheckedMul(stageCount, next.udmaChunkBytes, &payloadBytes)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - next.udmaStateOffset = 0; - next.udmaOutboundOffset = cursor; - if (!CheckedAdd(cursor, payloadBytes, &cursor) || - !AlignUp(cursor, kReduceGradUdmaAlignment, &cursor)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - next.udmaInboundOffset = cursor; - if (!CheckedAdd(cursor, payloadBytes, &cursor) || - !AlignUp(cursor, kReduceGradUdmaWorkspaceAlignment, &next.workspaceBytes)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { - if (next.transports[q] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA && - !DivideRoundUp(next.rowBytes[q], next.udmaChunkBytes, - &next.udmaChunkCounts[q])) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - } - uint64_t sequenceCount = 0; - for (uint32_t q = 0; q < kReduceGradProjectionCount; ++q) { - uint64_t projectionSequences = 0; - if (!CheckedMul(static_cast(next.prefetchSlots), - next.udmaChunkCounts[q], &projectionSequences) || - !CheckedAdd(sequenceCount, projectionSequences, &sequenceCount)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } - } - if (sequenceCount > UINT32_MAX) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } + uint64_t stateBytes = 0; + uint64_t payloadBytes = 0; + uint64_t cursor = 0; + if (!CheckedMul(next.laneCount, kReduceGradLaneStateStrideBytes, &stateBytes) || + !AlignUp(stateBytes, kReduceGradUdmaAlignment, &next.laneStateBytes) || + !CheckedMul(static_cast(rankSize), next.chunkBytes, + &next.bankStrideBytes) || + !CheckedMul(kReduceGradBankCount, next.bankStrideBytes, + &next.laneStrideBytes) || + !CheckedMul(next.laneCount, next.laneStrideBytes, &payloadBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.stagingOffset = next.laneStateBytes; + cursor = next.stagingOffset; + if (!CheckedAdd(cursor, payloadBytes, &cursor) || + !AlignUp(cursor, kReduceGradWorkspaceAlignment, &next.workspaceBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } *out = next; @@ -238,14 +226,13 @@ int TileXRMoonEpBuildReduceGradLayout(int64_t rank, int64_t rankSize, int TileXRMoonEpBuildReduceGradLayout(int64_t rank, int64_t rankSize, int64_t expertCount, const uint64_t rowElements[kReduceGradProjectionCount], - uint64_t peerWindowBytes, uint64_t requestedUdmaChunkBytes, + uint32_t transportQpCount, uint64_t requestedChunkBytes, ReduceGradLayout *out) { const int64_t prefetchSlots = rankSize > 0 && expertCount > 0 ? expertCount / rankSize : 0; return TileXRMoonEpBuildReduceGradLayout(rank, rankSize, expertCount, - prefetchSlots, rowElements, peerWindowBytes, - requestedUdmaChunkBytes, out); + prefetchSlots, rowElements, transportQpCount, requestedChunkBytes, out); } } // namespace TileXRMoonEp diff --git a/src/moonep/reduce_grad/host/reduce_grad_layout.h b/src/moonep/reduce_grad/host/reduce_grad_layout.h index 8691c80e..155d1899 100644 --- a/src/moonep/reduce_grad/host/reduce_grad_layout.h +++ b/src/moonep/reduce_grad/host/reduce_grad_layout.h @@ -7,17 +7,15 @@ namespace TileXRMoonEp { -uint64_t TileXRMoonEpReduceGradPeerWindowBytes(); - int TileXRMoonEpBuildReduceGradLayout(int64_t rank, int64_t rankSize, int64_t expertCount, int64_t prefetchSlots, const uint64_t rowElements[kReduceGradProjectionCount], - uint64_t peerWindowBytes, uint64_t requestedUdmaChunkBytes, + uint32_t transportQpCount, uint64_t requestedChunkBytes, ReduceGradLayout *out); int TileXRMoonEpBuildReduceGradLayout(int64_t rank, int64_t rankSize, int64_t expertCount, const uint64_t rowElements[kReduceGradProjectionCount], - uint64_t peerWindowBytes, uint64_t requestedUdmaChunkBytes, + uint32_t transportQpCount, uint64_t requestedChunkBytes, ReduceGradLayout *out); } // namespace TileXRMoonEp diff --git a/src/moonep/reduce_grad/kernels/tilexr_moonep_reduce_grad_kernel.cpp b/src/moonep/reduce_grad/kernels/tilexr_moonep_reduce_grad_kernel.cpp index da7a755e..114e78dd 100644 --- a/src/moonep/reduce_grad/kernels/tilexr_moonep_reduce_grad_kernel.cpp +++ b/src/moonep/reduce_grad/kernels/tilexr_moonep_reduce_grad_kernel.cpp @@ -3,22 +3,33 @@ #include "comm_args.h" #include "reduce_grad_common.h" -#include "tilexr_data_as_flag.h" +#include "tilexr_sync.h" #include "tilexr_udma.h" namespace TileXRMoonEp { namespace Kernel { -constexpr uint32_t kPeerBatchRecords = 64; -constexpr uint32_t kIoBufferBytes = 32 * 1024; -constexpr uint32_t kFlagBufferBytes = 32 * 1024; -constexpr uint32_t kClearBufferBytes = 4 * 1024; -constexpr uint32_t kUdmaPollBackoffMax = 256; +constexpr uint32_t kSyncBufferBytes = + TileXR::TILEXR_MAX_RANK_SIZE * 4U * sizeof(int64_t); +constexpr uint32_t kWqeBufferBytes = TileXR::TILEXR_UDMA_WQE_SCRATCH_BYTES; +constexpr uint32_t kItemBufferBytes = sizeof(ReduceGradBankItem); + +struct ReduceGradPendingWave { + bool used[TileXR::TILEXR_MAX_RANK_SIZE]; + uint32_t frontiers[TileXR::TILEXR_MAX_RANK_SIZE]; + uint64_t chunkIndex; + uint64_t itemChunkBytes; + uint32_t sequence; + uint32_t projection; + uint32_t localExpert; + uint32_t waveStart; + uint32_t contributorCount; + uint32_t remoteContributorCount; +}; template -__aicore__ inline void SyncEvent() +__aicore__ inline void SyncEvent(AscendC::TEventID eventId) { - const AscendC::TEventID eventId = GetTPipePtr()->FetchEventID(event); AscendC::SetFlag(eventId); AscendC::WaitFlag(eventId); } @@ -30,676 +41,1151 @@ __aicore__ inline uint64_t MinU64(uint64_t lhs, uint64_t rhs) __aicore__ inline uint64_t CeilDivU64(uint64_t value, uint64_t divisor) { - return divisor == 0 ? 0 : value / divisor + (value % divisor == 0 ? 0 : 1); + return divisor == 0U ? 0U : value / divisor + (value % divisor == 0U ? 0U : 1U); } class ReduceGradKernel { public: - __aicore__ inline void Init(GM_ADDR commArgs, GM_ADDR expertsToCopy, - GM_ADDR gate, GM_ADDR up, GM_ADDR down, GM_ADDR workspace, GM_ADDR status, + __aicore__ inline void Init(GM_ADDR commArgs, GM_ADDR profileInfo, + GM_ADDR profileRegistry, GM_ADDR expertsToCopy, + GM_ADDR gateGradient, GM_ADDR upGradient, GM_ADDR downGradient, + GM_ADDR gateSource, GM_ADDR upSource, GM_ADDR downSource, + GM_ADDR workspace, GM_ADDR status, int64_t rank, int64_t rankSize, int64_t expertCount, - int64_t expertsPerRank, int64_t prefetchSlots, int64_t controlBlockCount, + int64_t expertsPerRank, int64_t prefetchSlots, uint64_t gateRowElements, uint64_t upRowElements, - uint64_t downRowElements, uint64_t gateRowBytes, uint64_t upRowBytes, - uint64_t downRowBytes, uint32_t gateTransport, uint32_t upTransport, - uint32_t downTransport, uint32_t udmaQpCount, uint64_t peerRecordBaseOffset, - uint64_t peerHalfBytes, uint64_t peerSlotStrideBytes, - uint64_t peerChunkPayloadBytes, uint64_t udmaStateOffset, - uint64_t udmaOutboundOffset, uint64_t udmaInboundOffset, - uint64_t udmaChunkBytes, uint64_t workspaceBytes, + uint64_t downRowElements, uint64_t gateRowBytes, + uint64_t upRowBytes, uint64_t downRowBytes, + uint64_t gateChunkCount, uint64_t upChunkCount, + uint64_t downChunkCount, uint32_t gateQpBase, + uint32_t upQpBase, uint32_t downQpBase, + uint32_t gateQpCount, uint32_t upQpCount, + uint32_t downQpCount, uint32_t lane0PhysicalQp, + uint32_t lane1PhysicalQp, uint32_t lane2PhysicalQp, + uint32_t lane3PhysicalQp, uint32_t lane4PhysicalQp, + uint32_t lane5PhysicalQp, uint32_t lane6PhysicalQp, + uint32_t lane7PhysicalQp, uint32_t transportQpCount, + uint32_t qpCount, uint32_t laneCount, + uint64_t laneStateBytes, uint64_t stagingOffset, + uint64_t bankStrideBytes, uint64_t laneStrideBytes, + uint64_t chunkBytes, uint64_t workspaceBytes, uint64_t waitIterations, int64_t magic) { args_ = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgs); + profileInfo_ = reinterpret_cast<__gm__ TileXR::UDMAInfo *>(profileInfo); + profileRegistry_ = reinterpret_cast<__gm__ TileXR::TileXRUDMAProfileRegistry *>( + profileRegistry); expertsToCopy_ = reinterpret_cast<__gm__ int32_t *>(expertsToCopy); - gradients_[kReduceGradGate] = gate; - gradients_[kReduceGradUp] = up; - gradients_[kReduceGradDown] = down; + gradients_[kReduceGradGate] = gateGradient; + gradients_[kReduceGradUp] = upGradient; + gradients_[kReduceGradDown] = downGradient; + sources_[kReduceGradGate] = gateSource; + sources_[kReduceGradUp] = upSource; + sources_[kReduceGradDown] = downSource; workspace_ = workspace; - status_ = reinterpret_cast<__gm__ int32_t *>(status); + status_ = reinterpret_cast<__gm__ uint32_t *>(status); rank_ = rank; rankSize_ = rankSize; expertCount_ = expertCount; expertsPerRank_ = expertsPerRank; prefetchSlots_ = prefetchSlots; - controlBlockCount_ = controlBlockCount; rowElements_[kReduceGradGate] = gateRowElements; rowElements_[kReduceGradUp] = upRowElements; rowElements_[kReduceGradDown] = downRowElements; rowBytes_[kReduceGradGate] = gateRowBytes; rowBytes_[kReduceGradUp] = upRowBytes; rowBytes_[kReduceGradDown] = downRowBytes; - transports_[kReduceGradGate] = gateTransport; - transports_[kReduceGradUp] = upTransport; - transports_[kReduceGradDown] = downTransport; - udmaQpCount_ = udmaQpCount; - peerRecordBaseOffset_ = peerRecordBaseOffset; - peerHalfBytes_ = peerHalfBytes; - peerSlotStrideBytes_ = peerSlotStrideBytes; - peerChunkPayloadBytes_ = peerChunkPayloadBytes; - udmaStateOffset_ = udmaStateOffset; - udmaOutboundOffset_ = udmaOutboundOffset; - udmaInboundOffset_ = udmaInboundOffset; - udmaChunkBytes_ = udmaChunkBytes; + chunkCounts_[kReduceGradGate] = gateChunkCount; + chunkCounts_[kReduceGradUp] = upChunkCount; + chunkCounts_[kReduceGradDown] = downChunkCount; + projectionQpBase_[kReduceGradGate] = gateQpBase; + projectionQpBase_[kReduceGradUp] = upQpBase; + projectionQpBase_[kReduceGradDown] = downQpBase; + projectionQpCounts_[kReduceGradGate] = gateQpCount; + projectionQpCounts_[kReduceGradUp] = upQpCount; + projectionQpCounts_[kReduceGradDown] = downQpCount; + lanePhysicalQps_[0] = lane0PhysicalQp; + lanePhysicalQps_[1] = lane1PhysicalQp; + lanePhysicalQps_[2] = lane2PhysicalQp; + lanePhysicalQps_[3] = lane3PhysicalQp; + lanePhysicalQps_[4] = lane4PhysicalQp; + lanePhysicalQps_[5] = lane5PhysicalQp; + lanePhysicalQps_[6] = lane6PhysicalQp; + lanePhysicalQps_[7] = lane7PhysicalQp; + transportQpCount_ = transportQpCount; + qpCount_ = qpCount; + laneCount_ = laneCount; + laneStateBytes_ = laneStateBytes; + stagingOffset_ = stagingOffset; + bankStrideBytes_ = bankStrideBytes; + laneStrideBytes_ = laneStrideBytes; + chunkBytes_ = chunkBytes; workspaceBytes_ = workspaceBytes; - waitIterations_ = waitIterations; + waitIterations_ = waitIterations == 0U ? 1U : waitIterations; magic_ = magic; + blockIdx_ = static_cast(AscendC::GetBlockIdx()); + blockCount_ = static_cast(AscendC::GetBlockNum()); - const int64_t subBlockCount = static_cast(get_subblockdim()); - blockIdx_ = static_cast(get_block_idx()) * subBlockCount + - static_cast(get_subblockid()); - blockCount_ = static_cast(get_block_num()) * subBlockCount; - receiverCount_ = blockCount_ - controlBlockCount_; - - pipe_.InitBuffer(ioBuf_, kIoBufferBytes); - pipe_.InitBuffer(accumBuf_, kIoBufferBytes); - pipe_.InitBuffer(flagBuf_, kFlagBufferBytes); - pipe_.InitBuffer(clearBuf_, kClearBufferBytes); - pipe_.InitBuffer(udmaWqeBuf_, TileXR::TILEXR_UDMA_WQE_SCRATCH_BYTES); + pipe_.InitBuffer(accumBuf_, kReduceGradKernelTileBytes); + pipe_.InitBuffer(accumPongBuf_, kReduceGradKernelTileBytes); + pipe_.InitBuffer(inputPingBuf_, kReduceGradKernelTileBytes); + pipe_.InitBuffer(inputPongBuf_, kReduceGradKernelTileBytes); + pipe_.InitBuffer(syncBuf_, kSyncBufferBytes); + pipe_.InitBuffer(wqeBuf_, kWqeBufferBytes); + pipe_.InitBuffer(itemBuf_, kItemBufferBytes); } __aicore__ inline void Process() { - if (args_ == nullptr || expertsToCopy_ == nullptr || status_ == nullptr || - rank_ < 0 || rank_ >= rankSize_ || rankSize_ <= 0 || expertCount_ <= 0 || - expertsPerRank_ <= 0 || prefetchSlots_ <= 0 || - prefetchSlots_ > expertsPerRank_ || controlBlockCount_ < 0 || - (rankSize_ > 1 && controlBlockCount_ == 0) || receiverCount_ <= 0 || - magic_ <= 0 || - (UsesUdma() && (udmaQpCount_ == 0 || - udmaQpCount_ > kReduceGradMaxUdmaQpCount || - TileXR::UDMAQpCount(args_) != udmaQpCount_))) { - SetStatus(kReduceGradDeviceInvalidState); + if (args_ == nullptr || status_ == nullptr) { + return; + } + if (blockIdx_ == 0U) { + AscendC::AtomicExch(status_, static_cast(0U)); + } + AscendC::SyncAll(); + + if (!CollectiveConfigurationValid()) { + PublishStatus(kReduceGradDeviceInvalidState); return; } - if (blockIdx_ < controlBlockCount_) { - RunSender(); + const bool configured = ConfigurationValid(); + if (!configured) { + PublishStatus(kReduceGradDeviceInvalidState); } else { - RunReceiver(); + ValidatePlanEntries(); + } + AscendC::SyncAll(); + if (DeviceStatus() == 0U) { + if (blockIdx_ < laneCount_) { + RunLaneLeader(blockIdx_); + } else { + RunLaneHelper(); + } + } + + AscendC::SyncAll(); + if (blockIdx_ == 0U) { + CrossRankBarrier(); + } + AscendC::SyncAll(); + if (DeviceStatus() == 0U) { + ClearLocalSources(); } } private: - __aicore__ inline void SetStatus(int32_t value) + __aicore__ inline bool CollectiveConfigurationValid() const { - status_[0] = value; + return rank_ >= 0 && rank_ < rankSize_ && + rankSize_ >= kReduceGradMinRankCount && + rankSize_ <= TileXR::TILEXR_MAX_RANK_SIZE && blockCount_ != 0U && + blockIdx_ < blockCount_ && magic_ > 0 && args_->peerMems[rank_] != nullptr; } - __aicore__ inline int32_t ExpertForSlot(int64_t sourceRank, int64_t slot) const + __aicore__ inline bool DataPointersValid() const { - return expertsToCopy_[sourceRank * prefetchSlots_ + slot]; + if (expertsToCopy_ == nullptr) { + return false; + } + for (uint32_t projection = 0U; + projection < kReduceGradProjectionCount; ++projection) { + if (gradients_[projection] == nullptr || sources_[projection] == nullptr) { + return false; + } + } + return true; } - __aicore__ inline int64_t OwnerForExpert(int32_t expert) const + __aicore__ inline bool ConfigurationValid() const { - return expert < 0 ? -1 : static_cast(expert) / expertsPerRank_; + if (!DataPointersValid() || expertCount_ <= 0 || + expertsPerRank_ <= 0 || expertCount_ != expertsPerRank_ * rankSize_ || + expertCount_ > INT32_MAX || prefetchSlots_ <= 0 || + prefetchSlots_ > INT32_MAX || prefetchSlots_ > INT32_MAX / rankSize_ || + blockCount_ < 2U * laneCount_) { + return false; + } + for (uint32_t projection = 0U; + projection < kReduceGradProjectionCount; ++projection) { + if (rowElements_[projection] == 0U || rowBytes_[projection] == 0U || + rowBytes_[projection] != rowElements_[projection] * sizeof(float)) { + return false; + } + } + if (workspace_ == nullptr || profileInfo_ == nullptr || + profileRegistry_ == nullptr || + transportQpCount_ < kReduceGradMinMultiRankQpCount || + transportQpCount_ > kReduceGradMaxTransportQpCount || + qpCount_ < kReduceGradProjectionCount || + qpCount_ > kReduceGradMaxUdmaQpCount || laneCount_ != qpCount_ || + chunkBytes_ == 0U || + chunkBytes_ > UINT32_MAX || laneStateBytes_ < + laneCount_ * kReduceGradLaneStateStrideBytes || + stagingOffset_ != laneStateBytes_ || + bankStrideBytes_ != static_cast(rankSize_) * chunkBytes_ || + laneStrideBytes_ != kReduceGradBankCount * bankStrideBytes_ || + workspaceBytes_ < stagingOffset_ + laneCount_ * laneStrideBytes_ || + !TileXR::UDMAProfileRegistryValid(args_, profileInfo_, profileRegistry_) || + profileRegistry_->regionCount != kReduceGradProfileRegionCount || + profileRegistry_->qpCount != transportQpCount_) { + return false; + } + for (uint32_t lane = 0U; lane < laneCount_; ++lane) { + if (lanePhysicalQps_[lane] >= transportQpCount_) { + return false; + } + for (uint32_t prior = 0U; prior < lane; ++prior) { + if (lanePhysicalQps_[prior] == lanePhysicalQps_[lane]) { + return false; + } + } + } + uint32_t cursor = 0U; + for (uint32_t projection = 0U; + projection < kReduceGradProjectionCount; ++projection) { + if (projectionQpCounts_[projection] == 0U || + projectionQpBase_[projection] != cursor || + chunkCounts_[projection] != CeilDivU64( + rowBytes_[projection], chunkBytes_)) { + return false; + } + cursor += projectionQpCounts_[projection]; + } + return cursor == qpCount_; } - __aicore__ inline uint64_t ChunkBytes(uint32_t projection) const + __attribute__((always_inline)) inline __aicore__ void PublishStatus(uint32_t value) { - return transports_[projection] == kReduceGradTransportPeer ? - peerChunkPayloadBytes_ : udmaChunkBytes_; + if (value != 0U) { + AscendC::AtomicCas(status_, static_cast(0U), value); + } } - __aicore__ inline uint64_t ChunkCount(uint32_t projection) const + __aicore__ inline uint32_t DeviceStatus() const { - return CeilDivU64(rowBytes_[projection], ChunkBytes(projection)); + return AscendC::AtomicAdd(status_, static_cast(0U)); } - __aicore__ inline uint64_t Sequence(uint32_t projection, int64_t slot, - uint64_t chunk) const + __aicore__ inline void ValidatePlanEntries() { - uint64_t ordinal = 1; - for (uint32_t q = 0; q < projection; ++q) { - if (transports_[q] == kReduceGradTransportUdma) { - ordinal += static_cast(prefetchSlots_) * ChunkCount(q); + const uint64_t entries = static_cast(rankSize_) * + static_cast(prefetchSlots_); + for (uint64_t entry = blockIdx_; entry < entries; entry += blockCount_) { + const int32_t expert = expertsToCopy_[entry]; + if (expert < -1 || expert >= expertCount_) { + PublishStatus(kReduceGradDeviceInvalidState); } } - ordinal += static_cast(slot) * ChunkCount(projection) + chunk; - return (static_cast(static_cast(magic_)) << 32) | - static_cast(ordinal); } - __aicore__ inline bool UsesUdma() const + __aicore__ inline uint64_t LaneToken(uint32_t sequence) const { - for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { - if (transports_[projection] == kReduceGradTransportUdma) { - return true; + return (static_cast(static_cast(magic_)) << 32U) | + static_cast(sequence); + } + + __aicore__ inline uint32_t ProjectionForLane(uint32_t lane) const + { + for (uint32_t projection = 0U; + projection < kReduceGradProjectionCount; ++projection) { + const uint32_t begin = projectionQpBase_[projection]; + if (lane >= begin && lane < begin + projectionQpCounts_[projection]) { + return projection; } } - return false; + return kReduceGradProjectionCount; } - __aicore__ inline uint32_t UDMAQpFor(uint32_t projection, int64_t slot, - uint64_t chunk) const + __attribute__((always_inline)) inline __aicore__ uint32_t PhysicalQp( + uint32_t lane) const { - const uint32_t ordinal = static_cast(Sequence(projection, slot, chunk)); - return (ordinal - 1U) % udmaQpCount_; + return lanePhysicalQps_[lane]; } - __aicore__ inline int64_t ControlPeer(int64_t controlIndex) const + __aicore__ inline uint64_t LoadLaneToken(__gm__ uint64_t *token) const { - return controlIndex >= rank_ ? controlIndex + 1 : controlIndex; + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(token), sizeof(uint64_t)); + return *token; } - __aicore__ inline GM_ADDR PeerRecord(int64_t owner, int64_t source, - int64_t slot, uint64_t chunk) const + __aicore__ inline void StoreLaneToken( + __gm__ uint64_t *token, uint64_t value) const { - const uint64_t stage = chunk & 1U; - const uint64_t slotIndex = static_cast(source * prefetchSlots_ + slot); - return args_->peerMems[owner] + TileXR::IPC_DATA_OFFSET + - peerRecordBaseOffset_ + stage * peerHalfBytes_ + - slotIndex * peerSlotStrideBytes_; + *token = value; + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(token), sizeof(uint64_t)); } - __aicore__ inline uint64_t UDMAPeerStateOffset(int64_t peer) const + __aicore__ inline __gm__ uint64_t *BankReady(uint32_t lane, uint32_t bank) const { - return udmaStateOffset_ + static_cast(peer) * kReduceGradUdmaPeerStateBytes; + return reinterpret_cast<__gm__ uint64_t *>(workspace_ + + static_cast(lane) * kReduceGradLaneStateStrideBytes + + kReduceGradBankReadyOffset + + static_cast(bank) * kReduceGradLaneFlagStrideBytes); } - __aicore__ inline uint64_t UDMAReadyOffset(int64_t source, uint64_t stage) const + __aicore__ inline __gm__ uint64_t *BankDoneToken( + uint32_t lane, uint32_t bank, uint32_t helperIndex) const { - return UDMAPeerStateOffset(source) + kReduceGradUdmaReadyOffset + - stage * kReduceGradUdmaSignalStageStride; + return reinterpret_cast<__gm__ uint64_t *>(workspace_ + + static_cast(lane) * kReduceGradLaneStateStrideBytes + + kReduceGradBankDoneOffset + + static_cast(bank) * kReduceGradDoneBankStrideBytes + + static_cast(helperIndex) * kReduceGradDoneTokenBytes); } - __aicore__ inline uint64_t UDMACompletionOffset(int64_t source, uint64_t stage) const + __aicore__ inline __gm__ ReduceGradBankItem *BankItem( + uint32_t lane, uint32_t bank) const { - return UDMAPeerStateOffset(source) + kReduceGradUdmaCompletionOffset + - stage * kReduceGradUdmaSignalStageStride; + return reinterpret_cast<__gm__ ReduceGradBankItem *>(workspace_ + + static_cast(lane) * kReduceGradLaneStateStrideBytes + + kReduceGradBankItemOffset + + static_cast(bank) * kReduceGradLaneFlagStrideBytes); } - __aicore__ inline uint64_t UDMAPollScratchOffset(int64_t peer, uint64_t stage) const + __aicore__ inline __gm__ uint64_t *LaneAbort(uint32_t lane) const { - return UDMAPeerStateOffset(peer) + kReduceGradUdmaPollScratchOffset + - stage * kReduceGradUdmaSignalStageStride; + return reinterpret_cast<__gm__ uint64_t *>(workspace_ + + static_cast(lane) * kReduceGradLaneStateStrideBytes + + kReduceGradLaneErrorOffset); } - __aicore__ inline uint64_t UDMAOutboundOffset(int64_t target, uint64_t stage) const + __aicore__ inline uint64_t LaneAbortToken() const { - return udmaOutboundOffset_ + - (static_cast(target) * 2 + stage) * udmaChunkBytes_; + return LaneToken(UINT32_MAX); } - __aicore__ inline uint64_t UDMAInboundOffset(int64_t source, uint64_t stage) const + __aicore__ inline bool LaneAborted(uint32_t lane) const { - return udmaInboundOffset_ + - (static_cast(source) * 2 + stage) * udmaChunkBytes_; + return LoadLaneToken(LaneAbort(lane)) == LaneAbortToken(); } - __aicore__ inline uint64_t LoadSignal(GM_ADDR address) + __aicore__ inline void AbortLane(uint32_t lane) { - TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(address), sizeof(uint64_t)); - AscendC::PipeBarrier(); - return reinterpret_cast<__gm__ uint64_t *>(address)[0]; + StoreLaneToken(LaneAbort(lane), LaneAbortToken()); } - __aicore__ inline void UdmaPollBackoff(uint64_t attempt) + __aicore__ inline void FailLane(uint32_t lane) { - const uint32_t shift = static_cast(attempt >> 4) > 8U ? - 8U : static_cast(attempt >> 4); - const uint32_t spins = 1U << shift; - for (volatile uint32_t spin = 0; spin < spins && spin < kUdmaPollBackoffMax; - ++spin) { - __asm__ __volatile__(""); - } + AbortLane(lane); } - __aicore__ inline void StoreSignal(GM_ADDR address, uint64_t value) + __aicore__ inline GM_ADDR BankPayload(uint32_t lane, uint32_t bank) const { - reinterpret_cast<__gm__ uint64_t *>(address)[0] = value; - AscendC::PipeBarrier(); - TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(address), sizeof(uint64_t)); - AscendC::PipeBarrier(); + return workspace_ + stagingOffset_ + + static_cast(lane) * laneStrideBytes_ + + static_cast(bank) * bankStrideBytes_; } - __aicore__ inline bool WaitUdmaReady(int64_t source, uint64_t stage, uint64_t sequence) + __aicore__ inline void HelperShape(uint32_t lane, + uint32_t &helperCount, uint32_t &helperBegin) const { - GM_ADDR ready = workspace_ + UDMAReadyOffset(source, stage); - for (uint64_t attempt = 0; attempt < waitIterations_; ++attempt) { - if (LoadSignal(ready) == sequence) { - return true; - } + const uint32_t totalHelpers = blockCount_ - laneCount_; + const uint32_t base = totalHelpers / laneCount_; + const uint32_t extra = totalHelpers % laneCount_; + helperCount = base + (lane < extra ? 1U : 0U); + helperBegin = laneCount_ + lane * base + (lane < extra ? lane : extra); + } + + __aicore__ inline void HelperLane(uint32_t &lane, + uint32_t &helperIndex, uint32_t &helperCount) const + { + const uint32_t ordinal = blockIdx_ - laneCount_; + const uint32_t totalHelpers = blockCount_ - laneCount_; + const uint32_t base = totalHelpers / laneCount_; + const uint32_t extra = totalHelpers % laneCount_; + const uint32_t wider = (base + 1U) * extra; + if (ordinal < wider) { + lane = ordinal / (base + 1U); + helperIndex = ordinal % (base + 1U); + helperCount = base + 1U; + } else { + const uint32_t tailOrdinal = ordinal - wider; + lane = extra + tailOrdinal / base; + helperIndex = tailOrdinal % base; + helperCount = base; } - SetStatus(kReduceGradDevicePeerTimeout); - return false; } - __aicore__ inline bool WaitUdmaCompletion(int64_t target, uint32_t qpIdx, - uint64_t stage, uint64_t sequence) + __aicore__ inline uint32_t CountContributors(int32_t globalExpert) const { - const uint64_t remoteOffset = UDMACompletionOffset(rank_, stage); - GM_ADDR localScratch = workspace_ + UDMAPollScratchOffset(target, stage); - AscendC::LocalTensor wqeScratch = udmaWqeBuf_.Get(); - for (uint64_t attempt = 0; attempt < waitIterations_; ++attempt) { - const uint32_t getStatus = TileXR::UDMAGetNbiOnQp(args_, wqeScratch, - static_cast(target), qpIdx, - reinterpret_cast<__gm__ uint8_t *>(localScratch), remoteOffset, - static_cast(sizeof(uint64_t))); - if (getStatus != TileXR::TILEXR_UDMA_STATUS_SUCCESS || - TileXR::UDMAQuietStatusOnQp(args_, static_cast(target), qpIdx) != - TileXR::TILEXR_UDMA_STATUS_SUCCESS) { - SetStatus(kReduceGradDeviceUdmaCqError); - return false; - } - if (LoadSignal(localScratch) == sequence) { - return true; + uint32_t count = 0U; + const uint64_t entries = static_cast(rankSize_) * + static_cast(prefetchSlots_); + for (uint64_t entry = 0U; entry < entries; ++entry) { + if (expertsToCopy_[entry] == globalExpert) { + ++count; } - UdmaPollBackoff(attempt); } - SetStatus(kReduceGradDevicePeerTimeout); - return false; + return count; } - __aicore__ inline bool WaitPeerRecords(GM_ADDR recordBase, uint64_t payloadBytes, - bool ready) + __attribute__((always_inline)) inline __aicore__ bool IssueWave( + uint32_t lane, uint32_t projection, + uint32_t bank, int32_t globalExpert, uint64_t chunkIndex, + uint32_t waveStart, uint32_t contributorCount, + uint32_t sequence, ReduceGradPendingWave &pending) { - AscendC::LocalTensor flags = flagBuf_.Get(); - const uint32_t totalRecords = TileXR::DataAsFlagBlockCountForPayloadBytes(payloadBytes); - const uint32_t capacity = TileXR::DataAsFlagMaxCheckBlocks(kFlagBufferBytes); - if (totalRecords == 0 || capacity == 0) { - SetStatus(kReduceGradDeviceInvalidState); - return false; + AscendC::LocalTensor wqeScratch = wqeBuf_.Get(); + const uint64_t chunkOffset = chunkIndex * chunkBytes_; + const uint64_t bytes = MinU64( + rowBytes_[projection] - chunkOffset, chunkBytes_); + for (int32_t sourceRank = 0; sourceRank < rankSize_; ++sourceRank) { + pending.used[sourceRank] = false; + pending.frontiers[sourceRank] = 0U; } - uint32_t record = 0; - while (record < totalRecords) { - const uint32_t batch = MinU64(totalRecords - record, capacity); - bool matched = false; - for (uint64_t attempt = 0; attempt < waitIterations_; ++attempt) { - matched = ready ? TileXR::DataAsFlagCheckBatch( - reinterpret_cast<__gm__ uint8_t *>(recordBase), record, batch, flags) : - TileXR::DataAsFlagCheckBatchCleared( - reinterpret_cast<__gm__ uint8_t *>(recordBase), record, batch, flags); - if (matched) { - break; - } + pending.chunkIndex = chunkIndex; + pending.itemChunkBytes = bytes; + pending.sequence = sequence; + pending.projection = projection; + pending.localExpert = static_cast( + globalExpert - rank_ * expertsPerRank_); + pending.waveStart = waveStart; + pending.contributorCount = contributorCount; + pending.remoteContributorCount = 0U; + uint32_t matching = 0U; + uint32_t collected = 0U; + bool ok = true; + + const uint64_t entries = static_cast(rankSize_) * + static_cast(prefetchSlots_); + for (uint64_t entry = 0U; entry < entries && collected < contributorCount; + ++entry) { + if (expertsToCopy_[entry] != globalExpert) { + continue; } - if (!matched) { - SetStatus(kReduceGradDevicePeerTimeout); - return false; + if (matching++ < waveStart) { + continue; + } + const int32_t sourceRank = static_cast( + entry / static_cast(prefetchSlots_)); + const uint32_t slot = static_cast( + entry % static_cast(prefetchSlots_)); + if (sourceRank != rank_) { + const uint64_t localOffset = stagingOffset_ + + static_cast(lane) * laneStrideBytes_ + + static_cast(bank) * bankStrideBytes_ + + static_cast(pending.remoteContributorCount) * chunkBytes_; + const uint64_t remoteOffset = static_cast(slot) * + rowBytes_[projection] + chunkOffset; + const uint32_t post = TileXR::UDMAProfileGetNbiOnQpDeferred( + args_, profileInfo_, profileRegistry_, wqeScratch, + sourceRank, PhysicalQp(lane), kReduceGradStagingRegion, localOffset, + projection + 1U, remoteOffset, bytes); + if (post != TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + ok = false; + PublishStatus(kReduceGradDeviceUdmaCqError); + } else { + pending.used[sourceRank] = true; + } + ++pending.remoteContributorCount; } - record += batch; + ++collected; + } + if (collected != contributorCount) { + ok = false; + PublishStatus(kReduceGradDeviceInvalidState); } - return true; - } - __aicore__ inline void CopyDense(GM_ADDR destination, GM_ADDR source, uint64_t bytes) - { - AscendC::LocalTensor local = ioBuf_.Get(); - uint64_t offset = 0; - while (offset < bytes) { - const uint32_t batch = static_cast(MinU64(bytes - offset, kIoBufferBytes)); - AscendC::GlobalTensor sourceGlobal; - sourceGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(source + offset)); - AscendC::DataCopyExtParams inParams {1U, batch, 0U, 0U, 0U}; - AscendC::DataCopyPadExtParams pad {false, 0U, 0U, 0U}; - AscendC::DataCopyPad(local, sourceGlobal, inParams, pad); - SyncEvent(); - AscendC::GlobalTensor destinationGlobal; - destinationGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(destination + offset)); - AscendC::DataCopyPad(destinationGlobal, local, inParams); - SyncEvent(); - offset += batch; - } - } - - __aicore__ inline void ZeroDense(GM_ADDR destination, uint64_t bytes) - { - AscendC::LocalTensor local = ioBuf_.Get(); - uint64_t offset = 0; - while (offset < bytes) { - const uint32_t batch = static_cast(MinU64(bytes - offset, kIoBufferBytes)); - AscendC::Duplicate(local, 0.0f, batch / sizeof(float)); - SyncEvent(); - AscendC::GlobalTensor destinationGlobal; - destinationGlobal.SetGlobalBuffer( - reinterpret_cast<__gm__ float *>(destination + offset)); - AscendC::DataCopyExtParams outParams {1U, batch, 0U, 0U, 0U}; - AscendC::DataCopyPad(destinationGlobal, local, outParams); - SyncEvent(); - offset += batch; - } - } - - __aicore__ inline void AddDense(GM_ADDR destination, GM_ADDR source, uint64_t bytes) - { - AscendC::LocalTensor input = ioBuf_.Get(); - AscendC::LocalTensor output = accumBuf_.Get(); - uint64_t offset = 0; - while (offset < bytes) { - const uint32_t batch = static_cast(MinU64(bytes - offset, kIoBufferBytes)); - AscendC::GlobalTensor inputGlobal; - inputGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(source + offset)); - AscendC::GlobalTensor outputGlobal; - outputGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(destination + offset)); - AscendC::DataCopyExtParams inParams {1U, batch, 0U, 0U, 0U}; - AscendC::DataCopyPadExtParams pad {false, 0U, 0U, 0U}; - AscendC::DataCopyPad(input, inputGlobal, inParams, pad); - AscendC::DataCopyPad(output, outputGlobal, inParams, pad); - SyncEvent(); - AscendC::Add(output, output, input, batch / sizeof(float)); - SyncEvent(); - AscendC::DataCopyPad(outputGlobal, output, inParams); - SyncEvent(); - offset += batch; - } - } - - __aicore__ inline bool AddPeerPacked(GM_ADDR destination, GM_ADDR recordBase, - uint64_t payloadBytes) - { - if (!WaitPeerRecords(recordBase, payloadBytes, true)) { - return false; + for (int32_t sourceRank = 0; sourceRank < rankSize_; ++sourceRank) { + if (pending.used[sourceRank]) { + pending.frontiers[sourceRank] = TileXR::UDMAProfileCompletionFrontier( + args_, profileInfo_, profileRegistry_, sourceRank, + PhysicalQp(lane)); + } } - AscendC::LocalTensor input = ioBuf_.Get(); - AscendC::LocalTensor output = accumBuf_.Get(); - AscendC::LocalTensor clear = clearBuf_.Get(); - const uint32_t totalRecords = TileXR::DataAsFlagBlockCountForPayloadBytes(payloadBytes); - uint32_t record = 0; - while (record < totalRecords) { - const uint32_t records = static_cast(MinU64( - totalRecords - record, kPeerBatchRecords)); - const uint64_t payloadOffset = static_cast(record) * - TileXR::DATA_AS_FLAG_PAYLOAD_BYTES; - const uint32_t batchBytes = static_cast(MinU64( - payloadBytes - payloadOffset, - static_cast(records) * TileXR::DATA_AS_FLAG_PAYLOAD_BYTES)); - - AscendC::GlobalTensor packedGlobal; - packedGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( - recordBase + static_cast(record) * TileXR::DATA_AS_FLAG_BLOCK_BYTES)); - AscendC::DataCopyExtParams packedParams { - static_cast(records), TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, - TileXR::DATA_AS_FLAG_FLAG_BYTES, 0U, 0U}; - AscendC::DataCopyPadExtParams pad {false, 0U, 0U, 0U}; - AscendC::DataCopyPad(input, packedGlobal, packedParams, pad); - - AscendC::GlobalTensor outputGlobal; - outputGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( - destination + payloadOffset)); - AscendC::DataCopyExtParams outputParams {1U, batchBytes, 0U, 0U, 0U}; - AscendC::DataCopyPad(output, outputGlobal, outputParams, pad); - SyncEvent(); - AscendC::Add(output, output, input, batchBytes / sizeof(float)); - SyncEvent(); - AscendC::DataCopyPad(outputGlobal, output, outputParams); - SyncEvent(); - - AscendC::Duplicate(clear, 0.0f, - records * TileXR::DATA_AS_FLAG_FLAG_FLOATS); - SyncEvent(); - AscendC::GlobalTensor flagsGlobal; - flagsGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( - recordBase + static_cast(record) * TileXR::DATA_AS_FLAG_BLOCK_BYTES + - TileXR::DATA_AS_FLAG_FLAG_OFFSET_BYTES)); - AscendC::DataCopyExtParams clearParams { - static_cast(records), TileXR::DATA_AS_FLAG_FLAG_BYTES, - 0U, TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, 0U}; - AscendC::DataCopyPad(flagsGlobal, clear, clearParams); - SyncEvent(); - record += records; + for (int32_t sourceRank = 0; sourceRank < rankSize_; ++sourceRank) { + if (!pending.used[sourceRank]) { + continue; + } + const uint32_t flush = TileXR::UDMAProfileFlushQpDoorbell( + args_, profileInfo_, profileRegistry_, sourceRank, + PhysicalQp(lane)); + if (flush != TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + ok = false; + PublishStatus(kReduceGradDeviceUdmaCqError); + } } - return true; + return ok; } - __aicore__ inline bool SendPeerChunk(int64_t target, int64_t slot, - uint32_t projection, uint64_t chunk, GM_ADDR source, uint64_t bytes) + __attribute__((always_inline)) inline __aicore__ bool CompleteWave( + uint32_t lane, const ReduceGradPendingWave &pending) { - GM_ADDR record = PeerRecord(target, rank_, slot, chunk); - if (!WaitPeerRecords(record, bytes, false)) { - return false; + bool ok = true; + for (int32_t sourceRank = 0; sourceRank < rankSize_; ++sourceRank) { + if (!pending.used[sourceRank]) { + continue; + } + const uint32_t quiet = TileXR::UDMAProfileQuietStatusOnQpUntil( + args_, profileInfo_, profileRegistry_, sourceRank, + PhysicalQp(lane), + pending.frontiers[sourceRank]); + if (quiet != TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + ok = false; + PublishStatus(quiet == TileXR::TILEXR_UDMA_STATUS_CQ_TIMEOUT ? + kReduceGradDeviceUdmaTimeout : kReduceGradDeviceUdmaCqError); + } } - AscendC::LocalTensor local = ioBuf_.Get(); - TileXR::DataAsFlagInit(local); - return TileXR::DataAsFlagSend(reinterpret_cast<__gm__ uint8_t *>(record), - reinterpret_cast<__gm__ uint8_t *>(source), bytes, local) != 0; + return ok; } - __aicore__ inline bool PostUdmaChunk(int64_t target, int64_t slot, - uint32_t projection, uint64_t chunk, GM_ADDR source, uint64_t bytes) + __aicore__ inline bool WaitForHelpers(uint32_t lane, uint32_t bank, + uint32_t helperCount, uint32_t sequence) { - const uint64_t stage = chunk & 1U; - const uint64_t outboundOffset = UDMAOutboundOffset(target, stage); - GM_ADDR outbound = workspace_ + outboundOffset; - CopyDense(outbound, source, bytes); - TileXR::UDMACleanCacheLines( - reinterpret_cast<__gm__ uint8_t *>(outbound), bytes); - AscendC::PipeBarrier(); - const uint64_t sequence = Sequence(projection, slot, chunk); - const uint32_t qpIdx = UDMAQpFor(projection, slot, chunk); - AscendC::LocalTensor wqeScratch = udmaWqeBuf_.Get(); - const uint32_t putStatus = TileXR::UDMAPutRegisteredSignalNbiOnQp( - args_, wqeScratch, static_cast(target), qpIdx, - reinterpret_cast<__gm__ uint8_t *>(outbound), - UDMAInboundOffset(rank_, stage), static_cast(bytes), - UDMAReadyOffset(rank_, stage), sequence); - if (putStatus != TileXR::TILEXR_UDMA_STATUS_SUCCESS || - TileXR::UDMAQuietStatusOnQp(args_, static_cast(target), qpIdx) != - TileXR::TILEXR_UDMA_STATUS_SUCCESS) { - SetStatus(kReduceGradDeviceUdmaCqError); - return false; + const uint64_t expected = LaneToken(sequence); + for (uint32_t helper = 0U; helper < helperCount; ++helper) { + uint64_t attempts = 0U; + while (LoadLaneToken(BankDoneToken(lane, bank, helper)) != expected) { + if (LaneAborted(lane)) { + return false; + } + if (++attempts >= waitIterations_) { + PublishStatus(kReduceGradDeviceLeaderTimeout); + AbortLane(lane); + return false; + } + } } return true; } - __aicore__ inline bool CompleteChunk(int64_t target, int64_t slot, - uint32_t projection, uint64_t chunk, GM_ADDR source, uint64_t bytes) + __attribute__((always_inline)) inline __aicore__ void PublishItem( + uint32_t lane, uint32_t bank, + uint64_t token, uint64_t chunkIndex, uint64_t itemChunkBytes, + uint32_t kind, uint32_t projection, uint32_t localExpert, + uint32_t waveStart, uint32_t contributorCount, + uint32_t remoteContributorCount) { - bool complete = false; - if (transports_[projection] == kReduceGradTransportPeer) { - complete = WaitPeerRecords( - PeerRecord(target, rank_, slot, chunk), bytes, false); - } else { - const uint32_t qpIdx = UDMAQpFor(projection, slot, chunk); - complete = WaitUdmaCompletion( - target, qpIdx, chunk & 1U, Sequence(projection, slot, chunk)); - } - if (!complete) { - return false; - } - ZeroDense(source, bytes); - return true; + AscendC::LocalTensor local = itemBuf_.Get(); + __ubuf__ uint64_t *localWords = reinterpret_cast<__ubuf__ uint64_t *>( + local.GetPhyAddr()); + localWords[0] = token; + localWords[1] = chunkIndex; + localWords[2] = itemChunkBytes; + localWords[3] = PackUint32Pair(kind, projection); + localWords[4] = PackUint32Pair(localExpert, waveStart); + localWords[5] = PackUint32Pair( + contributorCount, remoteContributorCount); + localWords[6] = 0U; + localWords[7] = 0U; + AscendC::GlobalTensor destination; + destination.SetGlobalBuffer( + reinterpret_cast<__gm__ uint64_t *>(BankItem(lane, bank)), + sizeof(ReduceGradBankItem) / sizeof(uint64_t)); + const AscendC::DataCopyExtParams params { + 1U, sizeof(ReduceGradBankItem), 0U, 0U, 0U}; + SyncEvent(EVENT_ID3); + AscendC::DataCopyPad(destination, local, params); + SyncEvent(EVENT_ID3); + StoreLaneToken(BankReady(lane, bank), token); } - __aicore__ inline bool RunSenderTo(int64_t target) + __aicore__ inline uint64_t PackUint32Pair(uint32_t low, uint32_t high) const { - if (target < 0 || target >= rankSize_) { - SetStatus(kReduceGradDeviceInvalidState); - return false; + return static_cast(low) | + (static_cast(high) << 32U); + } + + __attribute__((always_inline)) inline __aicore__ void LoadItem( + uint32_t lane, uint32_t bank, uint64_t &token, uint64_t &chunkIndex, + uint64_t &itemChunkBytes, uint32_t &kind, uint32_t &projection, + uint32_t &localExpert, uint32_t &waveStart, uint32_t &contributorCount, + uint32_t &remoteContributorCount) + { + AscendC::LocalTensor local = itemBuf_.Get(); + __ubuf__ uint64_t *localWords = reinterpret_cast<__ubuf__ uint64_t *>( + local.GetPhyAddr()); + AscendC::GlobalTensor source; + source.SetGlobalBuffer( + reinterpret_cast<__gm__ uint64_t *>(BankItem(lane, bank)), + sizeof(ReduceGradBankItem) / sizeof(uint64_t)); + const AscendC::DataCopyExtParams params { + 1U, sizeof(ReduceGradBankItem), 0U, 0U, 0U}; + const AscendC::DataCopyPadExtParams pad {false, 0U, 0U, 0U}; + AscendC::DataCopyPad(local, source, params, pad); + SyncEvent(EVENT_ID3); + + token = localWords[0]; + chunkIndex = localWords[1]; + itemChunkBytes = localWords[2]; + const uint64_t kindProjection = localWords[3]; + const uint64_t expertWave = localWords[4]; + const uint64_t contributors = localWords[5]; + kind = static_cast(kindProjection); + projection = static_cast(kindProjection >> 32U); + localExpert = static_cast(expertWave); + waveStart = static_cast(expertWave >> 32U); + contributorCount = static_cast(contributors); + remoteContributorCount = static_cast(contributors >> 32U); + } + + __aicore__ inline void RunLaneLeader(uint32_t lane) + { + uint32_t helperCount = 0U; + uint32_t helperBegin = 0U; + HelperShape(lane, helperCount, helperBegin); + (void)helperBegin; + const uint32_t projection = ProjectionForLane(lane); + if (projection >= kReduceGradProjectionCount || helperCount == 0U) { + PublishStatus(kReduceGradDeviceInvalidState); + FailLane(lane); + return; } - for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { - const uint64_t chunkBytes = ChunkBytes(projection); - const uint64_t chunks = ChunkCount(projection); - if (chunkBytes == 0 || chunks == 0) { - SetStatus(kReduceGradDeviceInvalidState); - return false; + const uint32_t projectionLane = lane - projectionQpBase_[projection]; + ReduceGradPendingWave pending[kReduceGradBankCount]; + bool bankIssued[kReduceGradBankCount] = {}; + bool bankPublished[kReduceGradBankCount] = {}; + uint32_t bankSequence[kReduceGradBankCount] = {}; + uint32_t sequence = 1U; + uint32_t nextReadySequence = 1U; + bool stop = false; + + for (int64_t localExpert = 0; localExpert < expertsPerRank_ && !stop; + ++localExpert) { + const int32_t globalExpert = static_cast( + rank_ * expertsPerRank_ + localExpert); + const uint32_t totalContributors = CountContributors(globalExpert); + if (totalContributors == 0U) { + continue; } - for (int64_t slot = 0; slot < prefetchSlots_; ++slot) { - const int32_t expert = ExpertForSlot(rank_, slot); - if (expert < 0) { - continue; - } - if (expert >= expertCount_) { - SetStatus(kReduceGradDeviceInvalidState); - return false; - } - if (OwnerForExpert(expert) != target) { + for (uint64_t chunk = 0U; + chunk < chunkCounts_[projection] && !stop; ++chunk) { + const uint64_t workOrdinal = + static_cast(localExpert) * chunkCounts_[projection] + chunk; + if (workOrdinal % projectionQpCounts_[projection] != projectionLane) { continue; } - GM_ADDR sourceRow = gradients_[projection] + - static_cast(expertCount_ + slot) * rowBytes_[projection]; - for (uint64_t chunk = 0; chunk < chunks; ++chunk) { - if (chunk >= 2) { - const uint64_t completedChunk = chunk - 2; - const uint64_t completedOffset = completedChunk * chunkBytes; - const uint64_t completedBytes = MinU64( - rowBytes_[projection] - completedOffset, chunkBytes); - if (!CompleteChunk(target, slot, projection, completedChunk, - sourceRow + completedOffset, completedBytes)) { - return false; + for (uint32_t waveStart = 0U; + waveStart < totalContributors && !stop; + waveStart += static_cast(rankSize_)) { + const uint32_t bank = (sequence - 1U) % kReduceGradBankCount; + if (bankPublished[bank]) { + if (!WaitForHelpers( + lane, bank, helperCount, bankSequence[bank])) { + return; } + bankPublished[bank] = false; } - const uint64_t offset = chunk * chunkBytes; - const uint64_t bytes = MinU64(rowBytes_[projection] - offset, chunkBytes); - const bool ok = transports_[projection] == - kReduceGradTransportPeer ? - SendPeerChunk(target, slot, projection, chunk, sourceRow + offset, bytes) : - PostUdmaChunk(target, slot, projection, chunk, sourceRow + offset, bytes); - if (!ok) { - return false; + if (bankIssued[bank]) { + PublishStatus(kReduceGradDeviceInvalidState); + FailLane(lane); + return; } - } - const uint64_t drainBegin = chunks > 2 ? chunks - 2 : 0; - for (uint64_t chunk = drainBegin; chunk < chunks; ++chunk) { - const uint64_t offset = chunk * chunkBytes; - const uint64_t bytes = MinU64(rowBytes_[projection] - offset, chunkBytes); - if (!CompleteChunk(target, slot, projection, chunk, - sourceRow + offset, bytes)) { - return false; + const uint32_t contributorCount = static_cast(MinU64( + totalContributors - waveStart, + static_cast(rankSize_))); + const bool issued = IssueWave(lane, projection, bank, + globalExpert, chunk, waveStart, contributorCount, + sequence, pending[bank]); + if (!issued) { + FailLane(lane); + return; + } + bankIssued[bank] = true; + + if (sequence >= kReduceGradBankCount) { + const uint32_t readyBank = + (nextReadySequence - 1U) % kReduceGradBankCount; + if (!bankIssued[readyBank] || + pending[readyBank].sequence != nextReadySequence) { + PublishStatus(kReduceGradDeviceInvalidState); + FailLane(lane); + return; + } + const ReduceGradPendingWave &ready = pending[readyBank]; + if (!CompleteWave(lane, pending[readyBank])) { + FailLane(lane); + return; + } + PublishItem(lane, readyBank, LaneToken(ready.sequence), + ready.chunkIndex, ready.itemChunkBytes, + kReduceGradBankWork, ready.projection, + ready.localExpert, ready.waveStart, + ready.contributorCount, ready.remoteContributorCount); + bankIssued[readyBank] = false; + bankPublished[readyBank] = true; + bankSequence[readyBank] = ready.sequence; + ++nextReadySequence; } + ++sequence; } } } - return true; + + while (nextReadySequence < sequence) { + const uint32_t readyBank = + (nextReadySequence - 1U) % kReduceGradBankCount; + if (!bankIssued[readyBank] || + pending[readyBank].sequence != nextReadySequence) { + PublishStatus(kReduceGradDeviceInvalidState); + FailLane(lane); + return; + } + const ReduceGradPendingWave &ready = pending[readyBank]; + if (!CompleteWave(lane, pending[readyBank])) { + FailLane(lane); + return; + } + PublishItem(lane, readyBank, LaneToken(ready.sequence), + ready.chunkIndex, ready.itemChunkBytes, + kReduceGradBankWork, ready.projection, + ready.localExpert, ready.waveStart, + ready.contributorCount, ready.remoteContributorCount); + bankIssued[readyBank] = false; + bankPublished[readyBank] = true; + bankSequence[readyBank] = ready.sequence; + ++nextReadySequence; + } + for (uint32_t bank = 0U; bank < kReduceGradBankCount; ++bank) { + if (bankPublished[bank]) { + if (!WaitForHelpers(lane, bank, helperCount, bankSequence[bank])) { + return; + } + } + } + const uint32_t terminalBank = (sequence - 1U) % kReduceGradBankCount; + PublishItem(lane, terminalBank, LaneToken(sequence), 0U, 0U, + kReduceGradBankTerminal, 0U, 0U, 0U, 0U, 0U); + } + + __attribute__((always_inline)) inline __aicore__ uint32_t CollectWaveContributors( + uint32_t localExpert, uint32_t waveStart, uint32_t contributorCount, + int32_t contributors[]) const + { + const int32_t globalExpert = static_cast( + rank_ * expertsPerRank_ + localExpert); + uint32_t matching = 0U; + uint32_t collected = 0U; + const uint64_t entries = static_cast(rankSize_) * + static_cast(prefetchSlots_); + for (uint64_t entry = 0U; + entry < entries && collected < contributorCount; ++entry) { + if (expertsToCopy_[entry] != globalExpert) { + continue; + } + if (matching++ < waveStart) { + continue; + } + contributors[collected++] = static_cast(entry); + } + return collected; + } + + __attribute__((always_inline)) inline __aicore__ GM_ADDR ContributorAddress( + uint32_t projection, + int32_t contributor, GM_ADDR bankPayload, uint32_t &remoteOrdinal, + uint64_t chunkOffset, uint64_t tileOffset) const + { + const int32_t sourceRank = contributor / static_cast(prefetchSlots_); + const uint32_t slot = static_cast( + contributor % static_cast(prefetchSlots_)); + if (sourceRank == rank_) { + return sources_[projection] + static_cast(slot) * + rowBytes_[projection] + chunkOffset + tileOffset; + } + return bankPayload + static_cast(remoteOrdinal++) * + chunkBytes_ + tileOffset; + } + + __aicore__ inline void CopyGmToUb(AscendC::LocalTensor destination, + GM_ADDR source, uint32_t bytes) + { + AscendC::GlobalTensor sourceGlobal; + sourceGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(source)); + const AscendC::DataCopyExtParams params {1U, bytes, 0U, 0U, 0U}; + const AscendC::DataCopyPadExtParams pad {false, 0U, 0U, 0U}; + AscendC::DataCopyPad(destination, sourceGlobal, params, pad); + } + + __aicore__ inline void CopyUbToGm(GM_ADDR destination, + AscendC::LocalTensor source, uint32_t bytes) + { + AscendC::GlobalTensor destinationGlobal; + destinationGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(destination)); + const AscendC::DataCopyExtParams params {1U, bytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(destinationGlobal, source, params); + } + + __attribute__((always_inline)) inline __aicore__ void PrefetchAccumulator( + GM_ADDR output, uint64_t tileOffset, uint32_t tileBytes, + uint32_t accumIndex, bool storePending[]) + { + AscendC::LocalTensor accum = accumIndex == 0U ? + accumBuf_.Get() : accumPongBuf_.Get(); + const AscendC::TEventID event = accumIndex == 0U ? EVENT_ID2 : EVENT_ID3; + if (storePending[accumIndex]) { + AscendC::WaitFlag(event); + storePending[accumIndex] = false; + } + CopyGmToUb(accum, output + tileOffset, tileBytes); + AscendC::SetFlag(event); + } + + __attribute__((always_inline)) inline __aicore__ void AccumulateTile( + uint32_t projection, + GM_ADDR output, GM_ADDR bankPayload, const int32_t contributors[], + uint32_t contributorCount, uint64_t chunkOffset, uint64_t tileOffset, + uint32_t tileBytes, uint32_t accumIndex, bool storePending[], + bool hasNextTile, uint64_t nextTileOffset, uint32_t nextTileBytes) + { + AscendC::LocalTensor accum = accumIndex == 0U ? + accumBuf_.Get() : accumPongBuf_.Get(); + const AscendC::TEventID storeEvent = accumIndex == 0U ? + EVENT_ID2 : EVENT_ID3; + AscendC::LocalTensor inputs[2] = { + inputPingBuf_.Get(), inputPongBuf_.Get()}; + AscendC::WaitFlag(storeEvent); + + uint32_t remoteOrdinal = 0U; + bool reusePending[2] = {false, false}; + GM_ADDR firstSource = ContributorAddress(projection, contributors[0], + bankPayload, remoteOrdinal, chunkOffset, tileOffset); + CopyGmToUb(inputs[0], firstSource, tileBytes); + AscendC::SetFlag(EVENT_ID0); + + for (uint32_t contributor = 0U; + contributor < contributorCount; ++contributor) { + const uint32_t current = contributor & 1U; + const AscendC::TEventID currentEvent = current == 0U ? EVENT_ID0 : EVENT_ID1; + AscendC::WaitFlag(currentEvent); + + if (contributor + 1U < contributorCount) { + const uint32_t next = current ^ 1U; + const AscendC::TEventID nextEvent = next == 0U ? EVENT_ID0 : EVENT_ID1; + if (reusePending[next]) { + AscendC::WaitFlag(nextEvent); + reusePending[next] = false; + } + GM_ADDR nextSource = ContributorAddress(projection, + contributors[contributor + 1U], bankPayload, remoteOrdinal, + chunkOffset, tileOffset); + CopyGmToUb(inputs[next], nextSource, tileBytes); + AscendC::SetFlag(nextEvent); + } else if (hasNextTile) { + PrefetchAccumulator(output, nextTileOffset, nextTileBytes, + accumIndex ^ 1U, storePending); + } + + AscendC::Add(accum, accum, inputs[current], tileBytes / sizeof(float)); + AscendC::SetFlag(currentEvent); + reusePending[current] = true; + } + if (reusePending[0]) { + AscendC::WaitFlag(EVENT_ID0); + } + if (reusePending[1]) { + AscendC::WaitFlag(EVENT_ID1); + } + SyncEvent(storeEvent); + CopyUbToGm(output + tileOffset, accum, tileBytes); + AscendC::SetFlag(storeEvent); + storePending[accumIndex] = true; + } + + __attribute__((always_inline)) inline __aicore__ void ProcessBankItem( + uint32_t lane, uint32_t bank, + uint32_t projection, uint32_t localExpert, uint64_t chunkIndex, + uint64_t itemChunkBytes, uint32_t waveStart, uint32_t contributorCount, + uint32_t remoteContributorCount, uint32_t helperIndex, uint32_t helperCount) + { + if (projection >= kReduceGradProjectionCount || + localExpert >= static_cast(expertsPerRank_) || + contributorCount == 0U || + contributorCount > static_cast(rankSize_) || + chunkIndex >= chunkCounts_[projection] || + itemChunkBytes == 0U || itemChunkBytes > chunkBytes_) { + PublishStatus(kReduceGradDeviceInvalidState); + return; + } + int32_t contributors[TileXR::TILEXR_MAX_RANK_SIZE] = {}; + if (CollectWaveContributors(localExpert, waveStart, contributorCount, + contributors) != contributorCount) { + PublishStatus(kReduceGradDeviceInvalidState); + return; + } + uint32_t observedRemoteContributors = 0U; + for (uint32_t contributor = 0U; + contributor < contributorCount; ++contributor) { + const int32_t sourceRank = contributors[contributor] / + static_cast(prefetchSlots_); + if (sourceRank != rank_) { + ++observedRemoteContributors; + } + } + if (observedRemoteContributors != remoteContributorCount) { + PublishStatus(kReduceGradDeviceInvalidState); + return; + } + const uint64_t chunkOffset = chunkIndex * chunkBytes_; + GM_ADDR output = gradients_[projection] + + static_cast(rank_ * expertsPerRank_ + localExpert) * + rowBytes_[projection] + chunkOffset; + GM_ADDR bankPayload = BankPayload(lane, bank); + const uint64_t tiles = CeilDivU64(itemChunkBytes, + kReduceGradKernelTileBytes); + if (helperIndex >= tiles) { + return; + } + bool storePending[2] = {false, false}; + uint32_t accumIndex = 0U; + const uint64_t firstTileOffset = + static_cast(helperIndex) * kReduceGradKernelTileBytes; + const uint32_t firstTileBytes = static_cast(MinU64( + itemChunkBytes - firstTileOffset, kReduceGradKernelTileBytes)); + PrefetchAccumulator(output, firstTileOffset, firstTileBytes, + accumIndex, storePending); + for (uint64_t tile = helperIndex; tile < tiles; tile += helperCount) { + const uint64_t tileOffset = tile * kReduceGradKernelTileBytes; + const uint32_t tileBytes = static_cast(MinU64( + itemChunkBytes - tileOffset, kReduceGradKernelTileBytes)); + const uint64_t nextTile = tile + helperCount; + const bool hasNextTile = nextTile < tiles; + const uint64_t nextTileOffset = nextTile * kReduceGradKernelTileBytes; + const uint32_t nextTileBytes = hasNextTile ? + static_cast(MinU64( + itemChunkBytes - nextTileOffset, kReduceGradKernelTileBytes)) : 0U; + AccumulateTile(projection, output, bankPayload, contributors, + contributorCount, chunkOffset, tileOffset, tileBytes, + accumIndex, storePending, hasNextTile, + nextTileOffset, nextTileBytes); + accumIndex ^= 1U; + } + for (uint32_t index = 0U; index < 2U; ++index) { + if (storePending[index]) { + const AscendC::TEventID storeEvent = index == 0U ? + EVENT_ID2 : EVENT_ID3; + AscendC::WaitFlag(storeEvent); + } + } + } + + __aicore__ inline bool WaitForBank(uint32_t lane, uint32_t bank, + uint32_t sequence) + { + const uint64_t expected = LaneToken(sequence); + uint64_t attempts = 0U; + while (true) { + const uint64_t observed = LoadLaneToken(BankReady(lane, bank)); + if (observed == expected) { + return true; + } + if (LaneAborted(lane)) { + return false; + } + if (++attempts >= waitIterations_) { + PublishStatus(kReduceGradDeviceHelperTimeout); + AbortLane(lane); + return false; + } + } } - __aicore__ inline void RunSender() + __aicore__ inline void RunLaneHelper() { - for (int64_t controlIndex = blockIdx_; controlIndex < rankSize_ - 1; - controlIndex += controlBlockCount_) { - if (!RunSenderTo(ControlPeer(controlIndex))) { + uint32_t lane = 0U; + uint32_t helperIndex = 0U; + uint32_t helperCount = 0U; + HelperLane(lane, helperIndex, helperCount); + uint32_t sequence = 1U; + while (true) { + const uint32_t bank = (sequence - 1U) % kReduceGradBankCount; + if (!WaitForBank(lane, bank, sequence)) { return; } + uint64_t token = 0U; + uint64_t chunkIndex = 0U; + uint64_t itemChunkBytes = 0U; + uint32_t kind = 0U; + uint32_t projection = 0U; + uint32_t localExpert = 0U; + uint32_t waveStart = 0U; + uint32_t contributorCount = 0U; + uint32_t remoteContributorCount = 0U; + LoadItem(lane, bank, token, chunkIndex, itemChunkBytes, kind, + projection, localExpert, waveStart, contributorCount, + remoteContributorCount); + if (token != LaneToken(sequence)) { + PublishStatus(kReduceGradDeviceInvalidState); + } else if (kind == kReduceGradBankTerminal) { + StoreLaneToken(BankDoneToken( + lane, bank, helperIndex), LaneToken(sequence)); + break; + } else if (kind == kReduceGradBankWork) { + ProcessBankItem(lane, bank, projection, localExpert, chunkIndex, + itemChunkBytes, waveStart, contributorCount, + remoteContributorCount, helperIndex, helperCount); + } else { + PublishStatus(kReduceGradDeviceInvalidState); + } + StoreLaneToken(BankDoneToken( + lane, bank, helperIndex), LaneToken(sequence)); + ++sequence; } } - __aicore__ inline void RunReceiver() + __aicore__ inline void CrossRankBarrier() { - const int64_t receiverIndex = blockIdx_ - controlBlockCount_; - for (uint32_t projection = 0; projection < kReduceGradProjectionCount; ++projection) { - const uint64_t chunkBytes = ChunkBytes(projection); - const uint64_t chunks = ChunkCount(projection); - for (int64_t source = 0; source < rankSize_; ++source) { - for (int64_t slot = 0; slot < prefetchSlots_; ++slot) { - const int32_t expert = ExpertForSlot(source, slot); - if (expert < 0) { - continue; - } - if (expert >= expertCount_) { - SetStatus(kReduceGradDeviceInvalidState); - return; - } - if (OwnerForExpert(expert) != rank_) { - continue; - } - GM_ADDR outputRow = gradients_[projection] + - static_cast(expert) * rowBytes_[projection]; - const uint64_t localExpert = static_cast( - expert - rank_ * expertsPerRank_); - for (uint64_t chunk = 0; chunk < chunks; ++chunk) { - // A slot's two stage records are reused across every projection and - // chunk, so one receiver must consume them in sender order. - const uint64_t workIndex = localExpert; - if (static_cast( - workIndex % static_cast(receiverCount_)) != - receiverIndex) { - continue; - } - const uint64_t offset = chunk * chunkBytes; - const uint64_t bytes = MinU64(rowBytes_[projection] - offset, chunkBytes); - if (source == rank_) { - GM_ADDR sourceRow = gradients_[projection] + - static_cast(expertCount_ + slot) * rowBytes_[projection]; - AddDense(outputRow + offset, sourceRow + offset, bytes); - ZeroDense(sourceRow + offset, bytes); - } else if (transports_[projection] == - kReduceGradTransportPeer) { - if (!AddPeerPacked(outputRow + offset, - PeerRecord(rank_, source, slot, chunk), bytes)) { - return; - } - } else { - const uint64_t stage = chunk & 1U; - const uint64_t sequence = Sequence(projection, slot, chunk); - if (!WaitUdmaReady(source, stage, sequence)) { - return; - } - GM_ADDR inbound = workspace_ + UDMAInboundOffset(source, stage); - TileXR::UDMACleanCacheLines( - reinterpret_cast<__gm__ uint8_t *>(inbound), bytes); - AscendC::PipeBarrier(); - AddDense(outputRow + offset, inbound, bytes); - StoreSignal(workspace_ + UDMACompletionOffset(source, stage), sequence); - } - } + GM_ADDR shareAddrs[TileXR::TILEXR_MAX_RANK_SIZE] = {}; + AscendC::GlobalTensor peerMemTable; + peerMemTable.SetGlobalBuffer(&(args_->peerMems[0]), + TileXR::TILEXR_MAX_RANK_SIZE); + for (int32_t peer = 0; peer < rankSize_; ++peer) { + shareAddrs[peer] = peerMemTable.GetValue(peer); + } + GM_ADDR localShareAddr = shareAddrs[rank_]; + for (int32_t peer = 0; peer < rankSize_; ++peer) { + if (shareAddrs[peer] == nullptr) { + PublishStatus(kReduceGradDeviceInvalidState); + shareAddrs[peer] = localShareAddr; + } + } + SyncCollectives sync; + sync.Init(static_cast(rank_), static_cast(rankSize_), + shareAddrs, syncBuf_); + const int32_t localStep = DeviceStatus() == 0U ? + kReduceGradBarrierStep : kReduceGradBarrierFailureStep; + sync.SetInnerFlag(static_cast(magic_), localStep); + for (int32_t peer = 0; peer < rankSize_; ++peer) { + sync.WaitInnerFlag(static_cast(magic_), + kReduceGradBarrierStep, peer, 0); + } + const int64_t successValue = + (static_cast(static_cast(magic_)) << MAGIC_OFFSET) | + static_cast(kReduceGradBarrierStep); + for (int32_t peer = 0; peer < rankSize_; ++peer) { + if (sync.GetInnerFlag(peer, 0) != successValue) { + PublishStatus(kReduceGradDeviceInvalidState); + } + } + } + + __aicore__ inline void ClearLocalSources() + { + AscendC::LocalTensor zeros = accumBuf_.Get(); + AscendC::Duplicate(zeros, 0.0F, + static_cast(kReduceGradKernelTileBytes / sizeof(float))); + SyncEvent(EVENT_ID2); + bool copyPending = false; + uint64_t rowTileBase = 0U; + for (uint32_t projection = 0U; + projection < kReduceGradProjectionCount; ++projection) { + const uint64_t tilesPerRow = CeilDivU64( + rowBytes_[projection], kReduceGradKernelTileBytes); + for (int64_t slot = 0; slot < prefetchSlots_; ++slot) { + const uint64_t rowBlock = rowTileBase % blockCount_; + const uint64_t firstTile = + (static_cast(blockIdx_) + blockCount_ - rowBlock) % + blockCount_; + rowTileBase += tilesPerRow; + if (expertsToCopy_[rank_ * prefetchSlots_ + slot] < 0) { + continue; + } + GM_ADDR destination = sources_[projection] + + static_cast(slot) * rowBytes_[projection]; + for (uint64_t tile = firstTile; tile < tilesPerRow; + tile += static_cast(blockCount_)) { + const uint64_t offset = tile * kReduceGradKernelTileBytes; + const uint32_t bytes = static_cast(MinU64( + rowBytes_[projection] - offset, + kReduceGradKernelTileBytes)); + CopyUbToGm(destination + offset, zeros, bytes); + copyPending = true; } } } + if (copyPending) { + SyncEvent(EVENT_ID2); + } } __gm__ TileXR::CommArgs *args_{nullptr}; + __gm__ TileXR::UDMAInfo *profileInfo_{nullptr}; + __gm__ TileXR::TileXRUDMAProfileRegistry *profileRegistry_{nullptr}; __gm__ int32_t *expertsToCopy_{nullptr}; GM_ADDR gradients_[kReduceGradProjectionCount] = {}; + GM_ADDR sources_[kReduceGradProjectionCount] = {}; GM_ADDR workspace_{nullptr}; - __gm__ int32_t *status_{nullptr}; + __gm__ uint32_t *status_{nullptr}; int64_t rank_{0}; int64_t rankSize_{0}; int64_t expertCount_{0}; int64_t expertsPerRank_{0}; int64_t prefetchSlots_{0}; - int64_t controlBlockCount_{0}; - int64_t blockIdx_{0}; - int64_t blockCount_{0}; - int64_t receiverCount_{0}; uint64_t rowElements_[kReduceGradProjectionCount] = {}; uint64_t rowBytes_[kReduceGradProjectionCount] = {}; - uint32_t transports_[kReduceGradProjectionCount] = {}; - uint32_t udmaQpCount_{0}; - uint64_t peerRecordBaseOffset_{0}; - uint64_t peerHalfBytes_{0}; - uint64_t peerSlotStrideBytes_{0}; - uint64_t peerChunkPayloadBytes_{0}; - uint64_t udmaStateOffset_{0}; - uint64_t udmaOutboundOffset_{0}; - uint64_t udmaInboundOffset_{0}; - uint64_t udmaChunkBytes_{0}; + uint64_t chunkCounts_[kReduceGradProjectionCount] = {}; + uint32_t projectionQpBase_[kReduceGradProjectionCount] = {}; + uint32_t projectionQpCounts_[kReduceGradProjectionCount] = {}; + uint32_t lanePhysicalQps_[kReduceGradMaxUdmaQpCount] = {}; + uint32_t transportQpCount_{0}; + uint32_t qpCount_{0}; + uint32_t laneCount_{0}; + uint64_t laneStateBytes_{0}; + uint64_t stagingOffset_{0}; + uint64_t bankStrideBytes_{0}; + uint64_t laneStrideBytes_{0}; + uint64_t chunkBytes_{0}; uint64_t workspaceBytes_{0}; - uint64_t waitIterations_{0}; + uint64_t waitIterations_{1}; int64_t magic_{0}; + uint32_t blockIdx_{0}; + uint32_t blockCount_{0}; AscendC::TPipe pipe_; - AscendC::TBuf ioBuf_; AscendC::TBuf accumBuf_; - AscendC::TBuf flagBuf_; - AscendC::TBuf clearBuf_; - AscendC::TBuf udmaWqeBuf_; + AscendC::TBuf accumPongBuf_; + AscendC::TBuf inputPingBuf_; + AscendC::TBuf inputPongBuf_; + AscendC::TBuf syncBuf_; + AscendC::TBuf wqeBuf_; + AscendC::TBuf itemBuf_; }; } // namespace Kernel } // namespace TileXRMoonEp extern "C" __global__ __aicore__ void tilexr_moonep_reduce_grad_kernel( - GM_ADDR commArgs, GM_ADDR expertsToCopy, GM_ADDR gate, GM_ADDR up, GM_ADDR down, - GM_ADDR workspace, GM_ADDR status, int64_t rank, int64_t rankSize, - int64_t expertCount, int64_t expertsPerRank, int64_t prefetchSlots, - int64_t controlBlockCount, - uint64_t gateRowElements, uint64_t upRowElements, uint64_t downRowElements, - uint64_t gateRowBytes, uint64_t upRowBytes, uint64_t downRowBytes, - uint32_t gateTransport, uint32_t upTransport, uint32_t downTransport, - uint32_t udmaQpCount, uint64_t peerRecordBaseOffset, uint64_t peerHalfBytes, - uint64_t peerSlotStrideBytes, uint64_t peerChunkPayloadBytes, - uint64_t udmaStateOffset, uint64_t udmaOutboundOffset, - uint64_t udmaInboundOffset, uint64_t udmaChunkBytes, uint64_t workspaceBytes, + GM_ADDR commArgs, GM_ADDR profileInfo, GM_ADDR profileRegistry, + GM_ADDR expertsToCopy, GM_ADDR gateGradient, GM_ADDR upGradient, + GM_ADDR downGradient, GM_ADDR gateSource, GM_ADDR upSource, + GM_ADDR downSource, GM_ADDR workspace, GM_ADDR status, + int64_t rank, int64_t rankSize, int64_t expertCount, + int64_t expertsPerRank, int64_t prefetchSlots, + uint64_t gateRowElements, uint64_t upRowElements, + uint64_t downRowElements, uint64_t gateRowBytes, + uint64_t upRowBytes, uint64_t downRowBytes, + uint64_t gateChunkCount, uint64_t upChunkCount, + uint64_t downChunkCount, uint32_t gateQpBase, + uint32_t upQpBase, uint32_t downQpBase, + uint32_t gateQpCount, uint32_t upQpCount, + uint32_t downQpCount, uint32_t lane0PhysicalQp, + uint32_t lane1PhysicalQp, uint32_t lane2PhysicalQp, + uint32_t lane3PhysicalQp, uint32_t lane4PhysicalQp, + uint32_t lane5PhysicalQp, uint32_t lane6PhysicalQp, + uint32_t lane7PhysicalQp, uint32_t transportQpCount, + uint32_t qpCount, uint32_t laneCount, + uint64_t laneStateBytes, uint64_t stagingOffset, + uint64_t bankStrideBytes, uint64_t laneStrideBytes, + uint64_t chunkBytes, uint64_t workspaceBytes, uint64_t waitIterations, int64_t magic) { if constexpr (g_coreType == AscendC::AIV) { TileXRMoonEp::Kernel::ReduceGradKernel op; - op.Init(commArgs, expertsToCopy, gate, up, down, workspace, status, rank, - rankSize, expertCount, expertsPerRank, prefetchSlots, controlBlockCount, - gateRowElements, upRowElements, downRowElements, gateRowBytes, - upRowBytes, downRowBytes, gateTransport, upTransport, downTransport, - udmaQpCount, - peerRecordBaseOffset, peerHalfBytes, peerSlotStrideBytes, - peerChunkPayloadBytes, udmaStateOffset, udmaOutboundOffset, - udmaInboundOffset, udmaChunkBytes, workspaceBytes, waitIterations, magic); + op.Init(commArgs, profileInfo, profileRegistry, expertsToCopy, + gateGradient, upGradient, downGradient, gateSource, upSource, + downSource, workspace, status, rank, rankSize, expertCount, + expertsPerRank, prefetchSlots, gateRowElements, upRowElements, + downRowElements, gateRowBytes, upRowBytes, downRowBytes, + gateChunkCount, upChunkCount, downChunkCount, gateQpBase, + upQpBase, downQpBase, gateQpCount, upQpCount, downQpCount, + lane0PhysicalQp, lane1PhysicalQp, lane2PhysicalQp, + lane3PhysicalQp, lane4PhysicalQp, lane5PhysicalQp, + lane6PhysicalQp, lane7PhysicalQp, transportQpCount, + qpCount, laneCount, laneStateBytes, stagingOffset, + bankStrideBytes, laneStrideBytes, chunkBytes, workspaceBytes, + waitIterations, magic); op.Process(); } } diff --git a/tests/moonep/CMakeLists.txt b/tests/moonep/CMakeLists.txt index 795263c1..4c5bdc8e 100644 --- a/tests/moonep/CMakeLists.txt +++ b/tests/moonep/CMakeLists.txt @@ -83,12 +83,15 @@ file(WRITE "${TILEXR_MOONEP_FAKE_INCLUDE_DIR}/tilexr_api.h" "#include \"tilexr_udma_reg.h\"\n" "extern \"C\" {\n" "typedef void *TileXRCommPtr;\n" +"typedef uint32_t TileXRUDMAProfileHandle;\n" "int TileXRGetCommArgsHost(TileXRCommPtr comm, TileXR::CommArgs *&commArgs);\n" "int TileXRGetCommArgsDev(TileXRCommPtr comm, GM_ADDR &commArgs);\n" "int TileXRGetUDMARegistryHost(TileXRCommPtr comm, const TileXR::TileXRUDMARegistry **registry);\n" "int TileXRUDMAGetQpCount(TileXRCommPtr comm, uint32_t *qpCount);\n" +"int TileXRUDMAProfileRegister(TileXRCommPtr comm, const TileXR::TileXRUDMAProfileDesc *desc, TileXRUDMAProfileHandle *handle);\n" +"int TileXRUDMAProfileUnregister(TileXRCommPtr comm, TileXRUDMAProfileHandle handle);\n" +"int TileXRUDMAProfileQuery(TileXRCommPtr comm, TileXRUDMAProfileHandle handle, TileXR::TileXRUDMAProfileView *view);\n" "int TileXRCommNextMagic(TileXRCommPtr comm, int64_t *magic);\n" -"int TileXRGetUDMARegistryHost(TileXRCommPtr comm, const TileXR::TileXRUDMARegistry **registry);\n" "}\n" "#endif\n" "#endif\n") diff --git a/tests/moonep/demo/run_a5.sh b/tests/moonep/demo/run_a5.sh index 13be9f67..cb3c089f 100644 --- a/tests/moonep/demo/run_a5.sh +++ b/tests/moonep/demo/run_a5.sh @@ -2,7 +2,7 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -export TILEXR_UDMA_QP_ROUTE_SPEC=port_count:6,port_count:2 +export TILEXR_UDMA_QP_ROUTE_SPEC=port_count:6,port_count:6,port_count:2 if [[ -x "${SCRIPT_DIR}/tilexr_moonep_flow_demo" ]]; then TILEXR_INSTALL_PREFIX="$(cd "${SCRIPT_DIR}/.." && pwd)" else @@ -17,7 +17,7 @@ experts="${4:-$((rank_size * 2))}" hidden="${5:-32}" physical_device_count="${6:-$((rank_size < 8 ? rank_size : 8))}" -if ((rank_size <= 0 || rank_size > 128 || physical_device_count <= 0 || +if ((rank_size < 4 || rank_size > 128 || physical_device_count <= 0 || s <= 0 || k < 2 || experts <= 0 || hidden <= 0 || hidden % 32 != 0 || experts % rank_size != 0)); then echo "invalid rank/dimension configuration" >&2 diff --git a/tests/moonep/demo/tilexr_moonep_flow_demo.cpp b/tests/moonep/demo/tilexr_moonep_flow_demo.cpp index b46ada74..9eedf9ec 100644 --- a/tests/moonep/demo/tilexr_moonep_flow_demo.cpp +++ b/tests/moonep/demo/tilexr_moonep_flow_demo.cpp @@ -32,6 +32,7 @@ constexpr int32_t kExpectedCombineStatusSuccess = 3000; constexpr int32_t kExpectedPrefetchStatusSuccess = 4000; constexpr uint64_t kPrefetchWeightAlignment = 64; constexpr int32_t kExpectedReduceGradStatusSuccess = 0; +constexpr int kMinReduceGradRankCount = 4; struct Options { int world = 8; @@ -59,6 +60,7 @@ struct RuntimeResources { TileXRCommPtr comm = nullptr; bool udmaRegistered = false; TileXRUDMAMemHandle udmaHandle = 0; + TileXRMoonEpReduceGradPreparedV2 reduceGradPrepared = nullptr; std::vector allocations; }; @@ -239,7 +241,7 @@ bool ValidateOptions(const Options &options, uint64_t *routeCount, { if (routeCount == nullptr || tokenHiddenCount == nullptr || routeHiddenCount == nullptr || expertHiddenCount == nullptr || - options.world <= 0 || options.world > 128 || options.rank < 0 || + options.world < kMinReduceGradRankCount || options.world > 128 || options.rank < 0 || options.rank >= options.world || options.physicalDeviceCount <= 0 || options.device < 0 || options.device >= options.physicalDeviceCount || options.s <= 0 || options.k < 2 || options.experts <= 0 || @@ -551,6 +553,26 @@ bool Allocate(RuntimeResources *resources, uint64_t bytes, const std::string &na return true; } +bool AllocateAligned(RuntimeResources *resources, uint64_t bytes, uint64_t alignment, + const std::string &name, DeviceBuffer *buffer) +{ + uint64_t allocationBytes = 0; + if (resources == nullptr || buffer == nullptr || bytes == 0 || alignment == 0 || + !CheckedAdd(bytes, alignment - 1, &allocationBytes)) { + return false; + } + DeviceBuffer allocation; + if (!Allocate(resources, allocationBytes, name, &allocation)) { + return false; + } + const uintptr_t address = reinterpret_cast(allocation.data); + const uint64_t remainder = static_cast(address % alignment); + const uint64_t offset = remainder == 0 ? 0 : alignment - remainder; + buffer->data = static_cast(allocation.data) + offset; + buffer->bytes = bytes; + return true; +} + bool AllocateInt32(RuntimeResources *resources, uint64_t elements, const std::string &name, DeviceBuffer *buffer) { @@ -566,6 +588,16 @@ bool AllocateTyped(RuntimeResources *resources, uint64_t elements, return CountTypedBytes(elements, &bytes) && Allocate(resources, bytes, name, buffer); } +template +bool AllocateTypedAligned(RuntimeResources *resources, uint64_t elements, + uint64_t alignment, const std::string &name, + DeviceBuffer *buffer) +{ + uint64_t bytes = 0; + return CountTypedBytes(elements, &bytes) && + AllocateAligned(resources, bytes, alignment, name, buffer); +} + bool CopyHostToDevice(int rank, const DeviceBuffer &buffer, const std::vector &values, const std::string &name) { @@ -632,6 +664,16 @@ bool Cleanup(RuntimeResources *resources) ok = CheckAcl(resources->rank, "cleanup stream synchronize", aclrtSynchronizeStream(resources->stream)) && ok; } + if (resources->reduceGradPrepared != nullptr) { + const int destroyRet = TileXRMoonEpReduceGradDestroyPreparedV2( + resources->reduceGradPrepared); + ok = CheckTileXR(resources->rank, + "TileXRMoonEpReduceGradDestroyPreparedV2", destroyRet) && ok; + if (destroyRet != TILEXR_MOONEP_SUCCESS) { + return false; + } + resources->reduceGradPrepared = nullptr; + } if (resources->udmaRegistered && resources->comm != nullptr) { const int unregisterRet = TileXRUDMAUnregister(resources->comm, resources->udmaHandle); @@ -1077,6 +1119,7 @@ bool RunFlow(const Options &options, RuntimeResources *resources, DeviceBuffer upGradDev; DeviceBuffer downGradDev; DeviceBuffer reduceGradStatus; + DeviceBuffer reduceGradWorkspace; if (!AllocateInt32(resources, routeCount, "topk", &topk) || !AllocateInt32(resources, options.experts, "tokens per expert", &tpe) || @@ -1108,9 +1151,15 @@ bool RunFlow(const Options &options, RuntimeResources *resources, "backward dispatch", &backwardDispatchDev) || !AllocateTyped(resources, tokenHiddenCount, "backward combine", &backwardCombineDev) || - !AllocateTyped(resources, gateFullElements, "gate grad", &gateGradDev) || - !AllocateTyped(resources, upFullElements, "up grad", &upGradDev) || - !AllocateTyped(resources, downFullElements, "down grad", &downGradDev) || + !AllocateTypedAligned(resources, gateFullElements, + TILEXR_MOONEP_REDUCE_GRAD_WORKSPACE_ALIGNMENT, + "gate grad", &gateGradDev) || + !AllocateTypedAligned(resources, upFullElements, + TILEXR_MOONEP_REDUCE_GRAD_WORKSPACE_ALIGNMENT, + "up grad", &upGradDev) || + !AllocateTypedAligned(resources, downFullElements, + TILEXR_MOONEP_REDUCE_GRAD_WORKSPACE_ALIGNMENT, + "down grad", &downGradDev) || !AllocateInt32(resources, 1, "ReduceGrad status", &reduceGradStatus)) { return false; } @@ -1296,7 +1345,6 @@ bool RunFlow(const Options &options, RuntimeResources *resources, TileXRMoonEpReduceGradArgsV2 reduceGrad {}; reduceGrad.structSize = sizeof(reduceGrad); reduceGrad.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; - reduceGrad.comm = resources->comm; reduceGrad.plan = &plan; reduceGrad.gate = &gateGradTensor; reduceGrad.up = &upGradTensor; @@ -1322,18 +1370,63 @@ bool RunFlow(const Options &options, RuntimeResources *resources, if (!CheckTileXR(rank, "TileXRMoonEpReduceGradGetWorkspaceSizeV2", TileXRMoonEpReduceGradGetWorkspaceSizeV2( &reduceGradQuery, &reduceGradInfo)) || - reduceGradInfo.workspaceBytes != 0 || - reduceGradInfo.transports[0] != TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER || - reduceGradInfo.transports[1] != TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER || - reduceGradInfo.transports[2] != TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER) { + reduceGradInfo.workspaceBytes == 0 || reduceGradInfo.workspaceAlignment == 0 || + reduceGradInfo.udmaChunkBytes == 0 || reduceGradInfo.qpCount < 3 || + !AllocateAligned(resources, reduceGradInfo.workspaceBytes, + reduceGradInfo.workspaceAlignment, "ReduceGrad workspace", + &reduceGradWorkspace)) { std::cerr << "[rank " << rank - << "] flow demo requires the peer-only ReduceGrad V2 path" + << "] ReduceGrad owner-pull workspace preparation failed" << " workspace_bytes=" << reduceGradInfo.workspaceBytes - << " transports=" << reduceGradInfo.transports[0] << "," - << reduceGradInfo.transports[1] << "," - << reduceGradInfo.transports[2] << std::endl; + << " alignment=" << reduceGradInfo.workspaceAlignment + << " chunk_bytes=" << reduceGradInfo.udmaChunkBytes + << " qp_count=" << reduceGradInfo.qpCount << std::endl; + return false; + } + DeviceBuffer *gradBuffers[3] = {&gateGradDev, &upGradDev, &downGradDev}; + TileXRMoonEpReduceGradSourceSliceV2 reduceGradSources[3] {}; + for (uint32_t projection = 0; projection < 3; ++projection) { + uint64_t sourceOffset = 0; + uint64_t sourceBytes = 0; + uint64_t sourceEnd = 0; + if (!CheckedMultiply(static_cast(options.experts), + reduceGradInfo.rowBytes[projection], &sourceOffset) || + !CheckedMultiply(static_cast(b), + reduceGradInfo.rowBytes[projection], &sourceBytes) || + !CheckedAdd(sourceOffset, sourceBytes, &sourceEnd) || + sourceEnd > gradBuffers[projection]->bytes) { + return false; + } + reduceGradSources[projection].data = + static_cast(gradBuffers[projection]->data) + sourceOffset; + reduceGradSources[projection].bytes = sourceBytes; + reduceGradSources[projection].registrationBase = + gradBuffers[projection]->data; + reduceGradSources[projection].registrationBytes = + gradBuffers[projection]->bytes; + reduceGrad.sources[projection] = reduceGradSources[projection]; + } + TileXRMoonEpReduceGradPrepareArgsV2 reduceGradPrepare {}; + reduceGradPrepare.structSize = sizeof(reduceGradPrepare); + reduceGradPrepare.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; + reduceGradPrepare.comm = resources->comm; + reduceGradPrepare.plan = &plan; + reduceGradPrepare.gate = &gateGradTensor; + reduceGradPrepare.up = &upGradTensor; + reduceGradPrepare.down = &downGradTensor; + reduceGradPrepare.workspace = reduceGradWorkspace.data; + reduceGradPrepare.workspaceBytes = reduceGradWorkspace.bytes; + reduceGradPrepare.requestedUdmaChunkBytes = reduceGradQuery.requestedUdmaChunkBytes; + reduceGradPrepare.flags = TILEXR_MOONEP_FLAG_NONE; + for (uint32_t projection = 0; projection < 3; ++projection) { + reduceGradPrepare.sources[projection] = reduceGradSources[projection]; + } + if (!CheckTileXR(rank, "TileXRMoonEpReduceGradPrepareV2", + TileXRMoonEpReduceGradPrepareV2( + &reduceGradPrepare, &resources->reduceGradPrepared))) { return false; } + reduceGrad.prepared = resources->reduceGradPrepared; if (!CheckTileXR(rank, "Dispatch forward", TileXRMoonEpDispatchV1(&forwardDispatch, resources->stream)) || !CheckStageCompletion(rank, "dispatch_forward", kExpectedDispatchStatusSuccess, @@ -1526,11 +1619,11 @@ bool RunFlow(const Options &options, RuntimeResources *resources, << " prefetch_weight_status=" << kExpectedPrefetchStatusSuccess << " combine_status=" << kExpectedCombineStatusSuccess << " reduce_grad_status=" << kExpectedReduceGradStatusSuccess - << " reduce_grad_transport=peer" + << " reduce_grad_transport=owner_pull_udma" << " cu_last=" << cuHost.back() << " planning=native" << " dispatch=native prefetch_weight=native_udma_registered" - << " combine=native reduce_grad=native" + << " combine=native reduce_grad=native_owner_pull_udma" << " torch_validated=false" << " transport_performance_valid=false" << std::endl; diff --git a/tests/moonep/python/fakes.py b/tests/moonep/python/fakes.py index ad50e843..01314a3d 100644 --- a/tests/moonep/python/fakes.py +++ b/tests/moonep/python/fakes.py @@ -271,6 +271,8 @@ def __init__( stub_mask=0, ) self.registered_workspace = None + self.reduce_grad_prepared = None + self._next_reduce_grad_prepared = 1 self.reduce_grad_query_calls = 0 self._reduce_grad_lock = threading.RLock() self._reduce_grad_owner_token = None @@ -432,49 +434,69 @@ def reduce_grad_workspace_info( int(getattr(gradients, name).numel() // getattr(gradients, name).shape[0]) * 4 for name in ("gate", "up", "down") ) - transports = tuple("udma" if value > (1 << 20) else "peer" for value in row_bytes) - workspace_bytes = 4096 if "udma" in transports else 0 + chunk_bytes = int(requested_udma_chunk_bytes) or (8 << 20) + qp_count = max(3, self._udma_qp_count) + qp_counts = (1, 1, qp_count - 2) + lane_state_bytes = qp_count * 4096 + bank_stride_bytes = self.world_size * chunk_bytes + lane_stride_bytes = 2 * bank_stride_bytes + payload_bytes = qp_count * lane_stride_bytes + workspace_bytes = ( + (lane_state_bytes + payload_bytes + (2 << 20) - 1) // (2 << 20) + ) * (2 << 20) self.reduce_grad_query_calls += 1 return ReduceGradWorkspaceInfo( workspace_bytes=workspace_bytes, workspace_alignment=2 << 20, - udma_chunk_bytes=2 << 20 if workspace_bytes else 0, - peer_window_bytes=100 << 20, - peer_half_bytes=49 << 20, - peer_slot_stride_bytes=1 << 20, + udma_chunk_bytes=chunk_bytes, + lane_state_bytes=lane_state_bytes, + lane_state_stride_bytes=4096, + bank_stride_bytes=bank_stride_bytes, + lane_stride_bytes=lane_stride_bytes, row_bytes=row_bytes, - transports=transports, + chunk_counts=tuple( + (value + chunk_bytes - 1) // chunk_bytes for value in row_bytes + ), + projection_qp_counts=qp_counts, + qp_count=qp_count, block_dim=64, ) - def register_reduce_grad_workspace( - self, workspace, required_bytes, *, owner_token=None + def prepare_reduce_grad( + self, + context, + plan, + gradients, + sources, + source_registrations, + workspace, + *, + requested_udma_chunk_bytes=0, ): - with self._reduce_grad_lock: - self._require_reduce_grad_workspace_owner( - owner_token, "register_reduce_grad_workspace" - ) - if ( - self._active_udma_owner == "reduce_grad" - and self.registered_workspace is workspace - ): - return - self.registered_workspace = workspace - self._active_udma_owner = "reduce_grad" - self._active_projection = None - self.calls.append(("register_reduce_grad_workspace", required_bytes)) - - def unregister_reduce_grad_workspace(self, *, owner_token=None): - with self._reduce_grad_lock: - self._require_reduce_grad_workspace_owner( - owner_token, "unregister_reduce_grad_workspace" + prepared = self._next_reduce_grad_prepared + self._next_reduce_grad_prepared += 1 + self.registered_workspace = workspace + self.reduce_grad_prepared = prepared + self.calls.append( + ( + "prepare_reduce_grad", + prepared, + plan, + gradients, + sources, + source_registrations, + workspace, + requested_udma_chunk_bytes, ) - if ( - self.registered_workspace is not None - and self._active_udma_owner == "reduce_grad" - ): - self.calls.append(("unregister_reduce_grad_workspace", None)) - self._active_udma_owner = None + ) + return prepared + + def destroy_reduce_grad(self, prepared): + if prepared is None: + return + self.calls.append(("destroy_reduce_grad", prepared)) + if self.reduce_grad_prepared == prepared: + self.reduce_grad_prepared = None self.registered_workspace = None def reduce_grad( @@ -482,21 +504,22 @@ def reduce_grad( context, plan, gradients, - workspace, + sources, + source_registrations, + prepared, stream, wait_iterations, - *, - requested_udma_chunk_bytes=0, ): self.calls.append( ( "reduce_grad", plan, gradients, - workspace, + sources, + source_registrations, + prepared, stream, wait_iterations, - requested_udma_chunk_bytes, ) ) plan.reduce_grad_status._item = 0 diff --git a/tests/moonep/python/test_ffi_unittest.py b/tests/moonep/python/test_ffi_unittest.py index cd93687e..960a090b 100644 --- a/tests/moonep/python/test_ffi_unittest.py +++ b/tests/moonep/python/test_ffi_unittest.py @@ -26,6 +26,8 @@ TileXRMoonEPPrefetchWeightArgsV1, TileXRMoonEPReduceGradArgsV1, TileXRMoonEPReduceGradArgsV2, + TileXRMoonEPReduceGradPrepareArgsV2, + TileXRMoonEPReduceGradSourceSliceV2, TileXRMoonEPReduceGradWorkspaceInfoV2, TileXRMoonEPReduceGradWorkspaceQueryV2, TileXRMoonEPTensorV1, @@ -57,6 +59,8 @@ def __init__(self): self.stage_records = [] self.combine_v2_workspace_records = [] self.reduce_grad_records = [] + self.reduce_grad_prepare_records = [] + self.reduce_grad_destroy_records = [] self.register_calls = [] self.unregister_calls = [] self.comm = self._comm_library() @@ -305,18 +309,44 @@ def reduce_grad_query(query_ptr, info_ptr): info_ptr, ctypes.POINTER(TileXRMoonEPReduceGradWorkspaceInfoV2) ).contents descriptors = (query.gate.contents, query.up.contents, query.down.contents) - info.workspaceBytes = 0 + info.workspaceBytes = 2 << 20 info.workspaceAlignment = 2 << 20 - info.udmaChunkBytes = 0 - info.peerWindowBytes = 100 << 20 - info.peerHalfBytes = 49 << 20 - info.peerSlotStrideBytes = 1 << 20 + info.udmaChunkBytes = 4 << 20 + info.laneStateBytes = 3 * 4096 + info.laneStateStrideBytes = 4096 + info.bankStrideBytes = 8 << 20 + info.laneStrideBytes = 16 << 20 + info.qpCount = 3 info.blockDim = 64 for index, tensor_value in enumerate(descriptors): info.rowBytes[index] = int( tensor_value.elementCount // tensor_value.shape[0] ) * 4 - info.transports[index] = 1 + info.chunkCounts[index] = 1 + info.projectionQpCounts[index] = 1 + return 0 + + def reduce_grad_prepare(args_ptr, prepared_ptr): + args = ctypes.cast( + args_ptr, ctypes.POINTER(TileXRMoonEPReduceGradPrepareArgsV2) + ).contents + self.reduce_grad_prepare_records.append({ + "workspace_bytes": args.workspaceBytes, + "requested_chunk_bytes": args.requestedUdmaChunkBytes, + "source_bytes": tuple(int(value.bytes) for value in args.sources), + "source_ptrs": tuple(int(value.data) for value in args.sources), + "registration_bytes": tuple( + int(value.registrationBytes) for value in args.sources + ), + "registration_ptrs": tuple( + int(value.registrationBase) for value in args.sources + ), + }) + ctypes.cast(prepared_ptr, ctypes.POINTER(ctypes.c_void_p)).contents.value = 0x5678 + return 0 + + def reduce_grad_destroy(prepared): + self.reduce_grad_destroy_records.append(int(prepared.value)) return 0 def reduce_grad(args_ptr, stream): @@ -327,7 +357,11 @@ def reduce_grad(args_ptr, stream): "stream": stream.value, "flags": args.flags, "wait_iterations": args.waitIterations, - "workspace_bytes": args.workspaceBytes, + "prepared": int(args.prepared), + "source_bytes": tuple(int(value.bytes) for value in args.sources), + "registration_bytes": tuple( + int(value.registrationBytes) for value in args.sources + ), "shapes": tuple( tuple(pointer.contents.shape[: pointer.contents.rank]) for pointer in (args.gate, args.up, args.down) @@ -339,6 +373,8 @@ def reduce_grad(args_ptr, stream): return 0 library.TileXRMoonEpReduceGradGetWorkspaceSizeV2 = FakeFunction(reduce_grad_query) + library.TileXRMoonEpReduceGradPrepareV2 = FakeFunction(reduce_grad_prepare) + library.TileXRMoonEpReduceGradDestroyPreparedV2 = FakeFunction(reduce_grad_destroy) library.TileXRMoonEpReduceGradV2 = FakeFunction(reduce_grad) return library @@ -398,8 +434,10 @@ def test_ctypes_layout_matches_tilexr_moonep_header(self): self.assertEqual(TileXRMoonEPCombineArgsV1.dstLocal.offset, 24) self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradArgsV1), 48) self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradWorkspaceQueryV2), 64) - self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradWorkspaceInfoV2), 96) - self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradArgsV2), 96) + self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradWorkspaceInfoV2), 136) + self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradSourceSliceV2), 32) + self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradPrepareArgsV2), 176) + self.assertEqual(ctypes.sizeof(TileXRMoonEPReduceGradArgsV2), 168) self.assertEqual(TileXRMoonEPTensorV1.shape.offset, 32) self.assertEqual(TileXRMoonEPPlanV1.dst.offset, 56) self.assertEqual(TileXRMoonEPPlanV1.status.offset, 112) @@ -409,6 +447,10 @@ def test_ctypes_layout_matches_tilexr_moonep_header(self): self.assertEqual(TileXRMoonEPDispatchArgsV1.registeredWorkspace.offset, 64) self.assertEqual(TileXRMoonEPDispatchArgsV1.registeredWorkspaceBytes.offset, 72) self.assertEqual(TileXRMoonEPDispatchArgsV2.registeredWorkspace.offset, 64) + self.assertEqual(TileXRMoonEPReduceGradWorkspaceInfoV2.rowBytes.offset, 64) + self.assertEqual(TileXRMoonEPReduceGradWorkspaceInfoV2.chunkCounts.offset, 88) + self.assertEqual(TileXRMoonEPReduceGradPrepareArgsV2.sources.offset, 48) + self.assertEqual(TileXRMoonEPReduceGradArgsV2.sources.offset, 48) def test_invalid_combine_version_fails_before_library_load(self): loader = FakeCDLLLoader() @@ -551,6 +593,7 @@ def test_fake_cdll_receives_v1_descriptors_and_combine_v2_pointers(self): hidden_sh, route_weights_sk, _ = buffer.combine( plan, hidden_nvsh, route_weights_nvs ) + plan.status._item = 4000 gradients = ProjectionBuffers( tensor((6, 2, 4), torch.float32), tensor((6, 4, 3), torch.float32), @@ -643,10 +686,30 @@ def test_fake_cdll_receives_v1_descriptors_and_combine_v2_pointers(self): "stream": 0xCAFE, "flags": 0, "wait_iterations": 1234, - "workspace_bytes": 0, + "prepared": 0x5678, + "source_bytes": (64, 96, 48), + "registration_bytes": (192, 288, 144), "shapes": ((6, 2, 4), (6, 4, 3), (6, 3, 2)), "status_shape": (1,), }]) + self.assertEqual(loader.reduce_grad_prepare_records, [{ + "workspace_bytes": 2 << 20, + "requested_chunk_bytes": 0, + "source_bytes": (64, 96, 48), + "source_ptrs": tuple( + tensor.data_ptr() + 4 * row_bytes + for tensor, row_bytes in zip( + (gradients.gate, gradients.up, gradients.down), + (32, 48, 24), + ) + ), + "registration_bytes": (192, 288, 144), + "registration_ptrs": tuple( + tensor.data_ptr() + for tensor in (gradients.gate, gradients.up, gradients.down) + ), + }]) + self.assertEqual(loader.reduce_grad_destroy_records, [0x5678]) self.assertEqual( [byte_count for _, byte_count in loader.register_calls], [2 * 1024 * 1024] * 5, diff --git a/tests/moonep/python/test_reduce_grad_benchmark.py b/tests/moonep/python/test_reduce_grad_benchmark.py new file mode 100644 index 00000000..d24cf81b --- /dev/null +++ b/tests/moonep/python/test_reduce_grad_benchmark.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import unittest + +from tools.moonep.reduce_grad_benchmark import ( + ReduceGradDimensions, + build_experts_to_copy, + cross_rank_statistics, + expected_expert_value, + gradient_source_regions, + plan_statistics, +) + + +class ReduceGradBenchmarkTests(unittest.TestCase): + def test_gradient_tail_uses_full_gradient_as_registration_backing(self): + class Gradient: + def __init__(self): + self.narrow_calls = [] + + def narrow(self, dimension, start, length): + source = object() + self.narrow_calls.append((dimension, start, length, source)) + return source + + gradients = [Gradient(), Gradient(), Gradient()] + sources, registrations = gradient_source_regions(gradients, 8, 4) + self.assertEqual(registrations, gradients) + self.assertEqual( + sources, [gradient.narrow_calls[0][3] for gradient in gradients] + ) + self.assertTrue( + all(gradient.narrow_calls[0][:3] == (0, 8, 4) for gradient in gradients) + ) + + def test_rank_count_below_four_is_rejected(self): + with self.assertRaisesRegex(ValueError, "at least 4 ranks"): + ReduceGradDimensions(3, 2, 2, 32, 32) + with self.assertRaisesRegex(ValueError, "at least 4 ranks"): + build_experts_to_copy(3, 2, 2, "full") + + def test_rank_by_slot_index_overflow_is_rejected(self): + with self.assertRaisesRegex(ValueError, "contributor indices"): + ReduceGradDimensions(8, 8, (1 << 31) // 8 + 1, 32, 32) + + def test_baseline_full_plan_preserves_source_rank_slot_order(self): + dimensions = ReduceGradDimensions(8, 8, 14, 3584, 3072) + plan = build_experts_to_copy( + dimensions.ranks, + dimensions.slots, + dimensions.experts_per_rank, + "full", + ) + self.assertEqual(plan[0], [-1] * 14) + self.assertEqual(plan[1], list(range(8)) + [-1] * 6) + self.assertEqual(plan[2], list(range(8)) + [-1] * 6) + self.assertEqual(plan[3], list(range(8)) + [-1] * 6) + self.assertEqual(plan[4], [-1] * 14) + stats = plan_statistics(plan) + self.assertEqual(stats["live_entries"], 24) + self.assertEqual(stats["active_experts"], 8) + self.assertEqual(stats["max_contributors_per_expert"], 3) + self.assertNotEqual(expected_expert_value(plan, 0, 0), 0.0) + + def test_balanced_full_uses_every_source_slot_and_owner(self): + plan = build_experts_to_copy(8, 14, 8, "balanced-full") + stats = plan_statistics(plan) + self.assertEqual(stats["live_entries"], 112) + self.assertEqual(stats["density"], 1.0) + self.assertTrue(all(0 <= expert < 64 for row in plan for expert in row)) + owners = {expert // 8 for row in plan for expert in row} + self.assertEqual(owners, set(range(8))) + + def test_cross_rank_max_is_aggregated_before_percentiles(self): + stats = cross_rank_statistics(((1.0, 7.0), (3.0, 4.0), (5.0, 2.0))) + self.assertEqual(stats["cross_rank_max_us"], [5.0, 7.0]) + self.assertEqual(stats["p50_us"], 6.0) + self.assertAlmostEqual(stats["p99_us"], 6.98) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/moonep/python/test_tilexr_correctness_adapter.py b/tests/moonep/python/test_tilexr_correctness_adapter.py index cc28b3b0..32ba87c4 100644 --- a/tests/moonep/python/test_tilexr_correctness_adapter.py +++ b/tests/moonep/python/test_tilexr_correctness_adapter.py @@ -187,10 +187,15 @@ def test_all_downstream_stages_map_in_place_and_optional_weights(self): self.assertIs(native_reduce.gate_reduce, reduce_buffers.gate) self.assertIs(native_reduce.down_reduce, reduce_buffers.down) for _, full_tensor in full_grads.items(): - self.assertEqual(len(full_tensor.copy_calls), 2) + self.assertEqual(full_tensor.copy_calls, []) for _, reduce_tensor in reduce_buffers.items(): self.assertEqual(reduce_tensor.masked_fill_calls, []) - self.assertEqual(len(reduce_tensor.zero_calls), 1) + self.assertEqual(reduce_tensor.zero_calls, []) + prepared = [ + call for call in runtime.calls if call[0] == "prepare_reduce_grad" + ][-1] + for source, (_, value) in zip(prepared[4], reduce_buffers.items()): + self.assertEqual(source.data_ptr(), value[dimensions().rank].data_ptr()) native_plan = backend._native_plans[id(plan)].native self.assertEqual(native_plan.reduce_grad_status.item(), 0) backend.synchronize() @@ -211,7 +216,7 @@ def test_prefetch_preserves_unused_projection_slots(self): self.assertEqual(projection.copy_calls, []) backend.close() - def test_reduce_grad_restores_legacy_tail_and_unused_slots(self): + def test_reduce_grad_uses_direct_sources_without_tail_copies(self): torch, _, _, backend = make_backend() plan = self._planned(backend, torch) backend.synchronize() @@ -229,10 +234,10 @@ def test_reduce_grad_restores_legacy_tail_and_unused_slots(self): backend.synchronize() for _, full_tensor in full_grads.items(): - self.assertEqual(len(full_tensor.copy_calls), 2) + self.assertEqual(full_tensor.copy_calls, []) for _, reduce_tensor in reduce_buffers.items(): - self.assertEqual(len(reduce_tensor.copy_calls), dimensions().prefetch_slots) - self.assertEqual(len(reduce_tensor.zero_calls), 1) + self.assertEqual(reduce_tensor.copy_calls, []) + self.assertEqual(reduce_tensor.zero_calls, []) backend.close() def test_unknown_or_cloned_plan_is_rejected(self): diff --git a/tests/moonep/python/test_upstream_compat.py b/tests/moonep/python/test_upstream_compat.py index 269f7c40..eca61fc3 100644 --- a/tests/moonep/python/test_upstream_compat.py +++ b/tests/moonep/python/test_upstream_compat.py @@ -548,13 +548,16 @@ def test_reduce_grad_maps_in_place_and_returns_optional_event(self): ) self.assertIsNone(result) for value in full: - self.assertGreaterEqual( - sum(call[2] == (2, 8, 8) for call in value.copy_calls), 2 - ) + self.assertEqual(value.copy_calls, []) for value in reduced: - self.assertEqual(len(value.zero_calls), 1) - self.assertEqual(len(value.copy_calls), 1) - self.assertEqual(len(value.masked_fill_calls), 1) + self.assertEqual(value.zero_calls, []) + self.assertEqual(value.copy_calls, []) + self.assertEqual(value.masked_fill_calls, []) + prepared = [ + call for call in runtime.calls if call[0] == "prepare_reduce_grad" + ][-1] + for source, value in zip(prepared[4], reduced): + self.assertEqual(source.data_ptr(), value[0].data_ptr()) event = buffer.reduce_grad( plan=plan, diff --git a/tests/moonep/python/unittest_smoke.py b/tests/moonep/python/unittest_smoke.py index a5a379fa..8c38747e 100644 --- a/tests/moonep/python/unittest_smoke.py +++ b/tests/moonep/python/unittest_smoke.py @@ -204,7 +204,9 @@ def expert_backward(dispatched, state): "dispatch", "expert_backward", "combine", + "prepare_reduce_grad", "reduce_grad", + "destroy_reduce_grad", "udma_unregister", "close", ], @@ -216,7 +218,7 @@ def expert_backward(dispatched, state): with self.assertRaises((AttributeError, TypeError)): forward.state.plan.epoch = 99 - def test_reduce_grad_udma_workspace_is_registered_once_and_released(self): + def test_reduce_grad_profile_is_prepared_once_and_released(self): torch, runtime, buffer = make_buffer() plan, _ = buffer.planning( tensor((4, 2), torch.int32), tensor((4,), torch.int32) @@ -229,7 +231,8 @@ def test_reduce_grad_udma_workspace_is_registered_once_and_released(self): ) info = buffer.prepare_reduce_grad(plan, gradients) - self.assertEqual(info.transports, ("peer", "udma", "peer")) + self.assertEqual(info.qp_count, 3) + self.assertEqual(info.projection_qp_counts, (1, 1, 1)) self.assertIsNotNone(runtime.registered_workspace) self.assertEqual(runtime.registered_workspace.numel(), info.workspace_bytes) self.assertEqual( @@ -237,7 +240,7 @@ def test_reduce_grad_udma_workspace_is_registered_once_and_released(self): ) self.assertEqual( [call[0] for call in runtime.calls].count( - "register_reduce_grad_workspace" + "prepare_reduce_grad" ), 1, ) @@ -251,7 +254,7 @@ def test_reduce_grad_udma_workspace_is_registered_once_and_released(self): self.assertEqual(runtime.reduce_grad_query_calls, 1) self.assertEqual( [call[0] for call in runtime.calls].count( - "register_reduce_grad_workspace" + "prepare_reduce_grad" ), 1, ) @@ -259,7 +262,7 @@ def test_reduce_grad_udma_workspace_is_registered_once_and_released(self): self.assertIsNone(runtime.registered_workspace) names = [call[0] for call in runtime.calls] self.assertLess( - names.index("unregister_reduce_grad_workspace"), names.index("close") + names.index("destroy_reduce_grad"), names.index("close") ) def test_reduce_grad_rejects_second_inflight_call_and_reuses_after_sync(self): @@ -321,6 +324,28 @@ def test_reduce_grad_device_failure_releases_runtime_owner(self): self.assertEqual([call[0] for call in runtime.calls].count("reduce_grad"), 2) buffer.close() + def test_reduce_grad_transport_failure_poisons_context(self): + torch, runtime, buffer = make_buffer() + plan, _ = buffer.planning( + tensor((4, 2), torch.int32), tensor((4,), torch.int32) + ) + gradients = ProjectionBuffers( + tensor((6, 8), torch.float32), + tensor((6, 8), torch.float32), + tensor((6, 8), torch.float32), + ) + buffer.reduce_grad( + plan, + full_gate_grad=gradients.gate, + full_up_grad=gradients.up, + full_down_grad=gradients.down, + ) + plan.reduce_grad_status._item = 3 + with self.assertRaisesRegex(RuntimeError, "actual 3"): + buffer.synchronize() + self.assertTrue(buffer.context._poisoned) + self.assertIsNone(runtime._reduce_grad_owner_token) + def test_strict_tensor_checks(self): torch, _, buffer = make_buffer() tpe = tensor((4,), torch.int32) diff --git a/tests/moonep/unit/test_tilexr_moonep_abi_layout.cpp b/tests/moonep/unit/test_tilexr_moonep_abi_layout.cpp index 0a9ba63c..39055967 100644 --- a/tests/moonep/unit/test_tilexr_moonep_abi_layout.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_abi_layout.cpp @@ -80,27 +80,45 @@ int main() "ReduceGrad workspace query must be standard layout"); static_assert(std::is_standard_layout::value, "ReduceGrad workspace info must be standard layout"); + static_assert(std::is_standard_layout::value, + "ReduceGrad source slice must be standard layout"); + static_assert(std::is_standard_layout::value, + "ReduceGrad prepare args must be standard layout"); static_assert(std::is_standard_layout::value, "ReduceGrad V2 args must be standard layout"); static_assert(sizeof(TileXRMoonEpReduceGradWorkspaceQueryV2) == 64, "Unexpected ReduceGrad workspace query size"); - static_assert(sizeof(TileXRMoonEpReduceGradWorkspaceInfoV2) == 96, + static_assert(sizeof(TileXRMoonEpReduceGradWorkspaceInfoV2) == 136, "Unexpected ReduceGrad workspace info size"); - static_assert(offsetof(TileXRMoonEpReduceGradWorkspaceInfoV2, rowBytes) == 56, + static_assert(offsetof(TileXRMoonEpReduceGradWorkspaceInfoV2, rowBytes) == 64, "Unexpected ReduceGrad rowBytes offset"); - static_assert(offsetof(TileXRMoonEpReduceGradWorkspaceInfoV2, transports) == 80, - "Unexpected ReduceGrad transports offset"); - static_assert(sizeof(TileXRMoonEpReduceGradArgsV2) == 96, + static_assert(offsetof(TileXRMoonEpReduceGradWorkspaceInfoV2, chunkCounts) == 88, + "Unexpected ReduceGrad chunkCounts offset"); + static_assert(offsetof(TileXRMoonEpReduceGradWorkspaceInfoV2, projectionQpCounts) == 112, + "Unexpected ReduceGrad projection QP offset"); + static_assert(sizeof(TileXRMoonEpReduceGradSourceSliceV2) == 32, + "Unexpected ReduceGrad source slice size"); + static_assert(sizeof(TileXRMoonEpReduceGradPrepareArgsV2) == 176, + "Unexpected ReduceGrad prepare args size"); + static_assert(offsetof(TileXRMoonEpReduceGradPrepareArgsV2, sources) == 48, + "Unexpected ReduceGrad prepare source offset"); + static_assert(offsetof(TileXRMoonEpReduceGradPrepareArgsV2, workspace) == 144, + "Unexpected ReduceGrad prepare workspace offset"); + static_assert(sizeof(TileXRMoonEpReduceGradArgsV2) == 168, "Unexpected ReduceGrad V2 args size"); - static_assert(offsetof(TileXRMoonEpReduceGradArgsV2, workspace) == 48, - "Unexpected ReduceGrad workspace offset"); - static_assert(offsetof(TileXRMoonEpReduceGradArgsV2, status) == 64, + static_assert(offsetof(TileXRMoonEpReduceGradArgsV2, prepared) == 8, + "Unexpected ReduceGrad prepared handle offset"); + static_assert(offsetof(TileXRMoonEpReduceGradArgsV2, sources) == 48, + "Unexpected ReduceGrad launch source offset"); + static_assert(offsetof(TileXRMoonEpReduceGradArgsV2, status) == 144, "Unexpected ReduceGrad status offset"); Check(TILEXR_MOONEP_ABI_VERSION_V1 == 1, "ABI version must be 1"); Check(TILEXR_MOONEP_ABI_VERSION_V2 == 2, "V2 ABI version must be 2"); - Check(TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES == UINT64_C(1048576), - "ReduceGrad UDMA threshold must be exactly 1 MiB"); + Check(TILEXR_MOONEP_REDUCE_GRAD_DEFAULT_CHUNK_BYTES == UINT64_C(8388608), + "ReduceGrad default owner-pull chunk must be 8 MiB"); + Check(TILEXR_MOONEP_REDUCE_GRAD_WORKSPACE_ALIGNMENT == UINT64_C(2097152), + "ReduceGrad workspace alignment must be 2 MiB"); Check((TILEXR_MOONEP_STAGE_PLANNING & TILEXR_MOONEP_STAGE_DISPATCH) == 0, "Stage capability bits must not overlap"); Check(TILEXR_MOONEP_MAX_TENSOR_RANK == 4, "Tensor rank must remain fixed at four"); diff --git a/tests/moonep/unit/test_tilexr_moonep_c_header.c b/tests/moonep/unit/test_tilexr_moonep_c_header.c index c8b880b0..c11e76df 100644 --- a/tests/moonep/unit/test_tilexr_moonep_c_header.c +++ b/tests/moonep/unit/test_tilexr_moonep_c_header.c @@ -13,6 +13,9 @@ int main(void) TileXRMoonEpReduceGradArgsV1 reduce = {0}; TileXRMoonEpReduceGradWorkspaceQueryV2 query = {0}; TileXRMoonEpReduceGradWorkspaceInfoV2 info = {0}; + TileXRMoonEpReduceGradSourceSliceV2 source = {0}; + TileXRMoonEpReduceGradPrepareArgsV2 prepare = {0}; + TileXRMoonEpReduceGradPreparedV2 prepared = 0; TileXRMoonEpReduceGradArgsV2 reduceV2 = {0}; tensor.structSize = (uint32_t)sizeof(tensor); @@ -34,7 +37,11 @@ int main(void) reduce.output = &tensor; query.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; info.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; + source.data = &tensor; + prepare.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; + prepare.sources[0] = source; reduceV2.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; + reduceV2.prepared = prepared; return tensor.dtype == TILEXR_MOONEP_DTYPE_FLOAT16 && plan.nvS == 1 && dispatch.hiddenSh == &tensor && @@ -44,5 +51,6 @@ int main(void) reduce.input == &tensor && reduce.output == &tensor && query.abiVersion == TILEXR_MOONEP_ABI_VERSION_V2 && info.abiVersion == TILEXR_MOONEP_ABI_VERSION_V2 && + prepare.sources[0].data == &tensor && reduceV2.abiVersion == TILEXR_MOONEP_ABI_VERSION_V2 ? 0 : 1; } diff --git a/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp b/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp index e4874560..c7efb634 100644 --- a/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp @@ -231,22 +231,63 @@ int main() Excludes("reduce launch", reduceLaunch, "LaunchMoonEpKernel("); Contains("reduce kernel", reduceKernel, "extern \"C\" __global__ __aicore__ void tilexr_moonep_reduce_grad_kernel"); - Contains("reduce kernel", reduceKernel, "RunSender"); - Contains("reduce kernel", reduceKernel, "RunSenderTo"); + Contains("reduce kernel", reduceKernel, "RunLaneLeader"); + Contains("reduce kernel", reduceKernel, "RunLaneHelper"); + Contains("reduce kernel", reduceKernel, "IssueWave"); Contains("reduce kernel", reduceKernel, - "controlIndex += controlBlockCount_"); - Contains("reduce kernel", reduceKernel, "RunReceiver"); + "__attribute__((always_inline)) inline __aicore__ bool IssueWave"); Contains("reduce kernel", reduceKernel, - "sourceRank * prefetchSlots_ + slot"); + "__attribute__((always_inline)) inline __aicore__ bool CompleteWave"); Contains("reduce kernel", reduceKernel, - "source * prefetchSlots_ + slot"); - Contains("reduce kernel", reduceKernel, "slot < prefetchSlots_"); + "__attribute__((always_inline)) inline __aicore__ void PublishItem"); + Contains("reduce kernel", reduceKernel, + "__attribute__((always_inline)) inline __aicore__ void LoadItem"); + Excludes("reduce kernel", reduceKernel, "const ReduceGradBankItem &item"); + Excludes("reduce kernel", reduceKernel, "ReduceGradBankItem item"); + Contains("reduce kernel", reduceKernel, + "entry / static_cast(prefetchSlots_)"); + Excludes("reduce kernel", reduceKernel, "prefetchSlots_ > expertsPerRank_"); + Contains("reduce kernel", reduceKernel, "CollectWaveContributors"); Excludes("reduce kernel", reduceKernel, "reduceBuffers_"); - Contains("reduce kernel", reduceKernel, "DataAsFlagSend"); - Contains("reduce kernel", reduceKernel, "DataAsFlagCheckBatchCleared"); Contains("reduce kernel", reduceKernel, "AscendC::Add"); - Contains("reduce kernel", reduceKernel, "UDMAPutRegisteredSignalNbiOnQp"); - Contains("reduce kernel", reduceKernel, "UDMAQuietStatusOnQp"); + Contains("reduce kernel", reduceKernel, "accumPongBuf_"); + Contains("reduce kernel", reduceKernel, "HardEvent::MTE3_MTE2"); + Contains("reduce kernel", reduceKernel, "UDMAProfileGetNbiOnQpDeferred"); + Contains("reduce kernel", reduceKernel, "UDMAProfileFlushQpDoorbell"); + Contains("reduce kernel", reduceKernel, "UDMAProfileQuietStatusOnQpUntil"); + Contains("reduce kernel", reduceKernel, "AscendC::AtomicAdd"); + Contains("reduce kernel", reduceKernel, "bankIssued[kReduceGradBankCount]"); + Contains("reduce kernel", reduceKernel, "bankPublished[kReduceGradBankCount]"); + Contains("reduce kernel", reduceKernel, "nextReadySequence"); + Contains("reduce kernel", reduceKernel, "sequence >= kReduceGradBankCount"); + Contains("reduce kernel", reduceKernel, "ProjectionForLane"); + Contains("reduce kernel", reduceKernel, + "workOrdinal % projectionQpCounts_[projection]"); + Excludes("reduce kernel", reduceKernel, "StageWave"); + Excludes("reduce kernel", reduceKernel, "PublishLaneReadDone"); + Excludes("reduce kernel", reduceKernel, "PublishLaneWorkDone"); + Excludes("reduce kernel", reduceKernel, "ReleaseSources"); + Contains("reduce kernel", reduceKernel, "AscendC::SyncAll()"); + Contains("reduce kernel", reduceKernel, "SyncCollectives"); + Contains("reduce kernel", reduceKernel, "kReduceGradBarrierStep"); + Contains("reduce kernel", reduceKernel, "kReduceGradBarrierFailureStep"); + Contains("reduce kernel", reduceKernel, "CollectiveConfigurationValid"); + Contains("reduce kernel", reduceKernel, + "const bool configured = ConfigurationValid()"); + Contains("reduce kernel", reduceKernel, "if (!configured)"); + Contains("reduce kernel", reduceKernel, "sync.GetInnerFlag(peer, 0)"); + Contains("reduce kernel", reduceKernel, "if (DeviceStatus() == 0U)"); + Contains("reduce kernel", reduceKernel, "ClearLocalSources"); + Contains("reduce kernel clear", reduceKernel, + "tile += static_cast(blockCount_)"); + Excludes("reduce kernel clear", reduceKernel, + "rowOrdinal % blockCount_ != blockIdx_"); + Excludes("reduce kernel", reduceKernel, "RunSender"); + Excludes("reduce kernel", reduceKernel, "RunReceiver"); + Excludes("reduce kernel", reduceKernel, "DataAsFlag"); + Excludes("reduce kernel", reduceKernel, "UDMAPut"); + Excludes("reduce kernel", reduceKernel, "Acknowledgement"); + Excludes("reduce kernel", reduceKernel, "peerRecord"); Excludes("reduce kernel", reduceKernel, "tilexr_moonep_reduce_grad_status_kernel"); Excludes("reduce kernel", reduceKernel, "kReduceGradDeviceStatusSuccess"); diff --git a/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp b/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp index 9b84859e..e279bed3 100644 --- a/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp @@ -113,6 +113,18 @@ void TestLaunch() seenContext.layout.blockDim == 4, "prefetch must cap workers without rejecting the shared-domain QP count"); + Reset(); + qpNum = 3; + plan = Plan(); gate = Weight(0x100000, 4, 8); + up = Weight(0x101000, 4, 16); down = Weight(0x102000, 8, 8); + args = Args(&plan, &gate, &up, &down); + Status("prefetch three QPs", + TileXRMoonEp::TileXRMoonEpRunPrefetchWeightV1(&args, stream), + TILEXR_MOONEP_SUCCESS); + Check(launchCalls == 1 && seenContext.layout.qpNum == 3 && + seenContext.layout.blockDim == 2, + "prefetch three-QP layout must use two workers"); + gate.dtype = TILEXR_MOONEP_DTYPE_FLOAT32; Status("prefetch dtype", TileXRMoonEp::TileXRMoonEpRunPrefetchWeightV1(&args, stream), TILEXR_MOONEP_ERROR_INVALID_ARGUMENT); gate = Weight(0x100000, 4, 8); args.flags = 1; diff --git a/tests/moonep/unit/test_tilexr_moonep_reduce_grad_host.cpp b/tests/moonep/unit/test_tilexr_moonep_reduce_grad_host.cpp index 7ca99210..77ce38e4 100644 --- a/tests/moonep/unit/test_tilexr_moonep_reduce_grad_host.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_reduce_grad_host.cpp @@ -20,19 +20,23 @@ namespace { int g_failures = 0; int g_launchCalls = 0; int g_memsetCalls = 0; -int g_synchronizeCalls = 0; -aclError g_synchronizeReturn = ACL_SUCCESS; -int g_commHostReturn = TileXR::TILEXR_SUCCESS; -int g_commDevReturn = TileXR::TILEXR_SUCCESS; -int g_registryReturn = TileXR::TILEXR_SUCCESS; +int g_syncCalls = 0; +int g_profileRegisterCalls = 0; +int g_profileQueryCalls = 0; +int g_profileUnregisterCalls = 0; +int g_qpCountCalls = 0; +int g_profileRegisterReturn = TileXR::TILEXR_SUCCESS; +int g_profileQueryReturn = TileXR::TILEXR_SUCCESS; +int g_profileUnregisterReturn = TileXR::TILEXR_SUCCESS; int g_qpCountReturn = TileXR::TILEXR_SUCCESS; int g_launchReturn = TileXR::TILEXR_SUCCESS; -int g_qpCountCalls = 0; -uint32_t g_qpCount = 3; +aclError g_syncReturn = ACL_SUCCESS; +uint32_t g_qpCount = 8; TileXR::CommArgs g_commArgs {}; -TileXR::TileXRUDMARegistry g_registry {}; +TileXR::TileXRUDMAProfileDesc g_profileDesc {}; +TileXR::TileXRUDMAProfileRegistry g_profileRegistry {}; +TileXR::TileXRUDMAProfileView g_profileView {}; GM_ADDR g_commDev = reinterpret_cast(UINTPTR_C(0x800000)); -TileXRMoonEp::ReduceGradLayout g_launchedLayout {}; void Check(bool condition, const std::string &message) { @@ -50,34 +54,32 @@ void CheckStatus(const std::string &label, int actual, int expected) } } -TileXRMoonEpPlanV1 Plan() +TileXRMoonEpPlanV1 Plan(int64_t rankSize = 8, int64_t expertsPerRank = 2, + int64_t prefetchSlots = 2) { TileXRMoonEpPlanV1 plan {}; plan.structSize = sizeof(plan); plan.abiVersion = TILEXR_MOONEP_ABI_VERSION_V1; plan.n = 4; plan.k = 2; - plan.e = 4; - plan.b = 2; - plan.r = 2; + plan.e = rankSize * expertsPerRank; + plan.b = prefetchSlots; + plan.r = rankSize; plan.nvS = 4; - plan.dst = reinterpret_cast(UINTPTR_C(0x11000)); plan.expertsToCopy = reinterpret_cast(UINTPTR_C(0x13000)); - plan.remoteStats = reinterpret_cast(UINTPTR_C(0x14000)); - plan.status = reinterpret_cast(UINTPTR_C(0x15000)); return plan; } -TileXRMoonEpTensorV1 Gradient(void *data, uint64_t rowElements) +TileXRMoonEpTensorV1 Gradient(void *data, int64_t rows, uint64_t rowElements) { TileXRMoonEpTensorV1 tensor {}; tensor.structSize = sizeof(tensor); tensor.abiVersion = TILEXR_MOONEP_ABI_VERSION_V1; tensor.data = data; - tensor.elementCount = 6 * rowElements; + tensor.elementCount = static_cast(rows) * rowElements; tensor.dtype = TILEXR_MOONEP_DTYPE_FLOAT32; tensor.rank = 2; - tensor.shape[0] = 6; + tensor.shape[0] = rows; tensor.shape[1] = static_cast(rowElements); return tensor; } @@ -95,36 +97,37 @@ TileXRMoonEpTensorV1 Status() return tensor; } -void Reset() +void Reset(int64_t rankSize = 8) { g_launchCalls = 0; g_memsetCalls = 0; - g_synchronizeCalls = 0; - g_synchronizeReturn = ACL_SUCCESS; - g_commHostReturn = TileXR::TILEXR_SUCCESS; - g_commDevReturn = TileXR::TILEXR_SUCCESS; - g_registryReturn = TileXR::TILEXR_SUCCESS; + g_syncCalls = 0; + g_profileRegisterCalls = 0; + g_profileQueryCalls = 0; + g_profileUnregisterCalls = 0; + g_qpCountCalls = 0; + g_profileRegisterReturn = TileXR::TILEXR_SUCCESS; + g_profileQueryReturn = TileXR::TILEXR_SUCCESS; + g_profileUnregisterReturn = TileXR::TILEXR_SUCCESS; g_qpCountReturn = TileXR::TILEXR_SUCCESS; g_launchReturn = TileXR::TILEXR_SUCCESS; - g_qpCountCalls = 0; - g_qpCount = 3; - g_launchedLayout = TileXRMoonEp::ReduceGradLayout {}; + g_syncReturn = ACL_SUCCESS; + g_qpCount = 8; g_commArgs = TileXR::CommArgs {}; g_commArgs.rank = 0; g_commArgs.localRank = 0; - g_commArgs.rankSize = 2; - g_commArgs.localRankSize = 2; + g_commArgs.rankSize = static_cast(rankSize); + g_commArgs.localRankSize = static_cast(rankSize); g_commArgs.extraFlag = TileXR::ExtraFlag::TOPO_910A5; - g_commArgs.peerMems[0] = reinterpret_cast(UINTPTR_C(0x200000)); - g_commArgs.peerMems[1] = reinterpret_cast(UINTPTR_C(0x300000)); - g_registry = TileXR::TileXRUDMARegistry {}; - g_registry.rankSize = 2; - g_registry.regionCount = 1; + g_commArgs.extraFlag |= TileXR::ExtraFlag::UDMA; + g_profileDesc = TileXR::TileXRUDMAProfileDesc {}; + g_profileRegistry = TileXR::TileXRUDMAProfileRegistry {}; + g_profileView = TileXR::TileXRUDMAProfileView {}; } -TileXRMoonEpReduceGradWorkspaceInfoV2 Query( - TileXRMoonEpPlanV1 *plan, TileXRMoonEpTensorV1 *gate, - TileXRMoonEpTensorV1 *up, TileXRMoonEpTensorV1 *down, int expected) +TileXRMoonEpReduceGradWorkspaceInfoV2 Query(TileXRMoonEpPlanV1 *plan, + TileXRMoonEpTensorV1 *gate, TileXRMoonEpTensorV1 *up, + TileXRMoonEpTensorV1 *down, int expected) { TileXRMoonEpReduceGradWorkspaceQueryV2 query {}; query.structSize = sizeof(query); @@ -137,15 +140,17 @@ TileXRMoonEpReduceGradWorkspaceInfoV2 Query( TileXRMoonEpReduceGradWorkspaceInfoV2 info {}; info.structSize = sizeof(info); info.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; - CheckStatus("workspace query", TileXRMoonEpReduceGradGetWorkspaceSizeV2(&query, &info), expected); + CheckStatus("workspace query", TileXRMoonEpReduceGradGetWorkspaceSizeV2( + &query, &info), expected); return info; } -TileXRMoonEpReduceGradArgsV2 Args(TileXRMoonEpPlanV1 *plan, +TileXRMoonEpReduceGradPrepareArgsV2 PrepareArgs(TileXRMoonEpPlanV1 *plan, TileXRMoonEpTensorV1 *gate, TileXRMoonEpTensorV1 *up, - TileXRMoonEpTensorV1 *down, TileXRMoonEpTensorV1 *status) + TileXRMoonEpTensorV1 *down, + const TileXRMoonEpReduceGradWorkspaceInfoV2 &info) { - TileXRMoonEpReduceGradArgsV2 args {}; + TileXRMoonEpReduceGradPrepareArgsV2 args {}; args.structSize = sizeof(args); args.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; args.comm = reinterpret_cast(UINTPTR_C(0x1000)); @@ -153,270 +158,241 @@ TileXRMoonEpReduceGradArgsV2 Args(TileXRMoonEpPlanV1 *plan, args.gate = gate; args.up = up; args.down = down; - args.status = status; - args.waitIterations = 1000; + args.workspace = info.workspaceBytes == 0 ? nullptr : + reinterpret_cast(UINTPTR_C(0x20000000)); + args.workspaceBytes = info.workspaceBytes; + TileXRMoonEpTensorV1 *gradients[] = {gate, up, down}; + for (uint32_t projection = 0; projection < 3; ++projection) { + args.sources[projection].registrationBase = reinterpret_cast( + UINTPTR_C(0x30000000) + projection * UINTPTR_C(0x10000000)); + args.sources[projection].data = static_cast( + args.sources[projection].registrationBase) + 0x1000; + args.sources[projection].bytes = static_cast(plan->b) * + info.rowBytes[projection]; + args.sources[projection].registrationBytes = + args.sources[projection].bytes + 0x2000; + Check(args.sources[projection].bytes == static_cast(plan->b) * + static_cast(gradients[projection]->shape[1]) * sizeof(float), + "test source byte construction mismatch"); + } return args; } -void TestPeerOnly() +TileXRMoonEpReduceGradArgsV2 LaunchArgs(TileXRMoonEpReduceGradPreparedV2 prepared, + const TileXRMoonEpReduceGradPrepareArgsV2 &prepare, + TileXRMoonEpTensorV1 *status) { - Reset(); - TileXRMoonEpPlanV1 plan = Plan(); - TileXRMoonEpTensorV1 gate = Gradient(reinterpret_cast(UINTPTR_C(0x400000)), 64); - TileXRMoonEpTensorV1 up = Gradient(reinterpret_cast(UINTPTR_C(0x500000)), 128); - TileXRMoonEpTensorV1 down = Gradient(reinterpret_cast(UINTPTR_C(0x600000)), 256); - TileXRMoonEpTensorV1 status = Status(); - const auto info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); - Check(info.workspaceBytes == 0 && info.udmaChunkBytes == 0, - "peer-only query must not request UDMA workspace"); - Check(info.transports[0] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - info.transports[1] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - info.transports[2] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER, - "peer-only query selected the wrong transport"); - - auto args = Args(&plan, &gate, &up, &down, &status); - aclrtStream stream = reinterpret_cast(UINTPTR_C(0x700000)); - CheckStatus("peer launch", TileXRMoonEpReduceGradV2(&args, stream), TileXR::TILEXR_SUCCESS); - Check(g_memsetCalls == 1 && g_launchCalls == 1, - "peer launch must initialize status and launch exactly once"); - Check(g_qpCountCalls == 0, - "peer-only query and launch must not query UDMA QPs"); + TileXRMoonEpReduceGradArgsV2 args {}; + args.structSize = sizeof(args); + args.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; + args.prepared = prepared; + args.plan = prepare.plan; + args.gate = prepare.gate; + args.up = prepare.up; + args.down = prepare.down; + for (uint32_t projection = 0; projection < 3; ++projection) { + args.sources[projection] = prepare.sources[projection]; + } + args.status = status; + args.waitIterations = 1000; + return args; } -void TestCompactPrefetchSlots() +void TestPreparedLifecycleAndHotLaunch() { Reset(); TileXRMoonEpPlanV1 plan = Plan(); - plan.b = 1; TileXRMoonEpTensorV1 gate = Gradient( - reinterpret_cast(UINTPTR_C(0x400000)), 64); + reinterpret_cast(UINTPTR_C(0x400000)), 18, 1024); TileXRMoonEpTensorV1 up = Gradient( - reinterpret_cast(UINTPTR_C(0x500000)), 128); + reinterpret_cast(UINTPTR_C(0x500000)), 18, 2048); TileXRMoonEpTensorV1 down = Gradient( - reinterpret_cast(UINTPTR_C(0x600000)), 256); - TileXRMoonEpTensorV1 *gradients[] = {&gate, &up, &down}; - for (TileXRMoonEpTensorV1 *tensor : gradients) { - tensor->shape[0] = plan.e + plan.b; - tensor->elementCount = static_cast(tensor->shape[0]) * - static_cast(tensor->shape[1]); - } - TileXRMoonEpTensorV1 status = Status(); + reinterpret_cast(UINTPTR_C(0x600000)), 18, 4096); const auto info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); - Check(info.workspaceBytes == 0, - "compact-B peer query must not request UDMA workspace"); + Check(info.workspaceBytes > 0 && info.qpCount == 8 && + info.projectionQpCounts[0] == 2 && info.projectionQpCounts[1] == 2 && + info.projectionQpCounts[2] == 4, + "workspace query did not expose weighted owner-pull layout"); + auto prepare = PrepareArgs(&plan, &gate, &up, &down, info); + TileXRMoonEpReduceGradPreparedV2 prepared = nullptr; + CheckStatus("prepare", TileXRMoonEpReduceGradPrepareV2(&prepare, &prepared), + TileXR::TILEXR_SUCCESS); + Check(prepared != nullptr && g_profileRegisterCalls == 1 && + g_profileQueryCalls == 1, + "prepare must register and query exactly one persistent profile"); + Check(g_profileDesc.regionCount == 4 && g_profileDesc.qpBindingCount == 8, + "profile must contain staging plus three source regions"); + Check(g_profileDesc.regions[1].base == prepare.sources[0].data && + g_profileDesc.regions[1].bytes == prepare.sources[0].bytes && + g_profileDesc.regions[1].registrationBase == + prepare.sources[0].registrationBase && + g_profileDesc.regions[1].registrationBytes == + prepare.sources[0].registrationBytes, + "profile must keep the logical source view separate from its backing MR"); + for (uint32_t qp = 0; qp < 8; ++qp) { + Check(g_profileDesc.qpBindings[qp].localRegion == 0, + "every QP must write into the staging MR"); + } - auto args = Args(&plan, &gate, &up, &down, &status); + TileXRMoonEpTensorV1 status = Status(); + auto launch = LaunchArgs(prepared, prepare, &status); aclrtStream stream = reinterpret_cast(UINTPTR_C(0x700000)); - CheckStatus("compact-B peer launch", TileXRMoonEpReduceGradV2(&args, stream), + CheckStatus("hot launch", TileXRMoonEpReduceGradV2(&launch, stream), + TileXR::TILEXR_SUCCESS); + Check(g_profileRegisterCalls == 1 && g_profileQueryCalls == 2 && + g_memsetCalls == 1 && g_launchCalls == 1, + "hot launch must only validate the persistent profile and enqueue work"); + + launch.sources[0].data = reinterpret_cast(UINTPTR_C(0xdead0000)); + CheckStatus("source identity mismatch", TileXRMoonEpReduceGradV2(&launch, stream), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + launch.sources[0] = prepare.sources[0]; + launch.sources[0].registrationBase = + reinterpret_cast(UINTPTR_C(0xdead0000)); + CheckStatus("source registration identity mismatch", + TileXRMoonEpReduceGradV2(&launch, stream), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + launch.sources[0] = prepare.sources[0]; + gate.data = reinterpret_cast(UINTPTR_C(0xbeef0000)); + CheckStatus("gradient identity mismatch", TileXRMoonEpReduceGradV2(&launch, stream), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + gate.data = reinterpret_cast(UINTPTR_C(0x400000)); + Check(g_memsetCalls == 1 && g_launchCalls == 1, + "pointer mismatches must fail before enqueuing work"); + + g_launchReturn = -91; + CheckStatus("launch failure", TileXRMoonEpReduceGradV2(&launch, stream), -91); + Check(g_memsetCalls == 2 && g_launchCalls == 2 && g_syncCalls == 1, + "launch failure must drain the enqueued status memset"); + g_syncReturn = 1; + CheckStatus("launch failure drain failure", + TileXRMoonEpReduceGradV2(&launch, stream), TileXR::TILEXR_ERROR_MKIRT); + Check(g_memsetCalls == 3 && g_launchCalls == 3 && g_syncCalls == 2, + "failed launch drain must be attempted exactly once"); + g_launchReturn = TileXR::TILEXR_SUCCESS; + g_syncReturn = ACL_SUCCESS; + + g_profileView.infoDev = reinterpret_cast(UINTPTR_C(0xdeadbeef)); + CheckStatus("profile mismatch", TileXRMoonEpReduceGradV2(&launch, stream), + TileXR::TILEXR_ERROR_NOT_INITIALIZED); + g_profileView.infoDev = reinterpret_cast(UINTPTR_C(0x900000)); + + CheckStatus("destroy", TileXRMoonEpReduceGradDestroyPreparedV2(prepared), TileXR::TILEXR_SUCCESS); - Check(g_launchCalls == 1, - "compact-B ReduceGrad must launch exactly once"); + Check(g_profileUnregisterCalls == 1, + "destroy must unregister the persistent profile exactly once"); } -void TestMixedUdma() +void TestPreparationValidationAndCapability() { Reset(); - g_commArgs.localRankSize = 1; TileXRMoonEpPlanV1 plan = Plan(); - const uint64_t thresholdElements = TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES / sizeof(float); - TileXRMoonEpTensorV1 gate = Gradient(reinterpret_cast(UINTPTR_C(0x400000)), thresholdElements); - TileXRMoonEpTensorV1 up = Gradient(reinterpret_cast(UINTPTR_C(0x500000)), thresholdElements + 1); - TileXRMoonEpTensorV1 down = Gradient(reinterpret_cast(UINTPTR_C(0x600000)), 256); - TileXRMoonEpTensorV1 status = Status(); - const auto info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); - Check(info.transports[0] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - info.transports[1] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA, - "mixed query threshold selection mismatch"); - Check(info.workspaceBytes > 0 && - info.workspaceBytes % TileXRMoonEp::kReduceGradUdmaWorkspaceAlignment == 0 && - info.workspaceAlignment == TileXRMoonEp::kReduceGradUdmaWorkspaceAlignment, - "mixed query must return aligned UDMA workspace"); - Check(g_qpCountCalls == 1, - "mixed query must obtain the negotiated UDMA QP count"); - - void *workspace = reinterpret_cast(UINTPTR_C(0x1000000)); - g_commArgs.extraFlag |= TileXR::ExtraFlag::UDMA; - g_commArgs.udmaInfoPtr = reinterpret_cast(UINTPTR_C(0x17000)); - g_commArgs.udmaRegistryPtr = reinterpret_cast(UINTPTR_C(0x18000)); - for (int rank = 0; rank < 2; ++rank) { - g_registry.regions[rank].base = rank == 0 ? static_cast(workspace) : - reinterpret_cast(UINTPTR_C(0x2000000)); - g_registry.regions[rank].bytes = info.workspaceBytes; - } + TileXRMoonEpTensorV1 gate = Gradient( + reinterpret_cast(UINTPTR_C(0x400000)), 18, 1024); + TileXRMoonEpTensorV1 up = Gradient( + reinterpret_cast(UINTPTR_C(0x500000)), 18, 1024); + TileXRMoonEpTensorV1 down = Gradient( + reinterpret_cast(UINTPTR_C(0x600000)), 18, 1024); + auto info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); + auto prepare = PrepareArgs(&plan, &gate, &up, &down, info); + TileXRMoonEpReduceGradPreparedV2 prepared = nullptr; - auto args = Args(&plan, &gate, &up, &down, &status); - args.workspace = workspace; - args.workspaceBytes = info.workspaceBytes; - aclrtStream stream = reinterpret_cast(UINTPTR_C(0x700000)); - CheckStatus("mixed launch", TileXRMoonEpReduceGradV2(&args, stream), TileXR::TILEXR_SUCCESS); - Check(g_launchCalls == 1 && - g_launchedLayout.transports[TileXRMoonEp::kReduceGradUp] == - TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA && - g_launchedLayout.udmaQpCount == g_qpCount, - "cross-node mixed launch did not preserve transport selection"); - Check(g_qpCountCalls == 2, - "mixed launch must obtain the negotiated UDMA QP count"); - - g_registry.regions[1].base += 512; - CheckStatus("misaligned peer workspace", TileXRMoonEpReduceGradV2(&args, stream), + prepare.sources[1].bytes -= sizeof(float); + CheckStatus("source size", TileXRMoonEpReduceGradPrepareV2(&prepare, &prepared), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + prepare = PrepareArgs(&plan, &gate, &up, &down, info); + prepare.sources[1].registrationBase = prepare.sources[1].data; + prepare.sources[1].registrationBytes = prepare.sources[1].bytes - 1; + CheckStatus("source outside registration", + TileXRMoonEpReduceGradPrepareV2(&prepare, &prepared), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); - Check(g_launchCalls == 1, "misaligned peer workspace must not launch"); - g_registry.regions[1].base -= 512; - g_registry.regions[1].bytes += TileXRMoonEp::kReduceGradUdmaWorkspaceAlignment; - CheckStatus("mismatched peer workspace bytes", TileXRMoonEpReduceGradV2(&args, stream), + prepare = PrepareArgs(&plan, &gate, &up, &down, info); + prepare.workspace = reinterpret_cast(UINTPTR_C(0x20000200)); + CheckStatus("workspace alignment", TileXRMoonEpReduceGradPrepareV2(&prepare, &prepared), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); - Check(g_launchCalls == 1, "mismatched peer workspace bytes must not launch"); - g_registry.regions[1].bytes = info.workspaceBytes; g_commArgs.extraFlag &= ~TileXR::ExtraFlag::UDMA; - CheckStatus("missing UDMA", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_ERROR_NOT_INITIALIZED); - Check(g_launchCalls == 1, "missing UDMA must not launch"); - + (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_SUPPORT); g_commArgs.extraFlag |= TileXR::ExtraFlag::UDMA; - const int qpCallsBeforeFailureChecks = g_qpCountCalls; - g_qpCountReturn = -92; - (void)Query(&plan, &gate, &up, &down, -92); - Check(g_qpCountCalls == qpCallsBeforeFailureChecks + 1, - "UDMA QP query failure must be observed by workspace preparation"); + g_qpCount = 2; + (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_SUPPORT); + g_qpCount = 32; + (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_SUPPORT); + g_commArgs.extraFlag |= TileXR::ExtraFlag::UDMA_SHARED_QP; + info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); + Check(info.qpCount == 3, + "shared-domain query must report only the three active ReduceGrad lanes"); + prepare = PrepareArgs(&plan, &gate, &up, &down, info); + CheckStatus("shared-domain prepare", + TileXRMoonEpReduceGradPrepareV2(&prepare, &prepared), + TileXR::TILEXR_SUCCESS); + Check(g_profileDesc.qpBindingCount == 32, + "persistent profile must bind every transport QP in the shared domain"); + Check(g_profileDesc.qpBindings[0].remoteRegion == 1 && + g_profileDesc.qpBindings[1].remoteRegion == 2 && + g_profileDesc.qpBindings[16].remoteRegion == 3, + "active shared-domain QPs must bind gate/up/down to physical QPs 0/1/16"); + CheckStatus("shared-domain destroy", + TileXRMoonEpReduceGradDestroyPreparedV2(prepared), TileXR::TILEXR_SUCCESS); + prepared = nullptr; + g_qpCount = 33; + (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_SUPPORT); + Check(g_qpCountCalls == 9, + "UDMA-capable multi-rank queries must validate the current hardware QP count"); - g_qpCountReturn = TileXR::TILEXR_SUCCESS; - g_qpCount = 0; - (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_INITIALIZED); - g_qpCount = TileXRMoonEp::kReduceGradMaxUdmaQpCount + 1; - (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_INITIALIZED); - Check(g_qpCountCalls == qpCallsBeforeFailureChecks + 3, - "invalid UDMA QP counts must be validated after every query"); + Reset(); + plan = Plan(); + gate = Gradient(reinterpret_cast(UINTPTR_C(0x400000)), 18, 1024); + up = Gradient(reinterpret_cast(UINTPTR_C(0x500000)), 18, 1024); + down = Gradient(reinterpret_cast(UINTPTR_C(0x600000)), 18, 1024); + info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); + prepare = PrepareArgs(&plan, &gate, &up, &down, info); + g_profileRegisterReturn = -91; + CheckStatus("profile registration failure", + TileXRMoonEpReduceGradPrepareV2(&prepare, &prepared), -91); + Check(prepared == nullptr && g_profileQueryCalls == 0, + "failed registration must not expose a prepared handle"); } -void TestSingleRankLargeRowsDoNotRequireUdma() +void TestRanksBelowFourAreUnsupported() { - Reset(); - g_commArgs.rankSize = 1; - g_commArgs.localRankSize = 1; - g_registryReturn = TileXR::TILEXR_ERROR_NOT_INITIALIZED; - TileXRMoonEpPlanV1 plan = Plan(); - plan.r = 1; - plan.b = plan.e; - const uint64_t largeRow = - TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES / sizeof(float) + 1; + Reset(3); + TileXRMoonEpPlanV1 plan = Plan(3); TileXRMoonEpTensorV1 gate = Gradient( - reinterpret_cast(UINTPTR_C(0x400000)), largeRow); + reinterpret_cast(UINTPTR_C(0x400000)), 16, 1024); TileXRMoonEpTensorV1 up = Gradient( - reinterpret_cast(UINTPTR_C(0x500000)), largeRow + 1); + reinterpret_cast(UINTPTR_C(0x500000)), 16, 1024); TileXRMoonEpTensorV1 down = Gradient( - reinterpret_cast(UINTPTR_C(0x600000)), largeRow + 2); - TileXRMoonEpTensorV1 *gradients[] = {&gate, &up, &down}; - for (TileXRMoonEpTensorV1 *tensor : gradients) { - tensor->shape[0] = plan.e + plan.b; - tensor->elementCount = static_cast(tensor->shape[0]) * - static_cast(tensor->shape[1]); - } - TileXRMoonEpTensorV1 status = Status(); - const auto info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); - Check(info.workspaceBytes == 0 && info.udmaChunkBytes == 0, - "single-rank large-row query must not request UDMA workspace"); - Check(info.transports[0] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - info.transports[1] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - info.transports[2] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER, - "single-rank large-row query must select the local peer path"); - - auto args = Args(&plan, &gate, &up, &down, &status); - aclrtStream stream = reinterpret_cast(UINTPTR_C(0x700000)); - CheckStatus("single-rank large-row launch", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_SUCCESS); - Check(g_launchCalls == 1, - "single-rank large-row launch must not consult an unavailable UDMA registry"); + reinterpret_cast(UINTPTR_C(0x600000)), 16, 1024); + (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_NOT_SUPPORT); + Check(g_qpCountCalls == 0 && g_profileRegisterCalls == 0, + "rank counts below four must fail before UDMA profile work"); + + Reset(4); + plan = Plan(4); + gate = Gradient(reinterpret_cast(UINTPTR_C(0x400000)), 10, 1024); + up = Gradient(reinterpret_cast(UINTPTR_C(0x500000)), 10, 1024); + down = Gradient(reinterpret_cast(UINTPTR_C(0x600000)), 10, 1024); + (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); + Check(g_qpCountCalls == 1, + "four ranks must reach UDMA capability validation"); } -void TestLargeRankWorkspaceQuery() +void TestSlotsMayExceedExpertsPerRank() { Reset(); - g_commArgs.rankSize = TileXR::TILEXR_MAX_RANK_SIZE; - g_commArgs.localRankSize = 8; - for (int rank = 0; rank < g_commArgs.rankSize; ++rank) { - g_commArgs.peerMems[rank] = reinterpret_cast( - UINTPTR_C(0x200000) + static_cast(rank) * UINTPTR_C(0x100000)); - } - - TileXRMoonEpPlanV1 plan = Plan(); - plan.n = TileXR::TILEXR_MAX_RANK_SIZE; - plan.k = 1; - plan.e = TileXR::TILEXR_MAX_RANK_SIZE; - plan.b = 1; - plan.r = TileXR::TILEXR_MAX_RANK_SIZE; + TileXRMoonEpPlanV1 plan = Plan(8, 8, 14); TileXRMoonEpTensorV1 gate = Gradient( - reinterpret_cast(UINTPTR_C(0x400000)), 64); + reinterpret_cast(UINTPTR_C(0x400000)), 78, 1024); TileXRMoonEpTensorV1 up = Gradient( - reinterpret_cast(UINTPTR_C(0x500000)), 128); + reinterpret_cast(UINTPTR_C(0x500000)), 78, 1024); TileXRMoonEpTensorV1 down = Gradient( - reinterpret_cast(UINTPTR_C(0x600000)), 256); - TileXRMoonEpTensorV1 *gradients[] = {&gate, &up, &down}; - for (TileXRMoonEpTensorV1 *tensor : gradients) { - tensor->shape[0] = plan.e + plan.b; - tensor->elementCount = static_cast(tensor->shape[0]) * - static_cast(tensor->shape[1]); - } - + reinterpret_cast(UINTPTR_C(0x600000)), 78, 1024); const auto info = Query(&plan, &gate, &up, &down, TileXR::TILEXR_SUCCESS); - Check(info.blockDim == TileXRMoonEp::kReduceGradMaxAivBlockCount, - "128-rank workspace query must retain the 64-AIV launch limit"); -} - -void TestLaunchFailureDrainsEnqueuedStatusReset() -{ - Reset(); - TileXRMoonEpPlanV1 plan = Plan(); - TileXRMoonEpTensorV1 gate = Gradient(reinterpret_cast(UINTPTR_C(0x400000)), 64); - TileXRMoonEpTensorV1 up = Gradient(reinterpret_cast(UINTPTR_C(0x500000)), 64); - TileXRMoonEpTensorV1 down = Gradient(reinterpret_cast(UINTPTR_C(0x600000)), 64); - TileXRMoonEpTensorV1 status = Status(); - auto args = Args(&plan, &gate, &up, &down, &status); - aclrtStream stream = reinterpret_cast(UINTPTR_C(0x700000)); - - g_launchReturn = TileXR::TILEXR_ERROR_NOT_SUPPORT; - CheckStatus("launch failure", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_ERROR_NOT_SUPPORT); - Check(g_memsetCalls == 1 && g_launchCalls == 1 && g_synchronizeCalls == 1, - "launch failure must drain the enqueued status reset"); - - g_synchronizeReturn = -91; - CheckStatus("launch failure with drain failure", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_ERROR_MKIRT); - Check(g_memsetCalls == 2 && g_launchCalls == 2 && g_synchronizeCalls == 2, - "stream drain failure must be reported after a failed launch"); -} - -void TestValidation() -{ - Reset(); - TileXRMoonEpPlanV1 plan = Plan(); - TileXRMoonEpTensorV1 gate = Gradient(reinterpret_cast(UINTPTR_C(0x400000)), 64); - TileXRMoonEpTensorV1 up = Gradient(reinterpret_cast(UINTPTR_C(0x500000)), 64); - TileXRMoonEpTensorV1 down = Gradient(reinterpret_cast(UINTPTR_C(0x600000)), 64); - TileXRMoonEpTensorV1 status = Status(); - auto args = Args(&plan, &gate, &up, &down, &status); - aclrtStream stream = reinterpret_cast(UINTPTR_C(0x700000)); - - gate.dtype = TILEXR_MOONEP_DTYPE_FLOAT16; - CheckStatus("gradient dtype", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); - gate.dtype = TILEXR_MOONEP_DTYPE_FLOAT32; - g_commArgs.peerMems[1] = nullptr; - CheckStatus("missing peer", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_ERROR_NOT_INITIALIZED); - g_commArgs.peerMems[1] = reinterpret_cast(UINTPTR_C(0x300000)); - g_commArgs.extraFlag = 0; - CheckStatus("wrong architecture", TileXRMoonEpReduceGradV2(&args, stream), - TileXR::TILEXR_ERROR_NOT_SUPPORT); + Check(info.workspaceBytes > 0, + "workspace query must accept native dedicated-suite B greater than E/R"); - Reset(); - plan = Plan(); - plan.e = std::numeric_limits::max(); - plan.b = std::numeric_limits::max(); - plan.r = 1; - g_commArgs.rankSize = 1; - g_commArgs.localRankSize = 1; + plan.b = std::numeric_limits::max() / plan.r + 1; (void)Query(&plan, &gate, &up, &down, TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); } @@ -424,23 +400,14 @@ void TestValidation() extern "C" int TileXRGetCommArgsHost(TileXRCommPtr, TileXR::CommArgs *&commArgs) { - commArgs = g_commHostReturn == TileXR::TILEXR_SUCCESS ? &g_commArgs : nullptr; - return g_commHostReturn; + commArgs = &g_commArgs; + return TileXR::TILEXR_SUCCESS; } extern "C" int TileXRGetCommArgsDev(TileXRCommPtr, GM_ADDR &commArgs) { - commArgs = g_commDevReturn == TileXR::TILEXR_SUCCESS ? g_commDev : nullptr; - return g_commDevReturn; -} - -extern "C" int TileXRGetUDMARegistryHost(TileXRCommPtr, - const TileXR::TileXRUDMARegistry **registry) -{ - if (registry != nullptr) { - *registry = g_registryReturn == TileXR::TILEXR_SUCCESS ? &g_registry : nullptr; - } - return g_registryReturn; + commArgs = g_commDev; + return TileXR::TILEXR_SUCCESS; } extern "C" int TileXRUDMAGetQpCount(TileXRCommPtr, uint32_t *qpCount) @@ -452,6 +419,54 @@ extern "C" int TileXRUDMAGetQpCount(TileXRCommPtr, uint32_t *qpCount) return g_qpCountReturn; } +extern "C" int TileXRUDMAProfileRegister(TileXRCommPtr, + const TileXR::TileXRUDMAProfileDesc *desc, TileXRUDMAProfileHandle *handle) +{ + ++g_profileRegisterCalls; + if (g_profileRegisterReturn != TileXR::TILEXR_SUCCESS) { + return g_profileRegisterReturn; + } + g_profileDesc = *desc; + *handle = 7; + g_profileRegistry = TileXR::TileXRUDMAProfileRegistry {}; + g_profileRegistry.rankSize = static_cast(g_commArgs.rankSize); + g_profileRegistry.regionCount = desc->regionCount; + g_profileRegistry.qpCount = desc->qpBindingCount; + for (uint32_t qp = 0; qp < desc->qpBindingCount; ++qp) { + g_profileRegistry.qpBindings[qp] = desc->qpBindings[qp]; + } + for (int rank = 0; rank < g_commArgs.rankSize; ++rank) { + for (uint32_t region = 0; region < desc->regionCount; ++region) { + g_profileRegistry.regions[static_cast(rank) * + TileXR::TILEXR_UDMA_PROFILE_MAX_REGIONS + region] = desc->regions[region]; + } + } + g_profileView = TileXR::TileXRUDMAProfileView {}; + g_profileView.rankSize = static_cast(g_commArgs.rankSize); + g_profileView.regionCount = desc->regionCount; + g_profileView.qpCount = desc->qpBindingCount; + g_profileView.infoDev = reinterpret_cast(UINTPTR_C(0x900000)); + g_profileView.registryDev = reinterpret_cast(UINTPTR_C(0xa00000)); + g_profileView.registryHost = &g_profileRegistry; + return TileXR::TILEXR_SUCCESS; +} + +extern "C" int TileXRUDMAProfileQuery(TileXRCommPtr, TileXRUDMAProfileHandle, + TileXR::TileXRUDMAProfileView *view) +{ + ++g_profileQueryCalls; + if (g_profileQueryReturn == TileXR::TILEXR_SUCCESS && view != nullptr) { + *view = g_profileView; + } + return g_profileQueryReturn; +} + +extern "C" int TileXRUDMAProfileUnregister(TileXRCommPtr, TileXRUDMAProfileHandle) +{ + ++g_profileUnregisterCalls; + return g_profileUnregisterReturn; +} + extern "C" aclError aclrtMemsetAsync(void *, size_t, int32_t, size_t, aclrtStream) { ++g_memsetCalls; @@ -466,17 +481,16 @@ extern "C" aclError aclrtMemcpyAsync(void *, size_t, const void *, size_t, extern "C" aclError aclrtSynchronizeStream(aclrtStream) { - ++g_synchronizeCalls; - return g_synchronizeReturn; + ++g_syncCalls; + return g_syncReturn; } namespace TileXRMoonEp { -int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradParams &, - const ReduceGradLaunchContext &context) +int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradLaunchParams &, + const ReduceGradPreparedContext &) { ++g_launchCalls; - g_launchedLayout = context.layout; return g_launchReturn; } @@ -484,12 +498,9 @@ int TileXRMoonEpLaunchReduceGradKernel(const ReduceGradParams &, int main() { - TestPeerOnly(); - TestCompactPrefetchSlots(); - TestMixedUdma(); - TestSingleRankLargeRowsDoNotRequireUdma(); - TestLargeRankWorkspaceQuery(); - TestLaunchFailureDrainsEnqueuedStatusReset(); - TestValidation(); + TestPreparedLifecycleAndHotLaunch(); + TestPreparationValidationAndCapability(); + TestRanksBelowFourAreUnsupported(); + TestSlotsMayExceedExpertsPerRank(); return g_failures == 0 ? 0 : 1; } diff --git a/tests/moonep/unit/test_tilexr_moonep_reduce_grad_launch.cpp b/tests/moonep/unit/test_tilexr_moonep_reduce_grad_launch.cpp index 8220309a..5a3493cd 100644 --- a/tests/moonep/unit/test_tilexr_moonep_reduce_grad_launch.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_reduce_grad_launch.cpp @@ -32,18 +32,14 @@ int g_binaryRegisterCalls = 0; int g_binaryUnregisterCalls = 0; int g_functionRegisterCalls = 0; int g_launchCalls = 0; -bool g_binaryWasNull = false; rtDevBinary_t g_binary {}; void *g_binaryHandle = reinterpret_cast(UINTPTR_C(0xa000)); -void *g_functionBinaryHandle = nullptr; void *g_unregisteredBinaryHandle = nullptr; const void *g_functionStub = nullptr; std::string g_functionStubName; std::string g_functionKernelInfo; -uint32_t g_functionMode = UINT32_MAX; const void *g_launchKernel = nullptr; uint32_t g_blockDim = 0; -rtSmDesc_t *g_smDesc = nullptr; rtStream_t g_stream = nullptr; uint32_t g_launchFlags = UINT32_MAX; int g_schemMode = -1; @@ -74,26 +70,38 @@ void TestLaunch() TileXRMoonEpTensorV1 down = Tensor(reinterpret_cast(UINTPTR_C(0x5000))); TileXRMoonEpTensorV1 status = Tensor(reinterpret_cast(UINTPTR_C(0x6000))); - TileXRMoonEp::ReduceGradParams params {}; - params.comm = reinterpret_cast(UINTPTR_C(0x1000)); + TileXRMoonEp::ReduceGradLaunchParams params {}; params.plan = &plan; params.gradients[0] = &gate; params.gradients[1] = &up; params.gradients[2] = &down; - params.workspace = reinterpret_cast(UINTPTR_C(0x7000)); params.status = &status; params.waitIterations = 123; params.stream = reinterpret_cast(UINTPTR_C(0x8000)); - TileXRMoonEp::ReduceGradLaunchContext context {}; + TileXRMoonEp::ReduceGradPreparedContext context {}; + context.comm = reinterpret_cast(UINTPTR_C(0x1000)); context.devArgs = reinterpret_cast(UINTPTR_C(0x9000)); + context.profileView.infoDev = reinterpret_cast(UINTPTR_C(0xa000)); + context.profileView.registryDev = reinterpret_cast(UINTPTR_C(0xb000)); + context.expertsToCopy = plan.expertsToCopy; + context.workspace = reinterpret_cast(UINTPTR_C(0x7000)); + context.workspaceBytes = 900; + context.gradients[0] = gate; + context.gradients[1] = up; + context.gradients[2] = down; + for (uint32_t projection = 0; projection < 3; ++projection) { + context.sources[projection].data = reinterpret_cast( + UINTPTR_C(0xc000) + projection * UINTPTR_C(0x1000)); + context.sources[projection].bytes = 100 + projection; + } + auto &layout = context.layout; layout.rank = 1; layout.rankSize = 4; layout.expertCount = 8; layout.expertsPerRank = 2; layout.prefetchSlots = 1; - layout.controlBlockCount = 3; layout.blockDim = 64; layout.rowElements[0] = 10; layout.rowElements[1] = 20; @@ -101,27 +109,35 @@ void TestLaunch() layout.rowBytes[0] = 40; layout.rowBytes[1] = 80; layout.rowBytes[2] = 120; - layout.transports[0] = TileXRMoonEp::kReduceGradTransportPeer; - layout.transports[1] = TileXRMoonEp::kReduceGradTransportUdma; - layout.transports[2] = TileXRMoonEp::kReduceGradTransportPeer; - layout.udmaQpCount = 3; - layout.peerRecordBaseOffset = 1; - layout.peerHalfBytes = 2; - layout.peerSlotStrideBytes = 3; - layout.peerChunkPayloadBytes = 4; - layout.udmaStateOffset = 5; - layout.udmaOutboundOffset = 6; - layout.udmaInboundOffset = 7; - layout.udmaChunkBytes = 8; - layout.workspaceBytes = 9; + layout.chunkCounts[0] = 1; + layout.chunkCounts[1] = 2; + layout.chunkCounts[2] = 3; + layout.projectionQpBase[0] = 0; + layout.projectionQpBase[1] = 3; + layout.projectionQpBase[2] = 6; + layout.projectionQpCounts[0] = 3; + layout.projectionQpCounts[1] = 3; + layout.projectionQpCounts[2] = 2; + layout.transportQpCount = 32; + for (uint32_t lane = 0; lane < 8; ++lane) { + layout.lanePhysicalQps[lane] = lane == 7 ? 16 : lane; + } + layout.qpCount = 8; + layout.laneCount = 8; + layout.laneStateBytes = 32; + layout.stagingOffset = 32; + layout.bankStrideBytes = 64; + layout.laneStrideBytes = 128; + layout.chunkBytes = 16; + layout.workspaceBytes = 900; g_binaryRegisterReturn = -10; Check(TileXRMoonEp::TileXRMoonEpLaunchReduceGradKernel(params, context) == TileXR::TILEXR_ERROR_MKIRT, "ReduceGrad binary registration error was not translated"); Check(g_binaryRegisterCalls == 1 && g_functionRegisterCalls == 0 && - g_binaryUnregisterCalls == 0 && g_launchCalls == 0 && g_magicCalls == 0, - "ReduceGrad binary registration failure did not stop launch"); + g_launchCalls == 0 && g_magicCalls == 0, + "binary registration failure did not stop launch"); g_binaryRegisterReturn = RT_ERROR_NONE; g_functionRegisterReturn = -11; @@ -130,74 +146,73 @@ void TestLaunch() "ReduceGrad function registration error was not translated"); Check(g_binaryRegisterCalls == 2 && g_functionRegisterCalls == 1 && g_binaryUnregisterCalls == 1 && g_launchCalls == 0 && g_magicCalls == 0, - "ReduceGrad function registration failure did not stop launch"); + "function registration failure did not release the binary"); Check(g_unregisteredBinaryHandle == g_binaryHandle, - "ReduceGrad function registration failure did not release the binary handle"); + "function registration released the wrong binary handle"); g_functionRegisterReturn = RT_ERROR_NONE; Check(TileXRMoonEp::TileXRMoonEpLaunchReduceGradKernel(params, context) == TileXR::TILEXR_SUCCESS, "ReduceGrad runtime launch failed"); Check(g_binaryRegisterCalls == 3 && g_functionRegisterCalls == 2 && - g_binaryUnregisterCalls == 1 && g_launchCalls == 1 && g_magicCalls == 1, - "ReduceGrad successful registration/launch counts mismatch"); - Check(!g_binaryWasNull && - g_binary.data == TileXRMoonEpReduceGradKernelBinaryData && + g_launchCalls == 1 && g_magicCalls == 1, + "successful registration/launch counts mismatch"); + Check(g_binary.data == TileXRMoonEpReduceGradKernelBinaryData && g_binary.length == TileXRMoonEpReduceGradKernelBinarySize && - g_binary.magic == RT_DEV_BINARY_MAGIC_ELF_AIVEC && g_binary.version == 0, - "ReduceGrad binary descriptor ABI mismatch"); - Check(g_functionBinaryHandle == g_binaryHandle && g_functionStub != nullptr && + g_binary.magic == RT_DEV_BINARY_MAGIC_ELF_AIVEC, + "embedded pure-AIV binary descriptor mismatch"); + Check(g_functionStub != nullptr && g_functionStubName == "tilexr_moonep_reduce_grad_kernel" && - g_functionKernelInfo == "tilexr_moonep_reduce_grad_kernel" && - g_functionMode == 0, - "ReduceGrad function registration ABI mismatch"); - Check(g_launchKernel == g_functionStub && - std::string(static_cast(g_launchKernel)) == - "tilexr_moonep_reduce_grad_kernel", - "ReduceGrad launch did not reuse the registered stub"); - Check(g_blockDim == 64 && g_stream == params.stream && g_smDesc == nullptr && - g_launchFlags == 0 && g_schemMode == 1, - "ReduceGrad runtime launch metadata mismatch"); + g_functionKernelInfo == "tilexr_moonep_reduce_grad_kernel", + "registered function signature mismatch"); + Check(g_launchKernel == g_functionStub && g_blockDim == 64 && + g_stream == params.stream && g_launchFlags == 0 && g_schemMode == 1, + "runtime launch metadata mismatch"); Check(g_argsSize == sizeof(TileXRMoonEp::ReduceGradKernelArgs), - "ReduceGrad args size ABI mismatch"); + "kernel args size ABI mismatch"); Check(g_args.commArgs == context.devArgs && - g_args.expertsToCopy == plan.expertsToCopy && g_args.gate == gate.data && - g_args.up == up.data && g_args.down == down.data && - g_args.workspace == params.workspace && g_args.status == status.data, - "ReduceGrad pointer args ABI mismatch"); + g_args.profileInfo == context.profileView.infoDev && + g_args.profileRegistry == context.profileView.registryDev && + g_args.expertsToCopy == plan.expertsToCopy && + g_args.workspace == context.workspace && g_args.status == status.data, + "owner-pull pointer args mismatch"); + Check(g_args.gradients[0] == gate.data && g_args.gradients[1] == up.data && + g_args.gradients[2] == down.data && + g_args.sources[0] == context.sources[0].data && + g_args.sources[1] == context.sources[1].data && + g_args.sources[2] == context.sources[2].data, + "gradient/source slice args mismatch"); Check(g_args.rank == 1 && g_args.rankSize == 4 && g_args.expertCount == 8 && - g_args.expertsPerRank == 2 && g_args.prefetchSlots == 1 && - g_args.controlBlockCount == 3, - "ReduceGrad rank args ABI mismatch"); - Check(g_args.gateRowElements == 10 && g_args.upRowElements == 20 && - g_args.downRowElements == 30 && g_args.gateRowBytes == 40 && - g_args.upRowBytes == 80 && g_args.downRowBytes == 120, - "ReduceGrad row args ABI mismatch"); - Check(g_args.gateTransport == TileXRMoonEp::kReduceGradTransportPeer && - g_args.upTransport == TileXRMoonEp::kReduceGradTransportUdma && - g_args.downTransport == TileXRMoonEp::kReduceGradTransportPeer && - g_args.udmaQpCount == 3, - "ReduceGrad transport args ABI mismatch"); - Check(g_args.peerRecordBaseOffset == 1 && g_args.peerHalfBytes == 2 && - g_args.peerSlotStrideBytes == 3 && g_args.peerChunkPayloadBytes == 4 && - g_args.udmaStateOffset == 5 && g_args.udmaOutboundOffset == 6 && - g_args.udmaInboundOffset == 7 && g_args.udmaChunkBytes == 8 && - g_args.workspaceBytes == 9 && g_args.waitIterations == 123 && - g_args.magic == g_magic, - "ReduceGrad offset/wait args ABI mismatch"); + g_args.expertsPerRank == 2 && g_args.prefetchSlots == 1, + "rank args mismatch"); + Check(g_args.rowElements[0] == 10 && g_args.rowElements[1] == 20 && + g_args.rowElements[2] == 30 && g_args.rowBytes[0] == 40 && + g_args.rowBytes[1] == 80 && g_args.rowBytes[2] == 120 && + g_args.chunkCounts[2] == 3, + "projection shape args mismatch"); + Check(g_args.projectionQpBase[1] == 3 && + g_args.projectionQpCounts[2] == 2 && g_args.qpCount == 8 && + g_args.laneCount == 8 && g_args.transportQpCount == 32 && + g_args.lanePhysicalQps[7] == 16, + "QP allocation args mismatch"); + Check(g_args.laneStateBytes == 32 && g_args.stagingOffset == 32 && + g_args.bankStrideBytes == 64 && g_args.laneStrideBytes == 128 && + g_args.chunkBytes == 16 && g_args.workspaceBytes == 900 && + g_args.waitIterations == 123 && g_args.magic == g_magic, + "workspace/timing args mismatch"); g_launchReturn = -12; Check(TileXRMoonEp::TileXRMoonEpLaunchReduceGradKernel(params, context) == - TileXR::TILEXR_ERROR_MKIRT, "ReduceGrad launch error was not translated"); + TileXR::TILEXR_ERROR_MKIRT, "launch error was not translated"); Check(g_binaryRegisterCalls == 3 && g_functionRegisterCalls == 2 && - g_binaryUnregisterCalls == 1 && g_launchCalls == 2 && g_magicCalls == 2, - "ReduceGrad successful registration was not cached"); + g_launchCalls == 2 && g_magicCalls == 2, + "successful kernel registration was not cached"); g_launchReturn = RT_ERROR_NONE; g_magicReturn = -91; Check(TileXRMoonEp::TileXRMoonEpLaunchReduceGradKernel(params, context) == -91, - "ReduceGrad magic error was not propagated"); + "magic error was not propagated"); Check(g_launchCalls == 2 && g_magicCalls == 3, - "ReduceGrad magic failure did not stop launch"); + "magic failure did not stop launch"); } } // namespace @@ -214,7 +229,6 @@ extern "C" int TileXRCommNextMagic(TileXRCommPtr, int64_t *magic) extern "C" rtError_t rtDevBinaryRegister(const rtDevBinary_t *binary, void **handle) { ++g_binaryRegisterCalls; - g_binaryWasNull = binary == nullptr; if (binary != nullptr) { g_binary = *binary; } @@ -231,27 +245,24 @@ extern "C" rtError_t rtDevBinaryUnRegister(void *handle) return g_binaryUnregisterReturn; } -extern "C" rtError_t rtFunctionRegister(void *binaryHandle, const void *stubFunc, - const char_t *stubName, const void *kernelInfoExt, uint32_t functionMode) +extern "C" rtError_t rtFunctionRegister(void *, const void *stubFunc, + const char_t *stubName, const void *kernelInfoExt, uint32_t) { ++g_functionRegisterCalls; - g_functionBinaryHandle = binaryHandle; g_functionStub = stubFunc; g_functionStubName = stubName == nullptr ? "" : stubName; g_functionKernelInfo = kernelInfoExt == nullptr ? "" : static_cast(kernelInfoExt); - g_functionMode = functionMode; return g_functionRegisterReturn; } extern "C" rtError_t rtKernelLaunchWithFlagV2(const void *kernel, uint32_t blockDim, - rtArgsEx_t *argsInfo, rtSmDesc_t *smDesc, rtStream_t stream, uint32_t flags, + rtArgsEx_t *argsInfo, rtSmDesc_t *, rtStream_t stream, uint32_t flags, const rtTaskCfgInfo_t *cfgInfo) { ++g_launchCalls; g_launchKernel = kernel; g_blockDim = blockDim; - g_smDesc = smDesc; g_stream = stream; g_launchFlags = flags; g_schemMode = cfgInfo == nullptr ? -1 : cfgInfo->schemMode; diff --git a/tests/moonep/unit/test_tilexr_moonep_reduce_grad_layout.cpp b/tests/moonep/unit/test_tilexr_moonep_reduce_grad_layout.cpp index a964a19d..9c7a34b0 100644 --- a/tests/moonep/unit/test_tilexr_moonep_reduce_grad_layout.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_reduce_grad_layout.cpp @@ -1,13 +1,10 @@ #include -#include #include #include #include -#include "comm_args.h" #include "reduce_grad_common.h" #include "reduce_grad_layout.h" -#include "tilexr_moonep.h" #include "tilexr_types.h" namespace { @@ -30,148 +27,155 @@ void CheckStatus(const std::string &label, int actual, int expected) } } -void TestThresholdAndMixedLayout() +void TestEqualProjectionAllocationAndWorkspace() { const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = { - (TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES - sizeof(float)) / sizeof(float), - TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES / sizeof(float), - TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES / sizeof(float) + 1, - }; + UINT64_C(14) << 20, UINT64_C(14) << 20, UINT64_C(14) << 20}; TileXRMoonEp::ReduceGradLayout layout {}; - CheckStatus("mixed layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, rows, UINT64_C(512) << 20, 0, &layout), TileXR::TILEXR_SUCCESS); - Check(layout.transports[0] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER, - "below-threshold row must use peer memory"); - Check(layout.transports[1] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER, - "exactly 1 MiB row must use peer memory"); - Check(layout.transports[2] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_UDMA, - "row above 1 MiB must use UDMA"); - Check(layout.udmaQpCount == 0, - "layout builder must leave negotiated UDMA QP count to the host"); - Check(layout.rowBytes[1] == UINT64_C(1) << 20, - "exact-threshold row byte calculation mismatch"); - Check(layout.rowBytes[2] == (UINT64_C(1) << 20) + sizeof(float), - "above-threshold row byte calculation mismatch"); - Check(layout.peerRecordBaseOffset == UINT64_C(1) << 20, - "peer state reservation mismatch"); - Check(layout.peerHalfBytes > 0 && layout.peerSlotStrideBytes >= 512 && - layout.peerChunkPayloadBytes > 0, "peer layout must provide usable records"); - Check(layout.udmaChunkBytes >= (UINT64_C(1) << 20) && - layout.udmaChunkBytes <= UINT32_MAX, "UDMA chunk bounds mismatch"); - Check(layout.udmaOutboundOffset % TileXRMoonEp::kReduceGradUdmaAlignment == 0 && - layout.udmaInboundOffset % TileXRMoonEp::kReduceGradUdmaAlignment == 0 && - layout.workspaceBytes % TileXRMoonEp::kReduceGradUdmaWorkspaceAlignment == 0, - "UDMA workspace alignment mismatch"); - Check(layout.udmaInboundOffset > layout.udmaOutboundOffset && - layout.workspaceBytes > layout.udmaInboundOffset, - "UDMA workspace regions overlap or are empty"); - Check(TileXRMoonEp::kReduceGradUdmaSignalStageStride == - TileXR::TILEXR_UDMA_CACHE_LINE_SIZE && - TileXRMoonEp::kReduceGradUdmaCompletionOffset >= - 2 * TileXR::TILEXR_UDMA_CACHE_LINE_SIZE && - TileXRMoonEp::kReduceGradUdmaPollScratchOffset >= - TileXRMoonEp::kReduceGradUdmaCompletionOffset + - 2 * TileXR::TILEXR_UDMA_CACHE_LINE_SIZE && - TileXRMoonEp::kReduceGradUdmaPeerStateBytes >= - 6 * TileXR::TILEXR_UDMA_CACHE_LINE_SIZE, - "each UDMA stage signal must use a distinct cache line"); -} + CheckStatus("8-rank layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 3, 8, 384, 48, rows, 8, 0, &layout), TileXR::TILEXR_SUCCESS); -void TestCapacityInjection() -{ - const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = {64, 128, 256}; - TileXRMoonEp::ReduceGradLayout mainLayout {}; - TileXRMoonEp::ReduceGradLayout pr90Layout {}; - CheckStatus("main capacity", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 8, 64, rows, UINT64_C(100) << 20, 0, &mainLayout), TileXR::TILEXR_SUCCESS); - CheckStatus("PR90 capacity", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 8, 64, rows, UINT64_C(512) << 20, 0, &pr90Layout), TileXR::TILEXR_SUCCESS); - Check(pr90Layout.peerHalfBytes > mainLayout.peerHalfBytes, - "injected PR90 capacity must increase the peer half"); - Check(pr90Layout.peerSlotStrideBytes > mainLayout.peerSlotStrideBytes, - "injected PR90 capacity must increase the slot stride"); - Check(TileXRMoonEp::TileXRMoonEpReduceGradPeerWindowBytes() == - static_cast(TileXR::IPC_BUFF_MAX_SIZE), - "capacity resolver must follow IPC_BUFF_MAX_SIZE"); + Check(layout.projectionQpCounts[0] == 3 && + layout.projectionQpCounts[1] == 3 && + layout.projectionQpCounts[2] == 2, + "equal projection bytes must receive stable 3/3/2 QP allocation"); + Check(layout.projectionQpBase[0] == 0 && layout.projectionQpBase[1] == 3 && + layout.projectionQpBase[2] == 6, + "projection QP bases must be contiguous"); + for (uint32_t qp = 0; qp < 8; ++qp) { + const uint32_t expected = qp < 3 ? 0 : (qp < 6 ? 1 : 2); + Check(layout.qpProjection[qp] == expected, + "QP-to-projection mapping is not stable"); + } + + Check(layout.chunkBytes == TileXRMoonEp::kReduceGradDefaultChunkBytes, + "default chunk must be 8 MiB"); + Check(layout.laneStateBytes == + 8 * TileXRMoonEp::kReduceGradLaneStateStrideBytes && + layout.stagingOffset == layout.laneStateBytes, + "lane state sizing mismatch"); + Check(layout.bankStrideBytes == UINT64_C(64) << 20 && + layout.laneStrideBytes == UINT64_C(128) << 20, + "rank-sized bank strides mismatch"); + Check(layout.workspaceBytes == UINT64_C(1026) << 20, + "8-rank owner-pull workspace mismatch"); + Check(layout.workspaceBytes % TileXRMoonEp::kReduceGradWorkspaceAlignment == 0, + "workspace must retain 2 MiB registration alignment"); + Check(layout.blockDim == TileXRMoonEp::kReduceGradMaxAivBlockCount, + "default launch must use all AIV blocks"); } -void TestPureTransportLayouts() +void TestProportionalAllocationIsStable() { - const uint64_t peerRows[TileXRMoonEp::kReduceGradProjectionCount] = {1, 2, 3}; - const uint64_t udmaRows[TileXRMoonEp::kReduceGradProjectionCount] = { - (UINT64_C(1) << 18) + 1, - (UINT64_C(1) << 18) + 2, - (UINT64_C(1) << 18) + 3, - }; + const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = { + UINT64_C(16) << 20, UINT64_C(8) << 20, UINT64_C(4) << 20}; TileXRMoonEp::ReduceGradLayout layout {}; - CheckStatus("peer without capacity", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, peerRows, 0, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); - CheckStatus("UDMA without peer capacity", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, udmaRows, 0, UINT64_C(2) << 20, &layout), TileXR::TILEXR_SUCCESS); - Check(layout.peerHalfBytes == 0 && layout.workspaceBytes > 0, - "UDMA-only layout must not reserve peer records"); - Check(layout.controlBlockCount == 1, - "two-rank UDMA layout must assign one control block"); + CheckStatus("weighted layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 16, 384, 24, rows, 8, UINT64_C(2) << 20, &layout), + TileXR::TILEXR_SUCCESS); + Check(layout.projectionQpCounts[0] == 5 && + layout.projectionQpCounts[1] == 2 && + layout.projectionQpCounts[2] == 1, + "weighted QP allocation must follow projection bytes with stable ties"); + Check(layout.bankStrideBytes == UINT64_C(32) << 20 && + layout.laneStrideBytes == UINT64_C(64) << 20 && + layout.workspaceBytes == UINT64_C(514) << 20, + "16-rank 2 MiB workspace arithmetic mismatch"); + Check(layout.chunkCounts[0] == 32 && layout.chunkCounts[1] == 16 && + layout.chunkCounts[2] == 8, + "projection chunk counts mismatch"); } -void TestSingleRankLargeRowsStayLocal() +void TestSharedDomainUsesBoundedActiveLanes() { - const uint64_t largeRow = - TILEXR_MOONEP_REDUCE_GRAD_UDMA_THRESHOLD_BYTES / sizeof(float) + 1; const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = { - largeRow, largeRow + 1, largeRow + 2}; + UINT64_C(14) << 20, UINT64_C(14) << 20, UINT64_C(14) << 20}; TileXRMoonEp::ReduceGradLayout layout {}; - CheckStatus("single-rank large rows", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 1, 4, rows, UINT64_C(100) << 20, 0, &layout), TileXR::TILEXR_SUCCESS); - Check(layout.transports[0] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - layout.transports[1] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER && - layout.transports[2] == TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER, - "single-rank rows must stay on the local peer path regardless of row size"); - Check(layout.workspaceBytes == 0 && layout.udmaChunkBytes == 0, - "single-rank rows must not request a UDMA workspace"); - Check(layout.controlBlockCount == 0 && layout.peerChunkPayloadBytes > 0, - "single-rank layout must reserve all blocks for local reduction"); + CheckStatus("32-QP shared-domain layout", + TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 4, 4, 1, rows, 32, UINT64_C(8) << 20, &layout), + TileXR::TILEXR_SUCCESS); + Check(layout.transportQpCount == 32 && layout.qpCount == 3 && + layout.laneCount == 3, + "shared domain must expose all transport QPs and use three active lanes"); + Check(layout.lanePhysicalQps[0] == 0 && layout.lanePhysicalQps[1] == 1 && + layout.lanePhysicalQps[2] == 16, + "shared-domain lanes must preserve the measured 6/6/2 route mapping"); + Check(layout.projectionQpCounts[0] == 1 && + layout.projectionQpCounts[1] == 1 && + layout.projectionQpCounts[2] == 1, + "32-QP shared domain must allocate one active lane per projection"); + Check(layout.workspaceBytes == UINT64_C(194) << 20, + "inactive shared-domain QPs must not increase ReduceGrad workspace"); } -void TestLargeRankScheduling() +void TestMinimumRankCount() { const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = {64, 128, 256}; TileXRMoonEp::ReduceGradLayout layout {}; - CheckStatus("128-rank layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 128, 128, rows, UINT64_C(512) << 20, 0, &layout), TileXR::TILEXR_SUCCESS); - Check(layout.blockDim == TileXRMoonEp::kReduceGradMaxAivBlockCount, - "128-rank layout must use all available AIV blocks"); - Check(layout.controlBlockCount == TileXRMoonEp::kReduceGradMaxAivBlockCount - 1, - "128-rank layout must reserve one AIV block for receiving"); + CheckStatus("one-rank layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 1, 8, 8, rows, 3, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("three-rank layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 3, 6, 2, rows, 3, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("four-rank layout", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 4, 8, 2, rows, 3, 0, &layout), TileXR::TILEXR_SUCCESS); + Check(layout.bankStrideBytes == UINT64_C(32) << 20 && + layout.laneStrideBytes == UINT64_C(64) << 20 && + layout.workspaceBytes == UINT64_C(194) << 20, + "4-rank workspace arithmetic mismatch"); } -void TestInvalidInputs() +void TestSlotsMayExceedExpertsPerRank() +{ + const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = {1024, 1024, 1024}; + TileXRMoonEp::ReduceGradLayout layout {}; + CheckStatus("native dedicated-suite layout", + TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 64, 14, rows, 8, 0, &layout), TileXR::TILEXR_SUCCESS); + Check(layout.expertsPerRank == 8 && layout.prefetchSlots == 14, + "ReduceGrad must support source slots independently of local expert count"); +} + +void TestChunkAlignment() +{ + const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = {1024, 2048, 4096}; + TileXRMoonEp::ReduceGradLayout layout {}; + const uint64_t requested = (UINT64_C(2) << 20) + 1; + CheckStatus("unaligned chunk", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 64, 8, rows, 3, requested, &layout), TileXR::TILEXR_SUCCESS); + Check(layout.chunkBytes == requested + 511, + "chunk must align up to the 512-byte UDMA boundary"); +} + +void TestInvalidInputsAndOverflow() { const uint64_t rows[TileXRMoonEp::kReduceGradProjectionCount] = {1, 1, 1}; TileXRMoonEp::ReduceGradLayout layout {}; CheckStatus("null output", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, rows, UINT64_C(100) << 20, 0, nullptr), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + 0, 8, 16, 2, rows, 3, 0, nullptr), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); CheckStatus("bad rank", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 2, 2, 4, rows, UINT64_C(100) << 20, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + 8, 8, 16, 2, rows, 3, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); CheckStatus("nondivisible experts", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 3, rows, UINT64_C(100) << 20, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + 0, 8, 17, 2, rows, 3, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("too few QPs", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 16, 2, rows, 2, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("too many transport QPs", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 16, 2, rows, 33, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("oversized chunk", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 16, 2, rows, 3, UINT64_C(1) << 32, &layout), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("aligned chunk overflow", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 16, 2, rows, 3, UINT32_MAX, &layout), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckStatus("contributor index overflow", + TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( + 0, 8, 64, std::numeric_limits::max() / 8 + 1, + rows, 3, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); uint64_t overflowRows[TileXRMoonEp::kReduceGradProjectionCount] = { std::numeric_limits::max(), 1, 1}; CheckStatus("row byte overflow", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, overflowRows, UINT64_C(100) << 20, 0, &layout), - TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); - - const uint64_t udmaRows[TileXRMoonEp::kReduceGradProjectionCount] = { - (UINT64_C(1) << 18) + 1, - (UINT64_C(1) << 18) + 1, - (UINT64_C(1) << 18) + 1, - }; - CheckStatus("small UDMA chunk", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, udmaRows, 0, (UINT64_C(1) << 20) - 1, &layout), - TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); - CheckStatus("oversized UDMA chunk", TileXRMoonEp::TileXRMoonEpBuildReduceGradLayout( - 0, 2, 4, udmaRows, 0, UINT64_C(1) << 32, &layout), + 0, 8, 16, 2, overflowRows, 3, 0, &layout), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); } @@ -179,11 +183,12 @@ void TestInvalidInputs() int main() { - TestThresholdAndMixedLayout(); - TestCapacityInjection(); - TestPureTransportLayouts(); - TestSingleRankLargeRowsStayLocal(); - TestLargeRankScheduling(); - TestInvalidInputs(); + TestEqualProjectionAllocationAndWorkspace(); + TestProportionalAllocationIsStable(); + TestSharedDomainUsesBoundedActiveLanes(); + TestMinimumRankCount(); + TestSlotsMayExceedExpertsPerRank(); + TestChunkAlignment(); + TestInvalidInputsAndOverflow(); return g_failures == 0 ? 0 : 1; } diff --git a/tests/moonep/unit/test_tilexr_moonep_sources.cpp b/tests/moonep/unit/test_tilexr_moonep_sources.cpp index fcb875c6..aa1f8298 100644 --- a/tests/moonep/unit/test_tilexr_moonep_sources.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_sources.cpp @@ -230,6 +230,8 @@ int main() Contains("prefetch Host", prefetchHost, "TileXRGetUDMARegistryHost"); Contains("ReduceGrad Host", reduceHost, "TileXRMoonEpPrepareReduceGradLayout"); Contains("ReduceGrad Host", reduceHost, "TileXRMoonEpReduceGradV2"); + Contains("ReduceGrad launch-failure drain", reduceHost, + "aclrtSynchronizeStream(stream)"); Excludes("ReduceGrad Host", reduceHost, "moonep_stage_host.h"); Contains("test CMake", testCmake, "if(TARGET tilexr-moonep)"); @@ -241,7 +243,11 @@ int main() Contains("flow demo", flowDemo, "TileXRUDMAUnregister prefetch arena"); Contains("flow demo", flowDemo, "TileXRMoonEpGetCapabilitiesV2"); Contains("flow demo", flowDemo, "TileXRMoonEpReduceGradGetWorkspaceSizeV2"); - Contains("flow demo", flowDemo, "TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER"); + Contains("flow demo", flowDemo, "TileXRMoonEpReduceGradPrepareV2"); + Contains("flow demo", flowDemo, "TileXRMoonEpReduceGradDestroyPreparedV2"); + Contains("flow demo", flowDemo, "options.world < kMinReduceGradRankCount"); + Excludes("flow demo", flowDemo, ".transports"); + Excludes("flow demo", flowDemo, "TILEXR_MOONEP_REDUCE_GRAD_TRANSPORT_PEER"); Excludes("flow demo", flowDemo, "TileXRMoonEpReduceGradV1(&reduceGrad"); Contains("flow demo", flowDemo, "nvS != static_cast(routeCount)"); Contains("flow demo", flowDemo, "planning_status="); @@ -265,6 +271,7 @@ int main() CheckOrdered("native flow", flowDemo, { "TileXRMoonEpPlanningV1(&planning", "TileXRMoonEpReduceGradGetWorkspaceSizeV2(", + "TileXRMoonEpReduceGradPrepareV2(", "TileXRMoonEpDispatchV1(&forwardDispatch", "TileXRMoonEpPrefetchWeightV1(&prefetch", "TileXRMoonEpCombineV1(&forwardCombine", @@ -280,9 +287,12 @@ int main() Contains("flow demo", flowDemo, "TILEXR_MOONEP_FLOW_BARRIER_ADDR"); Contains("flow runner", flowRunner, "block_dim=$((64 / ranks_per_device))"); Contains("flow runner", flowRunner, "export TILEXR_ENABLE_UDMA=1"); + Contains("flow runner", flowRunner, + "TILEXR_UDMA_QP_ROUTE_SPEC=port_count:6,port_count:6,port_count:2"); Contains("flow runner", flowRunner, "${SCRIPT_DIR}/tilexr_moonep_flow_demo"); Contains("flow runner", flowRunner, "device=$((rank % physical_device_count))"); Contains("flow runner", flowRunner, "TILEXR_MOONEP_PLANNER_BLOCK_DIM"); + Contains("flow runner", flowRunner, "rank_size < 4"); Contains("flow runner", flowRunner, "torch_validated=false"); Contains("flow runner", flowRunner, "transport_performance_valid=false"); Excludes("flow demo", Lower(flowDemo), "hccl"); diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index dc214125..e809eff2 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -228,6 +228,93 @@ if(BUILD_TILEXR_UDMA_DEMO) -I${TILEXR_ROOT}/3rdparty -I${TILEXR_ROOT}/src/include ) + + if(TILEXR_UDMA_DEMO_SOC_TYPE STREQUAL "Ascend950") + find_program(TILEXR_UDMA_CCE_LINKER + NAMES ld.lld cce-ld + HINTS + "${ASCEND_HOME_PATH}/bin" + "${ASCEND_HOME_PATH}/tools/bisheng_compiler/bin" + ) + if(NOT TILEXR_UDMA_CCE_LINKER) + message(FATAL_ERROR "CCE linker not found in the selected CANN installation") + endif() + set(TILEXR_UDMA_PROFILE_PROBE_KERNEL_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_profile_probe_kernel.cpp") + set(TILEXR_UDMA_PROFILE_PROBE_KERNEL_RELOCATABLE + "${CMAKE_CURRENT_BINARY_DIR}/tilexr_udma_profile_probe_kernel_rel.o") + set(TILEXR_UDMA_PROFILE_PROBE_KERNEL_BINARY + "${CMAKE_CURRENT_BINARY_DIR}/tilexr_udma_profile_probe_kernel.o") + set(TILEXR_UDMA_PROFILE_PROBE_KERNEL_EMBED + "${CMAKE_CURRENT_BINARY_DIR}/tilexr_udma_profile_probe_kernel_embed.cpp") + set(TILEXR_UDMA_PROFILE_PROBE_KERNEL_INCLUDES + ${TILEXR_UDMA_DEMO_KERNEL_INCLUDES} + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include + -I${ASCEND_HOME_PATH}/${ARCH}-linux/ascendc/include + ) + add_custom_command( + OUTPUT "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_BINARY}" + COMMAND ${BISHENG_EXECUTABLE} + -xcce + -mllvm -cce-aicore-stack-size=0x8000 + -mllvm -cce-aicore-function-stack-size=0x8000 + -mllvm -cce-aicore-record-overflow=true + -mllvm -cce-aicore-addr-transform + -mllvm -cce-aicore-dcci-insert-for-scalar=false + --cce-aicore-arch=dav-c310-vec + --cce-aicore-only + -std=gnu++17 + ${TILEXR_UDMA_PROFILE_PROBE_KERNEL_INCLUDES} + -c "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_SOURCE}" + -o "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_RELOCATABLE}" + COMMAND ${TILEXR_UDMA_CCE_LINKER} + -m aicorelinux -Ttext=0 + "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_RELOCATABLE}" + --static --allow-multiple-definition + -o "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_BINARY}" + DEPENDS + "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_SOURCE}" + "${TILEXR_ROOT}/src/include/tilexr_udma.h" + "${TILEXR_ROOT}/src/include/tilexr_udma_reg.h" + VERBATIM + COMMENT "Building pure AICore TileXR UDMA profile probe kernel" + ) + add_custom_target(tilexr_udma_profile_probe_kernel ALL + DEPENDS "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_BINARY}") + add_custom_command( + OUTPUT "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_EMBED}" + COMMAND ${CMAKE_COMMAND} + -DTILEXR_MOONEP_KERNEL_BINARY=${TILEXR_UDMA_PROFILE_PROBE_KERNEL_BINARY} + -DTILEXR_MOONEP_KERNEL_EMBED_CPP=${TILEXR_UDMA_PROFILE_PROBE_KERNEL_EMBED} + -DTILEXR_MOONEP_KERNEL_DATA_SYMBOL=TileXRUDMAProfileProbeKernelBinaryData + -DTILEXR_MOONEP_KERNEL_SIZE_SYMBOL=TileXRUDMAProfileProbeKernelBinarySize + -P "${TILEXR_ROOT}/src/moonep/cmake/embed_moonep_kernel.cmake" + DEPENDS + "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_BINARY}" + "${TILEXR_ROOT}/src/moonep/cmake/embed_moonep_kernel.cmake" + VERBATIM + COMMENT "Embedding pure AICore TileXR UDMA profile probe kernel" + ) + + add_executable(tilexr_udma_profile_probe + demo/tilexr_udma_profile_probe.cpp + "${TILEXR_UDMA_PROFILE_PROBE_KERNEL_EMBED}" + ) + add_dependencies(tilexr_udma_profile_probe tilexr_udma_profile_probe_kernel) + target_link_libraries(tilexr_udma_profile_probe + ${TILEXR_LIB} + ascendcl + runtime + ascend_hal + ) + set_target_properties(tilexr_udma_profile_probe PROPERTIES + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH "$ORIGIN/../lib" + ) + list(APPEND INSTALL_TARGETS tilexr_udma_profile_probe) + message(STATUS "TileXR UDMA persistent-profile hardware probe enabled") + endif() + add_custom_command( OUTPUT "${TILEXR_UDMA_DEMO_KERNEL_SO}" COMMAND ${BISHENG_EXECUTABLE} diff --git a/tests/udma/build.sh b/tests/udma/build.sh index 02177a27..3fe07755 100755 --- a/tests/udma/build.sh +++ b/tests/udma/build.sh @@ -69,9 +69,14 @@ if [ -f "${INSTALL_DIR}/bin/tilexr_udma_demo" ]; then else echo " - tilexr_udma_demo : skipped (requires bisheng/AICore toolchain)" fi +if [ -f "${INSTALL_DIR}/bin/tilexr_udma_profile_probe" ]; then + echo " - tilexr_udma_profile_probe : persistent multi-MR deferred-READ hardware probe" +fi echo "" echo "Run tests with:" echo " bash run_tests.sh" echo "Run demo with:" echo " bash demo/run_tilexr_udma_demo.sh 0 2 16" +echo "Run the two-host profile probe with:" +echo " bash demo/run_tilexr_udma_profile_probe_mpi.sh --hosts host1:1,host2:1 --comm-id host1_data_ip:10067" echo "==========================================" diff --git a/tests/udma/demo/README.md b/tests/udma/demo/README.md index 5615d0a4..37237bd6 100644 --- a/tests/udma/demo/README.md +++ b/tests/udma/demo/README.md @@ -2,6 +2,32 @@ This demo shows TileXR-initialized UDMA communication with verbose diagnostics. UDMA means UnifiedBus DMA and this runtime path currently targets A5 / Ascend950 / 950 hardware. The host demo uses TileXR public APIs, registers ordinary `aclrtMalloc` memory with `TileXRUDMARegister`, and the AICore kernel uses `tilexr_udma.h`. +## Persistent Profile Hardware Probe + +`tilexr_udma_profile_probe` is the Ascend950 stop-gate for persistent multi-MR +profiles. It registers one legacy region alongside four 2 MiB-aligned profile +regions, binds every QP to local staging and one of three remote sources, and +issues four deferred READ WQEs per batch. The per-WQE sizes are 48 KiB, 256 KiB, +1 MiB, 2 MiB, 4 MiB, 8 MiB, and 16 MiB. + +The probe kernel is a pure AICore binary. The Host executable registers it with +`rtDevBinaryRegister` and `rtFunctionRegister`, then launches it with +`rtKernelLaunchWithFlagV2`. + +```bash +cd tests/udma +bash build.sh +bash demo/run_tilexr_udma_profile_probe_mpi.sh \ + --hosts 141.61.53.106:1,141.61.53.110:1 \ + --comm-id 172.27.12.106:10167 +``` + +The runner retains the combined log and JSONL timing records under +`tests/udma/logs/`. `wqe_bytes` is the size of one UDMA copy and `batch_bytes` +is the four-WQE aggregate. The launcher uses MPI when `mpirun` is available and +otherwise starts one bounded SSH process on each host; use `--launcher` to +select either mode explicitly. + ## Build ```bash diff --git a/tests/udma/demo/run_tilexr_udma_profile_probe_mpi.sh b/tests/udma/demo/run_tilexr_udma_profile_probe_mpi.sh new file mode 100644 index 00000000..a005eccd --- /dev/null +++ b/tests/udma/demo/run_tilexr_udma_profile_probe_mpi.sh @@ -0,0 +1,181 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +UDMA_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +TILEXR_ROOT="$(cd "${UDMA_DIR}/../.." && pwd)" +INSTALL_DIR="${UDMA_DIR}/install" + +HOSTS="" +RANK_SIZE=2 +COMM_ID="" +MPI_HOME="${MPI_HOME:-/usr/local/mpich}" +LAUNCHER="auto" +TIMEOUT_SECONDS=600 +WARMUP="${TILEXR_PROFILE_PROBE_WARMUP:-2}" +ITERATIONS="${TILEXR_PROFILE_PROBE_ITERATIONS:-5}" +OUTPUT_DIR="${UDMA_DIR}/logs/tilexr_udma_profile_probe_$(date +%Y%m%d_%H%M%S)" + +usage() { + cat < Two hosts, one rank per host + --comm-id TileXR rendezvous address on rank 0 + --launcher auto, mpi, or ssh (default: auto) + --mpi-home MPI installation (default: /usr/local/mpich) + --warmup Warmup launches per size/QP/mode (default: 2) + --iterations Timed launches per size/QP/mode (default: 5) + --output-dir Log and JSONL destination + --timeout Complete-run timeout (default: 600) +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --hosts) HOSTS="$2"; shift 2 ;; + --comm-id) COMM_ID="$2"; shift 2 ;; + --mpi-home) MPI_HOME="$2"; shift 2 ;; + --launcher) LAUNCHER="$2"; shift 2 ;; + --warmup) WARMUP="$2"; shift 2 ;; + --iterations) ITERATIONS="$2"; shift 2 ;; + --output-dir) OUTPUT_DIR="$2"; shift 2 ;; + --timeout) TIMEOUT_SECONDS="$2"; shift 2 ;; + --help|-h) usage; exit 0 ;; + *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;; + esac +done + +if [[ -z "${HOSTS}" || -z "${COMM_ID}" ]]; then + usage >&2 + exit 2 +fi +if [[ ! "${WARMUP}" =~ ^[0-9]+$ || ! "${ITERATIONS}" =~ ^[1-9][0-9]*$ || + ! "${TIMEOUT_SECONDS}" =~ ^[1-9][0-9]*$ ]]; then + echo "warmup must be non-negative; iterations and timeout must be positive" >&2 + exit 2 +fi +if [[ ! "${LAUNCHER}" =~ ^(auto|mpi|ssh)$ ]]; then + echo "launcher must be auto, mpi, or ssh" >&2 + exit 2 +fi +if [[ ! "${COMM_ID}" =~ ^[0-9]+(\.[0-9]+){3}:[1-9][0-9]*$ ]]; then + echo "comm-id must be an IPv4 address and port" >&2 + exit 2 +fi + +: "${ASCEND_HOME_PATH:=}" +: "${LD_LIBRARY_PATH:=}" +source "${TILEXR_ROOT}/scripts/common_env.sh" + +bin="${TILEXR_PROFILE_PROBE_BIN:-${INSTALL_DIR}/bin/tilexr_udma_profile_probe}" +if [[ ! -x "${bin}" ]]; then + echo "Missing profile probe binary: ${bin}" >&2 + echo "Build it with: cd ${UDMA_DIR} && bash build.sh" >&2 + exit 1 +fi + +mpi_bin="${MPI_HOME}/bin/mpirun" +if [[ ! -x "${mpi_bin}" ]]; then + mpi_bin="$(command -v mpirun || true)" +fi +if [[ "${LAUNCHER}" == "auto" ]]; then + if [[ -n "${mpi_bin}" && -x "${mpi_bin}" ]]; then + LAUNCHER="mpi" + else + LAUNCHER="ssh" + fi +fi +if [[ "${LAUNCHER}" == "mpi" && ( -z "${mpi_bin}" || ! -x "${mpi_bin}" ) ]]; then + echo "mpirun not found; set --mpi-home" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" +combined_log="${OUTPUT_DIR}/combined.log" +jsonl="${OUTPUT_DIR}/timings.jsonl" + +export TILEXR_COMM_ID="${COMM_ID}" +export TILEXR_UDMA_QP_ROUTE_SPEC="port_count:6,port_count:6,port_count:2" +export TILEXR_PROFILE_PROBE_WARMUP="${WARMUP}" +export TILEXR_PROFILE_PROBE_ITERATIONS="${ITERATIONS}" +export TILEXR_PROFILE_PROBE_DEVICE_BASE=0 +export TILEXR_ENABLE_IPC=0 +export TILEXR_ENABLE_SDMA=0 +export LD_LIBRARY_PATH="${TILEXR_ROOT}/install/lib64:${TILEXR_ROOT}/install/lib:${INSTALL_DIR}/lib64:${INSTALL_DIR}/lib:${LD_LIBRARY_PATH:-}" + +if [[ "${LAUNCHER}" == "mpi" ]]; then + set +e + timeout --signal=TERM --kill-after=30 "${TIMEOUT_SECONDS}" \ + "${mpi_bin}" -hosts "${HOSTS}" -n "${RANK_SIZE}" \ + -genv TILEXR_COMM_ID "${TILEXR_COMM_ID}" \ + -genv TILEXR_UDMA_QP_ROUTE_SPEC "${TILEXR_UDMA_QP_ROUTE_SPEC}" \ + -genv TILEXR_PROFILE_PROBE_WARMUP "${TILEXR_PROFILE_PROBE_WARMUP}" \ + -genv TILEXR_PROFILE_PROBE_ITERATIONS "${TILEXR_PROFILE_PROBE_ITERATIONS}" \ + -genv TILEXR_PROFILE_PROBE_DEVICE_BASE "${TILEXR_PROFILE_PROBE_DEVICE_BASE}" \ + -genv TILEXR_ENABLE_IPC "${TILEXR_ENABLE_IPC}" \ + -genv TILEXR_ENABLE_SDMA "${TILEXR_ENABLE_SDMA}" \ + -genv LD_LIBRARY_PATH "${LD_LIBRARY_PATH}" \ + "${bin}" 2>&1 | tee "${combined_log}" + run_status=${PIPESTATUS[0]} + set -e +else + IFS=',' read -r host0_spec host1_spec extra_host <<< "${HOSTS}" + if [[ -n "${extra_host:-}" || ! "${host0_spec:-}" =~ ^[A-Za-z0-9._-]+:1$ || + ! "${host1_spec:-}" =~ ^[A-Za-z0-9._-]+:1$ ]]; then + echo "ssh launcher requires --hosts host1:1,host2:1" >&2 + exit 2 + fi + host0="${host0_spec%:1}" + host1="${host1_spec%:1}" + for value in "${TILEXR_ROOT}" "${ASCEND_HOME_PATH}" "${LD_LIBRARY_PATH}" "${bin}"; do + if [[ ! "${value}" =~ ^[-A-Za-z0-9_./:]+$ ]]; then + echo "ssh launcher paths must not contain shell metacharacters: ${value}" >&2 + exit 2 + fi + done + + run_ssh_rank() { + local host="$1" + local rank="$2" + ssh -o BatchMode=yes "root@${host}" \ + "cd ${TILEXR_ROOT}; source ${ASCEND_HOME_PATH}/set_env.sh; \ + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}; \ + export TILEXR_COMM_ID=${TILEXR_COMM_ID}; \ + export TILEXR_UDMA_QP_ROUTE_SPEC=${TILEXR_UDMA_QP_ROUTE_SPEC}; \ + export TILEXR_PROFILE_PROBE_WARMUP=${TILEXR_PROFILE_PROBE_WARMUP}; \ + export TILEXR_PROFILE_PROBE_ITERATIONS=${TILEXR_PROFILE_PROBE_ITERATIONS}; \ + export TILEXR_PROFILE_PROBE_DEVICE_BASE=${TILEXR_PROFILE_PROBE_DEVICE_BASE}; \ + export TILEXR_ENABLE_IPC=${TILEXR_ENABLE_IPC}; \ + export TILEXR_ENABLE_SDMA=${TILEXR_ENABLE_SDMA}; \ + export RANK=${rank}; export RANK_SIZE=${RANK_SIZE}; \ + exec timeout --signal=TERM --kill-after=30 ${TIMEOUT_SECONDS} ${bin}" + } + + set +e + run_ssh_rank "${host0}" 0 > "${OUTPUT_DIR}/rank0.log" 2>&1 & + pid0=$! + run_ssh_rank "${host1}" 1 > "${OUTPUT_DIR}/rank1.log" 2>&1 & + pid1=$! + wait "${pid0}" + status0=$? + wait "${pid1}" + status1=$? + cat "${OUTPUT_DIR}/rank0.log" "${OUTPUT_DIR}/rank1.log" | tee "${combined_log}" + run_status=$((status0 != 0 || status1 != 0)) + set -e +fi + +sed -n 's/^TILEXR_UDMA_PROFILE_PROBE_JSON //p' "${combined_log}" > "${jsonl}" +success_count=$(grep -c 'TileXR UDMA profile probe success' "${combined_log}" || true) +if [[ "${run_status}" -ne 0 || "${success_count}" -ne "${RANK_SIZE}" ]]; then + echo "Profile probe failed: launcher=${LAUNCHER}, status=${run_status}, success_ranks=${success_count}/${RANK_SIZE}" >&2 + echo "Log: ${combined_log}" >&2 + exit 1 +fi + +echo "Profile probe passed on ${RANK_SIZE} ranks" +echo "Combined log: ${combined_log}" +echo "Machine-readable results: ${jsonl}" diff --git a/tests/udma/demo/tilexr_udma_profile_probe.cpp b/tests/udma/demo/tilexr_udma_profile_probe.cpp new file mode 100644 index 00000000..b2fe852a --- /dev/null +++ b/tests/udma/demo/tilexr_udma_profile_probe.cpp @@ -0,0 +1,767 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "acl/acl.h" +#include "runtime/kernel.h" +#include "tilexr_api.h" +#include "tilexr_types.h" + +extern "C" { +extern const unsigned char TileXRUDMAProfileProbeKernelBinaryData[]; +extern const std::size_t TileXRUDMAProfileProbeKernelBinarySize; +} + +namespace { + +constexpr uintptr_t kKernelSignatureValue = UINT64_C(0x5458505200001000); +constexpr const char* kKernelName = "tilexr_udma_profile_probe_kernel"; +constexpr uint64_t kProbeMagic = UINT64_C(0x5458505250524f42); +constexpr size_t kAlignment = 2U * 1024U * 1024U; +constexpr uint64_t kMaxTransferBytes = 16U * 1024U * 1024U; +constexpr uint32_t kBatchCount = 4U; +constexpr size_t kProfileRegionBytes = + static_cast(kMaxTransferBytes) * kBatchCount; +constexpr size_t kLegacyTransferBytes = 256U * 1024U; +constexpr size_t kLegacyRegionBytes = 2U * kLegacyTransferBytes; +constexpr size_t kLegacyDestinationOffset = kLegacyTransferBytes; +constexpr uint32_t kProfileRegionCount = 4U; +constexpr uint32_t kMinimumQpCount = 3U; +constexpr uint32_t kStatusWords = 11U; +constexpr uint32_t kModeProfileTransfer = 0U; +constexpr uint32_t kModeProfileConsume = 1U; +constexpr uint32_t kModeLegacyGet = 2U; +constexpr uint32_t kLegacyPatternRegion = 7U; + +constexpr uint64_t kTransferSizes[] = { + 48U * 1024U, + 256U * 1024U, + 1U * 1024U * 1024U, + 2U * 1024U * 1024U, + 4U * 1024U * 1024U, + 8U * 1024U * 1024U, + 16U * 1024U * 1024U, +}; + +struct DeviceRegion { + void* allocation = nullptr; + GM_ADDR base = nullptr; + size_t bytes = 0; +}; + +struct KernelRegistration { + void* binaryHandle = nullptr; + bool registered = false; +}; + +int GetEnvInt(const char* name, int defaultValue) +{ + const char* value = std::getenv(name); + return value == nullptr ? defaultValue : std::atoi(value); +} + +uint64_t GetEnvUint64(const char* name, uint64_t defaultValue) +{ + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return defaultValue; + } + char* end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 0); + return end != value && end != nullptr && end[0] == '\0' ? + static_cast(parsed) : defaultValue; +} + +int GetRank() +{ + const char* names[] = { + "PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK", "RANK"}; + for (const char* name : names) { + const char* value = std::getenv(name); + if (value != nullptr) { + return std::atoi(value); + } + } + return 0; +} + +int GetRankSize() +{ + const char* names[] = { + "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "RANK_SIZE"}; + for (const char* name : names) { + const char* value = std::getenv(name); + if (value != nullptr) { + return std::atoi(value); + } + } + return 1; +} + +int GetLocalRank() +{ + const char* names[] = { + "MPI_LOCALRANKID", "OMPI_COMM_WORLD_LOCAL_RANK", "MV2_COMM_WORLD_LOCAL_RANK", + "LOCAL_RANK"}; + for (const char* name : names) { + const char* value = std::getenv(name); + if (value != nullptr) { + return std::atoi(value); + } + } + return 0; +} + +void Log(int rank, const std::string& message) +{ + std::cout << "[rank " << rank << "] " << message << std::endl; +} + +bool CheckAcl(int rank, const char* step, aclError ret) +{ + if (ret == ACL_SUCCESS) { + return true; + } + std::cerr << "[rank " << rank << "] ERROR: " << step + << " failed, ret=" << ret << std::endl; + return false; +} + +bool CheckTileXR(int rank, const char* step, int ret) +{ + if (ret == TileXR::TILEXR_SUCCESS) { + return true; + } + std::cerr << "[rank " << rank << "] ERROR: " << step + << " failed, ret=" << ret << std::endl; + return false; +} + +bool AllocateAlignedRegion(int rank, size_t bytes, DeviceRegion& region) +{ + if (bytes == 0 || bytes > std::numeric_limits::max() - (kAlignment - 1U)) { + return false; + } + const size_t allocationBytes = bytes + kAlignment - 1U; + if (!CheckAcl(rank, "aclrtMalloc aligned region", + aclrtMalloc(®ion.allocation, allocationBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + return false; + } + const uintptr_t raw = reinterpret_cast(region.allocation); + const uintptr_t aligned = (raw + kAlignment - 1U) & ~(static_cast(kAlignment) - 1U); + region.base = reinterpret_cast(aligned); + region.bytes = bytes; + if ((aligned % kAlignment) != 0U || aligned < raw || + aligned - raw > allocationBytes - bytes) { + std::cerr << "[rank " << rank << "] ERROR: failed to derive a 2 MiB-aligned device region" + << std::endl; + return false; + } + return true; +} + +void FreeRegion(DeviceRegion& region) +{ + if (region.allocation != nullptr) { + (void)aclrtFree(region.allocation); + } + region = {}; +} + +uint8_t PatternByte(int rank, uint32_t region, uint64_t offset) +{ + const uint64_t mixed = offset * UINT64_C(1315423911) + + static_cast(rank + 1) * UINT64_C(2654435761) + + static_cast(region + 1U) * UINT64_C(2246822519); + return static_cast((mixed ^ (mixed >> 17U) ^ (mixed >> 31U)) & 0xFFU); +} + +bool InitializePattern(int rank, const DeviceRegion& region, uint32_t patternRegion, + size_t patternBytes, size_t offset = 0U) +{ + if (offset > region.bytes || patternBytes > region.bytes - offset) { + return false; + } + std::vector host(patternBytes); + for (size_t index = 0; index < patternBytes; ++index) { + host[index] = PatternByte(rank, patternRegion, index); + } + return CheckAcl(rank, "aclrtMemcpy pattern H2D", + aclrtMemcpy(region.base + offset, region.bytes - offset, + host.data(), host.size(), ACL_MEMCPY_HOST_TO_DEVICE)); +} + +bool ValidatePattern(int rank, const char* label, GM_ADDR device, size_t bytes, + int sourceRank, uint32_t sourceRegion, uint64_t sourceOffset = 0U) +{ + std::vector host(bytes); + if (!CheckAcl(rank, "aclrtMemcpy result D2H", + aclrtMemcpy(host.data(), host.size(), device, bytes, ACL_MEMCPY_DEVICE_TO_HOST))) { + return false; + } + for (size_t index = 0; index < bytes; ++index) { + const uint8_t expected = PatternByte( + sourceRank, sourceRegion, sourceOffset + index); + if (host[index] != expected) { + std::cerr << "[rank " << rank << "] ERROR: " << label + << " byte mismatch at offset=" << index + << " got=" << static_cast(host[index]) + << " expected=" << static_cast(expected) << std::endl; + return false; + } + } + return true; +} + +const void* KernelSignature() +{ + return reinterpret_cast(kKernelSignatureValue); +} + +bool RegisterKernel(int rank, KernelRegistration& registration) +{ + if (registration.registered) { + return true; + } + rtDevBinary_t binary {}; + binary.magic = RT_DEV_BINARY_MAGIC_ELF_AIVEC; + binary.version = 0; + binary.data = TileXRUDMAProfileProbeKernelBinaryData; + binary.length = static_cast(TileXRUDMAProfileProbeKernelBinarySize); + rtError_t ret = rtDevBinaryRegister(&binary, ®istration.binaryHandle); + if (ret != RT_ERROR_NONE) { + std::cerr << "[rank " << rank << "] ERROR: rtDevBinaryRegister failed, ret=" + << ret << std::endl; + return false; + } + ret = rtFunctionRegister(registration.binaryHandle, KernelSignature(), + kKernelName, kKernelName, 0); + if (ret != RT_ERROR_NONE) { + std::cerr << "[rank " << rank << "] ERROR: rtFunctionRegister failed, ret=" + << ret << std::endl; + (void)rtDevBinaryUnRegister(registration.binaryHandle); + registration.binaryHandle = nullptr; + return false; + } + registration.registered = true; + return true; +} + +void UnregisterKernel(KernelRegistration& registration) +{ + if (registration.binaryHandle != nullptr) { + (void)rtDevBinaryUnRegister(registration.binaryHandle); + } + registration = {}; +} + +bool LaunchKernel(int rank, KernelRegistration& registration, aclrtStream stream, + GM_ADDR commArgs, GM_ADDR profileInfo, GM_ADDR profileRegistry, + GM_ADDR legacy, GM_ADDR consumer, GM_ADDR status, int peer, + uint32_t qpIdx, uint32_t localRegion, uint32_t remoteRegion, + uint64_t transferBytes, uint32_t batchCount, uint32_t mode) +{ + if (!RegisterKernel(rank, registration)) { + return false; + } + struct KernelArgs { + GM_ADDR commArgs; + GM_ADDR profileInfo; + GM_ADDR profileRegistry; + GM_ADDR legacy; + GM_ADDR consumer; + GM_ADDR status; + int32_t peer; + uint32_t qpIdx; + uint32_t localRegion; + uint32_t remoteRegion; + uint64_t transferBytes; + uint32_t batchCount; + uint32_t mode; + } args { + commArgs, profileInfo, profileRegistry, legacy, consumer, status, + peer, qpIdx, localRegion, remoteRegion, transferBytes, batchCount, mode + }; + rtArgsEx_t argsInfo {}; + argsInfo.args = &args; + argsInfo.argsSize = sizeof(args); + rtTaskCfgInfo_t cfgInfo {}; + cfgInfo.schemMode = 1; + const rtError_t ret = rtKernelLaunchWithFlagV2(KernelSignature(), 1U, + &argsInfo, nullptr, static_cast(stream), 0U, &cfgInfo); + if (ret != RT_ERROR_NONE) { + std::cerr << "[rank " << rank << "] ERROR: rtKernelLaunchWithFlagV2 failed, ret=" + << ret << std::endl; + return false; + } + return true; +} + +bool ReadAndValidateStatus(int rank, GM_ADDR statusDev, uint32_t mode, + uint32_t qpIdx, uint32_t remoteRegion, uint64_t batchBytes) +{ + uint64_t status[kStatusWords] = {}; + if (!CheckAcl(rank, "aclrtMemcpy status D2H", + aclrtMemcpy(status, sizeof(status), statusDev, sizeof(status), + ACL_MEMCPY_DEVICE_TO_HOST))) { + return false; + } + const bool common = status[0] == kProbeMagic && + status[1] == TileXR::TILEXR_SUCCESS && + status[4] == TileXR::TILEXR_SUCCESS && + status[9] == batchBytes && status[10] == mode; + const bool profile = mode == kModeLegacyGet || + (status[2] != 0U && status[3] == TileXR::TILEXR_SUCCESS && + status[5] == qpIdx && status[6] == 0U && status[7] == remoteRegion); + if (!common || !profile) { + std::cerr << "[rank " << rank << "] ERROR: kernel status mismatch" + << " magic=" << std::hex << status[0] << std::dec + << " post=" << status[1] << " frontier=" << status[2] + << " flush=" << status[3] << " quiet=" << status[4] + << " qp=" << status[5] << " localRegion=" << status[6] + << " remoteRegion=" << status[7] << " batch=" << status[8] + << " bytes=" << status[9] << " mode=" << status[10] << std::endl; + return false; + } + return true; +} + +bool MeasureProfileCase(int rank, int device, int peer, + KernelRegistration& registration, aclrtStream stream, GM_ADDR commArgs, + const TileXR::TileXRUDMAProfileView& view, const DeviceRegion& staging, + const DeviceRegion& consumer, const DeviceRegion& status, + uint32_t qpIdx, uint32_t remoteRegion, uint64_t transferBytes, + uint32_t mode, uint32_t warmupIterations, uint32_t timedIterations, + uint64_t sourcePatternOffset = 0U) +{ + const uint64_t batchBytes = transferBytes * kBatchCount; + if (batchBytes > staging.bytes || batchBytes > consumer.bytes) { + return false; + } + if (!CheckAcl(rank, "aclrtMemset staging", + aclrtMemset(staging.base, staging.bytes, 0xA5, batchBytes)) || + !CheckAcl(rank, "aclrtMemset consumer", + aclrtMemset(consumer.base, consumer.bytes, 0x5A, batchBytes))) { + return false; + } + + for (uint32_t iteration = 0; iteration < warmupIterations; ++iteration) { + if (!LaunchKernel(rank, registration, stream, commArgs, view.infoDev, + view.registryDev, nullptr, consumer.base, status.base, peer, + qpIdx, 0U, remoteRegion, transferBytes, kBatchCount, mode)) { + return false; + } + } + if (!CheckAcl(rank, "aclrtSynchronizeStream warmup", aclrtSynchronizeStream(stream))) { + return false; + } + + for (uint32_t iteration = 0; iteration < timedIterations; ++iteration) { + aclrtEvent start = nullptr; + aclrtEvent stop = nullptr; + if (!CheckAcl(rank, "aclrtCreateEvent start", aclrtCreateEvent(&start)) || + !CheckAcl(rank, "aclrtCreateEvent stop", aclrtCreateEvent(&stop))) { + if (start != nullptr) { + (void)aclrtDestroyEvent(start); + } + if (stop != nullptr) { + (void)aclrtDestroyEvent(stop); + } + return false; + } + bool ok = CheckAcl(rank, "aclrtRecordEvent start", aclrtRecordEvent(start, stream)); + ok = ok && LaunchKernel(rank, registration, stream, commArgs, view.infoDev, + view.registryDev, nullptr, consumer.base, status.base, peer, + qpIdx, 0U, remoteRegion, transferBytes, kBatchCount, mode); + ok = ok && CheckAcl(rank, "aclrtRecordEvent stop", aclrtRecordEvent(stop, stream)); + ok = ok && CheckAcl(rank, "aclrtSynchronizeEvent stop", aclrtSynchronizeEvent(stop)); + float elapsedMs = 0.0F; + ok = ok && CheckAcl(rank, "aclrtEventElapsedTime", + aclrtEventElapsedTime(&elapsedMs, start, stop)); + (void)aclrtDestroyEvent(start); + (void)aclrtDestroyEvent(stop); + if (!ok || elapsedMs <= 0.0F) { + return false; + } + const double udmaGBps = static_cast(batchBytes) / + (static_cast(elapsedMs) * 1.0e6); + const uint64_t stagingIoBytes = mode == kModeProfileConsume ? 2U * batchBytes : 0U; + std::cout << "TILEXR_UDMA_PROFILE_PROBE_JSON {" + << "\"record\":\"timing\"," + << "\"rank\":" << rank << "," + << "\"device\":" << device << "," + << "\"mode\":\"" + << (mode == kModeProfileConsume ? "transfer_consume" : "transfer_only") << "\"," + << "\"qp\":" << qpIdx << "," + << "\"remote_region\":" << remoteRegion << "," + << "\"wqe_bytes\":" << transferBytes << "," + << "\"batch_count\":" << kBatchCount << "," + << "\"batch_bytes\":" << batchBytes << "," + << "\"staging_io_bytes\":" << stagingIoBytes << "," + << "\"iteration\":" << iteration << "," + << "\"elapsed_ms\":" << std::fixed << std::setprecision(6) << elapsedMs << "," + << "\"udma_gbps\":" << std::setprecision(6) << udmaGBps + << "}" << std::defaultfloat << std::endl; + } + + if (!ReadAndValidateStatus(rank, status.base, mode, qpIdx, remoteRegion, batchBytes) || + !ValidatePattern(rank, "staging", staging.base, batchBytes, + peer, remoteRegion, sourcePatternOffset)) { + return false; + } + if (mode == kModeProfileConsume && + !ValidatePattern(rank, "consumer", consumer.base, batchBytes, + peer, remoteRegion, sourcePatternOffset)) { + return false; + } + std::cout << "TILEXR_UDMA_PROFILE_PROBE_JSON {" + << "\"record\":\"correctness\",\"rank\":" << rank + << ",\"mode\":\"" + << (mode == kModeProfileConsume ? "transfer_consume" : "transfer_only") + << "\",\"qp\":" << qpIdx + << ",\"remote_region\":" << remoteRegion + << ",\"wqe_bytes\":" << transferBytes + << ",\"batch_bytes\":" << batchBytes + << ",\"byte_exact\":true}" << std::endl; + return true; +} + +bool RunLegacyGet(int rank, int peer, const char* phase, + KernelRegistration& registration, aclrtStream stream, GM_ADDR commArgs, + const DeviceRegion& legacy, const DeviceRegion& consumer, + const DeviceRegion& status) +{ + GM_ADDR destination = legacy.base + kLegacyDestinationOffset; + if (!CheckAcl(rank, "aclrtMemset legacy destination", + aclrtMemset(destination, legacy.bytes - kLegacyDestinationOffset, + 0xCC, kLegacyTransferBytes)) || + !LaunchKernel(rank, registration, stream, commArgs, nullptr, nullptr, + destination, consumer.base, status.base, peer, 0U, 0U, 0U, + kLegacyTransferBytes, 1U, kModeLegacyGet) || + !CheckAcl(rank, "aclrtSynchronizeStream legacy", aclrtSynchronizeStream(stream)) || + !ReadAndValidateStatus(rank, status.base, kModeLegacyGet, + 0U, 0U, kLegacyTransferBytes) || + !ValidatePattern(rank, "legacy destination", destination, + kLegacyTransferBytes, peer, kLegacyPatternRegion)) { + return false; + } + std::cout << "TILEXR_UDMA_PROFILE_PROBE_JSON {" + << "\"record\":\"lifecycle\",\"rank\":" << rank + << ",\"phase\":\"" << phase + << "\",\"legacy_get_byte_exact\":true}" << std::endl; + return true; +} + +void Cleanup(int rank, int device, TileXRCommPtr comm, aclrtStream stream, + KernelRegistration& registration, std::vector& profileRegions, + DeviceRegion& legacy, DeviceRegion& consumer, DeviceRegion& status) +{ + UnregisterKernel(registration); + if (comm != nullptr) { + (void)TileXRCommDestroy(comm); + } + for (DeviceRegion& region : profileRegions) { + FreeRegion(region); + } + FreeRegion(legacy); + FreeRegion(consumer); + FreeRegion(status); + if (stream != nullptr) { + (void)aclrtDestroyStream(stream); + } + (void)aclrtResetDevice(device); + (void)aclFinalize(); + Log(rank, "cleanup complete"); +} + +} // namespace + +int main() +{ + const int rank = GetRank(); + const int rankSize = GetRankSize(); + const int localRank = GetLocalRank(); + const int device = GetEnvInt("TILEXR_PROFILE_PROBE_DEVICE_BASE", 0) + localRank; + const uint32_t warmupIterations = static_cast( + std::max(0, GetEnvInt("TILEXR_PROFILE_PROBE_WARMUP", 2))); + const uint32_t timedIterations = static_cast( + std::max(1, GetEnvInt("TILEXR_PROFILE_PROBE_ITERATIONS", 5))); + const bool registrationOnly = + GetEnvInt("TILEXR_PROFILE_PROBE_REGISTRATION_ONLY", 0) != 0; + const uint64_t sourceRegistrationOffset = GetEnvUint64( + "TILEXR_PROFILE_PROBE_SOURCE_REG_OFFSET", 0U); + const uint64_t sourceRegistrationBytes = GetEnvUint64( + "TILEXR_PROFILE_PROBE_SOURCE_REG_BYTES", kProfileRegionBytes); + const uint64_t sourceViewOffset = GetEnvUint64( + "TILEXR_PROFILE_PROBE_SOURCE_VIEW_OFFSET", 0U); + const uint64_t defaultSourceViewBytes = + sourceViewOffset <= sourceRegistrationBytes ? + sourceRegistrationBytes - sourceViewOffset : 0U; + const uint64_t sourceViewBytes = GetEnvUint64( + "TILEXR_PROFILE_PROBE_SOURCE_VIEW_BYTES", defaultSourceViewBytes); + if (rankSize != 2 || rank < 0 || rank >= rankSize || device < 0) { + std::cerr << "ERROR: profile probe requires exactly two ranks and valid device ids" << std::endl; + return 2; + } + const int peer = 1 - rank; + + TileXRCommPtr comm = nullptr; + aclrtStream stream = nullptr; + KernelRegistration kernelRegistration; + std::vector profileRegions(kProfileRegionCount); + DeviceRegion legacy; + DeviceRegion consumer; + DeviceRegion status; + TileXRUDMAMemHandle legacyHandle = 0U; + TileXRUDMAProfileHandle profileHandle = 0U; + bool legacyRegistered = false; + bool profileRegistered = false; + bool initialized = false; + + if (!CheckAcl(rank, "aclInit", aclInit(nullptr)) || + !CheckAcl(rank, "aclrtSetDevice", aclrtSetDevice(device)) || + !CheckAcl(rank, "aclrtCreateStream", aclrtCreateStream(&stream))) { + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + initialized = true; + (void)initialized; + + if (!CheckTileXR(rank, "TileXRCommInitRankLocal", + TileXRCommInitRankLocal(rankSize, rank, &comm))) { + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + uint32_t qpCount = 0U; + GM_ADDR commArgs = nullptr; + if (!CheckTileXR(rank, "TileXRUDMAGetQpCount", TileXRUDMAGetQpCount(comm, &qpCount)) || + qpCount < kMinimumQpCount || + !CheckTileXR(rank, "TileXRGetCommArgsDev", TileXRGetCommArgsDev(comm, commArgs))) { + std::cerr << "[rank " << rank << "] ERROR: probe needs at least three UDMA QPs, got " + << qpCount << std::endl; + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + + bool ok = true; + for (DeviceRegion& region : profileRegions) { + ok = AllocateAlignedRegion(rank, kProfileRegionBytes, region) && ok; + } + ok = AllocateAlignedRegion(rank, kLegacyRegionBytes, legacy) && ok; + ok = AllocateAlignedRegion(rank, kProfileRegionBytes, consumer) && ok; + ok = AllocateAlignedRegion(rank, kAlignment, status) && ok; + if (!ok) { + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + if (sourceRegistrationBytes == 0U || sourceViewBytes == 0U || + sourceRegistrationOffset > kProfileRegionBytes || + sourceRegistrationBytes > kProfileRegionBytes - sourceRegistrationOffset || + sourceViewOffset > sourceRegistrationBytes || + sourceViewBytes > sourceRegistrationBytes - sourceViewOffset) { + std::cerr << "[rank " << rank + << "] ERROR: source registration range exceeds the aligned backing allocation" + << std::endl; + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 2; + } + for (uint32_t region = 1U; region < kProfileRegionCount; ++region) { + ok = InitializePattern(rank, profileRegions[region], region, + profileRegions[region].bytes) && ok; + } + ok = InitializePattern(rank, legacy, kLegacyPatternRegion, + kLegacyTransferBytes) && ok; + if (!ok) { + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + + if (!CheckTileXR(rank, "TileXRUDMARegister legacy", + TileXRUDMARegister(comm, legacy.base, legacy.bytes, &legacyHandle))) { + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + legacyRegistered = true; + TileXR::CommArgs* commArgsHost = nullptr; + if (!CheckTileXR(rank, "TileXRGetCommArgsHost", + TileXRGetCommArgsHost(comm, commArgsHost)) || commArgsHost == nullptr) { + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + const GM_ADDR legacyInfoBeforeProfile = commArgsHost->udmaInfoPtr; + const GM_ADDR legacyRegistryBeforeProfile = commArgsHost->udmaRegistryPtr; + + TileXR::TileXRUDMAProfileDesc desc {}; + desc.regionCount = kProfileRegionCount; + desc.qpBindingCount = qpCount; + for (uint32_t region = 0U; region < desc.regionCount; ++region) { + desc.regions[region].base = profileRegions[region].base; + desc.regions[region].bytes = profileRegions[region].bytes; + } + for (uint32_t region = 1U; region < desc.regionCount; ++region) { + desc.regions[region].registrationBase = + profileRegions[region].base + sourceRegistrationOffset; + desc.regions[region].registrationBytes = sourceRegistrationBytes; + desc.regions[region].base = + desc.regions[region].registrationBase + sourceViewOffset; + desc.regions[region].bytes = sourceViewBytes; + } + for (uint32_t qp = 0U; qp < qpCount; ++qp) { + desc.qpBindings[qp].localRegion = 0U; + desc.qpBindings[qp].remoteRegion = qp % 3U + 1U; + } + if (!CheckTileXR(rank, "TileXRUDMAProfileRegister", + TileXRUDMAProfileRegister(comm, &desc, &profileHandle))) { + (void)TileXRUDMAUnregister(comm, legacyHandle); + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + profileRegistered = true; + TileXR::TileXRUDMAProfileView view {}; + if (!CheckTileXR(rank, "TileXRUDMAProfileQuery", + TileXRUDMAProfileQuery(comm, profileHandle, &view)) || + view.infoDev == nullptr || view.registryDev == nullptr || + view.registryHost == nullptr || view.qpCount != qpCount || + commArgsHost->udmaInfoPtr != legacyInfoBeforeProfile || + commArgsHost->udmaRegistryPtr != legacyRegistryBeforeProfile) { + std::cerr << "[rank " << rank << "] ERROR: profile registration mutated legacy CommArgs" + << std::endl; + (void)TileXRUDMAProfileUnregister(comm, profileHandle); + (void)TileXRUDMAUnregister(comm, legacyHandle); + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return 1; + } + + if (registrationOnly) { + std::cout << "TILEXR_UDMA_PROFILE_PROBE_JSON {" + << "\"record\":\"registration\",\"rank\":" << rank + << ",\"source_offset\":" << sourceRegistrationOffset + << ",\"source_bytes\":" << sourceRegistrationBytes + << ",\"view_offset\":" << sourceViewOffset + << ",\"view_bytes\":" << sourceViewBytes + << ",\"source_base_mod_2m\":" + << (reinterpret_cast( + desc.regions[1].registrationBase) % kAlignment) + << ",\"passed\":true}" << std::endl; + const bool profileUnregistered = CheckTileXR(rank, + "TileXRUDMAProfileUnregister registration-only", + TileXRUDMAProfileUnregister(comm, profileHandle)); + profileRegistered = !profileUnregistered; + const bool legacyUnregistered = CheckTileXR(rank, + "TileXRUDMAUnregister registration-only", + TileXRUDMAUnregister(comm, legacyHandle)); + legacyRegistered = !legacyUnregistered; + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return profileUnregistered && legacyUnregistered ? 0 : 1; + } + + for (uint32_t qp = 0U; qp < qpCount && ok; ++qp) { + const uint32_t remoteRegion = desc.qpBindings[qp].remoteRegion; + for (uint64_t transferBytes : kTransferSizes) { + if (transferBytes > desc.regions[remoteRegion].bytes / kBatchCount) { + continue; + } + ok = MeasureProfileCase(rank, device, peer, kernelRegistration, + stream, commArgs, view, profileRegions[0], consumer, status, + qp, remoteRegion, transferBytes, kModeProfileTransfer, + warmupIterations, timedIterations, + sourceRegistrationOffset + sourceViewOffset) && ok; + ok = MeasureProfileCase(rank, device, peer, kernelRegistration, + stream, commArgs, view, profileRegions[0], consumer, status, + qp, remoteRegion, transferBytes, kModeProfileConsume, + warmupIterations, timedIterations, + sourceRegistrationOffset + sourceViewOffset) && ok; + } + } + ok = RunLegacyGet(rank, peer, "profile_and_legacy_coexist", + kernelRegistration, stream, commArgs, legacy, consumer, status) && ok; + + if (ok) { + ok = CheckTileXR(rank, "TileXRUDMAProfileUnregister", + TileXRUDMAProfileUnregister(comm, profileHandle)); + profileRegistered = !ok; + } + ok = ok && RunLegacyGet(rank, peer, "legacy_after_profile_unregister", + kernelRegistration, stream, commArgs, legacy, consumer, status); + + if (ok) { + profileHandle = 0U; + ok = CheckTileXR(rank, "TileXRUDMAProfileRegister second", + TileXRUDMAProfileRegister(comm, &desc, &profileHandle)); + profileRegistered = ok; + } + if (ok) { + view = {}; + ok = CheckTileXR(rank, "TileXRUDMAProfileQuery second", + TileXRUDMAProfileQuery(comm, profileHandle, &view)); + } + if (ok) { + ok = CheckTileXR(rank, "TileXRUDMAUnregister legacy", + TileXRUDMAUnregister(comm, legacyHandle)); + legacyRegistered = !ok; + } + if (ok) { + const uint32_t remoteRegion = desc.qpBindings[0].remoteRegion; + ok = MeasureProfileCase(rank, device, peer, kernelRegistration, + stream, commArgs, view, profileRegions[0], consumer, status, + 0U, remoteRegion, 2U * 1024U * 1024U, kModeProfileTransfer, + 0U, 1U, sourceRegistrationOffset + sourceViewOffset); + if (ok) { + std::cout << "TILEXR_UDMA_PROFILE_PROBE_JSON {" + << "\"record\":\"lifecycle\",\"rank\":" << rank + << ",\"phase\":\"profile_after_legacy_unregister\"," + << "\"profile_get_byte_exact\":true}" << std::endl; + } + } + + if (profileRegistered) { + const bool unregistered = CheckTileXR(rank, "TileXRUDMAProfileUnregister final", + TileXRUDMAProfileUnregister(comm, profileHandle)); + profileRegistered = !unregistered; + ok = unregistered && ok; + } + if (legacyRegistered) { + const bool unregistered = CheckTileXR(rank, "TileXRUDMAUnregister final", + TileXRUDMAUnregister(comm, legacyHandle)); + legacyRegistered = !unregistered; + ok = unregistered && ok; + } + + if (ok) { + std::cout << "TILEXR_UDMA_PROFILE_PROBE_JSON {" + << "\"record\":\"summary\",\"rank\":" << rank + << ",\"device\":" << device + << ",\"qp_count\":" << qpCount + << ",\"profile_regions\":" << kProfileRegionCount + << ",\"region_alignment_bytes\":" << kAlignment + << ",\"max_wqe_bytes\":" << kMaxTransferBytes + << ",\"max_batch_bytes\":" << kProfileRegionBytes + << ",\"passed\":true}" << std::endl; + std::cout << "[rank " << rank << "] TileXR UDMA profile probe success" << std::endl; + } + Cleanup(rank, device, comm, stream, kernelRegistration, + profileRegions, legacy, consumer, status); + return ok ? 0 : 1; +} diff --git a/tests/udma/demo/tilexr_udma_profile_probe_kernel.cpp b/tests/udma/demo/tilexr_udma_profile_probe_kernel.cpp new file mode 100644 index 00000000..530157e1 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_profile_probe_kernel.cpp @@ -0,0 +1,165 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#include "kernel_operator.h" +#include "tilexr_udma.h" + +namespace { + +constexpr uint64_t kProbeMagic = UINT64_C(0x5458505250524f42); // TXPRPROB +constexpr uint32_t kModeProfileTransfer = 0U; +constexpr uint32_t kModeProfileConsume = 1U; +constexpr uint32_t kModeLegacyGet = 2U; +constexpr uint32_t kConsumeTileBytes = 48U * 1024U; + +__aicore__ inline void StoreStatus(__gm__ uint64_t* status, uint32_t postStatus, + uint32_t frontier, uint32_t flushStatus, uint32_t quietStatus, + uint32_t qpIdx, uint32_t localRegion, uint32_t remoteRegion, + uint32_t batchCount, uint64_t batchBytes, uint32_t mode) +{ + if (status == nullptr) { + return; + } + status[0] = kProbeMagic; + status[1] = postStatus; + status[2] = frontier; + status[3] = flushStatus; + status[4] = quietStatus; + status[5] = qpIdx; + status[6] = localRegion; + status[7] = remoteRegion; + status[8] = batchCount; + status[9] = batchBytes; + status[10] = mode; +} + +__aicore__ inline uint32_t ConsumeStaging( + AscendC::TPipe& pipe, const __gm__ TileXR::CommArgs* args, + const __gm__ TileXR::TileXRUDMAProfileRegistry* registry, + uint32_t localRegion, __gm__ uint8_t* consumer, uint64_t bytes) +{ + if (!TileXR::UDMAProfileRegisteredRangeValid( + registry, args->rank, localRegion, 0U, bytes) || + consumer == nullptr || bytes == 0U || (bytes % sizeof(uint32_t)) != 0U) { + return TileXR::TILEXR_UDMA_STATUS_INVALID; + } + + __gm__ uint8_t* staging = TileXR::UDMAProfileRegisteredAddr( + registry, args->rank, localRegion, 0U); + if (staging == nullptr) { + return TileXR::TILEXR_UDMA_STATUS_INVALID; + } + + AscendC::TBuf copyBuf; + pipe.InitBuffer(copyBuf, kConsumeTileBytes); + AscendC::LocalTensor local = copyBuf.Get(); + AscendC::GlobalTensor source; + AscendC::GlobalTensor destination; + source.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(staging), bytes / sizeof(uint32_t)); + destination.SetGlobalBuffer( + reinterpret_cast<__gm__ uint32_t*>(consumer), bytes / sizeof(uint32_t)); + + const uint32_t tileElements = kConsumeTileBytes / sizeof(uint32_t); + const uint64_t totalElements = bytes / sizeof(uint32_t); + for (uint64_t offset = 0U; offset < totalElements; offset += tileElements) { + const uint64_t remaining = totalElements - offset; + const uint32_t elements = remaining < tileElements ? + static_cast(remaining) : tileElements; + AscendC::DataCopy(local, source[offset], elements); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(destination[offset], local, elements); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + return TileXR::TILEXR_UDMA_STATUS_SUCCESS; +} + +} // namespace + +extern "C" __global__ __aicore__ void tilexr_udma_profile_probe_kernel( + GM_ADDR commArgsGM, GM_ADDR profileInfoGM, GM_ADDR profileRegistryGM, + GM_ADDR legacyGM, GM_ADDR consumerGM, GM_ADDR statusGM, + int32_t peer, uint32_t qpIdx, uint32_t localRegion, uint32_t remoteRegion, + uint64_t transferBytes, uint32_t batchCount, uint32_t mode) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto profileInfo = reinterpret_cast<__gm__ TileXR::UDMAInfo*>(profileInfoGM); + auto profileRegistry = + reinterpret_cast<__gm__ TileXR::TileXRUDMAProfileRegistry*>(profileRegistryGM); + auto legacy = reinterpret_cast<__gm__ uint8_t*>(legacyGM); + auto consumer = reinterpret_cast<__gm__ uint8_t*>(consumerGM); + auto status = reinterpret_cast<__gm__ uint64_t*>(statusGM); + + AscendC::TPipe pipe; + AscendC::TBuf wqeBuf; + pipe.InitBuffer(wqeBuf, TileXR::TILEXR_UDMA_WQE_SCRATCH_BYTES); + auto wqeScratch = wqeBuf.Get(); + + uint32_t postStatus = TileXR::TILEXR_UDMA_STATUS_INVALID; + uint32_t frontier = 0U; + uint32_t flushStatus = TileXR::TILEXR_UDMA_STATUS_INVALID; + uint32_t quietStatus = TileXR::TILEXR_UDMA_STATUS_INVALID; + uint64_t batchBytes = 0U; + + if (args == nullptr || peer < 0 || peer >= args->rankSize || + peer == args->rank || transferBytes == 0U) { + StoreStatus(status, postStatus, frontier, flushStatus, quietStatus, + qpIdx, localRegion, remoteRegion, batchCount, batchBytes, mode); + return; + } + + if (mode == kModeLegacyGet) { + if (legacy != nullptr && transferBytes <= UINT32_MAX) { + postStatus = TileXR::UDMAGetNbiOnQp(args, wqeScratch, + peer, 0U, legacy, 0U, static_cast(transferBytes)); + if (postStatus == TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, 0U); + } + } + StoreStatus(status, postStatus, frontier, flushStatus, quietStatus, + 0U, 0U, 0U, 1U, transferBytes, mode); + return; + } + + if (batchCount != 0U) { + postStatus = TileXR::TILEXR_UDMA_STATUS_SUCCESS; + uint64_t offset = 0U; + for (uint32_t batch = 0U; batch < batchCount; ++batch) { + const uint32_t ret = TileXR::UDMAProfileGetNbiOnQpDeferred( + args, profileInfo, profileRegistry, wqeScratch, peer, qpIdx, + localRegion, offset, remoteRegion, offset, transferBytes); + if (ret != TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + postStatus = ret; + break; + } + if (transferBytes > UINT64_MAX - offset) { + postStatus = TileXR::TILEXR_UDMA_STATUS_INVALID; + break; + } + offset += transferBytes; + } + batchBytes = offset; + } + + if (postStatus == TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + frontier = TileXR::UDMAProfileCompletionFrontier( + args, profileInfo, profileRegistry, peer, qpIdx); + flushStatus = TileXR::UDMAProfileFlushQpDoorbell( + args, profileInfo, profileRegistry, peer, qpIdx); + } + if (flushStatus == TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + quietStatus = TileXR::UDMAProfileQuietStatusOnQpUntil( + args, profileInfo, profileRegistry, peer, qpIdx, frontier); + } + if (quietStatus == TileXR::TILEXR_UDMA_STATUS_SUCCESS && + mode == kModeProfileConsume) { + quietStatus = ConsumeStaging( + pipe, args, profileRegistry, localRegion, consumer, batchBytes); + } + + StoreStatus(status, postStatus, frontier, flushStatus, quietStatus, + qpIdx, localRegion, remoteRegion, batchCount, batchBytes, mode); +} diff --git a/tests/udma/unit/test_tilexr_udma_demo_sources.cpp b/tests/udma/unit/test_tilexr_udma_demo_sources.cpp index 48235a98..9b126885 100644 --- a/tests/udma/unit/test_tilexr_udma_demo_sources.cpp +++ b/tests/udma/unit/test_tilexr_udma_demo_sources.cpp @@ -179,6 +179,41 @@ int main() CheckNotContains(kernelPath, kernel, "aclshmem"); CheckNotContains(kernelPath, kernel, "shmem_"); + const std::string profileHostPath = "tests/udma/demo/tilexr_udma_profile_probe.cpp"; + const std::string profileKernelPath = + "tests/udma/demo/tilexr_udma_profile_probe_kernel.cpp"; + const std::string profileRunnerPath = + "tests/udma/demo/run_tilexr_udma_profile_probe_mpi.sh"; + const std::string profileHost = ReadFile(profileHostPath); + const std::string profileKernel = ReadFile(profileKernelPath); + const std::string profileRunner = ReadFile(profileRunnerPath); + CheckContains(profileHostPath, profileHost, "TileXRUDMAProfileRegister"); + CheckContains(profileHostPath, profileHost, "TileXRUDMAProfileUnregister"); + CheckContains(profileHostPath, profileHost, "TileXRUDMAProfileQuery"); + CheckContains(profileHostPath, profileHost, "TileXRUDMARegister"); + CheckContains(profileHostPath, profileHost, "TileXRUDMAUnregister"); + CheckContains(profileHostPath, profileHost, "rtDevBinaryRegister"); + CheckContains(profileHostPath, profileHost, "rtFunctionRegister"); + CheckContains(profileHostPath, profileHost, "rtKernelLaunchWithFlagV2"); + CheckNotContains(profileHostPath, profileHost, "<<<"); + CheckContains(profileKernelPath, profileKernel, "UDMAProfileGetNbiOnQpDeferred"); + CheckContains(profileKernelPath, profileKernel, "UDMAProfileCompletionFrontier"); + CheckContains(profileKernelPath, profileKernel, "UDMAProfileFlushQpDoorbell"); + CheckContains(profileKernelPath, profileKernel, "UDMAProfileQuietStatusOnQpUntil"); + CheckContains(profileKernelPath, profileKernel, "HardEvent::MTE2_MTE3"); + CheckContains(profileKernelPath, profileKernel, "HardEvent::MTE3_MTE2"); + CheckNotContains(profileKernelPath, profileKernel, "<<<"); + CheckNotContains(profileKernelPath, profileKernel, "launch_tilexr_udma_profile"); + CheckContains(profileRunnerPath, profileRunner, + "port_count:6,port_count:6,port_count:2"); + CheckContains(profileRunnerPath, profileRunner, "--launcher "); + CheckContains(profileRunnerPath, profileRunner, "LAUNCHER=\"ssh\""); + CheckContains(profileRunnerPath, profileRunner, "run_ssh_rank()"); + CheckContains(profileRunnerPath, profileRunner, "root@${host}"); + CheckContains(profileRunnerPath, profileRunner, "timings.jsonl"); + CheckContains(profileRunnerPath, profileRunner, + "timeout --signal=TERM --kill-after=30"); + CheckContains(hccpDefsPath, hccpDefs, "MEM_SEG_ACCESS_LOCAL_ONLY = 1"); CheckContains(hccpDefsPath, hccpDefs, "MEM_SEG_ACCESS_READ = (1 << 1)"); CheckContains(hccpDefsPath, hccpDefs, "MEM_SEG_ACCESS_WRITE = (1 << 2)"); diff --git a/tests/udma/unit/test_tilexr_udma_device_api.cpp b/tests/udma/unit/test_tilexr_udma_device_api.cpp index e35a5123..ec689410 100644 --- a/tests/udma/unit/test_tilexr_udma_device_api.cpp +++ b/tests/udma/unit/test_tilexr_udma_device_api.cpp @@ -49,6 +49,7 @@ struct Fixture { CommArgs args = {}; UDMAInfo info = {}; TileXRUDMARegistry registry = {}; + TileXRUDMAProfileRegistry profileRegistry = {}; std::array wq = {}; std::array cq = {}; std::array mem = {}; @@ -65,6 +66,8 @@ struct Fixture { std::array wqeScratch = {}; std::array localRegion = {}; std::array remoteRegion = {}; + std::array remoteRegionTwo = {}; + std::array remoteRegionThree = {}; Fixture() { @@ -86,6 +89,21 @@ struct Fixture { registry.regions[1].base = remoteRegion.data(); registry.regions[1].bytes = remoteRegion.size(); + profileRegistry.rankSize = kRankSize; + profileRegistry.regionCount = 4U; + profileRegistry.qpCount = kQpNum; + profileRegistry.qpBindings[0] = {0U, 1U}; + profileRegistry.qpBindings[1] = {0U, 2U}; + profileRegistry.regions[0].base = localRegion.data(); + profileRegistry.regions[0].bytes = localRegion.size(); + const size_t remoteBase = TILEXR_UDMA_PROFILE_MAX_REGIONS; + profileRegistry.regions[remoteBase + 1U].base = remoteRegion.data(); + profileRegistry.regions[remoteBase + 1U].bytes = remoteRegion.size(); + profileRegistry.regions[remoteBase + 2U].base = remoteRegionTwo.data(); + profileRegistry.regions[remoteBase + 2U].bytes = remoteRegionTwo.size(); + profileRegistry.regions[remoteBase + 3U].base = remoteRegionThree.data(); + profileRegistry.regions[remoteBase + 3U].bytes = remoteRegionThree.size(); + for (uint32_t qp = 0U; qp < kQpNum; ++qp) { const uint32_t entry = kQpNum + qp; sqBuffers[qp].resize(kSqBytes + kSqGuardBytes); @@ -360,6 +378,81 @@ void TestGetAndLegacyQp0Wrappers() "legacy PUT uses named ordered-completion flag"); } +void TestDeferredProfileGetAndCompletionFrontier() +{ + Fixture fixture; + auto scratch = fixture.Scratch(); + fixture.sqDoorbell[0] = 77U; + + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 0U, 0U, 8U, 1U, 16U, 0U) == + TILEXR_UDMA_STATUS_INVALID, "profile GET rejects zero bytes"); + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 0U, 1U, 8U, 1U, 16U, 32U) == + TILEXR_UDMA_STATUS_INVALID, "profile GET rejects wrong local region binding"); + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 0U, 0U, 8U, 2U, 16U, 32U) == + TILEXR_UDMA_STATUS_INVALID, "profile GET rejects wrong remote region binding"); + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 0U, 0U, 240U, 1U, 16U, 32U) == + TILEXR_UDMA_STATUS_INVALID, "profile GET rejects local staging overflow"); + + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 0U, 0U, 8U, 1U, 16U, 32U) == + TILEXR_UDMA_STATUS_SUCCESS, "first deferred profile GET succeeds"); + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 0U, 0U, 64U, 1U, 64U, 16U) == + TILEXR_UDMA_STATUS_SUCCESS, "second deferred profile GET succeeds"); + Check(fixture.sqHead[0] == 2U && fixture.wqeCount[0] == 2U, + "deferred profile GETs advance SQ and completion frontier"); + Check(fixture.sqDoorbell[0] == 77U, + "deferred profile GET batch does not ring the doorbell"); + + const UDMASqeCtx* firstSqe = fixture.Sqe(0U, 0U); + Check(firstSqe->opcode == static_cast(UDMAOpcode::READ), + "deferred profile GET emits READ"); + const uint64_t firstRemoteAddr = static_cast(firstSqe->rmtAddrLOrTokenId) | + (static_cast(firstSqe->rmtAddrHOrTokenValue) << 32U); + Check(firstRemoteAddr == AddressOf(fixture.remoteRegion.data() + 16U), + "profile GET uses bound remote region offset"); + const UDMASgeCtx* firstSge = reinterpret_cast( + fixture.sqBuffers[0].data() + sizeof(UDMASqeCtx)); + Check(firstSge->va == AddressOf(fixture.localRegion.data() + 8U) && + firstSge->len == 32U, "profile GET uses registered local staging address"); + Check(firstSge->tokenId == fixture.wq[Fixture::kQpNum].localTokenId, + "profile GET uses profile QP local token"); + + const uint32_t frontier = UDMAProfileCompletionFrontier(&fixture.args, + &fixture.info, &fixture.profileRegistry, 1, 0U); + Check(frontier == 2U, "profile completion frontier covers the full batch"); + Check(UDMAProfileFlushQpDoorbell(&fixture.args, &fixture.info, + &fixture.profileRegistry, 1, 0U) == TILEXR_UDMA_STATUS_SUCCESS, + "one profile doorbell flush succeeds"); + Check(fixture.sqDoorbell[0] == 2U, "one profile flush publishes the final SQ head"); + + fixture.Cqe(0U, 0U)->owner = 1U; + fixture.Cqe(0U, 0U)->entryIdx = 0U; + fixture.Cqe(0U, 1U)->owner = 1U; + fixture.Cqe(0U, 1U)->entryIdx = 1U; + Check(UDMAProfileQuietStatusOnQpUntil(&fixture.args, &fixture.info, + &fixture.profileRegistry, 1, 0U, frontier) == TILEXR_UDMA_STATUS_SUCCESS, + "one profile quiet consumes the batch frontier"); + Check(fixture.sqTail[0] == 2U && fixture.cqTail[0] == 2U, + "profile quiet reclaims the full deferred batch"); + Check(UDMAProfileQuietStatusOnQpUntil(&fixture.args, &fixture.info, + &fixture.profileRegistry, 1, 0U, frontier + 1U) == TILEXR_UDMA_STATUS_INVALID, + "profile quiet rejects a future completion frontier"); + + Check(UDMAProfileGetNbiOnQpDeferred(&fixture.args, &fixture.info, + &fixture.profileRegistry, scratch, 1, 1U, 0U, 0U, 2U, 24U, 8U) == + TILEXR_UDMA_STATUS_SUCCESS, "second QP uses its independent profile binding"); + const UDMASqeCtx* qpOneSqe = fixture.Sqe(1U, 0U); + const uint64_t qpOneRemoteAddr = static_cast(qpOneSqe->rmtAddrLOrTokenId) | + (static_cast(qpOneSqe->rmtAddrHOrTokenValue) << 32U); + Check(qpOneRemoteAddr == AddressOf(fixture.remoteRegionTwo.data() + 24U), + "second QP selects its bound remote region"); +} + void TestWriteNotifyWrapsWithinSqRing() { Fixture fixture; @@ -521,6 +614,7 @@ int main() TestDeferredPutAndFlushAreQpSpecific(); TestImmediatePutReclaimsCompletedFullSq(); TestGetAndLegacyQp0Wrappers(); + TestDeferredProfileGetAndCompletionFrontier(); TestWriteNotifyWrapsWithinSqRing(); TestWriteNotifyCompletionUsesFinalBbIndex(); TestLegacyQuietWithoutRegistry(); diff --git a/tests/udma/unit/test_tilexr_udma_registry.cpp b/tests/udma/unit/test_tilexr_udma_registry.cpp index 518c3a46..78fc904a 100644 --- a/tests/udma/unit/test_tilexr_udma_registry.cpp +++ b/tests/udma/unit/test_tilexr_udma_registry.cpp @@ -1,5 +1,6 @@ #include #include +#include #include "tilexr_udma_reg.h" @@ -47,11 +48,109 @@ void TestRemoteAddressCalculation() static_cast(0x200040)); } +TileXR::TileXRUDMAProfileDesc MakeProfileDesc() +{ + TileXR::TileXRUDMAProfileDesc desc = {}; + desc.regionCount = 4; + desc.qpBindingCount = 3; + for (uint32_t region = 0; region < desc.regionCount; ++region) { + desc.regions[region].base = reinterpret_cast(0x100000 + region * 0x10000); + desc.regions[region].bytes = 0x8000; + } + desc.qpBindings[0] = {0, 1}; + desc.qpBindings[1] = {0, 2}; + desc.qpBindings[2] = {0, 3}; + return desc; +} + +void TestProfileDescriptorValidation() +{ + auto desc = MakeProfileDesc(); + CHECK_TRUE(TileXR::UDMAProfileDescValid(&desc, 3)); + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&desc, 2)); + + auto invalid = desc; + invalid.regionCount = 0; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + invalid = desc; + invalid.regionCount = TileXR::TILEXR_UDMA_PROFILE_MAX_REGIONS + 1; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + invalid = desc; + invalid.qpBindings[2].remoteRegion = desc.regionCount; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + invalid = desc; + invalid.regions[1].base = nullptr; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + invalid = desc; + invalid.regions[1].bytes = 0; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + invalid = desc; + invalid.regions[1].base = reinterpret_cast( + std::numeric_limits::max() - 7U); + invalid.regions[1].bytes = 16; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + + auto backed = desc; + backed.regions[1].base = reinterpret_cast(0x210000); + backed.regions[1].bytes = 0x1000; + backed.regions[1].registrationBase = reinterpret_cast(0x200000); + backed.regions[1].registrationBytes = 0x20000; + CHECK_TRUE(TileXR::UDMAProfileDescValid(&backed, 3)); + CHECK_EQ(reinterpret_cast( + TileXR::UDMAProfileRegistrationBase(backed.regions[1])), + static_cast(0x200000)); + CHECK_EQ(TileXR::UDMAProfileRegistrationBytes(backed.regions[1]), + UINT64_C(0x20000)); + + invalid = backed; + invalid.regions[1].registrationBytes = 0; + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); + invalid = backed; + invalid.regions[1].base = reinterpret_cast(0x220000); + CHECK_TRUE(!TileXR::UDMAProfileDescValid(&invalid, 3)); +} + +void TestProfileContractAgreement() +{ + const auto desc = MakeProfileDesc(); + auto peer = desc; + peer.regions[0].base = reinterpret_cast(0x900000); + peer.regions[0].bytes = 0x4000; + CHECK_TRUE(TileXR::UDMAProfileContractsEqual(desc, peer)); + peer.qpBindings[1].remoteRegion = 3; + CHECK_TRUE(!TileXR::UDMAProfileContractsEqual(desc, peer)); +} + +void TestProfileRegistryAddressing() +{ + TileXR::TileXRUDMAProfileRegistry registry = {}; + registry.rankSize = 2; + registry.regionCount = 4; + registry.qpCount = 3; + registry.qpBindings[0] = {0, 1}; + registry.qpBindings[1] = {0, 2}; + registry.qpBindings[2] = {0, 3}; + const size_t rankOneGate = TileXR::TILEXR_UDMA_PROFILE_MAX_REGIONS + 1; + registry.regions[rankOneGate].base = reinterpret_cast(0x800000); + registry.regions[rankOneGate].bytes = 4096; + + CHECK_TRUE(TileXR::UDMAProfileRegistryValid(®istry, 2, 4, 3)); + CHECK_TRUE(TileXR::UDMAProfileRegionContains(®istry, 1, 1, 1024, 3072)); + CHECK_TRUE(!TileXR::UDMAProfileRegionContains(®istry, 1, 1, 1024, 3073)); + CHECK_TRUE(!TileXR::UDMAProfileRegionContains(®istry, 2, 1, 0, 1)); + CHECK_EQ(reinterpret_cast( + TileXR::UDMAProfileRemoteAddr(®istry, 1, 1, 64)), + static_cast(0x800040)); +} + } // namespace int main() { TestRemoteAddressCalculation(); + TestProfileDescriptorValidation(); + TestProfileContractAgreement(); + TestProfileRegistryAddressing(); if (g_failures != 0) { std::cerr << g_failures << " registry checks failed" << std::endl; return 1; diff --git a/tests/udma/unit/test_tilexr_udma_source_guard.cpp b/tests/udma/unit/test_tilexr_udma_source_guard.cpp index 66e3c0da..9730a359 100644 --- a/tests/udma/unit/test_tilexr_udma_source_guard.cpp +++ b/tests/udma/unit/test_tilexr_udma_source_guard.cpp @@ -174,15 +174,15 @@ void TestUDMAMemoryCleanupIsRetryable() const std::string transportPath = "src/comm/udma/tilexr_udma_transport.cpp"; const auto transport = ReadFile(transportPath); CheckContains(transportPath, transport, - "const int localRet = CleanupLocalRegistrations(registration.localRegistrations);"); + "const int localRet = CleanupLocalRegistrations(region.localRegistrations);"); CheckContains(transportPath, transport, "const int cleanupRet = CleanupAllMemory();"); CheckContains(transportPath, transport, "GM_ADDR TileXRUDMATransport::GetRegisteredMemoryPtr() const"); - CheckContains(transportPath, transport, "RaCtxRmemUnimport failed for peer"); + CheckContains(transportPath, transport, "RaCtxRmemUnimport failed for region"); CheckContains(transportPath, transport, "RaCtxLmemUnregister failed for eid"); CheckContains(transportPath, transport, "sq[index].localTokenId = registrationIt->second.tokenId;"); CheckContains(transportPath, transport, "it = byEid.erase(it);"); - CheckContains(transportPath, transport, "it = registration.remoteMemHandles.erase(it);"); + CheckContains(transportPath, transport, "it = handles.erase(it);"); const std::string contextPath = "src/comm/udma/tilexr_udma_context.cpp"; const auto context = ReadFile(contextPath); @@ -229,13 +229,13 @@ void TestUDMAUnregisterIsLocalAfterPublication() const std::string path = "src/comm/udma/tilexr_udma_context.cpp"; const auto text = ReadFile(path); const auto unregisterPos = text.find("int TileXRUDMAContext::UnregisterMemory"); - const auto getterPos = text.find("GM_ADDR TileXRUDMAContext::GetRegistryDev", unregisterPos); - if (unregisterPos == std::string::npos || getterPos == std::string::npos) { + const auto profileRegisterPos = text.find("int TileXRUDMAContext::RegisterProfile", unregisterPos); + if (unregisterPos == std::string::npos || profileRegisterPos == std::string::npos) { std::cerr << "failed to locate UDMA unregister body" << std::endl; ++g_failures; return; } - const auto body = text.substr(unregisterPos, getterPos - unregisterPos); + const auto body = text.substr(unregisterPos, profileRegisterPos - unregisterPos); CheckContains(path, body, "transport_->GetBaseUDMAInfoDev()"); CheckContains(path, body, "transport_->CleanupAllMemory()"); CheckNotContains(path, body, "AllGather("); @@ -334,8 +334,7 @@ void TestUDMAMultiQpHostTransportContract() "if (remoteClean && state.qpHandle != nullptr)"); CheckContains(transportPath, transport, "const bool hardwareClean = remoteClean && state.qpHandle == nullptr"); - CheckContains(transportPath, transport, - "registration.remoteMemHandles.count(importKey) == 0"); + CheckContains(transportPath, transport, "region.remoteMemHandles.count(importKey) == 0"); CheckContains(transportPath, transport, "BuildUDMAInfoImage(reinterpret_cast(registration.infoDev), qpCount_"); CheckContains(transportPath, transport, "return IsAvailable() ? qpCount_ : 0U;"); @@ -519,6 +518,69 @@ void TestUDMALocalRegistrationFlags() CheckContains(commPath, comm, "physicalInfo_.chipName == ChipName::CHIP_950PR"); } +void TestUDMAPersistentProfilesStayIndependent() +{ + const std::string apiPath = "src/include/tilexr_api.h"; + const auto api = ReadFile(apiPath); + CheckContains(apiPath, api, "typedef uint32_t TileXRUDMAProfileHandle;"); + CheckContains(apiPath, api, "int TileXRUDMAProfileRegister("); + CheckContains(apiPath, api, "int TileXRUDMAProfileUnregister("); + CheckContains(apiPath, api, "int TileXRUDMAProfileQuery("); + + const std::string registryPath = "src/include/tilexr_udma_reg.h"; + const auto registry = ReadFile(registryPath); + CheckContains(registryPath, registry, "TILEXR_UDMA_PROFILE_MAX_REGIONS = 8U"); + CheckContains(registryPath, registry, "TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS = 32U"); + CheckContains(registryPath, registry, "struct TileXRUDMAProfileQpBinding"); + CheckContains(registryPath, registry, "struct TileXRUDMAProfileView"); + CheckContains(registryPath, registry, "registrationBase"); + CheckContains(registryPath, registry, "registrationBytes"); + CheckContains(registryPath, registry, + "TileXRUDMAProfileQpBinding qpBindings[TILEXR_UDMA_PROFILE_MAX_QP_BINDINGS]"); + + const std::string contextPath = "src/comm/udma/tilexr_udma_context.cpp"; + const auto context = ReadFile(contextPath); + const auto registerPos = context.find("int TileXRUDMAContext::RegisterProfile"); + const auto unregisterPos = context.find("int TileXRUDMAContext::UnregisterProfile", registerPos); + if (registerPos == std::string::npos || unregisterPos == std::string::npos) { + std::cerr << "failed to locate persistent profile registration body" << std::endl; + ++g_failures; + } else { + const auto body = context.substr(registerPos, unregisterPos - registerPos); + CheckContains(contextPath, body, "transport_->PrepareProfile(desc)"); + CheckContains(contextPath, body, + "candidate->registry.regions[index].registrationBase = nullptr;"); + CheckContains(contextPath, body, "transport_->CommitPreparedProfile(candidateHandle)"); + CheckContains(contextPath, body, "profiles_.emplace(candidateHandle, std::move(candidate))"); + CheckNotContains(contextPath, body, "ApplyCommArgsState("); + CheckNotContains(contextPath, body, "udmaInfoDev_ ="); + } + + const std::string transportPath = "src/comm/udma/tilexr_udma_transport.cpp"; + const auto transport = ReadFile(transportPath); + CheckContains(transportPath, transport, + "registration.qpBindings[qpIdx].localRegion"); + CheckContains(transportPath, transport, + "registration.qpBindings[qpIdx].remoteRegion"); + CheckContains(transportPath, transport, + "UDMAProfileRegistrationBase(desc.regions[region])"); + CheckContains(transportPath, transport, + "profiles_.emplace(handle, std::move(preparedProfile_))"); + CheckContains(transportPath, transport, "registration.cleanupPending = true;"); + CheckContains(transportPath, transport, + "if (ret == TILEXR_SUCCESS && it->second->regions.empty()"); + + const std::string devicePath = "src/include/tilexr_udma.h"; + const auto device = ReadFile(devicePath); + CheckContains(devicePath, device, "UDMAProfileRegisteredRangeValid("); + CheckContains(devicePath, device, "UDMAProfileGetNbiOnQpDeferred("); + CheckContains(devicePath, device, + "UDMAPostSend(udmaInfo, wqeScratch"); + CheckContains(devicePath, device, "UDMAProfileCompletionFrontier("); + CheckContains(devicePath, device, "UDMAProfileFlushQpDoorbell("); + CheckContains(devicePath, device, "UDMAProfileQuietStatusOnQpUntil("); +} + } // namespace int main() @@ -530,6 +592,7 @@ int main() TestPublicHeadersDoNotExposeUDMAContext(); TestCommSourcesDoNotUseShmem(); TestUDMALocalRegistrationFlags(); + TestUDMAPersistentProfilesStayIndependent(); TestUDMAMemoryCleanupIsRetryable(); TestUDMARegistrationIsTransactional(); TestUDMAUnregisterIsLocalAfterPublication(); diff --git a/tools/moonep/reduce_grad_benchmark.py b/tools/moonep/reduce_grad_benchmark.py new file mode 100644 index 00000000..76cfdf7b --- /dev/null +++ b/tools/moonep/reduce_grad_benchmark.py @@ -0,0 +1,801 @@ +from __future__ import annotations + +import argparse +import datetime +import json +import math +import os +import platform +import socket +import struct +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Sequence + + +MIN_RANK_COUNT = 4 +INT32_MAX = (1 << 31) - 1 +UDMA_REGISTRATION_ALIGNMENT = 2 * 1024 * 1024 +NATIVE_BASELINE_COMMIT = "a49538a45e5c5bdc82aa6ae02548f99e72ec67eb" +PROJECTION_NAMES = ("gate", "up", "down") +PATTERNS = ("empty", "sparse", "mixed", "heavy", "full", "balanced-full") + + +@dataclass(frozen=True) +class ReduceGradDimensions: + ranks: int + experts_per_rank: int + slots: int + hidden: int + projection: int + + def __post_init__(self) -> None: + if self.ranks < MIN_RANK_COUNT: + raise ValueError(f"ReduceGrad requires at least {MIN_RANK_COUNT} ranks") + if min(self.experts_per_rank, self.slots, self.hidden, self.projection) <= 0: + raise ValueError("ReduceGrad dimensions must be positive") + if self.experts > INT32_MAX or self.slots > INT32_MAX: + raise ValueError("experts and slots must fit in int32") + if self.ranks * self.slots > INT32_MAX: + raise ValueError("rank-by-slot contributor indices must fit in int32") + + @property + def experts(self) -> int: + return self.ranks * self.experts_per_rank + + @property + def row_elements(self) -> int: + return self.hidden * self.projection + + @property + def row_bytes(self) -> int: + return self.row_elements * 4 + + +def build_experts_to_copy( + ranks: int, + slots: int, + experts_per_rank: int, + pattern: str, +) -> list[list[int]]: + if ranks < MIN_RANK_COUNT: + raise ValueError(f"ReduceGrad requires at least {MIN_RANK_COUNT} ranks") + if slots <= 0 or experts_per_rank <= 0: + raise ValueError("invalid ReduceGrad plan dimensions") + if ranks * experts_per_rank > INT32_MAX or slots > INT32_MAX: + raise ValueError("ReduceGrad expert and slot counts must fit in int32") + if ranks * slots > INT32_MAX: + raise ValueError("rank-by-slot contributor indices must fit in int32") + if pattern not in PATTERNS: + raise ValueError(f"unknown ReduceGrad pattern {pattern!r}") + + plan = [[-1 for _ in range(slots)] for _ in range(ranks)] + if pattern == "empty": + return plan + if pattern == "balanced-full": + for source in range(ranks): + for slot in range(slots): + owner = (source + slot) % ranks + plan[source][slot] = owner * experts_per_rank + slot % experts_per_rank + return plan + + counts_by_pattern = { + "sparse": (1,), + "mixed": (3, 0, 2, 1), + "heavy": (3, 3, 2, 3, 1, 0, 2, 3), + "full": (3,), + } + counts = counts_by_pattern[pattern] + for slot in range(min(slots, experts_per_rank)): + count = min(int(counts[slot % len(counts)]), ranks - 1) + for source in range(1, count + 1): + plan[source][slot] = slot + return plan + + +def flatten_plan(plan: Sequence[Sequence[int]]) -> list[int]: + return [int(value) for row in plan for value in row] + + +def plan_statistics(plan: Sequence[Sequence[int]]) -> dict[str, Any]: + flattened = flatten_plan(plan) + live = sum(value >= 0 for value in flattened) + counts: dict[int, int] = {} + for expert in flattened: + if expert >= 0: + counts[expert] = counts.get(expert, 0) + 1 + total = len(flattened) + return { + "entries": total, + "live_entries": live, + "density": (float(live) / total if total else 0.0), + "active_experts": len(counts), + "max_contributors_per_expert": max(counts.values(), default=0), + } + + +def percentile(values: Sequence[float], quantile: float) -> float: + if not values: + raise ValueError("percentile requires at least one value") + if quantile < 0.0 or quantile > 1.0: + raise ValueError("quantile must be in [0, 1]") + ordered = sorted(float(value) for value in values) + position = (len(ordered) - 1) * quantile + lower = int(math.floor(position)) + upper = int(math.ceil(position)) + if lower == upper: + return ordered[lower] + fraction = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * fraction + + +def cross_rank_statistics(samples_by_rank: Sequence[Sequence[float]]) -> dict[str, Any]: + if not samples_by_rank or not samples_by_rank[0]: + raise ValueError("cross-rank statistics require samples") + iterations = len(samples_by_rank[0]) + if any(len(samples) != iterations for samples in samples_by_rank): + raise ValueError("all ranks must report the same number of samples") + maxima = [ + max(float(samples[iteration]) for samples in samples_by_rank) + for iteration in range(iterations) + ] + return { + "cross_rank_max_us": maxima, + "p50_us": percentile(maxima, 0.50), + "p99_us": percentile(maxima, 0.99), + "min_us": min(maxima), + "max_us": max(maxima), + "mean_us": sum(maxima) / len(maxima), + } + + +def source_value(source_rank: int, projection: int, slot: int) -> float: + value = ( + (projection + 1) * 0.5 + + (source_rank + 1) * 0.03125 + + (slot + 1) * 0.0009765625 + ) + return struct.unpack("f", struct.pack("f", value))[0] + + +def fp32_add(lhs: float, rhs: float) -> float: + return struct.unpack("f", struct.pack("f", float(lhs) + float(rhs)))[0] + + +def expected_expert_value( + plan: Sequence[Sequence[int]], global_expert: int, projection: int +) -> float: + value = 0.0 + for source, row in enumerate(plan): + for slot, expert in enumerate(row): + if int(expert) == global_expert: + value = fp32_add(value, source_value(source, projection, slot)) + return value + + +def _align_up(value: int, alignment: int) -> int: + return ((int(value) + alignment - 1) // alignment) * alignment + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + temporary.replace(path) + + +def _command_output(command: Sequence[str]) -> str | None: + try: + completed = subprocess.run( + list(command), check=False, capture_output=True, text=True, timeout=10 + ) + except (OSError, subprocess.SubprocessError): + return None + output = completed.stdout.strip() or completed.stderr.strip() + return output if output else None + + +def environment_metadata(torch_module, torch_npu_module) -> dict[str, Any]: + npu_smi = _command_output(("npu-smi", "info")) + return { + "hostname": socket.gethostname(), + "platform": platform.platform(), + "python": sys.version, + "torch": str(torch_module.__version__), + "torch_npu": str(torch_npu_module.__version__), + "cann_home": os.environ.get("ASCEND_HOME_PATH"), + "cann_opp_path": os.environ.get("ASCEND_OPP_PATH"), + "driver_summary": None if npu_smi is None else npu_smi.splitlines()[:4], + } + + +def _current_stream_ptr(torch_module, device_index: int) -> int: + if int(torch_module.npu.current_device()) != int(device_index): + raise RuntimeError(f"current NPU device is not npu:{device_index}") + stream = torch_module.npu.current_stream() + value = getattr(stream, "npu_stream", None) + if value is None: + value = getattr(stream, "stream", None) + if value is None: + raise RuntimeError("torch.npu.current_stream() exposes no native stream pointer") + return int(value) + + +def _aligned_workspace(torch_module, device: str, size_bytes: int, alignment: int): + if size_bytes <= 0 or alignment <= 0: + raise ValueError("workspace size and alignment must be positive") + allocation = torch_module.empty( + (size_bytes + alignment - 1,), dtype=torch_module.uint8, device=device + ) + offset = (-int(allocation.data_ptr())) % alignment + workspace = allocation.narrow(0, offset, size_bytes) + if int(workspace.data_ptr()) % alignment: + raise RuntimeError("failed to align ReduceGrad UDMA workspace") + return workspace, allocation + + +def _aligned_float32_tensor(torch_module, device: str, shape: Sequence[int]): + elements = math.prod(int(value) for value in shape) + alignment_elements = UDMA_REGISTRATION_ALIGNMENT // 4 + allocation = torch_module.empty( + (elements + alignment_elements - 1,), + dtype=torch_module.float32, + device=device, + ) + offset_bytes = (-int(allocation.data_ptr())) % UDMA_REGISTRATION_ALIGNMENT + if offset_bytes % 4: + raise RuntimeError("float32 allocation cannot satisfy UDMA alignment") + tensor = allocation.narrow(0, offset_bytes // 4, elements).reshape(tuple(shape)) + if int(tensor.data_ptr()) % UDMA_REGISTRATION_ALIGNMENT: + raise RuntimeError("failed to align UDMA float32 tensor") + return tensor, allocation + + +def gradient_source_regions(full_gradients, expert_count: int, slots: int): + sources = [ + gradient.narrow(0, expert_count, slots) for gradient in full_gradients + ] + return sources, list(full_gradients) + + +class TileXRReduceGradRunner: + def __init__(self, torch_module, dimensions: ReduceGradDimensions, plan, args): + from tilexr_moonep import ( + MoonEPPlan, + ProjectionBuffers, + TileXRMoonEPRuntime, + ) + + self.torch = torch_module + self.dimensions = dimensions + self.plan_host = plan + self.rank = int(os.environ["RANK"]) + self.device_index = int(os.environ["LOCAL_RANK"]) + self.device = f"npu:{self.device_index}" + if args.num_sms is not None: + os.environ["TILEXR_MOONEP_REDUCE_GRAD_BLOCK_DIM"] = str(args.num_sms) + if args.qp_route_spec is not None: + os.environ["TILEXR_UDMA_QP_ROUTE_SPEC"] = args.qp_route_spec + self.runtime = TileXRMoonEPRuntime( + self.rank, + dimensions.ranks, + install_prefix=args.install_prefix, + ) + self.prepared = None + self.closed = False + plan_tensor = torch_module.tensor( + plan, dtype=torch_module.int32, device=self.device + ) + self.plan = MoonEPPlan( + dst=torch_module.zeros((1,), dtype=torch_module.int32, device=self.device), + experts_to_copy=plan_tensor, + zero_fill_ranges=torch_module.zeros( + (dimensions.experts + dimensions.slots, 2), + dtype=torch_module.int32, + device=self.device, + ), + remote_stats=torch_module.zeros( + (2,), dtype=torch_module.int32, device=self.device + ), + dup_groups=torch_module.zeros( + (1, 3), dtype=torch_module.int32, device=self.device + ), + dup_loffs=torch_module.zeros( + (1,), dtype=torch_module.int32, device=self.device + ), + dup_counts=torch_module.zeros( + (2,), dtype=torch_module.int32, device=self.device + ), + status=torch_module.zeros( + (1,), dtype=torch_module.int32, device=self.device + ), + reduce_grad_status=torch_module.zeros( + (1,), dtype=torch_module.int32, device=self.device + ), + workspace=torch_module.empty( + (1,), dtype=torch_module.uint8, device=self.device + ), + n=1, + tokens_per_rank=1, + topk=1, + expert_count=dimensions.experts, + rank_size=dimensions.ranks, + prefetch_slots=dimensions.slots, + nv_s=1, + token_padding=1, + epoch=1, + backend="tilexr", + runtime=self.runtime, + ) + rows = dimensions.experts + dimensions.slots + shapes = ( + (rows, dimensions.hidden, dimensions.projection), + (rows, dimensions.hidden, dimensions.projection), + (rows, dimensions.projection, dimensions.hidden), + ) + gradient_pairs = [ + _aligned_float32_tensor(torch_module, self.device, shape) + for shape in shapes + ] + self.full_gradients = [pair[0] for pair in gradient_pairs] + self.gradient_allocations = [pair[1] for pair in gradient_pairs] + for gradient in self.full_gradients: + gradient.zero_() + self.gradients = ProjectionBuffers(*self.full_gradients) + self.sources, self.source_registrations = gradient_source_regions( + self.full_gradients, dimensions.experts, dimensions.slots + ) + self.source_strategy = "gradient-backing-interior-slice" + self.reset_sources() + torch_module.npu.synchronize() + prepare_start = time.perf_counter() + self.info = self.runtime.reduce_grad_workspace_info( + None, + self.plan, + self.gradients, + requested_udma_chunk_bytes=args.chunk_bytes, + ) + self.workspace, self.workspace_allocation = _aligned_workspace( + torch_module, + self.device, + self.info.workspace_bytes, + self.info.workspace_alignment, + ) + self.prepared = self.runtime.prepare_reduce_grad( + None, + self.plan, + self.gradients, + self.sources, + self.source_registrations, + self.workspace, + requested_udma_chunk_bytes=args.chunk_bytes, + ) + torch_module.npu.synchronize() + self.prepare_ms = (time.perf_counter() - prepare_start) * 1000.0 + self.wait_iterations = args.wait_iterations + + def reset_sources(self) -> None: + for projection, source in enumerate(self.sources): + for slot in range(self.dimensions.slots): + source[slot].fill_(source_value(self.rank, projection, slot)) + + def launch(self) -> None: + self.runtime.reduce_grad( + None, + self.plan, + self.gradients, + self.sources, + self.source_registrations, + self.prepared, + _current_stream_ptr(self.torch, self.device_index), + self.wait_iterations, + ) + + def synchronize(self) -> None: + self.torch.npu.synchronize(self.device_index) + status = int(self.plan.reduce_grad_status.item()) + if status != 0: + raise RuntimeError(f"TileXR ReduceGrad device status is {status}, expected 0") + + def validate(self) -> dict[str, Any]: + d = self.dimensions + owner_begin = self.rank * d.experts_per_rank + owner_end = owner_begin + d.experts_per_rank + for projection, gradient in enumerate(self.full_gradients): + actual = gradient[owner_begin:owner_end] + expected = self.torch.empty_like(gradient[owner_begin:owner_end]) + for local_expert in range(d.experts_per_rank): + expected[local_expert].fill_(expected_expert_value( + self.plan_host, owner_begin + local_expert, projection + )) + if not self.torch.equal(actual, expected): + actual_flat = actual.reshape(-1) + expected_flat = expected.reshape(-1) + mismatch = actual_flat != expected_flat + mismatch_index = int(mismatch.to(self.torch.int32).argmax().item()) + staging_offset = int(self.info.lane_state_bytes) + staging_bytes = bytes( + int(value) for value in + self.workspace.narrow(0, staging_offset, 4).cpu().tolist() + ) + staging_first = struct.unpack("f", staging_bytes)[0] + raise RuntimeError( + f"TileXR projection {projection} owner rows differ from ordered " + f"FP32 reference at flat index {mismatch_index}: " + f"actual={float(actual_flat[mismatch_index].item())}, " + f"expected={float(expected_flat[mismatch_index].item())}, " + f"gate_lane0_bank0_staging={staging_first}" + ) + if owner_begin and bool(self.torch.count_nonzero(gradient[:owner_begin]).item()): + raise RuntimeError(f"TileXR projection {projection} modified lower non-owner rows") + if owner_end < d.experts and bool( + self.torch.count_nonzero(gradient[owner_end:d.experts]).item() + ): + raise RuntimeError(f"TileXR projection {projection} modified upper non-owner rows") + for slot in range(d.slots): + live = int(self.plan_host[self.rank][slot]) >= 0 + if live: + if bool(self.torch.count_nonzero(self.sources[projection][slot]).item()): + raise RuntimeError( + f"TileXR projection {projection} live source slot {slot} was not cleared" + ) + else: + expected_source = source_value(self.rank, projection, slot) + if not bool(self.torch.all( + self.sources[projection][slot] == expected_source + ).item()): + raise RuntimeError( + f"TileXR projection {projection} unused source slot {slot} changed" + ) + return {"passed": True, "mode": "exact_ordered_fp32"} + + def layout_metadata(self) -> dict[str, Any]: + metadata = self.info.as_dict() + metadata["source_strategy"] = self.source_strategy + metadata["source_alignment"] = UDMA_REGISTRATION_ALIGNMENT + metadata["configured_qp_route_spec"] = os.environ.get( + "TILEXR_UDMA_QP_ROUTE_SPEC" + ) + return metadata + + def close(self) -> None: + if self.closed: + return + try: + self.torch.npu.synchronize(self.device_index) + self.runtime.destroy_reduce_grad(self.prepared) + self.prepared = None + finally: + self.runtime.close() + self.closed = True + + +class NativeReduceGradRunner: + def __init__(self, torch_module, dimensions: ReduceGradDimensions, plan, args): + native_root = Path(args.native_root).resolve() + if str(native_root) not in sys.path: + sys.path.insert(0, str(native_root)) + from ascend_moonep.buffer_c import ( + create_sym_dist_tensor, + shmem_finalize, + shmem_init, + ) + from ascend_moonep.launchers import launch_grad_reduce + + self.torch = torch_module + self.dimensions = dimensions + self.plan_host = plan + self.rank = int(os.environ["RANK"]) + self.local_rank = int(os.environ["LOCAL_RANK"]) + self.device = f"npu:{self.local_rank}" + self._launch_grad_reduce = launch_grad_reduce + self._shmem_finalize = shmem_finalize + self.num_sms = ( + int(args.num_sms) if args.num_sms is not None else + int(os.environ.get("TILEXR_MOONEP_NATIVE_REDUCE_GRAD_BLOCK_DIM", "32")) + ) + projection_bytes = (dimensions.experts_per_rank + dimensions.slots) * ( + dimensions.row_bytes + ) + heap_bytes = _align_up(3 * projection_bytes + 256 * 1024 * 1024, 2 * 1024 * 1024) + prepare_start = time.perf_counter() + shmem_init(group=None, mem_size=heap_bytes) + shapes = ( + (dimensions.hidden, dimensions.projection), + (dimensions.hidden, dimensions.projection), + (dimensions.projection, dimensions.hidden), + ) + self.outputs = [ + create_sym_dist_tensor( + (dimensions.experts_per_rank, *shape), + torch_module.float32, + self.local_rank, + ) + for shape in shapes + ] + self.sources = [ + create_sym_dist_tensor( + (dimensions.slots, *shape), + torch_module.float32, + self.local_rank, + ) + for shape in shapes + ] + for output in self.outputs: + output.zero_() + self.plan = torch_module.tensor( + plan, dtype=torch_module.int32, device=self.device + ) + self.reset_sources() + torch_module.npu.synchronize() + self.prepare_ms = (time.perf_counter() - prepare_start) * 1000.0 + + def reset_sources(self) -> None: + for projection, source in enumerate(self.sources): + for slot in range(self.dimensions.slots): + source[slot].fill_(source_value(self.rank, projection, slot)) + + def launch(self) -> None: + for output, source in zip(self.outputs, self.sources): + self._launch_grad_reduce( + output, + source, + self.plan, + rank=self.rank, + num_sms=self.num_sms, + meta_buf=None, + meta_stride=0, + barrier_off=0, + grid_sync_bar=None, + ) + + def synchronize(self) -> None: + self.torch.npu.synchronize() + + def validate(self) -> dict[str, Any]: + d = self.dimensions + owner_begin = self.rank * d.experts_per_rank + for projection, output in enumerate(self.outputs): + expected = self.torch.empty_like(output) + for local_expert in range(d.experts_per_rank): + expected[local_expert].fill_(expected_expert_value( + self.plan_host, owner_begin + local_expert, projection + )) + if not self.torch.equal(output, expected): + raise RuntimeError( + f"native projection {projection} owner rows differ from ordered FP32 reference" + ) + for slot in range(d.slots): + live = int(self.plan_host[self.rank][slot]) >= 0 + if live: + if bool(self.torch.count_nonzero(self.sources[projection][slot]).item()): + raise RuntimeError( + f"native projection {projection} live source slot {slot} was not cleared" + ) + else: + expected_source = source_value(self.rank, projection, slot) + if not bool(self.torch.all( + self.sources[projection][slot] == expected_source + ).item()): + raise RuntimeError( + f"native projection {projection} unused source slot {slot} changed" + ) + return {"passed": True, "mode": "exact_ordered_fp32"} + + def layout_metadata(self) -> dict[str, Any]: + return { + "block_dim": self.num_sms, + "launch_count": 3, + "transport": "cann-shmem-mte-get", + } + + def close(self) -> None: + self._shmem_finalize() + + +def _time_iteration(torch_module, distributed, runner) -> float: + runner.reset_sources() + torch_module.npu.synchronize() + distributed.barrier() + start = torch_module.npu.Event(enable_timing=True) + end = torch_module.npu.Event(enable_timing=True) + start.record() + runner.launch() + end.record() + end.synchronize() + elapsed_us = float(start.elapsed_time(end)) * 1000.0 + runner.synchronize() + return elapsed_us + + +def _run(args: argparse.Namespace) -> int: + import torch + import torch.distributed as dist + import torch_npu + + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + dimensions = ReduceGradDimensions( + ranks=world_size, + experts_per_rank=args.experts_per_rank, + slots=args.slots, + hidden=args.hidden, + projection=args.projection, + ) + torch.npu.set_device(local_rank) + init_kwargs = { + "backend": args.coordination_backend, + "timeout": datetime.timedelta(seconds=args.coordination_timeout_seconds), + } + if args.coordination_backend == "hccl": + init_kwargs["device_id"] = torch.device(f"npu:{local_rank}") + dist.init_process_group(**init_kwargs) + dist.barrier() + plan = build_experts_to_copy( + world_size, dimensions.slots, dimensions.experts_per_rank, args.pattern + ) + output_dir = Path(args.output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + runner = None + result: dict[str, Any] = { + "schema_version": 1, + "status": "failed", + "backend": args.backend, + "rank": rank, + "world_size": world_size, + "coordination_backend": args.coordination_backend, + "case": { + **asdict(dimensions), + "experts": dimensions.experts, + "pattern": args.pattern, + "num_sms": args.num_sms, + "qp_route_spec": args.qp_route_spec, + "warmup": args.warmup, + "iterations": args.iterations, + }, + "plan": plan_statistics(plan), + "timing_boundary": { + "clock": "NPU device events on current stream", + "tilexr": "one fused TileXRMoonEpReduceGradV2 launch", + "native": "three sequential launch_grad_reduce launches", + "excluded": [ + "allocation", + "MR registration/import", + "prepared-profile creation", + "source reset", + "cross-rank pre-launch alignment", + "post-launch status synchronization", + ], + }, + } + failure = None + try: + setup_start = time.perf_counter() + if args.backend == "tilexr": + runner = TileXRReduceGradRunner(torch, dimensions, plan, args) + else: + runner = NativeReduceGradRunner(torch, dimensions, plan, args) + result["setup_ms"] = (time.perf_counter() - setup_start) * 1000.0 + result["prepare_ms"] = runner.prepare_ms + result["layout"] = runner.layout_metadata() + + correctness = {"passed": None, "mode": "disabled"} + if args.correctness: + required = 3 * (dimensions.experts + dimensions.slots) * dimensions.row_bytes + if required > args.correctness_max_bytes: + raise ValueError( + f"exact correctness would inspect {required} bytes, above " + f"--correctness-max-bytes={args.correctness_max_bytes}" + ) + runner.reset_sources() + torch.npu.synchronize() + dist.barrier() + runner.launch() + runner.synchronize() + correctness = runner.validate() + dist.barrier() + result["correctness"] = correctness + + for _ in range(args.warmup): + _time_iteration(torch, dist, runner) + samples = [ + _time_iteration(torch, dist, runner) for _ in range(args.iterations) + ] + gathered: list[Any] = [None for _ in range(world_size)] + dist.all_gather_object(gathered, samples) + result["local_samples_us"] = samples + result["status"] = "passed" + if rank == 0: + result["samples_by_rank_us"] = gathered + result["statistics"] = cross_rank_statistics(gathered) + result["commits"] = { + "tilexr_base": args.tilexr_commit, + "tilexr_source_snapshot_sha256": args.source_snapshot_sha256, + "native_moonep": NATIVE_BASELINE_COMMIT, + } + result["environment"] = environment_metadata(torch, torch_npu) + except Exception as exc: + result["failure_reason"] = f"{type(exc).__name__}: {exc}" + failure = (exc, exc.__traceback__) + finally: + try: + torch.npu.synchronize() + dist.barrier() + if runner is not None: + runner.close() + dist.barrier() + except Exception as cleanup_error: + result["cleanup_error"] = ( + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + if failure is None: + failure = (cleanup_error, cleanup_error.__traceback__) + result["status"] = "failed" + _write_json(output_dir / f"rank_{rank}.json", result) + if rank == 0: + _write_json(output_dir / "summary.json", result) + dist.destroy_process_group() + if failure is not None: + raise failure[0].with_traceback(failure[1]) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Isolated TileXR versus native MoonEP ReduceGrad benchmark" + ) + parser.add_argument("--backend", choices=("tilexr", "native"), required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--install-prefix", default=None) + parser.add_argument("--native-root", default="reference/ascend-moonep-dev") + parser.add_argument("--experts-per-rank", type=int, default=8) + parser.add_argument("--slots", type=int, default=8) + parser.add_argument("--hidden", type=int, default=32) + parser.add_argument("--projection", type=int, default=32) + parser.add_argument("--pattern", choices=PATTERNS, default="mixed") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument("--chunk-bytes", type=int, default=8 * 1024 * 1024) + parser.add_argument("--num-sms", type=int, default=None) + parser.add_argument( + "--qp-route-spec", + default=os.environ.get("TILEXR_UDMA_QP_ROUTE_SPEC"), + ) + parser.add_argument("--wait-iterations", type=int, default=1_000_000) + parser.add_argument( + "--coordination-backend", choices=("gloo", "hccl"), default="gloo" + ) + parser.add_argument("--coordination-timeout-seconds", type=int, default=120) + parser.add_argument("--correctness", action="store_true") + parser.add_argument("--correctness-max-bytes", type=int, default=512 * 1024 * 1024) + parser.add_argument("--tilexr-commit", default=os.environ.get("TILEXR_GIT_COMMIT")) + parser.add_argument( + "--source-snapshot-sha256", + default=os.environ.get("TILEXR_SOURCE_SNAPSHOT_SHA256"), + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + if args.warmup < 0 or args.iterations <= 0: + raise ValueError("warmup must be non-negative and iterations must be positive") + if args.chunk_bytes <= 0 or args.wait_iterations <= 0: + raise ValueError("chunk-bytes and wait-iterations must be positive") + if args.num_sms is not None and args.num_sms <= 0: + raise ValueError("num-sms must be positive") + if args.correctness_max_bytes <= 0: + raise ValueError("correctness-max-bytes must be positive") + if args.coordination_timeout_seconds <= 0: + raise ValueError("coordination-timeout-seconds must be positive") + return _run(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/moonep/tilexr_backend.py b/tools/moonep/tilexr_backend.py index b0cac6dc..2c2c05b9 100644 --- a/tools/moonep/tilexr_backend.py +++ b/tools/moonep/tilexr_backend.py @@ -233,19 +233,6 @@ def reduce_grad( dtype=self.torch.float32, reduce_buffers=True, ) - d = self.dimensions - full_tail = { - name: tensor.narrow(0, d.expert_count, d.prefetch_slots).clone() - for name, tensor in full_grads.items() - } - unused_reduce = { - name: { - slot: tensor[d.rank][slot].clone() - for slot in range(d.prefetch_slots) - if int(plan.experts_to_copy[d.rank][slot].item()) < 0 - } - for name, tensor in reduce_buffers.items() - } self.buffer.reduce_grad( entry.native, full_gate_grad=full_grads.gate, @@ -255,13 +242,6 @@ def reduce_grad( up_reduce_buffer=reduce_buffers.up, down_reduce_buffer=reduce_buffers.down, ) - # These restores follow the native launch on the same current stream. - for name, tensor in full_grads.items(): - tensor.narrow(0, d.expert_count, d.prefetch_slots).copy_(full_tail[name]) - for name, slots in unused_reduce.items(): - tensor = getattr(reduce_buffers, name) - for slot, snapshot in slots.items(): - tensor[d.rank][slot].copy_(snapshot) return ReduceGradResult(full_grads, reduce_buffers) def synchronize(self) -> None: