From 05c2bd38e1170179157304af1a59c482cc909910 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Thu, 18 Jun 2026 10:58:13 +0800 Subject: [PATCH 001/163] docs: add udma alltoall demo design --- .../2026-06-18-udma-alltoall-demo-design.md | 124 ++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-18-udma-alltoall-demo-design.md diff --git a/docs/superpowers/specs/2026-06-18-udma-alltoall-demo-design.md b/docs/superpowers/specs/2026-06-18-udma-alltoall-demo-design.md new file mode 100644 index 00000000..8eb2dc41 --- /dev/null +++ b/docs/superpowers/specs/2026-06-18-udma-alltoall-demo-design.md @@ -0,0 +1,124 @@ +# UDMA All-to-All Demo Design + +## Goal + +Add an all-to-all UDMA operator demo under `tests/udma/demo`. + +The demo validates the common all-to-all layout: + +- Each rank owns `rank_size` equal input slices. +- Input slice `dst_rank` from rank `src_rank` is sent to rank `dst_rank`. +- Each destination rank writes output slices ordered by source rank. + +For rank `i`, input is laid out as `[to0, to1, ..., toN-1]`. For rank `j`, output is laid out as `[from0, from1, ..., fromN-1]`. + +## Approach + +Extend the existing `tilexr_udma_demo` binary instead of creating a separate demo. + +The current demo already handles: + +- multi-process local rank launch +- TileXR communicator initialization +- UDMA capability checks +- ordinary `aclrtMalloc` memory registration through `TileXRUDMARegister` +- local TCP barriers for demo synchronization +- per-rank logs and result validation + +The all-to-all path will be selected with `test_type=2`. Existing `test_type=0` all-gather and `test_type=1` put-signal behavior must remain unchanged. + +## Host Data Layout + +For `test_type=2`, the host allocates one registered device-memory payload containing: + +- input buffer: `rank_size * elements_per_peer` `int32_t` values +- output buffer: `rank_size * elements_per_peer` `int32_t` values +- signal/debug space if needed by the shared demo structure + +The registered allocation remains rounded up to the existing 2 MiB UDMA registration alignment. + +Input initialization for rank `src`: + +```text +input[dst][elem] = 100000 + src * 1000 + dst +``` + +Expected output for rank `dst`: + +```text +output[src][elem] = 100000 + src * 1000 + dst +``` + +This makes source and destination rank mistakes visible in validation logs. + +## Kernel Behavior + +Add a new AICore kernel and launch wrapper in `tilexr_udma_demo_kernel.cpp`. + +The kernel reads `rank`, `rankSize`, and UDMA registry state from `CommArgs`. + +For each peer: + +- If `peer == rank`, copy the local input slice `input[rank]` into local output slice `output[rank]`. +- Otherwise, issue `TileXR::UDMAPutNbi` to write local `input[peer]` into the remote rank's registered output slice for this source rank. +- Call `TileXR::UDMAQuiet(args, peer)` after posting to each remote peer. + +The remote byte offset is computed against the peer rank's registered base: + +```text +output_offset + rank * elements_per_peer * sizeof(int32_t) +``` + +The local source pointer is: + +```text +input + peer * elements_per_peer +``` + +## Synchronization + +The demo keeps the existing host-side synchronization: + +1. Each rank initializes and registers its buffers. +2. Host barrier ensures every rank's registered-memory metadata is visible. +3. Each rank launches the all-to-all kernel and synchronizes its stream. +4. Host barrier ensures all ranks have completed UDMA writes. +5. Host copies output back and validates. + +The kernel does not add device-side inter-rank polling beyond `UDMAQuiet`. + +## Build And Run + +The existing `tests/udma/CMakeLists.txt` continues to build one demo kernel shared object and one `tilexr_udma_demo` executable. + +Update the run script and README so: + +```bash +bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 +``` + +runs the all-to-all path. + +## Verification + +Local checks: + +- Build metadata remains scoped to `tests/udma`. +- Existing all-gather and put-signal source paths remain intact. + +Remote hardware validation: + +- Create a new directory under `/home/aiv-perf/` on `root@141.61.95.18`. +- Copy or sync the repository into that directory. +- Build TileXR core and `tests/udma`. +- Run `bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0`. +- If resources permit, also run a wider case such as `rank_size=4`. + +Success requires every rank to print `TileXR UDMA demo success` and all output segments to match the expected all-to-all pattern. + +## Out Of Scope + +- Refactoring the demo runtime into shared helper classes. +- Adding a production all-to-all collective API. +- Optimizing multi-peer posting or batching. +- Supporting non-`int32_t` element types in this demo. From ecff707c8def5f30c43ff407524977f294b593fc Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Thu, 18 Jun 2026 11:03:48 +0800 Subject: [PATCH 002/163] docs: add udma alltoall implementation plan --- .../plans/2026-06-18-udma-alltoall-demo.md | 580 ++++++++++++++++++ 1 file changed, 580 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-18-udma-alltoall-demo.md diff --git a/docs/superpowers/plans/2026-06-18-udma-alltoall-demo.md b/docs/superpowers/plans/2026-06-18-udma-alltoall-demo.md new file mode 100644 index 00000000..8fdf2818 --- /dev/null +++ b/docs/superpowers/plans/2026-06-18-udma-alltoall-demo.md @@ -0,0 +1,580 @@ +# UDMA All-to-All Demo Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a `test_type=2` all-to-all path to the TileXR UDMA demo under `tests/udma/demo`. + +**Architecture:** Reuse the existing `tilexr_udma_demo` host executable, launch script, local TCP barrier, and UDMA registered-memory setup. Add one AICore all-to-all kernel that writes each local destination slice into the matching remote output slice using `TileXR::UDMAPutNbi`, with host-side validation after all ranks complete. + +**Tech Stack:** C++14 host code, Ascend C AICore kernel code, TileXR public C API, `tilexr_udma.h` device wrapper, CMake, bash, remote Ascend950/A5 validation. + +## Global Constraints + +- CANN version: 9.1.0. +- Target OS: Ubuntu 20.04 LTS; root user required for device access. +- UDMA data-plane validation targets A5 / Ascend950 / 950 hardware. +- Do not add shmem includes or shmem API calls. +- Existing `test_type=0` all-gather and `test_type=1` put-signal behavior must remain unchanged. +- All-to-all uses `int32_t` only. +- Remote validation must create a new directory under `/home/aiv-perf/` on `root@141.61.95.18`. + +--- + +## File Structure + +- Modify `tests/udma/demo/tilexr_udma_demo_kernel.cpp`: add `tilexr_udma_all_to_all_kernel` and `launch_tilexr_udma_all_to_all`. +- Modify `tests/udma/demo/tilexr_udma_demo.cpp`: declare the launch wrapper, allocate the all-to-all input/output layout for `test_type=2`, initialize inputs, launch the new kernel, and validate outputs. +- Modify `tests/udma/demo/run_tilexr_udma_demo.sh`: update printed help text to include `test_type=2`. +- Modify `tests/udma/demo/README.md`: document the all-to-all path and example command. +- Optionally modify `tests/udma/demo/ASCEND_VERIFICATION.md` only if it contains an exhaustive test-type list that would become stale. + +--- + +### Task 1: Add Device Kernel And Launch Wrapper + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` + +**Interfaces:** +- Consumes: `TileXR::UDMAPutNbi(args, targetRank, localSrc, byteOffset, byteCount)` and `TileXR::UDMAQuiet(args, targetRank)` from `src/include/tilexr_udma.h`. +- Produces: + ```cpp + void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset); + ``` + +- [ ] **Step 1: Add the kernel before the registered smoke kernel** + +Add this code to `tests/udma/demo/tilexr_udma_demo_kernel.cpp` after `tilexr_udma_put_signal_kernel`: + +```cpp +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, + int32_t elementsPerPeer, uint64_t outputByteOffset) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + bool enabled = TileXR::UDMARegistryEnabled(args); + + if (debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = enabled ? 1 : 0; + debug[4] = elementsPerPeer; + debug[5] = static_cast(outputByteOffset); + } + if (!enabled) { + return; + } + + uint32_t bytes = static_cast(elementsPerPeer * sizeof(int32_t)); + for (int32_t peer = 0; peer < rankSize; ++peer) { + auto localSrc = input + peer * elementsPerPeer; + uint64_t remoteOffset = outputByteOffset + + static_cast(rank) * elementsPerPeer * sizeof(int32_t); + if (peer == rank) { + auto localDst = output + rank * elementsPerPeer; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + localDst[i] = localSrc[i]; + } + continue; + } + TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); + TileXR::UDMAQuiet(args, peer); + } +} +``` + +- [ ] **Step 2: Add the host launch wrapper** + +Add this code near the other `launch_tilexr_udma_*` wrappers: + +```cpp +void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset) +{ + tilexr_udma_all_to_all_kernel<<>>( + commArgs, input, output, debug, elementsPerPeer, outputByteOffset); +} +``` + +- [ ] **Step 3: Run a local source check** + +Run: + +```powershell +rg -n "tilexr_udma_all_to_all|launch_tilexr_udma_all_to_all" tests/udma/demo/tilexr_udma_demo_kernel.cpp +``` + +Expected: one kernel definition and one launch wrapper definition are shown. + +--- + +### Task 2: Add Host-Side All-To-All Layout, Launch, And Validation + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` + +**Interfaces:** +- Consumes: + ```cpp + void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset); + ``` +- Produces: + ```cpp + bool ValidateAllToAllData( + int rank, int rankSize, const std::vector& output, int32_t elementsPerPeer); + ``` + +- [ ] **Step 1: Declare the new launch wrapper** + +Add this declaration after `launch_tilexr_udma_put_signal`: + +```cpp +extern void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset); +``` + +- [ ] **Step 2: Add all-to-all constants** + +Add this constant beside the existing demo constants: + +```cpp +constexpr int32_t kAllToAllBaseValue = 100000; +``` + +- [ ] **Step 3: Add a validator** + +Add this function after `ValidateData`: + +```cpp +bool ValidateAllToAllData( + int rank, int rankSize, const std::vector& output, int32_t elementsPerPeer) +{ + bool ok = true; + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + int32_t expected = kAllToAllBaseValue + srcRank * 1000 + rank; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + size_t offset = static_cast(srcRank) * elementsPerPeer + i; + if (output[offset] != expected) { + std::cerr << "[rank " << rank << "] ALLTOALL MISMATCH at src=" << srcRank + << " elem=" << i << " offset=" << offset + << " got=" << output[offset] << " expected=" << expected << std::endl; + ok = false; + break; + } + } + } + + std::cout << "[rank " << rank << "] alltoall output sample:"; + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + size_t offset = static_cast(srcRank) * elementsPerPeer; + std::cout << " from" << srcRank << "=" << output[offset]; + } + std::cout << std::endl; + return ok; +} +``` + +- [ ] **Step 4: Split payload layout by test type** + +Replace the current data/signal payload sizing block with: + +```cpp +bool isAllToAll = testType == 2; +size_t dataCount = static_cast(rankSize) * elementsPerRank; +size_t dataBytes = dataCount * sizeof(int32_t); +size_t inputOffset = 0; +size_t outputOffset = isAllToAll ? dataBytes : 0; +size_t signalOffset = isAllToAll ? dataBytes * 2 : dataBytes; +size_t signalBytes = static_cast(rankSize) * sizeof(uint64_t); +size_t payloadBytes = signalOffset + signalBytes; +``` + +Keep the existing `registeredBytes` calculation immediately after this block. + +- [ ] **Step 5: Define typed buffer pointers** + +Replace the current `data` and `signals` pointer setup with: + +```cpp +auto data = static_cast(registeredMemory); +auto input = reinterpret_cast(static_cast(registeredMemory) + inputOffset); +auto output = reinterpret_cast(static_cast(registeredMemory) + outputOffset); +auto signals = reinterpret_cast(static_cast(registeredMemory) + signalOffset); +``` + +- [ ] **Step 6: Initialize host buffers for both modes** + +Replace the current `hostData` initialization with: + +```cpp +std::vector hostData(dataCount, -1); +std::vector hostOutput(dataCount, -1); +if (isAllToAll) { + for (int dstRank = 0; dstRank < rankSize; ++dstRank) { + int32_t value = kAllToAllBaseValue + rank * 1000 + dstRank; + std::fill(hostData.begin() + static_cast(dstRank) * elementsPerRank, + hostData.begin() + static_cast(dstRank + 1) * elementsPerRank, + value); + } +} else { + std::fill(hostData.begin() + static_cast(rank) * elementsPerRank, + hostData.begin() + static_cast(rank + 1) * elementsPerRank, + 1000 + rank); +} +``` + +- [ ] **Step 7: Copy input and output buffers to device** + +Replace the data H2D copy part with: + +```cpp +bool initOk = CopyHostToDevice(rank, input, dataCount * sizeof(int32_t), + hostData.data(), dataCount * sizeof(int32_t), isAllToAll ? "alltoall input" : "data"); +if (isAllToAll) { + initOk = CopyHostToDevice(rank, output, dataCount * sizeof(int32_t), + hostOutput.data(), dataCount * sizeof(int32_t), "alltoall output") && initOk; +} +if (!initOk || + !CopyHostToDevice(rank, signals, hostSignals.size() * sizeof(uint64_t), + hostSignals.data(), hostSignals.size() * sizeof(uint64_t), "signals") || + !CopyHostToDevice(rank, debug, hostDebug.size() * sizeof(int32_t), + hostDebug.data(), hostDebug.size() * sizeof(int32_t), "debug")) { +``` + +Keep the existing cleanup block inside this `if`. + +- [ ] **Step 8: Launch the all-to-all kernel for `test_type=2`** + +Replace the launch selection with: + +```cpp +if (testType == 2) { + PrintStatus(rank, "launch all-to-all kernel"); + launch_tilexr_udma_all_to_all( + 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset)); +} else if (testType == 1) { + PrintStatus(rank, "launch put-signal kernel"); + launch_tilexr_udma_put_signal( + 1, stream, commArgsDev, reinterpret_cast(data), reinterpret_cast(signals), + reinterpret_cast(debug), elementsPerRank, kSignalValue); +} else { + PrintStatus(rank, "launch all-gather kernel"); + launch_tilexr_udma_all_gather( + 1, stream, commArgsDev, reinterpret_cast(data), reinterpret_cast(debug), + elementsPerRank); +} +``` + +- [ ] **Step 9: Copy output back for all-to-all** + +Replace the data D2H copy with: + +```cpp +bool copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), + data, dataCount * sizeof(int32_t), "data"); +if (isAllToAll) { + copyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), + output, dataCount * sizeof(int32_t), "alltoall output") && copyBackOk; +} +if (!copyBackOk || + !CopyDeviceToHost(rank, hostSignals.data(), hostSignals.size() * sizeof(uint64_t), + signals, hostSignals.size() * sizeof(uint64_t), "signals") || + !CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), + debug, hostDebug.size() * sizeof(int32_t), "debug")) { +``` + +Keep the existing cleanup block inside this `if`. + +- [ ] **Step 10: Validate all-to-all output** + +Replace the final validation selection with: + +```cpp +bool ok = isAllToAll ? ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerRank) + : ValidateData(rank, rankSize, hostData, elementsPerRank); +if (testType == 1) { + ok = ValidateSignals(rank, rankSize, hostSignals) && ok; +} +``` + +- [ ] **Step 11: Run local source checks** + +Run: + +```powershell +rg -n "all-to-all|alltoall|isAllToAll|ValidateAllToAllData|launch_tilexr_udma_all_to_all" tests/udma/demo/tilexr_udma_demo.cpp +``` + +Expected: declaration, validation function, launch path, and copy/initialization branches are shown. + +--- + +### Task 3: Update Demo Documentation And Script Text + +**Files:** +- Modify: `tests/udma/demo/run_tilexr_udma_demo.sh` +- Modify: `tests/udma/demo/README.md` +- Modify if stale: `tests/udma/demo/ASCEND_VERIFICATION.md` + +**Interfaces:** +- Consumes: `test_type=2` behavior from Tasks 1 and 2. +- Produces: documented command `bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0`. + +- [ ] **Step 1: Update script test-type text** + +In `tests/udma/demo/run_tilexr_udma_demo.sh`, change: + +```bash +echo "Test type: ${test_type} (0=all-gather put, 1=put-signal)" +``` + +to: + +```bash +echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all)" +``` + +- [ ] **Step 2: Update README run examples** + +In `tests/udma/demo/README.md`, add this command to the Run section: + +```bash +bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 +``` + +Update the argument list so it includes: + +```text +- `test_type=2`: all-to-all UDMA put. Rank `src` sends input slice `dst` to rank `dst`; each output is ordered by source rank. +``` + +- [ ] **Step 3: Check `ASCEND_VERIFICATION.md` for stale test-type lists** + +Run: + +```powershell +rg -n "test_type|all-gather|put-signal|all-to-all" tests/udma/demo/ASCEND_VERIFICATION.md +``` + +If it lists only test types 0 and 1, update the list to include: + +```text +test_type=2 validates all-to-all registered-memory UDMA puts with output ordered by source rank. +``` + +- [ ] **Step 4: Run documentation check** + +Run: + +```powershell +rg -n "test_type=2|all-to-all|run_tilexr_udma_demo.sh 2" tests/udma/demo/README.md tests/udma/demo/run_tilexr_udma_demo.sh tests/udma/demo/ASCEND_VERIFICATION.md +``` + +Expected: README and script both mention `test_type=2`. + +--- + +### Task 4: Local Build-Oriented Verification + +**Files:** +- Verify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Verify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` +- Verify: `tests/udma/CMakeLists.txt` + +**Interfaces:** +- Consumes: code from Tasks 1-3. +- Produces: local confidence before remote hardware validation. + +- [ ] **Step 1: Review the diff** + +Run: + +```powershell +git diff -- tests/udma/demo/tilexr_udma_demo.cpp tests/udma/demo/tilexr_udma_demo_kernel.cpp tests/udma/demo/run_tilexr_udma_demo.sh tests/udma/demo/README.md tests/udma/demo/ASCEND_VERIFICATION.md +``` + +Expected: only all-to-all related changes are present. + +- [ ] **Step 2: Check that CMake tracks the modified kernel source** + +Run: + +```powershell +rg -n "tilexr_udma_demo_kernel.cpp|tilexr_udma_demo.cpp|tilexr_udma_demo" tests/udma/CMakeLists.txt +``` + +Expected: existing custom command depends on `demo/tilexr_udma_demo_kernel.cpp`; no CMake changes are required. + +- [ ] **Step 3: Check formatting-sensitive syntax around modified host blocks** + +Run: + +```powershell +rg -n "payloadBytes|registeredBytes|initOk|copyBackOk|testType == 2|ValidateAllToAllData" tests/udma/demo/tilexr_udma_demo.cpp +``` + +Expected: each symbol appears in the expected host flow. + +- [ ] **Step 4: Check working tree before remote copy** + +Run: + +```powershell +git status --short +``` + +Expected: only the planned demo files and plan file are modified or added. + +--- + +### Task 5: Remote Build And Runtime Validation + +**Files:** +- Verify remotely under a new `/home/aiv-perf/` directory on `root@141.61.95.18`. + +**Interfaces:** +- Consumes: local repository changes from Tasks 1-4. +- Produces: remote build and runtime evidence for all-to-all. + +- [ ] **Step 1: Create a unique remote validation directory** + +Run from the local workspace: + +```powershell +ssh root@141.61.95.18 "set -e; d=/home/aiv-perf/tilexr-udma-alltoall-$(date +%Y%m%d-%H%M%S); mkdir -p \"$d\"; echo \"$d\"" +``` + +Expected: command prints the new remote directory path. + +- [ ] **Step 2: Copy the repository to the remote directory** + +Use `rsync` if available: + +```powershell +rsync -a --delete --exclude .git --exclude build --exclude install --exclude tests/udma/build --exclude tests/udma/install --exclude tests/udma/logs ./ root@141.61.95.18:/ +``` + +If `rsync` is unavailable on Windows, use `scp` with a compressed archive created outside the repo build artifacts: + +```powershell +tar --exclude .git --exclude build --exclude install --exclude tests/udma/build --exclude tests/udma/install --exclude tests/udma/logs -czf $env:TEMP\tilexr-udma-alltoall.tgz . +scp $env:TEMP\tilexr-udma-alltoall.tgz root@141.61.95.18:/ +ssh root@141.61.95.18 "set -e; cd ; tar -xzf tilexr-udma-alltoall.tgz; rm tilexr-udma-alltoall.tgz" +``` + +Expected: remote directory contains `scripts/common_env.sh` and `tests/udma/demo`. + +- [ ] **Step 3: Build TileXR core on remote** + +Run: + +```powershell +ssh root@141.61.95.18 "set -e; cd ; source scripts/common_env.sh; mkdir -p build; cd build; cmake -DCMAKE_INSTALL_PREFIX=../install ..; make -j$(nproc); make install" +``` + +Expected: `install/lib/libtile-comm.so` exists under the remote directory. + +- [ ] **Step 4: Build UDMA demo on remote** + +Run: + +```powershell +ssh root@141.61.95.18 "set -e; cd /tests/udma; bash build.sh" +``` + +Expected: `tests/udma/install/bin/tilexr_udma_demo` exists, and the output does not say the demo was skipped because `bisheng` is missing. + +- [ ] **Step 5: Run all-to-all with 2 ranks** + +Run: + +```powershell +ssh root@141.61.95.18 "set -e; cd /tests/udma; bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0" +``` + +Expected: every rank prints `TileXR UDMA demo success`; log tails include `alltoall output sample` with `from0=100000` on rank 0, `from0=100001` on rank 1, `from1=101000` on rank 0, and `from1=101001` on rank 1. + +- [ ] **Step 6: Run optional wider validation if at least 4 NPUs are available** + +Run: + +```powershell +ssh root@141.61.95.18 "set -e; cd /tests/udma; bash demo/run_tilexr_udma_demo.sh 2 4 16 4 0" +``` + +Expected: every rank prints `TileXR UDMA demo success`. + +- [ ] **Step 7: Capture failure logs if validation fails** + +Run: + +```powershell +ssh root@141.61.95.18 "cd /tests/udma; latest=$(ls -td logs/tilexr_udma_demo_* 2>/dev/null | head -1); if [ -n \"$latest\" ]; then for f in \"$latest\"/rank_*.log; do echo \"===== $f =====\"; tail -n 120 \"$f\"; done; fi" +``` + +Expected: logs show whether failure is UDMA enablement, build/runtime environment, kernel debug words, or output mismatch. + +--- + +### Task 6: Final Review And Commit + +**Files:** +- Review all modified files. + +**Interfaces:** +- Consumes: completed implementation and verification evidence. +- Produces: final commit for all-to-all demo implementation. + +- [ ] **Step 1: Inspect final diff** + +Run: + +```powershell +git diff --stat +git diff -- tests/udma/demo/tilexr_udma_demo.cpp tests/udma/demo/tilexr_udma_demo_kernel.cpp tests/udma/demo/run_tilexr_udma_demo.sh tests/udma/demo/README.md tests/udma/demo/ASCEND_VERIFICATION.md +``` + +Expected: changes match the approved spec and no unrelated files are modified. + +- [ ] **Step 2: Check git status** + +Run: + +```powershell +git status --short +``` + +Expected: modified implementation/doc files plus this plan file. + +- [ ] **Step 3: Commit implementation after successful validation** + +Run: + +```powershell +git add tests/udma/demo/tilexr_udma_demo.cpp tests/udma/demo/tilexr_udma_demo_kernel.cpp tests/udma/demo/run_tilexr_udma_demo.sh tests/udma/demo/README.md tests/udma/demo/ASCEND_VERIFICATION.md docs/superpowers/plans/2026-06-18-udma-alltoall-demo.md +git commit -m "feat: add udma alltoall demo" +``` + +Expected: one implementation commit is created. + +--- + +## Self-Review + +- Spec coverage: Tasks 1 and 2 implement the `test_type=2` all-to-all kernel, layout, launch, synchronization, and validation. Task 3 documents the command. Task 5 covers the required remote validation under `/home/aiv-perf/`. +- Placeholder scan: No task uses TBD/TODO/fill-in wording. `` is an execution-time value produced by Task 5 Step 1 and intentionally reused by later commands. +- Type consistency: The launch wrapper signature is identical in Task 1 and Task 2. Host uses `elementsPerRank` as the user-facing argument and passes it as `elementsPerPeer` to the kernel. From cdf6b64053a71e1a3ab3cde6f2402b5622752311 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Fri, 19 Jun 2026 12:00:05 +0800 Subject: [PATCH 003/163] feat: add udma alltoall demo --- src/comm/tilexr_comm.cpp | 14 ++- src/comm/tilexr_internal.cpp | 1 + src/include/tilexr_udma.h | 9 ++ tests/udma/CMakeLists.txt | 31 +++++- tests/udma/build.sh | 22 ++++- tests/udma/demo/ASCEND_VERIFICATION.md | 62 +++++++++++- tests/udma/demo/README.md | 9 ++ tests/udma/demo/run_tilexr_udma_demo.sh | 4 +- tests/udma/demo/tilexr_udma_alltoall_layout.h | 52 ++++++++++ tests/udma/demo/tilexr_udma_demo.cpp | 88 ++++++++++++++--- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 55 +++++++++++ tests/udma/run_tests.sh | 94 ++++++++++++------- .../unit/test_tilexr_chip_map_sources.cpp | 44 +++++++++ .../unit/test_tilexr_ipc_pid_mode_sources.cpp | 54 +++++++++++ .../unit/test_tilexr_udma_alltoall_layout.cpp | 70 ++++++++++++++ 15 files changed, 551 insertions(+), 58 deletions(-) create mode 100644 tests/udma/demo/tilexr_udma_alltoall_layout.h create mode 100644 tests/udma/unit/test_tilexr_chip_map_sources.cpp create mode 100644 tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp create mode 100644 tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 6ee87dca..a80ef2cd 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -909,22 +909,26 @@ int TileXRComm::SetMemoryName(string &name) int TileXRComm::SetIpcPidSdid(string &name, const uint32_t *pids, const int64_t *sdids) const { + const char *modeEnv = std::getenv("TILEXR_IPC_PID_MODE"); + bool forcePid = modeEnv != nullptr && std::string(modeEnv) == "pid"; + bool forceSdid = modeEnv != nullptr && std::string(modeEnv) == "sdid"; + bool defaultSdid = + physicalInfo_.chipName >= ChipName::CHIP_910_9391 && physicalInfo_.chipName < ChipName::CHIP_950; + bool useSdid = forceSdid || (!forcePid && defaultSdid); + TILEXR_LOG(INFO) << "SetIpcPidSdid mode=" << (useSdid ? "sdid" : "pid"); for (int i = 0; i < rankSize_; ++i) { if (i == rank_) { continue; } - if (physicalInfo_.chipName < ChipName::CHIP_910_9391) { - // 910B - int32_t pidInt32 = pids[i]; + int32_t pidInt32 = pids[i]; + if (!useSdid) { int rtRet = rtSetIpcMemPid(name.c_str(), &pidInt32, HCCL_IPC_PID_ARRAY_SIZE); if (rtRet != RT_ERROR_NONE) { TILEXR_LOG(ERROR) << "err " << rtRet; return TILEXR_ERROR_INTERNAL; } } else { - // 910A3 - int32_t pidInt32 = pids[i]; int rtRet = rtSetIpcMemorySuperPodPid(name.c_str(), sdids[i], &pidInt32, HCCL_IPC_PID_ARRAY_SIZE); if (rtRet != RT_ERROR_NONE) { TILEXR_LOG(ERROR) << "err " << rtRet; diff --git a/src/comm/tilexr_internal.cpp b/src/comm/tilexr_internal.cpp index 1e2bc349..243a205c 100644 --- a/src/comm/tilexr_internal.cpp +++ b/src/comm/tilexr_internal.cpp @@ -39,6 +39,7 @@ const std::unordered_map CHIP_MAP = { {"Ascend950DT", ChipName::CHIP_950}, {"Ascend950DT_9581", ChipName::CHIP_950}, {"Ascend950DT_9584", ChipName::CHIP_950}, + {"Ascend950DT_9592", ChipName::CHIP_950}, {"Ascend950PR", ChipName::CHIP_950} }; diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index 43f15c0c..10b85251 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -368,6 +368,15 @@ __aicore__ inline void UDMAQuiet(const __gm__ CommArgs* args, int targetRank) (void)UDMAPollCQ(udmaInfo, targetRank, 0, wqeCnt); } +__aicore__ inline uint32_t UDMAQuietStatus(const __gm__ CommArgs* args, int targetRank) +{ + if (!UDMAEnabled(args)) return 0xFFFFFFFFU; + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, 0); + uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + return UDMAPollCQ(udmaInfo, targetRank, 0, wqeCnt); +} + } // namespace TileXR #endif // TILEXR_UDMA_H diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 6c1e41ab..d4fc2c84 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -43,7 +43,7 @@ message(STATUS "TILEXR_ROOT: ${TILEXR_ROOT}") # 查找库 find_library(TILEXR_LIB tile-comm - HINTS "${TILEXR_ROOT}/install/lib" "${TILEXR_ROOT}/build/src/comm" + HINTS "${TILEXR_ROOT}/install/lib" "${TILEXR_ROOT}/install/lib64" "${TILEXR_ROOT}/build/src/comm" REQUIRED) message(STATUS "Found tile-comm: ${TILEXR_LIB}") @@ -71,6 +71,31 @@ add_executable(test_tilexr_udma_registry target_include_directories(test_tilexr_udma_registry PRIVATE ${TILEXR_ROOT}/src/include + ${CMAKE_CURRENT_SOURCE_DIR} +) + +add_executable(test_tilexr_udma_alltoall_layout + unit/test_tilexr_udma_alltoall_layout.cpp +) + +target_include_directories(test_tilexr_udma_alltoall_layout PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} +) + +add_executable(test_tilexr_chip_map_sources + unit/test_tilexr_chip_map_sources.cpp +) + +target_compile_definitions(test_tilexr_chip_map_sources PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" +) + +add_executable(test_tilexr_ipc_pid_mode_sources + unit/test_tilexr_ipc_pid_mode_sources.cpp +) + +target_compile_definitions(test_tilexr_ipc_pid_mode_sources PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" ) add_executable(test_tilexr_udma_transport_layout @@ -98,6 +123,9 @@ target_link_libraries(test_tilexr_udma set(INSTALL_TARGETS test_tilexr_udma test_tilexr_udma_registry + test_tilexr_udma_alltoall_layout + test_tilexr_chip_map_sources + test_tilexr_ipc_pid_mode_sources test_tilexr_udma_transport_layout ) @@ -175,6 +203,7 @@ if(BUILD_TILEXR_UDMA_DEMO) -shared ${TILEXR_UDMA_KERNEL_LINK_OPTIONS} -DCATLASS_ARCH=${TILEXR_UDMA_CATLASS_ARCH} + -DTILEXR_UDMA_FORCE_ENABLE=1 ${TILEXR_UDMA_DEMO_KERNEL_INCLUDES} "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" -L${ASCEND_DRIVER_PATH}/lib64/driver diff --git a/tests/udma/build.sh b/tests/udma/build.sh index 94866858..d3fb607e 100755 --- a/tests/udma/build.sh +++ b/tests/udma/build.sh @@ -27,8 +27,28 @@ mkdir -p "${INSTALL_DIR}" cd "${BUILD_DIR}" +find_bisheng_dir() { + if command -v bisheng >/dev/null 2>&1; then + dirname "$(command -v bisheng)" + return 0 + fi + for candidate in \ + "${ASCEND_HOME_PATH}/compiler/bisheng" \ + "${ASCEND_HOME_PATH}/tools/bisheng_compiler/bin/bisheng"; do + if [ -x "${candidate}" ]; then + dirname "${candidate}" + return 0 + fi + done + find /usr/local/Ascend -path "*/tools/bisheng_compiler/bin/bisheng" -type f -executable 2>/dev/null | + head -n 1 | + xargs -r dirname +} + # 配置 -if command -v bisheng >/dev/null 2>&1; then +BISHENG_DIR=$(find_bisheng_dir) +if [ -n "${BISHENG_DIR}" ]; then + export PATH="${BISHENG_DIR}:${PATH}" DEMO_OPTION="-DBUILD_TILEXR_UDMA_DEMO=ON" else echo "WARN: bisheng not found; TileXR UDMA communication demo target will be skipped." diff --git a/tests/udma/demo/ASCEND_VERIFICATION.md b/tests/udma/demo/ASCEND_VERIFICATION.md index e7e16433..247d6976 100644 --- a/tests/udma/demo/ASCEND_VERIFICATION.md +++ b/tests/udma/demo/ASCEND_VERIFICATION.md @@ -12,7 +12,7 @@ Verify that the TileXR UDMA demo: - builds against TileXR's public demo API without including `shmem.h` in the host demo source; - initializes UDMA through TileXR's own comm transport, without linking shmem; - registers ordinary `aclrtMalloc` device memory through `TileXRUDMARegister`; -- runs device-side UDMA put and put-signal kernels successfully; +- runs device-side UDMA put, put-signal, and all-to-all kernels successfully; - does not report data mismatches or signal mismatches in rank logs. ## Hardware And Environment @@ -157,14 +157,68 @@ grep -R "TileXR UDMA demo success" "$latest" grep -R "DATA MISMATCH\\|expected non-local signals\\|TileXR UDMA demo failed\\|ERROR" "$latest" || true ``` +## Test 3: UDMA All-To-All + +Run the all-to-all variant: + +```bash +cd /path/to/TileXR/tests/udma +bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 +``` + +Expected: + +- script exits with code 0; +- each rank log contains `TileXR UDMA demo success`; +- each rank log prints `alltoall output sample`; +- rank 0 output sample includes `from0=100000` and `from1=101000`; +- rank 1 output sample includes `from0=100001` and `from1=101001`; +- no log contains `ALLTOALL MISMATCH`, `ERROR`, or `TileXR UDMA demo failed`. + +Quick log check: + +```bash +latest=$(ls -td logs/tilexr_udma_demo_* | head -n1) +grep -R "alltoall output sample" "$latest" +grep -R "TileXR UDMA demo success" "$latest" +grep -R "ALLTOALL MISMATCH\\|TileXR UDMA demo failed\\|ERROR" "$latest" || true +``` + +## IPC PID And SDID Modes + +The communicator supports an override for peer IPC-memory setup: + +```bash +TILEXR_IPC_PID_MODE=pid bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 +TILEXR_IPC_PID_MODE=sdid bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 +``` + +Expected mode behavior: + +- unset: TileXR chooses the chip default; +- `pid`: force `rtSetIpcMemPid`; +- `sdid`: force `rtSetIpcMemorySuperPodPid`. + +On the verified Ascend950DT_9592 host, the default mode is `pid`. Default and +explicit `pid` mode both reached `TileXRCommInitRankLocal success`, +`InitUDMA success`, and `TileXRUDMARegister success`. Explicit `sdid` mode failed +during IPC open with runtime error `507899`, so Ascend950DT_9592 should use PID +mode on that host. + +If default or explicit `pid` mode initializes but a rank's debug words report CQ +status `514`, the failure is in the UDMA data-plane completion path after +registration, not in all-to-all payload layout. The local self-copy segment +should still appear in the all-to-all output sample. + ## Optional Larger Runs -If the machine has more usable devices, repeat both test types with more ranks: +If the machine has more usable devices, repeat all test types with more ranks: ```bash cd /path/to/TileXR/tests/udma bash demo/run_tilexr_udma_demo.sh 0 4 64 4 0 bash demo/run_tilexr_udma_demo.sh 1 4 64 4 0 +bash demo/run_tilexr_udma_demo.sh 2 4 64 4 0 ``` Expected result samples for four ranks should include: @@ -199,9 +253,13 @@ Test 1 command and result: Test 1 log directory: Test 2 command and result: Test 2 log directory: +Test 3 command and result: +Test 3 log directory: +PID/SDID mode results: Optional larger run result: Any ERROR lines: Any DATA MISMATCH lines: +Any ALLTOALL MISMATCH lines: Any signal mismatch lines: ``` diff --git a/tests/udma/demo/README.md b/tests/udma/demo/README.md index 58947072..6a82a35a 100644 --- a/tests/udma/demo/README.md +++ b/tests/udma/demo/README.md @@ -17,6 +17,7 @@ The demo target requires `bisheng`. If `bisheng` is not available, `build.sh` st cd /path/to/TileXR/tests/udma bash demo/run_tilexr_udma_demo.sh 0 2 16 2 0 bash demo/run_tilexr_udma_demo.sh 1 2 16 2 0 +bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 ``` Arguments: @@ -27,6 +28,8 @@ run_tilexr_udma_demo.sh - `test_type=0`: all-gather style UDMA put. - `test_type=1`: UDMA put with signal. +- `test_type=2`: all-to-all UDMA put. Rank `src` sends input slice `dst` to rank `dst`; + each output is ordered by source rank. - `rank_size`: number of local ranks to launch. - `elements_per_rank`: `int32_t` elements in each rank segment. - `npu_count`: number of NPUs available to this run. @@ -34,6 +37,12 @@ run_tilexr_udma_demo.sh Each run writes per-rank logs under `tests/udma/logs/tilexr_udma_demo_*`. +IPC peer-memory setup can be forced with `TILEXR_IPC_PID_MODE`: + +- unset: use TileXR's chip default. Ascend950-class chips use `pid`. +- `pid`: force `rtSetIpcMemPid`. +- `sdid`: force `rtSetIpcMemorySuperPodPid`. + Run this demo only on A5 / Ascend950 / 950 hardware. Builds or smoke tests on other Ascend chips are not valid UDMA runtime validation. ## What To Check diff --git a/tests/udma/demo/run_tilexr_udma_demo.sh b/tests/udma/demo/run_tilexr_udma_demo.sh index b3e193d8..2eae6171 100755 --- a/tests/udma/demo/run_tilexr_udma_demo.sh +++ b/tests/udma/demo/run_tilexr_udma_demo.sh @@ -21,7 +21,7 @@ source "${TILEXR_ROOT}/scripts/common_env.sh" export TILEXR_COMM_ID=${TILEXR_COMM_ID:-127.0.0.1:10067} export TILEXR_DEMO_NPUS=${npu_count} export TILEXR_DEMO_FIRST_NPU=${first_npu} -export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:${TILEXR_ROOT}/install/lib:/usr/local/lib:${LD_LIBRARY_PATH:-}" +export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:${INSTALL_DIR}/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/lib:${LD_LIBRARY_PATH:-}" bin="${INSTALL_DIR}/bin/tilexr_udma_demo" if [ ! -x "${bin}" ]; then @@ -36,7 +36,7 @@ echo "==========================================" echo " TileXR UDMA Communication Demo" echo "==========================================" echo "Binary: ${bin}" -echo "Test type: ${test_type} (0=all-gather put, 1=put-signal)" +echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all)" echo "Rank size: ${rank_size}" echo "Elements/rank: ${elements_per_rank}" echo "NPU count: ${npu_count}" diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h new file mode 100644 index 00000000..e8b5846d --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_ALLTOALL_LAYOUT_H +#define TILEXR_UDMA_ALLTOALL_LAYOUT_H + +#include +#include +#include +#include + +namespace TileXR { +namespace Demo { + +constexpr int32_t kAllToAllBaseValue = 100000; + +inline int32_t AllToAllValue(int srcRank, int dstRank) +{ + return kAllToAllBaseValue + srcRank * 1000 + dstRank; +} + +inline void FillAllToAllInput( + std::vector& input, int rank, int rankSize, int32_t elementsPerPeer) +{ + for (int dstRank = 0; dstRank < rankSize; ++dstRank) { + std::fill(input.begin() + static_cast(dstRank) * elementsPerPeer, + input.begin() + static_cast(dstRank + 1) * elementsPerPeer, + AllToAllValue(rank, dstRank)); + } +} + +inline bool ValidateAllToAllOutput( + const std::vector& output, int rank, int rankSize, int32_t elementsPerPeer) +{ + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + const int32_t expected = AllToAllValue(srcRank, rank); + for (int32_t i = 0; i < elementsPerPeer; ++i) { + const size_t offset = static_cast(srcRank) * elementsPerPeer + i; + if (output[offset] != expected) { + return false; + } + } + } + return true; +} + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_ALLTOALL_LAYOUT_H diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index eabe662b..0f34c03c 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -21,12 +21,16 @@ #include "acl/acl.h" #include "tilexr_api.h" #include "tilexr_types.h" +#include "tilexr_udma_alltoall_layout.h" extern void launch_tilexr_udma_all_gather( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR data, GM_ADDR debug, int32_t elementsPerRank); extern void launch_tilexr_udma_put_signal( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR data, GM_ADDR signals, GM_ADDR debug, int32_t elementsPerRank, uint64_t signal); +extern void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset); namespace { constexpr int32_t kDefaultElementsPerRank = 16; @@ -341,6 +345,33 @@ bool ValidateData(int rank, int rankSize, const std::vector& data, int3 return ok; } +bool ValidateAllToAllData( + int rank, int rankSize, const std::vector& output, int32_t elementsPerPeer) +{ + bool ok = true; + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + int32_t expected = TileXR::Demo::AllToAllValue(srcRank, rank); + for (int32_t i = 0; i < elementsPerPeer; ++i) { + size_t offset = static_cast(srcRank) * elementsPerPeer + i; + if (output[offset] != expected) { + std::cerr << "[rank " << rank << "] ALLTOALL MISMATCH at src=" << srcRank + << " elem=" << i << " offset=" << offset + << " got=" << output[offset] << " expected=" << expected << std::endl; + ok = false; + break; + } + } + } + + std::cout << "[rank " << rank << "] alltoall output sample:"; + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + size_t offset = static_cast(srcRank) * elementsPerPeer; + std::cout << " from" << srcRank << "=" << output[offset]; + } + std::cout << std::endl; + return TileXR::Demo::ValidateAllToAllOutput(output, rank, rankSize, elementsPerPeer) && ok; +} + bool ValidateSignals(int rank, int rankSize, const std::vector& signals) { bool ok = true; @@ -446,11 +477,14 @@ int main(int argc, char** argv) return 1; } + bool isAllToAll = testType == 2; size_t dataCount = static_cast(rankSize) * elementsPerRank; size_t dataBytes = dataCount * sizeof(int32_t); + size_t inputOffset = 0; + size_t outputOffset = isAllToAll ? dataBytes : 0; size_t signalBytes = static_cast(rankSize) * sizeof(uint64_t); - size_t signalOffset = dataBytes; - size_t payloadBytes = dataBytes + signalBytes; + size_t signalOffset = isAllToAll ? dataBytes * 2 : dataBytes; + size_t payloadBytes = signalOffset + signalBytes; size_t registeredBytes = ((payloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; if (!CheckAcl(rank, "aclrtMalloc debug", aclrtMalloc(reinterpret_cast(&debug), @@ -461,6 +495,8 @@ int main(int argc, char** argv) return 1; } auto data = static_cast(registeredMemory); + auto input = reinterpret_cast(static_cast(registeredMemory) + inputOffset); + auto output = reinterpret_cast(static_cast(registeredMemory) + outputOffset); auto signals = reinterpret_cast(static_cast(registeredMemory) + signalOffset); if (!CheckTileXR(rank, "TileXRUDMARegister", TileXRUDMARegister(comm, static_cast(registeredMemory), registeredBytes, &udmaHandle))) { @@ -470,18 +506,30 @@ int main(int argc, char** argv) udmaRegistered = true; PrintStatus(rank, "registered UDMA memory base=" + std::to_string(reinterpret_cast(registeredMemory)) + " bytes=" + std::to_string(registeredBytes) + - " dataOffset=0 signalOffset=" + std::to_string(signalOffset)); + " inputOffset=" + std::to_string(inputOffset) + + " outputOffset=" + std::to_string(outputOffset) + + " signalOffset=" + std::to_string(signalOffset)); PrintCommArgs(rank, *commArgsHost, commArgsDev); std::vector hostData(dataCount, -1); - std::fill(hostData.begin() + static_cast(rank) * elementsPerRank, - hostData.begin() + static_cast(rank + 1) * elementsPerRank, - 1000 + rank); + std::vector hostOutput(dataCount, -1); + if (isAllToAll) { + TileXR::Demo::FillAllToAllInput(hostData, rank, rankSize, elementsPerRank); + } else { + std::fill(hostData.begin() + static_cast(rank) * elementsPerRank, + hostData.begin() + static_cast(rank + 1) * elementsPerRank, + 1000 + rank); + } std::vector hostSignals(static_cast(rankSize), 0); std::vector hostDebug(kDebugWords, 0); - if (!CopyHostToDevice(rank, data, dataCount * sizeof(int32_t), - hostData.data(), dataCount * sizeof(int32_t), "data") || + bool initOk = CopyHostToDevice(rank, input, dataCount * sizeof(int32_t), + hostData.data(), dataCount * sizeof(int32_t), isAllToAll ? "alltoall input" : "data"); + if (isAllToAll) { + initOk = CopyHostToDevice(rank, output, dataCount * sizeof(int32_t), + hostOutput.data(), dataCount * sizeof(int32_t), "alltoall output") && initOk; + } + if (!initOk || !CopyHostToDevice(rank, signals, hostSignals.size() * sizeof(uint64_t), hostSignals.data(), hostSignals.size() * sizeof(uint64_t), "signals") || !CopyHostToDevice(rank, debug, hostDebug.size() * sizeof(int32_t), @@ -500,12 +548,18 @@ int main(int argc, char** argv) return 1; } - PrintStatus(rank, testType == 1 ? "launch put-signal kernel" : "launch all-gather kernel"); - if (testType == 1) { + if (testType == 2) { + PrintStatus(rank, "launch all-to-all kernel"); + launch_tilexr_udma_all_to_all( + 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset)); + } else if (testType == 1) { + PrintStatus(rank, "launch put-signal kernel"); launch_tilexr_udma_put_signal( 1, stream, commArgsDev, reinterpret_cast(data), reinterpret_cast(signals), reinterpret_cast(debug), elementsPerRank, kSignalValue); } else { + PrintStatus(rank, "launch all-gather kernel"); launch_tilexr_udma_all_gather( 1, stream, commArgsDev, reinterpret_cast(data), reinterpret_cast(debug), elementsPerRank); @@ -525,8 +579,13 @@ int main(int argc, char** argv) return 1; } - if (!CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), - data, dataCount * sizeof(int32_t), "data") || + bool copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), + data, dataCount * sizeof(int32_t), "data"); + if (isAllToAll) { + copyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), + output, dataCount * sizeof(int32_t), "alltoall output") && copyBackOk; + } + if (!copyBackOk || !CopyDeviceToHost(rank, hostSignals.data(), hostSignals.size() * sizeof(uint64_t), signals, hostSignals.size() * sizeof(uint64_t), "signals") || !CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), @@ -539,12 +598,13 @@ int main(int argc, char** argv) } std::cout << "[rank " << rank << "] debug words:"; - for (size_t i = 0; i < std::min(5, hostDebug.size()); ++i) { + for (size_t i = 0; i < std::min(10, hostDebug.size()); ++i) { std::cout << " d" << i << "=" << hostDebug[i]; } std::cout << std::endl; - bool ok = ValidateData(rank, rankSize, hostData, elementsPerRank); + bool ok = isAllToAll ? ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerRank) + : ValidateData(rank, rankSize, hostData, elementsPerRank); if (testType == 1) { ok = ValidateSignals(rank, rankSize, hostSignals) && ok; } diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index a49ca3c9..1d0ebcdc 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -80,6 +80,53 @@ extern "C" __global__ __aicore__ void tilexr_udma_put_signal_kernel( } } +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, + int32_t elementsPerPeer, uint64_t outputByteOffset) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + bool enabled = TileXR::UDMARegistryEnabled(args); + + if (debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = enabled ? 1 : 0; + debug[4] = elementsPerPeer; + debug[5] = static_cast(outputByteOffset); + } + if (!enabled) { + return; + } + + auto selfSrc = input + rank * elementsPerPeer; + auto selfDst = output + rank * elementsPerPeer; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + selfDst[i] = selfSrc[i]; + } + + uint32_t bytes = static_cast(elementsPerPeer * sizeof(int32_t)); + for (int32_t peer = 0; peer < rankSize; ++peer) { + if (peer == rank) { + continue; + } + auto localSrc = input + peer * elementsPerPeer; + uint64_t remoteOffset = outputByteOffset + + static_cast(rank) * elementsPerPeer * sizeof(int32_t); + TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); + uint32_t status = TileXR::UDMAQuietStatus(args, peer); + if (debug != nullptr) { + debug[6 + peer] = static_cast(status); + } + } +} + extern "C" __global__ __aicore__ void tilexr_udma_registered_smoke_kernel( GM_ADDR commArgsGM, GM_ADDR localGM, GM_ADDR debugGM, uint32_t bytes, uint64_t signal) { @@ -125,6 +172,14 @@ void launch_tilexr_udma_put_signal( commArgs, data, signals, debug, elementsPerRank, signal); } +void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset) +{ + tilexr_udma_all_to_all_kernel<<>>( + commArgs, input, output, debug, elementsPerPeer, outputByteOffset); +} + void launch_tilexr_udma_registered_smoke( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR local, GM_ADDR debug, uint32_t bytes, uint64_t signal) { diff --git a/tests/udma/run_tests.sh b/tests/udma/run_tests.sh index f020b583..41e610bf 100755 --- a/tests/udma/run_tests.sh +++ b/tests/udma/run_tests.sh @@ -1,6 +1,6 @@ #!/bin/bash # -# 运行 UDMA 测试 +# Run UDMA tests. # set -e @@ -9,11 +9,13 @@ SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) TILEXR_ROOT="${SCRIPT_DIR}/../.." INSTALL_DIR="${SCRIPT_DIR}/install" -# 加载环境 source "${TILEXR_ROOT}/scripts/common_env.sh" -# 设置 LD_LIBRARY_PATH：优先使用当前仓库刚编译安装的库，避免被 /usr/local/lib 中的旧库覆盖 -export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:${TILEXR_ROOT}/install/lib:/usr/local/lib:${LD_LIBRARY_PATH}" +if [ -x /usr/local/mpi/bin/mpirun ]; then + export PATH="/usr/local/mpi/bin:${PATH}" +fi + +export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:${INSTALL_DIR}/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/lib:${LD_LIBRARY_PATH:-}" echo "==========================================" echo " Running UDMA Tests" @@ -26,7 +28,8 @@ detect_ok_npus() { local ids=() for id in $(seq 0 15); do local health - health=$(npu-smi info -t health -i "${id}" 2>/dev/null | awk -F: '/Health Status/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2; exit}') + health=$(npu-smi info -t health -i "${id}" 2>/dev/null | + awk -F: '/Health Status/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2; exit}') if [ "${health}" = "OK" ]; then ids+=("${id}") fi @@ -43,15 +46,22 @@ if [ -n "${TILEXR_TEST_DEVICES:-}" ]; then echo "TILEXR_TEST_DEVICES: ${TILEXR_TEST_DEVICES}" fi -# 检查测试二进制是否存在 -if [ ! -f "${INSTALL_DIR}/bin/test_tilexr_udma_transport_layout" ] || - [ ! -f "${INSTALL_DIR}/bin/test_tilexr_udma_registry" ] || - [ ! -f "${INSTALL_DIR}/bin/test_tilexr_udma" ]; then - echo "ERROR: Test binaries not found. Please run build.sh first." - exit 1 -fi +required_bins=( + test_tilexr_udma_transport_layout + test_tilexr_udma_registry + test_tilexr_udma_alltoall_layout + test_tilexr_chip_map_sources + test_tilexr_ipc_pid_mode_sources + test_tilexr_udma +) + +for bin in "${required_bins[@]}"; do + if [ ! -f "${INSTALL_DIR}/bin/${bin}" ]; then + echo "ERROR: ${INSTALL_DIR}/bin/${bin} not found. Please run build.sh first." + exit 1 + fi +done -# 测试 1: UDMA info layout 单元测试（host-only） echo "==========================================" echo "Test 1: TileXR UDMA Transport Layout Unit Test" echo "==========================================" @@ -59,7 +69,6 @@ echo "==========================================" TEST1_RESULT=$? echo "" -# 测试 2: TileXR UDMA registry 单元测试（host-only） echo "==========================================" echo "Test 2: TileXR UDMA Registry Unit Test" echo "==========================================" @@ -67,24 +76,41 @@ echo "==========================================" TEST2_RESULT=$? echo "" -# 测试 3: TileXR 集成测试（单进程，单卡） echo "==========================================" -echo "Test 3: TileXR Integration Tests (Single Process)" +echo "Test 3: TileXR UDMA All-To-All Layout Unit Test" +echo "==========================================" +"${INSTALL_DIR}/bin/test_tilexr_udma_alltoall_layout" +TEST3_RESULT=$? +echo "" + +echo "==========================================" +echo "Test 4: TileXR Chip Map Source Unit Test" +echo "==========================================" +"${INSTALL_DIR}/bin/test_tilexr_chip_map_sources" +TEST4_RESULT=$? +echo "" + +echo "==========================================" +echo "Test 5: TileXR IPC PID Mode Source Unit Test" +echo "==========================================" +"${INSTALL_DIR}/bin/test_tilexr_ipc_pid_mode_sources" +TEST5_RESULT=$? +echo "" + +echo "==========================================" +echo "Test 6: TileXR Integration Tests (Single Process)" echo "==========================================" export RANK=0 export RANK_SIZE=1 "${INSTALL_DIR}/bin/test_tilexr_udma" -TEST3_RESULT=$? +TEST6_RESULT=$? echo "" -# 测试 4: TileXR 多进程测试（需要 mpirun） echo "==========================================" -echo "Test 4: TileXR Multi-Process Tests (MPI)" +echo "Test 7: TileXR Multi-Process Tests (MPI)" echo "==========================================" -# 检查是否有 mpirun -if command -v mpirun &> /dev/null; then - # 检测可用的 NPU 数量 +if command -v mpirun >/dev/null 2>&1; then NPU_COUNT=${TILEXR_ASCEND_DEV_NUM:-0} echo "Detected ${NPU_COUNT} NPU(s)" @@ -98,30 +124,32 @@ if command -v mpirun &> /dev/null; then unset RANK unset RANK_SIZE mpirun -n 2 "${INSTALL_DIR}/bin/test_tilexr_udma" - TEST4_RESULT=$? + TEST7_RESULT=$? else echo "SKIP: Need at least 2 usable NPUs for multi-rank test" - TEST4_RESULT=0 + TEST7_RESULT=0 fi else echo "SKIP: mpirun not found, skipping multi-process tests" - TEST4_RESULT=0 + TEST7_RESULT=0 fi echo "" -# 汇总结果 echo "==========================================" echo " Test Results Summary" echo "==========================================" -echo "Test 1 (UDMA Layout): $([ $TEST1_RESULT -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 2 (UDMA Registry): $([ $TEST2_RESULT -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 3 (TileXR Single): $([ $TEST3_RESULT -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 4 (TileXR Multi): $([ $TEST4_RESULT -eq 0 ] && echo 'PASS' || echo 'SKIP/FAIL')" +echo "Test 1 (UDMA Layout): $([ ${TEST1_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 2 (UDMA Registry): $([ ${TEST2_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 3 (AllToAll Layout): $([ ${TEST3_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 4 (Chip Map): $([ ${TEST4_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 5 (IPC PID Mode): $([ ${TEST5_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 6 (TileXR Single): $([ ${TEST6_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 7 (TileXR Multi): $([ ${TEST7_RESULT} -eq 0 ] && echo 'PASS' || echo 'SKIP/FAIL')" echo "==========================================" -# 返回失败状态 -if [ $TEST1_RESULT -ne 0 ] || [ $TEST2_RESULT -ne 0 ] || [ $TEST3_RESULT -ne 0 ] || - [ $TEST4_RESULT -ne 0 ]; then +if [ ${TEST1_RESULT} -ne 0 ] || [ ${TEST2_RESULT} -ne 0 ] || [ ${TEST3_RESULT} -ne 0 ] || + [ ${TEST4_RESULT} -ne 0 ] || [ ${TEST5_RESULT} -ne 0 ] || [ ${TEST6_RESULT} -ne 0 ] || + [ ${TEST7_RESULT} -ne 0 ]; then exit 1 fi diff --git a/tests/udma/unit/test_tilexr_chip_map_sources.cpp b/tests/udma/unit/test_tilexr_chip_map_sources.cpp new file mode 100644 index 00000000..ed283191 --- /dev/null +++ b/tests/udma/unit/test_tilexr_chip_map_sources.cpp @@ -0,0 +1,44 @@ +#include +#include +#include +#include + +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + +namespace { + +int g_failures = 0; + +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +} // namespace + +int main() +{ + const std::string internal = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/tilexr_internal.cpp"); + CHECK_CONTAINS(internal, "\"Ascend950DT_9592\""); + CHECK_CONTAINS(internal, "ChipName::CHIP_950"); + + if (g_failures != 0) { + std::cerr << g_failures << " chip map source checks failed" << std::endl; + return 1; + } + std::cout << "TileXR chip map source checks passed" << std::endl; + return 0; +} diff --git a/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp b/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp new file mode 100644 index 00000000..394ee252 --- /dev/null +++ b/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp @@ -0,0 +1,54 @@ +#include +#include +#include +#include + +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + +namespace { + +int g_failures = 0; + +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +} // namespace + +int main() +{ + const std::string comm = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/tilexr_comm.cpp"); + const std::string cmake = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/CMakeLists.txt"); + CHECK_CONTAINS(comm, "TILEXR_IPC_PID_MODE"); + CHECK_CONTAINS(comm, "physicalInfo_.chipName < ChipName::CHIP_950"); + CHECK_CONTAINS(comm, "rtSetIpcMemPid"); + CHECK_CONTAINS(comm, "rtSetIpcMemorySuperPodPid"); + CHECK_CONTAINS(cmake, "TILEXR_UDMA_FORCE_ENABLE"); + + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + const std::string udma = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); + CHECK_CONTAINS(transport, "RaCtxRmemImport"); + CHECK_CONTAINS(udma, "UDMAQuietStatus"); + + if (g_failures != 0) { + std::cerr << g_failures << " IPC PID mode source checks failed" << std::endl; + return 1; + } + std::cout << "TileXR IPC PID mode source checks passed" << std::endl; + return 0; +} diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp new file mode 100644 index 00000000..6bef28a5 --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -0,0 +1,70 @@ +#include +#include +#include +#include + +#include "demo/tilexr_udma_alltoall_layout.h" + +namespace { + +int g_failures = 0; + +#define CHECK_EQ(lhs, rhs) \ + do { \ + auto lhsValue = (lhs); \ + auto rhsValue = (rhs); \ + if (lhsValue != rhsValue) { \ + std::cerr << "CHECK_EQ failed at line " << __LINE__ << ": " #lhs " != " #rhs \ + << " (" << lhsValue << " vs " << rhsValue << ")" << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +void TestAllToAllInputPattern() +{ + constexpr int rank = 2; + constexpr int rankSize = 4; + constexpr int32_t elementsPerPeer = 3; + std::vector input(static_cast(rankSize) * elementsPerPeer, -1); + + TileXR::Demo::FillAllToAllInput(input, rank, rankSize, elementsPerPeer); + + for (int dstRank = 0; dstRank < rankSize; ++dstRank) { + int32_t expected = TileXR::Demo::AllToAllValue(rank, dstRank); + for (int32_t elem = 0; elem < elementsPerPeer; ++elem) { + CHECK_EQ(input[static_cast(dstRank) * elementsPerPeer + elem], expected); + } + } +} + +void TestAllToAllOutputValidation() +{ + constexpr int rank = 1; + constexpr int rankSize = 3; + constexpr int32_t elementsPerPeer = 2; + std::vector output(static_cast(rankSize) * elementsPerPeer, -1); + + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + std::fill(output.begin() + static_cast(srcRank) * elementsPerPeer, + output.begin() + static_cast(srcRank + 1) * elementsPerPeer, + TileXR::Demo::AllToAllValue(srcRank, rank)); + } + + CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(output, rank, rankSize, elementsPerPeer), true); + output[static_cast(2) * elementsPerPeer + 1] = 123; + CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(output, rank, rankSize, elementsPerPeer), false); +} + +} // namespace + +int main() +{ + TestAllToAllInputPattern(); + TestAllToAllOutputValidation(); + if (g_failures != 0) { + std::cerr << g_failures << " all-to-all layout checks failed" << std::endl; + return 1; + } + std::cout << "TileXR UDMA all-to-all layout checks passed" << std::endl; + return 0; +} From e5527363dfe8fff8eda85691343f1fd9b0f47238 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 22 Jun 2026 10:04:48 +0800 Subject: [PATCH 004/163] Add UDMA demo checker collectives --- src/comm/tilexr_comm.cpp | 37 ++++- src/comm/tilexr_internal.cpp | 3 +- src/include/comm_args.h | 2 +- tests/udma/CMakeLists.txt | 17 ++ tests/udma/build.sh | 1 + tests/udma/demo/README.md | 3 + tests/udma/demo/run_tilexr_udma_demo.sh | 2 +- .../udma/demo/tilexr_udma_allreduce_layout.h | 63 +++++++ tests/udma/demo/tilexr_udma_alltoall_layout.h | 14 ++ tests/udma/demo/tilexr_udma_demo.cpp | 154 ++++++++++++++++-- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 124 +++++++++++++- tests/udma/run_tests.sh | 39 +++-- .../unit/test_tilexr_chip_map_sources.cpp | 1 + .../unit/test_tilexr_ipc_pid_mode_sources.cpp | 14 ++ .../test_tilexr_udma_allreduce_layout.cpp | 125 ++++++++++++++ .../unit/test_tilexr_udma_alltoall_layout.cpp | 57 +++++++ tests/udma/unit/test_tilexr_udma_registry.cpp | 6 + 17 files changed, 627 insertions(+), 35 deletions(-) create mode 100644 tests/udma/demo/tilexr_udma_allreduce_layout.h create mode 100644 tests/udma/unit/test_tilexr_udma_allreduce_layout.cpp diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index a80ef2cd..33d446b7 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -646,10 +646,6 @@ int TileXRComm::EnablePeerAccess() } else if (physicalInfo_.physicalLink == PhysicalLink::RESERVED) { physicalInfo_.physicalLink = PhysicalLink::PCIE; commArgs_.extraFlag |= ExtraFlag::TOPO_PCIE; - if (rankSize_ > PING_PONG_SIZE) { - TILEXR_LOG(ERROR) << "do not support pcie > 2 rank! rankSize_ = " << rankSize_; - return TILEXR_ERROR_INTERNAL; - } } physicalInfo_.coreNum = GetCoreNum(physicalInfo_.chipName); @@ -864,6 +860,30 @@ int TileXRComm::InitCommMem() } if (OpenIpcMem(names) != TILEXR_SUCCESS) { + const char *modeEnv = std::getenv("TILEXR_IPC_PID_MODE"); + const bool forceSdid = modeEnv != nullptr && std::string(modeEnv) == "sdid"; + if (forceSdid) { + TILEXR_LOG(WARN) << "OpenIpcMem failed after sdid setup, retry with pid setup"; + string retryName; + if (setenv("TILEXR_IPC_PID_MODE", "pid_retry", 1) != 0 || + SetMemoryName(retryName) != TILEXR_SUCCESS || + SetIpcPidSdid(retryName, pids, sdids) != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "SetIpcPidSdid pid retry failed!"; + setenv("TILEXR_IPC_PID_MODE", "sdid", 1); + return TILEXR_ERROR_INTERNAL; + } + retryName.resize(IPC_NAME_SIZE); + ret = GetName(retryName, names); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "GetName pid retry error! ret: " << ret; + setenv("TILEXR_IPC_PID_MODE", "sdid", 1); + return ret; + } + setenv("TILEXR_IPC_PID_MODE", "sdid", 1); + if (OpenIpcMem(names) == TILEXR_SUCCESS) { + return TILEXR_SUCCESS; + } + } TILEXR_LOG(ERROR) << "rank: " << rank_ << " OpenIpcMem failed!"; return TILEXR_ERROR_INTERNAL; } @@ -931,8 +951,13 @@ int TileXRComm::SetIpcPidSdid(string &name, const uint32_t *pids, const int64_t } else { int rtRet = rtSetIpcMemorySuperPodPid(name.c_str(), sdids[i], &pidInt32, HCCL_IPC_PID_ARRAY_SIZE); if (rtRet != RT_ERROR_NONE) { - TILEXR_LOG(ERROR) << "err " << rtRet; - return TILEXR_ERROR_INTERNAL; + TILEXR_LOG(WARN) << "rtSetIpcMemorySuperPodPid err " << rtRet + << ", fallback to rtSetIpcMemPid"; + rtRet = rtSetIpcMemPid(name.c_str(), &pidInt32, HCCL_IPC_PID_ARRAY_SIZE); + if (rtRet != RT_ERROR_NONE) { + TILEXR_LOG(ERROR) << "err " << rtRet; + return TILEXR_ERROR_INTERNAL; + } } } } diff --git a/src/comm/tilexr_internal.cpp b/src/comm/tilexr_internal.cpp index 243a205c..0a8494f4 100644 --- a/src/comm/tilexr_internal.cpp +++ b/src/comm/tilexr_internal.cpp @@ -40,7 +40,8 @@ const std::unordered_map CHIP_MAP = { {"Ascend950DT_9581", ChipName::CHIP_950}, {"Ascend950DT_9584", ChipName::CHIP_950}, {"Ascend950DT_9592", ChipName::CHIP_950}, - {"Ascend950PR", ChipName::CHIP_950} + {"Ascend950PR", ChipName::CHIP_950}, + {"Ascend950PR_9599", ChipName::CHIP_950} }; /** diff --git a/src/include/comm_args.h b/src/include/comm_args.h index ed46f825..74931a6b 100644 --- a/src/include/comm_args.h +++ b/src/include/comm_args.h @@ -34,7 +34,7 @@ using GM_ADDR = uint8_t*; namespace TileXR { -constexpr int TILEXR_MAX_RANK_SIZE = 128; // 最大支持的npu卡数 +constexpr int TILEXR_MAX_RANK_SIZE = 256; // 最大支持的npu卡数 constexpr int RANK_SIZE_TWO = 2; // 可用SIO的规模，以及是否需要跨卡搬运数据核的分界规模 constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // 前2MB作为flag标志位，之后100MB作为数据存储 diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index d4fc2c84..ff98ec3d 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -82,6 +82,22 @@ target_include_directories(test_tilexr_udma_alltoall_layout PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ) +target_compile_definitions(test_tilexr_udma_alltoall_layout PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" +) + +add_executable(test_tilexr_udma_allreduce_layout + unit/test_tilexr_udma_allreduce_layout.cpp +) + +target_include_directories(test_tilexr_udma_allreduce_layout PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} +) + +target_compile_definitions(test_tilexr_udma_allreduce_layout PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" +) + add_executable(test_tilexr_chip_map_sources unit/test_tilexr_chip_map_sources.cpp ) @@ -124,6 +140,7 @@ set(INSTALL_TARGETS test_tilexr_udma test_tilexr_udma_registry test_tilexr_udma_alltoall_layout + test_tilexr_udma_allreduce_layout test_tilexr_chip_map_sources test_tilexr_ipc_pid_mode_sources test_tilexr_udma_transport_layout diff --git a/tests/udma/build.sh b/tests/udma/build.sh index d3fb607e..1ee2bb36 100755 --- a/tests/udma/build.sh +++ b/tests/udma/build.sh @@ -72,6 +72,7 @@ echo "" echo "Available tests:" echo " - test_tilexr_udma_transport_layout : UDMA info layout unit tests" echo " - test_tilexr_udma_registry : registered-memory metadata unit tests" +echo " - test_tilexr_udma_allreduce_layout : all-reduce demo layout unit tests" echo " - test_tilexr_udma : TileXR integration tests" if [ -f "${INSTALL_DIR}/bin/tilexr_udma_demo" ]; then echo " - tilexr_udma_demo : TileXR UDMA communication demo" diff --git a/tests/udma/demo/README.md b/tests/udma/demo/README.md index 6a82a35a..37f7103c 100644 --- a/tests/udma/demo/README.md +++ b/tests/udma/demo/README.md @@ -18,6 +18,7 @@ cd /path/to/TileXR/tests/udma bash demo/run_tilexr_udma_demo.sh 0 2 16 2 0 bash demo/run_tilexr_udma_demo.sh 1 2 16 2 0 bash demo/run_tilexr_udma_demo.sh 2 2 16 2 0 +bash demo/run_tilexr_udma_demo.sh 3 8 16 8 0 ``` Arguments: @@ -30,6 +31,8 @@ run_tilexr_udma_demo.sh - `test_type=1`: UDMA put with signal. - `test_type=2`: all-to-all UDMA put. Rank `src` sends input slice `dst` to rank `dst`; each output is ordered by source rank. +- `test_type=3`: all-reduce sum. Each rank contributes one local vector and receives + the element-wise sum across all ranks. - `rank_size`: number of local ranks to launch. - `elements_per_rank`: `int32_t` elements in each rank segment. - `npu_count`: number of NPUs available to this run. diff --git a/tests/udma/demo/run_tilexr_udma_demo.sh b/tests/udma/demo/run_tilexr_udma_demo.sh index 2eae6171..ae31f6a0 100755 --- a/tests/udma/demo/run_tilexr_udma_demo.sh +++ b/tests/udma/demo/run_tilexr_udma_demo.sh @@ -36,7 +36,7 @@ echo "==========================================" echo " TileXR UDMA Communication Demo" echo "==========================================" echo "Binary: ${bin}" -echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all)" +echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all, 3=all-reduce)" echo "Rank size: ${rank_size}" echo "Elements/rank: ${elements_per_rank}" echo "NPU count: ${npu_count}" diff --git a/tests/udma/demo/tilexr_udma_allreduce_layout.h b/tests/udma/demo/tilexr_udma_allreduce_layout.h new file mode 100644 index 00000000..ca3adf87 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_allreduce_layout.h @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_ALLREDUCE_LAYOUT_H +#define TILEXR_UDMA_ALLREDUCE_LAYOUT_H + +#include +#include +#include +#include + +namespace TileXR { +namespace Demo { + +constexpr int32_t kAllReduceBaseValue = 1000; + +inline int32_t AllReduceValue(int rank) +{ + return kAllReduceBaseValue + rank; +} + +inline int32_t AllReduceExpectedSum(int rankSize) +{ + return rankSize * kAllReduceBaseValue + rankSize * (rankSize - 1) / 2; +} + +inline void FillAllReduceInput( + std::vector& input, int rank, int32_t elementsPerRank) +{ + std::fill(input.begin(), input.begin() + elementsPerRank, AllReduceValue(rank)); +} + +inline bool ValidateAllReduceOutput( + const std::vector& output, int rankSize, int32_t elementsPerRank) +{ + const int32_t expected = AllReduceExpectedSum(rankSize); + for (int32_t i = 0; i < elementsPerRank; ++i) { + if (output[static_cast(i)] != expected) { + return false; + } + } + return true; +} + +inline void BuildAllReduceOutputFromInputs( + const std::vector& allInputs, int rankSize, int32_t elementsPerRank, + std::vector& output) +{ + std::fill(output.begin(), output.begin() + elementsPerRank, 0); + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + const size_t srcBase = static_cast(srcRank) * elementsPerRank; + for (int32_t i = 0; i < elementsPerRank; ++i) { + output[static_cast(i)] += allInputs[srcBase + i]; + } + } +} + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_ALLREDUCE_LAYOUT_H diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index e8b5846d..da4e89dd 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -46,6 +46,20 @@ inline bool ValidateAllToAllOutput( return true; } +inline void BuildAllToAllOutputFromInputs( + const std::vector& allInputs, int rank, int rankSize, int32_t elementsPerPeer, + std::vector& output) +{ + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + const size_t srcBase = static_cast(srcRank) * rankSize * elementsPerPeer + + static_cast(rank) * elementsPerPeer; + const size_t dstBase = static_cast(srcRank) * elementsPerPeer; + std::copy(allInputs.begin() + srcBase, + allInputs.begin() + srcBase + elementsPerPeer, + output.begin() + dstBase); + } +} + } // namespace Demo } // namespace TileXR diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 0f34c03c..ba3dfd94 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -21,6 +21,7 @@ #include "acl/acl.h" #include "tilexr_api.h" #include "tilexr_types.h" +#include "tilexr_udma_allreduce_layout.h" #include "tilexr_udma_alltoall_layout.h" extern void launch_tilexr_udma_all_gather( @@ -31,11 +32,24 @@ extern void launch_tilexr_udma_put_signal( extern void launch_tilexr_udma_all_to_all( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset); +extern void launch_tilexr_all_to_all_ipc_scatter( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); +extern void launch_tilexr_all_to_all_ipc_gather( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer); +extern void launch_tilexr_all_reduce_ipc_scatter( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerRank); +extern void launch_tilexr_all_reduce_ipc_sum( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerRank); namespace { constexpr int32_t kDefaultElementsPerRank = 16; constexpr uint64_t kSignalValue = 1000; -constexpr size_t kDebugWords = 16; +constexpr int kDebugUdmaStatusBase = 6; +constexpr int kDebugIpcScatter = kDebugUdmaStatusBase + TileXR::TILEXR_MAX_RANK_SIZE; +constexpr int kDebugIpcGather = kDebugIpcScatter + 1; +constexpr int kDebugAllReduceScatter = kDebugIpcGather + 1; +constexpr int kDebugAllReduceSum = kDebugAllReduceScatter + 1; +constexpr size_t kDebugWords = kDebugAllReduceSum + 1; constexpr int kDefaultCommPort = 10067; constexpr int kDemoBarrierPortOffset = 97; constexpr size_t kUdmaRegistrationAlignment = 2 * 1024 * 1024; @@ -372,6 +386,29 @@ bool ValidateAllToAllData( return TileXR::Demo::ValidateAllToAllOutput(output, rank, rankSize, elementsPerPeer) && ok; } +bool ValidateAllReduceData( + int rank, int rankSize, const std::vector& output, int32_t elementsPerRank) +{ + bool ok = true; + const int32_t expected = TileXR::Demo::AllReduceExpectedSum(rankSize); + for (int32_t i = 0; i < elementsPerRank; ++i) { + if (output[static_cast(i)] != expected) { + std::cerr << "[rank " << rank << "] ALLREDUCE MISMATCH at elem=" << i + << " got=" << output[static_cast(i)] + << " expected=" << expected << std::endl; + ok = false; + break; + } + } + + std::cout << "[rank " << rank << "] allreduce output sample:"; + for (int32_t i = 0; i < std::min(elementsPerRank, 8); ++i) { + std::cout << " elem" << i << "=" << output[static_cast(i)]; + } + std::cout << std::endl; + return TileXR::Demo::ValidateAllReduceOutput(output, rankSize, elementsPerRank) && ok; +} + bool ValidateSignals(int rank, int rankSize, const std::vector& signals) { bool ok = true; @@ -390,6 +427,16 @@ bool ValidateSignals(int rank, int rankSize, const std::vector& signal return ok; } +bool AllToAllUdmaComplete(int rankSize, const std::vector& debug) +{ + for (int peer = 0; peer < rankSize; ++peer) { + if (debug[kDebugUdmaStatusBase + peer] != 0) { + return false; + } + } + return true; +} + void Cleanup( TileXRCommPtr comm, aclrtStream stream, void* registeredMemory, int32_t* debug, int rank, int deviceId) { @@ -478,12 +525,14 @@ int main(int argc, char** argv) } bool isAllToAll = testType == 2; + bool isAllReduce = testType == 3; + bool hasOutput = isAllToAll || isAllReduce; size_t dataCount = static_cast(rankSize) * elementsPerRank; size_t dataBytes = dataCount * sizeof(int32_t); size_t inputOffset = 0; - size_t outputOffset = isAllToAll ? dataBytes : 0; + size_t outputOffset = hasOutput ? dataBytes : 0; size_t signalBytes = static_cast(rankSize) * sizeof(uint64_t); - size_t signalOffset = isAllToAll ? dataBytes * 2 : dataBytes; + size_t signalOffset = hasOutput ? dataBytes * 2 : dataBytes; size_t payloadBytes = signalOffset + signalBytes; size_t registeredBytes = ((payloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; @@ -515,6 +564,8 @@ int main(int argc, char** argv) std::vector hostOutput(dataCount, -1); if (isAllToAll) { TileXR::Demo::FillAllToAllInput(hostData, rank, rankSize, elementsPerRank); + } else if (isAllReduce) { + TileXR::Demo::FillAllReduceInput(hostData, rank, elementsPerRank); } else { std::fill(hostData.begin() + static_cast(rank) * elementsPerRank, hostData.begin() + static_cast(rank + 1) * elementsPerRank, @@ -523,11 +574,13 @@ int main(int argc, char** argv) std::vector hostSignals(static_cast(rankSize), 0); std::vector hostDebug(kDebugWords, 0); + const char* inputName = isAllToAll ? "alltoall input" : (isAllReduce ? "allreduce input" : "data"); bool initOk = CopyHostToDevice(rank, input, dataCount * sizeof(int32_t), - hostData.data(), dataCount * sizeof(int32_t), isAllToAll ? "alltoall input" : "data"); - if (isAllToAll) { + hostData.data(), dataCount * sizeof(int32_t), inputName); + if (hasOutput) { + const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; initOk = CopyHostToDevice(rank, output, dataCount * sizeof(int32_t), - hostOutput.data(), dataCount * sizeof(int32_t), "alltoall output") && initOk; + hostOutput.data(), dataCount * sizeof(int32_t), outputName) && initOk; } if (!initOk || !CopyHostToDevice(rank, signals, hostSignals.size() * sizeof(uint64_t), @@ -553,6 +606,11 @@ int main(int argc, char** argv) launch_tilexr_udma_all_to_all( 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset)); + } else if (testType == 3) { + PrintStatus(rank, "launch all-reduce IPC scatter kernel"); + launch_tilexr_all_reduce_ipc_scatter( + 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(debug), + elementsPerRank); } else if (testType == 1) { PrintStatus(rank, "launch put-signal kernel"); launch_tilexr_udma_put_signal( @@ -579,11 +637,71 @@ int main(int argc, char** argv) return 1; } + if (isAllReduce) { + PrintStatus(rank, "launch all-reduce IPC sum kernel"); + launch_tilexr_all_reduce_ipc_sum( + 1, stream, commArgsDev, reinterpret_cast(output), reinterpret_cast(debug), + elementsPerRank); + if (!CheckAcl(rank, "aclrtSynchronizeStream allreduce ipc sum", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed allreduce ipc sum")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } + + if (isAllToAll && + !CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), + debug, hostDebug.size() * sizeof(int32_t), "debug after alltoall udma")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + + bool usedIpcFallback = false; + if (isAllToAll && !AllToAllUdmaComplete(rankSize, hostDebug)) { + usedIpcFallback = true; + std::cout << "[rank " << rank << "] alltoall UDMA CQ incomplete, use IPC fallback:"; + for (int peer = 0; peer < rankSize; ++peer) { + std::cout << " peer" << peer << "=" << hostDebug[kDebugUdmaStatusBase + peer]; + } + std::cout << std::endl; + + launch_tilexr_all_to_all_ipc_scatter( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(debug), elementsPerRank); + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc scatter", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc scatter")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + + launch_tilexr_all_to_all_ipc_gather( + 1, stream, commArgsDev, reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank); + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc gather", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc gather")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } + bool copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), data, dataCount * sizeof(int32_t), "data"); - if (isAllToAll) { + if (hasOutput) { + const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; copyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), - output, dataCount * sizeof(int32_t), "alltoall output") && copyBackOk; + output, dataCount * sizeof(int32_t), outputName) && copyBackOk; } if (!copyBackOk || !CopyDeviceToHost(rank, hostSignals.data(), hostSignals.size() * sizeof(uint64_t), @@ -602,9 +720,25 @@ int main(int argc, char** argv) std::cout << " d" << i << "=" << hostDebug[i]; } std::cout << std::endl; + if (usedIpcFallback) { + std::cout << "[rank " << rank << "] alltoall IPC fallback completed" + << " scatter=" << hostDebug[kDebugIpcScatter] + << " gather=" << hostDebug[kDebugIpcGather] << std::endl; + } + if (isAllReduce) { + std::cout << "[rank " << rank << "] allreduce IPC completed" + << " scatter=" << hostDebug[kDebugAllReduceScatter] + << " sum=" << hostDebug[kDebugAllReduceSum] << std::endl; + } - bool ok = isAllToAll ? ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerRank) - : ValidateData(rank, rankSize, hostData, elementsPerRank); + bool ok = false; + if (isAllToAll) { + ok = ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerRank); + } else if (isAllReduce) { + ok = ValidateAllReduceData(rank, rankSize, hostOutput, elementsPerRank); + } else { + ok = ValidateData(rank, rankSize, hostData, elementsPerRank); + } if (testType == 1) { ok = ValidateSignals(rank, rankSize, hostSignals) && ok; } diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 1d0ebcdc..eeadbbe4 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -7,6 +7,13 @@ #include "tilexr_udma.h" constexpr int32_t TILEXR_UDMA_DEMO_MAGIC = 0x5444554d; // "TDUM" +constexpr uint64_t TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET = TileXR::IPC_DATA_OFFSET; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE = 6; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER = TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER + 1; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER = TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER + 1; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SUM = TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER + 1; extern "C" __global__ __aicore__ void tilexr_udma_all_gather_kernel( GM_ADDR commArgsGM, GM_ADDR dataGM, GM_ADDR debugGM, int32_t elementsPerRank) @@ -122,11 +129,98 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); uint32_t status = TileXR::UDMAQuietStatus(args, peer); if (debug != nullptr) { - debug[6 + peer] = static_cast(status); + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); } } } +extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_scatter_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR debugGM, int32_t elementsPerPeer) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + auto localSrc = input + dstRank * elementsPerPeer; + auto remoteBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[dstRank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + auto remoteDst = remoteBase + rank * elementsPerPeer; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + remoteDst[i] = localSrc[i]; + } + } + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER] = 1; + } +} + +extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_gather_kernel( + GM_ADDR commArgsGM, GM_ADDR outputGM, GM_ADDR debugGM, int32_t elementsPerPeer) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + auto localBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[args->rank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + for (int32_t srcRank = 0; srcRank < args->rankSize; ++srcRank) { + auto localSrc = localBase + srcRank * elementsPerPeer; + auto localDst = output + srcRank * elementsPerPeer; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + localDst[i] = localSrc[i]; + } + } + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER] = 1; + } +} + +extern "C" __global__ __aicore__ void tilexr_all_reduce_ipc_scatter_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR debugGM, int32_t elementsPerRank) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + auto remoteBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[dstRank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + auto remoteDst = remoteBase + rank * elementsPerRank; + for (int32_t i = 0; i < elementsPerRank; ++i) { + remoteDst[i] = input[i]; + } + } + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER] = 1; + } +} + +extern "C" __global__ __aicore__ void tilexr_all_reduce_ipc_sum_kernel( + GM_ADDR commArgsGM, GM_ADDR outputGM, GM_ADDR debugGM, int32_t elementsPerRank) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + auto localBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[args->rank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + for (int32_t i = 0; i < elementsPerRank; ++i) { + int32_t sum = 0; + for (int32_t srcRank = 0; srcRank < args->rankSize; ++srcRank) { + sum += localBase[srcRank * elementsPerRank + i]; + } + output[i] = sum; + } + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SUM] = 1; + } +} + extern "C" __global__ __aicore__ void tilexr_udma_registered_smoke_kernel( GM_ADDR commArgsGM, GM_ADDR localGM, GM_ADDR debugGM, uint32_t bytes, uint64_t signal) { @@ -180,6 +274,34 @@ void launch_tilexr_udma_all_to_all( commArgs, input, output, debug, elementsPerPeer, outputByteOffset); } +void launch_tilexr_all_to_all_ipc_scatter( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer) +{ + tilexr_all_to_all_ipc_scatter_kernel<<>>( + commArgs, input, debug, elementsPerPeer); +} + +void launch_tilexr_all_to_all_ipc_gather( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer) +{ + tilexr_all_to_all_ipc_gather_kernel<<>>( + commArgs, output, debug, elementsPerPeer); +} + +void launch_tilexr_all_reduce_ipc_scatter( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerRank) +{ + tilexr_all_reduce_ipc_scatter_kernel<<>>( + commArgs, input, debug, elementsPerRank); +} + +void launch_tilexr_all_reduce_ipc_sum( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerRank) +{ + tilexr_all_reduce_ipc_sum_kernel<<>>( + commArgs, output, debug, elementsPerRank); +} + void launch_tilexr_udma_registered_smoke( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR local, GM_ADDR debug, uint32_t bytes, uint64_t signal) { diff --git a/tests/udma/run_tests.sh b/tests/udma/run_tests.sh index 41e610bf..255bf7f5 100755 --- a/tests/udma/run_tests.sh +++ b/tests/udma/run_tests.sh @@ -50,6 +50,7 @@ required_bins=( test_tilexr_udma_transport_layout test_tilexr_udma_registry test_tilexr_udma_alltoall_layout + test_tilexr_udma_allreduce_layout test_tilexr_chip_map_sources test_tilexr_ipc_pid_mode_sources test_tilexr_udma @@ -84,30 +85,37 @@ TEST3_RESULT=$? echo "" echo "==========================================" -echo "Test 4: TileXR Chip Map Source Unit Test" +echo "Test 4: TileXR UDMA All-Reduce Layout Unit Test" echo "==========================================" -"${INSTALL_DIR}/bin/test_tilexr_chip_map_sources" +"${INSTALL_DIR}/bin/test_tilexr_udma_allreduce_layout" TEST4_RESULT=$? echo "" echo "==========================================" -echo "Test 5: TileXR IPC PID Mode Source Unit Test" +echo "Test 5: TileXR Chip Map Source Unit Test" echo "==========================================" -"${INSTALL_DIR}/bin/test_tilexr_ipc_pid_mode_sources" +"${INSTALL_DIR}/bin/test_tilexr_chip_map_sources" TEST5_RESULT=$? echo "" echo "==========================================" -echo "Test 6: TileXR Integration Tests (Single Process)" +echo "Test 6: TileXR IPC PID Mode Source Unit Test" +echo "==========================================" +"${INSTALL_DIR}/bin/test_tilexr_ipc_pid_mode_sources" +TEST6_RESULT=$? +echo "" + +echo "==========================================" +echo "Test 7: TileXR Integration Tests (Single Process)" echo "==========================================" export RANK=0 export RANK_SIZE=1 "${INSTALL_DIR}/bin/test_tilexr_udma" -TEST6_RESULT=$? +TEST7_RESULT=$? echo "" echo "==========================================" -echo "Test 7: TileXR Multi-Process Tests (MPI)" +echo "Test 8: TileXR Multi-Process Tests (MPI)" echo "==========================================" if command -v mpirun >/dev/null 2>&1; then @@ -124,14 +132,14 @@ if command -v mpirun >/dev/null 2>&1; then unset RANK unset RANK_SIZE mpirun -n 2 "${INSTALL_DIR}/bin/test_tilexr_udma" - TEST7_RESULT=$? + TEST8_RESULT=$? else echo "SKIP: Need at least 2 usable NPUs for multi-rank test" - TEST7_RESULT=0 + TEST8_RESULT=0 fi else echo "SKIP: mpirun not found, skipping multi-process tests" - TEST7_RESULT=0 + TEST8_RESULT=0 fi echo "" @@ -141,15 +149,16 @@ echo "==========================================" echo "Test 1 (UDMA Layout): $([ ${TEST1_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" echo "Test 2 (UDMA Registry): $([ ${TEST2_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" echo "Test 3 (AllToAll Layout): $([ ${TEST3_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 4 (Chip Map): $([ ${TEST4_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 5 (IPC PID Mode): $([ ${TEST5_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 6 (TileXR Single): $([ ${TEST6_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" -echo "Test 7 (TileXR Multi): $([ ${TEST7_RESULT} -eq 0 ] && echo 'PASS' || echo 'SKIP/FAIL')" +echo "Test 4 (AllReduce Layout): $([ ${TEST4_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 5 (Chip Map): $([ ${TEST5_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 6 (IPC PID Mode): $([ ${TEST6_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 7 (TileXR Single): $([ ${TEST7_RESULT} -eq 0 ] && echo 'PASS' || echo 'FAIL')" +echo "Test 8 (TileXR Multi): $([ ${TEST8_RESULT} -eq 0 ] && echo 'PASS' || echo 'SKIP/FAIL')" echo "==========================================" if [ ${TEST1_RESULT} -ne 0 ] || [ ${TEST2_RESULT} -ne 0 ] || [ ${TEST3_RESULT} -ne 0 ] || [ ${TEST4_RESULT} -ne 0 ] || [ ${TEST5_RESULT} -ne 0 ] || [ ${TEST6_RESULT} -ne 0 ] || - [ ${TEST7_RESULT} -ne 0 ]; then + [ ${TEST7_RESULT} -ne 0 ] || [ ${TEST8_RESULT} -ne 0 ]; then exit 1 fi diff --git a/tests/udma/unit/test_tilexr_chip_map_sources.cpp b/tests/udma/unit/test_tilexr_chip_map_sources.cpp index ed283191..25d06cc2 100644 --- a/tests/udma/unit/test_tilexr_chip_map_sources.cpp +++ b/tests/udma/unit/test_tilexr_chip_map_sources.cpp @@ -33,6 +33,7 @@ int main() { const std::string internal = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/tilexr_internal.cpp"); CHECK_CONTAINS(internal, "\"Ascend950DT_9592\""); + CHECK_CONTAINS(internal, "\"Ascend950PR_9599\""); CHECK_CONTAINS(internal, "ChipName::CHIP_950"); if (g_failures != 0) { diff --git a/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp b/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp index 394ee252..037647e1 100644 --- a/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp +++ b/tests/udma/unit/test_tilexr_ipc_pid_mode_sources.cpp @@ -19,6 +19,14 @@ int g_failures = 0; } \ } while (0) +#define CHECK_NOT_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) != std::string::npos) { \ + std::cerr << "CHECK_NOT_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + std::string ReadFile(const std::string& path) { std::ifstream in(path.c_str()); @@ -37,6 +45,12 @@ int main() CHECK_CONTAINS(comm, "physicalInfo_.chipName < ChipName::CHIP_950"); CHECK_CONTAINS(comm, "rtSetIpcMemPid"); CHECK_CONTAINS(comm, "rtSetIpcMemorySuperPodPid"); + CHECK_CONTAINS(comm, "fallback to rtSetIpcMemPid"); + CHECK_CONTAINS(comm, "OpenIpcMem failed after sdid setup, retry with pid setup"); + CHECK_CONTAINS(comm, "\"pid_retry\""); + CHECK_CONTAINS(comm, "SetMemoryName(retryName)"); + CHECK_CONTAINS(comm, "GetName(retryName, names)"); + CHECK_NOT_CONTAINS(comm, "do not support pcie > 2 rank"); CHECK_CONTAINS(cmake, "TILEXR_UDMA_FORCE_ENABLE"); const std::string transport = diff --git a/tests/udma/unit/test_tilexr_udma_allreduce_layout.cpp b/tests/udma/unit/test_tilexr_udma_allreduce_layout.cpp new file mode 100644 index 00000000..c0fcd390 --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_allreduce_layout.cpp @@ -0,0 +1,125 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "demo/tilexr_udma_allreduce_layout.h" + +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + +namespace { + +int g_failures = 0; + +#define CHECK_EQ(lhs, rhs) \ + do { \ + auto lhsValue = (lhs); \ + auto rhsValue = (rhs); \ + if (lhsValue != rhsValue) { \ + std::cerr << "CHECK_EQ failed at line " << __LINE__ << ": " #lhs " != " #rhs \ + << " (" << lhsValue << " vs " << rhsValue << ")" << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +void TestAllReduceInputPattern() +{ + constexpr int rank = 3; + constexpr int32_t elementsPerRank = 4; + std::vector input(elementsPerRank * 2, -1); + + TileXR::Demo::FillAllReduceInput(input, rank, elementsPerRank); + + for (int32_t i = 0; i < elementsPerRank; ++i) { + CHECK_EQ(input[static_cast(i)], TileXR::Demo::AllReduceValue(rank)); + } + CHECK_EQ(input[static_cast(elementsPerRank)], -1); +} + +void TestAllReduceExpectedEightRankSum() +{ + CHECK_EQ(TileXR::Demo::AllReduceExpectedSum(8), 8028); +} + +void TestAllReduceOutputValidation() +{ + constexpr int rankSize = 4; + constexpr int32_t elementsPerRank = 3; + std::vector output(elementsPerRank, TileXR::Demo::AllReduceExpectedSum(rankSize)); + + CHECK_EQ(TileXR::Demo::ValidateAllReduceOutput(output, rankSize, elementsPerRank), true); + output[1] = -1; + CHECK_EQ(TileXR::Demo::ValidateAllReduceOutput(output, rankSize, elementsPerRank), false); +} + +void TestBuildAllReduceOutput() +{ + constexpr int rankSize = 5; + constexpr int32_t elementsPerRank = 2; + std::vector allInputs(static_cast(rankSize) * elementsPerRank, -1); + + for (int rank = 0; rank < rankSize; ++rank) { + std::fill(allInputs.begin() + static_cast(rank) * elementsPerRank, + allInputs.begin() + static_cast(rank + 1) * elementsPerRank, + TileXR::Demo::AllReduceValue(rank)); + } + + std::vector output(elementsPerRank, -1); + TileXR::Demo::BuildAllReduceOutputFromInputs(allInputs, rankSize, elementsPerRank, output); + + CHECK_EQ(TileXR::Demo::ValidateAllReduceOutput(output, rankSize, elementsPerRank), true); +} + +void TestDemoAllReduceSourceHooks() +{ + const std::string demo = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); + const std::string kernel = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + const std::string script = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/run_tilexr_udma_demo.sh"); + + CHECK_CONTAINS(demo, "testType == 3"); + CHECK_CONTAINS(demo, "ValidateAllReduceData"); + CHECK_CONTAINS(kernel, "tilexr_all_reduce_ipc_scatter_kernel"); + CHECK_CONTAINS(kernel, "tilexr_all_reduce_ipc_sum_kernel"); + CHECK_CONTAINS(script, "3=all-reduce"); +} + +} // namespace + +int main() +{ + TestAllReduceInputPattern(); + TestAllReduceExpectedEightRankSum(); + TestAllReduceOutputValidation(); + TestBuildAllReduceOutput(); + TestDemoAllReduceSourceHooks(); + if (g_failures != 0) { + std::cerr << g_failures << " all-reduce layout checks failed" << std::endl; + return 1; + } + std::cout << "TileXR UDMA all-reduce layout checks passed" << std::endl; + return 0; +} diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 6bef28a5..6364c29f 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -1,10 +1,17 @@ #include #include +#include #include +#include +#include #include #include "demo/tilexr_udma_alltoall_layout.h" +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + namespace { int g_failures = 0; @@ -20,6 +27,22 @@ int g_failures = 0; } \ } while (0) +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + void TestAllToAllInputPattern() { constexpr int rank = 2; @@ -55,12 +78,46 @@ void TestAllToAllOutputValidation() CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(output, rank, rankSize, elementsPerPeer), false); } +void TestBuildAllToAllOutput() +{ + constexpr int rankSize = 3; + constexpr int32_t elementsPerPeer = 2; + std::vector allInputs(static_cast(rankSize) * rankSize * elementsPerPeer, -1); + + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + std::vector oneInput(static_cast(rankSize) * elementsPerPeer, -1); + TileXR::Demo::FillAllToAllInput(oneInput, srcRank, rankSize, elementsPerPeer); + std::copy(oneInput.begin(), oneInput.end(), + allInputs.begin() + static_cast(srcRank) * rankSize * elementsPerPeer); + } + + std::vector output(static_cast(rankSize) * elementsPerPeer, -1); + TileXR::Demo::BuildAllToAllOutputFromInputs(allInputs, 2, rankSize, elementsPerPeer, output); + + CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(output, 2, rankSize, elementsPerPeer), true); +} + +void TestDemoDebugLayoutSource() +{ + const std::string demo = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); + const std::string kernel = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + + CHECK_CONTAINS(demo, "kDebugUdmaStatusBase + TileXR::TILEXR_MAX_RANK_SIZE"); + CHECK_CONTAINS(demo, "kDebugIpcGather + 1"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer"); +} + } // namespace int main() { TestAllToAllInputPattern(); TestAllToAllOutputValidation(); + TestBuildAllToAllOutput(); + TestDemoDebugLayoutSource(); if (g_failures != 0) { std::cerr << g_failures << " all-to-all layout checks failed" << std::endl; return 1; diff --git a/tests/udma/unit/test_tilexr_udma_registry.cpp b/tests/udma/unit/test_tilexr_udma_registry.cpp index 518c3a46..0c4ffc5a 100644 --- a/tests/udma/unit/test_tilexr_udma_registry.cpp +++ b/tests/udma/unit/test_tilexr_udma_registry.cpp @@ -47,11 +47,17 @@ void TestRemoteAddressCalculation() static_cast(0x200040)); } +void TestRankScaleLimit() +{ + CHECK_EQ(TileXR::TILEXR_MAX_RANK_SIZE, 256); +} + } // namespace int main() { TestRemoteAddressCalculation(); + TestRankScaleLimit(); if (g_failures != 0) { std::cerr << g_failures << " registry checks failed" << std::endl; return 1; From 8d154336c4cfbe85dee70a5c33d6e027ba6b7066 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Wed, 24 Jun 2026 14:38:42 +0800 Subject: [PATCH 005/163] 256p alltoall udma checker ok --- src/comm/tilexr_comm.cpp | 8 + src/comm/udma/tilexr_udma_transport.cpp | 572 +++++++++++++----- src/comm/udma/tilexr_udma_transport.h | 2 + src/include/tilexr_udma.h | 21 +- src/include/tilexr_udma_types.h | 1 + tests/udma/CMakeLists.txt | 6 + tests/udma/demo/ALLTOALL_8P_RUNBOOK.md | 201 ++++++ tests/udma/demo/tilexr_udma_demo.cpp | 125 +++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 114 +++- .../unit/test_tilexr_udma_alltoall_layout.cpp | 55 ++ .../test_tilexr_udma_transport_layout.cpp | 106 ++++ 11 files changed, 1026 insertions(+), 185 deletions(-) create mode 100644 tests/udma/demo/ALLTOALL_8P_RUNBOOK.md diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 33d446b7..90324f14 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -346,6 +346,14 @@ int TileXRComm::RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMem TILEXR_LOG(ERROR) << "TileXR UDMA memory registration failed: " << ret; return TILEXR_ERROR_INTERNAL; } + udmaInfoDev_ = udmaTransport_->GetUDMAInfoDev(); + commArgs_.udmaInfoPtr = udmaInfoDev_; + ret = UpdateCommArgsDev(); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "TileXRUDMARegister failed to refresh CommArgs after UDMA info update: " << ret; + udmaTransport_->UnregisterMemory(localPtr); + return ret; + } if (socketExchange_ == nullptr) { TILEXR_LOG(ERROR) << "TileXRUDMARegister requires live socket exchange"; diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 7bafa5f1..be210836 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -12,8 +12,10 @@ #include #include #include +#include #include #include +#include #include #include @@ -33,6 +35,29 @@ uint32_t Log2Uint64(uint64_t value) return result; } +bool UDMADiagEnabled() +{ + const char* value = std::getenv("TILEXR_UDMA_DEBUG"); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + +std::string PtrToHex(uint64_t value) +{ + std::ostringstream os; + os << "0x" << std::hex << value; + return os.str(); +} + +std::string EidToHex(const HccpEid& eid) +{ + std::ostringstream os; + os << std::hex << std::setfill('0'); + for (uint8_t byte : eid.raw) { + os << std::setw(2) << static_cast(byte); + } + return os.str(); +} + HccpEid SwapEidForDevice(const HccpEid& hccpEid) { HccpEid swapped {}; @@ -289,24 +314,29 @@ bool ResolveLocalEidRoute( } // namespace struct TileXRUDMATransport::PerEidState { + struct PeerQueueState { + int peer = -1; + void* chanHandle = nullptr; + void* cqHandle = nullptr; + void* qpHandle = nullptr; + void* remoteQpHandle = nullptr; + CqInfoT cqInfo {}; + QpCreateInfo qpInfo {}; + uint32_t tpn = 0; + void* cqPiAddr = nullptr; + void* cqCiAddr = nullptr; + void* sqPiAddr = nullptr; + void* sqCiAddr = nullptr; + void* wqeCntAddr = nullptr; + void* amoAddr = nullptr; + UDMAWQCtx localWq {}; + UDMACQCtx localCq {}; + }; + uint32_t eidIndex = 0; void* ctxHandle = nullptr; void* tokenHandle = nullptr; - void* chanHandle = nullptr; - void* cqHandle = nullptr; - void* qpHandle = nullptr; - CqInfoT cqInfo {}; - QpCreateInfo qpInfo {}; - std::vector remoteQpHandles; - std::vector tpnList; - void* cqPiAddr = nullptr; - void* cqCiAddr = nullptr; - void* sqPiAddr = nullptr; - void* sqCiAddr = nullptr; - void* wqeCntAddr = nullptr; - void* amoAddr = nullptr; - UDMAWQCtx localWq {}; - UDMACQCtx localCq {}; + std::map peerQueues; }; TileXRUDMATransport::TileXRUDMATransport() = default; @@ -348,17 +378,7 @@ int TileXRUDMATransport::Init(const TileXRUDMATransportOptions& options) Shutdown(); return ret; } - ret = CreateQueues(); - if (ret != TILEXR_SUCCESS) { - Shutdown(); - return ret; - } - ret = ImportQueues(); - if (ret != TILEXR_SUCCESS) { - Shutdown(); - return ret; - } - ret = RefreshUDMAInfo(); + ret = EnsureUDMAInfoBuffer(); if (ret != TILEXR_SUCCESS) { Shutdown(); return ret; @@ -424,6 +444,24 @@ int TileXRUDMATransport::BuildRoutes() return TILEXR_ERROR_INTERNAL; } eidCount_ = eidNum; + const bool diag = UDMADiagEnabled(); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag BuildRoutes rank " << options_.rank + << " devId=" << options_.devId + << " logicDevId=" << logicDevId_ + << " deviceIdOffset=" << deviceIdOffset_ + << " phyId=" << (logicDevId_ + deviceIdOffset_) + << " runtimeEidCount=" << eidNum; + for (unsigned int i = 0; i < eidNum; ++i) { + TILEXR_LOG(INFO) << "UDMA diag local runtime eid rank " << options_.rank + << " idx=" << devEids[i].eidIndex + << " name=" << devEids[i].name + << " die=" << devEids[i].dieId + << " chip=" << devEids[i].chipId + << " func=" << devEids[i].funcId + << " eid=" << EidToHex(devEids[i].eid); + } + } uint32_t localId = static_cast(options_.devId); bool topoReady = false; @@ -456,6 +494,16 @@ int TileXRUDMATransport::BuildRoutes() if (ret != TILEXR_SUCCESS) { return ret; } + if (diag) { + std::ostringstream ids; + for (int rank = 0; rank < options_.rankSize; ++rank) { + ids << " rank" << rank << "=" << allLocalIds[rank]; + } + TILEXR_LOG(INFO) << "UDMA diag route ids rank " << options_.rank + << " localId=" << localId + << " topoReady=" << (topoReady ? 1 : 0) + << ids.str(); + } std::vector localRouteByPeer(options_.rankSize, -1); for (int peer = 0; peer < options_.rankSize; ++peer) { @@ -471,6 +519,13 @@ int TileXRUDMATransport::BuildRoutes() } peerLocalEid_[peer] = localEid; localRouteByPeer[peer] = static_cast(localEid); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag local route rank " << options_.rank + << " devLocalId=" << localId + << " peer=" << peer + << " peerLocalId=" << allLocalIds[peer] + << " localEid=" << localEid; + } } std::vector allRouteByPeer(options_.rankSize * options_.rankSize, -1); @@ -488,6 +543,12 @@ int TileXRUDMATransport::BuildRoutes() remoteEid = static_cast(devEids[0].eidIndex); } peerRemoteEid_[peer] = static_cast(remoteEid); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag remote route rank " << options_.rank + << " peer=" << peer + << " localEid=" << peerLocalEid_[peer] + << " remoteEid=" << peerRemoteEid_[peer]; + } } return TILEXR_SUCCESS; } @@ -517,6 +578,7 @@ int TileXRUDMATransport::CreateContexts() bool found = false; CtxInitAttr attr {}; auto targetEidIt = localEidByEid_.find(eidIndex); + const DevEidInfo* matchedEid = nullptr; for (unsigned int i = 0; i < eidNum; ++i) { bool matched = infoList[i].eidIndex == eidIndex; if (targetEidIt != localEidByEid_.end()) { @@ -529,6 +591,7 @@ int TileXRUDMATransport::CreateContexts() attr.ub.eid = infoList[i].eid; attr.ub.eidIndex = infoList[i].eidIndex; localEidByEid_[eidIndex] = infoList[i].eid; + matchedEid = &infoList[i]; found = true; break; } @@ -544,6 +607,17 @@ int TileXRUDMATransport::CreateContexts() TILEXR_LOG(WARN) << "TileXR UDMA RaCtxInit failed: " << ret; return TILEXR_ERROR_INTERNAL; } + if (UDMADiagEnabled() && matchedEid != nullptr) { + TILEXR_LOG(INFO) << "UDMA diag ctx init rank " << options_.rank + << " eid=" << eidIndex + << " runtimeIdx=" << matchedEid->eidIndex + << " name=" << matchedEid->name + << " die=" << matchedEid->dieId + << " chip=" << matchedEid->chipId + << " func=" << matchedEid->funcId + << " eidValue=" << EidToHex(matchedEid->eid) + << " ctx=" << ctxHandle; + } void* tokenHandle = nullptr; HccpTokenId tokenId {}; ret = loader_.RaCtxTokenIdAlloc(ctxHandle, &tokenId, &tokenHandle); @@ -586,127 +660,170 @@ int TileXRUDMATransport::CreateQueues() state.eidIndex = ctxEntry.first; state.ctxHandle = ctxEntry.second; state.tokenHandle = tokenHandleByEid_[ctxEntry.first]; - state.remoteQpHandles.assign(options_.rankSize, nullptr); - state.tpnList.assign(options_.rankSize, 0); - - ChanInfoT chanInfo {}; - chanInfo.in.dataPlaneFlag.bs.poolCqCstm = 1; - int ret = loader_.RaCtxChanCreate(state.ctxHandle, &chanInfo, &state.chanHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - - state.cqInfo.in.chanHandle = state.chanHandle; - state.cqInfo.in.depth = TILEXR_UDMA_CQ_DEPTH; - state.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; - ret = loader_.RaCtxCqCreate(state.ctxHandle, &state.cqInfo, &state.cqHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - state.localCq.cqn = 0; - state.localCq.bufAddr = state.cqInfo.out.bufAddr; - state.localCq.baseBkShift = Log2Uint64(state.cqInfo.out.cqeSize); - state.localCq.depth = state.cqInfo.in.depth; - if (AllocDeviceScalar(&state.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { - return TILEXR_ERROR_INTERNAL; - } - state.localCq.headAddr = reinterpret_cast(state.cqPiAddr); - state.localCq.tailAddr = reinterpret_cast(state.cqCiAddr); - state.localCq.dbMode = UDMADBMode::SW_DB; - state.localCq.dbAddr = state.cqInfo.out.swdbAddr; - - QpCreateAttr qpAttr {}; - qpAttr.scqHandle = state.cqHandle; - qpAttr.rcqHandle = state.cqHandle; - qpAttr.srqHandle = state.cqHandle; - qpAttr.sqDepth = TILEXR_UDMA_SQ_DEPTH; - qpAttr.rqDepth = TILEXR_UDMA_RQ_DEPTH_DEFAULT; - qpAttr.transportMode = CONN_RM; - qpAttr.ub.mode = JETTY_MODE_USER_CTL_NORMAL; - qpAttr.ub.flag.value = 1; - qpAttr.ub.jfsFlag.value = 2; - qpAttr.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - qpAttr.ub.rnrRetry = 7; - qpAttr.ub.extMode.piType = 0; - qpAttr.ub.extMode.cstmFlag.bs.sqCstm = 0; - qpAttr.ub.extMode.sqebbNum = TILEXR_UDMA_SQ_DEPTH; - qpAttr.ub.tokenIdHandle = state.tokenHandle; - ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &state.qpInfo, &state.qpHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - state.localWq.wqn = 0; - state.localWq.bufAddr = state.qpInfo.ub.sqBuffVa; - state.localWq.baseBkShift = Log2Uint64(state.qpInfo.ub.wqebbSize); - state.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; - if (AllocDeviceScalar(&state.sqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { - return TILEXR_ERROR_INTERNAL; + for (const auto& route : peerLocalEid_) { + if (route.second != state.eidIndex) { + continue; + } + int ret = CreatePeerQueue(state, route.first); + if (ret != TILEXR_SUCCESS) { + return ret; + } } - state.localWq.headAddr = reinterpret_cast(state.sqPiAddr); - state.localWq.tailAddr = reinterpret_cast(state.sqCiAddr); - state.localWq.dbMode = UDMADBMode::SW_DB; - state.localWq.dbAddr = state.qpInfo.ub.dbAddr; - state.localWq.wqeCntAddr = reinterpret_cast(state.wqeCntAddr); - state.localWq.amoAddr = reinterpret_cast(state.amoAddr); states_[state.eidIndex] = state; } return states_.empty() ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; } +int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer) +{ + const bool diag = UDMADiagEnabled(); + PerEidState::PeerQueueState queue {}; + queue.peer = peer; + + ChanInfoT chanInfo {}; + chanInfo.in.dataPlaneFlag.bs.poolCqCstm = 1; + int ret = loader_.RaCtxChanCreate(state.ctxHandle, &chanInfo, &queue.chanHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + + queue.cqInfo.in.chanHandle = queue.chanHandle; + queue.cqInfo.in.depth = TILEXR_UDMA_CQ_DEPTH; + queue.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; + ret = loader_.RaCtxCqCreate(state.ctxHandle, &queue.cqInfo, &queue.cqHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + queue.localCq.cqn = 0; + queue.localCq.bufAddr = queue.cqInfo.out.bufAddr; + queue.localCq.baseBkShift = Log2Uint64(queue.cqInfo.out.cqeSize); + queue.localCq.depth = queue.cqInfo.in.depth; + if (AllocDeviceScalar(&queue.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&queue.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + queue.localCq.headAddr = reinterpret_cast(queue.cqPiAddr); + queue.localCq.tailAddr = reinterpret_cast(queue.cqCiAddr); + queue.localCq.dbMode = UDMADBMode::SW_DB; + queue.localCq.dbAddr = queue.cqInfo.out.swdbAddr; + + QpCreateAttr qpAttr {}; + qpAttr.scqHandle = queue.cqHandle; + qpAttr.rcqHandle = queue.cqHandle; + qpAttr.srqHandle = queue.cqHandle; + qpAttr.sqDepth = TILEXR_UDMA_SQ_DEPTH; + qpAttr.rqDepth = TILEXR_UDMA_RQ_DEPTH_DEFAULT; + qpAttr.transportMode = CONN_RM; + qpAttr.ub.mode = JETTY_MODE_USER_CTL_NORMAL; + qpAttr.ub.flag.value = 1; + qpAttr.ub.jfsFlag.value = 2; + qpAttr.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + qpAttr.ub.rnrRetry = 7; + qpAttr.ub.extMode.piType = 0; + qpAttr.ub.extMode.cstmFlag.bs.sqCstm = 0; + qpAttr.ub.extMode.sqebbNum = TILEXR_UDMA_SQ_DEPTH; + qpAttr.ub.tokenIdHandle = state.tokenHandle; + ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &queue.qpInfo, &queue.qpHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + queue.localWq.wqn = 0; + queue.localWq.bufAddr = queue.qpInfo.ub.sqBuffVa; + queue.localWq.baseBkShift = Log2Uint64(queue.qpInfo.ub.wqebbSize); + queue.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; + if (AllocDeviceScalar(&queue.sqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&queue.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&queue.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&queue.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + queue.localWq.headAddr = reinterpret_cast(queue.sqPiAddr); + queue.localWq.tailAddr = reinterpret_cast(queue.sqCiAddr); + queue.localWq.dbMode = UDMADBMode::SW_DB; + queue.localWq.dbAddr = queue.qpInfo.ub.dbAddr; + queue.localWq.wqeCntAddr = reinterpret_cast(queue.wqeCntAddr); + queue.localWq.amoAddr = reinterpret_cast(queue.amoAddr); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag create peer queue rank " << options_.rank + << " peer=" << peer + << " eid=" << state.eidIndex + << " ctx=" << state.ctxHandle + << " chan=" << queue.chanHandle + << " qp=" << queue.qpHandle + << " cq=" << queue.cqHandle + << " sqBuf=" << PtrToHex(queue.localWq.bufAddr) + << " sqDb=" << PtrToHex(queue.localWq.dbAddr) + << " sqHead=" << PtrToHex(queue.localWq.headAddr) + << " sqTail=" << PtrToHex(queue.localWq.tailAddr) + << " wqeCnt=" << PtrToHex(queue.localWq.wqeCntAddr) + << " cqBuf=" << PtrToHex(queue.localCq.bufAddr) + << " cqDb=" << PtrToHex(queue.localCq.dbAddr) + << " cqHead=" << PtrToHex(queue.localCq.headAddr) + << " cqTail=" << PtrToHex(queue.localCq.tailAddr) + << " wqebbSize=" << queue.qpInfo.ub.wqebbSize + << " cqeSize=" << queue.cqInfo.out.cqeSize; + } + state.peerQueues[peer] = queue; + return TILEXR_SUCCESS; +} + int TileXRUDMATransport::ImportQueues() { - std::vector localImports(eidCount_); - std::vector localKeys(eidCount_); + const bool diag = UDMADiagEnabled(); + std::vector localPeerImports(options_.rankSize); + std::vector localPeerKeys(options_.rankSize); for (const auto& stateEntry : states_) { const auto& state = stateEntry.second; - if (state.eidIndex >= eidCount_) { - return TILEXR_ERROR_INTERNAL; + for (const auto& queueEntry : state.peerQueues) { + const int peer = queueEntry.first; + const auto& queue = queueEntry.second; + localPeerImports[peer].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + localPeerImports[peer].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + localPeerImports[peer].in.ub.policy = JETTY_GRP_POLICY_RR; + localPeerImports[peer].in.ub.type = TARGET_TYPE_JETTY; + localPeerImports[peer].in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + localPeerImports[peer].in.ub.tpType = 1; + localPeerImports[peer].in.key = queue.qpInfo.key; + localPeerKeys[peer] = queue.qpInfo.key; } - localImports[state.eidIndex].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; - localImports[state.eidIndex].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - localImports[state.eidIndex].in.ub.policy = JETTY_GRP_POLICY_RR; - localImports[state.eidIndex].in.ub.type = TARGET_TYPE_JETTY; - localImports[state.eidIndex].in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; - localImports[state.eidIndex].in.ub.tpType = 1; - localKeys[state.eidIndex] = state.qpInfo.key; } - std::vector allImports(options_.rankSize * eidCount_); - int ret = options_.exchange->AllGather(localImports.data(), localImports.size(), allImports.data()); + std::vector allImports(options_.rankSize * options_.rankSize); + int ret = options_.exchange->AllGather(localPeerImports.data(), localPeerImports.size(), allImports.data()); if (ret != TILEXR_SUCCESS) { return ret; } - std::vector allKeys(options_.rankSize * eidCount_); - ret = options_.exchange->AllGather(localKeys.data(), localKeys.size(), allKeys.data()); + std::vector allKeys(options_.rankSize * options_.rankSize); + ret = options_.exchange->AllGather(localPeerKeys.data(), localPeerKeys.size(), allKeys.data()); if (ret != TILEXR_SUCCESS) { return ret; } for (auto& stateEntry : states_) { auto& state = stateEntry.second; - for (int peer = 0; peer < options_.rankSize; ++peer) { - if (peer == options_.rank) { - continue; - } - const auto localRoute = peerLocalEid_.find(peer); - if (localRoute == peerLocalEid_.end() || localRoute->second != state.eidIndex) { - continue; - } + for (auto& queueEntry : state.peerQueues) { + const int peer = queueEntry.first; + auto& queue = queueEntry.second; const uint32_t remoteEid = peerRemoteEid_[peer]; if (remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - QpImportInfoT importInfo = allImports[peer * eidCount_ + remoteEid]; - importInfo.in.key = allKeys[peer * eidCount_ + remoteEid]; - ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &state.remoteQpHandles[peer]); + QpImportInfoT importInfo = allImports[(peer * options_.rankSize + options_.rank)]; + importInfo.in.key = allKeys[(peer * options_.rankSize + options_.rank)]; + ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &queue.remoteQpHandle); if (ret != 0) { return TILEXR_ERROR_INTERNAL; } - state.tpnList[peer] = importInfo.out.ub.tpn; + queue.tpn = importInfo.out.ub.tpn; + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag import qp rank " << options_.rank + << " peer=" << peer + << " localEid=" << state.eidIndex + << " remoteEid=" << remoteEid + << " remoteQp=" << queue.remoteQpHandle + << " tpn=" << queue.tpn + << " keySize=" << static_cast(importInfo.in.key.size); + } } } return TILEXR_SUCCESS; @@ -722,6 +839,7 @@ uint32_t TileXRUDMATransport::FallbackLocalEid() const int TileXRUDMATransport::RefreshUDMAInfo() { + const bool diag = UDMADiagEnabled(); if (eidCount_ == 0 || states_.empty()) { return TILEXR_ERROR_INTERNAL; } @@ -785,11 +903,36 @@ int TileXRUDMATransport::RefreshUDMAInfo() if (stateIt == states_.end()) { stateIt = fallbackIt; } - const auto& state = stateIt->second; - sq[rank] = state.localWq; - rq[rank] = state.localWq; - scq[rank] = state.localCq; - rcq[rank] = state.localCq; + auto& state = stateIt->second; + PerEidState::PeerQueueState* queuePtr = nullptr; + if (rank == options_.rank) { + if (!state.peerQueues.empty()) { + queuePtr = &state.peerQueues.begin()->second; + } else if (!fallbackIt->second.peerQueues.empty()) { + queuePtr = &fallbackIt->second.peerQueues.begin()->second; + } + } else { + const auto queueIt = state.peerQueues.find(rank); + if (queueIt == state.peerQueues.end()) { + return TILEXR_ERROR_INTERNAL; + } + queuePtr = &queueIt->second; + } + if (queuePtr == nullptr) { + return TILEXR_ERROR_INTERNAL; + } + auto& queue = *queuePtr; + if (!registeredMem_.empty()) { + const auto& localMrs = registeredMem_.begin()->second; + const auto localMrIt = localMrs.find(localEid); + if (localMrIt != localMrs.end()) { + queue.localWq.localTokenId = localMrIt->second.tokenId; + } + } + sq[rank] = queue.localWq; + rq[rank] = queue.localWq; + scq[rank] = queue.localCq; + rcq[rank] = queue.localCq; if (rank == options_.rank) { const auto localMemIt = localMemInfoByEid_.find(localEid); if (localMemIt != localMemInfoByEid_.end()) { @@ -797,17 +940,42 @@ int TileXRUDMATransport::RefreshUDMAInfo() } } else { mem[rank] = allMem[rank * eidCount_ + remoteEid]; - mem[rank].tpn = state.tpnList[rank]; + mem[rank].tpn = queue.tpn; } mem[rank].eidAddr = reinterpret_cast( eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag info image rank " << options_.rank + << " entryRank=" << rank + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " sqBuf=" << PtrToHex(sq[rank].bufAddr) + << " sqHead=" << PtrToHex(sq[rank].headAddr) + << " sqTail=" << PtrToHex(sq[rank].tailAddr) + << " localTokenId=" << sq[rank].localTokenId + << " wqeCnt=" << PtrToHex(sq[rank].wqeCntAddr) + << " cqBuf=" << PtrToHex(scq[rank].bufAddr) + << " cqTail=" << PtrToHex(scq[rank].tailAddr) + << " memAddr=" << PtrToHex(mem[rank].addr) + << " memLen=" << mem[rank].len + << " memTid=" << mem[rank].tid + << " memTpn=" << mem[rank].tpn + << " memEidAddr=" << PtrToHex(mem[rank].eidAddr); + } } - if (udmaInfoDev_ == nullptr) { - const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); - udmaInfoSize_ = static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize); + const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); + const uint32_t requiredInfoSize = + static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize); + if (udmaInfoDev_ == nullptr || udmaInfoSize_ < requiredInfoSize) { + if (udmaInfoDev_ != nullptr) { + aclrtFree(udmaInfoDev_); + udmaInfoDev_ = nullptr; + } + udmaInfoSize_ = requiredInfoSize; ret = aclrtMalloc(reinterpret_cast(&udmaInfoDev_), udmaInfoSize_, ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { + udmaInfoSize_ = 0; return TILEXR_ERROR_INTERNAL; } } @@ -825,28 +993,81 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_SUCCESS; } +int TileXRUDMATransport::EnsureUDMAInfoBuffer() +{ + if (udmaInfoDev_ != nullptr) { + return TILEXR_SUCCESS; + } + UDMAInfo info {}; + info.qpNum = 1; + udmaInfoSize_ = static_cast(sizeof(UDMAInfo)); + int ret = aclrtMalloc(reinterpret_cast(&udmaInfoDev_), udmaInfoSize_, ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + ret = aclrtMemcpy(udmaInfoDev_, udmaInfoSize_, &info, sizeof(info), ACL_MEMCPY_HOST_TO_DEVICE); + if (ret != ACL_SUCCESS) { + aclrtFree(udmaInfoDev_); + udmaInfoDev_ = nullptr; + udmaInfoSize_ = 0; + return TILEXR_ERROR_INTERNAL; + } + return TILEXR_SUCCESS; +} + int TileXRUDMATransport::RegisterMemory(GM_ADDR localPtr, size_t bytes) { if (!available_ || localPtr == nullptr || bytes == 0) { return TILEXR_ERROR_NOT_FOUND; } + const bool diag = UDMADiagEnabled(); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag register memory begin rank " << options_.rank + << " ptr=" << PtrToHex(reinterpret_cast(localPtr)) + << " bytes=" << bytes; + } + CleanupMemory(); + CleanupQueues(); + registeredPtr_ = nullptr; int ret = RegisterMemoryOnContexts(localPtr, bytes); if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA register memory on contexts failed rank " << options_.rank + << " ret=" << ret; return ret; } registeredPtr_ = localPtr; + ret = CreateQueues(); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA create queues after memory register failed rank " << options_.rank + << " ret=" << ret; + return ret; + } + ret = ImportQueues(); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA import queues after memory register failed rank " << options_.rank + << " ret=" << ret; + return ret; + } ret = ExchangeAndImportMemory(); if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA exchange/import memory failed rank " << options_.rank + << " ret=" << ret; + return ret; + } + ret = RefreshUDMAInfo(); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA refresh info after memory register failed rank " << options_.rank + << " ret=" << ret; return ret; } - return RefreshUDMAInfo(); + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag register memory end rank " << options_.rank; + } + return TILEXR_SUCCESS; } int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes) { - if (registeredPtr_ != nullptr) { - UnregisterMemory(registeredPtr_); - } std::map byEid; localMemInfoByEid_.clear(); for (const auto& ctxEntry : ctxHandleByEid_) { @@ -857,12 +1078,22 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes mrInfo.in.mem.size = 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 = 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) { + TILEXR_LOG(WARN) << "UDMA RaCtxLmemRegister failed rank " << options_.rank + << " eid=" << eidIndex + << " ctx=" << ctxEntry.second + << " ptr=" << PtrToHex(reinterpret_cast(localPtr)) + << " bytes=" << bytes + << " ret=" << ret + << " handle=" << lmemHandle; return TILEXR_ERROR_INTERNAL; } @@ -888,6 +1119,15 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes memInfo.len = static_cast(std::min(bytes, UINT32_MAX)); memInfo.addr = reinterpret_cast(localPtr); localMemInfoByEid_[eidIndex] = memInfo; + if (UDMADiagEnabled()) { + TILEXR_LOG(INFO) << "UDMA diag lmem registered rank " << options_.rank + << " eid=" << eidIndex + << " lmem=" << lmemHandle + << " tokenId=" << result.tokenId + << " tid=" << memInfo.tid + << " targetSeg=" << PtrToHex(result.targetSegHandle) + << " keySize=" << static_cast(result.key.size); + } } registeredMem_[reinterpret_cast(localPtr)] = byEid; return TILEXR_SUCCESS; @@ -903,10 +1143,13 @@ int TileXRUDMATransport::ExchangeAndImportMemory() std::vector allCounts(options_.rankSize); int ret = options_.exchange->AllGather(&localCount, 1, allCounts.data()); if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA memory count allgather failed rank " << options_.rank + << " ret=" << ret; return ret; } const uint32_t maxCount = *std::max_element(allCounts.begin(), allCounts.end()); if (maxCount == 0) { + TILEXR_LOG(WARN) << "UDMA memory exchange found zero max registration count rank " << options_.rank; return TILEXR_ERROR_INTERNAL; } @@ -927,6 +1170,9 @@ int TileXRUDMATransport::ExchangeAndImportMemory() std::vector all(options_.rankSize * maxCount); ret = options_.exchange->AllGather(local.data(), local.size(), all.data()); if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA memory info allgather failed rank " << options_.rank + << " ret=" << ret + << " localEntries=" << local.size(); return ret; } @@ -945,6 +1191,11 @@ int TileXRUDMATransport::ExchangeAndImportMemory() } } if (remote == nullptr) { + TILEXR_LOG(WARN) << "UDMA remote memory info missing rank " << options_.rank + << " peer=" << peer + << " remoteEid=" << remoteEid + << " peerCount=" << allCounts[peer] + << " maxCount=" << maxCount; return TILEXR_ERROR_INTERNAL; } const uint32_t localEid = peerLocalEid_[peer]; @@ -956,9 +1207,26 @@ int TileXRUDMATransport::ExchangeAndImportMemory() void* remoteHandle = nullptr; ret = loader_.RaCtxRmemImport(ctxHandleByEid_[localEid], &importInfo, &remoteHandle); if (ret != 0 || remoteHandle == nullptr) { + TILEXR_LOG(WARN) << "UDMA RaCtxRmemImport failed rank " << options_.rank + << " peer=" << peer + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " ctx=" << ctxHandleByEid_[localEid] + << " ret=" << ret + << " handle=" << remoteHandle + << " remoteToken=" << remote->mr.tokenValue + << " remoteTokenId=" << remote->mr.tokenId + << " remoteKeySize=" << static_cast(remote->mr.key.size); return TILEXR_ERROR_INTERNAL; } remoteMemHandles_[peer] = remoteHandle; + if (UDMADiagEnabled()) { + TILEXR_LOG(INFO) << "UDMA diag rmem imported rank " << options_.rank + << " peer=" << peer + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " remoteHandle=" << remoteHandle; + } } return TILEXR_SUCCESS; } @@ -1002,26 +1270,32 @@ void TileXRUDMATransport::CleanupQueues() { for (auto& stateEntry : states_) { auto& state = stateEntry.second; - for (void* remoteQp : state.remoteQpHandles) { - if (remoteQp != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxQpUnimport(state.ctxHandle, remoteQp); + for (auto& queueEntry : state.peerQueues) { + auto& queue = queueEntry.second; + if (queue.remoteQpHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxQpUnimport(state.ctxHandle, queue.remoteQpHandle); + queue.remoteQpHandle = nullptr; } + if (queue.qpHandle != nullptr) { + loader_.RaCtxQpDestroy(queue.qpHandle); + queue.qpHandle = nullptr; + } + if (queue.cqHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxCqDestroy(state.ctxHandle, queue.cqHandle); + queue.cqHandle = nullptr; + } + if (queue.chanHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxChanDestroy(state.ctxHandle, queue.chanHandle); + queue.chanHandle = nullptr; + } + FreeDeviceScalar(queue.cqPiAddr); + FreeDeviceScalar(queue.cqCiAddr); + FreeDeviceScalar(queue.sqPiAddr); + FreeDeviceScalar(queue.sqCiAddr); + FreeDeviceScalar(queue.wqeCntAddr); + FreeDeviceScalar(queue.amoAddr); } - if (state.qpHandle != nullptr) { - loader_.RaCtxQpDestroy(state.qpHandle); - } - if (state.cqHandle != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxCqDestroy(state.ctxHandle, state.cqHandle); - } - if (state.chanHandle != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxChanDestroy(state.ctxHandle, state.chanHandle); - } - FreeDeviceScalar(state.cqPiAddr); - FreeDeviceScalar(state.cqCiAddr); - FreeDeviceScalar(state.sqPiAddr); - FreeDeviceScalar(state.sqCiAddr); - FreeDeviceScalar(state.wqeCntAddr); - FreeDeviceScalar(state.amoAddr); + state.peerQueues.clear(); } states_.clear(); } diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index 0a787d12..d3adfac6 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -53,7 +53,9 @@ class TileXRUDMATransport { int BuildRoutes(); int CreateContexts(); int CreateQueues(); + int CreatePeerQueue(PerEidState& state, int peer); int ImportQueues(); + int EnsureUDMAInfoBuffer(); int RefreshUDMAInfo(); int RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes); int ExchangeAndImportMemory(); diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index 10b85251..c931e804 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -182,23 +182,28 @@ __aicore__ inline void UDMAFillNotifyData( __aicore__ inline void UDMAFillSqeCtx( __gm__ UDMASqeCtx* sqeCtx, __gm__ uint8_t* remoteAddr, __gm__ UDMAMemInfo* remoteMemInfo, - uint32_t curHead, UDMAOpcode opcode, const UDMASignalParams* signalParams) + uint32_t curHead, uint32_t depth, UDMAOpcode opcode, const UDMASignalParams* signalParams) { + sqeCtx->sqeBbIdx = curHead % depth; sqeCtx->opcode = static_cast(opcode); sqeCtx->flag = 0b00100010; + sqeCtx->rsv0 = 0; sqeCtx->nf = 0; sqeCtx->tokenEn = remoteMemInfo->tokenValueValid; sqeCtx->rmtJettyType = remoteMemInfo->rmtJettyType; - sqeCtx->owner = (curHead & TILEXR_UDMA_SQ_BB_COUNT) == 0 ? 1 : 0; + sqeCtx->owner = (curHead & depth) == 0 ? 1 : 0; sqeCtx->targetHint = remoteMemInfo->targetHint; + sqeCtx->rsv1 = 0; sqeCtx->inlineMsgLen = 0; sqeCtx->tpId = remoteMemInfo->tpn; sqeCtx->sgeNum = 1; sqeCtx->rmtJettyOrSegId = remoteMemInfo->tid; + sqeCtx->rsv2 = 0; sqeCtx->rmtTokenValue = remoteMemInfo->rmtTokenValue; sqeCtx->udfType = 0; sqeCtx->reduceDataType = 0; sqeCtx->reduceOpcode = 0; + sqeCtx->rsv3 = 0; uint64_t remoteAddrValue = reinterpret_cast(remoteAddr); sqeCtx->rmtAddrLOrTokenId = remoteAddrValue & 0xFFFFFFFF; sqeCtx->rmtAddrHOrTokenValue = (remoteAddrValue >> 32) & 0xFFFFFFFF; @@ -209,10 +214,11 @@ __aicore__ inline void UDMAFillSqeCtx( } __aicore__ inline void UDMAFillSgeCtx( - __gm__ UDMASgeCtx* sgeCtx, uint64_t messageLen, __gm__ uint8_t* localAddr) + __gm__ UDMASgeCtx* sgeCtx, uint64_t messageLen, __gm__ uint8_t* localAddr, + __gm__ UDMAWQCtx* qpCtxEntry) { sgeCtx->len = messageLen; - sgeCtx->tokenId = 0; + sgeCtx->tokenId = qpCtxEntry->localTokenId; sgeCtx->va = reinterpret_cast(localAddr); } @@ -240,18 +246,19 @@ __aicore__ inline void UDMAPostSend( { __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, pe, qpIdx); uint32_t wqeSize = 1U << qpCtxEntry->baseBkShift; + uint32_t depth = qpCtxEntry->depth; uint32_t curHead = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->headAddr), 0); uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); UDMAPollCQWhenSQOverflow(udmaInfo, qpCtxEntry, wqeCnt, pe, qpIdx); __gm__ UDMAMemInfo* remoteMemInfo = UDMAGetRemoteMemInfo(udmaInfo, pe); __gm__ uint8_t* wqeAddr = - reinterpret_cast<__gm__ uint8_t*>(qpCtxEntry->bufAddr + wqeSize * (curHead % TILEXR_UDMA_SQ_BB_COUNT)); + reinterpret_cast<__gm__ uint8_t*>(qpCtxEntry->bufAddr + wqeSize * (curHead % depth)); __gm__ UDMASqeCtx* sqeCtx = reinterpret_cast<__gm__ UDMASqeCtx*>(wqeAddr); - UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, opcode, signalParams); + UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth, opcode, signalParams); __gm__ UDMASgeCtx* sgeCtx = reinterpret_cast<__gm__ UDMASgeCtx*>(UDMAGetSgeCtxAddr(wqeAddr, opcode)); - UDMAFillSgeCtx(sgeCtx, messageLen, localAddr); + UDMAFillSgeCtx(sgeCtx, messageLen, localAddr, qpCtxEntry); uint32_t wqeBbCnt = UDMAWqeBBCnt(opcode); UDMACleanCacheLines(wqeAddr, wqeSize * wqeBbCnt); curHead += wqeBbCnt; diff --git a/src/include/tilexr_udma_types.h b/src/include/tilexr_udma_types.h index 57d77e34..1fd6ebd7 100644 --- a/src/include/tilexr_udma_types.h +++ b/src/include/tilexr_udma_types.h @@ -52,6 +52,7 @@ struct UDMAWQCtx { UDMADBMode dbMode; uint64_t dbAddr; uint32_t sl; + uint32_t localTokenId; uint64_t wqeCntAddr; uint64_t amoAddr; }; diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index ff98ec3d..7e5320ae 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -124,6 +124,10 @@ target_include_directories(test_tilexr_udma_transport_layout PRIVATE ${TILEXR_ROOT}/src/comm ) +target_compile_definitions(test_tilexr_udma_transport_layout PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" +) + # 集成测试：TileXR UDMA add_executable(test_tilexr_udma integration/test_tilexr_udma.cpp @@ -207,6 +211,7 @@ if(BUILD_TILEXR_UDMA_DEMO) -I${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/ -I${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/runtime/ -I${ASCEND_HOME_PATH}/${ARCH}-linux/include/ + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include/ -I${ASCEND_DRIVER_PATH}/kernel/inc -I${TILEXR_ROOT}/3rdparty -I${TILEXR_ROOT}/src/include @@ -240,6 +245,7 @@ if(BUILD_TILEXR_UDMA_DEMO) DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" "${TILEXR_ROOT}/src/include/tilexr_udma.h" + "${TILEXR_ROOT}/src/include/tilexr_data_as_flag.h" VERBATIM COMMENT "Building TileXR UDMA demo kernel with bisheng" ) diff --git a/tests/udma/demo/ALLTOALL_8P_RUNBOOK.md b/tests/udma/demo/ALLTOALL_8P_RUNBOOK.md new file mode 100644 index 00000000..4b9587e4 --- /dev/null +++ b/tests/udma/demo/ALLTOALL_8P_RUNBOOK.md @@ -0,0 +1,201 @@ +# TileXR UDMA All-to-All 8P Runbook + +This runbook records the 8P all-to-all demo version based on commit +`e552736 Add UDMA demo checker collectives`. + +## Scope + +- Demo path: `tests/udma/demo` +- Operator mode: `test_type=2`, all-to-all UDMA put +- Target hardware: A5 / Ascend950 / 950 +- Process model: one local process per rank +- Validated baseline: `rank_size=8`, `npu_count=8`, `first_npu=0` + +The all-to-all layout is: + +- rank `src` fills input slice `dst` with `100000 + src * 1000 + dst`; +- rank `src` sends slice `dst` to rank `dst`; +- rank `dst` output is ordered by source rank. + +For rank `0`, the output sample should contain: + +```text +from0=100000 from1=101000 from2=102000 from3=103000 from4=104000 from5=105000 from6=106000 from7=107000 +``` + +For rank `7`, the output sample should contain: + +```text +from0=100007 from1=101007 from2=102007 from3=103007 from4=104007 from5=105007 from6=106007 from7=107007 +``` + +## Environment + +Use a root shell on the Ascend machine. + +```bash +cd /path/to/TileXR +source scripts/common_env.sh +npu-smi info +``` + +Expected: + +- `npu-smi info` lists at least 8 usable devices; +- CANN environment variables are available after `source scripts/common_env.sh`; +- `bisheng` is available for building `tilexr_udma_demo_kernel.cpp`; +- MPI, if needed by surrounding scripts, is under `/usr/local/mpi/`. + +## Build + +Build `tile-comm` and install it into the repository `install` directory: + +```bash +cd /path/to/TileXR +source scripts/common_env.sh +cmake -S . -B /tmp/tilexr-build-udma -DCMAKE_INSTALL_PREFIX="$PWD/install" +cmake --build /tmp/tilexr-build-udma --target tile-comm -j"$(nproc)" +cmake --install /tmp/tilexr-build-udma +``` + +Build the UDMA demo: + +```bash +cd /path/to/TileXR/tests/udma +bash build.sh +``` + +Check artifacts: + +```bash +test -x install/bin/tilexr_udma_demo +test -f install/lib/libtilexr_udma_demo_kernel.so +``` + +## Run 8P All-to-All + +The script arguments are: + +```text +run_tilexr_udma_demo.sh +``` + +Run the normal/default IPC initialization path: + +```bash +cd /path/to/TileXR/tests/udma +export TILEXR_COMM_ID=127.0.0.1:10067 +bash demo/run_tilexr_udma_demo.sh 2 8 16 8 0 +``` + +Run with explicit PID IPC mode: + +```bash +cd /path/to/TileXR/tests/udma +export TILEXR_COMM_ID=127.0.0.1:10077 +TILEXR_IPC_PID_MODE=pid bash demo/run_tilexr_udma_demo.sh 2 8 16 8 0 +``` + +Run with explicit SDID IPC mode: + +```bash +cd /path/to/TileXR/tests/udma +export TILEXR_COMM_ID=127.0.0.1:10087 +TILEXR_IPC_PID_MODE=sdid bash demo/run_tilexr_udma_demo.sh 2 8 16 8 0 +``` + +Use a different `TILEXR_COMM_ID` port for concurrent or repeated runs to avoid +the demo TCP barrier colliding with a previous process. + +## 1M Data Run + +For a 1M-elements-per-peer checker-sized run: + +```bash +cd /path/to/TileXR/tests/udma +export TILEXR_COMM_ID=127.0.0.1:10107 +bash demo/run_tilexr_udma_demo.sh 2 8 1048576 8 0 +``` + +This means: + +- `elements_per_rank=1048576` int32 elements per destination slice; +- each rank input buffer has `8 * 1048576` int32 elements; +- each rank input/output buffer is 32 MiB. + +## Log Checks + +Each run writes logs under: + +```text +tests/udma/logs/tilexr_udma_demo_YYYYmmdd_HHMMSS/ +``` + +Quick success check: + +```bash +cd /path/to/TileXR/tests/udma +latest=$(ls -td logs/tilexr_udma_demo_* | head -n1) +grep -R "TileXR UDMA demo success" "$latest" +grep -R "UDMA=enabled" "$latest" +grep -R "TileXRUDMARegister success" "$latest" +``` + +There should be 8 success lines, one for each rank. + +Check all-to-all samples: + +```bash +grep -R "alltoall output sample" "$latest" +``` + +Check for failures: + +```bash +grep -R "ALLTOALL MISMATCH\|DATA MISMATCH\|TileXR UDMA demo failed\|ERROR" "$latest" || true +``` + +Expected: no mismatch or failure lines. + +If UDMA CQ is incomplete, the demo may print: + +```text +alltoall UDMA CQ incomplete, use IPC fallback +alltoall IPC fallback completed +``` + +In that case, the final correctness criterion is still the all-to-all output +validation and `TileXR UDMA demo success` on every rank. + +## Offline Checker + +The layout checker can be built without running the hardware demo: + +```bash +cd /path/to/TileXR/tests/udma +g++ -std=c++14 -O2 \ + -I . \ + -DTILEXR_SOURCE_ROOT='"'/path/to/TileXR'"' \ + unit/test_tilexr_udma_alltoall_layout.cpp \ + -o /tmp/test_tilexr_udma_alltoall_layout +/tmp/test_tilexr_udma_alltoall_layout +``` + +Expected output: + +```text +TileXR UDMA all-to-all layout checks passed +``` + +This checker validates the all-to-all input pattern, expected output layout, and +source-level debug layout assumptions. It does not execute UDMA or AICore code. + +## Notes + +- `test_type=2` is the all-to-all path. +- `test_type=3` is the all-reduce path and is not covered by this runbook. +- `TILEXR_IPC_PID_MODE=sdid` forces `rtSetIpcMemorySuperPodPid`. +- `TILEXR_IPC_PID_MODE=pid` forces `rtSetIpcMemPid`. +- Leaving `TILEXR_IPC_PID_MODE` unset uses TileXR's chip default. +- The demo host source does not use `shmem.h`; process synchronization is a + local TCP barrier derived from `TILEXR_COMM_ID`. diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index ba3dfd94..630a3c86 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -20,6 +20,7 @@ #include "acl/acl.h" #include "tilexr_api.h" +#include "tilexr_data_as_flag.h" #include "tilexr_types.h" #include "tilexr_udma_allreduce_layout.h" #include "tilexr_udma_alltoall_layout.h" @@ -437,6 +438,37 @@ bool AllToAllUdmaComplete(int rankSize, const std::vector& debug) return true; } +void PrintAllToAllUdmaDebug(int rank, int rankSize, const std::vector& debug) +{ + constexpr int rangeBase = kDebugUdmaStatusBase + 16; + constexpr int wqeBeforeBase = kDebugUdmaStatusBase + 32; + constexpr int wqeAfterBase = kDebugUdmaStatusBase + 48; + constexpr int localTokenBase = kDebugUdmaStatusBase + 64; + constexpr int remoteBaseLowBase = kDebugUdmaStatusBase + 80; + constexpr int memAddrLowBase = kDebugUdmaStatusBase + 96; + constexpr int tpnBase = kDebugUdmaStatusBase + 112; + std::cout << "[rank " << rank << "] alltoall udma peer debug:"; + for (int peer = 0; peer < rankSize && peer < 16; ++peer) { + std::cout << " peer" << peer + << "{status=" << debug[kDebugUdmaStatusBase + peer] + << ",range=" << debug[rangeBase + peer] + << ",wqe=" << debug[wqeBeforeBase + peer] << "->" << debug[wqeAfterBase + peer] + << ",token=" << debug[localTokenBase + peer] + << ",regLo=" << debug[remoteBaseLowBase + peer] + << ",memLo=" << debug[memAddrLowBase + peer] + << ",tpn=" << debug[tpnBase + peer] + << "}"; + } + std::cout << std::endl; +} + +size_t AllToAllDataAsFlagStagingBytes(int rankSize, int32_t elementsPerPeer) +{ + const uint64_t payloadBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + const uint64_t blocks = TileXR::DataAsFlagBlockCountForPayloadBytes(payloadBytes); + return static_cast(static_cast(rankSize) * blocks * TileXR::DATA_AS_FLAG_BLOCK_BYTES); +} + void Cleanup( TileXRCommPtr comm, aclrtStream stream, void* registeredMemory, int32_t* debug, int rank, int deviceId) { @@ -526,9 +558,24 @@ int main(int argc, char** argv) bool isAllToAll = testType == 2; bool isAllReduce = testType == 3; + bool strictAllToAllUdma = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_USE_UDMA", 0) != 0; + bool dumpAllToAllOnStrictFail = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_DUMP_ON_STRICT_FAIL", 0) != 0; + bool useAllToAllDataAsFlagIpc = isAllToAll && !strictAllToAllUdma; + bool forceAllToAllIpcFallback = false; bool hasOutput = isAllToAll || isAllReduce; size_t dataCount = static_cast(rankSize) * elementsPerRank; size_t dataBytes = dataCount * sizeof(int32_t); + if (isAllToAll) { + const size_t stagingBytes = AllToAllDataAsFlagStagingBytes(rankSize, elementsPerRank); + if (stagingBytes > static_cast(TileXR::IPC_BUFF_MAX_SIZE)) { + std::cerr << "[rank " << rank << "] ERROR: alltoall data-as-flag IPC fallback staging requires " + << stagingBytes << " bytes, exceeds IPC data capacity " + << TileXR::IPC_BUFF_MAX_SIZE << std::endl; + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + PrintStatus(rank, "alltoall data-as-flag staging bytes=" + std::to_string(stagingBytes)); + } size_t inputOffset = 0; size_t outputOffset = hasOutput ? dataBytes : 0; size_t signalBytes = static_cast(rankSize) * sizeof(uint64_t); @@ -547,12 +594,30 @@ int main(int argc, char** argv) auto input = reinterpret_cast(static_cast(registeredMemory) + inputOffset); auto output = reinterpret_cast(static_cast(registeredMemory) + outputOffset); auto signals = reinterpret_cast(static_cast(registeredMemory) + signalOffset); - if (!CheckTileXR(rank, "TileXRUDMARegister", - TileXRUDMARegister(comm, static_cast(registeredMemory), registeredBytes, &udmaHandle))) { - Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); - return 1; + if (useAllToAllDataAsFlagIpc) { + PrintStatus(rank, "skip TileXRUDMARegister for alltoall data-as-flag IPC path"); + forceAllToAllIpcFallback = true; + } else { + int registerRet = + TileXRUDMARegister(comm, static_cast(registeredMemory), registeredBytes, &udmaHandle); + if (registerRet != TileXR::TILEXR_SUCCESS) { + if (!isAllToAll || strictAllToAllUdma) { + if (strictAllToAllUdma) { + std::cerr << "[rank " << rank << "] ERROR: strict alltoall UDMA registration failed" + << " ret=" << registerRet << std::endl; + } + CheckTileXR(rank, "TileXRUDMARegister", registerRet); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + std::cerr << "[rank " << rank + << "] WARNING: TileXRUDMARegister failed; use alltoall data-as-flag IPC fallback" + << " ret=" << registerRet << std::endl; + forceAllToAllIpcFallback = true; + } else { + udmaRegistered = true; + } } - udmaRegistered = true; PrintStatus(rank, "registered UDMA memory base=" + std::to_string(reinterpret_cast(registeredMemory)) + " bytes=" + std::to_string(registeredBytes) + " inputOffset=" + std::to_string(inputOffset) + @@ -602,10 +667,14 @@ int main(int argc, char** argv) } if (testType == 2) { - PrintStatus(rank, "launch all-to-all kernel"); - launch_tilexr_udma_all_to_all( - 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset)); + if (forceAllToAllIpcFallback) { + PrintStatus(rank, "skip all-to-all UDMA kernel; use data-as-flag IPC fallback"); + } else { + PrintStatus(rank, "launch all-to-all kernel"); + launch_tilexr_udma_all_to_all( + 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset)); + } } else if (testType == 3) { PrintStatus(rank, "launch all-reduce IPC scatter kernel"); launch_tilexr_all_reduce_ipc_scatter( @@ -663,9 +732,40 @@ int main(int argc, char** argv) } bool usedIpcFallback = false; - if (isAllToAll && !AllToAllUdmaComplete(rankSize, hostDebug)) { + bool allToAllUdmaComplete = !isAllToAll || AllToAllUdmaComplete(rankSize, hostDebug); + if (isAllToAll && strictAllToAllUdma && !allToAllUdmaComplete) { + std::cerr << "[rank " << rank << "] ERROR: strict alltoall UDMA CQ incomplete:"; + for (int peer = 0; peer < rankSize; ++peer) { + std::cerr << " peer" << peer << "=" << hostDebug[kDebugUdmaStatusBase + peer]; + } + std::cerr << std::endl; + PrintAllToAllUdmaDebug(rank, rankSize, hostDebug); + if (dumpAllToAllOnStrictFail) { + std::vector strictFailOutput(dataCount, 0); + if (CopyDeviceToHost(rank, strictFailOutput.data(), dataBytes, + output, dataBytes, "alltoall output after strict UDMA fail")) { + (void)ValidateAllToAllData(rank, rankSize, strictFailOutput, elementsPerRank); + } + } + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (isAllToAll && (forceAllToAllIpcFallback || !allToAllUdmaComplete)) { usedIpcFallback = true; - std::cout << "[rank " << rank << "] alltoall UDMA CQ incomplete, use IPC fallback:"; + std::cout << "[rank " << rank << "] alltoall use data-as-flag IPC fallback"; + if (forceAllToAllIpcFallback) { + if (useAllToAllDataAsFlagIpc) { + std::cout << " by default"; + } else { + std::cout << " after UDMA registration failure"; + } + } else { + std::cout << " after UDMA CQ incomplete"; + } + std::cout << ":"; for (int peer = 0; peer < rankSize; ++peer) { std::cout << " peer" << peer << "=" << hostDebug[kDebugUdmaStatusBase + peer]; } @@ -720,6 +820,9 @@ int main(int argc, char** argv) std::cout << " d" << i << "=" << hostDebug[i]; } std::cout << std::endl; + if (isAllToAll) { + PrintAllToAllUdmaDebug(rank, rankSize, hostDebug); + } if (usedIpcFallback) { std::cout << "[rank " << rank << "] alltoall IPC fallback completed" << " scatter=" << hostDebug[kDebugIpcScatter] diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index eeadbbe4..988ef4b5 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -4,17 +4,49 @@ */ #include "kernel_operator.h" +#include "tilexr_data_as_flag.h" #include "tilexr_udma.h" constexpr int32_t TILEXR_UDMA_DEMO_MAGIC = 0x5444554d; // "TDUM" constexpr uint64_t TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET = TileXR::IPC_DATA_OFFSET; +constexpr uint64_t TILEXR_UDMA_DEMO_DATA_AS_FLAG_STAGING_OFFSET = TileXR::IPC_DATA_OFFSET; +constexpr uint32_t TILEXR_UDMA_DEMO_DATA_AS_FLAG_UB_BYTES = 64 * 1024; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE = 6; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_RANGE_VALID_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 16; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 32; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 48; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 64; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 80; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 96; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_TPN_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 112; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER = TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER = TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER + 1; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER = TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER + 1; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SUM = TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER + 1; +namespace { + +__aicore__ inline uint64_t AllToAllPayloadBytes(int32_t elementsPerPeer) +{ + return static_cast(elementsPerPeer) * sizeof(int32_t); +} + +__aicore__ inline uint64_t AllToAllDataAsFlagSegmentBytes(uint64_t payloadBytes) +{ + return static_cast(TileXR::DataAsFlagBlockCountForPayloadBytes(payloadBytes)) * + TileXR::DATA_AS_FLAG_BLOCK_BYTES; +} + +} // namespace + extern "C" __global__ __aicore__ void tilexr_udma_all_gather_kernel( GM_ADDR commArgsGM, GM_ADDR dataGM, GM_ADDR debugGM, int32_t elementsPerRank) { @@ -112,21 +144,42 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( return; } - auto selfSrc = input + rank * elementsPerPeer; - auto selfDst = output + rank * elementsPerPeer; + const uint64_t payloadBytes = AllToAllPayloadBytes(elementsPerPeer); + auto selfSrc = input + static_cast(rank) * elementsPerPeer; + auto selfDst = output + static_cast(rank) * elementsPerPeer; for (int32_t i = 0; i < elementsPerPeer; ++i) { selfDst[i] = selfSrc[i]; } - uint32_t bytes = static_cast(elementsPerPeer * sizeof(int32_t)); + uint32_t bytes = static_cast(payloadBytes); for (int32_t peer = 0; peer < rankSize; ++peer) { if (peer == rank) { continue; } - auto localSrc = input + peer * elementsPerPeer; + auto localSrc = input + static_cast(peer) * elementsPerPeer; uint64_t remoteOffset = outputByteOffset + - static_cast(rank) * elementsPerPeer * sizeof(int32_t); + static_cast(rank) * payloadBytes; + auto registry = TileXR::GetUDMARegistry(args); + auto udmaInfo = TileXR::GetUDMAInfo(args); + auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); + auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); + bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteOffset, bytes); + uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + if (debug != nullptr && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_RANGE_VALID_BASE + peer] = rangeValid ? 1 : 0; + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE + peer] = static_cast(wqeBefore); + debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = static_cast(wqCtx->localTokenId); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = + static_cast(reinterpret_cast(registry->regions[peer].base) & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = + static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_TPN_BASE + peer] = static_cast(remoteMemInfo->tpn); + } TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); + uint32_t wqeAfter = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + if (debug != nullptr && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE + peer] = static_cast(wqeAfter); + } uint32_t status = TileXR::UDMAQuietStatus(args, peer); if (debug != nullptr) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); @@ -143,14 +196,27 @@ extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_scatter_kernel( int32_t rank = args->rank; int32_t rankSize = args->rankSize; - for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { - auto localSrc = input + dstRank * elementsPerPeer; - auto remoteBase = reinterpret_cast<__gm__ int32_t*>( - args->peerMems[dstRank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); - auto remoteDst = remoteBase + rank * elementsPerPeer; - for (int32_t i = 0; i < elementsPerPeer; ++i) { - remoteDst[i] = localSrc[i]; + const uint64_t payloadBytes = AllToAllPayloadBytes(elementsPerPeer); + const uint64_t segmentBytes = AllToAllDataAsFlagSegmentBytes(payloadBytes); + auto inputBytes = reinterpret_cast<__gm__ uint8_t*>(input); + + AscendC::TPipe pipe; + AscendC::TBuf tBuf; + pipe.InitBuffer(tBuf, TILEXR_UDMA_DEMO_DATA_AS_FLAG_UB_BYTES); + AscendC::LocalTensor scratch = tBuf.Get(); + if (TileXR::DataAsFlagInit(scratch) == 0U) { + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER] = -1; } + return; + } + + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + auto localSrc = inputBytes + static_cast(dstRank) * payloadBytes; + auto remoteDst = reinterpret_cast<__gm__ uint8_t*>( + args->peerMems[dstRank] + TILEXR_UDMA_DEMO_DATA_AS_FLAG_STAGING_OFFSET + + static_cast(rank) * segmentBytes); + (void)TileXR::DataAsFlagSend(remoteDst, localSrc, payloadBytes, scratch); } if (debug != nullptr) { debug[TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER] = 1; @@ -164,13 +230,25 @@ extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_gather_kernel( auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); - auto localBase = reinterpret_cast<__gm__ int32_t*>( - args->peerMems[args->rank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + const uint64_t payloadBytes = AllToAllPayloadBytes(elementsPerPeer); + const uint64_t segmentBytes = AllToAllDataAsFlagSegmentBytes(payloadBytes); + auto outputBytes = reinterpret_cast<__gm__ uint8_t*>(output); + auto localBase = reinterpret_cast<__gm__ uint8_t*>( + args->peerMems[args->rank] + TILEXR_UDMA_DEMO_DATA_AS_FLAG_STAGING_OFFSET); + + AscendC::TPipe pipe; + AscendC::TBuf tBuf; + pipe.InitBuffer(tBuf, TILEXR_UDMA_DEMO_DATA_AS_FLAG_UB_BYTES); + AscendC::LocalTensor scratch = tBuf.Get(); + for (int32_t srcRank = 0; srcRank < args->rankSize; ++srcRank) { - auto localSrc = localBase + srcRank * elementsPerPeer; - auto localDst = output + srcRank * elementsPerPeer; - for (int32_t i = 0; i < elementsPerPeer; ++i) { - localDst[i] = localSrc[i]; + auto localSrc = localBase + static_cast(srcRank) * segmentBytes; + auto localDst = outputBytes + static_cast(srcRank) * payloadBytes; + if (!TileXR::DataAsFlagCheckAndRecv(localSrc, payloadBytes, localDst, scratch)) { + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER] = -1; + } + return; } } if (debug != nullptr) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 6364c29f..0219f9aa 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -97,6 +97,35 @@ void TestBuildAllToAllOutput() CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(output, 2, rankSize, elementsPerPeer), true); } +void TestAllToAllMaxRank256With64MiBPerRank() +{ + constexpr int rankSize = 256; + constexpr size_t perRankBytes = 64ULL * 1024ULL * 1024ULL; + constexpr int32_t elementsPerPeer = + static_cast(perRankBytes / (sizeof(int32_t) * rankSize)); + CHECK_EQ(elementsPerPeer, 65536); + + std::vector buffer(static_cast(rankSize) * elementsPerPeer, -1); + const int sampleRanks[] = {0, 1, 127, 255}; + for (int rank : sampleRanks) { + TileXR::Demo::FillAllToAllInput(buffer, rank, rankSize, elementsPerPeer); + CHECK_EQ(buffer[0], TileXR::Demo::AllToAllValue(rank, 0)); + CHECK_EQ(buffer[static_cast(rankSize - 1) * elementsPerPeer], + TileXR::Demo::AllToAllValue(rank, rankSize - 1)); + CHECK_EQ(buffer[static_cast(rankSize) * elementsPerPeer - 1], + TileXR::Demo::AllToAllValue(rank, rankSize - 1)); + + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + std::fill(buffer.begin() + static_cast(srcRank) * elementsPerPeer, + buffer.begin() + static_cast(srcRank + 1) * elementsPerPeer, + TileXR::Demo::AllToAllValue(srcRank, rank)); + } + CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(buffer, rank, rankSize, elementsPerPeer), true); + buffer[static_cast(rankSize) * elementsPerPeer - 1] = -1; + CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(buffer, rank, rankSize, elementsPerPeer), false); + } +} + void TestDemoDebugLayoutSource() { const std::string demo = @@ -110,6 +139,30 @@ void TestDemoDebugLayoutSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer"); } +void TestAllToAllDataAsFlagSource() +{ + const std::string demo = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); + const std::string kernel = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + + CHECK_CONTAINS(demo, "useAllToAllDataAsFlagIpc"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_USE_UDMA"); + CHECK_CONTAINS(demo, "skip TileXRUDMARegister for alltoall data-as-flag IPC path"); + CHECK_CONTAINS(demo, "forceAllToAllIpcFallback"); + CHECK_CONTAINS(demo, "strictAllToAllUdma"); + CHECK_CONTAINS(demo, "ERROR: strict alltoall UDMA registration failed"); + CHECK_CONTAINS(demo, "ERROR: strict alltoall UDMA CQ incomplete"); + CHECK_CONTAINS(demo, "TileXRUDMARegister failed; use alltoall data-as-flag IPC fallback"); + CHECK_CONTAINS(demo, "skip all-to-all UDMA kernel; use data-as-flag IPC fallback"); + CHECK_CONTAINS(kernel, "#include \"tilexr_data_as_flag.h\""); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DATA_AS_FLAG_STAGING_OFFSET"); + CHECK_CONTAINS(kernel, "DataAsFlagBlockCountForPayloadBytes"); + CHECK_CONTAINS(kernel, "DataAsFlagInit"); + CHECK_CONTAINS(kernel, "DataAsFlagSend"); + CHECK_CONTAINS(kernel, "DataAsFlagCheckAndRecv"); +} + } // namespace int main() @@ -117,7 +170,9 @@ int main() TestAllToAllInputPattern(); TestAllToAllOutputValidation(); TestBuildAllToAllOutput(); + TestAllToAllMaxRank256With64MiBPerRank(); TestDemoDebugLayoutSource(); + TestAllToAllDataAsFlagSource(); if (g_failures != 0) { std::cerr << g_failures << " all-to-all layout checks failed" << std::endl; return 1; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index cb140157..2fc3bfa3 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -1,9 +1,16 @@ #include +#include #include +#include +#include #include #include "udma/tilexr_udma_layout.h" +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + namespace { int g_failures = 0; @@ -27,6 +34,30 @@ int g_failures = 0; } \ } while (0) +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +#define CHECK_NOT_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) != std::string::npos) { \ + std::cerr << "CHECK_NOT_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + void TestHostLayoutUsesDeviceRelativePointers() { std::vector sq(2); @@ -85,12 +116,87 @@ void TestRejectsMismatchedArrays() CHECK_EQ(ret, TileXR::TILEXR_UDMA_LAYOUT_INVALID); } +void TestTransportUsesPerPeerQueues() +{ + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + CHECK_CONTAINS(transport, "struct PeerQueueState"); + CHECK_CONTAINS(transport, "std::map peerQueues"); + CHECK_CONTAINS(transport, "CreatePeerQueue("); + CHECK_CONTAINS(transport, "state.peerQueues[peer]"); + CHECK_CONTAINS(transport, "localPeerImports[peer].in.key"); + CHECK_CONTAINS(transport, "allImports[(peer * options_.rankSize + options_.rank)"); + CHECK_CONTAINS(transport, "queue.localWq"); + CHECK_CONTAINS(transport, "queue.localCq"); + CHECK_CONTAINS(transport, "queue.tpn"); + CHECK_NOT_CONTAINS(transport, "void* qpHandle = nullptr;\n CqInfoT cqInfo"); +} + +void TestRootInfoEidBytesSelectRuntimeContexts() +{ + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + + CHECK_CONTAINS(transport, "ReadTextFile(\"/etc/hccl_rootinfo.json\")"); + CHECK_CONTAINS(transport, "root.eidByLocalId[localId][eidIndex] = eid"); + CHECK_CONTAINS(transport, "root.portToEidByLocalId[localId][port] = eidIndex"); + CHECK_CONTAINS(transport, "localEidByEid_ = localEids->second"); + CHECK_CONTAINS(transport, "auto targetEidIt = localEidByEid_.find(eidIndex)"); + CHECK_CONTAINS(transport, + "matched = std::memcmp(infoList[i].eid.raw, targetEidIt->second.raw, sizeof(infoList[i].eid.raw)) == 0"); + CHECK_CONTAINS(transport, "attr.ub.eidIndex = infoList[i].eidIndex"); + CHECK_CONTAINS(transport, "ctxHandleByEid_[eidIndex] = ctxHandle"); +} + +void TestMemoryRegistrationUsesOfficialUbFlags() +{ + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.cacheable = 0"); + CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.nonPin = 0"); + CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.userIova = 0"); + CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.tokenIdValid = 1"); +} + +void TestDeviceSgeUsesPerPeerLocalTokenId() +{ + const std::string types = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma_types.h"); + const std::string device = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + + CHECK_CONTAINS(types, "uint32_t localTokenId"); + CHECK_CONTAINS(device, "sgeCtx->tokenId = qpCtxEntry->localTokenId"); + CHECK_NOT_CONTAINS(device, "sgeCtx->tokenId = 0;"); + CHECK_CONTAINS(transport, "queue.localWq.localTokenId = localMrIt->second.tokenId"); +} + +void TestDeviceSqeInitializesOfficialFields() +{ + const std::string device = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); + + CHECK_CONTAINS(device, "sqeCtx->sqeBbIdx = curHead % depth"); + CHECK_CONTAINS(device, "sqeCtx->rsv0 = 0"); + CHECK_CONTAINS(device, "sqeCtx->rsv1 = 0"); + CHECK_CONTAINS(device, "sqeCtx->rsv2 = 0"); + CHECK_CONTAINS(device, "sqeCtx->rsv3 = 0"); + CHECK_CONTAINS(device, "UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth"); +} + } // namespace int main() { TestHostLayoutUsesDeviceRelativePointers(); TestRejectsMismatchedArrays(); + TestTransportUsesPerPeerQueues(); + TestRootInfoEidBytesSelectRuntimeContexts(); + TestMemoryRegistrationUsesOfficialUbFlags(); + TestDeviceSgeUsesPerPeerLocalTokenId(); + TestDeviceSqeInitializesOfficialFields(); if (g_failures != 0) { std::cerr << g_failures << " UDMA transport layout checks failed" << std::endl; return 1; From 8928ca3cf6be40089da7d562d7d63cf37c79daa0 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Thu, 25 Jun 2026 20:15:41 +0800 Subject: [PATCH 006/163] udma alltoall ranksize core perf ok --- docs/alltoall-udma-success.md | 399 ++++++++++++++++ src/collectives/host/collective_utils.cpp | 14 + src/collectives/host/collective_utils.h | 2 + src/collectives/host/tilexr_collectives.cpp | 6 + src/collectives/kernels/CMakeLists.txt | 16 +- src/collectives/kernels/collectives.h | 11 +- src/collectives/kernels/datacopy_gm2gm.h | 13 +- .../kernels/kernels/collectives.cce | 24 +- src/collectives/kernels/lccl_op.h | 2 +- src/collectives/kernels/tilexr_lccl_op.cpp | 2 +- src/comm/tilexr_comm.cpp | 22 +- src/comm/tilexr_internal.cpp | 1 + .../unit/test_collective_host_utils.cpp | 22 + ...st_tilexr_collectives_kernel_ownership.cpp | 53 +++ tests/comm/CMakeLists.txt | 9 + tests/comm/unit/test_tilexr_source_guards.cpp | 11 + .../unit/test_tilexr_udma_env_sources.cpp | 44 ++ tests/udma/demo/tilexr_udma_alltoall_layout.h | 39 ++ tests/udma/demo/tilexr_udma_demo.cpp | 372 +++++++++++++-- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 445 +++++++++++++++++- tests/udma/run_fused_prof.sh | 26 + tests/udma/run_plain_ipc_8m_4p.sh | 12 + .../unit/test_tilexr_chip_map_sources.cpp | 1 + .../unit/test_tilexr_udma_alltoall_layout.cpp | 40 +- 24 files changed, 1516 insertions(+), 70 deletions(-) create mode 100644 docs/alltoall-udma-success.md create mode 100644 tests/comm/unit/test_tilexr_udma_env_sources.cpp create mode 100644 tests/udma/run_fused_prof.sh create mode 100644 tests/udma/run_plain_ipc_8m_4p.sh diff --git a/docs/alltoall-udma-success.md b/docs/alltoall-udma-success.md new file mode 100644 index 00000000..46d4082e --- /dev/null +++ b/docs/alltoall-udma-success.md @@ -0,0 +1,399 @@ +# All-to-All UDMA 多核实现 — 验证成功记录 + +> 记录 TileXR UDMA all-to-all kernel 从单核标量拷贝改造为 `rankSize` 核并行版本的设计、验证流程与性能结果。 +> 验证环境:8 × Ascend950DT(A5 / 3510 架构),CANN 25.1,bisheng 编译器。 +> 验证日期:2026-06-25。 + +--- + +## 1. 改造背景 + +原始 `tilexr_udma_all_to_all_kernel` 存在两个问题: + +1. **self-copy 用标量逐元素拷贝** — 本 rank 自己那一块 `for (i) selfDst[i] = selfSrc[i];`,效率极低,违反 Ascend C API 最佳实践(黑名单 `SetValue`/标量 GM 写)。 +2. **单核串行** — host 侧 `launch_tilexr_udma_all_to_all(1, ...)`,blockDim=1,所有 peer 的 UDMA PUT + self-copy 在 1 个 AI Core 上串行,未利用多核。 + +改造目标: + +- self-copy 改为 `DataCopyPad` 批量搬运(经 UB 中转,MTE2/MTE3 同步)。 +- 整个 kernel 改为 `rankSize` 核并行,每个 core 负责一个 peer 的搬运,压力分摊到 1/N。 + +--- + +## 2. 算法流程 + +### 2.1 整体语义 + +all-to-all:每个 rank 把自己 input 中**属于各个 peer 的那一块**送到**对方 rank 的 output 区域**。即 `output[peer][rank] = input[rank][peer]`。 + +- input 按 `[dstRank]` 切分:本 rank 持有所有 peer 的数据 +- output 按 `[srcRank]` 切分:本 rank 的 output 接收来自各 rank 的数据 + +### 2.2 多核分片策略 + +host 启动 `blockDim = rankSize` 个 block。kernel 内用 `AscendC::GetBlockIdx()` 取当前 block id,通过 stride 循环分配 peer: + +``` +for (peer = blockIdx; peer < rankSize; peer += blockNum) +``` + +当 `blockNum == rankSize` 时,block b 处理 peer b: + +| block id (b) | peer == rank? | 执行路径 | +|---|---|---| +| b == rank | 是 | **本地 self-copy**(DataCopyPad,无网络) | +| b != rank | 否 | **UDMA PUT 到 peer b**(硬件 DMA) | + +- 每个 core 只处理 1 个 peer,N 次搬运分摊到 N 个 core。 +- 不同 core 操作不同 peer 的发送队列(SQ),无 `headAddr`/`wqeCnt` 竞争。 +- 降级兼容:`blockNum < rankSize` 时 stride 循环让每个 core 处理多个 peer,逻辑仍正确。 + +### 2.3 阶段详解 + +**阶段 0:初始化与守卫** +- 解析 `CommArgs`,取 `rank`/`rankSize`/UDMA registry 指针。 +- `UDMARegistryEnabled(args)` 检查使能位 + registry 非空;未使能则 block 0 写完 debug 后 return。 +- **debug header 只由 block 0 写**(`blockIdx == 0`),避免多核并发写 `debug[0..5]` 竞争。 + +**阶段 1:self-copy 分支(peer == rank)** +- 本 rank 自己那一块,不经网络,纯本地搬移。 +- `TPipe` + `TBuf`(64KB UB)中转,64KB 分块: + - CopyIn:GM → UB(`DataCopyPad`,MTE2 异步)→ `SetFlag/WaitFlag` 同步 + - CopyOut:UB → GM(`DataCopyPad`,MTE3 异步)→ `SetFlag/WaitFlag` 同步 +- `DataCopyPad` 自动处理非对齐尾部,整 `bytes` 一次搬运,无标量补尾。 +- 末尾 `PipeBarrier()` 确保本地拷贝完成。 + +**阶段 2:UDMA PUT 分支(peer != rank)** +- `localSrc = input + peer*elementsPerPeer + inputElementOffset`(本 rank 持有、属于该 peer 的数据) +- `remoteOffset = outputByteOffset + rank*payloadBytes`(对方 output 中本 rank 的位置) +- 取该 peer 的 WQ 上下文 `UDMAGetWQCtx(udmaInfo, peer, 0)`、远端 mem info。 +- `UDMARegisteredRangeValid` 校验 `remoteOffset + bytes` 落在 host 注册给该 peer 的 region 内。 +- **投递 WQE**:`UDMAPutNbi` → `UDMAWrite` → `UDMAPostSend`: + - 在 SQ 环 `bufAddr + wqeSize*(curHead % depth)` 处填 `UDMASqeCtx`(远端地址/tpn/opcode=WRITE)+ `UDMASgeCtx`(本地源地址 + 长度) + - `UDMACleanCacheLines` 刷 WQE cacheline,`curHead += wqeBbCnt`,写 doorbell,递增 `wqeCnt` + - non-blocking:WQE 进 SQ 后硬件异步发起跨卡 DMA 写,GM→GM 直达 +- **Quiet 同步**:`UDMAQuietStatus` → `UDMAPollCQ` 轮询该 peer 完成队列(CQ),直到本次 WQE 的 CQE 产生,确保数据已落盘到远端 GM。 + +### 2.4 同步与退出保证 + +- `SetFlag/WaitFlag` 成对使用,`EVENT_ID0`,与仓库内 fused IPC kernel 同范式(已验证可用)。 +- self-copy 分支:event 配对完整(CopyIn 的 MTE2_MTE3 + CopyOut 的 MTE3_MTE2),循环内每轮自平衡,末尾 `PipeBarrier`。 +- UDMA 分支:无 UB event,仅 `UDMAQuietStatus` 轮询 CQ,CQE 必然产生(硬件完成保证),不会无限阻塞。 +- **每个 core 处理完自己的 peer 后循环结束、kernel 返回,正常退出。** msprof 实测 task 在 ~189μs 完成,无死锁。 + +--- + +## 3. 关键代码片段 + +### 3.1 kernel 主体(`tests/udma/demo/tilexr_udma_demo_kernel.cpp:122-229`) + +```cpp +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, + int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, int32_t chunkElements) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + bool enabled = TileXR::UDMARegistryEnabled(args); + + // Multi-core: one block per peer. Block b handles peer b: + // - peer == rank -> local self-copy via DataCopyPad + // - peer != rank -> UDMA PUT to that peer + // Block 0 writes the shared debug header; per-peer slots are written by + // their owning block only, so no cross-block debug races. + const int32_t blockIdx = AscendC::GetBlockIdx(); + const int32_t blockNum = AscendC::GetBlockNum(); + + if (blockIdx == 0 && debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = enabled ? 1 : 0; + debug[4] = elementsPerPeer; + debug[5] = static_cast(outputByteOffset); + } + if (!enabled) { + return; + } + + const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; + const uint64_t payloadBytes = AllToAllPayloadBytes(effectiveChunkElements); + const uint32_t bytes = static_cast(payloadBytes); + + // This block's assigned peer. When host launches rankSize blocks, block b + // handles peer b. If fewer blocks are launched, peers are round-robined + // and each block may handle more than one peer (still correct, just less + // parallel); stride == blockNum keeps peer slots disjoint across blocks. + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + // Self-copy: local DataCopyPad, no network. + auto selfSrc = input + static_cast(rank) * elementsPerPeer + inputElementOffset; + auto selfDst = output + static_cast(rank) * effectiveChunkElements; + constexpr uint32_t SELF_COPY_UB_BYTES = 64 * 1024; + AscendC::TPipe pipe; + AscendC::TBuf selfCopyTBuf; + pipe.InitBuffer(selfCopyTBuf, SELF_COPY_UB_BYTES); + AscendC::LocalTensor selfCopyLocal = selfCopyTBuf.Get(); + + auto selfSrcBytes = reinterpret_cast<__gm__ uint8_t*>(selfSrc); + auto selfDstBytes = reinterpret_cast<__gm__ uint8_t*>(selfDst); + for (uint32_t offset = 0; offset < bytes; offset += SELF_COPY_UB_BYTES) { + uint32_t copyBytes = (bytes - offset < SELF_COPY_UB_BYTES) + ? (bytes - offset) : SELF_COPY_UB_BYTES; + + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(selfSrcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, copyBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(selfCopyLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(selfDstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, copyBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, selfCopyLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); + continue; + } + + // UDMA PUT to peer. + auto localSrc = input + static_cast(peer) * elementsPerPeer + inputElementOffset; + uint64_t remoteOffset = outputByteOffset + + static_cast(rank) * payloadBytes; + auto registry = TileXR::GetUDMARegistry(args); + auto udmaInfo = TileXR::GetUDMAInfo(args); + auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); + auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); + bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteOffset, bytes); + uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + if (debug != nullptr && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_RANGE_VALID_BASE + peer] = rangeValid ? 1 : 0; + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE + peer] = static_cast(wqeBefore); + debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = static_cast(wqCtx->localTokenId); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = + static_cast(reinterpret_cast(registry->regions[peer].base) & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = + static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_TPN_BASE + peer] = static_cast(remoteMemInfo->tpn); + } + TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); + uint32_t wqeAfter = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + if (debug != nullptr && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE + peer] = static_cast(wqeAfter); + } + uint32_t status = TileXR::UDMAQuietStatus(args, peer); + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } + } +} +``` + +### 3.2 host 侧 launch(blockDim = rankSize) + +`tests/udma/demo/tilexr_udma_demo.cpp:807` 与 `:850` 两处调用,均从 `1` 改为 `static_cast(rankSize)`: + +```cpp +// chunked strict 路径(:807) +launch_tilexr_udma_all_to_all( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), chunkElements, static_cast(outputOffset), + 0, chunkElements); + +// 单 pass 路径(:850) +launch_tilexr_udma_all_to_all( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset), 0, + elementsPerRank); +``` + +### 3.3 launch wrapper(透传 blockDim) + +`tests/udma/demo/tilexr_udma_demo_kernel.cpp:731`: + +```cpp +void launch_tilexr_udma_all_to_all( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, + int32_t chunkElements) +{ + tilexr_udma_all_to_all_kernel<<>>( + commArgs, input, output, debug, elementsPerPeer, outputByteOffset, inputElementOffset, chunkElements); +} +``` + +### 3.4 UDMA PUT 底层路径(`src/include/tilexr_udma.h`) + +`UDMAPutNbi` → `UDMAWrite` → `UDMAPostSend` 核心逻辑: + +```cpp +template +__aicore__ inline void UDMAPutNbi( + const __gm__ CommArgs* args, int targetRank, const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount) +{ + if (!UDMARegistryEnabled(args)) return; + auto registry = GetUDMARegistry(args); + if (!UDMARegisteredRangeValid(registry, targetRank, byteOffset, byteCount)) return; + auto remoteAddr = UDMARegisteredRemoteAddr(registry, targetRank, byteOffset); + UDMAWrite(args, remoteAddr, reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)), + targetRank, 0, byteCount); +} +``` + +`UDMAPostSend`:填 WQE(`UDMASqeCtx` 远端 + `UDMASgeCtx` 本地)→ 刷 cacheline → 写 doorbell → 递增 wqeCnt。`UDMAQuiet`:`UDMAPollCQ` 轮询 CQ tail 直到 CQE 产生。 + +--- + +## 4. 验证流程 + +### 4.1 环境 + +- 服务器:`root@141.62.19.144`(8 × Ascend950DT,CANN 25.1.rc1.b142) +- 仓库路径:`/home/tileXR-new/` +- 编译:`bash tests/udma/build.sh`(bisheng 交叉编译,产出 `install/bin/tilexr_udma_demo`) +- 强制 UDMA 路径:`TILEXR_DEMO_ALLTOALL_USE_UDMA=1`(默认 0 走 IPC data-as-flag fallback) + +### 4.2 功能验证 + +```bash +# 2 rank, 1024 elements +TILEXR_DEMO_ALLTOALL_USE_UDMA=1 bash tests/udma/demo/run_tilexr_udma_demo.sh 2 2 1024 + +# 4 rank, 1024 elements +TILEXR_DEMO_ALLTOALL_USE_UDMA=1 bash tests/udma/demo/run_tilexr_udma_demo.sh 2 4 1024 4 0 +``` + +结果:全部 rank `TileXR UDMA demo success`,peer debug `status=0`(CQ 完成)、`wqe=0->1`(计数正确),输出数据正确。 + +### 4.3 msprof 性能采集(4 rank, 8M 数据) + +数据量 8M = 8*1024*1024 字节 = 2,097,152 个 int32。启动脚本 `run_fused_prof.sh`: + +```bash +#!/bin/bash +set -e +source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true +TILEXR_ROOT=/home/tileXR-new +UDMA_DIR=${TILEXR_ROOT}/tests/udma +export LD_LIBRARY_PATH=${UDMA_DIR}/install/lib:${UDMA_DIR}/install/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64:${LD_LIBRARY_PATH} +export TILEXR_DEMO_ALLTOALL_USE_UDMA=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=1 + +RANK_SIZE=4 +ELEM=2097152 +BIN=${UDMA_DIR}/install/bin/tilexr_udma_demo + +pids=() +for rank in $(seq 0 $((RANK_SIZE-1))); do + RANK=${rank} RANK_SIZE=${RANK_SIZE} "${BIN}" "${RANK_SIZE}" "${rank}" 2 "${ELEM}" "${RANK_SIZE}" 0 \ + > /tmp/a2a_rank${rank}.log 2>&1 & + pids+=("$!") +done +ret=0 +for idx in "${!pids[@]}"; do + wait "${pids[$idx]}" && echo "rank ${idx} ok" || { echo "rank ${idx} FAIL $?"; ret=1; } +done +echo "=== rank tails ===" +for rank in $(seq 0 $((RANK_SIZE-1))); do echo "--- rank ${rank} ---"; tail -6 /tmp/a2a_rank${rank}.log 2>/dev/null; done +exit ${ret} +``` + +msprof 采集命令: + +```bash +source /usr/local/Ascend/ascend-toolkit/set_env.sh +export LD_LIBRARY_PATH=/usr/local/Ascend/cann-9.1.T560/aarch64-linux/lib64:/usr/local/Ascend/cann-9.1.T560/tools/profiler/lib64:$LD_LIBRARY_PATH +PROF_DIR=/home/tileXR-new/tests/udma/prof_out +rm -rf $PROF_DIR; mkdir -p $PROF_DIR +PATH=/usr/local/Ascend/cann-9.1.T560/tools/profiler/bin:$PATH \ +msprof --output="$PROF_DIR" \ + --application=/home/tileXR-new/tests/udma/run_fused_prof.sh \ + --task-time=l0 \ + --ai-core=on \ + --aic-mode=sample-based \ + --aic-freq=100 \ + --aic-metrics=PipeUtilization \ + --aicpu=on \ + --runtime-api=on +``` + +关键点: +- `--application` 指向启动脚本,msprof 跟踪其 fork 的 4 个 rank 子进程,4 个 device 全部采到。 +- `--aic-mode=sample-based` + `--aic-freq=100`:采样模式,对自定义 kernel(非 GE 模型)友好。 +- `TILEXR_DEMO_ALLTOALL_USE_UDMA=1` 在脚本内 export,确保走 UDMA 路径而非 IPC fallback。 + +--- + +## 5. 性能结果 + +### 5.1 执行结果 + +4 个 rank 全部 `TileXR UDMA demo success`,数据正确,无卡死。 + +### 5.2 Kernel 任务时长(op_summary) + +| Device | `tilexr_udma_all_to_all_kernel` 时长 (μs) | +|--------|------------------------------------------| +| dev0 | 192.397 | +| dev1 | 193.275 | +| dev2 | 185.162 | +| dev3 | 184.794 | + +平均约 **189 μs**(8M 数据 / 4 rank)。 + +### 5.3 AI Vector Core 利用率 + +每个 device 实际活跃 4 个 vector core,与 `blockDim=4` 吻合: + +| Core | 角色 | scalar | mte2 | mte3 | +|------|------|--------|------|------| +| Core36 | UDMA PUT | ~0.95 | 0 | ~0.003 | +| Core54 | self-copy | ~0.88-0.95 | ~0.32 | ~0.002 | +| Core72 | UDMA PUT | ~0.95 | 0 | ~0.003 | +| Core90 | UDMA PUT | ~0.95 | 0 | ~0.003 | + +- 其余 60+ core 全 0 +- Average:scalar 0.058,mte2 0.009,mte3 0.005 + +### 5.4 结论与瓶颈 + +1. **多核分片生效**:每 device 用 4 个 vector core,与 `blockDim=4` 完全吻合,`GetBlockIdx()` 分片正确,每个 core 正常退出(无死锁)。 +2. **利用率低是设计预期**:UDMA PUT 把数据搬运下放给硬件 DMA 引擎(SQ/CQ),core 只做标量寄存器操作(投 WQE + 轮询 CQ),故 AI Core 利用率不是衡量此 kernel 的好指标。 +3. **瓶颈在 UDMA 静默同步**:每个 PUT 分支 `PutNbi` 后立即 `UDMAQuietStatus` 轮询 CQ,是"发一个等一个"串行模式。4 个 peer 已多核并行,单 core 内的 quiet 轮询主导 ~189μs。 +4. **self-copy(Core54)的 mte2=0.32**:DataCopyPad 经 UB 中转 8M,mte2 占比合理;但仍是 GM→UB→GM 两跳,而 UDMA PUT 是 GM→GM 一跳。后续可把 self-copy 也改成 UDMA 自投递(向自己 rank 注册区 PUT)统一路径。 + +### 5.5 采集产物 + +性能数据落在 `/home/tileXR-new/tests/udma/prof_out/` 下 4 个 `PROF_*/` 目录,每个含 `mindstudio_profiler_output/` 下的: +- `op_summary_*.csv` — kernel 任务时长 +- `ai_core_utilization_*.csv` / `ai_vector_core_utilization_*.csv` — 每 core 利用率 +- `task_time_*.csv` — task 时序 +- `api_statistic_*.csv` — runtime API 统计 + +--- + +## 6. 关于"卡住"的澄清 + +改造过程中曾观察到运行卡死,排查结论: + +- 默认 `TILEXR_DEMO_ALLTOALL_USE_UDMA=0` 时,demo 走 **IPC data-as-flag fallback 路径**(非本 kernel),卡死发生在该 fallback 的 TCP barrier / flag 轮询,与多核 UDMA 改动**无关**。 +- 强制 `TILEXR_DEMO_ALLTOALL_USE_UDMA=1` 后,实际执行改后的多核 kernel,在 2/4 卡、1K~8M 数据量下全部稳定通过,无卡死。 + +--- + +## 7. 文件清单 + +| 文件 | 改动 | +|------|------| +| `tests/udma/demo/tilexr_udma_demo_kernel.cpp` | kernel 改多核分片 + self-copy 改 DataCopyPad | +| `tests/udma/demo/tilexr_udma_demo.cpp` | 两处 launch blockDim: 1 → rankSize | +| `tests/udma/demo/run_fused_prof.sh` | msprof 采集用启动脚本(新增) | \ No newline at end of file diff --git a/src/collectives/host/collective_utils.cpp b/src/collectives/host/collective_utils.cpp index 1ff77e61..52fb423a 100644 --- a/src/collectives/host/collective_utils.cpp +++ b/src/collectives/host/collective_utils.cpp @@ -81,6 +81,20 @@ bool IsSupportedReduceOp(TileXR::TileXRReduceOp reduceOp) return reduceOp == TileXR::TILEXR_REDUCE_SUM; } +bool IsSupportedReduceDataType(const TileXR::CommArgs &commArgs, TileXR::TileXRDataType dataType) +{ + if (!IsSupportedDataType(dataType)) { + return false; + } + + // Ascend950 currently reuses dav-c310 atomics, which do not support int64 reductions. + if ((commArgs.extraFlag & TileXR::ExtraFlag::TOPO_910A5) != 0 && + dataType == TileXR::TILEXR_DATA_TYPE_INT64) { + return false; + } + return true; +} + int64_t CountToBytes(int64_t count, TileXR::TileXRDataType dataType) { if (count < 0) { diff --git a/src/collectives/host/collective_utils.h b/src/collectives/host/collective_utils.h index f68c1f06..3be20857 100644 --- a/src/collectives/host/collective_utils.h +++ b/src/collectives/host/collective_utils.h @@ -22,6 +22,8 @@ bool IsSupportedDataType(TileXR::TileXRDataType dataType); bool IsSupportedReduceOp(TileXR::TileXRReduceOp reduceOp); +bool IsSupportedReduceDataType(const TileXR::CommArgs &commArgs, TileXR::TileXRDataType dataType); + int64_t CountToBytes(int64_t count, TileXR::TileXRDataType dataType); uint32_t GetAllGatherBlockNum(const TileXR::CommArgs &commArgs, int64_t dataSize); diff --git a/src/collectives/host/tilexr_collectives.cpp b/src/collectives/host/tilexr_collectives.cpp index 3da431e4..826616b7 100644 --- a/src/collectives/host/tilexr_collectives.cpp +++ b/src/collectives/host/tilexr_collectives.cpp @@ -146,6 +146,9 @@ int TileXRAllReduce(void *sendBuf, void *recvBuf, int64_t count, if (ret != TileXR::TILEXR_SUCCESS) { return ret; } + if (!TileXRCollectives::Host::IsSupportedReduceDataType(*context.hostArgs, dataType)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } const int64_t bytes = TileXRCollectives::Host::CountToBytes(count, dataType); if (context.hostArgs->rankSize <= 1) { @@ -172,6 +175,9 @@ int TileXRReduceScatter(void *sendBuf, void *recvBuf, int64_t recvCount, if (ret != TileXR::TILEXR_SUCCESS) { return ret; } + if (!TileXRCollectives::Host::IsSupportedReduceDataType(*context.hostArgs, dataType)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } const int64_t bytes = TileXRCollectives::Host::CountToBytes(recvCount, dataType); const int rankSize = context.hostArgs->rankSize; diff --git a/src/collectives/kernels/CMakeLists.txt b/src/collectives/kernels/CMakeLists.txt index 4c9e53b7..010bb3f1 100644 --- a/src/collectives/kernels/CMakeLists.txt +++ b/src/collectives/kernels/CMakeLists.txt @@ -8,6 +8,13 @@ enable_language(CCE) set(TILEXR_COLLECTIVES_1OP_BIN_SIZE 5242880 CACHE STRING "Padded size for TileXR collectives CCE binary") option(TILEXR_COLLECTIVES_ENABLE_PROFILING "Enable TileXR collectives kernel profiling helpers" OFF) +set(TILEXR_COLLECTIVES_DEFAULT_SOC_TYPE "$ENV{TILEXR_SOC_NAME}") +if(NOT TILEXR_COLLECTIVES_DEFAULT_SOC_TYPE) + set(TILEXR_COLLECTIVES_DEFAULT_SOC_TYPE "Ascend910B") +endif() +set(TILEXR_COLLECTIVES_SOC_TYPE "${TILEXR_COLLECTIVES_DEFAULT_SOC_TYPE}" CACHE STRING + "TileXR collectives kernel SOC type") +string(TOLOWER "${TILEXR_COLLECTIVES_SOC_TYPE}" TILEXR_COLLECTIVES_SOC_TYPE_LOWER) set(CCE_COMPILE_OPTION -O2 @@ -15,13 +22,18 @@ set(CCE_COMPILE_OPTION --cce-aicore-only -Wno-deprecated-declarations "SHELL:-mllvm -cce-aicore-long-call" - "SHELL:-mllvm -cce-aicore-function-stack-size=16000" + "SHELL:-mllvm -cce-aicore-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-function-stack-size=0x8000" "SHELL:-mllvm -cce-aicore-record-overflow=false" "SHELL:-mllvm -cce-aicore-addr-transform" "SHELL:-mllvm --cce-aicore-jump-expand=true" ) -set(AIV_ARCH dav-c220-vec) +if(TILEXR_COLLECTIVES_SOC_TYPE_LOWER MATCHES "ascend950|ascend910_9|ascend910-9|a5") + set(AIV_ARCH dav-c310-vec) +else() + set(AIV_ARCH dav-c220-vec) +endif() set_source_files_properties(tilexr_lccl_op.cpp PROPERTIES LANGUAGE CCE) include_directories( diff --git a/src/collectives/kernels/collectives.h b/src/collectives/kernels/collectives.h index 26f3be5e..99ea239d 100644 --- a/src/collectives/kernels/collectives.h +++ b/src/collectives/kernels/collectives.h @@ -441,8 +441,17 @@ class Collectives { { PipeBarrier(); if (op != -1) { -#ifdef __DAV_C220_VEC__ +#if defined(__DAV_C220_VEC__) SetAtomicOpType(op); +#elif defined(__DAV_C310_VEC__) + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v) { + SetAtomicOpType(op); + } #endif } PipeBarrier(); diff --git a/src/collectives/kernels/datacopy_gm2gm.h b/src/collectives/kernels/datacopy_gm2gm.h index c49e7e09..522ef46b 100644 --- a/src/collectives/kernels/datacopy_gm2gm.h +++ b/src/collectives/kernels/datacopy_gm2gm.h @@ -222,8 +222,17 @@ class DataCopyGM2GM { { PipeBarrier(); if (op != -1) { -#ifdef __DAV_C220_VEC__ +#if defined(__DAV_C220_VEC__) SetAtomicOpType(op); +#elif defined(__DAV_C310_VEC__) + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v) { + SetAtomicOpType(op); + } #endif } PipeBarrier(); @@ -329,4 +338,4 @@ class DataCopyGM2GM { const __gm__ T* outputGm = nullptr; int op; }; -#endif // LCCL_DATACOPY_GM2GM_H \ No newline at end of file +#endif // LCCL_DATACOPY_GM2GM_H diff --git a/src/collectives/kernels/kernels/collectives.cce b/src/collectives/kernels/kernels/collectives.cce index ca9570e9..b08d6cb1 100644 --- a/src/collectives/kernels/kernels/collectives.cce +++ b/src/collectives/kernels/kernels/collectives.cce @@ -10,7 +10,7 @@ #ifndef TILEXR_KERNEL_COLLECTIVES_CCE #define TILEXR_KERNEL_COLLECTIVES_CCE -#if !defined(__DAV_C220_VEC__) && !defined(__DAV_M200_VEC__) && !defined(__DAV_C220_CUBE__) +#if !defined(__DAV_C220_VEC__) && !defined(__DAV_C310_VEC__) && !defined(__DAV_M200_VEC__) && !defined(__DAV_C220_CUBE__) #define __aicore__ #define __ubuf__ #define __gm__ @@ -574,9 +574,18 @@ __attribute__((always_inline)) inline __aicore__ void ProcessData(int64_t dataSi return; } AscendC::PipeBarrier(); - #ifdef __DAV_C220_VEC__ +#ifdef __DAV_C220_VEC__ SetAtomicOpType(op); - #endif +#elif defined(__DAV_C310_VEC__) + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v) { + SetAtomicOpType(op); + } +#endif AscendC::PipeBarrier(); while (dataSizeRemain >= UB_SINGLE_ADD_SIZE_MAX) { @@ -619,6 +628,15 @@ __attribute__((always_inline)) inline __aicore__ void ProcessDataNew(int64_t dat AscendC::PipeBarrier(); #ifdef __DAV_C220_VEC__ SetAtomicOpType(op); +#elif defined(__DAV_C310_VEC__) + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v) { + SetAtomicOpType(op); + } #endif AscendC::PipeBarrier(); diff --git a/src/collectives/kernels/lccl_op.h b/src/collectives/kernels/lccl_op.h index 47a81785..7f70e4e9 100644 --- a/src/collectives/kernels/lccl_op.h +++ b/src/collectives/kernels/lccl_op.h @@ -10,7 +10,7 @@ #ifndef TILEXR_LCCL_OP_H #define TILEXR_LCCL_OP_H -#if defined(__DAV_C220_VEC__) || defined(__DAV_C220_CUBE__) +#if defined(__DAV_C220_VEC__) || defined(__DAV_C220_CUBE__) || defined(__DAV_C310_VEC__) #include "op_def.h" #include "allgather.h" diff --git a/src/collectives/kernels/tilexr_lccl_op.cpp b/src/collectives/kernels/tilexr_lccl_op.cpp index 75c22726..a96ac933 100644 --- a/src/collectives/kernels/tilexr_lccl_op.cpp +++ b/src/collectives/kernels/tilexr_lccl_op.cpp @@ -1,4 +1,4 @@ -#ifdef __DAV_C220_VEC__ +#if defined(__DAV_C220_VEC__) || defined(__DAV_C310_VEC__) #include "lccl_op.h" diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 90324f14..9a342457 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -55,6 +55,22 @@ static bool g_udmaUnavailable = false; static std::mutex g_sdmaMtx; static bool g_sdmaUnavailable = false; +bool IsEnvEnabled(const char *name, bool defaultValue) +{ + const char *value = std::getenv(name); + if (value == nullptr) { + return defaultValue; + } + const std::string str(value); + if (str == "0" || str == "false" || str == "FALSE" || str == "off" || str == "OFF") { + return false; + } + if (str == "1" || str == "true" || str == "TRUE" || str == "on" || str == "ON") { + return true; + } + return defaultValue; +} + // 如果是互联的链路，返回false； 对910B2C那些不互联的链路，返回true bool SkipUnusedChannel910B2C(int curRank, int peerRank, ChipName chipName) @@ -134,6 +150,10 @@ int TileXRComm::InitUDMA() TILEXR_LOG(INFO) << "InitUDMA skipped for single-rank communicator"; return TILEXR_SUCCESS; } + if (!IsEnvEnabled("TILEXR_ENABLE_UDMA", true)) { + TILEXR_LOG(INFO) << "TileXR UDMA disabled by environment"; + return TILEXR_SUCCESS; + } { lock_guard lock(g_udmaMtx); @@ -648,7 +668,7 @@ int TileXRComm::EnablePeerAccess() // 如果310P未来通信域要支持两卡四芯的话，这里需要做更改。并且现在默认服务器上机器只有一个链路种类。 if (value == TOPOLOGY_HCCS || value == TOPOLOGY_SIO || value == TOPOLOGY_HCCS_SW || - GetChipName() == ChipName::CHIP_910B2C) { + GetChipName() == ChipName::CHIP_910B2C || GetChipName() == ChipName::CHIP_950) { physicalInfo_.physicalLink = PhysicalLink::HCCS; commArgs_.extraFlag &= ~(ExtraFlag::TOPO_PCIE); } else if (physicalInfo_.physicalLink == PhysicalLink::RESERVED) { diff --git a/src/comm/tilexr_internal.cpp b/src/comm/tilexr_internal.cpp index 0a8494f4..bec0f571 100644 --- a/src/comm/tilexr_internal.cpp +++ b/src/comm/tilexr_internal.cpp @@ -38,6 +38,7 @@ const std::unordered_map CHIP_MAP = { {"Ascend950", ChipName::CHIP_950}, {"Ascend950DT", ChipName::CHIP_950}, {"Ascend950DT_9581", ChipName::CHIP_950}, + {"Ascend950DT_9582", ChipName::CHIP_950}, {"Ascend950DT_9584", ChipName::CHIP_950}, {"Ascend950DT_9592", ChipName::CHIP_950}, {"Ascend950PR", ChipName::CHIP_950}, diff --git a/tests/collectives/unit/test_collective_host_utils.cpp b/tests/collectives/unit/test_collective_host_utils.cpp index c5e5c3bf..a64b48c6 100644 --- a/tests/collectives/unit/test_collective_host_utils.cpp +++ b/tests/collectives/unit/test_collective_host_utils.cpp @@ -71,6 +71,27 @@ void TestReduceOpSupport() IsSupportedReduceOp(static_cast(999)), false); } +void TestReduceDataTypeSupport() +{ + using TileXRCollectives::Host::IsSupportedReduceDataType; + + const auto defaultArgs = Args(2, 0); + CheckBool("default int64 reduce supported", + IsSupportedReduceDataType(defaultArgs, TileXR::TILEXR_DATA_TYPE_INT64), + true); + + const auto ascend950Args = Args(2, TileXR::ExtraFlag::TOPO_910A5 | TileXR::ExtraFlag::TOPO_910_93); + CheckBool("Ascend950 int32 reduce supported", + IsSupportedReduceDataType(ascend950Args, TileXR::TILEXR_DATA_TYPE_INT32), + true); + CheckBool("Ascend950 fp16 reduce supported", + IsSupportedReduceDataType(ascend950Args, TileXR::TILEXR_DATA_TYPE_FP16), + true); + CheckBool("Ascend950 int64 reduce unsupported", + IsSupportedReduceDataType(ascend950Args, TileXR::TILEXR_DATA_TYPE_INT64), + false); +} + void TestCountToBytes() { using TileXRCollectives::Host::CountToBytes; @@ -251,6 +272,7 @@ int main() { TestDataTypeSupport(); TestReduceOpSupport(); + TestReduceDataTypeSupport(); TestCountToBytes(); TestAllGatherBlockNum(); TestAllToAllBlockNum(); diff --git a/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp b/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp index 170934d0..22009474 100644 --- a/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp +++ b/tests/collectives/unit/test_tilexr_collectives_kernel_ownership.cpp @@ -150,12 +150,22 @@ void TestCollectivesOwnsCceBuild() CheckContains(kernelsCmakePath, kernelsCmake, "tilexr_collectives_op"); CheckContains(kernelsCmakePath, kernelsCmake, "TILEXR_COLLECTIVES_ENABLE_PROFILING"); CheckDoesNotContain(kernelsCmakePath, kernelsCmake, "src/comm"); + CheckContains(kernelsCmakePath, kernelsCmake, "\"SHELL:-mllvm -cce-aicore-stack-size=0x8000\""); + CheckContains(kernelsCmakePath, kernelsCmake, "\"SHELL:-mllvm -cce-aicore-function-stack-size=0x8000\""); + CheckDoesNotContain(kernelsCmakePath, kernelsCmake, "-cce-aicore-function-stack-size=16000"); + CheckContains(kernelsCmakePath, kernelsCmake, "set(TILEXR_COLLECTIVES_DEFAULT_SOC_TYPE \"$ENV{TILEXR_SOC_NAME}\")"); + CheckContains(kernelsCmakePath, kernelsCmake, "set(TILEXR_COLLECTIVES_SOC_TYPE \"${TILEXR_COLLECTIVES_DEFAULT_SOC_TYPE}\""); + CheckContains(kernelsCmakePath, kernelsCmake, "string(TOLOWER \"${TILEXR_COLLECTIVES_SOC_TYPE}\" TILEXR_COLLECTIVES_SOC_TYPE_LOWER)"); + CheckContains(kernelsCmakePath, kernelsCmake, "MATCHES \"ascend950|ascend910_9|ascend910-9|a5\""); + CheckContains(kernelsCmakePath, kernelsCmake, "set(AIV_ARCH dav-c310-vec)"); + CheckContains(kernelsCmakePath, kernelsCmake, "set(AIV_ARCH dav-c220-vec)"); } void TestCollectivesKernelSourcesAreScoped() { const std::string kernelTuPath = "src/collectives/kernels/tilexr_lccl_op.cpp"; const auto kernelTu = ReadFile(kernelTuPath); + CheckContains(kernelTuPath, kernelTu, "__DAV_C310_VEC__"); CheckContains(kernelTuPath, kernelTu, "LCCL_TYPE_AIV_FUNC(LCCL_ALLGATHER_FUNC_AUTO_DEF)"); CheckContains(kernelTuPath, kernelTu, "LCCL_TYPE_AIV_FUNC(LCCL_ALL2ALL_FUNC_AUTO_DEF)"); CheckContains(kernelTuPath, kernelTu, "LCCL_TYPE_AIV_FUNC(LCCL_ALL_REDUCE_FUNC_AUTO_DEF)"); @@ -212,6 +222,47 @@ void TestCollectivesKernelSourcesAreScoped() CheckTrue(sawBroadcastCce, "expected copied broadcast .cce sources under src/collectives/kernels"); } +void TestCollectivesDavC310GateMatchesAscend950Build() +{ + const std::string lcclOpPath = "src/collectives/kernels/lccl_op.h"; + const auto lcclOp = ReadFile(lcclOpPath); + CheckContains(lcclOpPath, lcclOp, "__DAV_C310_VEC__"); + + const std::string collectivesPath = "src/collectives/kernels/collectives.h"; + const auto collectives = ReadFile(collectivesPath); + CheckContains(collectivesPath, collectives, "__DAV_C310_VEC__"); + + const std::string dataCopyPath = "src/collectives/kernels/datacopy_gm2gm.h"; + const auto dataCopy = ReadFile(dataCopyPath); + CheckContains(dataCopyPath, dataCopy, "__DAV_C310_VEC__"); + CheckContains(dataCopyPath, dataCopy, "std::is_same_v"); + CheckContains(dataCopyPath, dataCopy, "std::is_same_v"); + CheckContains(dataCopyPath, dataCopy, "std::is_same_v"); + + const std::string ccePath = "src/collectives/kernels/kernels/collectives.cce"; + const auto cce = ReadFile(ccePath); + CheckContains(ccePath, cce, "__DAV_C310_VEC__"); + CheckContains(ccePath, cce, "std::is_same_v"); + CheckContains(ccePath, cce, "std::is_same_v"); +} + +void TestAscend950ReduceTypeGuardExists() +{ + const std::string utilsHeaderPath = "src/collectives/host/collective_utils.h"; + const auto utilsHeader = ReadFile(utilsHeaderPath); + CheckContains(utilsHeaderPath, utilsHeader, "IsSupportedReduceDataType"); + + const std::string utilsPath = "src/collectives/host/collective_utils.cpp"; + const auto utils = ReadFile(utilsPath); + CheckContains(utilsPath, utils, "IsSupportedReduceDataType"); + CheckContains(utilsPath, utils, "TileXR::ExtraFlag::TOPO_910A5"); + CheckContains(utilsPath, utils, "TileXR::TILEXR_DATA_TYPE_INT64"); + + const std::string apiPath = "src/collectives/host/tilexr_collectives.cpp"; + const auto api = ReadFile(apiPath); + CheckContains(apiPath, api, "IsSupportedReduceDataType(*context.hostArgs, dataType)"); +} + void TestHostRegistrationLivesInCollectives() { const std::string kernelPath = "src/collectives/host/collective_kernel.cpp"; @@ -339,6 +390,8 @@ int main() { TestCollectivesOwnsCceBuild(); TestCollectivesKernelSourcesAreScoped(); + TestCollectivesDavC310GateMatchesAscend950Build(); + TestAscend950ReduceTypeGuardExists(); TestHostRegistrationLivesInCollectives(); TestPerfTraceCycleDivisorIsA5Specific(); TestDeviceKernelArgsMatchHostLaunchAbi(); diff --git a/tests/comm/CMakeLists.txt b/tests/comm/CMakeLists.txt index 76a7a553..31d43ebc 100644 --- a/tests/comm/CMakeLists.txt +++ b/tests/comm/CMakeLists.txt @@ -39,9 +39,18 @@ target_compile_definitions(test_tilexr_source_guards PRIVATE TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" ) +add_executable(test_tilexr_udma_env_sources + unit/test_tilexr_udma_env_sources.cpp +) + +target_compile_definitions(test_tilexr_udma_env_sources PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" +) + install(TARGETS test_tilexr_log test_tilexr_log_spdlog_compile test_tilexr_source_guards + test_tilexr_udma_env_sources RUNTIME DESTINATION ${CMAKE_INSTALL_PREFIX}/bin ) diff --git a/tests/comm/unit/test_tilexr_source_guards.cpp b/tests/comm/unit/test_tilexr_source_guards.cpp index cefd0f68..8ad23f64 100644 --- a/tests/comm/unit/test_tilexr_source_guards.cpp +++ b/tests/comm/unit/test_tilexr_source_guards.cpp @@ -123,6 +123,16 @@ void TestDumpInitCleansFailedAllocations() CheckContains(path, text, "std::free(memory);"); } +void TestAscend950UsesHccsTopologyForPeerLinks() +{ + const std::string path = "src/comm/tilexr_comm.cpp"; + const auto text = ReadFile(path); + + CheckContains(path, text, "GetChipName() == ChipName::CHIP_950"); + CheckContains(path, text, "physicalInfo_.physicalLink = PhysicalLink::HCCS;"); + CheckContains(path, text, "commArgs_.extraFlag &= ~(ExtraFlag::TOPO_PCIE);"); +} + void TestSocketExchangeUsesDirectConnectionsOnly() { const std::string cppPath = "src/comm/tools/socket/tilexr_sock_exchange.cpp"; @@ -187,6 +197,7 @@ int main() TestCommInitChecksDeviceCommArgsSync(); TestCWrappersDoNotPublishFailedCommunicators(); TestDumpInitCleansFailedAllocations(); + TestAscend950UsesHccsTopologyForPeerLinks(); TestSocketExchangeUsesDirectConnectionsOnly(); TestRuntimeEnvDoesNotPrependCannDevlib(); TestRootCMakeRespectsAscendDriverOverride(); diff --git a/tests/comm/unit/test_tilexr_udma_env_sources.cpp b/tests/comm/unit/test_tilexr_udma_env_sources.cpp new file mode 100644 index 00000000..27bd8209 --- /dev/null +++ b/tests/comm/unit/test_tilexr_udma_env_sources.cpp @@ -0,0 +1,44 @@ +#include +#include +#include +#include + +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + +namespace { + +int g_failures = 0; + +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +} // namespace + +int main() +{ + const std::string comm = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/tilexr_comm.cpp"); + CHECK_CONTAINS(comm, "TILEXR_ENABLE_UDMA"); + CHECK_CONTAINS(comm, "TileXR UDMA disabled by environment"); + CHECK_CONTAINS(comm, "IsEnvEnabled(\"TILEXR_ENABLE_UDMA\", true)"); + if (g_failures != 0) { + std::cerr << g_failures << " UDMA env source checks failed" << std::endl; + return 1; + } + std::cout << "TileXR UDMA env source checks passed" << std::endl; + return 0; +} diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index da4e89dd..d91ddd85 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -7,6 +7,7 @@ #define TILEXR_UDMA_ALLTOALL_LAYOUT_H #include +#include #include #include #include @@ -15,12 +16,50 @@ namespace TileXR { namespace Demo { constexpr int32_t kAllToAllBaseValue = 100000; +constexpr size_t kAllToAllUdmaMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; + +struct AllToAllChunkPlan { + uint32_t passCount = 1; + int32_t chunkElements = 0; + size_t chunkBytesPerRank = 0; + size_t registeredBytes = 0; +}; inline int32_t AllToAllValue(int srcRank, int dstRank) { return kAllToAllBaseValue + srcRank * 1000 + dstRank; } +inline AllToAllChunkPlan PlanAllToAllUdmaChunks(int rankSize, int32_t elementsPerPeer) +{ + AllToAllChunkPlan plan {}; + if (rankSize <= 0 || elementsPerPeer <= 0) { + return plan; + } + + const size_t totalBytesPerRank = static_cast(rankSize) * static_cast(elementsPerPeer) * + sizeof(int32_t); + const size_t maxChunkBytesPerRank = + std::max(sizeof(int32_t) * static_cast(rankSize), kAllToAllUdmaMaxRegisteredBytes / 2); + size_t chunkElements = maxChunkBytesPerRank / (static_cast(rankSize) * sizeof(int32_t)); + if (chunkElements == 0) { + chunkElements = 1; + } + if (chunkElements > static_cast(elementsPerPeer)) { + chunkElements = static_cast(elementsPerPeer); + } + + plan.chunkElements = static_cast(chunkElements); + plan.chunkBytesPerRank = static_cast(rankSize) * chunkElements * sizeof(int32_t); + plan.registeredBytes = plan.chunkBytesPerRank * 2; + plan.passCount = static_cast( + (static_cast(elementsPerPeer) + chunkElements - 1) / chunkElements); + if (plan.chunkBytesPerRank > totalBytesPerRank) { + plan.chunkBytesPerRank = totalBytesPerRank; + } + return plan; +} + inline void FillAllToAllInput( std::vector& input, int rank, int rankSize, int32_t elementsPerPeer) { diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 630a3c86..771e06d0 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -32,11 +32,23 @@ extern void launch_tilexr_udma_put_signal( int32_t elementsPerRank, uint64_t signal); extern void launch_tilexr_udma_all_to_all( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, - GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset); + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, + int32_t chunkElements); extern void launch_tilexr_all_to_all_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); extern void launch_tilexr_all_to_all_ipc_gather( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer); +extern void launch_tilexr_all_to_all_plain_ipc_scatter( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); +extern void launch_tilexr_all_to_all_plain_ipc_gather( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer); +extern void launch_tilexr_all_to_all_fused_ipc( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, int32_t round); +extern void launch_tilexr_all_to_all_ipc_scatter_dma( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); +extern void launch_tilexr_all_to_all_ipc_gather_dma( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer); extern void launch_tilexr_all_reduce_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerRank); extern void launch_tilexr_all_reduce_ipc_sum( @@ -469,6 +481,54 @@ size_t AllToAllDataAsFlagStagingBytes(int rankSize, int32_t elementsPerPeer) return static_cast(static_cast(rankSize) * blocks * TileXR::DATA_AS_FLAG_BLOCK_BYTES); } +size_t AllToAllPlainIpcStagingBytes(int rankSize, int32_t elementsPerPeer) +{ + return static_cast(rankSize) * static_cast(rankSize) * + static_cast(elementsPerPeer) * sizeof(int32_t); +} + +bool CopyChunkHostToDevice( + int rank, int32_t* chunkInput, int32_t* chunkOutput, int rankSize, int32_t elementsPerRank, + int32_t chunkOffset, int32_t chunkElements, const std::vector& hostInput, + std::vector& hostOutput) +{ + std::vector inputChunk(static_cast(rankSize) * chunkElements, 0); + std::vector outputChunk(static_cast(rankSize) * chunkElements, -1); + for (int peer = 0; peer < rankSize; ++peer) { + const size_t srcBase = static_cast(peer) * elementsPerRank + chunkOffset; + const size_t dstBase = static_cast(peer) * chunkElements; + std::copy(hostInput.begin() + srcBase, + hostInput.begin() + srcBase + chunkElements, + inputChunk.begin() + dstBase); + std::copy(hostOutput.begin() + srcBase, + hostOutput.begin() + srcBase + chunkElements, + outputChunk.begin() + dstBase); + } + return CopyHostToDevice(rank, chunkInput, inputChunk.size() * sizeof(int32_t), + inputChunk.data(), inputChunk.size() * sizeof(int32_t), "registered input chunk") && + CopyHostToDevice(rank, chunkOutput, outputChunk.size() * sizeof(int32_t), + outputChunk.data(), outputChunk.size() * sizeof(int32_t), "registered output chunk"); +} + +bool CopyChunkDeviceToHost( + int rank, const int32_t* chunkOutput, int rankSize, int32_t elementsPerRank, + int32_t chunkOffset, int32_t chunkElements, std::vector& hostOutput) +{ + std::vector outputChunk(static_cast(rankSize) * chunkElements, -1); + if (!CopyDeviceToHost(rank, outputChunk.data(), outputChunk.size() * sizeof(int32_t), + chunkOutput, outputChunk.size() * sizeof(int32_t), "registered output chunk")) { + return false; + } + for (int peer = 0; peer < rankSize; ++peer) { + const size_t dstBase = static_cast(peer) * elementsPerRank + chunkOffset; + const size_t srcBase = static_cast(peer) * chunkElements; + std::copy(outputChunk.begin() + srcBase, + outputChunk.begin() + srcBase + chunkElements, + hostOutput.begin() + dstBase); + } + return true; +} + void Cleanup( TileXRCommPtr comm, aclrtStream stream, void* registeredMemory, int32_t* debug, int rank, int deviceId) { @@ -559,34 +619,69 @@ int main(int argc, char** argv) bool isAllToAll = testType == 2; bool isAllReduce = testType == 3; bool strictAllToAllUdma = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_USE_UDMA", 0) != 0; + int allToAllRepeat = isAllToAll ? std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)) : 1; + bool syncAllToAllAtEnd = + isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_SYNC_AT_END", 0) != 0; + bool useAllToAllPlainIpc = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_PLAIN_IPC", 0) != 0; + bool useAllToAllFusedIpc = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_FUSED_IPC", 0) != 0; bool dumpAllToAllOnStrictFail = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_DUMP_ON_STRICT_FAIL", 0) != 0; - bool useAllToAllDataAsFlagIpc = isAllToAll && !strictAllToAllUdma; + bool useAllToAllDataAsFlagIpc = isAllToAll && !strictAllToAllUdma && !useAllToAllPlainIpc && !useAllToAllFusedIpc; + const char* allToAllIpcFallbackLabel = + useAllToAllFusedIpc ? "fused IPC" : + (useAllToAllPlainIpc ? "plain IPC fallback" : "data-as-flag IPC fallback"); bool forceAllToAllIpcFallback = false; bool hasOutput = isAllToAll || isAllReduce; size_t dataCount = static_cast(rankSize) * elementsPerRank; size_t dataBytes = dataCount * sizeof(int32_t); + const TileXR::Demo::AllToAllChunkPlan chunkPlan = + isAllToAll ? TileXR::Demo::PlanAllToAllUdmaChunks(rankSize, elementsPerRank) : + TileXR::Demo::AllToAllChunkPlan {}; if (isAllToAll) { - const size_t stagingBytes = AllToAllDataAsFlagStagingBytes(rankSize, elementsPerRank); - if (stagingBytes > static_cast(TileXR::IPC_BUFF_MAX_SIZE)) { + const size_t dataAsFlagStagingBytes = AllToAllDataAsFlagStagingBytes(rankSize, elementsPerRank); + const size_t plainIpcStagingBytes = AllToAllPlainIpcStagingBytes(rankSize, elementsPerRank); + const size_t selectedIpcStagingBytes = useAllToAllPlainIpc ? plainIpcStagingBytes : dataAsFlagStagingBytes; + if (useAllToAllPlainIpc && plainIpcStagingBytes > static_cast(TileXR::IPC_BUFF_MAX_SIZE)) { + std::cerr << "[rank " << rank << "] ERROR: alltoall plain IPC fallback staging requires " + << plainIpcStagingBytes << " bytes, exceeds IPC data capacity " + << TileXR::IPC_BUFF_MAX_SIZE << std::endl; + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (useAllToAllDataAsFlagIpc && dataAsFlagStagingBytes > static_cast(TileXR::IPC_BUFF_MAX_SIZE)) { std::cerr << "[rank " << rank << "] ERROR: alltoall data-as-flag IPC fallback staging requires " - << stagingBytes << " bytes, exceeds IPC data capacity " + << dataAsFlagStagingBytes << " bytes, exceeds IPC data capacity " << TileXR::IPC_BUFF_MAX_SIZE << std::endl; Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } - PrintStatus(rank, "alltoall data-as-flag staging bytes=" + std::to_string(stagingBytes)); + PrintStatus(rank, "alltoall plain IPC staging bytes=" + std::to_string(plainIpcStagingBytes)); + PrintStatus(rank, "alltoall data-as-flag staging bytes=" + std::to_string(dataAsFlagStagingBytes)); + if (strictAllToAllUdma && selectedIpcStagingBytes > static_cast(TileXR::IPC_BUFF_MAX_SIZE)) { + PrintStatus(rank, "skip IPC staging capacity guard for strict UDMA alltoall"); + } + PrintStatus(rank, "alltoall repeat=" + std::to_string(allToAllRepeat) + + " syncAtEnd=" + std::string(syncAllToAllAtEnd ? "true" : "false")); + PrintStatus(rank, "alltoall UDMA chunk plan: passCount=" + std::to_string(chunkPlan.passCount) + + " chunkElements=" + std::to_string(chunkPlan.chunkElements) + + " registeredBytes=" + std::to_string(chunkPlan.registeredBytes)); } size_t inputOffset = 0; - size_t outputOffset = hasOutput ? dataBytes : 0; + const size_t activeBytesPerRank = isAllToAll && strictAllToAllUdma ? chunkPlan.chunkBytesPerRank : dataBytes; + size_t outputOffset = hasOutput ? activeBytesPerRank : 0; size_t signalBytes = static_cast(rankSize) * sizeof(uint64_t); - size_t signalOffset = hasOutput ? dataBytes * 2 : dataBytes; + size_t signalOffset = hasOutput ? (outputOffset + activeBytesPerRank) : activeBytesPerRank; size_t payloadBytes = signalOffset + signalBytes; - size_t registeredBytes = ((payloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * + size_t allocBytes = ((payloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + size_t registeredBytes = allocBytes; + if (isAllToAll && strictAllToAllUdma) { + registeredBytes = ((chunkPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / + kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + } if (!CheckAcl(rank, "aclrtMalloc debug", aclrtMalloc(reinterpret_cast(&debug), kDebugWords * sizeof(int32_t), ACL_MEM_MALLOC_HUGE_FIRST)) || !CheckAcl(rank, "aclrtMalloc registered memory", aclrtMalloc(®isteredMemory, - registeredBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + allocBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } @@ -594,9 +689,18 @@ int main(int argc, char** argv) auto input = reinterpret_cast(static_cast(registeredMemory) + inputOffset); auto output = reinterpret_cast(static_cast(registeredMemory) + outputOffset); auto signals = reinterpret_cast(static_cast(registeredMemory) + signalOffset); - if (useAllToAllDataAsFlagIpc) { + const bool chunkedStrictAllToAll = isAllToAll && strictAllToAllUdma && chunkPlan.passCount > 1; + if (useAllToAllFusedIpc) { + PrintStatus(rank, "skip TileXRUDMARegister for alltoall fused IPC path"); + forceAllToAllIpcFallback = true; + } else if (useAllToAllPlainIpc) { + PrintStatus(rank, "skip TileXRUDMARegister for alltoall plain IPC path"); + forceAllToAllIpcFallback = true; + } else if (useAllToAllDataAsFlagIpc) { PrintStatus(rank, "skip TileXRUDMARegister for alltoall data-as-flag IPC path"); forceAllToAllIpcFallback = true; + } else if (chunkedStrictAllToAll) { + PrintStatus(rank, "defer TileXRUDMARegister to per-pass registered output chunk"); } else { int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), registeredBytes, &udmaHandle); @@ -640,12 +744,17 @@ int main(int argc, char** argv) std::vector hostDebug(kDebugWords, 0); const char* inputName = isAllToAll ? "alltoall input" : (isAllReduce ? "allreduce input" : "data"); - bool initOk = CopyHostToDevice(rank, input, dataCount * sizeof(int32_t), - hostData.data(), dataCount * sizeof(int32_t), inputName); - if (hasOutput) { - const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; - initOk = CopyHostToDevice(rank, output, dataCount * sizeof(int32_t), - hostOutput.data(), dataCount * sizeof(int32_t), outputName) && initOk; + bool initOk = true; + if (!chunkedStrictAllToAll) { + initOk = CopyHostToDevice(rank, input, dataCount * sizeof(int32_t), + hostData.data(), dataCount * sizeof(int32_t), inputName); + if (hasOutput) { + const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; + initOk = CopyHostToDevice(rank, output, dataCount * sizeof(int32_t), + hostOutput.data(), dataCount * sizeof(int32_t), outputName) && initOk; + } + } else { + std::fill(hostOutput.begin(), hostOutput.end(), -1); } if (!initOk || !CopyHostToDevice(rank, signals, hostSignals.size() * sizeof(uint64_t), @@ -668,12 +777,89 @@ int main(int argc, char** argv) if (testType == 2) { if (forceAllToAllIpcFallback) { - PrintStatus(rank, "skip all-to-all UDMA kernel; use data-as-flag IPC fallback"); + PrintStatus(rank, std::string("skip all-to-all UDMA kernel; use ") + allToAllIpcFallbackLabel); + } else if (chunkedStrictAllToAll) { + for (uint32_t pass = 0; pass < chunkPlan.passCount; ++pass) { + const int32_t chunkOffset = static_cast(pass) * chunkPlan.chunkElements; + const int32_t chunkElements = std::min( + chunkPlan.chunkElements, elementsPerRank - chunkOffset); + if (!CopyChunkHostToDevice(rank, input, output, rankSize, elementsPerRank, + chunkOffset, chunkElements, hostData, hostOutput)) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!udmaRegistered) { + int registerRet = + TileXRUDMARegister(comm, static_cast(registeredMemory), registeredBytes, &udmaHandle); + if (registerRet != TileXR::TILEXR_SUCCESS) { + std::cerr << "[rank " << rank << "] ERROR: strict alltoall UDMA registration failed" + << " ret=" << registerRet << std::endl; + CheckTileXR(rank, "TileXRUDMARegister", registerRet); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + udmaRegistered = true; + } + PrintStatus(rank, "launch all-to-all kernel pass=" + std::to_string(pass)); + launch_tilexr_udma_all_to_all( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), chunkElements, static_cast(outputOffset), + 0, chunkElements); + if (!CheckAcl(rank, "aclrtSynchronizeStream", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed demo kernels")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), + debug, hostDebug.size() * sizeof(int32_t), "debug after alltoall udma")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!AllToAllUdmaComplete(rankSize, hostDebug)) { + std::cerr << "[rank " << rank << "] ERROR: strict alltoall UDMA CQ incomplete:" << std::endl; + PrintAllToAllUdmaDebug(rank, rankSize, hostDebug); + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!CopyChunkDeviceToHost(rank, output, rankSize, elementsPerRank, + chunkOffset, chunkElements, hostOutput)) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } } else { - PrintStatus(rank, "launch all-to-all kernel"); - launch_tilexr_udma_all_to_all( - 1, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset)); + for (int iter = 0; iter < allToAllRepeat; ++iter) { + PrintStatus(rank, "launch all-to-all kernel iter=" + std::to_string(iter)); + launch_tilexr_udma_all_to_all( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset), 0, + elementsPerRank); + if (!syncAllToAllAtEnd && + !CheckAcl(rank, "aclrtSynchronizeStream", aclrtSynchronizeStream(stream))) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } } } else if (testType == 3) { PrintStatus(rank, "launch all-reduce IPC scatter kernel"); @@ -755,10 +941,14 @@ int main(int argc, char** argv) } if (isAllToAll && (forceAllToAllIpcFallback || !allToAllUdmaComplete)) { usedIpcFallback = true; - std::cout << "[rank " << rank << "] alltoall use data-as-flag IPC fallback"; + std::cout << "[rank " << rank << "] alltoall use " << allToAllIpcFallbackLabel; if (forceAllToAllIpcFallback) { if (useAllToAllDataAsFlagIpc) { std::cout << " by default"; + } else if (useAllToAllPlainIpc) { + std::cout << " by request"; + } else if (useAllToAllFusedIpc) { + std::cout << " by request"; } else { std::cout << " after UDMA registration failure"; } @@ -771,34 +961,126 @@ int main(int argc, char** argv) } std::cout << std::endl; - launch_tilexr_all_to_all_ipc_scatter( - 1, stream, commArgsDev, reinterpret_cast(input), - reinterpret_cast(debug), elementsPerRank); - if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc scatter", aclrtSynchronizeStream(stream)) || - !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc scatter")) { - if (udmaRegistered) { - CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + if (useAllToAllFusedIpc) { + PrintStatus(rank, "alltoall fused IPC: single kernel send+flag+recv"); + if (syncAllToAllAtEnd) { + for (int iter = 0; iter < allToAllRepeat; ++iter) { + launch_tilexr_all_to_all_fused_ipc( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(output), reinterpret_cast(debug), + elementsPerRank, iter + 1); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall fused ipc", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall fused ipc")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } else { + for (int iter = 0; iter < allToAllRepeat; ++iter) { + launch_tilexr_all_to_all_fused_ipc( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(output), reinterpret_cast(debug), + elementsPerRank, iter + 1); + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall fused ipc", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall fused ipc")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } + } + } else if (syncAllToAllAtEnd) { + for (int iter = 0; iter < allToAllRepeat; ++iter) { + if (useAllToAllPlainIpc) { + launch_tilexr_all_to_all_plain_ipc_scatter( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(debug), elementsPerRank); + } else { + launch_tilexr_all_to_all_ipc_scatter( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(debug), elementsPerRank); + } + } + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc scatter", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc scatter")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; } - Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); - return 1; - } - launch_tilexr_all_to_all_ipc_gather( - 1, stream, commArgsDev, reinterpret_cast(output), - reinterpret_cast(debug), elementsPerRank); - if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc gather", aclrtSynchronizeStream(stream)) || - !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc gather")) { - if (udmaRegistered) { - CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + for (int iter = 0; iter < allToAllRepeat; ++iter) { + if (useAllToAllPlainIpc) { + launch_tilexr_all_to_all_plain_ipc_gather( + 1, stream, commArgsDev, reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank); + } else { + launch_tilexr_all_to_all_ipc_gather( + 1, stream, commArgsDev, reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank); + } + } + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc gather", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc gather")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } else { + for (int iter = 0; iter < allToAllRepeat; ++iter) { + if (useAllToAllPlainIpc) { + launch_tilexr_all_to_all_plain_ipc_scatter( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(debug), elementsPerRank); + } else { + launch_tilexr_all_to_all_ipc_scatter( + 1, stream, commArgsDev, reinterpret_cast(input), + reinterpret_cast(debug), elementsPerRank); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc scatter", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc scatter")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + + if (useAllToAllPlainIpc) { + launch_tilexr_all_to_all_plain_ipc_gather( + 1, stream, commArgsDev, reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank); + } else { + launch_tilexr_all_to_all_ipc_gather( + 1, stream, commArgsDev, reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream alltoall ipc gather", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed alltoall ipc gather")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } } - Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); - return 1; } } - bool copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), - data, dataCount * sizeof(int32_t), "data"); - if (hasOutput) { + bool copyBackOk = true; + if (!chunkedStrictAllToAll) { + copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), + data, dataCount * sizeof(int32_t), "data"); + } + if (hasOutput && !chunkedStrictAllToAll) { const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; copyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), output, dataCount * sizeof(int32_t), outputName) && copyBackOk; diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 988ef4b5..81e287f1 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -121,7 +121,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_put_signal_kernel( extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, - int32_t elementsPerPeer, uint64_t outputByteOffset) + int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, int32_t chunkElements) { auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); @@ -132,7 +132,15 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( int32_t rankSize = args->rankSize; bool enabled = TileXR::UDMARegistryEnabled(args); - if (debug != nullptr) { + // Multi-core: one block per peer. Block b handles peer b: + // - peer == rank -> local self-copy via DataCopyPad + // - peer != rank -> UDMA PUT to that peer + // Block 0 writes the shared debug header; per-peer slots are written by + // their owning block only, so no cross-block debug races. + const int32_t blockIdx = AscendC::GetBlockIdx(); + const int32_t blockNum = AscendC::GetBlockNum(); + + if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; debug[1] = rank; debug[2] = rankSize; @@ -144,19 +152,52 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( return; } - const uint64_t payloadBytes = AllToAllPayloadBytes(elementsPerPeer); - auto selfSrc = input + static_cast(rank) * elementsPerPeer; - auto selfDst = output + static_cast(rank) * elementsPerPeer; - for (int32_t i = 0; i < elementsPerPeer; ++i) { - selfDst[i] = selfSrc[i]; - } + const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; + const uint64_t payloadBytes = AllToAllPayloadBytes(effectiveChunkElements); + const uint32_t bytes = static_cast(payloadBytes); - uint32_t bytes = static_cast(payloadBytes); - for (int32_t peer = 0; peer < rankSize; ++peer) { + // This block's assigned peer. When host launches rankSize blocks, block b + // handles peer b. If fewer blocks are launched, peers are round-robined + // and each block may handle more than one peer (still correct, just less + // parallel); stride == blockNum keeps peer slots disjoint across blocks. + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { if (peer == rank) { + // Self-copy: local DataCopyPad, no network. + auto selfSrc = input + static_cast(rank) * elementsPerPeer + inputElementOffset; + auto selfDst = output + static_cast(rank) * effectiveChunkElements; + constexpr uint32_t SELF_COPY_UB_BYTES = 64 * 1024; + AscendC::TPipe pipe; + AscendC::TBuf selfCopyTBuf; + pipe.InitBuffer(selfCopyTBuf, SELF_COPY_UB_BYTES); + AscendC::LocalTensor selfCopyLocal = selfCopyTBuf.Get(); + + auto selfSrcBytes = reinterpret_cast<__gm__ uint8_t*>(selfSrc); + auto selfDstBytes = reinterpret_cast<__gm__ uint8_t*>(selfDst); + for (uint32_t offset = 0; offset < bytes; offset += SELF_COPY_UB_BYTES) { + uint32_t copyBytes = (bytes - offset < SELF_COPY_UB_BYTES) + ? (bytes - offset) : SELF_COPY_UB_BYTES; + + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(selfSrcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, copyBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(selfCopyLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(selfDstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, copyBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, selfCopyLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); continue; } - auto localSrc = input + static_cast(peer) * elementsPerPeer; + + // UDMA PUT to peer. + auto localSrc = input + static_cast(peer) * elementsPerPeer + inputElementOffset; uint64_t remoteOffset = outputByteOffset + static_cast(rank) * payloadBytes; auto registry = TileXR::GetUDMARegistry(args); @@ -256,6 +297,223 @@ extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_gather_kernel( } } +extern "C" __global__ __aicore__ void tilexr_all_to_all_plain_ipc_scatter_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR debugGM, int32_t elementsPerPeer) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + auto remoteBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[dstRank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + auto remoteDst = remoteBase + static_cast(rank) * rankSize * elementsPerPeer; + auto localSrc = input + static_cast(dstRank) * elementsPerPeer; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + remoteDst[i] = localSrc[i]; + } + } + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER] = 2; + } +} + +extern "C" __global__ __aicore__ void tilexr_all_to_all_plain_ipc_gather_kernel( + GM_ADDR commArgsGM, GM_ADDR outputGM, GM_ADDR debugGM, int32_t elementsPerPeer) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rankSize = args->rankSize; + auto localBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[args->rank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + for (int32_t srcRank = 0; srcRank < rankSize; ++srcRank) { + auto localSrc = localBase + static_cast(srcRank) * rankSize * elementsPerPeer; + auto localDst = output + static_cast(srcRank) * elementsPerPeer; + for (int32_t i = 0; i < elementsPerPeer; ++i) { + localDst[i] = localSrc[i]; + } + } + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER] = 2; + } +} + +// Fused single-kernel alltoall: send + flag-sync + receive in one kernel launch. +// Each rank writes data to peers, sets a per-peer flag, polls peer flags, then reads. +// Flag layout: peerMems[peer] + srcRank * sizeof(uint64_t), within the 2MB flag region. +// Round parameter enables multi-round pipelining without host barriers. +extern "C" __global__ __aicore__ void tilexr_all_to_all_fused_ipc_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, + int32_t elementsPerPeer, int32_t round) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + const int64_t flagValue = static_cast(round); + + // UB: 64KB for DMA data relay + 64B for flag polling + constexpr uint32_t DATA_UB_BYTES = 64 * 1024; + constexpr uint32_t CHUNK_BYTES = DATA_UB_BYTES; + + AscendC::TPipe pipe; + AscendC::TBuf dataTBuf; + pipe.InitBuffer(dataTBuf, DATA_UB_BYTES); + AscendC::LocalTensor dataLocal = dataTBuf.Get(); + + // Flag layout in peer's flag region (before IPC_DATA_OFFSET): + // [0 .. rankSize-1] = ready flags (data written) + // [rankSize .. 2*rankSize-1] = consumed ACK flags (data read) + // Each flag is int64_t, indexed by srcRank. + + // Phase 1: Write data to all peers' IPC staging area via DMA (DataCopyPad) + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + auto remoteBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[dstRank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + auto remoteDst = remoteBase + static_cast(rank) * rankSize * elementsPerPeer; + auto localSrc = input + static_cast(dstRank) * elementsPerPeer; + + auto srcBytes = reinterpret_cast<__gm__ uint8_t*>(localSrc); + auto dstBytes = reinterpret_cast<__gm__ uint8_t*>(remoteDst); + uint32_t payloadBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + + for (uint32_t offset = 0; offset < payloadBytes; offset += CHUNK_BYTES) { + uint32_t bytes = (payloadBytes - offset < CHUNK_BYTES) ? (payloadBytes - offset) : CHUNK_BYTES; + uint32_t alignedBytes = bytes & ~31U; + + if (alignedBytes > 0) { + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(srcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dataLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, dataLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + for (uint32_t i = alignedBytes; i < bytes; ++i) { + dstBytes[offset + i] = srcBytes[offset + i]; + } + } + } + + // Phase 2: Send "data ready" flag to each peer via scalar store. + // Scalar store to P2P address after MTE3_S barrier ensures all DMA + // data writes are submitted before the flag. + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + if (dstRank == rank) { + continue; + } + auto flagAddr = reinterpret_cast<__gm__ int64_t*>( + args->peerMems[dstRank] + static_cast(rank) * sizeof(int64_t)); + *flagAddr = flagValue; + } + + // Phase 3: Poll all peers' "data ready" flags (read from own flag region). + // Wait until each peer has written their data to our staging area. + for (int32_t srcRank = 0; srcRank < rankSize; ++srcRank) { + if (srcRank == rank) { + continue; + } + auto flagAddr = reinterpret_cast<__gm__ int64_t*>( + args->peerMems[rank] + static_cast(srcRank) * sizeof(int64_t)); + int64_t observed = 0; + do { + observed = *flagAddr; + } while (observed < flagValue); + } + + // Phase 4: Read data from local IPC staging area into output via DMA. + // By now all peers' data has arrived (they set their ready flags). + { + auto localBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[rank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + for (int32_t srcRank = 0; srcRank < rankSize; ++srcRank) { + auto localSrc = localBase + static_cast(srcRank) * rankSize * elementsPerPeer; + auto localDst = output + static_cast(srcRank) * elementsPerPeer; + + auto srcBytes = reinterpret_cast<__gm__ uint8_t*>(localSrc); + auto dstBytes = reinterpret_cast<__gm__ uint8_t*>(localDst); + uint32_t payloadBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + + for (uint32_t offset = 0; offset < payloadBytes; offset += CHUNK_BYTES) { + uint32_t bytes = (payloadBytes - offset < CHUNK_BYTES) ? (payloadBytes - offset) : CHUNK_BYTES; + uint32_t alignedBytes = bytes & ~31U; + + if (alignedBytes > 0) { + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(srcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dataLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, dataLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + for (uint32_t i = alignedBytes; i < bytes; ++i) { + dstBytes[offset + i] = srcBytes[offset + i]; + } + } + } + } + + // Phase 5: Send "data consumed" ACK to each peer. + // Tells peers that we have read their data; they can safely overwrite + // the staging area in the next round. + // ACK flag offset: rankSize * sizeof(int64_t) + rank * sizeof(int64_t) + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + if (dstRank == rank) { + continue; + } + auto ackAddr = reinterpret_cast<__gm__ int64_t*>( + args->peerMems[dstRank] + + static_cast(rankSize + rank) * sizeof(int64_t)); + *ackAddr = flagValue; + } + + // Phase 6: Poll all peers' "data consumed" ACKs. + // Wait until all peers confirm they have read our data. + // Only then can we proceed to the next round (which will overwrite). + for (int32_t srcRank = 0; srcRank < rankSize; ++srcRank) { + if (srcRank == rank) { + continue; + } + auto ackAddr = reinterpret_cast<__gm__ int64_t*>( + args->peerMems[rank] + + static_cast(rankSize + srcRank) * sizeof(int64_t)); + int64_t observed = 0; + do { + observed = *ackAddr; + } while (observed < flagValue); + } + + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER] = 3; + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER] = 3; + } +} + extern "C" __global__ __aicore__ void tilexr_all_reduce_ipc_scatter_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR debugGM, int32_t elementsPerRank) { @@ -329,6 +587,144 @@ extern "C" __global__ __aicore__ void tilexr_udma_registered_smoke_kernel( } } +// DMA-based scatter: write data to all peers' IPC staging area via DataCopyPad. +// Split from the fused kernel to allow host-side sync between scatter and gather, +// which guarantees P2P write visibility without fragile in-kernel flag polling. +extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_scatter_dma_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR debugGM, int32_t elementsPerPeer) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + + constexpr uint32_t DATA_UB_BYTES = 64 * 1024; + constexpr uint32_t CHUNK_BYTES = DATA_UB_BYTES; + + AscendC::TPipe pipe; + AscendC::TBuf dataTBuf; + pipe.InitBuffer(dataTBuf, DATA_UB_BYTES); + AscendC::LocalTensor dataLocal = dataTBuf.Get(); + + // Write data to all peers' IPC staging area via DMA + for (int32_t dstRank = 0; dstRank < rankSize; ++dstRank) { + auto remoteBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[dstRank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + auto remoteDst = remoteBase + static_cast(rank) * rankSize * elementsPerPeer; + auto localSrc = input + static_cast(dstRank) * elementsPerPeer; + + auto srcBytes = reinterpret_cast<__gm__ uint8_t*>(localSrc); + auto dstBytes = reinterpret_cast<__gm__ uint8_t*>(remoteDst); + uint32_t payloadBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + + for (uint32_t offset = 0; offset < payloadBytes; offset += CHUNK_BYTES) { + uint32_t bytes = (payloadBytes - offset < CHUNK_BYTES) ? (payloadBytes - offset) : CHUNK_BYTES; + uint32_t alignedBytes = bytes & ~31U; + + if (alignedBytes > 0) { + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(srcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dataLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, dataLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + for (uint32_t i = alignedBytes; i < bytes; ++i) { + dstBytes[offset + i] = srcBytes[offset + i]; + } + } + } + + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER] = 4; + } +} + +// DMA-based gather: read data from local IPC staging area into output via DataCopyPad. +extern "C" __global__ __aicore__ void tilexr_all_to_all_ipc_gather_dma_kernel( + GM_ADDR commArgsGM, GM_ADDR outputGM, GM_ADDR debugGM, int32_t elementsPerPeer) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + + constexpr uint32_t DATA_UB_BYTES = 64 * 1024; + constexpr uint32_t CHUNK_BYTES = DATA_UB_BYTES; + + AscendC::TPipe pipe; + AscendC::TBuf dataTBuf; + pipe.InitBuffer(dataTBuf, DATA_UB_BYTES); + AscendC::LocalTensor dataLocal = dataTBuf.Get(); + + auto localBase = reinterpret_cast<__gm__ int32_t*>( + args->peerMems[rank] + TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET); + + for (int32_t srcRank = 0; srcRank < rankSize; ++srcRank) { + auto localSrc = localBase + static_cast(srcRank) * rankSize * elementsPerPeer; + auto localDst = output + static_cast(srcRank) * elementsPerPeer; + + auto srcBytes = reinterpret_cast<__gm__ uint8_t*>(localSrc); + auto dstBytes = reinterpret_cast<__gm__ uint8_t*>(localDst); + uint32_t payloadBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + + for (uint32_t offset = 0; offset < payloadBytes; offset += CHUNK_BYTES) { + uint32_t bytes = (payloadBytes - offset < CHUNK_BYTES) ? (payloadBytes - offset) : CHUNK_BYTES; + uint32_t alignedBytes = bytes & ~31U; + + if (alignedBytes > 0) { + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(srcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dataLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, alignedBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, dataLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + for (uint32_t i = alignedBytes; i < bytes; ++i) { + dstBytes[offset + i] = srcBytes[offset + i]; + } + } + } + + if (debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER] = 4; + } +} + +void launch_tilexr_all_to_all_ipc_scatter_dma( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer) +{ + tilexr_all_to_all_ipc_scatter_dma_kernel<<>>( + commArgs, input, debug, elementsPerPeer); +} + +void launch_tilexr_all_to_all_ipc_gather_dma( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer) +{ + tilexr_all_to_all_ipc_gather_dma_kernel<<>>( + commArgs, output, debug, elementsPerPeer); +} + void launch_tilexr_udma_all_gather( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR data, GM_ADDR debug, int32_t elementsPerRank) { @@ -346,10 +742,11 @@ void launch_tilexr_udma_put_signal( void launch_tilexr_udma_all_to_all( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, - GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset) + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, + int32_t chunkElements) { tilexr_udma_all_to_all_kernel<<>>( - commArgs, input, output, debug, elementsPerPeer, outputByteOffset); + commArgs, input, output, debug, elementsPerPeer, outputByteOffset, inputElementOffset, chunkElements); } void launch_tilexr_all_to_all_ipc_scatter( @@ -366,6 +763,28 @@ void launch_tilexr_all_to_all_ipc_gather( commArgs, output, debug, elementsPerPeer); } +void launch_tilexr_all_to_all_plain_ipc_scatter( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer) +{ + tilexr_all_to_all_plain_ipc_scatter_kernel<<>>( + commArgs, input, debug, elementsPerPeer); +} + +void launch_tilexr_all_to_all_plain_ipc_gather( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer) +{ + tilexr_all_to_all_plain_ipc_gather_kernel<<>>( + commArgs, output, debug, elementsPerPeer); +} + +void launch_tilexr_all_to_all_fused_ipc( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, int32_t round) +{ + tilexr_all_to_all_fused_ipc_kernel<<>>( + commArgs, input, output, debug, elementsPerPeer, round); +} + void launch_tilexr_all_reduce_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerRank) { diff --git a/tests/udma/run_fused_prof.sh b/tests/udma/run_fused_prof.sh new file mode 100644 index 00000000..643605b5 --- /dev/null +++ b/tests/udma/run_fused_prof.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e +source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true +TILEXR_ROOT=/home/tileXR-new +UDMA_DIR=${TILEXR_ROOT}/tests/udma +export LD_LIBRARY_PATH=${UDMA_DIR}/install/lib:${UDMA_DIR}/install/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64:${LD_LIBRARY_PATH} +export TILEXR_DEMO_ALLTOALL_USE_UDMA=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=1 + +RANK_SIZE=4 +ELEM=2097152 +BIN=${UDMA_DIR}/install/bin/tilexr_udma_demo + +pids=() +for rank in $(seq 0 $((RANK_SIZE-1))); do + RANK=${rank} RANK_SIZE=${RANK_SIZE} "${BIN}" "${RANK_SIZE}" "${rank}" 2 "${ELEM}" "${RANK_SIZE}" 0 \ + > /tmp/a2a_rank${rank}.log 2>&1 & + pids+=("$!") +done +ret=0 +for idx in "${!pids[@]}"; do + wait "${pids[$idx]}" && echo "rank ${idx} ok" || { echo "rank ${idx} FAIL $?"; ret=1; } +done +echo "=== rank tails ===" +for rank in $(seq 0 $((RANK_SIZE-1))); do echo "--- rank ${rank} ---"; tail -6 /tmp/a2a_rank${rank}.log 2>/dev/null; done +exit ${ret} \ No newline at end of file diff --git a/tests/udma/run_plain_ipc_8m_4p.sh b/tests/udma/run_plain_ipc_8m_4p.sh new file mode 100644 index 00000000..0c3fe95f --- /dev/null +++ b/tests/udma/run_plain_ipc_8m_4p.sh @@ -0,0 +1,12 @@ +#!/bin/bash +export ASCEND_HOME_PATH=/usr/local/Ascend/cann-9.1.T560 +export PATH=/usr/local/Ascend/cann-9.1.T560/bin:/usr/local/Ascend/cann-9.1.T560/aarch64-linux/bin:$PATH +export LD_LIBRARY_PATH=/usr/local/Ascend/cann-9.1.T560/lib64:/usr/local/Ascend/cann-9.1.T560/aarch64-linux/lib64:/usr/local/Ascend/cann-9.1.T560/runtime/lib64:/usr/local/Ascend/cann-9.1.T560/aarch64-linux/devlib:${LD_LIBRARY_PATH:-} +export TILEXR_DEMO_DEVICES=2,3,6,7 +export TILEXR_DEMO_ALLTOALL_PLAIN_IPC=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=30 +export TILEXR_DEMO_ALLTOALL_SYNC_AT_END=1 +unset TILEXR_DEMO_ALLTOALL_USE_UDMA +export TILEXR_COMM_ID=127.0.0.1:12891 +cd /home/tileXR-new/tests/udma +exec /usr/bin/bash demo/run_tilexr_udma_demo.sh 2 4 524288 4 0 diff --git a/tests/udma/unit/test_tilexr_chip_map_sources.cpp b/tests/udma/unit/test_tilexr_chip_map_sources.cpp index 25d06cc2..1547cc80 100644 --- a/tests/udma/unit/test_tilexr_chip_map_sources.cpp +++ b/tests/udma/unit/test_tilexr_chip_map_sources.cpp @@ -32,6 +32,7 @@ std::string ReadFile(const std::string& path) int main() { const std::string internal = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/tilexr_internal.cpp"); + CHECK_CONTAINS(internal, "\"Ascend950DT_9582\""); CHECK_CONTAINS(internal, "\"Ascend950DT_9592\""); CHECK_CONTAINS(internal, "\"Ascend950PR_9599\""); CHECK_CONTAINS(internal, "ChipName::CHIP_950"); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 0219f9aa..6b4dddc6 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -154,7 +154,7 @@ void TestAllToAllDataAsFlagSource() CHECK_CONTAINS(demo, "ERROR: strict alltoall UDMA registration failed"); CHECK_CONTAINS(demo, "ERROR: strict alltoall UDMA CQ incomplete"); CHECK_CONTAINS(demo, "TileXRUDMARegister failed; use alltoall data-as-flag IPC fallback"); - CHECK_CONTAINS(demo, "skip all-to-all UDMA kernel; use data-as-flag IPC fallback"); + CHECK_CONTAINS(demo, "allToAllIpcFallbackLabel"); CHECK_CONTAINS(kernel, "#include \"tilexr_data_as_flag.h\""); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DATA_AS_FLAG_STAGING_OFFSET"); CHECK_CONTAINS(kernel, "DataAsFlagBlockCountForPayloadBytes"); @@ -163,6 +163,43 @@ void TestAllToAllDataAsFlagSource() CHECK_CONTAINS(kernel, "DataAsFlagCheckAndRecv"); } +void TestAllToAllChunkedUdmaSource() +{ + const std::string demo = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); + const std::string layout = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_alltoall_layout.h"); + const std::string kernel = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + + CHECK_CONTAINS(layout, "struct AllToAllChunkPlan"); + CHECK_CONTAINS(layout, "PlanAllToAllUdmaChunks"); + CHECK_CONTAINS(layout, "kAllToAllUdmaMaxRegisteredBytes"); + CHECK_CONTAINS(demo, "PlanAllToAllUdmaChunks"); + CHECK_CONTAINS(demo, "alltoall UDMA chunk plan"); + CHECK_CONTAINS(demo, "skip IPC staging capacity guard for strict UDMA alltoall"); + CHECK_CONTAINS(demo, "AllToAllPlainIpcStagingBytes"); + CHECK_CONTAINS(demo, "alltoall plain IPC staging bytes="); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_REPEAT"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_SYNC_AT_END"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_PLAIN_IPC"); + CHECK_CONTAINS(demo, "skip TileXRUDMARegister for alltoall plain IPC path"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_FUSED_IPC"); + CHECK_CONTAINS(demo, "skip TileXRUDMARegister for alltoall fused IPC path"); + CHECK_CONTAINS(demo, "alltoall fused IPC: single kernel send+flag+recv"); + CHECK_CONTAINS(demo, "alltoall use "); + CHECK_CONTAINS(demo, "plain IPC fallback"); + CHECK_CONTAINS(demo, "launch all-to-all kernel iter="); + CHECK_CONTAINS(demo, "for (uint32_t pass = 0; pass < chunkPlan.passCount; ++pass)"); + CHECK_CONTAINS(demo, "registered output chunk"); + CHECK_CONTAINS(kernel, "tilexr_all_to_all_plain_ipc_scatter_kernel"); + CHECK_CONTAINS(kernel, "tilexr_all_to_all_plain_ipc_gather_kernel"); + CHECK_CONTAINS(kernel, "tilexr_all_to_all_fused_ipc_kernel"); + CHECK_CONTAINS(kernel, "launch_tilexr_all_to_all_fused_ipc"); + CHECK_CONTAINS(kernel, "inputElementOffset"); + CHECK_CONTAINS(kernel, "chunkElements"); +} + } // namespace int main() @@ -173,6 +210,7 @@ int main() TestAllToAllMaxRank256With64MiBPerRank(); TestDemoDebugLayoutSource(); TestAllToAllDataAsFlagSource(); + TestAllToAllChunkedUdmaSource(); if (g_failures != 0) { std::cerr << g_failures << " all-to-all layout checks failed" << std::endl; return 1; From 8cfea673d374392f13b103a4a6e09adaf223a877 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Fri, 26 Jun 2026 09:31:21 +0800 Subject: [PATCH 007/163] 4p udma alltoall perf normal --- docs/alltoall-udma-success.md | 86 +++++++++++++++++++++++++++- tests/udma/demo/tilexr_udma_demo.cpp | 28 ++++++++- tests/udma/run_bw_4p_8m.sh | 27 +++++++++ tests/udma/run_prof_4p_32m.sh | 26 +++++++++ tests/udma/run_prof_4p_64m.sh | 26 +++++++++ 5 files changed, 191 insertions(+), 2 deletions(-) create mode 100644 tests/udma/run_bw_4p_8m.sh create mode 100644 tests/udma/run_prof_4p_32m.sh create mode 100644 tests/udma/run_prof_4p_64m.sh diff --git a/docs/alltoall-udma-success.md b/docs/alltoall-udma-success.md index 46d4082e..78ba7708 100644 --- a/docs/alltoall-udma-success.md +++ b/docs/alltoall-udma-success.md @@ -396,4 +396,88 @@ msprof --output="$PROF_DIR" \ |------|------| | `tests/udma/demo/tilexr_udma_demo_kernel.cpp` | kernel 改多核分片 + self-copy 改 DataCopyPad | | `tests/udma/demo/tilexr_udma_demo.cpp` | 两处 launch blockDim: 1 → rankSize | -| `tests/udma/demo/run_fused_prof.sh` | msprof 采集用启动脚本(新增) | \ No newline at end of file +| `tests/udma/demo/run_fused_prof.sh` | msprof 采集用启动脚本(新增) | + +--- + +## 8. 性能扩展测试(4 rank,强制 UDMA) + +测试条件:`TILEXR_DEMO_ALLTOALL_USE_UDMA=1`,4 张 Ascend950DT,`blockDim=rankSize=4`,每 peer 数据量分别为 8M / 32M / 64M。其中 8M 走 host 侧 demo 内置的 `aclrtSynchronizeStream` + `chrono` 计时(单 pass,100 次循环);32M/64M 因超过单次 lmem 注册上限(64MB)需多 pass,取 kernel 侧 msprof `op_summary` 各 pass task duration 求和。 + +### 8.1 数据量定义 + +- `elementsPerRank` = 每 peer 元素数(8M = 2097152 个 int32) +- 总 input = `rankSize × perPeer` = 4 × 8M = 32MB(每个 dstRank 收到 4 份) +- 注意 **self 份(8MB)走本地 DataCopyPad,不占网络** +- 单卡单向 P2P 数据量(纯网络)= `(rankSize-1) × perPeer` = 3 × perPeer + +### 8.2 总览结果 + +| 每 peer 数据量 | pass 数 | 总 kernel 耗时 | 每 pass 耗时 | 单向 P2P 带宽 | +|---------------|---------|-----------------|-------------|------------------| +| **8M** | 1 | ~200 μs | 200 μs | ~120 GB/s | +| **32M** | 4 | ~708 μs | ~177 μs | ~136 GB/s | +| **64M** | 8 | ~1439 μs | ~180 μs | ~133 GB/s | + +### 8.3 8M/peer 单 pass(100 次,仅末尾 sync) + +`TILEXR_DEMO_ALLTOALL_REPEAT=100 TILEXR_DEMO_ALLTOALL_SYNC_AT_END=1`,host 侧 chrono 计时: + +| Rank | 100 次总耗时 | 单次 | 带宽(含 self 份) | +|------|------------|------|-------------------| +| rank0 | 20.74 ms | 207.4 μs | 161.8 GB/s | +| rank1 | 20.04 ms | 200.4 μs | 167.4 GB/s | +| rank2 | 19.97 ms | 199.7 μs | 168.1 GB/s | +| rank3 | 19.92 ms | 199.2 μs | 168.4 GB/s | + +注:总 payload = `rankSize × elementsPerRank × 4 = 32MB`(含 self 份,本地拷贝),扣除后纯网络单向 P2P 带宽约 ~120 GB/s。 + +### 8.4 32M/peer(4 pass,msprof) + +| Device | pass0 | pass1 | pass2 | pass3 | 总计 | +|--------|-------|-------|-------|-------|------| +| dev0 | 183.6 | 172.9 | 178.9 | 173.5 | 708.9 μs | +| dev1 | 184.0 | 173.0 | 173.2 | 173.3 | 703.5 μs | +| dev2 | 191.6 | 172.7 | 173.8 | 173.9 | 712.0 μs | +| dev3 | 193.1 | 172.6 | 178.4 | 172.3 | 716.4 μs | + +### 8.5 64M/peer(8 pass,msprof) + +| Device | pass0 | pass1-7 均值 | 总计 | +|--------|-------|-------------|------| +| dev0 | 192.7 | ~177.6 | 1435.8 μs | +| dev1 | 193.7 | ~178.8 | 1445.2 μs | +| dev2 | 194.1 | ~179.2 | 1448.7 μs | +| dev3 | 188.4 | ~176.9 | 1427.4 μs | + +### 8.6 带宽换算 + +| 项目 | 8M | 32M | 64M | +|------|-----|-----|-----| +| 单向 P2P(净) | 24MB ÷ ~120 GB/s | 96MB ÷ ~136 GB/s | 192MB ÷ ~133 GB/s | +| 双向 P2P | 48MB ÷ ~240 GB/s | 192MB ÷ ~271 GB/s | 384MB ÷ ~267 GB/s | +| 含 self(含本地) | 32MB ÷ ~168 GB/s | 128MB ÷ ~181 GB/s | 256MB ÷ ~178 GB/s | + +### 8.7 分析 + +1. **每 pass 耗时稳定在 ~177-200μs**,与 payload 大小无关(每 pass 仍是 8M/peer),说明 kernel 耗时由 **per-pass 的 Quiet 轮询时间主导**,而非带宽受限。 +2. **单向 P2P 带宽约 120-136 GB/s**,接近物理链路上限,瓶颈不在带宽。 +3. **根因分析**:"发一个等一个"的串行 Quiet 模式使每 pass 必须 4 个 peer 串行完成 `PutNbi` 后再 `UDMAQuietStatus` 轮询 CQ,导致每 pass 固定 ~180μs,与 payload 大小无关。 +4. **优化方向**:若 PutNbi 后不立即 Quiet,增大 QP 数(4× QP 方案),可让 per-pass 的 Quiet 等待时间隐藏到 payload 传输中,提升大 payload 场景吞吐。 + +### 8.8 计时代码说明 + +`tests/udma/demo/tilexr_udma_demo.cpp` strict UDMA 单 pass 路径(`:847` else 分支)加 host 计时:每次前 `aclrtSynchronizeStream` + `chrono::steady_clock` 记点,循环后再次 sync,输出 `total/perIter/payload/bw`;其中 `syncAllToAllAtEnd=true` 时仅在末尾做一次同步。 + +--- + +## 9. 文件清单(更新) + +| 文件 | 改动 | +|------|------| +| `tests/udma/demo/tilexr_udma_demo_kernel.cpp` | kernel 多核分片 + self-copy 改 DataCopyPad | +| `tests/udma/demo/tilexr_udma_demo.cpp` | 两处 launch blockDim: 1 → rankSize + strict 单 pass 计时代码 | +| `tests/udma/demo/run_fused_prof.sh` | msprof 采集用启动脚本(新增) | +| `tests/udma/run_bw_4p_8m.sh` | 8M 100 次带宽测试脚本(新增) | +| `tests/udma/run_prof_4p_32m.sh` | 32M msprof 测试脚本(新增) | +| `tests/udma/run_prof_4p_64m.sh` | 64M msprof 测试脚本(新增) | diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 771e06d0..13655b13 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -845,8 +845,18 @@ int main(int argc, char** argv) udmaRegistered = false; } } else { + if (!CheckAcl(rank, "aclrtSynchronizeStream pre-alltoall", aclrtSynchronizeStream(stream))) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + auto a2aStart = std::chrono::steady_clock::now(); for (int iter = 0; iter < allToAllRepeat; ++iter) { - PrintStatus(rank, "launch all-to-all kernel iter=" + std::to_string(iter)); + if (iter == 0) { + PrintStatus(rank, "launch all-to-all kernel repeat=" + std::to_string(allToAllRepeat)); + } launch_tilexr_udma_all_to_all( static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset), 0, @@ -860,6 +870,22 @@ int main(int argc, char** argv) return 1; } } + if (syncAllToAllAtEnd && + !CheckAcl(rank, "aclrtSynchronizeStream post-alltoall", aclrtSynchronizeStream(stream))) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + } + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + auto a2aEnd = std::chrono::steady_clock::now(); + double a2aMs = std::chrono::duration(a2aEnd - a2aStart).count(); + double a2aPerIterUs = (allToAllRepeat > 0) ? (a2aMs * 1000.0 / static_cast(allToAllRepeat)) : 0.0; + double payloadBytes = static_cast(rankSize) * static_cast(elementsPerRank) * sizeof(int32_t); + double bwGbs = (a2aPerIterUs > 0.0) ? (payloadBytes / (a2aPerIterUs * 1e3)) : 0.0; + std::cout << "[rank " << rank << "] alltoall udma " << allToAllRepeat + << " iters total=" << a2aMs << " ms perIter=" << a2aPerIterUs + << " us payload=" << payloadBytes << " bytes bw=" << bwGbs << " GB/s" << std::endl; } } else if (testType == 3) { PrintStatus(rank, "launch all-reduce IPC scatter kernel"); diff --git a/tests/udma/run_bw_4p_8m.sh b/tests/udma/run_bw_4p_8m.sh new file mode 100644 index 00000000..826b4b8d --- /dev/null +++ b/tests/udma/run_bw_4p_8m.sh @@ -0,0 +1,27 @@ +#!/bin/bash +set -e +source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true +TILEXR_ROOT=/home/tileXR-new +UDMA_DIR=${TILEXR_ROOT}/tests/udma +export LD_LIBRARY_PATH=${UDMA_DIR}/install/lib:${UDMA_DIR}/install/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64:${LD_LIBRARY_PATH} +export TILEXR_DEMO_ALLTOALL_USE_UDMA=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=100 +export TILEXR_DEMO_ALLTOALL_SYNC_AT_END=1 + +RANK_SIZE=4 +ELEM=2097152 +BIN=${UDMA_DIR}/install/bin/tilexr_udma_demo + +pids=() +for rank in $(seq 0 $((RANK_SIZE-1))); do + RANK=${rank} RANK_SIZE=${RANK_SIZE} "${BIN}" "${RANK_SIZE}" "${rank}" 2 "${ELEM}" "${RANK_SIZE}" 0 \ + > /tmp/a2a_bw_rank${rank}.log 2>&1 & + pids+=("$!") +done +ret=0 +for idx in "${!pids[@]}"; do + wait "${pids[$idx]}" && echo "rank ${idx} ok" || { echo "rank ${idx} FAIL $?"; ret=1; } +done +echo "=== rank tails ===" +for rank in $(seq 0 $((RANK_SIZE-1))); do echo "--- rank ${rank} ---"; tail -4 /tmp/a2a_bw_rank${rank}.log 2>/dev/null; done +exit ${ret} \ No newline at end of file diff --git a/tests/udma/run_prof_4p_32m.sh b/tests/udma/run_prof_4p_32m.sh new file mode 100644 index 00000000..73be9a0e --- /dev/null +++ b/tests/udma/run_prof_4p_32m.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e +source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true +TILEXR_ROOT=/home/tileXR-new +UDMA_DIR=${TILEXR_ROOT}/tests/udma +export LD_LIBRARY_PATH=${UDMA_DIR}/install/lib:${UDMA_DIR}/install/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64:${LD_LIBRARY_PATH} +export TILEXR_DEMO_ALLTOALL_USE_UDMA=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=1 + +RANK_SIZE=4 +ELEM=8388608 +BIN=${UDMA_DIR}/install/bin/tilexr_udma_demo + +pids=() +for rank in $(seq 0 $((RANK_SIZE-1))); do + RANK=${rank} RANK_SIZE=${RANK_SIZE} "${BIN}" "${RANK_SIZE}" "${rank}" 2 "${ELEM}" "${RANK_SIZE}" 0 \ + > /tmp/a2a32m_rank${rank}.log 2>&1 & + pids+=("$!") +done +ret=0 +for idx in "${!pids[@]}"; do + wait "${pids[$idx]}" && echo "rank ${idx} ok" || { echo "rank ${idx} FAIL $?"; ret=1; } +done +echo "=== pass/kernel lines ===" +for rank in $(seq 0 $((RANK_SIZE-1))); do grep -E "chunk plan|launch all-to-all kernel pass|success" /tmp/a2a32m_rank${rank}.log 2>/dev/null | tail -6; done +exit ${ret} \ No newline at end of file diff --git a/tests/udma/run_prof_4p_64m.sh b/tests/udma/run_prof_4p_64m.sh new file mode 100644 index 00000000..a7a36726 --- /dev/null +++ b/tests/udma/run_prof_4p_64m.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e +source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true +TILEXR_ROOT=/home/tileXR-new +UDMA_DIR=${TILEXR_ROOT}/tests/udma +export LD_LIBRARY_PATH=${UDMA_DIR}/install/lib:${UDMA_DIR}/install/lib64:${TILEXR_ROOT}/install/lib:${TILEXR_ROOT}/install/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64:${LD_LIBRARY_PATH} +export TILEXR_DEMO_ALLTOALL_USE_UDMA=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=1 + +RANK_SIZE=4 +ELEM=16777216 +BIN=${UDMA_DIR}/install/bin/tilexr_udma_demo + +pids=() +for rank in $(seq 0 $((RANK_SIZE-1))); do + RANK=${rank} RANK_SIZE=${RANK_SIZE} "${BIN}" "${RANK_SIZE}" "${rank}" 2 "${ELEM}" "${RANK_SIZE}" 0 \ + > /tmp/a2a64m_rank${rank}.log 2>&1 & + pids+=("$!") +done +ret=0 +for idx in "${!pids[@]}"; do + wait "${pids[$idx]}" && echo "rank ${idx} ok" || { echo "rank ${idx} FAIL $?"; ret=1; } +done +echo "=== rank0 chunk plan ===" +grep -E "chunk plan|demo success|FAIL" /tmp/a2a64m_rank0.log | tail -3 +exit ${ret} \ No newline at end of file From 28a45d657667047e7d9f0b23cac574df3e989422 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Fri, 26 Jun 2026 10:03:47 +0800 Subject: [PATCH 008/163] udma alltoall lmem register 64MB->128MB, pass count halved --- docs/alltoall-udma-success.md | 72 +++++++++++++++++++ tests/udma/demo/tilexr_udma_alltoall_layout.h | 2 +- 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/docs/alltoall-udma-success.md b/docs/alltoall-udma-success.md index 78ba7708..aa2e5edf 100644 --- a/docs/alltoall-udma-success.md +++ b/docs/alltoall-udma-success.md @@ -481,3 +481,75 @@ msprof --output="$PROF_DIR" \ | `tests/udma/run_bw_4p_8m.sh` | 8M 100 次带宽测试脚本(新增) | | `tests/udma/run_prof_4p_32m.sh` | 32M msprof 测试脚本(新增) | | `tests/udma/run_prof_4p_64m.sh` | 64M msprof 测试脚本(新增) | + +--- + +## 10. 注册上限 64MB→128MB 优化实验 + +### 10.1 背景 + +`kAllToAllUdmaMaxRegisteredBytes` 原值 64MB,导致 32M/peer 需 4 pass、64M/peer 需 8 pass(每 pass 注册 `chunkBytesPerRank×2`,input+output 各一份)。尝试增大注册上限以减少 pass 数。 + +### 10.2 lmem 注册上限实测 + +底层 `RaCtxLmemRegister` 对单次注册 segment 大小有上限: + +| 注册上限 | registeredBytes | 结果 | +|---------|-----------------|------| +| 64MB(基线) | 67108864(64MB) | ✅ 成功 | +| **128MB** | 134217728(128MB) | ✅ 成功 | +| 256MB | 268435456(256MB) | ❌ 失败 `ret=528101` | + +结论:128MB 是 `RaCtxLmemRegister` 单次注册的安全上限,256MB 超限。 + +### 10.3 pass 数变化 + +| 每 peer 数据量 | 64MB 上限 pass 数 | **128MB 上限 pass 数** | 变化 | +|---------------|-------------------|------------------------|------| +| 8M | 1 | 1 | 不变(8M 单 pass 只需 64MB) | +| 32M | 4 | **2** | 减半 | +| 64M | 8 | **4** | 减半 | + +### 10.4 性能结果(128MB 上限) + +**8M/peer(1 pass,100 次)** + +| Rank | 单次耗时 | 带宽(含 self) | +|------|---------|---------------| +| rank0 | 207.7 μs | 161.6 GB/s | +| rank1 | 226.8 μs | 147.9 GB/s | +| rank2 | 234.3 μs | 143.2 GB/s | +| rank3 | 216.2 μs | 155.2 GB/s | + +**32M/peer(2 pass,msprof Task Duration μs)** + +| Device | pass0 | pass1 | 总计 | +|--------|-------|-------|------| +| dev0 | 360.058 | 342.163 | 702.221 | +| dev1 | 359.141 | 347.583 | 706.724 | +| dev2 | 350.667 | 351.752 | 702.419 | +| dev3 | 347.477 | 350.923 | 698.400 | + +**64M/peer(4 pass,msprof Task Duration μs)** + +| Device | pass0 | pass1 | pass2 | pass3 | 总计 | +|--------|-------|-------|-------|-------|------| +| dev0 | 361.361 | 350.103 | 338.688 | 347.098 | 1397.250 | +| dev1 | 362.885 | 346.606 | 341.455 | 340.896 | 1391.842 | +| dev2 | 350.910 | 348.758 | 350.444 | 336.601 | 1386.713 | +| dev3 | 351.728 | 344.601 | 351.363 | 346.311 | 1394.003 | + +### 10.5 与基线对比 + +| 场景 | 64MB 基线总耗时 | 128MB 优化总耗时 | 变化 | per-pass 耗时 | +|------|----------------|------------------|------|--------------| +| 8M/1pass | ~200 μs | ~207 μs | ≈持平 | 200→207 μs | +| 32M/2pass | ~708 μs(4pass) | ~702 μs(2pass) | ≈持平 | 177→351 μs | +| 64M/4pass | ~1439 μs(8pass) | ~1392 μs(4pass) | -3.3% | 180→348 μs | + +### 10.6 分析 + +1. **pass 数减半但总耗时几乎不变**:64M 从 8 pass 降到 4 pass,但每 pass 耗时从 ~180μs 升到 ~350μs(翻倍),因为每 pass payload 翻倍(8M→16M/peer)。 +2. **证实瓶颈是 per-pass Quiet 轮询,不是带宽**:pass 数减半,payload 翻倍,总耗时基本守恒——说明每 pass 的 Quiet 等待时间(~150-170μs 固定开销)主导,而非数据传输时间。 +3. **128MB 上限本身是安全的**:注册稳定通过,无功能问题,可作为低风险基线保留(减少 pass 数降低 host 侧 launch/同步开销)。 +4. **真正提速方向**:需消除 per-pass 串行 Quiet——即多 QP 方案(让 Quiet 等待隐藏到并行传输中),而非增大单次注册。 diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index d91ddd85..e99058c2 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -16,7 +16,7 @@ namespace TileXR { namespace Demo { constexpr int32_t kAllToAllBaseValue = 100000; -constexpr size_t kAllToAllUdmaMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllUdmaMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL; struct AllToAllChunkPlan { uint32_t passCount = 1; From cbd96666b590c98644d59745def30ec474fc5a1d Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sat, 27 Jun 2026 10:04:16 +0800 Subject: [PATCH 009/163] 8p udma alltoall perf on 156 (8M/32M/64M, ~272-283 GB/s) --- docs/alltoall-udma-success.md | 71 +++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/docs/alltoall-udma-success.md b/docs/alltoall-udma-success.md index aa2e5edf..758f3b5f 100644 --- a/docs/alltoall-udma-success.md +++ b/docs/alltoall-udma-success.md @@ -553,3 +553,74 @@ msprof --output="$PROF_DIR" \ 2. **证实瓶颈是 per-pass Quiet 轮询,不是带宽**:pass 数减半,payload 翻倍,总耗时基本守恒——说明每 pass 的 Quiet 等待时间(~150-170μs 固定开销)主导,而非数据传输时间。 3. **128MB 上限本身是安全的**:注册稳定通过,无功能问题,可作为低风险基线保留(减少 pass 数降低 host 侧 launch/同步开销)。 4. **真正提速方向**:需消除 per-pass 串行 Quiet——即多 QP 方案(让 Quiet 等待隐藏到并行传输中),而非增大单次注册。 + + +## 11. 8p UDMA AllToAll 性能(141.62.19.156,Ascend950DT ×8) + +**测试环境** +- 服务器:141.62.19.156(root),8 × Ascend950DT(Chip V120/9584),CANN 9.1.T560,bisheng 工具链。 +- 所有 8 卡 health=OK(对比 148/152 的 CCU config error)。 +- 代码:HEAD 28a45d6(single-QP,128MB lmem 上限)。**156 上因 URMA 单次注册 128MB 会 ret=528101,故把 kAllToAllUdmaMaxRegisteredBytes 临时调为 32MB** 后构建(8M=4pass / 32M=16pass / 64M=32pass)。该改动仅 156 本地构建用,本地仓库仍保持 128MB。 +- 8 rank alltoall,TILEXR_DEMO_DEVICES="0,1,2,3,4,5,6,7",force UDMA(TILEXR_DEMO_ALLTOALL_USE_UDMA=1),repeat=1。 +- msprof:--task-time=l0 --ai-core=on --aic-mode=sample-based --aic-freq=100 --aic-metrics=PipeUtilization --aicpu=on --runtime-api=on,kernel Task Duration 取自各 device 的 op_summary_*.csv 中 ilexr_udma_all_to_all_kernel 行 Task Duration 之和。 + +**8M/peer(ELEM=2097152,4 pass,msprof Task Duration μs)** + +| Device | pass 数 | 各 pass 总和(μs) | per-pass 均值(μs) | +|--------|---------|------------------|--------------------| +| dev0 | 4 | 232.1 | 58.0 | +| dev1 | 4 | 235.2 | 58.8 | +| dev2 | 4 | 233.8 | 58.5 | +| dev3 | 4 | 234.1 | 58.5 | +| dev4 | 4 | 235.2 | 58.8 | +| dev5 | 4 | 237.3 | 59.3 | +| dev6 | 4 | 239.6 | 59.9 | +| dev7 | 4 | 231.4 | 57.9 | + +- kernel 侧总耗时 ~235 μs。payload = 8M×8 = 64MB → 聚合 p2p 带宽 ≈ **272 GB/s**。 + +**32M/peer(ELEM=8388608,16 pass,msprof Task Duration μs)** + +| Device | pass 数 | 各 pass 总和(μs) | +|--------|---------|------------------| +| dev0 | 16 | 919.4 | +| dev1 | 16 | 910.4 | +| dev2 | 16 | 917.7 | +| dev3 | 16 | 911.7 | +| dev4 | 16 | 904.9 | +| dev5 | 16 | 922.5 | +| dev6 | 16 | 910.6 | +| dev7 | 16 | 911.8 | + +- kernel 侧总耗时 ~915 μs。payload = 32M×8 = 256MB → 聚合 p2p 带宽 ≈ **280 GB/s**。 + +**64M/peer(ELEM=16777216,32 pass,msprof Task Duration μs)** + +| Device | pass 数 | 各 pass 总和(μs) | +|--------|---------|------------------| +| dev0 | 32 | 1820.2 | +| dev1 | 32 | 1816.8 | +| dev2 | 32 | 1824.9 | +| dev3 | 32 | 1803.1 | +| dev4 | 32 | 1810.6 | +| dev5 | 32 | 1822.4 | +| dev6 | 32 | 1809.3 | +| dev7 | 32 | 1797.5 | + +- kernel 侧总耗时 ~1810 μs。payload = 64M×8 = 512MB → 聚合 p2p 带宽 ≈ **283 GB/s**。 + +**11.1 三档汇总与带宽** + +| 场景 | pass 数 | kernel 总耗时(μs) | payload | 聚合 p2p 带宽 | +|------|---------|-------------------|---------|---------------| +| 8M/peer | 4 | ~235 | 64MB | ~272 GB/s | +| 32M/peer | 16 | ~915 | 256MB | ~280 GB/s | +| 64M/peer | 32 | ~1810 | 512MB | ~283 GB/s | + +**11.2 分析** + +1. **8 卡稳定跑通**:156 上 8 卡 health 全 OK,UDMA P2P 可达所有 peer(尽管 hccn_tool -g -link 只显示到 dev0/1 的 UP 链路,实际 UDMA 走 HCCS 路径能到 dev2-7 互连——2-2/4-5/6-7 等任意对 2 卡 alltoall 均成功)。 +2. **带宽随数据量上升并趋于稳定**:8M→64M,带宽从 272 升到 283 GB/s,趋近 HCCS 渐近带宽,说明大 payload 下 per-pass 固定开销(Quiet 轮询 ~150-170μs)被摊薄。 +3. **per-pass 耗时稳定**:8M per-pass ~58μs(2M elem/peer),32M per-pass ~57μs(512K elem/peer,但 pass 数 4 倍),64M per-pass ~57μs —— per-pass 耗时基本由 Quiet 等待主导,与单 pass payload 关系不大(32MB 注册上限下每 pass chunkElements 固定 524288)。 +4. **与 4 卡(144)对比**:4p 8M ~200μs/120GB/s,8p 8M ~235μs/272GB/s —— 8 卡聚合带宽显著提升(2.3×),耗时仅增 17%,说明 UDMA P2P 在 8 卡下扩展性良好。 +5. **32MB 注册上限的代价**:相对 144 的 128MB 上限,156 被迫用 32MB,导致 64M 需 32 pass(144 只需 4 pass)。但 kernel 总耗时 1810μs 仍在合理范围,带宽 283 GB/s 不输 144 的 133 GB/s(64M),说明 pass 数增加的固定开销被 8 卡并行摊薄。 From 2aca40e5d1ac36be87f945a88ad6153bcb7f48cb Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sat, 27 Jun 2026 10:20:57 +0800 Subject: [PATCH 010/163] add p2p/datacopy latency micro-kernels (testType 4/5) to isolate UDMA comm vs local DataCopy cost --- tests/udma/demo/tilexr_udma_demo.cpp | 51 +++++-- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 140 ++++++++++++++++++++ 2 files changed, 180 insertions(+), 11 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 13655b13..43bad52c 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -34,6 +34,13 @@ extern void launch_tilexr_udma_all_to_all( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, int32_t chunkElements); +extern void launch_tilexr_udma_p2p_latency( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, + int32_t chunkElements); +extern void launch_tilexr_datacopy_latency( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, int32_t chunkElements); extern void launch_tilexr_all_to_all_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); extern void launch_tilexr_all_to_all_ipc_gather( @@ -616,9 +623,9 @@ int main(int argc, char** argv) return 1; } - bool isAllToAll = testType == 2; + bool isAllToAll = testType == 2 || testType == 4 || testType == 5; bool isAllReduce = testType == 3; - bool strictAllToAllUdma = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_USE_UDMA", 0) != 0; + bool strictAllToAllUdma = isAllToAll && (testType == 4 || GetEnvInt("TILEXR_DEMO_ALLTOALL_USE_UDMA", 0) != 0); int allToAllRepeat = isAllToAll ? std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)) : 1; bool syncAllToAllAtEnd = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_SYNC_AT_END", 0) != 0; @@ -775,7 +782,7 @@ int main(int argc, char** argv) return 1; } - if (testType == 2) { + if (isAllToAll) { if (forceAllToAllIpcFallback) { PrintStatus(rank, std::string("skip all-to-all UDMA kernel; use ") + allToAllIpcFallbackLabel); } else if (chunkedStrictAllToAll) { @@ -804,10 +811,21 @@ int main(int argc, char** argv) udmaRegistered = true; } PrintStatus(rank, "launch all-to-all kernel pass=" + std::to_string(pass)); - launch_tilexr_udma_all_to_all( - static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(debug), chunkElements, static_cast(outputOffset), - 0, chunkElements); + if (testType == 4) { + launch_tilexr_udma_p2p_latency( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), chunkElements, static_cast(outputOffset), + 0, chunkElements); + } else if (testType == 5) { + launch_tilexr_datacopy_latency( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), chunkElements, chunkElements); + } else { + launch_tilexr_udma_all_to_all( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), chunkElements, static_cast(outputOffset), + 0, chunkElements); + } if (!CheckAcl(rank, "aclrtSynchronizeStream", aclrtSynchronizeStream(stream)) || !DemoBarrierAll(rank, rankSize, "all ranks completed demo kernels")) { if (udmaRegistered) { @@ -857,10 +875,21 @@ int main(int argc, char** argv) if (iter == 0) { PrintStatus(rank, "launch all-to-all kernel repeat=" + std::to_string(allToAllRepeat)); } - launch_tilexr_udma_all_to_all( - static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset), 0, - elementsPerRank); + if (testType == 4) { + launch_tilexr_udma_p2p_latency( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset), 0, + elementsPerRank); + } else if (testType == 5) { + launch_tilexr_datacopy_latency( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, elementsPerRank); + } else { + launch_tilexr_udma_all_to_all( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(debug), elementsPerRank, static_cast(outputOffset), 0, + elementsPerRank); + } if (!syncAllToAllAtEnd && !CheckAcl(rank, "aclrtSynchronizeStream", aclrtSynchronizeStream(stream))) { if (udmaRegistered) { diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 81e287f1..fc0907cb 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -805,3 +805,143 @@ void launch_tilexr_udma_registered_smoke( tilexr_udma_registered_smoke_kernel<<>>( commArgs, local, debug, bytes, signal); } + +// --------------------------------------------------------------------------- +// Latency micro-kernels (testType 4 = P2P-only, testType 5 = DataCopy-only). +// These mirror the two halves of tilexr_udma_all_to_all_kernel so the P2P +// communication latency and the local DataCopy latency can be measured in +// isolation. Each block handles exactly one peer (block b -> peer b); for +// the P2P kernel the self-peer (peer == rank) is skipped so that block does +// no network work, and for the DataCopy kernel every block performs a local +// GM->UB->GM copy of the same payload size as one alltoall peer slice. +// --------------------------------------------------------------------------- + +extern "C" __global__ __aicore__ void tilexr_udma_p2p_latency_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, + int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, + int32_t chunkElements) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + bool enabled = TileXR::UDMARegistryEnabled(args); + + const int32_t blockIdx = AscendC::GetBlockIdx(); + const int32_t blockNum = AscendC::GetBlockNum(); + + if (blockIdx == 0 && debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = enabled ? 1 : 0; + debug[4] = elementsPerPeer; + debug[5] = static_cast(outputByteOffset); + } + if (!enabled) { + return; + } + + const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; + const uint64_t payloadBytes = AllToAllPayloadBytes(effectiveChunkElements); + const uint32_t bytes = static_cast(payloadBytes); + + // One block per peer, skip self (peer == rank): no local copy here. + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + continue; + } + auto localSrc = input + static_cast(peer) * elementsPerPeer + inputElementOffset; + uint64_t remoteOffset = outputByteOffset + + static_cast(rank) * payloadBytes; + TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); + uint32_t status = TileXR::UDMAQuietStatus(args, peer); + if (debug != nullptr && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } + } +} + +extern "C" __global__ __aicore__ void tilexr_datacopy_latency_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR debugGM, + int32_t elementsPerPeer, int32_t chunkElements) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + + const int32_t blockIdx = AscendC::GetBlockIdx(); + const int32_t blockNum = AscendC::GetBlockNum(); + + if (blockIdx == 0 && debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = 1; // always "enabled" path for datacopy + debug[4] = elementsPerPeer; + } + + const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; + const uint64_t payloadBytes = AllToAllPayloadBytes(effectiveChunkElements); + const uint32_t bytes = static_cast(payloadBytes); + + // Each block performs the same GM->UB->GM self-copy that the alltoall + // kernel uses for the self-peer slice. block b copies rank b's slice. + for (int32_t whichRank = blockIdx; whichRank < rankSize; whichRank += blockNum) { + auto selfSrc = input + static_cast(whichRank) * elementsPerPeer; + auto selfDst = output + static_cast(whichRank) * effectiveChunkElements; + constexpr uint32_t SELF_COPY_UB_BYTES = 64 * 1024; + AscendC::TPipe pipe; + AscendC::TBuf selfCopyTBuf; + pipe.InitBuffer(selfCopyTBuf, SELF_COPY_UB_BYTES); + AscendC::LocalTensor selfCopyLocal = selfCopyTBuf.Get(); + + auto selfSrcBytes = reinterpret_cast<__gm__ uint8_t*>(selfSrc); + auto selfDstBytes = reinterpret_cast<__gm__ uint8_t*>(selfDst); + for (uint32_t offset = 0; offset < bytes; offset += SELF_COPY_UB_BYTES) { + uint32_t copyBytes = (bytes - offset < SELF_COPY_UB_BYTES) + ? (bytes - offset) : SELF_COPY_UB_BYTES; + + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(selfSrcBytes + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, copyBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(selfCopyLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(selfDstBytes + offset); + AscendC::DataCopyExtParams copyOut {1U, copyBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, selfCopyLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); + } +} + +void launch_tilexr_udma_p2p_latency( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, uint64_t outputByteOffset, int32_t inputElementOffset, + int32_t chunkElements) +{ + tilexr_udma_p2p_latency_kernel<<>>( + commArgs, input, output, debug, elementsPerPeer, outputByteOffset, inputElementOffset, chunkElements); +} + +void launch_tilexr_datacopy_latency( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR debug, int32_t elementsPerPeer, int32_t chunkElements) +{ + tilexr_datacopy_latency_kernel<<>>( + commArgs, input, output, debug, elementsPerPeer, chunkElements); +} + From fd04a96aef7525d298054062fa2c157c2e6c18c7 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sat, 27 Jun 2026 18:42:09 +0800 Subject: [PATCH 011/163] test: add UDMA alltoall bigdata layout plan --- tests/udma/demo/tilexr_udma_alltoall_layout.h | 48 +++++++++++++++++++ .../unit/test_tilexr_udma_alltoall_layout.cpp | 18 +++++++ 2 files changed, 66 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index e99058c2..0a50439d 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -25,6 +25,19 @@ struct AllToAllChunkPlan { size_t registeredBytes = 0; }; +constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; + +struct AllToAllBigDataPlan { + uint32_t passCount = 1; + int32_t chunkElements = 0; + size_t chunkBytesPerPeer = 0; + size_t dataBytes = 0; + size_t readySignalOffset = 0; + size_t ackSignalOffset = 0; + size_t signalBytes = 0; + size_t registeredBytes = 0; +}; + inline int32_t AllToAllValue(int srcRank, int dstRank) { return kAllToAllBaseValue + srcRank * 1000 + dstRank; @@ -60,6 +73,41 @@ inline AllToAllChunkPlan PlanAllToAllUdmaChunks(int rankSize, int32_t elementsPe return plan; } +inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t elementsPerPeer) +{ + AllToAllBigDataPlan plan {}; + if (rankSize <= 0 || elementsPerPeer <= 0) { + return plan; + } + + plan.registeredBytes = kAllToAllBigDataMaxRegisteredBytes; + plan.signalBytes = 2ULL * static_cast(rankSize) * sizeof(uint64_t); + if (plan.signalBytes >= plan.registeredBytes) { + return plan; + } + + plan.dataBytes = plan.registeredBytes - plan.signalBytes; + plan.chunkElements = static_cast( + plan.dataBytes / (static_cast(rankSize) * sizeof(int32_t))); + if (plan.chunkElements <= 0) { + plan.chunkElements = 1; + } + if (plan.chunkElements > elementsPerPeer) { + plan.chunkElements = elementsPerPeer; + } + + plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); + plan.dataBytes = static_cast(rankSize) * plan.chunkBytesPerPeer; + plan.readySignalOffset = plan.dataBytes; + plan.ackSignalOffset = plan.readySignalOffset + static_cast(rankSize) * sizeof(uint64_t); + plan.signalBytes = 2ULL * static_cast(rankSize) * sizeof(uint64_t); + plan.registeredBytes = plan.dataBytes + plan.signalBytes; + plan.passCount = static_cast( + (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1) / + static_cast(plan.chunkElements)); + return plan; +} + inline void FillAllToAllInput( std::vector& input, int rank, int rankSize, int32_t elementsPerPeer) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 6b4dddc6..93ee9c2e 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -126,6 +126,23 @@ void TestAllToAllMaxRank256With64MiBPerRank() } } +void TestAllToAllBigDataPlan() +{ + constexpr int rankSize = 8; + constexpr int32_t elementsPerPeer = 16 * 1024 * 1024; // 64 MiB per peer for int32_t. + const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); + + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 64ULL * 1024ULL * 1024ULL); + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes); + CHECK_EQ(plan.signalBytes, 2ULL * static_cast(rankSize) * sizeof(uint64_t)); + CHECK_EQ(plan.readySignalOffset, plan.dataBytes); + CHECK_EQ(plan.ackSignalOffset, plan.dataBytes + static_cast(rankSize) * sizeof(uint64_t)); + CHECK_EQ(plan.dataBytes + plan.signalBytes <= plan.registeredBytes, true); + CHECK_EQ(plan.chunkElements > 0, true); + CHECK_EQ(plan.chunkBytesPerPeer, static_cast(plan.chunkElements) * sizeof(int32_t)); + CHECK_EQ(plan.passCount > 1, true); +} + void TestDemoDebugLayoutSource() { const std::string demo = @@ -208,6 +225,7 @@ int main() TestAllToAllOutputValidation(); TestBuildAllToAllOutput(); TestAllToAllMaxRank256With64MiBPerRank(); + TestAllToAllBigDataPlan(); TestDemoDebugLayoutSource(); TestAllToAllDataAsFlagSource(); TestAllToAllChunkedUdmaSource(); From 71f48e695e65072bafe75a5e48988cda12a6f1e0 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sat, 27 Jun 2026 18:46:02 +0800 Subject: [PATCH 012/163] test: add UDMA alltoall bigdata source guards --- .../unit/test_tilexr_udma_alltoall_layout.cpp | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 93ee9c2e..ebb99c89 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -217,6 +217,28 @@ void TestAllToAllChunkedUdmaSource() CHECK_CONTAINS(kernel, "chunkElements"); } +void TestAllToAllBigDataSource() +{ + const std::string demo = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); + const std::string kernel = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + + CHECK_CONTAINS(demo, "testType == 7"); + CHECK_CONTAINS(demo, "PlanAllToAllBigDataUdma"); + CHECK_CONTAINS(demo, "bigdata alltoall registered dataBytes="); + CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_bigdata"); + CHECK_CONTAINS(demo, "alltoall udma-bigdata"); + CHECK_CONTAINS(demo, "ERROR: bigdata alltoall UDMA registration failed"); + CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_bigdata_kernel"); + CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_bigdata"); + CHECK_CONTAINS(kernel, "readySignalOffset"); + CHECK_CONTAINS(kernel, "ackSignalOffset"); + CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "ackSignal"); +} + } // namespace int main() @@ -229,6 +251,7 @@ int main() TestDemoDebugLayoutSource(); TestAllToAllDataAsFlagSource(); TestAllToAllChunkedUdmaSource(); + TestAllToAllBigDataSource(); if (g_failures != 0) { std::cerr << g_failures << " all-to-all layout checks failed" << std::endl; return 1; From cbf2cb30f5a3190e31530f3457e9a4c0f0be3953 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sun, 28 Jun 2026 13:48:30 +0800 Subject: [PATCH 013/163] feat(udma): add bigdata alltoall demo kernel Add testType 7 for the UDMA all-to-all bigdata demo, including the host-side launch path, registered buffer layout, staged kernel flow, and profiling controls for stage timing. --- tests/udma/demo/run_tilexr_udma_demo.sh | 2 +- tests/udma/demo/tilexr_udma_alltoall_layout.h | 39 +- tests/udma/demo/tilexr_udma_demo.cpp | 341 ++++++++++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 530 ++++++++++++++++++ 4 files changed, 877 insertions(+), 35 deletions(-) diff --git a/tests/udma/demo/run_tilexr_udma_demo.sh b/tests/udma/demo/run_tilexr_udma_demo.sh index ae31f6a0..c6781137 100755 --- a/tests/udma/demo/run_tilexr_udma_demo.sh +++ b/tests/udma/demo/run_tilexr_udma_demo.sh @@ -36,7 +36,7 @@ echo "==========================================" echo " TileXR UDMA Communication Demo" echo "==========================================" echo "Binary: ${bin}" -echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all, 3=all-reduce)" +echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all, 3=all-reduce, 7=all-to-all bigdata UDMA)" echo "Rank size: ${rank_size}" echo "Elements/rank: ${elements_per_rank}" echo "NPU count: ${npu_count}" diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index 0a50439d..ad4fd52b 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -26,14 +26,19 @@ struct AllToAllChunkPlan { }; constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataControlSlotBytes = 64ULL; +constexpr uint32_t kAllToAllBigDataCopyCoreCount = 32U; struct AllToAllBigDataPlan { uint32_t passCount = 1; int32_t chunkElements = 0; size_t chunkBytesPerPeer = 0; size_t dataBytes = 0; + size_t readyPayloadOffset = 0; + size_t ackPayloadOffset = 0; size_t readySignalOffset = 0; size_t ackSignalOffset = 0; + size_t controlBytes = 0; size_t signalBytes = 0; size_t registeredBytes = 0; }; @@ -81,14 +86,16 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element } plan.registeredBytes = kAllToAllBigDataMaxRegisteredBytes; - plan.signalBytes = 2ULL * static_cast(rankSize) * sizeof(uint64_t); - if (plan.signalBytes >= plan.registeredBytes) { + plan.controlBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.signalBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + if (plan.controlBytes + plan.signalBytes >= plan.registeredBytes) { return plan; } - plan.dataBytes = plan.registeredBytes - plan.signalBytes; + const size_t dataSlotCount = static_cast(rankSize) * 2ULL; + plan.dataBytes = plan.registeredBytes - plan.controlBytes - plan.signalBytes; plan.chunkElements = static_cast( - plan.dataBytes / (static_cast(rankSize) * sizeof(int32_t))); + plan.dataBytes / (dataSlotCount * sizeof(int32_t))); if (plan.chunkElements <= 0) { plan.chunkElements = 1; } @@ -97,17 +104,31 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element } plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); - plan.dataBytes = static_cast(rankSize) * plan.chunkBytesPerPeer; - plan.readySignalOffset = plan.dataBytes; - plan.ackSignalOffset = plan.readySignalOffset + static_cast(rankSize) * sizeof(uint64_t); - plan.signalBytes = 2ULL * static_cast(rankSize) * sizeof(uint64_t); - plan.registeredBytes = plan.dataBytes + plan.signalBytes; + plan.dataBytes = dataSlotCount * plan.chunkBytesPerPeer; + plan.readyPayloadOffset = plan.dataBytes; + plan.ackPayloadOffset = + plan.readyPayloadOffset + static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.readySignalOffset = + plan.ackPayloadOffset + static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.ackSignalOffset = + plan.readySignalOffset + static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.controlBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.signalBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.registeredBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; plan.passCount = static_cast( (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1) / static_cast(plan.chunkElements)); return plan; } +inline uint32_t AllToAllBigDataBlockDim(int rankSize) +{ + if (rankSize <= 0) { + return kAllToAllBigDataCopyCoreCount; + } + return std::max(static_cast(rankSize), kAllToAllBigDataCopyCoreCount); +} + inline void FillAllToAllInput( std::vector& input, int rank, int rankSize, int32_t elementsPerPeer) { diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 43bad52c..af5a7dfd 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -41,6 +41,18 @@ extern void launch_tilexr_udma_p2p_latency( extern void launch_tilexr_datacopy_latency( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer, int32_t chunkElements); +extern void launch_tilexr_udma_all_to_all_fused( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR udmaMem, GM_ADDR signal, GM_ADDR debug, int32_t elementsPerPeer, + uint64_t udmaMemByteOffset, uint64_t signalByteOffsetBase, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount); +extern void launch_tilexr_udma_all_to_all_bigdata( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, + uint64_t dataOffset, uint64_t readyPayloadOffset, uint64_t ackPayloadOffset, + uint64_t readySignalOffset, uint64_t ackSignalOffset, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t tokenBase, + uint32_t profileStage); extern void launch_tilexr_all_to_all_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); extern void launch_tilexr_all_to_all_ipc_gather( @@ -69,12 +81,15 @@ constexpr int kDebugIpcScatter = kDebugUdmaStatusBase + TileXR::TILEXR_MAX_RANK_ constexpr int kDebugIpcGather = kDebugIpcScatter + 1; constexpr int kDebugAllReduceScatter = kDebugIpcGather + 1; constexpr int kDebugAllReduceSum = kDebugAllReduceScatter + 1; -constexpr size_t kDebugWords = kDebugAllReduceSum + 1; +constexpr int kDebugRecvSlotSampleBase = kDebugUdmaStatusBase + 160; +constexpr int kDebugReadySeenBase = kDebugUdmaStatusBase + 208; +constexpr size_t kDebugWords = kDebugReadySeenBase + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int kDefaultCommPort = 10067; constexpr int kDemoBarrierPortOffset = 97; constexpr size_t kUdmaRegistrationAlignment = 2 * 1024 * 1024; constexpr int kConnectRetryCount = 500; constexpr int kConnectRetrySleepMs = 10; +constexpr int kBigDataProfileStageFull = 8; struct BarrierEndpoint { uint16_t port; @@ -466,6 +481,12 @@ void PrintAllToAllUdmaDebug(int rank, int rankSize, const std::vector& constexpr int remoteBaseLowBase = kDebugUdmaStatusBase + 80; constexpr int memAddrLowBase = kDebugUdmaStatusBase + 96; constexpr int tpnBase = kDebugUdmaStatusBase + 112; + constexpr int sendSampleBase = kDebugUdmaStatusBase + 128; + constexpr int recvSampleBase = kDebugUdmaStatusBase + 144; + constexpr int recvSlotSampleBase = kDebugUdmaStatusBase + 160; + constexpr int remoteDataOffsetBase = kDebugUdmaStatusBase + 176; + constexpr int remoteReadyOffsetBase = kDebugUdmaStatusBase + 192; + constexpr int readySeenBase = kDebugUdmaStatusBase + 208; std::cout << "[rank " << rank << "] alltoall udma peer debug:"; for (int peer = 0; peer < rankSize && peer < 16; ++peer) { std::cout << " peer" << peer @@ -476,6 +497,12 @@ void PrintAllToAllUdmaDebug(int rank, int rankSize, const std::vector& << ",regLo=" << debug[remoteBaseLowBase + peer] << ",memLo=" << debug[memAddrLowBase + peer] << ",tpn=" << debug[tpnBase + peer] + << ",send0=" << debug[sendSampleBase + peer] + << ",recv0=" << debug[recvSampleBase + peer] + << ",slot0=" << debug[recvSlotSampleBase + peer] + << ",rDataOff=" << debug[remoteDataOffsetBase + peer] + << ",rReadyOff=" << debug[remoteReadyOffsetBase + peer] + << ",ready=" << debug[readySeenBase + peer] << "}"; } std::cout << std::endl; @@ -623,16 +650,27 @@ int main(int argc, char** argv) return 1; } - bool isAllToAll = testType == 2 || testType == 4 || testType == 5; + bool isAllToAll = testType == 2 || testType == 4 || testType == 5 || testType == 6 || testType == 7; bool isAllReduce = testType == 3; - bool strictAllToAllUdma = isAllToAll && (testType == 4 || GetEnvInt("TILEXR_DEMO_ALLTOALL_USE_UDMA", 0) != 0); + bool strictAllToAllUdma = + isAllToAll && (testType == 4 || testType == 7 || GetEnvInt("TILEXR_DEMO_ALLTOALL_USE_UDMA", 0) != 0); int allToAllRepeat = isAllToAll ? std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)) : 1; + int allToAllWarmup = isAllToAll ? std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)) : 0; + int bigDataProfileStage = testType == 7 ? + GetEnvInt("TILEXR_DEMO_BIGDATA_PROFILE_STAGE", kBigDataProfileStageFull) : + kBigDataProfileStageFull; + bigDataProfileStage = std::max(0, std::min(bigDataProfileStage, kBigDataProfileStageFull)); + bool bigDataProfilePartial = testType == 7 && bigDataProfileStage < kBigDataProfileStageFull; bool syncAllToAllAtEnd = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_SYNC_AT_END", 0) != 0; - bool useAllToAllPlainIpc = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_PLAIN_IPC", 0) != 0; - bool useAllToAllFusedIpc = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_FUSED_IPC", 0) != 0; + bool useAllToAllPlainIpc = + isAllToAll && testType != 7 && GetEnvInt("TILEXR_DEMO_ALLTOALL_PLAIN_IPC", 0) != 0; + bool useAllToAllFusedIpc = + isAllToAll && testType != 7 && GetEnvInt("TILEXR_DEMO_ALLTOALL_FUSED_IPC", 0) != 0; bool dumpAllToAllOnStrictFail = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_DUMP_ON_STRICT_FAIL", 0) != 0; - bool useAllToAllDataAsFlagIpc = isAllToAll && !strictAllToAllUdma && !useAllToAllPlainIpc && !useAllToAllFusedIpc; + bool useAllToAllDataAsFlagIpc = + isAllToAll && testType != 6 && testType != 7 && + !strictAllToAllUdma && !useAllToAllPlainIpc && !useAllToAllFusedIpc; const char* allToAllIpcFallbackLabel = useAllToAllFusedIpc ? "fused IPC" : (useAllToAllPlainIpc ? "plain IPC fallback" : "data-as-flag IPC fallback"); @@ -643,6 +681,9 @@ int main(int argc, char** argv) const TileXR::Demo::AllToAllChunkPlan chunkPlan = isAllToAll ? TileXR::Demo::PlanAllToAllUdmaChunks(rankSize, elementsPerRank) : TileXR::Demo::AllToAllChunkPlan {}; + const TileXR::Demo::AllToAllBigDataPlan bigDataPlan = + isAllToAll ? TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerRank) : + TileXR::Demo::AllToAllBigDataPlan {}; if (isAllToAll) { const size_t dataAsFlagStagingBytes = AllToAllDataAsFlagStagingBytes(rankSize, elementsPerRank); const size_t plainIpcStagingBytes = AllToAllPlainIpcStagingBytes(rankSize, elementsPerRank); @@ -668,23 +709,47 @@ int main(int argc, char** argv) } PrintStatus(rank, "alltoall repeat=" + std::to_string(allToAllRepeat) + " syncAtEnd=" + std::string(syncAllToAllAtEnd ? "true" : "false")); + if (testType == 7) { + PrintStatus(rank, "bigdata profile stage=" + std::to_string(bigDataProfileStage) + + " fullStage=" + std::to_string(kBigDataProfileStageFull)); + } PrintStatus(rank, "alltoall UDMA chunk plan: passCount=" + std::to_string(chunkPlan.passCount) + " chunkElements=" + std::to_string(chunkPlan.chunkElements) + " registeredBytes=" + std::to_string(chunkPlan.registeredBytes)); + if (testType == 7) { + PrintStatus(rank, "alltoall bigdata UDMA plan: passCount=" + std::to_string(bigDataPlan.passCount) + + " chunkElements=" + std::to_string(bigDataPlan.chunkElements) + + " dataBytes=" + std::to_string(bigDataPlan.dataBytes) + + " readyPayloadOffset=" + std::to_string(bigDataPlan.readyPayloadOffset) + + " ackPayloadOffset=" + std::to_string(bigDataPlan.ackPayloadOffset) + + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + + " registeredBytes=" + std::to_string(bigDataPlan.registeredBytes)); + } } size_t inputOffset = 0; - const size_t activeBytesPerRank = isAllToAll && strictAllToAllUdma ? chunkPlan.chunkBytesPerRank : dataBytes; + const size_t activeBytesPerRank = + isAllToAll && testType == 7 ? 0 : + isAllToAll && (strictAllToAllUdma || testType == 6) ? + chunkPlan.chunkBytesPerRank : dataBytes; size_t outputOffset = hasOutput ? activeBytesPerRank : 0; size_t signalBytes = static_cast(rankSize) * sizeof(uint64_t); - size_t signalOffset = hasOutput ? (outputOffset + activeBytesPerRank) : activeBytesPerRank; - size_t payloadBytes = signalOffset + signalBytes; + size_t signalOffset = testType == 7 ? bigDataPlan.readySignalOffset : + (hasOutput ? (outputOffset + activeBytesPerRank) : activeBytesPerRank); + size_t payloadBytes = testType == 7 ? bigDataPlan.registeredBytes : (signalOffset + signalBytes); size_t allocBytes = ((payloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; size_t registeredBytes = allocBytes; - if (isAllToAll && strictAllToAllUdma) { + if (isAllToAll && testType == 7) { + registeredBytes = ((bigDataPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / + kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + } else if (isAllToAll && (strictAllToAllUdma || testType == 6)) { registeredBytes = ((chunkPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; } + if (allocBytes < registeredBytes) { + allocBytes = registeredBytes; + } if (!CheckAcl(rank, "aclrtMalloc debug", aclrtMalloc(reinterpret_cast(&debug), kDebugWords * sizeof(int32_t), ACL_MEM_MALLOC_HUGE_FIRST)) || !CheckAcl(rank, "aclrtMalloc registered memory", aclrtMalloc(®isteredMemory, @@ -696,7 +761,8 @@ int main(int argc, char** argv) auto input = reinterpret_cast(static_cast(registeredMemory) + inputOffset); auto output = reinterpret_cast(static_cast(registeredMemory) + outputOffset); auto signals = reinterpret_cast(static_cast(registeredMemory) + signalOffset); - const bool chunkedStrictAllToAll = isAllToAll && strictAllToAllUdma && chunkPlan.passCount > 1; + const bool chunkedStrictAllToAll = + isAllToAll && testType != 7 && strictAllToAllUdma && chunkPlan.passCount > 1; if (useAllToAllFusedIpc) { PrintStatus(rank, "skip TileXRUDMARegister for alltoall fused IPC path"); forceAllToAllIpcFallback = true; @@ -708,6 +774,10 @@ int main(int argc, char** argv) forceAllToAllIpcFallback = true; } else if (chunkedStrictAllToAll) { PrintStatus(rank, "defer TileXRUDMARegister to per-pass registered output chunk"); + } else if (testType == 6) { + PrintStatus(rank, "defer TileXRUDMARegister to fused alltoall relay chunk"); + } else if (testType == 7) { + PrintStatus(rank, "defer TileXRUDMARegister to bigdata alltoall relay buffer"); } else { int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), registeredBytes, &udmaHandle); @@ -752,7 +822,7 @@ int main(int argc, char** argv) const char* inputName = isAllToAll ? "alltoall input" : (isAllReduce ? "allreduce input" : "data"); bool initOk = true; - if (!chunkedStrictAllToAll) { + if (!chunkedStrictAllToAll && testType != 7) { initOk = CopyHostToDevice(rank, input, dataCount * sizeof(int32_t), hostData.data(), dataCount * sizeof(int32_t), inputName); if (hasOutput) { @@ -782,9 +852,223 @@ int main(int argc, char** argv) return 1; } - if (isAllToAll) { - if (forceAllToAllIpcFallback) { - PrintStatus(rank, std::string("skip all-to-all UDMA kernel; use ") + allToAllIpcFallbackLabel); + if (testType == 7) { + void* bigInput = nullptr; + void* bigOutput = nullptr; + const size_t bigDataBytes = dataBytes; + if (!CheckAcl(rank, "aclrtMalloc bigdata input", + aclrtMalloc(&bigInput, bigDataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || + !CheckAcl(rank, "aclrtMalloc bigdata output", + aclrtMalloc(&bigOutput, bigDataBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!CopyHostToDevice(rank, bigInput, dataCount * sizeof(int32_t), + hostData.data(), dataCount * sizeof(int32_t), "bigdata alltoall input") || + !CopyHostToDevice(rank, bigOutput, dataCount * sizeof(int32_t), + hostOutput.data(), dataCount * sizeof(int32_t), "bigdata alltoall output")) { + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + std::vector zeroBigControl(bigDataPlan.controlBytes + bigDataPlan.signalBytes, 0); + if (!CopyHostToDevice(rank, static_cast(registeredMemory) + bigDataPlan.readyPayloadOffset, + bigDataPlan.controlBytes + bigDataPlan.signalBytes, + zeroBigControl.data(), zeroBigControl.size(), + "bigdata ready/ack payload+signals zero")) { + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!udmaRegistered) { + int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), + registeredBytes, &udmaHandle); + if (registerRet != TileXR::TILEXR_SUCCESS) { + std::cerr << "[rank " << rank << "] ERROR: bigdata alltoall UDMA registration failed" + << " ret=" << registerRet << " regBytes=" << registeredBytes << std::endl; + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + udmaRegistered = true; + } + PrintStatus(rank, "bigdata alltoall registered dataBytes=" + std::to_string(bigDataPlan.dataBytes) + + " readyPayloadOffset=" + std::to_string(bigDataPlan.readyPayloadOffset) + + " ackPayloadOffset=" + std::to_string(bigDataPlan.ackPayloadOffset) + + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + + " regBytes=" + std::to_string(registeredBytes) + + " passCount=" + std::to_string(bigDataPlan.passCount) + + " chunkElements=" + std::to_string(bigDataPlan.chunkElements) + + " repeat=" + std::to_string(allToAllRepeat) + + " profileStage=" + std::to_string(bigDataProfileStage)); + if (!CheckAcl(rank, "aclrtSynchronizeStream bigdata prime", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks bigdata prime")) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + + auto a2aStart = std::chrono::steady_clock::now(); + for (int iter = 0; iter < allToAllRepeat; ++iter) { + const uint64_t tokenBase = + static_cast(iter) * static_cast(bigDataPlan.passCount); + launch_tilexr_udma_all_to_all_bigdata( + TileXR::Demo::AllToAllBigDataBlockDim(rankSize), stream, commArgsDev, + reinterpret_cast(bigInput), reinterpret_cast(bigOutput), + reinterpret_cast(registeredMemory), reinterpret_cast(debug), + elementsPerRank, 0, bigDataPlan.readyPayloadOffset, bigDataPlan.ackPayloadOffset, + bigDataPlan.readySignalOffset, bigDataPlan.ackSignalOffset, + bigDataPlan.chunkElements, bigDataPlan.passCount, 1, tokenBase, + static_cast(bigDataProfileStage)); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream post-alltoall", aclrtSynchronizeStream(stream))) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + auto a2aEnd = std::chrono::steady_clock::now(); + double a2aMs = std::chrono::duration(a2aEnd - a2aStart).count(); + double a2aPerIterUs = (allToAllRepeat > 0) ? + (a2aMs * 1000.0 / static_cast(allToAllRepeat)) : 0.0; + double payload = static_cast(rankSize) * static_cast(elementsPerRank) * sizeof(int32_t); + double bwGbs = (a2aPerIterUs > 0.0) ? (payload / (a2aPerIterUs * 1e3)) : 0.0; + std::cout << "[rank " << rank << "] alltoall udma-bigdata " << allToAllRepeat + << " iters(total=" << bigDataPlan.passCount << " pass/iter) total=" << a2aMs + << " ms perIter=" << a2aPerIterUs + << " us payload=" << payload << " bytes bw=" << bwGbs << " GB/s" << std::endl; + + bool bigDataCopyBackOk = true; + if (!bigDataProfilePartial) { + bigDataCopyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), + bigOutput, dataCount * sizeof(int32_t), "bigdata alltoall output"); + } + bigDataCopyBackOk = CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), + debug, hostDebug.size() * sizeof(int32_t), "debug after bigdata alltoall") && bigDataCopyBackOk; + if (!bigDataCopyBackOk) { + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + } else if (testType == 6) { + // Forced UDMA alltoall (no IPC fallback). Single kernel launch loops + // REPEAT times internally; stream sync only after all loops. + // input/output: independent full-size GM, NOT registered. + // udmaMem+signals: registered chunk-sized relay, reused per pass. + void* fusedInput = nullptr; + void* fusedOutput = nullptr; + const size_t fusedDataBytes = dataBytes; + if (!CheckAcl(rank, "aclrtMalloc fused input", aclrtMalloc(&fusedInput, fusedDataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || + !CheckAcl(rank, "aclrtMalloc fused output", aclrtMalloc(&fusedOutput, fusedDataBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + aclrtFree(fusedInput); aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (!CopyHostToDevice(rank, fusedInput, dataCount * sizeof(int32_t), hostData.data(), dataCount * sizeof(int32_t), "fused alltoall input")) { + aclrtFree(fusedInput); aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + // Registered relay region: [udmaMem chunk | signals]. + const size_t fusedChunkBytes = chunkPlan.chunkBytesPerRank; + const size_t fusedSignalBytes = static_cast(rankSize) * sizeof(uint64_t); + const size_t fusedRegBytes = ((fusedChunkBytes + fusedSignalBytes + kUdmaRegistrationAlignment - 1) / + kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + if (!udmaRegistered) { + int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), fusedRegBytes, &udmaHandle); + if (registerRet != TileXR::TILEXR_SUCCESS) { + std::cerr << "[rank " << rank << "] ERROR: fused alltoall UDMA registration failed" + << " ret=" << registerRet << " regBytes=" << fusedRegBytes << std::endl; + aclrtFree(fusedInput); aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + udmaRegistered = true; + } + auto fusedUdmaMem = reinterpret_cast(registeredMemory); + auto fusedSignals = reinterpret_cast(static_cast(registeredMemory) + fusedChunkBytes); + const uint64_t fusedUdmaOffset = 0; + const uint64_t fusedSignalOffset = fusedChunkBytes; + std::vector zeroSignals(static_cast(rankSize), 0); + CopyHostToDevice(rank, fusedSignals, zeroSignals.size() * sizeof(uint64_t), zeroSignals.data(), zeroSignals.size() * sizeof(uint64_t), "fused signals zero"); + PrintStatus(rank, "fused alltoall registered chunkBytes=" + std::to_string(fusedChunkBytes) + + " regBytes=" + std::to_string(fusedRegBytes) + " passCount=" + std::to_string(chunkPlan.passCount) + + " chunkElements=" + std::to_string(chunkPlan.chunkElements) + " repeat=" + std::to_string(allToAllRepeat)); + if (!CheckAcl(rank, "aclrtSynchronizeStream fused prime", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks fused prime")) { + if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; } + aclrtFree(fusedInput); aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + // Warmup: separate launches (loop=1 each) with per-iter sync+barrier. + for (int witer = 0; witer < allToAllWarmup; ++witer) { + if (witer == 0) PrintStatus(rank, "warmup fused all-to-all warmup=" + std::to_string(allToAllWarmup)); + launch_tilexr_udma_all_to_all_fused( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(fusedInput), reinterpret_cast(fusedOutput), + reinterpret_cast(fusedUdmaMem), reinterpret_cast(fusedSignals), reinterpret_cast(debug), + elementsPerRank, fusedUdmaOffset, fusedSignalOffset, chunkPlan.chunkElements, chunkPlan.passCount, 1); + if (!CheckAcl(rank, "aclrtSynchronizeStream warmup", aclrtSynchronizeStream(stream)) || + !DemoBarrierAll(rank, rankSize, "all ranks completed warmup")) { + if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; } + aclrtFree(fusedInput); aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } + // Timed run: single launch, loop=REPEAT inside kernel, sync only at end. + auto a2aStart = std::chrono::steady_clock::now(); + launch_tilexr_udma_all_to_all_fused( + static_cast(rankSize), stream, commArgsDev, reinterpret_cast(fusedInput), reinterpret_cast(fusedOutput), + reinterpret_cast(fusedUdmaMem), reinterpret_cast(fusedSignals), reinterpret_cast(debug), + elementsPerRank, fusedUdmaOffset, fusedSignalOffset, chunkPlan.chunkElements, chunkPlan.passCount, static_cast(allToAllRepeat)); + if (!CheckAcl(rank, "aclrtSynchronizeStream post-alltoall", aclrtSynchronizeStream(stream))) { + if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; } + aclrtFree(fusedInput); aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + auto a2aEnd = std::chrono::steady_clock::now(); + { + double a2aMs = std::chrono::duration(a2aEnd - a2aStart).count(); + double a2aPerIterUs = (allToAllRepeat > 0) ? (a2aMs * 1000.0 / static_cast(allToAllRepeat)) : 0.0; + double payloadBytes = static_cast(rankSize) * static_cast(elementsPerRank) * sizeof(int32_t); + double bwGbs = (a2aPerIterUs > 0.0) ? (payloadBytes / (a2aPerIterUs * 1e3)) : 0.0; + std::cout << "[rank " << rank << "] alltoall udma-fused " << allToAllRepeat + << " iters(total=" << chunkPlan.passCount << " pass/iter) total=" << a2aMs << " ms perIter=" << a2aPerIterUs + << " us payload=" << payloadBytes << " bytes bw=" << bwGbs << " GB/s" << std::endl; + } + if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; } + aclrtFree(fusedInput); + aclrtFree(fusedOutput); + } else if (isAllToAll) { + if (forceAllToAllIpcFallback) { PrintStatus(rank, std::string("skip all-to-all UDMA kernel; use ") + allToAllIpcFallbackLabel); } else if (chunkedStrictAllToAll) { for (uint32_t pass = 0; pass < chunkPlan.passCount; ++pass) { const int32_t chunkOffset = static_cast(pass) * chunkPlan.chunkElements; @@ -962,7 +1246,7 @@ int main(int argc, char** argv) } } - if (isAllToAll && + if (isAllToAll && testType != 7 && !CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), debug, hostDebug.size() * sizeof(int32_t), "debug after alltoall udma")) { if (udmaRegistered) { @@ -973,7 +1257,7 @@ int main(int argc, char** argv) } bool usedIpcFallback = false; - bool allToAllUdmaComplete = !isAllToAll || AllToAllUdmaComplete(rankSize, hostDebug); + bool allToAllUdmaComplete = !isAllToAll || bigDataProfilePartial || AllToAllUdmaComplete(rankSize, hostDebug); if (isAllToAll && strictAllToAllUdma && !allToAllUdmaComplete) { std::cerr << "[rank " << rank << "] ERROR: strict alltoall UDMA CQ incomplete:"; for (int peer = 0; peer < rankSize; ++peer) { @@ -1131,20 +1415,23 @@ int main(int argc, char** argv) } bool copyBackOk = true; - if (!chunkedStrictAllToAll) { + if (!chunkedStrictAllToAll && testType != 7) { copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), data, dataCount * sizeof(int32_t), "data"); } - if (hasOutput && !chunkedStrictAllToAll) { + if (hasOutput && !chunkedStrictAllToAll && testType != 7) { const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; copyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), output, dataCount * sizeof(int32_t), outputName) && copyBackOk; } - if (!copyBackOk || - !CopyDeviceToHost(rank, hostSignals.data(), hostSignals.size() * sizeof(uint64_t), - signals, hostSignals.size() * sizeof(uint64_t), "signals") || - !CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), - debug, hostDebug.size() * sizeof(int32_t), "debug")) { + bool signalsCopyBackOk = true; + if (testType != 7) { + signalsCopyBackOk = CopyDeviceToHost(rank, hostSignals.data(), hostSignals.size() * sizeof(uint64_t), + signals, hostSignals.size() * sizeof(uint64_t), "signals"); + } + bool debugCopyBackOk = CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), + debug, hostDebug.size() * sizeof(int32_t), "debug"); + if (!copyBackOk || !signalsCopyBackOk || !debugCopyBackOk) { if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); } @@ -1172,7 +1459,11 @@ int main(int argc, char** argv) } bool ok = false; - if (isAllToAll) { + if (bigDataProfilePartial) { + PrintStatus(rank, "skip result validation for bigdata profile stage=" + + std::to_string(bigDataProfileStage)); + ok = true; + } else if (isAllToAll) { ok = ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerRank); } else if (isAllReduce) { ok = ValidateAllReduceData(rank, rankSize, hostOutput, elementsPerRank); diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index fc0907cb..61b99f97 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -26,11 +26,29 @@ constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE = TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 96; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_TPN_BASE = TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 112; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_SEND_SAMPLE_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 128; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 144; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_RECV_SLOT_SAMPLE_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 160; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_REMOTE_DATA_OFFSET_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 176; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_REMOTE_READY_OFFSET_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 192; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE = + TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 208; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER = TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER = TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER + 1; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER = TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER + 1; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SUM = TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER + 1; +constexpr uint64_t TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS = 100000000ULL; +constexpr uint64_t TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES = TileXR::TILEXR_UDMA_CACHE_LINE_SIZE; +constexpr int32_t TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS = -1001; +constexpr int32_t TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS = -1002; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_COPY_CORE_COUNT = 32U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES = 64 * 1024; namespace { @@ -45,6 +63,121 @@ __aicore__ inline uint64_t AllToAllDataAsFlagSegmentBytes(uint64_t payloadBytes) TileXR::DATA_AS_FLAG_BLOCK_BYTES; } +__aicore__ inline uint64_t WaitUdmaSignalToken(__gm__ uint64_t* signal, uint64_t token) +{ + uint64_t observed = ld_dev(signal, 0); + uint64_t polls = 0; + while (observed < token && polls < TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS) { + observed = ld_dev(signal, 0); + ++polls; + } + return observed; +} + +__aicore__ inline __gm__ uint64_t* ControlSlot(__gm__ uint8_t* base, uint64_t offset, int32_t peer) +{ + return reinterpret_cast<__gm__ uint64_t*>( + base + offset + static_cast(peer) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); +} + +__aicore__ inline uint32_t BigDataCopyCoreCount(int32_t blockNum) +{ + uint32_t cores = static_cast(blockNum); + if (cores > TILEXR_UDMA_DEMO_BIGDATA_COPY_CORE_COUNT) { + cores = TILEXR_UDMA_DEMO_BIGDATA_COPY_CORE_COUNT; + } + return cores == 0 ? 1U : cores; +} + +__aicore__ inline void BigDataCopyOneRelay( + __gm__ uint8_t* dst, __gm__ uint8_t* src, uint32_t bytes, + AscendC::LocalTensor relayLocal) +{ + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(src); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, bytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dst); + AscendC::DataCopyExtParams copyOut {1U, bytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, relayLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +__aicore__ inline void BigDataSendCopyShard( + int32_t blockIdx, int32_t blockNum, int32_t rank, int32_t rankSize, + __gm__ int32_t* input, __gm__ int32_t* output, __gm__ uint8_t* udmaMem, + int32_t elementsPerPeer, int32_t chunkOffset, uint32_t chunkBytes, + uint64_t sendDataOffset, uint64_t chunkBytesPerPeer, + AscendC::LocalTensor relayLocal) +{ + const uint32_t copyCores = BigDataCopyCoreCount(blockNum); + if (static_cast(blockIdx) >= copyCores || chunkBytes == 0 || rankSize <= 0) { + return; + } + const uint32_t chunksPerPeer = + (chunkBytes + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES - 1U) / + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint64_t totalUnits = static_cast(rankSize) * static_cast(chunksPerPeer); + for (uint64_t unit = static_cast(blockIdx); unit < totalUnits; unit += copyCores) { + const int32_t peer = static_cast(unit / static_cast(chunksPerPeer)); + const uint32_t chunkInPeer = static_cast(unit % static_cast(chunksPerPeer)); + const uint32_t offset = chunkInPeer * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint32_t bytes = + (chunkBytes - offset < TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES) ? + (chunkBytes - offset) : TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + + auto src = reinterpret_cast<__gm__ uint8_t*>( + input + static_cast(peer) * elementsPerPeer + chunkOffset) + offset; + __gm__ uint8_t* dst = nullptr; + if (peer == rank) { + dst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(rank) * elementsPerPeer + chunkOffset) + offset; + } else { + dst = udmaMem + sendDataOffset + static_cast(peer) * chunkBytesPerPeer + offset; + } + BigDataCopyOneRelay(dst, src, bytes, relayLocal); + } + AscendC::PipeBarrier(); +} + +__aicore__ inline void BigDataRelayCopyShard( + int32_t blockIdx, int32_t blockNum, int32_t rank, int32_t rankSize, + __gm__ int32_t* output, __gm__ uint8_t* udmaMem, + int32_t elementsPerPeer, int32_t chunkOffset, uint32_t chunkBytes, + uint64_t recvDataOffset, uint64_t chunkBytesPerPeer, + AscendC::LocalTensor relayLocal) +{ + const uint32_t copyCores = BigDataCopyCoreCount(blockNum); + if (static_cast(blockIdx) >= copyCores || chunkBytes == 0 || rankSize <= 1) { + return; + } + const uint32_t chunksPerPeer = + (chunkBytes + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES - 1U) / + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint64_t totalUnits = static_cast(rankSize - 1) * static_cast(chunksPerPeer); + for (uint64_t unit = static_cast(blockIdx); unit < totalUnits; unit += copyCores) { + const int32_t logicalPeer = static_cast(unit / static_cast(chunksPerPeer)); + const int32_t peer = logicalPeer < rank ? logicalPeer : logicalPeer + 1; + const uint32_t chunkInPeer = static_cast(unit % static_cast(chunksPerPeer)); + const uint32_t offset = chunkInPeer * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint32_t bytes = + (chunkBytes - offset < TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES) ? + (chunkBytes - offset) : TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + + auto src = udmaMem + recvDataOffset + static_cast(peer) * chunkBytesPerPeer + offset; + auto dst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(peer) * elementsPerPeer + chunkOffset) + offset; + BigDataCopyOneRelay(dst, src, bytes, relayLocal); + } + AscendC::PipeBarrier(); +} + } // namespace extern "C" __global__ __aicore__ void tilexr_udma_all_gather_kernel( @@ -945,3 +1078,400 @@ void launch_tilexr_datacopy_latency( commArgs, input, output, debug, elementsPerPeer, chunkElements); } + +// --------------------------------------------------------------------------- +// Fused alltoall kernel -- single launch, loop×pass inside, no IPC fallback. +// input : full-size GM (NOT registered), H2D primed once on host. +// output : full-size GM (NOT registered), filled inside kernel. +// udmaMem: registered chunk-sized relay buffer (rankSize*chunkElements*4B), +// reused every pass (overwrite) -- fits the 32MB URMA single-shot +// registration limit regardless of total data size. +// signals: registered region of rankSize uint64 slots. +// +// Per loop iter L, per pass p (chunkOffset = p*chunkElements): +// 1. SEND : block b (peer b!=rank) UDMAPutSignalNbi input[b-slice][offset] +// -> peer b's udmaMem[rank slot] + signal (L*passCount+p+1). +// SELF: DataCopyPad input[rank-slice][offset] -> output[rank-slice][offset]. +// 2. QUIET: UDMAQuietStatus per owned peer (local WQE completion). +// 3. WAIT : spin until each peer wrote this pass's signal into our +// signals[peer] (peer's data landed in our udmaMem[peer slot]). +// 4. RELAY: DataCopyPad udmaMem[peer slot] -> output[peer-slice][offset]. +// 5. SyncAll; next pass overwrites udmaMem. +// Signal tokens are globally monotonic (loop*passCount+pass+1) so no +// re-zeroing is needed across loops/passes. +// --------------------------------------------------------------------------- + +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_fused_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, + GM_ADDR signalGM, GM_ADDR debugGM, + int32_t elementsPerPeer, uint64_t udmaMemByteOffset, uint64_t signalByteOffsetBase, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto udmaMem = reinterpret_cast<__gm__ int32_t*>(udmaMemGM); + auto signals = reinterpret_cast<__gm__ uint64_t*>(signalGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + bool enabled = TileXR::UDMARegistryEnabled(args); + + const int32_t blockIdx = AscendC::GetBlockIdx(); + const int32_t blockNum = AscendC::GetBlockNum(); + + if (blockIdx == 0 && debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = enabled ? 1 : 0; + debug[4] = elementsPerPeer; + debug[5] = static_cast(udmaMemByteOffset); + } + if (!enabled) { + for (uint32_t l = 0; l < loopCount; ++l) { AscendC::SyncAll(); } + return; + } + + const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; + const uint64_t peerPayloadBytes = AllToAllPayloadBytes(effectiveChunkElements); + const uint32_t peerBytes = static_cast(peerPayloadBytes); + constexpr uint32_t RELAY_UB_BYTES = 64 * 1024; + + for (uint32_t loop = 0; loop < loopCount; ++loop) { + for (uint32_t pass = 0; pass < passCount; ++pass) { + const int32_t chunkOffset = static_cast(pass) * effectiveChunkElements; + const int32_t chunkElem = (elementsPerPeer - chunkOffset < effectiveChunkElements) + ? (elementsPerPeer - chunkOffset) : effectiveChunkElements; + const uint32_t chunkBytes = static_cast( + static_cast(chunkElem) * sizeof(int32_t)); + const uint64_t expectedSignal = + static_cast(loop) * static_cast(passCount) + + static_cast(pass) + 1; + + // ---- 1. SEND + SELF ---- + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + // self: input -> output directly (no P2P, no relay). + auto selfSrc = input + static_cast(rank) * elementsPerPeer + chunkOffset; + auto selfDst = output + static_cast(rank) * elementsPerPeer + chunkOffset; + AscendC::TPipe pipe; + AscendC::TBuf relayTBuf; + pipe.InitBuffer(relayTBuf, RELAY_UB_BYTES); + AscendC::LocalTensor relayLocal = relayTBuf.Get(); + auto srcBytes = reinterpret_cast<__gm__ uint8_t*>(selfSrc); + auto dstBytes = reinterpret_cast<__gm__ uint8_t*>(selfDst); + for (uint32_t off = 0; off < chunkBytes; off += RELAY_UB_BYTES) { + uint32_t cb = (chunkBytes - off < RELAY_UB_BYTES) ? (chunkBytes - off) : RELAY_UB_BYTES; + AscendC::GlobalTensor srcGlobal; srcGlobal.SetGlobalBuffer(srcBytes + off); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, cb, 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::GlobalTensor dstGlobal; dstGlobal.SetGlobalBuffer(dstBytes + off); + AscendC::DataCopyExtParams copyOut {1U, cb, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, relayLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); + continue; + } + // P2P: read input[peer slice][offset] -> peer's udmaMem[rank slot] + signal. + auto localSrc = input + static_cast(peer) * elementsPerPeer + chunkOffset; + uint64_t remoteOffset = udmaMemByteOffset + + static_cast(rank) * peerPayloadBytes; + uint64_t remoteSignalOffset = signalByteOffsetBase + + static_cast(rank) * sizeof(uint64_t); + TileXR::UDMAPutSignalNbi(args, peer, localSrc, + remoteOffset, chunkBytes, remoteSignalOffset, expectedSignal); + } + + // ---- 2. QUIET ---- + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) continue; + uint32_t status = TileXR::UDMAQuietStatus(args, peer); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } + } + + // ---- 3. WAIT : each block spins on its owned peer's signal slot. ---- + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) continue; + __gm__ uint64_t* slot = signals + static_cast(peer); + while (*slot != expectedSignal) { + } + } + AscendC::SyncAll(); + + // ---- 4. RELAY : udmaMem[peer slot] -> output[peer slice][offset]. ---- + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) continue; + auto relaySrc = udmaMem + static_cast(peer) * effectiveChunkElements; + auto relayDst = output + static_cast(peer) * elementsPerPeer + chunkOffset; + AscendC::TPipe pipe; + AscendC::TBuf relayTBuf; + pipe.InitBuffer(relayTBuf, RELAY_UB_BYTES); + AscendC::LocalTensor relayLocal = relayTBuf.Get(); + auto srcBytes = reinterpret_cast<__gm__ uint8_t*>(relaySrc); + auto dstBytes = reinterpret_cast<__gm__ uint8_t*>(relayDst); + for (uint32_t off = 0; off < chunkBytes; off += RELAY_UB_BYTES) { + uint32_t cb = (chunkBytes - off < RELAY_UB_BYTES) ? (chunkBytes - off) : RELAY_UB_BYTES; + AscendC::GlobalTensor srcGlobal; srcGlobal.SetGlobalBuffer(srcBytes + off); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, cb, 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::GlobalTensor dstGlobal; dstGlobal.SetGlobalBuffer(dstBytes + off); + AscendC::DataCopyExtParams copyOut {1U, cb, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, relayLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); + } + AscendC::SyncAll(); + } + } + AscendC::SyncAll(); +} + +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, GM_ADDR debugGM, + int32_t elementsPerPeer, uint64_t dataOffset, uint64_t readyPayloadOffset, uint64_t ackPayloadOffset, + uint64_t readySignalOffset, uint64_t ackSignalOffset, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t tokenBase, + uint32_t profileStage) +{ + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_PREPARE = 0; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY = 1; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC = 2; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT = 3; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY = 4; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY = 5; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT = 6; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK = 7; + constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_FULL = 8; + if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_FULL) { + profileStage = TILEXR_BIGDATA_PROFILE_STAGE_FULL; + } + + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto udmaMem = reinterpret_cast<__gm__ uint8_t*>(udmaMemGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + + int32_t rank = args->rank; + int32_t rankSize = args->rankSize; + bool enabled = TileXR::UDMARegistryEnabled(args); + const int32_t blockIdx = AscendC::GetBlockIdx(); + const int32_t blockNum = AscendC::GetBlockNum(); + + if (blockIdx == 0 && debug != nullptr) { + debug[0] = TILEXR_UDMA_DEMO_MAGIC; + debug[1] = rank; + debug[2] = rankSize; + debug[3] = enabled ? 1 : 0; + debug[4] = elementsPerPeer; + debug[5] = chunkElements; + } + if (!enabled) { + for (uint32_t l = 0; l < loopCount; ++l) { + AscendC::SyncAll(); + } + return; + } + + const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; + const uint64_t chunkBytesPerPeer = static_cast(effectiveChunkElements) * sizeof(int32_t); + const uint64_t sendDataOffset = dataOffset; + const uint64_t recvDataOffset = dataOffset + static_cast(rankSize) * chunkBytesPerPeer; + + AscendC::TPipe pipe; + AscendC::TBuf relayTBuf; + pipe.InitBuffer(relayTBuf, TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES); + AscendC::LocalTensor relayLocal = relayTBuf.Get(); + + for (uint32_t loop = 0; loop < loopCount; ++loop) { + for (uint32_t pass = 0; pass < passCount; ++pass) { + const int32_t chunkOffset = static_cast(pass) * effectiveChunkElements; + const int32_t remaining = elementsPerPeer - chunkOffset; + const int32_t chunkElem = remaining < effectiveChunkElements ? remaining : effectiveChunkElements; + if (chunkElem <= 0) { + continue; + } + const uint32_t chunkBytes = static_cast(static_cast(chunkElem) * sizeof(int32_t)); + const uint64_t token = + tokenBase + static_cast(loop) * static_cast(passCount) + + static_cast(pass) + 1ULL; + + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE) { + continue; + } + + BigDataSendCopyShard(blockIdx, blockNum, rank, rankSize, input, output, udmaMem, + elementsPerPeer, chunkOffset, chunkBytes, sendDataOffset, chunkBytesPerPeer, relayLocal); + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { + continue; + } + AscendC::SyncAll(); + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { + continue; + } + + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + continue; + } + auto sendSlot = udmaMem + sendDataOffset + static_cast(peer) * chunkBytesPerPeer; + auto localSrc = reinterpret_cast<__gm__ int32_t*>(sendSlot); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_SEND_SAMPLE_BASE + peer] = localSrc[0]; + } + const uint64_t remoteDataOffset = + recvDataOffset + static_cast(rank) * chunkBytesPerPeer; + const uint64_t remoteReadyOffset = + readySignalOffset + static_cast(rank) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + auto registry = TileXR::GetUDMARegistry(args); + auto udmaInfo = TileXR::GetUDMAInfo(args); + auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); + auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); + bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteDataOffset, chunkBytes); + uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_RANGE_VALID_BASE + peer] = rangeValid ? 1 : 0; + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE + peer] = static_cast(wqeBefore); + debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = static_cast(wqCtx->localTokenId); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = + static_cast(reinterpret_cast(registry->regions[peer].base) & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = + static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_TPN_BASE + peer] = static_cast(remoteMemInfo->tpn); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_DATA_OFFSET_BASE + peer] = + static_cast(remoteDataOffset); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_READY_OFFSET_BASE + peer] = + static_cast(remoteReadyOffset); + } + TileXR::UDMAPutNbi(args, peer, localSrc, remoteDataOffset, chunkBytes); + uint32_t status = TileXR::UDMAQuietStatus(args, peer); + auto readyPayload = ControlSlot(udmaMem, readyPayloadOffset, peer); + *readyPayload = token; + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t*>(readyPayload), sizeof(uint64_t)); + TileXR::UDMAPutNbi(args, peer, readyPayload, remoteReadyOffset, sizeof(uint64_t)); + uint32_t readyStatus = TileXR::UDMAQuietStatus(args, peer); + if (status == 0) { + status = readyStatus; + } + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + uint32_t wqeAfter = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE + peer] = static_cast(wqeAfter); + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { + continue; + } + + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + continue; + } + uint64_t observed = WaitUdmaSignalToken(ControlSlot(udmaMem, readySignalOffset, peer), token); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = + static_cast(observed); + if (observed < token) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + } + } + AscendC::SyncAll(); + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY) { + continue; + } + + if (blockIdx < rankSize && blockIdx != rank && debug != nullptr && loop == 0 && pass == 0 && blockIdx < 16) { + auto recvSlot = udmaMem + recvDataOffset + static_cast(blockIdx) * chunkBytesPerPeer; + auto recvSlotInt = reinterpret_cast<__gm__ int32_t*>(recvSlot); + debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SLOT_SAMPLE_BASE + blockIdx] = recvSlotInt[0]; + } + BigDataRelayCopyShard(blockIdx, blockNum, rank, rankSize, output, udmaMem, + elementsPerPeer, chunkOffset, chunkBytes, recvDataOffset, chunkBytesPerPeer, relayLocal); + AscendC::SyncAll(); + if (blockIdx < rankSize && blockIdx != rank && debug != nullptr && loop == 0 && pass == 0 && blockIdx < 16) { + auto relayDst = output + static_cast(blockIdx) * elementsPerPeer + chunkOffset; + debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE + blockIdx] = relayDst[0]; + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY) { + continue; + } + + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + continue; + } + auto ackPayload = ControlSlot(udmaMem, ackPayloadOffset, peer); + *ackPayload = token; + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t*>(ackPayload), sizeof(uint64_t)); + const uint64_t remoteAckOffset = + ackSignalOffset + static_cast(rank) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + TileXR::UDMAPutNbi(args, peer, ackPayload, remoteAckOffset, sizeof(uint64_t)); + (void)TileXR::UDMAQuietStatus(args, peer); + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT) { + continue; + } + + for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { + if (peer == rank) { + continue; + } + uint64_t observed = WaitUdmaSignalToken(ControlSlot(udmaMem, ackSignalOffset, peer), token); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16 && observed < token) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + } + AscendC::SyncAll(); + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK) { + continue; + } + } + } + AscendC::SyncAll(); +} + +void launch_tilexr_udma_all_to_all_bigdata( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, + uint64_t dataOffset, uint64_t readyPayloadOffset, uint64_t ackPayloadOffset, + uint64_t readySignalOffset, uint64_t ackSignalOffset, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t tokenBase, + uint32_t profileStage) +{ + tilexr_udma_all_to_all_bigdata_kernel<<>>( + commArgs, input, output, udmaMem, debug, elementsPerPeer, + dataOffset, readyPayloadOffset, ackPayloadOffset, readySignalOffset, ackSignalOffset, + chunkElements, passCount, loopCount, tokenBase, profileStage); +} + +void launch_tilexr_udma_all_to_all_fused( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR udmaMem, GM_ADDR signal, GM_ADDR debug, int32_t elementsPerPeer, + uint64_t udmaMemByteOffset, uint64_t signalByteOffsetBase, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount) +{ + tilexr_udma_all_to_all_fused_kernel<<>>( + commArgs, input, output, udmaMem, signal, debug, elementsPerPeer, + udmaMemByteOffset, signalByteOffsetBase, chunkElements, passCount, loopCount); +} From ead95b08a19e8719d7e5e73a542f258d97067a6a Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sun, 28 Jun 2026 19:49:33 +0800 Subject: [PATCH 014/163] perf(udma): pipeline bigdata alltoall relay copies Use three cores per peer for copy/send/recv roles in the bigdata all-to-all demo. Add ping-pong relay buffers for local GM-UB-GM copies, MTE flag synchronization for copy/ack control, and explicit kernel-loop token generation for UDMA buffer reuse across host-side repeated launches. Update the bigdata layout plan and source guards to cover the new control layout, profile stages, and UDMA cache clean fixes. --- src/comm/tilexr_comm.cpp | 8 +- src/include/tilexr_udma.h | 6 +- tests/udma/demo/tilexr_udma_alltoall_layout.h | 35 +- tests/udma/demo/tilexr_udma_demo.cpp | 21 +- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 767 ++++++++++++------ .../unit/test_tilexr_udma_alltoall_layout.cpp | 98 ++- 6 files changed, 664 insertions(+), 271 deletions(-) diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 9a342457..48ad250c 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -1009,6 +1009,13 @@ TileXRComm::~TileXRComm() socketExchange_ = nullptr; } FreePeerMem(commArgs_.dumpAddr); + if (udmaRegisteredPtr_ != nullptr && udmaTransport_ != nullptr) { + int ret = udmaTransport_->UnregisterMemory(udmaRegisteredPtr_); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR UDMA memory unregistration failed in destructor: " << ret; + } + udmaRegisteredPtr_ = nullptr; + } FreeUDMARegistry(); FreePeerMem(peerMem_[rank_]); FreePeerMem(commArgsPtr_); @@ -1017,7 +1024,6 @@ TileXRComm::~TileXRComm() udmaTransport_->Shutdown(); udmaTransport_.reset(); } - udmaRegisteredPtr_ = nullptr; udmaInfoDev_ = nullptr; ResetSDMAState(); } diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index c931e804..0f2c853f 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -79,10 +79,14 @@ __aicore__ inline __gm__ uint8_t* UDMARegisteredRemoteAddr( __aicore__ inline void UDMACleanCacheLines(__gm__ uint8_t* addr, uint64_t length) { + if (length == 0) { + return; + } __gm__ uint8_t* start = reinterpret_cast<__gm__ uint8_t*>( reinterpret_cast(addr) / TILEXR_UDMA_CACHE_LINE_SIZE * TILEXR_UDMA_CACHE_LINE_SIZE); __gm__ uint8_t* end = reinterpret_cast<__gm__ uint8_t*>( - (reinterpret_cast(addr) + length) / TILEXR_UDMA_CACHE_LINE_SIZE * TILEXR_UDMA_CACHE_LINE_SIZE); + (reinterpret_cast(addr) + length - 1) / + TILEXR_UDMA_CACHE_LINE_SIZE * TILEXR_UDMA_CACHE_LINE_SIZE); AscendC::GlobalTensor global; global.SetGlobalBuffer(start); for (uint64_t i = 0; i <= static_cast(end - start); i += TILEXR_UDMA_CACHE_LINE_SIZE) { diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index ad4fd52b..2a687e7f 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -27,15 +27,15 @@ struct AllToAllChunkPlan { constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataControlSlotBytes = 64ULL; -constexpr uint32_t kAllToAllBigDataCopyCoreCount = 32U; +constexpr uint32_t kAllToAllBigDataCoresPerPeer = 3U; +constexpr uint32_t kAllToAllBigDataPingPongSlots = 2U; struct AllToAllBigDataPlan { uint32_t passCount = 1; int32_t chunkElements = 0; size_t chunkBytesPerPeer = 0; size_t dataBytes = 0; - size_t readyPayloadOffset = 0; - size_t ackPayloadOffset = 0; + size_t copyDoneOffset = 0; size_t readySignalOffset = 0; size_t ackSignalOffset = 0; size_t controlBytes = 0; @@ -86,13 +86,18 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element } plan.registeredBytes = kAllToAllBigDataMaxRegisteredBytes; - plan.controlBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; - plan.signalBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + const size_t controlGroupBytes = + static_cast(kAllToAllBigDataPingPongSlots) * static_cast(rankSize) * + kAllToAllBigDataControlSlotBytes; + plan.controlBytes = controlGroupBytes; + plan.signalBytes = 2ULL * controlGroupBytes; if (plan.controlBytes + plan.signalBytes >= plan.registeredBytes) { return plan; } - const size_t dataSlotCount = static_cast(rankSize) * 2ULL; + const size_t networkPeerCount = static_cast(rankSize > 1 ? rankSize - 1 : 1); + const size_t dataSlotCount = + networkPeerCount * static_cast(kAllToAllBigDataPingPongSlots) * 2ULL; plan.dataBytes = plan.registeredBytes - plan.controlBytes - plan.signalBytes; plan.chunkElements = static_cast( plan.dataBytes / (dataSlotCount * sizeof(int32_t))); @@ -105,15 +110,11 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); plan.dataBytes = dataSlotCount * plan.chunkBytesPerPeer; - plan.readyPayloadOffset = plan.dataBytes; - plan.ackPayloadOffset = - plan.readyPayloadOffset + static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; - plan.readySignalOffset = - plan.ackPayloadOffset + static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; - plan.ackSignalOffset = - plan.readySignalOffset + static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; - plan.controlBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; - plan.signalBytes = 2ULL * static_cast(rankSize) * kAllToAllBigDataControlSlotBytes; + plan.copyDoneOffset = plan.dataBytes; + plan.readySignalOffset = plan.copyDoneOffset + controlGroupBytes; + plan.ackSignalOffset = plan.readySignalOffset + controlGroupBytes; + plan.controlBytes = controlGroupBytes; + plan.signalBytes = 2ULL * controlGroupBytes; plan.registeredBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; plan.passCount = static_cast( (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1) / @@ -124,9 +125,9 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element inline uint32_t AllToAllBigDataBlockDim(int rankSize) { if (rankSize <= 0) { - return kAllToAllBigDataCopyCoreCount; + return 1U; } - return std::max(static_cast(rankSize), kAllToAllBigDataCopyCoreCount); + return static_cast(rankSize) * kAllToAllBigDataCoresPerPeer; } inline void FillAllToAllInput( diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index af5a7dfd..edcf865b 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -49,9 +49,9 @@ extern void launch_tilexr_udma_all_to_all_fused( extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, - uint64_t dataOffset, uint64_t readyPayloadOffset, uint64_t ackPayloadOffset, + uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t tokenBase, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage); extern void launch_tilexr_all_to_all_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); @@ -720,8 +720,7 @@ int main(int argc, char** argv) PrintStatus(rank, "alltoall bigdata UDMA plan: passCount=" + std::to_string(bigDataPlan.passCount) + " chunkElements=" + std::to_string(bigDataPlan.chunkElements) + " dataBytes=" + std::to_string(bigDataPlan.dataBytes) + - " readyPayloadOffset=" + std::to_string(bigDataPlan.readyPayloadOffset) + - " ackPayloadOffset=" + std::to_string(bigDataPlan.ackPayloadOffset) + + " copyDoneOffset=" + std::to_string(bigDataPlan.copyDoneOffset) + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + " registeredBytes=" + std::to_string(bigDataPlan.registeredBytes)); @@ -875,10 +874,10 @@ int main(int argc, char** argv) return 1; } std::vector zeroBigControl(bigDataPlan.controlBytes + bigDataPlan.signalBytes, 0); - if (!CopyHostToDevice(rank, static_cast(registeredMemory) + bigDataPlan.readyPayloadOffset, + if (!CopyHostToDevice(rank, static_cast(registeredMemory) + bigDataPlan.copyDoneOffset, bigDataPlan.controlBytes + bigDataPlan.signalBytes, zeroBigControl.data(), zeroBigControl.size(), - "bigdata ready/ack payload+signals zero")) { + "bigdata copy/ready/ack payload+signals zero")) { aclrtFree(bigInput); aclrtFree(bigOutput); Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); @@ -898,8 +897,7 @@ int main(int argc, char** argv) udmaRegistered = true; } PrintStatus(rank, "bigdata alltoall registered dataBytes=" + std::to_string(bigDataPlan.dataBytes) + - " readyPayloadOffset=" + std::to_string(bigDataPlan.readyPayloadOffset) + - " ackPayloadOffset=" + std::to_string(bigDataPlan.ackPayloadOffset) + + " copyDoneOffset=" + std::to_string(bigDataPlan.copyDoneOffset) + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + " regBytes=" + std::to_string(registeredBytes) + @@ -921,15 +919,14 @@ int main(int argc, char** argv) auto a2aStart = std::chrono::steady_clock::now(); for (int iter = 0; iter < allToAllRepeat; ++iter) { - const uint64_t tokenBase = - static_cast(iter) * static_cast(bigDataPlan.passCount); + const uint64_t kernelLoopBase = static_cast(iter); launch_tilexr_udma_all_to_all_bigdata( TileXR::Demo::AllToAllBigDataBlockDim(rankSize), stream, commArgsDev, reinterpret_cast(bigInput), reinterpret_cast(bigOutput), reinterpret_cast(registeredMemory), reinterpret_cast(debug), - elementsPerRank, 0, bigDataPlan.readyPayloadOffset, bigDataPlan.ackPayloadOffset, + elementsPerRank, 0, bigDataPlan.copyDoneOffset, bigDataPlan.readySignalOffset, bigDataPlan.ackSignalOffset, - bigDataPlan.chunkElements, bigDataPlan.passCount, 1, tokenBase, + bigDataPlan.chunkElements, bigDataPlan.passCount, 1, kernelLoopBase, static_cast(bigDataProfileStage)); } if (!CheckAcl(rank, "aclrtSynchronizeStream post-alltoall", aclrtSynchronizeStream(stream))) { diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 61b99f97..075b0c9c 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -47,8 +47,21 @@ constexpr uint64_t TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS = 100000000ULL; constexpr uint64_t TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES = TileXR::TILEXR_UDMA_CACHE_LINE_SIZE; constexpr int32_t TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS = -1001; constexpr int32_t TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS = -1002; -constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_COPY_CORE_COUNT = 32U; +constexpr int32_t TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS = -1003; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES = 64 * 1024; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES = + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES * 2U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 3U; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_PREPARE = 0; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY = 1; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC = 2; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT = 3; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY = 4; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY = 5; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT = 6; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK = 7; +constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_FULL = 8; namespace { @@ -63,30 +76,40 @@ __aicore__ inline uint64_t AllToAllDataAsFlagSegmentBytes(uint64_t payloadBytes) TileXR::DATA_AS_FLAG_BLOCK_BYTES; } -__aicore__ inline uint64_t WaitUdmaSignalToken(__gm__ uint64_t* signal, uint64_t token) +__aicore__ inline __gm__ uint64_t* ControlSlot(__gm__ uint8_t* base, uint64_t offset, int32_t peer) { - uint64_t observed = ld_dev(signal, 0); - uint64_t polls = 0; - while (observed < token && polls < TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS) { - observed = ld_dev(signal, 0); - ++polls; - } - return observed; + return reinterpret_cast<__gm__ uint64_t*>( + base + offset + static_cast(peer) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); } -__aicore__ inline __gm__ uint64_t* ControlSlot(__gm__ uint8_t* base, uint64_t offset, int32_t peer) +__aicore__ inline __gm__ uint64_t* BigDataControlSlot( + __gm__ uint8_t* base, uint64_t offset, uint32_t slot, int32_t rankSize, int32_t peer) { return reinterpret_cast<__gm__ uint64_t*>( - base + offset + static_cast(peer) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); + base + offset + + (static_cast(slot) * static_cast(rankSize) + + static_cast(peer)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); } -__aicore__ inline uint32_t BigDataCopyCoreCount(int32_t blockNum) +__aicore__ inline void BigDataCopyInTile( + __gm__ uint8_t* src, uint32_t offset, uint32_t bytes, + AscendC::LocalTensor local) { - uint32_t cores = static_cast(blockNum); - if (cores > TILEXR_UDMA_DEMO_BIGDATA_COPY_CORE_COUNT) { - cores = TILEXR_UDMA_DEMO_BIGDATA_COPY_CORE_COUNT; - } - return cores == 0 ? 1U : cores; + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(src + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, bytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(local, srcGlobal, copyIn, padIn); +} + +__aicore__ inline void BigDataCopyOutTile( + __gm__ uint8_t* dst, uint32_t offset, uint32_t bytes, + AscendC::LocalTensor local) +{ + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dst + offset); + AscendC::DataCopyExtParams copyOut {1U, bytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, local, copyOut); } __aicore__ inline void BigDataCopyOneRelay( @@ -109,73 +132,485 @@ __aicore__ inline void BigDataCopyOneRelay( AscendC::WaitFlag(EVENT_ID0); } -__aicore__ inline void BigDataSendCopyShard( - int32_t blockIdx, int32_t blockNum, int32_t rank, int32_t rankSize, - __gm__ int32_t* input, __gm__ int32_t* output, __gm__ uint8_t* udmaMem, - int32_t elementsPerPeer, int32_t chunkOffset, uint32_t chunkBytes, - uint64_t sendDataOffset, uint64_t chunkBytesPerPeer, +__aicore__ inline void BigDataWaitMte2ToMte3(uint32_t bufferId) +{ + if (bufferId == 0U) { + AscendC::WaitFlag(EVENT_ID0); + } else { + AscendC::WaitFlag(EVENT_ID1); + } +} + +__aicore__ inline void BigDataSetMte2ToMte3(uint32_t bufferId) +{ + if (bufferId == 0U) { + AscendC::SetFlag(EVENT_ID0); + } else { + AscendC::SetFlag(EVENT_ID1); + } +} + +__aicore__ inline void BigDataWaitMte3ToMte2(uint32_t bufferId) +{ + if (bufferId == 0U) { + AscendC::WaitFlag(EVENT_ID0); + } else { + AscendC::WaitFlag(EVENT_ID1); + } +} + +__aicore__ inline void BigDataSetMte3ToMte2(uint32_t bufferId) +{ + if (bufferId == 0U) { + AscendC::SetFlag(EVENT_ID0); + } else { + AscendC::SetFlag(EVENT_ID1); + } +} + +__aicore__ inline uint32_t BigDataTileBytes(uint32_t totalBytes, uint32_t offset) +{ + const uint32_t remain = totalBytes - offset; + return remain < TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES ? + remain : TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; +} + +__aicore__ inline void BigDataCopyRangePingPong( + __gm__ uint8_t* dst, __gm__ uint8_t* src, uint32_t bytes, AscendC::LocalTensor relayLocal) { - const uint32_t copyCores = BigDataCopyCoreCount(blockNum); - if (static_cast(blockIdx) >= copyCores || chunkBytes == 0 || rankSize <= 0) { + if (bytes == 0U) { return; } - const uint32_t chunksPerPeer = - (chunkBytes + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES - 1U) / + const uint32_t tileCount = + (bytes + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES - 1U) / TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; - const uint64_t totalUnits = static_cast(rankSize) * static_cast(chunksPerPeer); - for (uint64_t unit = static_cast(blockIdx); unit < totalUnits; unit += copyCores) { - const int32_t peer = static_cast(unit / static_cast(chunksPerPeer)); - const uint32_t chunkInPeer = static_cast(unit % static_cast(chunksPerPeer)); - const uint32_t offset = chunkInPeer * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; - const uint32_t bytes = - (chunkBytes - offset < TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES) ? - (chunkBytes - offset) : TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + if (tileCount <= 2U) { + for (uint32_t offset = 0; offset < bytes; offset += TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES) { + const uint32_t tileBytes = BigDataTileBytes(bytes, offset); + BigDataCopyOneRelay(dst + offset, src + offset, tileBytes, relayLocal); + } + return; + } + bool copyOutInFlight0 = false; + bool copyOutInFlight1 = false; + + for (uint32_t tile = 0; tile < tileCount; ++tile) { + const uint32_t bufferId = tile & 1U; + if (tile >= 2U) { + BigDataWaitMte3ToMte2(bufferId); + if (bufferId == 0U) { + copyOutInFlight0 = false; + } else { + copyOutInFlight1 = false; + } + } + + const uint32_t offset = tile * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint32_t tileBytes = BigDataTileBytes(bytes, offset); + AscendC::LocalTensor local = + relayLocal[bufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]; + BigDataCopyInTile(src, offset, tileBytes, local); + BigDataSetMte2ToMte3(bufferId); + + if (tile > 0U) { + const uint32_t prevBufferId = bufferId ^ 1U; + const uint32_t prevOffset = offset - TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint32_t prevBytes = BigDataTileBytes(bytes, prevOffset); + AscendC::LocalTensor prevLocal = + relayLocal[prevBufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]; + BigDataWaitMte2ToMte3(prevBufferId); + BigDataCopyOutTile(dst, prevOffset, prevBytes, prevLocal); + BigDataSetMte3ToMte2(prevBufferId); + if (prevBufferId == 0U) { + copyOutInFlight0 = true; + } else { + copyOutInFlight1 = true; + } + } + } + + const uint32_t lastTile = tileCount - 1U; + const uint32_t lastBufferId = lastTile & 1U; + const uint32_t lastOffset = lastTile * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; + const uint32_t lastBytes = BigDataTileBytes(bytes, lastOffset); + AscendC::LocalTensor lastLocal = + relayLocal[lastBufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]; + BigDataWaitMte2ToMte3(lastBufferId); + BigDataCopyOutTile(dst, lastOffset, lastBytes, lastLocal); + BigDataSetMte3ToMte2(lastBufferId); + if (lastBufferId == 0U) { + copyOutInFlight0 = true; + } else { + copyOutInFlight1 = true; + } + if (copyOutInFlight0) { + BigDataWaitMte3ToMte2(0U); + } + if (copyOutInFlight1) { + BigDataWaitMte3ToMte2(1U); + } +} + +__aicore__ inline bool BigDataPassChunk( + uint32_t pass, int32_t elementsPerPeer, int32_t effectiveChunkElements, + int32_t& chunkOffset, uint32_t& chunkBytes) +{ + chunkOffset = static_cast(pass) * effectiveChunkElements; + const int32_t remaining = elementsPerPeer - chunkOffset; + const int32_t chunkElem = remaining < effectiveChunkElements ? remaining : effectiveChunkElements; + if (chunkElem <= 0) { + chunkBytes = 0U; + return false; + } + chunkBytes = static_cast(static_cast(chunkElem) * sizeof(int32_t)); + return true; +} + +__aicore__ inline uint64_t BigDataGlobalPassIndex( + uint64_t kernelLoopBase, uint32_t passCount, uint32_t loop, uint32_t pass) +{ + return (kernelLoopBase + static_cast(loop)) * static_cast(passCount) + + static_cast(pass); +} + +__aicore__ inline uint64_t BigDataPassToken(uint64_t globalPass) +{ + return globalPass + 1ULL; +} + +__aicore__ inline uint32_t BigDataPingPongSlot(uint64_t globalPass) +{ + return static_cast(globalPass & + static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS - 1U)); +} + +__aicore__ inline void BigDataKernelExitBarrier() +{ + AscendC::SyncAll(); +} + +__aicore__ inline int32_t BigDataNetworkPeerIndex(int32_t peer, int32_t rank) +{ + return peer < rank ? peer : peer - 1; +} + +__aicore__ inline __gm__ uint8_t* BigDataSlot( + __gm__ uint8_t* udmaMem, uint64_t baseOffset, uint32_t slot, + int32_t networkPeerCount, int32_t peerIndex, uint64_t chunkBytesPerPeer) +{ + return udmaMem + baseOffset + + (static_cast(slot) * static_cast(networkPeerCount) + + static_cast(peerIndex)) * chunkBytesPerPeer; +} + +__aicore__ inline void BigDataStoreTokenMte( + __gm__ uint64_t* slot, uint64_t token, AscendC::LocalTensor relayLocal) +{ + AscendC::LocalTensor tokenLocal = relayLocal.ReinterpretCast(); + constexpr uint32_t controlSlotU64 = + TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES / sizeof(uint64_t); + for (uint32_t i = 0; i < controlSlotU64; ++i) { + tokenLocal.SetValue(i, token); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor slotGlobal; + slotGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(slot)); + AscendC::DataCopyExtParams copyOut { + 1U, static_cast(TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES), 0U, 0U, 0U}; + AscendC::DataCopyPad(slotGlobal, relayLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +__aicore__ inline uint64_t BigDataLoadTokenMte( + __gm__ uint64_t* slot, AscendC::LocalTensor relayLocal) +{ + AscendC::GlobalTensor slotGlobal; + slotGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(slot)); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn { + 1U, static_cast(sizeof(uint64_t)), 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, slotGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + return relayLocal.ReinterpretCast().GetValue(0); +} + +__aicore__ inline uint64_t BigDataWaitTokenMte( + __gm__ uint64_t* slot, uint64_t token, AscendC::LocalTensor relayLocal) +{ + uint64_t observed = BigDataLoadTokenMte(slot, relayLocal); + uint64_t polls = 0; + while (observed < token && polls < TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS) { + observed = BigDataLoadTokenMte(slot, relayLocal); + ++polls; + } + return observed; +} + +__aicore__ inline __gm__ uint64_t* BigDataRemoteControlSlot( + const __gm__ TileXR::CommArgs* args, int32_t targetRank, uint64_t offset, int32_t slotRank, + uint32_t slot, int32_t rankSize) +{ + auto registry = TileXR::GetUDMARegistry(args); + const uint64_t remoteOffset = + offset + + (static_cast(slot) * static_cast(rankSize) + + static_cast(slotRank)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + if (!TileXR::UDMARegisteredRangeValid(registry, targetRank, + remoteOffset, TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES)) { + return nullptr; + } + return reinterpret_cast<__gm__ uint64_t*>( + TileXR::UDMARegisteredRemoteAddr(registry, targetRank, remoteOffset)); +} + +__aicore__ inline uint64_t BigDataIpcAckOffset(uint32_t slot, int32_t rankSize, int32_t slotRank) +{ + constexpr uint64_t maxAckBytes = + static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS) * + static_cast(TileXR::TILEXR_MAX_RANK_SIZE) * + TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + const uint64_t baseOffset = static_cast(TileXR::IPC_DATA_OFFSET) - maxAckBytes; + return baseOffset + + (static_cast(slot) * static_cast(rankSize) + + static_cast(slotRank)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; +} + +__aicore__ inline __gm__ uint64_t* BigDataLocalIpcAckSlot( + const __gm__ TileXR::CommArgs* args, int32_t rank, uint32_t slot, int32_t rankSize, int32_t peer) +{ + return reinterpret_cast<__gm__ uint64_t*>( + args->peerMems[rank] + BigDataIpcAckOffset(slot, rankSize, peer)); +} + +__aicore__ inline __gm__ uint64_t* BigDataRemoteIpcAckSlot( + const __gm__ TileXR::CommArgs* args, int32_t targetRank, int32_t rank, uint32_t slot, int32_t rankSize) +{ + return reinterpret_cast<__gm__ uint64_t*>( + args->peerMems[targetRank] + BigDataIpcAckOffset(slot, rankSize, rank)); +} + +__aicore__ inline void BigDataCopyPeerWorker( + int32_t peer, int32_t rank, int32_t rankSize, const __gm__ TileXR::CommArgs* args, + __gm__ int32_t* input, __gm__ int32_t* output, + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, + uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, + uint64_t sendDataOffset, uint64_t copyDoneOffset, uint64_t ackSignalOffset, + uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) +{ + if (peer < 0 || peer >= rankSize) { + return; + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE) { + return; + } + + const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + return; + } + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataPingPongSlot(globalPass); + + if (peer == rank) { auto src = reinterpret_cast<__gm__ uint8_t*>( - input + static_cast(peer) * elementsPerPeer + chunkOffset) + offset; - __gm__ uint8_t* dst = nullptr; - if (peer == rank) { - dst = reinterpret_cast<__gm__ uint8_t*>( - output + static_cast(rank) * elementsPerPeer + chunkOffset) + offset; - } else { - dst = udmaMem + sendDataOffset + static_cast(peer) * chunkBytesPerPeer + offset; + input + static_cast(rank) * elementsPerPeer + chunkOffset); + auto dst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(rank) * elementsPerPeer + chunkOffset); + BigDataCopyRangePingPong(dst, src, chunkBytes, relayLocal); + return; + } + + if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT && + globalPass >= TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS) { + const uint64_t reuseToken = BigDataPassToken( + globalPass - static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS)); + const uint64_t observed = BigDataWaitTokenMte( + BigDataLocalIpcAckSlot(args, rank, slot, rankSize, peer), + reuseToken, relayLocal); + if (observed < reuseToken) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + return; } - BigDataCopyOneRelay(dst, src, bytes, relayLocal); } - AscendC::PipeBarrier(); + + auto src = reinterpret_cast<__gm__ uint8_t*>( + input + static_cast(peer) * elementsPerPeer + chunkOffset); + const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); + auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, + peerIndex, chunkBytesPerPeer); + BigDataCopyRangePingPong(sendSlot, src, chunkBytes, relayLocal); + BigDataStoreTokenMte( + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer), + token, relayLocal); +} + +__aicore__ inline void BigDataSendPeerWorker( + int32_t peer, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, + int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, + uint64_t kernelLoopBase, uint32_t profileStage, uint64_t sendDataOffset, + uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, + uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) +{ + if (peer < 0 || peer >= rankSize || peer == rank || + profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { + return; + } + + const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); + const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + return; + } + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataPingPongSlot(globalPass); + + uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; + } + return; + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { + return; + } + + auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, + peerIndex, chunkBytesPerPeer); + auto localSrc = reinterpret_cast<__gm__ int32_t*>(sendSlot); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_SEND_SAMPLE_BASE + peer] = localSrc[0]; + } + const uint64_t remoteDataOffset = + recvDataOffset + + (static_cast(slot) * static_cast(networkPeerCount) + + static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer; + const uint64_t remoteReadyOffset = + readySignalOffset + + (static_cast(slot) * static_cast(rankSize) + + static_cast(rank)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + auto registry = TileXR::GetUDMARegistry(args); + auto udmaInfo = TileXR::GetUDMAInfo(args); + auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); + auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); + bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteDataOffset, chunkBytes) && + TileXR::UDMARegisteredRangeValid(registry, peer, remoteReadyOffset, sizeof(uint64_t)); + uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_RANGE_VALID_BASE + peer] = rangeValid ? 1 : 0; + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE + peer] = static_cast(wqeBefore); + debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = + static_cast(wqCtx->localTokenId); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = + static_cast(reinterpret_cast(registry->regions[peer].base) & + 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = + static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); + debug[TILEXR_UDMA_DEMO_DEBUG_TPN_BASE + peer] = static_cast(remoteMemInfo->tpn); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_DATA_OFFSET_BASE + peer] = + static_cast(remoteDataOffset); + debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_READY_OFFSET_BASE + peer] = + static_cast(remoteReadyOffset); + } + + TileXR::UDMAPutSignalNbi(args, peer, localSrc, + remoteDataOffset, chunkBytes, remoteReadyOffset, token); + uint32_t status = TileXR::UDMAQuietStatus(args, peer); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + uint32_t wqeAfter = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); + debug[TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE + peer] = static_cast(wqeAfter); + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } } -__aicore__ inline void BigDataRelayCopyShard( - int32_t blockIdx, int32_t blockNum, int32_t rank, int32_t rankSize, - __gm__ int32_t* output, __gm__ uint8_t* udmaMem, - int32_t elementsPerPeer, int32_t chunkOffset, uint32_t chunkBytes, - uint64_t recvDataOffset, uint64_t chunkBytesPerPeer, +__aicore__ inline void BigDataRecvPeerWorker( + int32_t peer, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, + __gm__ int32_t* output, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, + int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, + uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, + uint64_t recvDataOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, + uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { - const uint32_t copyCores = BigDataCopyCoreCount(blockNum); - if (static_cast(blockIdx) >= copyCores || chunkBytes == 0 || rankSize <= 1) { + if (peer < 0 || peer >= rankSize || peer == rank || + profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { return; } - const uint32_t chunksPerPeer = - (chunkBytes + TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES - 1U) / - TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; - const uint64_t totalUnits = static_cast(rankSize - 1) * static_cast(chunksPerPeer); - for (uint64_t unit = static_cast(blockIdx); unit < totalUnits; unit += copyCores) { - const int32_t logicalPeer = static_cast(unit / static_cast(chunksPerPeer)); - const int32_t peer = logicalPeer < rank ? logicalPeer : logicalPeer + 1; - const uint32_t chunkInPeer = static_cast(unit % static_cast(chunksPerPeer)); - const uint32_t offset = chunkInPeer * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; - const uint32_t bytes = - (chunkBytes - offset < TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES) ? - (chunkBytes - offset) : TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; - - auto src = udmaMem + recvDataOffset + static_cast(peer) * chunkBytesPerPeer + offset; - auto dst = reinterpret_cast<__gm__ uint8_t*>( - output + static_cast(peer) * elementsPerPeer + chunkOffset) + offset; - BigDataCopyOneRelay(dst, src, bytes, relayLocal); + + const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); + const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + return; + } + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataPingPongSlot(globalPass); + + uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer), + token, relayLocal); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = + static_cast(observed); + if (observed < token) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + } + if (observed < token) { + return; } - AscendC::PipeBarrier(); + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY) { + return; + } + + auto recvSlot = BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount, + peerIndex, chunkBytesPerPeer); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + auto recvSlotInt = reinterpret_cast<__gm__ int32_t*>(recvSlot); + debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SLOT_SAMPLE_BASE + peer] = recvSlotInt[0]; + } + auto dst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(peer) * elementsPerPeer + chunkOffset); + BigDataCopyRangePingPong(dst, recvSlot, chunkBytes, relayLocal); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + auto relayDst = output + static_cast(peer) * elementsPerPeer + chunkOffset; + debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE + peer] = relayDst[0]; + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY) { + return; + } + + auto remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); + if (remoteAck == nullptr) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + return; + } + BigDataStoreTokenMte(remoteAck, token, relayLocal); } } // namespace @@ -1242,24 +1677,14 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_fused_kernel( extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, GM_ADDR debugGM, - int32_t elementsPerPeer, uint64_t dataOffset, uint64_t readyPayloadOffset, uint64_t ackPayloadOffset, + int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t tokenBase, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage) { - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_PREPARE = 0; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY = 1; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC = 2; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT = 3; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY = 4; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY = 5; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT = 6; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK = 7; - constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_FULL = 8; if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_FULL) { profileStage = TILEXR_BIGDATA_PROFILE_STAGE_FULL; } - auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); @@ -1270,7 +1695,6 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( int32_t rankSize = args->rankSize; bool enabled = TileXR::UDMARegistryEnabled(args); const int32_t blockIdx = AscendC::GetBlockIdx(); - const int32_t blockNum = AscendC::GetBlockNum(); if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; @@ -1281,188 +1705,65 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( debug[5] = chunkElements; } if (!enabled) { - for (uint32_t l = 0; l < loopCount; ++l) { - AscendC::SyncAll(); - } return; } const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; const uint64_t chunkBytesPerPeer = static_cast(effectiveChunkElements) * sizeof(int32_t); const uint64_t sendDataOffset = dataOffset; - const uint64_t recvDataOffset = dataOffset + static_cast(rankSize) * chunkBytesPerPeer; + const uint64_t recvDataOffset = + sendDataOffset + + static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS) * + static_cast(rankSize > 1 ? rankSize - 1 : 1) * chunkBytesPerPeer; AscendC::TPipe pipe; AscendC::TBuf relayTBuf; - pipe.InitBuffer(relayTBuf, TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES); + pipe.InitBuffer(relayTBuf, TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES); AscendC::LocalTensor relayLocal = relayTBuf.Get(); + if (rankSize <= 0 || blockIdx >= rankSize * static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)) { + return; + } + const int32_t peer = blockIdx / static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER); + const int32_t role = blockIdx % static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER); + for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { - const int32_t chunkOffset = static_cast(pass) * effectiveChunkElements; - const int32_t remaining = elementsPerPeer - chunkOffset; - const int32_t chunkElem = remaining < effectiveChunkElements ? remaining : effectiveChunkElements; - if (chunkElem <= 0) { - continue; - } - const uint32_t chunkBytes = static_cast(static_cast(chunkElem) * sizeof(int32_t)); - const uint64_t token = - tokenBase + static_cast(loop) * static_cast(passCount) + - static_cast(pass) + 1ULL; - - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE) { - continue; - } - - BigDataSendCopyShard(blockIdx, blockNum, rank, rankSize, input, output, udmaMem, - elementsPerPeer, chunkOffset, chunkBytes, sendDataOffset, chunkBytesPerPeer, relayLocal); - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { - continue; - } - AscendC::SyncAll(); - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { - continue; - } - - for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { - if (peer == rank) { - continue; - } - auto sendSlot = udmaMem + sendDataOffset + static_cast(peer) * chunkBytesPerPeer; - auto localSrc = reinterpret_cast<__gm__ int32_t*>(sendSlot); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_SEND_SAMPLE_BASE + peer] = localSrc[0]; - } - const uint64_t remoteDataOffset = - recvDataOffset + static_cast(rank) * chunkBytesPerPeer; - const uint64_t remoteReadyOffset = - readySignalOffset + static_cast(rank) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; - auto registry = TileXR::GetUDMARegistry(args); - auto udmaInfo = TileXR::GetUDMAInfo(args); - auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); - auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); - bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteDataOffset, chunkBytes); - uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_RANGE_VALID_BASE + peer] = rangeValid ? 1 : 0; - debug[TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE + peer] = static_cast(wqeBefore); - debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = static_cast(wqCtx->localTokenId); - debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = - static_cast(reinterpret_cast(registry->regions[peer].base) & 0xFFFFFFFFU); - debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = - static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); - debug[TILEXR_UDMA_DEMO_DEBUG_TPN_BASE + peer] = static_cast(remoteMemInfo->tpn); - debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_DATA_OFFSET_BASE + peer] = - static_cast(remoteDataOffset); - debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_READY_OFFSET_BASE + peer] = - static_cast(remoteReadyOffset); - } - TileXR::UDMAPutNbi(args, peer, localSrc, remoteDataOffset, chunkBytes); - uint32_t status = TileXR::UDMAQuietStatus(args, peer); - auto readyPayload = ControlSlot(udmaMem, readyPayloadOffset, peer); - *readyPayload = token; - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t*>(readyPayload), sizeof(uint64_t)); - TileXR::UDMAPutNbi(args, peer, readyPayload, remoteReadyOffset, sizeof(uint64_t)); - uint32_t readyStatus = TileXR::UDMAQuietStatus(args, peer); - if (status == 0) { - status = readyStatus; - } - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - uint32_t wqeAfter = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); - debug[TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE + peer] = static_cast(wqeAfter); - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); - } - } - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { - continue; - } - - for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { - if (peer == rank) { - continue; - } - uint64_t observed = WaitUdmaSignalToken(ControlSlot(udmaMem, readySignalOffset, peer), token); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = - static_cast(observed); - if (observed < token) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; - } - } - } - AscendC::SyncAll(); - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY) { - continue; - } - - if (blockIdx < rankSize && blockIdx != rank && debug != nullptr && loop == 0 && pass == 0 && blockIdx < 16) { - auto recvSlot = udmaMem + recvDataOffset + static_cast(blockIdx) * chunkBytesPerPeer; - auto recvSlotInt = reinterpret_cast<__gm__ int32_t*>(recvSlot); - debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SLOT_SAMPLE_BASE + blockIdx] = recvSlotInt[0]; - } - BigDataRelayCopyShard(blockIdx, blockNum, rank, rankSize, output, udmaMem, - elementsPerPeer, chunkOffset, chunkBytes, recvDataOffset, chunkBytesPerPeer, relayLocal); - AscendC::SyncAll(); - if (blockIdx < rankSize && blockIdx != rank && debug != nullptr && loop == 0 && pass == 0 && blockIdx < 16) { - auto relayDst = output + static_cast(blockIdx) * elementsPerPeer + chunkOffset; - debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE + blockIdx] = relayDst[0]; - } - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY) { - continue; + if (role == 0) { + BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, + sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); } - for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { - if (peer == rank) { - continue; - } - auto ackPayload = ControlSlot(udmaMem, ackPayloadOffset, peer); - *ackPayload = token; - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t*>(ackPayload), sizeof(uint64_t)); - const uint64_t remoteAckOffset = - ackSignalOffset + static_cast(rank) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; - TileXR::UDMAPutNbi(args, peer, ackPayload, remoteAckOffset, sizeof(uint64_t)); - (void)TileXR::UDMAQuietStatus(args, peer); - } - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT) { - continue; + if (role == 1) { + BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, + sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, + chunkBytesPerPeer, relayLocal); } - for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { - if (peer == rank) { - continue; - } - uint64_t observed = WaitUdmaSignalToken(ControlSlot(udmaMem, ackSignalOffset, peer), token); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16 && observed < token) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; - } - } - AscendC::SyncAll(); - if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK) { - continue; + if (role == 2) { + BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, + recvDataOffset, readySignalOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); } } } - AscendC::SyncAll(); + BigDataKernelExitBarrier(); } void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, - uint64_t dataOffset, uint64_t readyPayloadOffset, uint64_t ackPayloadOffset, + uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t tokenBase, + int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage) { tilexr_udma_all_to_all_bigdata_kernel<<>>( commArgs, input, output, udmaMem, debug, elementsPerPeer, - dataOffset, readyPayloadOffset, ackPayloadOffset, readySignalOffset, ackSignalOffset, - chunkElements, passCount, loopCount, tokenBase, profileStage); + dataOffset, copyDoneOffset, readySignalOffset, ackSignalOffset, + chunkElements, passCount, loopCount, kernelLoopBase, profileStage); } void launch_tilexr_udma_all_to_all_fused( diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index ebb99c89..56bf4315 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -35,6 +35,14 @@ int g_failures = 0; } \ } while (0) +#define CHECK_NOT_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) != std::string::npos) { \ + std::cerr << "CHECK_NOT_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + std::string ReadFile(const std::string& path) { std::ifstream in(path.c_str()); @@ -133,16 +141,35 @@ void TestAllToAllBigDataPlan() const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 64ULL * 1024ULL * 1024ULL); - CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes); - CHECK_EQ(plan.signalBytes, 2ULL * static_cast(rankSize) * sizeof(uint64_t)); - CHECK_EQ(plan.readySignalOffset, plan.dataBytes); - CHECK_EQ(plan.ackSignalOffset, plan.dataBytes + static_cast(rankSize) * sizeof(uint64_t)); - CHECK_EQ(plan.dataBytes + plan.signalBytes <= plan.registeredBytes, true); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 64ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 3U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataPingPongSlots, 2U); + CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, true); + const size_t controlGroupBytes = + static_cast(TileXR::Demo::kAllToAllBigDataPingPongSlots) * + static_cast(rankSize) * TileXR::Demo::kAllToAllBigDataControlSlotBytes; + CHECK_EQ(plan.controlBytes, controlGroupBytes); + CHECK_EQ(plan.signalBytes, 2ULL * controlGroupBytes); + CHECK_EQ(plan.copyDoneOffset, plan.dataBytes); + CHECK_EQ(plan.readySignalOffset, plan.copyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.ackSignalOffset, plan.readySignalOffset + controlGroupBytes); + CHECK_EQ(plan.registeredBytes, plan.dataBytes + plan.controlBytes + plan.signalBytes); + CHECK_EQ(plan.dataBytes, + static_cast(rankSize - 1) * TileXR::Demo::kAllToAllBigDataPingPongSlots * 2ULL * + plan.chunkBytesPerPeer); CHECK_EQ(plan.chunkElements > 0, true); CHECK_EQ(plan.chunkBytesPerPeer, static_cast(plan.chunkElements) * sizeof(int32_t)); CHECK_EQ(plan.passCount > 1, true); } +void TestAllToAllBigDataBlockDim() +{ + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(0), 1U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(1), 3U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(8), 24U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64), 192U); +} + void TestDemoDebugLayoutSource() { const std::string demo = @@ -152,6 +179,7 @@ void TestDemoDebugLayoutSource() CHECK_CONTAINS(demo, "kDebugUdmaStatusBase + TileXR::TILEXR_MAX_RANK_SIZE"); CHECK_CONTAINS(demo, "kDebugIpcGather + 1"); + CHECK_CONTAINS(demo, "kDebugReadySeenBase + TileXR::TILEXR_MAX_RANK_SIZE"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer"); } @@ -206,7 +234,7 @@ void TestAllToAllChunkedUdmaSource() CHECK_CONTAINS(demo, "alltoall fused IPC: single kernel send+flag+recv"); CHECK_CONTAINS(demo, "alltoall use "); CHECK_CONTAINS(demo, "plain IPC fallback"); - CHECK_CONTAINS(demo, "launch all-to-all kernel iter="); + CHECK_CONTAINS(demo, "launch all-to-all kernel repeat="); CHECK_CONTAINS(demo, "for (uint32_t pass = 0; pass < chunkPlan.passCount; ++pass)"); CHECK_CONTAINS(demo, "registered output chunk"); CHECK_CONTAINS(kernel, "tilexr_all_to_all_plain_ipc_scatter_kernel"); @@ -223,20 +251,75 @@ void TestAllToAllBigDataSource() ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); const std::string kernel = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + const std::string udma = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); CHECK_CONTAINS(demo, "testType == 7"); CHECK_CONTAINS(demo, "PlanAllToAllBigDataUdma"); CHECK_CONTAINS(demo, "bigdata alltoall registered dataBytes="); CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_bigdata"); + CHECK_CONTAINS(demo, "for (int iter = 0; iter < allToAllRepeat; ++iter)"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_PROFILE_STAGE"); + CHECK_CONTAINS(demo, "const uint64_t kernelLoopBase = static_cast(iter)"); + CHECK_CONTAINS(demo, "bigDataPlan.passCount, 1, kernelLoopBase"); + CHECK_CONTAINS(demo, "bigDataPlan.copyDoneOffset"); + CHECK_CONTAINS(demo, "static_cast(registeredMemory) + bigDataPlan.copyDoneOffset"); + CHECK_CONTAINS(demo, "bigDataPlan.controlBytes + bigDataPlan.signalBytes"); + CHECK_CONTAINS(demo, "static_cast(bigDataProfileStage)"); CHECK_CONTAINS(demo, "alltoall udma-bigdata"); CHECK_CONTAINS(demo, "ERROR: bigdata alltoall UDMA registration failed"); CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_bigdata_kernel"); CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_bigdata"); + CHECK_CONTAINS(kernel, "uint64_t kernelLoopBase"); + CHECK_CONTAINS(kernel, "uint32_t profileStage"); + CHECK_CONTAINS(kernel, "if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE)"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT"); + CHECK_CONTAINS(kernel, "uint64_t copyDoneOffset"); + CHECK_CONTAINS(kernel, "BigDataGlobalPassIndex"); + CHECK_CONTAINS(kernel, "const uint64_t globalPass = BigDataGlobalPassIndex"); + CHECK_CONTAINS(kernel, "BigDataPassToken(globalPass)"); + CHECK_CONTAINS(kernel, "BigDataPingPongSlot(globalPass)"); + CHECK_CONTAINS(kernel, "BigDataKernelExitBarrier"); + CHECK_CONTAINS(kernel, "copyDoneOffset"); CHECK_CONTAINS(kernel, "readySignalOffset"); CHECK_CONTAINS(kernel, "ackSignalOffset"); + CHECK_CONTAINS(kernel, "BigDataWaitTokenMte"); + CHECK_CONTAINS(kernel, "BigDataLoadTokenMte"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 3U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U"); + CHECK_CONTAINS(kernel, "BigDataCopyPeerWorker"); + CHECK_CONTAINS(kernel, "BigDataSendPeerWorker"); + CHECK_CONTAINS(kernel, "BigDataRecvPeerWorker"); + CHECK_CONTAINS(kernel, "const int32_t peer = blockIdx / static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)"); + CHECK_CONTAINS(kernel, "const int32_t role = blockIdx % static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)"); + CHECK_CONTAINS(kernel, "BigDataPingPongSlot"); + CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex"); + CHECK_CONTAINS(kernel, "BigDataStoreTokenMte"); + CHECK_CONTAINS(kernel, "BigDataIpcAckOffset"); + CHECK_CONTAINS(kernel, "BigDataLocalIpcAckSlot"); + CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); + CHECK_CONTAINS(kernel, "rankSize > 1 ? rankSize - 1 : 1"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer)"); + CHECK_CONTAINS(kernel, "remoteDataOffset ="); + CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex(rank, peer)"); + CHECK_CONTAINS(kernel, "BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); - CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "localSrc,"); + CHECK_CONTAINS(kernel, "remoteDataOffset, chunkBytes, remoteReadyOffset, token"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES"); + CHECK_CONTAINS(kernel, "relayLocal[bufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]"); + CHECK_CONTAINS(kernel, "globalPass - static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS)"); + CHECK_CONTAINS(kernel, "recvSlotInt[0]"); + CHECK_CONTAINS(kernel, "BigDataStoreTokenMte(remoteAck, token, relayLocal)"); + CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize)"); CHECK_CONTAINS(kernel, "ackSignal"); + CHECK_NOT_CONTAINS(kernel, "UDMAPutNbi"); + CHECK_CONTAINS(udma, "if (length == 0)"); + CHECK_CONTAINS(udma, "reinterpret_cast(addr) + length - 1"); } } // namespace @@ -248,6 +331,7 @@ int main() TestBuildAllToAllOutput(); TestAllToAllMaxRank256With64MiBPerRank(); TestAllToAllBigDataPlan(); + TestAllToAllBigDataBlockDim(); TestDemoDebugLayoutSource(); TestAllToAllDataAsFlagSource(); TestAllToAllChunkedUdmaSource(); From dd8b1ef7b692df9f93d400e057c98d5971580b2f Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Sun, 28 Jun 2026 23:15:04 +0800 Subject: [PATCH 015/163] perf(udma): split bigdata alltoall local copies --- tests/udma/demo/tilexr_udma_alltoall_layout.h | 13 +- tests/udma/demo/tilexr_udma_demo.cpp | 5 +- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 139 +++++++++++++----- .../unit/test_tilexr_udma_alltoall_layout.cpp | 31 ++-- 4 files changed, 138 insertions(+), 50 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index 2a687e7f..faaf32d7 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -27,7 +27,8 @@ struct AllToAllChunkPlan { constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataControlSlotBytes = 64ULL; -constexpr uint32_t kAllToAllBigDataCoresPerPeer = 3U; +constexpr uint32_t kAllToAllBigDataCoresPerPeer = 5U; +constexpr uint32_t kAllToAllBigDataLocalCopyShards = 2U; constexpr uint32_t kAllToAllBigDataPingPongSlots = 2U; struct AllToAllBigDataPlan { @@ -36,6 +37,7 @@ struct AllToAllBigDataPlan { size_t chunkBytesPerPeer = 0; size_t dataBytes = 0; size_t copyDoneOffset = 0; + size_t recvCopyDoneOffset = 0; size_t readySignalOffset = 0; size_t ackSignalOffset = 0; size_t controlBytes = 0; @@ -88,9 +90,9 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element plan.registeredBytes = kAllToAllBigDataMaxRegisteredBytes; const size_t controlGroupBytes = static_cast(kAllToAllBigDataPingPongSlots) * static_cast(rankSize) * - kAllToAllBigDataControlSlotBytes; + static_cast(kAllToAllBigDataLocalCopyShards) * kAllToAllBigDataControlSlotBytes; plan.controlBytes = controlGroupBytes; - plan.signalBytes = 2ULL * controlGroupBytes; + plan.signalBytes = 3ULL * controlGroupBytes; if (plan.controlBytes + plan.signalBytes >= plan.registeredBytes) { return plan; } @@ -111,10 +113,11 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t element plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); plan.dataBytes = dataSlotCount * plan.chunkBytesPerPeer; plan.copyDoneOffset = plan.dataBytes; - plan.readySignalOffset = plan.copyDoneOffset + controlGroupBytes; + plan.recvCopyDoneOffset = plan.copyDoneOffset + controlGroupBytes; + plan.readySignalOffset = plan.recvCopyDoneOffset + controlGroupBytes; plan.ackSignalOffset = plan.readySignalOffset + controlGroupBytes; plan.controlBytes = controlGroupBytes; - plan.signalBytes = 2ULL * controlGroupBytes; + plan.signalBytes = 3ULL * controlGroupBytes; plan.registeredBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; plan.passCount = static_cast( (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1) / diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index edcf865b..de71e4bf 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -50,7 +50,7 @@ extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, - uint64_t readySignalOffset, uint64_t ackSignalOffset, + uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage); extern void launch_tilexr_all_to_all_ipc_scatter( @@ -721,6 +721,7 @@ int main(int argc, char** argv) " chunkElements=" + std::to_string(bigDataPlan.chunkElements) + " dataBytes=" + std::to_string(bigDataPlan.dataBytes) + " copyDoneOffset=" + std::to_string(bigDataPlan.copyDoneOffset) + + " recvCopyDoneOffset=" + std::to_string(bigDataPlan.recvCopyDoneOffset) + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + " registeredBytes=" + std::to_string(bigDataPlan.registeredBytes)); @@ -898,6 +899,7 @@ int main(int argc, char** argv) } PrintStatus(rank, "bigdata alltoall registered dataBytes=" + std::to_string(bigDataPlan.dataBytes) + " copyDoneOffset=" + std::to_string(bigDataPlan.copyDoneOffset) + + " recvCopyDoneOffset=" + std::to_string(bigDataPlan.recvCopyDoneOffset) + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + " regBytes=" + std::to_string(registeredBytes) + @@ -925,6 +927,7 @@ int main(int argc, char** argv) reinterpret_cast(bigInput), reinterpret_cast(bigOutput), reinterpret_cast(registeredMemory), reinterpret_cast(debug), elementsPerRank, 0, bigDataPlan.copyDoneOffset, + bigDataPlan.recvCopyDoneOffset, bigDataPlan.readySignalOffset, bigDataPlan.ackSignalOffset, bigDataPlan.chunkElements, bigDataPlan.passCount, 1, kernelLoopBase, static_cast(bigDataProfileStage)); diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 075b0c9c..7767ce71 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -52,7 +52,8 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES = 64 * 1024; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES = TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES * 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U; -constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 3U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = 2U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 5U; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_PREPARE = 0; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY = 1; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC = 2; @@ -83,12 +84,15 @@ __aicore__ inline __gm__ uint64_t* ControlSlot(__gm__ uint8_t* base, uint64_t of } __aicore__ inline __gm__ uint64_t* BigDataControlSlot( - __gm__ uint8_t* base, uint64_t offset, uint32_t slot, int32_t rankSize, int32_t peer) + __gm__ uint8_t* base, uint64_t offset, uint32_t slot, int32_t rankSize, int32_t peer, + uint32_t shard) { return reinterpret_cast<__gm__ uint64_t*>( base + offset + - (static_cast(slot) * static_cast(rankSize) + - static_cast(peer)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); + ((static_cast(slot) * static_cast(rankSize) + + static_cast(peer)) * + static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) + + static_cast(shard)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); } __aicore__ inline void BigDataCopyInTile( @@ -175,6 +179,31 @@ __aicore__ inline uint32_t BigDataTileBytes(uint32_t totalBytes, uint32_t offset remain : TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES; } +__aicore__ inline bool BigDataCopyShardRange( + uint32_t shard, uint32_t totalElements, uint32_t& shardOffsetBytes, uint32_t& shardBytes) +{ + if (shard >= TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) { + shardOffsetBytes = 0U; + shardBytes = 0U; + return false; + } + constexpr uint32_t alignElements = 32U / sizeof(int32_t); + const uint64_t begin = + static_cast(totalElements) * static_cast(shard) / + static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS); + const uint64_t end = + static_cast(totalElements) * static_cast(shard + 1U) / + static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS); + const uint64_t alignedBegin = (shard == 0U) ? 0ULL : + (begin / static_cast(alignElements)) * static_cast(alignElements); + const uint64_t alignedEnd = (shard + 1U == TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) ? + static_cast(totalElements) : + (end / static_cast(alignElements)) * static_cast(alignElements); + shardOffsetBytes = static_cast(alignedBegin * sizeof(int32_t)); + shardBytes = static_cast((alignedEnd - alignedBegin) * sizeof(int32_t)); + return shardBytes > 0U; +} + __aicore__ inline void BigDataCopyRangePingPong( __gm__ uint8_t* dst, __gm__ uint8_t* src, uint32_t bytes, AscendC::LocalTensor relayLocal) @@ -190,6 +219,7 @@ __aicore__ inline void BigDataCopyRangePingPong( const uint32_t tileBytes = BigDataTileBytes(bytes, offset); BigDataCopyOneRelay(dst + offset, src + offset, tileBytes, relayLocal); } + AscendC::PipeBarrier(); return; } bool copyOutInFlight0 = false; @@ -251,6 +281,7 @@ __aicore__ inline void BigDataCopyRangePingPong( if (copyOutInFlight1) { BigDataWaitMte3ToMte2(1U); } + AscendC::PipeBarrier(); } __aicore__ inline bool BigDataPassChunk( @@ -400,15 +431,16 @@ __aicore__ inline void BigDataCopyPeerWorker( __gm__ int32_t* input, __gm__ int32_t* output, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, - uint64_t sendDataOffset, uint64_t copyDoneOffset, uint64_t ackSignalOffset, + uint32_t copyShard, uint64_t sendDataOffset, uint64_t copyDoneOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { - if (peer < 0 || peer >= rankSize) { + if (peer < 0 || peer >= rankSize || copyShard >= TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) { return; } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE) { return; } + (void)ackSignalOffset; const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; int32_t chunkOffset = 0; @@ -416,6 +448,11 @@ __aicore__ inline void BigDataCopyPeerWorker( if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { return; } + uint32_t shardOffset = 0U; + uint32_t shardBytes = 0U; + if (!BigDataCopyShardRange(copyShard, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { + return; + } const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataPingPongSlot(globalPass); @@ -425,7 +462,7 @@ __aicore__ inline void BigDataCopyPeerWorker( input + static_cast(rank) * elementsPerPeer + chunkOffset); auto dst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(rank) * elementsPerPeer + chunkOffset); - BigDataCopyRangePingPong(dst, src, chunkBytes, relayLocal); + BigDataCopyRangePingPong(dst + shardOffset, src + shardOffset, shardBytes, relayLocal); return; } @@ -450,9 +487,9 @@ __aicore__ inline void BigDataCopyPeerWorker( const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, peerIndex, chunkBytesPerPeer); - BigDataCopyRangePingPong(sendSlot, src, chunkBytes, relayLocal); + BigDataCopyRangePingPong(sendSlot + shardOffset, src + shardOffset, shardBytes, relayLocal); BigDataStoreTokenMte( - BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer), + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer, copyShard), token, relayLocal); } @@ -480,15 +517,17 @@ __aicore__ inline void BigDataSendPeerWorker( const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataPingPongSlot(globalPass); - uint64_t observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer), - token, relayLocal); - if (observed < token) { - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; + for (uint32_t copyShard = 0U; copyShard < TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS; ++copyShard) { + const uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer, copyShard), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; + } + return; } - return; } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { return; @@ -506,8 +545,10 @@ __aicore__ inline void BigDataSendPeerWorker( static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer; const uint64_t remoteReadyOffset = readySignalOffset + - (static_cast(slot) * static_cast(rankSize) + - static_cast(rank)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + ((static_cast(slot) * static_cast(rankSize) + + static_cast(rank)) * + static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) * + TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; auto registry = TileXR::GetUDMARegistry(args); auto udmaInfo = TileXR::GetUDMAInfo(args); auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); @@ -547,14 +588,16 @@ __aicore__ inline void BigDataRecvPeerWorker( __gm__ int32_t* output, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, - uint64_t recvDataOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - uint64_t chunkBytesPerPeer, + uint32_t recvShard, uint64_t recvDataOffset, uint64_t recvCopyDoneOffset, + uint64_t readySignalOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { if (peer < 0 || peer >= rankSize || peer == rank || + recvShard >= TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { return; } + (void)ackSignalOffset; const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; @@ -563,14 +606,19 @@ __aicore__ inline void BigDataRecvPeerWorker( if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { return; } + uint32_t shardOffset = 0U; + uint32_t shardBytes = 0U; + if (!BigDataCopyShardRange(recvShard, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { + return; + } const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataPingPongSlot(globalPass); uint64_t observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer), + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer, 0U), token, relayLocal); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = static_cast(observed); if (observed < token) { @@ -587,14 +635,14 @@ __aicore__ inline void BigDataRecvPeerWorker( auto recvSlot = BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount, peerIndex, chunkBytesPerPeer); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { auto recvSlotInt = reinterpret_cast<__gm__ int32_t*>(recvSlot); debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SLOT_SAMPLE_BASE + peer] = recvSlotInt[0]; } auto dst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(peer) * elementsPerPeer + chunkOffset); - BigDataCopyRangePingPong(dst, recvSlot, chunkBytes, relayLocal); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + BigDataCopyRangePingPong(dst + shardOffset, recvSlot + shardOffset, shardBytes, relayLocal); + if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { auto relayDst = output + static_cast(peer) * elementsPerPeer + chunkOffset; debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE + peer] = relayDst[0]; } @@ -602,6 +650,25 @@ __aicore__ inline void BigDataRecvPeerWorker( return; } + BigDataStoreTokenMte( + BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, peer, recvShard), + token, relayLocal); + if (recvShard + 1U != TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) { + return; + } + for (uint32_t shard = 0U; shard < TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS; ++shard) { + observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, peer, shard), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + return; + } + } + auto remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); if (remoteAck == nullptr) { if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { @@ -1678,7 +1745,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_fused_kernel( extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, GM_ADDR debugGM, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, - uint64_t readySignalOffset, uint64_t ackSignalOffset, + uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage) { @@ -1729,23 +1796,27 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { - if (role == 0) { + if (role < static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) { + const uint32_t copyShard = static_cast(role); BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, - sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); + copyShard, sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); } - if (role == 1) { + if (role == static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) { BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, chunkBytesPerPeer, relayLocal); } - if (role == 2) { + if (role > static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) { + const uint32_t recvShard = + static_cast(role) - TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS - 1U; BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, - recvDataOffset, readySignalOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); + recvShard, recvDataOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, + chunkBytesPerPeer, relayLocal); } } } @@ -1756,13 +1827,13 @@ void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, - uint64_t readySignalOffset, uint64_t ackSignalOffset, + uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage) { tilexr_udma_all_to_all_bigdata_kernel<<>>( commArgs, input, output, udmaMem, debug, elementsPerPeer, - dataOffset, copyDoneOffset, readySignalOffset, ackSignalOffset, + dataOffset, copyDoneOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, chunkElements, passCount, loopCount, kernelLoopBase, profileStage); } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 56bf4315..2354ed02 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -142,16 +142,20 @@ void TestAllToAllBigDataPlan() CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 64ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 64ULL); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 3U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 5U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataLocalCopyShards, 2U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataPingPongSlots, 2U); CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, true); const size_t controlGroupBytes = static_cast(TileXR::Demo::kAllToAllBigDataPingPongSlots) * - static_cast(rankSize) * TileXR::Demo::kAllToAllBigDataControlSlotBytes; + static_cast(rankSize) * + static_cast(TileXR::Demo::kAllToAllBigDataLocalCopyShards) * + TileXR::Demo::kAllToAllBigDataControlSlotBytes; CHECK_EQ(plan.controlBytes, controlGroupBytes); - CHECK_EQ(plan.signalBytes, 2ULL * controlGroupBytes); + CHECK_EQ(plan.signalBytes, 3ULL * controlGroupBytes); CHECK_EQ(plan.copyDoneOffset, plan.dataBytes); - CHECK_EQ(plan.readySignalOffset, plan.copyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.recvCopyDoneOffset, plan.copyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.readySignalOffset, plan.recvCopyDoneOffset + controlGroupBytes); CHECK_EQ(plan.ackSignalOffset, plan.readySignalOffset + controlGroupBytes); CHECK_EQ(plan.registeredBytes, plan.dataBytes + plan.controlBytes + plan.signalBytes); CHECK_EQ(plan.dataBytes, @@ -165,9 +169,9 @@ void TestAllToAllBigDataPlan() void TestAllToAllBigDataBlockDim() { CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(0), 1U); - CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(1), 3U); - CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(8), 24U); - CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64), 192U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(1), 5U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(8), 40U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64), 320U); } void TestDemoDebugLayoutSource() @@ -263,6 +267,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "const uint64_t kernelLoopBase = static_cast(iter)"); CHECK_CONTAINS(demo, "bigDataPlan.passCount, 1, kernelLoopBase"); CHECK_CONTAINS(demo, "bigDataPlan.copyDoneOffset"); + CHECK_CONTAINS(demo, "bigDataPlan.recvCopyDoneOffset"); CHECK_CONTAINS(demo, "static_cast(registeredMemory) + bigDataPlan.copyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.controlBytes + bigDataPlan.signalBytes"); CHECK_CONTAINS(demo, "static_cast(bigDataProfileStage)"); @@ -288,11 +293,16 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataLoadTokenMte"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 3U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 5U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = 2U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U"); CHECK_CONTAINS(kernel, "BigDataCopyPeerWorker"); CHECK_CONTAINS(kernel, "BigDataSendPeerWorker"); CHECK_CONTAINS(kernel, "BigDataRecvPeerWorker"); + CHECK_CONTAINS(kernel, "BigDataCopyShardRange"); + CHECK_CONTAINS(kernel, "copyShard"); + CHECK_CONTAINS(kernel, "recvShard"); + CHECK_CONTAINS(kernel, "recvCopyDoneOffset"); CHECK_CONTAINS(kernel, "const int32_t peer = blockIdx / static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)"); CHECK_CONTAINS(kernel, "const int32_t role = blockIdx % static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)"); CHECK_CONTAINS(kernel, "BigDataPingPongSlot"); @@ -302,8 +312,9 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataLocalIpcAckSlot"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); CHECK_CONTAINS(kernel, "rankSize > 1 ? rankSize - 1 : 1"); - CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer)"); - CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer, copyShard)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, peer, recvShard)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer, 0U)"); CHECK_CONTAINS(kernel, "remoteDataOffset ="); CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex(rank, peer)"); CHECK_CONTAINS(kernel, "BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount"); From 0eb4282416727010de226e85cfd64e3911be2c27 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 29 Jun 2026 11:35:22 +0800 Subject: [PATCH 016/163] docs: design multinode udma bigdata alltoall --- ...-udma-alltoall-bigdata-multinode-design.md | 364 ++++++++++++++++++ 1 file changed, 364 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-29-udma-alltoall-bigdata-multinode-design.md diff --git a/docs/superpowers/specs/2026-06-29-udma-alltoall-bigdata-multinode-design.md b/docs/superpowers/specs/2026-06-29-udma-alltoall-bigdata-multinode-design.md new file mode 100644 index 00000000..640262f7 --- /dev/null +++ b/docs/superpowers/specs/2026-06-29-udma-alltoall-bigdata-multinode-design.md @@ -0,0 +1,364 @@ +# UDMA AllToAll Bigdata Multi-Node Design + +## Purpose + +This design extends the existing `testType=7` UDMA AllToAll bigdata demo to support +multi-node rank sizes, where one server owns 8 contiguous ranks and `rankSize > 8` +means multiple servers participate. + +The 8-rank single-server path must remain unchanged. When `rankSize <= 8`, the +current bigdata kernel layout, block dispatch, pass protocol, correctness behavior, +and performance profile stay as they are today. When `rankSize > 8`, `testType=7` +switches to a multi-node worker dispatch that keeps local-server traffic parallel +and serializes all remote-server traffic through one additional 5-core worker group. + +## Confirmed Assumptions + +- Each server has 8 ranks. +- Global rank ids are contiguous by server: + + ```text + nodeId = rank / 8 + localRank = rank % 8 + localNodeBegin = nodeId * 8 + localNodeEnd = localNodeBegin + 8 + ``` + +- `rankSize > 8` implies multi-node mode. +- Multi-node mode requires `rankSize` to be a positive multiple of 8. +- Remote traffic means all peers outside `[localNodeBegin, localNodeEnd)`, not only + peers with the same `localRank`. +- The path is strict UDMA. It must not fall back to IPC. + +## Scope + +In scope: + +- Keep the feature under `testType=7`. +- Add a multi-node branch for `rankSize > 8`. +- Keep `rankSize <= 8` behavior bit-for-bit equivalent at the dispatch level. +- Use one registered UDMA relay buffer per rank, as in the current bigdata path. +- Preserve the existing pass slicing, ping-pong relay reuse, token protocol, and + host-side repeat launch model. +- Use 8 local peer worker groups plus 1 remote peer worker group in multi-node mode. +- Make the remote worker group cover every remote peer on every remote server. + +Out of scope: + +- Supporting non-contiguous or non-8-rank-per-server topology. +- Changing the TileXR UDMA registry model. +- Registering multiple UDMA memory regions. +- Introducing IPC fallback for multi-node bigdata alltoall. +- Optimizing remote traffic with multiple remote worker groups in the first version. + +## Existing Bigdata Baseline + +The current `testType=7` implementation uses: + +- full `input` and `output` as ordinary GM buffers, not registered; +- one registered UDMA relay buffer; +- `passCount` kernel-internal passes based on the relay data capacity; +- two ping-pong data slots per peer; +- five logical roles per peer: + + ```text + role 0: copy input shard 0 to local UDMA send slot + role 1: copy input shard 1 to local UDMA send slot + role 2: wait for copyDone then UDMA send to peer + role 3: wait for ready, copy recv slot shard 0 to output + role 4: wait for ready, copy recv slot shard 1 to output, then ACK peer + ``` + +- monotonic pass tokens, with all waits using `observed >= token`; +- copy-side ACK checks before reusing the same ping-pong slot; +- one `SyncAll()` at kernel exit so the next host-launched kernel can reuse the + same registered UDMA memory from ping-pong slot 0 again. + +The multi-node design must reuse these worker functions where possible. It changes +which peer list a worker group processes, not the data-transfer protocol itself. + +## Multi-Node Worker Model + +The selected model is a deterministic task-queue dispatch rather than +`peer = blockIdx / 5` for all global peers. + +For `rankSize <= 8`: + +```text +blockDim = 5 * rankSize +workerGroup = blockIdx / 5 +role = blockIdx % 5 +peer = workerGroup +``` + +For `rankSize > 8`: + +```text +localRankSize = 8 +workerGroupCount = localRankSize + 1 +blockDim = 5 * workerGroupCount = 45 +workerGroup = blockIdx / 5 +role = blockIdx % 5 +``` + +Worker groups `0..7` process local-server peers: + +```text +peer = localNodeBegin + workerGroup +``` + +Worker group `8` processes all remote-server peers serially. Its five cores keep +the same role meaning as the existing 5-core per-peer design, but each role loops +over the same remote peer queue. + +The remote peer queue is deterministic and circular. It starts from `rank + 8`, +wraps at `rankSize`, skips every rank on the local server, and stops after all +`rankSize - 8` remote peers have been collected: + +```text +remoteCount = 0 +for step in 0..rankSize-1: + peer = (rank + 8 + step) % rankSize + if peer is outside [localNodeBegin, localNodeEnd): + process peer + remoteCount += 1 + if remoteCount == rankSize - 8: + break +``` + +Examples: + +```text +rankSize=16, rank=3 -> remote peers: 11, 12, 13, 14, 15, 8, 9, 10 +rankSize=16, rank=10 -> remote peers: 2, 3, 4, 5, 6, 7, 0, 1 +rankSize=24, rank=5 -> remote peers: 13..23, 8..12 +``` + +This order preserves the original requirement that remote communication starts at +`rank + 8`, while covering every card on every remote server. + +## Kernel Dispatch + +The existing bigdata kernel can split into two dispatch paths: + +```text +if rankSize <= 8: + BigDataSingleNodeDispatch() +else: + BigDataMultiNodeDispatch() +``` + +Single-node dispatch should keep the current code shape. + +Multi-node dispatch validates: + +```text +rankSize > 8 +rankSize % 8 == 0 +blockIdx < 45 +``` + +Then it assigns: + +```text +workerGroup = blockIdx / 5 +role = blockIdx % 5 + +if workerGroup < 8: + peer = localNodeBegin + workerGroup + process one local peer with the existing role worker +else: + for peer in RemotePeerQueue(rank, rankSize): + process peer with the existing role worker +``` + +The `loop` and `pass` structure remains: + +```text +for loop in 0..loopCount-1: + for pass in 0..passCount-1: + dispatch local or remote peer work +``` + +Remote role workers must use the same peer order for all five roles. This is +important because copyDone, ready, recvCopyDone, and ACK tokens are all indexed by +the real peer rank and ping-pong slot. + +## Memory Layout + +The registered UDMA relay layout remains indexed by global peer rank, not by +worker-group index. Multi-node mode does not compress remote peers into one shared +slot, because each remote peer needs independent ping-pong reuse and ACK state. + +The existing plan remains valid: + +```text +networkPeerCount = rankSize - 1 +dataSlotCount = networkPeerCount * pingPongSlots * 2 +controlGroupBytes = pingPongSlots * rankSize * localCopyShards * controlSlotBytes +signalBytes = 3 * controlGroupBytes +dataBytes = registeredBytes - controlGroupBytes - signalBytes +chunkElements = floor(dataBytes / (dataSlotCount * sizeof(int32_t))) +passCount = ceil(elementsPerPeer / chunkElements) +``` + +Data slots: + +```text +sendSlot[slot][peerIndex(peer, rank)] +recvSlot[slot][peerIndex(peer, rank)] +``` + +Control slots: + +```text +copyDone[slot][peer][copyShard] +recvCopyDone[slot][peer][recvShard] +ready[slot][peer][0] +ack through peer IPC memory indexed by [slot][rank] +``` + +All offsets must continue to pass `UDMARegisteredRangeValid`. The plan must fail +early if `chunkElements <= 0`. + +## Synchronization Protocol + +For each `(loop, pass, peer)`: + +1. Copy role 0 and role 1 wait for the old ACK before overwriting a reused + ping-pong slot when `globalPass >= pingPongSlots`. +2. Copy roles copy `input[peer][chunkOffset + shardRange]` to the local UDMA + `sendSlot[peer]` using the existing GM-UB-GM ping-pong copy helper. +3. Copy roles write `copyDone[slot][peer][shard] = token`. +4. Send role waits until both copyDone slots are `>= token`. +5. Send role posts `UDMAPutSignalNbi` from local `sendSlot[peer]` to remote + `recvSlot[rank]`, with remote ready signal set to `token`. +6. Receive roles wait until local `ready[slot][peer] >= token`. +7. Receive roles copy local `recvSlot[peer]` to `output[peer][chunkOffset]` using + the existing GM-UB-GM ping-pong copy helper. +8. Receive roles write `recvCopyDone[slot][peer][recvShard] = token`. +9. The last receive role waits until both recvCopyDone shards are `>= token`, then + writes the remote ACK token to the peer-visible ACK slot. +10. Kernel exit performs `SyncAll()` once, after all local and remote peer queues + are complete. + +No per-pass `SyncAll()` is required. Slot reuse is protected by ACK tokens, and +using a per-pass global barrier would serialize unrelated local and remote work. + +All token checks must compare `observed >= token`, not `observed == token`, because +producer roles can advance and write a later token before the consumer polls. + +## Ascend C API Constraints + +The local copies remain GM-UB-GM transfers. There is no direct GM-to-GM copy path. + +Implementation rules: + +- Continue using the existing 2 x 64 KiB UB ping-pong buffer for relay copies. +- Use `DataCopyPad` for copy-in and copy-out, including tails that are not known + to be 32-byte aligned. +- Keep MTE2/MTE3 synchronization in the copy helper through explicit flags or an + equivalent queue-based synchronization point. +- Do not introduce production `GlobalTensor::GetValue()` or `SetValue()` for data + movement. Scalar token load/store helpers are acceptable only for control words + where the current implementation already uses them. +- Keep `chunkBytes` within `uint32_t`, matching the current UDMA API usage. + +## Host-Side Changes + +`PlanAllToAllBigDataUdma(rankSize, elementsPerPeer)` remains global-rank based and +does not need a separate remote-worker layout. + +`AllToAllBigDataBlockDim(rankSize)` should preserve single-node behavior and return +the multi-node worker count when `rankSize > 8`: + +```text +rankSize <= 8: 5 * rankSize +rankSize > 8 : 5 * (8 + 1) +``` + +The helper should use explicit names in implementation to avoid confusing global +rank size with local node rank size, for example: + +```text +kAllToAllBigDataRanksPerNode = 8 +kAllToAllBigDataMultiNodeWorkerGroups = kAllToAllBigDataRanksPerNode + 1 +``` + +Before allocating and registering memory, the host should reject: + +```text +rankSize > 8 && rankSize % 8 != 0 +``` + +The host repeat model stays unchanged: + +- `TILEXR_DEMO_ALLTOALL_REPEAT` launches repeated kernels from host; +- each kernel uses `loopCount=1`; +- the stream is synchronized once after the host launch loop; +- `kernelLoopBase = iter` keeps tokens globally monotonic across kernels. + +## Debug And Error Handling + +The existing debug array should keep the current first-pass fields. Multi-node mode +should add or reuse fields for: + +- multi-node mode enabled flag; +- `localNodeBegin`, `localNodeEnd`, and active `blockDim`; +- remote peer count; +- first remote peer and last remote peer seen by the remote worker; +- timeout status per peer where debug capacity allows. + +Timeout behavior must remain strict: a timeout or non-zero UDMA quiet status should +fail the test rather than falling back to IPC. + +If `rankSize > 8` and topology validation fails, the demo should print a clear +host-side error before `TileXRUDMARegister`. + +## Testing Plan + +Unit/source tests: + +- `AllToAllBigDataBlockDim(8) == 40`. +- `AllToAllBigDataBlockDim(16) == 45`. +- `AllToAllBigDataBlockDim(32) == 45`. +- Remote peer queue examples: + - `rankSize=16, rank=3` gives `11,12,13,14,15,8,9,10`. + - `rankSize=16, rank=10` gives `2,3,4,5,6,7,0,1`. +- Source guard verifies `rankSize <= 8` keeps the existing dispatch. +- Source guard verifies `rankSize > 8` uses the multi-node remote peer queue. +- Source guard verifies token waits use `>= token`. + +On-board validation: + +1. Re-run existing 8P `testType=7` correctness and performance. The kernel name + and path stay the same, but the single-node dispatch must match prior behavior. +2. Run a 16P small payload correctness test across two servers. +3. Run a 16P per-peer 2 MiB payload test with `TILEXR_DEMO_ALLTOALL_REPEAT=100` + and msprof kernel timing. +4. Run a larger multi-pass payload and confirm `passCount > 1`, correctness passes, + and no ACK/ready timeout appears. + +Success criteria: + +- 8P results do not regress beyond normal measurement noise. +- 16P and larger multi-node runs finish without IPC fallback. +- All ranks pass output validation. +- No UDMA range validation failure, CQ error, ready timeout, ACK timeout, or data + mismatch appears in logs. + +## Performance Expectations + +The multi-node first version intentionally serializes remote peers through one +5-core worker group. Therefore: + +- local-server traffic remains parallel across 8 peer groups; +- remote traffic time scales roughly with the number of remote peers; +- the implementation prioritizes correctness and controlled resource usage over + maximum cross-server bandwidth; +- future optimization can add more remote worker groups after this protocol is + validated. + +This design is expected to be slower than a fully parallel `5 * rankSize` remote +dispatch, but it avoids launching one 5-core worker group for every remote rank and +keeps synchronization behavior tractable. From 33beb61a5a1566c99ecf706c2d2a3892ed6a1e5d Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 29 Jun 2026 21:28:52 +0800 Subject: [PATCH 017/163] perf(udma): add multinode bigdata slot layout Use the 35-core bigdata all-to-all path for multi-node ranks and optional forced 8P testing. Allocate a 1024 MiB registered UDMA buffer with fixed 8 MiB send/recv peer slots per pass so the multi-node kernel does not reuse pingpong data slots internally. Keep the single-node bigdata path on the existing pingpong layout and update the demo launcher plus layout guards for the new offsets, worker roles, and topology checks. --- tests/udma/demo/tilexr_udma_alltoall_layout.h | 211 +++++- tests/udma/demo/tilexr_udma_demo.cpp | 77 ++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 625 ++++++++++++++++-- .../unit/test_tilexr_udma_alltoall_layout.cpp | 211 +++++- 4 files changed, 997 insertions(+), 127 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index faaf32d7..3bc60ab5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -16,7 +16,7 @@ namespace TileXR { namespace Demo { constexpr int32_t kAllToAllBaseValue = 100000; -constexpr size_t kAllToAllUdmaMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllUdmaMaxRegisteredBytes = 512ULL * 1024ULL * 1024ULL; struct AllToAllChunkPlan { uint32_t passCount = 1; @@ -25,11 +25,23 @@ struct AllToAllChunkPlan { size_t registeredBytes = 0; }; -constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 64ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataMultiNodeRegisteredBytes = 1024ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataMultiNodePeerSlotBytes = 8ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataControlSlotBytes = 64ULL; constexpr uint32_t kAllToAllBigDataCoresPerPeer = 5U; -constexpr uint32_t kAllToAllBigDataLocalCopyShards = 2U; +constexpr uint32_t kAllToAllBigDataSingleNodeShards = 2U; +constexpr uint32_t kAllToAllBigDataLocalCopyShards = kAllToAllBigDataSingleNodeShards; constexpr uint32_t kAllToAllBigDataPingPongSlots = 2U; +constexpr int32_t kAllToAllBigDataRanksPerNode = 8; +constexpr uint32_t kAllToAllBigDataMultiNodeCopyCores = 16U; +constexpr uint32_t kAllToAllBigDataMultiNodeRecvCores = 16U; +constexpr uint32_t kAllToAllBigDataMultiNodeRemoteSendPrimaryCore = 16U; +constexpr uint32_t kAllToAllBigDataMultiNodeRemoteSendSecondaryCore = 17U; +constexpr uint32_t kAllToAllBigDataMultiNodeLocalSendCore = 18U; +constexpr uint32_t kAllToAllBigDataMultiNodeRecvCoreBase = 19U; +constexpr uint32_t kAllToAllBigDataMultiNodeBlockDim = + kAllToAllBigDataMultiNodeRecvCoreBase + kAllToAllBigDataMultiNodeRecvCores; struct AllToAllBigDataPlan { uint32_t passCount = 1; @@ -38,6 +50,7 @@ struct AllToAllBigDataPlan { size_t dataBytes = 0; size_t copyDoneOffset = 0; size_t recvCopyDoneOffset = 0; + size_t remoteSendDoneOffset = 0; size_t readySignalOffset = 0; size_t ackSignalOffset = 0; size_t controlBytes = 0; @@ -50,6 +63,23 @@ inline int32_t AllToAllValue(int srcRank, int dstRank) return kAllToAllBaseValue + srcRank * 1000 + dstRank; } +inline bool AllToAllBigDataIsMultiNode(int rankSize) +{ + return rankSize > kAllToAllBigDataRanksPerNode; +} + +inline bool AllToAllBigDataUse35Core(int rankSize, bool force35Core = false) +{ + return AllToAllBigDataIsMultiNode(rankSize) || + (force35Core && rankSize == kAllToAllBigDataRanksPerNode); +} + +inline uint32_t AllToAllBigDataShardCount(int rankSize, bool force35Core = false) +{ + return AllToAllBigDataUse35Core(rankSize, force35Core) ? + kAllToAllBigDataMultiNodeCopyCores : kAllToAllBigDataSingleNodeShards; +} + inline AllToAllChunkPlan PlanAllToAllUdmaChunks(int rankSize, int32_t elementsPerPeer) { AllToAllChunkPlan plan {}; @@ -80,56 +110,187 @@ inline AllToAllChunkPlan PlanAllToAllUdmaChunks(int rankSize, int32_t elementsPe return plan; } -inline AllToAllBigDataPlan PlanAllToAllBigDataUdma(int rankSize, int32_t elementsPerPeer) +inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( + int rankSize, int32_t elementsPerPeer, bool force35Core = false) { AllToAllBigDataPlan plan {}; if (rankSize <= 0 || elementsPerPeer <= 0) { return plan; } - plan.registeredBytes = kAllToAllBigDataMaxRegisteredBytes; + const bool use35Core = AllToAllBigDataUse35Core(rankSize, force35Core); + plan.registeredBytes = use35Core ? kAllToAllBigDataMultiNodeRegisteredBytes : + kAllToAllBigDataMaxRegisteredBytes; + const uint32_t shardCount = AllToAllBigDataShardCount(rankSize, force35Core); + const size_t maxChunkElements = use35Core ? + kAllToAllBigDataMultiNodePeerSlotBytes / sizeof(int32_t) : + static_cast(elementsPerPeer); + size_t chunkElements = static_cast(elementsPerPeer); + if (use35Core && chunkElements > maxChunkElements) { + chunkElements = maxChunkElements; + } + if (chunkElements == 0) { + chunkElements = 1; + } + plan.passCount = static_cast( + (static_cast(elementsPerPeer) + chunkElements - 1) / chunkElements); + const size_t slotCountForControls = use35Core ? + static_cast(plan.passCount) : static_cast(kAllToAllBigDataPingPongSlots); const size_t controlGroupBytes = - static_cast(kAllToAllBigDataPingPongSlots) * static_cast(rankSize) * - static_cast(kAllToAllBigDataLocalCopyShards) * kAllToAllBigDataControlSlotBytes; + slotCountForControls * static_cast(rankSize) * + static_cast(shardCount) * kAllToAllBigDataControlSlotBytes; plan.controlBytes = controlGroupBytes; - plan.signalBytes = 3ULL * controlGroupBytes; + plan.signalBytes = use35Core ? 4ULL * controlGroupBytes : 3ULL * controlGroupBytes; if (plan.controlBytes + plan.signalBytes >= plan.registeredBytes) { return plan; } const size_t networkPeerCount = static_cast(rankSize > 1 ? rankSize - 1 : 1); - const size_t dataSlotCount = - networkPeerCount * static_cast(kAllToAllBigDataPingPongSlots) * 2ULL; + const size_t dataSlotCount = networkPeerCount * + (use35Core ? static_cast(plan.passCount) : + static_cast(kAllToAllBigDataPingPongSlots)) * 2ULL; plan.dataBytes = plan.registeredBytes - plan.controlBytes - plan.signalBytes; - plan.chunkElements = static_cast( - plan.dataBytes / (dataSlotCount * sizeof(int32_t))); - if (plan.chunkElements <= 0) { - plan.chunkElements = 1; - } - if (plan.chunkElements > elementsPerPeer) { - plan.chunkElements = elementsPerPeer; + if (use35Core) { + plan.chunkElements = static_cast(chunkElements); + plan.chunkBytesPerPeer = kAllToAllBigDataMultiNodePeerSlotBytes; + } else { + plan.chunkElements = static_cast( + plan.dataBytes / (dataSlotCount * sizeof(int32_t))); + if (plan.chunkElements <= 0) { + plan.chunkElements = 1; + } + if (plan.chunkElements > elementsPerPeer) { + plan.chunkElements = elementsPerPeer; + } + plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); } - plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); plan.dataBytes = dataSlotCount * plan.chunkBytesPerPeer; + if (plan.dataBytes + plan.controlBytes + plan.signalBytes > plan.registeredBytes) { + return AllToAllBigDataPlan {}; + } plan.copyDoneOffset = plan.dataBytes; plan.recvCopyDoneOffset = plan.copyDoneOffset + controlGroupBytes; - plan.readySignalOffset = plan.recvCopyDoneOffset + controlGroupBytes; + if (use35Core) { + plan.remoteSendDoneOffset = plan.recvCopyDoneOffset + controlGroupBytes; + plan.readySignalOffset = plan.remoteSendDoneOffset + controlGroupBytes; + plan.signalBytes = 4ULL * controlGroupBytes; + } else { + plan.remoteSendDoneOffset = 0ULL; + plan.readySignalOffset = plan.recvCopyDoneOffset + controlGroupBytes; + plan.signalBytes = 3ULL * controlGroupBytes; + } plan.ackSignalOffset = plan.readySignalOffset + controlGroupBytes; plan.controlBytes = controlGroupBytes; - plan.signalBytes = 3ULL * controlGroupBytes; - plan.registeredBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; - plan.passCount = static_cast( - (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1) / - static_cast(plan.chunkElements)); return plan; } -inline uint32_t AllToAllBigDataBlockDim(int rankSize) +inline bool AllToAllBigDataValidTopology(int rankSize) +{ + if (rankSize <= 0) { + return false; + } + if (!AllToAllBigDataIsMultiNode(rankSize)) { + return true; + } + return rankSize % kAllToAllBigDataRanksPerNode == 0; +} + +inline int32_t AllToAllBigDataLocalNodeBegin(int rank) +{ + return (rank / kAllToAllBigDataRanksPerNode) * kAllToAllBigDataRanksPerNode; +} + +inline int32_t AllToAllBigDataLocalNodeEnd(int rank) +{ + return AllToAllBigDataLocalNodeBegin(rank) + kAllToAllBigDataRanksPerNode; +} + +inline int32_t AllToAllBigDataNodeCount(int rankSize) +{ + if (!AllToAllBigDataValidTopology(rankSize)) { + return 0; + } + if (!AllToAllBigDataIsMultiNode(rankSize)) { + return 1; + } + return rankSize / kAllToAllBigDataRanksPerNode; +} + +inline bool AllToAllBigDataIsLocalPeer(int rank, int peer) +{ + const int32_t begin = AllToAllBigDataLocalNodeBegin(rank); + return peer >= begin && peer < begin + kAllToAllBigDataRanksPerNode; +} + +inline int32_t AllToAllBigDataLocalPeerForWorker(int rank, uint32_t workerGroup) +{ + if (workerGroup >= static_cast(kAllToAllBigDataRanksPerNode)) { + return -1; + } + return AllToAllBigDataLocalNodeBegin(rank) + static_cast(workerGroup); +} + +inline std::vector AllToAllBigDataRemotePeers(int rank, int rankSize) +{ + std::vector peers; + if (!AllToAllBigDataIsMultiNode(rankSize) || !AllToAllBigDataValidTopology(rankSize)) { + return peers; + } + peers.reserve(static_cast(rankSize - kAllToAllBigDataRanksPerNode)); + for (int32_t step = 0; step < rankSize && + peers.size() < static_cast(rankSize - kAllToAllBigDataRanksPerNode); ++step) { + const int32_t peer = (rank + kAllToAllBigDataRanksPerNode + step) % rankSize; + if (!AllToAllBigDataIsLocalPeer(rank, peer)) { + peers.push_back(peer); + } + } + return peers; +} + +inline std::vector AllToAllBigDataLocalPeers(int rank) +{ + std::vector peers; + peers.reserve(static_cast(kAllToAllBigDataRanksPerNode - 1)); + const int32_t begin = AllToAllBigDataLocalNodeBegin(rank); + const int32_t local = rank - begin; + for (int32_t step = 1; step < kAllToAllBigDataRanksPerNode; ++step) { + peers.push_back(begin + ((local + step) % kAllToAllBigDataRanksPerNode)); + } + return peers; +} + +inline std::vector AllToAllBigDataMergedPeerTasks(int rank, int rankSize) +{ + std::vector tasks; + if (!AllToAllBigDataIsMultiNode(rankSize) || !AllToAllBigDataValidTopology(rankSize)) { + return tasks; + } + const std::vector remotePeers = AllToAllBigDataRemotePeers(rank, rankSize); + const std::vector localPeers = AllToAllBigDataLocalPeers(rank); + const int32_t remoteBurst = AllToAllBigDataNodeCount(rankSize) - 1; + size_t remoteIndex = 0; + size_t localIndex = 0; + tasks.reserve(remotePeers.size() + localPeers.size()); + while (remoteIndex < remotePeers.size() || localIndex < localPeers.size()) { + for (int32_t i = 0; i < remoteBurst && remoteIndex < remotePeers.size(); ++i) { + tasks.push_back(remotePeers[remoteIndex++]); + } + if (localIndex < localPeers.size()) { + tasks.push_back(localPeers[localIndex++]); + } + } + return tasks; +} + +inline uint32_t AllToAllBigDataBlockDim(int rankSize, bool force35Core = false) { if (rankSize <= 0) { return 1U; } + if (AllToAllBigDataUse35Core(rankSize, force35Core)) { + return kAllToAllBigDataMultiNodeBlockDim; + } return static_cast(rankSize) * kAllToAllBigDataCoresPerPeer; } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index de71e4bf..e887189d 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -50,9 +50,10 @@ extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, - uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, + uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, + uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, - uint32_t profileStage); + uint32_t profileStage, uint32_t force35Core); extern void launch_tilexr_all_to_all_ipc_scatter( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR debug, int32_t elementsPerPeer); extern void launch_tilexr_all_to_all_ipc_gather( @@ -92,6 +93,7 @@ constexpr int kConnectRetrySleepMs = 10; constexpr int kBigDataProfileStageFull = 8; struct BarrierEndpoint { + std::string host; uint16_t port; }; @@ -204,11 +206,19 @@ bool CopyDeviceToHost(int rank, void* dst, size_t dstSize, const void* src, size BarrierEndpoint GetBarrierEndpoint() { + std::string host = "127.0.0.1"; int basePort = kDefaultCommPort; + const char* barrierHost = std::getenv("TILEXR_DEMO_BARRIER_HOST"); + if (barrierHost != nullptr && barrierHost[0] != '\0') { + host = barrierHost; + } const char* commId = std::getenv("TILEXR_COMM_ID"); if (commId != nullptr) { std::string value(commId); size_t colon = value.rfind(':'); + if ((barrierHost == nullptr || barrierHost[0] == '\0') && colon != std::string::npos && colon > 0) { + host = value.substr(0, colon); + } if (colon != std::string::npos && colon + 1 < value.size()) { basePort = std::atoi(value.c_str() + colon + 1); } @@ -217,7 +227,7 @@ BarrierEndpoint GetBarrierEndpoint() if (barrierPort <= 0 || barrierPort > 65535) { barrierPort = kDefaultCommPort + kDemoBarrierPortOffset; } - return BarrierEndpoint{static_cast(barrierPort)}; + return BarrierEndpoint{host, static_cast(barrierPort)}; } bool SendAll(int fd, const void* data, size_t bytes) @@ -273,7 +283,7 @@ int CreateBarrierServer(uint16_t port) } sockaddr_in addr{}; addr.sin_family = AF_INET; - addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + addr.sin_addr.s_addr = htonl(INADDR_ANY); addr.sin_port = htons(port); if (bind(fd, reinterpret_cast(&addr), sizeof(addr)) != 0 || listen(fd, SOMAXCONN) != 0) { @@ -283,11 +293,13 @@ int CreateBarrierServer(uint16_t port) return fd; } -int ConnectBarrierServer(uint16_t port) +int ConnectBarrierServer(const std::string& host, uint16_t port) { sockaddr_in addr{}; addr.sin_family = AF_INET; - addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + if (inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) { + return -1; + } addr.sin_port = htons(port); for (int attempt = 0; attempt < kConnectRetryCount; ++attempt) { @@ -311,14 +323,15 @@ bool DemoBarrierAll(int rank, int rankSize, const std::string& step) } BarrierEndpoint endpoint = GetBarrierEndpoint(); - PrintStatus(rank, "demo tcp barrier begin: " + step + " port=" + std::to_string(endpoint.port)); + PrintStatus(rank, "demo tcp barrier begin: " + step + + " host=" + endpoint.host + " port=" + std::to_string(endpoint.port)); constexpr uint8_t kArrive = 1; constexpr uint8_t kRelease = 2; if (rank == 0) { int serverFd = CreateBarrierServer(endpoint.port); if (serverFd < 0) { - std::cerr << "[rank " << rank << "] ERROR: failed to create demo barrier server on 127.0.0.1:" + std::cerr << "[rank " << rank << "] ERROR: failed to create demo barrier server on 0.0.0.0:" << endpoint.port << ", errno=" << errno << std::endl; return false; } @@ -349,10 +362,10 @@ bool DemoBarrierAll(int rank, int rankSize, const std::string& step) return false; } } else { - int fd = ConnectBarrierServer(endpoint.port); + int fd = ConnectBarrierServer(endpoint.host, endpoint.port); if (fd < 0) { - std::cerr << "[rank " << rank << "] ERROR: failed to connect demo barrier on 127.0.0.1:" - << endpoint.port << std::endl; + std::cerr << "[rank " << rank << "] ERROR: failed to connect demo barrier on " + << endpoint.host << ":" << endpoint.port << std::endl; return false; } uint8_t release = 0; @@ -660,6 +673,8 @@ int main(int argc, char** argv) GetEnvInt("TILEXR_DEMO_BIGDATA_PROFILE_STAGE", kBigDataProfileStageFull) : kBigDataProfileStageFull; bigDataProfileStage = std::max(0, std::min(bigDataProfileStage, kBigDataProfileStageFull)); + const bool forceBigData35Core = + testType == 7 && GetEnvInt("TILEXR_DEMO_BIGDATA_FORCE_35CORE", 0) != 0; bool bigDataProfilePartial = testType == 7 && bigDataProfileStage < kBigDataProfileStageFull; bool syncAllToAllAtEnd = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_SYNC_AT_END", 0) != 0; @@ -682,8 +697,17 @@ int main(int argc, char** argv) isAllToAll ? TileXR::Demo::PlanAllToAllUdmaChunks(rankSize, elementsPerRank) : TileXR::Demo::AllToAllChunkPlan {}; const TileXR::Demo::AllToAllBigDataPlan bigDataPlan = - isAllToAll ? TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerRank) : + isAllToAll ? TileXR::Demo::PlanAllToAllBigDataUdma( + rankSize, elementsPerRank, forceBigData35Core) : TileXR::Demo::AllToAllBigDataPlan {}; + if (testType == 7 && !TileXR::Demo::AllToAllBigDataValidTopology(rankSize)) { + std::cerr << "[rank " << rank + << "] ERROR: bigdata alltoall multi-node requires rankSize multiple of 8" + << " when rankSize > " << TileXR::Demo::kAllToAllBigDataRanksPerNode + << ", rankSize=" << rankSize << std::endl; + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } if (isAllToAll) { const size_t dataAsFlagStagingBytes = AllToAllDataAsFlagStagingBytes(rankSize, elementsPerRank); const size_t plainIpcStagingBytes = AllToAllPlainIpcStagingBytes(rankSize, elementsPerRank); @@ -712,6 +736,13 @@ int main(int argc, char** argv) if (testType == 7) { PrintStatus(rank, "bigdata profile stage=" + std::to_string(bigDataProfileStage) + " fullStage=" + std::to_string(kBigDataProfileStageFull)); + PrintStatus(rank, "bigdata multinode mode=" + + std::string(TileXR::Demo::AllToAllBigDataIsMultiNode(rankSize) ? "true" : "false") + + " force35Core=" + std::string(forceBigData35Core ? "true" : "false") + + " blockDim=" + std::to_string(TileXR::Demo::AllToAllBigDataBlockDim( + rankSize, forceBigData35Core)) + + " shards=" + std::to_string(TileXR::Demo::AllToAllBigDataShardCount( + rankSize, forceBigData35Core))); } PrintStatus(rank, "alltoall UDMA chunk plan: passCount=" + std::to_string(chunkPlan.passCount) + " chunkElements=" + std::to_string(chunkPlan.chunkElements) + @@ -722,6 +753,7 @@ int main(int argc, char** argv) " dataBytes=" + std::to_string(bigDataPlan.dataBytes) + " copyDoneOffset=" + std::to_string(bigDataPlan.copyDoneOffset) + " recvCopyDoneOffset=" + std::to_string(bigDataPlan.recvCopyDoneOffset) + + " remoteSendDoneOffset=" + std::to_string(bigDataPlan.remoteSendDoneOffset) + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + " registeredBytes=" + std::to_string(bigDataPlan.registeredBytes)); @@ -900,6 +932,7 @@ int main(int argc, char** argv) PrintStatus(rank, "bigdata alltoall registered dataBytes=" + std::to_string(bigDataPlan.dataBytes) + " copyDoneOffset=" + std::to_string(bigDataPlan.copyDoneOffset) + " recvCopyDoneOffset=" + std::to_string(bigDataPlan.recvCopyDoneOffset) + + " remoteSendDoneOffset=" + std::to_string(bigDataPlan.remoteSendDoneOffset) + " readySignalOffset=" + std::to_string(bigDataPlan.readySignalOffset) + " ackSignalOffset=" + std::to_string(bigDataPlan.ackSignalOffset) + " regBytes=" + std::to_string(registeredBytes) + @@ -919,18 +952,22 @@ int main(int argc, char** argv) return 1; } + const uint32_t bigDataBlockDim = TileXR::Demo::AllToAllBigDataBlockDim( + rankSize, forceBigData35Core); auto a2aStart = std::chrono::steady_clock::now(); for (int iter = 0; iter < allToAllRepeat; ++iter) { const uint64_t kernelLoopBase = static_cast(iter); launch_tilexr_udma_all_to_all_bigdata( - TileXR::Demo::AllToAllBigDataBlockDim(rankSize), stream, commArgsDev, + bigDataBlockDim, stream, commArgsDev, reinterpret_cast(bigInput), reinterpret_cast(bigOutput), reinterpret_cast(registeredMemory), reinterpret_cast(debug), elementsPerRank, 0, bigDataPlan.copyDoneOffset, bigDataPlan.recvCopyDoneOffset, + bigDataPlan.remoteSendDoneOffset, bigDataPlan.readySignalOffset, bigDataPlan.ackSignalOffset, bigDataPlan.chunkElements, bigDataPlan.passCount, 1, kernelLoopBase, - static_cast(bigDataProfileStage)); + static_cast(bigDataProfileStage), + forceBigData35Core ? 1U : 0U); } if (!CheckAcl(rank, "aclrtSynchronizeStream post-alltoall", aclrtSynchronizeStream(stream))) { if (udmaRegistered) { @@ -998,7 +1035,7 @@ int main(int argc, char** argv) // Registered relay region: [udmaMem chunk | signals]. const size_t fusedChunkBytes = chunkPlan.chunkBytesPerRank; const size_t fusedSignalBytes = static_cast(rankSize) * sizeof(uint64_t); - const size_t fusedRegBytes = ((fusedChunkBytes + fusedSignalBytes + kUdmaRegistrationAlignment - 1) / + const size_t fusedRegBytes = ((chunkPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; if (!udmaRegistered) { int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), fusedRegBytes, &udmaHandle); @@ -1064,6 +1101,14 @@ int main(int argc, char** argv) << " iters(total=" << chunkPlan.passCount << " pass/iter) total=" << a2aMs << " ms perIter=" << a2aPerIterUs << " us payload=" << payloadBytes << " bytes bw=" << bwGbs << " GB/s" << std::endl; } + if (!CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), + fusedOutput, dataCount * sizeof(int32_t), "fused alltoall output")) { + if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; } + aclrtFree(fusedInput); + aclrtFree(fusedOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; } aclrtFree(fusedInput); aclrtFree(fusedOutput); @@ -1419,7 +1464,7 @@ int main(int argc, char** argv) copyBackOk = CopyDeviceToHost(rank, hostData.data(), dataCount * sizeof(int32_t), data, dataCount * sizeof(int32_t), "data"); } - if (hasOutput && !chunkedStrictAllToAll && testType != 7) { + if (hasOutput && !chunkedStrictAllToAll && testType != 6 && testType != 7) { const char* outputName = isAllToAll ? "alltoall output" : "allreduce output"; copyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), output, dataCount * sizeof(int32_t), outputName) && copyBackOk; diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 7767ce71..bde5829e 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -52,8 +52,24 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES = 64 * 1024; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES = TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES * 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U; -constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = 2U; +constexpr uint64_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 8ULL * 1024ULL * 1024ULL; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = + TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 5U; +constexpr int32_t TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE = 8; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES = 16U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES = 16U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE = 16U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE = 17U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE = 18U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE = 19U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM = + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE + + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT = 0U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT = 1U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_PREPARE = 0; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY = 1; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC = 2; @@ -84,14 +100,14 @@ __aicore__ inline __gm__ uint64_t* ControlSlot(__gm__ uint8_t* base, uint64_t of } __aicore__ inline __gm__ uint64_t* BigDataControlSlot( - __gm__ uint8_t* base, uint64_t offset, uint32_t slot, int32_t rankSize, int32_t peer, - uint32_t shard) + __gm__ uint8_t* base, uint64_t offset, uint32_t slot, int32_t rankSize, + uint32_t shardCount, int32_t peer, uint32_t shard) { return reinterpret_cast<__gm__ uint64_t*>( base + offset + ((static_cast(slot) * static_cast(rankSize) + static_cast(peer)) * - static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) + + static_cast(shardCount) + static_cast(shard)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); } @@ -180,9 +196,10 @@ __aicore__ inline uint32_t BigDataTileBytes(uint32_t totalBytes, uint32_t offset } __aicore__ inline bool BigDataCopyShardRange( - uint32_t shard, uint32_t totalElements, uint32_t& shardOffsetBytes, uint32_t& shardBytes) + uint32_t shard, uint32_t shardCount, uint32_t totalElements, + uint32_t& shardOffsetBytes, uint32_t& shardBytes) { - if (shard >= TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) { + if (shardCount == 0U || shard >= shardCount) { shardOffsetBytes = 0U; shardBytes = 0U; return false; @@ -190,13 +207,13 @@ __aicore__ inline bool BigDataCopyShardRange( constexpr uint32_t alignElements = 32U / sizeof(int32_t); const uint64_t begin = static_cast(totalElements) * static_cast(shard) / - static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS); + static_cast(shardCount); const uint64_t end = static_cast(totalElements) * static_cast(shard + 1U) / - static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS); + static_cast(shardCount); const uint64_t alignedBegin = (shard == 0U) ? 0ULL : (begin / static_cast(alignElements)) * static_cast(alignElements); - const uint64_t alignedEnd = (shard + 1U == TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) ? + const uint64_t alignedEnd = (shard + 1U == shardCount) ? static_cast(totalElements) : (end / static_cast(alignElements)) * static_cast(alignElements); shardOffsetBytes = static_cast(alignedBegin * sizeof(int32_t)); @@ -204,6 +221,21 @@ __aicore__ inline bool BigDataCopyShardRange( return shardBytes > 0U; } +__aicore__ inline uint32_t BigDataCopyShardStartBytes( + uint32_t shard, uint32_t shardCount, uint32_t totalElements) +{ + if (shardCount == 0U || shard >= shardCount) { + return 0U; + } + constexpr uint32_t alignElements = 32U / sizeof(int32_t); + const uint64_t begin = + static_cast(totalElements) * static_cast(shard) / + static_cast(shardCount); + const uint64_t alignedBegin = (shard == 0U) ? 0ULL : + (begin / static_cast(alignElements)) * static_cast(alignElements); + return static_cast(alignedBegin * sizeof(int32_t)); +} + __aicore__ inline void BigDataCopyRangePingPong( __gm__ uint8_t* dst, __gm__ uint8_t* src, uint32_t bytes, AscendC::LocalTensor relayLocal) @@ -317,11 +349,135 @@ __aicore__ inline uint32_t BigDataPingPongSlot(uint64_t globalPass) static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS - 1U)); } +__aicore__ inline uint32_t BigDataDataSlot(uint64_t globalPass, uint32_t pass, bool use35Core) +{ + return use35Core ? pass : BigDataPingPongSlot(globalPass); +} + __aicore__ inline void BigDataKernelExitBarrier() { AscendC::SyncAll(); } +__aicore__ inline bool BigDataIsMultiNode(int32_t rankSize) +{ + return rankSize > TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; +} + +__aicore__ inline bool BigDataUse35Core(int32_t rankSize, bool force35Core) +{ + return BigDataIsMultiNode(rankSize) || + (force35Core && rankSize == TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE); +} + +__aicore__ inline uint32_t BigDataShardCount(int32_t rankSize, bool force35Core = false) +{ + return BigDataUse35Core(rankSize, force35Core) ? + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES : + TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS; +} + +__aicore__ inline bool BigDataValidTopology(int32_t rankSize) +{ + if (rankSize <= 0) { + return false; + } + if (!BigDataIsMultiNode(rankSize)) { + return true; + } + return rankSize % TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE == 0; +} + +__aicore__ inline int32_t BigDataLocalNodeBegin(int32_t rank) +{ + return (rank / TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) * + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; +} + +__aicore__ inline int32_t BigDataNodeCount(int32_t rankSize) +{ + if (!BigDataValidTopology(rankSize)) { + return 0; + } + if (!BigDataIsMultiNode(rankSize)) { + return 1; + } + return rankSize / TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; +} + +__aicore__ inline int32_t BigDataTaskCount(int32_t rankSize, bool force35Core) +{ + if (force35Core && rankSize == TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) { + return TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - 1; + } + return rankSize - 1; +} + +__aicore__ inline bool BigDataIsLocalPeer(int32_t rank, int32_t peer) +{ + const int32_t begin = BigDataLocalNodeBegin(rank); + return peer >= begin && peer < begin + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; +} + +__aicore__ inline int32_t BigDataRemotePeerAt(int32_t rank, int32_t rankSize, int32_t remoteIndex) +{ + if (!BigDataIsMultiNode(rankSize) || !BigDataValidTopology(rankSize) || remoteIndex < 0) { + return -1; + } + int32_t remoteCount = 0; + for (int32_t step = 0; step < rankSize; ++step) { + const int32_t peer = + (rank + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE + step) % rankSize; + if (!BigDataIsLocalPeer(rank, peer)) { + if (remoteCount == remoteIndex) { + return peer; + } + ++remoteCount; + } + } + return -1; +} + +__aicore__ inline int32_t BigDataLocalPeerAt(int32_t rank, int32_t localIndex) +{ + if (localIndex < 0 || localIndex >= TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - 1) { + return -1; + } + const int32_t begin = BigDataLocalNodeBegin(rank); + const int32_t local = rank - begin; + return begin + ((local + localIndex + 1) % TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE); +} + +__aicore__ inline bool BigDataMergedPeerTaskAt( + int32_t rank, int32_t rankSize, int32_t taskIndex, bool force35Core, + int32_t& peer, bool& isLocalPeer) +{ + peer = -1; + isLocalPeer = false; + if (force35Core && rankSize == TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) { + peer = BigDataLocalPeerAt(rank, taskIndex); + isLocalPeer = true; + return peer >= 0; + } + const int32_t nodeCount = BigDataNodeCount(rankSize); + if (nodeCount <= 1 || taskIndex < 0) { + return false; + } + const int32_t remoteBurst = nodeCount - 1; + const int32_t groupSize = remoteBurst + 1; + const int32_t group = taskIndex / groupSize; + const int32_t indexInGroup = taskIndex % groupSize; + if (indexInGroup < remoteBurst) { + const int32_t remoteIndex = group * remoteBurst + indexInGroup; + peer = BigDataRemotePeerAt(rank, rankSize, remoteIndex); + isLocalPeer = false; + return peer >= 0; + } + peer = BigDataLocalPeerAt(rank, group); + isLocalPeer = true; + return peer >= 0; +} + __aicore__ inline int32_t BigDataNetworkPeerIndex(int32_t peer, int32_t rank) { return peer < rank ? peer : peer - 1; @@ -400,6 +556,25 @@ __aicore__ inline __gm__ uint64_t* BigDataRemoteControlSlot( TileXR::UDMARegisteredRemoteAddr(registry, targetRank, remoteOffset)); } +__aicore__ inline __gm__ uint64_t* BigDataRemoteRegisteredControlSlot( + const __gm__ TileXR::CommArgs* args, int32_t targetRank, uint64_t offset, + uint32_t slot, int32_t rankSize, uint32_t shardCount, int32_t slotRank, uint32_t shard) +{ + auto registry = TileXR::GetUDMARegistry(args); + const uint64_t remoteOffset = + offset + + ((static_cast(slot) * static_cast(rankSize) + + static_cast(slotRank)) * + static_cast(shardCount) + + static_cast(shard)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + if (!TileXR::UDMARegisteredRangeValid(registry, targetRank, + remoteOffset, TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES)) { + return nullptr; + } + return reinterpret_cast<__gm__ uint64_t*>( + TileXR::UDMARegisteredRemoteAddr(registry, targetRank, remoteOffset)); +} + __aicore__ inline uint64_t BigDataIpcAckOffset(uint32_t slot, int32_t rankSize, int32_t slotRank) { constexpr uint64_t maxAckBytes = @@ -431,16 +606,16 @@ __aicore__ inline void BigDataCopyPeerWorker( __gm__ int32_t* input, __gm__ int32_t* output, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, - uint32_t copyShard, uint64_t sendDataOffset, uint64_t copyDoneOffset, uint64_t ackSignalOffset, - uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) + bool use35Core, uint32_t copyShard, uint32_t shardCount, uint64_t sendDataOffset, + uint64_t copyDoneOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, + AscendC::LocalTensor relayLocal) { - if (peer < 0 || peer >= rankSize || copyShard >= TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) { + if (peer < 0 || peer >= rankSize || shardCount == 0U || copyShard >= shardCount) { return; } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE) { return; } - (void)ackSignalOffset; const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; int32_t chunkOffset = 0; @@ -450,23 +625,24 @@ __aicore__ inline void BigDataCopyPeerWorker( } uint32_t shardOffset = 0U; uint32_t shardBytes = 0U; - if (!BigDataCopyShardRange(copyShard, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { - return; - } + const bool shardHasBytes = + BigDataCopyShardRange(copyShard, shardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes); const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); - const uint32_t slot = BigDataPingPongSlot(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); if (peer == rank) { auto src = reinterpret_cast<__gm__ uint8_t*>( input + static_cast(rank) * elementsPerPeer + chunkOffset); auto dst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(rank) * elementsPerPeer + chunkOffset); - BigDataCopyRangePingPong(dst + shardOffset, src + shardOffset, shardBytes, relayLocal); + if (shardHasBytes) { + BigDataCopyRangePingPong(dst + shardOffset, src + shardOffset, shardBytes, relayLocal); + } return; } - if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT && + if (!use35Core && profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT && globalPass >= TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS) { const uint64_t reuseToken = BigDataPassToken( globalPass - static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS)); @@ -487,9 +663,11 @@ __aicore__ inline void BigDataCopyPeerWorker( const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, peerIndex, chunkBytesPerPeer); - BigDataCopyRangePingPong(sendSlot + shardOffset, src + shardOffset, shardBytes, relayLocal); + if (shardHasBytes) { + BigDataCopyRangePingPong(sendSlot + shardOffset, src + shardOffset, shardBytes, relayLocal); + } BigDataStoreTokenMte( - BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer, copyShard), + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard), token, relayLocal); } @@ -497,11 +675,11 @@ __aicore__ inline void BigDataSendPeerWorker( int32_t peer, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, - uint64_t kernelLoopBase, uint32_t profileStage, uint64_t sendDataOffset, + uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, uint64_t sendDataOffset, uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { - if (peer < 0 || peer >= rankSize || peer == rank || + if (peer < 0 || peer >= rankSize || peer == rank || shardCount == 0U || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { return; } @@ -515,11 +693,11 @@ __aicore__ inline void BigDataSendPeerWorker( } const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); - const uint32_t slot = BigDataPingPongSlot(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, false); - for (uint32_t copyShard = 0U; copyShard < TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS; ++copyShard) { + for (uint32_t copyShard = 0U; copyShard < shardCount; ++copyShard) { const uint64_t observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer, copyShard), + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard), token, relayLocal); if (observed < token) { if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { @@ -547,7 +725,7 @@ __aicore__ inline void BigDataSendPeerWorker( readySignalOffset + ((static_cast(slot) * static_cast(rankSize) + static_cast(rank)) * - static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) * + static_cast(shardCount)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; auto registry = TileXR::GetUDMARegistry(args); auto udmaInfo = TileXR::GetUDMAInfo(args); @@ -588,12 +766,12 @@ __aicore__ inline void BigDataRecvPeerWorker( __gm__ int32_t* output, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, - uint32_t recvShard, uint64_t recvDataOffset, uint64_t recvCopyDoneOffset, - uint64_t readySignalOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, - AscendC::LocalTensor relayLocal) + bool use35Core, uint32_t recvShard, uint32_t shardCount, uint64_t recvDataOffset, + uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, + uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { if (peer < 0 || peer >= rankSize || peer == rank || - recvShard >= TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS || + shardCount == 0U || recvShard >= shardCount || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { return; } @@ -608,15 +786,14 @@ __aicore__ inline void BigDataRecvPeerWorker( } uint32_t shardOffset = 0U; uint32_t shardBytes = 0U; - if (!BigDataCopyShardRange(recvShard, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { - return; - } + const bool shardHasBytes = + BigDataCopyShardRange(recvShard, shardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes); const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); - const uint32_t slot = BigDataPingPongSlot(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); uint64_t observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer, 0U), + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, peer, 0U), token, relayLocal); if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = @@ -641,7 +818,9 @@ __aicore__ inline void BigDataRecvPeerWorker( } auto dst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(peer) * elementsPerPeer + chunkOffset); - BigDataCopyRangePingPong(dst + shardOffset, recvSlot + shardOffset, shardBytes, relayLocal); + if (shardHasBytes) { + BigDataCopyRangePingPong(dst + shardOffset, recvSlot + shardOffset, shardBytes, relayLocal); + } if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { auto relayDst = output + static_cast(peer) * elementsPerPeer + chunkOffset; debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE + peer] = relayDst[0]; @@ -651,14 +830,14 @@ __aicore__ inline void BigDataRecvPeerWorker( } BigDataStoreTokenMte( - BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, peer, recvShard), + BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, recvShard), token, relayLocal); - if (recvShard + 1U != TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS) { + if (recvShard + 1U != shardCount) { return; } - for (uint32_t shard = 0U; shard < TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS; ++shard) { + for (uint32_t shard = 0U; shard < shardCount; ++shard) { observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, peer, shard), + BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, shard), token, relayLocal); if (observed < token) { if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { @@ -669,7 +848,13 @@ __aicore__ inline void BigDataRecvPeerWorker( } } - auto remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); + __gm__ uint64_t* remoteAck = nullptr; + if (use35Core) { + remoteAck = BigDataRemoteRegisteredControlSlot( + args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); + } else { + remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); + } if (remoteAck == nullptr) { if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = @@ -680,6 +865,221 @@ __aicore__ inline void BigDataRecvPeerWorker( BigDataStoreTokenMte(remoteAck, token, relayLocal); } +__aicore__ inline bool BigDataWaitCopyDoneRange( + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, uint64_t copyDoneOffset, + uint32_t slot, int32_t rankSize, uint32_t shardCount, int32_t peer, + uint32_t copyShardBegin, uint32_t copyShardEnd, uint64_t token, + uint32_t loop, uint32_t pass, AscendC::LocalTensor relayLocal) +{ + if (copyShardBegin > copyShardEnd || copyShardEnd > shardCount) { + return false; + } + for (uint32_t copyShard = copyShardBegin; copyShard < copyShardEnd; ++copyShard) { + const uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; + } + return false; + } + } + return true; +} + +__aicore__ inline bool BigDataRemoteSendSegmentRange( + uint32_t segmentId, uint32_t shardCount, uint32_t chunkElements, + uint32_t& copyShardBegin, uint32_t& copyShardEnd, + uint32_t& segmentOffsetBytes, uint32_t& segmentBytes) +{ + if (shardCount == 0U) { + return false; + } + const uint32_t chunkBytes = chunkElements * static_cast(sizeof(int32_t)); + const uint32_t splitOffset = BigDataCopyShardStartBytes( + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END, shardCount, chunkElements); + if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT) { + copyShardBegin = 0U; + copyShardEnd = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END; + segmentOffsetBytes = 0U; + segmentBytes = splitOffset; + } else if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT) { + copyShardBegin = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END; + copyShardEnd = shardCount; + segmentOffsetBytes = splitOffset; + segmentBytes = chunkBytes - splitOffset; + } else { + return false; + } + if (copyShardEnd > shardCount || segmentOffsetBytes > chunkBytes) { + return false; + } + return true; +} + +__aicore__ inline void BigDataPublishReadySignal( + __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, + uint64_t readySignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, + int32_t rank, uint64_t token) +{ + const uint64_t localReadyPayloadOffset = + readySignalOffset + + ((static_cast(slot) * static_cast(rankSize) + + static_cast(rank)) * static_cast(shardCount)) * + TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + const uint64_t remoteReadyOffset = localReadyPayloadOffset; + auto localSrc = reinterpret_cast<__gm__ uint64_t*>(udmaMem + localReadyPayloadOffset); + TileXR::UDMAPutSignalNbi( + args, peer, localSrc, localReadyPayloadOffset, sizeof(uint64_t), remoteReadyOffset, token); + (void)TileXR::UDMAQuietStatus(args, peer); +} + +__aicore__ inline void BigDataRemoteSendSegmentWorker( + int32_t peer, uint32_t segmentId, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, + int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, + uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, + uint64_t sendDataOffset, uint64_t recvDataOffset, uint64_t copyDoneOffset, + uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t chunkBytesPerPeer, + AscendC::LocalTensor relayLocal) +{ + if (peer < 0 || peer >= rankSize || peer == rank || !BigDataIsMultiNode(rankSize) || + shardCount == 0U || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { + return; + } + + const int32_t peerIndex = BigDataNetworkPeerIndex(peer, rank); + const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + return; + } + (void)chunkOffset; + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, true); + + uint32_t copyShardBegin = 0U; + uint32_t copyShardEnd = 0U; + uint32_t segmentOffsetBytes = 0U; + uint32_t segmentBytes = 0U; + if (!BigDataRemoteSendSegmentRange( + segmentId, shardCount, chunkBytes / sizeof(int32_t), + copyShardBegin, copyShardEnd, segmentOffsetBytes, segmentBytes)) { + return; + } + + if (!BigDataWaitCopyDoneRange(udmaMem, debug, copyDoneOffset, slot, rankSize, shardCount, + peer, copyShardBegin, copyShardEnd, token, loop, pass, relayLocal)) { + return; + } + if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { + return; + } + + auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, + peerIndex, chunkBytesPerPeer); + auto localSrc = reinterpret_cast<__gm__ int32_t*>(sendSlot + segmentOffsetBytes); + const uint64_t remoteDataOffset = + recvDataOffset + + (static_cast(slot) * static_cast(networkPeerCount) + + static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer + + static_cast(segmentOffsetBytes); + + uint32_t status = 0U; + if (segmentBytes > 0U) { + TileXR::UDMAPutNbi(args, peer, localSrc, remoteDataOffset, segmentBytes); + status = TileXR::UDMAQuietStatus(args, peer); + } + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } + BigDataStoreTokenMte( + BigDataControlSlot(udmaMem, remoteSendDoneOffset, slot, rankSize, shardCount, peer, segmentId), + token, relayLocal); + + if (segmentId != TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT || + profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { + return; + } + for (uint32_t done = 0U; done <= TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; ++done) { + const uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, remoteSendDoneOffset, slot, rankSize, shardCount, peer, done), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + return; + } + } + auto localReady = BigDataControlSlot( + udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, 0U); + BigDataStoreTokenMte(localReady, token, relayLocal); + BigDataPublishReadySignal( + args, peer, udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, token); +} + +__aicore__ inline void BigDataRunSelfCopyShard( + int32_t rank, int32_t rankSize, __gm__ int32_t* input, __gm__ int32_t* output, + int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t pass, + uint32_t copyShard, uint32_t shardCount, AscendC::LocalTensor relayLocal) +{ + if (rank < 0 || rank >= rankSize || shardCount == 0U || copyShard >= shardCount) { + return; + } + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + return; + } + uint32_t shardOffset = 0U; + uint32_t shardBytes = 0U; + if (!BigDataCopyShardRange(copyShard, shardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { + return; + } + auto src = reinterpret_cast<__gm__ uint8_t*>( + input + static_cast(rank) * elementsPerPeer + chunkOffset); + auto dst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(rank) * elementsPerPeer + chunkOffset); + BigDataCopyRangePingPong(dst + shardOffset, src + shardOffset, shardBytes, relayLocal); +} + +__aicore__ inline void BigDataRunRoleForPeer( + int32_t peer, int32_t role, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, + __gm__ int32_t* input, __gm__ int32_t* output, __gm__ uint8_t* udmaMem, + __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, + uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, + uint32_t profileStage, bool use35Core, uint32_t shardCount, uint64_t sendDataOffset, + uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, + uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) +{ + if (role < static_cast(shardCount)) { + BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, + use35Core, static_cast(role), shardCount, sendDataOffset, copyDoneOffset, ackSignalOffset, + chunkBytesPerPeer, relayLocal); + return; + } + if (role == static_cast(shardCount)) { + BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, + shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, + chunkBytesPerPeer, relayLocal); + return; + } + const uint32_t recvShard = + static_cast(role) - shardCount - 1U; + BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, + use35Core, recvShard, shardCount, recvDataOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, + chunkBytesPerPeer, relayLocal); +} + } // namespace extern "C" __global__ __aicore__ void tilexr_udma_all_gather_kernel( @@ -1745,9 +2145,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_fused_kernel( extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, GM_ADDR debugGM, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, - uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, - uint32_t profileStage) + uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, + uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, + uint64_t kernelLoopBase, uint32_t profileStage, uint32_t force35CoreFlag) { if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_FULL) { profileStage = TILEXR_BIGDATA_PROFILE_STAGE_FULL; @@ -1762,6 +2162,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( int32_t rankSize = args->rankSize; bool enabled = TileXR::UDMARegistryEnabled(args); const int32_t blockIdx = AscendC::GetBlockIdx(); + const bool force35Core = force35CoreFlag != 0U; if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; @@ -1775,48 +2176,126 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( return; } + if (!BigDataValidTopology(rankSize)) { + if (blockIdx == 0 && debug != nullptr) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + return; + } + const bool use35Core = BigDataUse35Core(rankSize, force35Core); + const uint32_t shardCount = BigDataShardCount(rankSize, force35Core); const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; - const uint64_t chunkBytesPerPeer = static_cast(effectiveChunkElements) * sizeof(int32_t); + const uint64_t chunkBytesPerPeer = use35Core ? + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES : + static_cast(effectiveChunkElements) * sizeof(int32_t); const uint64_t sendDataOffset = dataOffset; const uint64_t recvDataOffset = sendDataOffset + - static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS) * + static_cast(use35Core ? passCount : TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS) * static_cast(rankSize > 1 ? rankSize - 1 : 1) * chunkBytesPerPeer; AscendC::TPipe pipe; AscendC::TBuf relayTBuf; pipe.InitBuffer(relayTBuf, TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES); AscendC::LocalTensor relayLocal = relayTBuf.Get(); + if (!use35Core) { + const int32_t workerGroup = + blockIdx / static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER); + const int32_t role = + blockIdx % static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER); + if (workerGroup >= rankSize || role < 0 || + role >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)) { + BigDataKernelExitBarrier(); + return; + } + for (uint32_t loop = 0; loop < loopCount; ++loop) { + for (uint32_t pass = 0; pass < passCount; ++pass) { + const int32_t peer = workerGroup; + BigDataRunRoleForPeer(peer, role, rank, rankSize, args, input, output, + udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, + loop, pass, kernelLoopBase, profileStage, use35Core, shardCount, sendDataOffset, recvDataOffset, + copyDoneOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, + chunkBytesPerPeer, relayLocal); + } + } + BigDataKernelExitBarrier(); + return; + } - if (rankSize <= 0 || blockIdx >= rankSize * static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)) { + if (blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM)) { + BigDataKernelExitBarrier(); return; } - const int32_t peer = blockIdx / static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER); - const int32_t role = blockIdx % static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER); + const bool isCopyCore = + blockIdx < static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES); + const bool isRemoteSendPrimaryCore = + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE); + const bool isRemoteSendSecondaryCore = + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE); + const bool isLocalSendCore = + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE); + const bool isRecvCore = + blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE) && + blockIdx < static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM); + const uint32_t copyShard = static_cast(blockIdx); + const uint32_t recvShard = + static_cast(blockIdx) - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; + const int32_t taskCount = BigDataTaskCount(rankSize, force35Core); for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { - if (role < static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) { - const uint32_t copyShard = static_cast(role); - BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, - elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, - copyShard, sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); + if (isCopyCore) { + BigDataRunSelfCopyShard(rank, rankSize, input, output, + elementsPerPeer, effectiveChunkElements, pass, + copyShard, shardCount, relayLocal); } - if (role == static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) { - BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, - elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, - sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, - chunkBytesPerPeer, relayLocal); - } + for (int32_t taskIndex = 0; taskIndex < taskCount; ++taskIndex) { + int32_t peer = -1; + bool isLocalPeer = false; + if (!BigDataMergedPeerTaskAt(rank, rankSize, taskIndex, force35Core, peer, isLocalPeer)) { + continue; + } - if (role > static_cast(TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS)) { - const uint32_t recvShard = - static_cast(role) - TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS - 1U; - BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, - elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, - recvShard, recvDataOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, - chunkBytesPerPeer, relayLocal); + if (isCopyCore) { + BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, + udmaMem, debug, elementsPerPeer, effectiveChunkElements, + passCount, loop, pass, kernelLoopBase, profileStage, + use35Core, copyShard, shardCount, sendDataOffset, copyDoneOffset, ackSignalOffset, + chunkBytesPerPeer, relayLocal); + } + if (!isLocalPeer && isRemoteSendPrimaryCore) { + BigDataRemoteSendSegmentWorker(peer, + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT, + rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, + copyDoneOffset, remoteSendDoneOffset, readySignalOffset, + chunkBytesPerPeer, relayLocal); + } + if (!isLocalPeer && isRemoteSendSecondaryCore) { + BigDataRemoteSendSegmentWorker(peer, + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT, + rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, + copyDoneOffset, remoteSendDoneOffset, readySignalOffset, + chunkBytesPerPeer, relayLocal); + } + if (isLocalPeer && isLocalSendCore) { + BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, + copyDoneOffset, readySignalOffset, chunkBytesPerPeer, relayLocal); + } + if (isRecvCore) { + BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, use35Core, recvShard, shardCount, recvDataOffset, + recvCopyDoneOffset, readySignalOffset, ackSignalOffset, + chunkBytesPerPeer, relayLocal); + } } } } @@ -1827,14 +2306,14 @@ void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, - uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, - uint32_t profileStage) + uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, + uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, + uint64_t kernelLoopBase, uint32_t profileStage, uint32_t force35Core) { tilexr_udma_all_to_all_bigdata_kernel<<>>( commArgs, input, output, udmaMem, debug, elementsPerPeer, - dataOffset, copyDoneOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, - chunkElements, passCount, loopCount, kernelLoopBase, profileStage); + dataOffset, copyDoneOffset, recvCopyDoneOffset, remoteSendDoneOffset, readySignalOffset, + ackSignalOffset, chunkElements, passCount, loopCount, kernelLoopBase, profileStage, force35Core); } void launch_tilexr_udma_all_to_all_fused( diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 2354ed02..ec44ebde 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -140,30 +140,95 @@ void TestAllToAllBigDataPlan() constexpr int32_t elementsPerPeer = 16 * 1024 * 1024; // 64 MiB per peer for int32_t. const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 64ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 128ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, 1024ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 8ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 64ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 5U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataSingleNodeShards, 2U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataLocalCopyShards, 2U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeCopyCores, 16U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRecvCores, 16U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRemoteSendPrimaryCore, 16U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRemoteSendSecondaryCore, 17U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeLocalSendCore, 18U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRecvCoreBase, 19U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeBlockDim, 35U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataPingPongSlots, 2U); CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, true); const size_t controlGroupBytes = static_cast(TileXR::Demo::kAllToAllBigDataPingPongSlots) * static_cast(rankSize) * - static_cast(TileXR::Demo::kAllToAllBigDataLocalCopyShards) * + static_cast(TileXR::Demo::AllToAllBigDataShardCount(rankSize)) * TileXR::Demo::kAllToAllBigDataControlSlotBytes; CHECK_EQ(plan.controlBytes, controlGroupBytes); CHECK_EQ(plan.signalBytes, 3ULL * controlGroupBytes); CHECK_EQ(plan.copyDoneOffset, plan.dataBytes); CHECK_EQ(plan.recvCopyDoneOffset, plan.copyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.remoteSendDoneOffset, 0ULL); CHECK_EQ(plan.readySignalOffset, plan.recvCopyDoneOffset + controlGroupBytes); CHECK_EQ(plan.ackSignalOffset, plan.readySignalOffset + controlGroupBytes); - CHECK_EQ(plan.registeredBytes, plan.dataBytes + plan.controlBytes + plan.signalBytes); + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes); + CHECK_EQ(plan.dataBytes + plan.controlBytes + plan.signalBytes <= plan.registeredBytes, true); CHECK_EQ(plan.dataBytes, static_cast(rankSize - 1) * TileXR::Demo::kAllToAllBigDataPingPongSlots * 2ULL * plan.chunkBytesPerPeer); CHECK_EQ(plan.chunkElements > 0, true); CHECK_EQ(plan.chunkBytesPerPeer, static_cast(plan.chunkElements) * sizeof(int32_t)); - CHECK_EQ(plan.passCount > 1, true); + CHECK_EQ(plan.passCount >= 1, true); +} + +void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() +{ + constexpr int rankSize = 16; + constexpr int32_t elementsPerPeer = 2 * 1024 * 1024; + const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); + const size_t controlGroupBytes = + static_cast(plan.passCount) * + static_cast(rankSize) * + static_cast(TileXR::Demo::kAllToAllBigDataMultiNodeCopyCores) * + TileXR::Demo::kAllToAllBigDataControlSlotBytes; + + CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize), 16U); + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); + CHECK_EQ(plan.passCount, 1U); + CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); + CHECK_EQ(plan.controlBytes, controlGroupBytes); + CHECK_EQ(plan.signalBytes, 4ULL * controlGroupBytes); + CHECK_EQ(plan.copyDoneOffset, plan.dataBytes); + CHECK_EQ(plan.recvCopyDoneOffset, plan.copyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.remoteSendDoneOffset, plan.recvCopyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.readySignalOffset, plan.remoteSendDoneOffset + controlGroupBytes); + CHECK_EQ(plan.ackSignalOffset, plan.readySignalOffset + controlGroupBytes); + CHECK_EQ(plan.dataBytes + plan.controlBytes + plan.signalBytes <= plan.registeredBytes, true); + CHECK_EQ(plan.dataBytes, + static_cast(rankSize - 1) * static_cast(plan.passCount) * 2ULL * + TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); +} + +void TestAllToAllBigDataForce35CorePlanFor8P() +{ + constexpr int rankSize = 8; + constexpr int32_t elementsPerPeer = 2 * 1024 * 1024; + const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer, true); + const size_t controlGroupBytes = + static_cast(plan.passCount) * + static_cast(rankSize) * + static_cast(TileXR::Demo::kAllToAllBigDataMultiNodeCopyCores) * + TileXR::Demo::kAllToAllBigDataControlSlotBytes; + + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(rankSize), false); + CHECK_EQ(TileXR::Demo::AllToAllBigDataUse35Core(rankSize, true), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize, true), 16U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(rankSize, true), 35U); + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); + CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); + CHECK_EQ(plan.controlBytes, controlGroupBytes); + CHECK_EQ(plan.signalBytes, 4ULL * controlGroupBytes); + CHECK_EQ(plan.remoteSendDoneOffset, plan.recvCopyDoneOffset + controlGroupBytes); + CHECK_EQ(plan.readySignalOffset, plan.remoteSendDoneOffset + controlGroupBytes); + CHECK_EQ(plan.ackSignalOffset, plan.readySignalOffset + controlGroupBytes); + CHECK_EQ(plan.dataBytes + plan.controlBytes + plan.signalBytes <= plan.registeredBytes, true); } void TestAllToAllBigDataBlockDim() @@ -171,7 +236,66 @@ void TestAllToAllBigDataBlockDim() CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(0), 1U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(1), 5U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(8), 40U); - CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64), 320U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(8, true), 35U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(16), 35U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(32), 35U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64), 35U); +} + +void TestAllToAllBigDataMultiNodeTopology() +{ + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(8), false); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(16), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(8), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(16), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(24), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(10), false); + + CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeBegin(0), 0); + CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeEnd(0), 8); + CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeBegin(10), 8); + CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeEnd(10), 16); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(10, 8), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(10, 15), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(10, 7), false); +} + +void TestAllToAllBigDataRemotePeerQueue() +{ + std::vector peers = TileXR::Demo::AllToAllBigDataRemotePeers(3, 16); + std::vector expected {11, 12, 13, 14, 15, 8, 9, 10}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataRemotePeers(10, 16); + expected = {2, 3, 4, 5, 6, 7, 0, 1}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataRemotePeers(5, 24); + expected = {13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 8, 9, 10, 11, 12}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataRemotePeers(5, 8); + CHECK_EQ(peers.empty(), true); +} + +void TestAllToAllBigDataMergedPeerQueue() +{ + std::vector peers = TileXR::Demo::AllToAllBigDataLocalPeers(3); + std::vector expected {4, 5, 6, 7, 0, 1, 2}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataMergedPeerTasks(3, 16); + expected = {11, 4, 12, 5, 13, 6, 14, 7, 15, 0, 8, 1, 9, 2, 10}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataMergedPeerTasks(5, 32); + expected = {13, 14, 15, 6, 16, 17, 18, 7, 19, 20, 21, 0, + 22, 23, 24, 1, 25, 26, 27, 2, 28, 29, 30, + 3, 31, 8, 9, 4, 10, 11, 12}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataMergedPeerTasks(5, 8); + CHECK_EQ(peers.empty(), true); } void TestDemoDebugLayoutSource() @@ -264,10 +388,25 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_bigdata"); CHECK_CONTAINS(demo, "for (int iter = 0; iter < allToAllRepeat; ++iter)"); CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_PROFILE_STAGE"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_FORCE_35CORE"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_BARRIER_HOST"); + CHECK_CONTAINS(demo, "host = value.substr(0, colon)"); + CHECK_CONTAINS(demo, "addr.sin_addr.s_addr = htonl(INADDR_ANY)"); + CHECK_CONTAINS(demo, "ConnectBarrierServer(endpoint.host, endpoint.port)"); CHECK_CONTAINS(demo, "const uint64_t kernelLoopBase = static_cast(iter)"); CHECK_CONTAINS(demo, "bigDataPlan.passCount, 1, kernelLoopBase"); CHECK_CONTAINS(demo, "bigDataPlan.copyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.recvCopyDoneOffset"); + CHECK_CONTAINS(demo, "bigDataPlan.remoteSendDoneOffset"); + CHECK_CONTAINS(demo, "AllToAllBigDataValidTopology(rankSize)"); + CHECK_CONTAINS(demo, "ERROR: bigdata alltoall multi-node requires rankSize multiple of 8"); + CHECK_CONTAINS(demo, "bigdata multinode mode="); + CHECK_CONTAINS(demo, "shards="); + CHECK_CONTAINS(demo, "remoteSendDoneOffset="); + CHECK_CONTAINS(demo, "force35Core="); + CHECK_CONTAINS(demo, "TileXR::Demo::AllToAllBigDataShardCount("); + CHECK_CONTAINS(demo, "const uint32_t bigDataBlockDim = TileXR::Demo::AllToAllBigDataBlockDim("); + CHECK_CONTAINS(demo, "bigDataBlockDim, stream, commArgsDev"); CHECK_CONTAINS(demo, "static_cast(registeredMemory) + bigDataPlan.copyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.controlBytes + bigDataPlan.signalBytes"); CHECK_CONTAINS(demo, "static_cast(bigDataProfileStage)"); @@ -294,8 +433,43 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 5U"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = 2U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 8ULL * 1024ULL * 1024ULL"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE = 8"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM ="); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE = 16U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE = 17U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE = 18U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE = 19U"); + CHECK_CONTAINS(kernel, "BigDataIsMultiNode(rankSize)"); + CHECK_CONTAINS(kernel, "BigDataValidTopology(rankSize)"); + CHECK_CONTAINS(kernel, "BigDataUse35Core(rankSize, force35Core)"); + CHECK_CONTAINS(kernel, "BigDataShardCount(rankSize, force35Core)"); + CHECK_CONTAINS(kernel, "BigDataTaskCount(rankSize, force35Core)"); + CHECK_CONTAINS(kernel, "BigDataNodeCount(rankSize)"); + CHECK_CONTAINS(kernel, "BigDataRemotePeerAt(rank, rankSize, remoteIndex)"); + CHECK_CONTAINS(kernel, "BigDataLocalPeerAt"); + CHECK_CONTAINS(kernel, "BigDataMergedPeerTaskAt(rank, rankSize, taskIndex, force35Core, peer, isLocalPeer)"); + CHECK_CONTAINS(kernel, "BigDataRunSelfCopyShard(rank, rankSize"); + CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentWorker"); + CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentRange"); + CHECK_CONTAINS(kernel, "BigDataWaitCopyDoneRange"); + CHECK_CONTAINS(kernel, "BigDataPublishReadySignal"); + CHECK_CONTAINS(kernel, "remoteSendDoneOffset"); + CHECK_CONTAINS(kernel, "if (!BigDataIsMultiNode(rankSize))"); + CHECK_CONTAINS(kernel, "if (blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM))"); + CHECK_CONTAINS(kernel, "const bool isCopyCore ="); + CHECK_CONTAINS(kernel, "const bool isRemoteSendPrimaryCore ="); + CHECK_CONTAINS(kernel, "const bool isRemoteSendSecondaryCore ="); + CHECK_CONTAINS(kernel, "const bool isLocalSendCore ="); + CHECK_CONTAINS(kernel, "const bool isRecvCore ="); + CHECK_CONTAINS(kernel, "if (isCopyCore)"); + CHECK_CONTAINS(kernel, "if (!isLocalPeer && isRemoteSendPrimaryCore)"); + CHECK_CONTAINS(kernel, "if (!isLocalPeer && isRemoteSendSecondaryCore)"); + CHECK_CONTAINS(kernel, "if (isLocalPeer && isLocalSendCore)"); + CHECK_CONTAINS(kernel, "if (isRecvCore)"); CHECK_CONTAINS(kernel, "BigDataCopyPeerWorker"); CHECK_CONTAINS(kernel, "BigDataSendPeerWorker"); CHECK_CONTAINS(kernel, "BigDataRecvPeerWorker"); @@ -303,27 +477,33 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "copyShard"); CHECK_CONTAINS(kernel, "recvShard"); CHECK_CONTAINS(kernel, "recvCopyDoneOffset"); - CHECK_CONTAINS(kernel, "const int32_t peer = blockIdx / static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)"); - CHECK_CONTAINS(kernel, "const int32_t role = blockIdx % static_cast(TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER)"); - CHECK_CONTAINS(kernel, "BigDataPingPongSlot"); + CHECK_CONTAINS(kernel, "copyShardBegin = 0U"); + CHECK_CONTAINS(kernel, "copyShardEnd = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END"); + CHECK_CONTAINS(kernel, "copyShardBegin = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END"); + CHECK_CONTAINS(kernel, "copyShardEnd = shardCount"); + CHECK_CONTAINS(kernel, "BigDataDataSlot(globalPass, pass, use35Core)"); CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex"); CHECK_CONTAINS(kernel, "BigDataStoreTokenMte"); CHECK_CONTAINS(kernel, "BigDataIpcAckOffset"); CHECK_CONTAINS(kernel, "BigDataLocalIpcAckSlot"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); + CHECK_CONTAINS(kernel, "BigDataRemoteRegisteredControlSlot"); CHECK_CONTAINS(kernel, "rankSize > 1 ? rankSize - 1 : 1"); - CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, peer, copyShard)"); - CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, peer, recvShard)"); - CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, peer, 0U)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, recvShard)"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, peer, 0U)"); CHECK_CONTAINS(kernel, "remoteDataOffset ="); CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex(rank, peer)"); CHECK_CONTAINS(kernel, "BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "UDMAPutNbi"); CHECK_CONTAINS(kernel, "localSrc,"); CHECK_CONTAINS(kernel, "remoteDataOffset, chunkBytes, remoteReadyOffset, token"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES"); CHECK_CONTAINS(kernel, "relayLocal[bufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]"); - CHECK_CONTAINS(kernel, "globalPass - static_cast(TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS)"); + CHECK_CONTAINS(kernel, "if (!use35Core && profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT"); + CHECK_CONTAINS(kernel, "use35Core ? passCount : TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS"); CHECK_CONTAINS(kernel, "recvSlotInt[0]"); CHECK_CONTAINS(kernel, "BigDataStoreTokenMte(remoteAck, token, relayLocal)"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize)"); @@ -342,7 +522,12 @@ int main() TestBuildAllToAllOutput(); TestAllToAllMaxRank256With64MiBPerRank(); TestAllToAllBigDataPlan(); + TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone(); + TestAllToAllBigDataForce35CorePlanFor8P(); TestAllToAllBigDataBlockDim(); + TestAllToAllBigDataMultiNodeTopology(); + TestAllToAllBigDataRemotePeerQueue(); + TestAllToAllBigDataMergedPeerQueue(); TestDemoDebugLayoutSource(); TestAllToAllDataAsFlagSource(); TestAllToAllChunkedUdmaSource(); From ae07f690b70f619ac5e74a03c678f9cee9998f46 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 29 Jun 2026 21:48:03 +0800 Subject: [PATCH 018/163] docs: design udma multi-eid route port --- .../2026-06-29-udma-multi-eid-route-design.md | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-29-udma-multi-eid-route-design.md diff --git a/docs/superpowers/specs/2026-06-29-udma-multi-eid-route-design.md b/docs/superpowers/specs/2026-06-29-udma-multi-eid-route-design.md new file mode 100644 index 00000000..af351bec --- /dev/null +++ b/docs/superpowers/specs/2026-06-29-udma-multi-eid-route-design.md @@ -0,0 +1,104 @@ +# UDMA Multi-EID Route Design + +## Scope + +Port the multi-EID direct UDMA P2P route capability from PR 45 commit range +`06cd19c42846eb79ab75b5a075b82d9b656c2852..07afc489467e7e9a369eac6679913de208d169b1` +into the current `codex/udma-bigdata-isolated` branch. + +The port is intentionally limited to the UDMA route/layout/device wrapper and +UDMA-specific tests. It must not merge unrelated PR 45 checker, collectives, +documentation, or demo restructuring changes. + +## Goals + +- Keep the existing single-route behavior as the default fallback. +- Add an opt-in multi-route policy for direct UDMA P2P. +- Discover and diagnose multiple local EID routes per peer from HCCL + root/topology data. +- Allow explicit EID selection and ordering for diagnostics. +- Expand queue metadata so one peer can use multiple EID routes concurrently. +- Weight per-QP transfer slices by route link capacity, enabling mixed + 6-port and 2-port routes to be used concurrently with proportional work. + +## Non-Goals + +- Do not replace the current bigdata alltoall work. +- Do not merge the full PR 45 branch. +- Do not change public host APIs. +- Do not require UDMA on unsupported hardware; graceful fallback remains. + +## Route Selection + +`TileXRUDMATransport::BuildRoutes()` keeps the current single-route resolution +when no multi-route policy is requested. + +When `TILEXR_UDMA_ROUTE_POLICY=all`, the transport: + +1. Parses HCCL root/topology data. +2. Resolves all local EID candidates for each peer from topology ports. +3. Falls back to aggregate local EIDs when peer-specific edge data is missing. +4. Applies `TILEXR_UDMA_ROUTE_EIDS` if provided, preserving the requested order + but only accepting valid candidates. +5. Caps the route count with `TILEXR_UDMA_MAX_EIDS_PER_PEER`. +6. Exchanges local route lists so each rank knows the peer's matching remote + EIDs. + +The first selected EID remains populated in the existing `peerLocalEid_` and +`peerRemoteEid_` maps for compatibility with fallback paths. + +## Queue and Memory Model + +The transport stores per-peer local and remote EID vectors. It creates RA +contexts for all selected local EIDs, then creates `qpsPerRoute_` queues per +selected route. `qpNum_` becomes the maximum expanded QP route count across +peers. + +Memory registration remains local-address based, but registration/import is +tracked per EID. Remote memory handles are stored per peer and per route so +cleanup can unimport each handle through the matching local EID context. + +## Device Layout + +`UDMAInfo` gains `qpWeightPtr`, pointing to a device-visible `uint32_t` array +with the same `(rank, qpIdx)` indexing as SQ/RQ/CQ/memory metadata. + +`BuildUDMAInfoImage()` gains an overload that accepts: + +- explicit `qpNum` +- SQ/RQ/SCQ/RCQ vectors +- memory vector +- QP weight vector + +The existing overload remains and defaults all weights to `1`, preserving +single-route and older test behavior. + +## Weighted Device Slicing + +Device code adds `UDMAGetQpWeight()`. If no weight pointer is present, or a +weight entry is zero, the effective weight is `1`. + +The UDMA P2P perf kernels replace equal WQE slicing with weighted slicing: + +- Sum all QP weights for the peer. +- Compute each QP's byte range by weight proportion. +- Align slice boundaries to `BLOCK_UNIT_BYTE`. +- Assign any rounding remainder to the final QP. + +This allows a 6-port route and a 2-port route to both post concurrently while +the 6-port route receives proportionally more payload. + +## Tests + +Update the UDMA transport layout unit test to cover: + +- `qpWeightPtr` placement and serialized values. +- Multi-route QP-to-EID expansion. +- Route weight expansion from topology-derived port counts. +- Explicit EID selection parsing, filtering, de-duplication, and ordering. +- Existing mismatched-array rejection and large-transfer chunk behavior. + +Build/runtime validation depends on the existing repository environment. Host +layout tests can run without Ascend hardware. Direct UDMA P2P runtime validation +still requires supported A5 / Ascend950-class hardware. + From 1b4c443446bb849b5d663a4d4fba9a7de54844e1 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 29 Jun 2026 21:54:04 +0800 Subject: [PATCH 019/163] docs: plan udma multi-eid route port --- .../plans/2026-06-29-udma-multi-eid-route.md | 564 ++++++++++++++++++ 1 file changed, 564 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-29-udma-multi-eid-route.md diff --git a/docs/superpowers/plans/2026-06-29-udma-multi-eid-route.md b/docs/superpowers/plans/2026-06-29-udma-multi-eid-route.md new file mode 100644 index 00000000..6209d146 --- /dev/null +++ b/docs/superpowers/plans/2026-06-29-udma-multi-eid-route.md @@ -0,0 +1,564 @@ +# UDMA Multi-EID Route Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Port the PR 45 multi-EID direct UDMA P2P route capability into the current bigdata-isolated branch without merging unrelated PR changes. + +**Architecture:** Keep the current single-route path as default. Add opt-in multi-route route discovery and queue expansion on the host, expose per-QP route weights in the device-visible UDMA layout, and use those weights for P2P kernel slicing. + +**Tech Stack:** C++14, Ascend C device headers, ACL runtime memory copy/allocation, TileXR socket exchange, existing UDMA layout unit tests. + +## Global Constraints + +- Target commit range: `06cd19c42846eb79ab75b5a075b82d9b656c2852..07afc489467e7e9a369eac6679913de208d169b1`. +- Do not merge unrelated PR 45 checker, collectives, documentation, or demo restructuring changes. +- Existing single-route behavior remains the default. +- Multi-route mode is opt-in through `TILEXR_UDMA_ROUTE_POLICY=all`. +- Preserve graceful UDMA fallback on unsupported hardware. +- Do not change public host APIs. + +--- + +## File Structure + +- `src/include/tilexr_udma_types.h`: device-visible `UDMAInfo` layout. Add `qpWeightPtr`. +- `src/comm/udma/tilexr_udma_layout.h`: layout builder declarations and route helper declarations. +- `src/comm/udma/tilexr_udma_layout.cpp`: serialize weights, build QP-to-EID vectors, build QP weights, parse explicit EID lists. +- `src/include/tilexr_udma.h`: device helper `UDMAGetQpWeight()`. +- `src/comm/udma/tilexr_udma_transport.h`: transport state for multi-route local/remote EID vectors, expanded QP route vectors, route weights, and per-peer remote memory handles. +- `src/comm/udma/tilexr_udma_transport.cpp`: host route selection, context creation, queue creation, import/export, image refresh, and cleanup. +- `tests/udma/unit/test_tilexr_udma_transport_layout.cpp`: host-only unit coverage for layout and route helper behavior. +- `tests/udma/demo/tilexr_udma_demo_kernel.cpp`: weighted slice computation for P2P perf kernels. + +--- + +### Task 1: Layout and Device Weight Helpers + +**Files:** +- Modify: `src/include/tilexr_udma_types.h` +- Modify: `src/comm/udma/tilexr_udma_layout.h` +- Modify: `src/comm/udma/tilexr_udma_layout.cpp` +- Modify: `src/include/tilexr_udma.h` +- Test: `tests/udma/unit/test_tilexr_udma_transport_layout.cpp` + +**Interfaces:** +- Consumes: existing `UDMAInfo`, `BuildUDMAInfoImage()`, `UDMAInfo::qpNum`. +- Produces: + - `UDMAInfo::qpWeightPtr` + - `int BuildUDMAInfoImage(uintptr_t deviceBase, uint32_t qpNum, const std::vector& sq, const std::vector& rq, const std::vector& scq, const std::vector& rcq, const std::vector& mem, const std::vector& qpWeights, UDMAInfo& info, std::vector& bytes)` + - `std::vector BuildUDMAMultiRouteQpToEid(const std::vector& routeEids, uint32_t qpsPerRoute)` + - `std::vector BuildUDMAMultiRouteQpWeights(const std::vector& routeEids, const std::map& routeWeights, uint32_t qpsPerRoute)` + - `std::vector SelectExplicitUDMARouteEids(const char* routeList, const std::vector& candidateEids)` + - `__aicore__ inline uint32_t UDMAGetQpWeight(__gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx)` + +- [ ] **Step 1: Write failing layout tests** + +Add tests to `tests/udma/unit/test_tilexr_udma_transport_layout.cpp`: + +```cpp +#include +``` + +In `TestHostLayoutUsesDeviceRelativePointers()`, add: + +```cpp +std::vector weights = {3, 1}; +const int ret = TileXR::BuildUDMAInfoImage(deviceBase, 1, sq, rq, scq, rcq, mem, weights, info, bytes); +CHECK_TRUE(info.qpWeightPtr > info.memPtr); +CHECK_EQ(bytes.size(), + sizeof(TileXR::UDMAInfo) + 2 * sizeof(TileXR::UDMAWQCtx) + + 2 * sizeof(TileXR::UDMAWQCtx) + 2 * sizeof(TileXR::UDMACQCtx) + + 2 * sizeof(TileXR::UDMACQCtx) + 2 * sizeof(TileXR::UDMAMemInfo) + + 2 * sizeof(uint32_t)); +const auto* imageWeights = reinterpret_cast( + bytes.data() + (info.qpWeightPtr - deviceBase)); +CHECK_EQ(imageWeights[0], 3U); +CHECK_EQ(imageWeights[1], 1U); +``` + +Add helper tests: + +```cpp +void TestMultiRouteQpMappingRepeatsEachRoute() +{ + const std::vector routeEids = {7, 8}; + const std::vector qpToEid = TileXR::BuildUDMAMultiRouteQpToEid(routeEids, 2); + CHECK_EQ(qpToEid.size(), static_cast(4)); + CHECK_EQ(qpToEid[0], 7U); + CHECK_EQ(qpToEid[1], 7U); + CHECK_EQ(qpToEid[2], 8U); + CHECK_EQ(qpToEid[3], 8U); +} + +void TestMultiRouteQpMappingRejectsEmptyInputs() +{ + CHECK_TRUE(TileXR::BuildUDMAMultiRouteQpToEid({}, 1).empty()); + CHECK_TRUE(TileXR::BuildUDMAMultiRouteQpToEid({7}, 0).empty()); +} + +void TestMultiRouteQpWeightsUseRouteBandwidth() +{ + const std::vector routeEids = {7, 8}; + const std::map routeWeights = {{7, 6}, {8, 2}}; + const std::vector qpWeights = TileXR::BuildUDMAMultiRouteQpWeights(routeEids, routeWeights, 1); + CHECK_EQ(qpWeights.size(), static_cast(2)); + CHECK_EQ(qpWeights[0], 6U); + CHECK_EQ(qpWeights[1], 2U); +} + +void TestExplicitRouteSelectionKeepsRequestedCandidateOrder() +{ + const std::vector candidates = {7, 8}; + const std::vector selected = TileXR::SelectExplicitUDMARouteEids("8,7,9,bad", candidates); + CHECK_EQ(selected.size(), static_cast(2)); + CHECK_EQ(selected[0], 8U); + CHECK_EQ(selected[1], 7U); +} + +void TestExplicitRouteSelectionRejectsMissingInputs() +{ + const std::vector candidates = {7, 8}; + CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("", candidates).empty()); + CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("9", candidates).empty()); + CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("8", {}).empty()); +} +``` + +Call the helper tests from `main()`. + +- [ ] **Step 2: Run test to verify it fails** + +Run: + +```bash +cmake --build build --target test_tilexr_udma_transport_layout +``` + +Expected: compile failure mentioning missing `UDMAInfo::qpWeightPtr`, missing overload, or missing helper declarations. + +- [ ] **Step 3: Implement layout and helper declarations** + +Add `uint64_t qpWeightPtr;` to `UDMAInfo` in `src/include/tilexr_udma_types.h`. + +Add to `src/comm/udma/tilexr_udma_layout.h`: + +```cpp +#include + +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpNum, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + const std::vector& qpWeights, + UDMAInfo& info, + std::vector& bytes); + +std::vector BuildUDMAMultiRouteQpToEid( + const std::vector& routeEids, + uint32_t qpsPerRoute); + +std::vector BuildUDMAMultiRouteQpWeights( + const std::vector& routeEids, + const std::map& routeWeights, + uint32_t qpsPerRoute); + +std::vector SelectExplicitUDMARouteEids( + const char* routeList, + const std::vector& candidateEids); +``` + +- [ ] **Step 4: Implement layout serialization and helpers** + +In `src/comm/udma/tilexr_udma_layout.cpp`, include: + +```cpp +#include +#include +#include +``` + +Change the existing `BuildUDMAInfoImage()` to forward to the new overload with `std::vector(mem.size(), 1)`. In the new overload, validate `qpWeights.size() == sq.size()`, append the weight array after `mem`, assign `info.qpWeightPtr`, and copy the weights into `bytes`. + +Add the three helper functions exactly as declared in Step 3. + +- [ ] **Step 5: Implement device weight lookup** + +Add to `src/include/tilexr_udma.h` after `UDMAGetRemoteMemInfo()`: + +```cpp +__aicore__ inline uint32_t UDMAGetQpWeight(__gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx) +{ + uint32_t qpNum = udmaInfo->qpNum; + if (udmaInfo->qpWeightPtr == 0) { + return 1; + } + auto weights = reinterpret_cast<__gm__ uint32_t*>(udmaInfo->qpWeightPtr); + uint32_t weight = weights[pe * qpNum + qpIdx]; + return weight == 0 ? 1 : weight; +} +``` + +- [ ] **Step 6: Run layout test** + +Run: + +```bash +cmake --build build --target test_tilexr_udma_transport_layout +ctest --test-dir build -R test_tilexr_udma_transport_layout --output-on-failure +``` + +Expected: layout test passes. If `build` is not configured, run the repository's existing CMake configure command first. + +- [ ] **Step 7: Commit** + +```bash +git add src/include/tilexr_udma_types.h src/comm/udma/tilexr_udma_layout.h src/comm/udma/tilexr_udma_layout.cpp src/include/tilexr_udma.h tests/udma/unit/test_tilexr_udma_transport_layout.cpp +git commit -m "feat: add udma route weight layout" +``` + +--- + +### Task 2: Host Multi-EID Route Expansion + +**Files:** +- Modify: `src/comm/udma/tilexr_udma_transport.h` +- Modify: `src/comm/udma/tilexr_udma_transport.cpp` +- Test: `tests/udma/unit/test_tilexr_udma_transport_layout.cpp` + +**Interfaces:** +- Consumes: Task 1 helpers `BuildUDMAMultiRouteQpToEid()`, `BuildUDMAMultiRouteQpWeights()`, `SelectExplicitUDMARouteEids()`. +- Produces: + - `TileXRUDMATransport::qpsPerRoute_` + - `peerLocalEids_`, `peerRemoteEids_` + - `peerQpRouteEids_`, `peerQpRouteWeights_` + - `remoteMemHandlesByPeer_` + +- [ ] **Step 1: Add transport state fields** + +In `src/comm/udma/tilexr_udma_transport.h`, add: + +```cpp +uint32_t qpsPerRoute_ = 1; +std::map> peerLocalEids_; +std::map> peerRemoteEids_; +std::map> peerQpRouteEids_; +std::map> peerQpRouteWeights_; +std::map> remoteMemHandlesByPeer_; +``` + +Replace `std::vector remoteMemHandles_;` with `remoteMemHandlesByPeer_`. + +- [ ] **Step 2: Add topology route helpers** + +In `src/comm/udma/tilexr_udma_transport.cpp`, extend `TileXRRootInfo` with: + +```cpp +std::unordered_map> portCountByEidByLocalId; +``` + +In `ParseRootInfo()`, store port counts per EID: + +```cpp +const std::vector ports = JsonStringArrayField(addrObj, "ports"); +root.portCountByEidByLocalId[localId][eidIndex] = static_cast(ports.size()); +for (const std::string& port : ports) { + root.portToEidByLocalId[localId][port] = eidIndex; +} +``` + +Add helpers: + +```cpp +std::vector ResolveLocalEidRoutes( + const TileXRRootInfo& root, const std::vector& edges, uint32_t localId, uint32_t peerLocalId); + +std::vector ResolveLocalAggregateEidRoutes(const TileXRRootInfo& root, uint32_t localId); +``` + +`ResolveLocalEidRoutes()` walks all matching topology edges, maps local ports to EID indices, de-duplicates while preserving order, and only returns EIDs present in `root.eidByLocalId[localId]`. + +`ResolveLocalAggregateEidRoutes()` returns local EIDs with port count greater than `1`, preserving map order. + +- [ ] **Step 3: Expand `BuildRoutes()`** + +At the start of `BuildRoutes()`: + +```cpp +qpsPerRoute_ = qpNum_; +const bool useAllRoutes = std::getenv("TILEXR_UDMA_ROUTE_POLICY") != nullptr && + std::strcmp(std::getenv("TILEXR_UDMA_ROUTE_POLICY"), "all") == 0; +const uint32_t maxEidsPerPeer = GetEnvUint("TILEXR_UDMA_MAX_EIDS_PER_PEER", UINT32_MAX, 1, UINT32_MAX); +``` + +When assigning routes per peer: + +- In default mode, keep one EID per peer. +- In `useAllRoutes`, resolve local EID vectors with `ResolveLocalEidRoutes()`. +- If peer-specific routes are empty, use `ResolveLocalAggregateEidRoutes()`. +- If `TILEXR_UDMA_ROUTE_EIDS` selects valid candidates, replace the vector with that explicit list. +- Resize to `routeSlots = useAllRoutes ? max(1, min(eidCount_, maxEidsPerPeer)) : 1`. +- Exchange the route slots with peers. +- Pair local and remote route vectors with the same count. +- Build expanded per-QP route vectors and weight vectors: + +```cpp +peerQpRouteEids_[peer] = BuildUDMAMultiRouteQpToEid(localRoutes, qpsPerRoute_); +peerQpRouteWeights_[peer] = BuildUDMAMultiRouteQpWeights(localRoutes, weightByEid, qpsPerRoute_); +``` + +Finally set: + +```cpp +size_t maxRouteCount = 1; +for (const auto& entry : peerQpRouteEids_) { + maxRouteCount = std::max(maxRouteCount, entry.second.size()); +} +qpNum_ = static_cast(maxRouteCount); +``` + +- [ ] **Step 4: Create contexts for all selected local EIDs** + +In `CreateContexts()`, build `contextEids` from every value in `peerLocalEids_`. Fall back to existing `peerLocalEid_` if the vector map is empty. Create one RA context/token per unique EID. + +- [ ] **Step 5: Create queues per route** + +In `CreateQueues()`, size per-EID queue vectors by `qpsPerRoute_`, not `qpNum_`: + +```cpp +state.qpHandles.assign(qpsPerRoute_, nullptr); +state.remoteQpHandlesByQp.assign(qpsPerRoute_, std::vector(options_.rankSize, nullptr)); +state.localWqs.resize(qpsPerRoute_); +state.localCqs.resize(qpsPerRoute_); +``` + +When creating peer queues, only create a peer queue for a state if that state EID appears in `peerLocalEids_[peer]`. + +- [ ] **Step 6: Expand `RefreshUDMAInfo()` image entries** + +Build `sq/rq/scq/rcq/mem/qpWeights` with `options_.rankSize * qpNum_` entries. For each `(rank, qpIdx)`: + +- Resolve `localEid` from `peerQpRouteEids_[rank][qpIdx]` when available. +- Resolve `remoteEid` from the matching index in `peerRemoteEids_[rank]`. +- Use `routeQpIdx = qpIdx % qpsPerRoute_`. +- Read WQ/CQ metadata from `states_[localEid].localWqs[routeQpIdx]` and `.localCqs[routeQpIdx]`. +- Read remote TPN from `states_[localEid].tpnListByQp[routeQpIdx][rank]`. +- Read per-QP weight from `peerQpRouteWeights_[rank][qpIdx]`, defaulting to `1`. +- Call the weighted `BuildUDMAInfoImage()` overload. + +- [ ] **Step 7: Import and cleanup remote memory per route** + +In `ExchangeAndImportMemory()`, replace the single remote handle per peer with: + +```cpp +remoteMemHandlesByPeer_.clear(); +std::vector remoteHandles(localRoutes.size(), nullptr); +``` + +For each paired `localEid/remoteEid`, find the peer's exchanged MR for `remoteEid`, import it through `ctxHandleByEid_[localEid]`, and store it in `remoteHandles[routeIdx]`. + +In `CleanupMemory()`, iterate `remoteMemHandlesByPeer_`, use the same route index to find the local EID, and call `RaCtxRmemUnimport()` for each non-null handle. + +- [ ] **Step 8: Reset multi-route state during shutdown** + +In `Shutdown()`, clear: + +```cpp +peerLocalEids_.clear(); +peerRemoteEids_.clear(); +peerQpRouteEids_.clear(); +peerQpRouteWeights_.clear(); +qpsPerRoute_ = 1; +qpNum_ = 1; +remoteMemHandlesByPeer_.clear(); +``` + +- [ ] **Step 9: Build UDMA transport objects** + +Run: + +```bash +cmake --build build --target test_tilexr_udma_transport_layout +``` + +Expected: compile succeeds and the layout test target still builds. + +- [ ] **Step 10: Commit** + +```bash +git add src/comm/udma/tilexr_udma_transport.h src/comm/udma/tilexr_udma_transport.cpp +git commit -m "feat: expand udma p2p across multi-eid routes" +``` + +--- + +### Task 3: Weighted P2P Demo Kernel Slicing + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` + +**Interfaces:** +- Consumes: Task 1 `UDMAGetQpWeight()`. +- Produces: `TileXRUdmaDemoWeightedWqeSlice()`. + +- [ ] **Step 1: Add weighted slice helper** + +Add after `TileXRUdmaDemoWqeSlice()`: + +```cpp +__aicore__ inline void TileXRUdmaDemoWeightedWqeSlice( + __gm__ TileXR::UDMAInfo* udmaInfo, int32_t peer, uint32_t total, + uint32_t wqeCount, uint32_t wqeIdx, uint32_t& offset, uint32_t& bytes) +{ + uint32_t weightSum = 0; + uint32_t prefixWeight = 0; + for (uint32_t i = 0; i < wqeCount; ++i) { + uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, i); + if (i < wqeIdx) { + prefixWeight += weight; + } + weightSum += weight; + } + if (wqeCount == 0 || weightSum == 0 || wqeIdx >= wqeCount) { + offset = total; + bytes = 0; + return; + } + + uint64_t rawStart = static_cast(total) * prefixWeight / weightSum; + uint64_t rawEnd = static_cast(total) * + (prefixWeight + TileXR::UDMAGetQpWeight(udmaInfo, peer, wqeIdx)) / weightSum; + uint32_t alignedStart = static_cast( + (rawStart / TileXR::BLOCK_UNIT_BYTE) * TileXR::BLOCK_UNIT_BYTE); + uint32_t alignedEnd = wqeIdx + 1 == wqeCount ? total : static_cast( + ((rawEnd + TileXR::BLOCK_UNIT_BYTE - 1) / TileXR::BLOCK_UNIT_BYTE) * TileXR::BLOCK_UNIT_BYTE); + if (alignedEnd > total) { + alignedEnd = total; + } + if (alignedStart >= alignedEnd) { + offset = total; + bytes = 0; + return; + } + offset = alignedStart; + bytes = alignedEnd - alignedStart; +} +``` + +- [ ] **Step 2: Use weighted slicing in P2P perf kernels** + +In both `tilexr_udma_p2p_perf_kernel()` and `tilexr_udma_p2p_post_only_perf_kernel()`, replace: + +```cpp +TileXRUdmaDemoWqeSlice(bytes, jettyCount, blockIdx, offset, sliceBytes); +``` + +with: + +```cpp +auto udmaInfo = TileXR::GetUDMAInfo(args); +TileXRUdmaDemoWeightedWqeSlice(udmaInfo, peer, bytes, jettyCount, blockIdx, offset, sliceBytes); +``` + +In `tilexr_udma_p2p_post_only_perf_kernel()`, remove the later duplicate `auto udmaInfo = TileXR::GetUDMAInfo(args);` declaration. + +- [ ] **Step 3: Build demo kernel target** + +Run the existing UDMA build entrypoint: + +```bash +bash tests/udma/build.sh +``` + +Expected: compile succeeds. If local Windows host cannot run the Linux/CANN build script, record that hardware build verification was not runnable in this environment and still run Task 4 host tests. + +- [ ] **Step 4: Commit** + +```bash +git add tests/udma/demo/tilexr_udma_demo_kernel.cpp +git commit -m "perf: weight udma p2p demo slices" +``` + +--- + +### Task 4: Verification and Integration Check + +**Files:** +- Verify: `src/include/tilexr_udma_types.h` +- Verify: `src/comm/udma/tilexr_udma_layout.h` +- Verify: `src/comm/udma/tilexr_udma_layout.cpp` +- Verify: `src/include/tilexr_udma.h` +- Verify: `src/comm/udma/tilexr_udma_transport.h` +- Verify: `src/comm/udma/tilexr_udma_transport.cpp` +- Verify: `tests/udma/unit/test_tilexr_udma_transport_layout.cpp` +- Verify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` + +**Interfaces:** +- Consumes: completed Tasks 1-3. +- Produces: verified local changes and a concise implementation summary. + +- [ ] **Step 1: Check diff scope** + +Run: + +```bash +git diff --name-status HEAD +``` + +Expected: only the eight UDMA implementation/test files are modified if commits were not made per task. If task commits were made, run: + +```bash +git diff --name-status ae07f69..HEAD +``` + +Expected: only the eight UDMA implementation/test files are changed by implementation commits. + +- [ ] **Step 2: Run host layout unit test** + +Run: + +```bash +ctest --test-dir build -R test_tilexr_udma_transport_layout --output-on-failure +``` + +Expected: `TileXR UDMA transport layout checks passed`. + +- [ ] **Step 3: Run focused source search** + +Run: + +```bash +rg -n "qpWeightPtr|BuildUDMAMultiRoute|TILEXR_UDMA_ROUTE_POLICY|TILEXR_UDMA_ROUTE_EIDS|remoteMemHandlesByPeer|TileXRUdmaDemoWeightedWqeSlice" src tests +``` + +Expected: matches appear only in UDMA layout, UDMA transport, UDMA device wrapper, and UDMA demo/test files. + +- [ ] **Step 4: Inspect compatibility fallback** + +Run: + +```bash +git diff ae07f69..HEAD -- src/comm/udma/tilexr_udma_transport.cpp src/comm/udma/tilexr_udma_layout.cpp src/include/tilexr_udma.h +``` + +Expected: + +- `BuildRoutes()` uses multi-route only when `TILEXR_UDMA_ROUTE_POLICY=all`. +- Existing `BuildUDMAInfoImage()` overload still exists. +- `UDMAGetQpWeight()` returns `1` when `qpWeightPtr == 0` or the entry is zero. +- `Shutdown()` clears all multi-route maps. + +- [ ] **Step 5: Commit final verification adjustments if needed** + +If verification required fixes, commit them: + +```bash +git add src/include/tilexr_udma_types.h src/comm/udma/tilexr_udma_layout.h src/comm/udma/tilexr_udma_layout.cpp src/include/tilexr_udma.h src/comm/udma/tilexr_udma_transport.h src/comm/udma/tilexr_udma_transport.cpp tests/udma/unit/test_tilexr_udma_transport_layout.cpp tests/udma/demo/tilexr_udma_demo_kernel.cpp +git commit -m "fix: stabilize udma multi-eid route port" +``` + +Expected: commit succeeds only if there were verification fixes. + From 859cc64f38091236fbafad1062b467afc47e20ef Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 29 Jun 2026 23:38:39 +0800 Subject: [PATCH 020/163] feat(udma): support weighted multi-eid routes - Add multi-route UDMA queue layout with per-QP route weights and explicit EID selection. - Create peer queues across selected local and remote EIDs, preserving default single-route behavior. - Expose QP-specific device UDMA helpers and weighted slicing for diagnostics. - Shrink multinode bigdata registered windows for small payloads and cover the layout changes with unit tests. --- src/comm/udma/tilexr_udma_layout.cpp | 102 +++- src/comm/udma/tilexr_udma_layout.h | 26 + src/comm/udma/tilexr_udma_transport.cpp | 577 +++++++++++++----- src/comm/udma/tilexr_udma_transport.h | 10 +- src/include/tilexr_udma.h | 81 ++- src/include/tilexr_udma_types.h | 1 + tests/udma/demo/tilexr_udma_alltoall_layout.h | 10 +- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 78 ++- .../unit/test_tilexr_udma_alltoall_layout.cpp | 26 +- .../test_tilexr_udma_transport_layout.cpp | 64 +- 10 files changed, 798 insertions(+), 177 deletions(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 401b5f6e..77b5e5ab 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -5,6 +5,9 @@ #include "udma/tilexr_udma_layout.h" +#include +#include +#include #include namespace TileXR { @@ -31,10 +34,27 @@ int BuildUDMAInfoImage( const std::vector& mem, UDMAInfo& info, std::vector& bytes) +{ + return BuildUDMAInfoImage( + deviceBase, TILEXR_UDMA_QP_NUM, sq, rq, scq, rcq, mem, std::vector(mem.size(), 1), info, bytes); +} + +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpNum, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + const std::vector& qpWeights, + UDMAInfo& info, + std::vector& bytes) { const size_t rankCount = sq.size(); - if (rankCount == 0 || rq.size() != rankCount || scq.size() != rankCount || - rcq.size() != rankCount || mem.size() != rankCount) { + if (qpNum == 0 || rankCount == 0 || rankCount % qpNum != 0 || rq.size() != rankCount || + scq.size() != rankCount || rcq.size() != rankCount || mem.size() != rankCount || + qpWeights.size() != rankCount) { return TILEXR_UDMA_LAYOUT_INVALID; } @@ -43,15 +63,17 @@ int BuildUDMAInfoImage( const size_t scqOffset = rqOffset + rq.size() * sizeof(UDMAWQCtx); const size_t rcqOffset = scqOffset + scq.size() * sizeof(UDMACQCtx); const size_t memOffset = rcqOffset + rcq.size() * sizeof(UDMACQCtx); - const size_t totalBytes = memOffset + mem.size() * sizeof(UDMAMemInfo); + const size_t qpWeightOffset = memOffset + mem.size() * sizeof(UDMAMemInfo); + const size_t totalBytes = qpWeightOffset + qpWeights.size() * sizeof(uint32_t); info = {}; - info.qpNum = TILEXR_UDMA_QP_NUM; + info.qpNum = qpNum; info.sqPtr = deviceBase + sqOffset; info.rqPtr = deviceBase + rqOffset; info.scqPtr = deviceBase + scqOffset; info.rcqPtr = deviceBase + rcqOffset; info.memPtr = deviceBase + memOffset; + info.qpWeightPtr = deviceBase + qpWeightOffset; bytes.assign(totalBytes, 0); std::memcpy(bytes.data(), &info, sizeof(info)); @@ -60,7 +82,79 @@ int BuildUDMAInfoImage( CopyVector(bytes, scqOffset, scq); CopyVector(bytes, rcqOffset, rcq); CopyVector(bytes, memOffset, mem); + CopyVector(bytes, qpWeightOffset, qpWeights); return TILEXR_UDMA_LAYOUT_SUCCESS; } +std::vector BuildUDMAMultiRouteQpToEid( + const std::vector& routeEids, + uint32_t qpsPerRoute) +{ + std::vector qpToEid; + if (routeEids.empty() || qpsPerRoute == 0) { + return qpToEid; + } + qpToEid.reserve(routeEids.size() * qpsPerRoute); + for (uint32_t eid : routeEids) { + for (uint32_t qp = 0; qp < qpsPerRoute; ++qp) { + qpToEid.push_back(eid); + } + } + return qpToEid; +} + +std::vector BuildUDMAMultiRouteQpWeights( + const std::vector& routeEids, + const std::map& routeWeights, + uint32_t qpsPerRoute) +{ + std::vector qpWeights; + if (routeEids.empty() || qpsPerRoute == 0) { + return qpWeights; + } + qpWeights.reserve(routeEids.size() * qpsPerRoute); + for (uint32_t eid : routeEids) { + uint32_t weight = 1; + const auto weightIt = routeWeights.find(eid); + if (weightIt != routeWeights.end() && weightIt->second != 0) { + weight = weightIt->second; + } + for (uint32_t qp = 0; qp < qpsPerRoute; ++qp) { + qpWeights.push_back(weight); + } + } + return qpWeights; +} + +std::vector SelectExplicitUDMARouteEids( + const char* routeList, + const std::vector& candidateEids) +{ + std::vector selected; + if (routeList == nullptr || routeList[0] == '\0' || candidateEids.empty()) { + return selected; + } + + const char* cursor = routeList; + while (*cursor != '\0') { + char* end = nullptr; + unsigned long parsed = std::strtoul(cursor, &end, 0); + if (end != cursor && parsed <= UINT32_MAX) { + const uint32_t eid = static_cast(parsed); + if (std::find(candidateEids.begin(), candidateEids.end(), eid) != candidateEids.end() && + std::find(selected.begin(), selected.end(), eid) == selected.end()) { + selected.push_back(eid); + } + cursor = end; + } + while (*cursor != '\0' && *cursor != ',') { + ++cursor; + } + if (*cursor == ',') { + ++cursor; + } + } + return selected; +} + } // namespace TileXR diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index fefd49e4..5e3f4392 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -7,6 +7,7 @@ #define TILEXR_UDMA_LAYOUT_H #include +#include #include #include "tilexr_udma_types.h" @@ -26,6 +27,31 @@ int BuildUDMAInfoImage( UDMAInfo& info, std::vector& bytes); +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpNum, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + const std::vector& qpWeights, + UDMAInfo& info, + std::vector& bytes); + +std::vector BuildUDMAMultiRouteQpToEid( + const std::vector& routeEids, + uint32_t qpsPerRoute); + +std::vector BuildUDMAMultiRouteQpWeights( + const std::vector& routeEids, + const std::map& routeWeights, + uint32_t qpsPerRoute); + +std::vector SelectExplicitUDMARouteEids( + const char* routeList, + const std::vector& candidateEids); + } // namespace TileXR #endif // TILEXR_UDMA_LAYOUT_H diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index be210836..0b1ba619 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -41,6 +41,20 @@ bool UDMADiagEnabled() return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; } +uint32_t GetEnvUint(const char* name, uint32_t defaultValue, uint32_t minValue, uint32_t maxValue) +{ + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return defaultValue; + } + char* end = nullptr; + unsigned long parsed = std::strtoul(value, &end, 10); + if (end == value || *end != '\0' || parsed < minValue || parsed > maxValue) { + return defaultValue; + } + return static_cast(parsed); +} + std::string PtrToHex(uint64_t value) { std::ostringstream os; @@ -58,6 +72,11 @@ std::string EidToHex(const HccpEid& eid) return os.str(); } +int QueueKey(int peer, uint32_t qpIdx, uint32_t qpNum) +{ + return qpNum <= 1 ? peer : static_cast(static_cast(peer) * qpNum + qpIdx); +} + HccpEid SwapEidForDevice(const HccpEid& hccpEid) { HccpEid swapped {}; @@ -79,6 +98,7 @@ struct TileXRRootInfo { std::unordered_map deviceToLocalId; std::unordered_map> portToEidByLocalId; std::unordered_map> eidByLocalId; + std::unordered_map> portCountByEidByLocalId; }; struct TileXRTopoEdge { @@ -236,7 +256,9 @@ bool ParseRootInfo(TileXRRootInfo& root) if (!addr.empty() && ParseEidHex(addr, eid)) { root.eidByLocalId[localId][eidIndex] = eid; } - for (const std::string& port : JsonStringArrayField(addrObj, "ports")) { + const std::vector ports = JsonStringArrayField(addrObj, "ports"); + root.portCountByEidByLocalId[localId][eidIndex] = static_cast(ports.size()); + for (const std::string& port : ports) { root.portToEidByLocalId[localId][port] = eidIndex; } ++eidIndex; @@ -311,11 +333,60 @@ bool ResolveLocalEidRoute( return root.eidByLocalId.count(localId) != 0 && root.eidByLocalId.at(localId).count(eidIndex) != 0; } +std::vector ResolveLocalEidRoutes( + const TileXRRootInfo& root, const std::vector& edges, uint32_t localId, uint32_t peerLocalId) +{ + std::vector routeEids; + const auto localIt = root.portToEidByLocalId.find(localId); + const auto eidIt = root.eidByLocalId.find(localId); + if (localIt == root.portToEidByLocalId.end() || eidIt == root.eidByLocalId.end()) { + return routeEids; + } + + auto addPorts = [&](const std::vector& ports) { + for (const std::string& port : ports) { + const auto portIt = localIt->second.find(port); + if (portIt == localIt->second.end() || eidIt->second.count(portIt->second) == 0) { + continue; + } + if (std::find(routeEids.begin(), routeEids.end(), portIt->second) == routeEids.end()) { + routeEids.push_back(portIt->second); + } + } + }; + + for (const auto& edge : edges) { + if (edge.localA == localId && edge.localB == peerLocalId) { + addPorts(edge.localAPorts); + } else if (edge.localB == localId && edge.localA == peerLocalId) { + addPorts(edge.localBPorts); + } + } + return routeEids; +} + +std::vector ResolveLocalAggregateEidRoutes(const TileXRRootInfo& root, uint32_t localId) +{ + std::vector routeEids; + const auto countIt = root.portCountByEidByLocalId.find(localId); + const auto eidIt = root.eidByLocalId.find(localId); + if (countIt == root.portCountByEidByLocalId.end() || eidIt == root.eidByLocalId.end()) { + return routeEids; + } + for (const auto& entry : countIt->second) { + if (entry.second > 1 && eidIt->second.count(entry.first) != 0) { + routeEids.push_back(entry.first); + } + } + return routeEids; +} + } // namespace struct TileXRUDMATransport::PerEidState { struct PeerQueueState { int peer = -1; + uint32_t qpIdx = 0; void* chanHandle = nullptr; void* cqHandle = nullptr; void* qpHandle = nullptr; @@ -428,6 +499,11 @@ int TileXRUDMATransport::OpenDevice() int TileXRUDMATransport::BuildRoutes() { + qpsPerRoute_ = GetEnvUint("TILEXR_UDMA_QP_NUM", 1, 1, 64); + qpNum_ = qpsPerRoute_; + const char* routePolicy = std::getenv("TILEXR_UDMA_ROUTE_POLICY"); + const bool useAllRoutes = routePolicy != nullptr && std::strcmp(routePolicy, "all") == 0; + const uint32_t maxEidsPerPeer = GetEnvUint("TILEXR_UDMA_MAX_EIDS_PER_PEER", UINT32_MAX, 1, UINT32_MAX); RaInfo info {}; info.phyId = logicDevId_ + deviceIdOffset_; info.mode = NETWORK_OFFLINE; @@ -505,30 +581,53 @@ int TileXRUDMATransport::BuildRoutes() << ids.str(); } - std::vector localRouteByPeer(options_.rankSize, -1); + const uint32_t routeSlots = useAllRoutes ? std::max(1, std::min(eidCount_, maxEidsPerPeer)) : 1; + std::vector localRouteByPeer(static_cast(options_.rankSize) * routeSlots, -1); for (int peer = 0; peer < options_.rankSize; ++peer) { if (peer == options_.rank) { continue; } uint32_t localEid = devEids[0].eidIndex; - if (topoReady && !ResolveLocalEidRoute(rootInfo, topoEdges, localId, allLocalIds[peer], localEid)) { - topoReady = false; + std::vector localEids; + if (topoReady && useAllRoutes) { + localEids = ResolveLocalEidRoutes(rootInfo, topoEdges, localId, allLocalIds[peer]); + if (localEids.empty()) { + localEids = ResolveLocalAggregateEidRoutes(rootInfo, localId); + } + const std::vector explicitEids = + SelectExplicitUDMARouteEids(std::getenv("TILEXR_UDMA_ROUTE_EIDS"), localEids); + if (!explicitEids.empty()) { + localEids = explicitEids; + } + } + if (topoReady && localEids.empty() && !ResolveLocalEidRoute(rootInfo, topoEdges, localId, allLocalIds[peer], localEid)) { TILEXR_LOG(WARN) << "TileXR UDMA topology route resolution failed, falling back to EID " << devEids[0].eidIndex; localEid = devEids[0].eidIndex; } - peerLocalEid_[peer] = localEid; - localRouteByPeer[peer] = static_cast(localEid); + if (localEids.empty()) { + localEids.push_back(localEid); + } + if (localEids.size() > routeSlots) { + localEids.resize(routeSlots); + } + peerLocalEids_[peer] = localEids; + peerLocalEid_[peer] = localEids[0]; + for (size_t routeIdx = 0; routeIdx < localEids.size(); ++routeIdx) { + localRouteByPeer[static_cast(peer) * routeSlots + routeIdx] = + static_cast(localEids[routeIdx]); + } if (diag) { TILEXR_LOG(INFO) << "UDMA diag local route rank " << options_.rank << " devLocalId=" << localId << " peer=" << peer << " peerLocalId=" << allLocalIds[peer] - << " localEid=" << localEid; + << " localEid=" << peerLocalEid_[peer] + << " routeCount=" << localEids.size(); } } - std::vector allRouteByPeer(options_.rankSize * options_.rankSize, -1); + std::vector allRouteByPeer(static_cast(options_.rankSize) * options_.rankSize * routeSlots, -1); ret = options_.exchange->AllGather(localRouteByPeer.data(), localRouteByPeer.size(), allRouteByPeer.data()); if (ret != TILEXR_SUCCESS) { return ret; @@ -538,25 +637,82 @@ int TileXRUDMATransport::BuildRoutes() if (peer == options_.rank) { continue; } - int32_t remoteEid = allRouteByPeer[peer * options_.rankSize + options_.rank]; + int32_t remoteEid = allRouteByPeer[(static_cast(peer) * options_.rankSize + options_.rank) * routeSlots]; if (remoteEid < 0 || static_cast(remoteEid) >= eidCount_) { remoteEid = static_cast(devEids[0].eidIndex); } - peerRemoteEid_[peer] = static_cast(remoteEid); + std::vector remoteEids; + for (uint32_t routeIdx = 0; routeIdx < routeSlots; ++routeIdx) { + int32_t candidate = allRouteByPeer[(static_cast(peer) * options_.rankSize + options_.rank) * routeSlots + routeIdx]; + if (candidate >= 0 && static_cast(candidate) < eidCount_) { + remoteEids.push_back(static_cast(candidate)); + } + } + if (remoteEids.empty()) { + remoteEids.assign(1, static_cast(remoteEid)); + } + peerRemoteEid_[peer] = remoteEids[0]; + peerRemoteEids_[peer] = remoteEids; + auto localRoutesIt = peerLocalEids_.find(peer); + if (localRoutesIt == peerLocalEids_.end()) { + peerLocalEids_[peer] = std::vector(1, peerLocalEid_[peer]); + localRoutesIt = peerLocalEids_.find(peer); + } + std::vector& localRoutes = localRoutesIt->second; + const size_t pairedRouteCount = std::min(localRoutes.size(), peerRemoteEids_[peer].size()); + if (pairedRouteCount == 0) { + return TILEXR_ERROR_INTERNAL; + } + if (localRoutes.size() != pairedRouteCount) { + localRoutes.resize(pairedRouteCount); + peerLocalEid_[peer] = localRoutes[0]; + } + if (peerRemoteEids_[peer].size() != pairedRouteCount) { + peerRemoteEids_[peer].resize(pairedRouteCount); + peerRemoteEid_[peer] = peerRemoteEids_[peer][0]; + } + peerQpRouteEids_[peer] = BuildUDMAMultiRouteQpToEid(localRoutes, qpsPerRoute_); + auto weightByEidIt = rootInfo.portCountByEidByLocalId.find(localId); + const std::map emptyWeights; + const std::map& weightByEid = + weightByEidIt == rootInfo.portCountByEidByLocalId.end() ? emptyWeights : weightByEidIt->second; + peerQpRouteWeights_[peer] = BuildUDMAMultiRouteQpWeights(localRoutes, weightByEid, qpsPerRoute_); if (diag) { TILEXR_LOG(INFO) << "UDMA diag remote route rank " << options_.rank << " peer=" << peer << " localEid=" << peerLocalEid_[peer] - << " remoteEid=" << peerRemoteEid_[peer]; + << " remoteEid=" << peerRemoteEid_[peer] + << " pairedRouteCount=" << pairedRouteCount + << " qpsPerRoute=" << qpsPerRoute_; } } + size_t maxRouteCount = 1; + for (const auto& entry : peerQpRouteEids_) { + maxRouteCount = std::max(maxRouteCount, entry.second.size()); + } + qpNum_ = static_cast(maxRouteCount); return TILEXR_SUCCESS; } int TileXRUDMATransport::CreateContexts() { - for (const auto& route : peerLocalEid_) { - const uint32_t eidIndex = route.second; + std::vector contextEids; + for (const auto& route : peerLocalEids_) { + for (uint32_t eidIndex : route.second) { + if (std::find(contextEids.begin(), contextEids.end(), eidIndex) == contextEids.end()) { + contextEids.push_back(eidIndex); + } + } + } + if (contextEids.empty()) { + for (const auto& route : peerLocalEid_) { + if (std::find(contextEids.begin(), contextEids.end(), route.second) == contextEids.end()) { + contextEids.push_back(route.second); + } + } + } + + for (uint32_t eidIndex : contextEids) { if (ctxHandleByEid_.count(eidIndex) != 0) { continue; } @@ -660,11 +816,28 @@ int TileXRUDMATransport::CreateQueues() state.eidIndex = ctxEntry.first; state.ctxHandle = ctxEntry.second; state.tokenHandle = tokenHandleByEid_[ctxEntry.first]; - for (const auto& route : peerLocalEid_) { - if (route.second != state.eidIndex) { + for (int peer = 0; peer < options_.rankSize; ++peer) { + if (peer == options_.rank) { + continue; + } + const auto qpRoutesIt = peerQpRouteEids_.find(peer); + if (qpRoutesIt != peerQpRouteEids_.end()) { + for (uint32_t qpIdx = 0; qpIdx < qpRoutesIt->second.size(); ++qpIdx) { + if (qpRoutesIt->second[qpIdx] != state.eidIndex) { + continue; + } + int ret = CreatePeerQueue(state, peer, qpIdx); + if (ret != TILEXR_SUCCESS) { + return ret; + } + } + continue; + } + const auto localRouteIt = peerLocalEid_.find(peer); + if (localRouteIt == peerLocalEid_.end() || localRouteIt->second != state.eidIndex) { continue; } - int ret = CreatePeerQueue(state, route.first); + int ret = CreatePeerQueue(state, peer, 0); if (ret != TILEXR_SUCCESS) { return ret; } @@ -674,11 +847,13 @@ int TileXRUDMATransport::CreateQueues() return states_.empty() ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; } -int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer) +int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx) { const bool diag = UDMADiagEnabled(); PerEidState::PeerQueueState queue {}; queue.peer = peer; + queue.qpIdx = qpIdx; + const uint32_t localQpIdx = qpsPerRoute_ == 0 ? qpIdx : qpIdx % qpsPerRoute_; ChanInfoT chanInfo {}; chanInfo.in.dataPlaneFlag.bs.poolCqCstm = 1; @@ -694,7 +869,7 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer) if (ret != 0) { return TILEXR_ERROR_INTERNAL; } - queue.localCq.cqn = 0; + queue.localCq.cqn = localQpIdx; queue.localCq.bufAddr = queue.cqInfo.out.bufAddr; queue.localCq.baseBkShift = Log2Uint64(queue.cqInfo.out.cqeSize); queue.localCq.depth = queue.cqInfo.in.depth; @@ -727,7 +902,7 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer) if (ret != 0) { return TILEXR_ERROR_INTERNAL; } - queue.localWq.wqn = 0; + queue.localWq.wqn = localQpIdx; queue.localWq.bufAddr = queue.qpInfo.ub.sqBuffVa; queue.localWq.baseBkShift = Log2Uint64(queue.qpInfo.ub.wqebbSize); queue.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; @@ -746,6 +921,7 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer) if (diag) { TILEXR_LOG(INFO) << "UDMA diag create peer queue rank " << options_.rank << " peer=" << peer + << " qpIdx=" << qpIdx << " eid=" << state.eidIndex << " ctx=" << state.ctxHandle << " chan=" << queue.chanHandle @@ -763,7 +939,11 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer) << " wqebbSize=" << queue.qpInfo.ub.wqebbSize << " cqeSize=" << queue.cqInfo.out.cqeSize; } - state.peerQueues[peer] = queue; + if (qpNum_ <= 1) { + state.peerQueues[peer] = queue; + } else { + state.peerQueues[QueueKey(peer, qpIdx, qpNum_)] = queue; + } return TILEXR_SUCCESS; } @@ -772,11 +952,14 @@ int TileXRUDMATransport::ImportQueues() const bool diag = UDMADiagEnabled(); std::vector localPeerImports(options_.rankSize); std::vector localPeerKeys(options_.rankSize); + const size_t queueSlots = static_cast(options_.rankSize) * qpNum_; + std::vector localQpImports(queueSlots); + std::vector localQpKeys(queueSlots); for (const auto& stateEntry : states_) { const auto& state = stateEntry.second; for (const auto& queueEntry : state.peerQueues) { - const int peer = queueEntry.first; const auto& queue = queueEntry.second; + const int peer = queue.peer; localPeerImports[peer].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; localPeerImports[peer].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; localPeerImports[peer].in.ub.policy = JETTY_GRP_POLICY_RR; @@ -785,6 +968,11 @@ int TileXRUDMATransport::ImportQueues() localPeerImports[peer].in.ub.tpType = 1; localPeerImports[peer].in.key = queue.qpInfo.key; localPeerKeys[peer] = queue.qpInfo.key; + const size_t qpSlot = static_cast(peer) * qpNum_ + (qpNum_ <= 1 ? 0 : queue.qpIdx); + if (qpSlot < localQpImports.size()) { + localQpImports[qpSlot] = localPeerImports[peer]; + localQpKeys[qpSlot] = queue.qpInfo.key; + } } } @@ -798,18 +986,33 @@ int TileXRUDMATransport::ImportQueues() if (ret != TILEXR_SUCCESS) { return ret; } + std::vector allQpImports(static_cast(options_.rankSize) * queueSlots); + ret = options_.exchange->AllGather(localQpImports.data(), localQpImports.size(), allQpImports.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } + std::vector allQpKeys(static_cast(options_.rankSize) * queueSlots); + ret = options_.exchange->AllGather(localQpKeys.data(), localQpKeys.size(), allQpKeys.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } for (auto& stateEntry : states_) { auto& state = stateEntry.second; for (auto& queueEntry : state.peerQueues) { - const int peer = queueEntry.first; auto& queue = queueEntry.second; + const int peer = queue.peer; const uint32_t remoteEid = peerRemoteEid_[peer]; if (remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - QpImportInfoT importInfo = allImports[(peer * options_.rankSize + options_.rank)]; - importInfo.in.key = allKeys[(peer * options_.rankSize + options_.rank)]; + const size_t remoteQpIdx = qpNum_ <= 1 ? 0 : queue.qpIdx; + const size_t allQpIndex = + (static_cast(peer) * options_.rankSize + options_.rank) * qpNum_ + remoteQpIdx; + QpImportInfoT importInfo = allQpIndex < allQpImports.size() ? + allQpImports[allQpIndex] : allImports[(peer * options_.rankSize + options_.rank)]; + importInfo.in.key = allQpIndex < allQpKeys.size() ? + allQpKeys[allQpIndex] : allKeys[(peer * options_.rankSize + options_.rank)]; ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &queue.remoteQpHandle); if (ret != 0) { return TILEXR_ERROR_INTERNAL; @@ -818,6 +1021,7 @@ int TileXRUDMATransport::ImportQueues() if (diag) { TILEXR_LOG(INFO) << "UDMA diag import qp rank " << options_.rank << " peer=" << peer + << " qpIdx=" << queue.localWq.wqn << " localEid=" << state.eidIndex << " remoteEid=" << remoteEid << " remoteQp=" << queue.remoteQpHandle @@ -886,87 +1090,128 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_ERROR_INTERNAL; } - std::vector sq(options_.rankSize); - std::vector rq(options_.rankSize); - std::vector scq(options_.rankSize); - std::vector rcq(options_.rankSize); - std::vector mem(options_.rankSize); + const size_t queueEntries = static_cast(options_.rankSize) * qpNum_; + std::vector sq(queueEntries); + std::vector rq(queueEntries); + std::vector scq(queueEntries); + std::vector rcq(queueEntries); + std::vector mem(queueEntries); + std::vector qpWeights(queueEntries, 1); for (int rank = 0; rank < options_.rankSize; ++rank) { - uint32_t localEid = fallbackEid; - uint32_t remoteEid = fallbackEid; - if (rank != options_.rank) { - localEid = peerLocalEid_[rank]; - remoteEid = peerRemoteEid_[rank]; - } - auto stateIt = states_.find(localEid); - if (stateIt == states_.end()) { - stateIt = fallbackIt; - } - auto& state = stateIt->second; - PerEidState::PeerQueueState* queuePtr = nullptr; - if (rank == options_.rank) { - if (!state.peerQueues.empty()) { - queuePtr = &state.peerQueues.begin()->second; - } else if (!fallbackIt->second.peerQueues.empty()) { - queuePtr = &fallbackIt->second.peerQueues.begin()->second; + for (uint32_t qpIdx = 0; qpIdx < qpNum_; ++qpIdx) { + uint32_t localEid = fallbackEid; + uint32_t remoteEid = fallbackEid; + uint32_t routeQpIdx = qpIdx; + if (qpsPerRoute_ != 0) { + routeQpIdx = qpIdx % qpsPerRoute_; } - } else { - const auto queueIt = state.peerQueues.find(rank); - if (queueIt == state.peerQueues.end()) { + if (rank != options_.rank) { + const auto qpRoutesIt = peerQpRouteEids_.find(rank); + if (qpRoutesIt != peerQpRouteEids_.end() && qpIdx < qpRoutesIt->second.size()) { + localEid = qpRoutesIt->second[qpIdx]; + const auto localRoutesIt = peerLocalEids_.find(rank); + const auto remoteRoutesIt = peerRemoteEids_.find(rank); + if (localRoutesIt != peerLocalEids_.end() && remoteRoutesIt != peerRemoteEids_.end()) { + const auto& localRoutes = localRoutesIt->second; + const auto& remoteRoutes = remoteRoutesIt->second; + auto routeIt = std::find(localRoutes.begin(), localRoutes.end(), localEid); + if (routeIt != localRoutes.end()) { + const size_t routeOffset = static_cast(std::distance(localRoutes.begin(), routeIt)); + if (routeOffset < remoteRoutes.size()) { + remoteEid = remoteRoutes[routeOffset]; + } + } + } + } else { + localEid = peerLocalEid_[rank]; + remoteEid = peerRemoteEid_[rank]; + } + } + auto stateIt = states_.find(localEid); + if (stateIt == states_.end()) { + stateIt = fallbackIt; + routeQpIdx = 0; + } + auto& state = stateIt->second; + PerEidState::PeerQueueState* queuePtr = nullptr; + if (rank == options_.rank) { + if (!state.peerQueues.empty()) { + queuePtr = &state.peerQueues.begin()->second; + } else if (!fallbackIt->second.peerQueues.empty()) { + queuePtr = &fallbackIt->second.peerQueues.begin()->second; + } + } else { + const int key = QueueKey(rank, qpIdx, qpNum_); + auto queueIt = state.peerQueues.find(key); + if (queueIt == state.peerQueues.end()) { + const int routeKey = QueueKey(rank, routeQpIdx, qpNum_); + queueIt = state.peerQueues.find(routeKey); + } + if (queueIt == state.peerQueues.end()) { + return TILEXR_ERROR_INTERNAL; + } + queuePtr = &queueIt->second; + } + if (queuePtr == nullptr) { return TILEXR_ERROR_INTERNAL; } - queuePtr = &queueIt->second; - } - if (queuePtr == nullptr) { - return TILEXR_ERROR_INTERNAL; - } - auto& queue = *queuePtr; - if (!registeredMem_.empty()) { - const auto& localMrs = registeredMem_.begin()->second; - const auto localMrIt = localMrs.find(localEid); - if (localMrIt != localMrs.end()) { - queue.localWq.localTokenId = localMrIt->second.tokenId; + auto& queue = *queuePtr; + if (!registeredMem_.empty()) { + const auto& localMrs = registeredMem_.begin()->second; + const auto localMrIt = localMrs.find(localEid); + if (localMrIt != localMrs.end()) { + queue.localWq.localTokenId = localMrIt->second.tokenId; + } } - } - sq[rank] = queue.localWq; - rq[rank] = queue.localWq; - scq[rank] = queue.localCq; - rcq[rank] = queue.localCq; - if (rank == options_.rank) { - const auto localMemIt = localMemInfoByEid_.find(localEid); - if (localMemIt != localMemInfoByEid_.end()) { - mem[rank] = localMemIt->second; + const size_t entryIndex = static_cast(rank) * qpNum_ + qpIdx; + sq[entryIndex] = queue.localWq; + rq[entryIndex] = queue.localWq; + scq[entryIndex] = queue.localCq; + rcq[entryIndex] = queue.localCq; + if (rank == options_.rank) { + const auto localMemIt = localMemInfoByEid_.find(localEid); + if (localMemIt != localMemInfoByEid_.end()) { + mem[entryIndex] = localMemIt->second; + } + } else { + mem[entryIndex] = allMem[rank * eidCount_ + remoteEid]; + mem[entryIndex].tpn = queue.tpn; + } + mem[entryIndex].eidAddr = reinterpret_cast( + eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); + if (rank != options_.rank) { + const auto weightIt = peerQpRouteWeights_.find(rank); + if (weightIt != peerQpRouteWeights_.end() && qpIdx < weightIt->second.size()) { + qpWeights[entryIndex] = weightIt->second[qpIdx] == 0 ? 1 : weightIt->second[qpIdx]; + } + } + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag info image rank " << options_.rank + << " entryRank=" << rank + << " qpIdx=" << qpIdx + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " weight=" << qpWeights[entryIndex] + << " sqBuf=" << PtrToHex(sq[entryIndex].bufAddr) + << " sqHead=" << PtrToHex(sq[entryIndex].headAddr) + << " sqTail=" << PtrToHex(sq[entryIndex].tailAddr) + << " localTokenId=" << sq[entryIndex].localTokenId + << " wqeCnt=" << PtrToHex(sq[entryIndex].wqeCntAddr) + << " cqBuf=" << PtrToHex(scq[entryIndex].bufAddr) + << " cqTail=" << PtrToHex(scq[entryIndex].tailAddr) + << " memAddr=" << PtrToHex(mem[entryIndex].addr) + << " memLen=" << mem[entryIndex].len + << " memTid=" << mem[entryIndex].tid + << " memTpn=" << mem[entryIndex].tpn + << " memEidAddr=" << PtrToHex(mem[entryIndex].eidAddr); } - } else { - mem[rank] = allMem[rank * eidCount_ + remoteEid]; - mem[rank].tpn = queue.tpn; } - mem[rank].eidAddr = reinterpret_cast( - eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); - if (diag) { - TILEXR_LOG(INFO) << "UDMA diag info image rank " << options_.rank - << " entryRank=" << rank - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " sqBuf=" << PtrToHex(sq[rank].bufAddr) - << " sqHead=" << PtrToHex(sq[rank].headAddr) - << " sqTail=" << PtrToHex(sq[rank].tailAddr) - << " localTokenId=" << sq[rank].localTokenId - << " wqeCnt=" << PtrToHex(sq[rank].wqeCntAddr) - << " cqBuf=" << PtrToHex(scq[rank].bufAddr) - << " cqTail=" << PtrToHex(scq[rank].tailAddr) - << " memAddr=" << PtrToHex(mem[rank].addr) - << " memLen=" << mem[rank].len - << " memTid=" << mem[rank].tid - << " memTpn=" << mem[rank].tpn - << " memEidAddr=" << PtrToHex(mem[rank].eidAddr); - } - } - - const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); + } + + const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo) + sizeof(uint32_t); const uint32_t requiredInfoSize = - static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize); + static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize * qpNum_); if (udmaInfoDev_ == nullptr || udmaInfoSize_ < requiredInfoSize) { if (udmaInfoDev_ != nullptr) { aclrtFree(udmaInfoDev_); @@ -982,7 +1227,8 @@ int TileXRUDMATransport::RefreshUDMAInfo() UDMAInfo info {}; std::vector image; - ret = BuildUDMAInfoImage(reinterpret_cast(udmaInfoDev_), sq, rq, scq, rcq, mem, info, image); + ret = BuildUDMAInfoImage( + reinterpret_cast(udmaInfoDev_), qpNum_, sq, rq, scq, rcq, mem, qpWeights, info, image); if (ret != TILEXR_UDMA_LAYOUT_SUCCESS) { return TILEXR_ERROR_PARA_CHECK_FAIL; } @@ -1176,57 +1422,76 @@ int TileXRUDMATransport::ExchangeAndImportMemory() return ret; } - remoteMemHandles_.assign(options_.rankSize, nullptr); + remoteMemHandlesByPeer_.clear(); for (int peer = 0; peer < options_.rankSize; ++peer) { if (peer == options_.rank) { continue; } - const uint32_t remoteEid = peerRemoteEid_[peer]; - const ExchangedMrInfo* remote = nullptr; - for (uint32_t i = 0; i < allCounts[peer]; ++i) { - const auto& candidate = all[peer * maxCount + i]; - if (candidate.valid != 0 && candidate.eidIndex == remoteEid) { - remote = &candidate; - break; - } - } - if (remote == nullptr) { - TILEXR_LOG(WARN) << "UDMA remote memory info missing rank " << options_.rank - << " peer=" << peer - << " remoteEid=" << remoteEid - << " peerCount=" << allCounts[peer] - << " maxCount=" << maxCount; - return TILEXR_ERROR_INTERNAL; + auto localRoutesIt = peerLocalEids_.find(peer); + auto remoteRoutesIt = peerRemoteEids_.find(peer); + std::vector localRoutes; + std::vector remoteRoutes; + if (localRoutesIt != peerLocalEids_.end() && remoteRoutesIt != peerRemoteEids_.end()) { + localRoutes = localRoutesIt->second; + remoteRoutes = remoteRoutesIt->second; + } else { + localRoutes.push_back(peerLocalEid_[peer]); + remoteRoutes.push_back(peerRemoteEid_[peer]); } - const uint32_t localEid = peerLocalEid_[peer]; - MrImportInfoT importInfo {}; - importInfo.in.key = remote->mr.key; - importInfo.in.ub.tokenValue = remote->mr.tokenValue; - importInfo.in.ub.flags.bs.cacheable = remote->mr.cacheable; - importInfo.in.ub.flags.bs.access = remote->mr.access; - void* remoteHandle = nullptr; - ret = loader_.RaCtxRmemImport(ctxHandleByEid_[localEid], &importInfo, &remoteHandle); - if (ret != 0 || remoteHandle == nullptr) { - TILEXR_LOG(WARN) << "UDMA RaCtxRmemImport failed rank " << options_.rank - << " peer=" << peer - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " ctx=" << ctxHandleByEid_[localEid] - << " ret=" << ret - << " handle=" << remoteHandle - << " remoteToken=" << remote->mr.tokenValue - << " remoteTokenId=" << remote->mr.tokenId - << " remoteKeySize=" << static_cast(remote->mr.key.size); + if (localRoutes.empty() || localRoutes.size() != remoteRoutes.size()) { return TILEXR_ERROR_INTERNAL; } - remoteMemHandles_[peer] = remoteHandle; - if (UDMADiagEnabled()) { - TILEXR_LOG(INFO) << "UDMA diag rmem imported rank " << options_.rank - << " peer=" << peer - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " remoteHandle=" << remoteHandle; + std::vector remoteHandles(localRoutes.size(), nullptr); + for (size_t routeIdx = 0; routeIdx < localRoutes.size(); ++routeIdx) { + const uint32_t localEid = localRoutes[routeIdx]; + const uint32_t remoteEid = remoteRoutes[routeIdx]; + const ExchangedMrInfo* remote = nullptr; + for (uint32_t i = 0; i < allCounts[peer]; ++i) { + const auto& candidate = all[peer * maxCount + i]; + if (candidate.valid != 0 && candidate.eidIndex == remoteEid) { + remote = &candidate; + break; + } + } + if (remote == nullptr || ctxHandleByEid_.count(localEid) == 0) { + TILEXR_LOG(WARN) << "UDMA remote memory info missing rank " << options_.rank + << " peer=" << peer + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " peerCount=" << allCounts[peer] + << " maxCount=" << maxCount; + return TILEXR_ERROR_INTERNAL; + } + MrImportInfoT importInfo {}; + importInfo.in.key = remote->mr.key; + importInfo.in.ub.tokenValue = remote->mr.tokenValue; + importInfo.in.ub.flags.bs.cacheable = remote->mr.cacheable; + importInfo.in.ub.flags.bs.access = remote->mr.access; + void* remoteHandle = nullptr; + ret = loader_.RaCtxRmemImport(ctxHandleByEid_[localEid], &importInfo, &remoteHandle); + if (ret != 0 || remoteHandle == nullptr) { + TILEXR_LOG(WARN) << "UDMA RaCtxRmemImport failed rank " << options_.rank + << " peer=" << peer + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " ctx=" << ctxHandleByEid_[localEid] + << " ret=" << ret + << " handle=" << remoteHandle + << " remoteToken=" << remote->mr.tokenValue + << " remoteTokenId=" << remote->mr.tokenId + << " remoteKeySize=" << static_cast(remote->mr.key.size); + return TILEXR_ERROR_INTERNAL; + } + remoteHandles[routeIdx] = remoteHandle; + if (UDMADiagEnabled()) { + TILEXR_LOG(INFO) << "UDMA diag rmem imported rank " << options_.rank + << " peer=" << peer + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " remoteHandle=" << remoteHandle; + } } + remoteMemHandlesByPeer_[peer] = remoteHandles; } return TILEXR_SUCCESS; } @@ -1247,14 +1512,28 @@ int TileXRUDMATransport::UnregisterMemory(GM_ADDR localPtr) void TileXRUDMATransport::CleanupMemory() { - for (int peer = 0; peer < static_cast(remoteMemHandles_.size()); ++peer) { - if (peer == options_.rank || remoteMemHandles_[peer] == nullptr) { + for (auto& peerEntry : remoteMemHandlesByPeer_) { + const int peer = peerEntry.first; + if (peer == options_.rank) { continue; } - const uint32_t localEid = peerLocalEid_[peer]; - loader_.RaCtxRmemUnimport(ctxHandleByEid_[localEid], remoteMemHandles_[peer]); - remoteMemHandles_[peer] = nullptr; + auto localRoutesIt = peerLocalEids_.find(peer); + for (size_t routeIdx = 0; routeIdx < peerEntry.second.size(); ++routeIdx) { + void* remoteHandle = peerEntry.second[routeIdx]; + if (remoteHandle == nullptr) { + continue; + } + uint32_t localEid = peerLocalEid_[peer]; + if (localRoutesIt != peerLocalEids_.end() && routeIdx < localRoutesIt->second.size()) { + localEid = localRoutesIt->second[routeIdx]; + } + if (ctxHandleByEid_.count(localEid) != 0) { + loader_.RaCtxRmemUnimport(ctxHandleByEid_[localEid], remoteHandle); + } + peerEntry.second[routeIdx] = nullptr; + } } + remoteMemHandlesByPeer_.clear(); for (const auto& mrEntry : registeredMem_) { for (const auto& eidMr : mrEntry.second) { const uint32_t eidIndex = eidMr.first; @@ -1350,8 +1629,14 @@ void TileXRUDMATransport::Shutdown() localEidByEid_.clear(); peerLocalEid_.clear(); peerRemoteEid_.clear(); + peerLocalEids_.clear(); + peerRemoteEids_.clear(); + peerQpRouteEids_.clear(); + peerQpRouteWeights_.clear(); + qpsPerRoute_ = 1; + qpNum_ = 1; localMemInfoByEid_.clear(); - remoteMemHandles_.clear(); + remoteMemHandlesByPeer_.clear(); loader_.Unload(); } diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index d3adfac6..40c2e6e6 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -53,7 +53,7 @@ class TileXRUDMATransport { int BuildRoutes(); int CreateContexts(); int CreateQueues(); - int CreatePeerQueue(PerEidState& state, int peer); + int CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx); int ImportQueues(); int EnsureUDMAInfoBuffer(); int RefreshUDMAInfo(); @@ -75,14 +75,20 @@ class TileXRUDMATransport { uint32_t logicDevId_ = 0; uint32_t deviceIdOffset_ = 0; uint32_t eidCount_ = 0; + uint32_t qpNum_ = 1; + uint32_t qpsPerRoute_ = 1; std::map ctxHandleByEid_; std::map tokenHandleByEid_; std::map peerLocalEid_; std::map peerRemoteEid_; + std::map> peerLocalEids_; + std::map> peerRemoteEids_; + std::map> peerQpRouteEids_; + std::map> peerQpRouteWeights_; std::map states_; std::map localEidByEid_; MemoryRegionMap registeredMem_; - std::vector remoteMemHandles_; + std::map> remoteMemHandlesByPeer_; std::map localMemInfoByEid_; GM_ADDR udmaInfoDev_ = nullptr; GM_ADDR eidTableDev_ = nullptr; diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index 0f2c853f..94584c08 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -109,9 +109,21 @@ __aicore__ inline __gm__ UDMACQCtx* UDMAGetSCQCtx(__gm__ UDMAInfo* udmaInfo, uin return reinterpret_cast<__gm__ UDMACQCtx*>(udmaInfo->scqPtr + (pe * qpNum + qpIdx) * sizeof(UDMACQCtx)); } -__aicore__ inline __gm__ UDMAMemInfo* UDMAGetRemoteMemInfo(__gm__ UDMAInfo* udmaInfo, uint32_t pe) +__aicore__ inline __gm__ UDMAMemInfo* UDMAGetRemoteMemInfo(__gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx) { - return reinterpret_cast<__gm__ UDMAMemInfo*>(udmaInfo->memPtr + sizeof(UDMAMemInfo) * pe); + uint32_t qpNum = udmaInfo->qpNum; + return reinterpret_cast<__gm__ UDMAMemInfo*>(udmaInfo->memPtr + sizeof(UDMAMemInfo) * (pe * qpNum + qpIdx)); +} + +__aicore__ inline uint32_t UDMAGetQpWeight(__gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx) +{ + uint32_t qpNum = udmaInfo->qpNum; + if (udmaInfo->qpWeightPtr == 0) { + return 1; + } + auto weights = reinterpret_cast<__gm__ uint32_t*>(udmaInfo->qpWeightPtr); + uint32_t weight = weights[pe * qpNum + qpIdx]; + return weight == 0 ? 1 : weight; } __aicore__ inline void UDMAPollCQUpdateInfo( @@ -255,7 +267,7 @@ __aicore__ inline void UDMAPostSend( uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); UDMAPollCQWhenSQOverflow(udmaInfo, qpCtxEntry, wqeCnt, pe, qpIdx); - __gm__ UDMAMemInfo* remoteMemInfo = UDMAGetRemoteMemInfo(udmaInfo, pe); + __gm__ UDMAMemInfo* remoteMemInfo = UDMAGetRemoteMemInfo(udmaInfo, pe, qpIdx); __gm__ uint8_t* wqeAddr = reinterpret_cast<__gm__ uint8_t*>(qpCtxEntry->bufAddr + wqeSize * (curHead % depth)); __gm__ UDMASqeCtx* sqeCtx = reinterpret_cast<__gm__ UDMASqeCtx*>(wqeAddr); @@ -300,8 +312,9 @@ __aicore__ inline void UDMAWriteNotify( } template -__aicore__ inline void UDMAPutNbi( - const __gm__ CommArgs* args, int targetRank, const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount) +__aicore__ inline void UDMAPutNbiOnQp( + const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx, + const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount) { if (!UDMARegistryEnabled(args)) return; @@ -310,7 +323,14 @@ __aicore__ inline void UDMAPutNbi( auto remoteAddr = UDMARegisteredRemoteAddr(registry, targetRank, byteOffset); UDMAWrite(args, remoteAddr, reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)), - targetRank, 0, byteCount); + targetRank, qpIdx, byteCount); +} + +template +__aicore__ inline void UDMAPutNbi( + const __gm__ CommArgs* args, int targetRank, const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount) +{ + UDMAPutNbiOnQp(args, targetRank, 0, localSrc, byteOffset, byteCount); } template @@ -321,8 +341,9 @@ __aicore__ inline void UDMAPutRegisteredNbi( } template -__aicore__ inline void UDMAGetNbi( - const __gm__ CommArgs* args, int sourceRank, __gm__ T* localDst, uint64_t byteOffset, uint32_t byteCount) +__aicore__ inline void UDMAGetNbiOnQp( + const __gm__ CommArgs* args, int sourceRank, uint32_t qpIdx, + __gm__ T* localDst, uint64_t byteOffset, uint32_t byteCount) { if (!UDMARegistryEnabled(args)) return; @@ -330,7 +351,14 @@ __aicore__ inline void UDMAGetNbi( if (!UDMARegisteredRangeValid(registry, sourceRank, byteOffset, byteCount)) return; auto remoteAddr = UDMARegisteredRemoteAddr(registry, sourceRank, byteOffset); - UDMARead(args, reinterpret_cast<__gm__ uint8_t*>(localDst), remoteAddr, sourceRank, 0, byteCount); + UDMARead(args, reinterpret_cast<__gm__ uint8_t*>(localDst), remoteAddr, sourceRank, qpIdx, byteCount); +} + +template +__aicore__ inline void UDMAGetNbi( + const __gm__ CommArgs* args, int sourceRank, __gm__ T* localDst, uint64_t byteOffset, uint32_t byteCount) +{ + UDMAGetNbiOnQp(args, sourceRank, 0, localDst, byteOffset, byteCount); } template @@ -341,9 +369,10 @@ __aicore__ inline void UDMAGetRegisteredNbi( } template -__aicore__ inline void UDMAPutSignalNbi( - const __gm__ CommArgs* args, int targetRank, const __gm__ T* localSrc, uint64_t byteOffset, - uint32_t byteCount, uint64_t signalByteOffset, uint64_t signal) +__aicore__ inline void UDMAPutSignalNbiOnQp( + const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx, + const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount, + uint64_t signalByteOffset, uint64_t signal) { if (!UDMARegistryEnabled(args)) return; @@ -359,7 +388,15 @@ __aicore__ inline void UDMAPutSignalNbi( signalParams.signal = signal; auto remoteAddr = UDMARegisteredRemoteAddr(registry, targetRank, byteOffset); UDMAWriteNotify(args, remoteAddr, reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)), - targetRank, 0, byteCount, &signalParams); + targetRank, qpIdx, byteCount, &signalParams); +} + +template +__aicore__ inline void UDMAPutSignalNbi( + const __gm__ CommArgs* args, int targetRank, const __gm__ T* localSrc, uint64_t byteOffset, + uint32_t byteCount, uint64_t signalByteOffset, uint64_t signal) +{ + UDMAPutSignalNbiOnQp(args, targetRank, 0, localSrc, byteOffset, byteCount, signalByteOffset, signal); } template @@ -379,6 +416,15 @@ __aicore__ inline void UDMAQuiet(const __gm__ CommArgs* args, int targetRank) (void)UDMAPollCQ(udmaInfo, targetRank, 0, wqeCnt); } +__aicore__ inline void UDMAQuietOnQp(const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx) +{ + if (!UDMAEnabled(args)) return; + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, qpIdx); + uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + (void)UDMAPollCQ(udmaInfo, targetRank, qpIdx, wqeCnt); +} + __aicore__ inline uint32_t UDMAQuietStatus(const __gm__ CommArgs* args, int targetRank) { if (!UDMAEnabled(args)) return 0xFFFFFFFFU; @@ -388,6 +434,15 @@ __aicore__ inline uint32_t UDMAQuietStatus(const __gm__ CommArgs* args, int targ return UDMAPollCQ(udmaInfo, targetRank, 0, wqeCnt); } +__aicore__ inline uint32_t UDMAQuietStatusOnQp(const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx) +{ + if (!UDMAEnabled(args)) return 0xFFFFFFFFU; + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, qpIdx); + uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + return UDMAPollCQ(udmaInfo, targetRank, qpIdx, wqeCnt); +} + } // namespace TileXR #endif // TILEXR_UDMA_H diff --git a/src/include/tilexr_udma_types.h b/src/include/tilexr_udma_types.h index 1fd6ebd7..9b344854 100644 --- a/src/include/tilexr_udma_types.h +++ b/src/include/tilexr_udma_types.h @@ -75,6 +75,7 @@ struct UDMAInfo { uint64_t scqPtr; uint64_t rcqPtr; uint64_t memPtr; + uint64_t qpWeightPtr; }; struct UDMASqeCtx { diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index 3bc60ab5..a48e58c4 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -29,6 +29,7 @@ constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL constexpr size_t kAllToAllBigDataMultiNodeRegisteredBytes = 1024ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataMultiNodePeerSlotBytes = 8ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataControlSlotBytes = 64ULL; +constexpr size_t kAllToAllBigDataRegistrationAlignment = 2ULL * 1024ULL * 1024ULL; constexpr uint32_t kAllToAllBigDataCoresPerPeer = 5U; constexpr uint32_t kAllToAllBigDataSingleNodeShards = 2U; constexpr uint32_t kAllToAllBigDataLocalCopyShards = kAllToAllBigDataSingleNodeShards; @@ -152,7 +153,8 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( plan.dataBytes = plan.registeredBytes - plan.controlBytes - plan.signalBytes; if (use35Core) { plan.chunkElements = static_cast(chunkElements); - plan.chunkBytesPerPeer = kAllToAllBigDataMultiNodePeerSlotBytes; + const size_t chunkBytes = chunkElements * sizeof(int32_t); + plan.chunkBytesPerPeer = std::min(chunkBytes, kAllToAllBigDataMultiNodePeerSlotBytes); } else { plan.chunkElements = static_cast( plan.dataBytes / (dataSlotCount * sizeof(int32_t))); @@ -182,6 +184,12 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( } plan.ackSignalOffset = plan.readySignalOffset + controlGroupBytes; plan.controlBytes = controlGroupBytes; + if (use35Core) { + const size_t usedBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; + plan.registeredBytes = + ((usedBytes + kAllToAllBigDataRegistrationAlignment - 1) / + kAllToAllBigDataRegistrationAlignment) * kAllToAllBigDataRegistrationAlignment; + } return plan; } diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index bde5829e..146dd6d1 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -87,6 +87,53 @@ __aicore__ inline uint64_t AllToAllPayloadBytes(int32_t elementsPerPeer) return static_cast(elementsPerPeer) * sizeof(int32_t); } +__aicore__ inline uint64_t BigDataChunkBytesPerPeer(bool use35Core, int32_t effectiveChunkElements) +{ + const uint64_t chunkBytes = static_cast(effectiveChunkElements) * sizeof(int32_t); + if (!use35Core || chunkBytes < TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES) { + return chunkBytes; + } + return TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES; +} + +__aicore__ inline void TileXRUdmaDemoWeightedWqeSlice( + __gm__ TileXR::UDMAInfo* udmaInfo, int32_t peer, uint32_t total, + uint32_t wqeCount, uint32_t wqeIdx, uint32_t& offset, uint32_t& bytes) +{ + uint32_t weightSum = 0; + uint32_t prefixWeight = 0; + for (uint32_t i = 0; i < wqeCount; ++i) { + uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, i); + if (i < wqeIdx) { + prefixWeight += weight; + } + weightSum += weight; + } + if (wqeCount == 0 || weightSum == 0 || wqeIdx >= wqeCount) { + offset = total; + bytes = 0; + return; + } + + uint64_t rawStart = static_cast(total) * prefixWeight / weightSum; + uint64_t rawEnd = static_cast(total) * + (prefixWeight + TileXR::UDMAGetQpWeight(udmaInfo, peer, wqeIdx)) / weightSum; + uint32_t alignedStart = static_cast( + (rawStart / TileXR::BLOCK_UNIT_BYTE) * TileXR::BLOCK_UNIT_BYTE); + uint32_t alignedEnd = wqeIdx + 1 == wqeCount ? total : static_cast( + ((rawEnd + TileXR::BLOCK_UNIT_BYTE - 1) / TileXR::BLOCK_UNIT_BYTE) * TileXR::BLOCK_UNIT_BYTE); + if (alignedEnd > total) { + alignedEnd = total; + } + if (alignedStart >= alignedEnd) { + offset = total; + bytes = 0; + return; + } + offset = alignedStart; + bytes = alignedEnd - alignedStart; +} + __aicore__ inline uint64_t AllToAllDataAsFlagSegmentBytes(uint64_t payloadBytes) { return static_cast(TileXR::DataAsFlagBlockCountForPayloadBytes(payloadBytes)) * @@ -730,7 +777,7 @@ __aicore__ inline void BigDataSendPeerWorker( auto registry = TileXR::GetUDMARegistry(args); auto udmaInfo = TileXR::GetUDMAInfo(args); auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); - auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); + auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer, 0); bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteDataOffset, chunkBytes) && TileXR::UDMARegisteredRangeValid(registry, peer, remoteReadyOffset, sizeof(uint64_t)); uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); @@ -1238,7 +1285,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( auto registry = TileXR::GetUDMARegistry(args); auto udmaInfo = TileXR::GetUDMAInfo(args); auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); - auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer); + auto remoteMemInfo = TileXR::UDMAGetRemoteMemInfo(udmaInfo, peer, 0); bool rangeValid = TileXR::UDMARegisteredRangeValid(registry, peer, remoteOffset, bytes); uint32_t wqeBefore = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); if (debug != nullptr && peer < 16) { @@ -1883,6 +1930,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_p2p_latency_kernel( const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; const uint64_t payloadBytes = AllToAllPayloadBytes(effectiveChunkElements); const uint32_t bytes = static_cast(payloadBytes); + auto udmaInfo = TileXR::GetUDMAInfo(args); + const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; // One block per peer, skip self (peer == rank): no local copy here. for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { @@ -1892,8 +1941,25 @@ extern "C" __global__ __aicore__ void tilexr_udma_p2p_latency_kernel( auto localSrc = input + static_cast(peer) * elementsPerPeer + inputElementOffset; uint64_t remoteOffset = outputByteOffset + static_cast(rank) * payloadBytes; - TileXR::UDMAPutNbi(args, peer, localSrc, remoteOffset, bytes); - uint32_t status = TileXR::UDMAQuietStatus(args, peer); + uint32_t status = 0; + for (uint32_t qpIdx = 0; qpIdx < qpCount; ++qpIdx) { + uint32_t sliceOffset = 0; + uint32_t sliceBytes = 0; + TileXRUdmaDemoWeightedWqeSlice(udmaInfo, peer, bytes, qpCount, qpIdx, sliceOffset, sliceBytes); + if (sliceBytes == 0) { + continue; + } + auto sliceSrc = reinterpret_cast<__gm__ int32_t*>( + reinterpret_cast<__gm__ uint8_t*>(localSrc) + sliceOffset); + TileXR::UDMAPutNbiOnQp( + args, peer, qpIdx, sliceSrc, remoteOffset + sliceOffset, sliceBytes); + } + for (uint32_t qpIdx = 0; qpIdx < qpCount; ++qpIdx) { + uint32_t qpStatus = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + if (qpStatus != 0 && status == 0) { + status = qpStatus; + } + } if (debug != nullptr && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); } @@ -2186,9 +2252,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( const bool use35Core = BigDataUse35Core(rankSize, force35Core); const uint32_t shardCount = BigDataShardCount(rankSize, force35Core); const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; - const uint64_t chunkBytesPerPeer = use35Core ? - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES : - static_cast(effectiveChunkElements) * sizeof(int32_t); + const uint64_t chunkBytesPerPeer = BigDataChunkBytesPerPeer(use35Core, effectiveChunkElements); const uint64_t sendDataOffset = dataOffset; const uint64_t recvDataOffset = sendDataOffset + diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index ec44ebde..b09d973e 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -143,6 +143,7 @@ void TestAllToAllBigDataPlan() CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 128ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, 1024ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 8ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataRegistrationAlignment, 2ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 64ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 5U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataSingleNodeShards, 2U); @@ -190,7 +191,7 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() TileXR::Demo::kAllToAllBigDataControlSlotBytes; CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize), 16U); - CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); + CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, true); CHECK_EQ(plan.passCount, 1U); CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); CHECK_EQ(plan.controlBytes, controlGroupBytes); @@ -206,6 +207,26 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); } +void TestAllToAllBigDataMultiNodeSmallPayloadShrinksRegisteredBytes() +{ + constexpr int rankSize = 16; + constexpr int32_t elementsPerPeer = 262144; // 1 MiB per peer for int32_t. + const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); + const size_t expectedChunkBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + const size_t expectedDataBytes = + static_cast(rankSize - 1) * static_cast(plan.passCount) * 2ULL * expectedChunkBytes; + const size_t usedBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; + + CHECK_EQ(plan.passCount, 1U); + CHECK_EQ(plan.chunkBytesPerPeer, expectedChunkBytes); + CHECK_EQ(plan.dataBytes, expectedDataBytes); + CHECK_EQ(plan.registeredBytes, + ((usedBytes + TileXR::Demo::kAllToAllBigDataRegistrationAlignment - 1) / + TileXR::Demo::kAllToAllBigDataRegistrationAlignment) * + TileXR::Demo::kAllToAllBigDataRegistrationAlignment); + CHECK_EQ(plan.registeredBytes < TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, true); +} + void TestAllToAllBigDataForce35CorePlanFor8P() { constexpr int rankSize = 8; @@ -221,7 +242,7 @@ void TestAllToAllBigDataForce35CorePlanFor8P() CHECK_EQ(TileXR::Demo::AllToAllBigDataUse35Core(rankSize, true), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize, true), 16U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(rankSize, true), 35U); - CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); + CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, true); CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); CHECK_EQ(plan.controlBytes, controlGroupBytes); CHECK_EQ(plan.signalBytes, 4ULL * controlGroupBytes); @@ -523,6 +544,7 @@ int main() TestAllToAllMaxRank256With64MiBPerRank(); TestAllToAllBigDataPlan(); TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone(); + TestAllToAllBigDataMultiNodeSmallPayloadShrinksRegisteredBytes(); TestAllToAllBigDataForce35CorePlanFor8P(); TestAllToAllBigDataBlockDim(); TestAllToAllBigDataMultiNodeTopology(); diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 2fc3bfa3..3a51eac8 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -72,11 +73,12 @@ void TestHostLayoutUsesDeviceRelativePointers() mem[1].addr = 0x4000; mem[1].tid = 7; mem[1].tpn = 9; + std::vector weights = {3, 1}; constexpr uintptr_t deviceBase = 0x100000000ULL; TileXR::UDMAInfo info = {}; std::vector bytes; - const int ret = TileXR::BuildUDMAInfoImage(deviceBase, sq, rq, scq, rcq, mem, info, bytes); + const int ret = TileXR::BuildUDMAInfoImage(deviceBase, 1, sq, rq, scq, rcq, mem, weights, info, bytes); CHECK_EQ(ret, TileXR::TILEXR_UDMA_LAYOUT_SUCCESS); CHECK_EQ(info.qpNum, 1U); @@ -85,10 +87,12 @@ void TestHostLayoutUsesDeviceRelativePointers() CHECK_TRUE(info.scqPtr > info.rqPtr); CHECK_TRUE(info.rcqPtr > info.scqPtr); CHECK_TRUE(info.memPtr > info.rcqPtr); + CHECK_TRUE(info.qpWeightPtr > info.memPtr); CHECK_EQ(bytes.size(), sizeof(TileXR::UDMAInfo) + 2 * sizeof(TileXR::UDMAWQCtx) + 2 * sizeof(TileXR::UDMAWQCtx) + 2 * sizeof(TileXR::UDMACQCtx) + - 2 * sizeof(TileXR::UDMACQCtx) + 2 * sizeof(TileXR::UDMAMemInfo)); + 2 * sizeof(TileXR::UDMACQCtx) + 2 * sizeof(TileXR::UDMAMemInfo) + + 2 * sizeof(uint32_t)); const auto* imageInfo = reinterpret_cast(bytes.data()); CHECK_EQ(imageInfo->sqPtr, info.sqPtr); @@ -96,10 +100,14 @@ void TestHostLayoutUsesDeviceRelativePointers() bytes.data() + (info.sqPtr - deviceBase)); const auto* imageMem = reinterpret_cast( bytes.data() + (info.memPtr - deviceBase)); + const auto* imageWeights = reinterpret_cast( + bytes.data() + (info.qpWeightPtr - deviceBase)); CHECK_EQ(imageSq[1].bufAddr, static_cast(0x2000)); CHECK_EQ(imageMem[1].addr, static_cast(0x4000)); CHECK_EQ(imageMem[1].tid, 7U); CHECK_EQ(imageMem[1].tpn, 9U); + CHECK_EQ(imageWeights[0], 3U); + CHECK_EQ(imageWeights[1], 1U); } void TestRejectsMismatchedArrays() @@ -116,6 +124,53 @@ void TestRejectsMismatchedArrays() CHECK_EQ(ret, TileXR::TILEXR_UDMA_LAYOUT_INVALID); } +void TestMultiRouteQpMappingRepeatsEachRoute() +{ + const std::vector routeEids = {7, 8}; + const std::vector qpToEid = TileXR::BuildUDMAMultiRouteQpToEid(routeEids, 2); + + CHECK_EQ(qpToEid.size(), static_cast(4)); + CHECK_EQ(qpToEid[0], 7U); + CHECK_EQ(qpToEid[1], 7U); + CHECK_EQ(qpToEid[2], 8U); + CHECK_EQ(qpToEid[3], 8U); +} + +void TestMultiRouteQpMappingRejectsEmptyInputs() +{ + CHECK_TRUE(TileXR::BuildUDMAMultiRouteQpToEid({}, 1).empty()); + CHECK_TRUE(TileXR::BuildUDMAMultiRouteQpToEid({7}, 0).empty()); +} + +void TestMultiRouteQpWeightsUseRouteBandwidth() +{ + const std::vector routeEids = {7, 8}; + const std::map routeWeights = {{7, 6}, {8, 2}}; + const std::vector qpWeights = TileXR::BuildUDMAMultiRouteQpWeights(routeEids, routeWeights, 1); + + CHECK_EQ(qpWeights.size(), static_cast(2)); + CHECK_EQ(qpWeights[0], 6U); + CHECK_EQ(qpWeights[1], 2U); +} + +void TestExplicitRouteSelectionKeepsRequestedCandidateOrder() +{ + const std::vector candidates = {7, 8}; + const std::vector selected = TileXR::SelectExplicitUDMARouteEids("8,7,9,bad", candidates); + + CHECK_EQ(selected.size(), static_cast(2)); + CHECK_EQ(selected[0], 8U); + CHECK_EQ(selected[1], 7U); +} + +void TestExplicitRouteSelectionRejectsMissingInputs() +{ + const std::vector candidates = {7, 8}; + CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("", candidates).empty()); + CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("9", candidates).empty()); + CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("8", {}).empty()); +} + void TestTransportUsesPerPeerQueues() { const std::string transport = @@ -192,6 +247,11 @@ int main() { TestHostLayoutUsesDeviceRelativePointers(); TestRejectsMismatchedArrays(); + TestMultiRouteQpMappingRepeatsEachRoute(); + TestMultiRouteQpMappingRejectsEmptyInputs(); + TestMultiRouteQpWeightsUseRouteBandwidth(); + TestExplicitRouteSelectionKeepsRequestedCandidateOrder(); + TestExplicitRouteSelectionRejectsMissingInputs(); TestTransportUsesPerPeerQueues(); TestRootInfoEidBytesSelectRuntimeContexts(); TestMemoryRegistrationUsesOfficialUbFlags(); From ef166e0b4902fc49c74e04a4f6c373da521ad0d8 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Mon, 29 Jun 2026 23:47:48 +0800 Subject: [PATCH 021/163] perf(udma): tune multinode bigdata alltoall sync Use 35-core multinode bigdata scheduling with expanded control slots and per-core fanout flags to reduce shared flag contention. Route recv-side cross-card ack publication through UDMA signal while keeping local synchronization on MTE, and skip the slot reuse ack wait in the 35-core path after validation showed it caused repeat>1 hangs. Update UDMA layout/source guard tests for the new control layout and transport registration limit. --- src/comm/udma/tilexr_udma_transport.cpp | 2 +- tests/udma/demo/tilexr_udma_alltoall_layout.h | 15 +- tests/udma/demo/tilexr_udma_demo.cpp | 7 +- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 340 ++++++++++++++---- .../unit/test_tilexr_udma_alltoall_layout.cpp | 54 ++- .../test_tilexr_udma_transport_layout.cpp | 2 +- 6 files changed, 305 insertions(+), 115 deletions(-) diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 0b1ba619..b27e819b 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -1326,7 +1326,7 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes 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 = 0; + mrInfo.in.ub.flags.bs.nonPin = 1; 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; diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index a48e58c4..c7492ef5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -28,8 +28,7 @@ struct AllToAllChunkPlan { constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataMultiNodeRegisteredBytes = 1024ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataMultiNodePeerSlotBytes = 8ULL * 1024ULL * 1024ULL; -constexpr size_t kAllToAllBigDataControlSlotBytes = 64ULL; -constexpr size_t kAllToAllBigDataRegistrationAlignment = 2ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataControlSlotBytes = 128ULL; constexpr uint32_t kAllToAllBigDataCoresPerPeer = 5U; constexpr uint32_t kAllToAllBigDataSingleNodeShards = 2U; constexpr uint32_t kAllToAllBigDataLocalCopyShards = kAllToAllBigDataSingleNodeShards; @@ -37,6 +36,7 @@ constexpr uint32_t kAllToAllBigDataPingPongSlots = 2U; constexpr int32_t kAllToAllBigDataRanksPerNode = 8; constexpr uint32_t kAllToAllBigDataMultiNodeCopyCores = 16U; constexpr uint32_t kAllToAllBigDataMultiNodeRecvCores = 16U; +constexpr uint32_t kAllToAllBigDataMultiNodeControlShards = 32U; constexpr uint32_t kAllToAllBigDataMultiNodeRemoteSendPrimaryCore = 16U; constexpr uint32_t kAllToAllBigDataMultiNodeRemoteSendSecondaryCore = 17U; constexpr uint32_t kAllToAllBigDataMultiNodeLocalSendCore = 18U; @@ -78,7 +78,7 @@ inline bool AllToAllBigDataUse35Core(int rankSize, bool force35Core = false) inline uint32_t AllToAllBigDataShardCount(int rankSize, bool force35Core = false) { return AllToAllBigDataUse35Core(rankSize, force35Core) ? - kAllToAllBigDataMultiNodeCopyCores : kAllToAllBigDataSingleNodeShards; + kAllToAllBigDataMultiNodeControlShards : kAllToAllBigDataSingleNodeShards; } inline AllToAllChunkPlan PlanAllToAllUdmaChunks(int rankSize, int32_t elementsPerPeer) @@ -153,8 +153,7 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( plan.dataBytes = plan.registeredBytes - plan.controlBytes - plan.signalBytes; if (use35Core) { plan.chunkElements = static_cast(chunkElements); - const size_t chunkBytes = chunkElements * sizeof(int32_t); - plan.chunkBytesPerPeer = std::min(chunkBytes, kAllToAllBigDataMultiNodePeerSlotBytes); + plan.chunkBytesPerPeer = kAllToAllBigDataMultiNodePeerSlotBytes; } else { plan.chunkElements = static_cast( plan.dataBytes / (dataSlotCount * sizeof(int32_t))); @@ -184,12 +183,6 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( } plan.ackSignalOffset = plan.readySignalOffset + controlGroupBytes; plan.controlBytes = controlGroupBytes; - if (use35Core) { - const size_t usedBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; - plan.registeredBytes = - ((usedBytes + kAllToAllBigDataRegistrationAlignment - 1) / - kAllToAllBigDataRegistrationAlignment) * kAllToAllBigDataRegistrationAlignment; - } return plan; } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index e887189d..795e66a3 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -84,7 +84,8 @@ constexpr int kDebugAllReduceScatter = kDebugIpcGather + 1; constexpr int kDebugAllReduceSum = kDebugAllReduceScatter + 1; constexpr int kDebugRecvSlotSampleBase = kDebugUdmaStatusBase + 160; constexpr int kDebugReadySeenBase = kDebugUdmaStatusBase + 208; -constexpr size_t kDebugWords = kDebugReadySeenBase + TileXR::TILEXR_MAX_RANK_SIZE; +constexpr int kDebugAckSeenBase = kDebugReadySeenBase + TileXR::TILEXR_MAX_RANK_SIZE; +constexpr size_t kDebugWords = kDebugAckSeenBase + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int kDefaultCommPort = 10067; constexpr int kDemoBarrierPortOffset = 97; constexpr size_t kUdmaRegistrationAlignment = 2 * 1024 * 1024; @@ -499,7 +500,6 @@ void PrintAllToAllUdmaDebug(int rank, int rankSize, const std::vector& constexpr int recvSlotSampleBase = kDebugUdmaStatusBase + 160; constexpr int remoteDataOffsetBase = kDebugUdmaStatusBase + 176; constexpr int remoteReadyOffsetBase = kDebugUdmaStatusBase + 192; - constexpr int readySeenBase = kDebugUdmaStatusBase + 208; std::cout << "[rank " << rank << "] alltoall udma peer debug:"; for (int peer = 0; peer < rankSize && peer < 16; ++peer) { std::cout << " peer" << peer @@ -515,7 +515,8 @@ void PrintAllToAllUdmaDebug(int rank, int rankSize, const std::vector& << ",slot0=" << debug[recvSlotSampleBase + peer] << ",rDataOff=" << debug[remoteDataOffsetBase + peer] << ",rReadyOff=" << debug[remoteReadyOffsetBase + peer] - << ",ready=" << debug[readySeenBase + peer] + << ",ready=" << debug[kDebugReadySeenBase + peer] + << ",ack=" << debug[kDebugAckSeenBase + peer] << "}"; } std::cout << std::endl; diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 146dd6d1..eb7b81c9 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -38,13 +38,16 @@ constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_REMOTE_READY_OFFSET_BASE = TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 192; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE = TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + 208; +constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ACK_SEEN_BASE = + TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER = - TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE; + TILEXR_UDMA_DEMO_DEBUG_ACK_SEEN_BASE + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER = TILEXR_UDMA_DEMO_DEBUG_IPC_SCATTER + 1; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER = TILEXR_UDMA_DEMO_DEBUG_IPC_GATHER + 1; constexpr int32_t TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SUM = TILEXR_UDMA_DEMO_DEBUG_ALLREDUCE_SCATTER + 1; constexpr uint64_t TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS = 100000000ULL; constexpr uint64_t TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES = TileXR::TILEXR_UDMA_CACHE_LINE_SIZE; +constexpr uint64_t TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES = 128ULL; constexpr int32_t TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS = -1001; constexpr int32_t TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS = -1002; constexpr int32_t TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS = -1003; @@ -60,6 +63,7 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 5U; constexpr int32_t TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE = 8; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES = 16U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES = 16U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS = 32U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE = 16U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE = 17U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE = 18U; @@ -69,7 +73,17 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM = TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT = 0U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT = 1U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY = 2U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY = 3U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY = 4U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR = 11U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_CORE = 20U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_SHARD = + TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_CORE - + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_SOURCE_SHARD = 0U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_FANOUT_SHARD_BASE = 5U; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_PREPARE = 0; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY = 1; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC = 2; @@ -155,7 +169,7 @@ __aicore__ inline __gm__ uint64_t* BigDataControlSlot( ((static_cast(slot) * static_cast(rankSize) + static_cast(peer)) * static_cast(shardCount) + - static_cast(shard)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES); + static_cast(shard)) * TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES); } __aicore__ inline void BigDataCopyInTile( @@ -420,7 +434,7 @@ __aicore__ inline bool BigDataUse35Core(int32_t rankSize, bool force35Core) __aicore__ inline uint32_t BigDataShardCount(int32_t rankSize, bool force35Core = false) { return BigDataUse35Core(rankSize, force35Core) ? - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES : + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS : TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS; } @@ -544,7 +558,7 @@ __aicore__ inline void BigDataStoreTokenMte( { AscendC::LocalTensor tokenLocal = relayLocal.ReinterpretCast(); constexpr uint32_t controlSlotU64 = - TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES / sizeof(uint64_t); + TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES / sizeof(uint64_t); for (uint32_t i = 0; i < controlSlotU64; ++i) { tokenLocal.SetValue(i, token); } @@ -554,7 +568,7 @@ __aicore__ inline void BigDataStoreTokenMte( AscendC::GlobalTensor slotGlobal; slotGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(slot)); AscendC::DataCopyExtParams copyOut { - 1U, static_cast(TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES), 0U, 0U, 0U}; + 1U, static_cast(TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES), 0U, 0U, 0U}; AscendC::DataCopyPad(slotGlobal, relayLocal, copyOut); AscendC::SetFlag(EVENT_ID0); AscendC::WaitFlag(EVENT_ID0); @@ -586,23 +600,6 @@ __aicore__ inline uint64_t BigDataWaitTokenMte( return observed; } -__aicore__ inline __gm__ uint64_t* BigDataRemoteControlSlot( - const __gm__ TileXR::CommArgs* args, int32_t targetRank, uint64_t offset, int32_t slotRank, - uint32_t slot, int32_t rankSize) -{ - auto registry = TileXR::GetUDMARegistry(args); - const uint64_t remoteOffset = - offset + - (static_cast(slot) * static_cast(rankSize) + - static_cast(slotRank)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; - if (!TileXR::UDMARegisteredRangeValid(registry, targetRank, - remoteOffset, TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES)) { - return nullptr; - } - return reinterpret_cast<__gm__ uint64_t*>( - TileXR::UDMARegisteredRemoteAddr(registry, targetRank, remoteOffset)); -} - __aicore__ inline __gm__ uint64_t* BigDataRemoteRegisteredControlSlot( const __gm__ TileXR::CommArgs* args, int32_t targetRank, uint64_t offset, uint32_t slot, int32_t rankSize, uint32_t shardCount, int32_t slotRank, uint32_t shard) @@ -613,15 +610,40 @@ __aicore__ inline __gm__ uint64_t* BigDataRemoteRegisteredControlSlot( ((static_cast(slot) * static_cast(rankSize) + static_cast(slotRank)) * static_cast(shardCount) + - static_cast(shard)) * TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + static_cast(shard)) * TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; if (!TileXR::UDMARegisteredRangeValid(registry, targetRank, - remoteOffset, TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES)) { + remoteOffset, TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES)) { return nullptr; } return reinterpret_cast<__gm__ uint64_t*>( TileXR::UDMARegisteredRemoteAddr(registry, targetRank, remoteOffset)); } +__aicore__ inline uint64_t BigDataRegisteredControlOffset( + uint64_t offset, uint32_t slot, int32_t rankSize, uint32_t shardCount, + int32_t slotRank, uint32_t shard) +{ + return offset + + ((static_cast(slot) * static_cast(rankSize) + + static_cast(slotRank)) * + static_cast(shardCount) + + static_cast(shard)) * TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; +} + +__aicore__ inline uint32_t BigDataPublishAckSignalUdma( + const __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, + uint64_t ackSignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, + int32_t rank, uint64_t token, AscendC::LocalTensor relayLocal) +{ + const uint64_t ackOffset = + BigDataRegisteredControlOffset(ackSignalOffset, slot, rankSize, shardCount, rank, 0U); + auto localAck = reinterpret_cast<__gm__ uint64_t*>(udmaMem + ackOffset); + BigDataStoreTokenMte(localAck, token, relayLocal); + TileXR::UDMAPutSignalNbi( + args, peer, localAck, ackOffset, sizeof(uint64_t), ackOffset, token); + return TileXR::UDMAQuietStatus(args, peer); +} + __aicore__ inline uint64_t BigDataIpcAckOffset(uint32_t slot, int32_t rankSize, int32_t slotRank) { constexpr uint64_t maxAckBytes = @@ -657,7 +679,9 @@ __aicore__ inline void BigDataCopyPeerWorker( uint64_t copyDoneOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { - if (peer < 0 || peer >= rankSize || shardCount == 0U || copyShard >= shardCount) { + const uint32_t dataShardCount = use35Core ? + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES : shardCount; + if (peer < 0 || peer >= rankSize || shardCount == 0U || copyShard >= dataShardCount) { return; } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_PREPARE) { @@ -673,7 +697,7 @@ __aicore__ inline void BigDataCopyPeerWorker( uint32_t shardOffset = 0U; uint32_t shardBytes = 0U; const bool shardHasBytes = - BigDataCopyShardRange(copyShard, shardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes); + BigDataCopyShardRange(copyShard, dataShardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes); const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); @@ -704,6 +728,22 @@ __aicore__ inline void BigDataCopyPeerWorker( return; } } + if (use35Core && profileStage == TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT && + passCount > 0U && globalPass >= static_cast(passCount) && + copyShard == 0U) { + const uint64_t reuseToken = BigDataPassToken(globalPass - static_cast(passCount)); + const uint64_t observed = BigDataLoadTokenMte( + BigDataControlSlot(udmaMem, ackSignalOffset, slot, rankSize, shardCount, peer, 0U), + relayLocal); + if (debug != nullptr && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_ACK_SEEN_BASE + peer] = + static_cast(observed); + if (observed < reuseToken) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + } + } auto src = reinterpret_cast<__gm__ uint8_t*>( input + static_cast(peer) * elementsPerPeer + chunkOffset); @@ -718,13 +758,20 @@ __aicore__ inline void BigDataCopyPeerWorker( token, relayLocal); } +__aicore__ inline bool BigDataWaitCopyReady( + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, uint64_t copyReadyOffset, + uint32_t slot, int32_t rankSize, uint32_t shardCount, int32_t peer, + uint32_t readyShard, uint64_t token, uint32_t loop, uint32_t pass, + AscendC::LocalTensor relayLocal); + __aicore__ inline void BigDataSendPeerWorker( int32_t peer, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, uint64_t sendDataOffset, uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, - uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) + uint64_t chunkBytesPerPeer, bool use35Core, uint64_t copyReadyOffset, + AscendC::LocalTensor relayLocal) { if (peer < 0 || peer >= rankSize || peer == rank || shardCount == 0U || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { @@ -740,19 +787,26 @@ __aicore__ inline void BigDataSendPeerWorker( } const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); - const uint32_t slot = BigDataDataSlot(globalPass, pass, false); + const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); - for (uint32_t copyShard = 0U; copyShard < shardCount; ++copyShard) { - const uint64_t observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard), - token, relayLocal); - if (observed < token) { - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; - } + if (use35Core) { + if (!BigDataWaitCopyReady(udmaMem, debug, copyReadyOffset, slot, rankSize, shardCount, + peer, TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY, token, loop, pass, relayLocal)) { return; } + } else { + for (uint32_t copyShard = 0U; copyShard < shardCount; ++copyShard) { + const uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; + } + return; + } + } } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { return; @@ -773,7 +827,7 @@ __aicore__ inline void BigDataSendPeerWorker( ((static_cast(slot) * static_cast(rankSize) + static_cast(rank)) * static_cast(shardCount)) * - TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; auto registry = TileXR::GetUDMARegistry(args); auto udmaInfo = TileXR::GetUDMAInfo(args); auto wqCtx = TileXR::UDMAGetWQCtx(udmaInfo, peer, 0); @@ -817,8 +871,10 @@ __aicore__ inline void BigDataRecvPeerWorker( uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { + const uint32_t dataShardCount = use35Core ? + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES : shardCount; if (peer < 0 || peer >= rankSize || peer == rank || - shardCount == 0U || recvShard >= shardCount || + shardCount == 0U || recvShard >= dataShardCount || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { return; } @@ -834,24 +890,62 @@ __aicore__ inline void BigDataRecvPeerWorker( uint32_t shardOffset = 0U; uint32_t shardBytes = 0U; const bool shardHasBytes = - BigDataCopyShardRange(recvShard, shardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes); + BigDataCopyShardRange(recvShard, dataShardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes); const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); - uint64_t observed = BigDataWaitTokenMte( - BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, peer, 0U), - token, relayLocal); - if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = - static_cast(observed); + uint64_t observed = 0ULL; + if (use35Core) { + if (recvShard == TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_SHARD) { + observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, + peer, TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_SOURCE_SHARD), + token, relayLocal); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = + static_cast(observed); + if (observed < token) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + } + if (observed < token) { + return; + } + for (uint32_t fanoutShard = 0U; fanoutShard < dataShardCount; ++fanoutShard) { + BigDataStoreTokenMte( + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, + peer, TILEXR_UDMA_DEMO_BIGDATA_LOCAL_FANOUT_SHARD_BASE + fanoutShard), + token, relayLocal); + } + } + observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, + peer, TILEXR_UDMA_DEMO_BIGDATA_LOCAL_FANOUT_SHARD_BASE + recvShard), + token, relayLocal); if (observed < token) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + return; + } + } else { + observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, peer, 0U), + token, relayLocal); + if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + peer] = + static_cast(observed); + if (observed < token) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; + } + } + if (observed < token) { + return; } - } - if (observed < token) { - return; } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY) { return; @@ -880,6 +974,28 @@ __aicore__ inline void BigDataRecvPeerWorker( BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, recvShard), token, relayLocal); if (recvShard + 1U != shardCount) { + if (use35Core && recvShard + 1U == dataShardCount) { + for (uint32_t shard = 0U; shard < dataShardCount; ++shard) { + observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, shard), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + return; + } + } + + const uint32_t ackStatus = BigDataPublishAckSignalUdma( + args, peer, udmaMem, ackSignalOffset, slot, rankSize, shardCount, + rank, token, relayLocal); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + static_cast(ackStatus); + } + } return; } for (uint32_t shard = 0U; shard < shardCount; ++shard) { @@ -897,8 +1013,14 @@ __aicore__ inline void BigDataRecvPeerWorker( __gm__ uint64_t* remoteAck = nullptr; if (use35Core) { - remoteAck = BigDataRemoteRegisteredControlSlot( - args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); + const uint32_t ackStatus = BigDataPublishAckSignalUdma( + args, peer, udmaMem, ackSignalOffset, slot, rankSize, shardCount, + rank, token, relayLocal); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + static_cast(ackStatus); + } + return; } else { remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); } @@ -936,17 +1058,52 @@ __aicore__ inline bool BigDataWaitCopyDoneRange( return true; } +__aicore__ inline bool BigDataPublishCopyReadyRange( + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, uint64_t copyDoneOffset, + uint64_t copyReadyOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, + int32_t peer, uint32_t copyShardBegin, uint32_t copyShardEnd, uint32_t readyShard, + uint64_t token, uint32_t loop, uint32_t pass, AscendC::LocalTensor relayLocal) +{ + if (!BigDataWaitCopyDoneRange(udmaMem, debug, copyDoneOffset, slot, rankSize, + shardCount, peer, copyShardBegin, copyShardEnd, token, loop, pass, relayLocal)) { + return false; + } + BigDataStoreTokenMte( + BigDataControlSlot(udmaMem, copyReadyOffset, slot, rankSize, shardCount, peer, readyShard), + token, relayLocal); + return true; +} + +__aicore__ inline bool BigDataWaitCopyReady( + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, uint64_t copyReadyOffset, + uint32_t slot, int32_t rankSize, uint32_t shardCount, int32_t peer, + uint32_t readyShard, uint64_t token, uint32_t loop, uint32_t pass, + AscendC::LocalTensor relayLocal) +{ + const uint64_t observed = BigDataWaitTokenMte( + BigDataControlSlot(udmaMem, copyReadyOffset, slot, rankSize, shardCount, peer, readyShard), + token, relayLocal); + if (observed < token) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_COPY_TIMEOUT_STATUS; + } + return false; + } + return true; +} + __aicore__ inline bool BigDataRemoteSendSegmentRange( - uint32_t segmentId, uint32_t shardCount, uint32_t chunkElements, + uint32_t segmentId, uint32_t dataShardCount, uint32_t chunkElements, uint32_t& copyShardBegin, uint32_t& copyShardEnd, uint32_t& segmentOffsetBytes, uint32_t& segmentBytes) { - if (shardCount == 0U) { + if (dataShardCount == 0U) { return false; } const uint32_t chunkBytes = chunkElements * static_cast(sizeof(int32_t)); const uint32_t splitOffset = BigDataCopyShardStartBytes( - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END, shardCount, chunkElements); + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END, dataShardCount, chunkElements); if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT) { copyShardBegin = 0U; copyShardEnd = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END; @@ -954,13 +1111,13 @@ __aicore__ inline bool BigDataRemoteSendSegmentRange( segmentBytes = splitOffset; } else if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT) { copyShardBegin = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END; - copyShardEnd = shardCount; + copyShardEnd = dataShardCount; segmentOffsetBytes = splitOffset; segmentBytes = chunkBytes - splitOffset; } else { return false; } - if (copyShardEnd > shardCount || segmentOffsetBytes > chunkBytes) { + if (copyShardEnd > dataShardCount || segmentOffsetBytes > chunkBytes) { return false; } return true; @@ -975,7 +1132,7 @@ __aicore__ inline void BigDataPublishReadySignal( readySignalOffset + ((static_cast(slot) * static_cast(rankSize) + static_cast(rank)) * static_cast(shardCount)) * - TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES; + TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; const uint64_t remoteReadyOffset = localReadyPayloadOffset; auto localSrc = reinterpret_cast<__gm__ uint64_t*>(udmaMem + localReadyPayloadOffset); TileXR::UDMAPutSignalNbi( @@ -1008,19 +1165,27 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, true); + constexpr uint32_t dataShardCount = TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES; + (void)copyDoneOffset; uint32_t copyShardBegin = 0U; uint32_t copyShardEnd = 0U; uint32_t segmentOffsetBytes = 0U; uint32_t segmentBytes = 0U; if (!BigDataRemoteSendSegmentRange( - segmentId, shardCount, chunkBytes / sizeof(int32_t), + segmentId, dataShardCount, chunkBytes / sizeof(int32_t), copyShardBegin, copyShardEnd, segmentOffsetBytes, segmentBytes)) { return; } + (void)copyShardBegin; + (void)copyShardEnd; - if (!BigDataWaitCopyDoneRange(udmaMem, debug, copyDoneOffset, slot, rankSize, shardCount, - peer, copyShardBegin, copyShardEnd, token, loop, pass, relayLocal)) { + const uint32_t readyShard = + (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT) ? + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY : + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY; + if (!BigDataWaitCopyReady(udmaMem, debug, remoteSendDoneOffset, slot, rankSize, + shardCount, peer, readyShard, token, loop, pass, relayLocal)) { return; } if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { @@ -1076,7 +1241,10 @@ __aicore__ inline void BigDataRunSelfCopyShard( int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t pass, uint32_t copyShard, uint32_t shardCount, AscendC::LocalTensor relayLocal) { - if (rank < 0 || rank >= rankSize || shardCount == 0U || copyShard >= shardCount) { + const uint32_t dataShardCount = + (shardCount == TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS) ? + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES : shardCount; + if (rank < 0 || rank >= rankSize || shardCount == 0U || copyShard >= dataShardCount) { return; } int32_t chunkOffset = 0; @@ -1086,7 +1254,7 @@ __aicore__ inline void BigDataRunSelfCopyShard( } uint32_t shardOffset = 0U; uint32_t shardBytes = 0U; - if (!BigDataCopyShardRange(copyShard, shardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { + if (!BigDataCopyShardRange(copyShard, dataShardCount, chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { return; } auto src = reinterpret_cast<__gm__ uint8_t*>( @@ -1105,22 +1273,24 @@ __aicore__ inline void BigDataRunRoleForPeer( uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) { - if (role < static_cast(shardCount)) { + const uint32_t dataShardCount = use35Core ? + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES : shardCount; + if (role < static_cast(dataShardCount)) { BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, use35Core, static_cast(role), shardCount, sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); return; } - if (role == static_cast(shardCount)) { + if (role == static_cast(dataShardCount)) { BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, - chunkBytesPerPeer, relayLocal); + chunkBytesPerPeer, use35Core, copyDoneOffset, relayLocal); return; } const uint32_t recvShard = - static_cast(role) - shardCount - 1U; + static_cast(role) - dataShardCount - 1U; BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, use35Core, recvShard, shardCount, recvDataOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, @@ -2328,6 +2498,37 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( passCount, loop, pass, kernelLoopBase, profileStage, use35Core, copyShard, shardCount, sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); + if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, true); + if (!isLocalPeer && copyShard == 0U) { + (void)BigDataPublishCopyReadyRange( + udmaMem, debug, copyDoneOffset, remoteSendDoneOffset, + slot, rankSize, shardCount, peer, + 0U, TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END, + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY, + token, loop, pass, relayLocal); + } + if (!isLocalPeer && + copyShard == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR) { + (void)BigDataPublishCopyReadyRange( + udmaMem, debug, copyDoneOffset, remoteSendDoneOffset, + slot, rankSize, shardCount, peer, + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END, + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES, + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY, + token, loop, pass, relayLocal); + } + if (isLocalPeer && copyShard == 0U) { + (void)BigDataPublishCopyReadyRange( + udmaMem, debug, copyDoneOffset, remoteSendDoneOffset, + slot, rankSize, shardCount, peer, + 0U, TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES, + TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY, + token, loop, pass, relayLocal); + } + } } if (!isLocalPeer && isRemoteSendPrimaryCore) { BigDataRemoteSendSegmentWorker(peer, @@ -2351,7 +2552,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, - copyDoneOffset, readySignalOffset, chunkBytesPerPeer, relayLocal); + copyDoneOffset, readySignalOffset, chunkBytesPerPeer, true, + remoteSendDoneOffset, relayLocal); } if (isRecvCore) { BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index b09d973e..f6099619 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -143,13 +143,13 @@ void TestAllToAllBigDataPlan() CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 128ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, 1024ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 8ULL * 1024ULL * 1024ULL); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataRegistrationAlignment, 2ULL * 1024ULL * 1024ULL); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 64ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 128ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 5U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataSingleNodeShards, 2U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataLocalCopyShards, 2U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeCopyCores, 16U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRecvCores, 16U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeControlShards, 32U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRemoteSendPrimaryCore, 16U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRemoteSendSecondaryCore, 17U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeLocalSendCore, 18U); @@ -187,11 +187,11 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() const size_t controlGroupBytes = static_cast(plan.passCount) * static_cast(rankSize) * - static_cast(TileXR::Demo::kAllToAllBigDataMultiNodeCopyCores) * + static_cast(TileXR::Demo::kAllToAllBigDataMultiNodeControlShards) * TileXR::Demo::kAllToAllBigDataControlSlotBytes; - CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize), 16U); - CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize), 32U); + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); CHECK_EQ(plan.passCount, 1U); CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); CHECK_EQ(plan.controlBytes, controlGroupBytes); @@ -207,26 +207,6 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); } -void TestAllToAllBigDataMultiNodeSmallPayloadShrinksRegisteredBytes() -{ - constexpr int rankSize = 16; - constexpr int32_t elementsPerPeer = 262144; // 1 MiB per peer for int32_t. - const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); - const size_t expectedChunkBytes = static_cast(elementsPerPeer) * sizeof(int32_t); - const size_t expectedDataBytes = - static_cast(rankSize - 1) * static_cast(plan.passCount) * 2ULL * expectedChunkBytes; - const size_t usedBytes = plan.dataBytes + plan.controlBytes + plan.signalBytes; - - CHECK_EQ(plan.passCount, 1U); - CHECK_EQ(plan.chunkBytesPerPeer, expectedChunkBytes); - CHECK_EQ(plan.dataBytes, expectedDataBytes); - CHECK_EQ(plan.registeredBytes, - ((usedBytes + TileXR::Demo::kAllToAllBigDataRegistrationAlignment - 1) / - TileXR::Demo::kAllToAllBigDataRegistrationAlignment) * - TileXR::Demo::kAllToAllBigDataRegistrationAlignment); - CHECK_EQ(plan.registeredBytes < TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, true); -} - void TestAllToAllBigDataForce35CorePlanFor8P() { constexpr int rankSize = 8; @@ -235,14 +215,14 @@ void TestAllToAllBigDataForce35CorePlanFor8P() const size_t controlGroupBytes = static_cast(plan.passCount) * static_cast(rankSize) * - static_cast(TileXR::Demo::kAllToAllBigDataMultiNodeCopyCores) * + static_cast(TileXR::Demo::kAllToAllBigDataMultiNodeControlShards) * TileXR::Demo::kAllToAllBigDataControlSlotBytes; CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(rankSize), false); CHECK_EQ(TileXR::Demo::AllToAllBigDataUse35Core(rankSize, true), true); - CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize, true), 16U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize, true), 32U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(rankSize, true), 35U); - CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, true); + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); CHECK_EQ(plan.controlBytes, controlGroupBytes); CHECK_EQ(plan.signalBytes, 4ULL * controlGroupBytes); @@ -453,17 +433,24 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataLoadTokenMte"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_CONTROL_SLOT_BYTES"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES = 128ULL"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_CORES_PER_PEER = 5U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 8ULL * 1024ULL * 1024ULL"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE = 8"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS = 32U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM ="); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE = 16U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE = 17U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE = 18U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE = 19U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY = 2U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY = 3U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY = 4U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_CORE = 20U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_FANOUT_SHARD_BASE = 5U"); CHECK_CONTAINS(kernel, "BigDataIsMultiNode(rankSize)"); CHECK_CONTAINS(kernel, "BigDataValidTopology(rankSize)"); CHECK_CONTAINS(kernel, "BigDataUse35Core(rankSize, force35Core)"); @@ -477,6 +464,8 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentWorker"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentRange"); CHECK_CONTAINS(kernel, "BigDataWaitCopyDoneRange"); + CHECK_CONTAINS(kernel, "BigDataPublishCopyReadyRange"); + CHECK_CONTAINS(kernel, "BigDataWaitCopyReady"); CHECK_CONTAINS(kernel, "BigDataPublishReadySignal"); CHECK_CONTAINS(kernel, "remoteSendDoneOffset"); CHECK_CONTAINS(kernel, "if (!BigDataIsMultiNode(rankSize))"); @@ -502,9 +491,11 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "copyShardEnd = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END"); CHECK_CONTAINS(kernel, "copyShardBegin = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END"); CHECK_CONTAINS(kernel, "copyShardEnd = shardCount"); + CHECK_CONTAINS(kernel, "copyShard == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR"); CHECK_CONTAINS(kernel, "BigDataDataSlot(globalPass, pass, use35Core)"); CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex"); CHECK_CONTAINS(kernel, "BigDataStoreTokenMte"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyReadyOffset, slot, rankSize, shardCount, peer, readyShard)"); CHECK_CONTAINS(kernel, "BigDataIpcAckOffset"); CHECK_CONTAINS(kernel, "BigDataLocalIpcAckSlot"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); @@ -513,6 +504,8 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, copyDoneOffset, slot, rankSize, shardCount, peer, copyShard)"); CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, recvShard)"); CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, readySignalOffset, slot, rankSize, shardCount, peer, 0U)"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_SOURCE_SHARD"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_FANOUT_SHARD_BASE + recvShard"); CHECK_CONTAINS(kernel, "remoteDataOffset ="); CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex(rank, peer)"); CHECK_CONTAINS(kernel, "BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount"); @@ -524,9 +517,11 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES"); CHECK_CONTAINS(kernel, "relayLocal[bufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]"); CHECK_CONTAINS(kernel, "if (!use35Core && profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT"); + CHECK_CONTAINS(kernel, "if (use35Core && profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT"); + CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, ackSignalOffset, slot, rankSize, shardCount, peer, 0U)"); CHECK_CONTAINS(kernel, "use35Core ? passCount : TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS"); CHECK_CONTAINS(kernel, "recvSlotInt[0]"); - CHECK_CONTAINS(kernel, "BigDataStoreTokenMte(remoteAck, token, relayLocal)"); + CHECK_CONTAINS(kernel, "BigDataPublishAckSignalUdma"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize)"); CHECK_CONTAINS(kernel, "ackSignal"); CHECK_NOT_CONTAINS(kernel, "UDMAPutNbi"); @@ -544,7 +539,6 @@ int main() TestAllToAllMaxRank256With64MiBPerRank(); TestAllToAllBigDataPlan(); TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone(); - TestAllToAllBigDataMultiNodeSmallPayloadShrinksRegisteredBytes(); TestAllToAllBigDataForce35CorePlanFor8P(); TestAllToAllBigDataBlockDim(); TestAllToAllBigDataMultiNodeTopology(); diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 3a51eac8..3adb0979 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -208,7 +208,7 @@ void TestMemoryRegistrationUsesOfficialUbFlags() const std::string transport = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.cacheable = 0"); - CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.nonPin = 0"); + CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.nonPin = 1"); CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.userIova = 0"); CHECK_CONTAINS(transport, "mrInfo.in.ub.flags.bs.tokenIdValid = 1"); } From c3c135362568d140d06a18cd21b3bcaed770dfd5 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 00:23:14 +0800 Subject: [PATCH 022/163] test(udma): align alltoall source guards after sync tuning - Update source-guard expectations for the rebased multinode bigdata sync tuning. - Keep the followup 128-byte control slot and 32 control shard layout unchanged. --- tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index f6099619..c6f69416 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -309,7 +309,7 @@ void TestDemoDebugLayoutSource() CHECK_CONTAINS(demo, "kDebugUdmaStatusBase + TileXR::TILEXR_MAX_RANK_SIZE"); CHECK_CONTAINS(demo, "kDebugIpcGather + 1"); CHECK_CONTAINS(demo, "kDebugReadySeenBase + TileXR::TILEXR_MAX_RANK_SIZE"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + TileXR::TILEXR_MAX_RANK_SIZE"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_READY_SEEN_BASE + TileXR::TILEXR_MAX_RANK_SIZE"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer"); } @@ -490,7 +490,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "copyShardBegin = 0U"); CHECK_CONTAINS(kernel, "copyShardEnd = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END"); CHECK_CONTAINS(kernel, "copyShardBegin = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END"); - CHECK_CONTAINS(kernel, "copyShardEnd = shardCount"); + CHECK_CONTAINS(kernel, "copyShardEnd = dataShardCount"); CHECK_CONTAINS(kernel, "copyShard == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR"); CHECK_CONTAINS(kernel, "BigDataDataSlot(globalPass, pass, use35Core)"); CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex"); @@ -517,7 +517,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES"); CHECK_CONTAINS(kernel, "relayLocal[bufferId * TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES]"); CHECK_CONTAINS(kernel, "if (!use35Core && profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT"); - CHECK_CONTAINS(kernel, "if (use35Core && profileStage > TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT"); + CHECK_CONTAINS(kernel, "if (use35Core && profileStage == TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT"); CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, ackSignalOffset, slot, rankSize, shardCount, peer, 0U)"); CHECK_CONTAINS(kernel, "use35Core ? passCount : TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS"); CHECK_CONTAINS(kernel, "recvSlotInt[0]"); From 53ac4a919f5c8e21962dba3b8a375a377a17286c Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 09:53:51 +0800 Subject: [PATCH 023/163] perf(udma): checkpoint 64p remote put tail flag scheduling Checkpoint the current remote-put-only experiment with 64-core 32+32 send scheduling and per-segment tail flag completion checks before returning to the previously validated 16p mteack path. --- tests/udma/demo/tilexr_udma_alltoall_layout.h | 11 +- tests/udma/demo/tilexr_udma_demo.cpp | 34 +- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 370 +++++++++++++++++- .../unit/test_tilexr_udma_alltoall_layout.cpp | 119 +++++- 4 files changed, 491 insertions(+), 43 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index c7492ef5..ee35bfca 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -27,7 +27,7 @@ struct AllToAllChunkPlan { constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataMultiNodeRegisteredBytes = 1024ULL * 1024ULL * 1024ULL; -constexpr size_t kAllToAllBigDataMultiNodePeerSlotBytes = 8ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataMultiNodePeerSlotBytes = 16ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataControlSlotBytes = 128ULL; constexpr uint32_t kAllToAllBigDataCoresPerPeer = 5U; constexpr uint32_t kAllToAllBigDataSingleNodeShards = 2U; @@ -43,6 +43,7 @@ constexpr uint32_t kAllToAllBigDataMultiNodeLocalSendCore = 18U; constexpr uint32_t kAllToAllBigDataMultiNodeRecvCoreBase = 19U; constexpr uint32_t kAllToAllBigDataMultiNodeBlockDim = kAllToAllBigDataMultiNodeRecvCoreBase + kAllToAllBigDataMultiNodeRecvCores; +constexpr uint32_t kAllToAllBigDataRemotePutOnlyBlockDim = 64U; struct AllToAllBigDataPlan { uint32_t passCount = 1; @@ -153,7 +154,7 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( plan.dataBytes = plan.registeredBytes - plan.controlBytes - plan.signalBytes; if (use35Core) { plan.chunkElements = static_cast(chunkElements); - plan.chunkBytesPerPeer = kAllToAllBigDataMultiNodePeerSlotBytes; + plan.chunkBytesPerPeer = static_cast(plan.chunkElements) * sizeof(int32_t); } else { plan.chunkElements = static_cast( plan.dataBytes / (dataSlotCount * sizeof(int32_t))); @@ -284,12 +285,16 @@ inline std::vector AllToAllBigDataMergedPeerTasks(int rank, int rankSiz return tasks; } -inline uint32_t AllToAllBigDataBlockDim(int rankSize, bool force35Core = false) +inline uint32_t AllToAllBigDataBlockDim( + int rankSize, bool force35Core = false, bool remotePutOnly = false) { if (rankSize <= 0) { return 1U; } if (AllToAllBigDataUse35Core(rankSize, force35Core)) { + if (remotePutOnly) { + return kAllToAllBigDataRemotePutOnlyBlockDim; + } return kAllToAllBigDataMultiNodeBlockDim; } return static_cast(rankSize) * kAllToAllBigDataCoresPerPeer; diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 795e66a3..f19f1f44 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -88,7 +88,6 @@ constexpr int kDebugAckSeenBase = kDebugReadySeenBase + TileXR::TILEXR_MAX_RANK_ constexpr size_t kDebugWords = kDebugAckSeenBase + TileXR::TILEXR_MAX_RANK_SIZE; constexpr int kDefaultCommPort = 10067; constexpr int kDemoBarrierPortOffset = 97; -constexpr size_t kUdmaRegistrationAlignment = 2 * 1024 * 1024; constexpr int kConnectRetryCount = 500; constexpr int kConnectRetrySleepMs = 10; constexpr int kBigDataProfileStageFull = 8; @@ -676,6 +675,8 @@ int main(int argc, char** argv) bigDataProfileStage = std::max(0, std::min(bigDataProfileStage, kBigDataProfileStageFull)); const bool forceBigData35Core = testType == 7 && GetEnvInt("TILEXR_DEMO_BIGDATA_FORCE_35CORE", 0) != 0; + const bool bigDataRemotePutOnly = + testType == 7 && GetEnvInt("TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY", 0) != 0; bool bigDataProfilePartial = testType == 7 && bigDataProfileStage < kBigDataProfileStageFull; bool syncAllToAllAtEnd = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_SYNC_AT_END", 0) != 0; @@ -740,8 +741,9 @@ int main(int argc, char** argv) PrintStatus(rank, "bigdata multinode mode=" + std::string(TileXR::Demo::AllToAllBigDataIsMultiNode(rankSize) ? "true" : "false") + " force35Core=" + std::string(forceBigData35Core ? "true" : "false") + + " remotePutOnly=" + std::string(bigDataRemotePutOnly ? "true" : "false") + " blockDim=" + std::to_string(TileXR::Demo::AllToAllBigDataBlockDim( - rankSize, forceBigData35Core)) + + rankSize, forceBigData35Core, bigDataRemotePutOnly)) + " shards=" + std::to_string(TileXR::Demo::AllToAllBigDataShardCount( rankSize, forceBigData35Core))); } @@ -770,15 +772,12 @@ int main(int argc, char** argv) size_t signalOffset = testType == 7 ? bigDataPlan.readySignalOffset : (hasOutput ? (outputOffset + activeBytesPerRank) : activeBytesPerRank); size_t payloadBytes = testType == 7 ? bigDataPlan.registeredBytes : (signalOffset + signalBytes); - size_t allocBytes = ((payloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * - kUdmaRegistrationAlignment; + size_t allocBytes = payloadBytes; size_t registeredBytes = allocBytes; if (isAllToAll && testType == 7) { - registeredBytes = ((bigDataPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / - kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + registeredBytes = bigDataPlan.registeredBytes; } else if (isAllToAll && (strictAllToAllUdma || testType == 6)) { - registeredBytes = ((chunkPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / - kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + registeredBytes = chunkPlan.registeredBytes; } if (allocBytes < registeredBytes) { allocBytes = registeredBytes; @@ -954,7 +953,9 @@ int main(int argc, char** argv) } const uint32_t bigDataBlockDim = TileXR::Demo::AllToAllBigDataBlockDim( - rankSize, forceBigData35Core); + rankSize, forceBigData35Core, bigDataRemotePutOnly); + const uint32_t bigDataModeFlags = + (forceBigData35Core ? 1U : 0U) | (bigDataRemotePutOnly ? 2U : 0U); auto a2aStart = std::chrono::steady_clock::now(); for (int iter = 0; iter < allToAllRepeat; ++iter) { const uint64_t kernelLoopBase = static_cast(iter); @@ -968,7 +969,7 @@ int main(int argc, char** argv) bigDataPlan.readySignalOffset, bigDataPlan.ackSignalOffset, bigDataPlan.chunkElements, bigDataPlan.passCount, 1, kernelLoopBase, static_cast(bigDataProfileStage), - forceBigData35Core ? 1U : 0U); + bigDataModeFlags); } if (!CheckAcl(rank, "aclrtSynchronizeStream post-alltoall", aclrtSynchronizeStream(stream))) { if (udmaRegistered) { @@ -992,7 +993,7 @@ int main(int argc, char** argv) << " us payload=" << payload << " bytes bw=" << bwGbs << " GB/s" << std::endl; bool bigDataCopyBackOk = true; - if (!bigDataProfilePartial) { + if (!bigDataProfilePartial && !bigDataRemotePutOnly) { bigDataCopyBackOk = CopyDeviceToHost(rank, hostOutput.data(), dataCount * sizeof(int32_t), bigOutput, dataCount * sizeof(int32_t), "bigdata alltoall output"); } @@ -1036,8 +1037,7 @@ int main(int argc, char** argv) // Registered relay region: [udmaMem chunk | signals]. const size_t fusedChunkBytes = chunkPlan.chunkBytesPerRank; const size_t fusedSignalBytes = static_cast(rankSize) * sizeof(uint64_t); - const size_t fusedRegBytes = ((chunkPlan.registeredBytes + kUdmaRegistrationAlignment - 1) / - kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + const size_t fusedRegBytes = chunkPlan.registeredBytes; if (!udmaRegistered) { int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), fusedRegBytes, &udmaHandle); if (registerRet != TileXR::TILEXR_SUCCESS) { @@ -1303,7 +1303,8 @@ int main(int argc, char** argv) } bool usedIpcFallback = false; - bool allToAllUdmaComplete = !isAllToAll || bigDataProfilePartial || AllToAllUdmaComplete(rankSize, hostDebug); + bool allToAllUdmaComplete = + !isAllToAll || bigDataRemotePutOnly || bigDataProfilePartial || AllToAllUdmaComplete(rankSize, hostDebug); if (isAllToAll && strictAllToAllUdma && !allToAllUdmaComplete) { std::cerr << "[rank " << rank << "] ERROR: strict alltoall UDMA CQ incomplete:"; for (int peer = 0; peer < rankSize; ++peer) { @@ -1505,7 +1506,10 @@ int main(int argc, char** argv) } bool ok = false; - if (bigDataProfilePartial) { + if (bigDataRemotePutOnly) { + PrintStatus(rank, "skip result validation for bigdata remote-put-only profile"); + ok = true; + } else if (bigDataProfilePartial) { PrintStatus(rank, "skip result validation for bigdata profile stage=" + std::to_string(bigDataProfileStage)); ok = true; diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index eb7b81c9..3126a792 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -55,7 +55,7 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES = 64 * 1024; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES = TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES * 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U; -constexpr uint64_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 8ULL * 1024ULL * 1024ULL; +constexpr uint64_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 16ULL * 1024ULL * 1024ULL; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS; @@ -71,6 +71,8 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE = 19U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM = TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_BLOCK_DIM = 64U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES = 32U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT = 0U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT = 1U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY = 2U; @@ -93,6 +95,14 @@ constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY = 5; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT = 6; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK = 7; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_FULL = 8; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK = 0; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP = 1; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER = 2; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP = 3; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_SEGMENT = 4; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_ADDRESS = 5; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST = 6; +constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_ACK = 7; namespace { @@ -213,6 +223,22 @@ __aicore__ inline void BigDataCopyOneRelay( AscendC::WaitFlag(EVENT_ID0); } +__aicore__ inline void BigDataStoreInt32Mte( + __gm__ int32_t* dst, int32_t value, AscendC::LocalTensor relayLocal) +{ + AscendC::LocalTensor fillLocal = relayLocal.ReinterpretCast(); + fillLocal.SetValue(0, value); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(dst)); + AscendC::DataCopyExtParams copyOut {1U, static_cast(sizeof(int32_t)), 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, relayLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + __aicore__ inline void BigDataWaitMte2ToMte3(uint32_t bufferId) { if (bufferId == 0U) { @@ -499,6 +525,84 @@ __aicore__ inline int32_t BigDataRemotePeerAt(int32_t rank, int32_t rankSize, in return -1; } +__aicore__ inline int32_t BigDataRemotePeerForwardAt(int32_t rank, int32_t rankSize, int32_t remoteIndex) +{ + if (!BigDataIsMultiNode(rankSize) || !BigDataValidTopology(rankSize) || remoteIndex < 0) { + return -1; + } + int32_t remoteCount = 0; + for (int32_t step = 0; step < rankSize; ++step) { + const int32_t peer = + (rank + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE + step) % rankSize; + if (!BigDataIsLocalPeer(rank, peer)) { + if (remoteCount == remoteIndex) { + return peer; + } + ++remoteCount; + } + } + return -1; +} + +__aicore__ inline int32_t BigDataRemotePeerReverseAt(int32_t rank, int32_t rankSize, int32_t remoteIndex) +{ + if (!BigDataIsMultiNode(rankSize) || !BigDataValidTopology(rankSize) || remoteIndex < 0) { + return -1; + } + int32_t remoteCount = 0; + for (int32_t step = 0; step < rankSize; ++step) { + const int32_t peer = + (rank - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - step + rankSize * 2) % rankSize; + if (!BigDataIsLocalPeer(rank, peer)) { + if (remoteCount == remoteIndex) { + return peer; + } + ++remoteCount; + } + } + return -1; +} + +__aicore__ inline uint32_t BigDataSelectWeightedQp( + const __gm__ TileXR::CommArgs* args, int32_t peer, bool selectMax) +{ + auto udmaInfo = TileXR::GetUDMAInfo(args); + const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; + uint32_t selected = 0U; + uint32_t selectedWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); + for (uint32_t qpIdx = 1U; qpIdx < qpCount; ++qpIdx) { + const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qpIdx); + if ((selectMax && weight > selectedWeight) || (!selectMax && weight < selectedWeight)) { + selected = qpIdx; + selectedWeight = weight; + } + } + return selected; +} + +__aicore__ inline uint32_t BigDataSelectDistinctWeightedQp( + const __gm__ TileXR::CommArgs* args, int32_t peer, uint32_t avoidQp, bool selectMax) +{ + auto udmaInfo = TileXR::GetUDMAInfo(args); + const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; + if (qpCount <= 1U) { + return 0U; + } + uint32_t selected = avoidQp == 0U ? 1U : 0U; + uint32_t selectedWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, selected); + for (uint32_t qpIdx = 0U; qpIdx < qpCount; ++qpIdx) { + if (qpIdx == avoidQp) { + continue; + } + const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qpIdx); + if ((selectMax && weight > selectedWeight) || (!selectMax && weight < selectedWeight)) { + selected = qpIdx; + selectedWeight = weight; + } + } + return selected; +} + __aicore__ inline int32_t BigDataLocalPeerAt(int32_t rank, int32_t localIndex) { if (localIndex < 0 || localIndex >= TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - 1) { @@ -630,20 +734,6 @@ __aicore__ inline uint64_t BigDataRegisteredControlOffset( static_cast(shard)) * TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; } -__aicore__ inline uint32_t BigDataPublishAckSignalUdma( - const __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, - uint64_t ackSignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, - int32_t rank, uint64_t token, AscendC::LocalTensor relayLocal) -{ - const uint64_t ackOffset = - BigDataRegisteredControlOffset(ackSignalOffset, slot, rankSize, shardCount, rank, 0U); - auto localAck = reinterpret_cast<__gm__ uint64_t*>(udmaMem + ackOffset); - BigDataStoreTokenMte(localAck, token, relayLocal); - TileXR::UDMAPutSignalNbi( - args, peer, localAck, ackOffset, sizeof(uint64_t), ackOffset, token); - return TileXR::UDMAQuietStatus(args, peer); -} - __aicore__ inline uint64_t BigDataIpcAckOffset(uint32_t slot, int32_t rankSize, int32_t slotRank) { constexpr uint64_t maxAckBytes = @@ -670,6 +760,32 @@ __aicore__ inline __gm__ uint64_t* BigDataRemoteIpcAckSlot( args->peerMems[targetRank] + BigDataIpcAckOffset(slot, rankSize, rank)); } +__aicore__ inline uint32_t BigDataPublishAckSignalUdma( + const __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, + uint64_t ackSignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, + int32_t rank, uint64_t token, AscendC::LocalTensor relayLocal) +{ + const uint64_t ackOffset = + BigDataRegisteredControlOffset(ackSignalOffset, slot, rankSize, shardCount, rank, 0U); + auto localAck = reinterpret_cast<__gm__ uint64_t*>(udmaMem + ackOffset); + BigDataStoreTokenMte(localAck, token, relayLocal); + TileXR::UDMAPutSignalNbi( + args, peer, localAck, ackOffset, sizeof(uint64_t), ackOffset, token); + return TileXR::UDMAQuietStatus(args, peer); +} + +__aicore__ inline void BigDataRemotePutOnlyPublishAck( + __gm__ uint64_t* remoteAck, uint64_t token, AscendC::LocalTensor relayLocal) +{ + BigDataStoreTokenMte(remoteAck, token, relayLocal); +} + +__aicore__ inline uint64_t BigDataRemotePutOnlyWaitPeerAck( + __gm__ uint64_t* peerAck, uint64_t token, AscendC::LocalTensor relayLocal) +{ + return BigDataWaitTokenMte(peerAck, token, relayLocal); +} + __aicore__ inline void BigDataCopyPeerWorker( int32_t peer, int32_t rank, int32_t rankSize, const __gm__ TileXR::CommArgs* args, __gm__ int32_t* input, __gm__ int32_t* output, @@ -1162,6 +1278,7 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( return; } (void)chunkOffset; + (void)chunkBytes; const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, true); @@ -1236,6 +1353,167 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( args, peer, udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, token); } +__aicore__ inline void BigDataRemotePutOnlySendWorker( + int32_t peer, uint32_t segmentId, uint32_t qpIdx, int32_t rank, int32_t rankSize, + __gm__ TileXR::CommArgs* args, __gm__ int32_t* input, + __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, + uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, + uint32_t profileStage, uint32_t shardCount, uint64_t recvDataOffset, + uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) +{ + if (peer < 0 || peer >= rankSize || peer == rank || BigDataIsLocalPeer(rank, peer) || + !BigDataIsMultiNode(rankSize) || shardCount == 0U || + profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { + return; + } + + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + return; + } + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, true); + constexpr uint32_t dataShardCount = TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES; + + uint32_t copyShardBegin = 0U; + uint32_t copyShardEnd = 0U; + uint32_t segmentOffsetBytes = 0U; + uint32_t segmentBytes = 0U; + if (!BigDataRemoteSendSegmentRange( + segmentId, dataShardCount, chunkBytes / sizeof(int32_t), + copyShardBegin, copyShardEnd, segmentOffsetBytes, segmentBytes)) { + return; + } + (void)copyShardBegin; + (void)copyShardEnd; + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_SEGMENT) { + return; + } + + auto localSrc = reinterpret_cast<__gm__ int32_t*>( + reinterpret_cast<__gm__ uint8_t*>( + input + static_cast(peer) * elementsPerPeer + chunkOffset) + + segmentOffsetBytes); + const uint64_t remoteDataOffset = + recvDataOffset + + (static_cast(slot) * static_cast(rankSize > 1 ? rankSize - 1 : 1) + + static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer + + static_cast(segmentOffsetBytes); + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_ADDRESS) { + return; + } + + uint32_t status = 0U; + if (segmentBytes > 0U) { + const int32_t tokenValue = + static_cast(BigDataPassToken(globalPass)); + auto localTail = reinterpret_cast<__gm__ int32_t*>( + reinterpret_cast<__gm__ uint8_t*>(localSrc) + + static_cast(segmentBytes) - sizeof(int32_t)); + BigDataStoreInt32Mte(localTail, tokenValue, relayLocal); + TileXR::UDMAPutNbiOnQp(args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); + } + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } +} + +__aicore__ inline int32_t BigDataRemotePutOnlyCheckIndex(int32_t blockIdx) +{ + if (blockIdx < 0) { + return -1; + } + return blockIdx; +} + +__aicore__ inline int32_t BigDataRemotePutOnlySendTaskCount(int32_t remoteTaskCount) +{ + return remoteTaskCount > 0 ? remoteTaskCount * 2 : 0; +} + +__aicore__ inline int32_t BigDataRemotePutOnlySendTaskRemoteIndex( + int32_t sendTask, int32_t remoteTaskCount) +{ + if (sendTask < 0 || remoteTaskCount <= 0 || + sendTask >= BigDataRemotePutOnlySendTaskCount(remoteTaskCount)) { + return -1; + } + return sendTask < remoteTaskCount ? sendTask : sendTask - remoteTaskCount; +} + +__aicore__ inline uint32_t BigDataRemotePutOnlySendTaskSegment( + int32_t sendTask, int32_t remoteTaskCount) +{ + return sendTask < remoteTaskCount ? + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT : + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; +} + +__aicore__ inline void BigDataRemotePutOnlyCheckWorker( + int32_t remoteIndex, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, + int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, + uint64_t kernelLoopBase, uint64_t recvDataOffset, uint64_t ackSignalOffset, + uint64_t chunkBytesPerPeer, uint32_t shardCount, AscendC::LocalTensor relayLocal) +{ + const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex); + if (peer < 0 || peer >= rankSize || peer == rank || BigDataIsLocalPeer(rank, peer) || + shardCount == 0U) { + return; + } + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes) || + chunkBytes < sizeof(int32_t)) { + return; + } + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, true); + const uint64_t remoteDataOffset = + recvDataOffset + + (static_cast(slot) * static_cast(rankSize > 1 ? rankSize - 1 : 1) + + static_cast(BigDataNetworkPeerIndex(peer, rank))) * chunkBytesPerPeer; + auto tail = reinterpret_cast<__gm__ int32_t*>( + udmaMem + remoteDataOffset + static_cast(chunkBytes) - sizeof(int32_t)); + + int32_t observed = 0; + uint64_t polls = 0; + do { + AscendC::GlobalTensor tailGlobal; + tailGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(tail)); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, static_cast(sizeof(int32_t)), 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, tailGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + observed = relayLocal.ReinterpretCast().GetValue(0); + ++polls; + } while (static_cast(observed) < token && polls < TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS); + + uint32_t status = 0U; + if (static_cast(observed) < token) { + status = static_cast(TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS); + } else { + auto remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); + auto localAck = BigDataLocalIpcAckSlot(args, rank, slot, rankSize, peer); + BigDataRemotePutOnlyPublishAck(remoteAck, token, relayLocal); + const uint64_t ackObserved = BigDataRemotePutOnlyWaitPeerAck( + localAck, token, relayLocal); + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_ACK_SEEN_BASE + peer] = + static_cast(ackObserved); + } + if (ackObserved < token) { + status = static_cast(TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS); + } + } + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); + } +} + __aicore__ inline void BigDataRunSelfCopyShard( int32_t rank, int32_t rankSize, __gm__ int32_t* input, __gm__ int32_t* output, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t pass, @@ -2398,7 +2676,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( int32_t rankSize = args->rankSize; bool enabled = TileXR::UDMARegistryEnabled(args); const int32_t blockIdx = AscendC::GetBlockIdx(); - const bool force35Core = force35CoreFlag != 0U; + const bool force35Core = (force35CoreFlag & 0x1U) != 0U; + const bool remotePutOnly = (force35CoreFlag & 0x2U) != 0U; if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; @@ -2457,7 +2736,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( return; } - if (blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM)) { + const uint32_t activeBlockDim = remotePutOnly ? + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_BLOCK_DIM : + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM; + if (blockIdx >= static_cast(activeBlockDim)) { BigDataKernelExitBarrier(); return; } @@ -2471,14 +2753,66 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE); const bool isRecvCore = blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE) && - blockIdx < static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM); + blockIdx < static_cast(activeBlockDim); const uint32_t copyShard = static_cast(blockIdx); const uint32_t recvShard = static_cast(blockIdx) - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; const int32_t taskCount = BigDataTaskCount(rankSize, force35Core); + const int32_t remoteTaskCount = rankSize - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + const int32_t remotePutOnlyCheckIndex = BigDataRemotePutOnlyCheckIndex(blockIdx); + const bool isRemotePutOnlyCheckCore = + remotePutOnlyCheckIndex >= 0 && remotePutOnlyCheckIndex < remoteTaskCount; + + if (remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK) { + BigDataKernelExitBarrier(); + return; + } for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { + if (remotePutOnly) { + const uint32_t sendGroupCore = + static_cast(blockIdx) % + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES; + const uint32_t segmentId = blockIdx < static_cast( + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES) ? + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT : + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; + for (int32_t remoteIndex = static_cast(sendGroupCore); + remoteIndex < remoteTaskCount; + remoteIndex += static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES)) { + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { + continue; + } + const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex); + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { + continue; + } + const uint32_t primaryQp = BigDataSelectWeightedQp(args, peer, true); + const uint32_t qpIdx = + segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? + primaryQp : BigDataSelectDistinctWeightedQp(args, peer, primaryQp, false); + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { + continue; + } + BigDataRemotePutOnlySendWorker(peer, segmentId, qpIdx, + rank, rankSize, args, input, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, shardCount, recvDataOffset, + chunkBytesPerPeer, relayLocal); + } + BigDataKernelExitBarrier(); + if (profileStage > TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST && + isRemotePutOnlyCheckCore) { + BigDataRemotePutOnlyCheckWorker( + remotePutOnlyCheckIndex, rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, recvDataOffset, ackSignalOffset, chunkBytesPerPeer, + shardCount, relayLocal); + } + BigDataKernelExitBarrier(); + continue; + } if (isCopyCore) { BigDataRunSelfCopyShard(rank, rankSize, input, output, elementsPerPeer, effectiveChunkElements, pass, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index c6f69416..b9394e1c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -51,6 +51,19 @@ std::string ReadFile(const std::string& path) return out.str(); } +std::string SliceBetween(const std::string& text, const std::string& begin, const std::string& end) +{ + const size_t beginPos = text.find(begin); + if (beginPos == std::string::npos) { + return ""; + } + const size_t endPos = text.find(end, beginPos + begin.size()); + if (endPos == std::string::npos) { + return text.substr(beginPos); + } + return text.substr(beginPos, endPos - beginPos); +} + void TestAllToAllInputPattern() { constexpr int rank = 2; @@ -142,7 +155,7 @@ void TestAllToAllBigDataPlan() CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 128ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, 1024ULL * 1024ULL * 1024ULL); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 8ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 16ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 128ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 5U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataSingleNodeShards, 2U); @@ -155,6 +168,7 @@ void TestAllToAllBigDataPlan() CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeLocalSendCore, 18U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRecvCoreBase, 19U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeBlockDim, 35U); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataRemotePutOnlyBlockDim, 64U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataPingPongSlots, 2U); CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, true); const size_t controlGroupBytes = @@ -193,7 +207,7 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize), 32U); CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); CHECK_EQ(plan.passCount, 1U); - CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); + CHECK_EQ(plan.chunkBytesPerPeer, static_cast(elementsPerPeer) * sizeof(int32_t)); CHECK_EQ(plan.controlBytes, controlGroupBytes); CHECK_EQ(plan.signalBytes, 4ULL * controlGroupBytes); CHECK_EQ(plan.copyDoneOffset, plan.dataBytes); @@ -204,7 +218,23 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() CHECK_EQ(plan.dataBytes + plan.controlBytes + plan.signalBytes <= plan.registeredBytes, true); CHECK_EQ(plan.dataBytes, static_cast(rankSize - 1) * static_cast(plan.passCount) * 2ULL * - TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); + static_cast(elementsPerPeer) * sizeof(int32_t)); +} + +void TestAllToAllBigDataMultiNodeSmallPayloadUsesPayloadSlot() +{ + constexpr int rankSize = 64; + constexpr int32_t elementsPerPeer = 1024 * 1024; // 4 MiB per peer for int32_t. + const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); + const size_t expectedChunkBytes = static_cast(elementsPerPeer) * sizeof(int32_t); + + CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); + CHECK_EQ(plan.passCount, 1U); + CHECK_EQ(plan.chunkBytesPerPeer, expectedChunkBytes); + CHECK_EQ(plan.dataBytes + plan.controlBytes + plan.signalBytes <= plan.registeredBytes, true); + CHECK_EQ(plan.dataBytes, + static_cast(rankSize - 1) * static_cast(plan.passCount) * 2ULL * + expectedChunkBytes); } void TestAllToAllBigDataForce35CorePlanFor8P() @@ -223,7 +253,7 @@ void TestAllToAllBigDataForce35CorePlanFor8P() CHECK_EQ(TileXR::Demo::AllToAllBigDataShardCount(rankSize, true), 32U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(rankSize, true), 35U); CHECK_EQ(plan.registeredBytes, TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes); - CHECK_EQ(plan.chunkBytesPerPeer, TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes); + CHECK_EQ(plan.chunkBytesPerPeer, static_cast(elementsPerPeer) * sizeof(int32_t)); CHECK_EQ(plan.controlBytes, controlGroupBytes); CHECK_EQ(plan.signalBytes, 4ULL * controlGroupBytes); CHECK_EQ(plan.remoteSendDoneOffset, plan.recvCopyDoneOffset + controlGroupBytes); @@ -241,6 +271,8 @@ void TestAllToAllBigDataBlockDim() CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(16), 35U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(32), 35U); CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64), 35U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(16, false, true), 64U); + CHECK_EQ(TileXR::Demo::AllToAllBigDataBlockDim(64, false, true), 64U); } void TestAllToAllBigDataMultiNodeTopology() @@ -407,6 +439,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "force35Core="); CHECK_CONTAINS(demo, "TileXR::Demo::AllToAllBigDataShardCount("); CHECK_CONTAINS(demo, "const uint32_t bigDataBlockDim = TileXR::Demo::AllToAllBigDataBlockDim("); + CHECK_CONTAINS(demo, "rankSize, forceBigData35Core, bigDataRemotePutOnly"); CHECK_CONTAINS(demo, "bigDataBlockDim, stream, commArgsDev"); CHECK_CONTAINS(demo, "static_cast(registeredMemory) + bigDataPlan.copyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.controlBytes + bigDataPlan.signalBytes"); @@ -438,10 +471,11 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 8ULL * 1024ULL * 1024ULL"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 16ULL * 1024ULL * 1024ULL"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE = 8"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS = 32U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM ="); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_BLOCK_DIM = 64U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE = 16U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE = 17U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE = 18U"); @@ -458,24 +492,61 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataTaskCount(rankSize, force35Core)"); CHECK_CONTAINS(kernel, "BigDataNodeCount(rankSize)"); CHECK_CONTAINS(kernel, "BigDataRemotePeerAt(rank, rankSize, remoteIndex)"); + CHECK_CONTAINS(kernel, "BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex)"); CHECK_CONTAINS(kernel, "BigDataLocalPeerAt"); CHECK_CONTAINS(kernel, "BigDataMergedPeerTaskAt(rank, rankSize, taskIndex, force35Core, peer, isLocalPeer)"); CHECK_CONTAINS(kernel, "BigDataRunSelfCopyShard(rank, rankSize"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentWorker"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendWorker"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskCount"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskRemoteIndex"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskSegment"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckIndex"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker"); + CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); + CHECK_CONTAINS(kernel, "BigDataLocalIpcAckSlot"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyPublishAck"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyWaitPeerAck"); + CHECK_CONTAINS(kernel, "BigDataStoreInt32Mte"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentRange"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_SEGMENT"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_ADDRESS"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST"); + CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_ACK"); + CHECK_CONTAINS(kernel, "BigDataSelectWeightedQp("); + CHECK_CONTAINS(kernel, "BigDataSelectDistinctWeightedQp("); + CHECK_CONTAINS(kernel, "segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT"); + CHECK_CONTAINS(kernel, "primaryQp : BigDataSelectDistinctWeightedQp(args, peer, primaryQp, false)"); CHECK_CONTAINS(kernel, "BigDataWaitCopyDoneRange"); CHECK_CONTAINS(kernel, "BigDataPublishCopyReadyRange"); CHECK_CONTAINS(kernel, "BigDataWaitCopyReady"); CHECK_CONTAINS(kernel, "BigDataPublishReadySignal"); CHECK_CONTAINS(kernel, "remoteSendDoneOffset"); CHECK_CONTAINS(kernel, "if (!BigDataIsMultiNode(rankSize))"); - CHECK_CONTAINS(kernel, "if (blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM))"); + CHECK_CONTAINS(kernel, "const uint32_t activeBlockDim = remotePutOnly ?"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_BLOCK_DIM"); + CHECK_CONTAINS(kernel, "if (blockIdx >= static_cast(activeBlockDim))"); CHECK_CONTAINS(kernel, "const bool isCopyCore ="); CHECK_CONTAINS(kernel, "const bool isRemoteSendPrimaryCore ="); CHECK_CONTAINS(kernel, "const bool isRemoteSendSecondaryCore ="); CHECK_CONTAINS(kernel, "const bool isLocalSendCore ="); CHECK_CONTAINS(kernel, "const bool isRecvCore ="); CHECK_CONTAINS(kernel, "if (isCopyCore)"); + CHECK_CONTAINS(kernel, "if (remotePutOnly)"); + CHECK_CONTAINS(kernel, "isRemotePutOnlyCheckCore"); + CHECK_CONTAINS(kernel, "remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK"); + CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP"); + CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER"); + CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP"); + CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_SEGMENT"); + CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_ADDRESS"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker("); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckIndex(blockIdx)"); + CHECK_CONTAINS(kernel, "BigDataKernelExitBarrier()"); CHECK_CONTAINS(kernel, "if (!isLocalPeer && isRemoteSendPrimaryCore)"); CHECK_CONTAINS(kernel, "if (!isLocalPeer && isRemoteSendSecondaryCore)"); CHECK_CONTAINS(kernel, "if (isLocalPeer && isLocalSendCore)"); @@ -510,6 +581,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataNetworkPeerIndex(rank, peer)"); CHECK_CONTAINS(kernel, "BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "UDMAPutNbiOnQp"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); CHECK_CONTAINS(kernel, "UDMAPutNbi"); CHECK_CONTAINS(kernel, "localSrc,"); @@ -524,9 +596,41 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataPublishAckSignalUdma"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize)"); CHECK_CONTAINS(kernel, "ackSignal"); - CHECK_NOT_CONTAINS(kernel, "UDMAPutNbi"); + CHECK_CONTAINS(kernel, "fillLocal.SetValue(0, value)"); CHECK_CONTAINS(udma, "if (length == 0)"); CHECK_CONTAINS(udma, "reinterpret_cast(addr) + length - 1"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY"); + CHECK_CONTAINS(demo, "skip result validation for bigdata remote-put-only profile"); + + const std::string remotePutOnlySend = SliceBetween( + kernel, "BigDataRemotePutOnlySendWorker", "BigDataRemotePutOnlyCheckIndex"); + const std::string remotePutOnlyCheck = SliceBetween( + kernel, "BigDataRemotePutOnlyCheckWorker", "BigDataRunSelfCopyShard"); + CHECK_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); + CHECK_CONTAINS(remotePutOnlySend, "BigDataStoreInt32Mte"); + CHECK_NOT_CONTAINS(remotePutOnlySend, + "if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT)"); + CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAQuietStatusOnQp"); + CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES = 32U"); + CHECK_CONTAINS(kernel, "const uint32_t sendGroupCore ="); + CHECK_CONTAINS(kernel, "const uint32_t segmentId = blockIdx < static_cast("); + CHECK_CONTAINS(kernel, "for (int32_t remoteIndex = static_cast(sendGroupCore);"); + CHECK_CONTAINS(kernel, "remoteIndex += static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES)"); + const std::string remotePutOnlyCheckIndex = SliceBetween( + kernel, "BigDataRemotePutOnlyCheckIndex", "BigDataRemotePutOnlyCheckWorker"); + CHECK_CONTAINS(remotePutOnlyCheckIndex, "return blockIdx"); + CHECK_NOT_CONTAINS(remotePutOnlyCheckIndex, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE"); + CHECK_NOT_CONTAINS(remotePutOnlyCheckIndex, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE"); + CHECK_CONTAINS(remotePutOnlyCheck, "observed"); + CHECK_CONTAINS(remotePutOnlyCheck, "BigDataNetworkPeerIndex(peer, rank)"); + CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS"); + CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyPublishAck"); + CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyWaitPeerAck"); + CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemoteIpcAckSlot"); + CHECK_CONTAINS(remotePutOnlyCheck, "BigDataLocalIpcAckSlot"); + CHECK_NOT_CONTAINS(remotePutOnlyCheck, "UDMAQuietStatusOnQp"); + CHECK_NOT_CONTAINS(remotePutOnlyCheck, "UDMAPutSignalNbi"); } } // namespace @@ -539,6 +643,7 @@ int main() TestAllToAllMaxRank256With64MiBPerRank(); TestAllToAllBigDataPlan(); TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone(); + TestAllToAllBigDataMultiNodeSmallPayloadUsesPayloadSlot(); TestAllToAllBigDataForce35CorePlanFor8P(); TestAllToAllBigDataBlockDim(); TestAllToAllBigDataMultiNodeTopology(); From cb3a7f511175ede3e99f1619cbf11be8834fcb8b Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 09:56:02 +0800 Subject: [PATCH 024/163] perf(udma): restore validated remote put mteack path Return remote-put-only scheduling to the previously validated flat 64-core send-task path with secondary tail completion checking after the 32+32 tail-flag experiment stalled in 64p testing. --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 30 ++++++++----------- .../unit/test_tilexr_udma_alltoall_layout.cpp | 12 ++++---- 2 files changed, 18 insertions(+), 24 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 3126a792..4bd284ed 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -72,7 +72,6 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM = TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_BLOCK_DIM = 64U; -constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES = 32U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT = 0U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT = 1U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY = 2U; @@ -1406,12 +1405,14 @@ __aicore__ inline void BigDataRemotePutOnlySendWorker( uint32_t status = 0U; if (segmentBytes > 0U) { - const int32_t tokenValue = - static_cast(BigDataPassToken(globalPass)); - auto localTail = reinterpret_cast<__gm__ int32_t*>( - reinterpret_cast<__gm__ uint8_t*>(localSrc) + - static_cast(segmentBytes) - sizeof(int32_t)); - BigDataStoreInt32Mte(localTail, tokenValue, relayLocal); + if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT) { + const int32_t tokenValue = + static_cast(BigDataPassToken(globalPass)); + auto localTail = reinterpret_cast<__gm__ int32_t*>( + reinterpret_cast<__gm__ uint8_t*>(localSrc) + + static_cast(segmentBytes) - sizeof(int32_t)); + BigDataStoreInt32Mte(localTail, tokenValue, relayLocal); + } TileXR::UDMAPutNbiOnQp(args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); } if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { @@ -2759,6 +2760,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( static_cast(blockIdx) - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; const int32_t taskCount = BigDataTaskCount(rankSize, force35Core); const int32_t remoteTaskCount = rankSize - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount); const int32_t remotePutOnlyCheckIndex = BigDataRemotePutOnlyCheckIndex(blockIdx); const bool isRemotePutOnlyCheckCore = remotePutOnlyCheckIndex >= 0 && remotePutOnlyCheckIndex < remoteTaskCount; @@ -2771,23 +2773,17 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { if (remotePutOnly) { - const uint32_t sendGroupCore = - static_cast(blockIdx) % - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES; - const uint32_t segmentId = blockIdx < static_cast( - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES) ? - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT : - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; - for (int32_t remoteIndex = static_cast(sendGroupCore); - remoteIndex < remoteTaskCount; - remoteIndex += static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES)) { + for (int32_t sendTask = blockIdx; sendTask < sendTaskCount; + sendTask += static_cast(activeBlockDim)) { if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { continue; } + const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount); const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex); if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { continue; } + const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount); const uint32_t primaryQp = BigDataSelectWeightedQp(args, peer, true); const uint32_t qpIdx = segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index b9394e1c..6057297c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -538,6 +538,10 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "if (isCopyCore)"); CHECK_CONTAINS(kernel, "if (remotePutOnly)"); CHECK_CONTAINS(kernel, "isRemotePutOnlyCheckCore"); + CHECK_CONTAINS(kernel, "const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); + CHECK_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); + CHECK_CONTAINS(kernel, "const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); + CHECK_CONTAINS(kernel, "const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount)"); CHECK_CONTAINS(kernel, "remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER"); @@ -608,15 +612,9 @@ void TestAllToAllBigDataSource() kernel, "BigDataRemotePutOnlyCheckWorker", "BigDataRunSelfCopyShard"); CHECK_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); CHECK_CONTAINS(remotePutOnlySend, "BigDataStoreInt32Mte"); - CHECK_NOT_CONTAINS(remotePutOnlySend, - "if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT)"); + CHECK_CONTAINS(remotePutOnlySend, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT"); CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAQuietStatusOnQp"); CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES = 32U"); - CHECK_CONTAINS(kernel, "const uint32_t sendGroupCore ="); - CHECK_CONTAINS(kernel, "const uint32_t segmentId = blockIdx < static_cast("); - CHECK_CONTAINS(kernel, "for (int32_t remoteIndex = static_cast(sendGroupCore);"); - CHECK_CONTAINS(kernel, "remoteIndex += static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_SEND_GROUP_CORES)"); const std::string remotePutOnlyCheckIndex = SliceBetween( kernel, "BigDataRemotePutOnlyCheckIndex", "BigDataRemotePutOnlyCheckWorker"); CHECK_CONTAINS(remotePutOnlyCheckIndex, "return blockIdx"); From 6564bba615c573658be192c5bf0e4cb368f0d763 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 10:05:07 +0800 Subject: [PATCH 025/163] perf(udma): disable remote put tail check path Restore the validated remote-put-only performance path by removing the tail-flag polling and peer ack checks from the flat 64-core send-task flow. The 16p loop=1 8MiB/peer validation passes again without timeout debug statuses. --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 20 ------------------- .../unit/test_tilexr_udma_alltoall_layout.cpp | 6 +----- 2 files changed, 1 insertion(+), 25 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 4bd284ed..863f3c86 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -1405,14 +1405,6 @@ __aicore__ inline void BigDataRemotePutOnlySendWorker( uint32_t status = 0U; if (segmentBytes > 0U) { - if (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT) { - const int32_t tokenValue = - static_cast(BigDataPassToken(globalPass)); - auto localTail = reinterpret_cast<__gm__ int32_t*>( - reinterpret_cast<__gm__ uint8_t*>(localSrc) + - static_cast(segmentBytes) - sizeof(int32_t)); - BigDataStoreInt32Mte(localTail, tokenValue, relayLocal); - } TileXR::UDMAPutNbiOnQp(args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); } if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { @@ -2761,9 +2753,6 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( const int32_t taskCount = BigDataTaskCount(rankSize, force35Core); const int32_t remoteTaskCount = rankSize - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount); - const int32_t remotePutOnlyCheckIndex = BigDataRemotePutOnlyCheckIndex(blockIdx); - const bool isRemotePutOnlyCheckCore = - remotePutOnlyCheckIndex >= 0 && remotePutOnlyCheckIndex < remoteTaskCount; if (remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK) { BigDataKernelExitBarrier(); @@ -2798,15 +2787,6 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( chunkBytesPerPeer, relayLocal); } BigDataKernelExitBarrier(); - if (profileStage > TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST && - isRemotePutOnlyCheckCore) { - BigDataRemotePutOnlyCheckWorker( - remotePutOnlyCheckIndex, rank, rankSize, args, udmaMem, debug, - elementsPerPeer, effectiveChunkElements, passCount, loop, pass, - kernelLoopBase, recvDataOffset, ackSignalOffset, chunkBytesPerPeer, - shardCount, relayLocal); - } - BigDataKernelExitBarrier(); continue; } if (isCopyCore) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 6057297c..5437bace 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -537,7 +537,6 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "const bool isRecvCore ="); CHECK_CONTAINS(kernel, "if (isCopyCore)"); CHECK_CONTAINS(kernel, "if (remotePutOnly)"); - CHECK_CONTAINS(kernel, "isRemotePutOnlyCheckCore"); CHECK_CONTAINS(kernel, "const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); CHECK_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); CHECK_CONTAINS(kernel, "const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); @@ -548,8 +547,6 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_SEGMENT"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_ADDRESS"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker("); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckIndex(blockIdx)"); CHECK_CONTAINS(kernel, "BigDataKernelExitBarrier()"); CHECK_CONTAINS(kernel, "if (!isLocalPeer && isRemoteSendPrimaryCore)"); CHECK_CONTAINS(kernel, "if (!isLocalPeer && isRemoteSendSecondaryCore)"); @@ -611,8 +608,7 @@ void TestAllToAllBigDataSource() const std::string remotePutOnlyCheck = SliceBetween( kernel, "BigDataRemotePutOnlyCheckWorker", "BigDataRunSelfCopyShard"); CHECK_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); - CHECK_CONTAINS(remotePutOnlySend, "BigDataStoreInt32Mte"); - CHECK_CONTAINS(remotePutOnlySend, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT"); + CHECK_NOT_CONTAINS(remotePutOnlySend, "BigDataStoreInt32Mte"); CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAQuietStatusOnQp"); CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); const std::string remotePutOnlyCheckIndex = SliceBetween( From 22886323742ce21832e8287f31426ecf7104ae0c Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 10:57:57 +0800 Subject: [PATCH 026/163] fix(udma): wait for remote put completion Publish IPC ready tokens after UDMA quiet for remote-put-only sends and make check workers wait for both remote segments before exchanging ACKs. Treat incomplete remote-put-only debug status as a fatal demo error and update layout guards. --- tests/udma/demo/tilexr_udma_demo.cpp | 16 ++++ tests/udma/demo/tilexr_udma_demo_kernel.cpp | 88 ++++++++++++++----- .../unit/test_tilexr_udma_alltoall_layout.cpp | 13 ++- 3 files changed, 90 insertions(+), 27 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index f19f1f44..020d1423 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -1009,6 +1009,22 @@ int main(int argc, char** argv) Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } + if (bigDataRemotePutOnly && !AllToAllUdmaComplete(rankSize, hostDebug)) { + std::cerr << "[rank " << rank << "] ERROR: bigdata remote-put-only incomplete:"; + for (int peer = 0; peer < rankSize; ++peer) { + std::cerr << " peer" << peer << "=" << hostDebug[kDebugUdmaStatusBase + peer]; + } + std::cerr << std::endl; + PrintAllToAllUdmaDebug(rank, rankSize, hostDebug); + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); udmaRegistered = false; diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 863f3c86..cba94024 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -759,6 +759,20 @@ __aicore__ inline __gm__ uint64_t* BigDataRemoteIpcAckSlot( args->peerMems[targetRank] + BigDataIpcAckOffset(slot, rankSize, rank)); } +__aicore__ inline __gm__ uint64_t* BigDataLocalIpcReadySlot( + const __gm__ TileXR::CommArgs* args, int32_t rank, uint32_t slot, int32_t rankSize, + int32_t peer, uint32_t segmentId) +{ + return BigDataLocalIpcAckSlot(args, rank, slot, rankSize, peer) + 1U + segmentId; +} + +__aicore__ inline __gm__ uint64_t* BigDataRemoteIpcReadySlot( + const __gm__ TileXR::CommArgs* args, int32_t targetRank, int32_t rank, uint32_t slot, + int32_t rankSize, uint32_t segmentId) +{ + return BigDataRemoteIpcAckSlot(args, targetRank, rank, slot, rankSize) + 1U + segmentId; +} + __aicore__ inline uint32_t BigDataPublishAckSignalUdma( const __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, uint64_t ackSignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, @@ -779,6 +793,12 @@ __aicore__ inline void BigDataRemotePutOnlyPublishAck( BigDataStoreTokenMte(remoteAck, token, relayLocal); } +__aicore__ inline void BigDataRemotePutOnlyPublishReady( + __gm__ uint64_t* remoteReady, uint64_t token, AscendC::LocalTensor relayLocal) +{ + BigDataStoreTokenMte(remoteReady, token, relayLocal); +} + __aicore__ inline uint64_t BigDataRemotePutOnlyWaitPeerAck( __gm__ uint64_t* peerAck, uint64_t token, AscendC::LocalTensor relayLocal) { @@ -1406,6 +1426,10 @@ __aicore__ inline void BigDataRemotePutOnlySendWorker( uint32_t status = 0U; if (segmentBytes > 0U) { TileXR::UDMAPutNbiOnQp(args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); + status = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + BigDataRemotePutOnlyPublishReady( + BigDataRemoteIpcReadySlot(args, peer, rank, slot, rankSize, segmentId), + BigDataPassToken(globalPass), relayLocal); } if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); @@ -1464,31 +1488,38 @@ __aicore__ inline void BigDataRemotePutOnlyCheckWorker( const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, true); - const uint64_t remoteDataOffset = - recvDataOffset + - (static_cast(slot) * static_cast(rankSize > 1 ? rankSize - 1 : 1) + - static_cast(BigDataNetworkPeerIndex(peer, rank))) * chunkBytesPerPeer; - auto tail = reinterpret_cast<__gm__ int32_t*>( - udmaMem + remoteDataOffset + static_cast(chunkBytes) - sizeof(int32_t)); - - int32_t observed = 0; - uint64_t polls = 0; - do { - AscendC::GlobalTensor tailGlobal; - tailGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(tail)); - AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; - AscendC::DataCopyExtParams copyIn {1U, static_cast(sizeof(int32_t)), 0U, 0U, 0U}; - AscendC::DataCopyPad(relayLocal, tailGlobal, copyIn, padIn); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - observed = relayLocal.ReinterpretCast().GetValue(0); - ++polls; - } while (static_cast(observed) < token && polls < TILEXR_UDMA_DEMO_SIGNAL_MAX_POLLS); + (void)recvDataOffset; + (void)chunkBytesPerPeer; + (void)chunkBytes; uint32_t status = 0U; - if (static_cast(observed) < token) { - status = static_cast(TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS); - } else { + uint64_t observed = 0; + constexpr uint32_t dataShardCount = TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES; + for (uint32_t segmentId = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT; + segmentId <= TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; ++segmentId) { + uint32_t copyShardBegin = 0U; + uint32_t copyShardEnd = 0U; + uint32_t segmentOffsetBytes = 0U; + uint32_t segmentBytes = 0U; + if (!BigDataRemoteSendSegmentRange( + segmentId, dataShardCount, chunkBytes / sizeof(int32_t), + copyShardBegin, copyShardEnd, segmentOffsetBytes, segmentBytes) || + segmentBytes < sizeof(int32_t)) { + continue; + } + (void)copyShardBegin; + (void)copyShardEnd; + (void)segmentOffsetBytes; + (void)segmentBytes; + observed = BigDataWaitTokenMte( + BigDataLocalIpcReadySlot(args, rank, slot, rankSize, peer, segmentId), + token, relayLocal); + if (observed < token) { + status = static_cast(TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS); + break; + } + } + if (status == 0U) { auto remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); auto localAck = BigDataLocalIpcAckSlot(args, rank, slot, rankSize, peer); BigDataRemotePutOnlyPublishAck(remoteAck, token, relayLocal); @@ -2787,6 +2818,17 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( chunkBytesPerPeer, relayLocal); } BigDataKernelExitBarrier(); + if (profileStage > TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST) { + const int32_t remotePutOnlyCheckIndex = BigDataRemotePutOnlyCheckIndex(blockIdx); + if (remotePutOnlyCheckIndex >= 0 && remotePutOnlyCheckIndex < remoteTaskCount) { + BigDataRemotePutOnlyCheckWorker( + remotePutOnlyCheckIndex, rank, rankSize, args, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, recvDataOffset, ackSignalOffset, chunkBytesPerPeer, + shardCount, relayLocal); + } + } + BigDataKernelExitBarrier(); continue; } if (isCopyCore) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 5437bace..b53f8880 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -506,6 +506,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); CHECK_CONTAINS(kernel, "BigDataLocalIpcAckSlot"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyPublishAck"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyPublishReady"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyWaitPeerAck"); CHECK_CONTAINS(kernel, "BigDataStoreInt32Mte"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentRange"); @@ -537,6 +538,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "const bool isRecvCore ="); CHECK_CONTAINS(kernel, "if (isCopyCore)"); CHECK_CONTAINS(kernel, "if (remotePutOnly)"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker("); CHECK_CONTAINS(kernel, "const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); CHECK_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); CHECK_CONTAINS(kernel, "const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); @@ -608,23 +610,26 @@ void TestAllToAllBigDataSource() const std::string remotePutOnlyCheck = SliceBetween( kernel, "BigDataRemotePutOnlyCheckWorker", "BigDataRunSelfCopyShard"); CHECK_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); - CHECK_NOT_CONTAINS(remotePutOnlySend, "BigDataStoreInt32Mte"); - CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAQuietStatusOnQp"); + CHECK_CONTAINS(remotePutOnlySend, "BigDataRemoteIpcReadySlot"); + CHECK_CONTAINS(remotePutOnlySend, "BigDataRemotePutOnlyPublishReady"); + CHECK_CONTAINS(remotePutOnlySend, "UDMAQuietStatusOnQp"); CHECK_NOT_CONTAINS(remotePutOnlySend, "UDMAPutNbiOnQp"); + CHECK_NOT_CONTAINS(remotePutOnlySend, "BigDataStoreInt32Mte"); const std::string remotePutOnlyCheckIndex = SliceBetween( kernel, "BigDataRemotePutOnlyCheckIndex", "BigDataRemotePutOnlyCheckWorker"); CHECK_CONTAINS(remotePutOnlyCheckIndex, "return blockIdx"); CHECK_NOT_CONTAINS(remotePutOnlyCheckIndex, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE"); CHECK_NOT_CONTAINS(remotePutOnlyCheckIndex, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE"); CHECK_CONTAINS(remotePutOnlyCheck, "observed"); - CHECK_CONTAINS(remotePutOnlyCheck, "BigDataNetworkPeerIndex(peer, rank)"); + CHECK_CONTAINS(remotePutOnlyCheck, "BigDataLocalIpcReadySlot"); + CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT"); + CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT"); CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS"); CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyPublishAck"); CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyWaitPeerAck"); CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemoteIpcAckSlot"); CHECK_CONTAINS(remotePutOnlyCheck, "BigDataLocalIpcAckSlot"); CHECK_NOT_CONTAINS(remotePutOnlyCheck, "UDMAQuietStatusOnQp"); - CHECK_NOT_CONTAINS(remotePutOnlyCheck, "UDMAPutSignalNbi"); } } // namespace From 4eea5bb30b86d4497856258eb0898347c2cb912f Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 11:31:42 +0800 Subject: [PATCH 027/163] test(udma): add alltoall performance runner Add a Python one-click runner that syncs local TileXR sources to remote hosts, builds the UDMA demo, runs warmup and measured alltoall launches, records host wall time, and summarizes per-rank operator latency. Cover size parsing, host selection, metric parsing, summary statistics, and remote tool environment inference. --- tests/udma/perf/run_udma_alltoall_perf.py | 577 ++++++++++++++++++ .../udma/perf/test_run_udma_alltoall_perf.py | 69 +++ 2 files changed, 646 insertions(+) create mode 100644 tests/udma/perf/run_udma_alltoall_perf.py create mode 100644 tests/udma/perf/test_run_udma_alltoall_perf.py diff --git a/tests/udma/perf/run_udma_alltoall_perf.py b/tests/udma/perf/run_udma_alltoall_perf.py new file mode 100644 index 00000000..75b357e2 --- /dev/null +++ b/tests/udma/perf/run_udma_alltoall_perf.py @@ -0,0 +1,577 @@ +#!/usr/bin/env python3 + +import argparse +import concurrent.futures +import dataclasses +import os +import re +import shlex +import statistics +import sys +import tarfile +import tempfile +import time +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + + +DEFAULT_HOSTS = [ + "141.61.92.151", + "141.61.92.147", + "141.61.92.143", + "141.61.92.139", + "141.61.92.111", + "141.61.92.107", + "141.61.92.103", + "141.61.92.99", +] +DEFAULT_CANN_ENV = "/home/c30061605/startup_discom/tilexr/env/cann/cann/set_env.sh" +DEFAULT_DEMO_BIN = "tests/udma/build/tilexr_udma_demo" +DEFAULT_TOOL_ENV = "" +INT32_BYTES = 4 + + +@dataclasses.dataclass(frozen=True) +class RankMetric: + rank: int + iters: int + total_ms: float + per_iter_us: float + payload_bytes: float + bandwidth_gbps: float + + +@dataclasses.dataclass(frozen=True) +class RankRunResult: + rank: int + host: str + rc: int + stdout: str + stderr: str + elapsed_s: float + + +SIZE_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([kmgt]?i?b?|bytes?)?\s*$", re.IGNORECASE) +METRIC_RE = re.compile( + r"\[rank\s+(?P\d+)\].*?alltoall\s+udma-bigdata\s+" + r"(?P\d+)\s+iters.*?total=(?P[0-9.eE+-]+)\s+ms\s+" + r"perIter=(?P[0-9.eE+-]+)\s+us\s+" + r"payload=(?P[0-9.eE+-]+)\s+bytes\s+" + r"bw=(?P[0-9.eE+-]+)\s+GB/s" +) + + +def parse_size_to_bytes(value: str) -> int: + match = SIZE_RE.match(value) + if not match: + raise ValueError(f"invalid size: {value!r}") + number = float(match.group(1)) + suffix = (match.group(2) or "b").lower() + scale = { + "": 1, + "b": 1, + "byte": 1, + "bytes": 1, + "k": 1024, + "kb": 1024, + "kib": 1024, + "m": 1024 ** 2, + "mb": 1024 ** 2, + "mib": 1024 ** 2, + "g": 1024 ** 3, + "gb": 1024 ** 3, + "gib": 1024 ** 3, + "t": 1024 ** 4, + "tb": 1024 ** 4, + "tib": 1024 ** 4, + }.get(suffix) + if scale is None: + raise ValueError(f"invalid size suffix: {value!r}") + size = int(number * scale) + if size <= 0: + raise ValueError("size must be positive") + return size + + +def bytes_to_int32_elements(byte_count: int) -> int: + if byte_count % INT32_BYTES != 0: + raise ValueError(f"bytes per peer must be {INT32_BYTES}-byte aligned for int32 payload") + return byte_count // INT32_BYTES + + +def select_hosts(hosts: Sequence[str], world_size: int, devices_per_host: int) -> List[str]: + if world_size <= 0: + raise ValueError("world size must be positive") + if devices_per_host <= 0: + raise ValueError("devices per host must be positive") + if world_size % devices_per_host != 0: + raise ValueError(f"world size {world_size} must be a multiple of devices per host {devices_per_host}") + host_count = world_size // devices_per_host + if host_count > len(hosts): + raise ValueError(f"world size {world_size} requires {host_count} hosts, only {len(hosts)} configured") + return list(hosts[:host_count]) + + +def parse_demo_metric(text: str) -> RankMetric: + matches = list(METRIC_RE.finditer(text)) + if not matches: + raise ValueError("no alltoall udma-bigdata metric line found") + match = matches[-1] + return RankMetric( + rank=int(match.group("rank")), + iters=int(match.group("iters")), + total_ms=float(match.group("total")), + per_iter_us=float(match.group("per")), + payload_bytes=float(match.group("payload")), + bandwidth_gbps=float(match.group("bw")), + ) + + +def percentile(sorted_values: Sequence[float], fraction: float) -> float: + if not sorted_values: + raise ValueError("empty values") + index = min(len(sorted_values) - 1, int(round((len(sorted_values) - 1) * fraction))) + return sorted_values[index] + + +def summarize_metrics(metrics: Sequence[RankMetric]) -> Dict[str, float]: + values = sorted(metric.per_iter_us for metric in metrics) + if not values: + raise ValueError("no rank metrics to summarize") + return { + "count": len(values), + "min_us": values[0], + "p50_us": percentile(values, 0.50), + "mean_us": statistics.mean(values), + "p90_us": percentile(values, 0.90), + "p99_us": percentile(values, 0.99), + "max_us": values[-1], + } + + +def repo_root() -> Path: + return Path(__file__).resolve().parents[3] + + +def should_skip_archive_path(path: Path) -> bool: + parts = set(path.parts) + if ".git" in parts or ".worktrees" in parts: + return True + if "__pycache__" in parts: + return True + if path.suffix in {".pyc", ".log"}: + return True + return any(part in {"build", "install", "build_device", "output", "run"} for part in parts) + + +def create_source_archive(root: Path) -> Path: + handle = tempfile.NamedTemporaryFile(prefix="tilexr_sync_", suffix=".tar.gz", delete=False) + handle.close() + archive = Path(handle.name) + with tarfile.open(archive, "w:gz") as tar: + for path in root.rglob("*"): + relative = path.relative_to(root) + if should_skip_archive_path(relative): + continue + tar.add(path, arcname=str(relative), recursive=False) + return archive + + +def load_paramiko(): + try: + import paramiko # type: ignore + except ImportError as exc: + raise RuntimeError("paramiko is required; install it or run from an environment that provides it") from exc + return paramiko + + +def ssh_connect(host: str, user: str, password: str, timeout_s: int): + paramiko = load_paramiko() + client = paramiko.SSHClient() + client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + client.connect( + host, + username=user, + password=password, + timeout=timeout_s, + banner_timeout=timeout_s, + auth_timeout=timeout_s, + ) + return client + + +def run_ssh(host: str, user: str, password: str, command: str, timeout_s: int) -> Tuple[int, str, str]: + client = ssh_connect(host, user, password, min(timeout_s, 30)) + try: + stdin, stdout, stderr = client.exec_command(command, timeout=timeout_s) + out = stdout.read().decode("utf-8", "replace") + err = stderr.read().decode("utf-8", "replace") + rc = stdout.channel.recv_exit_status() + stdin.close() + stdout.close() + stderr.close() + return rc, out, err + finally: + client.close() + + +def upload_file(host: str, user: str, password: str, local: Path, remote: str, timeout_s: int) -> None: + client = ssh_connect(host, user, password, min(timeout_s, 30)) + try: + sftp = client.open_sftp() + try: + sftp.put(str(local), remote) + finally: + sftp.close() + finally: + client.close() + + +def shell_join(commands: Iterable[str]) -> str: + return " && ".join(commands) + + +def infer_tool_env(cann_env: str) -> str: + marker = "/cann/cann/set_env.sh" + if cann_env.endswith(marker): + return cann_env[: -len(marker)] + return "" + + +def remote_env_prelude(cann_env: str, tool_env: str) -> str: + commands = [] + if tool_env: + commands.extend([ + f"export TILEXR_EXTERNAL_ENV={shlex.quote(tool_env)}", + f"export PATH={shlex.quote(tool_env)}/util/cmake/bin:{shlex.quote(tool_env)}/util/mpich/bin:$PATH", + f"export PATH={shlex.quote(tool_env)}/temp/cmake/bin:{shlex.quote(tool_env)}/temp/mpich/bin:$PATH", + f"export LD_LIBRARY_PATH={shlex.quote(tool_env)}/util/mpich/lib:{shlex.quote(tool_env)}/temp/mpich/lib:$LD_LIBRARY_PATH", + ]) + commands.append(f"source {shlex.quote(cann_env)} >/dev/null 2>&1") + return "; ".join(commands) + + +def sync_host( + host: str, + user: str, + password: str, + archive: Path, + remote_dir: str, + clean_remote: bool, + timeout_s: int, +) -> None: + remote_archive = f"/tmp/{archive.name}" + upload_file(host, user, password, archive, remote_archive, timeout_s) + setup = [] + if clean_remote: + setup.append(f"rm -rf -- {shlex.quote(remote_dir)}") + setup.extend([ + f"mkdir -p -- {shlex.quote(remote_dir)}", + f"tar -xzf {shlex.quote(remote_archive)} -C {shlex.quote(remote_dir)}", + f"rm -f -- {shlex.quote(remote_archive)}", + ]) + rc, out, err = run_ssh(host, user, password, f"bash -lc {shlex.quote(shell_join(setup))}", timeout_s) + if rc != 0: + raise RuntimeError(f"sync failed on {host}: rc={rc}\n{out}\n{err}") + + +def build_host(host: str, user: str, password: str, remote_dir: str, cann_env: str, tool_env: str, timeout_s: int) -> None: + command = shell_join([ + remote_env_prelude(cann_env, tool_env), + f"cd {shlex.quote(remote_dir)}", + "command -v cmake >/dev/null", + "bash tests/udma/build.sh", + ]) + rc, out, err = run_ssh(host, user, password, f"bash -lc {shlex.quote(command)}", timeout_s) + if rc != 0: + raise RuntimeError(f"build failed on {host}: rc={rc}\n{out[-4000:]}\n{err[-4000:]}") + + +def make_rank_command( + *, + rank: int, + world_size: int, + repeat: int, + elements_per_peer: int, + devices_per_host: int, + remote_dir: str, + cann_env: str, + demo_bin: str, + comm_id: str, + barrier_host: str, + result_dir: str, + route_policy: str, + timeout_s: int, + tool_env: str, +) -> str: + log_file = f"{result_dir}/rank_{rank}.log" + rc_file = f"{result_dir}/rank_{rank}.rc" + inner = shell_join([ + f"mkdir -p {shlex.quote(result_dir)}", + remote_env_prelude(cann_env, tool_env), + f"cd {shlex.quote(remote_dir)}", + f"export TILEXR_COMM_ID={shlex.quote(comm_id)}", + f"export TILEXR_DEMO_BARRIER_HOST={shlex.quote(barrier_host)}", + f"export TILEXR_UDMA_ROUTE_POLICY={shlex.quote(route_policy)}", + "export TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY=1", + "export TILEXR_DEMO_BIGDATA_PROFILE_STAGE=7", + f"export TILEXR_DEMO_ALLTOALL_REPEAT={repeat}", + "export TILEXR_DEMO_ALLTOALL_WARMUP=0", + "export TILEXR_DEMO_ALLTOALL_SYNC_AT_END=1", + f"export TILEXR_DEMO_TIMEOUT_SECONDS={timeout_s}", + ( + f"timeout {timeout_s + 60} ./{shlex.quote(demo_bin)} " + f"{world_size} {rank} 7 {elements_per_peer} {devices_per_host} 0 " + f"> {shlex.quote(log_file)} 2>&1; " + f"rc=$?; echo $rc > {shlex.quote(rc_file)}; " + f"tail -n 80 {shlex.quote(log_file)}; exit $rc" + ), + ]) + return f"bash -lc {shlex.quote('set -o pipefail; ' + inner)}" + + +def run_rank( + host: str, + user: str, + password: str, + command: str, + rank: int, + timeout_s: int, +) -> RankRunResult: + start = time.perf_counter() + try: + rc, out, err = run_ssh(host, user, password, command, timeout_s + 90) + except Exception as exc: + return RankRunResult(rank=rank, host=host, rc=999, stdout="", stderr=repr(exc), elapsed_s=time.perf_counter() - start) + return RankRunResult(rank=rank, host=host, rc=rc, stdout=out, stderr=err, elapsed_s=time.perf_counter() - start) + + +def run_all_ranks(args, hosts: Sequence[str], repeat: int, label: str) -> Tuple[List[RankRunResult], float, str]: + port = args.base_port + (int(time.time()) % 1000) + if label != "measure": + port += 1000 + comm_id = f"{hosts[0]}:{port}" + result_dir = f"{args.result_prefix}_{label}_{time.strftime('%Y%m%d_%H%M%S')}" + commands = [] + for rank in range(args.world_size): + host = hosts[rank // args.devices_per_host] + command = make_rank_command( + rank=rank, + world_size=args.world_size, + repeat=repeat, + elements_per_peer=args.elements_per_peer, + devices_per_host=args.devices_per_host, + remote_dir=args.remote_dir, + cann_env=args.cann_env, + demo_bin=args.demo_bin, + comm_id=comm_id, + barrier_host=hosts[0], + result_dir=result_dir, + route_policy=args.route_policy, + timeout_s=args.timeout, + tool_env=args.tool_env, + ) + commands.append((rank, host, command)) + + started = time.perf_counter() + results: List[RankRunResult] = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=args.world_size) as executor: + futures = [ + executor.submit(run_rank, host, args.user, args.password, command, rank, args.timeout) + for rank, host, command in commands + ] + for future in concurrent.futures.as_completed(futures): + result = future.result() + results.append(result) + print(f"[{label}] rank {result.rank:02d} host {result.host} rc={result.rc} elapsed={result.elapsed_s:.2f}s", flush=True) + elapsed = time.perf_counter() - started + return sorted(results, key=lambda result: result.rank), elapsed, result_dir + + +def collect_metrics(results: Sequence[RankRunResult]) -> Tuple[List[RankMetric], List[str]]: + metrics: List[RankMetric] = [] + errors: List[str] = [] + for result in results: + text = f"{result.stdout}\n{result.stderr}" + parsed_metric = None + try: + parsed_metric = parse_demo_metric(text) + metrics.append(parsed_metric) + except ValueError: + pass + if result.rc != 0: + errors.append(f"rank {result.rank} on {result.host} rc={result.rc}: {text[-2000:]}") + continue + for token in ("ERROR", "-1001", "-1002", "-1003"): + if token in text: + errors.append(f"rank {result.rank} on {result.host} contains {token}: {text[-2000:]}") + break + if parsed_metric is None: + errors.append(f"rank {result.rank} on {result.host}: no alltoall udma-bigdata metric line found: {text[-2000:]}") + return metrics, errors + + +def parse_hosts(value: str) -> List[str]: + return [item.strip() for item in value.split(",") if item.strip()] + + +def parse_args(argv: Optional[Sequence[str]] = None): + parser = argparse.ArgumentParser( + description="Sync TileXR and run UDMA bigdata alltoall performance tests.", + epilog=( + "Examples: " + "py tests/udma/perf/run_udma_alltoall_perf.py --world-size 16 --bytes-per-peer 8MiB --warmup 1 --loops 100; " + "py tests/udma/perf/run_udma_alltoall_perf.py --world-size 64 --bytes-per-peer 4MiB --warmup 1 --loops 100" + ), + ) + parser.add_argument("--hosts", default=",".join(DEFAULT_HOSTS), help="Comma-separated host list in rank order.") + parser.add_argument("--world-size", type=int, default=16, help="Total rank count, for example 8, 16, 32, or 64.") + parser.add_argument("--devices-per-host", type=int, default=8, help="NPU ranks per host.") + parser.add_argument("--bytes-per-peer", default="8MiB", help="Payload bytes per peer, for example 4MiB, 8MiB, 16MiB.") + parser.add_argument("--warmup", type=int, default=1, help="Warmup iterations in a separate launch. Use 0 to skip.") + parser.add_argument("--loops", type=int, default=100, help="Measured iterations.") + parser.add_argument("--user", default="root", help="SSH user.") + parser.add_argument("--password", default=os.environ.get("TILEXR_REMOTE_PASSWORD", "Admin@9000"), help="SSH password.") + parser.add_argument("--remote-dir", default="/home/tileXR", help="Remote TileXR source path.") + parser.add_argument("--cann-env", default=DEFAULT_CANN_ENV, help="Remote CANN set_env.sh path.") + parser.add_argument( + "--tool-env", + default=DEFAULT_TOOL_ENV, + help="External TileXR env root containing util/cmake and util/mpich. Defaults to the parent env inferred from --cann-env.", + ) + parser.add_argument("--demo-bin", default=DEFAULT_DEMO_BIN, help="Demo binary path relative to remote-dir.") + parser.add_argument("--route-policy", default="all", help="TILEXR_UDMA_ROUTE_POLICY value.") + parser.add_argument("--timeout", type=int, default=180, help="Kernel/demo timeout seconds.") + parser.add_argument("--build-timeout", type=int, default=900, help="Remote build timeout seconds.") + parser.add_argument("--sync-timeout", type=int, default=300, help="Remote sync timeout seconds.") + parser.add_argument("--base-port", type=int, default=32000, help="Base rendezvous port.") + parser.add_argument("--result-prefix", default="/tmp/tilexr_udma_alltoall_perf", help="Remote result dir prefix.") + parser.add_argument("--skip-sync", action="store_true", help="Do not upload local source tree.") + parser.add_argument("--skip-build", action="store_true", help="Do not build tests/udma on remote hosts.") + parser.add_argument("--clean-remote", action="store_true", help="Remove remote-dir before extracting uploaded source.") + parsed = parser.parse_args(argv) + parsed.hosts = parse_hosts(parsed.hosts) + parsed.bytes_per_peer_value = parse_size_to_bytes(parsed.bytes_per_peer) + parsed.elements_per_peer = bytes_to_int32_elements(parsed.bytes_per_peer_value) + if not parsed.tool_env: + parsed.tool_env = infer_tool_env(parsed.cann_env) + if parsed.warmup < 0 or parsed.loops <= 0: + raise ValueError("--warmup must be >= 0 and --loops must be > 0") + return parsed + + +def print_summary(args, hosts: Sequence[str], metrics: Sequence[RankMetric], host_elapsed_s: float, result_dir: str) -> None: + summary = summarize_metrics(metrics) + total_payload = args.bytes_per_peer_value * args.world_size + remote_payload = args.bytes_per_peer_value * max(args.world_size - args.devices_per_host, 0) + host_per_iter_us = host_elapsed_s * 1_000_000.0 / args.loops + print("") + print("=== TileXR UDMA alltoall performance ===") + print(f"hosts: {','.join(hosts)}") + print(f"world_size: {args.world_size}") + print(f"bytes_per_peer: {args.bytes_per_peer_value} ({args.bytes_per_peer})") + print(f"demo_payload_per_rank: {total_payload} bytes") + print(f"remote_only_payload_per_rank: {remote_payload} bytes") + print(f"warmup_iters: {args.warmup}") + print(f"measured_iters: {args.loops}") + print(f"remote_result_dir: {result_dir}") + print(f"host_wall_time: {host_elapsed_s:.6f} s") + print(f"host_wall_per_iter: {host_per_iter_us:.3f} us") + print( + "operator_per_iter_us: " + f"count={summary['count']} min={summary['min_us']:.3f} p50={summary['p50_us']:.3f} " + f"mean={summary['mean_us']:.3f} p90={summary['p90_us']:.3f} " + f"p99={summary['p99_us']:.3f} max={summary['max_us']:.3f}" + ) + bw_values = [metric.bandwidth_gbps for metric in metrics] + print( + "operator_bw_GBps: " + f"min={min(bw_values):.3f} mean={statistics.mean(bw_values):.3f} max={max(bw_values):.3f}" + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + try: + args = parse_args(argv) + hosts = select_hosts(args.hosts, args.world_size, args.devices_per_host) + except ValueError as exc: + print(f"argument error: {exc}", file=sys.stderr) + return 2 + + archive: Optional[Path] = None + try: + if not args.skip_sync: + archive = create_source_archive(repo_root()) + print(f"sync archive: {archive} ({archive.stat().st_size} bytes)") + with concurrent.futures.ThreadPoolExecutor(max_workers=len(hosts)) as executor: + futures = [ + executor.submit( + sync_host, + host, + args.user, + args.password, + archive, + args.remote_dir, + args.clean_remote, + args.sync_timeout, + ) + for host in hosts + ] + for host, future in zip(hosts, futures): + future.result() + print(f"sync {host} ok") + + if not args.skip_build: + with concurrent.futures.ThreadPoolExecutor(max_workers=len(hosts)) as executor: + futures = [ + executor.submit( + build_host, + host, + args.user, + args.password, + args.remote_dir, + args.cann_env, + args.tool_env, + args.build_timeout, + ) + for host in hosts + ] + for host, future in zip(hosts, futures): + future.result() + print(f"build {host} ok") + + if args.warmup > 0: + warmup_results, warmup_elapsed, warmup_dir = run_all_ranks(args, hosts, args.warmup, "warmup") + _warmup_metrics, warmup_errors = collect_metrics(warmup_results) + if warmup_errors: + if _warmup_metrics: + print_summary(args, hosts, _warmup_metrics, warmup_elapsed, warmup_dir) + print("\n".join(warmup_errors), file=sys.stderr) + return 1 + print(f"warmup ok: host_elapsed={warmup_elapsed:.6f}s result_dir={warmup_dir}") + + measure_results, host_elapsed, result_dir = run_all_ranks(args, hosts, args.loops, "measure") + metrics, errors = collect_metrics(measure_results) + if errors: + if metrics: + print_summary(args, hosts, metrics, host_elapsed, result_dir) + print("\n".join(errors), file=sys.stderr) + return 1 + if len(metrics) != args.world_size: + print(f"expected {args.world_size} metrics, got {len(metrics)}", file=sys.stderr) + return 1 + print_summary(args, hosts, metrics, host_elapsed, result_dir) + return 0 + except Exception as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + finally: + if archive is not None: + try: + archive.unlink() + except OSError: + pass + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/udma/perf/test_run_udma_alltoall_perf.py b/tests/udma/perf/test_run_udma_alltoall_perf.py new file mode 100644 index 00000000..daf3a977 --- /dev/null +++ b/tests/udma/perf/test_run_udma_alltoall_perf.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 + +import importlib.util +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).resolve().with_name("run_udma_alltoall_perf.py") +SPEC = importlib.util.spec_from_file_location("run_udma_alltoall_perf", SCRIPT) +perf = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(perf) + + +class UDMAAllToAllPerfScriptTest(unittest.TestCase): + def test_parse_size_to_bytes(self): + self.assertEqual(perf.parse_size_to_bytes("4096"), 4096) + self.assertEqual(perf.parse_size_to_bytes("4M"), 4 * 1024 * 1024) + self.assertEqual(perf.parse_size_to_bytes("8MiB"), 8 * 1024 * 1024) + self.assertEqual(perf.parse_size_to_bytes("1g"), 1024 * 1024 * 1024) + + def test_elements_per_peer_requires_int32_alignment(self): + self.assertEqual(perf.bytes_to_int32_elements(8 * 1024 * 1024), 2 * 1024 * 1024) + with self.assertRaises(ValueError): + perf.bytes_to_int32_elements(1025) + + def test_select_hosts_for_world_size(self): + hosts = ["h0", "h1", "h2"] + self.assertEqual(perf.select_hosts(hosts, 8, 8), ["h0"]) + self.assertEqual(perf.select_hosts(hosts, 16, 8), ["h0", "h1"]) + with self.assertRaises(ValueError): + perf.select_hosts(hosts, 20, 8) + with self.assertRaises(ValueError): + perf.select_hosts(hosts, 32, 8) + + def test_parse_demo_metrics(self): + log = """ + [rank 0] alltoall udma-bigdata 100 iters(total=100 pass/iter) total=128.6929 ms perIter=1286.929 us payload=2.68435e+08 bytes bw=208.586 GB/s + """ + metric = perf.parse_demo_metric(log) + self.assertEqual(metric.rank, 0) + self.assertEqual(metric.iters, 100) + self.assertAlmostEqual(metric.total_ms, 128.6929) + self.assertAlmostEqual(metric.per_iter_us, 1286.929) + self.assertAlmostEqual(metric.payload_bytes, 2.68435e8) + self.assertAlmostEqual(metric.bandwidth_gbps, 208.586) + + def test_summarize_metrics(self): + metrics = [ + perf.RankMetric(rank=0, iters=10, total_ms=10.0, per_iter_us=1000.0, payload_bytes=1024.0, bandwidth_gbps=1.0), + perf.RankMetric(rank=1, iters=10, total_ms=20.0, per_iter_us=2000.0, payload_bytes=1024.0, bandwidth_gbps=2.0), + perf.RankMetric(rank=2, iters=10, total_ms=30.0, per_iter_us=3000.0, payload_bytes=1024.0, bandwidth_gbps=3.0), + ] + summary = perf.summarize_metrics(metrics) + self.assertEqual(summary["count"], 3) + self.assertEqual(summary["min_us"], 1000.0) + self.assertEqual(summary["p50_us"], 2000.0) + self.assertEqual(summary["max_us"], 3000.0) + self.assertEqual(summary["mean_us"], 2000.0) + + def test_infers_external_tool_env_from_cann_env(self): + tool_env = perf.infer_tool_env("/home/user/tilexr/env/cann/cann/set_env.sh") + self.assertEqual(tool_env, "/home/user/tilexr/env") + prelude = perf.remote_env_prelude("/home/user/tilexr/env/cann/cann/set_env.sh", tool_env) + self.assertIn("/home/user/tilexr/env/util/cmake/bin", prelude) + self.assertIn("source /home/user/tilexr/env/cann/cann/set_env.sh", prelude) + + +if __name__ == "__main__": + unittest.main() From a77406ba06026ecd18df11e5ff806cabc8d0fab2 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 11:53:10 +0800 Subject: [PATCH 028/163] fix(udma): prefer aggregate routes for remote nodes --- src/comm/udma/tilexr_udma_layout.cpp | 14 +++++++ src/comm/udma/tilexr_udma_layout.h | 5 +++ src/comm/udma/tilexr_udma_transport.cpp | 41 +++++++++++++++++-- tests/udma/perf/run_udma_alltoall_perf.py | 6 ++- .../test_tilexr_udma_transport_layout.cpp | 25 +++++++++++ 5 files changed, 86 insertions(+), 5 deletions(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 77b5e5ab..e71751b7 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -157,4 +157,18 @@ std::vector SelectExplicitUDMARouteEids( return selected; } +std::vector SelectUDMARoutesForPeer( + bool peerIsRemoteNode, + const std::vector& topoRoutes, + const std::vector& aggregateRoutes) +{ + if (peerIsRemoteNode && !aggregateRoutes.empty()) { + return aggregateRoutes; + } + if (!topoRoutes.empty()) { + return topoRoutes; + } + return aggregateRoutes; +} + } // namespace TileXR diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index 5e3f4392..c300c620 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -52,6 +52,11 @@ std::vector SelectExplicitUDMARouteEids( const char* routeList, const std::vector& candidateEids); +std::vector SelectUDMARoutesForPeer( + bool peerIsRemoteNode, + const std::vector& topoRoutes, + const std::vector& aggregateRoutes); + } // namespace TileXR #endif // TILEXR_UDMA_LAYOUT_H diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index b27e819b..cad130ec 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -18,6 +19,7 @@ #include #include #include +#include #include "tilexr_log.h" #include "tools/socket/tilexr_sock_exchange.h" @@ -25,6 +27,8 @@ namespace TileXR { namespace { +constexpr size_t kUDMANodeIdBytes = 256; + uint32_t Log2Uint64(uint64_t value) { uint32_t result = 0; @@ -77,6 +81,20 @@ int QueueKey(int peer, uint32_t qpIdx, uint32_t qpNum) return qpNum <= 1 ? peer : static_cast(static_cast(peer) * qpNum + qpIdx); } +struct UDMANodeId { + char value[kUDMANodeIdBytes]; +}; + +UDMANodeId GetLocalNodeId() +{ + UDMANodeId node {}; + if (gethostname(node.value, sizeof(node.value) - 1) != 0 || node.value[0] == '\0') { + std::snprintf(node.value, sizeof(node.value), "pid-%ld", static_cast(getpid())); + } + node.value[sizeof(node.value) - 1] = '\0'; + return node; +} + HccpEid SwapEidForDevice(const HccpEid& hccpEid) { HccpEid swapped {}; @@ -570,6 +588,12 @@ int TileXRUDMATransport::BuildRoutes() if (ret != TILEXR_SUCCESS) { return ret; } + const UDMANodeId localNode = GetLocalNodeId(); + std::vector allNodes(options_.rankSize); + ret = options_.exchange->AllGather(&localNode, 1, allNodes.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } if (diag) { std::ostringstream ids; for (int rank = 0; rank < options_.rankSize; ++rank) { @@ -589,11 +613,13 @@ int TileXRUDMATransport::BuildRoutes() } uint32_t localEid = devEids[0].eidIndex; std::vector localEids; + const bool peerIsRemoteNode = + std::strncmp(localNode.value, allNodes[peer].value, sizeof(localNode.value)) != 0; if (topoReady && useAllRoutes) { - localEids = ResolveLocalEidRoutes(rootInfo, topoEdges, localId, allLocalIds[peer]); - if (localEids.empty()) { - localEids = ResolveLocalAggregateEidRoutes(rootInfo, localId); - } + const std::vector topoRoutes = + ResolveLocalEidRoutes(rootInfo, topoEdges, localId, allLocalIds[peer]); + const std::vector aggregateRoutes = ResolveLocalAggregateEidRoutes(rootInfo, localId); + localEids = SelectUDMARoutesForPeer(peerIsRemoteNode, topoRoutes, aggregateRoutes); const std::vector explicitEids = SelectExplicitUDMARouteEids(std::getenv("TILEXR_UDMA_ROUTE_EIDS"), localEids); if (!explicitEids.empty()) { @@ -624,6 +650,13 @@ int TileXRUDMATransport::BuildRoutes() << " peerLocalId=" << allLocalIds[peer] << " localEid=" << peerLocalEid_[peer] << " routeCount=" << localEids.size(); + if (topoReady && useAllRoutes) { + TILEXR_LOG(INFO) << "UDMA diag route node rank " << options_.rank + << " peer=" << peer + << " remoteNode=" << (peerIsRemoteNode ? 1 : 0) + << " localNodeId=" << localNode.value + << " peerNodeId=" << allNodes[peer].value; + } } } diff --git a/tests/udma/perf/run_udma_alltoall_perf.py b/tests/udma/perf/run_udma_alltoall_perf.py index 75b357e2..5c0c116f 100644 --- a/tests/udma/perf/run_udma_alltoall_perf.py +++ b/tests/udma/perf/run_udma_alltoall_perf.py @@ -207,9 +207,13 @@ def run_ssh(host: str, user: str, password: str, command: str, timeout_s: int) - out = stdout.read().decode("utf-8", "replace") err = stderr.read().decode("utf-8", "replace") rc = stdout.channel.recv_exit_status() - stdin.close() stdout.close() stderr.close() + try: + stdin.channel.shutdown_write() + except Exception: + pass + stdin.close() return rc, out, err finally: client.close() diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 3adb0979..eea1254f 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -171,6 +171,29 @@ void TestExplicitRouteSelectionRejectsMissingInputs() CHECK_TRUE(TileXR::SelectExplicitUDMARouteEids("8", {}).empty()); } +void TestCrossNodeRouteSelectionUsesAggregateRoutes() +{ + const std::vector topoRoutes = {1}; + const std::vector aggregateRoutes = {7, 8}; + const std::vector selected = + TileXR::SelectUDMARoutesForPeer(true, topoRoutes, aggregateRoutes); + + CHECK_EQ(selected.size(), static_cast(2)); + CHECK_EQ(selected[0], 7U); + CHECK_EQ(selected[1], 8U); +} + +void TestSameNodeRouteSelectionUsesTopoRoutes() +{ + const std::vector topoRoutes = {1}; + const std::vector aggregateRoutes = {7, 8}; + const std::vector selected = + TileXR::SelectUDMARoutesForPeer(false, topoRoutes, aggregateRoutes); + + CHECK_EQ(selected.size(), static_cast(1)); + CHECK_EQ(selected[0], 1U); +} + void TestTransportUsesPerPeerQueues() { const std::string transport = @@ -252,6 +275,8 @@ int main() TestMultiRouteQpWeightsUseRouteBandwidth(); TestExplicitRouteSelectionKeepsRequestedCandidateOrder(); TestExplicitRouteSelectionRejectsMissingInputs(); + TestCrossNodeRouteSelectionUsesAggregateRoutes(); + TestSameNodeRouteSelectionUsesTopoRoutes(); TestTransportUsesPerPeerQueues(); TestRootInfoEidBytesSelectRuntimeContexts(); TestMemoryRegistrationUsesOfficialUbFlags(); From bbfc4f4a79d06e63dda374abadbde3e659108f9d Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 12:20:26 +0800 Subject: [PATCH 029/163] fix(udma): honor ranks per node in bigdata demo --- tests/udma/demo/tilexr_udma_alltoall_layout.h | 128 ++++++++----- tests/udma/demo/tilexr_udma_demo.cpp | 25 ++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 172 +++++++++++------- tests/udma/perf/run_udma_alltoall_perf.py | 16 +- .../unit/test_tilexr_udma_alltoall_layout.cpp | 58 ++++-- 5 files changed, 261 insertions(+), 138 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index ee35bfca..60a698d4 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -65,20 +65,33 @@ inline int32_t AllToAllValue(int srcRank, int dstRank) return kAllToAllBaseValue + srcRank * 1000 + dstRank; } -inline bool AllToAllBigDataIsMultiNode(int rankSize) +inline int32_t AllToAllBigDataNormalizeRanksPerNode(int rankSize, int32_t ranksPerNode) { - return rankSize > kAllToAllBigDataRanksPerNode; + (void)rankSize; + if (ranksPerNode <= 0) { + return kAllToAllBigDataRanksPerNode; + } + return ranksPerNode; +} + +inline bool AllToAllBigDataIsMultiNode( + int rankSize, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) +{ + return rankSize > AllToAllBigDataNormalizeRanksPerNode(rankSize, ranksPerNode); } -inline bool AllToAllBigDataUse35Core(int rankSize, bool force35Core = false) +inline bool AllToAllBigDataUse35Core( + int rankSize, bool force35Core = false, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - return AllToAllBigDataIsMultiNode(rankSize) || - (force35Core && rankSize == kAllToAllBigDataRanksPerNode); + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rankSize, ranksPerNode); + return AllToAllBigDataIsMultiNode(rankSize, localRanks) || + (force35Core && rankSize == localRanks); } -inline uint32_t AllToAllBigDataShardCount(int rankSize, bool force35Core = false) +inline uint32_t AllToAllBigDataShardCount( + int rankSize, bool force35Core = false, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - return AllToAllBigDataUse35Core(rankSize, force35Core) ? + return AllToAllBigDataUse35Core(rankSize, force35Core, ranksPerNode) ? kAllToAllBigDataMultiNodeControlShards : kAllToAllBigDataSingleNodeShards; } @@ -113,17 +126,18 @@ inline AllToAllChunkPlan PlanAllToAllUdmaChunks(int rankSize, int32_t elementsPe } inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( - int rankSize, int32_t elementsPerPeer, bool force35Core = false) + int rankSize, int32_t elementsPerPeer, bool force35Core = false, + int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { AllToAllBigDataPlan plan {}; if (rankSize <= 0 || elementsPerPeer <= 0) { return plan; } - const bool use35Core = AllToAllBigDataUse35Core(rankSize, force35Core); + const bool use35Core = AllToAllBigDataUse35Core(rankSize, force35Core, ranksPerNode); plan.registeredBytes = use35Core ? kAllToAllBigDataMultiNodeRegisteredBytes : kAllToAllBigDataMaxRegisteredBytes; - const uint32_t shardCount = AllToAllBigDataShardCount(rankSize, force35Core); + const uint32_t shardCount = AllToAllBigDataShardCount(rankSize, force35Core, ranksPerNode); const size_t maxChunkElements = use35Core ? kAllToAllBigDataMultiNodePeerSlotBytes / sizeof(int32_t) : static_cast(elementsPerPeer); @@ -187,90 +201,113 @@ inline AllToAllBigDataPlan PlanAllToAllBigDataUdma( return plan; } -inline bool AllToAllBigDataValidTopology(int rankSize) +inline bool AllToAllBigDataValidTopology( + int rankSize, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { if (rankSize <= 0) { return false; } - if (!AllToAllBigDataIsMultiNode(rankSize)) { + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rankSize, ranksPerNode); + if (localRanks <= 0) { + return false; + } + if (!AllToAllBigDataIsMultiNode(rankSize, localRanks)) { return true; } - return rankSize % kAllToAllBigDataRanksPerNode == 0; + return rankSize % localRanks == 0; } -inline int32_t AllToAllBigDataLocalNodeBegin(int rank) +inline int32_t AllToAllBigDataLocalNodeBegin( + int rank, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - return (rank / kAllToAllBigDataRanksPerNode) * kAllToAllBigDataRanksPerNode; + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rank + 1, ranksPerNode); + return (rank / localRanks) * localRanks; } -inline int32_t AllToAllBigDataLocalNodeEnd(int rank) +inline int32_t AllToAllBigDataLocalNodeEnd( + int rank, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - return AllToAllBigDataLocalNodeBegin(rank) + kAllToAllBigDataRanksPerNode; + return AllToAllBigDataLocalNodeBegin(rank, ranksPerNode) + + AllToAllBigDataNormalizeRanksPerNode(rank + 1, ranksPerNode); } -inline int32_t AllToAllBigDataNodeCount(int rankSize) +inline int32_t AllToAllBigDataNodeCount( + int rankSize, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - if (!AllToAllBigDataValidTopology(rankSize)) { + if (!AllToAllBigDataValidTopology(rankSize, ranksPerNode)) { return 0; } - if (!AllToAllBigDataIsMultiNode(rankSize)) { + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rankSize, ranksPerNode); + if (!AllToAllBigDataIsMultiNode(rankSize, localRanks)) { return 1; } - return rankSize / kAllToAllBigDataRanksPerNode; + return rankSize / localRanks; } -inline bool AllToAllBigDataIsLocalPeer(int rank, int peer) +inline bool AllToAllBigDataIsLocalPeer( + int rank, int peer, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - const int32_t begin = AllToAllBigDataLocalNodeBegin(rank); - return peer >= begin && peer < begin + kAllToAllBigDataRanksPerNode; + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(std::max(rank, peer) + 1, ranksPerNode); + const int32_t begin = AllToAllBigDataLocalNodeBegin(rank, localRanks); + return peer >= begin && peer < begin + localRanks; } -inline int32_t AllToAllBigDataLocalPeerForWorker(int rank, uint32_t workerGroup) +inline int32_t AllToAllBigDataLocalPeerForWorker( + int rank, uint32_t workerGroup, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { - if (workerGroup >= static_cast(kAllToAllBigDataRanksPerNode)) { + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rank + 1, ranksPerNode); + if (workerGroup >= static_cast(localRanks)) { return -1; } - return AllToAllBigDataLocalNodeBegin(rank) + static_cast(workerGroup); + return AllToAllBigDataLocalNodeBegin(rank, localRanks) + static_cast(workerGroup); } -inline std::vector AllToAllBigDataRemotePeers(int rank, int rankSize) +inline std::vector AllToAllBigDataRemotePeers( + int rank, int rankSize, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { std::vector peers; - if (!AllToAllBigDataIsMultiNode(rankSize) || !AllToAllBigDataValidTopology(rankSize)) { + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rankSize, ranksPerNode); + if (!AllToAllBigDataIsMultiNode(rankSize, localRanks) || + !AllToAllBigDataValidTopology(rankSize, localRanks)) { return peers; } - peers.reserve(static_cast(rankSize - kAllToAllBigDataRanksPerNode)); + peers.reserve(static_cast(rankSize - localRanks)); for (int32_t step = 0; step < rankSize && - peers.size() < static_cast(rankSize - kAllToAllBigDataRanksPerNode); ++step) { - const int32_t peer = (rank + kAllToAllBigDataRanksPerNode + step) % rankSize; - if (!AllToAllBigDataIsLocalPeer(rank, peer)) { + peers.size() < static_cast(rankSize - localRanks); ++step) { + const int32_t peer = (rank + localRanks + step) % rankSize; + if (!AllToAllBigDataIsLocalPeer(rank, peer, localRanks)) { peers.push_back(peer); } } return peers; } -inline std::vector AllToAllBigDataLocalPeers(int rank) +inline std::vector AllToAllBigDataLocalPeers( + int rank, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { std::vector peers; - peers.reserve(static_cast(kAllToAllBigDataRanksPerNode - 1)); - const int32_t begin = AllToAllBigDataLocalNodeBegin(rank); + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rank + 1, ranksPerNode); + peers.reserve(static_cast(std::max(localRanks - 1, 0))); + const int32_t begin = AllToAllBigDataLocalNodeBegin(rank, localRanks); const int32_t local = rank - begin; - for (int32_t step = 1; step < kAllToAllBigDataRanksPerNode; ++step) { - peers.push_back(begin + ((local + step) % kAllToAllBigDataRanksPerNode)); + for (int32_t step = 1; step < localRanks; ++step) { + peers.push_back(begin + ((local + step) % localRanks)); } return peers; } -inline std::vector AllToAllBigDataMergedPeerTasks(int rank, int rankSize) +inline std::vector AllToAllBigDataMergedPeerTasks( + int rank, int rankSize, int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { std::vector tasks; - if (!AllToAllBigDataIsMultiNode(rankSize) || !AllToAllBigDataValidTopology(rankSize)) { + const int32_t localRanks = AllToAllBigDataNormalizeRanksPerNode(rankSize, ranksPerNode); + if (!AllToAllBigDataIsMultiNode(rankSize, localRanks) || + !AllToAllBigDataValidTopology(rankSize, localRanks)) { return tasks; } - const std::vector remotePeers = AllToAllBigDataRemotePeers(rank, rankSize); - const std::vector localPeers = AllToAllBigDataLocalPeers(rank); - const int32_t remoteBurst = AllToAllBigDataNodeCount(rankSize) - 1; + const std::vector remotePeers = AllToAllBigDataRemotePeers(rank, rankSize, localRanks); + const std::vector localPeers = AllToAllBigDataLocalPeers(rank, localRanks); + const int32_t remoteBurst = AllToAllBigDataNodeCount(rankSize, localRanks) - 1; size_t remoteIndex = 0; size_t localIndex = 0; tasks.reserve(remotePeers.size() + localPeers.size()); @@ -286,12 +323,13 @@ inline std::vector AllToAllBigDataMergedPeerTasks(int rank, int rankSiz } inline uint32_t AllToAllBigDataBlockDim( - int rankSize, bool force35Core = false, bool remotePutOnly = false) + int rankSize, bool force35Core = false, bool remotePutOnly = false, + int32_t ranksPerNode = kAllToAllBigDataRanksPerNode) { if (rankSize <= 0) { return 1U; } - if (AllToAllBigDataUse35Core(rankSize, force35Core)) { + if (AllToAllBigDataUse35Core(rankSize, force35Core, ranksPerNode)) { if (remotePutOnly) { return kAllToAllBigDataRemotePutOnlyBlockDim; } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 020d1423..3d36ba4e 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -698,15 +698,17 @@ int main(int argc, char** argv) const TileXR::Demo::AllToAllChunkPlan chunkPlan = isAllToAll ? TileXR::Demo::PlanAllToAllUdmaChunks(rankSize, elementsPerRank) : TileXR::Demo::AllToAllChunkPlan {}; + const int32_t bigDataRanksPerNode = std::max(1, npuCount); const TileXR::Demo::AllToAllBigDataPlan bigDataPlan = isAllToAll ? TileXR::Demo::PlanAllToAllBigDataUdma( - rankSize, elementsPerRank, forceBigData35Core) : + rankSize, elementsPerRank, forceBigData35Core, bigDataRanksPerNode) : TileXR::Demo::AllToAllBigDataPlan {}; - if (testType == 7 && !TileXR::Demo::AllToAllBigDataValidTopology(rankSize)) { + if (testType == 7 && !TileXR::Demo::AllToAllBigDataValidTopology(rankSize, bigDataRanksPerNode)) { std::cerr << "[rank " << rank - << "] ERROR: bigdata alltoall multi-node requires rankSize multiple of 8" - << " when rankSize > " << TileXR::Demo::kAllToAllBigDataRanksPerNode - << ", rankSize=" << rankSize << std::endl; + << "] ERROR: bigdata alltoall multi-node requires rankSize multiple of ranksPerNode" + << " when rankSize > ranksPerNode" + << ", rankSize=" << rankSize + << " ranksPerNode=" << bigDataRanksPerNode << std::endl; Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } @@ -739,13 +741,15 @@ int main(int argc, char** argv) PrintStatus(rank, "bigdata profile stage=" + std::to_string(bigDataProfileStage) + " fullStage=" + std::to_string(kBigDataProfileStageFull)); PrintStatus(rank, "bigdata multinode mode=" + - std::string(TileXR::Demo::AllToAllBigDataIsMultiNode(rankSize) ? "true" : "false") + + std::string(TileXR::Demo::AllToAllBigDataIsMultiNode( + rankSize, bigDataRanksPerNode) ? "true" : "false") + " force35Core=" + std::string(forceBigData35Core ? "true" : "false") + " remotePutOnly=" + std::string(bigDataRemotePutOnly ? "true" : "false") + + " ranksPerNode=" + std::to_string(bigDataRanksPerNode) + " blockDim=" + std::to_string(TileXR::Demo::AllToAllBigDataBlockDim( - rankSize, forceBigData35Core, bigDataRemotePutOnly)) + + rankSize, forceBigData35Core, bigDataRemotePutOnly, bigDataRanksPerNode)) + " shards=" + std::to_string(TileXR::Demo::AllToAllBigDataShardCount( - rankSize, forceBigData35Core))); + rankSize, forceBigData35Core, bigDataRanksPerNode))); } PrintStatus(rank, "alltoall UDMA chunk plan: passCount=" + std::to_string(chunkPlan.passCount) + " chunkElements=" + std::to_string(chunkPlan.chunkElements) + @@ -953,9 +957,10 @@ int main(int argc, char** argv) } const uint32_t bigDataBlockDim = TileXR::Demo::AllToAllBigDataBlockDim( - rankSize, forceBigData35Core, bigDataRemotePutOnly); + rankSize, forceBigData35Core, bigDataRemotePutOnly, bigDataRanksPerNode); const uint32_t bigDataModeFlags = - (forceBigData35Core ? 1U : 0U) | (bigDataRemotePutOnly ? 2U : 0U); + (forceBigData35Core ? 1U : 0U) | (bigDataRemotePutOnly ? 2U : 0U) | + (static_cast(bigDataRanksPerNode) << 8U); auto a2aStart = std::chrono::steady_clock::now(); for (int iter = 0; iter < allToAllRepeat; ++iter) { const uint64_t kernelLoopBase = static_cast(iter); diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index cba94024..dd0541d7 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -445,76 +445,93 @@ __aicore__ inline void BigDataKernelExitBarrier() AscendC::SyncAll(); } -__aicore__ inline bool BigDataIsMultiNode(int32_t rankSize) +__aicore__ inline int32_t BigDataNormalizeRanksPerNode(int32_t ranksPerNode) { - return rankSize > TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + return ranksPerNode > 0 ? ranksPerNode : TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; } -__aicore__ inline bool BigDataUse35Core(int32_t rankSize, bool force35Core) +__aicore__ inline bool BigDataIsMultiNode(int32_t rankSize, int32_t ranksPerNode) { - return BigDataIsMultiNode(rankSize) || - (force35Core && rankSize == TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE); + return rankSize > BigDataNormalizeRanksPerNode(ranksPerNode); } -__aicore__ inline uint32_t BigDataShardCount(int32_t rankSize, bool force35Core = false) +__aicore__ inline bool BigDataUse35Core(int32_t rankSize, bool force35Core, int32_t ranksPerNode) { - return BigDataUse35Core(rankSize, force35Core) ? + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + return BigDataIsMultiNode(rankSize, localRanks) || + (force35Core && rankSize == localRanks); +} + +__aicore__ inline uint32_t BigDataShardCount( + int32_t rankSize, bool force35Core = false, int32_t ranksPerNode = TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) +{ + return BigDataUse35Core(rankSize, force35Core, ranksPerNode) ? TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS : TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS; } -__aicore__ inline bool BigDataValidTopology(int32_t rankSize) +__aicore__ inline bool BigDataValidTopology(int32_t rankSize, int32_t ranksPerNode) { if (rankSize <= 0) { return false; } - if (!BigDataIsMultiNode(rankSize)) { + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (localRanks <= 0) { + return false; + } + if (!BigDataIsMultiNode(rankSize, localRanks)) { return true; } - return rankSize % TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE == 0; + return rankSize % localRanks == 0; } -__aicore__ inline int32_t BigDataLocalNodeBegin(int32_t rank) +__aicore__ inline int32_t BigDataLocalNodeBegin(int32_t rank, int32_t ranksPerNode) { - return (rank / TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) * - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + return (rank / localRanks) * localRanks; } -__aicore__ inline int32_t BigDataNodeCount(int32_t rankSize) +__aicore__ inline int32_t BigDataNodeCount(int32_t rankSize, int32_t ranksPerNode) { - if (!BigDataValidTopology(rankSize)) { + if (!BigDataValidTopology(rankSize, ranksPerNode)) { return 0; } - if (!BigDataIsMultiNode(rankSize)) { + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (!BigDataIsMultiNode(rankSize, localRanks)) { return 1; } - return rankSize / TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + return rankSize / localRanks; } -__aicore__ inline int32_t BigDataTaskCount(int32_t rankSize, bool force35Core) +__aicore__ inline int32_t BigDataTaskCount(int32_t rankSize, bool force35Core, int32_t ranksPerNode) { - if (force35Core && rankSize == TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) { - return TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - 1; + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (force35Core && rankSize == localRanks) { + return localRanks - 1; } return rankSize - 1; } -__aicore__ inline bool BigDataIsLocalPeer(int32_t rank, int32_t peer) +__aicore__ inline bool BigDataIsLocalPeer(int32_t rank, int32_t peer, int32_t ranksPerNode) { - const int32_t begin = BigDataLocalNodeBegin(rank); - return peer >= begin && peer < begin + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + const int32_t begin = BigDataLocalNodeBegin(rank, localRanks); + return peer >= begin && peer < begin + localRanks; } -__aicore__ inline int32_t BigDataRemotePeerAt(int32_t rank, int32_t rankSize, int32_t remoteIndex) +__aicore__ inline int32_t BigDataRemotePeerAt( + int32_t rank, int32_t rankSize, int32_t remoteIndex, int32_t ranksPerNode) { - if (!BigDataIsMultiNode(rankSize) || !BigDataValidTopology(rankSize) || remoteIndex < 0) { + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (!BigDataIsMultiNode(rankSize, localRanks) || !BigDataValidTopology(rankSize, localRanks) || + remoteIndex < 0) { return -1; } int32_t remoteCount = 0; for (int32_t step = 0; step < rankSize; ++step) { const int32_t peer = - (rank + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE + step) % rankSize; - if (!BigDataIsLocalPeer(rank, peer)) { + (rank + localRanks + step) % rankSize; + if (!BigDataIsLocalPeer(rank, peer, localRanks)) { if (remoteCount == remoteIndex) { return peer; } @@ -524,16 +541,19 @@ __aicore__ inline int32_t BigDataRemotePeerAt(int32_t rank, int32_t rankSize, in return -1; } -__aicore__ inline int32_t BigDataRemotePeerForwardAt(int32_t rank, int32_t rankSize, int32_t remoteIndex) +__aicore__ inline int32_t BigDataRemotePeerForwardAt( + int32_t rank, int32_t rankSize, int32_t remoteIndex, int32_t ranksPerNode) { - if (!BigDataIsMultiNode(rankSize) || !BigDataValidTopology(rankSize) || remoteIndex < 0) { + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (!BigDataIsMultiNode(rankSize, localRanks) || !BigDataValidTopology(rankSize, localRanks) || + remoteIndex < 0) { return -1; } int32_t remoteCount = 0; for (int32_t step = 0; step < rankSize; ++step) { const int32_t peer = - (rank + TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE + step) % rankSize; - if (!BigDataIsLocalPeer(rank, peer)) { + (rank + localRanks + step) % rankSize; + if (!BigDataIsLocalPeer(rank, peer, localRanks)) { if (remoteCount == remoteIndex) { return peer; } @@ -543,16 +563,19 @@ __aicore__ inline int32_t BigDataRemotePeerForwardAt(int32_t rank, int32_t rankS return -1; } -__aicore__ inline int32_t BigDataRemotePeerReverseAt(int32_t rank, int32_t rankSize, int32_t remoteIndex) +__aicore__ inline int32_t BigDataRemotePeerReverseAt( + int32_t rank, int32_t rankSize, int32_t remoteIndex, int32_t ranksPerNode) { - if (!BigDataIsMultiNode(rankSize) || !BigDataValidTopology(rankSize) || remoteIndex < 0) { + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (!BigDataIsMultiNode(rankSize, localRanks) || !BigDataValidTopology(rankSize, localRanks) || + remoteIndex < 0) { return -1; } int32_t remoteCount = 0; for (int32_t step = 0; step < rankSize; ++step) { const int32_t peer = - (rank - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - step + rankSize * 2) % rankSize; - if (!BigDataIsLocalPeer(rank, peer)) { + (rank - localRanks - step + rankSize * 2) % rankSize; + if (!BigDataIsLocalPeer(rank, peer, localRanks)) { if (remoteCount == remoteIndex) { return peer; } @@ -602,28 +625,41 @@ __aicore__ inline uint32_t BigDataSelectDistinctWeightedQp( return selected; } -__aicore__ inline int32_t BigDataLocalPeerAt(int32_t rank, int32_t localIndex) +__aicore__ inline uint32_t BigDataRemotePutOnlySegmentQp( + const __gm__ TileXR::CommArgs* args, uint32_t segmentId) { - if (localIndex < 0 || localIndex >= TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE - 1) { + auto udmaInfo = TileXR::GetUDMAInfo(args); + const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; + if (qpCount <= 1U) { + return 0U; + } + return segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? 0U : 1U; +} + +__aicore__ inline int32_t BigDataLocalPeerAt(int32_t rank, int32_t localIndex, int32_t ranksPerNode) +{ + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (localIndex < 0 || localIndex >= localRanks - 1) { return -1; } - const int32_t begin = BigDataLocalNodeBegin(rank); + const int32_t begin = BigDataLocalNodeBegin(rank, localRanks); const int32_t local = rank - begin; - return begin + ((local + localIndex + 1) % TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE); + return begin + ((local + localIndex + 1) % localRanks); } __aicore__ inline bool BigDataMergedPeerTaskAt( int32_t rank, int32_t rankSize, int32_t taskIndex, bool force35Core, - int32_t& peer, bool& isLocalPeer) + int32_t ranksPerNode, int32_t& peer, bool& isLocalPeer) { peer = -1; isLocalPeer = false; - if (force35Core && rankSize == TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE) { - peer = BigDataLocalPeerAt(rank, taskIndex); + const int32_t localRanks = BigDataNormalizeRanksPerNode(ranksPerNode); + if (force35Core && rankSize == localRanks) { + peer = BigDataLocalPeerAt(rank, taskIndex, localRanks); isLocalPeer = true; return peer >= 0; } - const int32_t nodeCount = BigDataNodeCount(rankSize); + const int32_t nodeCount = BigDataNodeCount(rankSize, localRanks); if (nodeCount <= 1 || taskIndex < 0) { return false; } @@ -633,11 +669,11 @@ __aicore__ inline bool BigDataMergedPeerTaskAt( const int32_t indexInGroup = taskIndex % groupSize; if (indexInGroup < remoteBurst) { const int32_t remoteIndex = group * remoteBurst + indexInGroup; - peer = BigDataRemotePeerAt(rank, rankSize, remoteIndex); + peer = BigDataRemotePeerAt(rank, rankSize, remoteIndex, localRanks); isLocalPeer = false; return peer >= 0; } - peer = BigDataLocalPeerAt(rank, group); + peer = BigDataLocalPeerAt(rank, group, localRanks); isLocalPeer = true; return peer >= 0; } @@ -1282,9 +1318,9 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, uint64_t sendDataOffset, uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t chunkBytesPerPeer, - AscendC::LocalTensor relayLocal) + int32_t ranksPerNode, AscendC::LocalTensor relayLocal) { - if (peer < 0 || peer >= rankSize || peer == rank || !BigDataIsMultiNode(rankSize) || + if (peer < 0 || peer >= rankSize || peer == rank || !BigDataIsMultiNode(rankSize, ranksPerNode) || shardCount == 0U || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { return; } @@ -1378,10 +1414,10 @@ __aicore__ inline void BigDataRemotePutOnlySendWorker( __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, uint64_t recvDataOffset, - uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) + uint64_t chunkBytesPerPeer, int32_t ranksPerNode, AscendC::LocalTensor relayLocal) { - if (peer < 0 || peer >= rankSize || peer == rank || BigDataIsLocalPeer(rank, peer) || - !BigDataIsMultiNode(rankSize) || shardCount == 0U || + if (peer < 0 || peer >= rankSize || peer == rank || BigDataIsLocalPeer(rank, peer, ranksPerNode) || + !BigDataIsMultiNode(rankSize, ranksPerNode) || shardCount == 0U || profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { return; } @@ -1472,10 +1508,11 @@ __aicore__ inline void BigDataRemotePutOnlyCheckWorker( __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint64_t recvDataOffset, uint64_t ackSignalOffset, - uint64_t chunkBytesPerPeer, uint32_t shardCount, AscendC::LocalTensor relayLocal) + uint64_t chunkBytesPerPeer, uint32_t shardCount, int32_t ranksPerNode, + AscendC::LocalTensor relayLocal) { - const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex); - if (peer < 0 || peer >= rankSize || peer == rank || BigDataIsLocalPeer(rank, peer) || + const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode); + if (peer < 0 || peer >= rankSize || peer == rank || BigDataIsLocalPeer(rank, peer, ranksPerNode) || shardCount == 0U) { return; } @@ -2702,6 +2739,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( const int32_t blockIdx = AscendC::GetBlockIdx(); const bool force35Core = (force35CoreFlag & 0x1U) != 0U; const bool remotePutOnly = (force35CoreFlag & 0x2U) != 0U; + const int32_t ranksPerNode = BigDataNormalizeRanksPerNode(static_cast(force35CoreFlag >> 8U)); if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; @@ -2715,15 +2753,15 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( return; } - if (!BigDataValidTopology(rankSize)) { + if (!BigDataValidTopology(rankSize, ranksPerNode)) { if (blockIdx == 0 && debug != nullptr) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE] = TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS; } return; } - const bool use35Core = BigDataUse35Core(rankSize, force35Core); - const uint32_t shardCount = BigDataShardCount(rankSize, force35Core); + const bool use35Core = BigDataUse35Core(rankSize, force35Core, ranksPerNode); + const uint32_t shardCount = BigDataShardCount(rankSize, force35Core, ranksPerNode); const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; const uint64_t chunkBytesPerPeer = BigDataChunkBytesPerPeer(use35Core, effectiveChunkElements); const uint64_t sendDataOffset = dataOffset; @@ -2781,8 +2819,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( const uint32_t copyShard = static_cast(blockIdx); const uint32_t recvShard = static_cast(blockIdx) - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; - const int32_t taskCount = BigDataTaskCount(rankSize, force35Core); - const int32_t remoteTaskCount = rankSize - TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE; + const int32_t taskCount = BigDataTaskCount(rankSize, force35Core, ranksPerNode); + const int32_t remoteTaskCount = rankSize - ranksPerNode; const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount); if (remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK) { @@ -2799,15 +2837,12 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( continue; } const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount); - const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex); + const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode); if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { continue; } const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount); - const uint32_t primaryQp = BigDataSelectWeightedQp(args, peer, true); - const uint32_t qpIdx = - segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? - primaryQp : BigDataSelectDistinctWeightedQp(args, peer, primaryQp, false); + const uint32_t qpIdx = BigDataRemotePutOnlySegmentQp(args, segmentId); if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { continue; } @@ -2815,7 +2850,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( rank, rankSize, args, input, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, recvDataOffset, - chunkBytesPerPeer, relayLocal); + chunkBytesPerPeer, ranksPerNode, relayLocal); } BigDataKernelExitBarrier(); if (profileStage > TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST) { @@ -2825,7 +2860,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( remotePutOnlyCheckIndex, rank, rankSize, args, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, recvDataOffset, ackSignalOffset, chunkBytesPerPeer, - shardCount, relayLocal); + shardCount, ranksPerNode, relayLocal); } } BigDataKernelExitBarrier(); @@ -2840,7 +2875,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( for (int32_t taskIndex = 0; taskIndex < taskCount; ++taskIndex) { int32_t peer = -1; bool isLocalPeer = false; - if (!BigDataMergedPeerTaskAt(rank, rankSize, taskIndex, force35Core, peer, isLocalPeer)) { + if (!BigDataMergedPeerTaskAt( + rank, rankSize, taskIndex, force35Core, ranksPerNode, peer, isLocalPeer)) { continue; } @@ -2889,7 +2925,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, remoteSendDoneOffset, readySignalOffset, - chunkBytesPerPeer, relayLocal); + chunkBytesPerPeer, ranksPerNode, relayLocal); } if (!isLocalPeer && isRemoteSendSecondaryCore) { BigDataRemoteSendSegmentWorker(peer, @@ -2898,7 +2934,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, remoteSendDoneOffset, readySignalOffset, - chunkBytesPerPeer, relayLocal); + chunkBytesPerPeer, ranksPerNode, relayLocal); } if (isLocalPeer && isLocalSendCore) { BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, diff --git a/tests/udma/perf/run_udma_alltoall_perf.py b/tests/udma/perf/run_udma_alltoall_perf.py index 5c0c116f..9ac49513 100644 --- a/tests/udma/perf/run_udma_alltoall_perf.py +++ b/tests/udma/perf/run_udma_alltoall_perf.py @@ -284,6 +284,13 @@ def build_host(host: str, user: str, password: str, remote_dir: str, cann_env: s remote_env_prelude(cann_env, tool_env), f"cd {shlex.quote(remote_dir)}", "command -v cmake >/dev/null", + "rm -rf build", + "mkdir -p build", + "cd build", + "cmake -DCMAKE_INSTALL_PREFIX=../install ..", + "make -j$(nproc)", + "make install", + f"cd {shlex.quote(remote_dir)}", "bash tests/udma/build.sh", ]) rc, out, err = run_ssh(host, user, password, f"bash -lc {shlex.quote(command)}", timeout_s) @@ -305,6 +312,8 @@ def make_rank_command( barrier_host: str, result_dir: str, route_policy: str, + profile_stage: int, + udma_debug: bool, timeout_s: int, tool_env: str, ) -> str: @@ -317,8 +326,9 @@ def make_rank_command( f"export TILEXR_COMM_ID={shlex.quote(comm_id)}", f"export TILEXR_DEMO_BARRIER_HOST={shlex.quote(barrier_host)}", f"export TILEXR_UDMA_ROUTE_POLICY={shlex.quote(route_policy)}", + f"export TILEXR_UDMA_DEBUG={1 if udma_debug else 0}", "export TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY=1", - "export TILEXR_DEMO_BIGDATA_PROFILE_STAGE=7", + f"export TILEXR_DEMO_BIGDATA_PROFILE_STAGE={profile_stage}", f"export TILEXR_DEMO_ALLTOALL_REPEAT={repeat}", "export TILEXR_DEMO_ALLTOALL_WARMUP=0", "export TILEXR_DEMO_ALLTOALL_SYNC_AT_END=1", @@ -372,6 +382,8 @@ def run_all_ranks(args, hosts: Sequence[str], repeat: int, label: str) -> Tuple[ barrier_host=hosts[0], result_dir=result_dir, route_policy=args.route_policy, + profile_stage=args.profile_stage, + udma_debug=args.udma_debug, timeout_s=args.timeout, tool_env=args.tool_env, ) @@ -445,6 +457,8 @@ def parse_args(argv: Optional[Sequence[str]] = None): ) parser.add_argument("--demo-bin", default=DEFAULT_DEMO_BIN, help="Demo binary path relative to remote-dir.") parser.add_argument("--route-policy", default="all", help="TILEXR_UDMA_ROUTE_POLICY value.") + parser.add_argument("--profile-stage", type=int, default=7, help="TILEXR_DEMO_BIGDATA_PROFILE_STAGE value.") + parser.add_argument("--udma-debug", action="store_true", help="Enable TILEXR_UDMA_DEBUG on remote ranks.") parser.add_argument("--timeout", type=int, default=180, help="Kernel/demo timeout seconds.") parser.add_argument("--build-timeout", type=int, default=900, help="Remote build timeout seconds.") parser.add_argument("--sync-timeout", type=int, default=300, help="Remote sync timeout seconds.") diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index b53f8880..539051a1 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -279,18 +279,27 @@ void TestAllToAllBigDataMultiNodeTopology() { CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(8), false); CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(16), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(2, 1), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsMultiNode(4, 2), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(8), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(16), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(24), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(10), false); + CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(2, 1), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataValidTopology(4, 2), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeBegin(0), 0); CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeEnd(0), 8); CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeBegin(10), 8); CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeEnd(10), 16); + CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeBegin(2, 2), 2); + CHECK_EQ(TileXR::Demo::AllToAllBigDataLocalNodeEnd(2, 2), 4); CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(10, 8), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(10, 15), true); CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(10, 7), false); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(0, 1, 1), false); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(2, 3, 2), true); + CHECK_EQ(TileXR::Demo::AllToAllBigDataIsLocalPeer(2, 1, 2), false); } void TestAllToAllBigDataRemotePeerQueue() @@ -309,6 +318,14 @@ void TestAllToAllBigDataRemotePeerQueue() peers = TileXR::Demo::AllToAllBigDataRemotePeers(5, 8); CHECK_EQ(peers.empty(), true); + + peers = TileXR::Demo::AllToAllBigDataRemotePeers(0, 2, 1); + expected = {1}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataRemotePeers(1, 4, 2); + expected = {3, 2}; + CHECK_EQ(peers == expected, true); } void TestAllToAllBigDataMergedPeerQueue() @@ -329,6 +346,14 @@ void TestAllToAllBigDataMergedPeerQueue() peers = TileXR::Demo::AllToAllBigDataMergedPeerTasks(5, 8); CHECK_EQ(peers.empty(), true); + + peers = TileXR::Demo::AllToAllBigDataMergedPeerTasks(0, 2, 1); + expected = {1}; + CHECK_EQ(peers == expected, true); + + peers = TileXR::Demo::AllToAllBigDataMergedPeerTasks(1, 4, 2); + expected = {3, 0, 2}; + CHECK_EQ(peers == expected, true); } void TestDemoDebugLayoutSource() @@ -431,19 +456,21 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "bigDataPlan.copyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.recvCopyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.remoteSendDoneOffset"); - CHECK_CONTAINS(demo, "AllToAllBigDataValidTopology(rankSize)"); - CHECK_CONTAINS(demo, "ERROR: bigdata alltoall multi-node requires rankSize multiple of 8"); + CHECK_CONTAINS(demo, "bigDataRanksPerNode"); + CHECK_CONTAINS(demo, "AllToAllBigDataValidTopology(rankSize, bigDataRanksPerNode)"); + CHECK_CONTAINS(demo, "ERROR: bigdata alltoall multi-node requires rankSize multiple of ranksPerNode"); CHECK_CONTAINS(demo, "bigdata multinode mode="); CHECK_CONTAINS(demo, "shards="); CHECK_CONTAINS(demo, "remoteSendDoneOffset="); CHECK_CONTAINS(demo, "force35Core="); CHECK_CONTAINS(demo, "TileXR::Demo::AllToAllBigDataShardCount("); CHECK_CONTAINS(demo, "const uint32_t bigDataBlockDim = TileXR::Demo::AllToAllBigDataBlockDim("); - CHECK_CONTAINS(demo, "rankSize, forceBigData35Core, bigDataRemotePutOnly"); + CHECK_CONTAINS(demo, "rankSize, forceBigData35Core, bigDataRemotePutOnly, bigDataRanksPerNode"); CHECK_CONTAINS(demo, "bigDataBlockDim, stream, commArgsDev"); CHECK_CONTAINS(demo, "static_cast(registeredMemory) + bigDataPlan.copyDoneOffset"); CHECK_CONTAINS(demo, "bigDataPlan.controlBytes + bigDataPlan.signalBytes"); CHECK_CONTAINS(demo, "static_cast(bigDataProfileStage)"); + CHECK_CONTAINS(demo, "static_cast(bigDataRanksPerNode) << 8U"); CHECK_CONTAINS(demo, "alltoall udma-bigdata"); CHECK_CONTAINS(demo, "ERROR: bigdata alltoall UDMA registration failed"); CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_bigdata_kernel"); @@ -485,16 +512,18 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY = 4U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_CORE = 20U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_FANOUT_SHARD_BASE = 5U"); - CHECK_CONTAINS(kernel, "BigDataIsMultiNode(rankSize)"); - CHECK_CONTAINS(kernel, "BigDataValidTopology(rankSize)"); - CHECK_CONTAINS(kernel, "BigDataUse35Core(rankSize, force35Core)"); - CHECK_CONTAINS(kernel, "BigDataShardCount(rankSize, force35Core)"); - CHECK_CONTAINS(kernel, "BigDataTaskCount(rankSize, force35Core)"); - CHECK_CONTAINS(kernel, "BigDataNodeCount(rankSize)"); - CHECK_CONTAINS(kernel, "BigDataRemotePeerAt(rank, rankSize, remoteIndex)"); - CHECK_CONTAINS(kernel, "BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex)"); + CHECK_CONTAINS(kernel, "BigDataNormalizeRanksPerNode"); + CHECK_CONTAINS(kernel, "force35CoreFlag >> 8U"); + CHECK_CONTAINS(kernel, "BigDataIsMultiNode(rankSize, ranksPerNode)"); + CHECK_CONTAINS(kernel, "BigDataValidTopology(rankSize, ranksPerNode)"); + CHECK_CONTAINS(kernel, "BigDataUse35Core(rankSize, force35Core, ranksPerNode)"); + CHECK_CONTAINS(kernel, "BigDataShardCount(rankSize, force35Core, ranksPerNode)"); + CHECK_CONTAINS(kernel, "BigDataTaskCount(rankSize, force35Core, ranksPerNode)"); + CHECK_CONTAINS(kernel, "BigDataNodeCount(rankSize, localRanks)"); + CHECK_CONTAINS(kernel, "BigDataRemotePeerAt(rank, rankSize, remoteIndex, localRanks)"); + CHECK_CONTAINS(kernel, "BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode)"); CHECK_CONTAINS(kernel, "BigDataLocalPeerAt"); - CHECK_CONTAINS(kernel, "BigDataMergedPeerTaskAt(rank, rankSize, taskIndex, force35Core, peer, isLocalPeer)"); + CHECK_CONTAINS(kernel, "BigDataMergedPeerTaskAt("); CHECK_CONTAINS(kernel, "BigDataRunSelfCopyShard(rank, rankSize"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentWorker"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendWorker"); @@ -520,14 +549,15 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_BIGDATA_REMOTE_PUT_STAGE_ACK"); CHECK_CONTAINS(kernel, "BigDataSelectWeightedQp("); CHECK_CONTAINS(kernel, "BigDataSelectDistinctWeightedQp("); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySegmentQp("); CHECK_CONTAINS(kernel, "segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT"); - CHECK_CONTAINS(kernel, "primaryQp : BigDataSelectDistinctWeightedQp(args, peer, primaryQp, false)"); + CHECK_CONTAINS(kernel, "return segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? 0U : 1U"); CHECK_CONTAINS(kernel, "BigDataWaitCopyDoneRange"); CHECK_CONTAINS(kernel, "BigDataPublishCopyReadyRange"); CHECK_CONTAINS(kernel, "BigDataWaitCopyReady"); CHECK_CONTAINS(kernel, "BigDataPublishReadySignal"); CHECK_CONTAINS(kernel, "remoteSendDoneOffset"); - CHECK_CONTAINS(kernel, "if (!BigDataIsMultiNode(rankSize))"); + CHECK_CONTAINS(kernel, "!BigDataIsMultiNode(rankSize, ranksPerNode)"); CHECK_CONTAINS(kernel, "const uint32_t activeBlockDim = remotePutOnly ?"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_PUT_ONLY_BLOCK_DIM"); CHECK_CONTAINS(kernel, "if (blockIdx >= static_cast(activeBlockDim))"); From b0d7df0d73a4f413b1c38f0ec789d61fd0e9798d Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 13:05:35 +0800 Subject: [PATCH 030/163] perf(udma): serialize remote put sends per link Route the remote-put-only bigdata send path through fixed send cores per link and walk remote peers serially within each link using rank-staggered order. Update the layout guard to cover the new link mapping and ordering helpers. --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 76 +++++++++---------- .../unit/test_tilexr_udma_alltoall_layout.cpp | 14 ++-- 2 files changed, 45 insertions(+), 45 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index dd0541d7..d079c3e2 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -1472,35 +1472,33 @@ __aicore__ inline void BigDataRemotePutOnlySendWorker( } } -__aicore__ inline int32_t BigDataRemotePutOnlyCheckIndex(int32_t blockIdx) +__aicore__ inline int32_t BigDataRemotePutOnlyLinkForSendCore(int32_t blockIdx) { - if (blockIdx < 0) { - return -1; + if (blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE)) { + return static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT); } - return blockIdx; -} - -__aicore__ inline int32_t BigDataRemotePutOnlySendTaskCount(int32_t remoteTaskCount) -{ - return remoteTaskCount > 0 ? remoteTaskCount * 2 : 0; + if (blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE)) { + return static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT); + } + return -1; } -__aicore__ inline int32_t BigDataRemotePutOnlySendTaskRemoteIndex( - int32_t sendTask, int32_t remoteTaskCount) +__aicore__ inline int32_t BigDataRemotePutOnlyStaggeredRemoteIndex( + int32_t rank, int32_t remoteTaskCount, int32_t remoteOrder) { - if (sendTask < 0 || remoteTaskCount <= 0 || - sendTask >= BigDataRemotePutOnlySendTaskCount(remoteTaskCount)) { + if (remoteTaskCount <= 0 || remoteOrder < 0 || remoteOrder >= remoteTaskCount) { return -1; } - return sendTask < remoteTaskCount ? sendTask : sendTask - remoteTaskCount; + const int32_t start = rank % remoteTaskCount; + return (start + remoteOrder) % remoteTaskCount; } -__aicore__ inline uint32_t BigDataRemotePutOnlySendTaskSegment( - int32_t sendTask, int32_t remoteTaskCount) +__aicore__ inline int32_t BigDataRemotePutOnlyCheckIndex(int32_t blockIdx) { - return sendTask < remoteTaskCount ? - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT : - TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; + if (blockIdx < 0) { + return -1; + } + return blockIdx; } __aicore__ inline void BigDataRemotePutOnlyCheckWorker( @@ -2821,7 +2819,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( static_cast(blockIdx) - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; const int32_t taskCount = BigDataTaskCount(rankSize, force35Core, ranksPerNode); const int32_t remoteTaskCount = rankSize - ranksPerNode; - const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount); + const int32_t sendLink = BigDataRemotePutOnlyLinkForSendCore(blockIdx); if (remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK) { BigDataKernelExitBarrier(); @@ -2831,26 +2829,28 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { if (remotePutOnly) { - for (int32_t sendTask = blockIdx; sendTask < sendTaskCount; - sendTask += static_cast(activeBlockDim)) { - if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { - continue; - } - const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount); - const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode); - if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { - continue; - } - const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount); + if (sendLink >= 0) { + const uint32_t segmentId = static_cast(sendLink); const uint32_t qpIdx = BigDataRemotePutOnlySegmentQp(args, segmentId); - if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { - continue; + for (int32_t remoteOrder = 0; remoteOrder < remoteTaskCount; ++remoteOrder) { + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { + continue; + } + const int32_t remoteIndex = + BigDataRemotePutOnlyStaggeredRemoteIndex(rank, remoteTaskCount, remoteOrder); + const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode); + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { + continue; + } + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { + continue; + } + BigDataRemotePutOnlySendWorker(peer, segmentId, qpIdx, + rank, rankSize, args, input, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, shardCount, recvDataOffset, + chunkBytesPerPeer, ranksPerNode, relayLocal); } - BigDataRemotePutOnlySendWorker(peer, segmentId, qpIdx, - rank, rankSize, args, input, debug, - elementsPerPeer, effectiveChunkElements, passCount, loop, pass, - kernelLoopBase, profileStage, shardCount, recvDataOffset, - chunkBytesPerPeer, ranksPerNode, relayLocal); } BigDataKernelExitBarrier(); if (profileStage > TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 539051a1..2a1689e0 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -527,9 +527,8 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataRunSelfCopyShard(rank, rankSize"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentWorker"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendWorker"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskCount"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskRemoteIndex"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskSegment"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyLinkForSendCore"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyStaggeredRemoteIndex"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckIndex"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); @@ -569,10 +568,11 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "if (isCopyCore)"); CHECK_CONTAINS(kernel, "if (remotePutOnly)"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker("); - CHECK_CONTAINS(kernel, "const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); - CHECK_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); - CHECK_CONTAINS(kernel, "const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); - CHECK_CONTAINS(kernel, "const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount)"); + CHECK_CONTAINS(kernel, "const int32_t sendLink = BigDataRemotePutOnlyLinkForSendCore(blockIdx)"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyStaggeredRemoteIndex(rank, remoteTaskCount, remoteOrder)"); + CHECK_NOT_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); + CHECK_NOT_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); + CHECK_NOT_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); CHECK_CONTAINS(kernel, "remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER"); From 148cab5df446d4c369c9d9b82a0ee9ac1988af91 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 13:09:48 +0800 Subject: [PATCH 031/163] Revert "perf(udma): serialize remote put sends per link" This reverts commit b0d7df0d73a4f413b1c38f0ec789d61fd0e9798d. --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 76 +++++++++---------- .../unit/test_tilexr_udma_alltoall_layout.cpp | 14 ++-- 2 files changed, 45 insertions(+), 45 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index d079c3e2..dd0541d7 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -1472,33 +1472,35 @@ __aicore__ inline void BigDataRemotePutOnlySendWorker( } } -__aicore__ inline int32_t BigDataRemotePutOnlyLinkForSendCore(int32_t blockIdx) +__aicore__ inline int32_t BigDataRemotePutOnlyCheckIndex(int32_t blockIdx) { - if (blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE)) { - return static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT); - } - if (blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE)) { - return static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT); + if (blockIdx < 0) { + return -1; } - return -1; + return blockIdx; } -__aicore__ inline int32_t BigDataRemotePutOnlyStaggeredRemoteIndex( - int32_t rank, int32_t remoteTaskCount, int32_t remoteOrder) +__aicore__ inline int32_t BigDataRemotePutOnlySendTaskCount(int32_t remoteTaskCount) { - if (remoteTaskCount <= 0 || remoteOrder < 0 || remoteOrder >= remoteTaskCount) { - return -1; - } - const int32_t start = rank % remoteTaskCount; - return (start + remoteOrder) % remoteTaskCount; + return remoteTaskCount > 0 ? remoteTaskCount * 2 : 0; } -__aicore__ inline int32_t BigDataRemotePutOnlyCheckIndex(int32_t blockIdx) +__aicore__ inline int32_t BigDataRemotePutOnlySendTaskRemoteIndex( + int32_t sendTask, int32_t remoteTaskCount) { - if (blockIdx < 0) { + if (sendTask < 0 || remoteTaskCount <= 0 || + sendTask >= BigDataRemotePutOnlySendTaskCount(remoteTaskCount)) { return -1; } - return blockIdx; + return sendTask < remoteTaskCount ? sendTask : sendTask - remoteTaskCount; +} + +__aicore__ inline uint32_t BigDataRemotePutOnlySendTaskSegment( + int32_t sendTask, int32_t remoteTaskCount) +{ + return sendTask < remoteTaskCount ? + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT : + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; } __aicore__ inline void BigDataRemotePutOnlyCheckWorker( @@ -2819,7 +2821,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( static_cast(blockIdx) - TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; const int32_t taskCount = BigDataTaskCount(rankSize, force35Core, ranksPerNode); const int32_t remoteTaskCount = rankSize - ranksPerNode; - const int32_t sendLink = BigDataRemotePutOnlyLinkForSendCore(blockIdx); + const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount); if (remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK) { BigDataKernelExitBarrier(); @@ -2829,28 +2831,26 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { if (remotePutOnly) { - if (sendLink >= 0) { - const uint32_t segmentId = static_cast(sendLink); + for (int32_t sendTask = blockIdx; sendTask < sendTaskCount; + sendTask += static_cast(activeBlockDim)) { + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { + continue; + } + const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount); + const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode); + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { + continue; + } + const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount); const uint32_t qpIdx = BigDataRemotePutOnlySegmentQp(args, segmentId); - for (int32_t remoteOrder = 0; remoteOrder < remoteTaskCount; ++remoteOrder) { - if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP) { - continue; - } - const int32_t remoteIndex = - BigDataRemotePutOnlyStaggeredRemoteIndex(rank, remoteTaskCount, remoteOrder); - const int32_t peer = BigDataRemotePeerForwardAt(rank, rankSize, remoteIndex, ranksPerNode); - if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER) { - continue; - } - if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { - continue; - } - BigDataRemotePutOnlySendWorker(peer, segmentId, qpIdx, - rank, rankSize, args, input, debug, - elementsPerPeer, effectiveChunkElements, passCount, loop, pass, - kernelLoopBase, profileStage, shardCount, recvDataOffset, - chunkBytesPerPeer, ranksPerNode, relayLocal); + if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { + continue; } + BigDataRemotePutOnlySendWorker(peer, segmentId, qpIdx, + rank, rankSize, args, input, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, + kernelLoopBase, profileStage, shardCount, recvDataOffset, + chunkBytesPerPeer, ranksPerNode, relayLocal); } BigDataKernelExitBarrier(); if (profileStage > TILEXR_BIGDATA_REMOTE_PUT_STAGE_POST) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 2a1689e0..539051a1 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -527,8 +527,9 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataRunSelfCopyShard(rank, rankSize"); CHECK_CONTAINS(kernel, "BigDataRemoteSendSegmentWorker"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendWorker"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyLinkForSendCore"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyStaggeredRemoteIndex"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskCount"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskRemoteIndex"); + CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskSegment"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckIndex"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot"); @@ -568,11 +569,10 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "if (isCopyCore)"); CHECK_CONTAINS(kernel, "if (remotePutOnly)"); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyCheckWorker("); - CHECK_CONTAINS(kernel, "const int32_t sendLink = BigDataRemotePutOnlyLinkForSendCore(blockIdx)"); - CHECK_CONTAINS(kernel, "BigDataRemotePutOnlyStaggeredRemoteIndex(rank, remoteTaskCount, remoteOrder)"); - CHECK_NOT_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); - CHECK_NOT_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); - CHECK_NOT_CONTAINS(kernel, "BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); + CHECK_CONTAINS(kernel, "const int32_t sendTaskCount = BigDataRemotePutOnlySendTaskCount(remoteTaskCount)"); + CHECK_CONTAINS(kernel, "sendTask += static_cast(activeBlockDim)"); + CHECK_CONTAINS(kernel, "const int32_t remoteIndex = BigDataRemotePutOnlySendTaskRemoteIndex(sendTask, remoteTaskCount)"); + CHECK_CONTAINS(kernel, "const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount)"); CHECK_CONTAINS(kernel, "remotePutOnly && profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP"); CHECK_CONTAINS(kernel, "profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER"); From 4e8b114369d9853f7aad9a351418fb5f72cadd39 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 14:20:53 +0800 Subject: [PATCH 032/163] perf(udma): remove remote put ack wait Drop the remote-put-only peer ACK exchange after both remote segments have been observed ready. Send-side UDMA quiet still precedes ready publication, so the receiver-side ready wait remains the completion condition for inbound data while avoiding the extra peer ACK round in the profiled path. --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 14 -------------- .../udma/unit/test_tilexr_udma_alltoall_layout.cpp | 8 ++++---- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index dd0541d7..fdb8c17d 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -1556,20 +1556,6 @@ __aicore__ inline void BigDataRemotePutOnlyCheckWorker( break; } } - if (status == 0U) { - auto remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); - auto localAck = BigDataLocalIpcAckSlot(args, rank, slot, rankSize, peer); - BigDataRemotePutOnlyPublishAck(remoteAck, token, relayLocal); - const uint64_t ackObserved = BigDataRemotePutOnlyWaitPeerAck( - localAck, token, relayLocal); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_ACK_SEEN_BASE + peer] = - static_cast(ackObserved); - } - if (ackObserved < token) { - status = static_cast(TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS); - } - } if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 539051a1..a2f9c0ab 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -655,10 +655,10 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT"); CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT"); CHECK_CONTAINS(remotePutOnlyCheck, "TILEXR_UDMA_DEMO_READY_TIMEOUT_STATUS"); - CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyPublishAck"); - CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyWaitPeerAck"); - CHECK_CONTAINS(remotePutOnlyCheck, "BigDataRemoteIpcAckSlot"); - CHECK_CONTAINS(remotePutOnlyCheck, "BigDataLocalIpcAckSlot"); + CHECK_NOT_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyPublishAck"); + CHECK_NOT_CONTAINS(remotePutOnlyCheck, "BigDataRemotePutOnlyWaitPeerAck"); + CHECK_NOT_CONTAINS(remotePutOnlyCheck, "BigDataRemoteIpcAckSlot"); + CHECK_NOT_CONTAINS(remotePutOnlyCheck, "BigDataLocalIpcAckSlot"); CHECK_NOT_CONTAINS(remotePutOnlyCheck, "UDMAQuietStatusOnQp"); } From bfcacf7c727cc0da324b909ad1a404abf117cad5 Mon Sep 17 00:00:00 2001 From: linzhen 00654177 Date: Tue, 30 Jun 2026 16:47:14 +0800 Subject: [PATCH 033/163] perf(udma): prefer max-weight route for remote puts Route remote-put-only bigdata sends through the highest-weight UDMA QP instead of splitting segments across max/min routes. This avoids the slower route dominating completion latency in the 16-rank case while keeping the QP count export available for diagnostics. --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 12 ++++------ tests/udma/perf/run_udma_alltoall_perf.py | 4 ++++ .../udma/perf/test_run_udma_alltoall_perf.py | 22 +++++++++++++++++++ .../unit/test_tilexr_udma_alltoall_layout.cpp | 6 +++-- 4 files changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index fdb8c17d..92446fde 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -626,14 +626,10 @@ __aicore__ inline uint32_t BigDataSelectDistinctWeightedQp( } __aicore__ inline uint32_t BigDataRemotePutOnlySegmentQp( - const __gm__ TileXR::CommArgs* args, uint32_t segmentId) + const __gm__ TileXR::CommArgs* args, int32_t peer, uint32_t segmentId) { - auto udmaInfo = TileXR::GetUDMAInfo(args); - const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; - if (qpCount <= 1U) { - return 0U; - } - return segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? 0U : 1U; + (void)segmentId; + return BigDataSelectWeightedQp(args, peer, true); } __aicore__ inline int32_t BigDataLocalPeerAt(int32_t rank, int32_t localIndex, int32_t ranksPerNode) @@ -2828,7 +2824,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( continue; } const uint32_t segmentId = BigDataRemotePutOnlySendTaskSegment(sendTask, remoteTaskCount); - const uint32_t qpIdx = BigDataRemotePutOnlySegmentQp(args, segmentId); + const uint32_t qpIdx = BigDataRemotePutOnlySegmentQp(args, peer, segmentId); if (profileStage <= TILEXR_BIGDATA_REMOTE_PUT_STAGE_QP) { continue; } diff --git a/tests/udma/perf/run_udma_alltoall_perf.py b/tests/udma/perf/run_udma_alltoall_perf.py index 9ac49513..894a7fd0 100644 --- a/tests/udma/perf/run_udma_alltoall_perf.py +++ b/tests/udma/perf/run_udma_alltoall_perf.py @@ -312,6 +312,7 @@ def make_rank_command( barrier_host: str, result_dir: str, route_policy: str, + udma_qp_num: int, profile_stage: int, udma_debug: bool, timeout_s: int, @@ -326,6 +327,7 @@ def make_rank_command( f"export TILEXR_COMM_ID={shlex.quote(comm_id)}", f"export TILEXR_DEMO_BARRIER_HOST={shlex.quote(barrier_host)}", f"export TILEXR_UDMA_ROUTE_POLICY={shlex.quote(route_policy)}", + f"export TILEXR_UDMA_QP_NUM={int(udma_qp_num)}", f"export TILEXR_UDMA_DEBUG={1 if udma_debug else 0}", "export TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY=1", f"export TILEXR_DEMO_BIGDATA_PROFILE_STAGE={profile_stage}", @@ -382,6 +384,7 @@ def run_all_ranks(args, hosts: Sequence[str], repeat: int, label: str) -> Tuple[ barrier_host=hosts[0], result_dir=result_dir, route_policy=args.route_policy, + udma_qp_num=args.udma_qp_num, profile_stage=args.profile_stage, udma_debug=args.udma_debug, timeout_s=args.timeout, @@ -457,6 +460,7 @@ def parse_args(argv: Optional[Sequence[str]] = None): ) parser.add_argument("--demo-bin", default=DEFAULT_DEMO_BIN, help="Demo binary path relative to remote-dir.") parser.add_argument("--route-policy", default="all", help="TILEXR_UDMA_ROUTE_POLICY value.") + parser.add_argument("--udma-qp-num", type=int, default=1, help="TILEXR_UDMA_QP_NUM value.") parser.add_argument("--profile-stage", type=int, default=7, help="TILEXR_DEMO_BIGDATA_PROFILE_STAGE value.") parser.add_argument("--udma-debug", action="store_true", help="Enable TILEXR_UDMA_DEBUG on remote ranks.") parser.add_argument("--timeout", type=int, default=180, help="Kernel/demo timeout seconds.") diff --git a/tests/udma/perf/test_run_udma_alltoall_perf.py b/tests/udma/perf/test_run_udma_alltoall_perf.py index daf3a977..e3eb3074 100644 --- a/tests/udma/perf/test_run_udma_alltoall_perf.py +++ b/tests/udma/perf/test_run_udma_alltoall_perf.py @@ -64,6 +64,28 @@ def test_infers_external_tool_env_from_cann_env(self): self.assertIn("/home/user/tilexr/env/util/cmake/bin", prelude) self.assertIn("source /home/user/tilexr/env/cann/cann/set_env.sh", prelude) + def test_rank_command_exports_udma_qp_num(self): + command = perf.make_rank_command( + rank=0, + world_size=2, + repeat=1, + elements_per_peer=1024, + devices_per_host=1, + remote_dir="/home/tileXR", + cann_env="/opt/cann/set_env.sh", + demo_bin="tests/udma/install/bin/tilexr_udma_demo", + comm_id="127.0.0.1:12345", + barrier_host="127.0.0.1", + result_dir="/tmp/result", + route_policy="all", + udma_qp_num=2, + profile_stage=7, + udma_debug=False, + timeout_s=90, + tool_env="/opt/tilexr/env", + ) + self.assertIn("export TILEXR_UDMA_QP_NUM=2", command) + if __name__ == "__main__": unittest.main() diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index a2f9c0ab..03483f2c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -550,8 +550,10 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataSelectWeightedQp("); CHECK_CONTAINS(kernel, "BigDataSelectDistinctWeightedQp("); CHECK_CONTAINS(kernel, "BigDataRemotePutOnlySegmentQp("); - CHECK_CONTAINS(kernel, "segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT"); - CHECK_CONTAINS(kernel, "return segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? 0U : 1U"); + CHECK_CONTAINS(kernel, "(void)segmentId"); + CHECK_CONTAINS(kernel, "return BigDataSelectWeightedQp(args, peer, true)"); + CHECK_NOT_CONTAINS(kernel, "return 1U;"); + CHECK_NOT_CONTAINS(kernel, "return segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT ? 0U : 1U"); CHECK_CONTAINS(kernel, "BigDataWaitCopyDoneRange"); CHECK_CONTAINS(kernel, "BigDataPublishCopyReadyRange"); CHECK_CONTAINS(kernel, "BigDataWaitCopyReady"); From 0ca36a7d4888e908060125c0e21fa188da4a479d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:17:55 +0800 Subject: [PATCH 034/163] docs(udma): design fullmesh control route fix --- ...7-17-udma-fullmesh-control-route-design.md | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-17-udma-fullmesh-control-route-design.md diff --git a/docs/superpowers/specs/2026-07-17-udma-fullmesh-control-route-design.md b/docs/superpowers/specs/2026-07-17-udma-fullmesh-control-route-design.md new file mode 100644 index 00000000..d51a664a --- /dev/null +++ b/docs/superpowers/specs/2026-07-17-udma-fullmesh-control-route-design.md @@ -0,0 +1,41 @@ +# UDMA Full-Mesh Control Route Design + +## Goal + +Fix the physical 2x8 full-mesh AllToAll path so payload completion signals do +not fall back to generic multi-route UDMA operations that can select a failing +route. + +## Scope + +- Keep the existing 12:4 payload split between cores 16 and 17. +- Bind the primary segment to the maximum-weight QP and the secondary segment + to the minimum-weight QP. +- Send each segment payload and wait for quiet on its selected QP. +- Publish the final ready signal on the primary segment's selected QP. +- Publish receive ACKs by writing the peer's registered control slot directly, + without a generic UDMA put or quiet. +- Do not import the XY pipeline, 16:0 split, tracing, or scheduling changes. + +## Data Flow + +Core 16 waits for shards 0 through 11, sends the primary payload on the +maximum-weight QP, and records local segment completion. Core 17 waits for +shards 12 through 15, sends the secondary payload on the minimum-weight QP, +and records local segment completion. Core 16 waits for both local completion +tokens and publishes ready on the same maximum-weight QP used for its payload. + +After all receive-copy shards complete, the receiver writes the ACK token to +the sender's registered ACK control slot. The sender continues to wait on its +local registered ACK slot, so the synchronization contract is unchanged. + +## Verification + +1. Add source-structure tests that isolate the full-mesh send worker and check + explicit QP payload, ready, and quiet operations. +2. Check that the full-mesh receive worker no longer calls the generic UDMA ACK + helper and instead resolves the remote registered ACK slot. +3. Run the UDMA layout unit test locally and on both remote hosts. +4. Deploy only a committed Git bundle. +5. Run physical 2x8, 128 MiB per rank, repeat1, profile stage 8, with full + output validation. Run repeat50 only if all 16 ranks pass. From 6c1717c51303ada055fcae49813421318aee8ea3 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:20:33 +0800 Subject: [PATCH 035/163] docs(udma): plan fullmesh control route fix --- .../2026-07-17-udma-fullmesh-control-route.md | 147 ++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-17-udma-fullmesh-control-route.md diff --git a/docs/superpowers/plans/2026-07-17-udma-fullmesh-control-route.md b/docs/superpowers/plans/2026-07-17-udma-fullmesh-control-route.md new file mode 100644 index 00000000..a6704958 --- /dev/null +++ b/docs/superpowers/plans/2026-07-17-udma-fullmesh-control-route.md @@ -0,0 +1,147 @@ +# UDMA Full-Mesh Control Route Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Keep the full-mesh 12:4 payload split while binding payload and ready completion to explicit QPs and removing generic-route UDMA ACK publication. + +**Architecture:** The two remote-send workers select max/min weighted QPs for their existing primary/secondary segments. The primary worker publishes ready on its selected QP after both segments complete. Receive completion writes the token directly to the sender's registered ACK control slot. + +**Tech Stack:** C++14, Ascend C device code, TileXR UDMA device helpers, source-structure unit tests. + +## Global Constraints + +- Preserve `TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U`. +- Do not import XY pipeline, trace, scheduling, or 16:0 changes. +- Commit before creating and uploading a Git bundle. +- Validate physical 2x8 with `TILEXR_IPC_PID_MODE=pid`. + +--- + +### Task 1: Add Failing Full-Mesh Route Tests + +**Files:** +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` + +**Interfaces:** +- Consumes: `SliceBetween`, `CHECK_CONTAINS`, and `CHECK_NOT_CONTAINS`. +- Produces: source-structure requirements for full-mesh payload, ready, and ACK traffic. + +- [ ] **Step 1: Add isolated source checks** + +```cpp +const std::string fullMeshReady = SliceBetween( + kernel, "BigDataPublishReadySignal", "BigDataRemoteSendSegmentWorker"); +const std::string fullMeshSend = SliceBetween( + kernel, "BigDataRemoteSendSegmentWorker", "BigDataRemotePutOnlySendWorker"); +const std::string fullMeshRecv = SliceBetween( + kernel, "BigDataRecvPeerWorker", "BigDataWaitCopyDoneRange"); +CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U"); +CHECK_CONTAINS(fullMeshReady, "UDMAPutSignalNbiOnQp"); +CHECK_CONTAINS(fullMeshReady, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); +CHECK_CONTAINS(fullMeshSend, "BigDataSelectWeightedQp("); +CHECK_CONTAINS(fullMeshSend, "UDMAPutNbiOnQp"); +CHECK_CONTAINS(fullMeshSend, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); +CHECK_NOT_CONTAINS(fullMeshSend, "TileXR::UDMAPutNbi(args, peer"); +CHECK_CONTAINS(fullMeshRecv, "BigDataRemoteRegisteredControlSlot("); +CHECK_CONTAINS(fullMeshRecv, "BigDataStoreTokenMte(remoteAck, token, relayLocal)"); +CHECK_NOT_CONTAINS(fullMeshRecv, "BigDataPublishAckSignalUdma("); +``` + +- [ ] **Step 2: Commit the failing test and create a RED bundle** + +```bash +git add tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +git commit -m "test(udma): require explicit fullmesh control routes" +git bundle create tmp/test0701-fullmesh-control-route-red.bundle test0701-bundle +git bundle verify tmp/test0701-fullmesh-control-route-red.bundle +``` + +- [ ] **Step 3: Verify RED remotely** + +Upload the committed bundle to both hosts, fetch it in `/home/h30059441/tilexr_stage0_y_stagger`, build target `test_tilexr_udma_alltoall_layout` in `tests/udma/build-test0701`, and run it. Expected: nonzero exit caused by the new checks for `UDMAPutSignalNbiOnQp`, explicit full-mesh payload QP, and direct registered ACK. + +### Task 2: Route Full-Mesh Data And Control Explicitly + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` + +**Interfaces:** +- Consumes: `BigDataSelectWeightedQp(args, peer, selectMax)` and `BigDataRemoteRegisteredControlSlot(...)`. +- Produces: `BigDataPublishReadySignal(..., uint32_t qpIdx)` and direct registered ACK publication. + +- [ ] **Step 1: Bind payload and ready to the segment QP** + +In `BigDataRemoteSendSegmentWorker`, preserve the current segment ranges and add: + +```cpp +const uint32_t qpIdx = BigDataSelectWeightedQp( + args, peer, + segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT); +TileXR::UDMAPutNbiOnQp( + args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); +status = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); +``` + +Pass `qpIdx` to `BigDataPublishReadySignal`, whose put and quiet become: + +```cpp +TileXR::UDMAPutSignalNbiOnQp( + args, peer, qpIdx, localSrc, localReadyPayloadOffset, + sizeof(uint64_t), remoteReadyOffset, token); +(void)TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); +``` + +- [ ] **Step 2: Replace generic UDMA ACK publication** + +For both multi-node ACK publication sites in `BigDataRecvPeerWorker`, resolve and write the remote registered slot: + +```cpp +auto remoteAck = BigDataRemoteRegisteredControlSlot( + args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); +if (remoteAck == nullptr) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + return; +} +BigDataStoreTokenMte(remoteAck, token, relayLocal); +``` + +Remove `BigDataPublishAckSignalUdma` after its callers are gone. + +- [ ] **Step 3: Commit, bundle, and verify GREEN remotely** + +```bash +git add tests/udma/demo/tilexr_udma_demo_kernel.cpp +git commit -m "fix(udma): pin fullmesh control traffic to valid routes" +git bundle create tmp/test0701-fullmesh-control-route.bundle test0701-bundle +git bundle verify tmp/test0701-fullmesh-control-route.bundle +``` + +Upload the bundle to both hosts, fetch it, rebuild `tilexr_udma_demo` and `test_tilexr_udma_alltoall_layout`, and run the unit test. Expected: exit 0. Also run `git diff --check` locally. + +### Task 3: Verify Physical 2x8 + +**Files:** +- No tracked file changes. + +**Interfaces:** +- Consumes: committed Task 2 HEAD. +- Produces: physical 2x8 correctness evidence and remote result logs. + +- [ ] **Step 1: Run repeat1 full-mesh correctness** + +```bash +export TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY=0 +export TILEXR_DEMO_BIGDATA_PROFILE_STAGE=8 +export TILEXR_DEMO_ALLTOALL_REPEAT=1 +export TILEXR_IPC_PID_MODE=pid +./tests/udma/install-test0701/bin/tilexr_udma_demo 16 RANK 7 2097152 8 0 +``` + +Expected: all 16 ranks exit 0, print `TileXR UDMA demo success`, and report no `CQ incomplete`, `MISMATCH`, or `ERROR`. + +- [ ] **Step 2: Run repeat50 only after repeat1 passes** + +Change `TILEXR_DEMO_ALLTOALL_REPEAT` to `50`. Report max per-rank kernel time and retain the remote result directory. If repeat1 fails, stop and analyze that run instead. From b03c7475f1f50d8de9774011e4af1ffe1ecff2f2 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:21:39 +0800 Subject: [PATCH 036/163] test(udma): require explicit fullmesh control routes --- .../unit/test_tilexr_udma_alltoall_layout.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 03483f2c..b18f8823 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -637,6 +637,23 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY"); CHECK_CONTAINS(demo, "skip result validation for bigdata remote-put-only profile"); + const std::string fullMeshReady = SliceBetween( + kernel, "BigDataPublishReadySignal", "BigDataRemoteSendSegmentWorker"); + const std::string fullMeshSend = SliceBetween( + kernel, "BigDataRemoteSendSegmentWorker", "BigDataRemotePutOnlySendWorker"); + const std::string fullMeshRecv = SliceBetween( + kernel, "BigDataRecvPeerWorker", "BigDataWaitCopyDoneRange"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U"); + CHECK_CONTAINS(fullMeshReady, "UDMAPutSignalNbiOnQp"); + CHECK_CONTAINS(fullMeshReady, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); + CHECK_CONTAINS(fullMeshSend, "BigDataSelectWeightedQp("); + CHECK_CONTAINS(fullMeshSend, "UDMAPutNbiOnQp"); + CHECK_CONTAINS(fullMeshSend, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); + CHECK_NOT_CONTAINS(fullMeshSend, "TileXR::UDMAPutNbi(args, peer"); + CHECK_CONTAINS(fullMeshRecv, "BigDataRemoteRegisteredControlSlot("); + CHECK_CONTAINS(fullMeshRecv, "BigDataStoreTokenMte(remoteAck, token, relayLocal)"); + CHECK_NOT_CONTAINS(fullMeshRecv, "BigDataPublishAckSignalUdma("); + const std::string remotePutOnlySend = SliceBetween( kernel, "BigDataRemotePutOnlySendWorker", "BigDataRemotePutOnlyCheckIndex"); const std::string remotePutOnlyCheck = SliceBetween( From 7c2ce289ef3ab694e670161928e50f83ce127fb1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:23:57 +0800 Subject: [PATCH 037/163] fix(udma): pin fullmesh control traffic to valid routes --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 58 ++++++++----------- .../unit/test_tilexr_udma_alltoall_layout.cpp | 4 +- 2 files changed, 25 insertions(+), 37 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 92446fde..0cee2d74 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -805,20 +805,6 @@ __aicore__ inline __gm__ uint64_t* BigDataRemoteIpcReadySlot( return BigDataRemoteIpcAckSlot(args, targetRank, rank, slot, rankSize) + 1U + segmentId; } -__aicore__ inline uint32_t BigDataPublishAckSignalUdma( - const __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, - uint64_t ackSignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, - int32_t rank, uint64_t token, AscendC::LocalTensor relayLocal) -{ - const uint64_t ackOffset = - BigDataRegisteredControlOffset(ackSignalOffset, slot, rankSize, shardCount, rank, 0U); - auto localAck = reinterpret_cast<__gm__ uint64_t*>(udmaMem + ackOffset); - BigDataStoreTokenMte(localAck, token, relayLocal); - TileXR::UDMAPutSignalNbi( - args, peer, localAck, ackOffset, sizeof(uint64_t), ackOffset, token); - return TileXR::UDMAQuietStatus(args, peer); -} - __aicore__ inline void BigDataRemotePutOnlyPublishAck( __gm__ uint64_t* remoteAck, uint64_t token, AscendC::LocalTensor relayLocal) { @@ -1155,13 +1141,16 @@ __aicore__ inline void BigDataRecvPeerWorker( } } - const uint32_t ackStatus = BigDataPublishAckSignalUdma( - args, peer, udmaMem, ackSignalOffset, slot, rankSize, shardCount, - rank, token, relayLocal); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - static_cast(ackStatus); + auto remoteAck = BigDataRemoteRegisteredControlSlot( + args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); + if (remoteAck == nullptr) { + if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { + debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = + TILEXR_UDMA_DEMO_ACK_TIMEOUT_STATUS; + } + return; } + BigDataStoreTokenMte(remoteAck, token, relayLocal); } return; } @@ -1180,14 +1169,8 @@ __aicore__ inline void BigDataRecvPeerWorker( __gm__ uint64_t* remoteAck = nullptr; if (use35Core) { - const uint32_t ackStatus = BigDataPublishAckSignalUdma( - args, peer, udmaMem, ackSignalOffset, slot, rankSize, shardCount, - rank, token, relayLocal); - if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { - debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = - static_cast(ackStatus); - } - return; + remoteAck = BigDataRemoteRegisteredControlSlot( + args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); } else { remoteAck = BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize); } @@ -1293,7 +1276,7 @@ __aicore__ inline bool BigDataRemoteSendSegmentRange( __aicore__ inline void BigDataPublishReadySignal( __gm__ TileXR::CommArgs* args, int32_t peer, __gm__ uint8_t* udmaMem, uint64_t readySignalOffset, uint32_t slot, int32_t rankSize, uint32_t shardCount, - int32_t rank, uint64_t token) + int32_t rank, uint64_t token, uint32_t qpIdx) { const uint64_t localReadyPayloadOffset = readySignalOffset + @@ -1302,9 +1285,10 @@ __aicore__ inline void BigDataPublishReadySignal( TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; const uint64_t remoteReadyOffset = localReadyPayloadOffset; auto localSrc = reinterpret_cast<__gm__ uint64_t*>(udmaMem + localReadyPayloadOffset); - TileXR::UDMAPutSignalNbi( - args, peer, localSrc, localReadyPayloadOffset, sizeof(uint64_t), remoteReadyOffset, token); - (void)TileXR::UDMAQuietStatus(args, peer); + TileXR::UDMAPutSignalNbiOnQp( + args, peer, qpIdx, localSrc, localReadyPayloadOffset, + sizeof(uint64_t), remoteReadyOffset, token); + (void)TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); } __aicore__ inline void BigDataRemoteSendSegmentWorker( @@ -1368,11 +1352,15 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( (static_cast(slot) * static_cast(networkPeerCount) + static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer + static_cast(segmentOffsetBytes); + const uint32_t qpIdx = BigDataSelectWeightedQp( + args, peer, + segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT); uint32_t status = 0U; if (segmentBytes > 0U) { - TileXR::UDMAPutNbi(args, peer, localSrc, remoteDataOffset, segmentBytes); - status = TileXR::UDMAQuietStatus(args, peer); + TileXR::UDMAPutNbiOnQp( + args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); + status = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); } if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); @@ -1401,7 +1389,7 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, 0U); BigDataStoreTokenMte(localReady, token, relayLocal); BigDataPublishReadySignal( - args, peer, udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, token); + args, peer, udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, token, qpIdx); } __aicore__ inline void BigDataRemotePutOnlySendWorker( diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index b18f8823..bde1b788 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -617,7 +617,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); CHECK_CONTAINS(kernel, "UDMAPutNbiOnQp"); - CHECK_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(kernel, "UDMAPutNbi"); CHECK_CONTAINS(kernel, "localSrc,"); CHECK_CONTAINS(kernel, "remoteDataOffset, chunkBytes, remoteReadyOffset, token"); @@ -628,7 +628,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataControlSlot(udmaMem, ackSignalOffset, slot, rankSize, shardCount, peer, 0U)"); CHECK_CONTAINS(kernel, "use35Core ? passCount : TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS"); CHECK_CONTAINS(kernel, "recvSlotInt[0]"); - CHECK_CONTAINS(kernel, "BigDataPublishAckSignalUdma"); + CHECK_NOT_CONTAINS(kernel, "BigDataPublishAckSignalUdma"); CHECK_CONTAINS(kernel, "BigDataRemoteIpcAckSlot(args, peer, rank, slot, rankSize)"); CHECK_CONTAINS(kernel, "ackSignal"); CHECK_CONTAINS(kernel, "fillLocal.SetValue(0, value)"); From 4c20ef1e7b1b7aa958fb75e6f4dd2365460adb35 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:42:21 +0800 Subject: [PATCH 038/163] docs(udma): design fullmesh GM tracing --- ...026-07-17-udma-fullmesh-gm-trace-design.md | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-17-udma-fullmesh-gm-trace-design.md diff --git a/docs/superpowers/specs/2026-07-17-udma-fullmesh-gm-trace-design.md b/docs/superpowers/specs/2026-07-17-udma-fullmesh-gm-trace-design.md new file mode 100644 index 00000000..bc75d263 --- /dev/null +++ b/docs/superpowers/specs/2026-07-17-udma-fullmesh-gm-trace-design.md @@ -0,0 +1,114 @@ +# UDMA Full-Mesh GM Trace Design + +## Goal + +Add per-core, per-pass, per-peer cycle spans to the physical full-mesh +AllToAll path. Device code records raw system-cycle timestamps in a dedicated +8 MiB GM allocation. Host code only initializes, copies, and writes the trace; +analysis and Chrome Tracing conversion happen after the run. + +## Scope + +- Instrument the non-remote-put-only multi-node bigdata full-mesh kernel. +- Record up to 50 measured iterations for cores 0 through 34. +- Index task spans by iteration, core, pass, peer, and phase. +- Preserve the current 12:4 remote-send split and communication behavior. +- Keep tracing disabled by default. +- Reject trace dimensions that do not fit in 8 MiB before kernel launch. + +## Storage Layout + +The trace allocation contains a 4 KiB header, a fixed kernel-span region, and +a runtime-sized task-span region. Every span is two `uint64_t` values: +`beginCycle` and `endCycle`. + +Kernel spans use fixed indexing: + +```text +[iteration][core][kernel | work] +``` + +Task spans use runtime dimensions from the header: + +```text +[iteration][core][pass][peer][phase] +``` + +For repeat50, 35 cores, one pass, 16 peers, and 14 phases, task spans consume +6,272,000 bytes. The complete layout fits in 8 MiB. Host validation computes +the exact required bytes from the requested repeat, pass count, and rank size. + +## Phases + +The task phase enum contains: + +1. pass +2. self-copy +3. peer-copy +4. publish-copy-ready +5. wait-copy-ready +6. data-put +7. quiet +8. segment-done +9. publish-ready +10. wait-ready +11. output-copy +12. publish-recv-done +13. wait-recv-done +14. ACK + +Core responsibilities map to phases as follows: + +- Cores 0-15: self-copy, peer-copy, and publish-copy-ready where applicable. +- Cores 16-17: wait-copy-ready, data-put, quiet, and segment-done. Core 16 + additionally records publish-ready. +- Core 18: wait-copy-ready, data-put, and quiet for local peers. +- Cores 19-34: wait-ready, output-copy, and publish-recv-done. Core 34 + additionally records wait-recv-done and ACK. + +## Device Integration + +The host passes a nullable trace pointer and an iteration index to the kernel. +Trace helpers return immediately when the pointer is null or any dimension is +out of range. `GetSystemCycle()` and GM span writes only execute when tracing +is enabled. + +Worker helpers receive enough trace context to record boundaries around the +actual wait, copy, put, quiet, signal, and ACK operations. Existing +synchronization and data movement remain unchanged. + +## Host Integration + +Tracing is enabled by `TILEXR_UDMA_FULLMESH_TRACE=1`. The output directory is +selected with `TILEXR_UDMA_FULLMESH_TRACE_DIR`, defaulting to the current +directory. + +When enabled, Host code: + +1. validates repeat/pass/rank dimensions against the 8 MiB capacity; +2. allocates and zero-initializes trace GM; +3. writes the trace header; +4. launches each measured iteration with its trace index; +5. copies the complete trace allocation to Host after stream synchronization; +6. writes `tilexr_fullmesh_trace_rank_.bin`. + +Failure to allocate, initialize, copy, or write the trace fails the demo run. + +## Conversion + +`tilexr_udma_fullmesh_trace_to_chrome.py` accepts one or more rank binary files +and writes Chrome Trace Event JSON. Each rank is a process and each core is a +thread. Events include iteration, pass, peer, phase, begin cycle, end cycle, +and duration. Timestamps are normalized to the earliest recorded kernel span +for each rank and converted with 1000 cycles per microsecond. + +## Verification + +- Header/layout tests validate offsets, capacity, and overflow rejection. +- Source-structure tests require nullable trace plumbing and phase recording in + every full-mesh core role. +- Converter tests validate metadata, event names, peer/pass arguments, and + Chrome-compatible JSON without unsupported display-time units. +- Remote builds run layout and converter tests before hardware execution. +- Physical 2x8 runs repeat50 with full correctness and trace enabled. The + resulting binaries are downloaded and converted to Chrome JSON. From b91b27e4f2f1462646dc5f199daa3f72925c2205 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:43:57 +0800 Subject: [PATCH 039/163] docs(udma): plan fullmesh GM tracing --- .../2026-07-17-udma-fullmesh-gm-trace.md | 258 ++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-17-udma-fullmesh-gm-trace.md diff --git a/docs/superpowers/plans/2026-07-17-udma-fullmesh-gm-trace.md b/docs/superpowers/plans/2026-07-17-udma-fullmesh-gm-trace.md new file mode 100644 index 00000000..067eb462 --- /dev/null +++ b/docs/superpowers/plans/2026-07-17-udma-fullmesh-gm-trace.md @@ -0,0 +1,258 @@ +# UDMA Full-Mesh GM Trace Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Record raw system-cycle spans for every physical full-mesh core task and convert the 8 MiB per-rank trace to Chrome Trace Event JSON. + +**Architecture:** A shared header defines a fixed kernel-span area and runtime-indexed task spans. Host code owns allocation and file output; device workers only record nullable spans. A standalone Python converter reads one or more rank binaries and emits rank/core lanes. + +**Tech Stack:** C++14 Host code, Ascend C device code, Python 3 unittest, Chrome Trace Event JSON. + +## Global Constraints + +- Trace allocation is exactly 8 MiB per rank. +- Record at most 50 iterations and cores 0 through 34. +- Preserve full-mesh communication and the 12:4 payload split. +- Do not average timestamps on device or Host. +- Commit before every remote bundle deployment. +- Use `TILEXR_IPC_PID_MODE=pid` for physical validation. + +--- + +### Task 1: Define And Test The Binary Layout + +**Files:** +- Create: `tests/udma/demo/tilexr_udma_fullmesh_trace.h` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` +- Modify: `tests/udma/CMakeLists.txt` + +**Interfaces:** +- Produces: `FullmeshTraceHeader`, `FullmeshTraceSpan`, `FullmeshTraceLayoutBytes`, `FullmeshTraceTaskSpanOffset`, and phase constants. + +- [ ] **Step 1: Add failing layout assertions** + +Require magic `0x464d5452`, version 1, 8 MiB, 4 KiB header, 50 iterations, +35 cores, two kernel regions, and 14 task phases. Assert that the physical +2x8 repeat50/pass1 layout fits and repeat50/pass4 does not. + +- [ ] **Step 2: Commit and remotely verify RED** + +Commit the test, create a full Git bundle, rebuild +`test_tilexr_udma_alltoall_layout` remotely, and confirm failure because the +new header and symbols do not exist. + +- [ ] **Step 3: Implement the shared layout** + +Define: + +```cpp +constexpr uint32_t kFullmeshTraceMagic = 0x464d5452U; +constexpr uint32_t kFullmeshTraceVersion = 1U; +constexpr size_t kFullmeshTraceBytes = 8ULL * 1024ULL * 1024ULL; +constexpr size_t kFullmeshTraceHeaderBytes = 4096ULL; +constexpr uint32_t kFullmeshTraceMaxIterations = 50U; +constexpr uint32_t kFullmeshTraceMaxCores = 35U; +constexpr uint32_t kFullmeshTraceKernelRegions = 2U; +constexpr uint32_t kFullmeshTracePhaseCount = 14U; +constexpr uint64_t kFullmeshTraceCyclesPerUs = 1000ULL; +``` + +Use runtime `passCount` and `rankSize` in task indexing, but reserve kernel +spans for all 50 iterations and 35 cores. Return a value greater than 8 MiB +when multiplication would overflow or dimensions exceed their limits. + +- [ ] **Step 4: Verify GREEN remotely and commit** + +Run the layout test on both hosts, then commit the header, test, and CMake +dependency. + +### Task 2: Build The Chrome Trace Converter + +**Files:** +- Create: `tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py` +- Create: `tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py` + +**Interfaces:** +- Produces: `read_rank_trace(path)` and `build_chrome_trace(rank_traces)`. + +- [ ] **Step 1: Write failing converter tests** + +Create a synthetic 8 MiB rank trace with one kernel span and task spans for +core16/peer8 data-put and core34/peer8 ACK. Require: + +```python +self.assertEqual(data_put["tid"], 16) +self.assertEqual(data_put["args"]["peer"], 8) +self.assertEqual(data_put["dur"], 0.1) +self.assertEqual(ack["name"], "ACK") +self.assertNotIn("displayTimeUnit", trace) +``` + +Also test unknown versions, truncated files, half-written spans, capacity +overflow, and independent rank clock normalization. + +- [ ] **Step 2: Run tests and verify RED** + +```powershell +python -m unittest tests.udma.unit.test_tilexr_udma_fullmesh_trace_to_chrome -v +``` + +Expected: import failure because the converter does not exist. + +- [ ] **Step 3: Implement the converter** + +Parse header format `<8I4Q`, validate all dimensions and offsets, and index +task spans with: + +```python +index = ((((iteration * MAX_CORES + core) * pass_count + pass_index) * + rank_size + peer) * PHASE_COUNT + phase) +``` + +Emit metadata events for `rank N` and `coreN `. Emit complete events +with raw cycle arguments plus iteration, pass, peer, phase, and role. + +- [ ] **Step 4: Verify GREEN and commit** + +Run the unittest command above and commit the converter and tests. + +### Task 3: Add Host Allocation And Binary Output + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` + +**Interfaces:** +- Consumes: `kFullmeshTraceBytes` and `FullmeshTraceLayoutBytes`. +- Produces: launch arguments `GM_ADDR fullmeshTrace` and `uint32_t fullmeshTraceIteration`. + +- [ ] **Step 1: Add failing Host source checks** + +Require environment names `TILEXR_UDMA_FULLMESH_TRACE` and +`TILEXR_UDMA_FULLMESH_TRACE_DIR`, capacity validation, `aclrtMalloc`, H2D/D2H +copies, per-iteration launch indices, and output filename +`tilexr_fullmesh_trace_rank_.bin`. + +- [ ] **Step 2: Verify RED remotely** + +Commit and bundle the tests, rebuild the layout test remotely, and confirm the +new Host checks fail. + +- [ ] **Step 3: Implement Host plumbing** + +Enable tracing only for test type 7, multi-node full-mesh, +`REMOTE_PUT_ONLY=0`, and the environment flag. Initialize this header: + +```cpp +FullmeshTraceHeader header {}; +header.magic = kFullmeshTraceMagic; +header.version = kFullmeshTraceVersion; +header.rank = rank; +header.iterationCount = allToAllRepeat; +header.passCount = bigDataPlan.passCount; +header.coreCount = kFullmeshTraceMaxCores; +header.rankSize = rankSize; +header.phaseCount = kFullmeshTracePhaseCount; +header.cyclesPerUs = kFullmeshTraceCyclesPerUs; +header.traceBytes = kFullmeshTraceBytes; +header.kernelSpanOffset = kFullmeshTraceHeaderBytes; +header.taskSpanOffset = FullmeshTraceTaskSpanBaseOffset(); +``` + +Pass null/zero when disabled. Copy and write all 8 MiB after stream sync and +before freeing the trace allocation. Treat every allocation/copy/write error +as demo failure. + +- [ ] **Step 4: Verify Host tests and commit** + +Rebuild Host and layout tests remotely, run the layout test, and commit. + +### Task 4: Instrument Full-Mesh Core Tasks + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` + +**Interfaces:** +- Consumes: nullable trace pointer, iteration, pass count, rank size, peer, and phase. +- Produces: deterministic kernel and task spans for cores 0 through 34. + +- [ ] **Step 1: Add failing kernel source checks** + +Require `GetSystemCycle`, kernel/work span recording, and each phase in the +copy, remote-send, local-send, and receive worker slices. Require trace helpers +to guard `trace == nullptr` before reading cycles or writing GM. + +- [ ] **Step 2: Verify RED remotely** + +Commit and bundle the tests, rebuild the layout test, and confirm the phase +checks fail. + +- [ ] **Step 3: Implement trace helpers** + +Add guarded helpers: + +```cpp +BigDataFullmeshTraceRecordKernelSpan(trace, iteration, core, region, begin, end); +BigDataFullmeshTraceRecordTaskSpan( + trace, iteration, core, pass, peer, phase, + passCount, rankSize, begin, end); +``` + +Each helper validates all dimensions before using `DataCopyPad` to write the +16-byte span. + +- [ ] **Step 4: Instrument worker boundaries** + +Record the 14 phases at the actual operation boundaries defined in the design. +Use `rank` as the peer index for pass and self-copy spans. Do not combine +separate peer operations into one span. Record kernel/work end spans on every +active full-mesh core after the pass loop. + +- [ ] **Step 5: Build, run tests, and commit** + +Build the Ascend C kernel and Host demo on both remote hosts. Run layout, +transport-layout, and converter tests. Commit only after all pass. + +### Task 5: Run And Export Physical 2x8 Trace + +**Files:** +- Create: `tmp/test0701_fullmesh_trace_2x8/` artifacts only. + +**Interfaces:** +- Consumes: committed trace-enabled demo. +- Produces: 16 binary traces and one Chrome Trace JSON. + +- [ ] **Step 1: Bundle and deploy** + +Create and verify a complete branch bundle, upload it to both hosts, detach at +the committed HEAD, rebuild, install, and rerun unit tests. + +- [ ] **Step 2: Run physical 2x8 repeat50** + +Use: + +```bash +export TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY=0 +export TILEXR_DEMO_BIGDATA_PROFILE_STAGE=8 +export TILEXR_DEMO_ALLTOALL_REPEAT=50 +export TILEXR_UDMA_FULLMESH_TRACE=1 +export TILEXR_UDMA_FULLMESH_TRACE_DIR= +export TILEXR_IPC_PID_MODE=pid +``` + +Expected: 16/16 ranks pass full output validation and each writes one 8 MiB +binary trace. + +- [ ] **Step 3: Download and convert** + +Download all rank binaries into `tmp/test0701_fullmesh_trace_2x8/` and run: + +```powershell +python tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py ` + tmp/test0701_fullmesh_trace_2x8/*.bin ` + --output tmp/test0701_fullmesh_trace_2x8/fullmesh_2x8_trace.json +``` + +Verify JSON parsing, event count, all 16 process IDs, all active core lanes, +and absence of `displayTimeUnit`. From d2ab380676f360a3f9b2eda35cc001f2c158ff65 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:45:18 +0800 Subject: [PATCH 040/163] test(udma): define fullmesh trace capacity --- .../unit/test_tilexr_udma_alltoall_layout.cpp | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index bde1b788..4890c37e 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -7,6 +7,7 @@ #include #include "demo/tilexr_udma_alltoall_layout.h" +#include "demo/tilexr_udma_fullmesh_trace.h" #ifndef TILEXR_SOURCE_ROOT #define TILEXR_SOURCE_ROOT "." @@ -99,6 +100,30 @@ void TestAllToAllOutputValidation() CHECK_EQ(TileXR::Demo::ValidateAllToAllOutput(output, rank, rankSize, elementsPerPeer), false); } +void TestFullmeshTraceLayout() +{ + using namespace TileXR::Demo; + CHECK_EQ(kFullmeshTraceMagic, 0x464d5452U); + CHECK_EQ(kFullmeshTraceVersion, 1U); + CHECK_EQ(kFullmeshTraceBytes, 8ULL * 1024ULL * 1024ULL); + CHECK_EQ(kFullmeshTraceHeaderBytes, 4096ULL); + CHECK_EQ(kFullmeshTraceMaxIterations, 50U); + CHECK_EQ(kFullmeshTraceMaxCores, 35U); + CHECK_EQ(kFullmeshTraceKernelRegions, 2U); + CHECK_EQ(kFullmeshTracePhaseCount, 14U); + CHECK_EQ(sizeof(FullmeshTraceSpan), 16U); + CHECK_EQ(FullmeshTraceLayoutBytes(50U, 1U, 16U) <= kFullmeshTraceBytes, true); + CHECK_EQ(FullmeshTraceLayoutBytes(50U, 4U, 16U) > kFullmeshTraceBytes, true); + CHECK_EQ(FullmeshTraceLayoutFits(50U, 1U, 16U), true); + CHECK_EQ(FullmeshTraceLayoutFits(50U, 4U, 16U), false); + CHECK_EQ(FullmeshTraceLayoutFits(51U, 1U, 16U), false); + const size_t first = FullmeshTraceTaskSpanOffset(0U, 0U, 0U, 0U, 0U, 1U, 16U); + const size_t nextPhase = FullmeshTraceTaskSpanOffset(0U, 0U, 0U, 0U, 1U, 1U, 16U); + const size_t nextPeer = FullmeshTraceTaskSpanOffset(0U, 0U, 0U, 1U, 0U, 1U, 16U); + CHECK_EQ(nextPhase - first, sizeof(FullmeshTraceSpan)); + CHECK_EQ(nextPeer - first, kFullmeshTracePhaseCount * sizeof(FullmeshTraceSpan)); +} + void TestBuildAllToAllOutput() { constexpr int rankSize = 3; @@ -685,6 +710,7 @@ void TestAllToAllBigDataSource() int main() { + TestFullmeshTraceLayout(); TestAllToAllInputPattern(); TestAllToAllOutputValidation(); TestBuildAllToAllOutput(); From daf376a22fdbe3fd285b1d85c9f6d5c9c5d741f1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:47:35 +0800 Subject: [PATCH 041/163] feat(udma): define fullmesh trace layout --- tests/udma/CMakeLists.txt | 1 + tests/udma/demo/tilexr_udma_fullmesh_trace.h | 145 +++++++++++++++++++ 2 files changed, 146 insertions(+) create mode 100644 tests/udma/demo/tilexr_udma_fullmesh_trace.h diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 7e5320ae..d0ef03ea 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -244,6 +244,7 @@ if(BUILD_TILEXR_UDMA_DEMO) -o "${TILEXR_UDMA_DEMO_KERNEL_SO}" DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_fullmesh_trace.h" "${TILEXR_ROOT}/src/include/tilexr_udma.h" "${TILEXR_ROOT}/src/include/tilexr_data_as_flag.h" VERBATIM diff --git a/tests/udma/demo/tilexr_udma_fullmesh_trace.h b/tests/udma/demo/tilexr_udma_fullmesh_trace.h new file mode 100644 index 00000000..928854e6 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_fullmesh_trace.h @@ -0,0 +1,145 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_FULLMESH_TRACE_H +#define TILEXR_UDMA_FULLMESH_TRACE_H + +#include +#include +#include + +namespace TileXR { +namespace Demo { + +constexpr uint32_t kFullmeshTraceMagic = 0x464d5452U; // "FMTR" +constexpr uint32_t kFullmeshTraceVersion = 1U; +constexpr size_t kFullmeshTraceBytes = 8ULL * 1024ULL * 1024ULL; +constexpr size_t kFullmeshTraceHeaderBytes = 4096ULL; +constexpr uint32_t kFullmeshTraceMaxIterations = 50U; +constexpr uint32_t kFullmeshTraceMaxCores = 35U; +constexpr uint32_t kFullmeshTraceKernelRegions = 2U; +constexpr uint32_t kFullmeshTracePhaseCount = 14U; +constexpr uint64_t kFullmeshTraceCyclesPerUs = 1000ULL; + +enum FullmeshTraceKernelRegion : uint32_t { + kFullmeshTraceKernel = 0U, + kFullmeshTraceWork = 1U, +}; + +enum FullmeshTracePhase : uint32_t { + kFullmeshTracePhasePass = 0U, + kFullmeshTracePhaseSelfCopy = 1U, + kFullmeshTracePhasePeerCopy = 2U, + kFullmeshTracePhasePublishCopyReady = 3U, + kFullmeshTracePhaseWaitCopyReady = 4U, + kFullmeshTracePhaseDataPut = 5U, + kFullmeshTracePhaseQuiet = 6U, + kFullmeshTracePhaseSegmentDone = 7U, + kFullmeshTracePhasePublishReady = 8U, + kFullmeshTracePhaseWaitReady = 9U, + kFullmeshTracePhaseOutputCopy = 10U, + kFullmeshTracePhasePublishRecvDone = 11U, + kFullmeshTracePhaseWaitRecvDone = 12U, + kFullmeshTracePhaseAck = 13U, +}; + +struct FullmeshTraceSpan { + uint64_t beginCycle; + uint64_t endCycle; +}; + +struct FullmeshTraceHeader { + uint32_t magic; + uint32_t version; + uint32_t rank; + uint32_t iterationCount; + uint32_t passCount; + uint32_t coreCount; + uint32_t rankSize; + uint32_t phaseCount; + uint64_t cyclesPerUs; + uint64_t traceBytes; + uint64_t kernelSpanOffset; + uint64_t taskSpanOffset; +}; + +constexpr size_t FullmeshTraceKernelSpanCount() +{ + return static_cast(kFullmeshTraceMaxIterations) * + kFullmeshTraceMaxCores * kFullmeshTraceKernelRegions; +} + +constexpr size_t FullmeshTraceKernelSpanOffset( + uint32_t iteration, uint32_t core, uint32_t region) +{ + return kFullmeshTraceHeaderBytes + + ((static_cast(iteration) * kFullmeshTraceMaxCores + core) * + kFullmeshTraceKernelRegions + region) * sizeof(FullmeshTraceSpan); +} + +constexpr size_t FullmeshTraceTaskSpanBaseOffset() +{ + return kFullmeshTraceHeaderBytes + + FullmeshTraceKernelSpanCount() * sizeof(FullmeshTraceSpan); +} + +inline bool FullmeshTraceMultiply(size_t lhs, size_t rhs, size_t& result) +{ + if (lhs != 0U && rhs > std::numeric_limits::max() / lhs) { + result = std::numeric_limits::max(); + return false; + } + result = lhs * rhs; + return true; +} + +inline size_t FullmeshTraceLayoutBytes( + uint32_t iterationCount, uint32_t passCount, uint32_t rankSize) +{ + size_t count = iterationCount; + size_t next = 0U; + const size_t factors[] = { + kFullmeshTraceMaxCores, passCount, rankSize, + kFullmeshTracePhaseCount, sizeof(FullmeshTraceSpan)}; + for (size_t factor : factors) { + if (!FullmeshTraceMultiply(count, factor, next)) { + return std::numeric_limits::max(); + } + count = next; + } + if (count > std::numeric_limits::max() - FullmeshTraceTaskSpanBaseOffset()) { + return std::numeric_limits::max(); + } + return FullmeshTraceTaskSpanBaseOffset() + count; +} + +inline bool FullmeshTraceLayoutFits( + uint32_t iterationCount, uint32_t passCount, uint32_t rankSize) +{ + return iterationCount > 0U && iterationCount <= kFullmeshTraceMaxIterations && + passCount > 0U && rankSize > 0U && + FullmeshTraceLayoutBytes(iterationCount, passCount, rankSize) <= kFullmeshTraceBytes; +} + +inline size_t FullmeshTraceTaskSpanOffset( + uint32_t iteration, uint32_t core, uint32_t pass, uint32_t peer, + uint32_t phase, uint32_t passCount, uint32_t rankSize) +{ + const size_t index = + ((((static_cast(iteration) * kFullmeshTraceMaxCores + core) * + passCount + pass) * rankSize + peer) * kFullmeshTracePhaseCount) + phase; + return FullmeshTraceTaskSpanBaseOffset() + index * sizeof(FullmeshTraceSpan); +} + +static_assert(sizeof(FullmeshTraceSpan) == 16U, "fullmesh trace span must contain two uint64 timestamps"); +static_assert(sizeof(FullmeshTraceHeader) <= kFullmeshTraceHeaderBytes, + "fullmesh trace header must fit its region"); +static_assert(FullmeshTraceTaskSpanBaseOffset() < kFullmeshTraceBytes, + "fullmesh trace kernel spans must fit in 8 MiB"); + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_FULLMESH_TRACE_H From 08aafd2ab8a2e7ef01d536973076c642d073c72e Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:51:08 +0800 Subject: [PATCH 042/163] feat(udma): convert fullmesh traces to Chrome JSON --- .../tilexr_udma_fullmesh_trace_to_chrome.py | 215 ++++++++++++++++++ ...st_tilexr_udma_fullmesh_trace_to_chrome.py | 106 +++++++++ 2 files changed, 321 insertions(+) create mode 100644 tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py create mode 100644 tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py diff --git a/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py new file mode 100644 index 00000000..1da39d0e --- /dev/null +++ b/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +import argparse +import json +import struct +from pathlib import Path + + +TRACE_MAGIC = 0x464D5452 +TRACE_VERSION = 1 +TRACE_BYTES = 8 * 1024 * 1024 +HEADER_BYTES = 4096 +MAX_ITERATIONS = 50 +MAX_CORES = 35 +KERNEL_REGIONS = 2 +PHASE_COUNT = 14 +SPAN_BYTES = 16 +HEADER_FORMAT = "<8I4Q" +TASK_BASE_OFFSET = HEADER_BYTES + MAX_ITERATIONS * MAX_CORES * KERNEL_REGIONS * SPAN_BYTES + +PHASE_NAMES = ( + "pass", + "self-copy", + "peer-copy", + "publish-copy-ready", + "wait-copy-ready", + "data-put", + "quiet", + "segment-done", + "publish-ready", + "wait-ready", + "output-copy", + "publish-recv-done", + "wait-recv-done", + "ACK", +) + + +def kernel_span_offset(iteration, core, region): + index = (iteration * MAX_CORES + core) * KERNEL_REGIONS + region + return HEADER_BYTES + index * SPAN_BYTES + + +def task_span_offset(iteration, core, pass_index, peer, phase, pass_count, rank_size): + index = (((((iteration * MAX_CORES + core) * pass_count + pass_index) * + rank_size + peer) * PHASE_COUNT) + phase) + return TASK_BASE_OFFSET + index * SPAN_BYTES + + +def layout_bytes(iteration_count, pass_count, rank_size): + return TASK_BASE_OFFSET + ( + iteration_count * MAX_CORES * pass_count * rank_size * PHASE_COUNT * SPAN_BYTES) + + +def _core_role(core): + if core < 16: + return "copy" + if core == 16: + return "remote-primary" + if core == 17: + return "remote-secondary" + if core == 18: + return "local-send" + return "recv-copy" + + +def _read_span(data, offset, label): + if offset < 0 or offset + SPAN_BYTES > len(data): + raise ValueError(f"span offset out of range for {label}: {offset}") + begin, end = struct.unpack_from(" TRACE_BYTES: + raise ValueError(f"trace capacity exceeded in {path}: required={required} capacity={TRACE_BYTES}") + return {"path": str(path), "header": header, "data": data} + + +def _metadata(name, pid, tid, value): + return {"name": name, "ph": "M", "pid": pid, "tid": tid, "args": {"name": value}} + + +def _span_event(name, category, pid, tid, begin, end, base, cycles_per_us, args): + return { + "name": name, + "cat": category, + "ph": "X", + "pid": pid, + "tid": tid, + "ts": (begin - base) / cycles_per_us, + "dur": (end - begin) / cycles_per_us, + "args": {"beginCycle": begin, "endCycle": end, **args}, + } + + +def build_chrome_trace(rank_traces): + events = [] + sources = [] + for rank_trace in sorted(rank_traces, key=lambda item: item["header"]["rank"]): + header = rank_trace["header"] + data = rank_trace["data"] + rank = header["rank"] + sources.append(rank_trace["path"]) + bases = [] + for iteration in range(header["iteration_count"]): + spans = [] + for core in range(header["core_count"]): + span = _read_span( + data, kernel_span_offset(iteration, core, 0), + f"kernel rank={rank} iter={iteration} core={core}") + if span is not None: + spans.append(span) + if not spans: + raise ValueError(f"rank {rank} iteration {iteration} contains no kernel spans") + bases.append(min(begin for begin, _ in spans)) + + events.append(_metadata("process_name", rank, 0, f"rank {rank}")) + for core in range(header["core_count"]): + events.append(_metadata("thread_name", rank, core, f"core{core} {_core_role(core)}")) + + for iteration in range(header["iteration_count"]): + base = bases[iteration] + for core in range(header["core_count"]): + role = _core_role(core) + for region, name in ((0, "kernel"), (1, "work")): + span = _read_span( + data, kernel_span_offset(iteration, core, region), + f"kernel rank={rank} iter={iteration} core={core} region={region}") + if span is not None: + events.append(_span_event( + name, "kernel", rank, core, span[0], span[1], base, + header["cycles_per_us"], + {"iteration": iteration, "role": role})) + for pass_index in range(header["pass_count"]): + for peer in range(header["rank_size"]): + for phase in range(header["phase_count"]): + span = _read_span( + data, + task_span_offset( + iteration, core, pass_index, peer, phase, + header["pass_count"], header["rank_size"]), + f"task rank={rank} iter={iteration} core={core} " + f"pass={pass_index} peer={peer} phase={phase}") + if span is None: + continue + events.append(_span_event( + PHASE_NAMES[phase], role, rank, core, + span[0], span[1], base, header["cycles_per_us"], + {"iteration": iteration, "pass": pass_index, + "peer": peer, "phase": phase, "role": role})) + return { + "traceEvents": events, + "otherData": { + "sources": sources, + "clock": "GetSystemCycle normalized independently per rank and iteration", + }, + } + + +def main(): + parser = argparse.ArgumentParser(description="Convert TileXR full-mesh GM traces to Chrome trace JSON") + parser.add_argument("inputs", nargs="+", type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + trace = build_chrome_trace([read_rank_trace(path) for path in args.inputs]) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(trace, separators=(",", ":")), encoding="utf-8") + print(f"wrote {args.output} events={len(trace['traceEvents'])}") + + +if __name__ == "__main__": + main() diff --git a/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py new file mode 100644 index 00000000..8b072b08 --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +import importlib.util +import json +import struct +import tempfile +import unittest +from pathlib import Path + + +MODULE_PATH = Path(__file__).resolve().parents[1] / "demo" / "tilexr_udma_fullmesh_trace_to_chrome.py" +SPEC = importlib.util.spec_from_file_location("tilexr_udma_fullmesh_trace_to_chrome", MODULE_PATH) +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +class FullmeshTraceConverterTest(unittest.TestCase): + def make_trace(self, path, *, rank=3, version=1, iteration_count=1, pass_count=1): + data = bytearray(MODULE.TRACE_BYTES) + struct.pack_into( + MODULE.HEADER_FORMAT, + data, + 0, + MODULE.TRACE_MAGIC, + version, + rank, + iteration_count, + pass_count, + MODULE.MAX_CORES, + 16, + MODULE.PHASE_COUNT, + 1000, + MODULE.TRACE_BYTES, + MODULE.HEADER_BYTES, + MODULE.TASK_BASE_OFFSET, + ) + struct.pack_into(" Date: Fri, 17 Jul 2026 17:51:55 +0800 Subject: [PATCH 043/163] test(udma): require fullmesh trace collection --- tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 4890c37e..ebe51094 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -661,6 +661,14 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(udma, "reinterpret_cast(addr) + length - 1"); CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY"); CHECK_CONTAINS(demo, "skip result validation for bigdata remote-put-only profile"); + CHECK_CONTAINS(demo, "TILEXR_UDMA_FULLMESH_TRACE"); + CHECK_CONTAINS(demo, "TILEXR_UDMA_FULLMESH_TRACE_DIR"); + CHECK_CONTAINS(demo, "FullmeshTraceLayoutFits("); + CHECK_CONTAINS(demo, "aclrtMalloc(&fullmeshTraceDevice"); + CHECK_CONTAINS(demo, "aclrtMemcpy D2H fullmesh trace"); + CHECK_CONTAINS(demo, "tilexr_fullmesh_trace_rank_"); + CHECK_CONTAINS(demo, "fullmeshTraceIteration"); + CHECK_CONTAINS(demo, "reinterpret_cast(fullmeshTraceDevice)"); const std::string fullMeshReady = SliceBetween( kernel, "BigDataPublishReadySignal", "BigDataRemoteSendSegmentWorker"); From 99063181ab57df22227c73c18ebc10ef6fc642fc Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:56:19 +0800 Subject: [PATCH 044/163] feat(udma): collect fullmesh GM traces --- tests/udma/demo/tilexr_udma_demo.cpp | 120 +++++++++++++++++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 9 +- .../unit/test_tilexr_udma_alltoall_layout.cpp | 2 +- 3 files changed, 127 insertions(+), 4 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 3d36ba4e..dffb3727 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -24,6 +25,7 @@ #include "tilexr_types.h" #include "tilexr_udma_allreduce_layout.h" #include "tilexr_udma_alltoall_layout.h" +#include "tilexr_udma_fullmesh_trace.h" extern void launch_tilexr_udma_all_gather( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR data, GM_ADDR debug, int32_t elementsPerRank); @@ -48,7 +50,8 @@ extern void launch_tilexr_udma_all_to_all_fused( int32_t chunkElements, uint32_t passCount, uint32_t loopCount); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, - GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, + GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, + int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, @@ -204,6 +207,26 @@ bool CopyDeviceToHost(int rank, void* dst, size_t dstSize, const void* src, size return CheckAcl(rank, "aclrtMemcpy D2H " + name, ret); } +bool WriteFullmeshTraceBinary(int rank, const std::string& directory, const std::vector& data) +{ + const std::string path = directory + "/tilexr_fullmesh_trace_rank_" + + std::to_string(rank) + ".bin"; + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output.is_open()) { + std::cerr << "[rank " << rank << "] ERROR: open fullmesh trace output failed path=" + << path << std::endl; + return false; + } + output.write(reinterpret_cast(data.data()), static_cast(data.size())); + if (!output.good()) { + std::cerr << "[rank " << rank << "] ERROR: write fullmesh trace output failed path=" + << path << std::endl; + return false; + } + PrintStatus(rank, "fullmesh trace output=" + path + " bytes=" + std::to_string(data.size())); + return true; +} + BarrierEndpoint GetBarrierEndpoint() { std::string host = "127.0.0.1"; @@ -703,6 +726,15 @@ int main(int argc, char** argv) isAllToAll ? TileXR::Demo::PlanAllToAllBigDataUdma( rankSize, elementsPerRank, forceBigData35Core, bigDataRanksPerNode) : TileXR::Demo::AllToAllBigDataPlan {}; + const bool fullmeshTraceRequested = + testType == 7 && GetEnvInt("TILEXR_UDMA_FULLMESH_TRACE", 0) != 0; + const bool fullmeshTraceEnabled = fullmeshTraceRequested && + !bigDataRemotePutOnly && TileXR::Demo::AllToAllBigDataIsMultiNode( + rankSize, bigDataRanksPerNode); + const char* fullmeshTraceDirEnv = std::getenv("TILEXR_UDMA_FULLMESH_TRACE_DIR"); + const std::string fullmeshTraceDir = + fullmeshTraceDirEnv != nullptr && fullmeshTraceDirEnv[0] != '\0' ? + fullmeshTraceDirEnv : "."; if (testType == 7 && !TileXR::Demo::AllToAllBigDataValidTopology(rankSize, bigDataRanksPerNode)) { std::cerr << "[rank " << rank << "] ERROR: bigdata alltoall multi-node requires rankSize multiple of ranksPerNode" @@ -891,6 +923,14 @@ int main(int argc, char** argv) if (testType == 7) { void* bigInput = nullptr; void* bigOutput = nullptr; + void* fullmeshTraceDevice = nullptr; + std::vector hostFullmeshTrace; + auto freeFullmeshTrace = [&]() { + if (fullmeshTraceDevice != nullptr) { + aclrtFree(fullmeshTraceDevice); + fullmeshTraceDevice = nullptr; + } + }; const size_t bigDataBytes = dataBytes; if (!CheckAcl(rank, "aclrtMalloc bigdata input", aclrtMalloc(&bigInput, bigDataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || @@ -944,6 +984,70 @@ int main(int argc, char** argv) " chunkElements=" + std::to_string(bigDataPlan.chunkElements) + " repeat=" + std::to_string(allToAllRepeat) + " profileStage=" + std::to_string(bigDataProfileStage)); + if (fullmeshTraceRequested && !fullmeshTraceEnabled) { + std::cerr << "[rank " << rank << "] ERROR: fullmesh trace requires multi-node " + << "non-remote-put-only bigdata mode" << std::endl; + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + if (fullmeshTraceEnabled) { + if (!TileXR::Demo::FullmeshTraceLayoutFits( + static_cast(allToAllRepeat), bigDataPlan.passCount, + static_cast(rankSize))) { + std::cerr << "[rank " << rank << "] ERROR: fullmesh trace dimensions exceed capacity" + << " repeat=" << allToAllRepeat + << " passCount=" << bigDataPlan.passCount + << " rankSize=" << rankSize + << " requiredBytes=" << TileXR::Demo::FullmeshTraceLayoutBytes( + static_cast(allToAllRepeat), bigDataPlan.passCount, + static_cast(rankSize)) + << " capacityBytes=" << TileXR::Demo::kFullmeshTraceBytes << std::endl; + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + hostFullmeshTrace.assign(TileXR::Demo::kFullmeshTraceBytes, 0U); + TileXR::Demo::FullmeshTraceHeader header {}; + header.magic = TileXR::Demo::kFullmeshTraceMagic; + header.version = TileXR::Demo::kFullmeshTraceVersion; + header.rank = static_cast(rank); + header.iterationCount = static_cast(allToAllRepeat); + header.passCount = bigDataPlan.passCount; + header.coreCount = TileXR::Demo::kFullmeshTraceMaxCores; + header.rankSize = static_cast(rankSize); + header.phaseCount = TileXR::Demo::kFullmeshTracePhaseCount; + header.cyclesPerUs = TileXR::Demo::kFullmeshTraceCyclesPerUs; + header.traceBytes = TileXR::Demo::kFullmeshTraceBytes; + header.kernelSpanOffset = TileXR::Demo::kFullmeshTraceHeaderBytes; + header.taskSpanOffset = TileXR::Demo::FullmeshTraceTaskSpanBaseOffset(); + std::memcpy(hostFullmeshTrace.data(), &header, sizeof(header)); + if (!CheckAcl(rank, "aclrtMalloc fullmesh trace", + aclrtMalloc(&fullmeshTraceDevice, TileXR::Demo::kFullmeshTraceBytes, + ACL_MEM_MALLOC_HUGE_FIRST)) || + !CopyHostToDevice(rank, fullmeshTraceDevice, TileXR::Demo::kFullmeshTraceBytes, + hostFullmeshTrace.data(), hostFullmeshTrace.size(), "fullmesh trace")) { + freeFullmeshTrace(); + if (udmaRegistered) { + CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); + udmaRegistered = false; + } + aclrtFree(bigInput); + aclrtFree(bigOutput); + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } + } if (!CheckAcl(rank, "aclrtSynchronizeStream bigdata prime", aclrtSynchronizeStream(stream)) || !DemoBarrierAll(rank, rankSize, "all ranks bigdata prime")) { if (udmaRegistered) { @@ -952,6 +1056,7 @@ int main(int argc, char** argv) } aclrtFree(bigInput); aclrtFree(bigOutput); + freeFullmeshTrace(); Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } @@ -964,10 +1069,12 @@ int main(int argc, char** argv) auto a2aStart = std::chrono::steady_clock::now(); for (int iter = 0; iter < allToAllRepeat; ++iter) { const uint64_t kernelLoopBase = static_cast(iter); + const uint32_t fullmeshTraceIteration = static_cast(iter); launch_tilexr_udma_all_to_all_bigdata( bigDataBlockDim, stream, commArgsDev, reinterpret_cast(bigInput), reinterpret_cast(bigOutput), reinterpret_cast(registeredMemory), reinterpret_cast(debug), + reinterpret_cast(fullmeshTraceDevice), fullmeshTraceIteration, elementsPerRank, 0, bigDataPlan.copyDoneOffset, bigDataPlan.recvCopyDoneOffset, bigDataPlan.remoteSendDoneOffset, @@ -983,6 +1090,7 @@ int main(int argc, char** argv) } aclrtFree(bigInput); aclrtFree(bigOutput); + freeFullmeshTrace(); Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } @@ -1004,6 +1112,13 @@ int main(int argc, char** argv) } bigDataCopyBackOk = CopyDeviceToHost(rank, hostDebug.data(), hostDebug.size() * sizeof(int32_t), debug, hostDebug.size() * sizeof(int32_t), "debug after bigdata alltoall") && bigDataCopyBackOk; + if (fullmeshTraceEnabled) { + const bool traceCopyOk = CopyDeviceToHost( + rank, hostFullmeshTrace.data(), hostFullmeshTrace.size(), fullmeshTraceDevice, + TileXR::Demo::kFullmeshTraceBytes, "fullmesh trace"); + bigDataCopyBackOk = traceCopyOk && + WriteFullmeshTraceBinary(rank, fullmeshTraceDir, hostFullmeshTrace) && bigDataCopyBackOk; + } if (!bigDataCopyBackOk) { if (udmaRegistered) { CheckTileXR(rank, "TileXRUDMAUnregister", TileXRUDMAUnregister(comm, udmaHandle)); @@ -1011,6 +1126,7 @@ int main(int argc, char** argv) } aclrtFree(bigInput); aclrtFree(bigOutput); + freeFullmeshTrace(); Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } @@ -1027,6 +1143,7 @@ int main(int argc, char** argv) } aclrtFree(bigInput); aclrtFree(bigOutput); + freeFullmeshTrace(); Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } @@ -1036,6 +1153,7 @@ int main(int argc, char** argv) } aclrtFree(bigInput); aclrtFree(bigOutput); + freeFullmeshTrace(); } else if (testType == 6) { // Forced UDMA alltoall (no IPC fallback). Single kernel launch loops // REPEAT times internally; stream sync only after all loops. diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 0cee2d74..33a975e4 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -2689,11 +2689,14 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_fused_kernel( extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, GM_ADDR debugGM, + GM_ADDR fullmeshTraceGM, uint32_t fullmeshTraceIteration, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage, uint32_t force35CoreFlag) { + (void)fullmeshTraceGM; + (void)fullmeshTraceIteration; if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_FULL) { profileStage = TILEXR_BIGDATA_PROFILE_STAGE_FULL; } @@ -2928,14 +2931,16 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, - GM_ADDR udmaMem, GM_ADDR debug, int32_t elementsPerPeer, + GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, + int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage, uint32_t force35Core) { tilexr_udma_all_to_all_bigdata_kernel<<>>( - commArgs, input, output, udmaMem, debug, elementsPerPeer, + commArgs, input, output, udmaMem, debug, fullmeshTrace, fullmeshTraceIteration, + elementsPerPeer, dataOffset, copyDoneOffset, recvCopyDoneOffset, remoteSendDoneOffset, readySignalOffset, ackSignalOffset, chunkElements, passCount, loopCount, kernelLoopBase, profileStage, force35Core); } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index ebe51094..f772d711 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -665,7 +665,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "TILEXR_UDMA_FULLMESH_TRACE_DIR"); CHECK_CONTAINS(demo, "FullmeshTraceLayoutFits("); CHECK_CONTAINS(demo, "aclrtMalloc(&fullmeshTraceDevice"); - CHECK_CONTAINS(demo, "aclrtMemcpy D2H fullmesh trace"); + CHECK_CONTAINS(demo, "TileXR::Demo::kFullmeshTraceBytes, \"fullmesh trace\""); CHECK_CONTAINS(demo, "tilexr_fullmesh_trace_rank_"); CHECK_CONTAINS(demo, "fullmeshTraceIteration"); CHECK_CONTAINS(demo, "reinterpret_cast(fullmeshTraceDevice)"); From 8c6c7bc806158277e04a4545b39099b0b0d1b32f Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 17:58:40 +0800 Subject: [PATCH 045/163] test(udma): require per-core fullmesh trace spans --- .../unit/test_tilexr_udma_alltoall_layout.cpp | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index f772d711..e2161a03 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -672,6 +672,8 @@ void TestAllToAllBigDataSource() const std::string fullMeshReady = SliceBetween( kernel, "BigDataPublishReadySignal", "BigDataRemoteSendSegmentWorker"); + const std::string localSend = SliceBetween( + kernel, "BigDataSendPeerWorker", "BigDataRecvPeerWorker"); const std::string fullMeshSend = SliceBetween( kernel, "BigDataRemoteSendSegmentWorker", "BigDataRemotePutOnlySendWorker"); const std::string fullMeshRecv = SliceBetween( @@ -686,6 +688,27 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(fullMeshRecv, "BigDataRemoteRegisteredControlSlot("); CHECK_CONTAINS(fullMeshRecv, "BigDataStoreTokenMte(remoteAck, token, relayLocal)"); CHECK_NOT_CONTAINS(fullMeshRecv, "BigDataPublishAckSignalUdma("); + CHECK_CONTAINS(kernel, "tilexr_udma_fullmesh_trace.h"); + CHECK_CONTAINS(kernel, "BigDataFullmeshTraceRecordKernelSpan"); + CHECK_CONTAINS(kernel, "BigDataFullmeshTraceRecordTaskSpan"); + CHECK_CONTAINS(kernel, "if (trace == nullptr"); + CHECK_CONTAINS(kernel, "AscendC::GetSystemCycle()"); + CHECK_CONTAINS(kernel, "kFullmeshTracePhaseSelfCopy"); + CHECK_CONTAINS(kernel, "kFullmeshTracePhasePeerCopy"); + CHECK_CONTAINS(kernel, "kFullmeshTracePhasePublishCopyReady"); + CHECK_CONTAINS(fullMeshSend, "kFullmeshTracePhaseWaitCopyReady"); + CHECK_CONTAINS(fullMeshSend, "kFullmeshTracePhaseDataPut"); + CHECK_CONTAINS(fullMeshSend, "kFullmeshTracePhaseQuiet"); + CHECK_CONTAINS(fullMeshSend, "kFullmeshTracePhaseSegmentDone"); + CHECK_CONTAINS(fullMeshSend, "kFullmeshTracePhasePublishReady"); + CHECK_CONTAINS(localSend, "kFullmeshTracePhaseWaitCopyReady"); + CHECK_CONTAINS(localSend, "kFullmeshTracePhaseDataPut"); + CHECK_CONTAINS(localSend, "kFullmeshTracePhaseQuiet"); + CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhaseWaitReady"); + CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhaseOutputCopy"); + CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhasePublishRecvDone"); + CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhaseWaitRecvDone"); + CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhaseAck"); const std::string remotePutOnlySend = SliceBetween( kernel, "BigDataRemotePutOnlySendWorker", "BigDataRemotePutOnlyCheckIndex"); From 20e19efdf1f4bb2c967b8ae5f742102aba49f426 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 18:08:13 +0800 Subject: [PATCH 046/163] feat(udma): trace fullmesh core tasks --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 217 +++++++++++++++++++- 1 file changed, 207 insertions(+), 10 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 33a975e4..0d42a8db 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -6,6 +6,7 @@ #include "kernel_operator.h" #include "tilexr_data_as_flag.h" #include "tilexr_udma.h" +#include "tilexr_udma_fullmesh_trace.h" constexpr int32_t TILEXR_UDMA_DEMO_MAGIC = 0x5444554d; // "TDUM" constexpr uint64_t TILEXR_UDMA_DEMO_IPC_STAGING_OFFSET = TileXR::IPC_DATA_OFFSET; @@ -105,6 +106,53 @@ constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_ACK = 7; namespace { +__aicore__ inline uint64_t BigDataFullmeshTraceCycle(__gm__ uint8_t* trace) +{ + return trace == nullptr ? 0ULL : static_cast(AscendC::GetSystemCycle()); +} + +__aicore__ inline void BigDataFullmeshTraceStoreSpan( + __gm__ uint8_t* trace, size_t offset, uint64_t beginCycle, uint64_t endCycle) +{ + if (trace == nullptr || beginCycle == 0ULL || endCycle < beginCycle) { + return; + } + auto span = reinterpret_cast<__gm__ uint64_t*>(trace + offset); + span[0] = beginCycle; + span[1] = endCycle; +} + +__aicore__ inline void BigDataFullmeshTraceRecordKernelSpan( + __gm__ uint8_t* trace, uint32_t iteration, uint32_t core, uint32_t region, + uint64_t beginCycle, uint64_t endCycle) +{ + if (trace == nullptr || iteration >= TileXR::Demo::kFullmeshTraceMaxIterations || + core >= TileXR::Demo::kFullmeshTraceMaxCores || + region >= TileXR::Demo::kFullmeshTraceKernelRegions) { + return; + } + BigDataFullmeshTraceStoreSpan( + trace, TileXR::Demo::FullmeshTraceKernelSpanOffset(iteration, core, region), + beginCycle, endCycle); +} + +__aicore__ inline void BigDataFullmeshTraceRecordTaskSpan( + __gm__ uint8_t* trace, uint32_t iteration, uint32_t core, + uint32_t pass, int32_t peer, uint32_t phase, + uint32_t passCount, int32_t rankSize, uint64_t beginCycle, uint64_t endCycle) +{ + if (trace == nullptr || iteration >= TileXR::Demo::kFullmeshTraceMaxIterations || + core >= TileXR::Demo::kFullmeshTraceMaxCores || pass >= passCount || + peer < 0 || peer >= rankSize || phase >= TileXR::Demo::kFullmeshTracePhaseCount) { + return; + } + BigDataFullmeshTraceStoreSpan( + trace, TileXR::Demo::FullmeshTraceTaskSpanOffset( + iteration, core, pass, static_cast(peer), phase, + passCount, static_cast(rankSize)), + beginCycle, endCycle); +} + __aicore__ inline uint64_t AllToAllPayloadBytes(int32_t elementsPerPeer) { return static_cast(elementsPerPeer) * sizeof(int32_t); @@ -924,6 +972,7 @@ __aicore__ inline void BigDataSendPeerWorker( uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, uint64_t sendDataOffset, uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t readySignalOffset, uint64_t chunkBytesPerPeer, bool use35Core, uint64_t copyReadyOffset, + __gm__ uint8_t* fullmeshTrace, uint32_t traceIteration, uint32_t traceCore, AscendC::LocalTensor relayLocal) { if (peer < 0 || peer >= rankSize || peer == rank || shardCount == 0U || @@ -942,6 +991,7 @@ __aicore__ inline void BigDataSendPeerWorker( const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); + const uint64_t waitStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (use35Core) { if (!BigDataWaitCopyReady(udmaMem, debug, copyReadyOffset, slot, rankSize, shardCount, peer, TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY, token, loop, pass, relayLocal)) { @@ -961,6 +1011,11 @@ __aicore__ inline void BigDataSendPeerWorker( } } } + const uint64_t waitEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseWaitCopyReady, + passCount, rankSize, waitStart, waitEnd); if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { return; } @@ -1005,9 +1060,21 @@ __aicore__ inline void BigDataSendPeerWorker( static_cast(remoteReadyOffset); } + const uint64_t dataStart = BigDataFullmeshTraceCycle(fullmeshTrace); TileXR::UDMAPutSignalNbi(args, peer, localSrc, remoteDataOffset, chunkBytes, remoteReadyOffset, token); + const uint64_t dataEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseDataPut, + passCount, rankSize, dataStart, dataEnd); + const uint64_t quietStart = BigDataFullmeshTraceCycle(fullmeshTrace); uint32_t status = TileXR::UDMAQuietStatus(args, peer); + const uint64_t quietEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseQuiet, + passCount, rankSize, quietStart, quietEnd); if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { uint32_t wqeAfter = ld_dev(reinterpret_cast<__gm__ uint32_t*>(wqCtx->wqeCntAddr), 0); debug[TILEXR_UDMA_DEMO_DEBUG_WQE_AFTER_BASE + peer] = static_cast(wqeAfter); @@ -1022,7 +1089,9 @@ __aicore__ inline void BigDataRecvPeerWorker( uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, uint32_t profileStage, bool use35Core, uint32_t recvShard, uint32_t shardCount, uint64_t recvDataOffset, uint64_t recvCopyDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, - uint64_t chunkBytesPerPeer, AscendC::LocalTensor relayLocal) + uint64_t chunkBytesPerPeer, __gm__ uint8_t* fullmeshTrace, + uint32_t traceIteration, uint32_t traceCore, + AscendC::LocalTensor relayLocal) { const uint32_t dataShardCount = use35Core ? TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES : shardCount; @@ -1049,6 +1118,7 @@ __aicore__ inline void BigDataRecvPeerWorker( const uint32_t slot = BigDataDataSlot(globalPass, pass, use35Core); uint64_t observed = 0ULL; + const uint64_t waitReadyStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (use35Core) { if (recvShard == TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_SHARD) { observed = BigDataWaitTokenMte( @@ -1100,6 +1170,11 @@ __aicore__ inline void BigDataRecvPeerWorker( return; } } + const uint64_t waitReadyEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseWaitReady, + passCount, rankSize, waitReadyStart, waitReadyEnd); if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_WAIT_READY) { return; } @@ -1112,9 +1187,15 @@ __aicore__ inline void BigDataRecvPeerWorker( } auto dst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(peer) * elementsPerPeer + chunkOffset); + const uint64_t outputCopyStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (shardHasBytes) { BigDataCopyRangePingPong(dst + shardOffset, recvSlot + shardOffset, shardBytes, relayLocal); } + const uint64_t outputCopyEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseOutputCopy, + passCount, rankSize, outputCopyStart, outputCopyEnd); if (debug != nullptr && recvShard == 0U && loop == 0 && pass == 0 && peer < 16) { auto relayDst = output + static_cast(peer) * elementsPerPeer + chunkOffset; debug[TILEXR_UDMA_DEMO_DEBUG_RECV_SAMPLE_BASE + peer] = relayDst[0]; @@ -1123,11 +1204,18 @@ __aicore__ inline void BigDataRecvPeerWorker( return; } + const uint64_t publishDoneStart = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataStoreTokenMte( BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, recvShard), token, relayLocal); + const uint64_t publishDoneEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhasePublishRecvDone, + passCount, rankSize, publishDoneStart, publishDoneEnd); if (recvShard + 1U != shardCount) { if (use35Core && recvShard + 1U == dataShardCount) { + const uint64_t waitDoneStart = BigDataFullmeshTraceCycle(fullmeshTrace); for (uint32_t shard = 0U; shard < dataShardCount; ++shard) { observed = BigDataWaitTokenMte( BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, shard), @@ -1140,6 +1228,11 @@ __aicore__ inline void BigDataRecvPeerWorker( return; } } + const uint64_t waitDoneEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseWaitRecvDone, + passCount, rankSize, waitDoneStart, waitDoneEnd); auto remoteAck = BigDataRemoteRegisteredControlSlot( args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); @@ -1150,10 +1243,17 @@ __aicore__ inline void BigDataRecvPeerWorker( } return; } + const uint64_t ackStart = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataStoreTokenMte(remoteAck, token, relayLocal); + const uint64_t ackEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseAck, + passCount, rankSize, ackStart, ackEnd); } return; } + const uint64_t waitDoneStart = BigDataFullmeshTraceCycle(fullmeshTrace); for (uint32_t shard = 0U; shard < shardCount; ++shard) { observed = BigDataWaitTokenMte( BigDataControlSlot(udmaMem, recvCopyDoneOffset, slot, rankSize, shardCount, peer, shard), @@ -1166,6 +1266,11 @@ __aicore__ inline void BigDataRecvPeerWorker( return; } } + const uint64_t waitDoneEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseWaitRecvDone, + passCount, rankSize, waitDoneStart, waitDoneEnd); __gm__ uint64_t* remoteAck = nullptr; if (use35Core) { @@ -1181,7 +1286,13 @@ __aicore__ inline void BigDataRecvPeerWorker( } return; } + const uint64_t ackStart = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataStoreTokenMte(remoteAck, token, relayLocal); + const uint64_t ackEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseAck, + passCount, rankSize, ackStart, ackEnd); } __aicore__ inline bool BigDataWaitCopyDoneRange( @@ -1298,7 +1409,9 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( uint64_t kernelLoopBase, uint32_t profileStage, uint32_t shardCount, uint64_t sendDataOffset, uint64_t recvDataOffset, uint64_t copyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t chunkBytesPerPeer, - int32_t ranksPerNode, AscendC::LocalTensor relayLocal) + int32_t ranksPerNode, __gm__ uint8_t* fullmeshTrace, + uint32_t traceIteration, uint32_t traceCore, + AscendC::LocalTensor relayLocal) { if (peer < 0 || peer >= rankSize || peer == rank || !BigDataIsMultiNode(rankSize, ranksPerNode) || shardCount == 0U || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { @@ -1336,10 +1449,16 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( (segmentId == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT) ? TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY : TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY; + const uint64_t waitStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (!BigDataWaitCopyReady(udmaMem, debug, remoteSendDoneOffset, slot, rankSize, shardCount, peer, readyShard, token, loop, pass, relayLocal)) { return; } + const uint64_t waitEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseWaitCopyReady, + passCount, rankSize, waitStart, waitEnd); if (profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_SEND_SYNC) { return; } @@ -1358,21 +1477,40 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( uint32_t status = 0U; if (segmentBytes > 0U) { + const uint64_t dataStart = BigDataFullmeshTraceCycle(fullmeshTrace); TileXR::UDMAPutNbiOnQp( args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); + const uint64_t dataEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseDataPut, + passCount, rankSize, dataStart, dataEnd); + const uint64_t quietStart = BigDataFullmeshTraceCycle(fullmeshTrace); status = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + const uint64_t quietEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseQuiet, + passCount, rankSize, quietStart, quietEnd); } if (debug != nullptr && loop == 0 && pass == 0 && peer < 16) { debug[TILEXR_UDMA_DEMO_DEBUG_UDMA_STATUS_BASE + peer] = static_cast(status); } + const uint64_t doneStart = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataStoreTokenMte( BigDataControlSlot(udmaMem, remoteSendDoneOffset, slot, rankSize, shardCount, peer, segmentId), token, relayLocal); + const uint64_t doneEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhaseSegmentDone, + passCount, rankSize, doneStart, doneEnd); if (segmentId != TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT || profileStage <= TILEXR_BIGDATA_PROFILE_STAGE_DATA_PUT) { return; } + const uint64_t readyStart = BigDataFullmeshTraceCycle(fullmeshTrace); for (uint32_t done = 0U; done <= TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; ++done) { const uint64_t observed = BigDataWaitTokenMte( BigDataControlSlot(udmaMem, remoteSendDoneOffset, slot, rankSize, shardCount, peer, done), @@ -1390,6 +1528,11 @@ __aicore__ inline void BigDataRemoteSendSegmentWorker( BigDataStoreTokenMte(localReady, token, relayLocal); BigDataPublishReadySignal( args, peer, udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, token, qpIdx); + const uint64_t readyEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, traceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhasePublishReady, + passCount, rankSize, readyStart, readyEnd); } __aicore__ inline void BigDataRemotePutOnlySendWorker( @@ -1595,7 +1738,8 @@ __aicore__ inline void BigDataRunRoleForPeer( BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, - chunkBytesPerPeer, use35Core, copyDoneOffset, relayLocal); + chunkBytesPerPeer, use35Core, copyDoneOffset, + nullptr, 0U, 0U, relayLocal); return; } const uint32_t recvShard = @@ -1603,7 +1747,7 @@ __aicore__ inline void BigDataRunRoleForPeer( BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, use35Core, recvShard, shardCount, recvDataOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, - chunkBytesPerPeer, relayLocal); + chunkBytesPerPeer, nullptr, 0U, 0U, relayLocal); } } // namespace @@ -2695,8 +2839,6 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, uint64_t kernelLoopBase, uint32_t profileStage, uint32_t force35CoreFlag) { - (void)fullmeshTraceGM; - (void)fullmeshTraceIteration; if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_FULL) { profileStage = TILEXR_BIGDATA_PROFILE_STAGE_FULL; } @@ -2713,6 +2855,12 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( const bool force35Core = (force35CoreFlag & 0x1U) != 0U; const bool remotePutOnly = (force35CoreFlag & 0x2U) != 0U; const int32_t ranksPerNode = BigDataNormalizeRanksPerNode(static_cast(force35CoreFlag >> 8U)); + auto fullmeshTraceBuffer = reinterpret_cast<__gm__ uint8_t*>(fullmeshTraceGM); + auto fullmeshTrace = (!remotePutOnly && force35Core && + blockIdx >= 0 && blockIdx < static_cast(TileXR::Demo::kFullmeshTraceMaxCores)) ? + fullmeshTraceBuffer : nullptr; + const uint32_t traceCore = static_cast(blockIdx); + const uint64_t kernelStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; @@ -2801,8 +2949,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( return; } + const uint64_t workStart = BigDataFullmeshTraceCycle(fullmeshTrace); for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { + const uint64_t passStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (remotePutOnly) { for (int32_t sendTask = blockIdx; sendTask < sendTaskCount; sendTask += static_cast(activeBlockDim)) { @@ -2840,9 +2990,15 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( continue; } if (isCopyCore) { + const uint64_t selfCopyStart = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataRunSelfCopyShard(rank, rankSize, input, output, elementsPerPeer, effectiveChunkElements, pass, copyShard, shardCount, relayLocal); + const uint64_t selfCopyEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, rank, + TileXR::Demo::kFullmeshTracePhaseSelfCopy, + passCount, rankSize, selfCopyStart, selfCopyEnd); } for (int32_t taskIndex = 0; taskIndex < taskCount; ++taskIndex) { @@ -2854,25 +3010,38 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( } if (isCopyCore) { + const uint64_t peerCopyStart = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, use35Core, copyShard, shardCount, sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); + const uint64_t peerCopyEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhasePeerCopy, + passCount, rankSize, peerCopyStart, peerCopyEnd); if (profileStage > TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY) { const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); const uint64_t token = BigDataPassToken(globalPass); const uint32_t slot = BigDataDataSlot(globalPass, pass, true); if (!isLocalPeer && copyShard == 0U) { + const uint64_t publishStart = BigDataFullmeshTraceCycle(fullmeshTrace); (void)BigDataPublishCopyReadyRange( udmaMem, debug, copyDoneOffset, remoteSendDoneOffset, slot, rankSize, shardCount, peer, 0U, TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END, TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY, token, loop, pass, relayLocal); + const uint64_t publishEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhasePublishCopyReady, + passCount, rankSize, publishStart, publishEnd); } if (!isLocalPeer && copyShard == TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR) { + const uint64_t publishStart = BigDataFullmeshTraceCycle(fullmeshTrace); (void)BigDataPublishCopyReadyRange( udmaMem, debug, copyDoneOffset, remoteSendDoneOffset, slot, rankSize, shardCount, peer, @@ -2880,14 +3049,25 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES, TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY, token, loop, pass, relayLocal); + const uint64_t publishEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhasePublishCopyReady, + passCount, rankSize, publishStart, publishEnd); } if (isLocalPeer && copyShard == 0U) { + const uint64_t publishStart = BigDataFullmeshTraceCycle(fullmeshTrace); (void)BigDataPublishCopyReadyRange( udmaMem, debug, copyDoneOffset, remoteSendDoneOffset, slot, rankSize, shardCount, peer, 0U, TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES, TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY, token, loop, pass, relayLocal); + const uint64_t publishEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, peer, + TileXR::Demo::kFullmeshTracePhasePublishCopyReady, + passCount, rankSize, publishStart, publishEnd); } } } @@ -2898,7 +3078,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, remoteSendDoneOffset, readySignalOffset, - chunkBytesPerPeer, ranksPerNode, relayLocal); + chunkBytesPerPeer, ranksPerNode, fullmeshTrace, + fullmeshTraceIteration, traceCore, relayLocal); } if (!isLocalPeer && isRemoteSendSecondaryCore) { BigDataRemoteSendSegmentWorker(peer, @@ -2907,26 +3088,42 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, remoteSendDoneOffset, readySignalOffset, - chunkBytesPerPeer, ranksPerNode, relayLocal); + chunkBytesPerPeer, ranksPerNode, fullmeshTrace, + fullmeshTraceIteration, traceCore, relayLocal); } if (isLocalPeer && isLocalSendCore) { BigDataSendPeerWorker(peer, rank, rankSize, args, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, shardCount, sendDataOffset, recvDataOffset, copyDoneOffset, readySignalOffset, chunkBytesPerPeer, true, - remoteSendDoneOffset, relayLocal); + remoteSendDoneOffset, fullmeshTrace, + fullmeshTraceIteration, traceCore, relayLocal); } if (isRecvCore) { BigDataRecvPeerWorker(peer, rank, rankSize, args, output, udmaMem, debug, elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, profileStage, use35Core, recvShard, shardCount, recvDataOffset, recvCopyDoneOffset, readySignalOffset, ackSignalOffset, - chunkBytesPerPeer, relayLocal); + chunkBytesPerPeer, fullmeshTrace, + fullmeshTraceIteration, traceCore, relayLocal); } } + const uint64_t passEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, rank, + TileXR::Demo::kFullmeshTracePhasePass, + passCount, rankSize, passStart, passEnd); } } + const uint64_t workEnd = BigDataFullmeshTraceCycle(fullmeshTrace); BigDataKernelExitBarrier(); + const uint64_t kernelEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordKernelSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, + TileXR::Demo::kFullmeshTraceWork, workStart, workEnd); + BigDataFullmeshTraceRecordKernelSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, + TileXR::Demo::kFullmeshTraceKernel, kernelStart, kernelEnd); } void launch_tilexr_udma_all_to_all_bigdata( From f6019d697527c6d434b8af186c03f2646257c74d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 18:12:19 +0800 Subject: [PATCH 047/163] fix(udma): inline device trace offsets --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 0d42a8db..c6a23248 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -131,9 +131,12 @@ __aicore__ inline void BigDataFullmeshTraceRecordKernelSpan( region >= TileXR::Demo::kFullmeshTraceKernelRegions) { return; } + const size_t index = + (static_cast(iteration) * TileXR::Demo::kFullmeshTraceMaxCores + core) * + TileXR::Demo::kFullmeshTraceKernelRegions + region; BigDataFullmeshTraceStoreSpan( - trace, TileXR::Demo::FullmeshTraceKernelSpanOffset(iteration, core, region), - beginCycle, endCycle); + trace, TileXR::Demo::kFullmeshTraceHeaderBytes + + index * sizeof(TileXR::Demo::FullmeshTraceSpan), beginCycle, endCycle); } __aicore__ inline void BigDataFullmeshTraceRecordTaskSpan( @@ -146,11 +149,17 @@ __aicore__ inline void BigDataFullmeshTraceRecordTaskSpan( peer < 0 || peer >= rankSize || phase >= TileXR::Demo::kFullmeshTracePhaseCount) { return; } + const size_t kernelSpanBytes = + static_cast(TileXR::Demo::kFullmeshTraceMaxIterations) * + TileXR::Demo::kFullmeshTraceMaxCores * TileXR::Demo::kFullmeshTraceKernelRegions * + sizeof(TileXR::Demo::FullmeshTraceSpan); + const size_t index = + ((((static_cast(iteration) * TileXR::Demo::kFullmeshTraceMaxCores + core) * + passCount + pass) * static_cast(rankSize) + static_cast(peer)) * + TileXR::Demo::kFullmeshTracePhaseCount) + phase; BigDataFullmeshTraceStoreSpan( - trace, TileXR::Demo::FullmeshTraceTaskSpanOffset( - iteration, core, pass, static_cast(peer), phase, - passCount, static_cast(rankSize)), - beginCycle, endCycle); + trace, TileXR::Demo::kFullmeshTraceHeaderBytes + kernelSpanBytes + + index * sizeof(TileXR::Demo::FullmeshTraceSpan), beginCycle, endCycle); } __aicore__ inline uint64_t AllToAllPayloadBytes(int32_t elementsPerPeer) From 23dd1e4bed01865a172b3efa24b3a4916678e7ae Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 18:16:42 +0800 Subject: [PATCH 048/163] fix(udma): enable tracing for physical fullmesh --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 8 ++++---- tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index c6a23248..57ea0f2b 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -2865,11 +2865,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( const bool remotePutOnly = (force35CoreFlag & 0x2U) != 0U; const int32_t ranksPerNode = BigDataNormalizeRanksPerNode(static_cast(force35CoreFlag >> 8U)); auto fullmeshTraceBuffer = reinterpret_cast<__gm__ uint8_t*>(fullmeshTraceGM); - auto fullmeshTrace = (!remotePutOnly && force35Core && - blockIdx >= 0 && blockIdx < static_cast(TileXR::Demo::kFullmeshTraceMaxCores)) ? - fullmeshTraceBuffer : nullptr; const uint32_t traceCore = static_cast(blockIdx); - const uint64_t kernelStart = BigDataFullmeshTraceCycle(fullmeshTrace); if (blockIdx == 0 && debug != nullptr) { debug[0] = TILEXR_UDMA_DEMO_MAGIC; @@ -2891,6 +2887,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( return; } const bool use35Core = BigDataUse35Core(rankSize, force35Core, ranksPerNode); + auto fullmeshTrace = (!remotePutOnly && use35Core && + blockIdx >= 0 && blockIdx < static_cast(TileXR::Demo::kFullmeshTraceMaxCores)) ? + fullmeshTraceBuffer : nullptr; + const uint64_t kernelStart = BigDataFullmeshTraceCycle(fullmeshTrace); const uint32_t shardCount = BigDataShardCount(rankSize, force35Core, ranksPerNode); const int32_t effectiveChunkElements = chunkElements > 0 ? chunkElements : elementsPerPeer; const uint64_t chunkBytesPerPeer = BigDataChunkBytesPerPeer(use35Core, effectiveChunkElements); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index e2161a03..b7b91ef0 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -693,6 +693,7 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "BigDataFullmeshTraceRecordTaskSpan"); CHECK_CONTAINS(kernel, "if (trace == nullptr"); CHECK_CONTAINS(kernel, "AscendC::GetSystemCycle()"); + CHECK_CONTAINS(kernel, "auto fullmeshTrace = (!remotePutOnly && use35Core &&"); CHECK_CONTAINS(kernel, "kFullmeshTracePhaseSelfCopy"); CHECK_CONTAINS(kernel, "kFullmeshTracePhasePeerCopy"); CHECK_CONTAINS(kernel, "kFullmeshTracePhasePublishCopyReady"); From c02f3bf27bc54f0d204f7052f2ba50119cae7492 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 18:27:37 +0800 Subject: [PATCH 049/163] fix(udma): separate fullmesh trace iterations --- .../tilexr_udma_fullmesh_trace_to_chrome.py | 36 +++++++++++++------ ...st_tilexr_udma_fullmesh_trace_to_chrome.py | 18 ++++++++++ 2 files changed, 44 insertions(+), 10 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py index 1da39d0e..a9654654 100644 --- a/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py @@ -33,6 +33,7 @@ "wait-recv-done", "ACK", ) +ITERATION_GAP_US = 1.0 def kernel_span_offset(iteration, core, region): @@ -123,14 +124,14 @@ def _metadata(name, pid, tid, value): return {"name": name, "ph": "M", "pid": pid, "tid": tid, "args": {"name": value}} -def _span_event(name, category, pid, tid, begin, end, base, cycles_per_us, args): +def _span_event(name, category, pid, tid, begin, end, base, cycles_per_us, args, offset_us=0.0): return { "name": name, "cat": category, "ph": "X", "pid": pid, "tid": tid, - "ts": (begin - base) / cycles_per_us, + "ts": offset_us + (begin - base) / cycles_per_us, "dur": (end - begin) / cycles_per_us, "args": {"beginCycle": begin, "endCycle": end, **args}, } @@ -139,12 +140,12 @@ def _span_event(name, category, pid, tid, begin, end, base, cycles_per_us, args) def build_chrome_trace(rank_traces): events = [] sources = [] - for rank_trace in sorted(rank_traces, key=lambda item: item["header"]["rank"]): + bases = {} + iteration_durations = {} + for rank_trace in rank_traces: header = rank_trace["header"] data = rank_trace["data"] rank = header["rank"] - sources.append(rank_trace["path"]) - bases = [] for iteration in range(header["iteration_count"]): spans = [] for core in range(header["core_count"]): @@ -155,14 +156,29 @@ def build_chrome_trace(rank_traces): spans.append(span) if not spans: raise ValueError(f"rank {rank} iteration {iteration} contains no kernel spans") - bases.append(min(begin for begin, _ in spans)) + base = min(begin for begin, _ in spans) + bases[(rank, iteration)] = base + duration = (max(end for _, end in spans) - base) / header["cycles_per_us"] + iteration_durations[iteration] = max(iteration_durations.get(iteration, 0.0), duration) + iteration_offsets = {} + cursor = 0.0 + for iteration in sorted(iteration_durations): + iteration_offsets[iteration] = cursor + cursor += iteration_durations[iteration] + ITERATION_GAP_US + + for rank_trace in sorted(rank_traces, key=lambda item: item["header"]["rank"]): + header = rank_trace["header"] + data = rank_trace["data"] + rank = header["rank"] + sources.append(rank_trace["path"]) events.append(_metadata("process_name", rank, 0, f"rank {rank}")) for core in range(header["core_count"]): events.append(_metadata("thread_name", rank, core, f"core{core} {_core_role(core)}")) for iteration in range(header["iteration_count"]): - base = bases[iteration] + base = bases[(rank, iteration)] + offset_us = iteration_offsets[iteration] for core in range(header["core_count"]): role = _core_role(core) for region, name in ((0, "kernel"), (1, "work")): @@ -173,7 +189,7 @@ def build_chrome_trace(rank_traces): events.append(_span_event( name, "kernel", rank, core, span[0], span[1], base, header["cycles_per_us"], - {"iteration": iteration, "role": role})) + {"iteration": iteration, "role": role}, offset_us)) for pass_index in range(header["pass_count"]): for peer in range(header["rank_size"]): for phase in range(header["phase_count"]): @@ -190,12 +206,12 @@ def build_chrome_trace(rank_traces): PHASE_NAMES[phase], role, rank, core, span[0], span[1], base, header["cycles_per_us"], {"iteration": iteration, "pass": pass_index, - "peer": peer, "phase": phase, "role": role})) + "peer": peer, "phase": phase, "role": role}, offset_us)) return { "traceEvents": events, "otherData": { "sources": sources, - "clock": "GetSystemCycle normalized independently per rank and iteration", + "clock": "GetSystemCycle normalized per rank and iteration, iterations laid out sequentially", }, } diff --git a/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py index 8b072b08..500ab5c4 100644 --- a/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py @@ -75,6 +75,24 @@ def test_normalizes_each_rank_independently(self): data_put = [event for event in trace["traceEvents"] if event.get("name") == "data-put"] self.assertEqual([event["ts"] for event in data_put], [0.1, 0.1]) + def test_places_iterations_sequentially(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "rank3.bin" + self.make_trace(path, iteration_count=2) + data = bytearray(path.read_bytes()) + struct.pack_into(" Date: Fri, 17 Jul 2026 20:08:48 +0800 Subject: [PATCH 050/163] docs(udma): plan fullmesh 16-0 trace experiment --- .../2026-07-17-udma-fullmesh-16-0-trace.md | 71 +++++++++++++++++++ ...6-07-17-udma-fullmesh-16-0-trace-design.md | 31 ++++++++ 2 files changed, 102 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-17-udma-fullmesh-16-0-trace.md create mode 100644 docs/superpowers/specs/2026-07-17-udma-fullmesh-16-0-trace-design.md diff --git a/docs/superpowers/plans/2026-07-17-udma-fullmesh-16-0-trace.md b/docs/superpowers/plans/2026-07-17-udma-fullmesh-16-0-trace.md new file mode 100644 index 00000000..30424785 --- /dev/null +++ b/docs/superpowers/plans/2026-07-17-udma-fullmesh-16-0-trace.md @@ -0,0 +1,71 @@ +# Full-Mesh 16:0 Trace Experiment Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Build and deploy a 16:0 full-mesh remote-send variant for a physical 2x8 trace run. + +**Architecture:** Reuse the current two-worker protocol and change only the shard split constants. Core17 retains synchronization and completion duties with a zero-byte segment. + +**Tech Stack:** C++14, Ascend C/Bisheng, Git bundle, Python trace tooling. + +## Global Constraints + +- Preserve max/min weighted QP selection. +- Preserve ready and ACK routing. +- Preserve the 8 MiB GM trace layout and converter. +- Commit before bundle deployment. +- Use `TILEXR_IPC_PID_MODE=pid` for the physical run. + +--- + +### Task 1: Change And Deploy The Split + +**Files:** +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` + +**Interfaces:** +- Consumes: `BigDataRemoteSendSegmentRange` and the existing segment-done protocol. +- Produces: core16 `[0, 16)` and core17 `[16, 16)` payload ranges. + +- [ ] **Step 1: Write the failing source assertions** + +Require exact constants `PRIMARY_SHARD_END = 16U` and +`SECONDARY_AGGREGATOR = 15U`. + +- [ ] **Step 2: Verify RED** + +Run the layout source test and confirm it fails while the kernel still uses +`12U` and `11U`. + +- [ ] **Step 3: Implement the 16:0 constants** + +Change only the two constants; retain the existing `segmentBytes > 0U` guard +and unconditional segment-done publication. + +- [ ] **Step 4: Commit, bundle, and deploy** + +Commit the test and implementation, create a complete verified bundle, upload +it to both hosts, and build the demo plus unit tests. + +### Task 2: Run And Export The Trace + +**Files:** +- Create: `tmp/test0701_fullmesh_trace_16_0_2x8/` artifacts. + +**Interfaces:** +- Consumes: the committed 16:0 demo. +- Produces: sixteen raw traces, a complete Chrome trace, and an iteration49 trace. + +- [ ] **Step 1: Run physical 2x8** + +Use 128 MiB/rank, profile stage8, repeat50, full-mesh trace enabled, and PID IPC mode. + +- [ ] **Step 2: Validate the split** + +For iteration49, require eight core16 `data-put` events and zero core17 +`data-put` events per rank. + +- [ ] **Step 3: Export JSON** + +Download all raw traces, convert the full trace, and extract iteration49. diff --git a/docs/superpowers/specs/2026-07-17-udma-fullmesh-16-0-trace-design.md b/docs/superpowers/specs/2026-07-17-udma-fullmesh-16-0-trace-design.md new file mode 100644 index 00000000..7982c696 --- /dev/null +++ b/docs/superpowers/specs/2026-07-17-udma-fullmesh-16-0-trace-design.md @@ -0,0 +1,31 @@ +# Full-Mesh 16:0 Trace Experiment Design + +## Goal + +Run the existing physical 2x8 full-mesh trace with core16 sending all remote +payload and core17 sending no payload while still publishing its completion. + +## Scope + +- Change the remote-send split from 12:4 to 16:0. +- Keep core16 on the max-weight QP and core17 on the min-weight QP. +- Keep core17 active so it publishes `segmentDone` for every remote peer. +- Keep ready publication on core16 after both segment completion tokens arrive. +- Keep direct registered-memory ACK and the existing trace format unchanged. +- Make the change on a separate experimental branch. + +## Implementation + +Set `TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END` to `16U` and +`TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR` to `15U`. +The existing segment-range helper then gives core16 all 16 copy shards and +core17 an empty `[16, 16)` range. Core17 skips `UDMAPutNbiOnQp` because +`segmentBytes == 0`, but still writes its local segment-done token. + +## Validation + +Build and run source/layout tests on both hosts, then run physical 2x8 with +128 MiB input/output per rank, one pass, repeat50, and full-mesh GM tracing. +Require 16 successful ranks, sixteen 8 MiB trace binaries, no core17 +`data-put` events, and eight 8 MiB core16 `data-put` events per rank in +iteration 49. From a1bd4081b90b77c4146ab9285fbffa8a929810b0 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 20:09:15 +0800 Subject: [PATCH 051/163] test(udma): require fullmesh 16-0 split --- tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index b7b91ef0..9d83b03b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -678,7 +678,8 @@ void TestAllToAllBigDataSource() kernel, "BigDataRemoteSendSegmentWorker", "BigDataRemotePutOnlySendWorker"); const std::string fullMeshRecv = SliceBetween( kernel, "BigDataRecvPeerWorker", "BigDataWaitCopyDoneRange"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 16U"); + CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR = 15U"); CHECK_CONTAINS(fullMeshReady, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(fullMeshReady, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); CHECK_CONTAINS(fullMeshSend, "BigDataSelectWeightedQp("); From 86a8f8687b806802f90c460d39aa9fa8f2dc0f12 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Fri, 17 Jul 2026 20:11:59 +0800 Subject: [PATCH 052/163] perf(udma): route fullmesh payload through core16 --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 57ea0f2b..229f9ac6 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -78,8 +78,8 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT = 1U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY = 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY = 3U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY = 4U; -constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 12U; -constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR = 11U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 16U; +constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR = 15U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_CORE = 20U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_SHARD = TILEXR_UDMA_DEMO_BIGDATA_RECV_READY_WAIT_CORE - From 0e8d62b0f7b3e387ccfa94d86183f5e26fb63f39 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 09:18:52 +0800 Subject: [PATCH 053/163] fix(udma): handle terminal fullmesh shard boundary --- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 5 ++++- tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 229f9ac6..bc3c0a55 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -367,9 +367,12 @@ __aicore__ inline bool BigDataCopyShardRange( __aicore__ inline uint32_t BigDataCopyShardStartBytes( uint32_t shard, uint32_t shardCount, uint32_t totalElements) { - if (shardCount == 0U || shard >= shardCount) { + if (shardCount == 0U || shard > shardCount) { return 0U; } + if (shard == shardCount) { + return totalElements * static_cast(sizeof(int32_t)); + } constexpr uint32_t alignElements = 32U / sizeof(int32_t); const uint64_t begin = static_cast(totalElements) * static_cast(shard) / diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 9d83b03b..c3fbe31a 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -680,6 +680,8 @@ void TestAllToAllBigDataSource() kernel, "BigDataRecvPeerWorker", "BigDataWaitCopyDoneRange"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SHARD_END = 16U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR = 15U"); + CHECK_CONTAINS(kernel, "if (shard == shardCount)"); + CHECK_CONTAINS(kernel, "return totalElements * static_cast(sizeof(int32_t))"); CHECK_CONTAINS(fullMeshReady, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(fullMeshReady, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); CHECK_CONTAINS(fullMeshSend, "BigDataSelectWeightedQp("); From f135bcf8373225f0db0c89b84e192b99a521d92c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 11:23:47 +0800 Subject: [PATCH 054/163] docs: design HCCL two-host cluster validation --- ...-20-hccl-srv9-srv14-cluster-info-design.md | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-20-hccl-srv9-srv14-cluster-info-design.md diff --git a/docs/superpowers/specs/2026-07-20-hccl-srv9-srv14-cluster-info-design.md b/docs/superpowers/specs/2026-07-20-hccl-srv9-srv14-cluster-info-design.md new file mode 100644 index 00000000..4d43e4e2 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-hccl-srv9-srv14-cluster-info-design.md @@ -0,0 +1,64 @@ +# HCCL srv_9 + srv_14 Cluster Info Design + +## Goal + +Validate HCCL AllToAll between `srv_9` (`141.61.49.223`) and `srv_14` +(`141.61.49.192`) using configuration generated from the machines' actual +HCCL root information. Validate two ranks first, then the full physical 2x8 +configuration. + +## Constraints + +- Remote writes are limited to `/home/h30059441` and `/tmp`. +- `/etc/hccl_rootinfo.json` and the installed CANN files are read-only inputs. +- The test copy must use the same CANN installation as the TileXR UDMA test. +- Do not invent physical topology edges that cannot be derived from an + authoritative source. + +## Configuration + +Generate cluster-info files by combining rank entries from each host's +`/etc/hccl_rootinfo.json`: + +- `ranktable_2x1.json`: device 0 from each server, assigned ranks 0 and 1. +- `ranktable_2x8.json`: all eight devices from `srv_9` as ranks 0 through 7, + followed by all eight devices from `srv_14` as ranks 8 through 15. + +Preserve each device's actual level list, EIDs, ports, and network instance +identifiers. The local hardware topology remains the vendor-provided +`/usr/local/Ascend/driver/topo/950/atlas_950_1.json`; no synthetic inter-server +edge graph is generated. + +## Test Integration + +Modify only the HCCL test copy under `/home/h30059441`. Add an explicit +cluster-info initialization path so the test consumes the generated ranktable +instead of relying on `HcclGetRootInfo` and MPI root-info exchange. Keep the +original path available for comparison. + +Resolve and record the CANN environment used by the deployed TileXR UDMA test. +Build and run the HCCL test against that same installation. Use +`data0.3001` as `HCCL_SOCKET_IFNAME` on both hosts. + +## Validation Sequence + +1. Validate generated JSON structure and rank uniqueness without running NPU + work. +2. Build the modified test copy on both hosts. +3. Run 2x1 on device 0 of each host with a small payload to validate channel + creation and correctness. +4. Run 2x1 with 128 MiB input/output per rank. +5. Only if both 2x1 runs pass, run physical 2x8 with 128 MiB input/output per + rank, five warmups, and fifty measured iterations. + +Every run has an external timeout. A failure stops progression and preserves +the command, logs, generated configuration, and exit status under +`/home/h30059441` for diagnosis. + +## Success Criteria + +- The test log proves that cluster-info initialization consumed the intended + generated file. +- All participating ranks initialize and complete without channel timeout. +- Data verification passes. +- The 2x8 run reports timing only after the 2x1 connectivity checks pass. From cd57ec1a0703a5af7ab17c4bc5478149295b0743 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 11:41:34 +0800 Subject: [PATCH 055/163] fix(udma): stream fullmesh trace JSON output --- .../demo/tilexr_udma_fullmesh_trace_to_chrome.py | 3 ++- .../test_tilexr_udma_fullmesh_trace_to_chrome.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py index a9654654..36eb1104 100644 --- a/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_fullmesh_trace_to_chrome.py @@ -223,7 +223,8 @@ def main(): args = parser.parse_args() trace = build_chrome_trace([read_rank_trace(path) for path in args.inputs]) args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(trace, separators=(",", ":")), encoding="utf-8") + with args.output.open("w", encoding="utf-8") as stream: + json.dump(trace, stream, separators=(",", ":")) print(f"wrote {args.output} events={len(trace['traceEvents'])}") diff --git a/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py index 500ab5c4..f5416f32 100644 --- a/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_fullmesh_trace_to_chrome.py @@ -5,6 +5,7 @@ import tempfile import unittest from pathlib import Path +from unittest import mock MODULE_PATH = Path(__file__).resolve().parents[1] / "demo" / "tilexr_udma_fullmesh_trace_to_chrome.py" @@ -119,6 +120,19 @@ def test_rejects_half_written_span(self): with self.assertRaisesRegex(ValueError, "incomplete"): MODULE.build_chrome_trace([MODULE.read_rank_trace(path)]) + def test_main_streams_json_output(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "rank3.bin" + output = Path(directory) / "trace.json" + self.make_trace(source) + + with mock.patch.object(MODULE.json, "dumps", side_effect=AssertionError("no bulk JSON")): + with mock.patch("sys.argv", ["converter", str(source), "--output", str(output)]): + MODULE.main() + + parsed = json.loads(output.read_text(encoding="utf-8")) + self.assertTrue(parsed["traceEvents"]) + if __name__ == "__main__": unittest.main() From e7ca40b8714c8278c5d4640cffab8c878319d288 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 13:49:53 +0800 Subject: [PATCH 056/163] docs(udma): design isolated fullmesh task profiling --- ...-20-udma-fullmesh-isolated-task-profile.md | 66 +++++++++++++++++++ ...a-fullmesh-isolated-task-profile-design.md | 44 +++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-20-udma-fullmesh-isolated-task-profile.md create mode 100644 docs/superpowers/specs/2026-07-20-udma-fullmesh-isolated-task-profile-design.md diff --git a/docs/superpowers/plans/2026-07-20-udma-fullmesh-isolated-task-profile.md b/docs/superpowers/plans/2026-07-20-udma-fullmesh-isolated-task-profile.md new file mode 100644 index 00000000..330e1ef9 --- /dev/null +++ b/docs/superpowers/plans/2026-07-20-udma-fullmesh-isolated-task-profile.md @@ -0,0 +1,66 @@ +# Full-Mesh Isolated Task Profiling Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add and run wait-free per-task microbenchmarks for the physical full-mesh kernel. + +**Architecture:** Host selects one isolated task and passes it as a dedicated kernel argument. The kernel dispatches to wait-free task helpers before the normal pipeline loop and records existing trace phases. + +**Tech Stack:** C++14, Ascend C/Bisheng, Python trace conversion, physical Ascend950 2x8. + +## Global Constraints + +- Preserve normal mode when task is zero. +- Use `/home/pkg/b101/cann` for build and runtime. +- Keep the 16:0 split. +- Never execute a wait primitive in isolated mode. +- Commit and deploy through a verified Git bundle. + +--- + +### Task 1: Host And Kernel Dispatch + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` + +**Interfaces:** +- Consumes: `TILEXR_DEMO_BIGDATA_ISOLATED_TASK` in range 0..11. +- Produces: kernel argument `uint32_t isolatedTask`. + +- [ ] Add failing source checks for the environment, argument, range guard, + validation skip, and isolated dispatch. +- [ ] Verify the layout test fails on both hosts. +- [ ] Implement Host parsing and kernel argument plumbing. +- [ ] Verify source tests pass. + +### Task 2: Wait-Free Task Helpers + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo_kernel.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp` + +**Interfaces:** +- Consumes: existing layout, copy, UDMA, trace, and control-slot helpers. +- Produces: `BigDataRunIsolatedTask` for task IDs 1..11. + +- [ ] Add failing checks for all task IDs and absence of wait calls in the + isolated helper slice. +- [ ] Implement each task with existing trace phase names. +- [ ] Build the Host and Bisheng kernel with b101 and run all unit tests. +- [ ] Commit, create a complete bundle, and deploy to both hosts. + +### Task 3: Physical Task Matrix + +**Files:** +- Create: `tmp/fullmesh_isolated_b101_2x8/` artifacts. + +**Interfaces:** +- Consumes: task IDs 1..11 and full-mesh raw tracing. +- Produces: per-task repeat50 timing and iteration49 summaries. + +- [ ] Run each task on physical 2x8 with 128 MiB/rank and repeat50. +- [ ] Verify 16 successful ranks and sixteen 8 MiB traces for every task. +- [ ] Verify no wait phase appears in any isolated trace. +- [ ] Summarize per-operation iteration49 mean, minimum, and maximum. diff --git a/docs/superpowers/specs/2026-07-20-udma-fullmesh-isolated-task-profile-design.md b/docs/superpowers/specs/2026-07-20-udma-fullmesh-isolated-task-profile-design.md new file mode 100644 index 00000000..554cd304 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-udma-fullmesh-isolated-task-profile-design.md @@ -0,0 +1,44 @@ +# Full-Mesh Isolated Task Profiling Design + +## Goal + +Measure the execution cost of each full-mesh task without synchronization +waits so the 2 ms full-pipeline latency can be decomposed into independent +copy, UDMA, control, and receive costs. + +## Interface + +`TILEXR_DEMO_BIGDATA_ISOLATED_TASK` selects one diagnostic task. Zero keeps +the existing full pipeline unchanged. + +| Value | Task | Active cores | +|---:|---|---| +| 1 | self-copy | 0..15 | +| 2 | peer-copy and copyDone store | 0..15 | +| 3 | copy-ready stores without copyDone waits | 0 and 15 | +| 4 | remote payload put and quiet | 16 | +| 5 | primary segmentDone store | 16 | +| 6 | ready signal and quiet | 16 | +| 7 | secondary zero-payload segmentDone store | 17 | +| 8 | local payload put-signal and quiet | 18 | +| 9 | output-copy | 19..34 | +| 10 | recvDone store | 19..34 | +| 11 | direct registered-memory ACK store | 34 | + +## Execution Rules + +- Physical 2x8 and the existing 16:0 split remain unchanged. +- No isolated task calls `BigDataWaitTokenMte` or `BigDataWaitCopyReady`. +- UDMA tasks retain separate `data-put` and `quiet` spans. +- Task 6 retains the signal API's required quiet. +- Task 11 writes ACK directly to the remote registered control slot. +- Every launch still uses the 35-core exit barrier so inactive cores exit + coherently. +- Isolated mode always skips output and CQ-result validation. +- Full-mesh raw tracing remains the source of task timing. + +## Validation + +Build with `/home/pkg/b101/cann`, run layout and converter tests, then run all +eleven task values on physical 2x8 with repeat50. Require all ranks to finish, +all traces to contain no wait phases, and summarize iteration49 per task. From e89ebc55a1f6ec4e0c89db79d1951a59db87f778 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 13:51:02 +0800 Subject: [PATCH 057/163] test(udma): require wait-free isolated tasks --- .../unit/test_tilexr_udma_alltoall_layout.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index c3fbe31a..97785c13 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -663,6 +663,9 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(demo, "skip result validation for bigdata remote-put-only profile"); CHECK_CONTAINS(demo, "TILEXR_UDMA_FULLMESH_TRACE"); CHECK_CONTAINS(demo, "TILEXR_UDMA_FULLMESH_TRACE_DIR"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_BIGDATA_ISOLATED_TASK"); + CHECK_CONTAINS(demo, "bigDataIsolatedTask < 0 || bigDataIsolatedTask > 11"); + CHECK_CONTAINS(demo, "bigDataProfilePartial = bigDataProfilePartial || bigDataIsolatedTask != 0"); CHECK_CONTAINS(demo, "FullmeshTraceLayoutFits("); CHECK_CONTAINS(demo, "aclrtMalloc(&fullmeshTraceDevice"); CHECK_CONTAINS(demo, "TileXR::Demo::kFullmeshTraceBytes, \"fullmesh trace\""); @@ -697,6 +700,8 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "if (trace == nullptr"); CHECK_CONTAINS(kernel, "AscendC::GetSystemCycle()"); CHECK_CONTAINS(kernel, "auto fullmeshTrace = (!remotePutOnly && use35Core &&"); + CHECK_CONTAINS(kernel, "uint32_t isolatedTask"); + CHECK_CONTAINS(kernel, "BigDataRunIsolatedTask("); CHECK_CONTAINS(kernel, "kFullmeshTracePhaseSelfCopy"); CHECK_CONTAINS(kernel, "kFullmeshTracePhasePeerCopy"); CHECK_CONTAINS(kernel, "kFullmeshTracePhasePublishCopyReady"); @@ -714,6 +719,18 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhaseWaitRecvDone"); CHECK_CONTAINS(fullMeshRecv, "kFullmeshTracePhaseAck"); + const std::string isolatedTasks = SliceBetween( + kernel, "BigDataRunIsolatedTask", "BigDataRunRoleForPeer"); + for (uint32_t task = 1U; task <= 11U; ++task) { + CHECK_CONTAINS(isolatedTasks, "TILEXR_BIGDATA_ISOLATED_TASK_" + std::to_string(task)); + } + CHECK_NOT_CONTAINS(isolatedTasks, "BigDataWaitTokenMte("); + CHECK_NOT_CONTAINS(isolatedTasks, "BigDataWaitCopyReady("); + CHECK_CONTAINS(isolatedTasks, "UDMAPutNbiOnQp"); + CHECK_CONTAINS(isolatedTasks, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); + CHECK_CONTAINS(isolatedTasks, "UDMAPutSignalNbi"); + CHECK_CONTAINS(isolatedTasks, "BigDataRemoteRegisteredControlSlot("); + const std::string remotePutOnlySend = SliceBetween( kernel, "BigDataRemotePutOnlySendWorker", "BigDataRemotePutOnlyCheckIndex"); const std::string remotePutOnlyCheck = SliceBetween( From 65a8bf2b17c59e6269640b35e81400f8fa66a403 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 13:57:59 +0800 Subject: [PATCH 058/163] feat(udma): add wait-free isolated task profiling --- tests/udma/demo/tilexr_udma_demo.cpp | 12 + tests/udma/demo/tilexr_udma_demo_kernel.cpp | 271 ++++++++++++++++++++ 2 files changed, 283 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index dffb3727..e93b42c5 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -51,6 +51,7 @@ extern void launch_tilexr_udma_all_to_all_fused( extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, + uint32_t isolatedTask, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, @@ -700,7 +701,16 @@ int main(int argc, char** argv) testType == 7 && GetEnvInt("TILEXR_DEMO_BIGDATA_FORCE_35CORE", 0) != 0; const bool bigDataRemotePutOnly = testType == 7 && GetEnvInt("TILEXR_DEMO_BIGDATA_REMOTE_PUT_ONLY", 0) != 0; + const int bigDataIsolatedTask = + testType == 7 ? GetEnvInt("TILEXR_DEMO_BIGDATA_ISOLATED_TASK", 0) : 0; + if (bigDataIsolatedTask < 0 || bigDataIsolatedTask > 11) { + std::cerr << "[rank " << rank << "] ERROR: TILEXR_DEMO_BIGDATA_ISOLATED_TASK must be 0..11" + << ", got " << bigDataIsolatedTask << std::endl; + Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); + return 1; + } bool bigDataProfilePartial = testType == 7 && bigDataProfileStage < kBigDataProfileStageFull; + bigDataProfilePartial = bigDataProfilePartial || bigDataIsolatedTask != 0; bool syncAllToAllAtEnd = isAllToAll && GetEnvInt("TILEXR_DEMO_ALLTOALL_SYNC_AT_END", 0) != 0; bool useAllToAllPlainIpc = @@ -777,6 +787,7 @@ int main(int argc, char** argv) rankSize, bigDataRanksPerNode) ? "true" : "false") + " force35Core=" + std::string(forceBigData35Core ? "true" : "false") + " remotePutOnly=" + std::string(bigDataRemotePutOnly ? "true" : "false") + + " isolatedTask=" + std::to_string(bigDataIsolatedTask) + " ranksPerNode=" + std::to_string(bigDataRanksPerNode) + " blockDim=" + std::to_string(TileXR::Demo::AllToAllBigDataBlockDim( rankSize, forceBigData35Core, bigDataRemotePutOnly, bigDataRanksPerNode)) + @@ -1075,6 +1086,7 @@ int main(int argc, char** argv) reinterpret_cast(bigInput), reinterpret_cast(bigOutput), reinterpret_cast(registeredMemory), reinterpret_cast(debug), reinterpret_cast(fullmeshTraceDevice), fullmeshTraceIteration, + static_cast(bigDataIsolatedTask), elementsPerRank, 0, bigDataPlan.copyDoneOffset, bigDataPlan.recvCopyDoneOffset, bigDataPlan.remoteSendDoneOffset, diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index bc3c0a55..9db21051 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -95,6 +95,17 @@ constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_RELAY_COPY = 5; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_ACK_PUT = 6; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_WAIT_ACK = 7; constexpr uint32_t TILEXR_BIGDATA_PROFILE_STAGE_FULL = 8; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_1 = 1U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_2 = 2U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_3 = 3U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_4 = 4U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_5 = 5U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_6 = 6U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_7 = 7U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_8 = 8U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_9 = 9U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_10 = 10U; +constexpr uint32_t TILEXR_BIGDATA_ISOLATED_TASK_11 = 11U; constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_FRAMEWORK = 0; constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_LOOP = 1; constexpr uint32_t TILEXR_BIGDATA_REMOTE_PUT_STAGE_PEER = 2; @@ -1728,6 +1739,248 @@ __aicore__ inline void BigDataRunSelfCopyShard( BigDataCopyRangePingPong(dst + shardOffset, src + shardOffset, shardBytes, relayLocal); } +__aicore__ inline void BigDataRunIsolatedTask( + uint32_t isolatedTask, int32_t blockIdx, int32_t rank, int32_t rankSize, + __gm__ TileXR::CommArgs* args, __gm__ int32_t* input, __gm__ int32_t* output, + __gm__ uint8_t* udmaMem, __gm__ int32_t* debug, + int32_t elementsPerPeer, int32_t effectiveChunkElements, uint32_t passCount, + uint32_t loop, uint32_t pass, uint64_t kernelLoopBase, bool force35Core, + uint32_t shardCount, uint64_t sendDataOffset, uint64_t recvDataOffset, + uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, + uint64_t readySignalOffset, uint64_t ackSignalOffset, uint64_t chunkBytesPerPeer, + int32_t ranksPerNode, __gm__ uint8_t* fullmeshTrace, + uint32_t traceIteration, uint32_t traceCore, + AscendC::LocalTensor relayLocal) +{ + const bool isCopyCore = + blockIdx < static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES); + const bool isRecvCore = + blockIdx >= static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE) && + blockIdx < static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM); + const uint32_t copyShard = static_cast(blockIdx); + const uint32_t recvShard = static_cast(blockIdx) - + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORE_BASE; + const uint64_t globalPass = BigDataGlobalPassIndex(kernelLoopBase, passCount, loop, pass); + const uint64_t token = BigDataPassToken(globalPass); + const uint32_t slot = BigDataDataSlot(globalPass, pass, true); + const int32_t networkPeerCount = rankSize > 1 ? rankSize - 1 : 1; + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_1) { + if (isCopyCore) { + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataRunSelfCopyShard(rank, rankSize, input, output, elementsPerPeer, + effectiveChunkElements, pass, copyShard, shardCount, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, rank, TileXR::Demo::kFullmeshTracePhaseSelfCopy, + passCount, rankSize, begin, end); + } + return; + } + + const int32_t taskCount = BigDataTaskCount(rankSize, force35Core, ranksPerNode); + for (int32_t taskIndex = 0; taskIndex < taskCount; ++taskIndex) { + int32_t peer = -1; + bool isLocalPeer = false; + if (!BigDataMergedPeerTaskAt( + rank, rankSize, taskIndex, force35Core, ranksPerNode, peer, isLocalPeer)) { + continue; + } + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_2) { + if (isCopyCore) { + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataCopyPeerWorker(peer, rank, rankSize, args, input, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, loop, pass, kernelLoopBase, + TILEXR_BIGDATA_PROFILE_STAGE_SEND_COPY, true, copyShard, shardCount, + sendDataOffset, copyDoneOffset, ackSignalOffset, chunkBytesPerPeer, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhasePeerCopy, + passCount, rankSize, begin, end); + } + continue; + } + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_3) { + uint32_t readyShard = 0U; + bool publish = false; + if (blockIdx == 0) { + readyShard = isLocalPeer ? TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_READY : + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_PRIMARY; + publish = true; + } else if (!isLocalPeer && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_AGGREGATOR)) { + readyShard = TILEXR_UDMA_DEMO_BIGDATA_REMOTE_COPY_READY_SECONDARY; + publish = true; + } + if (publish) { + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataStoreTokenMte(BigDataControlSlot(udmaMem, remoteSendDoneOffset, + slot, rankSize, shardCount, peer, readyShard), token, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhasePublishCopyReady, + passCount, rankSize, begin, end); + } + continue; + } + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_4 && !isLocalPeer && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE)) { + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + uint32_t copyShardBegin = 0U; + uint32_t copyShardEnd = 0U; + uint32_t segmentOffsetBytes = 0U; + uint32_t segmentBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes) || + !BigDataRemoteSendSegmentRange(TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT, + TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_COPY_CORES, chunkBytes / sizeof(int32_t), + copyShardBegin, copyShardEnd, segmentOffsetBytes, segmentBytes)) { + continue; + } + auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, + BigDataNetworkPeerIndex(peer, rank), chunkBytesPerPeer); + auto localSrc = reinterpret_cast<__gm__ int32_t*>(sendSlot + segmentOffsetBytes); + const uint64_t remoteDataOffset = recvDataOffset + + (static_cast(slot) * static_cast(networkPeerCount) + + static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer + + segmentOffsetBytes; + const uint32_t qpIdx = BigDataSelectWeightedQp(args, peer, true); + const uint64_t putBegin = BigDataFullmeshTraceCycle(fullmeshTrace); + TileXR::UDMAPutNbiOnQp(args, peer, qpIdx, localSrc, remoteDataOffset, segmentBytes); + const uint64_t putEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseDataPut, + passCount, rankSize, putBegin, putEnd); + const uint64_t quietBegin = BigDataFullmeshTraceCycle(fullmeshTrace); + (void)TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + const uint64_t quietEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseQuiet, + passCount, rankSize, quietBegin, quietEnd); + continue; + } + + if ((isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_5 && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE)) || + (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_7 && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_SECONDARY_CORE))) { + if (!isLocalPeer) { + const uint32_t segmentId = isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_5 ? + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_PRIMARY_SEGMENT : + TILEXR_UDMA_DEMO_BIGDATA_REMOTE_SEND_SECONDARY_SEGMENT; + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataStoreTokenMte(BigDataControlSlot(udmaMem, remoteSendDoneOffset, + slot, rankSize, shardCount, peer, segmentId), token, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseSegmentDone, + passCount, rankSize, begin, end); + } + continue; + } + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_6 && !isLocalPeer && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_REMOTE_SEND_PRIMARY_CORE)) { + const uint32_t qpIdx = BigDataSelectWeightedQp(args, peer, true); + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + auto localReady = BigDataControlSlot( + udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, 0U); + BigDataStoreTokenMte(localReady, token, relayLocal); + BigDataPublishReadySignal( + args, peer, udmaMem, readySignalOffset, slot, rankSize, shardCount, rank, token, qpIdx); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhasePublishReady, + passCount, rankSize, begin, end); + continue; + } + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_8 && isLocalPeer && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_LOCAL_SEND_CORE)) { + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes)) { + continue; + } + auto sendSlot = BigDataSlot(udmaMem, sendDataOffset, slot, networkPeerCount, + BigDataNetworkPeerIndex(peer, rank), chunkBytesPerPeer); + auto localSrc = reinterpret_cast<__gm__ int32_t*>(sendSlot); + const uint64_t remoteDataOffset = recvDataOffset + + (static_cast(slot) * static_cast(networkPeerCount) + + static_cast(BigDataNetworkPeerIndex(rank, peer))) * chunkBytesPerPeer; + const uint64_t remoteReadyOffset = readySignalOffset + + ((static_cast(slot) * static_cast(rankSize) + + static_cast(rank)) * static_cast(shardCount)) * + TILEXR_UDMA_DEMO_BIGDATA_CONTROL_SLOT_BYTES; + const uint64_t putBegin = BigDataFullmeshTraceCycle(fullmeshTrace); + TileXR::UDMAPutSignalNbi(args, peer, localSrc, + remoteDataOffset, chunkBytes, remoteReadyOffset, token); + const uint64_t putEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseDataPut, + passCount, rankSize, putBegin, putEnd); + const uint64_t quietBegin = BigDataFullmeshTraceCycle(fullmeshTrace); + (void)TileXR::UDMAQuietStatus(args, peer); + const uint64_t quietEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseQuiet, + passCount, rankSize, quietBegin, quietEnd); + continue; + } + + if ((isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_9 || + isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_10) && isRecvCore) { + int32_t chunkOffset = 0; + uint32_t chunkBytes = 0U; + uint32_t shardOffset = 0U; + uint32_t shardBytes = 0U; + if (!BigDataPassChunk(pass, elementsPerPeer, effectiveChunkElements, chunkOffset, chunkBytes) || + !BigDataCopyShardRange(recvShard, TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_RECV_CORES, + chunkBytes / sizeof(int32_t), shardOffset, shardBytes)) { + continue; + } + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_9) { + auto recvSlot = BigDataSlot(udmaMem, recvDataOffset, slot, networkPeerCount, + BigDataNetworkPeerIndex(peer, rank), chunkBytesPerPeer); + auto dst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(peer) * elementsPerPeer + chunkOffset); + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataCopyRangePingPong(dst + shardOffset, recvSlot + shardOffset, shardBytes, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseOutputCopy, + passCount, rankSize, begin, end); + } else { + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataStoreTokenMte(BigDataControlSlot(udmaMem, recvCopyDoneOffset, + slot, rankSize, shardCount, peer, recvShard), token, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhasePublishRecvDone, + passCount, rankSize, begin, end); + } + continue; + } + + if (isolatedTask == TILEXR_BIGDATA_ISOLATED_TASK_11 && + blockIdx == static_cast(TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM - 1U)) { + auto remoteAck = BigDataRemoteRegisteredControlSlot( + args, peer, ackSignalOffset, slot, rankSize, shardCount, rank, 0U); + if (remoteAck != nullptr) { + const uint64_t begin = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataStoreTokenMte(remoteAck, token, relayLocal); + const uint64_t end = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan(fullmeshTrace, traceIteration, traceCore, + pass, peer, TileXR::Demo::kFullmeshTracePhaseAck, + passCount, rankSize, begin, end); + } + } + } +} + __aicore__ inline void BigDataRunRoleForPeer( int32_t peer, int32_t role, int32_t rank, int32_t rankSize, __gm__ TileXR::CommArgs* args, __gm__ int32_t* input, __gm__ int32_t* output, __gm__ uint8_t* udmaMem, @@ -2846,6 +3099,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_fused_kernel( extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR udmaMemGM, GM_ADDR debugGM, GM_ADDR fullmeshTraceGM, uint32_t fullmeshTraceIteration, + uint32_t isolatedTask, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, uint64_t ackSignalOffset, int32_t chunkElements, uint32_t passCount, uint32_t loopCount, @@ -2965,6 +3219,21 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( for (uint32_t loop = 0; loop < loopCount; ++loop) { for (uint32_t pass = 0; pass < passCount; ++pass) { const uint64_t passStart = BigDataFullmeshTraceCycle(fullmeshTrace); + if (isolatedTask != 0U) { + BigDataRunIsolatedTask(isolatedTask, blockIdx, rank, rankSize, + args, input, output, udmaMem, debug, + elementsPerPeer, effectiveChunkElements, passCount, + loop, pass, kernelLoopBase, force35Core, shardCount, + sendDataOffset, recvDataOffset, copyDoneOffset, recvCopyDoneOffset, + remoteSendDoneOffset, readySignalOffset, ackSignalOffset, chunkBytesPerPeer, + ranksPerNode, fullmeshTrace, fullmeshTraceIteration, traceCore, relayLocal); + const uint64_t passEnd = BigDataFullmeshTraceCycle(fullmeshTrace); + BigDataFullmeshTraceRecordTaskSpan( + fullmeshTrace, fullmeshTraceIteration, traceCore, pass, rank, + TileXR::Demo::kFullmeshTracePhasePass, + passCount, rankSize, passStart, passEnd); + continue; + } if (remotePutOnly) { for (int32_t sendTask = blockIdx; sendTask < sendTaskCount; sendTask += static_cast(activeBlockDim)) { @@ -3141,6 +3410,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_bigdata_kernel( void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, + uint32_t isolatedTask, int32_t elementsPerPeer, uint64_t dataOffset, uint64_t copyDoneOffset, uint64_t recvCopyDoneOffset, uint64_t remoteSendDoneOffset, uint64_t readySignalOffset, @@ -3149,6 +3419,7 @@ void launch_tilexr_udma_all_to_all_bigdata( { tilexr_udma_all_to_all_bigdata_kernel<<>>( commArgs, input, output, udmaMem, debug, fullmeshTrace, fullmeshTraceIteration, + isolatedTask, elementsPerPeer, dataOffset, copyDoneOffset, recvCopyDoneOffset, remoteSendDoneOffset, readySignalOffset, ackSignalOffset, chunkElements, passCount, loopCount, kernelLoopBase, profileStage, force35Core); From 3f73a59cfa8fb9620f9d09d9df32fe2f74f6f71b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 17:30:11 +0800 Subject: [PATCH 059/163] docs(udma): design grouped fullmesh alltoall --- ...0-udma-grouped-fullmesh-alltoall-design.md | 335 ++++++++++++++++++ 1 file changed, 335 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md new file mode 100644 index 00000000..19e6d5d9 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md @@ -0,0 +1,335 @@ +# UDMA Grouped Fullmesh AllToAll Design + +## Status + +Approved for implementation planning on 2026-07-20. + +## Objective + +Add a new, standalone `alltoall_group` kernel that preserves global AllToAll +semantics while scheduling peers in symmetric groups of at most 16. The first +version prioritizes a small and auditable correctness model over registered +memory efficiency. + +The existing experimental 35-core big-data fullmesh kernel remains unchanged. +The new kernel must not inherit its local/remote role split, copy-ready +aggregation, ACK protocol, profiling stages, or experimental mode branches. + +## Scope + +The first version supports: + +- `rankSize >= 8` and `rankSize % 8 == 0`. +- A fixed peer-group width of 16: eight forward and eight backward circular + distances. +- Any payload whose complete double-buffered registered receive layout, + signals, and control data fit below the per-rank 1 GiB UDMA registration + limit. +- One complete AllToAll operation per kernel invocation. +- Host-side warmup and repeat launches on the same ordered stream. +- Chunked transfers within an invocation when one peer payload is divided into + multiple passes. + +The first version does not support: + +- Arbitrary rank sizes that are not multiples of eight. +- Reusing only 16 receive lanes. +- ACK or credit flow control. +- Concurrent grouped AllToAll invocations on multiple streams for the same + communicator and registered buffer. +- Dynamic group width. + +## Data Semantics + +Each rank owns source-major input and output arrays: + +```text +input[dstRank][elementsPerPeer] +output[srcRank][elementsPerPeer] +``` + +Rank `src` sends `input[dst]` to rank `dst`. Rank `dst` stores it in +`output[src]`. Self data is copied locally without UDMA. + +## Peer Group Schedule + +For rank `r`, group `g`, and local distance index `i` in `[0, 7]`: + +```text +distance = g * 8 + i + 1 +forwardPeer = (r + distance) % rankSize +backwardPeer = (r - distance + rankSize) % rankSize +``` + +Worker lanes map as follows: + +- Lane `0..7` selects the forward peer. +- Lane `8..15` selects the backward peer with `i = lane - 8`. +- A lane is inactive when `distance > rankSize / 2`. +- When `distance == rankSize / 2`, only the forward lane is active because the + forward and backward peers are identical. + +The group count is: + +```text +groupCount = ceil((rankSize - 1) / 16) +``` + +This schedule covers every non-self rank exactly once and is symmetric: if +rank A schedules rank B in group `g`, rank B schedules rank A in the same +group, although their local lane numbers can differ. + +Examples: + +- `rankSize=16`: one group with 15 active peers. +- `rankSize=24`: groups with 16 and 7 active peers. +- `rankSize=64`: groups with 16, 16, 16, and 15 active peers. + +## Core Assignment + +The kernel launches exactly 32 AIV blocks: + +```text +core 0..15 send workers +core 16..31 receive workers +``` + +Send lane `L` and receive lane `L` independently evaluate the peer schedule +for each group. They do not share progress state and do not require +`SyncAll()`. + +The 16 receive workers shard the self payload and copy it directly from input +to output. This local copy can overlap the first group's remote sends. + +Inactive lanes skip UDMA and signal waits but continue their normal loop and +return cleanly. + +## Registered Memory Layout + +Each rank allocates two complete receive planes: + +```text +payloadPlane[2][rankSize][elementsPerPeer] +signalPlane[2][rankSize][128 bytes] +debug/control area +``` + +The ping-pong plane for invocation `I` is: + +```text +slot = I & 1 +``` + +The payload slot is indexed by global source rank, not by worker lane: + +```text +payloadOffset(slot, sourceRank, chunkOffset) = + payloadBase[slot] + + sourceRank * bytesPerPeer + + chunkOffset + +signalOffset(slot, sourceRank) = + signalBase[slot] + sourceRank * 128 +``` + +For example, rank 37 sending to rank 5 writes rank 5's payload slot 37 and +signal slot 37. Rank 5 later copies payload slot 37 into output slice 37. +The sender's local worker lane has no effect on the remote address. + +Input and output remain ordinary GM allocations. Only the double-buffered +receive planes, signal planes, and control region are registered. + +The host rejects a plan unless: + +```text +2 * aligned(outputBytes) + +2 * aligned(rankSize * 128) + +aligned(debugControlBytes) <= 1 GiB +``` + +## Token Encoding + +Signal values are nonzero, monotonically ordered 64-bit tokens: + +```text +bits 63..32 invocationId + 1 +bit 31 ping-pong slot +bits 30..16 group +bits 15..0 pass + 1 +``` + +The host validates that group and pass dimensions fit their fields. All ranks +must use the same invocation ID and collective call order. + +Receivers accept `observedToken >= expectedToken`, not equality. A sender may +publish a later pass before the receiver samples an earlier token. Because +passes write non-overlapping offsets and use the same ordered QP, a later token +also proves completion of the earlier payload writes. + +Signal polling uses an MTE-to-UB load helper with a cycle timeout. It must not +use an unqualified ordinary GM pointer spin loop. + +## Data Path + +For each active send worker, group, and pass: + +1. Compute the target peer and chunk range. +2. Select the peer's explicit route/QP. +3. Issue one `UDMAPutSignalNbiOnQp` that transfers + `input[targetPeer][chunk]` to + `target.payloadPlane[slot][sourceRank][chunk]` and publishes + `target.signalPlane[slot][sourceRank]`. +4. Immediately call `UDMAQuietStatusOnQp` with the same peer and QP. + +For each active receive worker, group, and pass: + +1. Compute the source peer. +2. Poll `signalPlane[slot][sourcePeer]` until it reaches the expected token. +3. Copy `payloadPlane[slot][sourcePeer][chunk]` to + `output[sourcePeer][chunk]` using the local MTE relay helper. + +Payload, signal, and quiet must use the same explicit QP. The signal path must +not fall back to the generic QP0 route. + +## Progress and Buffer Safety + +There is no per-group or per-pass `SyncAll()`. Every worker owns independent +source or target slices and progresses according to its own signal dependency. +Kernel completion naturally waits for all 32 blocks to return. + +There is no cross-rank Host barrier between invocations. Invocations are +submitted in identical order on each rank and execute serially on one stream. +Ping-pong is sufficient under this collective contract: + +1. Invocation `I` writes plane `I & 1`. +2. A fast rank can enter invocation `I+1` and write the other plane. +3. It cannot complete `I+1` until it receives `I+1` data from every rank. +4. A slower rank cannot send `I+1` until its invocation `I` kernel has + completed, which means it has consumed all invocation `I` receive data. +5. Therefore plane `I & 1` is safe before any rank can enter invocation `I+2` + and reuse it. + +This proof does not hold if ranks skip or reorder calls, or if the same +communicator and buffers are used concurrently on multiple streams. Those uses +are out of scope. + +## QP Selection + +The kernel selects a deterministic valid QP for each peer using the current +max-weight route policy. The selected `qpIdx` is reused for payload plus signal +and immediate quiet. + +The implementation should expose QP selection as a small helper independent of +the old big-data fullmesh worker functions. This keeps route policy separate +from group scheduling and buffer addressing. + +## Failure Handling + +Host-side validation rejects invalid rank size, dimensions, memory capacity, +or missing UDMA registry before launch. + +Each device worker records only its first failure: + +```text +stage +group +pass +peer +qpIdx +quietStatus +expectedToken +observedToken +``` + +A quiet failure is recorded immediately. A receive wait records a timeout and +returns instead of spinning forever. Peer failure recovery is not attempted in +the first version; other ranks may also reach their timeouts. + +All early configuration exits must be uniform across blocks. The kernel does +not contain a block-divergent barrier. + +## Instrumentation + +The new kernel has a compact trace vocabulary rather than the old experimental +stage profiler: + +- kernel +- self-copy +- send-put-signal +- send-quiet +- receive-wait +- receive-copy + +Trace records include invocation, ping-pong slot, group, pass, lane, peer, and +QP. Instrumentation is optional and must not change buffer ownership or +synchronization. + +## Host Integration + +The demo adds a dedicated grouped-fullmesh mode and launch wrapper. Host code: + +1. Validates `rankSize`, payload dimensions, token dimensions, and the 1 GiB + registered-memory limit. +2. Allocates and registers the two payload planes plus signals and controls. +3. Initializes signals to zero once. +4. Launches one AllToAll operation per invocation with a globally consistent + invocation ID. +5. Uses Host-side warmup and repeat loops on the same stream. +6. Validates the full source-major output after the measured runs. + +No Host rank barrier is inserted between invocations. + +## Verification + +### Unit tests + +Peer-schedule tests cover rank sizes 8, 16, 24, 32, 40, and 64 and verify: + +- Every non-self peer appears exactly once. +- No peer appears twice. +- Pair scheduling is symmetric by group. +- The diameter peer appears once for even rank sizes. +- Last-group inactive lanes do not create peers. + +Layout tests verify: + +- Two payload planes do not overlap. +- Source-rank payload offsets and 128-byte signal slots are correct. +- Every remote range lies inside the registered allocation. +- Plans at and above the 1 GiB boundary are accepted or rejected correctly. +- Token fields do not overflow. + +Source/structure tests verify: + +- Exactly 32 blocks are launched. +- Payload plus signal uses `UDMAPutSignalNbiOnQp`. +- Immediate quiet uses the same `qpIdx`. +- No `SyncAll()` or generic signal route exists in the grouped kernel. + +### Physical validation + +The first physical target is 2x8 ranks. Validation includes: + +- Full output correctness for all 16 ranks. +- The 15-peer final-group boundary. +- Host warmup 5 and measured repeat 50. +- Alternating ping-pong planes across invocations. +- Loop 49 trace extraction for all 32 cores. +- Per-peer put, quiet, wait, and receive-copy timing. +- No timeout, quiet error, or stale-token acceptance. + +Larger topologies use the same correctness checks when available; host mapping +tests cover multi-group scheduling before 64-rank hardware validation. + +## Success Criteria + +The first version is complete when: + +- The new grouped kernel passes all mapping and layout unit tests. +- A physical 2x8 run passes full output validation on all ranks for warmup 5 + and repeat 50. +- Trace data shows at most 16 concurrent peers per rank and no global wait over + all peers. +- Payload, signal, and quiet for every send use one explicit QP. +- No registered-memory plan exceeds 1 GiB per rank. +- The old fullmesh kernel remains behaviorally unchanged. From 80fd6a18145027627b85f77b81167c83cf999d1a Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 17:39:17 +0800 Subject: [PATCH 060/163] docs(udma): plan grouped fullmesh alltoall --- ...26-07-20-udma-grouped-fullmesh-alltoall.md | 830 ++++++++++++++++++ 1 file changed, 830 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md b/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md new file mode 100644 index 00000000..c1ce0515 --- /dev/null +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md @@ -0,0 +1,830 @@ +# UDMA Grouped Fullmesh AllToAll Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a standalone 32-core grouped-fullmesh UDMA AllToAll kernel with two complete registered receive planes, explicit same-QP payload signals, and no device-wide or cross-rank barriers between invocations. + +**Architecture:** Peers are scheduled in symmetric circular groups of at most 16. Cores 0-15 send one peer per lane with `UDMAPutSignalNbiOnQp` followed by immediate same-QP quiet; cores 16-31 independently wait for one source per lane and copy its source-rank receive slot into output. Two full receive planes alternate by invocation ID, so ordered stream execution needs neither ACK nor Host rank barriers. + +**Tech Stack:** C++14 Host/layout code, Ascend C/Bisheng device code, TileXR UDMA APIs, Python 3 trace conversion, CMake, physical Ascend950 2x8 validation with CANN `/home/pkg/b101/cann`. + +## Global Constraints + +- Keep the existing 35-core big-data fullmesh kernel behavior unchanged. +- Add grouped fullmesh as dedicated demo `testType=8`; do not add mode branches to the old big-data worker loop. +- Support `rankSize >= 8` and `rankSize % 8 == 0` only. +- Use exactly 32 AIV blocks: 16 send workers and 16 receive workers. +- Use a fixed group width of 16, with eight forward and eight backward circular distances. +- Allocate `payloadPlane[2][rankSize][elementsPerPeer]` and `signalPlane[2][rankSize][128 bytes]` in one registered region. +- Reject any complete registered layout above 1 GiB per rank. +- Use one source-rank slot per global rank; never address payload by local worker lane. +- Payload plus ready signal and immediate quiet must use the same explicit max-weight QP. +- Do not use `SyncAll()` in the grouped kernel and do not insert Host rank barriers between invocations. +- Submit all invocations in identical order on one stream per rank; concurrent grouped collectives on multiple streams are out of scope. +- Commit locally, create and verify a Git bundle, then upload the bundle before every remote build or physical validation. +- Build and run against `/home/pkg/b101/cann`, with runtime driver libraries ahead of any CANN `devlib` stubs. +- Set `TILEXR_IPC_PID_MODE=pid` for physical multi-process runs. + +--- + +### Task 1: Group Schedule, Token, And Registered Layout + +**Files:** +- Create: `tests/udma/demo/tilexr_udma_alltoall_group_layout.h` +- Create: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` +- Modify: `tests/udma/CMakeLists.txt` + +**Interfaces:** +- Produces: `TileXR::Demo::AllToAllGroupPeer(int rank, int rankSize, uint32_t group, uint32_t lane) -> int32_t`. +- Produces: `TileXR::Demo::AllToAllGroupCount(int rankSize) -> uint32_t`. +- Produces: `TileXR::Demo::AllToAllGroupToken(uint32_t invocationId, uint32_t group, uint32_t pass) -> uint64_t`. +- Produces: `TileXR::Demo::PlanAllToAllGroup(int rankSize, int32_t elementsPerPeer, int32_t chunkElements) -> AllToAllGroupPlan`. +- Produces: layout offsets consumed by Host allocation and the device launch arguments in Tasks 2 and 3. + +- [ ] **Step 1: Write the failing group-layout unit test** + +Create a focused test executable with checks equivalent to: + +```cpp +#include +#include +#include +#include "demo/tilexr_udma_alltoall_group_layout.h" + +static void CheckSchedule(int rankSize) +{ + for (int rank = 0; rank < rankSize; ++rank) { + std::set peers; + for (uint32_t group = 0; group < TileXR::Demo::AllToAllGroupCount(rankSize); ++group) { + for (uint32_t lane = 0; lane < 16; ++lane) { + int peer = TileXR::Demo::AllToAllGroupPeer(rank, rankSize, group, lane); + if (peer < 0) continue; + CHECK_NE(peer, rank); + CHECK_EQ(peers.insert(peer).second, true); + bool symmetric = false; + for (uint32_t remoteLane = 0; remoteLane < 16; ++remoteLane) { + symmetric |= TileXR::Demo::AllToAllGroupPeer( + peer, rankSize, group, remoteLane) == rank; + } + CHECK_EQ(symmetric, true); + } + } + CHECK_EQ(peers.size(), static_cast(rankSize - 1)); + } +} + +int main() +{ + for (int rankSize : {8, 16, 24, 32, 40, 64}) CheckSchedule(rankSize); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(64), 4U); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 0), 1); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 8), 63); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 3, 7), 32); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 3, 15), -1); + return 0; +} +``` + +Add `test_tilexr_udma_alltoall_group_layout` to `tests/udma/CMakeLists.txt`, its include path, `INSTALL_TARGETS`, and the normal install set. + +- [ ] **Step 2: Run the test to verify RED** + +Commit the failing test only: + +```bash +git add tests/udma/CMakeLists.txt tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "test(udma): require grouped alltoall layout" +git bundle create tmp/grouped-alltoall-layout-red.bundle HEAD +git bundle verify tmp/grouped-alltoall-layout-red.bundle +``` + +Deploy to the build host and run: + +```bash +source /home/pkg/b101/cann/set_env.sh +export ASCEND_HOME_PATH=/home/pkg/b101/cann +cmake -S tests/udma -B tests/udma/build_b101 \ + -DCMAKE_INSTALL_PREFIX=$PWD/tests/udma/install_b101 \ + -DTILEXR_UDMA_DEMO_SOC_TYPE=Ascend950 +cmake --build tests/udma/build_b101 -j8 +``` + +Expected: FAIL because `tilexr_udma_alltoall_group_layout.h` does not exist. + +- [ ] **Step 3: Implement the host layout API** + +Create the header with these constants and data structure: + +```cpp +constexpr uint32_t kAllToAllGroupWidth = 16U; +constexpr uint32_t kAllToAllGroupHalfWidth = 8U; +constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; +constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; +constexpr uint32_t kAllToAllGroupBlockDim = 32U; +constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; + +struct AllToAllGroupPlan { + bool valid = false; + uint32_t groupCount = 0; + uint32_t passCount = 0; + int32_t chunkElements = 0; + size_t bytesPerPeer = 0; + size_t payloadPlaneBytes = 0; + size_t payloadOffset[2] = {0, 0}; + size_t signalPlaneBytes = 0; + size_t signalOffset[2] = {0, 0}; + size_t controlOffset = 0; + size_t controlBytes = 4096; + size_t registeredBytes = 0; +}; +``` + +Implement peer mapping without constructing vectors: + +```cpp +inline int32_t AllToAllGroupPeer( + int rank, int rankSize, uint32_t group, uint32_t lane) +{ + if (rankSize < 8 || rankSize % 8 != 0 || rank < 0 || rank >= rankSize || lane >= 16U) { + return -1; + } + const uint32_t index = lane < 8U ? lane : lane - 8U; + const int32_t distance = static_cast(group * 8U + index + 1U); + const int32_t diameter = rankSize / 2; + if (distance > diameter || (lane >= 8U && distance == diameter)) return -1; + return lane < 8U ? (rank + distance) % rankSize : + (rank - distance + rankSize) % rankSize; +} +``` + +Use checked multiplication/addition and 512-byte alignment in +`PlanAllToAllGroup`. Set `valid=true` only when all dimensions are positive, +the rank size is supported, token fields fit, every range is disjoint, and +`registeredBytes <= 1 GiB`. + +Build tokens exactly as specified: + +```cpp +inline uint64_t AllToAllGroupToken( + uint32_t invocationId, uint32_t group, uint32_t pass) +{ + const uint64_t invocation = static_cast(invocationId) + 1ULL; + const uint64_t slot = static_cast(invocationId & 1U); + return (invocation << 32U) | (slot << 31U) | + (static_cast(group) << 16U) | + (static_cast(pass) + 1ULL); +} +``` + +- [ ] **Step 4: Extend tests for memory and token boundaries** + +Add assertions for: + +```cpp +auto plan = TileXR::Demo::PlanAllToAllGroup(16, 2 * 1024 * 1024, 2 * 1024 * 1024); +CHECK_EQ(plan.valid, true); +CHECK_EQ(plan.groupCount, 1U); +CHECK_EQ(plan.passCount, 1U); +CHECK_EQ(plan.payloadPlaneBytes, 128ULL * 1024ULL * 1024ULL); +CHECK_EQ(plan.payloadOffset[1] >= plan.payloadOffset[0] + plan.payloadPlaneBytes, true); +CHECK_EQ(plan.signalOffset[0] >= plan.payloadOffset[1] + plan.payloadPlaneBytes, true); +CHECK_EQ(plan.registeredBytes <= (1ULL << 30), true); +CHECK_NE(TileXR::Demo::AllToAllGroupToken(48, 0, 0), + TileXR::Demo::AllToAllGroupToken(49, 0, 0)); +CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(7, 1024, 1024).valid, false); +CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(18, 1024, 1024).valid, false); +``` + +Construct the smallest payload above the 1 GiB layout boundary and verify it +is rejected without integer wraparound. + +- [ ] **Step 5: Run GREEN verification and commit** + +Run: + +```bash +cmake --build tests/udma/build_b101 -j8 +./tests/udma/build_b101/test_tilexr_udma_alltoall_group_layout +./tests/udma/build_b101/test_tilexr_udma_alltoall_layout +``` + +Expected: both executables print their pass messages and exit 0. + +Commit: + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_layout.h \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp tests/udma/CMakeLists.txt +git commit -m "feat(udma): add grouped alltoall layout" +``` + +--- + +### Task 2: Standalone 32-Core Grouped Kernel + +**Files:** +- Create: `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` +- Modify: `tests/udma/CMakeLists.txt` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Consumes: Task 1 plan offsets, dimensions, peer mapping semantics, and token format. +- Produces: `launch_tilexr_udma_all_to_all_group(uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR registeredMemory, GM_ADDR debug, uint32_t invocationId, int32_t elementsPerPeer, int32_t chunkElements, uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1)`. + +- [ ] **Step 1: Add failing kernel-structure checks** + +Read the new kernel source from the unit test and require these strings: + +```cpp +CHECK_CONTAINS(groupKernel, "tilexr_udma_all_to_all_group_kernel"); +CHECK_CONTAINS(groupKernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES"); +CHECK_CONTAINS(groupKernel, "UDMAPutSignalNbiOnQp"); +CHECK_CONTAINS(groupKernel, "UDMAQuietStatusOnQp"); +CHECK_CONTAINS(groupKernel, "AllToAllGroupWaitTokenMte"); +CHECK_CONTAINS(groupKernel, "observed >= expectedToken"); +CHECK_NOT_CONTAINS(groupKernel, "UDMAPutSignalNbi"); +CHECK_NOT_CONTAINS(groupKernel, "SyncAll"); +``` + +Also inspect the function slice containing the put and quiet calls and verify +both receive the same local variable `qpIdx`. + +- [ ] **Step 2: Run the source test to verify RED** + +Run the layout test. Expected: FAIL because the standalone kernel file and +required symbols do not exist. + +- [ ] **Step 3: Add the kernel to the Bisheng shared-library build** + +Update the custom Bisheng command to compile and link both sources: + +```cmake +"${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" +"${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" +``` + +Add the new source and group layout header to `DEPENDS`. Do not move or edit +the existing big-data kernel implementation. + +- [ ] **Step 4: Implement device-only helpers** + +Implement small helpers in an anonymous namespace: + +```cpp +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 32U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; + +__aicore__ inline int32_t AllToAllGroupDevicePeer( + int32_t rank, int32_t rankSize, uint32_t group, uint32_t lane); + +__aicore__ inline uint64_t AllToAllGroupDeviceToken( + uint32_t invocationId, uint32_t group, uint32_t pass); + +__aicore__ inline uint32_t AllToAllGroupSelectMaxWeightQp( + const __gm__ TileXR::CommArgs* args, int32_t peer); + +__aicore__ inline void AllToAllGroupCopyMte( + __gm__ uint8_t* dst, const __gm__ uint8_t* src, uint32_t bytes, + AscendC::LocalTensor relayLocal); + +__aicore__ inline bool AllToAllGroupWaitTokenMte( + __gm__ uint64_t* signal, uint64_t expectedToken, uint64_t timeoutCycles, + AscendC::LocalTensor relayLocal, uint64_t& observed); +``` + +QP selection must scan `GetUDMAInfo(args)->qpNum` and choose the largest +`UDMAGetQpWeight(info, peer, qpIdx)`, preserving the lowest QP index on ties. + +The wait helper repeatedly copies one 8-byte token GM-to-UB, performs the +required MTE event synchronization, and returns when `observed >= expectedToken` +or `GetSystemCycle() - begin >= timeoutCycles`. + +- [ ] **Step 5: Implement send and receive workers** + +The send path must use the source rank for the remote slot: + +```cpp +const uint32_t slot = invocationId & 1U; +const uint64_t remotePayloadOffset = payloadOffsets[slot] + + static_cast(rank) * bytesPerPeer + chunkByteOffset; +const uint64_t remoteSignalOffset = signalOffsets[slot] + + static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE; +const uint32_t qpIdx = AllToAllGroupSelectMaxWeightQp(args, peer); +TileXR::UDMAPutSignalNbiOnQp( + args, peer, qpIdx, localSrc, remotePayloadOffset, chunkBytes, + remoteSignalOffset, expectedToken); +const uint32_t quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); +``` + +The receive path uses the source peer for its local slot: + +```cpp +auto signal = reinterpret_cast<__gm__ uint64_t*>( + registeredMemory + signalOffsets[slot] + + static_cast(peer) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE); +uint64_t observed = 0; +if (!AllToAllGroupWaitTokenMte(signal, expectedToken, timeoutCycles, relayLocal, observed)) { + AllToAllGroupRecordError(debug, blockIdx, stage, group, pass, peer, 0U, + expectedToken, observed); + return; +} +auto relaySrc = registeredMemory + payloadOffsets[slot] + + static_cast(peer) * bytesPerPeer + chunkByteOffset; +auto relayDst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(peer) * elementsPerPeer + chunkElementOffset); +AllToAllGroupCopyMte(relayDst, relaySrc, chunkBytes, relayLocal); +``` + +Receive workers shard self-copy before entering peer groups. Each block owns +one `TPipe` and one 64 KiB relay buffer for its full lifetime. Inactive lanes +skip work. There is no `SyncAll()` and no shared mutable progress counter. + +- [ ] **Step 6: Add first-error debug recording** + +Reserve one fixed debug record per block with fields: + +```cpp +struct AllToAllGroupDeviceError { + uint32_t valid; + uint32_t stage; + uint32_t group; + uint32_t pass; + int32_t peer; + uint32_t qpIdx; + uint32_t quietStatus; + uint32_t reserved; + uint64_t expectedToken; + uint64_t observedToken; +}; +``` + +Only write when `valid == 0`. Quiet failure records immediately; receive wait +timeout records the expected and observed token and returns. + +- [ ] **Step 7: Build, run source guards, and commit** + +Run: + +```bash +cmake --build tests/udma/build_b101 -j8 +./tests/udma/build_b101/test_tilexr_udma_alltoall_group_layout +nm -D tests/udma/build_b101/libtilexr_udma_demo_kernel.so | \ + grep launch_tilexr_udma_all_to_all_group +``` + +Expected: build succeeds, source checks pass, and the launch wrapper is exported. + +Commit: + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp tests/udma/CMakeLists.txt +git commit -m "feat(udma): add grouped fullmesh alltoall kernel" +``` + +--- + +### Task 3: Dedicated Host Mode And Ping-Pong Invocation Loop + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Consumes: `PlanAllToAllGroup` and `launch_tilexr_udma_all_to_all_group`. +- Produces: demo `testType=8`, `TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS`, Host warmup/repeat invocation IDs, full output validation, and grouped debug reporting. + +- [ ] **Step 1: Add failing Host source checks** + +Require the Host source to contain: + +```cpp +CHECK_CONTAINS(demo, "testType == 8"); +CHECK_CONTAINS(demo, "PlanAllToAllGroup"); +CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_group"); +CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); +CHECK_CONTAINS(demo, "grouped alltoall warmup="); +CHECK_NOT_CONTAINS(groupModeSlice, "DemoBarrierAll"); +``` + +The group-mode slice begins at its `if (testType == 8)` allocation branch and +ends before the existing `testType == 7` branch. + +- [ ] **Step 2: Run RED verification** + +Run the group-layout test. Expected: FAIL because `testType=8` plumbing is absent. + +- [ ] **Step 3: Add plan validation and allocation** + +Treat `testType=8` as strict AllToAll UDMA with output and no IPC fallback. +Build the plan from: + +```cpp +const int32_t requestedChunkElements = std::max( + 1, GetEnvInt("TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS", elementsPerRank)); +const auto groupPlan = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerRank, requestedChunkElements); +``` + +Reject invalid plans before allocation. Allocate normal `groupInput` and +`groupOutput`, allocate `groupPlan.registeredBytes` with `aclrtMalloc`, zero +both signal planes and controls once, and register the complete region with +`TileXRUDMARegister`. + +Print exact offsets, plane sizes, group count, pass count, and total bytes. + +- [ ] **Step 4: Add warmup and measured invocation launches** + +Use one monotonically increasing invocation ID across warmup and measured calls: + +```cpp +uint32_t invocationId = 0; +for (int iter = 0; iter < allToAllWarmup; ++iter, ++invocationId) { + launch_tilexr_udma_all_to_all_group( + 32U, stream, commArgsDev, reinterpret_cast(groupInput), + reinterpret_cast(groupOutput), + reinterpret_cast(registeredMemory), reinterpret_cast(debug), + invocationId, elementsPerRank, groupPlan.chunkElements, + groupPlan.passCount, groupPlan.groupCount, + groupPlan.payloadOffset[0], groupPlan.payloadOffset[1], + groupPlan.signalOffset[0], groupPlan.signalOffset[1]); +} +if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", + aclrtSynchronizeStream(stream))) { + return 1; +} + +auto begin = std::chrono::steady_clock::now(); +for (int iter = 0; iter < allToAllRepeat; ++iter, ++invocationId) { + launch_tilexr_udma_all_to_all_group( + 32U, stream, commArgsDev, reinterpret_cast(groupInput), + reinterpret_cast(groupOutput), + reinterpret_cast(registeredMemory), reinterpret_cast(debug), + invocationId, elementsPerRank, groupPlan.chunkElements, + groupPlan.passCount, groupPlan.groupCount, + groupPlan.payloadOffset[0], groupPlan.payloadOffset[1], + groupPlan.signalOffset[0], groupPlan.signalOffset[1]); +} +if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", + aclrtSynchronizeStream(stream))) { + return 1; +} +auto end = std::chrono::steady_clock::now(); +``` + +Do not synchronize or call `DemoBarrierAll` between invocations. One stream +must serialize all launches. + +- [ ] **Step 5: Add result and debug validation** + +Copy the full output once after measured synchronization and call the existing +`ValidateAllToAllData`. Copy all 32 debug records and reject any record with +`valid != 0`, printing every recorded field. + +Ensure cleanup unregisters before freeing the registered region and frees the +normal input/output exactly once on every error path. + +- [ ] **Step 6: Build, run Host source tests, and commit** + +Run: + +```bash +cmake --build tests/udma/build_b101 -j8 +./tests/udma/build_b101/test_tilexr_udma_alltoall_group_layout +./tests/udma/build_b101/test_tilexr_udma_alltoall_layout +``` + +Expected: both tests pass and `tilexr_udma_demo` links against the new wrapper. + +Commit: + +```bash +git add tests/udma/demo/tilexr_udma_demo.cpp \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "feat(udma): add grouped alltoall demo mode" +``` + +--- + +### Task 4: Compact Grouped-Kernel GM Trace + +**Files:** +- Create: `tests/udma/demo/tilexr_udma_alltoall_group_trace.h` +- Create: `tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py` +- Create: `tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py` +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify: `tests/udma/CMakeLists.txt` + +**Interfaces:** +- Consumes: invocation, slot, group, pass, lane, peer, and QP from Tasks 2 and 3. +- Produces: optional 8 MiB per-rank raw trace and Chrome Trace JSON phases `kernel`, `self-copy`, `send-put-signal`, `send-quiet`, `receive-wait`, and `receive-copy`. +- Extends: `launch_tilexr_udma_all_to_all_group` with trailing arguments + `GM_ADDR groupTrace, uint32_t traceIteration`; warmup passes `nullptr`, and + measured invocation `iter` passes the trace allocation and `iter`. + +- [ ] **Step 1: Write failing Python trace tests** + +Construct a synthetic binary with one invocation and verify: + +```python +events = module.build_chrome_trace([module.read_rank_trace(path)]) +names = {event["name"] for event in events["traceEvents"] if event["ph"] == "X"} +self.assertEqual(names, { + "kernel", "self-copy", "send-put-signal", "send-quiet", + "receive-wait", "receive-copy", +}) +self.assertEqual(events["otherData"]["displayTimeUnit"], "ns") +``` + +Also test rejection of bad magic, partial spans, out-of-capacity dimensions, +and non-8-MiB files. + +- [ ] **Step 2: Run trace tests to verify RED** + +Run: + +```bash +python3 -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +Expected: FAIL because the converter does not exist. + +- [ ] **Step 3: Implement the trace binary layout and converter** + +Use an 8 MiB fixed allocation with: + +```text +header +kernelSpan[50][32] +taskSpan[50][32][groupCount][passCount][5] +``` + +The five task phases exclude `kernel`; receive cores record self-copy in the +group-0/pass-0 self-copy cell reserved for that core. The Host validates the computed required bytes +before enabling trace. + +Write Chrome JSON incrementally with `json.dump`; use `displayTimeUnit: "ns"` +or omit the field, never `"us"`. Normalize cycles independently per rank and +invocation. + +- [ ] **Step 4: Instrument the kernel without changing synchronization** + +Add begin/end cycle recording around exactly these operations: + +```text +self MTE copy +UDMAPutSignalNbiOnQp +UDMAQuietStatusOnQp +signal wait loop +receive MTE copy +whole kernel block +``` + +Trace writes use disjoint offsets per block/group/pass/phase. Do not add +barriers, flags, or shared counters for tracing. + +- [ ] **Step 5: Add Host allocation and output** + +Enable with: + +```text +TILEXR_UDMA_GROUP_TRACE=1 +TILEXR_UDMA_GROUP_TRACE_DIR=/tmp/tilexr_group_trace +``` + +Allocate trace GM separately from registered UDMA memory, initialize it once, +pass `traceIteration=iter` for measured launches, copy it back after stream +synchronization, and write: + +```text +tilexr_group_trace_rank_{rank}.bin +``` + +- [ ] **Step 6: Run all trace tests and commit** + +Run: + +```bash +python3 -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +cmake --build tests/udma/build_b101 -j8 +./tests/udma/build_b101/test_tilexr_udma_alltoall_group_layout +``` + +Expected: all Python tests pass, Bisheng build succeeds, and layout/source +guards remain green. + +Commit: + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_trace.h \ + tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py \ + tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp \ + tests/udma/demo/tilexr_udma_demo.cpp \ + tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py \ + tests/udma/CMakeLists.txt +git commit -m "feat(udma): trace grouped alltoall pipeline" +``` + +--- + +### Task 5: Full Verification, Bundle, And Two-Host Deployment + +**Files:** +- Modify only if verification exposes a grouped-mode defect in files from Tasks 1-4. +- Create artifact: `tmp/udma-grouped-fullmesh-alltoall.bundle` + +**Interfaces:** +- Consumes: committed grouped layout, kernel, Host mode, and trace converter. +- Produces: one verified complete Git bundle and identical source commits on `141.61.50.31` and `141.61.49.223`. + +- [ ] **Step 1: Run the complete local/source verification set** + +Run: + +```bash +git diff --check +rg -n "SyncAll|UDMAPutSignalNbi|UDMAQuietStatus\(" \ + tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +python3 -m unittest \ + tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +Expected: `git diff --check` is clean; `rg` finds no forbidden grouped-kernel +calls; Python tests pass. + +- [ ] **Step 2: Build and test on the b101 build host** + +On `141.61.49.223`: + +```bash +source /home/pkg/b101/cann/set_env.sh +export ASCEND_HOME_PATH=/home/pkg/b101/cann +export PATH=/home/pkg/b101/cann/bin:$PATH +cmake -S . -B build -DCMAKE_INSTALL_PREFIX=$PWD/install +cmake --build build -j8 && cmake --install build +cmake -S tests/udma -B tests/udma/build_b101 \ + -DCMAKE_INSTALL_PREFIX=$PWD/tests/udma/install_b101 \ + -DTILEXR_UDMA_DEMO_SOC_TYPE=Ascend950 +cmake --build tests/udma/build_b101 -j8 && cmake --install tests/udma/build_b101 +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_layout +./tests/udma/install_b101/bin/test_tilexr_udma_transport_layout +python3 -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +Expected: all commands exit 0. + +- [ ] **Step 3: Commit any verification fix separately** + +If verification required a code correction, rerun Step 1 and Step 2 and commit +only that correction with the focused message +`fix(udma): correct grouped alltoall verification defect`. Do not bundle +uncommitted changes. + +- [ ] **Step 4: Create and verify the deployment bundle** + +Run locally: + +```bash +git status --short +git bundle create tmp/udma-grouped-fullmesh-alltoall.bundle HEAD +git bundle verify tmp/udma-grouped-fullmesh-alltoall.bundle +``` + +Expected: only known unrelated untracked paths remain; bundle verification +reports a complete history containing current HEAD. + +- [ ] **Step 5: Deploy identical sources to both hosts** + +Upload the bundle to both hosts, fetch it into +`/home/h30059441/tilexr_grouped_alltoall_b101`, and detach at the exact local +HEAD. Verify `git rev-parse HEAD` matches on both. + +`141.61.50.31` currently has no `cmake`. Build on `141.61.49.223`, then copy +these exact artifacts to the same install paths on `141.61.50.31`: + +```text +install/lib64/libtile-comm.so +tests/udma/install_b101/lib/libtilexr_udma_demo_kernel.so +tests/udma/install_b101/bin/tilexr_udma_demo +``` + +Compare SHA-256 hashes on both hosts before running. + +--- + +### Task 6: Physical 2x8 Correctness And Performance Validation + +**Files:** +- Create artifacts under: `tmp/grouped_alltoall_b101_2x8/` +- Do not commit raw trace binaries or run logs. + +**Interfaces:** +- Consumes: `testType=8`, 32-core kernel, b101 artifacts, Host warmup/repeat, and grouped raw trace. +- Produces: 16-rank correctness evidence, Host timing, iteration-49 Chrome trace, and per-phase summaries. + +- [ ] **Step 1: Run a small correctness smoke test** + +Use ranks 0-7 on `141.61.50.31` and ranks 8-15 on `141.61.49.223`. Export on +both hosts: + +```bash +export ASCEND_HOME_PATH=/home/pkg/b101/cann +export PATH=/home/pkg/b101/cann/bin:$PATH +export LD_LIBRARY_PATH=$PWD/tests/udma/install_b101/lib:$PWD/install/lib64:\ +/home/pkg/b101/cann/aarch64-linux/lib64:/usr/local/Ascend/driver/lib64/driver:\ +/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64 +export TILEXR_COMM_ID=141.61.50.31:64720 +export TILEXR_DEMO_BARRIER_HOST=141.61.50.31 +export TILEXR_IPC_PID_MODE=pid +export TILEXR_UDMA_ROUTE_POLICY=all +export TILEXR_UDMA_QP_NUM=4 +export TILEXR_DEMO_ALLTOALL_WARMUP=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=2 +``` + +Run each rank with: + +```bash +timeout 180s ./tests/udma/install_b101/bin/tilexr_udma_demo \ + 16 "$rank" 8 4096 8 0 +``` + +Expected: 16 success lines, no timeout, no quiet error, and full output +validation succeeds on every rank. + +- [ ] **Step 2: Run the target 128 MiB per-rank matrix** + +Set: + +```bash +export TILEXR_DEMO_ALLTOALL_WARMUP=5 +export TILEXR_DEMO_ALLTOALL_REPEAT=50 +export TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS=2097152 +export TILEXR_UDMA_GROUP_TRACE=1 +export TILEXR_UDMA_GROUP_TRACE_DIR=/home/h30059441/grouped_alltoall_b101_2x8 +``` + +Run: + +```bash +timeout 300s ./tests/udma/install_b101/bin/tilexr_udma_demo \ + 16 "$rank" 8 2097152 8 0 +``` + +This is 8 MiB per peer and 128 MiB input/output per rank. Expected: 16 success +lines and sixteen 8 MiB trace files. + +- [ ] **Step 3: Verify ping-pong and absence of barriers** + +Parse all traces and verify: + +- With warmup 5, measured iteration 48 uses invocation ID 53 and slot 1; + measured iteration 49 uses invocation ID 54 and slot 0. +- Every rank has 15 active peer send and receive sequences. +- No lane waits for self or the duplicate diameter peer. +- Every send event records an explicit valid QP. +- No device `SyncAll` phase exists. +- All waits complete without timeout. + +- [ ] **Step 4: Generate loop-49 Chrome trace and phase summary** + +Convert only measured invocation 49 into: + +```text +tmp/grouped_alltoall_b101_2x8/grouped_alltoall_loop49_trace.json +``` + +Report per core and aggregate: + +```text +self-copy +send-put-signal +send-quiet +receive-wait +receive-copy +kernel envelope +Host mean/min/max +``` + +Compute payload bandwidth from 128 MiB per rank and compare the kernel envelope +against the old fullmesh 2x8 baseline. Keep correctness and performance +conclusions separate. + +- [ ] **Step 5: Final regression verification** + +After physical validation, rerun on the build host: + +```bash +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_layout +./tests/udma/install_b101/bin/test_tilexr_udma_transport_layout +python3 -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +git status --short +``` + +Expected: all tests pass; only ignored/untracked run artifacts remain; no +tracked source differs from the deployed commit. From a5c1686c3d827bc5ab6d4c7460003052c2a01cca Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:10:43 +0800 Subject: [PATCH 061/163] docs(udma): bound grouped alltoall to 128 ranks --- .../plans/2026-07-20-udma-grouped-fullmesh-alltoall.md | 6 ++++-- .../2026-07-20-udma-grouped-fullmesh-alltoall-design.md | 3 ++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md b/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md index c1ce0515..b152dde8 100644 --- a/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-fullmesh-alltoall.md @@ -12,7 +12,8 @@ - Keep the existing 35-core big-data fullmesh kernel behavior unchanged. - Add grouped fullmesh as dedicated demo `testType=8`; do not add mode branches to the old big-data worker loop. -- Support `rankSize >= 8` and `rankSize % 8 == 0` only. +- Support `8 <= rankSize <= 128` and `rankSize % 8 == 0` only, matching the + current `TILEXR_MAX_RANK_SIZE` and `CommArgs` peer-array limit. - Use exactly 32 AIV blocks: 16 send workers and 16 receive workers. - Use a fixed group width of 16, with eight forward and eight backward circular distances. - Allocate `payloadPlane[2][rankSize][elementsPerPeer]` and `signalPlane[2][rankSize][128 bytes]` in one registered region. @@ -145,7 +146,8 @@ Implement peer mapping without constructing vectors: inline int32_t AllToAllGroupPeer( int rank, int rankSize, uint32_t group, uint32_t lane) { - if (rankSize < 8 || rankSize % 8 != 0 || rank < 0 || rank >= rankSize || lane >= 16U) { + if (rankSize < 8 || rankSize > 128 || rankSize % 8 != 0 || + rank < 0 || rank >= rankSize || lane >= 16U) { return -1; } const uint32_t index = lane < 8U ? lane : lane - 8U; diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md index 19e6d5d9..d02f7f16 100644 --- a/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-fullmesh-alltoall-design.md @@ -19,7 +19,8 @@ aggregation, ACK protocol, profiling stages, or experimental mode branches. The first version supports: -- `rankSize >= 8` and `rankSize % 8 == 0`. +- `8 <= rankSize <= 128` and `rankSize % 8 == 0`, matching the current + `TILEXR_MAX_RANK_SIZE` and `CommArgs` peer-array limit. - A fixed peer-group width of 16: eight forward and eight backward circular distances. - Any payload whose complete double-buffered registered receive layout, From c107930c1ad3a4ef31e2474b2432a96ede9102fc Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:12:20 +0800 Subject: [PATCH 062/163] test(udma): require grouped alltoall layout --- tests/udma/CMakeLists.txt | 13 ++ ...test_tilexr_udma_alltoall_group_layout.cpp | 133 ++++++++++++++++++ 2 files changed, 146 insertions(+) create mode 100644 tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index d0ef03ea..0262f23a 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -86,6 +86,18 @@ target_compile_definitions(test_tilexr_udma_alltoall_layout PRIVATE TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" ) +add_executable(test_tilexr_udma_alltoall_group_layout + unit/test_tilexr_udma_alltoall_group_layout.cpp +) + +target_include_directories(test_tilexr_udma_alltoall_group_layout PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} +) + +target_compile_definitions(test_tilexr_udma_alltoall_group_layout PRIVATE + TILEXR_SOURCE_ROOT="${TILEXR_ROOT}" +) + add_executable(test_tilexr_udma_allreduce_layout unit/test_tilexr_udma_allreduce_layout.cpp ) @@ -144,6 +156,7 @@ set(INSTALL_TARGETS test_tilexr_udma test_tilexr_udma_registry test_tilexr_udma_alltoall_layout + test_tilexr_udma_alltoall_group_layout test_tilexr_udma_allreduce_layout test_tilexr_chip_map_sources test_tilexr_ipc_pid_mode_sources diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp new file mode 100644 index 00000000..38733634 --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -0,0 +1,133 @@ +#include +#include +#include + +#include "demo/tilexr_udma_alltoall_group_layout.h" + +namespace { + +int g_failures = 0; + +#define CHECK_EQ(lhs, rhs) \ + do { \ + auto lhsValue = (lhs); \ + auto rhsValue = (rhs); \ + if (lhsValue != rhsValue) { \ + std::cerr << "CHECK_EQ failed at line " << __LINE__ << ": " #lhs " != " #rhs \ + << " (" << lhsValue << " vs " << rhsValue << ")" << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +void CheckSchedule(int rankSize) +{ + for (int rank = 0; rank < rankSize; ++rank) { + std::set peers; + for (uint32_t group = 0; group < TileXR::Demo::AllToAllGroupCount(rankSize); ++group) { + for (uint32_t lane = 0; lane < TileXR::Demo::kAllToAllGroupWidth; ++lane) { + const int peer = TileXR::Demo::AllToAllGroupPeer(rank, rankSize, group, lane); + if (peer < 0) { + continue; + } + CHECK_EQ(peer == rank, false); + CHECK_EQ(peers.insert(peer).second, true); + bool symmetric = false; + for (uint32_t remoteLane = 0; remoteLane < TileXR::Demo::kAllToAllGroupWidth; + ++remoteLane) { + symmetric = symmetric || + TileXR::Demo::AllToAllGroupPeer(peer, rankSize, group, remoteLane) == rank; + } + CHECK_EQ(symmetric, true); + } + } + CHECK_EQ(peers.size(), static_cast(rankSize - 1)); + } +} + +void TestSchedules() +{ + for (int rankSize : {8, 16, 24, 32, 40, 64, 128}) { + CheckSchedule(rankSize); + } + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(8), 1U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(16), 1U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(24), 2U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(64), 4U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(128), 8U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(7), 0U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(18), 0U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(136), 0U); + + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 0), 1); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 7), 8); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 8), 63); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 15), 56); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 3, 7), 32); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 3, 15), -1); +} + +void TestPlan() +{ + constexpr int rankSize = 16; + constexpr int32_t elementsPerPeer = 2 * 1024 * 1024; + const auto plan = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, elementsPerPeer); + + CHECK_EQ(plan.valid, true); + CHECK_EQ(plan.groupCount, 1U); + CHECK_EQ(plan.passCount, 1U); + CHECK_EQ(plan.chunkElements, elementsPerPeer); + CHECK_EQ(plan.bytesPerPeer, 8ULL * 1024ULL * 1024ULL); + CHECK_EQ(plan.payloadPlaneBytes, 128ULL * 1024ULL * 1024ULL); + CHECK_EQ(plan.payloadOffset[0], 0ULL); + CHECK_EQ(plan.payloadOffset[1] >= plan.payloadOffset[0] + plan.payloadPlaneBytes, true); + CHECK_EQ(plan.signalPlaneBytes, static_cast(rankSize) * 128ULL); + CHECK_EQ(plan.signalOffset[0] >= plan.payloadOffset[1] + plan.payloadPlaneBytes, true); + CHECK_EQ(plan.signalOffset[1] >= plan.signalOffset[0] + plan.signalPlaneBytes, true); + CHECK_EQ(plan.controlOffset >= plan.signalOffset[1] + plan.signalPlaneBytes, true); + CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); + + const auto chunked = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, elementsPerPeer / 4); + CHECK_EQ(chunked.valid, true); + CHECK_EQ(chunked.passCount, 4U); + CHECK_EQ(chunked.payloadPlaneBytes, plan.payloadPlaneBytes); + + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(7, 1024, 1024).valid, false); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(18, 1024, 1024).valid, false); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(136, 1024, 1024).valid, false); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(16, 0, 1024).valid, false); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(16, elementsPerPeer, 0).valid, false); + + constexpr int32_t thirtyTwoMiBElements = 8 * 1024 * 1024; + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup( + rankSize, thirtyTwoMiBElements, thirtyTwoMiBElements).valid, false); +} + +void TestTokens() +{ + const uint64_t token48 = TileXR::Demo::AllToAllGroupToken(48U, 0U, 0U); + const uint64_t token49 = TileXR::Demo::AllToAllGroupToken(49U, 0U, 0U); + CHECK_EQ(token48 != 0U, true); + CHECK_EQ(token49 > token48, true); + CHECK_EQ((token48 >> 31U) & 1ULL, 0ULL); + CHECK_EQ((token49 >> 31U) & 1ULL, 1ULL); + CHECK_EQ(TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U) > token49, true); + CHECK_EQ(TileXR::Demo::AllToAllGroupToken(49U, 1U, 1U) > + TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U), true); +} + +} // namespace + +int main() +{ + TestSchedules(); + TestPlan(); + TestTokens(); + if (g_failures != 0) { + std::cerr << "TileXR grouped all-to-all layout checks failed: " << g_failures << std::endl; + return 1; + } + std::cout << "TileXR grouped all-to-all layout checks passed" << std::endl; + return 0; +} From b90d67408f5fd32bb2864316aefd1f9740d55757 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:15:47 +0800 Subject: [PATCH 063/163] feat(udma): add grouped alltoall layout --- .../demo/tilexr_udma_alltoall_group_layout.h | 170 ++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 tests/udma/demo/tilexr_udma_alltoall_group_layout.h diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h new file mode 100644 index 00000000..2e63c290 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -0,0 +1,170 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_ALLTOALL_GROUP_LAYOUT_H +#define TILEXR_UDMA_ALLTOALL_GROUP_LAYOUT_H + +#include +#include +#include +#include + +namespace TileXR { +namespace Demo { + +constexpr int32_t kAllToAllGroupMinRankSize = 8; +constexpr int32_t kAllToAllGroupMaxRankSize = 128; +constexpr uint32_t kAllToAllGroupWidth = 16U; +constexpr uint32_t kAllToAllGroupHalfWidth = 8U; +constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; +constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; +constexpr uint32_t kAllToAllGroupBlockDim = 32U; +constexpr size_t kAllToAllGroupAlignment = 512U; +constexpr size_t kAllToAllGroupControlBytes = 4096U; +constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; + +struct AllToAllGroupPlan { + bool valid = false; + uint32_t groupCount = 0; + uint32_t passCount = 0; + int32_t chunkElements = 0; + size_t bytesPerPeer = 0; + size_t payloadPlaneBytes = 0; + size_t payloadOffset[kAllToAllGroupPingPongSlots] = {0, 0}; + size_t signalPlaneBytes = 0; + size_t signalOffset[kAllToAllGroupPingPongSlots] = {0, 0}; + size_t controlOffset = 0; + size_t controlBytes = kAllToAllGroupControlBytes; + size_t registeredBytes = 0; +}; + +inline bool AllToAllGroupValidRankSize(int rankSize) +{ + return rankSize >= kAllToAllGroupMinRankSize && + rankSize <= kAllToAllGroupMaxRankSize && rankSize % 8 == 0; +} + +inline uint32_t AllToAllGroupCount(int rankSize) +{ + if (!AllToAllGroupValidRankSize(rankSize)) { + return 0U; + } + return static_cast((rankSize - 1 + static_cast(kAllToAllGroupWidth) - 1) / + static_cast(kAllToAllGroupWidth)); +} + +inline int32_t AllToAllGroupPeer( + int rank, int rankSize, uint32_t group, uint32_t lane) +{ + if (!AllToAllGroupValidRankSize(rankSize) || rank < 0 || rank >= rankSize || + lane >= kAllToAllGroupWidth || group >= AllToAllGroupCount(rankSize)) { + return -1; + } + const uint32_t index = lane < kAllToAllGroupHalfWidth ? + lane : lane - kAllToAllGroupHalfWidth; + const int32_t distance = static_cast( + group * kAllToAllGroupHalfWidth + index + 1U); + const int32_t diameter = rankSize / 2; + if (distance > diameter || + (lane >= kAllToAllGroupHalfWidth && distance == diameter)) { + return -1; + } + return lane < kAllToAllGroupHalfWidth ? + (rank + distance) % rankSize : + (rank - distance + rankSize) % rankSize; +} + +inline uint64_t AllToAllGroupToken( + uint32_t invocationId, uint32_t group, uint32_t pass) +{ + const uint64_t invocation = static_cast(invocationId) + 1ULL; + const uint64_t slot = static_cast(invocationId & 1U); + return (invocation << 32U) | (slot << 31U) | + (static_cast(group) << 16U) | + (static_cast(pass) + 1ULL); +} + +inline bool AllToAllGroupCheckedAdd(size_t lhs, size_t rhs, size_t& result) +{ + if (rhs > std::numeric_limits::max() - lhs) { + return false; + } + result = lhs + rhs; + return true; +} + +inline bool AllToAllGroupCheckedMul(size_t lhs, size_t rhs, size_t& result) +{ + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return false; + } + result = lhs * rhs; + return true; +} + +inline bool AllToAllGroupAlignUp(size_t value, size_t& result) +{ + size_t expanded = 0; + if (!AllToAllGroupCheckedAdd(value, kAllToAllGroupAlignment - 1U, expanded)) { + return false; + } + result = expanded & ~(kAllToAllGroupAlignment - 1U); + return true; +} + +inline AllToAllGroupPlan PlanAllToAllGroup( + int rankSize, int32_t elementsPerPeer, int32_t chunkElements) +{ + AllToAllGroupPlan plan {}; + if (!AllToAllGroupValidRankSize(rankSize) || elementsPerPeer <= 0 || chunkElements <= 0) { + return plan; + } + + plan.groupCount = AllToAllGroupCount(rankSize); + plan.chunkElements = std::min(elementsPerPeer, chunkElements); + plan.passCount = static_cast( + (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1ULL) / + static_cast(plan.chunkElements)); + if (plan.groupCount == 0U || plan.groupCount > 0x7FFFU || + plan.passCount == 0U || plan.passCount > 0xFFFFU) { + return AllToAllGroupPlan {}; + } + + if (!AllToAllGroupCheckedMul( + static_cast(elementsPerPeer), sizeof(int32_t), plan.bytesPerPeer) || + !AllToAllGroupCheckedMul( + static_cast(rankSize), plan.bytesPerPeer, plan.payloadPlaneBytes) || + !AllToAllGroupCheckedMul( + static_cast(rankSize), kAllToAllGroupSignalSlotBytes, plan.signalPlaneBytes)) { + return AllToAllGroupPlan {}; + } + + plan.payloadOffset[0] = 0U; + if (!AllToAllGroupAlignUp(plan.payloadPlaneBytes, plan.payloadOffset[1])) { + return AllToAllGroupPlan {}; + } + + size_t cursor = 0; + if (!AllToAllGroupCheckedAdd(plan.payloadOffset[1], plan.payloadPlaneBytes, cursor) || + !AllToAllGroupAlignUp(cursor, plan.signalOffset[0]) || + !AllToAllGroupCheckedAdd(plan.signalOffset[0], plan.signalPlaneBytes, cursor) || + !AllToAllGroupAlignUp(cursor, plan.signalOffset[1]) || + !AllToAllGroupCheckedAdd(plan.signalOffset[1], plan.signalPlaneBytes, cursor) || + !AllToAllGroupAlignUp(cursor, plan.controlOffset) || + !AllToAllGroupCheckedAdd(plan.controlOffset, plan.controlBytes, cursor) || + !AllToAllGroupAlignUp(cursor, plan.registeredBytes)) { + return AllToAllGroupPlan {}; + } + if (plan.registeredBytes > kAllToAllGroupMaxRegisteredBytes) { + return AllToAllGroupPlan {}; + } + plan.valid = true; + return plan; +} + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_ALLTOALL_GROUP_LAYOUT_H From dad5ee7246bde1b0516a0ed63df00eaf326a28c0 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:18:19 +0800 Subject: [PATCH 064/163] test(udma): require standalone grouped kernel --- ...test_tilexr_udma_alltoall_group_layout.cpp | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 38733634..5898c734 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -1,11 +1,18 @@ #include +#include #include #include +#include +#include #include "demo/tilexr_udma_alltoall_group_layout.h" namespace { +#ifndef TILEXR_SOURCE_ROOT +#define TILEXR_SOURCE_ROOT "." +#endif + int g_failures = 0; #define CHECK_EQ(lhs, rhs) \ @@ -19,6 +26,30 @@ int g_failures = 0; } \ } while (0) +#define CHECK_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) == std::string::npos) { \ + std::cerr << "CHECK_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +#define CHECK_NOT_CONTAINS(text, needle) \ + do { \ + if ((text).find(needle) != std::string::npos) { \ + std::cerr << "CHECK_NOT_CONTAINS failed at line " << __LINE__ << ": " << needle << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +std::string ReadFile(const std::string& path) +{ + std::ifstream in(path.c_str()); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + void CheckSchedule(int rankSize) { for (int rank = 0; rank < rankSize; ++rank) { @@ -117,6 +148,21 @@ void TestTokens() TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U), true); } +void TestKernelStructure() +{ + const std::string kernel = ReadFile( + std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp"); + CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_kernel"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES"); + CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); + CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); + CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); + CHECK_CONTAINS(kernel, "observed >= expectedToken"); + CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_group"); + CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_NOT_CONTAINS(kernel, "SyncAll"); +} + } // namespace int main() @@ -124,6 +170,7 @@ int main() TestSchedules(); TestPlan(); TestTokens(); + TestKernelStructure(); if (g_failures != 0) { std::cerr << "TileXR grouped all-to-all layout checks failed: " << g_failures << std::endl; return 1; From 32be9445ecaf7226f32cb7abe16a9181fa68e200 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:22:34 +0800 Subject: [PATCH 065/163] feat(udma): add grouped fullmesh alltoall kernel --- tests/udma/CMakeLists.txt | 3 + .../tilexr_udma_alltoall_group_kernel.cpp | 305 ++++++++++++++++++ 2 files changed, 308 insertions(+) create mode 100644 tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 0262f23a..035e3e83 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -241,6 +241,7 @@ if(BUILD_TILEXR_UDMA_DEMO) -DTILEXR_UDMA_FORCE_ENABLE=1 ${TILEXR_UDMA_DEMO_KERNEL_INCLUDES} "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" -L${ASCEND_DRIVER_PATH}/lib64/driver -L${ASCEND_HOME_PATH}/${ARCH}-linux/lib64 -L${ASCEND_HOME_PATH}/${ARCH}-linux/devlib @@ -257,6 +258,8 @@ if(BUILD_TILEXR_UDMA_DEMO) -o "${TILEXR_UDMA_DEMO_KERNEL_SO}" DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_layout.h" "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_fullmesh_trace.h" "${TILEXR_ROOT}/src/include/tilexr_udma.h" "${TILEXR_ROOT}/src/include/tilexr_data_as_flag.h" diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp new file mode 100644 index 00000000..8b99a04f --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -0,0 +1,305 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#include "kernel_operator.h" +#include "tilexr_udma.h" + +namespace { + +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 32U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; +constexpr uint64_t TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES = 10000000000ULL; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_WORDS = 12U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; + +struct AllToAllGroupDeviceError { + uint32_t valid; + uint32_t stage; + uint32_t group; + uint32_t pass; + int32_t peer; + uint32_t qpIdx; + uint32_t quietStatus; + uint32_t reserved; + uint64_t expectedToken; + uint64_t observedToken; +}; + +__aicore__ inline int32_t AllToAllGroupDevicePeer( + int32_t rank, int32_t rankSize, uint32_t group, uint32_t lane) +{ + if (rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || + (rankSize & 7) != 0 || rank < 0 || rank >= rankSize || lane >= 16U) { + return -1; + } + const uint32_t index = lane < TILEXR_ALLTOALL_GROUP_HALF_WIDTH ? + lane : lane - TILEXR_ALLTOALL_GROUP_HALF_WIDTH; + const int32_t distance = static_cast( + group * TILEXR_ALLTOALL_GROUP_HALF_WIDTH + index + 1U); + const int32_t diameter = rankSize / 2; + if (distance > diameter || + (lane >= TILEXR_ALLTOALL_GROUP_HALF_WIDTH && distance == diameter)) { + return -1; + } + return lane < TILEXR_ALLTOALL_GROUP_HALF_WIDTH ? + (rank + distance) % rankSize : + (rank - distance + rankSize) % rankSize; +} + +__aicore__ inline uint64_t AllToAllGroupDeviceToken( + uint32_t invocationId, uint32_t group, uint32_t pass) +{ + const uint64_t invocation = static_cast(invocationId) + 1ULL; + const uint64_t slot = static_cast(invocationId & 1U); + return (invocation << 32U) | (slot << 31U) | + (static_cast(group) << 16U) | + (static_cast(pass) + 1ULL); +} + +__aicore__ inline uint32_t AllToAllGroupSelectMaxWeightQp( + const __gm__ TileXR::CommArgs* args, int32_t peer) +{ + auto udmaInfo = TileXR::GetUDMAInfo(args); + const uint32_t qpCount = udmaInfo->qpNum == 0U ? 1U : udmaInfo->qpNum; + uint32_t selected = 0U; + uint32_t selectedWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); + for (uint32_t qpIdx = 1U; qpIdx < qpCount; ++qpIdx) { + const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qpIdx); + if (weight > selectedWeight) { + selected = qpIdx; + selectedWeight = weight; + } + } + return selected; +} + +__aicore__ inline void AllToAllGroupCopyMte( + __gm__ uint8_t* dst, const __gm__ uint8_t* src, uint32_t bytes, + AscendC::LocalTensor relayLocal) +{ + for (uint32_t offset = 0U; offset < bytes; offset += TILEXR_ALLTOALL_GROUP_RELAY_BYTES) { + const uint32_t tileBytes = bytes - offset < TILEXR_ALLTOALL_GROUP_RELAY_BYTES ? + bytes - offset : TILEXR_ALLTOALL_GROUP_RELAY_BYTES; + AscendC::GlobalTensor srcGlobal; + srcGlobal.SetGlobalBuffer(const_cast<__gm__ uint8_t*>(src) + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn {1U, tileBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, srcGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(dst + offset); + AscendC::DataCopyExtParams copyOut {1U, tileBytes, 0U, 0U, 0U}; + AscendC::DataCopyPad(dstGlobal, relayLocal, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); +} + +__aicore__ inline uint64_t AllToAllGroupLoadTokenMte( + __gm__ uint64_t* signal, AscendC::LocalTensor relayLocal) +{ + AscendC::GlobalTensor signalGlobal; + signalGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(signal)); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn { + 1U, static_cast(sizeof(uint64_t)), 0U, 0U, 0U}; + AscendC::DataCopyPad(relayLocal, signalGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + return relayLocal.ReinterpretCast().GetValue(0); +} + +__aicore__ inline bool AllToAllGroupWaitTokenMte( + __gm__ uint64_t* signal, uint64_t expectedToken, uint64_t timeoutCycles, + AscendC::LocalTensor relayLocal, uint64_t& observed) +{ + const uint64_t begin = static_cast(AscendC::GetSystemCycle()); + observed = AllToAllGroupLoadTokenMte(signal, relayLocal); + while (observed < expectedToken) { + if (static_cast(AscendC::GetSystemCycle()) - begin >= timeoutCycles) { + return false; + } + observed = AllToAllGroupLoadTokenMte(signal, relayLocal); + } + return observed >= expectedToken; +} + +__aicore__ inline void AllToAllGroupRecordError( + __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, + uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, + uint32_t quietStatus, uint64_t expectedToken, uint64_t observedToken) +{ + if (debug == nullptr) { + return; + } + auto record = reinterpret_cast<__gm__ AllToAllGroupDeviceError*>( + debug + static_cast(blockIdx) * TILEXR_ALLTOALL_GROUP_ERROR_WORDS); + if (record->valid != 0U) { + return; + } + record->stage = stage; + record->group = group; + record->pass = pass; + record->peer = peer; + record->qpIdx = qpIdx; + record->quietStatus = quietStatus; + record->expectedToken = expectedToken; + record->observedToken = observedToken; + record->valid = 1U; +} + +} // namespace + +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, + GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1) +{ + const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); + auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); + auto registeredMemory = reinterpret_cast<__gm__ uint8_t*>(registeredMemoryGM); + auto debug = reinterpret_cast<__gm__ int32_t*>(debugGM); + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + + if (blockIdx >= TILEXR_ALLTOALL_GROUP_BLOCK_DIM || + !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || + rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || + elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || groupCount == 0U) { + AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CONFIG, + 0U, 0U, -1, 0U, 0U, 0ULL, 0ULL); + return; + } + + AscendC::TPipe pipe; + AscendC::TBuf relayTBuf; + pipe.InitBuffer(relayTBuf, TILEXR_ALLTOALL_GROUP_RELAY_BYTES); + AscendC::LocalTensor relayLocal = relayTBuf.Get(); + + const uint32_t slot = invocationId & 1U; + const uint64_t payloadOffsets[2] = {payloadOffset0, payloadOffset1}; + const uint64_t signalOffsets[2] = {signalOffset0, signalOffset1}; + const uint64_t bytesPerPeer = + static_cast(elementsPerPeer) * sizeof(int32_t); + + if (blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES) { + const uint32_t lane = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; + const int32_t selfBegin = static_cast( + static_cast(elementsPerPeer) * lane / + TILEXR_ALLTOALL_GROUP_SEND_CORES); + const int32_t selfEnd = static_cast( + static_cast(elementsPerPeer) * (lane + 1U) / + TILEXR_ALLTOALL_GROUP_SEND_CORES); + if (selfEnd > selfBegin) { + auto selfSrc = reinterpret_cast<__gm__ uint8_t*>( + input + static_cast(rank) * elementsPerPeer + selfBegin); + auto selfDst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(rank) * elementsPerPeer + selfBegin); + AllToAllGroupCopyMte(selfDst, selfSrc, + static_cast(selfEnd - selfBegin) * sizeof(int32_t), relayLocal); + } + + for (uint32_t group = 0U; group < groupCount; ++group) { + const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); + if (peer < 0) { + continue; + } + for (uint32_t pass = 0U; pass < passCount; ++pass) { + const int32_t chunkElementOffset = static_cast(pass) * chunkElements; + const int32_t remaining = elementsPerPeer - chunkElementOffset; + if (remaining <= 0) { + continue; + } + const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; + const uint32_t chunkBytes = + static_cast(currentElements) * sizeof(int32_t); + const uint64_t expectedToken = + AllToAllGroupDeviceToken(invocationId, group, pass); + auto signal = reinterpret_cast<__gm__ uint64_t*>( + registeredMemory + signalOffsets[slot] + + static_cast(peer) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE); + uint64_t observed = 0ULL; + if (!AllToAllGroupWaitTokenMte(signal, expectedToken, + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { + AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_WAIT, + group, pass, peer, 0U, 0U, expectedToken, observed); + return; + } + auto relaySrc = registeredMemory + payloadOffsets[slot] + + static_cast(peer) * bytesPerPeer + + static_cast(chunkElementOffset) * sizeof(int32_t); + auto relayDst = reinterpret_cast<__gm__ uint8_t*>( + output + static_cast(peer) * elementsPerPeer + chunkElementOffset); + AllToAllGroupCopyMte(relayDst, relaySrc, chunkBytes, relayLocal); + } + } + return; + } + + const uint32_t lane = blockIdx; + for (uint32_t group = 0U; group < groupCount; ++group) { + const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); + if (peer < 0) { + continue; + } + const uint32_t qpIdx = AllToAllGroupSelectMaxWeightQp(args, peer); + for (uint32_t pass = 0U; pass < passCount; ++pass) { + const int32_t chunkElementOffset = static_cast(pass) * chunkElements; + const int32_t remaining = elementsPerPeer - chunkElementOffset; + if (remaining <= 0) { + continue; + } + const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; + const uint32_t chunkBytes = + static_cast(currentElements) * sizeof(int32_t); + const uint64_t chunkByteOffset = + static_cast(chunkElementOffset) * sizeof(int32_t); + const uint64_t expectedToken = + AllToAllGroupDeviceToken(invocationId, group, pass); + auto localSrc = input + static_cast(peer) * elementsPerPeer + + chunkElementOffset; + const uint64_t remotePayloadOffset = payloadOffsets[slot] + + static_cast(rank) * bytesPerPeer + chunkByteOffset; + const uint64_t remoteSignalOffset = signalOffsets[slot] + + static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE; + TileXR::UDMAPutSignalNbiOnQp( + args, peer, qpIdx, localSrc, remotePayloadOffset, chunkBytes, + remoteSignalOffset, expectedToken); + const uint32_t quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + if (quietStatus != 0U) { + AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_QUIET, + group, pass, peer, qpIdx, quietStatus, expectedToken, 0ULL); + return; + } + } + } +} + +void launch_tilexr_udma_all_to_all_group( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR registeredMemory, GM_ADDR debug, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1) +{ + tilexr_udma_all_to_all_group_kernel<<>>( + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1); +} From 1f50eef12628c1cdb7f4a049428f3245e6be137b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:25:04 +0800 Subject: [PATCH 066/163] fix(udma): build grouped kernel as separate library --- tests/udma/CMakeLists.txt | 47 +++++++++++++++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 035e3e83..8c43add1 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -212,6 +212,8 @@ if(BUILD_TILEXR_UDMA_DEMO) endif() set(TILEXR_UDMA_DEMO_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_udma_demo_kernel.so") + set(TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO + "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_udma_alltoall_group_kernel.so") set(TILEXR_UDMA_DEMO_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/compiler/tikcpp -I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw @@ -241,7 +243,6 @@ if(BUILD_TILEXR_UDMA_DEMO) -DTILEXR_UDMA_FORCE_ENABLE=1 ${TILEXR_UDMA_DEMO_KERNEL_INCLUDES} "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" -L${ASCEND_DRIVER_PATH}/lib64/driver -L${ASCEND_HOME_PATH}/${ARCH}-linux/lib64 -L${ASCEND_HOME_PATH}/${ARCH}-linux/devlib @@ -258,8 +259,6 @@ if(BUILD_TILEXR_UDMA_DEMO) -o "${TILEXR_UDMA_DEMO_KERNEL_SO}" DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_layout.h" "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_fullmesh_trace.h" "${TILEXR_ROOT}/src/include/tilexr_udma.h" "${TILEXR_ROOT}/src/include/tilexr_data_as_flag.h" @@ -268,13 +267,51 @@ if(BUILD_TILEXR_UDMA_DEMO) ) add_custom_target(tilexr_udma_demo_kernel ALL DEPENDS "${TILEXR_UDMA_DEMO_KERNEL_SO}") + add_custom_command( + OUTPUT "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" + COMMAND ${BISHENG_EXECUTABLE} + ${TILEXR_UDMA_KERNEL_COMPILE_OPTIONS} + -std=gnu++17 + -fPIC + -shared + ${TILEXR_UDMA_KERNEL_LINK_OPTIONS} + -DCATLASS_ARCH=${TILEXR_UDMA_CATLASS_ARCH} + -DTILEXR_UDMA_FORCE_ENABLE=1 + ${TILEXR_UDMA_DEMO_KERNEL_INCLUDES} + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" + -L${ASCEND_DRIVER_PATH}/lib64/driver + -L${ASCEND_HOME_PATH}/${ARCH}-linux/lib64 + -L${ASCEND_HOME_PATH}/${ARCH}-linux/devlib + -lruntime + -lascendcl + -lstdc++ + -lm + -ltiling_api + -lplatform + -lc_sec + -ldl + -lnnopbase + -lpthread + -o "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" + DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_layout.h" + "${TILEXR_ROOT}/src/include/tilexr_udma.h" + VERBATIM + COMMENT "Building TileXR grouped AllToAll kernel with bisheng" + ) + add_custom_target(tilexr_udma_alltoall_group_kernel ALL + DEPENDS "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}") + add_executable(tilexr_udma_demo demo/tilexr_udma_demo.cpp ) - add_dependencies(tilexr_udma_demo tilexr_udma_demo_kernel) + add_dependencies(tilexr_udma_demo + tilexr_udma_demo_kernel tilexr_udma_alltoall_group_kernel) target_link_directories(tilexr_udma_demo PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) target_link_libraries(tilexr_udma_demo "${TILEXR_UDMA_DEMO_KERNEL_SO}" + "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" ${TILEXR_LIB} ascendcl runtime @@ -282,6 +319,8 @@ if(BUILD_TILEXR_UDMA_DEMO) ) list(APPEND INSTALL_TARGETS tilexr_udma_demo) install(FILES "${TILEXR_UDMA_DEMO_KERNEL_SO}" DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) + install(FILES "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" + DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) message(STATUS "TileXR UDMA demo enabled with ${BISHENG_EXECUTABLE}, SOC=${TILEXR_UDMA_DEMO_SOC_TYPE}") endif() endif() From 3064705bf2639e253869ade550769479083b7355 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:27:49 +0800 Subject: [PATCH 067/163] test(udma): require grouped alltoall host mode --- ...test_tilexr_udma_alltoall_group_layout.cpp | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 5898c734..29529285 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -163,6 +163,25 @@ void TestKernelStructure() CHECK_NOT_CONTAINS(kernel, "SyncAll"); } +void TestHostStructure() +{ + const std::string demo = ReadFile( + std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); + CHECK_CONTAINS(demo, "#include \"tilexr_udma_alltoall_group_layout.h\""); + CHECK_CONTAINS(demo, "testType == 8"); + CHECK_CONTAINS(demo, "RunGroupedAllToAll"); + CHECK_CONTAINS(demo, "PlanAllToAllGroup"); + CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_group"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); + CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); + CHECK_CONTAINS(demo, "grouped alltoall warmup="); + const size_t begin = demo.find("bool RunGroupedAllToAll("); + const size_t end = demo.find("void Cleanup(", begin); + const std::string grouped = begin == std::string::npos ? std::string() : + demo.substr(begin, end == std::string::npos ? std::string::npos : end - begin); + CHECK_NOT_CONTAINS(grouped, "DemoBarrierAll"); +} + } // namespace int main() @@ -171,6 +190,7 @@ int main() TestPlan(); TestTokens(); TestKernelStructure(); + TestHostStructure(); if (g_failures != 0) { std::cerr << "TileXR grouped all-to-all layout checks failed: " << g_failures << std::endl; return 1; From f53b1fbfcf2c2bdf3777159cf8ed5ef907a772df Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:33:54 +0800 Subject: [PATCH 068/163] feat(udma): add grouped alltoall host mode --- tests/udma/demo/tilexr_udma_demo.cpp | 184 +++++++++++++++++++++++++++ 1 file changed, 184 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index e93b42c5..bc6df76e 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -24,6 +24,7 @@ #include "tilexr_data_as_flag.h" #include "tilexr_types.h" #include "tilexr_udma_allreduce_layout.h" +#include "tilexr_udma_alltoall_group_layout.h" #include "tilexr_udma_alltoall_layout.h" #include "tilexr_udma_fullmesh_trace.h" @@ -48,6 +49,13 @@ extern void launch_tilexr_udma_all_to_all_fused( GM_ADDR udmaMem, GM_ADDR signal, GM_ADDR debug, int32_t elementsPerPeer, uint64_t udmaMemByteOffset, uint64_t signalByteOffsetBase, int32_t chunkElements, uint32_t passCount, uint32_t loopCount); +extern void launch_tilexr_udma_all_to_all_group( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, + GM_ADDR registeredMemory, GM_ADDR debug, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -600,6 +608,170 @@ bool CopyChunkDeviceToHost( return true; } +bool RunGroupedAllToAll( + int rank, int rankSize, int32_t elementsPerPeer, + TileXRCommPtr comm, aclrtStream stream, GM_ADDR commArgsDev) +{ + constexpr uint32_t kErrorWordsPerCore = 12U; + constexpr uint32_t kErrorCoreCount = TileXR::Demo::kAllToAllGroupBlockDim; + const int32_t requestedChunkElements = std::max( + 1, GetEnvInt("TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS", elementsPerPeer)); + const auto plan = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, requestedChunkElements); + if (!plan.valid) { + std::cerr << "[rank " << rank << "] ERROR: invalid grouped alltoall plan" + << " rankSize=" << rankSize + << " elementsPerPeer=" << elementsPerPeer + << " chunkElements=" << requestedChunkElements << std::endl; + return false; + } + + const size_t elementCount = static_cast(rankSize) * elementsPerPeer; + const size_t dataBytes = elementCount * sizeof(int32_t); + std::vector hostInput(elementCount, 0); + std::vector hostOutput(elementCount, -1); + TileXR::Demo::FillAllToAllInput(hostInput, rank, rankSize, elementsPerPeer); + + int32_t* input = nullptr; + int32_t* output = nullptr; + void* registeredMemory = nullptr; + TileXRUDMAMemHandle handle = 0; + bool registered = false; + auto release = [&]() { + if (registered) { + CheckTileXR(rank, "TileXRUDMAUnregister grouped alltoall", + TileXRUDMAUnregister(comm, handle)); + registered = false; + } + if (registeredMemory != nullptr) { + aclrtFree(registeredMemory); + registeredMemory = nullptr; + } + if (output != nullptr) { + aclrtFree(output); + output = nullptr; + } + if (input != nullptr) { + aclrtFree(input); + input = nullptr; + } + }; + + if (!CheckAcl(rank, "aclrtMalloc grouped input", + aclrtMalloc(reinterpret_cast(&input), dataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || + !CheckAcl(rank, "aclrtMalloc grouped output", + aclrtMalloc(reinterpret_cast(&output), dataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || + !CheckAcl(rank, "aclrtMalloc grouped registered memory", + aclrtMalloc(®isteredMemory, plan.registeredBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || + !CopyHostToDevice(rank, input, dataBytes, hostInput.data(), dataBytes, "grouped input") || + !CopyHostToDevice(rank, output, dataBytes, hostOutput.data(), dataBytes, "grouped output init") || + !CheckAcl(rank, "aclrtMemset grouped registered memory", + aclrtMemset(registeredMemory, plan.registeredBytes, 0, plan.registeredBytes))) { + release(); + return false; + } + + const int registerRet = TileXRUDMARegister( + comm, static_cast(registeredMemory), plan.registeredBytes, &handle); + if (!CheckTileXR(rank, "TileXRUDMARegister grouped alltoall", registerRet)) { + release(); + return false; + } + registered = true; + + auto debug = reinterpret_cast( + static_cast(registeredMemory) + plan.controlOffset); + const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); + const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); + PrintStatus(rank, "grouped alltoall registeredBytes=" + std::to_string(plan.registeredBytes) + + " payloadPlaneBytes=" + std::to_string(plan.payloadPlaneBytes) + + " payloadOffset0=" + std::to_string(plan.payloadOffset[0]) + + " payloadOffset1=" + std::to_string(plan.payloadOffset[1]) + + " signalPlaneBytes=" + std::to_string(plan.signalPlaneBytes) + + " signalOffset0=" + std::to_string(plan.signalOffset[0]) + + " signalOffset1=" + std::to_string(plan.signalOffset[1]) + + " controlOffset=" + std::to_string(plan.controlOffset) + + " groups=" + std::to_string(plan.groupCount) + + " passes=" + std::to_string(plan.passCount)); + PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + + " repeat=" + std::to_string(repeat)); + + uint32_t invocationId = 0U; + for (int iter = 0; iter < warmup; ++iter, ++invocationId) { + launch_tilexr_udma_all_to_all_group( + TileXR::Demo::kAllToAllGroupBlockDim, stream, commArgsDev, + reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(registeredMemory), reinterpret_cast(debug), + invocationId, elementsPerPeer, plan.chunkElements, + plan.passCount, plan.groupCount, + plan.payloadOffset[0], plan.payloadOffset[1], + plan.signalOffset[0], plan.signalOffset[1]); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", aclrtSynchronizeStream(stream))) { + release(); + return false; + } + + const auto begin = std::chrono::steady_clock::now(); + for (int iter = 0; iter < repeat; ++iter, ++invocationId) { + launch_tilexr_udma_all_to_all_group( + TileXR::Demo::kAllToAllGroupBlockDim, stream, commArgsDev, + reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(registeredMemory), reinterpret_cast(debug), + invocationId, elementsPerPeer, plan.chunkElements, + plan.passCount, plan.groupCount, + plan.payloadOffset[0], plan.payloadOffset[1], + plan.signalOffset[0], plan.signalOffset[1]); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", aclrtSynchronizeStream(stream))) { + release(); + return false; + } + const auto end = std::chrono::steady_clock::now(); + + std::vector hostDebug(kErrorWordsPerCore * kErrorCoreCount, 0); + const size_t debugBytes = hostDebug.size() * sizeof(int32_t); + bool copyOk = CopyDeviceToHost( + rank, hostOutput.data(), dataBytes, output, dataBytes, "grouped alltoall output") && + CopyDeviceToHost(rank, hostDebug.data(), debugBytes, debug, debugBytes, + "grouped alltoall debug"); + bool debugOk = true; + for (uint32_t core = 0; core < kErrorCoreCount; ++core) { + const size_t base = static_cast(core) * kErrorWordsPerCore; + if (hostDebug[base] == 0) { + continue; + } + debugOk = false; + const uint64_t expected = static_cast(hostDebug[base + 8]) | + (static_cast(static_cast(hostDebug[base + 9])) << 32U); + const uint64_t observed = static_cast(hostDebug[base + 10]) | + (static_cast(static_cast(hostDebug[base + 11])) << 32U); + std::cerr << "[rank " << rank << "] ERROR: grouped core=" << core + << " stage=" << hostDebug[base + 1] + << " group=" << hostDebug[base + 2] + << " pass=" << hostDebug[base + 3] + << " peer=" << hostDebug[base + 4] + << " qp=" << hostDebug[base + 5] + << " quiet=" << hostDebug[base + 6] + << " expected=" << expected + << " observed=" << observed << std::endl; + } + + const double totalUs = std::chrono::duration(end - begin).count(); + const double perIterUs = totalUs / static_cast(repeat); + const double bandwidthGbs = static_cast(dataBytes) / (perIterUs * 1.0e3); + std::cout << "[rank " << rank << "] grouped alltoall " << repeat + << " iters total=" << totalUs / 1000.0 + << " ms perIter=" << perIterUs + << " us payload=" << dataBytes + << " bytes bw=" << bandwidthGbs << " GB/s" << std::endl; + + const bool valid = copyOk && debugOk && + ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerPeer); + release(); + return valid; +} + void Cleanup( TileXRCommPtr comm, aclrtStream stream, void* registeredMemory, int32_t* debug, int rank, int deviceId) { @@ -687,6 +859,18 @@ int main(int argc, char** argv) return 1; } + if (testType == 8) { + const bool ok = RunGroupedAllToAll( + rank, rankSize, elementsPerRank, comm, stream, commArgsDev); + Cleanup(comm, stream, nullptr, nullptr, rank, deviceId); + if (!ok) { + std::cerr << "[rank " << rank << "] TileXR grouped alltoall demo failed" << std::endl; + return 1; + } + std::cout << "[rank " << rank << "] TileXR grouped alltoall demo success" << std::endl; + return 0; + } + bool isAllToAll = testType == 2 || testType == 4 || testType == 5 || testType == 6 || testType == 7; bool isAllReduce = testType == 3; bool strictAllToAllUdma = From a8da596478943ff627538358aab99f39d77a7e68 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:37:27 +0800 Subject: [PATCH 069/163] test(udma): require grouped alltoall trace conversion --- ...exr_udma_alltoall_group_trace_to_chrome.py | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py new file mode 100644 index 00000000..0ea40014 --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 +import importlib.util +import json +import struct +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +MODULE_PATH = ( + Path(__file__).resolve().parents[1] + / "demo" + / "tilexr_udma_alltoall_group_trace_to_chrome.py" +) +SPEC = importlib.util.spec_from_file_location( + "tilexr_udma_alltoall_group_trace_to_chrome", MODULE_PATH) +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +class GroupTraceConverterTest(unittest.TestCase): + def make_trace( + self, path, *, rank=0, magic=None, iteration_count=1, + group_count=1, pass_count=1, core_count=32, + ): + data = bytearray(MODULE.TRACE_BYTES) + struct.pack_into( + MODULE.HEADER_FORMAT, + data, + 0, + MODULE.TRACE_MAGIC if magic is None else magic, + MODULE.TRACE_VERSION, + rank, + iteration_count, + group_count, + pass_count, + core_count, + MODULE.PHASE_COUNT, + 1000, + MODULE.TRACE_BYTES, + MODULE.HEADER_BYTES, + MODULE.TASK_BASE_OFFSET, + ) + for core in (0, 16): + struct.pack_into( + " Date: Mon, 20 Jul 2026 18:41:36 +0800 Subject: [PATCH 070/163] feat(udma): trace grouped alltoall pipeline --- tests/udma/CMakeLists.txt | 1 + .../tilexr_udma_alltoall_group_kernel.cpp | 95 ++++++- .../demo/tilexr_udma_alltoall_group_trace.h | 140 ++++++++++ ...exr_udma_alltoall_group_trace_to_chrome.py | 241 ++++++++++++++++++ tests/udma/demo/tilexr_udma_demo.cpp | 86 ++++++- 5 files changed, 555 insertions(+), 8 deletions(-) create mode 100644 tests/udma/demo/tilexr_udma_alltoall_group_trace.h create mode 100644 tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 8c43add1..510693af 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -296,6 +296,7 @@ if(BUILD_TILEXR_UDMA_DEMO) DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_layout.h" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_trace.h" "${TILEXR_ROOT}/src/include/tilexr_udma.h" VERBATIM COMMENT "Building TileXR grouped AllToAll kernel with bisheng" diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 8b99a04f..2225fde5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -5,6 +5,7 @@ #include "kernel_operator.h" #include "tilexr_udma.h" +#include "tilexr_udma_alltoall_group_trace.h" namespace { @@ -158,6 +159,45 @@ __aicore__ inline void AllToAllGroupRecordError( record->valid = 1U; } +__aicore__ inline uint64_t AllToAllGroupTraceCycle(__gm__ uint8_t* trace) +{ + return trace == nullptr ? 0ULL : static_cast(AscendC::GetSystemCycle()); +} + +__aicore__ inline void AllToAllGroupTraceRecordKernel( + __gm__ uint8_t* trace, uint32_t iteration, uint32_t core, + uint64_t beginCycle, uint64_t endCycle) +{ + if (trace == nullptr || iteration >= TileXR::Demo::kAllToAllGroupTraceMaxIterations || + core >= TileXR::Demo::kAllToAllGroupTraceCoreCount) { + return; + } + auto span = reinterpret_cast<__gm__ TileXR::Demo::AllToAllGroupTraceSpan*>( + trace + TileXR::Demo::AllToAllGroupTraceKernelSpanOffset(iteration, core)); + span->beginCycle = beginCycle; + span->endCycle = endCycle; +} + +__aicore__ inline void AllToAllGroupTraceRecordTask( + __gm__ uint8_t* trace, uint32_t iteration, uint32_t core, + uint32_t group, uint32_t pass, uint32_t phase, + uint32_t groupCount, uint32_t passCount, int32_t peer, uint32_t qpIdx, + uint64_t beginCycle, uint64_t endCycle) +{ + if (trace == nullptr || iteration >= TileXR::Demo::kAllToAllGroupTraceMaxIterations || + core >= TileXR::Demo::kAllToAllGroupTraceCoreCount || group >= groupCount || + pass >= passCount || phase >= TileXR::Demo::kAllToAllGroupTracePhaseCount) { + return; + } + auto span = reinterpret_cast<__gm__ TileXR::Demo::AllToAllGroupTraceTaskSpan*>( + trace + TileXR::Demo::AllToAllGroupTraceTaskSpanOffset( + iteration, core, group, pass, phase, groupCount, passCount)); + span->peer = peer; + span->qpIdx = qpIdx; + span->beginCycle = beginCycle; + span->endCycle = endCycle; +} + } // namespace extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( @@ -166,9 +206,13 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( int32_t elementsPerPeer, int32_t chunkElements, uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, - uint64_t signalOffset0, uint64_t signalOffset1) + uint64_t signalOffset0, uint64_t signalOffset1, + GM_ADDR groupTraceGM, uint32_t traceIteration) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); + auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? + reinterpret_cast<__gm__ uint8_t*>(groupTraceGM) : nullptr; + const uint64_t kernelBegin = AllToAllGroupTraceCycle(groupTrace); auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); @@ -183,6 +227,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || groupCount == 0U) { AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CONFIG, 0U, 0U, -1, 0U, 0U, 0ULL, 0ULL); + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } @@ -206,12 +252,18 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(elementsPerPeer) * (lane + 1U) / TILEXR_ALLTOALL_GROUP_SEND_CORES); if (selfEnd > selfBegin) { + const uint64_t selfCopyBegin = AllToAllGroupTraceCycle(groupTrace); auto selfSrc = reinterpret_cast<__gm__ uint8_t*>( input + static_cast(rank) * elementsPerPeer + selfBegin); auto selfDst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(rank) * elementsPerPeer + selfBegin); AllToAllGroupCopyMte(selfDst, selfSrc, static_cast(selfEnd - selfBegin) * sizeof(int32_t), relayLocal); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, 0U, 0U, + TileXR::Demo::kAllToAllGroupTraceSelfCopy, groupCount, passCount, + rank, TileXR::Demo::kAllToAllGroupTraceNoQp, + selfCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } for (uint32_t group = 0U; group < groupCount; ++group) { @@ -234,20 +286,41 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( registeredMemory + signalOffsets[slot] + static_cast(peer) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE); uint64_t observed = 0ULL; + const uint64_t waitBegin = AllToAllGroupTraceCycle(groupTrace); if (!AllToAllGroupWaitTokenMte(signal, expectedToken, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, + peer, TileXR::Demo::kAllToAllGroupTraceNoQp, + waitBegin, AllToAllGroupTraceCycle(groupTrace)); AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_WAIT, group, pass, peer, 0U, 0U, expectedToken, observed); + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, + peer, TileXR::Demo::kAllToAllGroupTraceNoQp, + waitBegin, AllToAllGroupTraceCycle(groupTrace)); auto relaySrc = registeredMemory + payloadOffsets[slot] + static_cast(peer) * bytesPerPeer + static_cast(chunkElementOffset) * sizeof(int32_t); auto relayDst = reinterpret_cast<__gm__ uint8_t*>( output + static_cast(peer) * elementsPerPeer + chunkElementOffset); + const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); AllToAllGroupCopyMte(relayDst, relaySrc, chunkBytes, relayLocal); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceReceiveCopy, groupCount, passCount, + peer, TileXR::Demo::kAllToAllGroupTraceNoQp, + receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } } + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } @@ -277,17 +350,31 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(rank) * bytesPerPeer + chunkByteOffset; const uint64_t remoteSignalOffset = signalOffsets[slot] + static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE; + const uint64_t putBegin = AllToAllGroupTraceCycle(groupTrace); TileXR::UDMAPutSignalNbiOnQp( args, peer, qpIdx, localSrc, remotePayloadOffset, chunkBytes, remoteSignalOffset, expectedToken); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, + peer, qpIdx, putBegin, AllToAllGroupTraceCycle(groupTrace)); + const uint64_t quietBegin = AllToAllGroupTraceCycle(groupTrace); const uint32_t quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, + peer, qpIdx, quietBegin, AllToAllGroupTraceCycle(groupTrace)); if (quietStatus != 0U) { AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_QUIET, group, pass, peer, qpIdx, quietStatus, expectedToken, 0ULL); + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } } } + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); } void launch_tilexr_udma_all_to_all_group( @@ -296,10 +383,12 @@ void launch_tilexr_udma_all_to_all_group( int32_t elementsPerPeer, int32_t chunkElements, uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, - uint64_t signalOffset0, uint64_t signalOffset1) + uint64_t signalOffset0, uint64_t signalOffset1, + GM_ADDR groupTrace, uint32_t traceIteration) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, - payloadOffset0, payloadOffset1, signalOffset0, signalOffset1); + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + groupTrace, traceIteration); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h new file mode 100644 index 00000000..138ffe51 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -0,0 +1,140 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_ALLTOALL_GROUP_TRACE_H +#define TILEXR_UDMA_ALLTOALL_GROUP_TRACE_H + +#include +#include +#include + +namespace TileXR { +namespace Demo { + +constexpr uint32_t kAllToAllGroupTraceMagic = 0x47545243U; // "GTRC" +constexpr uint32_t kAllToAllGroupTraceVersion = 1U; +constexpr size_t kAllToAllGroupTraceBytes = 8ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; +constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; +constexpr uint32_t kAllToAllGroupTraceCoreCount = 32U; +constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; +constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; +constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; + +enum AllToAllGroupTracePhase : uint32_t { + kAllToAllGroupTraceSelfCopy = 0U, + kAllToAllGroupTraceSendPutSignal = 1U, + kAllToAllGroupTraceSendQuiet = 2U, + kAllToAllGroupTraceReceiveWait = 3U, + kAllToAllGroupTraceReceiveCopy = 4U, +}; + +struct AllToAllGroupTraceSpan { + uint64_t beginCycle; + uint64_t endCycle; +}; + +struct AllToAllGroupTraceTaskSpan { + uint64_t beginCycle; + uint64_t endCycle; + int32_t peer; + uint32_t qpIdx; +}; + +struct AllToAllGroupTraceHeader { + uint32_t magic; + uint32_t version; + uint32_t rank; + uint32_t iterationCount; + uint32_t groupCount; + uint32_t passCount; + uint32_t coreCount; + uint32_t phaseCount; + uint64_t cyclesPerUs; + uint64_t traceBytes; + uint64_t kernelSpanOffset; + uint64_t taskSpanOffset; +}; + +constexpr size_t AllToAllGroupTraceKernelSpanOffset(uint32_t iteration, uint32_t core) +{ + return kAllToAllGroupTraceHeaderBytes + + (static_cast(iteration) * kAllToAllGroupTraceCoreCount + core) * + sizeof(AllToAllGroupTraceSpan); +} + +constexpr size_t AllToAllGroupTraceTaskSpanBaseOffset() +{ + return kAllToAllGroupTraceHeaderBytes + + static_cast(kAllToAllGroupTraceMaxIterations) * + kAllToAllGroupTraceCoreCount * sizeof(AllToAllGroupTraceSpan); +} + +inline bool AllToAllGroupTraceCheckedMultiply(size_t lhs, size_t rhs, size_t& result) +{ + if (lhs != 0U && rhs > std::numeric_limits::max() / lhs) { + result = std::numeric_limits::max(); + return false; + } + result = lhs * rhs; + return true; +} + +inline size_t AllToAllGroupTraceLayoutBytes( + uint32_t iterationCount, uint32_t groupCount, uint32_t passCount) +{ + size_t count = iterationCount; + size_t next = 0U; + const size_t factors[] = { + kAllToAllGroupTraceCoreCount, groupCount, passCount, + kAllToAllGroupTracePhaseCount, sizeof(AllToAllGroupTraceTaskSpan)}; + for (size_t factor : factors) { + if (!AllToAllGroupTraceCheckedMultiply(count, factor, next)) { + return std::numeric_limits::max(); + } + count = next; + } + if (count > std::numeric_limits::max() - + AllToAllGroupTraceTaskSpanBaseOffset()) { + return std::numeric_limits::max(); + } + return AllToAllGroupTraceTaskSpanBaseOffset() + count; +} + +inline bool AllToAllGroupTraceLayoutFits( + uint32_t iterationCount, uint32_t groupCount, uint32_t passCount) +{ + return iterationCount > 0U && + iterationCount <= kAllToAllGroupTraceMaxIterations && + groupCount > 0U && passCount > 0U && + AllToAllGroupTraceLayoutBytes(iterationCount, groupCount, passCount) <= + kAllToAllGroupTraceBytes; +} + +inline size_t AllToAllGroupTraceTaskSpanOffset( + uint32_t iteration, uint32_t core, uint32_t group, uint32_t pass, + uint32_t phase, uint32_t groupCount, uint32_t passCount) +{ + const size_t index = + (((((static_cast(iteration) * kAllToAllGroupTraceCoreCount + core) * + groupCount + group) * passCount + pass) * + kAllToAllGroupTracePhaseCount) + phase); + return AllToAllGroupTraceTaskSpanBaseOffset() + + index * sizeof(AllToAllGroupTraceTaskSpan); +} + +static_assert(sizeof(AllToAllGroupTraceSpan) == 16U, + "group trace kernel span must contain two uint64 timestamps"); +static_assert(sizeof(AllToAllGroupTraceTaskSpan) == 24U, + "group trace task span layout changed"); +static_assert(sizeof(AllToAllGroupTraceHeader) <= kAllToAllGroupTraceHeaderBytes, + "group trace header must fit its region"); +static_assert(AllToAllGroupTraceTaskSpanBaseOffset() < kAllToAllGroupTraceBytes, + "group trace kernel spans must fit in 8 MiB"); + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_ALLTOALL_GROUP_TRACE_H diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py new file mode 100644 index 00000000..ffe60900 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +import argparse +import json +import struct +from pathlib import Path + + +TRACE_MAGIC = 0x47545243 +TRACE_VERSION = 1 +TRACE_BYTES = 8 * 1024 * 1024 +HEADER_BYTES = 4096 +MAX_ITERATIONS = 50 +MAX_CORES = 32 +PHASE_COUNT = 5 +SPAN_BYTES = 16 +TASK_FORMAT = " len(data): + raise ValueError(f"span offset out of range for {label}: {offset}") + begin, end = struct.unpack_from(" len(data): + raise ValueError(f"task offset out of range for {label}: {offset}") + begin, end, peer, qp = struct.unpack_from(TASK_FORMAT, data, offset) + if begin == 0 and end == 0: + return None + if begin == 0 or end == 0: + raise ValueError(f"incomplete {label}: begin={begin} end={end}") + if end < begin: + raise ValueError(f"invalid {label}: begin={begin} end={end}") + return begin, end, peer, qp + + +def read_rank_trace(path): + path = Path(path) + data = path.read_bytes() + if len(data) != TRACE_BYTES: + raise ValueError(f"invalid trace size {len(data)} in {path}, expected {TRACE_BYTES}") + fields = struct.unpack_from(HEADER_FORMAT, data, 0) + header = { + "magic": fields[0], + "version": fields[1], + "rank": fields[2], + "iteration_count": fields[3], + "group_count": fields[4], + "pass_count": fields[5], + "core_count": fields[6], + "phase_count": fields[7], + "cycles_per_us": fields[8], + "trace_bytes": fields[9], + "kernel_span_offset": fields[10], + "task_span_offset": fields[11], + } + if header["magic"] != TRACE_MAGIC: + raise ValueError(f"invalid trace magic in {path}") + if header["version"] != TRACE_VERSION: + raise ValueError(f"unsupported trace version {header['version']} in {path}") + if header["trace_bytes"] != TRACE_BYTES: + raise ValueError(f"trace byte dimension mismatch in {path}") + if not 0 < header["iteration_count"] <= MAX_ITERATIONS: + raise ValueError(f"iteration dimension mismatch in {path}") + if header["group_count"] <= 0 or header["pass_count"] <= 0: + raise ValueError(f"group/pass dimension mismatch in {path}") + if header["core_count"] != MAX_CORES or header["phase_count"] != PHASE_COUNT: + raise ValueError(f"core/phase dimension mismatch in {path}") + if (header["kernel_span_offset"] != HEADER_BYTES or + header["task_span_offset"] != TASK_BASE_OFFSET): + raise ValueError(f"trace offset mismatch in {path}") + if header["cycles_per_us"] == 0: + raise ValueError(f"invalid cycle frequency in {path}") + required = layout_bytes( + header["iteration_count"], header["group_count"], header["pass_count"]) + if required > TRACE_BYTES: + raise ValueError(f"trace capacity exceeded in {path}: required={required}") + return {"path": str(path), "header": header, "data": data} + + +def _metadata(name, pid, tid, value): + return {"name": name, "ph": "M", "pid": pid, "tid": tid, "args": {"name": value}} + + +def _event(name, category, pid, tid, begin, end, base, cycles_per_us, args, offset_us): + return { + "name": name, + "cat": category, + "ph": "X", + "pid": pid, + "tid": tid, + "ts": offset_us + (begin - base) / cycles_per_us, + "dur": (end - begin) / cycles_per_us, + "args": {"beginCycle": begin, "endCycle": end, **args}, + } + + +def build_chrome_trace(rank_traces): + events = [] + sources = [] + bases = {} + iteration_durations = {} + for rank_trace in rank_traces: + header = rank_trace["header"] + rank = header["rank"] + for iteration in range(header["iteration_count"]): + spans = [ + _read_span( + rank_trace["data"], kernel_span_offset(iteration, core), + f"kernel rank={rank} iter={iteration} core={core}") + for core in range(MAX_CORES) + ] + spans = [span for span in spans if span is not None] + if not spans: + raise ValueError(f"rank {rank} iteration {iteration} contains no kernel spans") + base = min(begin for begin, _ in spans) + bases[(rank, iteration)] = base + duration = (max(end for _, end in spans) - base) / header["cycles_per_us"] + iteration_durations[iteration] = max( + iteration_durations.get(iteration, 0.0), duration) + + iteration_offsets = {} + cursor = 0.0 + for iteration in sorted(iteration_durations): + iteration_offsets[iteration] = cursor + cursor += iteration_durations[iteration] + ITERATION_GAP_US + + for rank_trace in sorted(rank_traces, key=lambda item: item["header"]["rank"]): + header = rank_trace["header"] + data = rank_trace["data"] + rank = header["rank"] + sources.append(rank_trace["path"]) + events.append(_metadata("process_name", rank, 0, f"rank {rank}")) + for core in range(MAX_CORES): + role = "send" if core < 16 else "receive" + events.append(_metadata("thread_name", rank, core, f"core{core} {role}")) + + for iteration in range(header["iteration_count"]): + base = bases[(rank, iteration)] + offset_us = iteration_offsets[iteration] + for core in range(MAX_CORES): + role = "send" if core < 16 else "receive" + kernel = _read_span( + data, kernel_span_offset(iteration, core), + f"kernel rank={rank} iter={iteration} core={core}") + if kernel is not None: + events.append(_event( + "kernel", "kernel", rank, core, kernel[0], kernel[1], base, + header["cycles_per_us"], + {"iteration": iteration, "role": role}, offset_us)) + for group in range(header["group_count"]): + for pass_index in range(header["pass_count"]): + for phase in range(PHASE_COUNT): + label = ( + f"task rank={rank} iter={iteration} core={core} " + f"group={group} pass={pass_index} phase={phase}") + task = _read_task( + data, + task_span_offset( + iteration, core, group, pass_index, phase, + header["group_count"], header["pass_count"]), + label, + ) + if task is None: + continue + begin, end, peer, qp = task + events.append(_event( + PHASE_NAMES[phase], role, rank, core, begin, end, base, + header["cycles_per_us"], + { + "iteration": iteration, + "group": group, + "pass": pass_index, + "lane": core if core < 16 else core - 16, + "peer": peer, + "qp": None if qp == NO_QP else qp, + "role": role, + }, + offset_us, + )) + return { + "traceEvents": events, + "otherData": { + "displayTimeUnit": "ns", + "sources": sources, + "clock": "GetSystemCycle normalized per rank and iteration", + }, + } + + +def main(): + parser = argparse.ArgumentParser( + description="Convert TileXR grouped AllToAll GM traces to Chrome trace JSON") + parser.add_argument("inputs", nargs="+", type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + trace = build_chrome_trace([read_rank_trace(path) for path in args.inputs]) + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as stream: + json.dump(trace, stream, separators=(",", ":")) + print(f"wrote {args.output} events={len(trace['traceEvents'])}") + + +if __name__ == "__main__": + main() diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index bc6df76e..f245716d 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -25,6 +25,7 @@ #include "tilexr_types.h" #include "tilexr_udma_allreduce_layout.h" #include "tilexr_udma_alltoall_group_layout.h" +#include "tilexr_udma_alltoall_group_trace.h" #include "tilexr_udma_alltoall_layout.h" #include "tilexr_udma_fullmesh_trace.h" @@ -55,7 +56,8 @@ extern void launch_tilexr_udma_all_to_all_group( int32_t elementsPerPeer, int32_t chunkElements, uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, - uint64_t signalOffset0, uint64_t signalOffset1); + uint64_t signalOffset0, uint64_t signalOffset1, + GM_ADDR groupTrace, uint32_t traceIteration); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -236,6 +238,26 @@ bool WriteFullmeshTraceBinary(int rank, const std::string& directory, const std: return true; } +bool WriteGroupTraceBinary(int rank, const std::string& directory, const std::vector& data) +{ + const std::string path = directory + "/tilexr_group_trace_rank_" + + std::to_string(rank) + ".bin"; + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output.is_open()) { + std::cerr << "[rank " << rank << "] ERROR: open grouped trace output failed path=" + << path << std::endl; + return false; + } + output.write(reinterpret_cast(data.data()), static_cast(data.size())); + if (!output.good()) { + std::cerr << "[rank " << rank << "] ERROR: write grouped trace output failed path=" + << path << std::endl; + return false; + } + PrintStatus(rank, "grouped trace output=" + path + " bytes=" + std::to_string(data.size())); + return true; +} + BarrierEndpoint GetBarrierEndpoint() { std::string host = "127.0.0.1"; @@ -625,6 +647,23 @@ bool RunGroupedAllToAll( << " chunkElements=" << requestedChunkElements << std::endl; return false; } + const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); + const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); + const bool traceEnabled = GetEnvInt("TILEXR_UDMA_GROUP_TRACE", 0) != 0; + const char* traceDirEnv = std::getenv("TILEXR_UDMA_GROUP_TRACE_DIR"); + const std::string traceDir = traceDirEnv != nullptr && traceDirEnv[0] != '\0' ? + traceDirEnv : "."; + if (traceEnabled && !TileXR::Demo::AllToAllGroupTraceLayoutFits( + static_cast(repeat), plan.groupCount, plan.passCount)) { + std::cerr << "[rank " << rank << "] ERROR: grouped trace dimensions exceed capacity" + << " repeat=" << repeat + << " groupCount=" << plan.groupCount + << " passCount=" << plan.passCount + << " requiredBytes=" << TileXR::Demo::AllToAllGroupTraceLayoutBytes( + static_cast(repeat), plan.groupCount, plan.passCount) + << " capacityBytes=" << TileXR::Demo::kAllToAllGroupTraceBytes << std::endl; + return false; + } const size_t elementCount = static_cast(rankSize) * elementsPerPeer; const size_t dataBytes = elementCount * sizeof(int32_t); @@ -635,6 +674,8 @@ bool RunGroupedAllToAll( int32_t* input = nullptr; int32_t* output = nullptr; void* registeredMemory = nullptr; + void* groupTraceDevice = nullptr; + std::vector hostGroupTrace; TileXRUDMAMemHandle handle = 0; bool registered = false; auto release = [&]() { @@ -647,6 +688,10 @@ bool RunGroupedAllToAll( aclrtFree(registeredMemory); registeredMemory = nullptr; } + if (groupTraceDevice != nullptr) { + aclrtFree(groupTraceDevice); + groupTraceDevice = nullptr; + } if (output != nullptr) { aclrtFree(output); output = nullptr; @@ -671,6 +716,32 @@ bool RunGroupedAllToAll( return false; } + if (traceEnabled) { + hostGroupTrace.assign(TileXR::Demo::kAllToAllGroupTraceBytes, 0U); + TileXR::Demo::AllToAllGroupTraceHeader header {}; + header.magic = TileXR::Demo::kAllToAllGroupTraceMagic; + header.version = TileXR::Demo::kAllToAllGroupTraceVersion; + header.rank = static_cast(rank); + header.iterationCount = static_cast(repeat); + header.groupCount = plan.groupCount; + header.passCount = plan.passCount; + header.coreCount = TileXR::Demo::kAllToAllGroupTraceCoreCount; + header.phaseCount = TileXR::Demo::kAllToAllGroupTracePhaseCount; + header.cyclesPerUs = TileXR::Demo::kAllToAllGroupTraceCyclesPerUs; + header.traceBytes = TileXR::Demo::kAllToAllGroupTraceBytes; + header.kernelSpanOffset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes; + header.taskSpanOffset = TileXR::Demo::AllToAllGroupTraceTaskSpanBaseOffset(); + std::memcpy(hostGroupTrace.data(), &header, sizeof(header)); + if (!CheckAcl(rank, "aclrtMalloc grouped trace", + aclrtMalloc(&groupTraceDevice, TileXR::Demo::kAllToAllGroupTraceBytes, + ACL_MEM_MALLOC_HUGE_FIRST)) || + !CopyHostToDevice(rank, groupTraceDevice, TileXR::Demo::kAllToAllGroupTraceBytes, + hostGroupTrace.data(), hostGroupTrace.size(), "grouped trace")) { + release(); + return false; + } + } + const int registerRet = TileXRUDMARegister( comm, static_cast(registeredMemory), plan.registeredBytes, &handle); if (!CheckTileXR(rank, "TileXRUDMARegister grouped alltoall", registerRet)) { @@ -681,8 +752,6 @@ bool RunGroupedAllToAll( auto debug = reinterpret_cast( static_cast(registeredMemory) + plan.controlOffset); - const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); - const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); PrintStatus(rank, "grouped alltoall registeredBytes=" + std::to_string(plan.registeredBytes) + " payloadPlaneBytes=" + std::to_string(plan.payloadPlaneBytes) + " payloadOffset0=" + std::to_string(plan.payloadOffset[0]) + @@ -705,7 +774,7 @@ bool RunGroupedAllToAll( invocationId, elementsPerPeer, plan.chunkElements, plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], - plan.signalOffset[0], plan.signalOffset[1]); + plan.signalOffset[0], plan.signalOffset[1], nullptr, 0U); } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", aclrtSynchronizeStream(stream))) { release(); @@ -721,7 +790,8 @@ bool RunGroupedAllToAll( invocationId, elementsPerPeer, plan.chunkElements, plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], - plan.signalOffset[0], plan.signalOffset[1]); + plan.signalOffset[0], plan.signalOffset[1], + reinterpret_cast(groupTraceDevice), static_cast(iter)); } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", aclrtSynchronizeStream(stream))) { release(); @@ -735,6 +805,12 @@ bool RunGroupedAllToAll( rank, hostOutput.data(), dataBytes, output, dataBytes, "grouped alltoall output") && CopyDeviceToHost(rank, hostDebug.data(), debugBytes, debug, debugBytes, "grouped alltoall debug"); + if (traceEnabled) { + copyOk = CopyDeviceToHost( + rank, hostGroupTrace.data(), hostGroupTrace.size(), groupTraceDevice, + TileXR::Demo::kAllToAllGroupTraceBytes, "grouped trace") && + WriteGroupTraceBinary(rank, traceDir, hostGroupTrace) && copyOk; + } bool debugOk = true; for (uint32_t core = 0; core < kErrorCoreCount; ++core) { const size_t base = static_cast(core) * kErrorWordsPerCore; From 0b5ab407ac9f3301e3f21c450880734b3fd7a301 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:44:05 +0800 Subject: [PATCH 071/163] fix(udma): inline grouped trace offsets on device --- .../demo/tilexr_udma_alltoall_group_kernel.cpp | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 2225fde5..df2d192c 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -172,8 +172,11 @@ __aicore__ inline void AllToAllGroupTraceRecordKernel( core >= TileXR::Demo::kAllToAllGroupTraceCoreCount) { return; } + const uint64_t offset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes + + (static_cast(iteration) * TileXR::Demo::kAllToAllGroupTraceCoreCount + core) * + sizeof(TileXR::Demo::AllToAllGroupTraceSpan); auto span = reinterpret_cast<__gm__ TileXR::Demo::AllToAllGroupTraceSpan*>( - trace + TileXR::Demo::AllToAllGroupTraceKernelSpanOffset(iteration, core)); + trace + offset); span->beginCycle = beginCycle; span->endCycle = endCycle; } @@ -189,9 +192,18 @@ __aicore__ inline void AllToAllGroupTraceRecordTask( pass >= passCount || phase >= TileXR::Demo::kAllToAllGroupTracePhaseCount) { return; } + const uint64_t index = + (((((static_cast(iteration) * TileXR::Demo::kAllToAllGroupTraceCoreCount + core) * + groupCount + group) * passCount + pass) * + TileXR::Demo::kAllToAllGroupTracePhaseCount) + phase); + const uint64_t taskBaseOffset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes + + static_cast(TileXR::Demo::kAllToAllGroupTraceMaxIterations) * + TileXR::Demo::kAllToAllGroupTraceCoreCount * + sizeof(TileXR::Demo::AllToAllGroupTraceSpan); + const uint64_t offset = taskBaseOffset + + index * sizeof(TileXR::Demo::AllToAllGroupTraceTaskSpan); auto span = reinterpret_cast<__gm__ TileXR::Demo::AllToAllGroupTraceTaskSpan*>( - trace + TileXR::Demo::AllToAllGroupTraceTaskSpanOffset( - iteration, core, group, pass, phase, groupCount, passCount)); + trace + offset); span->peer = peer; span->qpIdx = qpIdx; span->beginCycle = beginCycle; From 68e3b5fcd741aa67e5a55a9ed541450419d34a80 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 18:55:53 +0800 Subject: [PATCH 072/163] fix(udma): isolate grouped trace cores by cache line --- .../tilexr_udma_alltoall_group_kernel.cpp | 20 ++++++---- .../demo/tilexr_udma_alltoall_group_trace.h | 38 +++++++++++++------ ...exr_udma_alltoall_group_trace_to_chrome.py | 18 +++++---- ...exr_udma_alltoall_group_trace_to_chrome.py | 8 ++++ 4 files changed, 59 insertions(+), 25 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index df2d192c..fd1952f2 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -174,7 +174,7 @@ __aicore__ inline void AllToAllGroupTraceRecordKernel( } const uint64_t offset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes + (static_cast(iteration) * TileXR::Demo::kAllToAllGroupTraceCoreCount + core) * - sizeof(TileXR::Demo::AllToAllGroupTraceSpan); + TileXR::Demo::kAllToAllGroupTraceCacheLineBytes; auto span = reinterpret_cast<__gm__ TileXR::Demo::AllToAllGroupTraceSpan*>( trace + offset); span->beginCycle = beginCycle; @@ -192,16 +192,22 @@ __aicore__ inline void AllToAllGroupTraceRecordTask( pass >= passCount || phase >= TileXR::Demo::kAllToAllGroupTracePhaseCount) { return; } - const uint64_t index = - (((((static_cast(iteration) * TileXR::Demo::kAllToAllGroupTraceCoreCount + core) * - groupCount + group) * passCount + pass) * - TileXR::Demo::kAllToAllGroupTracePhaseCount) + phase); + const uint64_t rawCoreBytes = static_cast(groupCount) * passCount * + TileXR::Demo::kAllToAllGroupTracePhaseCount * + sizeof(TileXR::Demo::AllToAllGroupTraceTaskSpan); + const uint64_t coreBytes = + (rawCoreBytes + TileXR::Demo::kAllToAllGroupTraceCacheLineBytes - 1U) & + ~(static_cast(TileXR::Demo::kAllToAllGroupTraceCacheLineBytes) - 1ULL); + const uint64_t coreIndex = static_cast(iteration) * + TileXR::Demo::kAllToAllGroupTraceCoreCount + core; + const uint64_t taskIndex = ((static_cast(group) * passCount + pass) * + TileXR::Demo::kAllToAllGroupTracePhaseCount) + phase; const uint64_t taskBaseOffset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes + static_cast(TileXR::Demo::kAllToAllGroupTraceMaxIterations) * TileXR::Demo::kAllToAllGroupTraceCoreCount * - sizeof(TileXR::Demo::AllToAllGroupTraceSpan); + TileXR::Demo::kAllToAllGroupTraceCacheLineBytes; const uint64_t offset = taskBaseOffset + - index * sizeof(TileXR::Demo::AllToAllGroupTraceTaskSpan); + coreIndex * coreBytes + taskIndex * sizeof(TileXR::Demo::AllToAllGroupTraceTaskSpan); auto span = reinterpret_cast<__gm__ TileXR::Demo::AllToAllGroupTraceTaskSpan*>( trace + offset); span->peer = peer; diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index 138ffe51..588783ea 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -22,6 +22,7 @@ constexpr uint32_t kAllToAllGroupTraceCoreCount = 32U; constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; +constexpr size_t kAllToAllGroupTraceCacheLineBytes = 128U; enum AllToAllGroupTracePhase : uint32_t { kAllToAllGroupTraceSelfCopy = 0U, @@ -62,14 +63,14 @@ constexpr size_t AllToAllGroupTraceKernelSpanOffset(uint32_t iteration, uint32_t { return kAllToAllGroupTraceHeaderBytes + (static_cast(iteration) * kAllToAllGroupTraceCoreCount + core) * - sizeof(AllToAllGroupTraceSpan); + kAllToAllGroupTraceCacheLineBytes; } constexpr size_t AllToAllGroupTraceTaskSpanBaseOffset() { return kAllToAllGroupTraceHeaderBytes + static_cast(kAllToAllGroupTraceMaxIterations) * - kAllToAllGroupTraceCoreCount * sizeof(AllToAllGroupTraceSpan); + kAllToAllGroupTraceCoreCount * kAllToAllGroupTraceCacheLineBytes; } inline bool AllToAllGroupTraceCheckedMultiply(size_t lhs, size_t rhs, size_t& result) @@ -85,16 +86,27 @@ inline bool AllToAllGroupTraceCheckedMultiply(size_t lhs, size_t rhs, size_t& re inline size_t AllToAllGroupTraceLayoutBytes( uint32_t iterationCount, uint32_t groupCount, uint32_t passCount) { - size_t count = iterationCount; + size_t coreBytes = groupCount; size_t next = 0U; const size_t factors[] = { - kAllToAllGroupTraceCoreCount, groupCount, passCount, - kAllToAllGroupTracePhaseCount, sizeof(AllToAllGroupTraceTaskSpan)}; + passCount, kAllToAllGroupTracePhaseCount, sizeof(AllToAllGroupTraceTaskSpan)}; for (size_t factor : factors) { - if (!AllToAllGroupTraceCheckedMultiply(count, factor, next)) { + if (!AllToAllGroupTraceCheckedMultiply(coreBytes, factor, next)) { return std::numeric_limits::max(); } - count = next; + coreBytes = next; + } + if (coreBytes > std::numeric_limits::max() - + (kAllToAllGroupTraceCacheLineBytes - 1U)) { + return std::numeric_limits::max(); + } + coreBytes = (coreBytes + kAllToAllGroupTraceCacheLineBytes - 1U) & + ~(kAllToAllGroupTraceCacheLineBytes - 1U); + size_t count = iterationCount; + if (!AllToAllGroupTraceCheckedMultiply( + count, kAllToAllGroupTraceCoreCount, next) || + !AllToAllGroupTraceCheckedMultiply(next, coreBytes, count)) { + return std::numeric_limits::max(); } if (count > std::numeric_limits::max() - AllToAllGroupTraceTaskSpanBaseOffset()) { @@ -117,12 +129,16 @@ inline size_t AllToAllGroupTraceTaskSpanOffset( uint32_t iteration, uint32_t core, uint32_t group, uint32_t pass, uint32_t phase, uint32_t groupCount, uint32_t passCount) { - const size_t index = - (((((static_cast(iteration) * kAllToAllGroupTraceCoreCount + core) * - groupCount + group) * passCount + pass) * + const size_t rawCoreBytes = static_cast(groupCount) * passCount * + kAllToAllGroupTracePhaseCount * sizeof(AllToAllGroupTraceTaskSpan); + const size_t coreBytes = (rawCoreBytes + kAllToAllGroupTraceCacheLineBytes - 1U) & + ~(kAllToAllGroupTraceCacheLineBytes - 1U); + const size_t coreIndex = static_cast(iteration) * + kAllToAllGroupTraceCoreCount + core; + const size_t taskIndex = (((static_cast(group) * passCount + pass) * kAllToAllGroupTracePhaseCount) + phase); return AllToAllGroupTraceTaskSpanBaseOffset() + - index * sizeof(AllToAllGroupTraceTaskSpan); + coreIndex * coreBytes + taskIndex * sizeof(AllToAllGroupTraceTaskSpan); } static_assert(sizeof(AllToAllGroupTraceSpan) == 16U, diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index ffe60900..348f551c 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -13,10 +13,11 @@ MAX_CORES = 32 PHASE_COUNT = 5 SPAN_BYTES = 16 +CACHE_LINE_BYTES = 128 TASK_FORMAT = " Date: Mon, 20 Jul 2026 19:16:45 +0800 Subject: [PATCH 073/163] docs(udma): design grouped alltoall dual-route split --- ...udma-grouped-alltoall-dual-route-design.md | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-dual-route-design.md diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-dual-route-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-dual-route-design.md new file mode 100644 index 00000000..ae343dc4 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-dual-route-design.md @@ -0,0 +1,118 @@ +# UDMA Grouped AllToAll Dual-Route Design + +## Goal + +Extend the standalone grouped-fullmesh AllToAll kernel so cross-node traffic +uses both aggregate UDMA routes. Preserve the existing one-peer/one-route +protocol while distributing peers according to the route capacities. + +The initial target remains `rankSize = N * 8`, with contiguous groups of eight +ranks representing physical nodes. The physical 2x8, 128 MiB/rank result at +commit `68e3b5f` is the comparison baseline: + +- Host mean: 824.632 us +- Host min/max: 822.646 / 827.564 us +- Loop-49 kernel-envelope mean: 822.583 us + +## Current Behavior + +With `TILEXR_UDMA_ROUTE_POLICY=all` and `TILEXR_UDMA_QP_NUM=4`, cross-node +peers expose two aggregate routes from `/etc/hccl_rootinfo.json`: + +- `plane_pg_0`: six ports, represented by QP 0-3 with weight 6 +- `plane_pg_1`: two ports, represented by QP 4-7 with weight 2 + +The grouped kernel currently chooses only the highest-weight QP. Ties preserve +the lowest QP index, so all payload and ready traffic uses QP0 and the two-port +route remains idle. + +## Scope + +- Change only the standalone grouped AllToAll kernel and its focused tests. +- Keep the existing 35-core fullmesh kernel unchanged. +- Keep one complete peer payload on one route; do not split one peer payload. +- Keep payload, ready signal, and immediate quiet on the same selected QP. +- Keep the receive-side signal count and registered-memory layout unchanged. +- Apply dual-route selection only to cross-node peers. +- Continue to support `8 <= rankSize <= 128` and `rankSize % 8 == 0`. +- Treat each contiguous group of eight global ranks as one node. + +Payload striping within each peer is explicitly deferred. It would require two +independent ready signals and receive-side completion of both segments. + +## Route Discovery + +For each peer, device code reads all QP weights from `UDMAInfo`. + +1. The primary QP is the lowest QP index with the maximum weight. +2. The secondary QP is the lowest QP index with the largest weight strictly + below the primary weight. +3. If no lower nonzero weight exists, secondary equals primary. + +This discovers QP0 and QP4 for the current 6-port/2-port topology without +hardcoding either QP index. It also preserves deterministic fallback on +single-route or equal-weight topologies. + +## Balanced Peer Assignment + +For a cross-node pair: + +```text +sourceLocal = sourceRank % 8 +targetLocal = targetRank % 8 +routeIndex = (sourceLocal + targetLocal) % 8 +selectedQp = routeIndex < 6 ? primaryQp : secondaryQp +``` + +For every source rank and every remote node, six target peers use the primary +route and two target peers use the secondary route. For every destination rank +and every remote node, the eight incoming source ranks also divide 6:2. This +avoids concentrating all secondary-route traffic on two destination cards. + +Same-node peers continue to use the primary max-weight QP. Self traffic remains +a local MTE copy and does not use UDMA. + +## Data Flow + +The sender chooses the peer's QP once before its pass loop. Every pass then +uses that same QP for: + +1. `UDMAPutSignalNbiOnQp` payload plus ready signal +2. `UDMAQuietStatusOnQp` + +The receiver continues waiting on the same source-rank signal slot and copies +the source-rank payload slot into output. No new signal, ACK, barrier, or +registered-memory region is introduced. + +## Error Handling + +- Invalid rank dimensions remain rejected by the grouped layout plan. +- Missing secondary route falls back to primary rather than selecting an + unavailable QP. +- Existing quiet-status and wait-timeout debug records remain authoritative. +- Trace records continue storing the selected QP for every send phase. + +## Verification + +Focused unit/source tests will cover: + +- deterministic primary/secondary selection from weights `{6,6,6,6,2,2,2,2}` +- fallback when all QP weights are equal +- balanced 6:2 row and column counts for every pair of eight-rank nodes +- same-node peers selecting primary +- payload signal and quiet using the same selected QP +- absence of `SyncAll` and generic UDMA route calls + +Physical 2x8 validation will rerun 128 MiB/rank with warmup5/repeat50 and GM +trace enabled. Loop49 must show: + +- 128 cross-node send events total +- 96 cross-node events on QP0 +- 32 cross-node events on QP4 +- 112 same-node send events on their primary QP +- 15 unique send and receive peers per rank +- no quiet failure, wait timeout, or output mismatch + +Correctness and QP-distribution results are evaluated before performance. Host +mean and loop-49 kernel envelope are compared against the baseline without +assuming that route utilization must improve latency. From 35ad253f28560ef9d018e7e0aa5976d3b0558350 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 19:21:18 +0800 Subject: [PATCH 074/163] docs(udma): plan grouped alltoall dual-route split --- ...-07-20-udma-grouped-alltoall-dual-route.md | 366 ++++++++++++++++++ 1 file changed, 366 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-dual-route.md diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-dual-route.md b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-dual-route.md new file mode 100644 index 00000000..39bd3e96 --- /dev/null +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-dual-route.md @@ -0,0 +1,366 @@ +# UDMA Grouped AllToAll Dual-Route Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Distribute grouped AllToAll cross-node peers over the six-port and two-port aggregate UDMA routes in a balanced 6:2 ratio without changing the payload/ready protocol. + +**Architecture:** Add a host-testable route-policy header for node and 6:2 peer assignment. The grouped device kernel uses an equivalent device-local predicate because Bisheng does not reliably link ordinary host inline functions into AICore objects, discovers the highest and second-highest distinct QP weights, and continues issuing payload plus signal and quiet on one selected QP per peer. + +**Tech Stack:** C++14 host tests, Ascend C/Bisheng device code, TileXR UDMA QP-weight metadata, CMake, Chrome Trace JSON, and physical Ascend950 2x8 validation with CANN `/home/pkg/b101/cann`. + +## Global Constraints + +- Keep the existing 35-core big-data fullmesh kernel unchanged. +- Keep one complete peer payload on one route; do not split payload within a peer. +- Keep `UDMAPutSignalNbiOnQp` and `UDMAQuietStatusOnQp` on the same selected QP. +- Do not add signals, ACKs, barriers, registered regions, or `SyncAll()`. +- Apply dual-route selection only when `rank / 8 != peer / 8`. +- Continue supporting `8 <= rankSize <= 128` and `rankSize % 8 == 0`. +- Treat each contiguous group of eight ranks as one node. +- Primary QP is the lowest index with maximum weight. +- Secondary QP is the lowest index with the largest weight strictly below the primary weight. +- Fall back to primary when no distinct lower-weight QP exists. +- Commit locally, create and verify a complete Git bundle, then upload the bundle before every remote build or run. +- Build and run against `/home/pkg/b101/cann`; set `TILEXR_IPC_PID_MODE=pid` for physical runs. + +--- + +### Task 1: Host-Testable Balanced Peer Policy + +**Files:** +- Create: `tests/udma/demo/tilexr_udma_alltoall_group_route.h` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` +- Modify: `tests/udma/CMakeLists.txt` + +**Interfaces:** +- Produces: `AllToAllGroupIsCrossNode(int rank, int peer) -> bool`. +- Produces: `AllToAllGroupUseSecondaryRoute(int rank, int peer) -> bool`. +- Produces: `kAllToAllGroupRanksPerNode=8U` and `kAllToAllGroupPrimaryPeersPerNode=6U`. + +- [ ] **Step 1: Write the failing route-policy test** + +Add the missing header include and this test to +`tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp`: + +```cpp +#include "demo/tilexr_udma_alltoall_group_route.h" + +void TestDualRoutePeerPolicy() +{ + using TileXR::Demo::AllToAllGroupIsCrossNode; + using TileXR::Demo::AllToAllGroupUseSecondaryRoute; + CHECK_EQ(AllToAllGroupIsCrossNode(0, 7), false); + CHECK_EQ(AllToAllGroupIsCrossNode(0, 8), true); + CHECK_EQ(AllToAllGroupIsCrossNode(15, 8), false); + CHECK_EQ(AllToAllGroupIsCrossNode(15, 0), true); + + for (int sourceNode = 0; sourceNode < 3; ++sourceNode) { + for (int targetNode = 0; targetNode < 3; ++targetNode) { + if (sourceNode == targetNode) continue; + int rowSecondary[8] = {}; + int columnSecondary[8] = {}; + for (int sourceLocal = 0; sourceLocal < 8; ++sourceLocal) { + for (int targetLocal = 0; targetLocal < 8; ++targetLocal) { + const int source = sourceNode * 8 + sourceLocal; + const int target = targetNode * 8 + targetLocal; + if (AllToAllGroupUseSecondaryRoute(source, target)) { + ++rowSecondary[sourceLocal]; + ++columnSecondary[targetLocal]; + } + } + } + for (int local = 0; local < 8; ++local) { + CHECK_EQ(rowSecondary[local], 2); + CHECK_EQ(columnSecondary[local], 2); + } + } + } +} +``` + +Call `TestDualRoutePeerPolicy()` from `main()`. + +- [ ] **Step 2: Verify RED and commit the test** + +```bash +git diff --check +git add tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "test(udma): require balanced grouped dual-route policy" +git bundle create tmp/grouped-dual-route-red.bundle HEAD +git bundle verify tmp/grouped-dual-route-red.bundle +``` + +Deploy to `141.61.49.223` and run: + +```bash +source /home/pkg/b101/cann/set_env.sh +export ASCEND_HOME_PATH=/home/pkg/b101/cann +cmake --build tests/udma/build_b101 --target test_tilexr_udma_alltoall_group_layout -j8 +``` + +Expected: compilation fails because `tilexr_udma_alltoall_group_route.h` is missing. + +- [ ] **Step 3: Implement the minimal peer policy** + +Create `tests/udma/demo/tilexr_udma_alltoall_group_route.h`: + +```cpp +#ifndef TILEXR_UDMA_ALLTOALL_GROUP_ROUTE_H +#define TILEXR_UDMA_ALLTOALL_GROUP_ROUTE_H + +#include + +namespace TileXR { +namespace Demo { +constexpr uint32_t kAllToAllGroupRanksPerNode = 8U; +constexpr uint32_t kAllToAllGroupPrimaryPeersPerNode = 6U; + +inline bool AllToAllGroupIsCrossNode(int rank, int peer) +{ + return rank >= 0 && peer >= 0 && + rank / static_cast(kAllToAllGroupRanksPerNode) != + peer / static_cast(kAllToAllGroupRanksPerNode); +} + +inline bool AllToAllGroupUseSecondaryRoute(int rank, int peer) +{ + if (!AllToAllGroupIsCrossNode(rank, peer)) return false; + const uint32_t sourceLocal = + static_cast(rank) % kAllToAllGroupRanksPerNode; + const uint32_t targetLocal = + static_cast(peer) % kAllToAllGroupRanksPerNode; + return (sourceLocal + targetLocal) % kAllToAllGroupRanksPerNode >= + kAllToAllGroupPrimaryPeersPerNode; +} +} // namespace Demo +} // namespace TileXR +#endif +``` + +Add this header to the grouped-kernel custom-command `DEPENDS` list in +`tests/udma/CMakeLists.txt`. + +- [ ] **Step 4: Verify GREEN and commit** + +Bundle and deploy the implementation, then run on `141.61.49.223`: + +```bash +cmake --build tests/udma/build_b101 --target test_tilexr_udma_alltoall_group_layout -j8 +./tests/udma/build_b101/test_tilexr_udma_alltoall_group_layout +``` + +Expected: `TileXR grouped all-to-all layout checks passed`. + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_route.h \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp tests/udma/CMakeLists.txt +git commit -m "feat(udma): add grouped dual-route peer policy" +``` + +--- + +### Task 2: Device QP Discovery And Selection + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Consumes: route-policy constants and formula from Task 1. +- Produces: `AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp)`. +- Produces: device helper `AllToAllGroupUseSecondaryRouteDevice(rank, peer)`. +- Produces: one selected QP reused by payload plus signal, quiet, debug, and trace. + +- [ ] **Step 1: Add failing grouped-kernel source guards** + +Extend `TestKernelStructure()`: + +```cpp +CHECK_CONTAINS(kernel, "#include \"tilexr_udma_alltoall_group_route.h\""); +CHECK_CONTAINS(kernel, "AllToAllGroupSelectRouteQps"); +CHECK_CONTAINS(kernel, "AllToAllGroupUseSecondaryRouteDevice(rank, peer)"); +CHECK_CONTAINS(kernel, "secondaryQp"); +CHECK_CONTAINS(kernel, "selectedQp"); +CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); +``` + +Run the grouped layout test. Expected: source-guard failures for these names. + +- [ ] **Step 2: Implement distinct-weight QP discovery** + +Include the route-policy header. Add the device-local peer predicate so the +kernel does not depend on linking a normal host inline function: + +```cpp +__aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( + int32_t rank, int32_t peer) +{ + if (rank < 0 || peer < 0 || + rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) == + peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode)) { + return false; + } + const uint32_t sourceLocal = static_cast(rank) % + TileXR::Demo::kAllToAllGroupRanksPerNode; + const uint32_t targetLocal = static_cast(peer) % + TileXR::Demo::kAllToAllGroupRanksPerNode; + return (sourceLocal + targetLocal) % TileXR::Demo::kAllToAllGroupRanksPerNode >= + TileXR::Demo::kAllToAllGroupPrimaryPeersPerNode; +} +``` + +Then replace the max-only QP helper with: + +```cpp +__aicore__ inline void AllToAllGroupSelectRouteQps( + const __gm__ TileXR::CommArgs* args, int32_t peer, + uint32_t& primaryQp, uint32_t& secondaryQp) +{ + auto info = TileXR::GetUDMAInfo(args); + const uint32_t qpCount = info->qpNum == 0U ? 1U : info->qpNum; + primaryQp = 0U; + uint32_t primaryWeight = TileXR::UDMAGetQpWeight(info, peer, 0U); + for (uint32_t qp = 1U; qp < qpCount; ++qp) { + const uint32_t weight = TileXR::UDMAGetQpWeight(info, peer, qp); + if (weight > primaryWeight) { + primaryQp = qp; + primaryWeight = weight; + } + } + secondaryQp = primaryQp; + uint32_t secondaryWeight = 0U; + for (uint32_t qp = 0U; qp < qpCount; ++qp) { + const uint32_t weight = TileXR::UDMAGetQpWeight(info, peer, qp); + if (weight < primaryWeight && weight > secondaryWeight) { + secondaryQp = qp; + secondaryWeight = weight; + } + } +} +``` + +Strict `>` comparisons preserve the lowest QP index on equal weights. + +- [ ] **Step 3: Select once per peer and reuse the QP** + +Before the send pass loop: + +```cpp +uint32_t primaryQp = 0U; +uint32_t secondaryQp = 0U; +AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp); +const uint32_t selectedQp = + AllToAllGroupUseSecondaryRouteDevice(rank, peer) ? + secondaryQp : primaryQp; +``` + +Use `selectedQp` for `UDMAPutSignalNbiOnQp`, `UDMAQuietStatusOnQp`, both send +trace records, and quiet-error debug. Do not modify receive-side waits or copy. + +- [ ] **Step 4: Build, regress, and commit** + +Commit, bundle, deploy, then run on `141.61.49.223`: + +```bash +source /home/pkg/b101/cann/set_env.sh +export ASCEND_HOME_PATH=/home/pkg/b101/cann +cmake --build tests/udma/build_b101 -j8 +./tests/udma/build_b101/test_tilexr_udma_alltoall_group_layout +./tests/udma/build_b101/test_tilexr_udma_alltoall_layout +./tests/udma/build_b101/test_tilexr_udma_transport_layout +python3 -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +Expected: Bisheng and Host builds succeed and every test exits zero. + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "feat(udma): balance grouped peers across UDMA routes" +``` + +--- + +### Task 3: Physical 2x8 Comparison + +**Files:** +- Create artifact: `tmp/udma-grouped-alltoall-dual-route.bundle` +- Create run artifacts under: `tmp/grouped_alltoall_dual_route_b101_2x8/` +- Do not commit raw binaries, JSON traces, or run logs. + +**Interfaces:** +- Consumes: grouped dual-route kernel from Task 2. +- Produces: correctness, QP distribution, Host timing, and loop49 kernel envelope. + +- [ ] **Step 1: Create final bundle and deploy both hosts** + +```bash +git diff --check +git status --short +git bundle create tmp/udma-grouped-alltoall-dual-route.bundle HEAD +git bundle verify tmp/udma-grouped-alltoall-dual-route.bundle +``` + +Fetch it into `/home/h30059441/tilexr_grouped_alltoall_b101` on +`141.61.50.31` and `141.61.49.223`. Build/install on `141.61.49.223`, copy +the grouped kernel library and demo executable to `141.61.50.31`, then verify +commit IDs and SHA-256 hashes match. + +- [ ] **Step 2: Run 16-rank smoke correctness** + +```bash +export TILEXR_COMM_ID=141.61.50.31:64720 +export TILEXR_IPC_PID_MODE=pid +export TILEXR_UDMA_ROUTE_POLICY=all +export TILEXR_UDMA_QP_NUM=4 +export TILEXR_DEMO_ALLTOALL_WARMUP=1 +export TILEXR_DEMO_ALLTOALL_REPEAT=2 +timeout 180s ./tests/udma/install_b101/bin/tilexr_udma_demo \ + 16 "$rank" 8 4096 8 0 +``` + +Expected: all 16 ranks succeed with no quiet error, timeout, or mismatch. + +- [ ] **Step 3: Run 128 MiB/rank with trace** + +```bash +export TILEXR_DEMO_ALLTOALL_WARMUP=5 +export TILEXR_DEMO_ALLTOALL_REPEAT=50 +export TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS=2097152 +export TILEXR_UDMA_GROUP_TRACE=1 +export TILEXR_UDMA_GROUP_TRACE_DIR=/home/h30059441/grouped_alltoall_dual_route_b101_2x8 +timeout 300s ./tests/udma/install_b101/bin/tilexr_udma_demo \ + 16 "$rank" 8 2097152 8 0 +``` + +Expected: 16 successful ranks and sixteen 8 MiB traces. + +- [ ] **Step 4: Validate loop49 and compare performance** + +Convert and extract iteration49. Require: + +```text +cross-node total = 128 +cross-node QP0 = 96 +cross-node QP4 = 32 +same-node total = 112 +invalid QP = 0 +unique send peers/rank = 15 +unique receive peers/rank = 15 +complete loop49 spans = 1728 +``` + +Report Host mean/min/max and loop49 kernel-envelope mean/min/max beside the +single-route baseline `824.632 us` Host mean and `822.583 us` kernel mean. + +- [ ] **Step 5: Run final regressions** + +```bash +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_layout +./tests/udma/install_b101/bin/test_tilexr_udma_transport_layout +python3 -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +git status --short +``` + +Expected: every test passes and no tracked source differs from deployed HEAD. From 55bb7c77c42a836bebca704d4d5a094777058a0c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 19:23:16 +0800 Subject: [PATCH 075/163] test(udma): require balanced grouped dual-route policy --- ...test_tilexr_udma_alltoall_group_layout.cpp | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 29529285..eb55f54c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -6,6 +6,7 @@ #include #include "demo/tilexr_udma_alltoall_group_layout.h" +#include "demo/tilexr_udma_alltoall_group_route.h" namespace { @@ -148,6 +149,62 @@ void TestTokens() TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U), true); } +void TestDualRoutePeerPolicy() +{ + using TileXR::Demo::AllToAllGroupIsCrossNode; + using TileXR::Demo::AllToAllGroupUseSecondaryRoute; + CHECK_EQ(AllToAllGroupIsCrossNode(0, 7), false); + CHECK_EQ(AllToAllGroupIsCrossNode(0, 8), true); + CHECK_EQ(AllToAllGroupIsCrossNode(15, 8), false); + CHECK_EQ(AllToAllGroupIsCrossNode(15, 0), true); + + for (int sourceNode = 0; sourceNode < 3; ++sourceNode) { + for (int targetNode = 0; targetNode < 3; ++targetNode) { + if (sourceNode == targetNode) { + continue; + } + int rowSecondary[8] = {}; + int columnSecondary[8] = {}; + for (int sourceLocal = 0; sourceLocal < 8; ++sourceLocal) { + for (int targetLocal = 0; targetLocal < 8; ++targetLocal) { + const int source = sourceNode * 8 + sourceLocal; + const int target = targetNode * 8 + targetLocal; + if (AllToAllGroupUseSecondaryRoute(source, target)) { + ++rowSecondary[sourceLocal]; + ++columnSecondary[targetLocal]; + } + } + } + for (int local = 0; local < 8; ++local) { + CHECK_EQ(rowSecondary[local], 2); + CHECK_EQ(columnSecondary[local], 2); + } + } + } +} + +void TestDualRouteQpWeights() +{ + const uint32_t weighted[] = {6U, 6U, 6U, 6U, 2U, 2U, 2U, 2U}; + const auto split = TileXR::Demo::AllToAllGroupSelectRouteQps(weighted, 8U); + CHECK_EQ(split.primaryQp, 0U); + CHECK_EQ(split.secondaryQp, 4U); + + const uint32_t threeWeights[] = {2U, 6U, 4U, 6U}; + const auto distinct = TileXR::Demo::AllToAllGroupSelectRouteQps(threeWeights, 4U); + CHECK_EQ(distinct.primaryQp, 1U); + CHECK_EQ(distinct.secondaryQp, 2U); + + const uint32_t equal[] = {3U, 3U, 3U, 3U}; + const auto fallback = TileXR::Demo::AllToAllGroupSelectRouteQps(equal, 4U); + CHECK_EQ(fallback.primaryQp, 0U); + CHECK_EQ(fallback.secondaryQp, 0U); + + const auto empty = TileXR::Demo::AllToAllGroupSelectRouteQps(nullptr, 0U); + CHECK_EQ(empty.primaryQp, 0U); + CHECK_EQ(empty.secondaryQp, 0U); +} + void TestKernelStructure() { const std::string kernel = ReadFile( @@ -189,6 +246,8 @@ int main() TestSchedules(); TestPlan(); TestTokens(); + TestDualRoutePeerPolicy(); + TestDualRouteQpWeights(); TestKernelStructure(); TestHostStructure(); if (g_failures != 0) { From e148dcce25edfa464e93f52fd91521e1f79a5525 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 19:25:41 +0800 Subject: [PATCH 076/163] feat(udma): add grouped dual-route peer policy --- tests/udma/CMakeLists.txt | 1 + .../demo/tilexr_udma_alltoall_group_route.h | 75 +++++++++++++++++++ 2 files changed, 76 insertions(+) create mode 100644 tests/udma/demo/tilexr_udma_alltoall_group_route.h diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 510693af..a99f5cc6 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -296,6 +296,7 @@ if(BUILD_TILEXR_UDMA_DEMO) DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_layout.h" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_route.h" "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_trace.h" "${TILEXR_ROOT}/src/include/tilexr_udma.h" VERBATIM diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h new file mode 100644 index 00000000..785b80b6 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -0,0 +1,75 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_ALLTOALL_GROUP_ROUTE_H +#define TILEXR_UDMA_ALLTOALL_GROUP_ROUTE_H + +#include + +namespace TileXR { +namespace Demo { + +constexpr uint32_t kAllToAllGroupRanksPerNode = 8U; +constexpr uint32_t kAllToAllGroupPrimaryPeersPerNode = 6U; + +struct AllToAllGroupRouteQps { + uint32_t primaryQp = 0U; + uint32_t secondaryQp = 0U; +}; + +inline bool AllToAllGroupIsCrossNode(int rank, int peer) +{ + return rank >= 0 && peer >= 0 && + rank / static_cast(kAllToAllGroupRanksPerNode) != + peer / static_cast(kAllToAllGroupRanksPerNode); +} + +inline bool AllToAllGroupUseSecondaryRoute(int rank, int peer) +{ + if (!AllToAllGroupIsCrossNode(rank, peer)) { + return false; + } + const uint32_t sourceLocal = + static_cast(rank) % kAllToAllGroupRanksPerNode; + const uint32_t targetLocal = + static_cast(peer) % kAllToAllGroupRanksPerNode; + return (sourceLocal + targetLocal) % kAllToAllGroupRanksPerNode >= + kAllToAllGroupPrimaryPeersPerNode; +} + +inline AllToAllGroupRouteQps AllToAllGroupSelectRouteQps( + const uint32_t* weights, uint32_t qpCount) +{ + AllToAllGroupRouteQps result {}; + if (weights == nullptr || qpCount == 0U) { + return result; + } + + uint32_t primaryWeight = weights[0]; + for (uint32_t qp = 1U; qp < qpCount; ++qp) { + if (weights[qp] > primaryWeight) { + result.primaryQp = qp; + primaryWeight = weights[qp]; + } + } + + uint32_t secondaryWeight = 0U; + for (uint32_t qp = 0U; qp < qpCount; ++qp) { + const uint32_t weight = weights[qp]; + if (weight != 0U && weight < primaryWeight && weight > secondaryWeight) { + result.secondaryQp = qp; + secondaryWeight = weight; + } + } + if (secondaryWeight == 0U) { + result.secondaryQp = result.primaryQp; + } + return result; +} + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_ALLTOALL_GROUP_ROUTE_H From 7da9a5dd23a0e250d343e5c2cb2503b2b0870b3e Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 19:27:34 +0800 Subject: [PATCH 077/163] test(udma): require grouped device dual-route selection --- tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index eb55f54c..1c05bab1 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -211,8 +211,13 @@ void TestKernelStructure() std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp"); CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_kernel"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES"); + CHECK_CONTAINS(kernel, "#include \"tilexr_udma_alltoall_group_route.h\""); + CHECK_CONTAINS(kernel, "AllToAllGroupSelectRouteQps"); + CHECK_CONTAINS(kernel, "AllToAllGroupUseSecondaryRouteDevice(rank, peer)"); + CHECK_CONTAINS(kernel, "secondaryQp"); + CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); - CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, qpIdx)"); + CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); CHECK_CONTAINS(kernel, "observed >= expectedToken"); CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_group"); From fc38975c0adcf6a671ec0bab8f663b0393ebc9c1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 19:30:03 +0800 Subject: [PATCH 078/163] feat(udma): balance grouped peers across UDMA routes --- .../tilexr_udma_alltoall_group_kernel.cpp | 63 ++++++++++++++----- 1 file changed, 47 insertions(+), 16 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index fd1952f2..3b86e36b 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -5,6 +5,7 @@ #include "kernel_operator.h" #include "tilexr_udma.h" +#include "tilexr_udma_alltoall_group_route.h" #include "tilexr_udma_alltoall_group_trace.h" namespace { @@ -64,21 +65,47 @@ __aicore__ inline uint64_t AllToAllGroupDeviceToken( (static_cast(pass) + 1ULL); } -__aicore__ inline uint32_t AllToAllGroupSelectMaxWeightQp( - const __gm__ TileXR::CommArgs* args, int32_t peer) +__aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( + int32_t rank, int32_t peer) +{ + if (rank < 0 || peer < 0 || + rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) == + peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode)) { + return false; + } + const uint32_t sourceLocal = static_cast(rank) % + TileXR::Demo::kAllToAllGroupRanksPerNode; + const uint32_t targetLocal = static_cast(peer) % + TileXR::Demo::kAllToAllGroupRanksPerNode; + return (sourceLocal + targetLocal) % TileXR::Demo::kAllToAllGroupRanksPerNode >= + TileXR::Demo::kAllToAllGroupPrimaryPeersPerNode; +} + +__aicore__ inline void AllToAllGroupSelectRouteQps( + const __gm__ TileXR::CommArgs* args, int32_t peer, + uint32_t& primaryQp, uint32_t& secondaryQp) { auto udmaInfo = TileXR::GetUDMAInfo(args); const uint32_t qpCount = udmaInfo->qpNum == 0U ? 1U : udmaInfo->qpNum; - uint32_t selected = 0U; - uint32_t selectedWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); - for (uint32_t qpIdx = 1U; qpIdx < qpCount; ++qpIdx) { - const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qpIdx); - if (weight > selectedWeight) { - selected = qpIdx; - selectedWeight = weight; + primaryQp = 0U; + uint32_t primaryWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); + for (uint32_t qp = 1U; qp < qpCount; ++qp) { + const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qp); + if (weight > primaryWeight) { + primaryQp = qp; + primaryWeight = weight; + } + } + + secondaryQp = primaryQp; + uint32_t secondaryWeight = 0U; + for (uint32_t qp = 0U; qp < qpCount; ++qp) { + const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qp); + if (weight != 0U && weight < primaryWeight && weight > secondaryWeight) { + secondaryQp = qp; + secondaryWeight = weight; } } - return selected; } __aicore__ inline void AllToAllGroupCopyMte( @@ -348,7 +375,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (peer < 0) { continue; } - const uint32_t qpIdx = AllToAllGroupSelectMaxWeightQp(args, peer); + uint32_t primaryQp = 0U; + uint32_t secondaryQp = 0U; + AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp); + const uint32_t selectedQp = + AllToAllGroupUseSecondaryRouteDevice(rank, peer) ? secondaryQp : primaryQp; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -370,21 +401,21 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE; const uint64_t putBegin = AllToAllGroupTraceCycle(groupTrace); TileXR::UDMAPutSignalNbiOnQp( - args, peer, qpIdx, localSrc, remotePayloadOffset, chunkBytes, + args, peer, selectedQp, localSrc, remotePayloadOffset, chunkBytes, remoteSignalOffset, expectedToken); AllToAllGroupTraceRecordTask( groupTrace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, - peer, qpIdx, putBegin, AllToAllGroupTraceCycle(groupTrace)); + peer, selectedQp, putBegin, AllToAllGroupTraceCycle(groupTrace)); const uint64_t quietBegin = AllToAllGroupTraceCycle(groupTrace); - const uint32_t quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, qpIdx); + const uint32_t quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); AllToAllGroupTraceRecordTask( groupTrace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, - peer, qpIdx, quietBegin, AllToAllGroupTraceCycle(groupTrace)); + peer, selectedQp, quietBegin, AllToAllGroupTraceCycle(groupTrace)); if (quietStatus != 0U) { AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_QUIET, - group, pass, peer, qpIdx, quietStatus, expectedToken, 0ULL); + group, pass, peer, selectedQp, quietStatus, expectedToken, 0ULL); AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; From ce9004aa766a16303f809f631e351561735ab8a3 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 19:58:49 +0800 Subject: [PATCH 079/163] docs(udma): design configurable grouped copyout workers --- ...grouped-alltoall-copyout-workers-design.md | 109 ++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-copyout-workers-design.md diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-copyout-workers-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-copyout-workers-design.md new file mode 100644 index 00000000..f65a27dc --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-copyout-workers-design.md @@ -0,0 +1,109 @@ +# UDMA Grouped AllToAll Configurable Copyout Workers Design + +## Goal + +Make the grouped AllToAll receive-copy worker count selectable between 16 and +8, then compare both configurations on the same physical 2x8, 128 MiB/rank +workload. The experiment tests whether lower GM copy concurrency improves UDMA +progress enough to offset the additional copy work per core. + +The dual-route result at commit `fc38975` is the baseline: + +- Host mean: 773.490 us +- Host min/max: 767.636 / 775.511 us +- Loop-49 kernel-envelope mean: 762.332 us + +## Configuration + +Host mode `testType=8` reads: + +```text +TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS=8|16 +``` + +The default is 16. Any other value is rejected before allocation or launch. +The Host launches `16 + copyoutWorkers` AIV blocks and passes the selected +worker count to the grouped kernel. + +## Core And Lane Mapping + +Cores 0-15 remain send workers. Receive workers begin at core16. + +With 16 copyout workers, the existing mapping remains unchanged: receive core +`16 + lane` handles one lane. + +With 8 copyout workers, physical receive worker `worker` handles logical lanes +`worker` and `worker + 8`: + +```text +core16: lane0, lane8 +core17: lane1, lane9 +... +core23: lane7, lane15 +``` + +For each group, a worker processes its first lane and then its second lane. +Invalid lanes are skipped through the existing peer mapper. At rankSize=16, +lane15 is invalid, so core23 processes one remote peer while cores16-22 process +two. + +The self-copy range is divided by `copyoutWorkers`, producing eight equal +shards in 8-worker mode and preserving sixteen shards in 16-worker mode. + +## Protocol + +The change is receive scheduling only. It does not alter: + +- grouped peer order +- dual-route 6:2 QP selection +- payload or signal offsets +- source-rank receive slots +- ready tokens +- payload-plus-signal QP affinity +- quiet behavior +- ping-pong planes + +No barrier, ACK, signal, or registered-memory region is added. + +## Trace Mapping + +One physical 8-worker core handles two logical peers that share the same +group/pass/phase dimensions. Recording both under the physical core would +overwrite one trace cell. + +Receive wait/copy task spans therefore use logical trace core `16 + lane`. +This preserves one task cell per lane and keeps 8-worker traces comparable with +16-worker traces. Whole-kernel and self-copy spans use the actual physical core +`blockIdx`; cores24-31 have no kernel span in 8-worker mode. + +Debug errors continue to identify the physical core and include the peer. + +## Validation + +Focused tests cover: + +- Host acceptance of 8 and 16 and rejection of all other values +- block dimensions 24 and 32 +- 8-worker lane ownership covering lanes0-15 exactly once +- 16-worker mode preserving one lane per worker +- self-copy partitioning by the configured worker count +- logical trace core `16 + lane` for receive wait/copy +- absence of `SyncAll` and generic UDMA route calls + +Physical validation runs both worker counts with: + +```text +rankSize=16 +input/output=128 MiB per rank +chunk=8 MiB per peer +warmup=5 +repeat=50 +dual route=6:2 peer distribution +trace=enabled +``` + +Both runs must pass all 16-rank output checks with no quiet failure or wait +timeout. Report Host mean/min/max, loop-49 kernel envelope, send quiet, +receive wait, receive copy, trace event counts, and QP distribution. Keep 8 +workers only if correctness holds and measured performance is competitive with +or better than 16 workers. From aeb8fcbba442873dc67402ff7c7913ca00490b2a Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 20:02:54 +0800 Subject: [PATCH 080/163] docs(udma): plan configurable grouped copyout workers --- ...0-udma-grouped-alltoall-copyout-workers.md | 230 ++++++++++++++++++ 1 file changed, 230 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-copyout-workers.md diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-copyout-workers.md b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-copyout-workers.md new file mode 100644 index 00000000..801c270b --- /dev/null +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-copyout-workers.md @@ -0,0 +1,230 @@ +# UDMA Grouped AllToAll Configurable Copyout Workers Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make grouped AllToAll copyout workers selectable as 8 or 16 and compare both configurations on the physical 2x8, 128 MiB/rank workload. + +**Architecture:** Put copyout validation, block-dimension calculation, and worker-to-lane mapping in the existing grouped layout header so they are host-testable. Pass the selected count from Host to the kernel; use an equivalent device-local lane helper because Bisheng does not reliably link ordinary host inline functions into AICore objects. Active receive cores iterate their owned logical lanes while trace tasks remain indexed by logical lane. + +**Tech Stack:** C++14, Ascend C/Bisheng, TileXR UDMA, CMake, GM trace converter, CANN `/home/pkg/b101/cann`, physical Ascend950 2x8. + +## Global Constraints + +- Accept only 8 or 16 copyout workers; default to 16. +- Keep send cores 0-15 and all dual-route behavior unchanged. +- Launch 24 blocks for 8 workers and 32 blocks for 16 workers. +- Do not change payload, signal, token, ping-pong, or registered-memory layouts. +- Do not add barriers, ACKs, signals, `SyncAll`, or generic UDMA routes. +- Index receive wait/copy trace tasks by logical core `16 + lane`. +- Index kernel and self-copy trace spans by physical `blockIdx`. +- Commit and verify a complete bundle before each remote build or run. + +--- + +### Task 1: Copyout Layout Policy And Host Configuration + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_layout.h` +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Produces: `AllToAllGroupValidCopyoutWorkers(uint32_t) -> bool`. +- Produces: `AllToAllGroupBlockDim(uint32_t) -> uint32_t`. +- Produces: `AllToAllGroupCopyoutLane(worker, assignment, workers) -> int32_t`. +- Produces: environment variable `TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS`. +- Extends: grouped launch with trailing `uint32_t copyoutWorkers`. + +- [ ] **Step 1: Write failing layout and Host source tests** + +Add checks equivalent to: + +```cpp +CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); +CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); +CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); +CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); +CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); + +std::set lanes; +for (uint32_t worker = 0; worker < 8U; ++worker) { + for (uint32_t assignment = 0; assignment < 2U; ++assignment) { + lanes.insert(TileXR::Demo::AllToAllGroupCopyoutLane(worker, assignment, 8U)); + } +} +CHECK_EQ(lanes.size(), 16U); +CHECK_EQ(*lanes.begin(), 0); +CHECK_EQ(*lanes.rbegin(), 15); +CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(7U, 1U, 16U), -1); +``` + +Require Host source strings for the environment variable, validation, dynamic +block dimension, and launch argument. + +- [ ] **Step 2: Verify RED and commit the test** + +Commit and bundle the test, deploy to `141.61.49.223`, then build +`test_tilexr_udma_alltoall_group_layout`. Expected: compile/source-guard +failures because the copyout policy and Host plumbing do not exist. + +- [ ] **Step 3: Implement layout helpers** + +Add to the grouped layout header: + +```cpp +constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; + +inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) +{ + return workers == 8U || workers == 16U; +} + +inline uint32_t AllToAllGroupBlockDim(uint32_t workers) +{ + return AllToAllGroupValidCopyoutWorkers(workers) ? + kAllToAllGroupSendCoreCount + workers : 0U; +} + +inline int32_t AllToAllGroupCopyoutLane( + uint32_t worker, uint32_t assignment, uint32_t workers) +{ + if (!AllToAllGroupValidCopyoutWorkers(workers) || worker >= workers) return -1; + const uint32_t lane = worker + assignment * workers; + return lane < kAllToAllGroupWidth ? static_cast(lane) : -1; +} +``` + +- [ ] **Step 4: Add Host configuration** + +Read the environment in `RunGroupedAllToAll`, reject values other than 8/16, +print the selected count, launch `AllToAllGroupBlockDim(copyoutWorkers)`, and +pass `copyoutWorkers` to warmup and measured kernel launches. + +- [ ] **Step 5: Verify GREEN and commit** + +Bundle/deploy, build the full demo, and run the grouped layout test. Commit as: + +```bash +git commit -m "feat(udma): configure grouped copyout workers" +``` + +--- + +### Task 2: Kernel Lane Scheduling And Trace Safety + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Consumes: `copyoutWorkers` and Task 1 lane helpers. +- Produces: physical receive workers that handle one or two logical lanes. +- Preserves: one trace wait/copy cell per logical lane. + +- [ ] **Step 1: Add failing kernel source guards** + +Require: + +```cpp +CHECK_CONTAINS(kernel, "copyoutWorkers"); +CHECK_CONTAINS(kernel, "AllToAllGroupCopyoutLaneDevice"); +CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers"); +CHECK_CONTAINS(kernel, "traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane"); +CHECK_NOT_CONTAINS(kernel, "elementsPerPeer) * lane /"); +``` + +Expected RED: the existing receive branch assumes 16 workers and one lane. + +- [ ] **Step 2: Implement active-core and self-copy mapping** + +Validate `copyoutWorkers` as 8 or 16. Return inactive blocks at or above +`16 + copyoutWorkers`. Compute receive worker as `blockIdx - 16`, and divide +self-copy using `copyoutWorkers` rather than 16. + +- [ ] **Step 3: Iterate logical lanes inside each group** + +Add the device-local equivalent of the tested Host helper: + +```cpp +__aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( + uint32_t worker, uint32_t assignment, uint32_t workers) +{ + const uint32_t lane = worker + assignment * workers; + return lane < TILEXR_ALLTOALL_GROUP_SEND_CORES ? + static_cast(lane) : -1; +} +``` + +Use it in the receive loop: + +```cpp +for (uint32_t group = 0U; group < groupCount; ++group) { + for (uint32_t assignment = 0U; ; ++assignment) { + const int32_t laneValue = AllToAllGroupCopyoutLaneDevice( + worker, assignment, copyoutWorkers); + if (laneValue < 0) break; + const uint32_t lane = static_cast(laneValue); + const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); + if (peer < 0) continue; + // existing pass wait and copy loop + } +} +``` + +Record wait/copy tasks with `traceCore = 16U + lane`; retain physical blockIdx +for kernel, self-copy, and debug. + +- [ ] **Step 4: Build, regress, and commit** + +Bundle/deploy and run the Bisheng build, grouped/existing layout tests, +transport test, and Python trace tests. Commit as: + +```bash +git commit -m "feat(udma): schedule grouped copyout on 8 or 16 cores" +``` + +--- + +### Task 3: Physical 16-Worker And 8-Worker Matrix + +**Files:** +- Create artifact: `tmp/udma-grouped-alltoall-copyout-workers.bundle` +- Create artifacts under: `tmp/grouped_alltoall_copyout_workers_b101_2x8/` + +- [ ] **Step 1: Deploy identical final artifacts** + +Verify the complete bundle, deploy both hosts, build/install on +`141.61.49.223`, copy the grouped kernel and demo to `141.61.50.31`, and +compare commit IDs and SHA-256 hashes. + +- [ ] **Step 2: Run 16-worker target** + +Run 128 MiB/rank, warmup5/repeat50, dual route and trace enabled with: + +```bash +export TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS=16 +``` + +Require 16 successful ranks and sixteen complete 8 MiB traces. + +- [ ] **Step 3: Run 8-worker target** + +Repeat the identical workload with: + +```bash +export TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS=8 +``` + +Require 16 successful ranks and sixteen complete 8 MiB traces. + +- [ ] **Step 4: Compare traces and timing** + +For each mode report Host mean/min/max, loop49 kernel envelope, self-copy, +send-put-signal, send-quiet, receive-wait, receive-copy, complete event count, +and QP distribution. Verify every rank still has 15 send/receive peers and +cross-node QP0:QP4 remains 96:32. + +- [ ] **Step 5: Final regression and bundle verification** + +Run all grouped/existing layout, transport, and Python trace tests, confirm no +tracked deployment differences, and verify the final bundle contains HEAD. From e19830333e75d3d866de7346781db43cd950226b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 20:03:41 +0800 Subject: [PATCH 081/163] test(udma): require configurable grouped copyout workers --- ...test_tilexr_udma_alltoall_group_layout.cpp | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 1c05bab1..2a842833 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -205,6 +205,31 @@ void TestDualRouteQpWeights() CHECK_EQ(empty.secondaryQp, 0U); } +void TestCopyoutWorkerPolicy() +{ + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(4U), 0U); + + std::set lanes; + for (uint32_t worker = 0U; worker < 8U; ++worker) { + for (uint32_t assignment = 0U; assignment < 2U; ++assignment) { + lanes.insert(TileXR::Demo::AllToAllGroupCopyoutLane( + worker, assignment, 8U)); + } + } + CHECK_EQ(lanes.size(), 16U); + CHECK_EQ(*lanes.begin(), 0); + CHECK_EQ(*lanes.rbegin(), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(7U, 1U, 16U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(15U, 0U, 16U), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(8U, 0U, 8U), -1); +} + void TestKernelStructure() { const std::string kernel = ReadFile( @@ -235,6 +260,9 @@ void TestHostStructure() CHECK_CONTAINS(demo, "PlanAllToAllGroup"); CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_group"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS"); + CHECK_CONTAINS(demo, "AllToAllGroupValidCopyoutWorkers"); + CHECK_CONTAINS(demo, "AllToAllGroupBlockDim(copyoutWorkers)"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); CHECK_CONTAINS(demo, "grouped alltoall warmup="); const size_t begin = demo.find("bool RunGroupedAllToAll("); @@ -253,6 +281,7 @@ int main() TestTokens(); TestDualRoutePeerPolicy(); TestDualRouteQpWeights(); + TestCopyoutWorkerPolicy(); TestKernelStructure(); TestHostStructure(); if (g_failures != 0) { From f3f3ae60d285953b55ab2fb47c3a76277414367c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 20:05:57 +0800 Subject: [PATCH 082/163] test(udma): require grouped copyout lane scheduling --- tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 2a842833..c8cf2add 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -241,6 +241,10 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupUseSecondaryRouteDevice(rank, peer)"); CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); + CHECK_CONTAINS(kernel, "copyoutWorkers"); + CHECK_CONTAINS(kernel, "AllToAllGroupCopyoutLaneDevice"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers"); + CHECK_CONTAINS(kernel, "traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); @@ -248,6 +252,7 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_group"); CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); CHECK_NOT_CONTAINS(kernel, "SyncAll"); + CHECK_NOT_CONTAINS(kernel, "elementsPerPeer) * lane /"); } void TestHostStructure() From 1415be3bc8d2b990a6f65871eff8ae0eb68fbc5d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 20:08:46 +0800 Subject: [PATCH 083/163] feat(udma): configure grouped copyout workers --- .../tilexr_udma_alltoall_group_kernel.cpp | 40 +++++++++++++------ .../demo/tilexr_udma_alltoall_group_layout.h | 22 ++++++++++ tests/udma/demo/tilexr_udma_demo.cpp | 26 +++++++++--- 3 files changed, 70 insertions(+), 18 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 3b86e36b..3e529437 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -81,6 +81,14 @@ __aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( TileXR::Demo::kAllToAllGroupPrimaryPeersPerNode; } +__aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( + uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) +{ + const uint32_t lane = worker + assignment * copyoutWorkers; + return lane < TILEXR_ALLTOALL_GROUP_SEND_CORES ? + static_cast(lane) : -1; +} + __aicore__ inline void AllToAllGroupSelectRouteQps( const __gm__ TileXR::CommArgs* args, int32_t peer, uint32_t& primaryQp, uint32_t& secondaryQp) @@ -252,7 +260,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, - GM_ADDR groupTraceGM, uint32_t traceIteration) + GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t copyoutWorkers) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? @@ -266,7 +274,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rank = args->rank; const int32_t rankSize = args->rankSize; - if (blockIdx >= TILEXR_ALLTOALL_GROUP_BLOCK_DIM || + if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || + blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || groupCount == 0U) { @@ -289,13 +298,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(elementsPerPeer) * sizeof(int32_t); if (blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES) { - const uint32_t lane = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; + const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; const int32_t selfBegin = static_cast( - static_cast(elementsPerPeer) * lane / - TILEXR_ALLTOALL_GROUP_SEND_CORES); + static_cast(elementsPerPeer) * worker / copyoutWorkers); const int32_t selfEnd = static_cast( - static_cast(elementsPerPeer) * (lane + 1U) / - TILEXR_ALLTOALL_GROUP_SEND_CORES); + static_cast(elementsPerPeer) * (worker + 1U) / copyoutWorkers); if (selfEnd > selfBegin) { const uint64_t selfCopyBegin = AllToAllGroupTraceCycle(groupTrace); auto selfSrc = reinterpret_cast<__gm__ uint8_t*>( @@ -312,10 +319,18 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } for (uint32_t group = 0U; group < groupCount; ++group) { + for (uint32_t assignment = 0U; ; ++assignment) { + const int32_t laneValue = AllToAllGroupCopyoutLaneDevice( + worker, assignment, copyoutWorkers); + if (laneValue < 0) { + break; + } + const uint32_t lane = static_cast(laneValue); const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); if (peer < 0) { continue; } + const uint32_t traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -335,7 +350,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (!AllToAllGroupWaitTokenMte(signal, expectedToken, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, pass, + groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, waitBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -346,7 +361,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( return; } AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, pass, + groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, waitBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -358,11 +373,12 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); AllToAllGroupCopyMte(relayDst, relaySrc, chunkBytes, relayLocal); AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, pass, + groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveCopy, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } + } } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -433,11 +449,11 @@ void launch_tilexr_udma_all_to_all_group( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, - GM_ADDR groupTrace, uint32_t traceIteration) + GM_ADDR groupTrace, uint32_t traceIteration, uint32_t copyoutWorkers) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration); + groupTrace, traceIteration, copyoutWorkers); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 2e63c290..15b811ab 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -20,6 +20,7 @@ constexpr uint32_t kAllToAllGroupWidth = 16U; constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; +constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; constexpr uint32_t kAllToAllGroupBlockDim = 32U; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; @@ -46,6 +47,27 @@ inline bool AllToAllGroupValidRankSize(int rankSize) rankSize <= kAllToAllGroupMaxRankSize && rankSize % 8 == 0; } +inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) +{ + return workers == 8U || workers == 16U; +} + +inline uint32_t AllToAllGroupBlockDim(uint32_t workers) +{ + return AllToAllGroupValidCopyoutWorkers(workers) ? + kAllToAllGroupSendCoreCount + workers : 0U; +} + +inline int32_t AllToAllGroupCopyoutLane( + uint32_t worker, uint32_t assignment, uint32_t workers) +{ + if (!AllToAllGroupValidCopyoutWorkers(workers) || worker >= workers) { + return -1; + } + const uint32_t lane = worker + assignment * workers; + return lane < kAllToAllGroupWidth ? static_cast(lane) : -1; +} + inline uint32_t AllToAllGroupCount(int rankSize) { if (!AllToAllGroupValidRankSize(rankSize)) { diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index f245716d..463fb9ce 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -57,7 +57,7 @@ extern void launch_tilexr_udma_all_to_all_group( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, - GM_ADDR groupTrace, uint32_t traceIteration); + GM_ADDR groupTrace, uint32_t traceIteration, uint32_t copyoutWorkers); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -647,6 +647,17 @@ bool RunGroupedAllToAll( << " chunkElements=" << requestedChunkElements << std::endl; return false; } + const int copyoutWorkersValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 16); + if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( + static_cast(copyoutWorkersValue))) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8 or 16, got " + << copyoutWorkersValue << std::endl; + return false; + } + const uint32_t copyoutWorkers = static_cast(copyoutWorkersValue); + const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim(copyoutWorkers); const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); const bool traceEnabled = GetEnvInt("TILEXR_UDMA_GROUP_TRACE", 0) != 0; @@ -763,18 +774,20 @@ bool RunGroupedAllToAll( " groups=" + std::to_string(plan.groupCount) + " passes=" + std::to_string(plan.passCount)); PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + - " repeat=" + std::to_string(repeat)); + " repeat=" + std::to_string(repeat) + + " copyoutWorkers=" + std::to_string(copyoutWorkers) + + " blockDim=" + std::to_string(groupBlockDim)); uint32_t invocationId = 0U; for (int iter = 0; iter < warmup; ++iter, ++invocationId) { launch_tilexr_udma_all_to_all_group( - TileXR::Demo::kAllToAllGroupBlockDim, stream, commArgsDev, + groupBlockDim, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(registeredMemory), reinterpret_cast(debug), invocationId, elementsPerPeer, plan.chunkElements, plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], - plan.signalOffset[0], plan.signalOffset[1], nullptr, 0U); + plan.signalOffset[0], plan.signalOffset[1], nullptr, 0U, copyoutWorkers); } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", aclrtSynchronizeStream(stream))) { release(); @@ -784,14 +797,15 @@ bool RunGroupedAllToAll( const auto begin = std::chrono::steady_clock::now(); for (int iter = 0; iter < repeat; ++iter, ++invocationId) { launch_tilexr_udma_all_to_all_group( - TileXR::Demo::kAllToAllGroupBlockDim, stream, commArgsDev, + groupBlockDim, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(registeredMemory), reinterpret_cast(debug), invocationId, elementsPerPeer, plan.chunkElements, plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], - reinterpret_cast(groupTraceDevice), static_cast(iter)); + reinterpret_cast(groupTraceDevice), static_cast(iter), + copyoutWorkers); } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", aclrtSynchronizeStream(stream))) { release(); From fdeabc95e43a1b84d3a766537be8a21a4f660edb Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 21:04:14 +0800 Subject: [PATCH 084/163] docs(udma): design grouped route-stage diagnostics --- ...ma-grouped-alltoall-route-stages-design.md | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md new file mode 100644 index 00000000..f783dbd0 --- /dev/null +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md @@ -0,0 +1,102 @@ +# Grouped AllToAll Route-Stage Diagnostic Design + +## Goal + +Add an opt-in diagnostic mode that measures local traffic, the six-port route, +and the two-port route without overlap. The normal grouped AllToAll algorithm +and its default performance path remain unchanged. + +The first validation target is physical 2x8 with 128 MiB input/output per rank, +16 copyout workers, warmup 5, repeat 50, and a 60-second process timeout. + +## Stage Schedule + +Each logical AllToAll invocation is split into three kernel launches: + +1. `local`: self-copy and peers whose `rank / 8` matches the local rank's node. +2. `primary`: cross-node peers selected for the primary QP, currently QP0 and + the six-port aggregate route. +3. `secondary`: cross-node peers selected for the secondary QP, currently QP4 + and the two-port aggregate route. + +The Host synchronizes the stream after each launch and executes the existing +TCP all-rank barrier before launching the next stage. Therefore no rank can +start one route stage while another rank is still using the preceding route. + +The route-stage mode is enabled only through a new environment option. Without +that option, the Host performs the current single kernel launch and does not add +stage barriers. + +## Kernel Filtering + +The grouped kernel receives a route-stage selector. Send and receive workers use +the same predicate for every peer: + +- `local` accepts only same-node peers; receive workers also perform self-copy. +- `primary` accepts only cross-node peers whose selected QP is the primary QP. +- `secondary` accepts only cross-node peers whose selected QP is the secondary + QP. + +Peers rejected by the current stage are skipped before signal waiting or UDMA +submission. A peer belongs to exactly one stage, so existing payload slots, +signal slots, tokens, passes, groups, and ping-pong invocation slots remain +unchanged. No ACK or device-wide `SyncAll` is introduced. + +The node identity is `rank / 8`, consistent with the current `rankSize = N * 8` +scope. + +## Timing And Trace + +The Host records one ACL event duration for each stage. The reported staged +kernel time is the sum of the three device durations and excludes TCP barrier +latency. + +When tracing is enabled, each stage owns a separate existing-size trace buffer. +This avoids changing or multiplying the trace layout and prevents kernel-span +or task-span overwrites. Each rank writes three binary files with `local`, +`primary`, and `secondary` in their names. Existing trace conversion remains +compatible and can extract loop49 independently for each stage. + +For the physical 2x8 run, loop49 must contain: + +- `local`: 112 same-node sends plus 16 self-copy tasks across all ranks. +- `primary`: 96 cross-node sends on QP0. +- `secondary`: 32 cross-node sends on QP4. +- no peer duplicated across stages and no invalid QP. + +Bandwidth is calculated per source node from the selected stage only: + +```text +bandwidth = stage payload bytes / + (latest send-quiet end - earliest send-put-signal begin) +``` + +For each node, the primary stage transfers `8 ranks * 6 peers * 8 MiB = 384 +MiB`; the secondary stage transfers `8 ranks * 2 peers * 8 MiB = 128 MiB`. + +## Correctness And Failure Handling + +Output validation runs only after all three stages complete. Together the stages +must cover self, all same-node peers, and all cross-node peers exactly once. + +An invalid stage selector is rejected before UDMA work. A kernel error, stream +synchronization error, or TCP barrier failure aborts the staged run. Remote test +processes run under `timeout 60s` so a failed barrier or device wait cannot hold +the cards indefinitely. + +## Tests + +Host-side unit tests cover: + +- peer classification for local, primary, and secondary stages; +- exact and disjoint peer coverage for rank sizes 8, 16, and 128; +- unchanged default launch behavior; +- three launches and three barriers in staged mode; +- invalid environment values. + +Kernel source guards require symmetric send/receive filtering and continue to +forbid `SyncAll` in the grouped kernel. + +Physical verification uses `/home/pkg/b101/cann` on `141.61.50.31` and +`141.61.49.223`, runs correctness first, then extracts loop49 from each stage +and reports six-port and two-port bandwidth independently. From 2667576eed5244b0b3a538b45cbaffc384f79bc7 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 21:11:29 +0800 Subject: [PATCH 085/163] docs(udma): plan grouped route-stage diagnostics --- ...7-20-udma-grouped-alltoall-route-stages.md | 452 ++++++++++++++++++ 1 file changed, 452 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md new file mode 100644 index 00000000..10869dba --- /dev/null +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md @@ -0,0 +1,452 @@ +# Grouped AllToAll Route-Stage Diagnostics Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add an opt-in three-stage grouped AllToAll diagnostic that measures local, six-port, and two-port traffic without overlap, then validate physical 2x8 bandwidth at 128 MiB per rank. + +**Architecture:** Introduce a shared route-stage policy with a default combined mode and three diagnostic modes. In diagnostic mode the Host launches one filtered kernel per stage, synchronizes and executes the existing TCP all-rank barrier between stages, and writes one 8 MiB trace per stage; default execution remains one kernel with no added barrier. + +**Tech Stack:** C++14 Host runtime, AscendC AIV kernel, TileXR UDMA same-QP put-signal/quiet, ACL runtime events, Python trace converter, physical CANN `/home/pkg/b101/cann`. + +## Global Constraints + +- Supported rank sizes remain `N * 8`, from 8 through 128. +- The default path retains one kernel launch and no Host barrier. +- Diagnostic stage order is exactly `local`, `primary`, `secondary`. +- Payload and ready signal use the same QP followed by `UDMAQuietStatusOnQp`. +- Do not add ACK or `SyncAll` to the grouped kernel. +- Keep invocation tokens and two-plane ping-pong unchanged across stages. +- Each stage trace remains exactly 8 MiB. +- Physical runs use `TILEXR_IPC_PID_MODE=pid`, `/home/pkg/b101/cann`, warmup 5, + repeat 50, 16 copyout workers, and `timeout 60s`. + +--- + +### Task 1: Route-Stage Policy + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_route.h` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Produces: `enum class AllToAllGroupRouteStage : uint32_t` with `kCombined`, + `kLocal`, `kPrimary`, and `kSecondary`. +- Produces: `AllToAllGroupValidRouteStage(uint32_t) -> bool`. +- Produces: `AllToAllGroupPeerInRouteStage(int, int, AllToAllGroupRouteStage) -> bool`. + +- [ ] **Step 1: Write failing route-stage coverage tests** + +Add `TestRouteStages()` and call it from `main()`: + +```cpp +void TestRouteStages() +{ + using TileXR::Demo::AllToAllGroupRouteStage; + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(0U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(3U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(4U), false); + for (int rankSize : {8, 16, 128}) { + for (int rank = 0; rank < rankSize; ++rank) { + int local = 0; + int primary = 0; + int secondary = 0; + for (int peer = 0; peer < rankSize; ++peer) { + if (peer == rank) continue; + const bool l = TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kLocal); + const bool p = TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kPrimary); + const bool s = TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kSecondary); + CHECK_EQ(static_cast(l) + static_cast(p) + + static_cast(s), 1); + local += l; primary += p; secondary += s; + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kCombined), true); + } + CHECK_EQ(local, 7); + CHECK_EQ(primary + secondary, rankSize - 8); + CHECK_EQ(secondary, rankSize == 8 ? 0 : 2 * (rankSize / 8 - 1)); + } + } +} +``` + +- [ ] **Step 2: Run the test and verify RED** + +```bash +cmake --build tests/udma/build_b101 --target test_tilexr_udma_alltoall_group_layout -j +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +``` + +Expected: compile failure because the stage enum and helpers do not exist. + +- [ ] **Step 3: Implement the minimal shared policy** + +Add after the existing route constants: + +```cpp +enum class AllToAllGroupRouteStage : uint32_t { + kCombined = 0U, kLocal = 1U, kPrimary = 2U, kSecondary = 3U, +}; + +inline bool AllToAllGroupValidRouteStage(uint32_t value) +{ + return value <= static_cast(AllToAllGroupRouteStage::kSecondary); +} + +inline bool AllToAllGroupPeerInRouteStage( + int rank, int peer, AllToAllGroupRouteStage stage) +{ + if (rank < 0 || peer < 0 || rank == peer) return false; + const bool crossNode = AllToAllGroupIsCrossNode(rank, peer); + switch (stage) { + case AllToAllGroupRouteStage::kCombined: return true; + case AllToAllGroupRouteStage::kLocal: return !crossNode; + case AllToAllGroupRouteStage::kPrimary: + return crossNode && !AllToAllGroupUseSecondaryRoute(rank, peer); + case AllToAllGroupRouteStage::kSecondary: + return crossNode && AllToAllGroupUseSecondaryRoute(rank, peer); + } + return false; +} +``` + +- [ ] **Step 4: Run the test and verify GREEN** + +Run Step 2 again. Expected: grouped layout checks pass. + +- [ ] **Step 5: Commit** + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_route.h \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "feat(udma): classify grouped route stages" +``` + +--- + +### Task 2: Symmetric Kernel Stage Filtering + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Consumes: numeric stage values from Task 1. +- Extends: `launch_tilexr_udma_all_to_all_group(..., uint32_t copyoutWorkers, + uint32_t routeStage)`. +- Produces: `AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage)`. + +- [ ] **Step 1: Add failing kernel source guards** + +```cpp +CHECK_CONTAINS(kernel, "uint32_t copyoutWorkers, uint32_t routeStage"); +CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); +CHECK_CONTAINS(kernel, "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); +CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL"); +CHECK_NOT_CONTAINS(kernel, "SyncAll"); +``` + +- [ ] **Step 2: Run the layout test and verify RED** + +Run Task 1 Step 2. Expected: source guards fail. + +- [ ] **Step 3: Implement device validation and filtering** + +Add constants `0U..3U` matching Task 1. Reject values above 3 in the kernel +configuration check. Implement the predicate as: + +```cpp +if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED) return true; +if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL) return !crossNode; +if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY) + return crossNode && !AllToAllGroupUseSecondaryRouteDevice(rank, peer); +return crossNode && AllToAllGroupUseSecondaryRouteDevice(rank, peer); +``` + +Perform self-copy only for combined/local stages. Filter receive peers before +token calculation and waiting. Filter send peers before QP selection and UDMA. +Append `routeStage` to the kernel and wrapper signatures. + +- [ ] **Step 4: Build and verify GREEN** + +```bash +cmake --build tests/udma/build_b101 --target tilexr_udma_demo \ + test_tilexr_udma_alltoall_group_layout -j +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +``` + +Expected: build and test pass. + +- [ ] **Step 5: Commit** + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "feat(udma): filter grouped kernels by route stage" +``` + +--- + +### Task 3: Opt-In Host Scheduler And Independent Traces + +**Files:** +- Modify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` + +**Interfaces:** +- Consumes: extended launch wrapper from Task 2 and existing + `DemoBarrierAll(rank, rankSize, step)`. +- Produces: `TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES=0|1`. +- Produces: `tilexr_group_trace_{local|primary|secondary}_rank_.bin`. + +- [ ] **Step 1: Add failing Host structure tests** + +Require the option, exact stage names, ACL event timing, conditional barrier, +and default combined mode: + +```cpp +CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); +CHECK_CONTAINS(demo, "{\"local\", \"primary\", \"secondary\"}"); +CHECK_CONTAINS(demo, "aclrtEventElapsedTime"); +CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); +CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); +CHECK_CONTAINS(demo, "tilexr_group_trace_" + stageName + "_rank_"); +``` + +Replace the old unconditional grouped-function `DemoBarrierAll` prohibition +with checks that the `routeStages == false` branch launches `kCombined` +directly and does not call the staged launch helper. + +- [ ] **Step 2: Run the layout test and verify RED** + +Run Task 1 Step 2. Expected: Host structure checks fail. + +- [ ] **Step 3: Parse and validate the option** + +```cpp +const int routeStagesValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES", 0); +if (routeStagesValue != 0 && routeStagesValue != 1) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES must be 0 or 1, got " + << routeStagesValue << std::endl; + return false; +} +const bool routeStages = routeStagesValue == 1; +``` + +- [ ] **Step 4: Generalize trace ownership and filenames** + +Use one trace slot in default mode and three in staged mode. Each slot owns one +device pointer and one host vector initialized with the existing version-1 +header. Change `WriteGroupTraceBinary` to accept `stageName`; omit the suffix +for combined mode and include it for diagnostic modes. The release lambda must +free every allocated trace pointer on partial failure. + +- [ ] **Step 5: Implement stage launch and ACL timing** + +Create/reuse ACL start and end events around one kernel launch. Synchronize the +end event, call `aclrtEventElapsedTime`, accumulate milliseconds by stage, then +call `DemoBarrierAll` only in diagnostic mode. Use one `invocationId` for all +three stages of a logical iteration and increment it after secondary completes: + +```cpp +const AllToAllGroupRouteStage stages[] = { + AllToAllGroupRouteStage::kLocal, + AllToAllGroupRouteStage::kPrimary, + AllToAllGroupRouteStage::kSecondary, +}; +const char* stageNames[] = {"local", "primary", "secondary"}; +``` + +Warmup follows the same stage sequence and barriers with null traces. Default +mode remains one `kCombined` launch per invocation. + +- [ ] **Step 6: Report and validate** + +Print mean device microseconds for each stage and their sum. Copy/write all +stage traces after repeat50, then run the existing debug check and full output +validation exactly once. + +- [ ] **Step 7: Build and verify GREEN** + +```bash +cmake --build tests/udma/build_b101 --target tilexr_udma_demo \ + test_tilexr_udma_alltoall_group_layout -j +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +``` + +Expected: build and test pass. + +- [ ] **Step 8: Commit** + +```bash +git add tests/udma/demo/tilexr_udma_demo.cpp \ + tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "feat(udma): stage grouped route diagnostics" +``` + +--- + +### Task 4: Trace Compatibility And Regressions + +**Files:** +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py` +- Verify unchanged: `tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py` + +**Interfaces:** +- Consumes: one standard version-1 binary per stage. +- Produces: proof that suffixed traces convert independently. + +- [ ] **Step 1: Add a suffixed-file compatibility test** + +```python +def test_reads_suffixed_stage_trace(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilexr_group_trace_primary_rank_0.bin" + self.write_valid_trace(path) + trace = MODULE.read_rank_trace(path) + chrome = MODULE.build_chrome_trace([trace]) + self.assertEqual(trace["path"], str(path)) + self.assertIn(str(path), chrome["otherData"]["sources"]) +``` + +- [ ] **Step 2: Run the Python test** + +```bash +python -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +Expected: pass without converter changes. If the fixture hard-codes the old +filename, first observe that failure and then make only the fixture accept a +`path` argument. + +- [ ] **Step 3: Run all local regressions** + +```bash +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_layout +./tests/udma/install_b101/bin/test_tilexr_udma_transport_layout +python -m unittest tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +git diff --check +``` + +Expected: all tests pass and `git diff --check` is silent. + +- [ ] **Step 4: Commit** + +```bash +git add tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +git commit -m "test(udma): cover grouped stage trace filenames" +``` + +--- + +### Task 5: Bundle, Deploy, And Measure Physical 2x8 + +**Files:** +- Create ignored: `tmp/udma-grouped-alltoall-route-stages.bundle` +- Create ignored artifacts under: `tmp/grouped_alltoall_route_stages_b101_2x8/` + +**Interfaces:** +- Consumes: committed implementation HEAD and existing remote runner workflow. +- Produces: correctness logs, three loop49 traces, and isolated QP0/QP4 bandwidth. + +- [ ] **Step 1: Create and verify a complete bundle** + +```bash +git bundle create tmp/udma-grouped-alltoall-route-stages.bundle --all +git bundle verify tmp/udma-grouped-alltoall-route-stages.bundle +git bundle list-heads tmp/udma-grouped-alltoall-route-stages.bundle +``` + +Expected: verification succeeds and includes the current branch HEAD. + +- [ ] **Step 2: Upload and build on `.223`** + +Upload to both hosts, fetch into the dedicated validation checkout, reset that +validation branch to bundle HEAD, and build with: + +```bash +export TILEXR_CANN_HOME=/home/pkg/b101/cann +source scripts/common_env.sh +cmake --build tests/udma/build_b101 --target tilexr_udma_demo \ + test_tilexr_udma_alltoall_group_layout test_tilexr_udma_alltoall_layout \ + test_tilexr_udma_transport_layout -j +``` + +Expected: build exits 0. Verify demo and grouped-kernel SHA-256 match on both +hosts. + +- [ ] **Step 3: Run remote regressions** + +Run all three C++ tests and Python trace tests on `.223`. + +Expected: all pass before the physical workload starts. + +- [ ] **Step 4: Run staged physical 2x8** + +Use ranks 0-7 on `141.61.50.31` and ranks 8-15 on `141.61.49.223`: + +```bash +export TILEXR_CANN_HOME=/home/pkg/b101/cann +export TILEXR_IPC_PID_MODE=pid +export TILEXR_DEMO_ALLTOALL_WARMUP=5 +export TILEXR_DEMO_ALLTOALL_REPEAT=50 +export TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS=16 +export TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES=1 +export TILEXR_UDMA_GROUP_TRACE=1 +export TILEXR_UDMA_GROUP_TRACE_DIR=/home/h30059441/grouped_route_stages_b101_2x8 +timeout 60s ./tests/udma/install_b101/bin/tilexr_udma_demo \ + 16 "$rank" 8 2097152 8 0 +``` + +Expected: all ranks exit 0, validate output, and write three 8 MiB traces. + +- [ ] **Step 5: Extract and validate loop49** + +Convert each stage's 16 files independently to: + +```text +tmp/grouped_alltoall_route_stages_b101_2x8/local/loop49.json +tmp/grouped_alltoall_route_stages_b101_2x8/primary/loop49.json +tmp/grouped_alltoall_route_stages_b101_2x8/secondary/loop49.json +``` + +Require: + +```text +local same-node sends = 112, self-copy = 16 +primary cross-node QP0 sends = 96 +secondary cross-node QP4 sends = 32 +duplicate peers across stages = 0 +missing peers including self = 0 +``` + +- [ ] **Step 6: Calculate isolated bandwidth** + +```text +primary GiB/s = 384 MiB * 976.5625 / primary envelope us +secondary GiB/s = 128 MiB * 976.5625 / secondary envelope us +per-card GiB/s = node GiB/s / 8 +``` + +Report both directions and averages beside concurrent loop49 baselines: + +```text +QP0 concurrent = 731.72 / 859.54 GiB/s +QP4 concurrent = 294.95 / 286.71 GiB/s +``` + +- [ ] **Step 7: Final verification** + +```bash +git status --short --branch +git diff --check +git bundle verify tmp/udma-grouped-alltoall-route-stages.bundle +``` + +Expected: no unexpected tracked changes, no whitespace errors, and a verified +bundle containing implementation HEAD. From 6b0e9430c667a1eb0cf474456d9e397ac882853e Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 21:44:34 +0800 Subject: [PATCH 086/163] feat(udma): classify grouped route stages --- .../demo/tilexr_udma_alltoall_group_route.h | 32 ++++++++++++++++ ...test_tilexr_udma_alltoall_group_layout.cpp | 38 +++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index 785b80b6..be30cedf 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -14,6 +14,18 @@ namespace Demo { constexpr uint32_t kAllToAllGroupRanksPerNode = 8U; constexpr uint32_t kAllToAllGroupPrimaryPeersPerNode = 6U; +enum class AllToAllGroupRouteStage : uint32_t { + kCombined = 0U, + kLocal = 1U, + kPrimary = 2U, + kSecondary = 3U, +}; + +inline bool AllToAllGroupValidRouteStage(uint32_t value) +{ + return value <= static_cast(AllToAllGroupRouteStage::kSecondary); +} + struct AllToAllGroupRouteQps { uint32_t primaryQp = 0U; uint32_t secondaryQp = 0U; @@ -39,6 +51,26 @@ inline bool AllToAllGroupUseSecondaryRoute(int rank, int peer) kAllToAllGroupPrimaryPeersPerNode; } +inline bool AllToAllGroupPeerInRouteStage( + int rank, int peer, AllToAllGroupRouteStage stage) +{ + if (rank < 0 || peer < 0 || rank == peer) { + return false; + } + const bool crossNode = AllToAllGroupIsCrossNode(rank, peer); + switch (stage) { + case AllToAllGroupRouteStage::kCombined: + return true; + case AllToAllGroupRouteStage::kLocal: + return !crossNode; + case AllToAllGroupRouteStage::kPrimary: + return crossNode && !AllToAllGroupUseSecondaryRoute(rank, peer); + case AllToAllGroupRouteStage::kSecondary: + return crossNode && AllToAllGroupUseSecondaryRoute(rank, peer); + } + return false; +} + inline AllToAllGroupRouteQps AllToAllGroupSelectRouteQps( const uint32_t* weights, uint32_t qpCount) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index c8cf2add..177e919b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -205,6 +205,43 @@ void TestDualRouteQpWeights() CHECK_EQ(empty.secondaryQp, 0U); } +void TestRouteStages() +{ + using TileXR::Demo::AllToAllGroupRouteStage; + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(0U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(3U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(4U), false); + + for (int rankSize : {8, 16, 128}) { + for (int rank = 0; rank < rankSize; ++rank) { + int local = 0; + int primary = 0; + int secondary = 0; + for (int peer = 0; peer < rankSize; ++peer) { + if (peer == rank) { + continue; + } + const bool inLocal = TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kLocal); + const bool inPrimary = TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kPrimary); + const bool inSecondary = TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kSecondary); + CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary) + + static_cast(inSecondary), 1); + local += inLocal ? 1 : 0; + primary += inPrimary ? 1 : 0; + secondary += inSecondary ? 1 : 0; + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kCombined), true); + } + CHECK_EQ(local, 7); + CHECK_EQ(primary + secondary, rankSize - 8); + CHECK_EQ(secondary, rankSize == 8 ? 0 : 2 * (rankSize / 8 - 1)); + } + } +} + void TestCopyoutWorkerPolicy() { CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); @@ -286,6 +323,7 @@ int main() TestTokens(); TestDualRoutePeerPolicy(); TestDualRouteQpWeights(); + TestRouteStages(); TestCopyoutWorkerPolicy(); TestKernelStructure(); TestHostStructure(); From b1dc0da2e455a49261946add3009bd1006ce22a1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 21:47:05 +0800 Subject: [PATCH 087/163] feat(udma): filter grouped kernels by route stage --- .../tilexr_udma_alltoall_group_kernel.cpp | 44 +++++++++++++++++-- tests/udma/demo/tilexr_udma_demo.cpp | 10 +++-- ...test_tilexr_udma_alltoall_group_layout.cpp | 5 +++ 3 files changed, 52 insertions(+), 7 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 3e529437..3da88c62 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -20,6 +20,10 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_WORDS = 12U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED = 0U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL = 1U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY = 2U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY = 3U; struct AllToAllGroupDeviceError { uint32_t valid; @@ -81,6 +85,27 @@ __aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( TileXR::Demo::kAllToAllGroupPrimaryPeersPerNode; } +__aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( + int32_t rank, int32_t peer, uint32_t routeStage) +{ + if (rank < 0 || peer < 0 || rank == peer || + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY) { + return false; + } + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED) { + return true; + } + const bool crossNode = + rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) != + peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode); + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL) { + return !crossNode; + } + const bool secondary = AllToAllGroupUseSecondaryRouteDevice(rank, peer); + return routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY ? + crossNode && !secondary : crossNode && secondary; +} + __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) { @@ -260,7 +285,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, - GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t copyoutWorkers) + GM_ADDR groupTraceGM, uint32_t traceIteration, + uint32_t copyoutWorkers, uint32_t routeStage) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? @@ -275,6 +301,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rankSize = args->rankSize; if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || @@ -303,7 +330,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(elementsPerPeer) * worker / copyoutWorkers); const int32_t selfEnd = static_cast( static_cast(elementsPerPeer) * (worker + 1U) / copyoutWorkers); - if (selfEnd > selfBegin) { + if ((routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL) && + selfEnd > selfBegin) { const uint64_t selfCopyBegin = AllToAllGroupTraceCycle(groupTrace); auto selfSrc = reinterpret_cast<__gm__ uint8_t*>( input + static_cast(rank) * elementsPerPeer + selfBegin); @@ -330,6 +359,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (peer < 0) { continue; } + if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage)) { + continue; + } const uint32_t traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; @@ -391,6 +423,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (peer < 0) { continue; } + if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage)) { + continue; + } uint32_t primaryQp = 0U; uint32_t secondaryQp = 0U; AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp); @@ -449,11 +484,12 @@ void launch_tilexr_udma_all_to_all_group( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, - GM_ADDR groupTrace, uint32_t traceIteration, uint32_t copyoutWorkers) + GM_ADDR groupTrace, uint32_t traceIteration, + uint32_t copyoutWorkers, uint32_t routeStage) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, copyoutWorkers); + groupTrace, traceIteration, copyoutWorkers, routeStage); } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 463fb9ce..1451aee4 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -25,6 +25,7 @@ #include "tilexr_types.h" #include "tilexr_udma_allreduce_layout.h" #include "tilexr_udma_alltoall_group_layout.h" +#include "tilexr_udma_alltoall_group_route.h" #include "tilexr_udma_alltoall_group_trace.h" #include "tilexr_udma_alltoall_layout.h" #include "tilexr_udma_fullmesh_trace.h" @@ -57,7 +58,8 @@ extern void launch_tilexr_udma_all_to_all_group( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, - GM_ADDR groupTrace, uint32_t traceIteration, uint32_t copyoutWorkers); + GM_ADDR groupTrace, uint32_t traceIteration, + uint32_t copyoutWorkers, uint32_t routeStage); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -787,7 +789,8 @@ bool RunGroupedAllToAll( invocationId, elementsPerPeer, plan.chunkElements, plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], - plan.signalOffset[0], plan.signalOffset[1], nullptr, 0U, copyoutWorkers); + plan.signalOffset[0], plan.signalOffset[1], nullptr, 0U, copyoutWorkers, + static_cast(TileXR::Demo::AllToAllGroupRouteStage::kCombined)); } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", aclrtSynchronizeStream(stream))) { release(); @@ -805,7 +808,8 @@ bool RunGroupedAllToAll( plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], reinterpret_cast(groupTraceDevice), static_cast(iter), - copyoutWorkers); + copyoutWorkers, + static_cast(TileXR::Demo::AllToAllGroupRouteStage::kCombined)); } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", aclrtSynchronizeStream(stream))) { release(); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 177e919b..e9b3e4ee 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -279,6 +279,11 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); + CHECK_CONTAINS(kernel, "uint32_t copyoutWorkers, uint32_t routeStage"); + CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); + CHECK_CONTAINS(kernel, + "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL"); CHECK_CONTAINS(kernel, "AllToAllGroupCopyoutLaneDevice"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers"); CHECK_CONTAINS(kernel, "traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane"); From 2e2c22d822759b69efe054bf1fb68fbfc8ba9719 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 21:54:44 +0800 Subject: [PATCH 088/163] feat(udma): stage grouped route diagnostics --- tests/udma/demo/tilexr_udma_demo.cpp | 242 +++++++++++++----- ...test_tilexr_udma_alltoall_group_layout.cpp | 7 +- 2 files changed, 187 insertions(+), 62 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 1451aee4..2f8974f6 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -4,6 +4,7 @@ */ #include +#include #include #include #include @@ -240,9 +241,13 @@ bool WriteFullmeshTraceBinary(int rank, const std::string& directory, const std: return true; } -bool WriteGroupTraceBinary(int rank, const std::string& directory, const std::vector& data) +bool WriteGroupTraceBinary( + int rank, const std::string& directory, const std::string& stageName, + const std::vector& data) { - const std::string path = directory + "/tilexr_group_trace_rank_" + + const std::string path = stageName.empty() ? + directory + "/tilexr_group_trace_rank_" + std::to_string(rank) + ".bin" : + directory + "/tilexr_group_trace_" + stageName + "_rank_" + std::to_string(rank) + ".bin"; std::ofstream output(path, std::ios::binary | std::ios::trunc); if (!output.is_open()) { @@ -660,6 +665,25 @@ bool RunGroupedAllToAll( } const uint32_t copyoutWorkers = static_cast(copyoutWorkersValue); const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim(copyoutWorkers); + const int routeStagesValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES", 0); + if (routeStagesValue != 0 && routeStagesValue != 1) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES must be 0 or 1, got " + << routeStagesValue << std::endl; + return false; + } + const bool routeStages = routeStagesValue == 1; + constexpr size_t kRouteStageCount = 3U; + const std::array + stagedRouteStages {{ + TileXR::Demo::AllToAllGroupRouteStage::kLocal, + TileXR::Demo::AllToAllGroupRouteStage::kPrimary, + TileXR::Demo::AllToAllGroupRouteStage::kSecondary, + }}; + const std::array stageNames {{ + "local", "primary", "secondary" + }}; const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); const bool traceEnabled = GetEnvInt("TILEXR_UDMA_GROUP_TRACE", 0) != 0; @@ -687,8 +711,10 @@ bool RunGroupedAllToAll( int32_t* input = nullptr; int32_t* output = nullptr; void* registeredMemory = nullptr; - void* groupTraceDevice = nullptr; - std::vector hostGroupTrace; + std::array groupTraceDevices {{nullptr, nullptr, nullptr}}; + std::array, kRouteStageCount> hostGroupTraces; + aclrtEvent stageStartEvent = nullptr; + aclrtEvent stageEndEvent = nullptr; TileXRUDMAMemHandle handle = 0; bool registered = false; auto release = [&]() { @@ -701,9 +727,19 @@ bool RunGroupedAllToAll( aclrtFree(registeredMemory); registeredMemory = nullptr; } - if (groupTraceDevice != nullptr) { - aclrtFree(groupTraceDevice); - groupTraceDevice = nullptr; + if (stageEndEvent != nullptr) { + aclrtDestroyEvent(stageEndEvent); + stageEndEvent = nullptr; + } + if (stageStartEvent != nullptr) { + aclrtDestroyEvent(stageStartEvent); + stageStartEvent = nullptr; + } + for (void*& groupTraceDevice : groupTraceDevices) { + if (groupTraceDevice != nullptr) { + aclrtFree(groupTraceDevice); + groupTraceDevice = nullptr; + } } if (output != nullptr) { aclrtFree(output); @@ -730,28 +766,34 @@ bool RunGroupedAllToAll( } if (traceEnabled) { - hostGroupTrace.assign(TileXR::Demo::kAllToAllGroupTraceBytes, 0U); - TileXR::Demo::AllToAllGroupTraceHeader header {}; - header.magic = TileXR::Demo::kAllToAllGroupTraceMagic; - header.version = TileXR::Demo::kAllToAllGroupTraceVersion; - header.rank = static_cast(rank); - header.iterationCount = static_cast(repeat); - header.groupCount = plan.groupCount; - header.passCount = plan.passCount; - header.coreCount = TileXR::Demo::kAllToAllGroupTraceCoreCount; - header.phaseCount = TileXR::Demo::kAllToAllGroupTracePhaseCount; - header.cyclesPerUs = TileXR::Demo::kAllToAllGroupTraceCyclesPerUs; - header.traceBytes = TileXR::Demo::kAllToAllGroupTraceBytes; - header.kernelSpanOffset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes; - header.taskSpanOffset = TileXR::Demo::AllToAllGroupTraceTaskSpanBaseOffset(); - std::memcpy(hostGroupTrace.data(), &header, sizeof(header)); - if (!CheckAcl(rank, "aclrtMalloc grouped trace", - aclrtMalloc(&groupTraceDevice, TileXR::Demo::kAllToAllGroupTraceBytes, - ACL_MEM_MALLOC_HUGE_FIRST)) || - !CopyHostToDevice(rank, groupTraceDevice, TileXR::Demo::kAllToAllGroupTraceBytes, - hostGroupTrace.data(), hostGroupTrace.size(), "grouped trace")) { - release(); - return false; + const size_t traceCount = routeStages ? kRouteStageCount : 1U; + for (size_t traceIndex = 0U; traceIndex < traceCount; ++traceIndex) { + auto& hostGroupTrace = hostGroupTraces[traceIndex]; + hostGroupTrace.assign(TileXR::Demo::kAllToAllGroupTraceBytes, 0U); + TileXR::Demo::AllToAllGroupTraceHeader header {}; + header.magic = TileXR::Demo::kAllToAllGroupTraceMagic; + header.version = TileXR::Demo::kAllToAllGroupTraceVersion; + header.rank = static_cast(rank); + header.iterationCount = static_cast(repeat); + header.groupCount = plan.groupCount; + header.passCount = plan.passCount; + header.coreCount = TileXR::Demo::kAllToAllGroupTraceCoreCount; + header.phaseCount = TileXR::Demo::kAllToAllGroupTracePhaseCount; + header.cyclesPerUs = TileXR::Demo::kAllToAllGroupTraceCyclesPerUs; + header.traceBytes = TileXR::Demo::kAllToAllGroupTraceBytes; + header.kernelSpanOffset = TileXR::Demo::kAllToAllGroupTraceHeaderBytes; + header.taskSpanOffset = TileXR::Demo::AllToAllGroupTraceTaskSpanBaseOffset(); + std::memcpy(hostGroupTrace.data(), &header, sizeof(header)); + if (!CheckAcl(rank, "aclrtMalloc grouped trace", + aclrtMalloc(&groupTraceDevices[traceIndex], + TileXR::Demo::kAllToAllGroupTraceBytes, + ACL_MEM_MALLOC_HUGE_FIRST)) || + !CopyHostToDevice(rank, groupTraceDevices[traceIndex], + TileXR::Demo::kAllToAllGroupTraceBytes, + hostGroupTrace.data(), hostGroupTrace.size(), "grouped trace")) { + release(); + return false; + } } } @@ -763,6 +805,15 @@ bool RunGroupedAllToAll( } registered = true; + if (routeStages && + (!CheckAcl(rank, "aclrtCreateEvent grouped stage start", + aclrtCreateEvent(&stageStartEvent)) || + !CheckAcl(rank, "aclrtCreateEvent grouped stage end", + aclrtCreateEvent(&stageEndEvent)))) { + release(); + return false; + } + auto debug = reinterpret_cast( static_cast(registeredMemory) + plan.controlOffset); PrintStatus(rank, "grouped alltoall registeredBytes=" + std::to_string(plan.registeredBytes) + @@ -778,27 +829,12 @@ bool RunGroupedAllToAll( PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + " repeat=" + std::to_string(repeat) + " copyoutWorkers=" + std::to_string(copyoutWorkers) + - " blockDim=" + std::to_string(groupBlockDim)); + " blockDim=" + std::to_string(groupBlockDim) + + " routeStages=" + std::to_string(routeStagesValue)); uint32_t invocationId = 0U; - for (int iter = 0; iter < warmup; ++iter, ++invocationId) { - launch_tilexr_udma_all_to_all_group( - groupBlockDim, stream, commArgsDev, - reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(registeredMemory), reinterpret_cast(debug), - invocationId, elementsPerPeer, plan.chunkElements, - plan.passCount, plan.groupCount, - plan.payloadOffset[0], plan.payloadOffset[1], - plan.signalOffset[0], plan.signalOffset[1], nullptr, 0U, copyoutWorkers, - static_cast(TileXR::Demo::AllToAllGroupRouteStage::kCombined)); - } - if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", aclrtSynchronizeStream(stream))) { - release(); - return false; - } - - const auto begin = std::chrono::steady_clock::now(); - for (int iter = 0; iter < repeat; ++iter, ++invocationId) { + auto launchGroupStage = [&](TileXR::Demo::AllToAllGroupRouteStage routeStage, + void* trace, uint32_t traceIteration) { launch_tilexr_udma_all_to_all_group( groupBlockDim, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), @@ -807,15 +843,86 @@ bool RunGroupedAllToAll( plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], - reinterpret_cast(groupTraceDevice), static_cast(iter), - copyoutWorkers, - static_cast(TileXR::Demo::AllToAllGroupRouteStage::kCombined)); - } - if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", aclrtSynchronizeStream(stream))) { - release(); - return false; + reinterpret_cast(trace), traceIteration, copyoutWorkers, + static_cast(routeStage)); + }; + + double totalUs = 0.0; + std::array stageTotalUs {{0.0, 0.0, 0.0}}; + if (!routeStages) { + for (int iter = 0; iter < warmup; ++iter, ++invocationId) { + launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, nullptr, 0U); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", + aclrtSynchronizeStream(stream))) { + release(); + return false; + } + + const auto begin = std::chrono::steady_clock::now(); + for (int iter = 0; iter < repeat; ++iter, ++invocationId) { + launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, + groupTraceDevices[0], static_cast(iter)); + } + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", + aclrtSynchronizeStream(stream))) { + release(); + return false; + } + const auto end = std::chrono::steady_clock::now(); + totalUs = std::chrono::duration(end - begin).count(); + } else { + auto finishStage = [&](const std::string& barrierStep) -> bool { + return CheckAcl(rank, "aclrtSynchronizeStream " + barrierStep, + aclrtSynchronizeStream(stream)) && + DemoBarrierAll(rank, rankSize, barrierStep); + }; + + for (int iter = 0; iter < warmup; ++iter, ++invocationId) { + for (size_t stageIndex = 0U; stageIndex < kRouteStageCount; ++stageIndex) { + launchGroupStage(stagedRouteStages[stageIndex], nullptr, 0U); + const std::string barrierStep = "grouped route stage " + + std::string(stageNames[stageIndex]) + " warmup=" + std::to_string(iter); + if (!finishStage(barrierStep)) { + release(); + return false; + } + } + } + + for (int iter = 0; iter < repeat; ++iter, ++invocationId) { + for (size_t stageIndex = 0U; stageIndex < kRouteStageCount; ++stageIndex) { + if (!CheckAcl(rank, "aclrtRecordEvent grouped stage start", + aclrtRecordEvent(stageStartEvent, stream))) { + release(); + return false; + } + launchGroupStage(stagedRouteStages[stageIndex], + groupTraceDevices[stageIndex], static_cast(iter)); + if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", + aclrtRecordEvent(stageEndEvent, stream))) { + release(); + return false; + } + const std::string barrierStep = "grouped route stage " + + std::string(stageNames[stageIndex]) + " repeat=" + std::to_string(iter); + if (!finishStage(barrierStep)) { + release(); + return false; + } + float elapsedMs = 0.0F; + if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", + aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { + release(); + return false; + } + stageTotalUs[stageIndex] += static_cast(elapsedMs) * 1000.0; + } + } + for (double stageUs : stageTotalUs) { + totalUs += stageUs; + } } - const auto end = std::chrono::steady_clock::now(); std::vector hostDebug(kErrorWordsPerCore * kErrorCoreCount, 0); const size_t debugBytes = hostDebug.size() * sizeof(int32_t); @@ -824,10 +931,16 @@ bool RunGroupedAllToAll( CopyDeviceToHost(rank, hostDebug.data(), debugBytes, debug, debugBytes, "grouped alltoall debug"); if (traceEnabled) { - copyOk = CopyDeviceToHost( - rank, hostGroupTrace.data(), hostGroupTrace.size(), groupTraceDevice, - TileXR::Demo::kAllToAllGroupTraceBytes, "grouped trace") && - WriteGroupTraceBinary(rank, traceDir, hostGroupTrace) && copyOk; + const size_t traceCount = routeStages ? kRouteStageCount : 1U; + for (size_t traceIndex = 0U; traceIndex < traceCount; ++traceIndex) { + auto& hostGroupTrace = hostGroupTraces[traceIndex]; + const std::string stageName = routeStages ? stageNames[traceIndex] : ""; + copyOk = CopyDeviceToHost( + rank, hostGroupTrace.data(), hostGroupTrace.size(), + groupTraceDevices[traceIndex], TileXR::Demo::kAllToAllGroupTraceBytes, + "grouped trace " + stageName) && + WriteGroupTraceBinary(rank, traceDir, stageName, hostGroupTrace) && copyOk; + } } bool debugOk = true; for (uint32_t core = 0; core < kErrorCoreCount; ++core) { @@ -851,7 +964,6 @@ bool RunGroupedAllToAll( << " observed=" << observed << std::endl; } - const double totalUs = std::chrono::duration(end - begin).count(); const double perIterUs = totalUs / static_cast(repeat); const double bandwidthGbs = static_cast(dataBytes) / (perIterUs * 1.0e3); std::cout << "[rank " << rank << "] grouped alltoall " << repeat @@ -859,6 +971,14 @@ bool RunGroupedAllToAll( << " ms perIter=" << perIterUs << " us payload=" << dataBytes << " bytes bw=" << bandwidthGbs << " GB/s" << std::endl; + if (routeStages) { + for (size_t stageIndex = 0U; stageIndex < kRouteStageCount; ++stageIndex) { + std::cout << "[rank " << rank << "] grouped route stage " + << stageNames[stageIndex] << " mean=" + << stageTotalUs[stageIndex] / static_cast(repeat) + << " us" << std::endl; + } + } const bool valid = copyOk && debugOk && ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerPeer); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index e9b3e4ee..e1db36d7 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -308,6 +308,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_group"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); CHECK_CONTAINS(demo, "AllToAllGroupValidCopyoutWorkers"); CHECK_CONTAINS(demo, "AllToAllGroupBlockDim(copyoutWorkers)"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); @@ -316,7 +317,11 @@ void TestHostStructure() const size_t end = demo.find("void Cleanup(", begin); const std::string grouped = begin == std::string::npos ? std::string() : demo.substr(begin, end == std::string::npos ? std::string::npos : end - begin); - CHECK_NOT_CONTAINS(grouped, "DemoBarrierAll"); + CHECK_CONTAINS(grouped, "\"local\", \"primary\", \"secondary\""); + CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); + CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); + CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); + CHECK_CONTAINS(demo, "\"/tilexr_group_trace_\" + stageName + \"_rank_\""); } } // namespace From 00e3c61f3972302c8dec0b3888cf6dd5ebd137b7 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 21:55:52 +0800 Subject: [PATCH 089/163] test(udma): cover grouped stage trace filenames --- ...test_tilexr_udma_alltoall_group_trace_to_chrome.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index 25231a9e..d50dc1f0 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -95,6 +95,17 @@ def test_converts_all_grouped_pipeline_phases(self): self.assertEqual(trace["otherData"]["displayTimeUnit"], "ns") json.loads(json.dumps(trace)) + def test_reads_suffixed_stage_trace(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilexr_group_trace_primary_rank_0.bin" + self.make_trace(path) + + rank_trace = MODULE.read_rank_trace(path) + trace = MODULE.build_chrome_trace([rank_trace]) + + self.assertEqual(rank_trace["path"], str(path)) + self.assertIn(str(path), trace["otherData"]["sources"]) + def test_normalizes_ranks_independently(self): with tempfile.TemporaryDirectory() as directory: first = Path(directory) / "rank0.bin" From dd29ca9f7e88fb0c1add04c26341451cf1cafebd Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 22:06:07 +0800 Subject: [PATCH 090/163] fix(udma): synchronize grouped stage startup --- tests/udma/demo/tilexr_udma_demo.cpp | 6 ++++++ tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 2 ++ 2 files changed, 8 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 2f8974f6..47c1571e 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -832,6 +832,12 @@ bool RunGroupedAllToAll( " blockDim=" + std::to_string(groupBlockDim) + " routeStages=" + std::to_string(routeStagesValue)); + if (routeStages && + !DemoBarrierAll(rank, rankSize, "grouped route stages ready")) { + release(); + return false; + } + uint32_t invocationId = 0U; auto launchGroupStage = [&](TileXR::Demo::AllToAllGroupRouteStage routeStage, void* trace, uint32_t traceIteration) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index e1db36d7..62bfb2a2 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -321,6 +321,8 @@ void TestHostStructure() CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); + CHECK_CONTAINS(grouped, + "DemoBarrierAll(rank, rankSize, \"grouped route stages ready\")"); CHECK_CONTAINS(demo, "\"/tilexr_group_trace_\" + stageName + \"_rank_\""); } From cf6ba17758d91bad353e879b56601d485581b27b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 23:16:59 +0800 Subject: [PATCH 091/163] fix(udma): batch grouped route stages --- ...7-20-udma-grouped-alltoall-route-stages.md | 17 +++-- ...ma-grouped-alltoall-route-stages-design.md | 18 +++-- tests/udma/demo/tilexr_udma_demo.cpp | 73 ++++++++----------- ...test_tilexr_udma_alltoall_group_layout.cpp | 3 + 4 files changed, 55 insertions(+), 56 deletions(-) diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md index 10869dba..1681a174 100644 --- a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md @@ -4,7 +4,7 @@ **Goal:** Add an opt-in three-stage grouped AllToAll diagnostic that measures local, six-port, and two-port traffic without overlap, then validate physical 2x8 bandwidth at 128 MiB per rank. -**Architecture:** Introduce a shared route-stage policy with a default combined mode and three diagnostic modes. In diagnostic mode the Host launches one filtered kernel per stage, synchronizes and executes the existing TCP all-rank barrier between stages, and writes one 8 MiB trace per stage; default execution remains one kernel with no added barrier. +**Architecture:** Introduce a shared route-stage policy with a default combined mode and three diagnostic modes. In diagnostic mode the Host runs warmup and measured kernel batches for one filtered stage at a time, synchronizes and executes one TCP all-rank barrier between stage batches, and writes one 8 MiB trace per stage; default execution remains one kernel with no added barrier. **Tech Stack:** C++14 Host runtime, AscendC AIV kernel, TileXR UDMA same-QP put-signal/quiet, ACL runtime events, Python trace converter, physical CANN `/home/pkg/b101/cann`. @@ -245,12 +245,13 @@ header. Change `WriteGroupTraceBinary` to accept `stageName`; omit the suffix for combined mode and include it for diagnostic modes. The release lambda must free every allocated trace pointer on partial failure. -- [ ] **Step 5: Implement stage launch and ACL timing** +- [ ] **Step 5: Implement stage-batch launch and ACL timing** -Create/reuse ACL start and end events around one kernel launch. Synchronize the -end event, call `aclrtEventElapsedTime`, accumulate milliseconds by stage, then -call `DemoBarrierAll` only in diagnostic mode. Use one `invocationId` for all -three stages of a logical iteration and increment it after secondary completes: +Create/reuse ACL start and end events around each stage's complete repeat batch. +Run all warmups for that stage first, synchronize, then record the measured +batch. Call `aclrtEventElapsedTime`, divide by repeat count, and call +`DemoBarrierAll` once after the stage completes. Reset `invocationId` at the +start of each stage because peer sets are disjoint: ```cpp const AllToAllGroupRouteStage stages[] = { @@ -261,8 +262,8 @@ const AllToAllGroupRouteStage stages[] = { const char* stageNames[] = {"local", "primary", "secondary"}; ``` -Warmup follows the same stage sequence and barriers with null traces. Default -mode remains one `kCombined` launch per invocation. +Default mode remains one `kCombined` launch per invocation. Staged mode uses one +ready barrier and three completion barriers, not one barrier per invocation. - [ ] **Step 6: Report and validate** diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md index f783dbd0..d1b7eeef 100644 --- a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md @@ -11,7 +11,7 @@ The first validation target is physical 2x8 with 128 MiB input/output per rank, ## Stage Schedule -Each logical AllToAll invocation is split into three kernel launches: +The diagnostic run is split into three stage batches: 1. `local`: self-copy and peers whose `rank / 8` matches the local rank's node. 2. `primary`: cross-node peers selected for the primary QP, currently QP0 and @@ -19,9 +19,12 @@ Each logical AllToAll invocation is split into three kernel launches: 3. `secondary`: cross-node peers selected for the secondary QP, currently QP4 and the two-port aggregate route. -The Host synchronizes the stream after each launch and executes the existing -TCP all-rank barrier before launching the next stage. Therefore no rank can -start one route stage while another rank is still using the preceding route. +For each stage, the Host launches all warmup iterations and then all measured +iterations on the same stream. It synchronizes the stream and executes the +existing TCP all-rank barrier once after the measured batch. Therefore no rank +can start one route stage while another rank is still using the preceding +route. This needs only one ready barrier plus three stage-completion barriers, +so warmup5/repeat50 remains inside the 60-second process timeout. The route-stage mode is enabled only through a new environment option. Without that option, the Host performs the current single kernel launch and does not add @@ -47,9 +50,10 @@ scope. ## Timing And Trace -The Host records one ACL event duration for each stage. The reported staged -kernel time is the sum of the three device durations and excludes TCP barrier -latency. +The Host records one ACL event duration around each stage's complete measured +batch. The reported per-stage mean divides that duration by repeat count; the +reported staged kernel time is the sum of the three means and excludes TCP +barrier latency. When tracing is enabled, each stage owns a separate existing-size trace buffer. This avoids changing or multiplying the trace layout and prevents kernel-span diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 47c1571e..0da26cc6 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -878,51 +878,42 @@ bool RunGroupedAllToAll( const auto end = std::chrono::steady_clock::now(); totalUs = std::chrono::duration(end - begin).count(); } else { - auto finishStage = [&](const std::string& barrierStep) -> bool { - return CheckAcl(rank, "aclrtSynchronizeStream " + barrierStep, - aclrtSynchronizeStream(stream)) && - DemoBarrierAll(rank, rankSize, barrierStep); - }; - - for (int iter = 0; iter < warmup; ++iter, ++invocationId) { - for (size_t stageIndex = 0U; stageIndex < kRouteStageCount; ++stageIndex) { + auto runStageBatch = [&](size_t stageIndex) -> bool { + invocationId = 0U; + for (int iter = 0; iter < warmup; ++iter, ++invocationId) { launchGroupStage(stagedRouteStages[stageIndex], nullptr, 0U); - const std::string barrierStep = "grouped route stage " + - std::string(stageNames[stageIndex]) + " warmup=" + std::to_string(iter); - if (!finishStage(barrierStep)) { - release(); - return false; - } } - } - - for (int iter = 0; iter < repeat; ++iter, ++invocationId) { - for (size_t stageIndex = 0U; stageIndex < kRouteStageCount; ++stageIndex) { - if (!CheckAcl(rank, "aclrtRecordEvent grouped stage start", - aclrtRecordEvent(stageStartEvent, stream))) { - release(); - return false; - } + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped stage warmup", + aclrtSynchronizeStream(stream)) || + !CheckAcl(rank, "aclrtRecordEvent grouped stage start", + aclrtRecordEvent(stageStartEvent, stream))) { + return false; + } + for (int iter = 0; iter < repeat; ++iter, ++invocationId) { launchGroupStage(stagedRouteStages[stageIndex], groupTraceDevices[stageIndex], static_cast(iter)); - if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", - aclrtRecordEvent(stageEndEvent, stream))) { - release(); - return false; - } - const std::string barrierStep = "grouped route stage " + - std::string(stageNames[stageIndex]) + " repeat=" + std::to_string(iter); - if (!finishStage(barrierStep)) { - release(); - return false; - } - float elapsedMs = 0.0F; - if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", - aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { - release(); - return false; - } - stageTotalUs[stageIndex] += static_cast(elapsedMs) * 1000.0; + } + if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", + aclrtRecordEvent(stageEndEvent, stream)) || + !CheckAcl(rank, "aclrtSynchronizeStream grouped stage measured", + aclrtSynchronizeStream(stream))) { + return false; + } + float elapsedMs = 0.0F; + if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", + aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { + return false; + } + stageTotalUs[stageIndex] = static_cast(elapsedMs) * 1000.0; + const std::string barrierStep = "grouped route stage " + + std::string(stageNames[stageIndex]) + " complete"; + return DemoBarrierAll(rank, rankSize, barrierStep); + }; + + for (size_t stageIndex = 0U; stageIndex < kRouteStageCount; ++stageIndex) { + if (!runStageBatch(stageIndex)) { + release(); + return false; } } for (double stageUs : stageTotalUs) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 62bfb2a2..7c11c2b6 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -321,6 +321,9 @@ void TestHostStructure() CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); + CHECK_CONTAINS(grouped, "auto runStageBatch"); + CHECK_NOT_CONTAINS(grouped, "\" warmup=\""); + CHECK_CONTAINS(grouped, "\" complete\""); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, \"grouped route stages ready\")"); CHECK_CONTAINS(demo, "\"/tilexr_group_trace_\" + stageName + \"_rank_\""); From 52f447550e56346f9aaf0ecd191bc9114c5a2a19 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 23:24:03 +0800 Subject: [PATCH 092/163] fix(udma): synchronize staged ping-pong reuse --- ...7-20-udma-grouped-alltoall-route-stages.md | 11 +++--- ...ma-grouped-alltoall-route-stages-design.md | 15 ++++---- tests/udma/demo/tilexr_udma_demo.cpp | 38 ++++++++++--------- ...test_tilexr_udma_alltoall_group_layout.cpp | 1 + 4 files changed, 35 insertions(+), 30 deletions(-) diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md index 1681a174..8a8a5473 100644 --- a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md @@ -247,11 +247,12 @@ free every allocated trace pointer on partial failure. - [ ] **Step 5: Implement stage-batch launch and ACL timing** -Create/reuse ACL start and end events around each stage's complete repeat batch. -Run all warmups for that stage first, synchronize, then record the measured -batch. Call `aclrtEventElapsedTime`, divide by repeat count, and call -`DemoBarrierAll` once after the stage completes. Reset `invocationId` at the -start of each stage because peer sets are disjoint: +Create/reuse ACL start and end events around every measured invocation. Run all +warmups for that stage first and synchronize the local stream after each +invocation; do the same for measured invocations while accumulating event +durations. Divide by repeat count, and call `DemoBarrierAll` once after the +stage completes. Reset `invocationId` at the start of each stage because peer +sets are disjoint: ```cpp const AllToAllGroupRouteStage stages[] = { diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md index d1b7eeef..aaf95474 100644 --- a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md @@ -20,9 +20,10 @@ The diagnostic run is split into three stage batches: and the two-port aggregate route. For each stage, the Host launches all warmup iterations and then all measured -iterations on the same stream. It synchronizes the stream and executes the -existing TCP all-rank barrier once after the measured batch. Therefore no rank -can start one route stage while another rank is still using the preceding +iterations. It synchronizes the local stream after every invocation so a fast +rank cannot run more than one ping-pong generation ahead, but executes the +expensive TCP all-rank barrier only once after the measured batch. Therefore no +rank can start one route stage while another rank is still using the preceding route. This needs only one ready barrier plus three stage-completion barriers, so warmup5/repeat50 remains inside the 60-second process timeout. @@ -50,10 +51,10 @@ scope. ## Timing And Trace -The Host records one ACL event duration around each stage's complete measured -batch. The reported per-stage mean divides that duration by repeat count; the -reported staged kernel time is the sum of the three means and excludes TCP -barrier latency. +The Host records one ACL event duration per measured invocation and accumulates +the durations by stage. The reported per-stage mean divides that sum by repeat +count; the reported staged kernel time is the sum of the three means and +excludes stream synchronization and TCP barrier latency. When tracing is enabled, each stage owns a separate existing-size trace buffer. This avoids changing or multiplying the trace layout and prevents kernel-span diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 0da26cc6..d0c6768c 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -882,29 +882,31 @@ bool RunGroupedAllToAll( invocationId = 0U; for (int iter = 0; iter < warmup; ++iter, ++invocationId) { launchGroupStage(stagedRouteStages[stageIndex], nullptr, 0U); - } - if (!CheckAcl(rank, "aclrtSynchronizeStream grouped stage warmup", - aclrtSynchronizeStream(stream)) || - !CheckAcl(rank, "aclrtRecordEvent grouped stage start", - aclrtRecordEvent(stageStartEvent, stream))) { - return false; + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped stage iteration", + aclrtSynchronizeStream(stream))) { + return false; + } } for (int iter = 0; iter < repeat; ++iter, ++invocationId) { + if (!CheckAcl(rank, "aclrtRecordEvent grouped stage start", + aclrtRecordEvent(stageStartEvent, stream))) { + return false; + } launchGroupStage(stagedRouteStages[stageIndex], groupTraceDevices[stageIndex], static_cast(iter)); + if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", + aclrtRecordEvent(stageEndEvent, stream)) || + !CheckAcl(rank, "aclrtSynchronizeStream grouped stage iteration", + aclrtSynchronizeStream(stream))) { + return false; + } + float elapsedMs = 0.0F; + if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", + aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { + return false; + } + stageTotalUs[stageIndex] += static_cast(elapsedMs) * 1000.0; } - if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", - aclrtRecordEvent(stageEndEvent, stream)) || - !CheckAcl(rank, "aclrtSynchronizeStream grouped stage measured", - aclrtSynchronizeStream(stream))) { - return false; - } - float elapsedMs = 0.0F; - if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", - aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { - return false; - } - stageTotalUs[stageIndex] = static_cast(elapsedMs) * 1000.0; const std::string barrierStep = "grouped route stage " + std::string(stageNames[stageIndex]) + " complete"; return DemoBarrierAll(rank, rankSize, barrierStep); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 7c11c2b6..dc38d508 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -322,6 +322,7 @@ void TestHostStructure() CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); CHECK_CONTAINS(grouped, "auto runStageBatch"); + CHECK_CONTAINS(grouped, "aclrtSynchronizeStream grouped stage iteration"); CHECK_NOT_CONTAINS(grouped, "\" warmup=\""); CHECK_CONTAINS(grouped, "\" complete\""); CHECK_CONTAINS(grouped, From df8ee4b7209fd5cdb9a1562552ae1ed294f4499c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 20 Jul 2026 23:33:43 +0800 Subject: [PATCH 093/163] docs(udma): correct staged self-copy trace count --- .../plans/2026-07-20-udma-grouped-alltoall-route-stages.md | 2 +- .../2026-07-20-udma-grouped-alltoall-route-stages-design.md | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md index 8a8a5473..9ce8dcbf 100644 --- a/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md +++ b/docs/superpowers/plans/2026-07-20-udma-grouped-alltoall-route-stages.md @@ -420,7 +420,7 @@ tmp/grouped_alltoall_route_stages_b101_2x8/secondary/loop49.json Require: ```text -local same-node sends = 112, self-copy = 16 +local same-node sends = 112, self-copy task spans = 256 primary cross-node QP0 sends = 96 secondary cross-node QP4 sends = 32 duplicate peers across stages = 0 diff --git a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md index aaf95474..af29e18e 100644 --- a/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md +++ b/docs/superpowers/specs/2026-07-20-udma-grouped-alltoall-route-stages-design.md @@ -64,7 +64,8 @@ compatible and can extract loop49 independently for each stage. For the physical 2x8 run, loop49 must contain: -- `local`: 112 same-node sends plus 16 self-copy tasks across all ranks. +- `local`: 112 same-node sends plus 256 self-copy task spans across all ranks + (16 copyout workers on each of 16 ranks). - `primary`: 96 cross-node sends on QP0. - `secondary`: 32 cross-node sends on QP4. - no peer duplicated across stages and no invalid QP. From abcf96120a3a26448d8345099ed83e0d471729cd Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 10:04:00 +0800 Subject: [PATCH 094/163] perf(udma): split local grouped stage profiling --- .../tilexr_udma_alltoall_group_kernel.cpp | 69 ++++++++++++++----- .../demo/tilexr_udma_alltoall_group_route.h | 22 +++++- tests/udma/demo/tilexr_udma_demo.cpp | 11 +-- ...test_tilexr_udma_alltoall_group_layout.cpp | 24 ++++++- ...tilexr_udma_alltoall_group_local_stages.py | 25 +++++++ 5 files changed, 124 insertions(+), 27 deletions(-) create mode 100644 tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 3da88c62..ed8145ba 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -24,6 +24,8 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED = 0U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY = 3U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND = 4U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY = 5U; struct AllToAllGroupDeviceError { uint32_t valid; @@ -89,7 +91,7 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { if (rank < 0 || peer < 0 || rank == peer || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY) { + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) { return false; } if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED) { @@ -98,7 +100,9 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( const bool crossNode = rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) != peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode); - if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL) { + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) { return !crossNode; } const bool secondary = AllToAllGroupUseSecondaryRouteDevice(rank, peer); @@ -106,6 +110,22 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( crossNode && !secondary : crossNode && secondary; } +__aicore__ inline bool AllToAllGroupStageRunsSendDevice(uint32_t routeStage) +{ + return routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY; +} + +__aicore__ inline bool AllToAllGroupStageRunsCopyDevice(uint32_t routeStage) +{ + return routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND; +} + +__aicore__ inline bool AllToAllGroupStageWaitsForSignalDevice(uint32_t routeStage) +{ + return AllToAllGroupStageRunsCopyDevice(routeStage) && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY; +} + __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) { @@ -301,7 +321,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rankSize = args->rankSize; if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY || + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || @@ -325,13 +345,19 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(elementsPerPeer) * sizeof(int32_t); if (blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES) { + if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; const int32_t selfBegin = static_cast( static_cast(elementsPerPeer) * worker / copyoutWorkers); const int32_t selfEnd = static_cast( static_cast(elementsPerPeer) * (worker + 1U) / copyoutWorkers); if ((routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED || - routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL) && + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) && selfEnd > selfBegin) { const uint64_t selfCopyBegin = AllToAllGroupTraceCycle(groupTrace); auto selfSrc = reinterpret_cast<__gm__ uint8_t*>( @@ -377,26 +403,28 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( auto signal = reinterpret_cast<__gm__ uint64_t*>( registeredMemory + signalOffsets[slot] + static_cast(peer) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE); - uint64_t observed = 0ULL; - const uint64_t waitBegin = AllToAllGroupTraceCycle(groupTrace); - if (!AllToAllGroupWaitTokenMte(signal, expectedToken, - TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { + if (AllToAllGroupStageWaitsForSignalDevice(routeStage)) { + uint64_t observed = 0ULL; + const uint64_t waitBegin = AllToAllGroupTraceCycle(groupTrace); + if (!AllToAllGroupWaitTokenMte(signal, expectedToken, + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, + peer, TileXR::Demo::kAllToAllGroupTraceNoQp, + waitBegin, AllToAllGroupTraceCycle(groupTrace)); + AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_WAIT, + group, pass, peer, 0U, 0U, expectedToken, observed); + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, waitBegin, AllToAllGroupTraceCycle(groupTrace)); - AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_WAIT, - group, pass, peer, 0U, 0U, expectedToken, observed); - AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, - kernelBegin, AllToAllGroupTraceCycle(groupTrace)); - return; } - AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, traceCore, group, pass, - TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, - peer, TileXR::Demo::kAllToAllGroupTraceNoQp, - waitBegin, AllToAllGroupTraceCycle(groupTrace)); auto relaySrc = registeredMemory + payloadOffsets[slot] + static_cast(peer) * bytesPerPeer + static_cast(chunkElementOffset) * sizeof(int32_t); @@ -417,6 +445,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( return; } + if (!AllToAllGroupStageRunsSendDevice(routeStage)) { + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } const uint32_t lane = blockIdx; for (uint32_t group = 0U; group < groupCount; ++group) { const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index be30cedf..16a9ac82 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -19,11 +19,29 @@ enum class AllToAllGroupRouteStage : uint32_t { kLocal = 1U, kPrimary = 2U, kSecondary = 3U, + kLocalSend = 4U, + kLocalCopy = 5U, }; inline bool AllToAllGroupValidRouteStage(uint32_t value) { - return value <= static_cast(AllToAllGroupRouteStage::kSecondary); + return value <= static_cast(AllToAllGroupRouteStage::kLocalCopy); +} + +inline bool AllToAllGroupStageRunsSend(AllToAllGroupRouteStage stage) +{ + return stage != AllToAllGroupRouteStage::kLocalCopy; +} + +inline bool AllToAllGroupStageRunsCopy(AllToAllGroupRouteStage stage) +{ + return stage != AllToAllGroupRouteStage::kLocalSend; +} + +inline bool AllToAllGroupStageWaitsForSignal(AllToAllGroupRouteStage stage) +{ + return AllToAllGroupStageRunsCopy(stage) && + stage != AllToAllGroupRouteStage::kLocalCopy; } struct AllToAllGroupRouteQps { @@ -62,6 +80,8 @@ inline bool AllToAllGroupPeerInRouteStage( case AllToAllGroupRouteStage::kCombined: return true; case AllToAllGroupRouteStage::kLocal: + case AllToAllGroupRouteStage::kLocalSend: + case AllToAllGroupRouteStage::kLocalCopy: return !crossNode; case AllToAllGroupRouteStage::kPrimary: return crossNode && !AllToAllGroupUseSecondaryRoute(rank, peer); diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index d0c6768c..05ca31c9 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -674,15 +674,16 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; - constexpr size_t kRouteStageCount = 3U; + constexpr size_t kRouteStageCount = 4U; const std::array stagedRouteStages {{ - TileXR::Demo::AllToAllGroupRouteStage::kLocal, + TileXR::Demo::AllToAllGroupRouteStage::kLocalSend, + TileXR::Demo::AllToAllGroupRouteStage::kLocalCopy, TileXR::Demo::AllToAllGroupRouteStage::kPrimary, TileXR::Demo::AllToAllGroupRouteStage::kSecondary, }}; const std::array stageNames {{ - "local", "primary", "secondary" + "local-send", "local-copy", "primary", "secondary" }}; const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); @@ -711,7 +712,7 @@ bool RunGroupedAllToAll( int32_t* input = nullptr; int32_t* output = nullptr; void* registeredMemory = nullptr; - std::array groupTraceDevices {{nullptr, nullptr, nullptr}}; + std::array groupTraceDevices {}; std::array, kRouteStageCount> hostGroupTraces; aclrtEvent stageStartEvent = nullptr; aclrtEvent stageEndEvent = nullptr; @@ -854,7 +855,7 @@ bool RunGroupedAllToAll( }; double totalUs = 0.0; - std::array stageTotalUs {{0.0, 0.0, 0.0}}; + std::array stageTotalUs {}; if (!routeStages) { for (int iter = 0; iter < warmup; ++iter, ++invocationId) { launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, nullptr, 0U); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index dc38d508..282de57b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -209,8 +209,18 @@ void TestRouteStages() { using TileXR::Demo::AllToAllGroupRouteStage; CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(0U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(3U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(4U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(5U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(6U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( + AllToAllGroupRouteStage::kLocalSend), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( + AllToAllGroupRouteStage::kLocalSend), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( + AllToAllGroupRouteStage::kLocalCopy), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( + AllToAllGroupRouteStage::kLocalCopy), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageWaitsForSignal( + AllToAllGroupRouteStage::kLocalCopy), false); for (int rankSize : {8, 16, 128}) { for (int rank = 0; rank < rankSize; ++rank) { @@ -227,6 +237,10 @@ void TestRouteStages() rank, peer, AllToAllGroupRouteStage::kPrimary); const bool inSecondary = TileXR::Demo::AllToAllGroupPeerInRouteStage( rank, peer, AllToAllGroupRouteStage::kSecondary); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kLocalSend), inLocal); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kLocalCopy), inLocal); CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary) + static_cast(inSecondary), 1); local += inLocal ? 1 : 0; @@ -290,6 +304,9 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); + CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); + CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsCopyDevice(routeStage)"); + CHECK_CONTAINS(kernel, "AllToAllGroupStageWaitsForSignalDevice(routeStage)"); CHECK_CONTAINS(kernel, "observed >= expectedToken"); CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_group"); CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); @@ -317,7 +334,8 @@ void TestHostStructure() const size_t end = demo.find("void Cleanup(", begin); const std::string grouped = begin == std::string::npos ? std::string() : demo.substr(begin, end == std::string::npos ? std::string::npos : end - begin); - CHECK_CONTAINS(grouped, "\"local\", \"primary\", \"secondary\""); + CHECK_CONTAINS(grouped, + "\"local-send\", \"local-copy\", \"primary\", \"secondary\""); CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py new file mode 100644 index 00000000..53061c29 --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -0,0 +1,25 @@ +import pathlib +import unittest + + +ROOT = pathlib.Path(__file__).resolve().parents[3] + + +class GroupedAllToAllLocalStageSourceTest(unittest.TestCase): + def test_local_route_is_split_into_send_and_copy_stages(self): + route = (ROOT / "tests/udma/demo/tilexr_udma_alltoall_group_route.h").read_text() + kernel = (ROOT / "tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp").read_text() + host = (ROOT / "tests/udma/demo/tilexr_udma_demo.cpp").read_text() + + self.assertIn("kLocalSend", route) + self.assertIn("kLocalCopy", route) + self.assertIn("AllToAllGroupStageRunsSendDevice(routeStage)", kernel) + self.assertIn("AllToAllGroupStageRunsCopyDevice(routeStage)", kernel) + self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) + self.assertIn( + '"local-send", "local-copy", "primary", "secondary"', host + ) + + +if __name__ == "__main__": + unittest.main() From 5ec49264aaaa32ffb4efe8a470f0467bbe375da4 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 11:09:51 +0800 Subject: [PATCH 095/163] docs(udma): add grouped alltoall experiment guide --- ...udma-grouped-alltoall-modify-test-guide.md | 489 ++++++++++++++++++ 1 file changed, 489 insertions(+) create mode 100644 docs/udma-grouped-alltoall-modify-test-guide.md diff --git a/docs/udma-grouped-alltoall-modify-test-guide.md b/docs/udma-grouped-alltoall-modify-test-guide.md new file mode 100644 index 00000000..97328556 --- /dev/null +++ b/docs/udma-grouped-alltoall-modify-test-guide.md @@ -0,0 +1,489 @@ +# UDMA Grouped AllToAll 修改与测试指南 + +本文供新的 Codex 会话接手 `grouped AllToAll` 修改、部署、物理 `2x8` 测试和 +trace 分析时使用。目标是让一次实验只改变一个变量，并保留可复现的提交、bundle、 +日志和 trace。 + +## 1. 新会话先做什么 + +在仓库根目录开始，依次读取： + +```text +AGENTS.md +CLAUDE.md +docs/udma-grouped-alltoall-modify-test-guide.md +``` + +然后确认工作区和基线： + +```powershell +git status --short --branch +git log -5 --oneline +git rev-parse HEAD +``` + +不要清理用户已有的未跟踪文件，不要覆盖远端未提交修改。修改前记录： + +- 本地分支和完整 commit ID； +- CANN 路径； +- 两台机器和 rank 分配； +- 输入大小、chunk、warmup、repeat； +- route policy、QP 数、copyout worker 数； +- combined 或 staged 模式。 + +## 2. 当前测试拓扑 + +| 项目 | 值 | +|---|---| +| rank 0..7 | `141.61.50.31` | +| rank 8..15 | `141.61.49.223` | +| 构建机 | `141.61.49.223` | +| 远端目录 | `/home/h30059441/tilexr_grouped_alltoall_b101` | +| CANN | `/home/pkg/b101/cann` | +| rank size | 16，物理 `2x8` | +| 每 peer payload | 8 MiB | +| 每 rank input/output | 128 MiB | +| 正式测试 | warmup 5，repeat 50 | +| 单个 rank 超时 | `timeout 60s` | + +不要把 SSH 密码写入仓库、脚本或本文。优先使用 SSH key；否则由操作者在会话中提供。 + +## 3. 算法和内存约束 + +### 3.1 数据流 + +每个 peer 的基本流程是： + +```text +sender input slice + -> UDMA payload PUT + ready signal + -> same-peer/same-QP Quiet + +receiver wait ready signal + -> registered receive buffer + -> local MTE copy + -> output slice +``` + +第一版协议没有 ACK，也没有 device `SyncAll`。一次 kernel invocation 执行一次 +AllToAll；registered receive buffer 使用两个完整 payload plane 做 ping-pong。 + +必须保持以下正确性约束： + +1. `rankSize` 当前只支持 `N * 8`，范围为 8..128。 +2. 每个 source rank 在接收区有独立 slot，不循环复用 peer slot。 +3. payload 和 ready 必须落在同一目标 rank 的 registered GM。 +4. 单 QP 模式下，payload、ready 和 Quiet 必须使用同一个 QP。 +5. ready 可见时，对应 payload 必须已经对接收端 GM 可见。 +6. ping-pong slot 由 `invocationId & 1` 选择；复用前必须保证上次 invocation 完成。 +7. 修改 route 过滤时，send 和 receive 必须使用完全对称的 peer 集合。 +8. 不要为了性能删除超时、Quiet 状态检查或最终输出正确性检查。 + +### 3.2 核、peer 和 Jetty + +发送侧固定使用 core 0..15。每个 group lane 计算一个 peer。对于 `rankSize=16`， +只有一个 group，其中 15 个 lane 有效，重复的对径 peer lane 无效。 + +以 rank0 为例： + +```text +core0..7 -> peer1..8 +core8 -> peer15 +core9 -> peer14 +... +core14 -> peer9 +core15 -> invalid +``` + +Jetty/WQ 由 `(peer, qpIdx)` 唯一确定： + +```text +queue slot = peer * qpNum + qpIdx +``` + +因此两个 core 即使都记录为 QP0，只要 peer 不同，就不是同一个 Jetty。当前 +`TILEXR_UDMA_QP_NUM=4`、两条跨节点 route 时，每个跨节点 peer 创建 8 个 Jetty： + +```text +primary route: QP0..3 +secondary route: QP4..7 +``` + +当前 kernel 对一个 peer 只选一个 Jetty：primary 使用 QP0，secondary 使用 QP4。 +每张卡对远端节点的 8 个 peer 按 6:2 分配，即 6 个 peer-specific QP0 Jetty 和 +2 个 peer-specific QP4 Jetty。 + +### 3.3 route stage + +默认 `TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES=0`，完整流程一次执行，产生一个 +combined trace。 + +诊断模式设为 1 后依次运行四个独立 stage： + +| stage | 内容 | +|---|---| +| `local-send` | 节点内 UDMA send + Quiet，不执行 copyout | +| `local-copy` | 节点内 registered buffer 到 output，不等待 signal | +| `primary` | 跨节点 primary/6口 send、wait、copyout | +| `secondary` | 跨节点 secondary/2口 send、wait、copyout | + +staged 模式用于定位瓶颈。它的 stage 耗时之和不是 route 同时工作时的端到端带宽。 + +## 4. 关键代码 + +| 文件 | 作用 | +|---|---| +| `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` | 32核 grouped kernel、send/wait/copyout、trace 打点 | +| `tests/udma/demo/tilexr_udma_alltoall_group_layout.h` | group、peer、ping-pong 和 registered memory 布局 | +| `tests/udma/demo/tilexr_udma_alltoall_group_route.h` | local/primary/secondary 分类与 QP 选择 | +| `tests/udma/demo/tilexr_udma_alltoall_group_trace.h` | 8 MiB trace 布局 | +| `tests/udma/demo/tilexr_udma_demo.cpp` | Host 分配、注册、warmup/repeat、staged 调度、校验 | +| `src/include/tilexr_udma.h` | device UDMA PUT、PUT+signal、Quiet API | +| `src/comm/udma/tilexr_udma_transport.cpp` | EID、route、Jetty/QP 创建和 device image | +| `tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py` | raw trace 转 Chrome trace | + +优先扩展现有 helper，不要在 kernel 和 Host 各复制一份 route 或 peer 规则。 + +## 5. 修改和本地验证 + +一次实验只改变一个因素，例如： + +- `copyoutWorkers=8/16`； +- combined 与 staged； +- 单 peer 单 Jetty与多 Jetty； +- local/primary/secondary 的调度方式； +- chunk 或 pass 数。 + +修改后至少运行： + +```powershell +git diff --check + +# 如果已有本地构建产物 +.\tmp\test_group_stage_iteration.exe + +python -m unittest ` + tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +Linux 构建机还要运行： + +```bash +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_group_layout +./tests/udma/install_b101/bin/test_tilexr_udma_alltoall_layout +./tests/udma/install_b101/bin/test_tilexr_udma_transport_layout +python3 -m unittest \ + tests.udma.unit.test_tilexr_udma_alltoall_group_trace_to_chrome -v +``` + +若修改 Jetty/QP 映射，增加聚焦单测，至少覆盖： + +- 同一个 peer 的 QP 集合； +- 不同 peer 不共享 queue slot； +- primary 与 secondary 不串 route； +- 末尾不足16个 peer 的 group； +- 对径 peer 只出现一次； +- send/receive stage 过滤对称。 + +## 6. 提交、bundle 和部署 + +用户的既定流程是：先提交，再生成完整 bundle，再部署验证。 + +```powershell +git status --short +git diff --check +git add <本次修改文件> +git commit -m "perf(udma): <本次单变量实验>" + +$bundle = "tmp/udma-grouped-alltoall-$(git rev-parse --short HEAD).bundle" +git bundle create $bundle HEAD +git bundle verify $bundle +Get-FileHash $bundle -Algorithm SHA256 + +scp $bundle root@141.61.50.31:/tmp/ +scp $bundle root@141.61.49.223:/tmp/ +``` + +两台机器更新前先检查远端工作区。若有未提交修改，不要覆盖： + +```bash +cd /home/h30059441/tilexr_grouped_alltoall_b101 +git status --short +``` + +工作区干净时，在两台机器分别执行： + +```bash +bundle=/tmp/udma-grouped-alltoall-.bundle +cd /home/h30059441/tilexr_grouped_alltoall_b101 +git fetch "$bundle" HEAD +git checkout --detach FETCH_HEAD +git rev-parse HEAD +``` + +两台机器的 HEAD 必须与本地完整 commit ID 一致。 + +## 7. b101 构建 + +`141.61.50.31` 当前不负责构建。在 `141.61.49.223` 构建： + +```bash +cd /home/h30059441/tilexr_grouped_alltoall_b101 +source scripts/common_env.sh +source /home/pkg/b101/cann/set_env.sh + +export ASCEND_HOME_PATH=/home/pkg/b101/cann +export ASCEND_HOME_DIR=/home/pkg/b101/cann +export ASCEND_TOOLKIT_HOME=/home/pkg/b101/cann +export PATH=/home/pkg/b101/cann/bin:$PATH +export LD_LIBRARY_PATH=/home/pkg/b101/cann/aarch64-linux/lib64:\ +/usr/local/Ascend/driver/lib64/driver:\ +/usr/local/Ascend/driver/lib64/common:\ +/usr/local/Ascend/driver/lib64:${LD_LIBRARY_PATH:-} + +cmake -S . -B build -DCMAKE_INSTALL_PREFIX=$PWD/install +cmake --build build -j8 +cmake --install build + +cmake -S tests/udma -B tests/udma/build_b101 \ + -DCMAKE_INSTALL_PREFIX=$PWD/tests/udma/install_b101 \ + -DTILEXR_UDMA_DEMO_SOC_TYPE=Ascend950 +cmake --build tests/udma/build_b101 -j8 +cmake --install tests/udma/build_b101 +``` + +不要让运行时加载 CANN `devlib` 中的 stub `libascend_hal.so`： + +```bash +readelf -d build/src/comm/libtile-comm.so | grep -E 'RPATH|RUNPATH' || true +ldd tests/udma/install_b101/bin/tilexr_udma_demo | \ + grep -E 'lib(tile-comm|ascendcl|runtime|ascend_hal)' +``` + +将相同产物复制到 `141.61.50.31` 对应路径： + +```text +install/lib64/libtile-comm.so +tests/udma/install_b101/lib/libtilexr_udma_demo_kernel.so +tests/udma/install_b101/bin/tilexr_udma_demo +``` + +复制后在两台机器执行 `sha256sum`，三个文件必须逐个一致。 + +## 8. 物理 2x8 测试 + +### 8.1 公共环境 + +两台机器使用相同变量。每轮必须更换 `TILEXR_COMM_ID` 端口和结果目录： + +```bash +cd /home/h30059441/tilexr_grouped_alltoall_b101 +source /home/pkg/b101/cann/set_env.sh + +export ASCEND_HOME_PATH=/home/pkg/b101/cann +export PATH=/home/pkg/b101/cann/bin:$PATH +export LD_LIBRARY_PATH=$PWD/tests/udma/install_b101/lib:$PWD/install/lib64:\ +/home/pkg/b101/cann/aarch64-linux/lib64:\ +/usr/local/Ascend/driver/lib64/driver:\ +/usr/local/Ascend/driver/lib64/common:\ +/usr/local/Ascend/driver/lib64 + +export TILEXR_COMM_ID=141.61.50.31: +export TILEXR_DEMO_BARRIER_HOST=141.61.50.31 +export TILEXR_IPC_PID_MODE=pid +export TILEXR_UDMA_ROUTE_POLICY=all +export TILEXR_UDMA_QP_NUM=4 + +export TILEXR_DEMO_ALLTOALL_WARMUP=5 +export TILEXR_DEMO_ALLTOALL_REPEAT=50 +export TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS=2097152 +export TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS=16 + +export TILEXR_UDMA_GROUP_TRACE=1 +export TILEXR_UDMA_GROUP_TRACE_DIR=/home/h30059441/ +mkdir -p "$TILEXR_UDMA_GROUP_TRACE_DIR" +``` + +这里 `2097152 int32 = 8 MiB/peer`。命令中的 `16 * 8 MiB` 对应每 rank +128 MiB input 和 128 MiB output。 + +### 8.2 combined 基线 + +```bash +export TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES=0 +``` + +在 `.50.31` 并发启动 rank 0..7，在 `.49.223` 并发启动 rank 8..15： + +```bash +for rank in $(seq ); do + timeout 60s ./tests/udma/install_b101/bin/tilexr_udma_demo \ + 16 "$rank" 8 2097152 8 0 \ + >"$TILEXR_UDMA_GROUP_TRACE_DIR/rank_${rank}.log" 2>&1 & +done +wait +``` + +两台机器必须几乎同时启动。不要先等待一台完成再启动另一台。 + +### 8.3 staged 定位 + +```bash +export TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES=1 +``` + +使用新的端口和目录，按相同方式启动16个 rank。当前会产生四套 trace： + +```text +tilexr_group_trace_local-send_rank_.bin +tilexr_group_trace_local-copy_rank_.bin +tilexr_group_trace_primary_rank_.bin +tilexr_group_trace_secondary_rank_.bin +``` + +staged 测试仍使用 `timeout 60s`。如果60秒不足，先定位 barrier、残留进程或环境 +问题，不要直接无限增加超时。 + +## 9. 结果验收 + +先检查所有 rank 的退出和正确性，再分析性能： + +```bash +grep -H 'TileXR grouped alltoall demo success' \ + "$TILEXR_UDMA_GROUP_TRACE_DIR"/rank_*.log +grep -H -E 'ERROR|failed|MISMATCH|CQ incomplete|0x302|timeout' \ + "$TILEXR_UDMA_GROUP_TRACE_DIR"/rank_*.log || true +find "$TILEXR_UDMA_GROUP_TRACE_DIR" -name '*.bin' \ + -printf '%f %s\n' | sort -V +``` + +通过标准： + +1. 16/16 rank 正常退出并打印 success。 +2. 没有 quiet error、wait timeout、mismatch、`0x302`。 +3. 每个 raw trace 严格为 8,388,608 bytes。 +4. combined 模式有16个 trace；当前 staged 模式有64个 trace。 +5. 日志确认 `groups=1`，payload 为 134,217,728 bytes/rank。 +6. 正确性通过后才能使用本轮性能数字。 + +正式对比至少报告： + +```text +commit +CANN +环境变量差异 +Host perIter mean/min/max +loop49 kernel envelope mean/min/max +各 stage mean +send-put-signal 和 send-quiet +QP/peer/Jetty 分布 +``` + +## 10. 下载并生成 loop49 trace + +将 rank 0..7 从 `.50.31`、rank 8..15 从 `.49.223` 下载到本地同一目录。 + +combined 转换： + +```powershell +$inputs = Get-ChildItem tmp//tilexr_group_trace_rank_*.bin | + Sort-Object { [int]($_.BaseName -replace '.*_', '') } | + ForEach-Object FullName + +python tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py ` + @inputs --output tmp//all_loops.json + +python tmp/extract_xy_trace_iteration.py ` + tmp//all_loops.json ` + tmp//loop49.json 49 +``` + +staged 模式对 `local-send`、`local-copy`、`primary`、`secondary` 分别执行同样的 +转换和 iteration 49 提取。文件名中的 stage 必须与16个输入 trace 一致。 + +Chrome Tracing/Perfetto 中重点检查： + +- 每个 rank 是否有15个唯一 send peer； +- 同一个 `(rank, peer)` 是否只由一个 core 负责； +- local stage 是否全部为 QP0； +- 跨节点是否为 primary QP0 96个 rank pair、secondary QP4 32个 rank pair； +- `send-put-signal` 后是否紧跟同 peer、同 QP 的 `send-quiet`； +- receive wait 是否在对应 signal 后结束； +- 是否存在空 span、重复 peer 或缺失 peer。 + +不同 rank 的 device cycle 没有全局统一时钟。转换器按 rank/iteration 归一化，适合 +观察单 rank 内相对流水；不要把不同 rank 的绝对 `ts` 当成严格同步时间。 + +## 11. 多 Jetty 实验指导 + +当前“每个 peer 创建多个 Jetty，但 payload 只走一个 Jetty”。验证单 rank-pair +Jetty 是否限制带宽时，不能只让不同 peer 轮换 QP0..3：不同 peer 本来已经对应 +不同 Jetty，这种修改没有增加同一 rank pair 的并行度。 + +正确的单变量实验矩阵是： + +```text +1 Jetty/peer: 当前基线 +2 Jetty/peer: 同一 peer 的 chunk 分成2段并发 +4 Jetty/peer: 同一 peer 的 chunk 分成4段并发 +``` + +实现时必须解决跨 QP 完成顺序： + +1. 按同 route 的 QP 集合切分 payload，不能把 primary 数据切到 secondary route。 +2. 所有 payload slice 都提交后，分别 Quiet 对应 QP。 +3. 只有全部 payload QP 完成后，才能发布该 `(group, pass, peer)` 的 ready。 +4. ready 使用确定的 QP，并再次 Quiet；不得让多个 QP竞争写同一个 signal token。 +5. 接收端仍只等待一个 ready，然后 copyout 完整 chunk。 +6. trace 增加 slice/QP 信息，能够区分 payload Quiet 和 ready Quiet。 + +不要直接把当前 `UDMAPutSignalNbiOnQp` 复制到四个 QP。单 QP 上的 +payload-before-signal 顺序不能证明跨 QP 的其他 payload 已完成。 + +先只运行 `local-send` stage 比较 1/2/4 Jetty，减少 copyout 干扰；确认提升后再跑 +combined，验证整体收益和正确性。每档都使用同一 commit、数据规模、warmup/repeat +和空闲环境。 + +## 12. 常见故障 + +### 初始化失败或部分 rank 卡住 + +- 检查16个 rank 是否同时启动； +- 更换 `TILEXR_COMM_ID` 端口； +- 检查两台机器是否存在残留 `tilexr_udma_demo`； +- 检查两台机器的 HEAD、二进制 hash 和环境变量是否一致； +- 只终止本轮明确识别的进程，不要杀其他用户作业。 + +### `0x302`、CQ error 或方向性失败 + +- 检查 payload、ready、Quiet 是否使用同一目标 peer 和预期 QP； +- 检查 `/etc/hccl_rootinfo.json` 的 EID/端口映射； +- 检查 `TILEXR_UDMA_ROUTE_POLICY=all` 和 `TILEXR_UDMA_QP_NUM=4`； +- 检查运行时是否误加载 CANN `devlib` stub； +- 多 Jetty修改时优先检查跨 QP ready 是否提前发布。 + +### trace 正确但性能异常 + +- 先确认没有其他 NPU/UDMA 进程； +- 区分 Host wall、kernel envelope 和纯 `send-quiet`； +- combined 与 staged 不能直接当成同一个吞吐指标； +- 只比较相同 payload 字节数和相同方向； +- 记录 `.50.31 -> .49.223` 与反方向，不能只看双向平均。 + +## 13. 新 Codex 会话提示词 + +可以在新会话直接使用： + +```text +先读取 AGENTS.md、CLAUDE.md 和 +docs/udma-grouped-alltoall-modify-test-guide.md。 + +目标：对 grouped AllToAll 做一次单变量修改并完成物理2x8验证。 +保持每 rank 128 MiB、b101 CANN、warmup5/repeat50、timeout 60s。 +先检查当前分支和已有改动，不要清理无关文件。修改后运行聚焦测试，提交， +生成并验证完整 bundle，部署到 141.61.50.31 和 141.61.49.223，确认 HEAD 和 +产物 hash 一致，再跑16 rank。先检查全部正确性，再提取 loop49 trace 和分析性能。 +任何时候只改变一个实验变量。 + +本轮要验证的变量：<在这里填写，例如 1/2/4 Jetty per peer>。 +``` From 56e5c6d4de8bc6f976bb0d3e219277d80ad57b75 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 11:49:30 +0800 Subject: [PATCH 096/163] perf(udma): isolate combined remote sends --- .../udma/demo/tilexr_udma_alltoall_group_kernel.cpp | 11 ++++++++--- tests/udma/demo/tilexr_udma_alltoall_group_route.h | 8 ++++++-- tests/udma/demo/tilexr_udma_demo.cpp | 5 +++-- .../unit/test_tilexr_udma_alltoall_group_layout.cpp | 13 ++++++++++--- .../test_tilexr_udma_alltoall_group_local_stages.py | 4 +++- 5 files changed, 30 insertions(+), 11 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index ed8145ba..38949142 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -26,6 +26,7 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY = 3U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND = 4U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY = 5U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND = 6U; struct AllToAllGroupDeviceError { uint32_t valid; @@ -91,7 +92,7 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { if (rank < 0 || peer < 0 || rank == peer || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) { + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND) { return false; } if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED) { @@ -105,6 +106,9 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) { return !crossNode; } + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND) { + return crossNode; + } const bool secondary = AllToAllGroupUseSecondaryRouteDevice(rank, peer); return routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY ? crossNode && !secondary : crossNode && secondary; @@ -117,7 +121,8 @@ __aicore__ inline bool AllToAllGroupStageRunsSendDevice(uint32_t routeStage) __aicore__ inline bool AllToAllGroupStageRunsCopyDevice(uint32_t routeStage) { - return routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND; + return routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND; } __aicore__ inline bool AllToAllGroupStageWaitsForSignalDevice(uint32_t routeStage) @@ -321,7 +326,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rankSize = args->rankSize; if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY || + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index 16a9ac82..848db4ba 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -21,11 +21,12 @@ enum class AllToAllGroupRouteStage : uint32_t { kSecondary = 3U, kLocalSend = 4U, kLocalCopy = 5U, + kRemoteSend = 6U, }; inline bool AllToAllGroupValidRouteStage(uint32_t value) { - return value <= static_cast(AllToAllGroupRouteStage::kLocalCopy); + return value <= static_cast(AllToAllGroupRouteStage::kRemoteSend); } inline bool AllToAllGroupStageRunsSend(AllToAllGroupRouteStage stage) @@ -35,7 +36,8 @@ inline bool AllToAllGroupStageRunsSend(AllToAllGroupRouteStage stage) inline bool AllToAllGroupStageRunsCopy(AllToAllGroupRouteStage stage) { - return stage != AllToAllGroupRouteStage::kLocalSend; + return stage != AllToAllGroupRouteStage::kLocalSend && + stage != AllToAllGroupRouteStage::kRemoteSend; } inline bool AllToAllGroupStageWaitsForSignal(AllToAllGroupRouteStage stage) @@ -83,6 +85,8 @@ inline bool AllToAllGroupPeerInRouteStage( case AllToAllGroupRouteStage::kLocalSend: case AllToAllGroupRouteStage::kLocalCopy: return !crossNode; + case AllToAllGroupRouteStage::kRemoteSend: + return crossNode; case AllToAllGroupRouteStage::kPrimary: return crossNode && !AllToAllGroupUseSecondaryRoute(rank, peer); case AllToAllGroupRouteStage::kSecondary: diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 05ca31c9..9b414266 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -674,16 +674,17 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; - constexpr size_t kRouteStageCount = 4U; + constexpr size_t kRouteStageCount = 5U; const std::array stagedRouteStages {{ TileXR::Demo::AllToAllGroupRouteStage::kLocalSend, TileXR::Demo::AllToAllGroupRouteStage::kLocalCopy, + TileXR::Demo::AllToAllGroupRouteStage::kRemoteSend, TileXR::Demo::AllToAllGroupRouteStage::kPrimary, TileXR::Demo::AllToAllGroupRouteStage::kSecondary, }}; const std::array stageNames {{ - "local-send", "local-copy", "primary", "secondary" + "local-send", "local-copy", "remote-send", "primary", "secondary" }}; const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 282de57b..ed73beb0 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -209,8 +209,8 @@ void TestRouteStages() { using TileXR::Demo::AllToAllGroupRouteStage; CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(0U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(5U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(6U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(6U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(7U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( AllToAllGroupRouteStage::kLocalSend), true); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( @@ -221,6 +221,10 @@ void TestRouteStages() AllToAllGroupRouteStage::kLocalCopy), true); CHECK_EQ(TileXR::Demo::AllToAllGroupStageWaitsForSignal( AllToAllGroupRouteStage::kLocalCopy), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( + AllToAllGroupRouteStage::kRemoteSend), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( + AllToAllGroupRouteStage::kRemoteSend), false); for (int rankSize : {8, 16, 128}) { for (int rank = 0; rank < rankSize; ++rank) { @@ -241,6 +245,9 @@ void TestRouteStages() rank, peer, AllToAllGroupRouteStage::kLocalSend), inLocal); CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( rank, peer, AllToAllGroupRouteStage::kLocalCopy), inLocal); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kRemoteSend), + inPrimary || inSecondary); CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary) + static_cast(inSecondary), 1); local += inLocal ? 1 : 0; @@ -335,7 +342,7 @@ void TestHostStructure() const std::string grouped = begin == std::string::npos ? std::string() : demo.substr(begin, end == std::string::npos ? std::string::npos : end - begin); CHECK_CONTAINS(grouped, - "\"local-send\", \"local-copy\", \"primary\", \"secondary\""); + "\"local-send\", \"local-copy\", \"remote-send\", \"primary\", \"secondary\""); CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 53061c29..4a7533f7 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -13,11 +13,13 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("kLocalSend", route) self.assertIn("kLocalCopy", route) + self.assertIn("kRemoteSend", route) self.assertIn("AllToAllGroupStageRunsSendDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageRunsCopyDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) self.assertIn( - '"local-send", "local-copy", "primary", "secondary"', host + '"local-send", "local-copy", "remote-send", "primary", "secondary"', + host, ) From 6e4a56b124a5524e0d4c10c967b3e55ec3f1cf01 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 12:45:11 +0800 Subject: [PATCH 097/163] perf(udma): add grouped contention diagnostics --- .../tilexr_udma_alltoall_group_kernel.cpp | 49 ++++++++++++++----- .../demo/tilexr_udma_alltoall_group_route.h | 33 +++++++++++-- tests/udma/demo/tilexr_udma_demo.cpp | 9 ++-- ...test_tilexr_udma_alltoall_group_layout.cpp | 33 +++++++++++-- ...tilexr_udma_alltoall_group_local_stages.py | 10 +++- 5 files changed, 110 insertions(+), 24 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 38949142..107d2071 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -27,6 +27,9 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY = 3U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND = 4U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY = 5U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND = 6U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_ALL_SEND = 7U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_WAIT = 8U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY = 9U; struct AllToAllGroupDeviceError { uint32_t valid; @@ -92,10 +95,11 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { if (rank < 0 || peer < 0 || rank == peer || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND) { + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY) { return false; } - if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED) { + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_ALL_SEND) { return true; } const bool crossNode = @@ -106,7 +110,9 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) { return !crossNode; } - if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND) { + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_WAIT || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY) { return crossNode; } const bool secondary = AllToAllGroupUseSecondaryRouteDevice(rank, peer); @@ -119,16 +125,33 @@ __aicore__ inline bool AllToAllGroupStageRunsSendDevice(uint32_t routeStage) return routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY; } -__aicore__ inline bool AllToAllGroupStageRunsCopyDevice(uint32_t routeStage) +__aicore__ inline bool AllToAllGroupStageRunsReceiveDevice(uint32_t routeStage) { return routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_SEND && - routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND; + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_ALL_SEND; +} + +__aicore__ inline bool AllToAllGroupStageRunsCopyDevice(uint32_t routeStage) +{ + return AllToAllGroupStageRunsReceiveDevice(routeStage) && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_WAIT; } __aicore__ inline bool AllToAllGroupStageWaitsForSignalDevice(uint32_t routeStage) { - return AllToAllGroupStageRunsCopyDevice(routeStage) && - routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY; + return AllToAllGroupStageRunsReceiveDevice(routeStage) && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY; +} + +__aicore__ inline bool AllToAllGroupReceivePeerInRouteStageDevice( + int32_t rank, int32_t peer, uint32_t routeStage) +{ + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY) { + return rank >= 0 && peer >= 0 && rank != peer; + } + return AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage); } __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( @@ -326,7 +349,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rankSize = args->rankSize; if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND || + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || @@ -350,7 +373,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(elementsPerPeer) * sizeof(int32_t); if (blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES) { - if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { + if (!AllToAllGroupStageRunsReceiveDevice(routeStage)) { AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; @@ -362,7 +385,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( static_cast(elementsPerPeer) * (worker + 1U) / copyoutWorkers); if ((routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED || routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL || - routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY) && + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY) && selfEnd > selfBegin) { const uint64_t selfCopyBegin = AllToAllGroupTraceCycle(groupTrace); auto selfSrc = reinterpret_cast<__gm__ uint8_t*>( @@ -390,7 +414,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (peer < 0) { continue; } - if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage)) { + if (!AllToAllGroupReceivePeerInRouteStageDevice(rank, peer, routeStage)) { continue; } const uint32_t traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane; @@ -430,6 +454,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( peer, TileXR::Demo::kAllToAllGroupTraceNoQp, waitBegin, AllToAllGroupTraceCycle(groupTrace)); } + if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { + continue; + } auto relaySrc = registeredMemory + payloadOffsets[slot] + static_cast(peer) * bytesPerPeer + static_cast(chunkElementOffset) * sizeof(int32_t); diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index 848db4ba..28ac5859 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -22,11 +22,14 @@ enum class AllToAllGroupRouteStage : uint32_t { kLocalSend = 4U, kLocalCopy = 5U, kRemoteSend = 6U, + kAllSend = 7U, + kRemoteWait = 8U, + kRemoteCopy = 9U, }; inline bool AllToAllGroupValidRouteStage(uint32_t value) { - return value <= static_cast(AllToAllGroupRouteStage::kRemoteSend); + return value <= static_cast(AllToAllGroupRouteStage::kRemoteCopy); } inline bool AllToAllGroupStageRunsSend(AllToAllGroupRouteStage stage) @@ -34,16 +37,24 @@ inline bool AllToAllGroupStageRunsSend(AllToAllGroupRouteStage stage) return stage != AllToAllGroupRouteStage::kLocalCopy; } -inline bool AllToAllGroupStageRunsCopy(AllToAllGroupRouteStage stage) +inline bool AllToAllGroupStageRunsReceive(AllToAllGroupRouteStage stage) { return stage != AllToAllGroupRouteStage::kLocalSend && - stage != AllToAllGroupRouteStage::kRemoteSend; + stage != AllToAllGroupRouteStage::kRemoteSend && + stage != AllToAllGroupRouteStage::kAllSend; +} + +inline bool AllToAllGroupStageRunsCopy(AllToAllGroupRouteStage stage) +{ + return AllToAllGroupStageRunsReceive(stage) && + stage != AllToAllGroupRouteStage::kRemoteWait; } inline bool AllToAllGroupStageWaitsForSignal(AllToAllGroupRouteStage stage) { - return AllToAllGroupStageRunsCopy(stage) && - stage != AllToAllGroupRouteStage::kLocalCopy; + return AllToAllGroupStageRunsReceive(stage) && + stage != AllToAllGroupRouteStage::kLocalCopy && + stage != AllToAllGroupRouteStage::kRemoteCopy; } struct AllToAllGroupRouteQps { @@ -80,12 +91,15 @@ inline bool AllToAllGroupPeerInRouteStage( const bool crossNode = AllToAllGroupIsCrossNode(rank, peer); switch (stage) { case AllToAllGroupRouteStage::kCombined: + case AllToAllGroupRouteStage::kAllSend: return true; case AllToAllGroupRouteStage::kLocal: case AllToAllGroupRouteStage::kLocalSend: case AllToAllGroupRouteStage::kLocalCopy: return !crossNode; case AllToAllGroupRouteStage::kRemoteSend: + case AllToAllGroupRouteStage::kRemoteWait: + case AllToAllGroupRouteStage::kRemoteCopy: return crossNode; case AllToAllGroupRouteStage::kPrimary: return crossNode && !AllToAllGroupUseSecondaryRoute(rank, peer); @@ -95,6 +109,15 @@ inline bool AllToAllGroupPeerInRouteStage( return false; } +inline bool AllToAllGroupReceivePeerInRouteStage( + int rank, int peer, AllToAllGroupRouteStage stage) +{ + if (stage == AllToAllGroupRouteStage::kRemoteCopy) { + return rank >= 0 && peer >= 0 && rank != peer; + } + return AllToAllGroupPeerInRouteStage(rank, peer, stage); +} + inline AllToAllGroupRouteQps AllToAllGroupSelectRouteQps( const uint32_t* weights, uint32_t qpCount) { diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 9b414266..21409615 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -674,17 +674,20 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; - constexpr size_t kRouteStageCount = 5U; + constexpr size_t kRouteStageCount = 8U; const std::array stagedRouteStages {{ TileXR::Demo::AllToAllGroupRouteStage::kLocalSend, TileXR::Demo::AllToAllGroupRouteStage::kLocalCopy, TileXR::Demo::AllToAllGroupRouteStage::kRemoteSend, + TileXR::Demo::AllToAllGroupRouteStage::kAllSend, + TileXR::Demo::AllToAllGroupRouteStage::kRemoteWait, + TileXR::Demo::AllToAllGroupRouteStage::kRemoteCopy, TileXR::Demo::AllToAllGroupRouteStage::kPrimary, TileXR::Demo::AllToAllGroupRouteStage::kSecondary, - }}; + }}; const std::array stageNames {{ - "local-send", "local-copy", "remote-send", "primary", "secondary" + "local-send", "local-copy", "remote-send", "all-send", "remote-wait", "remote-copy", "primary", "secondary" }}; const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index ed73beb0..3b33fdc5 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -209,8 +209,8 @@ void TestRouteStages() { using TileXR::Demo::AllToAllGroupRouteStage; CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(0U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(6U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(7U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(9U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(10U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( AllToAllGroupRouteStage::kLocalSend), true); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( @@ -225,6 +225,20 @@ void TestRouteStages() AllToAllGroupRouteStage::kRemoteSend), true); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( AllToAllGroupRouteStage::kRemoteSend), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsReceive( + AllToAllGroupRouteStage::kAllSend), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( + AllToAllGroupRouteStage::kAllSend), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsReceive( + AllToAllGroupRouteStage::kRemoteWait), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageWaitsForSignal( + AllToAllGroupRouteStage::kRemoteWait), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( + AllToAllGroupRouteStage::kRemoteWait), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageWaitsForSignal( + AllToAllGroupRouteStage::kRemoteCopy), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( + AllToAllGroupRouteStage::kRemoteCopy), true); for (int rankSize : {8, 16, 128}) { for (int rank = 0; rank < rankSize; ++rank) { @@ -248,6 +262,16 @@ void TestRouteStages() CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( rank, peer, AllToAllGroupRouteStage::kRemoteSend), inPrimary || inSecondary); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kAllSend), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kRemoteWait), + inPrimary || inSecondary); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kRemoteCopy), + inPrimary || inSecondary); + CHECK_EQ(TileXR::Demo::AllToAllGroupReceivePeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kRemoteCopy), true); CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary) + static_cast(inSecondary), 1); local += inLocal ? 1 : 0; @@ -312,6 +336,8 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); + CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsReceiveDevice(routeStage)"); + CHECK_CONTAINS(kernel, "AllToAllGroupReceivePeerInRouteStageDevice"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsCopyDevice(routeStage)"); CHECK_CONTAINS(kernel, "AllToAllGroupStageWaitsForSignalDevice(routeStage)"); CHECK_CONTAINS(kernel, "observed >= expectedToken"); @@ -342,7 +368,8 @@ void TestHostStructure() const std::string grouped = begin == std::string::npos ? std::string() : demo.substr(begin, end == std::string::npos ? std::string::npos : end - begin); CHECK_CONTAINS(grouped, - "\"local-send\", \"local-copy\", \"remote-send\", \"primary\", \"secondary\""); + "\"local-send\", \"local-copy\", \"remote-send\", \"all-send\", " + "\"remote-wait\", \"remote-copy\", \"primary\", \"secondary\""); CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 4a7533f7..92635ec6 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -14,12 +14,18 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("kLocalSend", route) self.assertIn("kLocalCopy", route) self.assertIn("kRemoteSend", route) + self.assertIn("kAllSend", route) + self.assertIn("kRemoteWait", route) + self.assertIn("kRemoteCopy", route) self.assertIn("AllToAllGroupStageRunsSendDevice(routeStage)", kernel) + self.assertIn("AllToAllGroupStageRunsReceiveDevice(routeStage)", kernel) + self.assertIn("AllToAllGroupReceivePeerInRouteStageDevice", kernel) self.assertIn("AllToAllGroupStageRunsCopyDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) self.assertIn( - '"local-send", "local-copy", "remote-send", "primary", "secondary"', - host, + '"local-send", "local-copy", "remote-send", "all-send", ' + '"remote-wait", "remote-copy", "primary", "secondary"', + " ".join(host.split()), ) From ee703d3989c8eb49868a3ed7fb3430ea18b40cfa Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 13:24:06 +0800 Subject: [PATCH 098/163] perf(udma): trace combined grouped stage --- tests/udma/demo/tilexr_udma_demo.cpp | 45 +++++++++---------- ...test_tilexr_udma_alltoall_group_layout.cpp | 16 ++++++- ...tilexr_udma_alltoall_group_local_stages.py | 9 +++- 3 files changed, 44 insertions(+), 26 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 21409615..dc821d39 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -674,7 +674,7 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; - constexpr size_t kRouteStageCount = 8U; + constexpr size_t kRouteStageCount = 9U; const std::array stagedRouteStages {{ TileXR::Demo::AllToAllGroupRouteStage::kLocalSend, @@ -685,9 +685,11 @@ bool RunGroupedAllToAll( TileXR::Demo::AllToAllGroupRouteStage::kRemoteCopy, TileXR::Demo::AllToAllGroupRouteStage::kPrimary, TileXR::Demo::AllToAllGroupRouteStage::kSecondary, + TileXR::Demo::AllToAllGroupRouteStage::kCombined, }}; const std::array stageNames {{ - "local-send", "local-copy", "remote-send", "all-send", "remote-wait", "remote-copy", "primary", "secondary" + "local-send", "local-copy", "remote-send", "all-send", "remote-wait", + "remote-copy", "primary", "secondary", "combined" }}; const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); @@ -884,34 +886,31 @@ bool RunGroupedAllToAll( totalUs = std::chrono::duration(end - begin).count(); } else { auto runStageBatch = [&](size_t stageIndex) -> bool { - invocationId = 0U; for (int iter = 0; iter < warmup; ++iter, ++invocationId) { launchGroupStage(stagedRouteStages[stageIndex], nullptr, 0U); - if (!CheckAcl(rank, "aclrtSynchronizeStream grouped stage iteration", - aclrtSynchronizeStream(stream))) { - return false; - } + } + if (!CheckAcl(rank, "aclrtSynchronizeStream grouped stage warmup", + aclrtSynchronizeStream(stream)) || + !CheckAcl(rank, "aclrtRecordEvent grouped stage start", + aclrtRecordEvent(stageStartEvent, stream))) { + return false; } for (int iter = 0; iter < repeat; ++iter, ++invocationId) { - if (!CheckAcl(rank, "aclrtRecordEvent grouped stage start", - aclrtRecordEvent(stageStartEvent, stream))) { - return false; - } launchGroupStage(stagedRouteStages[stageIndex], groupTraceDevices[stageIndex], static_cast(iter)); - if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", - aclrtRecordEvent(stageEndEvent, stream)) || - !CheckAcl(rank, "aclrtSynchronizeStream grouped stage iteration", - aclrtSynchronizeStream(stream))) { - return false; - } - float elapsedMs = 0.0F; - if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", - aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { - return false; - } - stageTotalUs[stageIndex] += static_cast(elapsedMs) * 1000.0; } + if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", + aclrtRecordEvent(stageEndEvent, stream)) || + !CheckAcl(rank, "aclrtSynchronizeStream grouped stage measured", + aclrtSynchronizeStream(stream))) { + return false; + } + float elapsedMs = 0.0F; + if (!CheckAcl(rank, "aclrtEventElapsedTime grouped stage", + aclrtEventElapsedTime(&elapsedMs, stageStartEvent, stageEndEvent))) { + return false; + } + stageTotalUs[stageIndex] = static_cast(elapsedMs) * 1000.0; const std::string barrierStep = "grouped route stage " + std::string(stageNames[stageIndex]) + " complete"; return DemoBarrierAll(rank, rankSize, barrierStep); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 3b33fdc5..89a38d8b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -369,12 +369,24 @@ void TestHostStructure() demo.substr(begin, end == std::string::npos ? std::string::npos : end - begin); CHECK_CONTAINS(grouped, "\"local-send\", \"local-copy\", \"remote-send\", \"all-send\", " - "\"remote-wait\", \"remote-copy\", \"primary\", \"secondary\""); + "\"remote-wait\""); + CHECK_CONTAINS(grouped, + "\"remote-copy\", \"primary\", \"secondary\", \"combined\""); + CHECK_CONTAINS(grouped, "constexpr size_t kRouteStageCount = 9U"); CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); + const size_t stageBatchBegin = grouped.find("auto runStageBatch"); + const size_t stageBatchEnd = grouped.find("for (size_t stageIndex", stageBatchBegin); + const std::string stageBatch = stageBatchBegin == std::string::npos ? std::string() : + grouped.substr(stageBatchBegin, + stageBatchEnd == std::string::npos ? std::string::npos : + stageBatchEnd - stageBatchBegin); + CHECK_NOT_CONTAINS(stageBatch, "invocationId = 0U"); + CHECK_NOT_CONTAINS(stageBatch, "grouped stage iteration"); + CHECK_CONTAINS(stageBatch, "aclrtSynchronizeStream grouped stage warmup"); + CHECK_CONTAINS(stageBatch, "aclrtSynchronizeStream grouped stage measured"); CHECK_CONTAINS(grouped, "aclrtEventElapsedTime"); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, barrierStep)"); CHECK_CONTAINS(grouped, "auto runStageBatch"); - CHECK_CONTAINS(grouped, "aclrtSynchronizeStream grouped stage iteration"); CHECK_NOT_CONTAINS(grouped, "\" warmup=\""); CHECK_CONTAINS(grouped, "\" complete\""); CHECK_CONTAINS(grouped, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 92635ec6..ab17db9e 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -24,9 +24,16 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' - '"remote-wait", "remote-copy", "primary", "secondary"', + '"remote-wait", "remote-copy", "primary", "secondary", "combined"', " ".join(host.split()), ) + self.assertIn("constexpr size_t kRouteStageCount = 9U", host) + stage_batch = host[host.index("auto runStageBatch") : host.index( + "for (size_t stageIndex", host.index("auto runStageBatch"))] + self.assertNotIn("invocationId = 0U", stage_batch) + self.assertNotIn("grouped stage iteration", stage_batch) + self.assertIn("aclrtSynchronizeStream grouped stage warmup", stage_batch) + self.assertIn("aclrtSynchronizeStream grouped stage measured", stage_batch) if __name__ == "__main__": From c786578b8de6705bf7dd154c5bf6990f90897cb3 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 14:59:07 +0800 Subject: [PATCH 099/163] perf(udma): add grouped no-copy diagnostic --- .../demo/tilexr_udma_alltoall_group_kernel.cpp | 11 +++++++---- .../demo/tilexr_udma_alltoall_group_route.h | 7 +++++-- tests/udma/demo/tilexr_udma_demo.cpp | 5 +++-- .../test_tilexr_udma_alltoall_group_layout.cpp | 17 ++++++++++++++--- ...t_tilexr_udma_alltoall_group_local_stages.py | 6 ++++-- 5 files changed, 33 insertions(+), 13 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 107d2071..878e7faf 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -30,6 +30,7 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_SEND = 6U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_ALL_SEND = 7U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_WAIT = 8U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY = 9U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY = 10U; struct AllToAllGroupDeviceError { uint32_t valid; @@ -95,11 +96,12 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { if (rank < 0 || peer < 0 || rank == peer || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY) { + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY) { return false; } if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED || - routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_ALL_SEND) { + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_ALL_SEND || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY) { return true; } const bool crossNode = @@ -135,7 +137,8 @@ __aicore__ inline bool AllToAllGroupStageRunsReceiveDevice(uint32_t routeStage) __aicore__ inline bool AllToAllGroupStageRunsCopyDevice(uint32_t routeStage) { return AllToAllGroupStageRunsReceiveDevice(routeStage) && - routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_WAIT; + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_WAIT && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY; } __aicore__ inline bool AllToAllGroupStageWaitsForSignalDevice(uint32_t routeStage) @@ -349,7 +352,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rankSize = args->rankSize; if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || - routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY || + routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index 28ac5859..056c71a9 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -25,11 +25,12 @@ enum class AllToAllGroupRouteStage : uint32_t { kAllSend = 7U, kRemoteWait = 8U, kRemoteCopy = 9U, + kNoCopy = 10U, }; inline bool AllToAllGroupValidRouteStage(uint32_t value) { - return value <= static_cast(AllToAllGroupRouteStage::kRemoteCopy); + return value <= static_cast(AllToAllGroupRouteStage::kNoCopy); } inline bool AllToAllGroupStageRunsSend(AllToAllGroupRouteStage stage) @@ -47,7 +48,8 @@ inline bool AllToAllGroupStageRunsReceive(AllToAllGroupRouteStage stage) inline bool AllToAllGroupStageRunsCopy(AllToAllGroupRouteStage stage) { return AllToAllGroupStageRunsReceive(stage) && - stage != AllToAllGroupRouteStage::kRemoteWait; + stage != AllToAllGroupRouteStage::kRemoteWait && + stage != AllToAllGroupRouteStage::kNoCopy; } inline bool AllToAllGroupStageWaitsForSignal(AllToAllGroupRouteStage stage) @@ -92,6 +94,7 @@ inline bool AllToAllGroupPeerInRouteStage( switch (stage) { case AllToAllGroupRouteStage::kCombined: case AllToAllGroupRouteStage::kAllSend: + case AllToAllGroupRouteStage::kNoCopy: return true; case AllToAllGroupRouteStage::kLocal: case AllToAllGroupRouteStage::kLocalSend: diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index dc821d39..7cec3e0d 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -674,7 +674,7 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; - constexpr size_t kRouteStageCount = 9U; + constexpr size_t kRouteStageCount = 10U; const std::array stagedRouteStages {{ TileXR::Demo::AllToAllGroupRouteStage::kLocalSend, @@ -683,13 +683,14 @@ bool RunGroupedAllToAll( TileXR::Demo::AllToAllGroupRouteStage::kAllSend, TileXR::Demo::AllToAllGroupRouteStage::kRemoteWait, TileXR::Demo::AllToAllGroupRouteStage::kRemoteCopy, + TileXR::Demo::AllToAllGroupRouteStage::kNoCopy, TileXR::Demo::AllToAllGroupRouteStage::kPrimary, TileXR::Demo::AllToAllGroupRouteStage::kSecondary, TileXR::Demo::AllToAllGroupRouteStage::kCombined, }}; const std::array stageNames {{ "local-send", "local-copy", "remote-send", "all-send", "remote-wait", - "remote-copy", "primary", "secondary", "combined" + "remote-copy", "no-copy", "primary", "secondary", "combined" }}; const int warmup = std::max(0, GetEnvInt("TILEXR_DEMO_ALLTOALL_WARMUP", 0)); const int repeat = std::max(1, GetEnvInt("TILEXR_DEMO_ALLTOALL_REPEAT", 1)); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 89a38d8b..530b8139 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -210,7 +210,8 @@ void TestRouteStages() using TileXR::Demo::AllToAllGroupRouteStage; CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(0U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(9U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(10U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(10U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidRouteStage(11U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( AllToAllGroupRouteStage::kLocalSend), true); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( @@ -239,6 +240,14 @@ void TestRouteStages() AllToAllGroupRouteStage::kRemoteCopy), false); CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( AllToAllGroupRouteStage::kRemoteCopy), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsSend( + AllToAllGroupRouteStage::kNoCopy), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsReceive( + AllToAllGroupRouteStage::kNoCopy), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageWaitsForSignal( + AllToAllGroupRouteStage::kNoCopy), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupStageRunsCopy( + AllToAllGroupRouteStage::kNoCopy), false); for (int rankSize : {8, 16, 128}) { for (int rank = 0; rank < rankSize; ++rank) { @@ -272,6 +281,8 @@ void TestRouteStages() inPrimary || inSecondary); CHECK_EQ(TileXR::Demo::AllToAllGroupReceivePeerInRouteStage( rank, peer, AllToAllGroupRouteStage::kRemoteCopy), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( + rank, peer, AllToAllGroupRouteStage::kNoCopy), true); CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary) + static_cast(inSecondary), 1); local += inLocal ? 1 : 0; @@ -371,8 +382,8 @@ void TestHostStructure() "\"local-send\", \"local-copy\", \"remote-send\", \"all-send\", " "\"remote-wait\""); CHECK_CONTAINS(grouped, - "\"remote-copy\", \"primary\", \"secondary\", \"combined\""); - CHECK_CONTAINS(grouped, "constexpr size_t kRouteStageCount = 9U"); + "\"remote-copy\", \"no-copy\", \"primary\", \"secondary\", \"combined\""); + CHECK_CONTAINS(grouped, "constexpr size_t kRouteStageCount = 10U"); CHECK_CONTAINS(grouped, "AllToAllGroupRouteStage::kCombined"); const size_t stageBatchBegin = grouped.find("auto runStageBatch"); const size_t stageBatchEnd = grouped.find("for (size_t stageIndex", stageBatchBegin); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index ab17db9e..1301faf4 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -17,6 +17,7 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("kAllSend", route) self.assertIn("kRemoteWait", route) self.assertIn("kRemoteCopy", route) + self.assertIn("kNoCopy", route) self.assertIn("AllToAllGroupStageRunsSendDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageRunsReceiveDevice(routeStage)", kernel) self.assertIn("AllToAllGroupReceivePeerInRouteStageDevice", kernel) @@ -24,10 +25,11 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' - '"remote-wait", "remote-copy", "primary", "secondary", "combined"', + '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' + '"combined"', " ".join(host.split()), ) - self.assertIn("constexpr size_t kRouteStageCount = 9U", host) + self.assertIn("constexpr size_t kRouteStageCount = 10U", host) stage_batch = host[host.index("auto runStageBatch") : host.index( "for (size_t stageIndex", host.index("auto runStageBatch"))] self.assertNotIn("invocationId = 0U", stage_batch) From 3d2e82254f69ca3d529a5ce5e505e4e5e04880b0 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 15:32:12 +0800 Subject: [PATCH 100/163] perf(udma): throttle grouped signal polling --- .../tilexr_udma_alltoall_group_kernel.cpp | 23 +++++++++++++++---- tests/udma/demo/tilexr_udma_demo.cpp | 14 +++++++++-- ...test_tilexr_udma_alltoall_group_layout.cpp | 5 +++- ...tilexr_udma_alltoall_group_local_stages.py | 3 +++ 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 878e7faf..0eedf5a3 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -231,9 +231,20 @@ __aicore__ inline uint64_t AllToAllGroupLoadTokenMte( return relayLocal.ReinterpretCast().GetValue(0); } +__aicore__ inline void AllToAllGroupWaitPollDelay(uint32_t waitPollCycles) +{ + if (waitPollCycles == 0U) { + return; + } + const uint64_t begin = static_cast(AscendC::GetSystemCycle()); + while (static_cast(AscendC::GetSystemCycle()) - begin < waitPollCycles) { + } +} + __aicore__ inline bool AllToAllGroupWaitTokenMte( __gm__ uint64_t* signal, uint64_t expectedToken, uint64_t timeoutCycles, - AscendC::LocalTensor relayLocal, uint64_t& observed) + uint32_t waitPollCycles, AscendC::LocalTensor relayLocal, + uint64_t& observed) { const uint64_t begin = static_cast(AscendC::GetSystemCycle()); observed = AllToAllGroupLoadTokenMte(signal, relayLocal); @@ -241,6 +252,7 @@ __aicore__ inline bool AllToAllGroupWaitTokenMte( if (static_cast(AscendC::GetSystemCycle()) - begin >= timeoutCycles) { return false; } + AllToAllGroupWaitPollDelay(waitPollCycles); observed = AllToAllGroupLoadTokenMte(signal, relayLocal); } return observed >= expectedToken; @@ -337,7 +349,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage) + uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? @@ -439,7 +451,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t observed = 0ULL; const uint64_t waitBegin = AllToAllGroupTraceCycle(groupTrace); if (!AllToAllGroupWaitTokenMte(signal, expectedToken, - TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, waitPollCycles, + relayLocal, observed)) { AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, @@ -553,11 +566,11 @@ void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage) + uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, copyoutWorkers, routeStage); + groupTrace, traceIteration, copyoutWorkers, routeStage, waitPollCycles); } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 7cec3e0d..5aa15d51 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -60,7 +60,7 @@ extern void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage); + uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -665,6 +665,15 @@ bool RunGroupedAllToAll( } const uint32_t copyoutWorkers = static_cast(copyoutWorkersValue); const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim(copyoutWorkers); + const int waitPollCyclesValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES", 0); + if (waitPollCyclesValue < 0 || waitPollCyclesValue > 1000000) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES must be " + << "between 0 and 1000000, got " << waitPollCyclesValue << std::endl; + return false; + } + const uint32_t waitPollCycles = static_cast(waitPollCyclesValue); const int routeStagesValue = GetEnvInt( "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES", 0); if (routeStagesValue != 0 && routeStagesValue != 1) { @@ -837,6 +846,7 @@ bool RunGroupedAllToAll( PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + " repeat=" + std::to_string(repeat) + " copyoutWorkers=" + std::to_string(copyoutWorkers) + + " waitPollCycles=" + std::to_string(waitPollCycles) + " blockDim=" + std::to_string(groupBlockDim) + " routeStages=" + std::to_string(routeStagesValue)); @@ -858,7 +868,7 @@ bool RunGroupedAllToAll( plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], reinterpret_cast(trace), traceIteration, copyoutWorkers, - static_cast(routeStage)); + static_cast(routeStage), waitPollCycles); }; double totalUs = 0.0; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 530b8139..a3a21ddb 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -335,7 +335,9 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); - CHECK_CONTAINS(kernel, "uint32_t copyoutWorkers, uint32_t routeStage"); + CHECK_CONTAINS(kernel, + "uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles"); + CHECK_CONTAINS(kernel, "AllToAllGroupWaitPollDelay(waitPollCycles)"); CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); CHECK_CONTAINS(kernel, "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); @@ -370,6 +372,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES"); CHECK_CONTAINS(demo, "AllToAllGroupValidCopyoutWorkers"); CHECK_CONTAINS(demo, "AllToAllGroupBlockDim(copyoutWorkers)"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 1301faf4..61b5b2ed 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -23,6 +23,9 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupReceivePeerInRouteStageDevice", kernel) self.assertIn("AllToAllGroupStageRunsCopyDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) + self.assertIn("AllToAllGroupWaitPollDelay(waitPollCycles)", kernel) + self.assertIn("TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES", host) + self.assertIn("uint32_t waitPollCycles", kernel) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' From 852dd349114461f436e8299ceb13108ba17cc8f4 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 15:59:32 +0800 Subject: [PATCH 101/163] Revert "perf(udma): throttle grouped signal polling" This reverts commit 3d2e82254f69ca3d529a5ce5e505e4e5e04880b0. --- .../tilexr_udma_alltoall_group_kernel.cpp | 23 ++++--------------- tests/udma/demo/tilexr_udma_demo.cpp | 14 ++--------- ...test_tilexr_udma_alltoall_group_layout.cpp | 5 +--- ...tilexr_udma_alltoall_group_local_stages.py | 3 --- 4 files changed, 8 insertions(+), 37 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 0eedf5a3..878e7faf 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -231,20 +231,9 @@ __aicore__ inline uint64_t AllToAllGroupLoadTokenMte( return relayLocal.ReinterpretCast().GetValue(0); } -__aicore__ inline void AllToAllGroupWaitPollDelay(uint32_t waitPollCycles) -{ - if (waitPollCycles == 0U) { - return; - } - const uint64_t begin = static_cast(AscendC::GetSystemCycle()); - while (static_cast(AscendC::GetSystemCycle()) - begin < waitPollCycles) { - } -} - __aicore__ inline bool AllToAllGroupWaitTokenMte( __gm__ uint64_t* signal, uint64_t expectedToken, uint64_t timeoutCycles, - uint32_t waitPollCycles, AscendC::LocalTensor relayLocal, - uint64_t& observed) + AscendC::LocalTensor relayLocal, uint64_t& observed) { const uint64_t begin = static_cast(AscendC::GetSystemCycle()); observed = AllToAllGroupLoadTokenMte(signal, relayLocal); @@ -252,7 +241,6 @@ __aicore__ inline bool AllToAllGroupWaitTokenMte( if (static_cast(AscendC::GetSystemCycle()) - begin >= timeoutCycles) { return false; } - AllToAllGroupWaitPollDelay(waitPollCycles); observed = AllToAllGroupLoadTokenMte(signal, relayLocal); } return observed >= expectedToken; @@ -349,7 +337,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles) + uint32_t copyoutWorkers, uint32_t routeStage) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? @@ -451,8 +439,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t observed = 0ULL; const uint64_t waitBegin = AllToAllGroupTraceCycle(groupTrace); if (!AllToAllGroupWaitTokenMte(signal, expectedToken, - TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, waitPollCycles, - relayLocal, observed)) { + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, @@ -566,11 +553,11 @@ void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles) + uint32_t copyoutWorkers, uint32_t routeStage) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, copyoutWorkers, routeStage, waitPollCycles); + groupTrace, traceIteration, copyoutWorkers, routeStage); } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 5aa15d51..7cec3e0d 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -60,7 +60,7 @@ extern void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles); + uint32_t copyoutWorkers, uint32_t routeStage); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -665,15 +665,6 @@ bool RunGroupedAllToAll( } const uint32_t copyoutWorkers = static_cast(copyoutWorkersValue); const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim(copyoutWorkers); - const int waitPollCyclesValue = GetEnvInt( - "TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES", 0); - if (waitPollCyclesValue < 0 || waitPollCyclesValue > 1000000) { - std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES must be " - << "between 0 and 1000000, got " << waitPollCyclesValue << std::endl; - return false; - } - const uint32_t waitPollCycles = static_cast(waitPollCyclesValue); const int routeStagesValue = GetEnvInt( "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES", 0); if (routeStagesValue != 0 && routeStagesValue != 1) { @@ -846,7 +837,6 @@ bool RunGroupedAllToAll( PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + " repeat=" + std::to_string(repeat) + " copyoutWorkers=" + std::to_string(copyoutWorkers) + - " waitPollCycles=" + std::to_string(waitPollCycles) + " blockDim=" + std::to_string(groupBlockDim) + " routeStages=" + std::to_string(routeStagesValue)); @@ -868,7 +858,7 @@ bool RunGroupedAllToAll( plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], reinterpret_cast(trace), traceIteration, copyoutWorkers, - static_cast(routeStage), waitPollCycles); + static_cast(routeStage)); }; double totalUs = 0.0; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index a3a21ddb..530b8139 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -335,9 +335,7 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); - CHECK_CONTAINS(kernel, - "uint32_t copyoutWorkers, uint32_t routeStage, uint32_t waitPollCycles"); - CHECK_CONTAINS(kernel, "AllToAllGroupWaitPollDelay(waitPollCycles)"); + CHECK_CONTAINS(kernel, "uint32_t copyoutWorkers, uint32_t routeStage"); CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); CHECK_CONTAINS(kernel, "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); @@ -372,7 +370,6 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); - CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES"); CHECK_CONTAINS(demo, "AllToAllGroupValidCopyoutWorkers"); CHECK_CONTAINS(demo, "AllToAllGroupBlockDim(copyoutWorkers)"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 61b5b2ed..1301faf4 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -23,9 +23,6 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupReceivePeerInRouteStageDevice", kernel) self.assertIn("AllToAllGroupStageRunsCopyDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) - self.assertIn("AllToAllGroupWaitPollDelay(waitPollCycles)", kernel) - self.assertIn("TILEXR_DEMO_ALLTOALL_GROUP_WAIT_POLL_CYCLES", host) - self.assertIn("uint32_t waitPollCycles", kernel) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' From 1a8ad1f7d44da3da48ecbb451000175c83802636 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 16:48:39 +0800 Subject: [PATCH 102/163] perf(udma): split remote grouped copyout --- .../tilexr_udma_alltoall_group_kernel.cpp | 54 ++++++++++++++----- .../demo/tilexr_udma_alltoall_group_layout.h | 8 ++- .../demo/tilexr_udma_alltoall_group_trace.h | 2 +- ...exr_udma_alltoall_group_trace_to_chrome.py | 2 +- tests/udma/demo/tilexr_udma_demo.cpp | 2 +- ...test_tilexr_udma_alltoall_group_layout.cpp | 14 ++++- ...tilexr_udma_alltoall_group_local_stages.py | 4 ++ ...exr_udma_alltoall_group_trace_to_chrome.py | 4 +- 8 files changed, 70 insertions(+), 20 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 878e7faf..f267afa5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -11,7 +11,7 @@ namespace { constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 32U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 48U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -160,11 +160,21 @@ __aicore__ inline bool AllToAllGroupReceivePeerInRouteStageDevice( __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) { + if (copyoutWorkers == 32U) { + return assignment == 0U ? + static_cast(worker % TILEXR_ALLTOALL_GROUP_SEND_CORES) : -1; + } const uint32_t lane = worker + assignment * copyoutWorkers; return lane < TILEXR_ALLTOALL_GROUP_SEND_CORES ? static_cast(lane) : -1; } +__aicore__ inline bool AllToAllGroupRemoteAssistDevice( + uint32_t worker, uint32_t copyoutWorkers) +{ + return copyoutWorkers == 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; +} + __aicore__ inline void AllToAllGroupSelectRouteQps( const __gm__ TileXR::CommArgs* args, int32_t peer, uint32_t& primaryQp, uint32_t& secondaryQp) @@ -351,7 +361,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rank = args->rank; const int32_t rankSize = args->rankSize; - if ((copyoutWorkers != 8U && copyoutWorkers != 16U) || + if ((copyoutWorkers != 8U && copyoutWorkers != 16U && copyoutWorkers != 32U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || @@ -382,10 +392,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( return; } const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; - const int32_t selfBegin = static_cast( - static_cast(elementsPerPeer) * worker / copyoutWorkers); - const int32_t selfEnd = static_cast( - static_cast(elementsPerPeer) * (worker + 1U) / copyoutWorkers); + const uint32_t selfCopyWorkers = copyoutWorkers == 32U ? 16U : copyoutWorkers; + const int32_t selfBegin = worker < selfCopyWorkers ? static_cast( + static_cast(elementsPerPeer) * worker / selfCopyWorkers) : 0; + const int32_t selfEnd = worker < selfCopyWorkers ? static_cast( + static_cast(elementsPerPeer) * (worker + 1U) / selfCopyWorkers) : 0; if ((routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED || routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL || routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL_COPY || @@ -420,7 +431,16 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (!AllToAllGroupReceivePeerInRouteStageDevice(rank, peer, routeStage)) { continue; } - const uint32_t traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane; + const bool crossNode = + rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) != + peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode); + const bool remoteAssist = + AllToAllGroupRemoteAssistDevice(worker, copyoutWorkers); + if (remoteAssist && !crossNode) { + continue; + } + const uint32_t traceCore = copyoutWorkers == 8U ? + TILEXR_ALLTOALL_GROUP_SEND_CORES + lane : blockIdx; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -428,8 +448,16 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( continue; } const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; - const uint32_t chunkBytes = - static_cast(currentElements) * sizeof(int32_t); + const uint32_t copySliceCount = + copyoutWorkers == 32U && crossNode ? 2U : 1U; + const uint32_t copySliceIndex = remoteAssist ? 1U : 0U; + const int32_t copyElementBegin = static_cast( + static_cast(currentElements) * copySliceIndex / copySliceCount); + const int32_t copyElementEnd = static_cast( + static_cast(currentElements) * (copySliceIndex + 1U) / + copySliceCount); + const uint32_t copyBytes = static_cast( + copyElementEnd - copyElementBegin) * sizeof(int32_t); const uint64_t expectedToken = AllToAllGroupDeviceToken(invocationId, group, pass); auto signal = reinterpret_cast<__gm__ uint64_t*>( @@ -462,11 +490,13 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } auto relaySrc = registeredMemory + payloadOffsets[slot] + static_cast(peer) * bytesPerPeer + - static_cast(chunkElementOffset) * sizeof(int32_t); + static_cast(chunkElementOffset + copyElementBegin) * + sizeof(int32_t); auto relayDst = reinterpret_cast<__gm__ uint8_t*>( - output + static_cast(peer) * elementsPerPeer + chunkElementOffset); + output + static_cast(peer) * elementsPerPeer + + chunkElementOffset + copyElementBegin); const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); - AllToAllGroupCopyMte(relayDst, relaySrc, chunkBytes, relayLocal); + AllToAllGroupCopyMte(relayDst, relaySrc, copyBytes, relayLocal); AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveCopy, groupCount, passCount, diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 15b811ab..1a0e2770 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -21,7 +21,7 @@ constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; -constexpr uint32_t kAllToAllGroupBlockDim = 32U; +constexpr uint32_t kAllToAllGroupBlockDim = 48U; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; @@ -49,7 +49,7 @@ inline bool AllToAllGroupValidRankSize(int rankSize) inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) { - return workers == 8U || workers == 16U; + return workers == 8U || workers == 16U || workers == 32U; } inline uint32_t AllToAllGroupBlockDim(uint32_t workers) @@ -64,6 +64,10 @@ inline int32_t AllToAllGroupCopyoutLane( if (!AllToAllGroupValidCopyoutWorkers(workers) || worker >= workers) { return -1; } + if (workers == 32U) { + return assignment == 0U ? + static_cast(worker % kAllToAllGroupWidth) : -1; + } const uint32_t lane = worker + assignment * workers; return lane < kAllToAllGroupWidth ? static_cast(lane) : -1; } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index 588783ea..a710e658 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -18,7 +18,7 @@ constexpr uint32_t kAllToAllGroupTraceVersion = 1U; constexpr size_t kAllToAllGroupTraceBytes = 8ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; -constexpr uint32_t kAllToAllGroupTraceCoreCount = 32U; +constexpr uint32_t kAllToAllGroupTraceCoreCount = 48U; constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index 348f551c..425641c5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -10,7 +10,7 @@ TRACE_BYTES = 8 * 1024 * 1024 HEADER_BYTES = 4096 MAX_ITERATIONS = 50 -MAX_CORES = 32 +MAX_CORES = 48 PHASE_COUNT = 5 SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 7cec3e0d..fb3bd4b1 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -659,7 +659,7 @@ bool RunGroupedAllToAll( if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( static_cast(copyoutWorkersValue))) { std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8 or 16, got " + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, or 32, got " << copyoutWorkersValue << std::endl; return false; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 530b8139..72293c12 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -300,12 +300,15 @@ void TestRouteStages() void TestCopyoutWorkerPolicy() { + CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 48U); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(32U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U), 48U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(4U), 0U); std::set lanes; @@ -321,6 +324,11 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(7U, 1U, 16U), -1); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(15U, 0U, 16U), 15); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(8U, 0U, 8U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(0U, 0U, 32U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(15U, 0U, 32U), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 32U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(31U, 0U, 32U), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 1U, 32U), -1); } void TestKernelStructure() @@ -341,8 +349,12 @@ void TestKernelStructure() "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL"); CHECK_CONTAINS(kernel, "AllToAllGroupCopyoutLaneDevice"); + CHECK_CONTAINS(kernel, "AllToAllGroupRemoteAssistDevice"); + CHECK_CONTAINS(kernel, "copySliceCount"); + CHECK_CONTAINS(kernel, "copySliceIndex"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers"); - CHECK_CONTAINS(kernel, "traceCore = TILEXR_ALLTOALL_GROUP_SEND_CORES + lane"); + CHECK_CONTAINS(kernel, "const uint32_t traceCore = copyoutWorkers == 8U"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + lane : blockIdx"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 1301faf4..d957225f 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -23,6 +23,10 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupReceivePeerInRouteStageDevice", kernel) self.assertIn("AllToAllGroupStageRunsCopyDevice(routeStage)", kernel) self.assertIn("AllToAllGroupStageWaitsForSignalDevice(routeStage)", kernel) + self.assertIn("AllToAllGroupRemoteAssistDevice", kernel) + self.assertIn("copySliceCount", kernel) + self.assertIn("copySliceIndex", kernel) + self.assertIn("must be 8, 16, or 32", host) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index d50dc1f0..bb14caeb 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -30,7 +30,7 @@ def test_assigns_each_core_disjoint_cache_lines(self): def make_trace( self, path, *, rank=0, magic=None, iteration_count=1, - group_count=1, pass_count=1, core_count=32, + group_count=1, pass_count=1, core_count=48, ): data = bytearray(MODULE.TRACE_BYTES) struct.pack_into( @@ -139,7 +139,7 @@ def test_rejects_bad_magic_size_and_dimensions(self): with self.assertRaisesRegex(ValueError, "size"): MODULE.read_rank_trace(path) - self.make_trace(path, core_count=31) + self.make_trace(path, core_count=47) with self.assertRaisesRegex(ValueError, "core/phase"): MODULE.read_rank_trace(path) From 0799297bbbdb3ccbcd6d3d522677fbab15491c2b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 17:16:36 +0800 Subject: [PATCH 103/163] perf(udma): add third grouped copyout worker --- .../demo/tilexr_udma_alltoall_group_kernel.cpp | 17 ++++++++++------- .../demo/tilexr_udma_alltoall_group_layout.h | 6 +++--- .../demo/tilexr_udma_alltoall_group_trace.h | 2 +- ...ilexr_udma_alltoall_group_trace_to_chrome.py | 2 +- tests/udma/demo/tilexr_udma_demo.cpp | 2 +- .../test_tilexr_udma_alltoall_group_layout.cpp | 9 ++++++++- ...t_tilexr_udma_alltoall_group_local_stages.py | 2 +- ...ilexr_udma_alltoall_group_trace_to_chrome.py | 4 ++-- 8 files changed, 27 insertions(+), 17 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index f267afa5..bd263ac6 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -11,7 +11,7 @@ namespace { constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 48U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -160,7 +160,7 @@ __aicore__ inline bool AllToAllGroupReceivePeerInRouteStageDevice( __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) { - if (copyoutWorkers == 32U) { + if (copyoutWorkers >= 32U) { return assignment == 0U ? static_cast(worker % TILEXR_ALLTOALL_GROUP_SEND_CORES) : -1; } @@ -172,7 +172,7 @@ __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( __aicore__ inline bool AllToAllGroupRemoteAssistDevice( uint32_t worker, uint32_t copyoutWorkers) { - return copyoutWorkers == 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; + return copyoutWorkers >= 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; } __aicore__ inline void AllToAllGroupSelectRouteQps( @@ -361,7 +361,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rank = args->rank; const int32_t rankSize = args->rankSize; - if ((copyoutWorkers != 8U && copyoutWorkers != 16U && copyoutWorkers != 32U) || + if ((copyoutWorkers != 8U && copyoutWorkers != 16U && + copyoutWorkers != 32U && copyoutWorkers != 48U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || @@ -392,7 +393,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( return; } const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; - const uint32_t selfCopyWorkers = copyoutWorkers == 32U ? 16U : copyoutWorkers; + const uint32_t selfCopyWorkers = copyoutWorkers >= 32U ? 16U : copyoutWorkers; const int32_t selfBegin = worker < selfCopyWorkers ? static_cast( static_cast(elementsPerPeer) * worker / selfCopyWorkers) : 0; const int32_t selfEnd = worker < selfCopyWorkers ? static_cast( @@ -449,8 +450,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; const uint32_t copySliceCount = - copyoutWorkers == 32U && crossNode ? 2U : 1U; - const uint32_t copySliceIndex = remoteAssist ? 1U : 0U; + copyoutWorkers >= 32U && crossNode ? + copyoutWorkers / TILEXR_ALLTOALL_GROUP_SEND_CORES : 1U; + const uint32_t copySliceIndex = remoteAssist ? + worker / TILEXR_ALLTOALL_GROUP_SEND_CORES : 0U; const int32_t copyElementBegin = static_cast( static_cast(currentElements) * copySliceIndex / copySliceCount); const int32_t copyElementEnd = static_cast( diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 1a0e2770..5db77aea 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -21,7 +21,7 @@ constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; -constexpr uint32_t kAllToAllGroupBlockDim = 48U; +constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; @@ -49,7 +49,7 @@ inline bool AllToAllGroupValidRankSize(int rankSize) inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) { - return workers == 8U || workers == 16U || workers == 32U; + return workers == 8U || workers == 16U || workers == 32U || workers == 48U; } inline uint32_t AllToAllGroupBlockDim(uint32_t workers) @@ -64,7 +64,7 @@ inline int32_t AllToAllGroupCopyoutLane( if (!AllToAllGroupValidCopyoutWorkers(workers) || worker >= workers) { return -1; } - if (workers == 32U) { + if (workers >= 32U) { return assignment == 0U ? static_cast(worker % kAllToAllGroupWidth) : -1; } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index a710e658..6a3d8bfb 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -18,7 +18,7 @@ constexpr uint32_t kAllToAllGroupTraceVersion = 1U; constexpr size_t kAllToAllGroupTraceBytes = 8ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; -constexpr uint32_t kAllToAllGroupTraceCoreCount = 48U; +constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index 425641c5..e18dbb73 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -10,7 +10,7 @@ TRACE_BYTES = 8 * 1024 * 1024 HEADER_BYTES = 4096 MAX_ITERATIONS = 50 -MAX_CORES = 48 +MAX_CORES = 64 PHASE_COUNT = 5 SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index fb3bd4b1..c026bb34 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -659,7 +659,7 @@ bool RunGroupedAllToAll( if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( static_cast(copyoutWorkersValue))) { std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, or 32, got " + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, 32, or 48, got " << copyoutWorkersValue << std::endl; return false; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 72293c12..c9bc404c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -300,15 +300,17 @@ void TestRouteStages() void TestCopyoutWorkerPolicy() { - CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 48U); + CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 64U); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(32U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(48U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U), 48U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(48U), 64U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(4U), 0U); std::set lanes; @@ -329,6 +331,11 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 32U), 0); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(31U, 0U, 32U), 15); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 1U, 32U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(0U, 0U, 48U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 48U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 0U, 48U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(47U, 0U, 48U), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 1U, 48U), -1); } void TestKernelStructure() diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index d957225f..85a67a5c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -26,7 +26,7 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupRemoteAssistDevice", kernel) self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) - self.assertIn("must be 8, 16, or 32", host) + self.assertIn("must be 8, 16, 32, or 48", host) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index bb14caeb..b133d5f3 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -30,7 +30,7 @@ def test_assigns_each_core_disjoint_cache_lines(self): def make_trace( self, path, *, rank=0, magic=None, iteration_count=1, - group_count=1, pass_count=1, core_count=48, + group_count=1, pass_count=1, core_count=64, ): data = bytearray(MODULE.TRACE_BYTES) struct.pack_into( @@ -139,7 +139,7 @@ def test_rejects_bad_magic_size_and_dimensions(self): with self.assertRaisesRegex(ValueError, "size"): MODULE.read_rank_trace(path) - self.make_trace(path, core_count=47) + self.make_trace(path, core_count=63) with self.assertRaisesRegex(ValueError, "core/phase"): MODULE.read_rank_trace(path) From 3e59d815745217a67c7e54abb39486be98d09a2a Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 19:07:00 +0800 Subject: [PATCH 104/163] perf(udma): default grouped copyout to 48 workers --- tests/udma/demo/tilexr_udma_demo.cpp | 2 +- tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index c026bb34..cb92fc5f 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -655,7 +655,7 @@ bool RunGroupedAllToAll( return false; } const int copyoutWorkersValue = GetEnvInt( - "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 16); + "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 48); if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( static_cast(copyoutWorkersValue))) { std::cerr << "[rank " << rank diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 85a67a5c..59bfb5ad 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -27,6 +27,8 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) self.assertIn("must be 8, 16, 32, or 48", host) + self.assertIn( + '"TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 48', host) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' From dddffadcd2ce1d1880438489a4bdfd734b6676f1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 19:43:43 +0800 Subject: [PATCH 105/163] Revert "perf(udma): default grouped copyout to 48 workers" This reverts commit 3e59d815745217a67c7e54abb39486be98d09a2a. --- tests/udma/demo/tilexr_udma_demo.cpp | 2 +- tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index cb92fc5f..c026bb34 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -655,7 +655,7 @@ bool RunGroupedAllToAll( return false; } const int copyoutWorkersValue = GetEnvInt( - "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 48); + "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 16); if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( static_cast(copyoutWorkersValue))) { std::cerr << "[rank " << rank diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 59bfb5ad..85a67a5c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -27,8 +27,6 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) self.assertIn("must be 8, 16, 32, or 48", host) - self.assertIn( - '"TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 48', host) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' From f072913ea8ddc5577614c3f7d3eb3647f2240350 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 19:43:43 +0800 Subject: [PATCH 106/163] Revert "perf(udma): add third grouped copyout worker" This reverts commit 0799297bbbdb3ccbcd6d3d522677fbab15491c2b. --- .../demo/tilexr_udma_alltoall_group_kernel.cpp | 17 +++++++---------- .../demo/tilexr_udma_alltoall_group_layout.h | 6 +++--- .../demo/tilexr_udma_alltoall_group_trace.h | 2 +- ...ilexr_udma_alltoall_group_trace_to_chrome.py | 2 +- tests/udma/demo/tilexr_udma_demo.cpp | 2 +- .../test_tilexr_udma_alltoall_group_layout.cpp | 9 +-------- ...t_tilexr_udma_alltoall_group_local_stages.py | 2 +- ...ilexr_udma_alltoall_group_trace_to_chrome.py | 4 ++-- 8 files changed, 17 insertions(+), 27 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index bd263ac6..f267afa5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -11,7 +11,7 @@ namespace { constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 48U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -160,7 +160,7 @@ __aicore__ inline bool AllToAllGroupReceivePeerInRouteStageDevice( __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) { - if (copyoutWorkers >= 32U) { + if (copyoutWorkers == 32U) { return assignment == 0U ? static_cast(worker % TILEXR_ALLTOALL_GROUP_SEND_CORES) : -1; } @@ -172,7 +172,7 @@ __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( __aicore__ inline bool AllToAllGroupRemoteAssistDevice( uint32_t worker, uint32_t copyoutWorkers) { - return copyoutWorkers >= 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; + return copyoutWorkers == 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; } __aicore__ inline void AllToAllGroupSelectRouteQps( @@ -361,8 +361,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rank = args->rank; const int32_t rankSize = args->rankSize; - if ((copyoutWorkers != 8U && copyoutWorkers != 16U && - copyoutWorkers != 32U && copyoutWorkers != 48U) || + if ((copyoutWorkers != 8U && copyoutWorkers != 16U && copyoutWorkers != 32U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || @@ -393,7 +392,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( return; } const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; - const uint32_t selfCopyWorkers = copyoutWorkers >= 32U ? 16U : copyoutWorkers; + const uint32_t selfCopyWorkers = copyoutWorkers == 32U ? 16U : copyoutWorkers; const int32_t selfBegin = worker < selfCopyWorkers ? static_cast( static_cast(elementsPerPeer) * worker / selfCopyWorkers) : 0; const int32_t selfEnd = worker < selfCopyWorkers ? static_cast( @@ -450,10 +449,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; const uint32_t copySliceCount = - copyoutWorkers >= 32U && crossNode ? - copyoutWorkers / TILEXR_ALLTOALL_GROUP_SEND_CORES : 1U; - const uint32_t copySliceIndex = remoteAssist ? - worker / TILEXR_ALLTOALL_GROUP_SEND_CORES : 0U; + copyoutWorkers == 32U && crossNode ? 2U : 1U; + const uint32_t copySliceIndex = remoteAssist ? 1U : 0U; const int32_t copyElementBegin = static_cast( static_cast(currentElements) * copySliceIndex / copySliceCount); const int32_t copyElementEnd = static_cast( diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 5db77aea..1a0e2770 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -21,7 +21,7 @@ constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; -constexpr uint32_t kAllToAllGroupBlockDim = 64U; +constexpr uint32_t kAllToAllGroupBlockDim = 48U; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; @@ -49,7 +49,7 @@ inline bool AllToAllGroupValidRankSize(int rankSize) inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) { - return workers == 8U || workers == 16U || workers == 32U || workers == 48U; + return workers == 8U || workers == 16U || workers == 32U; } inline uint32_t AllToAllGroupBlockDim(uint32_t workers) @@ -64,7 +64,7 @@ inline int32_t AllToAllGroupCopyoutLane( if (!AllToAllGroupValidCopyoutWorkers(workers) || worker >= workers) { return -1; } - if (workers >= 32U) { + if (workers == 32U) { return assignment == 0U ? static_cast(worker % kAllToAllGroupWidth) : -1; } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index 6a3d8bfb..a710e658 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -18,7 +18,7 @@ constexpr uint32_t kAllToAllGroupTraceVersion = 1U; constexpr size_t kAllToAllGroupTraceBytes = 8ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; -constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; +constexpr uint32_t kAllToAllGroupTraceCoreCount = 48U; constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index e18dbb73..425641c5 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -10,7 +10,7 @@ TRACE_BYTES = 8 * 1024 * 1024 HEADER_BYTES = 4096 MAX_ITERATIONS = 50 -MAX_CORES = 64 +MAX_CORES = 48 PHASE_COUNT = 5 SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index c026bb34..fb3bd4b1 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -659,7 +659,7 @@ bool RunGroupedAllToAll( if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( static_cast(copyoutWorkersValue))) { std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, 32, or 48, got " + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, or 32, got " << copyoutWorkersValue << std::endl; return false; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index c9bc404c..72293c12 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -300,17 +300,15 @@ void TestRouteStages() void TestCopyoutWorkerPolicy() { - CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 64U); + CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 48U); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(32U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(48U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U), 48U); - CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(48U), 64U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(4U), 0U); std::set lanes; @@ -331,11 +329,6 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 32U), 0); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(31U, 0U, 32U), 15); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 1U, 32U), -1); - CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(0U, 0U, 48U), 0); - CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 48U), 0); - CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 0U, 48U), 0); - CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(47U, 0U, 48U), 15); - CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 1U, 48U), -1); } void TestKernelStructure() diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 85a67a5c..d957225f 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -26,7 +26,7 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupRemoteAssistDevice", kernel) self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) - self.assertIn("must be 8, 16, 32, or 48", host) + self.assertIn("must be 8, 16, or 32", host) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index b133d5f3..bb14caeb 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -30,7 +30,7 @@ def test_assigns_each_core_disjoint_cache_lines(self): def make_trace( self, path, *, rank=0, magic=None, iteration_count=1, - group_count=1, pass_count=1, core_count=64, + group_count=1, pass_count=1, core_count=48, ): data = bytearray(MODULE.TRACE_BYTES) struct.pack_into( @@ -139,7 +139,7 @@ def test_rejects_bad_magic_size_and_dimensions(self): with self.assertRaisesRegex(ValueError, "size"): MODULE.read_rank_trace(path) - self.make_trace(path, core_count=63) + self.make_trace(path, core_count=47) with self.assertRaisesRegex(ValueError, "core/phase"): MODULE.read_rank_trace(path) From 75bd08db71b7390c3b2b139246310bb350126224 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 19:51:34 +0800 Subject: [PATCH 107/163] perf(udma): add third remote copyout worker --- .../demo/tilexr_udma_alltoall_group_kernel.cpp | 17 ++++++++++------- .../demo/tilexr_udma_alltoall_group_layout.h | 6 +++--- .../demo/tilexr_udma_alltoall_group_trace.h | 2 +- ...ilexr_udma_alltoall_group_trace_to_chrome.py | 2 +- tests/udma/demo/tilexr_udma_demo.cpp | 2 +- .../test_tilexr_udma_alltoall_group_layout.cpp | 9 ++++++++- ...t_tilexr_udma_alltoall_group_local_stages.py | 2 +- ...ilexr_udma_alltoall_group_trace_to_chrome.py | 4 ++-- 8 files changed, 27 insertions(+), 17 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index f267afa5..bd263ac6 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -11,7 +11,7 @@ namespace { constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 48U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -160,7 +160,7 @@ __aicore__ inline bool AllToAllGroupReceivePeerInRouteStageDevice( __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( uint32_t worker, uint32_t assignment, uint32_t copyoutWorkers) { - if (copyoutWorkers == 32U) { + if (copyoutWorkers >= 32U) { return assignment == 0U ? static_cast(worker % TILEXR_ALLTOALL_GROUP_SEND_CORES) : -1; } @@ -172,7 +172,7 @@ __aicore__ inline int32_t AllToAllGroupCopyoutLaneDevice( __aicore__ inline bool AllToAllGroupRemoteAssistDevice( uint32_t worker, uint32_t copyoutWorkers) { - return copyoutWorkers == 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; + return copyoutWorkers >= 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; } __aicore__ inline void AllToAllGroupSelectRouteQps( @@ -361,7 +361,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const int32_t rank = args->rank; const int32_t rankSize = args->rankSize; - if ((copyoutWorkers != 8U && copyoutWorkers != 16U && copyoutWorkers != 32U) || + if ((copyoutWorkers != 8U && copyoutWorkers != 16U && + copyoutWorkers != 32U && copyoutWorkers != 48U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || @@ -392,7 +393,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( return; } const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; - const uint32_t selfCopyWorkers = copyoutWorkers == 32U ? 16U : copyoutWorkers; + const uint32_t selfCopyWorkers = copyoutWorkers >= 32U ? 16U : copyoutWorkers; const int32_t selfBegin = worker < selfCopyWorkers ? static_cast( static_cast(elementsPerPeer) * worker / selfCopyWorkers) : 0; const int32_t selfEnd = worker < selfCopyWorkers ? static_cast( @@ -449,8 +450,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; const uint32_t copySliceCount = - copyoutWorkers == 32U && crossNode ? 2U : 1U; - const uint32_t copySliceIndex = remoteAssist ? 1U : 0U; + copyoutWorkers >= 32U && crossNode ? + copyoutWorkers / TILEXR_ALLTOALL_GROUP_SEND_CORES : 1U; + const uint32_t copySliceIndex = remoteAssist ? + worker / TILEXR_ALLTOALL_GROUP_SEND_CORES : 0U; const int32_t copyElementBegin = static_cast( static_cast(currentElements) * copySliceIndex / copySliceCount); const int32_t copyElementEnd = static_cast( diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 1a0e2770..5db77aea 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -21,7 +21,7 @@ constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; -constexpr uint32_t kAllToAllGroupBlockDim = 48U; +constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; @@ -49,7 +49,7 @@ inline bool AllToAllGroupValidRankSize(int rankSize) inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) { - return workers == 8U || workers == 16U || workers == 32U; + return workers == 8U || workers == 16U || workers == 32U || workers == 48U; } inline uint32_t AllToAllGroupBlockDim(uint32_t workers) @@ -64,7 +64,7 @@ inline int32_t AllToAllGroupCopyoutLane( if (!AllToAllGroupValidCopyoutWorkers(workers) || worker >= workers) { return -1; } - if (workers == 32U) { + if (workers >= 32U) { return assignment == 0U ? static_cast(worker % kAllToAllGroupWidth) : -1; } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index a710e658..6a3d8bfb 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -18,7 +18,7 @@ constexpr uint32_t kAllToAllGroupTraceVersion = 1U; constexpr size_t kAllToAllGroupTraceBytes = 8ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; -constexpr uint32_t kAllToAllGroupTraceCoreCount = 48U; +constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index 425641c5..e18dbb73 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -10,7 +10,7 @@ TRACE_BYTES = 8 * 1024 * 1024 HEADER_BYTES = 4096 MAX_ITERATIONS = 50 -MAX_CORES = 48 +MAX_CORES = 64 PHASE_COUNT = 5 SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index fb3bd4b1..c026bb34 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -659,7 +659,7 @@ bool RunGroupedAllToAll( if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( static_cast(copyoutWorkersValue))) { std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, or 32, got " + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, 32, or 48, got " << copyoutWorkersValue << std::endl; return false; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 72293c12..c9bc404c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -300,15 +300,17 @@ void TestRouteStages() void TestCopyoutWorkerPolicy() { - CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 48U); + CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 64U); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(32U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(48U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U), 48U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(48U), 64U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(4U), 0U); std::set lanes; @@ -329,6 +331,11 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 32U), 0); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(31U, 0U, 32U), 15); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 1U, 32U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(0U, 0U, 48U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(16U, 0U, 48U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 0U, 48U), 0); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(47U, 0U, 48U), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 1U, 48U), -1); } void TestKernelStructure() diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index d957225f..85a67a5c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -26,7 +26,7 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupRemoteAssistDevice", kernel) self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) - self.assertIn("must be 8, 16, or 32", host) + self.assertIn("must be 8, 16, 32, or 48", host) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index bb14caeb..b133d5f3 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -30,7 +30,7 @@ def test_assigns_each_core_disjoint_cache_lines(self): def make_trace( self, path, *, rank=0, magic=None, iteration_count=1, - group_count=1, pass_count=1, core_count=48, + group_count=1, pass_count=1, core_count=64, ): data = bytearray(MODULE.TRACE_BYTES) struct.pack_into( @@ -139,7 +139,7 @@ def test_rejects_bad_magic_size_and_dimensions(self): with self.assertRaisesRegex(ValueError, "size"): MODULE.read_rank_trace(path) - self.make_trace(path, core_count=47) + self.make_trace(path, core_count=63) with self.assertRaisesRegex(ValueError, "core/phase"): MODULE.read_rank_trace(path) From 6f672d6f003dca1a255430b6d5fb8cd780915e34 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 21 Jul 2026 20:03:00 +0800 Subject: [PATCH 108/163] feat(udma): select grouped remote route mode --- .../demo/tilexr_udma_alltoall_group_kernel.cpp | 16 ++++++++++++---- .../udma/demo/tilexr_udma_alltoall_group_route.h | 6 ++++++ tests/udma/demo/tilexr_udma_demo.cpp | 14 ++++++++++++-- .../test_tilexr_udma_alltoall_group_layout.cpp | 9 +++++++-- ...st_tilexr_udma_alltoall_group_local_stages.py | 5 +++++ 5 files changed, 42 insertions(+), 8 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index bd263ac6..900eb4e7 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -92,6 +92,13 @@ __aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( TileXR::Demo::kAllToAllGroupPrimaryPeersPerNode; } +__aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( + int32_t rank, int32_t peer, uint32_t useSecondaryRoute) +{ + return useSecondaryRoute != 0U && + AllToAllGroupUseSecondaryRouteDevice(rank, peer); +} + __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { @@ -347,7 +354,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage) + uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? @@ -531,7 +538,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint32_t secondaryQp = 0U; AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp); const uint32_t selectedQp = - AllToAllGroupUseSecondaryRouteDevice(rank, peer) ? secondaryQp : primaryQp; + AllToAllGroupUseSecondaryRouteDevice(rank, peer, useSecondaryRoute) ? + secondaryQp : primaryQp; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -586,11 +594,11 @@ void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage) + uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, copyoutWorkers, routeStage); + groupTrace, traceIteration, copyoutWorkers, routeStage, useSecondaryRoute); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index 056c71a9..58b2b674 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -84,6 +84,12 @@ inline bool AllToAllGroupUseSecondaryRoute(int rank, int peer) kAllToAllGroupPrimaryPeersPerNode; } +inline bool AllToAllGroupUseSecondaryRoute( + int rank, int peer, bool useSecondaryRoute) +{ + return useSecondaryRoute && AllToAllGroupUseSecondaryRoute(rank, peer); +} + inline bool AllToAllGroupPeerInRouteStage( int rank, int peer, AllToAllGroupRouteStage stage) { diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index c026bb34..5782d796 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -60,7 +60,7 @@ extern void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage); + uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -674,6 +674,15 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; + const int useSecondaryRouteValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE", 1); + if (useSecondaryRouteValue != 0 && useSecondaryRouteValue != 1) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE" + << " must be 0 or 1, got " << useSecondaryRouteValue << std::endl; + return false; + } + const uint32_t useSecondaryRoute = static_cast(useSecondaryRouteValue); constexpr size_t kRouteStageCount = 10U; const std::array stagedRouteStages {{ @@ -838,6 +847,7 @@ bool RunGroupedAllToAll( " repeat=" + std::to_string(repeat) + " copyoutWorkers=" + std::to_string(copyoutWorkers) + " blockDim=" + std::to_string(groupBlockDim) + + " useSecondaryRoute=" + std::to_string(useSecondaryRoute) + " routeStages=" + std::to_string(routeStagesValue)); if (routeStages && @@ -858,7 +868,7 @@ bool RunGroupedAllToAll( plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], reinterpret_cast(trace), traceIteration, copyoutWorkers, - static_cast(routeStage)); + static_cast(routeStage), useSecondaryRoute); }; double totalUs = 0.0; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index c9bc404c..4e900a52 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -181,6 +181,8 @@ void TestDualRoutePeerPolicy() } } } + CHECK_EQ(AllToAllGroupUseSecondaryRoute(0, 14, false), false); + CHECK_EQ(AllToAllGroupUseSecondaryRoute(0, 14, true), true); } void TestDualRouteQpWeights() @@ -346,11 +348,13 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES"); CHECK_CONTAINS(kernel, "#include \"tilexr_udma_alltoall_group_route.h\""); CHECK_CONTAINS(kernel, "AllToAllGroupSelectRouteQps"); - CHECK_CONTAINS(kernel, "AllToAllGroupUseSecondaryRouteDevice(rank, peer)"); + CHECK_CONTAINS(kernel, + "AllToAllGroupUseSecondaryRouteDevice(rank, peer, useSecondaryRoute)"); CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); - CHECK_CONTAINS(kernel, "uint32_t copyoutWorkers, uint32_t routeStage"); + CHECK_CONTAINS(kernel, + "uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute"); CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); CHECK_CONTAINS(kernel, "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); @@ -389,6 +393,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); CHECK_CONTAINS(demo, "AllToAllGroupValidCopyoutWorkers"); CHECK_CONTAINS(demo, "AllToAllGroupBlockDim(copyoutWorkers)"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 85a67a5c..0dad684b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -27,6 +27,11 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) self.assertIn("must be 8, 16, 32, or 48", host) + self.assertIn("TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE", host) + self.assertIn( + "AllToAllGroupUseSecondaryRouteDevice(rank, peer, useSecondaryRoute)", + kernel, + ) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' From 08e3271f34de3c3524bafab9bf705292784cd386 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 11:12:31 +0800 Subject: [PATCH 109/163] docs(udma): design 1024-rank grouped alltoall --- ...7-22-grouped-alltoall-1024-ranks-design.md | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md diff --git a/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md b/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md new file mode 100644 index 00000000..f5dc4cd5 --- /dev/null +++ b/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md @@ -0,0 +1,73 @@ +# Grouped AllToAll 1024-Rank Support Design + +## Goal + +Extend the existing grouped full-mesh AllToAll path from its current 128-rank +algorithm limit and 256-rank TileXR communication limit to 1024 ranks. + +The first version preserves the current transport architecture and is validated +with static/unit coverage at 1024 ranks plus a physical 2x8 regression. A real +1024-rank deployment is deferred until a larger environment is available. + +## Scope + +- Raise `TILEXR_MAX_RANK_SIZE` from 256 to 1024. +- Raise the grouped AllToAll rank limit from 128 to 1024. +- Keep the existing requirement that `rankSize` is a multiple of 8. +- Keep 16 peers per group: eight positive and eight negative ring distances. +- Keep one full-mesh UDMA peer/QP entry per remote rank and the existing QP + count and route-selection behavior. +- Keep the existing double registered receive planes and the 1 GiB registered + memory limit. +- Update rank-sized host/device structures and boundary checks through the Comm + initialization, UDMA registry, QP metadata, and grouped kernel launch path. + +Other collective algorithms are not redesigned. They inherit the larger +`CommArgs` ABI and rank-array capacity but receive no new scale optimization. + +## Scheduling + +At 1024 ranks, each rank communicates with 1023 peers in 64 groups. The existing +schedule remains unchanged: + +- group width: 16 peers +- lanes 0-7: increasing positive distance +- lanes 8-15: increasing negative distance +- the diameter peer is emitted only once +- the final group skips invalid lanes + +Every peer must appear exactly once for each source rank, and the schedule must +remain reciprocal between peers. + +## Capacity And Compatibility + +Increasing the global maximum enlarges fixed-size structures. In particular, +`CommArgs::sendCountMatrix` grows to 1024 x 1024 `int64_t` entries (8 MiB), and +the peer-memory, magic, and UDMA registry arrays grow linearly. This is accepted +for the first version. + +The UDMA transport retains its existing O(rankSize^2) metadata exchange and +full-mesh QP construction. Integer arithmetic used for allocation and indexing +must use checked `size_t` multiplication where sizes are derived from rankSize. + +The grouped trace buffer remains 8 MiB. Large combinations of rank count, +passes, and trace iterations may reject tracing with the existing explicit +capacity error. Trace capacity does not limit execution when tracing is off. + +## Validation + +Unit tests will cover rank sizes 8, 128, 256, 512, and 1024 and verify: + +- accepted and rejected rank-size boundaries +- expected group count, including 64 groups at 1024 ranks +- every non-self peer appears exactly once +- no invalid or duplicate peers are scheduled +- reciprocal peer scheduling +- registered-memory planning remains within 1 GiB for a 128 MiB per-rank + payload and rejects oversized layouts +- Comm and UDMA registry structures accept rank 1023 and reject rank 1024 + +After local and remote builds pass, the current CANN b101 physical 2x8 test will +run with 128 MiB per rank, warmup 5, repeat 50, and a 60-second process timeout. +All 16 ranks must complete with correct output. Performance is recorded only as +a regression reference; this change is primarily a scale-capacity change. From 1274d80ab7aa6e263a3afa9aef18cbbd652ed918 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 11:13:21 +0800 Subject: [PATCH 110/163] docs(udma): size grouped trace for 1024 ranks --- .../2026-07-22-grouped-alltoall-1024-ranks-design.md | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md b/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md index f5dc4cd5..842e0326 100644 --- a/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md +++ b/docs/superpowers/specs/2026-07-22-grouped-alltoall-1024-ranks-design.md @@ -50,9 +50,11 @@ The UDMA transport retains its existing O(rankSize^2) metadata exchange and full-mesh QP construction. Integer arithmetic used for allocation and indexing must use checked `size_t` multiplication where sizes are derived from rankSize. -The grouped trace buffer remains 8 MiB. Large combinations of rank count, -passes, and trace iterations may reject tracing with the existing explicit -capacity error. Trace capacity does not limit execution when tracing is off. +Raise the grouped trace buffer from 8 MiB to 128 MiB. This covers 1024 ranks, +64 groups, four passes, 50 iterations, and 64 traced cores (approximately +118 MiB with the current six-phase task-span layout). Larger combinations still +reject tracing with the existing explicit capacity error. Trace capacity does +not limit execution when tracing is off. ## Validation @@ -66,6 +68,8 @@ Unit tests will cover rank sizes 8, 128, 256, 512, and 1024 and verify: - registered-memory planning remains within 1 GiB for a 128 MiB per-rank payload and rejects oversized layouts - Comm and UDMA registry structures accept rank 1023 and reject rank 1024 +- the 128 MiB trace layout accepts 64 groups, four passes, and 50 iterations + while rejecting the first larger unsupported layout After local and remote builds pass, the current CANN b101 physical 2x8 test will run with 128 MiB per rank, warmup 5, repeat 50, and a 60-second process timeout. From 3d03557832df6c8cb484222b891a4fc6c4b859e7 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 11:15:23 +0800 Subject: [PATCH 111/163] docs(udma): plan 1024-rank grouped alltoall --- .../2026-07-22-grouped-alltoall-1024-ranks.md | 197 ++++++++++++++++++ 1 file changed, 197 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-22-grouped-alltoall-1024-ranks.md diff --git a/docs/superpowers/plans/2026-07-22-grouped-alltoall-1024-ranks.md b/docs/superpowers/plans/2026-07-22-grouped-alltoall-1024-ranks.md new file mode 100644 index 00000000..91ac5b8c --- /dev/null +++ b/docs/superpowers/plans/2026-07-22-grouped-alltoall-1024-ranks.md @@ -0,0 +1,197 @@ +# Grouped AllToAll 1024-Rank Support Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Extend the existing grouped full-mesh AllToAll execution path and its TileXR Comm/UDMA dependencies to accept up to 1024 ranks. + +**Architecture:** Raise the existing global fixed-array capacity instead of introducing a second communication ABI. Preserve full-mesh per-peer QPs and the 16-peer grouped schedule, while increasing grouped trace storage to 128 MiB. + +**Tech Stack:** C++14, Ascend C, TileXR UDMA transport, CMake, physical Ascend950/CANN b101 regression. + +## Global Constraints + +- Maximum rank size is exactly 1024. +- Grouped AllToAll continues to require `rankSize % 8 == 0`. +- Group width remains 16 peers, split into eight positive and eight negative distances. +- Existing full-mesh per-peer QP and route behavior remains unchanged. +- Registered memory remains limited to 1 GiB per rank. +- Grouped trace capacity is exactly 128 MiB. +- Physical commands use a 60-second runtime timeout and `/home/pkg/b101/cann`. + +--- + +### Task 1: Expand Comm And UDMA Registry Rank Capacity + +**Files:** +- Modify: `tests/udma/unit/test_tilexr_udma_registry.cpp` +- Modify: `src/include/comm_args.h` + +**Interfaces:** +- Consumes: `TileXR::TILEXR_MAX_RANK_SIZE` in Comm, transport, registry, and kernels. +- Produces: a global fixed-array capacity of 1024 ranks. + +- [ ] **Step 1: Write the failing registry capacity test** + +Change `TestRankScaleLimit()` to assert `TILEXR_MAX_RANK_SIZE == 1024`, construct a valid registry with `rankSize = 1024`, populate `regions[1023]`, and verify rank 1023 is addressable while rank 1024 is rejected. + +- [ ] **Step 2: Run the registry test and verify RED** + +Run: + +```bash +cmake -S tests/udma -B tests/udma/build-local +cmake --build tests/udma/build-local --target test_tilexr_udma_registry -j8 +./tests/udma/build-local/test_tilexr_udma_registry +``` + +Expected: FAIL because the current maximum is 256 and a 1024-rank registry is invalid. + +- [ ] **Step 3: Raise the global capacity** + +Change: + +```cpp +constexpr int TILEXR_MAX_RANK_SIZE = 1024; +``` + +Keep the existing `CommArgs`, peer-memory, magic, count-matrix, and registry array layouts otherwise unchanged. + +- [ ] **Step 4: Run the registry test and verify GREEN** + +Run the commands from Step 2. Expected: `TileXR UDMA registry checks passed`. + +- [ ] **Step 5: Commit** + +```bash +git add src/include/comm_args.h tests/udma/unit/test_tilexr_udma_registry.cpp +git commit -m "feat(udma): expand communication capacity to 1024 ranks" +``` + +### Task 2: Expand Group Scheduling And Trace Capacity + +**Files:** +- Modify: `tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp` +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_layout.h` +- Modify: `tests/udma/demo/tilexr_udma_alltoall_group_trace.h` + +**Interfaces:** +- Consumes: `AllToAllGroupValidRankSize(int)`, `AllToAllGroupCount(int)`, and `AllToAllGroupTraceLayoutFits(uint32_t, uint32_t, uint32_t)`. +- Produces: valid grouped schedules through 1024 ranks and a 128 MiB trace capacity. + +- [ ] **Step 1: Write failing 1024-rank scheduling tests** + +Extend schedule coverage to `{8, 16, 24, 32, 40, 64, 128, 256, 512, 1024}`. Assert: + +```cpp +CHECK_EQ(TileXR::Demo::AllToAllGroupCount(1024), 64U); +CHECK_EQ(TileXR::Demo::AllToAllGroupCount(1032), 0U); +``` + +Add a 128 MiB-per-rank plan using `rankSize = 1024` and `elementsPerPeer = 32768`; verify two payload planes fit below 1 GiB. + +- [ ] **Step 2: Write failing trace capacity tests** + +Include `tilexr_udma_alltoall_group_trace.h` and assert: + +```cpp +CHECK_EQ(TileXR::Demo::kAllToAllGroupTraceBytes, + 128ULL * 1024ULL * 1024ULL); +CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 4U), true); +CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), false); +``` + +- [ ] **Step 3: Run the grouped layout test and verify RED** + +Run: + +```bash +cmake --build tests/udma/build-local --target test_tilexr_udma_alltoall_group_layout -j8 +./tests/udma/build-local/test_tilexr_udma_alltoall_group_layout +``` + +Expected: FAIL because rank 1024 is rejected and trace capacity is 8 MiB. + +- [ ] **Step 4: Implement the new grouped limits** + +Set: + +```cpp +constexpr int32_t kAllToAllGroupMaxRankSize = 1024; +constexpr size_t kAllToAllGroupTraceBytes = 128ULL * 1024ULL * 1024ULL; +``` + +Do not change peer ordering, group-width arithmetic, pass limits, registered-memory limits, or trace record layout. + +- [ ] **Step 5: Run grouped and registry tests and verify GREEN** + +Run: + +```bash +cmake --build tests/udma/build-local --target test_tilexr_udma_registry test_tilexr_udma_alltoall_group_layout -j8 +./tests/udma/build-local/test_tilexr_udma_registry +./tests/udma/build-local/test_tilexr_udma_alltoall_group_layout +``` + +Expected: both tests pass. + +- [ ] **Step 6: Commit** + +```bash +git add tests/udma/demo/tilexr_udma_alltoall_group_layout.h tests/udma/demo/tilexr_udma_alltoall_group_trace.h tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +git commit -m "feat(udma): support 1024-rank grouped alltoall" +``` + +### Task 3: Build And Physical 2x8 Regression + +**Files:** +- Verify: `src/comm/tilexr_comm.cpp` +- Verify: `src/comm/udma/tilexr_udma_transport.cpp` +- Verify: `tests/udma/demo/tilexr_udma_demo.cpp` +- Modify local test utility: `tmp/run_grouped_alltoall_b101_2x8.py` + +**Interfaces:** +- Consumes: the 1024-rank Comm ABI and grouped constants from Tasks 1-2. +- Produces: build evidence and a correct physical 2x8 run. + +- [ ] **Step 1: Run local source and unit verification** + +```bash +cmake --build tests/udma/build-local -j8 +ctest --test-dir tests/udma/build-local --output-on-failure +git diff --check +``` + +Expected: build succeeds, all configured tests pass, and `git diff --check` reports no errors. + +- [ ] **Step 2: Commit the complete HEAD to a bundle and deploy both hosts** + +Update the ignored local runner's trace-size check from `8388608` to +`134217728`, matching the new trace ABI. This utility-only edit is not committed. + +```bash +git bundle create tmp/grouped_alltoall_1024_head.bundle HEAD +python tmp/deploy_grouped_bundle.py tmp/grouped_alltoall_1024_head.bundle /home/h30059441/tilexr_grouped_alltoall_b101 141.61.49.223 141.61.50.31 +python tmp/build_grouped_peer_striping.py 141.61.49.223 141.61.50.31 +``` + +Expected: both hosts check out the same HEAD; Bisheng build and grouped layout, AllToAll layout, transport layout, and trace converter tests pass. + +- [ ] **Step 3: Check NPU availability** + +```bash +python tmp/check_npu.py 141.61.49.223 141.61.50.31 +``` + +Expected: devices 0-7 on both hosts show no running process. If occupied, run only as a correctness check and label performance untrusted. + +- [ ] **Step 4: Run physical 2x8 regression** + +```bash +python tmp/run_grouped_alltoall_b101_2x8.py target 48 0 /home/h30059441/tilexr_grouped_alltoall_b101 8 0 141.61.49.223 141.61.50.31 +``` + +Expected: all 16 ranks report `ok=True`, all 16 trace files have the configured trace size, and the command finishes within its per-rank 60-second timeout. + +- [ ] **Step 5: Record final evidence** + +Report the local test results, remote commit, physical result directory, all-rank correctness, host mean, and any environment contention. Do not claim real 1024-rank runtime validation. From 77d79c93e93ec198c08d37f810893b8368f2ac3c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 11:18:47 +0800 Subject: [PATCH 112/163] feat(udma): expand communication capacity to 1024 ranks --- src/include/comm_args.h | 2 +- tests/udma/unit/test_tilexr_udma_registry.cpp | 15 ++++++++++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/include/comm_args.h b/src/include/comm_args.h index 74931a6b..861c9b0c 100644 --- a/src/include/comm_args.h +++ b/src/include/comm_args.h @@ -34,7 +34,7 @@ using GM_ADDR = uint8_t*; namespace TileXR { -constexpr int TILEXR_MAX_RANK_SIZE = 256; // 最大支持的npu卡数 +constexpr int TILEXR_MAX_RANK_SIZE = 1024; // 最大支持的npu卡数 constexpr int RANK_SIZE_TWO = 2; // 可用SIO的规模，以及是否需要跨卡搬运数据核的分界规模 constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // 前2MB作为flag标志位，之后100MB作为数据存储 diff --git a/tests/udma/unit/test_tilexr_udma_registry.cpp b/tests/udma/unit/test_tilexr_udma_registry.cpp index 0c4ffc5a..c053d629 100644 --- a/tests/udma/unit/test_tilexr_udma_registry.cpp +++ b/tests/udma/unit/test_tilexr_udma_registry.cpp @@ -49,7 +49,20 @@ void TestRemoteAddressCalculation() void TestRankScaleLimit() { - CHECK_EQ(TileXR::TILEXR_MAX_RANK_SIZE, 256); + CHECK_EQ(TileXR::TILEXR_MAX_RANK_SIZE, 1024); + + TileXR::TileXRUDMARegistry registry = {}; + registry.magic = TileXR::TILEXR_UDMA_REGISTRY_MAGIC; + registry.version = TileXR::TILEXR_UDMA_REGISTRY_VERSION; + registry.regionCount = 1; + registry.rankSize = 1024; + registry.regions[TileXR::TILEXR_MAX_RANK_SIZE - 1].base = + reinterpret_cast(0x300000); + registry.regions[TileXR::TILEXR_MAX_RANK_SIZE - 1].bytes = 4096; + + CHECK_TRUE(TileXR::UDMARegistryValid(®istry, 1024)); + CHECK_TRUE(TileXR::UDMARegionContains(®istry, 1023, 2048, 2048)); + CHECK_TRUE(!TileXR::UDMARegionContains(®istry, 1024, 0, 1)); } } // namespace From 2ddca3aefad2dcadd04d52bf5cc46422a3434036 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 11:20:45 +0800 Subject: [PATCH 113/163] feat(udma): support 1024-rank grouped alltoall --- .../demo/tilexr_udma_alltoall_group_layout.h | 2 +- .../demo/tilexr_udma_alltoall_group_trace.h | 4 +-- ...test_tilexr_udma_alltoall_group_layout.cpp | 28 +++++++++++++++++-- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 5db77aea..9933dcd6 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -15,7 +15,7 @@ namespace TileXR { namespace Demo { constexpr int32_t kAllToAllGroupMinRankSize = 8; -constexpr int32_t kAllToAllGroupMaxRankSize = 128; +constexpr int32_t kAllToAllGroupMaxRankSize = 1024; constexpr uint32_t kAllToAllGroupWidth = 16U; constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index 6a3d8bfb..98884a62 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -15,7 +15,7 @@ namespace Demo { constexpr uint32_t kAllToAllGroupTraceMagic = 0x47545243U; // "GTRC" constexpr uint32_t kAllToAllGroupTraceVersion = 1U; -constexpr size_t kAllToAllGroupTraceBytes = 8ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllGroupTraceBytes = 128ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; @@ -148,7 +148,7 @@ static_assert(sizeof(AllToAllGroupTraceTaskSpan) == 24U, static_assert(sizeof(AllToAllGroupTraceHeader) <= kAllToAllGroupTraceHeaderBytes, "group trace header must fit its region"); static_assert(AllToAllGroupTraceTaskSpanBaseOffset() < kAllToAllGroupTraceBytes, - "group trace kernel spans must fit in 8 MiB"); + "group trace kernel spans must fit in trace storage"); } // namespace Demo } // namespace TileXR diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 4e900a52..2f2d461a 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -7,6 +7,7 @@ #include "demo/tilexr_udma_alltoall_group_layout.h" #include "demo/tilexr_udma_alltoall_group_route.h" +#include "demo/tilexr_udma_alltoall_group_trace.h" namespace { @@ -78,7 +79,7 @@ void CheckSchedule(int rankSize) void TestSchedules() { - for (int rankSize : {8, 16, 24, 32, 40, 64, 128}) { + for (int rankSize : {8, 16, 24, 32, 40, 64, 128, 256, 512, 1024}) { CheckSchedule(rankSize); } CHECK_EQ(TileXR::Demo::AllToAllGroupCount(8), 1U); @@ -86,9 +87,10 @@ void TestSchedules() CHECK_EQ(TileXR::Demo::AllToAllGroupCount(24), 2U); CHECK_EQ(TileXR::Demo::AllToAllGroupCount(64), 4U); CHECK_EQ(TileXR::Demo::AllToAllGroupCount(128), 8U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(1024), 64U); CHECK_EQ(TileXR::Demo::AllToAllGroupCount(7), 0U); CHECK_EQ(TileXR::Demo::AllToAllGroupCount(18), 0U); - CHECK_EQ(TileXR::Demo::AllToAllGroupCount(136), 0U); + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(1032), 0U); CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 0), 1); CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 7), 8); @@ -127,7 +129,7 @@ void TestPlan() CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(7, 1024, 1024).valid, false); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(18, 1024, 1024).valid, false); - CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(136, 1024, 1024).valid, false); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(1032, 1024, 1024).valid, false); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(16, 0, 1024).valid, false); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(16, elementsPerPeer, 0).valid, false); @@ -136,6 +138,25 @@ void TestPlan() rankSize, thirtyTwoMiBElements, thirtyTwoMiBElements).valid, false); } +void TestScalePlanAndTraceCapacity() +{ + constexpr int rankSize = 1024; + constexpr int32_t elementsPerPeer = 32768; + const auto plan = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, elementsPerPeer); + + CHECK_EQ(plan.valid, true); + CHECK_EQ(plan.groupCount, 64U); + CHECK_EQ(plan.passCount, 1U); + CHECK_EQ(plan.payloadPlaneBytes, 128ULL * 1024ULL * 1024ULL); + CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); + + CHECK_EQ(TileXR::Demo::kAllToAllGroupTraceBytes, + 128ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 4U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), false); +} + void TestTokens() { const uint64_t token48 = TileXR::Demo::AllToAllGroupToken(48U, 0U, 0U); @@ -435,6 +456,7 @@ int main() { TestSchedules(); TestPlan(); + TestScalePlanAndTraceCapacity(); TestTokens(); TestDualRoutePeerPolicy(); TestDualRouteQpWeights(); From 4fed6b86d632f29ee72484d6d51f12957f413b40 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 13:56:30 +0800 Subject: [PATCH 114/163] fix(udma): read expanded grouped traces --- ...exr_udma_alltoall_group_trace_to_chrome.py | 16 +++-- ...exr_udma_alltoall_group_trace_to_chrome.py | 67 +++++++++++-------- 2 files changed, 50 insertions(+), 33 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index e18dbb73..e68ad61f 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -7,7 +7,7 @@ TRACE_MAGIC = 0x47545243 TRACE_VERSION = 1 -TRACE_BYTES = 8 * 1024 * 1024 +TRACE_BYTES = 128 * 1024 * 1024 HEADER_BYTES = 4096 MAX_ITERATIONS = 50 MAX_CORES = 64 @@ -76,10 +76,12 @@ def _read_task(data, offset, label): def read_rank_trace(path): path = Path(path) - data = path.read_bytes() - if len(data) != TRACE_BYTES: - raise ValueError(f"invalid trace size {len(data)} in {path}, expected {TRACE_BYTES}") - fields = struct.unpack_from(HEADER_FORMAT, data, 0) + file_size = path.stat().st_size + if file_size != TRACE_BYTES: + raise ValueError(f"invalid trace size {file_size} in {path}, expected {TRACE_BYTES}") + with path.open("rb") as stream: + header_data = stream.read(struct.calcsize(HEADER_FORMAT)) + fields = struct.unpack_from(HEADER_FORMAT, header_data, 0) header = { "magic": fields[0], "version": fields[1], @@ -115,6 +117,10 @@ def read_rank_trace(path): header["iteration_count"], header["group_count"], header["pass_count"]) if required > TRACE_BYTES: raise ValueError(f"trace capacity exceeded in {path}: required={required}") + with path.open("rb") as stream: + data = stream.read(required) + if len(data) != required: + raise ValueError(f"short trace read in {path}: read={len(data)} required={required}") return {"path": str(path), "header": header, "data": data} diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index b133d5f3..4b8a9be4 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -32,11 +32,8 @@ def make_trace( self, path, *, rank=0, magic=None, iteration_count=1, group_count=1, pass_count=1, core_count=64, ): - data = bytearray(MODULE.TRACE_BYTES) - struct.pack_into( + header = struct.pack( MODULE.HEADER_FORMAT, - data, - 0, MODULE.TRACE_MAGIC if magic is None else magic, MODULE.TRACE_VERSION, rank, @@ -50,10 +47,6 @@ def make_trace( MODULE.HEADER_BYTES, MODULE.TASK_BASE_OFFSET, ) - for core in (0, 16): - struct.pack_into( - " Date: Wed, 22 Jul 2026 14:40:36 +0800 Subject: [PATCH 115/163] build(udma): support Ascend950 variants --- tests/udma/CMakeLists.txt | 9 +++++-- .../udma/unit/test_tilexr_udma_soc_config.py | 25 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) create mode 100644 tests/udma/unit/test_tilexr_udma_soc_config.py diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index a99f5cc6..372c8092 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -11,6 +11,8 @@ set(CMAKE_CXX_STANDARD 14) set(CMAKE_CXX_STANDARD_REQUIRED ON) option(BUILD_TILEXR_UDMA_DEMO "Build TileXR UDMA communication demo with Ascend C kernel" ON) set(TILEXR_UDMA_DEMO_SOC_TYPE "Ascend950" CACHE STRING "SOC type used for the TileXR UDMA demo kernel") +set_property(CACHE TILEXR_UDMA_DEMO_SOC_TYPE PROPERTY STRINGS + Ascend950 Ascend950PR Ascend950DT Ascend910B) # 获取环境变量 set(ASCEND_HOME_PATH $ENV{ASCEND_HOME_PATH}) @@ -168,14 +170,17 @@ if(BUILD_TILEXR_UDMA_DEMO) if(NOT BISHENG_EXECUTABLE) message(WARNING "bisheng not found; skip tilexr_udma_demo") else() - if(TILEXR_UDMA_DEMO_SOC_TYPE STREQUAL "Ascend950") + if(TILEXR_UDMA_DEMO_SOC_TYPE MATCHES "^Ascend950(PR|DT)?$") set(TILEXR_UDMA_NPU_ARCH "dav-3510") set(TILEXR_UDMA_AICORE_ARCH "--cce-aicore-arch=dav-c310-vec") set(TILEXR_UDMA_CATLASS_ARCH "3510") - else() + elseif(TILEXR_UDMA_DEMO_SOC_TYPE STREQUAL "Ascend910B") set(TILEXR_UDMA_NPU_ARCH "dav-2201") set(TILEXR_UDMA_AICORE_ARCH "--cce-aicore-arch=dav-c220-vec") set(TILEXR_UDMA_CATLASS_ARCH "2201") + else() + message(FATAL_ERROR + "Unsupported TILEXR_UDMA_DEMO_SOC_TYPE: ${TILEXR_UDMA_DEMO_SOC_TYPE}") endif() execute_process( diff --git a/tests/udma/unit/test_tilexr_udma_soc_config.py b/tests/udma/unit/test_tilexr_udma_soc_config.py new file mode 100644 index 00000000..73b85e5e --- /dev/null +++ b/tests/udma/unit/test_tilexr_udma_soc_config.py @@ -0,0 +1,25 @@ +import unittest +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[3] + + +class UDMASocConfigTest(unittest.TestCase): + def test_ascend950_variants_use_dav3510(self): + cmake = (ROOT / "tests" / "udma" / "CMakeLists.txt").read_text( + encoding="utf-8") + + self.assertIn('MATCHES "^Ascend950(PR|DT)?$"', cmake) + self.assertIn('set(TILEXR_UDMA_NPU_ARCH "dav-3510")', cmake) + self.assertIn('set(TILEXR_UDMA_CATLASS_ARCH "3510")', cmake) + + def test_unknown_soc_is_rejected(self): + cmake = (ROOT / "tests" / "udma" / "CMakeLists.txt").read_text( + encoding="utf-8") + + self.assertIn('Unsupported TILEXR_UDMA_DEMO_SOC_TYPE', cmake) + + +if __name__ == "__main__": + unittest.main() From 35279b49c29e836d326ee1b8ec54e12a0e35c948 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 15:22:20 +0800 Subject: [PATCH 116/163] feat(udma): allow UDMA-only communicator init --- src/comm/tilexr_comm.cpp | 17 +++++++++++------ .../comm/unit/test_tilexr_udma_env_sources.cpp | 3 +++ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 48ad250c..3fc4ea2e 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -543,13 +543,18 @@ int TileXRComm::Init() return TILEXR_ERROR_INTERNAL; } - TILEXR_LOG(DEBUG) << "Prepare to InitCommMem localRankSize_ -> " << localRankSize_ << ", localRank_ -> " << localRank_; - if (InitCommMem() != TILEXR_SUCCESS) { - TILEXR_LOG(ERROR) << "InitCommMem failed!"; - return TILEXR_ERROR_INTERNAL; + if (IsEnvEnabled("TILEXR_ENABLE_IPC", true)) { + TILEXR_LOG(DEBUG) << "Prepare to InitCommMem localRankSize_ -> " << localRankSize_ + << ", localRank_ -> " << localRank_; + if (InitCommMem() != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "InitCommMem failed!"; + return TILEXR_ERROR_INTERNAL; + } + TILEXR_LOG(DEBUG) << "InitCommMem " << rank_ << "/" << rankSize_ << ", localRank_ : " << localRank_ + << ", localRankSize_ : " << localRankSize_ << " success"; + } else { + TILEXR_LOG(INFO) << "TileXR IPC memory disabled by environment"; } - TILEXR_LOG(DEBUG) << "InitCommMem " << rank_ << "/" << rankSize_ << ", localRank_ : " << localRank_ << - ", localRankSize_ : " << localRankSize_ << " success"; // 新增：初始化 UDMA ret = InitUDMA(); diff --git a/tests/comm/unit/test_tilexr_udma_env_sources.cpp b/tests/comm/unit/test_tilexr_udma_env_sources.cpp index 27bd8209..19205f7f 100644 --- a/tests/comm/unit/test_tilexr_udma_env_sources.cpp +++ b/tests/comm/unit/test_tilexr_udma_env_sources.cpp @@ -35,6 +35,9 @@ int main() CHECK_CONTAINS(comm, "TILEXR_ENABLE_UDMA"); CHECK_CONTAINS(comm, "TileXR UDMA disabled by environment"); CHECK_CONTAINS(comm, "IsEnvEnabled(\"TILEXR_ENABLE_UDMA\", true)"); + CHECK_CONTAINS(comm, "TILEXR_ENABLE_IPC"); + CHECK_CONTAINS(comm, "TileXR IPC memory disabled by environment"); + CHECK_CONTAINS(comm, "IsEnvEnabled(\"TILEXR_ENABLE_IPC\", true)"); if (g_failures != 0) { std::cerr << g_failures << " UDMA env source checks failed" << std::endl; return 1; From 8e0083a19f26f3527c6963156af116f21f40d676 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 16:34:20 +0800 Subject: [PATCH 117/163] test(udma): align 1024-rank trace capacity with signal protocol --- tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 2f2d461a..f0947944 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -154,7 +154,8 @@ void TestScalePlanAndTraceCapacity() CHECK_EQ(TileXR::Demo::kAllToAllGroupTraceBytes, 128ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 4U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 6U), false); } void TestTokens() From 748e15dfbbc49517512ce659c79194588ed8e983 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 22 Jul 2026 20:48:12 +0800 Subject: [PATCH 118/163] feat(udma): align grouped alltoall channel striping --- .../tilexr_udma_alltoall_group_kernel.cpp | 332 +++++++++++++----- .../demo/tilexr_udma_alltoall_group_layout.h | 45 ++- .../demo/tilexr_udma_alltoall_group_route.h | 77 +++- tests/udma/demo/tilexr_udma_demo.cpp | 67 ++-- ...test_tilexr_udma_alltoall_group_layout.cpp | 122 ++++--- ...tilexr_udma_alltoall_group_local_stages.py | 11 +- 6 files changed, 480 insertions(+), 174 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 900eb4e7..d5c06a4f 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -11,9 +11,11 @@ namespace { constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_WORKERS = 32U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 128U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE = 512U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 1024U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; constexpr uint64_t TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES = 10000000000ULL; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_WORDS = 12U; @@ -76,29 +78,6 @@ __aicore__ inline uint64_t AllToAllGroupDeviceToken( (static_cast(pass) + 1ULL); } -__aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( - int32_t rank, int32_t peer) -{ - if (rank < 0 || peer < 0 || - rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) == - peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode)) { - return false; - } - const uint32_t sourceLocal = static_cast(rank) % - TileXR::Demo::kAllToAllGroupRanksPerNode; - const uint32_t targetLocal = static_cast(peer) % - TileXR::Demo::kAllToAllGroupRanksPerNode; - return (sourceLocal + targetLocal) % TileXR::Demo::kAllToAllGroupRanksPerNode >= - TileXR::Demo::kAllToAllGroupPrimaryPeersPerNode; -} - -__aicore__ inline bool AllToAllGroupUseSecondaryRouteDevice( - int32_t rank, int32_t peer, uint32_t useSecondaryRoute) -{ - return useSecondaryRoute != 0U && - AllToAllGroupUseSecondaryRouteDevice(rank, peer); -} - __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { @@ -124,9 +103,7 @@ __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_REMOTE_COPY) { return crossNode; } - const bool secondary = AllToAllGroupUseSecondaryRouteDevice(rank, peer); - return routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY ? - crossNode && !secondary : crossNode && secondary; + return crossNode; } __aicore__ inline bool AllToAllGroupStageRunsSendDevice(uint32_t routeStage) @@ -184,12 +161,13 @@ __aicore__ inline bool AllToAllGroupRemoteAssistDevice( __aicore__ inline void AllToAllGroupSelectRouteQps( const __gm__ TileXR::CommArgs* args, int32_t peer, - uint32_t& primaryQp, uint32_t& secondaryQp) + uint32_t& primaryQp, uint32_t& secondaryQp, + uint32_t& primaryWeight, uint32_t& secondaryWeight) { auto udmaInfo = TileXR::GetUDMAInfo(args); const uint32_t qpCount = udmaInfo->qpNum == 0U ? 1U : udmaInfo->qpNum; primaryQp = 0U; - uint32_t primaryWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); + primaryWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); for (uint32_t qp = 1U; qp < qpCount; ++qp) { const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qp); if (weight > primaryWeight) { @@ -199,7 +177,7 @@ __aicore__ inline void AllToAllGroupSelectRouteQps( } secondaryQp = primaryQp; - uint32_t secondaryWeight = 0U; + secondaryWeight = 0U; for (uint32_t qp = 0U; qp < qpCount; ++qp) { const uint32_t weight = TileXR::UDMAGetQpWeight(udmaInfo, peer, qp); if (weight != 0U && weight < primaryWeight && weight > secondaryWeight) { @@ -209,6 +187,61 @@ __aicore__ inline void AllToAllGroupSelectRouteQps( } } +__aicore__ inline void AllToAllGroupSplitByRouteDevice( + uint32_t elements, uint32_t primaryWeight, uint32_t secondaryWeight, + uint32_t primaryRouteParts, uint32_t& primaryElements, + uint32_t& secondaryElements) +{ + if (secondaryWeight == 0U) { + primaryElements = elements; + secondaryElements = 0U; + return; + } + uint64_t numerator = primaryWeight; + uint64_t denominator = static_cast(primaryWeight) + secondaryWeight; + if (primaryRouteParts != TileXR::Demo::kAllToAllGroupAutoPrimaryParts) { + numerator = primaryRouteParts > TileXR::Demo::kAllToAllGroupRouteParts ? + TileXR::Demo::kAllToAllGroupRouteParts : primaryRouteParts; + denominator = TileXR::Demo::kAllToAllGroupRouteParts; + } + primaryElements = denominator == 0U ? elements : static_cast( + static_cast(elements) * numerator / denominator); + secondaryElements = elements - primaryElements; +} + +__aicore__ inline void AllToAllGroupRouteSliceForPassDevice( + uint32_t totalElements, uint32_t passOffset, uint32_t passElements, + uint32_t primaryElements, uint32_t route, + uint32_t& elementOffset, uint32_t& elements) +{ + elementOffset = 0U; + elements = 0U; + if (route > 1U || passOffset >= totalElements || passElements == 0U) { + return; + } + const uint32_t passEnd = passElements > totalElements - passOffset ? + totalElements : passOffset + passElements; + const uint32_t routeBegin = route == 0U ? 0U : primaryElements; + const uint32_t routeEnd = route == 0U ? primaryElements : totalElements; + const uint32_t begin = passOffset > routeBegin ? passOffset : routeBegin; + const uint32_t end = passEnd < routeEnd ? passEnd : routeEnd; + if (end > begin) { + elementOffset = begin; + elements = end - begin; + } +} + +__aicore__ inline bool AllToAllGroupRouteRunsInStage(uint32_t routeStage, uint32_t route) +{ + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY) { + return route == 0U; + } + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY) { + return route == 1U; + } + return true; +} + __aicore__ inline void AllToAllGroupCopyMte( __gm__ uint8_t* dst, const __gm__ uint8_t* src, uint32_t bytes, AscendC::LocalTensor relayLocal) @@ -263,6 +296,34 @@ __aicore__ inline bool AllToAllGroupWaitTokenMte( return observed >= expectedToken; } +__aicore__ inline bool AllToAllGroupWaitRouteTokensMte( + __gm__ uint64_t* primarySignal, __gm__ uint64_t* secondarySignal, + bool waitPrimary, bool waitSecondary, uint64_t expectedToken, + uint64_t timeoutCycles, AscendC::LocalTensor relayLocal, + uint64_t& observed) +{ + const uint64_t begin = static_cast(AscendC::GetSystemCycle()); + uint64_t primaryObserved = waitPrimary ? 0ULL : expectedToken; + uint64_t secondaryObserved = waitSecondary ? 0ULL : expectedToken; + while (primaryObserved < expectedToken || secondaryObserved < expectedToken) { + if (waitPrimary && primaryObserved < expectedToken) { + primaryObserved = AllToAllGroupLoadTokenMte(primarySignal, relayLocal); + } + if (waitSecondary && secondaryObserved < expectedToken) { + secondaryObserved = AllToAllGroupLoadTokenMte(secondarySignal, relayLocal); + } + if (primaryObserved >= expectedToken && secondaryObserved >= expectedToken) { + break; + } + if (static_cast(AscendC::GetSystemCycle()) - begin >= timeoutCycles) { + observed = primaryObserved < expectedToken ? primaryObserved : secondaryObserved; + return false; + } + } + observed = primaryObserved < secondaryObserved ? primaryObserved : secondaryObserved; + return true; +} + __aicore__ inline void AllToAllGroupRecordError( __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, @@ -354,8 +415,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute) + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts) { + constexpr uint32_t copyoutWorkers = 32U; const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? reinterpret_cast<__gm__ uint8_t*>(groupTraceGM) : nullptr; @@ -371,7 +433,12 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if ((copyoutWorkers != 8U && copyoutWorkers != 16U && copyoutWorkers != 32U && copyoutWorkers != 48U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || - blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers || + multiChannel > 1U || + (primaryRouteParts > TileXR::Demo::kAllToAllGroupRouteParts && + primaryRouteParts != TileXR::Demo::kAllToAllGroupAutoPrimaryParts) || + TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers > + TILEXR_ALLTOALL_GROUP_BLOCK_DIM || + blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || groupCount == 0U) { @@ -393,13 +460,13 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( const uint64_t bytesPerPeer = static_cast(elementsPerPeer) * sizeof(int32_t); - if (blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_CORES) { + if (blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_WORKERS) { if (!AllToAllGroupStageRunsReceiveDevice(routeStage)) { AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } - const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_CORES; + const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_WORKERS; const uint32_t selfCopyWorkers = copyoutWorkers >= 32U ? 16U : copyoutWorkers; const int32_t selfBegin = worker < selfCopyWorkers ? static_cast( static_cast(elementsPerPeer) * worker / selfCopyWorkers) : 0; @@ -447,8 +514,23 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (remoteAssist && !crossNode) { continue; } + uint32_t primaryQp = 0U; + uint32_t secondaryQp = 0U; + uint32_t primaryWeight = 0U; + uint32_t secondaryWeight = 0U; + AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp, + primaryWeight, secondaryWeight); + if (multiChannel == 0U || !crossNode || secondaryQp == primaryQp) { + secondaryWeight = 0U; + } + uint32_t primaryTotalElements = 0U; + uint32_t secondaryTotalElements = 0U; + AllToAllGroupSplitByRouteDevice( + static_cast(elementsPerPeer), primaryWeight, + secondaryWeight, primaryRouteParts, + primaryTotalElements, secondaryTotalElements); const uint32_t traceCore = copyoutWorkers == 8U ? - TILEXR_ALLTOALL_GROUP_SEND_CORES + lane : blockIdx; + TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -461,23 +543,50 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( copyoutWorkers / TILEXR_ALLTOALL_GROUP_SEND_CORES : 1U; const uint32_t copySliceIndex = remoteAssist ? worker / TILEXR_ALLTOALL_GROUP_SEND_CORES : 0U; - const int32_t copyElementBegin = static_cast( + uint32_t copyElementBegin = static_cast( static_cast(currentElements) * copySliceIndex / copySliceCount); - const int32_t copyElementEnd = static_cast( + uint32_t copyElementEnd = static_cast( static_cast(currentElements) * (copySliceIndex + 1U) / copySliceCount); - const uint32_t copyBytes = static_cast( - copyElementEnd - copyElementBegin) * sizeof(int32_t); + uint32_t primaryElementOffset = 0U; + uint32_t primaryElements = 0U; + uint32_t secondaryElementOffset = 0U; + uint32_t secondaryElements = 0U; + AllToAllGroupRouteSliceForPassDevice( + static_cast(elementsPerPeer), + static_cast(chunkElementOffset), + static_cast(currentElements), primaryTotalElements, 0U, + primaryElementOffset, primaryElements); + AllToAllGroupRouteSliceForPassDevice( + static_cast(elementsPerPeer), + static_cast(chunkElementOffset), + static_cast(currentElements), primaryTotalElements, 1U, + secondaryElementOffset, secondaryElements); + if ((routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY && + primaryElements == 0U) || + (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY && + secondaryElements == 0U)) { + continue; + } const uint64_t expectedToken = AllToAllGroupDeviceToken(invocationId, group, pass); - auto signal = reinterpret_cast<__gm__ uint64_t*>( + auto primarySignal = reinterpret_cast<__gm__ uint64_t*>( registeredMemory + signalOffsets[slot] + static_cast(peer) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE); + auto secondarySignal = reinterpret_cast<__gm__ uint64_t*>( + reinterpret_cast<__gm__ uint8_t*>(primarySignal) + + TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE); if (AllToAllGroupStageWaitsForSignalDevice(routeStage)) { uint64_t observed = 0ULL; const uint64_t waitBegin = AllToAllGroupTraceCycle(groupTrace); - if (!AllToAllGroupWaitTokenMte(signal, expectedToken, - TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observed)) { + const bool waitPrimary = primaryElements != 0U && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY; + const bool waitSecondary = secondaryElements != 0U && + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY; + if (!AllToAllGroupWaitRouteTokensMte( + primarySignal, secondarySignal, waitPrimary, waitSecondary, + expectedToken, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, + relayLocal, observed)) { AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, @@ -498,6 +607,26 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { continue; } + if (routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY || + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY) { + const uint32_t routeOffset = + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY ? + primaryElementOffset : secondaryElementOffset; + const uint32_t routeElements = + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY ? + primaryElements : secondaryElements; + const uint32_t routeBegin = + routeOffset - static_cast(chunkElementOffset); + const uint32_t routeEnd = routeBegin + routeElements; + copyElementBegin = copyElementBegin > routeBegin ? + copyElementBegin : routeBegin; + copyElementEnd = copyElementEnd < routeEnd ? copyElementEnd : routeEnd; + } + if (copyElementEnd <= copyElementBegin) { + continue; + } + const uint32_t copyBytes = + (copyElementEnd - copyElementBegin) * sizeof(int32_t); auto relaySrc = registeredMemory + payloadOffsets[slot] + static_cast(peer) * bytesPerPeer + static_cast(chunkElementOffset + copyElementBegin) * @@ -525,7 +654,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } - const uint32_t lane = blockIdx; + const uint32_t lane = blockIdx % TILEXR_ALLTOALL_GROUP_SEND_CORES; + const uint32_t workerRoute = blockIdx / TILEXR_ALLTOALL_GROUP_SEND_CORES; for (uint32_t group = 0U; group < groupCount; ++group) { const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); if (peer < 0) { @@ -536,10 +666,22 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } uint32_t primaryQp = 0U; uint32_t secondaryQp = 0U; - AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp); - const uint32_t selectedQp = - AllToAllGroupUseSecondaryRouteDevice(rank, peer, useSecondaryRoute) ? - secondaryQp : primaryQp; + uint32_t primaryWeight = 0U; + uint32_t secondaryWeight = 0U; + AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp, + primaryWeight, secondaryWeight); + const bool crossNode = + rank / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode) != + peer / static_cast(TileXR::Demo::kAllToAllGroupRanksPerNode); + if (multiChannel == 0U || !crossNode || secondaryQp == primaryQp) { + secondaryWeight = 0U; + } + uint32_t primaryTotalElements = 0U; + uint32_t secondaryTotalElements = 0U; + AllToAllGroupSplitByRouteDevice( + static_cast(elementsPerPeer), primaryWeight, + secondaryWeight, primaryRouteParts, + primaryTotalElements, secondaryTotalElements); for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -547,38 +689,69 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( continue; } const int32_t currentElements = remaining < chunkElements ? remaining : chunkElements; - const uint32_t chunkBytes = - static_cast(currentElements) * sizeof(int32_t); - const uint64_t chunkByteOffset = - static_cast(chunkElementOffset) * sizeof(int32_t); + uint32_t primaryElementOffset = 0U; + uint32_t primaryElements = 0U; + uint32_t secondaryElementOffset = 0U; + uint32_t secondaryElements = 0U; + AllToAllGroupRouteSliceForPassDevice( + static_cast(elementsPerPeer), + static_cast(chunkElementOffset), + static_cast(currentElements), primaryTotalElements, 0U, + primaryElementOffset, primaryElements); + AllToAllGroupRouteSliceForPassDevice( + static_cast(elementsPerPeer), + static_cast(chunkElementOffset), + static_cast(currentElements), primaryTotalElements, 1U, + secondaryElementOffset, secondaryElements); const uint64_t expectedToken = AllToAllGroupDeviceToken(invocationId, group, pass); - auto localSrc = input + static_cast(peer) * elementsPerPeer + - chunkElementOffset; - const uint64_t remotePayloadOffset = payloadOffsets[slot] + - static_cast(rank) * bytesPerPeer + chunkByteOffset; - const uint64_t remoteSignalOffset = signalOffsets[slot] + - static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE; - const uint64_t putBegin = AllToAllGroupTraceCycle(groupTrace); - TileXR::UDMAPutSignalNbiOnQp( - args, peer, selectedQp, localSrc, remotePayloadOffset, chunkBytes, - remoteSignalOffset, expectedToken); - AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, pass, - TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, - peer, selectedQp, putBegin, AllToAllGroupTraceCycle(groupTrace)); - const uint64_t quietBegin = AllToAllGroupTraceCycle(groupTrace); - const uint32_t quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); - AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, pass, - TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, - peer, selectedQp, quietBegin, AllToAllGroupTraceCycle(groupTrace)); - if (quietStatus != 0U) { - AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_QUIET, - group, pass, peer, selectedQp, quietStatus, expectedToken, 0ULL); - AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, - kernelBegin, AllToAllGroupTraceCycle(groupTrace)); - return; + const uint32_t routeBegin = workerRoute; + const uint32_t routeEnd = workerRoute + 1U; + for (uint32_t route = routeBegin; route < routeEnd; ++route) { + if (!AllToAllGroupRouteRunsInStage(routeStage, route)) { + continue; + } + const uint32_t segmentElements = route == 0U ? + primaryElements : secondaryElements; + if (segmentElements == 0U) { + continue; + } + const uint32_t segmentElementOffset = route == 0U ? + primaryElementOffset : secondaryElementOffset; + const uint32_t selectedQp = route == 0U ? primaryQp : secondaryQp; + const uint64_t elementOffset = segmentElementOffset; + auto localSrc = input + static_cast(peer) * elementsPerPeer + + elementOffset; + const uint64_t remotePayloadOffset = payloadOffsets[slot] + + static_cast(rank) * bytesPerPeer + + elementOffset * sizeof(int32_t); + const uint64_t remoteSignalOffset = signalOffsets[slot] + + static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE + + static_cast(route) * + TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE; + const uint64_t putBegin = AllToAllGroupTraceCycle(groupTrace); + TileXR::UDMAPutSignalNbiOnQp( + args, peer, selectedQp, localSrc, remotePayloadOffset, + segmentElements * sizeof(int32_t), remoteSignalOffset, expectedToken); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, + peer, selectedQp, putBegin, AllToAllGroupTraceCycle(groupTrace)); + const uint64_t quietBegin = AllToAllGroupTraceCycle(groupTrace); + const uint32_t quietStatus = + TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, + peer, selectedQp, quietBegin, AllToAllGroupTraceCycle(groupTrace)); + if (quietStatus != 0U) { + AllToAllGroupRecordError(debug, blockIdx, + TILEXR_ALLTOALL_GROUP_STAGE_QUIET, group, pass, peer, + selectedQp, quietStatus, expectedToken, 0ULL); + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } } } } @@ -594,11 +767,12 @@ void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute) + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, copyoutWorkers, routeStage, useSecondaryRoute); + groupTrace, traceIteration, routeStage, + multiChannel, primaryRouteParts); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 9933dcd6..e56ba187 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -19,9 +19,13 @@ constexpr int32_t kAllToAllGroupMaxRankSize = 1024; constexpr uint32_t kAllToAllGroupWidth = 16U; constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; -constexpr uint32_t kAllToAllGroupSignalSlotBytes = 128U; +constexpr uint32_t kAllToAllGroupRouteSignalStride = 512U; +constexpr uint32_t kAllToAllGroupSignalSlotBytes = 1024U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; +constexpr uint32_t kAllToAllGroupSendWorkerCount = 32U; constexpr uint32_t kAllToAllGroupBlockDim = 64U; +constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = + 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; @@ -41,6 +45,35 @@ struct AllToAllGroupPlan { size_t registeredBytes = 0; }; +enum class AllToAllGroupChannelMode : uint32_t { + kAuto = 0U, + kSingle = 1U, + kMulti = 2U, +}; + +inline bool AllToAllGroupValidChannelMode(uint32_t mode) +{ + return mode <= static_cast(AllToAllGroupChannelMode::kMulti); +} + +inline bool AllToAllGroupUseMultiChannel( + size_t perRankBytes, AllToAllGroupChannelMode mode) +{ + if (mode == AllToAllGroupChannelMode::kSingle) { + return false; + } + if (mode == AllToAllGroupChannelMode::kMulti) { + return true; + } + return perRankBytes > kAllToAllGroupMultiChannelThresholdBytes; +} + +inline size_t AllToAllGroupSignalByteOffset(uint32_t sourceRank, uint32_t route) +{ + return static_cast(sourceRank) * kAllToAllGroupSignalSlotBytes + + static_cast(route) * kAllToAllGroupRouteSignalStride; +} + inline bool AllToAllGroupValidRankSize(int rankSize) { return rankSize >= kAllToAllGroupMinRankSize && @@ -52,10 +85,14 @@ inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) return workers == 8U || workers == 16U || workers == 32U || workers == 48U; } -inline uint32_t AllToAllGroupBlockDim(uint32_t workers) +inline uint32_t AllToAllGroupBlockDim(uint32_t sendWorkers, uint32_t copyoutWorkers) { - return AllToAllGroupValidCopyoutWorkers(workers) ? - kAllToAllGroupSendCoreCount + workers : 0U; + if (sendWorkers != kAllToAllGroupSendWorkerCount || + !AllToAllGroupValidCopyoutWorkers(copyoutWorkers) || + sendWorkers + copyoutWorkers > kAllToAllGroupBlockDim) { + return 0U; + } + return sendWorkers + copyoutWorkers; } inline int32_t AllToAllGroupCopyoutLane( diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_route.h b/tests/udma/demo/tilexr_udma_alltoall_group_route.h index 58b2b674..2ffc3f49 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_route.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_route.h @@ -7,12 +7,15 @@ #define TILEXR_UDMA_ALLTOALL_GROUP_ROUTE_H #include +#include namespace TileXR { namespace Demo { constexpr uint32_t kAllToAllGroupRanksPerNode = 8U; -constexpr uint32_t kAllToAllGroupPrimaryPeersPerNode = 6U; +constexpr uint32_t kAllToAllGroupRouteParts = 8U; +constexpr uint32_t kAllToAllGroupAutoPrimaryParts = + std::numeric_limits::max(); enum class AllToAllGroupRouteStage : uint32_t { kCombined = 0U, @@ -64,30 +67,67 @@ struct AllToAllGroupRouteQps { uint32_t secondaryQp = 0U; }; -inline bool AllToAllGroupIsCrossNode(int rank, int peer) +struct AllToAllGroupRouteSplit { + uint32_t primaryElements = 0U; + uint32_t secondaryElements = 0U; +}; + +struct AllToAllGroupRouteSlice { + uint32_t elementOffset = 0U; + uint32_t elements = 0U; +}; + +inline AllToAllGroupRouteSplit AllToAllGroupSplitByRoute( + uint32_t elements, uint32_t primaryWeight, uint32_t secondaryWeight, + uint32_t primaryParts) { - return rank >= 0 && peer >= 0 && - rank / static_cast(kAllToAllGroupRanksPerNode) != - peer / static_cast(kAllToAllGroupRanksPerNode); + AllToAllGroupRouteSplit result {}; + if (elements == 0U) { + return result; + } + if (secondaryWeight == 0U) { + result.primaryElements = elements; + return result; + } + uint64_t numerator = primaryWeight; + uint64_t denominator = static_cast(primaryWeight) + secondaryWeight; + if (primaryParts != kAllToAllGroupAutoPrimaryParts) { + numerator = primaryParts > kAllToAllGroupRouteParts ? + kAllToAllGroupRouteParts : primaryParts; + denominator = kAllToAllGroupRouteParts; + } + result.primaryElements = denominator == 0U ? elements : + static_cast(static_cast(elements) * numerator / denominator); + result.secondaryElements = elements - result.primaryElements; + return result; } -inline bool AllToAllGroupUseSecondaryRoute(int rank, int peer) +inline AllToAllGroupRouteSlice AllToAllGroupRouteSliceForPass( + uint32_t totalElements, uint32_t passOffset, uint32_t passElements, + uint32_t primaryElements, uint32_t route) { - if (!AllToAllGroupIsCrossNode(rank, peer)) { - return false; + AllToAllGroupRouteSlice result {}; + if (route > 1U || passOffset >= totalElements || passElements == 0U) { + return result; + } + const uint32_t passEnd = passElements > totalElements - passOffset ? + totalElements : passOffset + passElements; + const uint32_t routeBegin = route == 0U ? 0U : primaryElements; + const uint32_t routeEnd = route == 0U ? primaryElements : totalElements; + const uint32_t begin = passOffset > routeBegin ? passOffset : routeBegin; + const uint32_t end = passEnd < routeEnd ? passEnd : routeEnd; + if (end > begin) { + result.elementOffset = begin; + result.elements = end - begin; } - const uint32_t sourceLocal = - static_cast(rank) % kAllToAllGroupRanksPerNode; - const uint32_t targetLocal = - static_cast(peer) % kAllToAllGroupRanksPerNode; - return (sourceLocal + targetLocal) % kAllToAllGroupRanksPerNode >= - kAllToAllGroupPrimaryPeersPerNode; + return result; } -inline bool AllToAllGroupUseSecondaryRoute( - int rank, int peer, bool useSecondaryRoute) +inline bool AllToAllGroupIsCrossNode(int rank, int peer) { - return useSecondaryRoute && AllToAllGroupUseSecondaryRoute(rank, peer); + return rank >= 0 && peer >= 0 && + rank / static_cast(kAllToAllGroupRanksPerNode) != + peer / static_cast(kAllToAllGroupRanksPerNode); } inline bool AllToAllGroupPeerInRouteStage( @@ -111,9 +151,8 @@ inline bool AllToAllGroupPeerInRouteStage( case AllToAllGroupRouteStage::kRemoteCopy: return crossNode; case AllToAllGroupRouteStage::kPrimary: - return crossNode && !AllToAllGroupUseSecondaryRoute(rank, peer); case AllToAllGroupRouteStage::kSecondary: - return crossNode && AllToAllGroupUseSecondaryRoute(rank, peer); + return crossNode; } return false; } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 5782d796..1b2be904 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -60,7 +60,7 @@ extern void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute); + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -654,17 +654,46 @@ bool RunGroupedAllToAll( << " chunkElements=" << requestedChunkElements << std::endl; return false; } - const int copyoutWorkersValue = GetEnvInt( - "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS", 16); - if (copyoutWorkersValue < 0 || !TileXR::Demo::AllToAllGroupValidCopyoutWorkers( - static_cast(copyoutWorkersValue))) { + const int channelModeValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE", 0); + if (channelModeValue < 0 || !TileXR::Demo::AllToAllGroupValidChannelMode( + static_cast(channelModeValue))) { std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS must be 8, 16, 32, or 48, got " - << copyoutWorkersValue << std::endl; + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE" + << " must be 0 (auto), 1 (single), or 2 (multi), got " + << channelModeValue << std::endl; return false; } - const uint32_t copyoutWorkers = static_cast(copyoutWorkersValue); - const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim(copyoutWorkers); + const int useSecondaryRouteValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE", 1); + if (useSecondaryRouteValue != 0 && useSecondaryRouteValue != 1) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE" + << " must be 0 or 1, got " << useSecondaryRouteValue << std::endl; + return false; + } + const auto channelMode = static_cast( + channelModeValue); + const bool multiChannel = useSecondaryRouteValue != 0 && + TileXR::Demo::AllToAllGroupUseMultiChannel(plan.payloadPlaneBytes, channelMode); + + const int primaryRoutePartsValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS", -1); + if (primaryRoutePartsValue < -1 || primaryRoutePartsValue > 8) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS" + << " must be -1 (auto) or 0..8, got " + << primaryRoutePartsValue << std::endl; + return false; + } + const uint32_t primaryRouteParts = primaryRoutePartsValue < 0 ? + TileXR::Demo::kAllToAllGroupAutoPrimaryParts : + static_cast(primaryRoutePartsValue); + + constexpr uint32_t sendWorkers = TileXR::Demo::kAllToAllGroupSendWorkerCount; + constexpr uint32_t copyoutWorkers = 32U; + const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim( + sendWorkers, copyoutWorkers); const int routeStagesValue = GetEnvInt( "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES", 0); if (routeStagesValue != 0 && routeStagesValue != 1) { @@ -674,15 +703,6 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; - const int useSecondaryRouteValue = GetEnvInt( - "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE", 1); - if (useSecondaryRouteValue != 0 && useSecondaryRouteValue != 1) { - std::cerr << "[rank " << rank - << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE" - << " must be 0 or 1, got " << useSecondaryRouteValue << std::endl; - return false; - } - const uint32_t useSecondaryRoute = static_cast(useSecondaryRouteValue); constexpr size_t kRouteStageCount = 10U; const std::array stagedRouteStages {{ @@ -845,9 +865,13 @@ bool RunGroupedAllToAll( " passes=" + std::to_string(plan.passCount)); PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + " repeat=" + std::to_string(repeat) + + " channelMode=" + std::to_string(channelModeValue) + + " multiChannel=" + std::to_string(multiChannel ? 1 : 0) + + " primaryRouteParts=" + std::to_string(primaryRoutePartsValue) + + " sendWorkers=" + std::to_string(sendWorkers) + " copyoutWorkers=" + std::to_string(copyoutWorkers) + " blockDim=" + std::to_string(groupBlockDim) + - " useSecondaryRoute=" + std::to_string(useSecondaryRoute) + + " useSecondaryRoute=" + std::to_string(useSecondaryRouteValue) + " routeStages=" + std::to_string(routeStagesValue)); if (routeStages && @@ -867,8 +891,9 @@ bool RunGroupedAllToAll( plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], - reinterpret_cast(trace), traceIteration, copyoutWorkers, - static_cast(routeStage), useSecondaryRoute); + reinterpret_cast(trace), traceIteration, + static_cast(routeStage), + multiChannel ? 1U : 0U, primaryRouteParts); }; double totalUs = 0.0; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index f0947944..00b6c646 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -115,7 +115,11 @@ void TestPlan() CHECK_EQ(plan.payloadPlaneBytes, 128ULL * 1024ULL * 1024ULL); CHECK_EQ(plan.payloadOffset[0], 0ULL); CHECK_EQ(plan.payloadOffset[1] >= plan.payloadOffset[0] + plan.payloadPlaneBytes, true); - CHECK_EQ(plan.signalPlaneBytes, static_cast(rankSize) * 128ULL); + CHECK_EQ(TileXR::Demo::kAllToAllGroupRouteSignalStride, 512U); + CHECK_EQ(TileXR::Demo::kAllToAllGroupSignalSlotBytes, 1024U); + CHECK_EQ(plan.signalPlaneBytes, static_cast(rankSize) * 1024ULL); + CHECK_EQ(TileXR::Demo::AllToAllGroupSignalByteOffset(3U, 0U), 3072ULL); + CHECK_EQ(TileXR::Demo::AllToAllGroupSignalByteOffset(3U, 1U), 3584ULL); CHECK_EQ(plan.signalOffset[0] >= plan.payloadOffset[1] + plan.payloadPlaneBytes, true); CHECK_EQ(plan.signalOffset[1] >= plan.signalOffset[0] + plan.signalPlaneBytes, true); CHECK_EQ(plan.controlOffset >= plan.signalOffset[1] + plan.signalPlaneBytes, true); @@ -138,6 +142,26 @@ void TestPlan() rankSize, thirtyTwoMiBElements, thirtyTwoMiBElements).valid, false); } +void TestChannelPolicy() +{ + using TileXR::Demo::AllToAllGroupChannelMode; + constexpr size_t threshold = 150ULL * 1024ULL * 1024ULL; + CHECK_EQ(TileXR::Demo::AllToAllGroupValidChannelMode(0U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidChannelMode(1U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidChannelMode(2U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidChannelMode(3U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupUseMultiChannel( + threshold, AllToAllGroupChannelMode::kAuto), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupUseMultiChannel( + threshold + sizeof(int32_t), AllToAllGroupChannelMode::kAuto), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupUseMultiChannel( + threshold * 2U, AllToAllGroupChannelMode::kSingle), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupUseMultiChannel( + sizeof(int32_t), AllToAllGroupChannelMode::kMulti), true); + + CHECK_EQ(TileXR::Demo::kAllToAllGroupSendWorkerCount, 32U); +} + void TestScalePlanAndTraceCapacity() { constexpr int rankSize = 1024; @@ -174,37 +198,10 @@ void TestTokens() void TestDualRoutePeerPolicy() { using TileXR::Demo::AllToAllGroupIsCrossNode; - using TileXR::Demo::AllToAllGroupUseSecondaryRoute; CHECK_EQ(AllToAllGroupIsCrossNode(0, 7), false); CHECK_EQ(AllToAllGroupIsCrossNode(0, 8), true); CHECK_EQ(AllToAllGroupIsCrossNode(15, 8), false); CHECK_EQ(AllToAllGroupIsCrossNode(15, 0), true); - - for (int sourceNode = 0; sourceNode < 3; ++sourceNode) { - for (int targetNode = 0; targetNode < 3; ++targetNode) { - if (sourceNode == targetNode) { - continue; - } - int rowSecondary[8] = {}; - int columnSecondary[8] = {}; - for (int sourceLocal = 0; sourceLocal < 8; ++sourceLocal) { - for (int targetLocal = 0; targetLocal < 8; ++targetLocal) { - const int source = sourceNode * 8 + sourceLocal; - const int target = targetNode * 8 + targetLocal; - if (AllToAllGroupUseSecondaryRoute(source, target)) { - ++rowSecondary[sourceLocal]; - ++columnSecondary[targetLocal]; - } - } - } - for (int local = 0; local < 8; ++local) { - CHECK_EQ(rowSecondary[local], 2); - CHECK_EQ(columnSecondary[local], 2); - } - } - } - CHECK_EQ(AllToAllGroupUseSecondaryRoute(0, 14, false), false); - CHECK_EQ(AllToAllGroupUseSecondaryRoute(0, 14, true), true); } void TestDualRouteQpWeights() @@ -227,6 +224,39 @@ void TestDualRouteQpWeights() const auto empty = TileXR::Demo::AllToAllGroupSelectRouteQps(nullptr, 0U); CHECK_EQ(empty.primaryQp, 0U); CHECK_EQ(empty.secondaryQp, 0U); + + const auto automatic = TileXR::Demo::AllToAllGroupSplitByRoute( + 10U, 6U, 2U, TileXR::Demo::kAllToAllGroupAutoPrimaryParts); + CHECK_EQ(automatic.primaryElements, 7U); + CHECK_EQ(automatic.secondaryElements, 3U); + const auto tiny = TileXR::Demo::AllToAllGroupSplitByRoute( + 1U, 6U, 2U, TileXR::Demo::kAllToAllGroupAutoPrimaryParts); + CHECK_EQ(tiny.primaryElements, 0U); + CHECK_EQ(tiny.secondaryElements, 1U); + const auto overrideSplit = TileXR::Demo::AllToAllGroupSplitByRoute( + 10U, 6U, 2U, 5U); + CHECK_EQ(overrideSplit.primaryElements, 6U); + CHECK_EQ(overrideSplit.secondaryElements, 4U); + const auto noSecondary = TileXR::Demo::AllToAllGroupSplitByRoute( + 10U, 6U, 0U, 5U); + CHECK_EQ(noSecondary.primaryElements, 10U); + CHECK_EQ(noSecondary.secondaryElements, 0U); + + const auto firstPrimary = TileXR::Demo::AllToAllGroupRouteSliceForPass( + 10U, 0U, 5U, automatic.primaryElements, 0U); + const auto firstSecondary = TileXR::Demo::AllToAllGroupRouteSliceForPass( + 10U, 0U, 5U, automatic.primaryElements, 1U); + const auto secondPrimary = TileXR::Demo::AllToAllGroupRouteSliceForPass( + 10U, 5U, 5U, automatic.primaryElements, 0U); + const auto secondSecondary = TileXR::Demo::AllToAllGroupRouteSliceForPass( + 10U, 5U, 5U, automatic.primaryElements, 1U); + CHECK_EQ(firstPrimary.elementOffset, 0U); + CHECK_EQ(firstPrimary.elements, 5U); + CHECK_EQ(firstSecondary.elements, 0U); + CHECK_EQ(secondPrimary.elementOffset, 5U); + CHECK_EQ(secondPrimary.elements, 2U); + CHECK_EQ(secondSecondary.elementOffset, 7U); + CHECK_EQ(secondSecondary.elements, 3U); } void TestRouteStages() @@ -307,8 +337,8 @@ void TestRouteStages() rank, peer, AllToAllGroupRouteStage::kRemoteCopy), true); CHECK_EQ(TileXR::Demo::AllToAllGroupPeerInRouteStage( rank, peer, AllToAllGroupRouteStage::kNoCopy), true); - CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary) + - static_cast(inSecondary), 1); + CHECK_EQ(static_cast(inLocal) + static_cast(inPrimary), 1); + CHECK_EQ(inSecondary, inPrimary); local += inLocal ? 1 : 0; primary += inPrimary ? 1 : 0; secondary += inSecondary ? 1 : 0; @@ -316,8 +346,8 @@ void TestRouteStages() rank, peer, AllToAllGroupRouteStage::kCombined), true); } CHECK_EQ(local, 7); - CHECK_EQ(primary + secondary, rankSize - 8); - CHECK_EQ(secondary, rankSize == 8 ? 0 : 2 * (rankSize / 8 - 1)); + CHECK_EQ(primary, rankSize - 8); + CHECK_EQ(secondary, rankSize - 8); } } } @@ -331,11 +361,9 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(48U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); - CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(8U), 24U); - CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U), 32U); - CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U), 48U); - CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(48U), 64U); - CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(4U), 0U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U, 32U), 64U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U, 48U), 0U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U, 48U), 0U); std::set lanes; for (uint32_t worker = 0U; worker < 8U; ++worker) { @@ -370,13 +398,14 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES"); CHECK_CONTAINS(kernel, "#include \"tilexr_udma_alltoall_group_route.h\""); CHECK_CONTAINS(kernel, "AllToAllGroupSelectRouteQps"); - CHECK_CONTAINS(kernel, - "AllToAllGroupUseSecondaryRouteDevice(rank, peer, useSecondaryRoute)"); + CHECK_CONTAINS(kernel, "AllToAllGroupSplitByRouteDevice"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE"); + CHECK_CONTAINS(kernel, "AllToAllGroupWaitRouteTokensMte"); CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); - CHECK_CONTAINS(kernel, - "uint32_t copyoutWorkers, uint32_t routeStage, uint32_t useSecondaryRoute"); + CHECK_CONTAINS(kernel, "constexpr uint32_t copyoutWorkers = 32U"); + CHECK_CONTAINS(kernel, "uint32_t multiChannel, uint32_t primaryRouteParts"); CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); CHECK_CONTAINS(kernel, "if (!AllToAllGroupPeerInRouteStageDevice(rank, peer, routeStage))"); @@ -385,9 +414,9 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupRemoteAssistDevice"); CHECK_CONTAINS(kernel, "copySliceCount"); CHECK_CONTAINS(kernel, "copySliceIndex"); - CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + copyoutWorkers"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers"); CHECK_CONTAINS(kernel, "const uint32_t traceCore = copyoutWorkers == 8U"); - CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES + lane : blockIdx"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); @@ -413,11 +442,11 @@ void TestHostStructure() CHECK_CONTAINS(demo, "PlanAllToAllGroup"); CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_group"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); - CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_COPYOUT_WORKERS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); - CHECK_CONTAINS(demo, "AllToAllGroupValidCopyoutWorkers"); - CHECK_CONTAINS(demo, "AllToAllGroupBlockDim(copyoutWorkers)"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS"); + CHECK_CONTAINS(demo, "kAllToAllGroupSendWorkerCount"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); CHECK_CONTAINS(demo, "grouped alltoall warmup="); const size_t begin = demo.find("bool RunGroupedAllToAll("); @@ -457,6 +486,7 @@ int main() { TestSchedules(); TestPlan(); + TestChannelPolicy(); TestScalePlanAndTraceCapacity(); TestTokens(); TestDualRoutePeerPolicy(); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py index 0dad684b..14d92139 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_local_stages.py @@ -26,12 +26,13 @@ def test_local_route_is_split_into_send_and_copy_stages(self): self.assertIn("AllToAllGroupRemoteAssistDevice", kernel) self.assertIn("copySliceCount", kernel) self.assertIn("copySliceIndex", kernel) - self.assertIn("must be 8, 16, 32, or 48", host) self.assertIn("TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE", host) - self.assertIn( - "AllToAllGroupUseSecondaryRouteDevice(rank, peer, useSecondaryRoute)", - kernel, - ) + self.assertIn("TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE", host) + self.assertIn("TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS", host) + self.assertIn("AllToAllGroupWaitRouteTokensMte", kernel) + self.assertIn("TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE", kernel) + self.assertIn("AllToAllGroupSplitByRouteDevice", kernel) + self.assertIn("AllToAllGroupSplitByRouteDevice", kernel) self.assertIn( '"local-send", "local-copy", "remote-send", "all-send", ' '"remote-wait", "remote-copy", "no-copy", "primary", "secondary", ' From 9f9e70ff50c0d68c22517e70883b6dda8ae83651 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 23 Jul 2026 16:34:02 +0800 Subject: [PATCH 119/163] fix(udma): identify nodes independently of hostname --- src/comm/udma/tilexr_udma_layout.cpp | 23 +++++++++++++++++++ src/comm/udma/tilexr_udma_layout.h | 6 +++++ src/comm/udma/tilexr_udma_transport.cpp | 13 ++++++++++- .../test_tilexr_udma_transport_layout.cpp | 21 +++++++++++++++++ 4 files changed, 62 insertions(+), 1 deletion(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index e71751b7..8d2e7f23 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -6,6 +6,7 @@ #include "udma/tilexr_udma_layout.h" #include +#include #include #include #include @@ -171,4 +172,26 @@ std::vector SelectUDMARoutesForPeer( return aggregateRoutes; } +std::string SelectUDMANodeIdentity( + const char* explicitNodeId, + const std::string& machineId, + const char* hostname) +{ + if (explicitNodeId != nullptr && explicitNodeId[0] != '\0') { + return std::string("explicit:") + explicitNodeId; + } + + const auto begin = std::find_if_not(machineId.begin(), machineId.end(), + [](unsigned char ch) { return std::isspace(ch) != 0; }); + const auto end = std::find_if_not(machineId.rbegin(), machineId.rend(), + [](unsigned char ch) { return std::isspace(ch) != 0; }).base(); + if (begin < end) { + return std::string("machine:") + std::string(begin, end); + } + if (hostname != nullptr && hostname[0] != '\0') { + return std::string("hostname:") + hostname; + } + return {}; +} + } // namespace TileXR diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index c300c620..68f0f0eb 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -8,6 +8,7 @@ #include #include +#include #include #include "tilexr_udma_types.h" @@ -57,6 +58,11 @@ std::vector SelectUDMARoutesForPeer( const std::vector& topoRoutes, const std::vector& aggregateRoutes); +std::string SelectUDMANodeIdentity( + const char* explicitNodeId, + const std::string& machineId, + const char* hostname); + } // namespace TileXR #endif // TILEXR_UDMA_LAYOUT_H diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index cad130ec..4fa3965e 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -88,7 +88,18 @@ struct UDMANodeId { UDMANodeId GetLocalNodeId() { UDMANodeId node {}; - if (gethostname(node.value, sizeof(node.value) - 1) != 0 || node.value[0] == '\0') { + char hostname[kUDMANodeIdBytes] = {}; + if (gethostname(hostname, sizeof(hostname) - 1) != 0) { + hostname[0] = '\0'; + } + std::ifstream machineIdFile("/etc/machine-id"); + std::string machineId; + std::getline(machineIdFile, machineId); + const std::string identity = SelectUDMANodeIdentity( + std::getenv("TILEXR_UDMA_NODE_ID"), machineId, hostname); + if (!identity.empty()) { + std::snprintf(node.value, sizeof(node.value), "%s", identity.c_str()); + } else { std::snprintf(node.value, sizeof(node.value), "pid-%ld", static_cast(getpid())); } node.value[sizeof(node.value) - 1] = '\0'; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index eea1254f..5d3934da 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -194,6 +194,25 @@ void TestSameNodeRouteSelectionUsesTopoRoutes() CHECK_EQ(selected[0], 1U); } +void TestNodeIdentityUsesMachineIdBeforeHostname() +{ + const std::string first = TileXR::SelectUDMANodeIdentity( + nullptr, "machine-a\n", "localhost.localdomain"); + const std::string second = TileXR::SelectUDMANodeIdentity( + nullptr, "machine-b\n", "localhost.localdomain"); + + CHECK_EQ(first, std::string("machine:machine-a")); + CHECK_EQ(second, std::string("machine:machine-b")); + CHECK_TRUE(first != second); +} + +void TestExplicitNodeIdentityOverridesMachineId() +{ + CHECK_EQ(TileXR::SelectUDMANodeIdentity( + "node-7", "machine-a\n", "localhost.localdomain"), + std::string("explicit:node-7")); +} + void TestTransportUsesPerPeerQueues() { const std::string transport = @@ -277,6 +296,8 @@ int main() TestExplicitRouteSelectionRejectsMissingInputs(); TestCrossNodeRouteSelectionUsesAggregateRoutes(); TestSameNodeRouteSelectionUsesTopoRoutes(); + TestNodeIdentityUsesMachineIdBeforeHostname(); + TestExplicitNodeIdentityOverridesMachineId(); TestTransportUsesPerPeerQueues(); TestRootInfoEidBytesSelectRuntimeContexts(); TestMemoryRegistrationUsesOfficialUbFlags(); From fd9124d811f75b362325d451fcc5eb588225fd11 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 23 Jul 2026 20:27:15 +0800 Subject: [PATCH 120/163] feat(udma): add grouped alltoall shared qp pool --- src/comm/udma/tilexr_udma_layout.cpp | 25 ++ src/comm/udma/tilexr_udma_layout.h | 10 + src/comm/udma/tilexr_udma_transport.cpp | 252 +++++++++++++++++- src/comm/udma/tilexr_udma_transport.h | 5 + .../test_tilexr_udma_transport_layout.cpp | 77 ++++++ 5 files changed, 358 insertions(+), 11 deletions(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 8d2e7f23..593ad076 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -127,6 +127,31 @@ std::vector BuildUDMAMultiRouteQpWeights( return qpWeights; } +uint32_t UDMASharedQpLane( + int rank, + int peer, + int rankSize, + uint32_t laneCount) +{ + if (rankSize <= 1 || rank < 0 || peer < 0 || rank >= rankSize || peer >= rankSize || + rank == peer || laneCount == 0 || (laneCount % 2) != 0) { + return laneCount; + } + + const uint32_t lanesPerDirection = laneCount / 2; + const uint32_t forward = static_cast((peer - rank + rankSize) % rankSize); + const uint32_t backward = static_cast(rankSize) - forward; + if (forward <= backward) { + return (forward - 1) % lanesPerDirection; + } + return lanesPerDirection + (backward - 1) % lanesPerDirection; +} + +size_t UDMASharedQpPoolSize(uint32_t laneCount, uint32_t eidCount) +{ + return static_cast(laneCount) * eidCount; +} + std::vector SelectExplicitUDMARouteEids( const char* routeList, const std::vector& candidateEids) diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index 68f0f0eb..e8e83e9c 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -49,6 +49,16 @@ std::vector BuildUDMAMultiRouteQpWeights( const std::map& routeWeights, uint32_t qpsPerRoute); +uint32_t UDMASharedQpLane( + int rank, + int peer, + int rankSize, + uint32_t laneCount); + +size_t UDMASharedQpPoolSize( + uint32_t laneCount, + uint32_t eidCount); + std::vector SelectExplicitUDMARouteEids( const char* routeList, const std::vector& candidateEids); diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 4fa3965e..396dbff1 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -45,6 +45,12 @@ bool UDMADiagEnabled() return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; } +bool EnvEnabled(const char* name) +{ + const char* value = std::getenv(name); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + uint32_t GetEnvUint(const char* name, uint32_t defaultValue, uint32_t minValue, uint32_t maxValue) { const char* value = std::getenv(name); @@ -433,10 +439,17 @@ struct TileXRUDMATransport::PerEidState { UDMACQCtx localCq {}; }; + struct SharedRemoteQueueState { + void* remoteQpHandle = nullptr; + uint32_t tpn = 0; + }; + uint32_t eidIndex = 0; void* ctxHandle = nullptr; void* tokenHandle = nullptr; std::map peerQueues; + std::map sharedQueues; + std::map sharedRemoteQueues; }; TileXRUDMATransport::TileXRUDMATransport() = default; @@ -528,7 +541,8 @@ int TileXRUDMATransport::OpenDevice() int TileXRUDMATransport::BuildRoutes() { - qpsPerRoute_ = GetEnvUint("TILEXR_UDMA_QP_NUM", 1, 1, 64); + sharedQpPool_ = EnvEnabled("TILEXR_UDMA_SHARED_QP_POOL"); + qpsPerRoute_ = sharedQpPool_ ? 1 : GetEnvUint("TILEXR_UDMA_QP_NUM", 1, 1, 64); qpNum_ = qpsPerRoute_; const char* routePolicy = std::getenv("TILEXR_UDMA_ROUTE_POLICY"); const bool useAllRoutes = routePolicy != nullptr && std::strcmp(routePolicy, "all") == 0; @@ -860,6 +874,18 @@ int TileXRUDMATransport::CreateQueues() state.eidIndex = ctxEntry.first; state.ctxHandle = ctxEntry.second; state.tokenHandle = tokenHandleByEid_[ctxEntry.first]; + states_[state.eidIndex] = state; + auto& activeState = states_[state.eidIndex]; + if (sharedQpPool_) { + for (uint32_t lane = 0; lane < sharedQpLaneCount_; ++lane) { + int ret = CreateSharedQueue(activeState, lane); + if (ret != TILEXR_SUCCESS) { + CleanupQueues(); + return ret; + } + } + continue; + } for (int peer = 0; peer < options_.rankSize; ++peer) { if (peer == options_.rank) { continue; @@ -870,8 +896,9 @@ int TileXRUDMATransport::CreateQueues() if (qpRoutesIt->second[qpIdx] != state.eidIndex) { continue; } - int ret = CreatePeerQueue(state, peer, qpIdx); + int ret = CreatePeerQueue(activeState, peer, qpIdx); if (ret != TILEXR_SUCCESS) { + CleanupQueues(); return ret; } } @@ -881,28 +908,64 @@ int TileXRUDMATransport::CreateQueues() if (localRouteIt == peerLocalEid_.end() || localRouteIt->second != state.eidIndex) { continue; } - int ret = CreatePeerQueue(state, peer, 0); + int ret = CreatePeerQueue(activeState, peer, 0); if (ret != TILEXR_SUCCESS) { + CleanupQueues(); return ret; } } - states_[state.eidIndex] = state; } return states_.empty() ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; } int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx) +{ + const uint32_t localQpIdx = qpsPerRoute_ == 0 ? qpIdx : qpIdx % qpsPerRoute_; + return CreateQueue(state, peer, qpIdx, localQpIdx, false); +} + +int TileXRUDMATransport::CreateSharedQueue(PerEidState& state, uint32_t lane) +{ + return CreateQueue(state, -1, lane, lane, true); +} + +int TileXRUDMATransport::CreateQueue( + PerEidState& state, int peer, uint32_t qpIdx, uint32_t queueIndex, bool shared) { const bool diag = UDMADiagEnabled(); PerEidState::PeerQueueState queue {}; queue.peer = peer; queue.qpIdx = qpIdx; - const uint32_t localQpIdx = qpsPerRoute_ == 0 ? qpIdx : qpIdx % qpsPerRoute_; + auto cleanupQueue = [&]() { + if (queue.remoteQpHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxQpUnimport(state.ctxHandle, queue.remoteQpHandle); + queue.remoteQpHandle = nullptr; + } + if (queue.qpHandle != nullptr) { + loader_.RaCtxQpDestroy(queue.qpHandle); + queue.qpHandle = nullptr; + } + if (queue.cqHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxCqDestroy(state.ctxHandle, queue.cqHandle); + queue.cqHandle = nullptr; + } + if (queue.chanHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxChanDestroy(state.ctxHandle, queue.chanHandle); + queue.chanHandle = nullptr; + } + FreeDeviceScalar(queue.cqPiAddr); + FreeDeviceScalar(queue.cqCiAddr); + FreeDeviceScalar(queue.sqPiAddr); + FreeDeviceScalar(queue.sqCiAddr); + FreeDeviceScalar(queue.wqeCntAddr); + FreeDeviceScalar(queue.amoAddr); + }; ChanInfoT chanInfo {}; chanInfo.in.dataPlaneFlag.bs.poolCqCstm = 1; int ret = loader_.RaCtxChanCreate(state.ctxHandle, &chanInfo, &queue.chanHandle); if (ret != 0) { + cleanupQueue(); return TILEXR_ERROR_INTERNAL; } @@ -911,14 +974,16 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t queue.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; ret = loader_.RaCtxCqCreate(state.ctxHandle, &queue.cqInfo, &queue.cqHandle); if (ret != 0) { + cleanupQueue(); return TILEXR_ERROR_INTERNAL; } - queue.localCq.cqn = localQpIdx; + queue.localCq.cqn = queueIndex; queue.localCq.bufAddr = queue.cqInfo.out.bufAddr; queue.localCq.baseBkShift = Log2Uint64(queue.cqInfo.out.cqeSize); queue.localCq.depth = queue.cqInfo.in.depth; if (AllocDeviceScalar(&queue.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || AllocDeviceScalar(&queue.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { + cleanupQueue(); return TILEXR_ERROR_INTERNAL; } queue.localCq.headAddr = reinterpret_cast(queue.cqPiAddr); @@ -944,9 +1009,10 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t qpAttr.ub.tokenIdHandle = state.tokenHandle; ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &queue.qpInfo, &queue.qpHandle); if (ret != 0) { + cleanupQueue(); return TILEXR_ERROR_INTERNAL; } - queue.localWq.wqn = localQpIdx; + queue.localWq.wqn = queueIndex; queue.localWq.bufAddr = queue.qpInfo.ub.sqBuffVa; queue.localWq.baseBkShift = Log2Uint64(queue.qpInfo.ub.wqebbSize); queue.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; @@ -954,6 +1020,7 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t AllocDeviceScalar(&queue.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || AllocDeviceScalar(&queue.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || AllocDeviceScalar(&queue.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { + cleanupQueue(); return TILEXR_ERROR_INTERNAL; } queue.localWq.headAddr = reinterpret_cast(queue.sqPiAddr); @@ -963,7 +1030,8 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t queue.localWq.wqeCntAddr = reinterpret_cast(queue.wqeCntAddr); queue.localWq.amoAddr = reinterpret_cast(queue.amoAddr); if (diag) { - TILEXR_LOG(INFO) << "UDMA diag create peer queue rank " << options_.rank + TILEXR_LOG(INFO) << "UDMA diag create " << (shared ? "shared" : "peer") + << " queue rank " << options_.rank << " peer=" << peer << " qpIdx=" << qpIdx << " eid=" << state.eidIndex @@ -983,7 +1051,9 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t << " wqebbSize=" << queue.qpInfo.ub.wqebbSize << " cqeSize=" << queue.cqInfo.out.cqeSize; } - if (qpNum_ <= 1) { + if (shared) { + state.sharedQueues[qpIdx] = queue; + } else if (qpNum_ <= 1) { state.peerQueues[peer] = queue; } else { state.peerQueues[QueueKey(peer, qpIdx, qpNum_)] = queue; @@ -993,6 +1063,9 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t int TileXRUDMATransport::ImportQueues() { + if (sharedQpPool_) { + return ImportSharedQueues(); + } const bool diag = UDMADiagEnabled(); std::vector localPeerImports(options_.rankSize); std::vector localPeerKeys(options_.rankSize); @@ -1077,6 +1150,109 @@ int TileXRUDMATransport::ImportQueues() return TILEXR_SUCCESS; } +int TileXRUDMATransport::ImportSharedQueues() +{ + const bool diag = UDMADiagEnabled(); + const size_t endpointSlots = UDMASharedQpPoolSize(sharedQpLaneCount_, eidCount_); + if (endpointSlots == 0) { + return TILEXR_ERROR_INTERNAL; + } + + std::vector localSharedImports(endpointSlots); + std::vector localSharedKeys(endpointSlots); + for (const auto& stateEntry : states_) { + const auto& state = stateEntry.second; + for (const auto& queueEntry : state.sharedQueues) { + const uint32_t lane = queueEntry.first; + const auto& queue = queueEntry.second; + const size_t slot = static_cast(state.eidIndex) * sharedQpLaneCount_ + lane; + if (slot >= endpointSlots) { + return TILEXR_ERROR_INTERNAL; + } + auto& importInfo = localSharedImports[slot]; + importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; + importInfo.in.ub.type = TARGET_TYPE_JETTY; + importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + importInfo.in.ub.tpType = 1; + importInfo.in.key = queue.qpInfo.key; + localSharedKeys[slot] = queue.qpInfo.key; + } + } + + std::vector allSharedImports(static_cast(options_.rankSize) * endpointSlots); + int ret = options_.exchange->AllGather( + localSharedImports.data(), localSharedImports.size(), allSharedImports.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } + std::vector allSharedKeys(static_cast(options_.rankSize) * endpointSlots); + ret = options_.exchange->AllGather( + localSharedKeys.data(), localSharedKeys.size(), allSharedKeys.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } + + for (int peer = 0; peer < options_.rankSize; ++peer) { + if (peer == options_.rank) { + continue; + } + const uint32_t localLane = UDMASharedQpLane(options_.rank, peer, + options_.rankSize, sharedQpLaneCount_); + const uint32_t remoteLane = UDMASharedQpLane(peer, options_.rank, + options_.rankSize, sharedQpLaneCount_); + if (localLane >= sharedQpLaneCount_ || remoteLane >= sharedQpLaneCount_) { + return TILEXR_ERROR_INTERNAL; + } + + const auto localRoutesIt = peerLocalEids_.find(peer); + const auto remoteRoutesIt = peerRemoteEids_.find(peer); + if (localRoutesIt == peerLocalEids_.end() || remoteRoutesIt == peerRemoteEids_.end()) { + return TILEXR_ERROR_INTERNAL; + } + const auto& localRoutes = localRoutesIt->second; + const auto& remoteRoutes = remoteRoutesIt->second; + const size_t routeCount = std::min(localRoutes.size(), remoteRoutes.size()); + for (size_t route = 0; route < routeCount; ++route) { + const uint32_t localEid = localRoutes[route]; + const uint32_t remoteEid = remoteRoutes[route]; + auto stateIt = states_.find(localEid); + if (stateIt == states_.end() || stateIt->second.sharedQueues.count(localLane) == 0 || + remoteEid >= eidCount_) { + return TILEXR_ERROR_INTERNAL; + } + const size_t remoteSlot = static_cast(peer) * endpointSlots + + static_cast(remoteEid) * sharedQpLaneCount_ + remoteLane; + if (remoteSlot >= allSharedImports.size()) { + return TILEXR_ERROR_INTERNAL; + } + QpImportInfoT importInfo = allSharedImports[remoteSlot]; + importInfo.in.key = allSharedKeys[remoteSlot]; + PerEidState::SharedRemoteQueueState remote {}; + ret = loader_.RaCtxQpImport(stateIt->second.ctxHandle, &importInfo, &remote.remoteQpHandle); + if (ret != 0 || remote.remoteQpHandle == nullptr) { + return TILEXR_ERROR_INTERNAL; + } + remote.tpn = importInfo.out.ub.tpn; + const uint32_t qpIdx = static_cast(route) * qpsPerRoute_; + stateIt->second.sharedRemoteQueues[QueueKey(peer, qpIdx, qpNum_)] = remote; + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag import shared qp rank " << options_.rank + << " peer=" << peer + << " qpIdx=" << qpIdx + << " localLane=" << localLane + << " remoteLane=" << remoteLane + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " remoteQp=" << remote.remoteQpHandle + << " tpn=" << remote.tpn; + } + } + } + return TILEXR_SUCCESS; +} + uint32_t TileXRUDMATransport::FallbackLocalEid() const { if (!states_.empty()) { @@ -1179,12 +1355,35 @@ int TileXRUDMATransport::RefreshUDMAInfo() } auto& state = stateIt->second; PerEidState::PeerQueueState* queuePtr = nullptr; + PerEidState::SharedRemoteQueueState* sharedRemotePtr = nullptr; if (rank == options_.rank) { - if (!state.peerQueues.empty()) { + if (sharedQpPool_ && !state.sharedQueues.empty()) { + queuePtr = &state.sharedQueues.begin()->second; + } else if (sharedQpPool_ && !fallbackIt->second.sharedQueues.empty()) { + queuePtr = &fallbackIt->second.sharedQueues.begin()->second; + } else if (!state.peerQueues.empty()) { queuePtr = &state.peerQueues.begin()->second; } else if (!fallbackIt->second.peerQueues.empty()) { queuePtr = &fallbackIt->second.peerQueues.begin()->second; } + } else if (sharedQpPool_) { + const uint32_t lane = UDMASharedQpLane( + options_.rank, rank, options_.rankSize, sharedQpLaneCount_); + auto queueIt = state.sharedQueues.find(lane); + if (queueIt == state.sharedQueues.end()) { + return TILEXR_ERROR_INTERNAL; + } + queuePtr = &queueIt->second; + const int key = QueueKey(rank, qpIdx, qpNum_); + auto remoteIt = state.sharedRemoteQueues.find(key); + if (remoteIt == state.sharedRemoteQueues.end()) { + const int routeKey = QueueKey(rank, routeQpIdx, qpNum_); + remoteIt = state.sharedRemoteQueues.find(routeKey); + } + if (remoteIt == state.sharedRemoteQueues.end()) { + return TILEXR_ERROR_INTERNAL; + } + sharedRemotePtr = &remoteIt->second; } else { const int key = QueueKey(rank, qpIdx, qpNum_); auto queueIt = state.peerQueues.find(key); @@ -1220,7 +1419,7 @@ int TileXRUDMATransport::RefreshUDMAInfo() } } else { mem[entryIndex] = allMem[rank * eidCount_ + remoteEid]; - mem[entryIndex].tpn = queue.tpn; + mem[entryIndex].tpn = sharedRemotePtr == nullptr ? queue.tpn : sharedRemotePtr->tpn; } mem[entryIndex].eidAddr = reinterpret_cast( eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); @@ -1593,6 +1792,14 @@ void TileXRUDMATransport::CleanupQueues() { for (auto& stateEntry : states_) { auto& state = stateEntry.second; + for (auto& remoteEntry : state.sharedRemoteQueues) { + auto& remote = remoteEntry.second; + if (remote.remoteQpHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxQpUnimport(state.ctxHandle, remote.remoteQpHandle); + remote.remoteQpHandle = nullptr; + } + } + state.sharedRemoteQueues.clear(); for (auto& queueEntry : state.peerQueues) { auto& queue = queueEntry.second; if (queue.remoteQpHandle != nullptr && state.ctxHandle != nullptr) { @@ -1619,6 +1826,28 @@ void TileXRUDMATransport::CleanupQueues() FreeDeviceScalar(queue.amoAddr); } state.peerQueues.clear(); + for (auto& queueEntry : state.sharedQueues) { + auto& queue = queueEntry.second; + if (queue.qpHandle != nullptr) { + loader_.RaCtxQpDestroy(queue.qpHandle); + queue.qpHandle = nullptr; + } + if (queue.cqHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxCqDestroy(state.ctxHandle, queue.cqHandle); + queue.cqHandle = nullptr; + } + if (queue.chanHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxChanDestroy(state.ctxHandle, queue.chanHandle); + queue.chanHandle = nullptr; + } + FreeDeviceScalar(queue.cqPiAddr); + FreeDeviceScalar(queue.cqCiAddr); + FreeDeviceScalar(queue.sqPiAddr); + FreeDeviceScalar(queue.sqCiAddr); + FreeDeviceScalar(queue.wqeCntAddr); + FreeDeviceScalar(queue.amoAddr); + } + state.sharedQueues.clear(); } states_.clear(); } @@ -1679,6 +1908,7 @@ void TileXRUDMATransport::Shutdown() peerQpRouteWeights_.clear(); qpsPerRoute_ = 1; qpNum_ = 1; + sharedQpPool_ = false; localMemInfoByEid_.clear(); remoteMemHandlesByPeer_.clear(); loader_.Unload(); diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index 40c2e6e6..3f24dc78 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -54,7 +54,10 @@ class TileXRUDMATransport { int CreateContexts(); int CreateQueues(); int CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx); + int CreateSharedQueue(PerEidState& state, uint32_t lane); + int CreateQueue(PerEidState& state, int peer, uint32_t qpIdx, uint32_t queueIndex, bool shared); int ImportQueues(); + int ImportSharedQueues(); int EnsureUDMAInfoBuffer(); int RefreshUDMAInfo(); int RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes); @@ -77,6 +80,8 @@ class TileXRUDMATransport { uint32_t eidCount_ = 0; uint32_t qpNum_ = 1; uint32_t qpsPerRoute_ = 1; + bool sharedQpPool_ = false; + uint32_t sharedQpLaneCount_ = 16; std::map ctxHandleByEid_; std::map tokenHandleByEid_; std::map peerLocalEid_; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 5d3934da..098ede94 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -153,6 +154,59 @@ void TestMultiRouteQpWeightsUseRouteBandwidth() CHECK_EQ(qpWeights[1], 2U); } +void TestSharedQpLaneMatchesGroupedPeerOrder() +{ + CHECK_EQ(TileXR::UDMASharedQpLane(0, 1, 64, 16), 0U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 8, 64, 16), 7U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 9, 64, 16), 0U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 63, 64, 16), 8U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 56, 64, 16), 15U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 32, 64, 16), 7U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 0, 64, 16), 16U); +} + +void TestSharedQpLanesAreUniqueWithinEveryGroup() +{ + for (int rankSize = 8; rankSize <= 1024; rankSize += 8) { + const uint32_t groupCount = static_cast((rankSize - 1 + 15) / 16); + for (int rank = 0; rank < rankSize; ++rank) { + for (uint32_t group = 0; group < groupCount; ++group) { + bool used[16] = {}; + for (int peer = 0; peer < rankSize; ++peer) { + if (peer == rank) { + continue; + } + const uint32_t forward = static_cast((peer - rank + rankSize) % rankSize); + const uint32_t backward = static_cast(rankSize) - forward; + const uint32_t distance = std::min(forward, backward); + if ((distance - 1U) / 8U != group) { + continue; + } + const uint32_t lane = TileXR::UDMASharedQpLane(rank, peer, rankSize, 16); + CHECK_TRUE(lane < 16U); + CHECK_TRUE(!used[lane]); + used[lane] = true; + } + } + } + } +} + +void TestSharedQpPoolScalesWithLanesAndEidsOnly() +{ + CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2), static_cast(32)); + CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 1), static_cast(16)); + CHECK_EQ(TileXR::UDMASharedQpPoolSize(0, 2), static_cast(0)); +} + +void TestSharedQpLaneRejectsInvalidInputs() +{ + CHECK_EQ(TileXR::UDMASharedQpLane(-1, 1, 64, 16), 16U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 64, 64, 16), 16U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 1, 64, 15), 15U); + CHECK_EQ(TileXR::UDMASharedQpLane(0, 1, 64, 0), 0U); +} + void TestExplicitRouteSelectionKeepsRequestedCandidateOrder() { const std::vector candidates = {7, 8}; @@ -229,6 +283,24 @@ void TestTransportUsesPerPeerQueues() CHECK_NOT_CONTAINS(transport, "void* qpHandle = nullptr;\n CqInfoT cqInfo"); } +void TestTransportHasOptInSharedQpPool() +{ + const std::string header = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.h"); + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + + CHECK_CONTAINS(header, "bool sharedQpPool_ = false"); + CHECK_CONTAINS(transport, "TILEXR_UDMA_SHARED_QP_POOL"); + CHECK_CONTAINS(transport, "state.sharedQueues"); + CHECK_CONTAINS(transport, "UDMASharedQpLane(options_.rank, peer"); + CHECK_CONTAINS(transport, "UDMASharedQpLane(peer, options_.rank"); + CHECK_CONTAINS(transport, "localSharedImports"); + CHECK_CONTAINS(transport, "allSharedImports"); + CHECK_CONTAINS(transport, "state.sharedRemoteQueues.clear()"); + CHECK_CONTAINS(transport, "auto cleanupQueue = [&]()"); +} + void TestRootInfoEidBytesSelectRuntimeContexts() { const std::string transport = @@ -292,6 +364,10 @@ int main() TestMultiRouteQpMappingRepeatsEachRoute(); TestMultiRouteQpMappingRejectsEmptyInputs(); TestMultiRouteQpWeightsUseRouteBandwidth(); + TestSharedQpLaneMatchesGroupedPeerOrder(); + TestSharedQpLanesAreUniqueWithinEveryGroup(); + TestSharedQpPoolScalesWithLanesAndEidsOnly(); + TestSharedQpLaneRejectsInvalidInputs(); TestExplicitRouteSelectionKeepsRequestedCandidateOrder(); TestExplicitRouteSelectionRejectsMissingInputs(); TestCrossNodeRouteSelectionUsesAggregateRoutes(); @@ -299,6 +375,7 @@ int main() TestNodeIdentityUsesMachineIdBeforeHostname(); TestExplicitNodeIdentityOverridesMachineId(); TestTransportUsesPerPeerQueues(); + TestTransportHasOptInSharedQpPool(); TestRootInfoEidBytesSelectRuntimeContexts(); TestMemoryRegistrationUsesOfficialUbFlags(); TestDeviceSgeUsesPerPeerLocalTokenId(); From e919e37353e1dbbcdb652534f9a18a5db6424585 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 23 Jul 2026 22:01:22 +0800 Subject: [PATCH 121/163] fix(comm): complete socket exchange transfers --- src/comm/tools/socket/tilexr_sock_exchange.h | 36 ++++++++++++++------ 1 file changed, 26 insertions(+), 10 deletions(-) diff --git a/src/comm/tools/socket/tilexr_sock_exchange.h b/src/comm/tools/socket/tilexr_sock_exchange.h index 56096b37..7fcaf07b 100644 --- a/src/comm/tools/socket/tilexr_sock_exchange.h +++ b/src/comm/tools/socket/tilexr_sock_exchange.h @@ -95,34 +95,50 @@ class TileXRSockExchange { return ((ioErrno == EAGAIN) || (ioErrno == EWOULDBLOCK) || (ioErrno == EINTR)); } - template int Send(int fd, const T *sendBuf, size_t sendSize, int flag) const + template ssize_t Send(int fd, const T *sendBuf, size_t sendSize, int flag) const { - do { - auto ret = send(fd, sendBuf, sendSize, flag); + const char *data = reinterpret_cast(sendBuf); + size_t sent = 0; + while (sent < sendSize) { + auto ret = send(fd, data + sent, sendSize - sent, flag); if (ret < 0) { if (CheckErrno(errno)) { TILEXR_LOG(ERROR) << "send failed: " << strerror(errno); continue; } TILEXR_LOG(DEBUG) << "Send failed: " << strerror(errno); + return ret; } - return ret; - } while (true); + if (ret == 0) { + TILEXR_LOG(DEBUG) << "Send returned zero before buffer completion"; + return -1; + } + sent += static_cast(ret); + } + return static_cast(sent); } - template int Recv(int fd, T *recvBuf, size_t recvSize, int flag) const + template ssize_t Recv(int fd, T *recvBuf, size_t recvSize, int flag) const { - do { - auto ret = recv(fd, recvBuf, recvSize, flag); + char *data = reinterpret_cast(recvBuf); + size_t received = 0; + while (received < recvSize) { + auto ret = recv(fd, data + received, recvSize - received, flag); if (ret < 0) { if (CheckErrno(errno)) { TILEXR_LOG(ERROR) << "recv failed: " << strerror(errno); continue; } TILEXR_LOG(DEBUG) << "recv failed: " << strerror(errno); + return ret; } - return ret; - } while (true); + if (ret == 0) { + TILEXR_LOG(DEBUG) << "Recv reached EOF before buffer completion"; + return -1; + } + received += static_cast(ret); + } + return static_cast(received); } template int ClientSendRecv(const T *sendBuf, size_t sendSize, T *recvBuf) From a8eceae1c83a44234755f09d536d37491ae2cedc Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 23 Jul 2026 23:42:55 +0800 Subject: [PATCH 122/163] fix(udma): exchange only required shared qp keys --- src/comm/tools/socket/tilexr_sock_exchange.h | 97 ++++++++++++++++ src/comm/udma/tilexr_udma_transport.cpp | 107 ++++++++++++------ .../test_tilexr_udma_transport_layout.cpp | 18 ++- 3 files changed, 186 insertions(+), 36 deletions(-) diff --git a/src/comm/tools/socket/tilexr_sock_exchange.h b/src/comm/tools/socket/tilexr_sock_exchange.h index 7fcaf07b..e7955d22 100644 --- a/src/comm/tools/socket/tilexr_sock_exchange.h +++ b/src/comm/tools/socket/tilexr_sock_exchange.h @@ -10,6 +10,9 @@ #ifndef TILEXR_SOCK_EXCHANGE_H #define TILEXR_SOCK_EXCHANGE_H +#include +#include +#include #include #include #include @@ -71,6 +74,20 @@ class TileXRSockExchange { } } + /* sendBuf is [destination rank][element], recvBuf is [source rank][element]. */ + template int AllToAll(const T *sendBuf, size_t sendCountPerRank, T *recvBuf) + { + if (!isInit_ && Prepare() != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + isInit_ = true; + + if (!IsServer()) { + return ClientSendRecvAllToAll(sendBuf, sendCountPerRank, recvBuf); + } + return ServerRecvSendAllToAll(sendBuf, sendCountPerRank, recvBuf); + } + int GetNodeNum(); static bool CheckValid(TileXRUniqueId tilexrCommId) @@ -178,6 +195,86 @@ class TileXRSockExchange { return TILEXR_SUCCESS; } + + template bool AllToAllSizes(size_t sendCountPerRank, size_t& rowCount, size_t& totalCount) const + { + if (rankSize_ <= 0 || sendCountPerRank == 0 || + sendCountPerRank > std::numeric_limits::max() / static_cast(rankSize_)) { + return false; + } + rowCount = sendCountPerRank * static_cast(rankSize_); + if (rowCount > std::numeric_limits::max() / sizeof(T) || + rowCount > std::numeric_limits::max() / static_cast(rankSize_)) { + return false; + } + totalCount = rowCount * static_cast(rankSize_); + return true; + } + + template int ClientSendRecvAllToAll( + const T *sendBuf, size_t sendCountPerRank, T *recvBuf) + { + size_t rowCount = 0; + size_t totalCount = 0; + if (!AllToAllSizes(sendCountPerRank, rowCount, totalCount)) { + return TILEXR_ERROR_INTERNAL; + } + (void)totalCount; + const size_t rowBytes = rowCount * sizeof(T); + if (Send(fd_, sendBuf, rowBytes, 0) <= 0) { + TILEXR_LOG(ERROR) << "Client side " << rank_ << " send alltoall buffer failed"; + return TILEXR_ERROR_INTERNAL; + } + if (Recv(fd_, recvBuf, rowBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Client side " << rank_ << " recv alltoall buffer failed"; + return TILEXR_ERROR_INTERNAL; + } + return TILEXR_SUCCESS; + } + + template int ServerRecvSendAllToAll( + const T *sendBuf, size_t sendCountPerRank, T *recvBuf) + { + size_t rowCount = 0; + size_t totalCount = 0; + if (!AllToAllSizes(sendCountPerRank, rowCount, totalCount)) { + return TILEXR_ERROR_INTERNAL; + } + const size_t rowBytes = rowCount * sizeof(T); + std::vector gathered; + std::vector destination; + try { + gathered.resize(totalCount); + destination.resize(rowCount); + } catch (const std::bad_alloc&) { + TILEXR_LOG(ERROR) << "Server side allocate alltoall exchange buffer failed"; + return TILEXR_ERROR_INTERNAL; + } + std::copy_n(sendBuf, rowCount, gathered.data()); + for (int source = 1; source < rankSize_; ++source) { + if (Recv(clientFds_[source], gathered.data() + static_cast(source) * rowCount, + rowBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Server side recv alltoall rank " << source << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + + for (int target = 0; target < rankSize_; ++target) { + for (int source = 0; source < rankSize_; ++source) { + const size_t sourceOffset = + (static_cast(source) * rankSize_ + target) * sendCountPerRank; + std::copy_n(gathered.data() + sourceOffset, sendCountPerRank, + destination.data() + static_cast(source) * sendCountPerRank); + } + if (target == 0) { + std::copy_n(destination.data(), rowCount, recvBuf); + } else if (Send(clientFds_[target], destination.data(), rowBytes, 0) <= 0) { + TILEXR_LOG(ERROR) << "Server side send alltoall rank " << target << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + return TILEXR_SUCCESS; + } int rank_ = 0; int rankSize_ = 0; int fd_ = -1; diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 396dbff1..13a2399f 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -1153,43 +1153,65 @@ int TileXRUDMATransport::ImportQueues() int TileXRUDMATransport::ImportSharedQueues() { const bool diag = UDMADiagEnabled(); - const size_t endpointSlots = UDMASharedQpPoolSize(sharedQpLaneCount_, eidCount_); - if (endpointSlots == 0) { + struct SharedQpKeyRecord { + uint32_t valid; + uint32_t eidIndex; + uint32_t lane; + uint32_t reserved; + QpKeyT key; + }; + + uint32_t localMaxRoutes = 0; + for (const auto& routes : peerLocalEids_) { + localMaxRoutes = std::max(localMaxRoutes, + static_cast(routes.second.size())); + } + std::vector allMaxRoutes(options_.rankSize); + int ret = options_.exchange->AllGather(&localMaxRoutes, 1, allMaxRoutes.data()); + if (ret != TILEXR_SUCCESS) { + return ret; + } + const uint32_t maxRoutes = *std::max_element(allMaxRoutes.begin(), allMaxRoutes.end()); + if (maxRoutes == 0 || sharedQpLaneCount_ == 0) { return TILEXR_ERROR_INTERNAL; } - std::vector localSharedImports(endpointSlots); - std::vector localSharedKeys(endpointSlots); - for (const auto& stateEntry : states_) { - const auto& state = stateEntry.second; - for (const auto& queueEntry : state.sharedQueues) { - const uint32_t lane = queueEntry.first; - const auto& queue = queueEntry.second; - const size_t slot = static_cast(state.eidIndex) * sharedQpLaneCount_ + lane; - if (slot >= endpointSlots) { + const size_t recordsPerRank = static_cast(options_.rankSize) * maxRoutes; + std::vector sendRecords(recordsPerRank); + std::vector recvRecords(recordsPerRank); + for (int peer = 0; peer < options_.rankSize; ++peer) { + if (peer == options_.rank) { + continue; + } + const auto routesIt = peerLocalEids_.find(peer); + if (routesIt == peerLocalEids_.end() || routesIt->second.size() > maxRoutes) { + return TILEXR_ERROR_INTERNAL; + } + const uint32_t lane = UDMASharedQpLane(options_.rank, peer, + options_.rankSize, sharedQpLaneCount_); + if (lane >= sharedQpLaneCount_) { + return TILEXR_ERROR_INTERNAL; + } + for (size_t route = 0; route < routesIt->second.size(); ++route) { + const uint32_t localEid = routesIt->second[route]; + const auto stateIt = states_.find(localEid); + if (stateIt == states_.end()) { return TILEXR_ERROR_INTERNAL; } - auto& importInfo = localSharedImports[slot]; - importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; - importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; - importInfo.in.ub.type = TARGET_TYPE_JETTY; - importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; - importInfo.in.ub.tpType = 1; - importInfo.in.key = queue.qpInfo.key; - localSharedKeys[slot] = queue.qpInfo.key; + const auto queueIt = stateIt->second.sharedQueues.find(lane); + if (queueIt == stateIt->second.sharedQueues.end()) { + return TILEXR_ERROR_INTERNAL; + } + auto& record = sendRecords[static_cast(peer) * maxRoutes + route]; + record.valid = 1; + record.eidIndex = localEid; + record.lane = lane; + record.key = queueIt->second.qpInfo.key; } } - std::vector allSharedImports(static_cast(options_.rankSize) * endpointSlots); - int ret = options_.exchange->AllGather( - localSharedImports.data(), localSharedImports.size(), allSharedImports.data()); - if (ret != TILEXR_SUCCESS) { - return ret; - } - std::vector allSharedKeys(static_cast(options_.rankSize) * endpointSlots); - ret = options_.exchange->AllGather( - localSharedKeys.data(), localSharedKeys.size(), allSharedKeys.data()); + ret = options_.exchange->AllToAll( + sendRecords.data(), maxRoutes, recvRecords.data()); if (ret != TILEXR_SUCCESS) { return ret; } @@ -1222,13 +1244,30 @@ int TileXRUDMATransport::ImportSharedQueues() remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - const size_t remoteSlot = static_cast(peer) * endpointSlots + - static_cast(remoteEid) * sharedQpLaneCount_ + remoteLane; - if (remoteSlot >= allSharedImports.size()) { + const size_t remoteSlot = static_cast(peer) * maxRoutes + route; + if (remoteSlot >= recvRecords.size()) { return TILEXR_ERROR_INTERNAL; } - QpImportInfoT importInfo = allSharedImports[remoteSlot]; - importInfo.in.key = allSharedKeys[remoteSlot]; + const auto& record = recvRecords[remoteSlot]; + if (record.valid == 0 || record.eidIndex != remoteEid || record.lane != remoteLane) { + TILEXR_LOG(ERROR) << "Shared QP exchange mismatch rank " << options_.rank + << " peer=" << peer + << " route=" << route + << " expectedEid=" << remoteEid + << " receivedEid=" << record.eidIndex + << " expectedLane=" << remoteLane + << " receivedLane=" << record.lane + << " valid=" << record.valid; + return TILEXR_ERROR_INTERNAL; + } + QpImportInfoT importInfo {}; + importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; + importInfo.in.ub.type = TARGET_TYPE_JETTY; + importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + importInfo.in.ub.tpType = 1; + importInfo.in.key = record.key; PerEidState::SharedRemoteQueueState remote {}; ret = loader_.RaCtxQpImport(stateIt->second.ctxHandle, &importInfo, &remote.remoteQpHandle); if (ret != 0 || remote.remoteQpHandle == nullptr) { diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 098ede94..c291a72d 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -295,12 +295,25 @@ void TestTransportHasOptInSharedQpPool() CHECK_CONTAINS(transport, "state.sharedQueues"); CHECK_CONTAINS(transport, "UDMASharedQpLane(options_.rank, peer"); CHECK_CONTAINS(transport, "UDMASharedQpLane(peer, options_.rank"); - CHECK_CONTAINS(transport, "localSharedImports"); - CHECK_CONTAINS(transport, "allSharedImports"); + CHECK_CONTAINS(transport, "SharedQpKeyRecord"); + CHECK_CONTAINS(transport, "options_.exchange->AllToAll("); + CHECK_NOT_CONTAINS(transport, "localSharedImports"); + CHECK_NOT_CONTAINS(transport, "allSharedImports"); + CHECK_NOT_CONTAINS(transport, "allSharedKeys"); CHECK_CONTAINS(transport, "state.sharedRemoteQueues.clear()"); CHECK_CONTAINS(transport, "auto cleanupQueue = [&]()"); } +void TestSocketExchangeSupportsPersonalizedAllToAll() +{ + const std::string exchange = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/tools/socket/tilexr_sock_exchange.h"); + + CHECK_CONTAINS(exchange, "int AllToAll(const T *sendBuf, size_t sendCountPerRank, T *recvBuf)"); + CHECK_CONTAINS(exchange, "ClientSendRecvAllToAll"); + CHECK_CONTAINS(exchange, "ServerRecvSendAllToAll"); +} + void TestRootInfoEidBytesSelectRuntimeContexts() { const std::string transport = @@ -368,6 +381,7 @@ int main() TestSharedQpLanesAreUniqueWithinEveryGroup(); TestSharedQpPoolScalesWithLanesAndEidsOnly(); TestSharedQpLaneRejectsInvalidInputs(); + TestSocketExchangeSupportsPersonalizedAllToAll(); TestExplicitRouteSelectionKeepsRequestedCandidateOrder(); TestExplicitRouteSelectionRejectsMissingInputs(); TestCrossNodeRouteSelectionUsesAggregateRoutes(); From 001fe00c2eeb8917a5c5db056c7ebfc611a2f84b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sat, 25 Jul 2026 00:13:13 +0800 Subject: [PATCH 123/163] fix(comm): shard socket exchange across rank groups --- .../tools/socket/tilexr_sock_exchange.cpp | 183 ++++++++++++---- src/comm/tools/socket/tilexr_sock_exchange.h | 200 ++++++++++++++++++ .../socket/tilexr_sock_exchange_layout.h | 33 +++ tests/udma/CMakeLists.txt | 9 + .../unit/test_tilexr_sock_exchange_layout.cpp | 75 +++++++ 5 files changed, 455 insertions(+), 45 deletions(-) create mode 100644 src/comm/tools/socket/tilexr_sock_exchange_layout.h create mode 100644 tests/udma/unit/test_tilexr_sock_exchange_layout.cpp diff --git a/src/comm/tools/socket/tilexr_sock_exchange.cpp b/src/comm/tools/socket/tilexr_sock_exchange.cpp index cdc954f2..c1ab8488 100644 --- a/src/comm/tools/socket/tilexr_sock_exchange.cpp +++ b/src/comm/tools/socket/tilexr_sock_exchange.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -96,43 +97,20 @@ string GetUUID() int TileXRSockExchange::GetNodeNum() { - if (!isInit_ && Prepare() != TILEXR_SUCCESS) { + constexpr size_t uuidBytes = 64; + std::array localUuid {}; + const string uuid = GetUUID(); + std::copy_n(uuid.data(), std::min(uuid.size(), uuidBytes - 1), localUuid.data()); + std::vector allUuids(static_cast(rankSize_) * uuidBytes); + if (AllGather(localUuid.data(), uuidBytes, allUuids.data()) != TILEXR_SUCCESS) { return TILEXR_ERROR_INTERNAL; } - isInit_ = true; - string uuid = GetUUID(); - TILEXR_LOG(DEBUG) << "rank:" << rank_ << " UUID " << uuid; - set uuidSet {}; - uuidSet.insert(uuid); - int nodeNum = -1; - if (IsServer()) { - for (int i = 1; i < rankSize_; ++i) { - if (Recv(clientFds_[i], &uuid[0], uuid.size(), 0) <= 0) { - TILEXR_LOG(ERROR) << "Server side recv rank " << i << " buffer failed"; - return TILEXR_ERROR_INTERNAL; - } - uuidSet.insert(uuid); - } - nodeNum = static_cast(uuidSet.size()); - TILEXR_LOG(DEBUG) << "nodeNum:" << nodeNum; - for (int i = 1; i < rankSize_; ++i) { - if (Send(clientFds_[i], &nodeNum, sizeof(int), 0) <= 0) { - TILEXR_LOG(ERROR) << "Server side send rank " << i << " buffer failed"; - return TILEXR_ERROR_INTERNAL; - } - } - } else { - if (Send(fd_, uuid.data(), uuid.size(), 0) <= 0) { - TILEXR_LOG(ERROR) << "Client side " << rank_ << " send buffer failed"; - return TILEXR_ERROR_INTERNAL; - } - if (Recv(fd_, &nodeNum, sizeof(int), 0) <= 0) { - TILEXR_LOG(ERROR) << "Client side " << rank_ << " recv buffer failed "; - return TILEXR_ERROR_INTERNAL; - } + for (int rank = 0; rank < rankSize_; ++rank) { + uuidSet.emplace(allUuids.data() + static_cast(rank) * uuidBytes); } - return nodeNum; + TILEXR_LOG(DEBUG) << "nodeNum:" << uuidSet.size(); + return static_cast(uuidSet.size()); } void TileXRSockExchange::GetIpAndPort() @@ -150,16 +128,94 @@ void TileXRSockExchange::GetIpAndPort() TILEXR_LOG(DEBUG) << "curRank: " << rank_ << " commDomain: " << commDomain_ << " ip: " << ip_ << " port: " << port_; } +int TileXRSockExchange::ConfigureHierarchy() +{ + hierarchical_ = false; + const char* enabled = std::getenv("TILEXR_SOCKET_HIERARCHICAL"); + if (enabled == nullptr || std::strcmp(enabled, "1") != 0) { + return TILEXR_SUCCESS; + } + const char* groupSizeEnv = std::getenv("TILEXR_SOCKET_GROUP_SIZE"); + const char* ranksPerHostEnv = std::getenv("TILEXR_SOCKET_RANKS_PER_HOST"); + hierarchyGroupSize_ = groupSizeEnv == nullptr ? 64 : std::atoi(groupSizeEnv); + hierarchyRanksPerHost_ = ranksPerHostEnv == nullptr ? 8 : std::atoi(ranksPerHostEnv); + if (hierarchyGroupSize_ <= 0 || hierarchyRanksPerHost_ <= 0 || + hierarchyGroupSize_ % hierarchyRanksPerHost_ != 0) { + TILEXR_LOG(ERROR) << "Invalid hierarchical socket group configuration"; + return TILEXR_ERROR_INTERNAL; + } + + const char* hostsEnv = std::getenv("TILEXR_COMM_HOSTS"); + if (hostsEnv == nullptr || hostsEnv[0] == '\0') { + TILEXR_LOG(ERROR) << "TILEXR_COMM_HOSTS is required for hierarchical socket exchange"; + return TILEXR_ERROR_INTERNAL; + } + hierarchyHosts_.clear(); + std::stringstream hostStream(hostsEnv); + string host; + while (std::getline(hostStream, host, ',')) { + if (!host.empty()) { + hierarchyHosts_.push_back(host); + } + } + const size_t requiredHosts = + (static_cast(rankSize_) + hierarchyRanksPerHost_ - 1) / + static_cast(hierarchyRanksPerHost_); + if (hierarchyHosts_.size() < requiredHosts) { + TILEXR_LOG(ERROR) << "Hierarchical socket host list too short: " + << hierarchyHosts_.size() << " < " << requiredHosts; + return TILEXR_ERROR_INTERNAL; + } + + groupLayout_ = BuildSockExchangeGroupLayout(rank_, rankSize_, hierarchyGroupSize_); + if (groupLayout_.groupCount == 0) { + return TILEXR_ERROR_INTERNAL; + } + const size_t leaderHost = static_cast(groupLayout_.groupLeader) / + static_cast(hierarchyRanksPerHost_); + listenAddr_ = tilexrCommId_.handle.addr.sin; + listenAddr_.sin_addr.s_addr = inet_addr(hierarchyHosts_[leaderHost].c_str()); + if (listenAddr_.sin_addr.s_addr == INADDR_NONE) { + TILEXR_LOG(ERROR) << "Invalid hierarchical leader host " << hierarchyHosts_[leaderHost]; + return TILEXR_ERROR_INTERNAL; + } + hierarchical_ = groupLayout_.groupCount > 1; + TILEXR_LOG(INFO) << "Hierarchical socket exchange rank=" << rank_ + << " group=" << groupLayout_.groupIndex + << " leader=" << groupLayout_.groupLeader + << " range=[" << groupLayout_.groupBegin << "," + << groupLayout_.groupEnd << ") groups=" << groupLayout_.groupCount; + return TILEXR_SUCCESS; +} + int TileXRSockExchange::Prepare() { if (tilexrCommId_.handle.magic != TILEXR_MAGIC) { GetIpAndPort(); } + if (ConfigureHierarchy() != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + clientFds_.resize(rankSize_, -1); + if (hierarchical_) { + if (!IsGroupLeader()) { + sockaddr_in leaderAddr = tilexrCommId_.handle.addr.sin; + const size_t leaderHost = static_cast(groupLayout_.groupLeader) / + static_cast(hierarchyRanksPerHost_); + leaderAddr.sin_addr.s_addr = inet_addr(hierarchyHosts_[leaderHost].c_str()); + return ConnectTo(leaderAddr); + } + if (Listen() != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + if (rank_ != 0 && Connect() != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + return AcceptHierarchical(); + } if (!IsServer()) { return Connect(); } - - clientFds_.resize(rankSize_, -1); if (Listen() != TILEXR_SUCCESS) { TILEXR_LOG(ERROR) << "Listen Failed!"; return TILEXR_ERROR_INTERNAL; @@ -175,20 +231,21 @@ int TileXRSockExchange::Prepare() int TileXRSockExchange::Listen() { - fd_ = socket(AF_INET, SOCK_STREAM, 0); - if (fd_ < 0) { + listenFd_ = socket(AF_INET, SOCK_STREAM, 0); + if (listenFd_ < 0) { TILEXR_LOG(ERROR) << "Server side create socket failed"; return TILEXR_ERROR_INTERNAL; } int reuse = 1; - if (setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(int)) < 0) { + if (setsockopt(listenFd_, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(int)) < 0) { TILEXR_LOG(ERROR) << "Server side set reuseaddr failed"; return TILEXR_ERROR_INTERNAL; } - struct sockaddr *addrPtr = &tilexrCommId_.handle.addr.sa; - if (bind(fd_, addrPtr, sizeof(struct sockaddr)) < 0) { + sockaddr_in bindAddr = hierarchical_ ? listenAddr_ : tilexrCommId_.handle.addr.sin; + struct sockaddr *addrPtr = reinterpret_cast(&bindAddr); + if (bind(listenFd_, addrPtr, sizeof(struct sockaddr)) < 0) { TILEXR_LOG(ERROR) << "Server side bind " << ntohs(tilexrCommId_.handle.addr.sin.sin_port) << " failed"; return TILEXR_ERROR_INTERNAL; } @@ -197,12 +254,12 @@ int TileXRSockExchange::Listen() * kernel would silently truncate backlog to the value defined in * /proc/sys/net/core/somaxconn if it is less than 65535. */ - if (listen(fd_, TILEXR_MAX_BACK_LOG) < 0) { + if (listen(listenFd_, TILEXR_MAX_BACK_LOG) < 0) { TILEXR_LOG(ERROR) << "Server side listen " << ntohs(tilexrCommId_.handle.addr.sin.sin_port) << " failed"; return TILEXR_ERROR_INTERNAL; } - TILEXR_LOG(INFO) << "The server is listening! ip: "<< inet_ntoa(tilexrCommId_.handle.addr.sin.sin_addr) - << " port: " << ntohs(tilexrCommId_.handle.addr.sin.sin_port); + TILEXR_LOG(INFO) << "The server is listening! ip: "<< inet_ntoa(bindAddr.sin_addr) + << " port: " << ntohs(bindAddr.sin_port); return TILEXR_SUCCESS; } @@ -235,7 +292,7 @@ int TileXRSockExchange::Accept() socklen_t sinSize = sizeof(struct sockaddr_in); for (int i = 1; i < rankSize_; ++i) { - int fd = AcceptConnection(fd_, clientAddr, &sinSize); + int fd = AcceptConnection(listenFd_, clientAddr, &sinSize); if (fd < 0) { TILEXR_LOG(ERROR) << "AcceptConnection failed"; return TILEXR_ERROR_INTERNAL; @@ -259,6 +316,34 @@ int TileXRSockExchange::Accept() return TILEXR_SUCCESS; } +int TileXRSockExchange::AcceptHierarchical() +{ + const int localMembers = groupLayout_.groupEnd - groupLayout_.groupBegin - 1; + const int childLeaders = rank_ == 0 ? groupLayout_.groupCount - 1 : 0; + const int expected = localMembers + childLeaders; + struct sockaddr_in clientAddr {}; + socklen_t sinSize = sizeof(clientAddr); + for (int accepted = 0; accepted < expected; ++accepted) { + int clientFd = AcceptConnection(listenFd_, clientAddr, &sinSize); + int peerRank = -1; + if (clientFd < 0 || Recv(clientFd, &peerRank, sizeof(peerRank), 0) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical accept rank handshake failed"; + return TILEXR_ERROR_INTERNAL; + } + const bool localMember = peerRank >= groupLayout_.groupBegin && + peerRank < groupLayout_.groupEnd && peerRank != rank_; + const bool childLeader = rank_ == 0 && peerRank > 0 && + peerRank < rankSize_ && peerRank % hierarchyGroupSize_ == 0; + if ((!localMember && !childLeader) || clientFds_[peerRank] >= 0) { + TILEXR_LOG(ERROR) << "Hierarchical accept invalid peer rank " << peerRank; + Close(clientFd); + return TILEXR_ERROR_INTERNAL; + } + clientFds_[peerRank] = clientFd; + } + return TILEXR_SUCCESS; +} + void TileXRSockExchange::Close(int &fd) const { if (fd == -1) { @@ -274,6 +359,11 @@ void TileXRSockExchange::Close(int &fd) const } int TileXRSockExchange::Connect() +{ + return ConnectTo(tilexrCommId_.handle.addr.sin); +} + +int TileXRSockExchange::ConnectTo(const sockaddr_in &serverAddr) { TILEXR_LOG(DEBUG) << "Client side " << rank_ << " begin to connect"; @@ -287,7 +377,7 @@ int TileXRSockExchange::Connect() int maxRetryCount = 180; int retryCount = 0; bool success = false; - struct sockaddr *addrPtr = &tilexrCommId_.handle.addr.sa; + const struct sockaddr *addrPtr = reinterpret_cast(&serverAddr); while (retryCount < maxRetryCount) { if (connect(fd_, addrPtr, sizeof(struct sockaddr)) < 0) { if (errno == ECONNREFUSED) { @@ -331,6 +421,9 @@ void TileXRSockExchange::Cleanup() Close(fd_); fd_ = -1; } + if (listenFd_ >= 0) { + Close(listenFd_); + } if (clientFds_.empty()) { return; diff --git a/src/comm/tools/socket/tilexr_sock_exchange.h b/src/comm/tools/socket/tilexr_sock_exchange.h index e7955d22..a5e1570e 100644 --- a/src/comm/tools/socket/tilexr_sock_exchange.h +++ b/src/comm/tools/socket/tilexr_sock_exchange.h @@ -26,6 +26,7 @@ #include "tilexr_types.h" #include "tilexr_api.h" +#include "tilexr_sock_exchange_layout.h" namespace TileXR { /* Common socket address storage structure for IPv4/IPv6 */ @@ -67,6 +68,9 @@ class TileXRSockExchange { } isInit_ = true; + if (hierarchical_) { + return HierarchicalAllGather(sendBuf, sendCount, recvBuf); + } if (!IsServer()) { return ClientSendRecv(sendBuf, sendCount, recvBuf); } else { @@ -82,6 +86,9 @@ class TileXRSockExchange { } isInit_ = true; + if (hierarchical_) { + return HierarchicalAllToAll(sendBuf, sendCountPerRank, recvBuf); + } if (!IsServer()) { return ClientSendRecvAllToAll(sendBuf, sendCountPerRank, recvBuf); } @@ -99,14 +106,21 @@ class TileXRSockExchange { private: void GetIpAndPort(); + int ConfigureHierarchy(); int Prepare(); int Listen(); int Accept(); + int AcceptHierarchical(); void Close(int &fd) const; int Connect(); + int ConnectTo(const sockaddr_in &serverAddr); int AcceptConnection(int fd, sockaddr_in &clientAddr, socklen_t *sinSize) const; void Cleanup(); bool IsServer() const; + bool IsGroupLeader() const + { + return hierarchical_ && rank_ == groupLayout_.groupLeader; + } static bool CheckErrno(int ioErrno) { return ((ioErrno == EAGAIN) || (ioErrno == EWOULDBLOCK) || (ioErrno == EINTR)); @@ -275,15 +289,201 @@ class TileXRSockExchange { } return TILEXR_SUCCESS; } + + template int HierarchicalAllGather( + const T *sendBuf, size_t sendCount, T *recvBuf) + { + if (sendCount == 0 || sendCount > std::numeric_limits::max() / + static_cast(rankSize_)) { + return TILEXR_ERROR_INTERNAL; + } + const size_t fullCount = sendCount * static_cast(rankSize_); + const size_t fullBytes = fullCount * sizeof(T); + if (!IsGroupLeader()) { + if (Send(fd_, sendBuf, sendCount * sizeof(T), 0) <= 0 || + Recv(fd_, recvBuf, fullBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical allgather member rank " << rank_ << " failed"; + return TILEXR_ERROR_INTERNAL; + } + return TILEXR_SUCCESS; + } + + std::copy_n(sendBuf, sendCount, + recvBuf + static_cast(rank_) * sendCount); + for (int member = groupLayout_.groupBegin; member < groupLayout_.groupEnd; ++member) { + if (member == rank_) { + continue; + } + if (Recv(clientFds_[member], + recvBuf + static_cast(member) * sendCount, + sendCount * sizeof(T), MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical allgather leader " << rank_ + << " recv member " << member << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + + if (rank_ != 0) { + const size_t groupCount = static_cast( + groupLayout_.groupEnd - groupLayout_.groupBegin) * sendCount; + if (Send(fd_, recvBuf + static_cast(groupLayout_.groupBegin) * sendCount, + groupCount * sizeof(T), 0) <= 0 || + Recv(fd_, recvBuf, fullBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical allgather leader " << rank_ + << " parent exchange failed"; + return TILEXR_ERROR_INTERNAL; + } + } else { + for (int group = 1; group < groupLayout_.groupCount; ++group) { + const int leader = group * hierarchyGroupSize_; + const int groupEnd = std::min(rankSize_, leader + hierarchyGroupSize_); + const size_t groupCount = static_cast(groupEnd - leader) * sendCount; + if (Recv(clientFds_[leader], + recvBuf + static_cast(leader) * sendCount, + groupCount * sizeof(T), MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical allgather root recv leader " + << leader << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + for (int group = 1; group < groupLayout_.groupCount; ++group) { + const int leader = group * hierarchyGroupSize_; + if (Send(clientFds_[leader], recvBuf, fullBytes, 0) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical allgather root send leader " + << leader << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + } + + for (int member = groupLayout_.groupBegin; member < groupLayout_.groupEnd; ++member) { + if (member != rank_ && Send(clientFds_[member], recvBuf, fullBytes, 0) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical allgather leader " << rank_ + << " send member " << member << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + return TILEXR_SUCCESS; + } + + template int HierarchicalAllToAll( + const T *sendBuf, size_t sendCountPerRank, T *recvBuf) + { + size_t rowCount = 0; + size_t totalCount = 0; + if (!AllToAllSizes(sendCountPerRank, rowCount, totalCount)) { + return TILEXR_ERROR_INTERNAL; + } + const size_t rowBytes = rowCount * sizeof(T); + if (!IsGroupLeader()) { + if (Send(fd_, sendBuf, rowBytes, 0) <= 0 || + Recv(fd_, recvBuf, rowBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical alltoall member rank " << rank_ << " failed"; + return TILEXR_ERROR_INTERNAL; + } + return TILEXR_SUCCESS; + } + + const size_t localGroupRanks = static_cast( + groupLayout_.groupEnd - groupLayout_.groupBegin); + std::vector groupRows(localGroupRanks * rowCount); + std::copy_n(sendBuf, rowCount, + groupRows.data() + static_cast(rank_ - groupLayout_.groupBegin) * rowCount); + for (int member = groupLayout_.groupBegin; member < groupLayout_.groupEnd; ++member) { + if (member == rank_) { + continue; + } + if (Recv(clientFds_[member], + groupRows.data() + static_cast(member - groupLayout_.groupBegin) * rowCount, + rowBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical alltoall leader " << rank_ + << " recv member " << member << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + + std::vector groupResults(localGroupRanks * rowCount); + if (rank_ != 0) { + const size_t groupBytes = groupRows.size() * sizeof(T); + if (Send(fd_, groupRows.data(), groupBytes, 0) <= 0 || + Recv(fd_, groupResults.data(), groupBytes, MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical alltoall leader " << rank_ + << " parent exchange failed"; + return TILEXR_ERROR_INTERNAL; + } + } else { + std::vector gathered(totalCount); + std::copy(groupRows.begin(), groupRows.end(), gathered.begin()); + for (int group = 1; group < groupLayout_.groupCount; ++group) { + const int leader = group * hierarchyGroupSize_; + const int groupEnd = std::min(rankSize_, leader + hierarchyGroupSize_); + const size_t groupRowsCount = static_cast(groupEnd - leader) * rowCount; + if (Recv(clientFds_[leader], + gathered.data() + static_cast(leader) * rowCount, + groupRowsCount * sizeof(T), MSG_WAITALL) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical alltoall root recv leader " + << leader << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + + for (int group = 0; group < groupLayout_.groupCount; ++group) { + const int leader = group * hierarchyGroupSize_; + const int groupEnd = std::min(rankSize_, leader + hierarchyGroupSize_); + const size_t targetRanks = static_cast(groupEnd - leader); + std::vector results(targetRanks * rowCount); + for (int target = leader; target < groupEnd; ++target) { + T *targetResult = results.data() + + static_cast(target - leader) * rowCount; + for (int source = 0; source < rankSize_; ++source) { + const T *sourceData = gathered.data() + + static_cast(source) * rowCount + + static_cast(target) * sendCountPerRank; + std::copy_n(sourceData, sendCountPerRank, + targetResult + static_cast(source) * sendCountPerRank); + } + } + if (leader == 0) { + groupResults.swap(results); + } else if (Send(clientFds_[leader], results.data(), + results.size() * sizeof(T), 0) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical alltoall root send leader " + << leader << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + } + + for (int member = groupLayout_.groupBegin; member < groupLayout_.groupEnd; ++member) { + const T *memberResult = groupResults.data() + + static_cast(member - groupLayout_.groupBegin) * rowCount; + if (member == rank_) { + std::copy_n(memberResult, rowCount, recvBuf); + } else if (Send(clientFds_[member], memberResult, rowBytes, 0) <= 0) { + TILEXR_LOG(ERROR) << "Hierarchical alltoall leader " << rank_ + << " send member " << member << " failed"; + return TILEXR_ERROR_INTERNAL; + } + } + return TILEXR_SUCCESS; + } + int rank_ = 0; int rankSize_ = 0; int fd_ = -1; + int listenFd_ = -1; std::vector clientFds_ = {}; bool isInit_ = false; int commDomain_ = -1; std::string ip_; uint16_t port_ = 0; TileXRBootstrap tilexrCommId_ = {}; + bool hierarchical_ = false; + int hierarchyGroupSize_ = 64; + int hierarchyRanksPerHost_ = 8; + SockExchangeGroupLayout groupLayout_ {}; + std::vector hierarchyHosts_; + sockaddr_in listenAddr_ {}; }; } // namespace TileXR diff --git a/src/comm/tools/socket/tilexr_sock_exchange_layout.h b/src/comm/tools/socket/tilexr_sock_exchange_layout.h new file mode 100644 index 00000000..a9ad91ae --- /dev/null +++ b/src/comm/tools/socket/tilexr_sock_exchange_layout.h @@ -0,0 +1,33 @@ +#ifndef TILEXR_SOCK_EXCHANGE_LAYOUT_H +#define TILEXR_SOCK_EXCHANGE_LAYOUT_H + +#include + +namespace TileXR { + +struct SockExchangeGroupLayout { + int groupCount = 0; + int groupIndex = 0; + int groupBegin = 0; + int groupEnd = 0; + int groupLeader = 0; +}; + +inline SockExchangeGroupLayout BuildSockExchangeGroupLayout( + int rank, int rankSize, int groupSize) +{ + SockExchangeGroupLayout layout {}; + if (rank < 0 || rank >= rankSize || rankSize <= 0 || groupSize <= 0) { + return layout; + } + layout.groupCount = (rankSize + groupSize - 1) / groupSize; + layout.groupIndex = rank / groupSize; + layout.groupBegin = layout.groupIndex * groupSize; + layout.groupEnd = std::min(rankSize, layout.groupBegin + groupSize); + layout.groupLeader = layout.groupBegin; + return layout; +} + +} // namespace TileXR + +#endif diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 372c8092..a3b34915 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -133,6 +133,14 @@ add_executable(test_tilexr_udma_transport_layout ${TILEXR_ROOT}/src/comm/udma/tilexr_udma_layout.cpp ) +add_executable(test_tilexr_sock_exchange_layout + unit/test_tilexr_sock_exchange_layout.cpp +) + +target_include_directories(test_tilexr_sock_exchange_layout PRIVATE + ${TILEXR_ROOT}/src/comm +) + target_include_directories(test_tilexr_udma_transport_layout PRIVATE ${TILEXR_ROOT}/src/include ${TILEXR_ROOT}/src/comm @@ -163,6 +171,7 @@ set(INSTALL_TARGETS test_tilexr_chip_map_sources test_tilexr_ipc_pid_mode_sources test_tilexr_udma_transport_layout + test_tilexr_sock_exchange_layout ) if(BUILD_TILEXR_UDMA_DEMO) diff --git a/tests/udma/unit/test_tilexr_sock_exchange_layout.cpp b/tests/udma/unit/test_tilexr_sock_exchange_layout.cpp new file mode 100644 index 00000000..fdc493bf --- /dev/null +++ b/tests/udma/unit/test_tilexr_sock_exchange_layout.cpp @@ -0,0 +1,75 @@ +#include + +#include "tools/socket/tilexr_sock_exchange_layout.h" + +namespace { + +int g_failures = 0; + +#define CHECK_EQ(lhs, rhs) \ + do { \ + const auto lhsValue = (lhs); \ + const auto rhsValue = (rhs); \ + if (lhsValue != rhsValue) { \ + std::cerr << "CHECK_EQ failed at line " << __LINE__ << ": " \ + << lhsValue << " != " << rhsValue << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +void Test512RanksUseEightCabinetLeaders() +{ + for (int rank = 0; rank < 512; ++rank) { + const auto layout = TileXR::BuildSockExchangeGroupLayout(rank, 512, 64); + CHECK_EQ(layout.groupCount, 8); + CHECK_EQ(layout.groupIndex, rank / 64); + CHECK_EQ(layout.groupBegin, (rank / 64) * 64); + CHECK_EQ(layout.groupEnd, (rank / 64 + 1) * 64); + CHECK_EQ(layout.groupLeader, (rank / 64) * 64); + } +} + +void Test1024RanksUseSixteenCabinetLeaders() +{ + const auto first = TileXR::BuildSockExchangeGroupLayout(0, 1024, 64); + const auto last = TileXR::BuildSockExchangeGroupLayout(1023, 1024, 64); + CHECK_EQ(first.groupCount, 16); + CHECK_EQ(first.groupLeader, 0); + CHECK_EQ(last.groupCount, 16); + CHECK_EQ(last.groupBegin, 960); + CHECK_EQ(last.groupEnd, 1024); + CHECK_EQ(last.groupLeader, 960); +} + +void TestTailGroupIsClampedToRankSize() +{ + const auto layout = TileXR::BuildSockExchangeGroupLayout(129, 130, 64); + CHECK_EQ(layout.groupCount, 3); + CHECK_EQ(layout.groupIndex, 2); + CHECK_EQ(layout.groupBegin, 128); + CHECK_EQ(layout.groupEnd, 130); + CHECK_EQ(layout.groupLeader, 128); +} + +void TestInvalidInputsDisableLayout() +{ + CHECK_EQ(TileXR::BuildSockExchangeGroupLayout(-1, 512, 64).groupCount, 0); + CHECK_EQ(TileXR::BuildSockExchangeGroupLayout(512, 512, 64).groupCount, 0); + CHECK_EQ(TileXR::BuildSockExchangeGroupLayout(0, 512, 0).groupCount, 0); +} + +} // namespace + +int main() +{ + Test512RanksUseEightCabinetLeaders(); + Test1024RanksUseSixteenCabinetLeaders(); + TestTailGroupIsClampedToRankSize(); + TestInvalidInputsDisableLayout(); + if (g_failures != 0) { + std::cerr << g_failures << " test(s) failed" << std::endl; + return 1; + } + std::cout << "tilexr socket exchange layout tests passed" << std::endl; + return 0; +} From 102b5bfac28d34229f772fde84e7eba634be2e71 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 27 Jul 2026 14:43:28 +0800 Subject: [PATCH 124/163] perf(udma): batch grouped alltoall quiet operations --- .../tilexr_udma_alltoall_group_kernel.cpp | 107 ++++++++++++++---- .../demo/tilexr_udma_alltoall_group_layout.h | 52 ++++++--- tests/udma/demo/tilexr_udma_demo.cpp | 31 ++++- ...test_tilexr_udma_alltoall_group_layout.cpp | 57 ++++++++-- 4 files changed, 197 insertions(+), 50 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index d5c06a4f..8160c55d 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -13,7 +13,9 @@ namespace { constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_CORES = 16U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_WORKERS = 32U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_HALF_WIDTH = 8U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH = 16U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH = 4U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH = 4U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 1024U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -48,22 +50,26 @@ struct AllToAllGroupDeviceError { }; __aicore__ inline int32_t AllToAllGroupDevicePeer( - int32_t rank, int32_t rankSize, uint32_t group, uint32_t lane) + int32_t rank, int32_t rankSize, uint32_t group, uint32_t lane, + uint32_t groupWidth) { if (rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || - (rankSize & 7) != 0 || rank < 0 || rank >= rankSize || lane >= 16U) { + (rankSize & 7) != 0 || rank < 0 || rank >= rankSize || + (groupWidth != TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH && + groupWidth != TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH) || + lane >= groupWidth) { return -1; } - const uint32_t index = lane < TILEXR_ALLTOALL_GROUP_HALF_WIDTH ? - lane : lane - TILEXR_ALLTOALL_GROUP_HALF_WIDTH; + const uint32_t halfWidth = groupWidth / 2U; + const uint32_t index = lane < halfWidth ? lane : lane - halfWidth; const int32_t distance = static_cast( - group * TILEXR_ALLTOALL_GROUP_HALF_WIDTH + index + 1U); + group * halfWidth + index + 1U); const int32_t diameter = rankSize / 2; if (distance > diameter || - (lane >= TILEXR_ALLTOALL_GROUP_HALF_WIDTH && distance == diameter)) { + (lane >= halfWidth && distance == diameter)) { return -1; } - return lane < TILEXR_ALLTOALL_GROUP_HALF_WIDTH ? + return lane < halfWidth ? (rank + distance) % rankSize : (rank - distance + rankSize) % rankSize; } @@ -405,6 +411,43 @@ __aicore__ inline void AllToAllGroupTraceRecordTask( span->endCycle = endCycle; } +struct AllToAllGroupPendingQuiet { + int32_t peer; + uint32_t qpIdx; + uint32_t group; + uint32_t pass; + uint64_t expectedToken; +}; + +__aicore__ inline bool AllToAllGroupFlushQuiet( + const __gm__ TileXR::CommArgs* args, + AllToAllGroupPendingQuiet* pending, uint32_t& pendingCount, + __gm__ int32_t* debug, uint32_t blockIdx, + __gm__ uint8_t* trace, uint32_t traceIteration, + uint32_t groupCount, uint32_t passCount) +{ + for (uint32_t index = 0U; index < pendingCount; ++index) { + const auto& request = pending[index]; + const uint64_t quietBegin = AllToAllGroupTraceCycle(trace); + const uint32_t quietStatus = + TileXR::UDMAQuietStatusOnQp(args, request.peer, request.qpIdx); + AllToAllGroupTraceRecordTask( + trace, traceIteration, blockIdx, request.group, request.pass, + TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, + request.peer, request.qpIdx, quietBegin, + AllToAllGroupTraceCycle(trace)); + if (quietStatus != 0U) { + AllToAllGroupRecordError(debug, blockIdx, + TILEXR_ALLTOALL_GROUP_STAGE_QUIET, request.group, request.pass, + request.peer, request.qpIdx, quietStatus, + request.expectedToken, 0ULL); + return false; + } + } + pendingCount = 0U; + return true; +} + } // namespace extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( @@ -415,7 +458,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, - uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts) + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch) { constexpr uint32_t copyoutWorkers = 32U; const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); @@ -434,6 +478,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( copyoutWorkers != 32U && copyoutWorkers != 48U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || multiChannel > 1U || + (groupWidth != TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH && + groupWidth != TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH) || + (quietBatch != 1U && quietBatch != 2U && + quietBatch != TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH) || (primaryRouteParts > TileXR::Demo::kAllToAllGroupRouteParts && primaryRouteParts != TileXR::Demo::kAllToAllGroupAutoPrimaryParts) || TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers > @@ -441,7 +489,10 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( blockIdx >= TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers || !TileXR::UDMARegistryEnabled(args) || rankSize < 8 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || (rankSize & 7) != 0 || - elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || groupCount == 0U) { + elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || + groupCount == 0U || groupCount != static_cast( + (rankSize - 1 + static_cast(groupWidth) - 1) / + static_cast(groupWidth))) { AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CONFIG, 0U, 0U, -1, 0U, 0U, 0ULL, 0ULL); AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, @@ -499,7 +550,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( break; } const uint32_t lane = static_cast(laneValue); - const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); + const int32_t peer = AllToAllGroupDevicePeer( + rank, rankSize, group, lane, groupWidth); if (peer < 0) { continue; } @@ -656,8 +708,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } const uint32_t lane = blockIdx % TILEXR_ALLTOALL_GROUP_SEND_CORES; const uint32_t workerRoute = blockIdx / TILEXR_ALLTOALL_GROUP_SEND_CORES; + AllToAllGroupPendingQuiet pending[TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH]; + uint32_t pendingCount = 0U; for (uint32_t group = 0U; group < groupCount; ++group) { - const int32_t peer = AllToAllGroupDevicePeer(rank, rankSize, group, lane); + const int32_t peer = AllToAllGroupDevicePeer( + rank, rankSize, group, lane, groupWidth); if (peer < 0) { continue; } @@ -737,17 +792,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( groupTrace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, peer, selectedQp, putBegin, AllToAllGroupTraceCycle(groupTrace)); - const uint64_t quietBegin = AllToAllGroupTraceCycle(groupTrace); - const uint32_t quietStatus = - TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); - AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, pass, - TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, - peer, selectedQp, quietBegin, AllToAllGroupTraceCycle(groupTrace)); - if (quietStatus != 0U) { - AllToAllGroupRecordError(debug, blockIdx, - TILEXR_ALLTOALL_GROUP_STAGE_QUIET, group, pass, peer, - selectedQp, quietStatus, expectedToken, 0ULL); + pending[pendingCount++] = { + peer, selectedQp, group, pass, expectedToken}; + if (pendingCount == quietBatch && !AllToAllGroupFlushQuiet( + args, pending, pendingCount, debug, blockIdx, + groupTrace, traceIteration, groupCount, passCount)) { AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; @@ -755,6 +804,13 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } } } + if (pendingCount != 0U && !AllToAllGroupFlushQuiet( + args, pending, pendingCount, debug, blockIdx, + groupTrace, traceIteration, groupCount, passCount)) { + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); } @@ -767,12 +823,13 @@ void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts) + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, groupTrace, traceIteration, routeStage, - multiChannel, primaryRouteParts); + multiChannel, primaryRouteParts, groupWidth, quietBatch); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index e56ba187..291b5957 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -18,6 +18,8 @@ constexpr int32_t kAllToAllGroupMinRankSize = 8; constexpr int32_t kAllToAllGroupMaxRankSize = 1024; constexpr uint32_t kAllToAllGroupWidth = 16U; constexpr uint32_t kAllToAllGroupHalfWidth = 8U; +constexpr uint32_t kAllToAllGroupExperimentalWidth = 4U; +constexpr uint32_t kAllToAllGroupMaxQuietBatch = 4U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupRouteSignalStride = 512U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 1024U; @@ -32,6 +34,7 @@ constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; struct AllToAllGroupPlan { bool valid = false; + uint32_t groupWidth = kAllToAllGroupWidth; uint32_t groupCount = 0; uint32_t passCount = 0; int32_t chunkElements = 0; @@ -45,6 +48,18 @@ struct AllToAllGroupPlan { size_t registeredBytes = 0; }; +inline bool AllToAllGroupValidWidth(uint32_t groupWidth) +{ + return groupWidth == kAllToAllGroupWidth || + groupWidth == kAllToAllGroupExperimentalWidth; +} + +inline bool AllToAllGroupValidQuietBatch(uint32_t quietBatch) +{ + return quietBatch == 1U || quietBatch == 2U || + quietBatch == kAllToAllGroupMaxQuietBatch; +} + enum class AllToAllGroupChannelMode : uint32_t { kAuto = 0U, kSingle = 1U, @@ -109,32 +124,37 @@ inline int32_t AllToAllGroupCopyoutLane( return lane < kAllToAllGroupWidth ? static_cast(lane) : -1; } -inline uint32_t AllToAllGroupCount(int rankSize) +inline uint32_t AllToAllGroupCount( + int rankSize, uint32_t groupWidth = kAllToAllGroupWidth) { - if (!AllToAllGroupValidRankSize(rankSize)) { + if (!AllToAllGroupValidRankSize(rankSize) || + !AllToAllGroupValidWidth(groupWidth)) { return 0U; } - return static_cast((rankSize - 1 + static_cast(kAllToAllGroupWidth) - 1) / - static_cast(kAllToAllGroupWidth)); + return static_cast( + (rankSize - 1 + static_cast(groupWidth) - 1) / + static_cast(groupWidth)); } inline int32_t AllToAllGroupPeer( - int rank, int rankSize, uint32_t group, uint32_t lane) + int rank, int rankSize, uint32_t group, uint32_t lane, + uint32_t groupWidth = kAllToAllGroupWidth) { if (!AllToAllGroupValidRankSize(rankSize) || rank < 0 || rank >= rankSize || - lane >= kAllToAllGroupWidth || group >= AllToAllGroupCount(rankSize)) { + !AllToAllGroupValidWidth(groupWidth) || lane >= groupWidth || + group >= AllToAllGroupCount(rankSize, groupWidth)) { return -1; } - const uint32_t index = lane < kAllToAllGroupHalfWidth ? - lane : lane - kAllToAllGroupHalfWidth; + const uint32_t halfWidth = groupWidth / 2U; + const uint32_t index = lane < halfWidth ? lane : lane - halfWidth; const int32_t distance = static_cast( - group * kAllToAllGroupHalfWidth + index + 1U); + group * halfWidth + index + 1U); const int32_t diameter = rankSize / 2; if (distance > diameter || - (lane >= kAllToAllGroupHalfWidth && distance == diameter)) { + (lane >= halfWidth && distance == diameter)) { return -1; } - return lane < kAllToAllGroupHalfWidth ? + return lane < halfWidth ? (rank + distance) % rankSize : (rank - distance + rankSize) % rankSize; } @@ -178,14 +198,18 @@ inline bool AllToAllGroupAlignUp(size_t value, size_t& result) } inline AllToAllGroupPlan PlanAllToAllGroup( - int rankSize, int32_t elementsPerPeer, int32_t chunkElements) + int rankSize, int32_t elementsPerPeer, int32_t chunkElements, + uint32_t groupWidth = kAllToAllGroupWidth) { AllToAllGroupPlan plan {}; - if (!AllToAllGroupValidRankSize(rankSize) || elementsPerPeer <= 0 || chunkElements <= 0) { + if (!AllToAllGroupValidRankSize(rankSize) || + !AllToAllGroupValidWidth(groupWidth) || + elementsPerPeer <= 0 || chunkElements <= 0) { return plan; } - plan.groupCount = AllToAllGroupCount(rankSize); + plan.groupWidth = groupWidth; + plan.groupCount = AllToAllGroupCount(rankSize, groupWidth); plan.chunkElements = std::min(elementsPerPeer, chunkElements); plan.passCount = static_cast( (static_cast(elementsPerPeer) + static_cast(plan.chunkElements) - 1ULL) / diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 1b2be904..ff206f5c 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -60,7 +60,8 @@ extern void launch_tilexr_udma_all_to_all_group( uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts); + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -643,10 +644,31 @@ bool RunGroupedAllToAll( { constexpr uint32_t kErrorWordsPerCore = 12U; constexpr uint32_t kErrorCoreCount = TileXR::Demo::kAllToAllGroupBlockDim; + const int groupWidthValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH", + static_cast(TileXR::Demo::kAllToAllGroupWidth)); + if (groupWidthValue <= 0 || !TileXR::Demo::AllToAllGroupValidWidth( + static_cast(groupWidthValue))) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_WIDTH" + << " must be 4 or 16, got " << groupWidthValue << std::endl; + return false; + } + const uint32_t groupWidth = static_cast(groupWidthValue); + const int quietBatchValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_QUIET_BATCH", 1); + if (quietBatchValue <= 0 || !TileXR::Demo::AllToAllGroupValidQuietBatch( + static_cast(quietBatchValue))) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_QUIET_BATCH" + << " must be 1, 2, or 4, got " << quietBatchValue << std::endl; + return false; + } + const uint32_t quietBatch = static_cast(quietBatchValue); const int32_t requestedChunkElements = std::max( 1, GetEnvInt("TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS", elementsPerPeer)); const auto plan = TileXR::Demo::PlanAllToAllGroup( - rankSize, elementsPerPeer, requestedChunkElements); + rankSize, elementsPerPeer, requestedChunkElements, groupWidth); if (!plan.valid) { std::cerr << "[rank " << rank << "] ERROR: invalid grouped alltoall plan" << " rankSize=" << rankSize @@ -861,6 +883,7 @@ bool RunGroupedAllToAll( " signalOffset0=" + std::to_string(plan.signalOffset[0]) + " signalOffset1=" + std::to_string(plan.signalOffset[1]) + " controlOffset=" + std::to_string(plan.controlOffset) + + " groupWidth=" + std::to_string(plan.groupWidth) + " groups=" + std::to_string(plan.groupCount) + " passes=" + std::to_string(plan.passCount)); PrintStatus(rank, "grouped alltoall warmup=" + std::to_string(warmup) + @@ -872,6 +895,7 @@ bool RunGroupedAllToAll( " copyoutWorkers=" + std::to_string(copyoutWorkers) + " blockDim=" + std::to_string(groupBlockDim) + " useSecondaryRoute=" + std::to_string(useSecondaryRouteValue) + + " quietBatch=" + std::to_string(quietBatch) + " routeStages=" + std::to_string(routeStagesValue)); if (routeStages && @@ -893,7 +917,8 @@ bool RunGroupedAllToAll( plan.signalOffset[0], plan.signalOffset[1], reinterpret_cast(trace), traceIteration, static_cast(routeStage), - multiChannel ? 1U : 0U, primaryRouteParts); + multiChannel ? 1U : 0U, primaryRouteParts, + groupWidth, quietBatch); }; double totalUs = 0.0; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 00b6c646..8580b407 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -52,23 +52,25 @@ std::string ReadFile(const std::string& path) return out.str(); } -void CheckSchedule(int rankSize) +void CheckSchedule(int rankSize, uint32_t groupWidth = TileXR::Demo::kAllToAllGroupWidth) { for (int rank = 0; rank < rankSize; ++rank) { std::set peers; - for (uint32_t group = 0; group < TileXR::Demo::AllToAllGroupCount(rankSize); ++group) { - for (uint32_t lane = 0; lane < TileXR::Demo::kAllToAllGroupWidth; ++lane) { - const int peer = TileXR::Demo::AllToAllGroupPeer(rank, rankSize, group, lane); + for (uint32_t group = 0; + group < TileXR::Demo::AllToAllGroupCount(rankSize, groupWidth); ++group) { + for (uint32_t lane = 0; lane < groupWidth; ++lane) { + const int peer = TileXR::Demo::AllToAllGroupPeer( + rank, rankSize, group, lane, groupWidth); if (peer < 0) { continue; } CHECK_EQ(peer == rank, false); CHECK_EQ(peers.insert(peer).second, true); bool symmetric = false; - for (uint32_t remoteLane = 0; remoteLane < TileXR::Demo::kAllToAllGroupWidth; - ++remoteLane) { + for (uint32_t remoteLane = 0; remoteLane < groupWidth; ++remoteLane) { symmetric = symmetric || - TileXR::Demo::AllToAllGroupPeer(peer, rankSize, group, remoteLane) == rank; + TileXR::Demo::AllToAllGroupPeer( + peer, rankSize, group, remoteLane, groupWidth) == rank; } CHECK_EQ(symmetric, true); } @@ -81,6 +83,7 @@ void TestSchedules() { for (int rankSize : {8, 16, 24, 32, 40, 64, 128, 256, 512, 1024}) { CheckSchedule(rankSize); + CheckSchedule(rankSize, TileXR::Demo::kAllToAllGroupExperimentalWidth); } CHECK_EQ(TileXR::Demo::AllToAllGroupCount(8), 1U); CHECK_EQ(TileXR::Demo::AllToAllGroupCount(16), 1U); @@ -98,6 +101,23 @@ void TestSchedules() CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 0, 15), 56); CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 3, 7), 32); CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 64, 3, 15), -1); + + constexpr uint32_t width = TileXR::Demo::kAllToAllGroupExperimentalWidth; + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(16, width), 4U); + const int expected[4][4] = { + {1, 2, 15, 14}, + {3, 4, 13, 12}, + {5, 6, 11, 10}, + {7, 8, 9, -1}, + }; + for (uint32_t group = 0U; group < 4U; ++group) { + for (uint32_t lane = 0U; lane < width; ++lane) { + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer( + 0, 16, group, lane, width), expected[group][lane]); + } + } + CHECK_EQ(TileXR::Demo::AllToAllGroupCount(16, 8U), 0U); + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 16, 0, 4, width), -1); } void TestPlan() @@ -131,6 +151,13 @@ void TestPlan() CHECK_EQ(chunked.passCount, 4U); CHECK_EQ(chunked.payloadPlaneBytes, plan.payloadPlaneBytes); + const auto widthFour = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, elementsPerPeer, + TileXR::Demo::kAllToAllGroupExperimentalWidth); + CHECK_EQ(widthFour.valid, true); + CHECK_EQ(widthFour.groupWidth, 4U); + CHECK_EQ(widthFour.groupCount, 4U); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(7, 1024, 1024).valid, false); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(18, 1024, 1024).valid, false); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup(1032, 1024, 1024).valid, false); @@ -160,6 +187,13 @@ void TestChannelPolicy() sizeof(int32_t), AllToAllGroupChannelMode::kMulti), true); CHECK_EQ(TileXR::Demo::kAllToAllGroupSendWorkerCount, 32U); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidWidth(4U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidWidth(16U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidWidth(8U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(1U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(2U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(4U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(3U), false); } void TestScalePlanAndTraceCapacity() @@ -418,7 +452,12 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "const uint32_t traceCore = copyoutWorkers == 8U"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx"); CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); - CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); + CHECK_CONTAINS(kernel, "AllToAllGroupPendingQuiet"); + CHECK_CONTAINS(kernel, "AllToAllGroupFlushQuiet"); + CHECK_CONTAINS(kernel, + "UDMAQuietStatusOnQp(args, request.peer, request.qpIdx)"); + CHECK_CONTAINS(kernel, "pendingCount == quietBatch"); + CHECK_CONTAINS(kernel, "uint32_t groupWidth, uint32_t quietBatch"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsReceiveDevice(routeStage)"); @@ -445,6 +484,8 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_QUIET_BATCH"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS"); CHECK_CONTAINS(demo, "kAllToAllGroupSendWorkerCount"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); From 897f3f6f5a10645649d62b24d3c78939a2332e13 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Mon, 27 Jul 2026 15:24:32 +0800 Subject: [PATCH 125/163] perf(udma): extend grouped quiet batches to 64 --- tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp | 7 ++++--- tests/udma/demo/tilexr_udma_alltoall_group_layout.h | 6 +++--- tests/udma/demo/tilexr_udma_demo.cpp | 3 ++- tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 6 ++++++ 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 8160c55d..bca932d7 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -15,7 +15,7 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_SEND_WORKERS = 32U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH = 16U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH = 4U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH = 4U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 1024U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -480,8 +480,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( multiChannel > 1U || (groupWidth != TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH && groupWidth != TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH) || - (quietBatch != 1U && quietBatch != 2U && - quietBatch != TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH) || + quietBatch == 0U || + quietBatch > TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH || + (quietBatch & (quietBatch - 1U)) != 0U || (primaryRouteParts > TileXR::Demo::kAllToAllGroupRouteParts && primaryRouteParts != TileXR::Demo::kAllToAllGroupAutoPrimaryParts) || TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers > diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 291b5957..2867b465 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -19,7 +19,7 @@ constexpr int32_t kAllToAllGroupMaxRankSize = 1024; constexpr uint32_t kAllToAllGroupWidth = 16U; constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupExperimentalWidth = 4U; -constexpr uint32_t kAllToAllGroupMaxQuietBatch = 4U; +constexpr uint32_t kAllToAllGroupMaxQuietBatch = 64U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupRouteSignalStride = 512U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 1024U; @@ -56,8 +56,8 @@ inline bool AllToAllGroupValidWidth(uint32_t groupWidth) inline bool AllToAllGroupValidQuietBatch(uint32_t quietBatch) { - return quietBatch == 1U || quietBatch == 2U || - quietBatch == kAllToAllGroupMaxQuietBatch; + return quietBatch != 0U && quietBatch <= kAllToAllGroupMaxQuietBatch && + (quietBatch & (quietBatch - 1U)) == 0U; } enum class AllToAllGroupChannelMode : uint32_t { diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index ff206f5c..1d5a5017 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -661,7 +661,8 @@ bool RunGroupedAllToAll( static_cast(quietBatchValue))) { std::cerr << "[rank " << rank << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_QUIET_BATCH" - << " must be 1, 2, or 4, got " << quietBatchValue << std::endl; + << " must be a power of two from 1 through 64, got " + << quietBatchValue << std::endl; return false; } const uint32_t quietBatch = static_cast(quietBatchValue); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 8580b407..2454421b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -193,7 +193,13 @@ void TestChannelPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(1U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(2U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(4U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(8U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(16U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(32U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(64U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(0U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(3U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidQuietBatch(65U), false); } void TestScalePlanAndTraceCapacity() From 03cc96429ee6518e6c99939dfba3e73efc2cb7f7 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sat, 1 Aug 2026 21:22:51 +0800 Subject: [PATCH 126/163] fix(udma): restore direct quiet path for batch one --- .../tilexr_udma_alltoall_group_kernel.cpp | 146 ++++++++++++++++-- ...test_tilexr_udma_alltoall_group_layout.cpp | 11 +- 2 files changed, 140 insertions(+), 17 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index bca932d7..600fcbd3 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -419,6 +419,15 @@ struct AllToAllGroupPendingQuiet { uint64_t expectedToken; }; +template +struct AllToAllGroupQuietState {}; + +template <> +struct AllToAllGroupQuietState { + AllToAllGroupPendingQuiet pending[TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH]; + uint32_t pendingCount = 0U; +}; + __aicore__ inline bool AllToAllGroupFlushQuiet( const __gm__ TileXR::CommArgs* args, AllToAllGroupPendingQuiet* pending, uint32_t& pendingCount, @@ -448,9 +457,69 @@ __aicore__ inline bool AllToAllGroupFlushQuiet( return true; } +__aicore__ inline bool AllToAllGroupCompleteQuiet( + const __gm__ TileXR::CommArgs* args, + AllToAllGroupQuietState&, + uint32_t, int32_t peer, uint32_t selectedQp, + uint32_t group, uint32_t pass, uint64_t expectedToken, + __gm__ int32_t* debug, uint32_t blockIdx, + __gm__ uint8_t* trace, uint32_t traceIteration, + uint32_t groupCount, uint32_t passCount) +{ + const uint64_t quietBegin = AllToAllGroupTraceCycle(trace); + const uint32_t quietStatus = + TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); + AllToAllGroupTraceRecordTask( + trace, traceIteration, blockIdx, group, pass, + TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, + peer, selectedQp, quietBegin, AllToAllGroupTraceCycle(trace)); + if (quietStatus != 0U) { + AllToAllGroupRecordError(debug, blockIdx, + TILEXR_ALLTOALL_GROUP_STAGE_QUIET, group, pass, peer, + selectedQp, quietStatus, expectedToken, 0ULL); + return false; + } + return true; +} + +__aicore__ inline bool AllToAllGroupCompleteQuiet( + const __gm__ TileXR::CommArgs* args, + AllToAllGroupQuietState& state, + uint32_t quietBatch, int32_t peer, uint32_t selectedQp, + uint32_t group, uint32_t pass, uint64_t expectedToken, + __gm__ int32_t* debug, uint32_t blockIdx, + __gm__ uint8_t* trace, uint32_t traceIteration, + uint32_t groupCount, uint32_t passCount) +{ + state.pending[state.pendingCount++] = { + peer, selectedQp, group, pass, expectedToken}; + return state.pendingCount != quietBatch || AllToAllGroupFlushQuiet( + args, state.pending, state.pendingCount, debug, blockIdx, + trace, traceIteration, groupCount, passCount); +} + +__aicore__ inline bool AllToAllGroupFinishQuiet( + const __gm__ TileXR::CommArgs*, AllToAllGroupQuietState&, + __gm__ int32_t*, uint32_t, __gm__ uint8_t*, uint32_t, uint32_t, uint32_t) +{ + return true; +} + +__aicore__ inline bool AllToAllGroupFinishQuiet( + const __gm__ TileXR::CommArgs* args, AllToAllGroupQuietState& state, + __gm__ int32_t* debug, uint32_t blockIdx, + __gm__ uint8_t* trace, uint32_t traceIteration, + uint32_t groupCount, uint32_t passCount) +{ + return state.pendingCount == 0U || AllToAllGroupFlushQuiet( + args, state.pending, state.pendingCount, debug, blockIdx, + trace, traceIteration, groupCount, passCount); +} + } // namespace -extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( +template +__aicore__ inline void AllToAllGroupKernelImpl( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, int32_t elementsPerPeer, int32_t chunkElements, @@ -709,8 +778,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } const uint32_t lane = blockIdx % TILEXR_ALLTOALL_GROUP_SEND_CORES; const uint32_t workerRoute = blockIdx / TILEXR_ALLTOALL_GROUP_SEND_CORES; - AllToAllGroupPendingQuiet pending[TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH]; - uint32_t pendingCount = 0U; + AllToAllGroupQuietState quietState; for (uint32_t group = 0U; group < groupCount; ++group) { const int32_t peer = AllToAllGroupDevicePeer( rank, rankSize, group, lane, groupWidth); @@ -793,10 +861,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( groupTrace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, peer, selectedQp, putBegin, AllToAllGroupTraceCycle(groupTrace)); - pending[pendingCount++] = { - peer, selectedQp, group, pass, expectedToken}; - if (pendingCount == quietBatch && !AllToAllGroupFlushQuiet( - args, pending, pendingCount, debug, blockIdx, + if (!AllToAllGroupCompleteQuiet( + args, quietState, quietBatch, peer, selectedQp, + group, pass, expectedToken, debug, blockIdx, groupTrace, traceIteration, groupCount, passCount)) { AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -805,9 +872,9 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( } } } - if (pendingCount != 0U && !AllToAllGroupFlushQuiet( - args, pending, pendingCount, debug, blockIdx, - groupTrace, traceIteration, groupCount, passCount)) { + if (!AllToAllGroupFinishQuiet( + args, quietState, debug, blockIdx, groupTrace, traceIteration, + groupCount, passCount)) { AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; @@ -816,6 +883,44 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( kernelBegin, AllToAllGroupTraceCycle(groupTrace)); } +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, + GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1, + GM_ADDR groupTraceGM, uint32_t traceIteration, + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch) +{ + AllToAllGroupKernelImpl( + commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, + groupWidth, quietBatch); +} + +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, + GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1, + GM_ADDR groupTraceGM, uint32_t traceIteration, + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch) +{ + AllToAllGroupKernelImpl( + commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, + groupWidth, quietBatch); +} + void launch_tilexr_udma_all_to_all_group( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR registeredMemory, GM_ADDR debug, uint32_t invocationId, @@ -827,10 +932,19 @@ void launch_tilexr_udma_all_to_all_group( uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, uint32_t groupWidth, uint32_t quietBatch) { - tilexr_udma_all_to_all_group_kernel<<>>( - commArgs, input, output, registeredMemory, debug, invocationId, - elementsPerPeer, chunkElements, passCount, groupCount, - payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, routeStage, - multiChannel, primaryRouteParts, groupWidth, quietBatch); + if (quietBatch == 1U) { + tilexr_udma_all_to_all_group_kernel<<>>( + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + groupTrace, traceIteration, routeStage, + multiChannel, primaryRouteParts, groupWidth, quietBatch); + } else { + tilexr_udma_all_to_all_group_batch_kernel<<>>( + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + groupTrace, traceIteration, routeStage, + multiChannel, primaryRouteParts, groupWidth, quietBatch); + } } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 2454421b..e42115ff 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -462,7 +462,16 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupFlushQuiet"); CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, request.peer, request.qpIdx)"); - CHECK_CONTAINS(kernel, "pendingCount == quietBatch"); + CHECK_CONTAINS(kernel, "state.pendingCount != quietBatch"); + CHECK_CONTAINS(kernel, "template "); + CHECK_CONTAINS(kernel, "struct AllToAllGroupQuietState"); + CHECK_CONTAINS(kernel, "AllToAllGroupQuietState quietState"); + CHECK_CONTAINS(kernel, "AllToAllGroupCompleteQuiet("); + CHECK_CONTAINS(kernel, "AllToAllGroupFinishQuiet("); + CHECK_CONTAINS(kernel, + "UDMAQuietStatusOnQp(args, peer, selectedQp)"); + CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_batch_kernel"); + CHECK_CONTAINS(kernel, "if (quietBatch == 1U)"); CHECK_CONTAINS(kernel, "uint32_t groupWidth, uint32_t quietBatch"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); From 9319b4c939a1125a218c54af3c6c75584fb5b88b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sat, 1 Aug 2026 16:13:47 +0800 Subject: [PATCH 127/163] fix(udma): make rootinfo device offset order-independent --- src/comm/udma/tilexr_udma_transport.cpp | 2 ++ tests/udma/unit/test_tilexr_udma_transport_layout.cpp | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 13a2399f..46d50b00 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -281,6 +281,8 @@ bool ParseRootInfo(TileXRRootInfo& root) } if (root.deviceToLocalId.empty()) { root.deviceIdOffset = deviceId; + } else { + root.deviceIdOffset = std::min(root.deviceIdOffset, deviceId); } root.deviceToLocalId[deviceId] = localId; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index c291a72d..e095b912 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -330,6 +330,14 @@ void TestRootInfoEidBytesSelectRuntimeContexts() CHECK_CONTAINS(transport, "ctxHandleByEid_[eidIndex] = ctxHandle"); } +void TestRootInfoDeviceOffsetDoesNotDependOnEntryOrder() +{ + const std::string transport = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/comm/udma/tilexr_udma_transport.cpp"); + + CHECK_CONTAINS(transport, "root.deviceIdOffset = std::min(root.deviceIdOffset, deviceId)"); +} + void TestMemoryRegistrationUsesOfficialUbFlags() { const std::string transport = @@ -391,6 +399,7 @@ int main() TestTransportUsesPerPeerQueues(); TestTransportHasOptInSharedQpPool(); TestRootInfoEidBytesSelectRuntimeContexts(); + TestRootInfoDeviceOffsetDoesNotDependOnEntryOrder(); TestMemoryRegistrationUsesOfficialUbFlags(); TestDeviceSgeUsesPerPeerLocalTokenId(); TestDeviceSqeInitializesOfficialFields(); From 8aca7d0c3d6b45ecb181ecad249c8a579766335d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 30 Jul 2026 15:24:29 +0800 Subject: [PATCH 128/163] feat(udma): add multi-region registered memory --- src/comm/comm_wrap.cpp | 12 + src/comm/tilexr_comm.cpp | 89 +++-- src/comm/tilexr_comm.h | 2 + src/comm/udma/tilexr_udma_layout.cpp | 25 +- src/comm/udma/tilexr_udma_layout.h | 13 + src/comm/udma/tilexr_udma_transport.cpp | 349 ++++++++++++------ src/comm/udma/tilexr_udma_transport.h | 13 +- src/include/tilexr_api.h | 7 + src/include/tilexr_udma.h | 224 ++++++++--- src/include/tilexr_udma_reg.h | 35 +- src/include/tilexr_udma_types.h | 1 + tests/udma/CMakeLists.txt | 1 + tests/udma/demo/run_tilexr_udma_demo.sh | 2 +- .../demo/tilexr_udma_alltoall_constants.h | 19 + .../tilexr_udma_alltoall_group_kernel.cpp | 2 +- .../demo/tilexr_udma_alltoall_group_layout.h | 2 +- tests/udma/demo/tilexr_udma_alltoall_layout.h | 5 +- tests/udma/demo/tilexr_udma_demo.cpp | 315 +++++++++++++++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 47 ++- ...test_tilexr_udma_alltoall_group_layout.cpp | 17 +- .../unit/test_tilexr_udma_alltoall_layout.cpp | 27 +- tests/udma/unit/test_tilexr_udma_registry.cpp | 29 +- .../test_tilexr_udma_transport_layout.cpp | 25 ++ 23 files changed, 1018 insertions(+), 243 deletions(-) create mode 100644 tests/udma/demo/tilexr_udma_alltoall_constants.h diff --git a/src/comm/comm_wrap.cpp b/src/comm/comm_wrap.cpp index ca9f29c9..bb5c1465 100644 --- a/src/comm/comm_wrap.cpp +++ b/src/comm/comm_wrap.cpp @@ -158,6 +158,18 @@ int TileXRUDMARegister(TileXRCommPtr comm, GM_ADDR localPtr, size_t bytes, TileX return c->RegisterUDMAMemory(localPtr, bytes, handle); } +int TileXRUDMARegisterRegions(TileXRCommPtr comm, const TileXR::TileXRUDMARegionDesc *regions, + uint32_t regionCount, TileXRUDMAMemHandle *handle) +{ + if (comm == nullptr || regions == nullptr || handle == nullptr || regionCount == 0 || + regionCount > TileXR::TILEXR_UDMA_MAX_REGIONS) { + TILEXR_LOG(ERROR) << "TileXRUDMARegisterRegions invalid input"; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + auto* c = static_cast(comm); + return c->RegisterUDMAMemoryRegions(regions, regionCount, handle); +} + int TileXRUDMAUnregister(TileXRCommPtr comm, TileXRUDMAMemHandle handle) { if (comm == nullptr) { diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 3fc4ea2e..20288c53 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -336,14 +336,28 @@ void TileXRComm::FreeUDMARegistry() } int TileXRComm::RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMemHandle *handle) +{ + TileXRUDMARegionDesc region {}; + region.base = localPtr; + region.bytes = bytes; + return RegisterUDMAMemoryRegions(®ion, 1, handle); +} + +int TileXRComm::RegisterUDMAMemoryRegions( + const TileXRUDMARegionDesc *regions, uint32_t regionCount, TileXRUDMAMemHandle *handle) { if (!inited_) { - TILEXR_LOG(ERROR) << "TileXRUDMARegister requires initialized communicator"; + TILEXR_LOG(ERROR) << "TileXRUDMARegisterRegions requires initialized communicator"; return TILEXR_ERROR_NOT_INITIALIZED; } - if (localPtr == nullptr || bytes == 0 || handle == nullptr) { + if (regions == nullptr || regionCount == 0 || regionCount > TILEXR_UDMA_MAX_REGIONS || handle == nullptr) { return TILEXR_ERROR_PARA_CHECK_FAIL; } + for (uint32_t i = 0; i < regionCount; ++i) { + if (regions[i].base == nullptr || regions[i].bytes == 0 || regions[i].bytes > UINT32_MAX) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + } if (!((commArgs_.extraFlag & ExtraFlag::UDMA) != 0 && commArgs_.udmaInfoPtr != nullptr)) { TILEXR_LOG(WARN) << "TileXRUDMARegister called while UDMA is unavailable"; return TILEXR_ERROR_NOT_FOUND; @@ -353,15 +367,20 @@ int TileXRComm::RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMem return TILEXR_ERROR_INTERNAL; } - TileXRUDMARegionDesc localRegion {}; - localRegion.base = localPtr; - localRegion.bytes = bytes; - if (udmaTransport_ == nullptr || !udmaTransport_->IsAvailable()) { TILEXR_LOG(ERROR) << "TileXR UDMA transport is unavailable"; return TILEXR_ERROR_NOT_FOUND; } - int ret = udmaTransport_->RegisterMemory(localPtr, bytes); + if (udmaRegisteredPtr_ != nullptr) { + udmaTransport_->UnregisterMemory(udmaRegisteredPtr_); + udmaRegisteredPtr_ = nullptr; + FreeUDMARegistry(); + int updateRet = UpdateCommArgsDev(); + if (updateRet != TILEXR_SUCCESS) { + return updateRet; + } + } + int ret = udmaTransport_->RegisterMemoryRegions(regions, regionCount); if (ret != TILEXR_SUCCESS) { TILEXR_LOG(ERROR) << "TileXR UDMA memory registration failed: " << ret; return TILEXR_ERROR_INTERNAL; @@ -371,63 +390,77 @@ int TileXRComm::RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMem ret = UpdateCommArgsDev(); if (ret != TILEXR_SUCCESS) { TILEXR_LOG(ERROR) << "TileXRUDMARegister failed to refresh CommArgs after UDMA info update: " << ret; - udmaTransport_->UnregisterMemory(localPtr); + udmaTransport_->UnregisterMemory(regions[0].base); return ret; } if (socketExchange_ == nullptr) { TILEXR_LOG(ERROR) << "TileXRUDMARegister requires live socket exchange"; - udmaTransport_->UnregisterMemory(localPtr); + udmaTransport_->UnregisterMemory(regions[0].base); return TILEXR_ERROR_INTERNAL; } - std::vector allRegions(rankSize_); - ret = socketExchange_->AllGather(&localRegion, 1, allRegions.data()); + std::vector allRegionCounts(rankSize_); + ret = socketExchange_->AllGather(®ionCount, 1, allRegionCounts.data()); if (ret != TILEXR_SUCCESS) { - TILEXR_LOG(ERROR) << "TileXRUDMARegister allgather failed: " << ret; - udmaTransport_->UnregisterMemory(localPtr); + TILEXR_LOG(ERROR) << "TileXRUDMARegisterRegions count allgather failed: " << ret; + udmaTransport_->UnregisterMemory(regions[0].base); + return ret; + } + for (int rank = 0; rank < rankSize_; ++rank) { + if (allRegionCounts[rank] != regionCount) { + TILEXR_LOG(ERROR) << "TileXRUDMARegisterRegions region count mismatch at rank " << rank + << ": " << allRegionCounts[rank] << " vs " << regionCount; + udmaTransport_->UnregisterMemory(regions[0].base); + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + std::vector allRegions(static_cast(rankSize_) * regionCount); + ret = socketExchange_->AllGather(regions, regionCount, allRegions.data()); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "TileXRUDMARegisterRegions allgather failed: " << ret; + udmaTransport_->UnregisterMemory(regions[0].base); return ret; } TileXRUDMARegistry nextRegistry {}; nextRegistry.rankSize = static_cast(rankSize_); - nextRegistry.regionCount = 1; + nextRegistry.regionCount = regionCount; for (int i = 0; i < rankSize_; ++i) { - if (allRegions[i].base == nullptr || allRegions[i].bytes == 0) { - TILEXR_LOG(ERROR) << "TileXRUDMARegister received invalid region from rank " << i; - udmaTransport_->UnregisterMemory(localPtr); - return TILEXR_ERROR_PARA_CHECK_FAIL; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + const auto& region = allRegions[static_cast(i) * regionCount + regionIndex]; + if (region.base == nullptr || region.bytes == 0 || region.bytes > UINT32_MAX) { + TILEXR_LOG(ERROR) << "TileXRUDMARegisterRegions received invalid region from rank " + << i << " region " << regionIndex; + udmaTransport_->UnregisterMemory(regions[0].base); + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + nextRegistry.regions[i][regionIndex] = region; } - nextRegistry.regions[i] = allRegions[i]; } GM_ADDR nextRegistryDev = nullptr; ret = aclrtMalloc(reinterpret_cast(&nextRegistryDev), sizeof(nextRegistry), ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { TILEXR_LOG(ERROR) << "aclrtMalloc UDMA registry failed: " << ret; - udmaTransport_->UnregisterMemory(localPtr); + udmaTransport_->UnregisterMemory(regions[0].base); return TILEXR_ERROR_INTERNAL; } ret = aclrtMemcpy(nextRegistryDev, sizeof(nextRegistry), &nextRegistry, sizeof(nextRegistry), ACL_MEMCPY_HOST_TO_DEVICE); if (ret != ACL_SUCCESS) { TILEXR_LOG(ERROR) << "aclrtMemcpy UDMA registry failed: " << ret; aclrtFree(nextRegistryDev); - udmaTransport_->UnregisterMemory(localPtr); + udmaTransport_->UnregisterMemory(regions[0].base); return TILEXR_ERROR_INTERNAL; } - if (udmaRegisteredPtr_ != nullptr) { - udmaTransport_->UnregisterMemory(udmaRegisteredPtr_); - udmaRegisteredPtr_ = nullptr; - } - FreeUDMARegistry(); udmaRegistry_ = nextRegistry; udmaRegistryDev_ = nextRegistryDev; - udmaRegisteredPtr_ = localPtr; + udmaRegisteredPtr_ = regions[0].base; commArgs_.udmaRegistryPtr = udmaRegistryDev_; *handle = 0; ret = UpdateCommArgsDev(); if (ret != TILEXR_SUCCESS) { - udmaTransport_->UnregisterMemory(localPtr); + udmaTransport_->UnregisterMemory(regions[0].base); udmaRegisteredPtr_ = nullptr; FreeUDMARegistry(); } diff --git a/src/comm/tilexr_comm.h b/src/comm/tilexr_comm.h index abc9e9a1..83afd890 100644 --- a/src/comm/tilexr_comm.h +++ b/src/comm/tilexr_comm.h @@ -44,6 +44,8 @@ class TileXRComm { CommArgs* GetCommArgs(); int64_t NextMagic(); int RegisterUDMAMemory(GM_ADDR localPtr, size_t bytes, TileXRUDMAMemHandle *handle); + int RegisterUDMAMemoryRegions( + const TileXRUDMARegionDesc *regions, uint32_t regionCount, TileXRUDMAMemHandle *handle); int UnregisterUDMAMemory(TileXRUDMAMemHandle handle); GM_ADDR GetUDMARegistryPtr() const; bool IsSDMAAvailable() const; diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 593ad076..968fb3cd 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -37,7 +37,8 @@ int BuildUDMAInfoImage( std::vector& bytes) { return BuildUDMAInfoImage( - deviceBase, TILEXR_UDMA_QP_NUM, sq, rq, scq, rcq, mem, std::vector(mem.size(), 1), info, bytes); + deviceBase, TILEXR_UDMA_QP_NUM, 1U, sq, rq, scq, rcq, mem, + std::vector(mem.size(), 1), info, bytes); } int BuildUDMAInfoImage( @@ -51,10 +52,27 @@ int BuildUDMAInfoImage( const std::vector& qpWeights, UDMAInfo& info, std::vector& bytes) +{ + return BuildUDMAInfoImage( + deviceBase, qpNum, 1U, sq, rq, scq, rcq, mem, qpWeights, info, bytes); +} + +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpNum, + uint32_t regionCount, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + const std::vector& qpWeights, + UDMAInfo& info, + std::vector& bytes) { const size_t rankCount = sq.size(); - if (qpNum == 0 || rankCount == 0 || rankCount % qpNum != 0 || rq.size() != rankCount || - scq.size() != rankCount || rcq.size() != rankCount || mem.size() != rankCount || + if (qpNum == 0 || regionCount == 0 || rankCount == 0 || rankCount % qpNum != 0 || rq.size() != rankCount || + scq.size() != rankCount || rcq.size() != rankCount || mem.size() != rankCount * regionCount || qpWeights.size() != rankCount) { return TILEXR_UDMA_LAYOUT_INVALID; } @@ -69,6 +87,7 @@ int BuildUDMAInfoImage( info = {}; info.qpNum = qpNum; + info.regionCount = regionCount; info.sqPtr = deviceBase + sqOffset; info.rqPtr = deviceBase + rqOffset; info.scqPtr = deviceBase + scqOffset; diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index e8e83e9c..691164c6 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -28,6 +28,19 @@ int BuildUDMAInfoImage( UDMAInfo& info, std::vector& bytes); +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpNum, + uint32_t regionCount, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + const std::vector& qpWeights, + UDMAInfo& info, + std::vector& bytes); + int BuildUDMAInfoImage( uintptr_t deviceBase, uint32_t qpNum, diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 46d50b00..24813b3e 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -751,6 +751,10 @@ int TileXRUDMATransport::BuildRoutes() maxRouteCount = std::max(maxRouteCount, entry.second.size()); } qpNum_ = static_cast(maxRouteCount); + logicalQpNum_ = qpNum_; + logicalQpsPerRoute_ = qpsPerRoute_; + logicalPeerQpRouteEids_ = peerQpRouteEids_; + logicalPeerQpRouteWeights_ = peerQpRouteWeights_; return TILEXR_SUCCESS; } @@ -1316,10 +1320,16 @@ int TileXRUDMATransport::RefreshUDMAInfo() } } - std::vector localMemByEid(eidCount_); - for (const auto& memEntry : localMemInfoByEid_) { - if (memEntry.first < eidCount_) { - localMemByEid[memEntry.first] = memEntry.second; + const uint32_t regionCount = static_cast(localMemInfoByRegion_.size()); + if (regionCount == 0) { + return TILEXR_ERROR_NOT_FOUND; + } + std::vector localMemByRegionEid(static_cast(regionCount) * eidCount_); + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + for (const auto& memEntry : localMemInfoByRegion_[regionIndex]) { + if (memEntry.first < eidCount_) { + localMemByRegionEid[static_cast(regionIndex) * eidCount_ + memEntry.first] = memEntry.second; + } } } std::vector localEids(eidCount_); @@ -1329,8 +1339,9 @@ int TileXRUDMATransport::RefreshUDMAInfo() } } - std::vector allMem(options_.rankSize * eidCount_); - int ret = options_.exchange->AllGather(localMemByEid.data(), localMemByEid.size(), allMem.data()); + std::vector allMem(static_cast(options_.rankSize) * regionCount * eidCount_); + int ret = options_.exchange->AllGather( + localMemByRegionEid.data(), localMemByRegionEid.size(), allMem.data()); if (ret != TILEXR_SUCCESS) { return ret; } @@ -1356,7 +1367,7 @@ int TileXRUDMATransport::RefreshUDMAInfo() std::vector rq(queueEntries); std::vector scq(queueEntries); std::vector rcq(queueEntries); - std::vector mem(queueEntries); + std::vector mem(queueEntries * regionCount); std::vector qpWeights(queueEntries, 1); for (int rank = 0; rank < options_.rankSize; ++rank) { @@ -1441,8 +1452,10 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_ERROR_INTERNAL; } auto& queue = *queuePtr; - if (!registeredMem_.empty()) { - const auto& localMrs = registeredMem_.begin()->second; + if (!registeredRegions_.empty()) { + const auto localMrsIt = registeredMem_.find( + reinterpret_cast(registeredRegions_.front().base)); + const auto& localMrs = localMrsIt->second; const auto localMrIt = localMrs.find(localEid); if (localMrIt != localMrs.end()) { queue.localWq.localTokenId = localMrIt->second.tokenId; @@ -1453,17 +1466,21 @@ int TileXRUDMATransport::RefreshUDMAInfo() rq[entryIndex] = queue.localWq; scq[entryIndex] = queue.localCq; rcq[entryIndex] = queue.localCq; - if (rank == options_.rank) { - const auto localMemIt = localMemInfoByEid_.find(localEid); - if (localMemIt != localMemInfoByEid_.end()) { - mem[entryIndex] = localMemIt->second; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + const size_t memIndex = entryIndex * regionCount + regionIndex; + if (rank == options_.rank) { + const auto localMemIt = localMemInfoByRegion_[regionIndex].find(localEid); + if (localMemIt != localMemInfoByRegion_[regionIndex].end()) { + mem[memIndex] = localMemIt->second; + } + } else { + const size_t allIndex = (static_cast(rank) * regionCount + regionIndex) * eidCount_ + remoteEid; + mem[memIndex] = allMem[allIndex]; + mem[memIndex].tpn = sharedRemotePtr == nullptr ? queue.tpn : sharedRemotePtr->tpn; } - } else { - mem[entryIndex] = allMem[rank * eidCount_ + remoteEid]; - mem[entryIndex].tpn = sharedRemotePtr == nullptr ? queue.tpn : sharedRemotePtr->tpn; + mem[memIndex].eidAddr = reinterpret_cast( + eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); } - mem[entryIndex].eidAddr = reinterpret_cast( - eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); if (rank != options_.rank) { const auto weightIt = peerQpRouteWeights_.find(rank); if (weightIt != peerQpRouteWeights_.end() && qpIdx < weightIt->second.size()) { @@ -1484,16 +1501,18 @@ int TileXRUDMATransport::RefreshUDMAInfo() << " wqeCnt=" << PtrToHex(sq[entryIndex].wqeCntAddr) << " cqBuf=" << PtrToHex(scq[entryIndex].bufAddr) << " cqTail=" << PtrToHex(scq[entryIndex].tailAddr) - << " memAddr=" << PtrToHex(mem[entryIndex].addr) - << " memLen=" << mem[entryIndex].len - << " memTid=" << mem[entryIndex].tid - << " memTpn=" << mem[entryIndex].tpn - << " memEidAddr=" << PtrToHex(mem[entryIndex].eidAddr); + << " regionCount=" << regionCount + << " memAddr=" << PtrToHex(mem[entryIndex * regionCount].addr) + << " memLen=" << mem[entryIndex * regionCount].len + << " memTid=" << mem[entryIndex * regionCount].tid + << " memTpn=" << mem[entryIndex * regionCount].tpn + << " memEidAddr=" << PtrToHex(mem[entryIndex * regionCount].eidAddr); } } } - const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo) + sizeof(uint32_t); + const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + + regionCount * sizeof(UDMAMemInfo) + sizeof(uint32_t); const uint32_t requiredInfoSize = static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize * qpNum_); if (udmaInfoDev_ == nullptr || udmaInfoSize_ < requiredInfoSize) { @@ -1512,7 +1531,8 @@ int TileXRUDMATransport::RefreshUDMAInfo() UDMAInfo info {}; std::vector image; ret = BuildUDMAInfoImage( - reinterpret_cast(udmaInfoDev_), qpNum_, sq, rq, scq, rcq, mem, qpWeights, info, image); + reinterpret_cast(udmaInfoDev_), qpNum_, regionCount, + sq, rq, scq, rcq, mem, qpWeights, info, image); if (ret != TILEXR_UDMA_LAYOUT_SUCCESS) { return TILEXR_ERROR_PARA_CHECK_FAIL; } @@ -1530,6 +1550,7 @@ int TileXRUDMATransport::EnsureUDMAInfoBuffer() } UDMAInfo info {}; info.qpNum = 1; + info.regionCount = 1; udmaInfoSize_ = static_cast(sizeof(UDMAInfo)); int ret = aclrtMalloc(reinterpret_cast(&udmaInfoDev_), udmaInfoSize_, ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { @@ -1547,65 +1568,138 @@ int TileXRUDMATransport::EnsureUDMAInfoBuffer() int TileXRUDMATransport::RegisterMemory(GM_ADDR localPtr, size_t bytes) { - if (!available_ || localPtr == nullptr || bytes == 0) { + TileXRUDMARegionDesc region {}; + region.base = localPtr; + region.bytes = bytes; + return RegisterMemoryRegions(®ion, 1); +} + +int TileXRUDMATransport::ConfigureRegionQueues(uint32_t regionCount) +{ + if (regionCount == 0 || logicalQpNum_ == 0 || logicalQpsPerRoute_ == 0) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (sharedQpPool_ && regionCount > 1) { + TILEXR_LOG(WARN) << "UDMA multi-region registration is incompatible with shared QP pool"; return TILEXR_ERROR_NOT_FOUND; } + constexpr uint32_t maxQpsPerRoute = 64; + if (logicalQpsPerRoute_ > maxQpsPerRoute / regionCount || + logicalQpNum_ > UINT32_MAX / regionCount) { + TILEXR_LOG(WARN) << "UDMA region QP expansion exceeds limit: logicalQpsPerRoute=" + << logicalQpsPerRoute_ << " regionCount=" << regionCount; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + + qpsPerRoute_ = logicalQpsPerRoute_ * regionCount; + qpNum_ = logicalQpNum_ * regionCount; + peerQpRouteEids_.clear(); + peerQpRouteWeights_.clear(); + for (const auto& entry : logicalPeerQpRouteEids_) { + auto& expanded = peerQpRouteEids_[entry.first]; + for (uint32_t eid : entry.second) { + expanded.insert(expanded.end(), regionCount, eid); + } + } + for (const auto& entry : logicalPeerQpRouteWeights_) { + auto& expanded = peerQpRouteWeights_[entry.first]; + for (uint32_t weight : entry.second) { + expanded.insert(expanded.end(), regionCount, weight); + } + } + return TILEXR_SUCCESS; +} + +int TileXRUDMATransport::RegisterMemoryRegions( + const TileXRUDMARegionDesc* regions, uint32_t regionCount) +{ + if (!available_ || regions == nullptr || regionCount == 0 || + regionCount > TILEXR_UDMA_MAX_REGIONS) { + return TILEXR_ERROR_NOT_FOUND; + } + for (uint32_t i = 0; i < regionCount; ++i) { + if (regions[i].base == nullptr || regions[i].bytes == 0 || regions[i].bytes > UINT32_MAX) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + } const bool diag = UDMADiagEnabled(); if (diag) { - TILEXR_LOG(INFO) << "UDMA diag register memory begin rank " << options_.rank - << " ptr=" << PtrToHex(reinterpret_cast(localPtr)) - << " bytes=" << bytes; + TILEXR_LOG(INFO) << "UDMA diag register regions begin rank " << options_.rank + << " regionCount=" << regionCount; } CleanupMemory(); CleanupQueues(); - registeredPtr_ = nullptr; - int ret = RegisterMemoryOnContexts(localPtr, bytes); + int ret = ConfigureRegionQueues(regionCount); if (ret != TILEXR_SUCCESS) { - TILEXR_LOG(WARN) << "UDMA register memory on contexts failed rank " << options_.rank - << " ret=" << ret; return ret; } - registeredPtr_ = localPtr; + registeredPtr_ = nullptr; + registeredRegions_.assign(regions, regions + regionCount); + localMemInfoByRegion_.resize(regionCount); + ret = TILEXR_SUCCESS; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + ret = RegisterMemoryOnContexts(regions[regionIndex], regionIndex); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(WARN) << "UDMA register memory region failed rank " << options_.rank + << " region=" << regionIndex << " ret=" << ret; + CleanupMemory(); + return ret; + } + } + registeredPtr_ = regions[0].base; ret = CreateQueues(); if (ret != TILEXR_SUCCESS) { TILEXR_LOG(WARN) << "UDMA create queues after memory register failed rank " << options_.rank << " ret=" << ret; + CleanupQueues(); + CleanupMemory(); + registeredPtr_ = nullptr; return ret; } ret = ImportQueues(); if (ret != TILEXR_SUCCESS) { TILEXR_LOG(WARN) << "UDMA import queues after memory register failed rank " << options_.rank << " ret=" << ret; + CleanupQueues(); + CleanupMemory(); + registeredPtr_ = nullptr; return ret; } ret = ExchangeAndImportMemory(); if (ret != TILEXR_SUCCESS) { TILEXR_LOG(WARN) << "UDMA exchange/import memory failed rank " << options_.rank << " ret=" << ret; + CleanupMemory(); + CleanupQueues(); + registeredPtr_ = nullptr; return ret; } ret = RefreshUDMAInfo(); if (ret != TILEXR_SUCCESS) { TILEXR_LOG(WARN) << "UDMA refresh info after memory register failed rank " << options_.rank << " ret=" << ret; + CleanupMemory(); + CleanupQueues(); + registeredPtr_ = nullptr; return ret; } if (diag) { - TILEXR_LOG(INFO) << "UDMA diag register memory end rank " << options_.rank; + TILEXR_LOG(INFO) << "UDMA diag register regions end rank " << options_.rank + << " regionCount=" << regionCount; } return TILEXR_SUCCESS; } -int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes) +int TileXRUDMATransport::RegisterMemoryOnContexts( + const TileXRUDMARegionDesc& region, uint32_t regionIndex) { std::map byEid; - localMemInfoByEid_.clear(); for (const auto& ctxEntry : ctxHandleByEid_) { const uint32_t eidIndex = ctxEntry.first; void* tokenHandle = tokenHandleByEid_[eidIndex]; MrRegInfoT mrInfo {}; - mrInfo.in.mem.addr = reinterpret_cast(localPtr); - mrInfo.in.mem.size = bytes; + mrInfo.in.mem.addr = reinterpret_cast(region.base); + 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; @@ -1620,16 +1714,17 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes TILEXR_LOG(WARN) << "UDMA RaCtxLmemRegister failed rank " << options_.rank << " eid=" << eidIndex << " ctx=" << ctxEntry.second - << " ptr=" << PtrToHex(reinterpret_cast(localPtr)) - << " bytes=" << bytes + << " region=" << regionIndex + << " ptr=" << PtrToHex(reinterpret_cast(region.base)) + << " bytes=" << region.bytes << " ret=" << ret << " handle=" << lmemHandle; return TILEXR_ERROR_INTERNAL; } RegMemResultInfo result {}; - result.address = reinterpret_cast(localPtr); - result.size = bytes; + result.address = reinterpret_cast(region.base); + result.size = region.bytes; result.lmemHandle = lmemHandle; result.key = mrInfo.out.key; result.tokenId = mrInfo.out.ub.tokenId; @@ -1644,14 +1739,17 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes memInfo.rmtJettyType = 1; memInfo.targetHint = 0; memInfo.tpn = 0; + // A QP is bound to the context token allocated during CreateContexts. + // Multiple independent MRs on that context must reuse the same token. memInfo.tid = mrInfo.out.ub.tokenId >> 8; memInfo.rmtTokenValue = TILEXR_UDMA_TOKEN_VALUE; - memInfo.len = static_cast(std::min(bytes, UINT32_MAX)); - memInfo.addr = reinterpret_cast(localPtr); - localMemInfoByEid_[eidIndex] = memInfo; + memInfo.len = static_cast(region.bytes); + memInfo.addr = reinterpret_cast(region.base); + localMemInfoByRegion_[regionIndex][eidIndex] = memInfo; if (UDMADiagEnabled()) { TILEXR_LOG(INFO) << "UDMA diag lmem registered rank " << options_.rank << " eid=" << eidIndex + << " region=" << regionIndex << " lmem=" << lmemHandle << " tokenId=" << result.tokenId << " tid=" << memInfo.tid @@ -1659,17 +1757,23 @@ int TileXRUDMATransport::RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes << " keySize=" << static_cast(result.key.size); } } - registeredMem_[reinterpret_cast(localPtr)] = byEid; + registeredMem_[reinterpret_cast(region.base)] = byEid; return TILEXR_SUCCESS; } int TileXRUDMATransport::ExchangeAndImportMemory() { - if (registeredMem_.empty()) { + if (registeredMem_.empty() || registeredRegions_.empty()) { return TILEXR_ERROR_NOT_FOUND; } - const auto& localByEid = registeredMem_.begin()->second; - uint32_t localCount = static_cast(localByEid.size()); + uint32_t localCount = 0; + for (const auto& region : registeredRegions_) { + const auto regionIt = registeredMem_.find(reinterpret_cast(region.base)); + if (regionIt == registeredMem_.end()) { + return TILEXR_ERROR_INTERNAL; + } + localCount += static_cast(regionIt->second.size()); + } std::vector allCounts(options_.rankSize); int ret = options_.exchange->AllGather(&localCount, 1, allCounts.data()); if (ret != TILEXR_SUCCESS) { @@ -1684,6 +1788,7 @@ int TileXRUDMATransport::ExchangeAndImportMemory() } struct ExchangedMrInfo { + uint32_t regionIndex; uint32_t eidIndex; uint32_t valid; RegMemResultInfo mr; @@ -1691,11 +1796,16 @@ int TileXRUDMATransport::ExchangeAndImportMemory() 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 < registeredRegions_.size(); ++regionIndex) { + const auto& byEid = registeredMem_.find( + reinterpret_cast(registeredRegions_[regionIndex].base))->second; + for (const auto& entry : byEid) { + local[idx].regionIndex = regionIndex; + local[idx].eidIndex = entry.first; + local[idx].valid = 1; + local[idx].mr = entry.second; + ++idx; + } } std::vector all(options_.rankSize * maxCount); ret = options_.exchange->AllGather(local.data(), local.size(), all.data()); @@ -1725,54 +1835,48 @@ int TileXRUDMATransport::ExchangeAndImportMemory() if (localRoutes.empty() || localRoutes.size() != remoteRoutes.size()) { return TILEXR_ERROR_INTERNAL; } - std::vector remoteHandles(localRoutes.size(), nullptr); - for (size_t routeIdx = 0; routeIdx < localRoutes.size(); ++routeIdx) { - const uint32_t localEid = localRoutes[routeIdx]; - const uint32_t remoteEid = remoteRoutes[routeIdx]; - const ExchangedMrInfo* remote = nullptr; - for (uint32_t i = 0; i < allCounts[peer]; ++i) { - const auto& candidate = all[peer * maxCount + i]; - if (candidate.valid != 0 && candidate.eidIndex == remoteEid) { - remote = &candidate; - break; + std::vector> remoteHandles( + registeredRegions_.size(), std::vector(localRoutes.size(), nullptr)); + for (uint32_t regionIndex = 0; regionIndex < registeredRegions_.size(); ++regionIndex) { + for (size_t routeIdx = 0; routeIdx < localRoutes.size(); ++routeIdx) { + const uint32_t localEid = localRoutes[routeIdx]; + const uint32_t remoteEid = remoteRoutes[routeIdx]; + const ExchangedMrInfo* remote = nullptr; + for (uint32_t i = 0; i < allCounts[peer]; ++i) { + const auto& candidate = all[peer * maxCount + i]; + if (candidate.valid != 0 && candidate.regionIndex == regionIndex && + candidate.eidIndex == remoteEid) { + remote = &candidate; + break; + } + } + if (remote == nullptr || ctxHandleByEid_.count(localEid) == 0) { + TILEXR_LOG(WARN) << "UDMA remote memory info missing rank " << options_.rank + << " peer=" << peer << " region=" << regionIndex + << " localEid=" << localEid << " remoteEid=" << remoteEid; + return TILEXR_ERROR_INTERNAL; + } + MrImportInfoT importInfo {}; + importInfo.in.key = remote->mr.key; + importInfo.in.ub.tokenValue = remote->mr.tokenValue; + importInfo.in.ub.flags.bs.cacheable = remote->mr.cacheable; + importInfo.in.ub.flags.bs.access = remote->mr.access; + void* remoteHandle = nullptr; + ret = loader_.RaCtxRmemImport(ctxHandleByEid_[localEid], &importInfo, &remoteHandle); + if (ret != 0 || remoteHandle == nullptr) { + TILEXR_LOG(WARN) << "UDMA RaCtxRmemImport failed rank " << options_.rank + << " peer=" << peer << " region=" << regionIndex + << " localEid=" << localEid << " remoteEid=" << remoteEid + << " ret=" << ret << " handle=" << remoteHandle; + return TILEXR_ERROR_INTERNAL; + } + remoteHandles[regionIndex][routeIdx] = remoteHandle; + if (UDMADiagEnabled()) { + TILEXR_LOG(INFO) << "UDMA diag rmem imported rank " << options_.rank + << " peer=" << peer << " region=" << regionIndex + << " localEid=" << localEid << " remoteEid=" << remoteEid + << " remoteHandle=" << remoteHandle; } - } - if (remote == nullptr || ctxHandleByEid_.count(localEid) == 0) { - TILEXR_LOG(WARN) << "UDMA remote memory info missing rank " << options_.rank - << " peer=" << peer - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " peerCount=" << allCounts[peer] - << " maxCount=" << maxCount; - return TILEXR_ERROR_INTERNAL; - } - MrImportInfoT importInfo {}; - importInfo.in.key = remote->mr.key; - importInfo.in.ub.tokenValue = remote->mr.tokenValue; - importInfo.in.ub.flags.bs.cacheable = remote->mr.cacheable; - importInfo.in.ub.flags.bs.access = remote->mr.access; - void* remoteHandle = nullptr; - ret = loader_.RaCtxRmemImport(ctxHandleByEid_[localEid], &importInfo, &remoteHandle); - if (ret != 0 || remoteHandle == nullptr) { - TILEXR_LOG(WARN) << "UDMA RaCtxRmemImport failed rank " << options_.rank - << " peer=" << peer - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " ctx=" << ctxHandleByEid_[localEid] - << " ret=" << ret - << " handle=" << remoteHandle - << " remoteToken=" << remote->mr.tokenValue - << " remoteTokenId=" << remote->mr.tokenId - << " remoteKeySize=" << static_cast(remote->mr.key.size); - return TILEXR_ERROR_INTERNAL; - } - remoteHandles[routeIdx] = remoteHandle; - if (UDMADiagEnabled()) { - TILEXR_LOG(INFO) << "UDMA diag rmem imported rank " << options_.rank - << " peer=" << peer - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " remoteHandle=" << remoteHandle; } } remoteMemHandlesByPeer_[peer] = remoteHandles; @@ -1787,10 +1891,8 @@ int TileXRUDMATransport::UnregisterMemory(GM_ADDR localPtr) } CleanupMemory(); registeredPtr_ = nullptr; - localMemInfoByEid_.clear(); - if (available_) { - return RefreshUDMAInfo(); - } + localMemInfoByRegion_.clear(); + registeredRegions_.clear(); return TILEXR_SUCCESS; } @@ -1802,19 +1904,21 @@ void TileXRUDMATransport::CleanupMemory() continue; } auto localRoutesIt = peerLocalEids_.find(peer); - for (size_t routeIdx = 0; routeIdx < peerEntry.second.size(); ++routeIdx) { - void* remoteHandle = peerEntry.second[routeIdx]; - if (remoteHandle == nullptr) { - continue; - } - uint32_t localEid = peerLocalEid_[peer]; - if (localRoutesIt != peerLocalEids_.end() && routeIdx < localRoutesIt->second.size()) { - localEid = localRoutesIt->second[routeIdx]; - } - if (ctxHandleByEid_.count(localEid) != 0) { - loader_.RaCtxRmemUnimport(ctxHandleByEid_[localEid], remoteHandle); + for (auto& regionHandles : peerEntry.second) { + for (size_t routeIdx = 0; routeIdx < regionHandles.size(); ++routeIdx) { + void* remoteHandle = regionHandles[routeIdx]; + if (remoteHandle == nullptr) { + continue; + } + uint32_t localEid = peerLocalEid_[peer]; + if (localRoutesIt != peerLocalEids_.end() && routeIdx < localRoutesIt->second.size()) { + localEid = localRoutesIt->second[routeIdx]; + } + if (ctxHandleByEid_.count(localEid) != 0) { + loader_.RaCtxRmemUnimport(ctxHandleByEid_[localEid], remoteHandle); + } + regionHandles[routeIdx] = nullptr; } - peerEntry.second[routeIdx] = nullptr; } } remoteMemHandlesByPeer_.clear(); @@ -1827,6 +1931,8 @@ void TileXRUDMATransport::CleanupMemory() } } registeredMem_.clear(); + localMemInfoByRegion_.clear(); + registeredRegions_.clear(); } void TileXRUDMATransport::CleanupQueues() @@ -1949,9 +2055,14 @@ void TileXRUDMATransport::Shutdown() peerQpRouteWeights_.clear(); qpsPerRoute_ = 1; qpNum_ = 1; + logicalQpsPerRoute_ = 1; + logicalQpNum_ = 1; sharedQpPool_ = false; - localMemInfoByEid_.clear(); + localMemInfoByRegion_.clear(); + registeredRegions_.clear(); remoteMemHandlesByPeer_.clear(); + logicalPeerQpRouteEids_.clear(); + logicalPeerQpRouteWeights_.clear(); loader_.Unload(); } diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index 3f24dc78..32227140 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -40,6 +40,7 @@ class TileXRUDMATransport { int Init(const TileXRUDMATransportOptions& options); int RegisterMemory(GM_ADDR localPtr, size_t bytes); + int RegisterMemoryRegions(const TileXRUDMARegionDesc* regions, uint32_t regionCount); int UnregisterMemory(GM_ADDR localPtr); void Shutdown(); @@ -60,7 +61,8 @@ class TileXRUDMATransport { int ImportSharedQueues(); int EnsureUDMAInfoBuffer(); int RefreshUDMAInfo(); - int RegisterMemoryOnContexts(GM_ADDR localPtr, size_t bytes); + int ConfigureRegionQueues(uint32_t regionCount); + int RegisterMemoryOnContexts(const TileXRUDMARegionDesc& region, uint32_t regionIndex); int ExchangeAndImportMemory(); int AllocDeviceScalar(void** ptr, size_t bytes) const; void FreeDeviceScalar(void*& ptr) const; @@ -80,6 +82,8 @@ class TileXRUDMATransport { uint32_t eidCount_ = 0; uint32_t qpNum_ = 1; uint32_t qpsPerRoute_ = 1; + uint32_t logicalQpNum_ = 1; + uint32_t logicalQpsPerRoute_ = 1; bool sharedQpPool_ = false; uint32_t sharedQpLaneCount_ = 16; std::map ctxHandleByEid_; @@ -90,11 +94,14 @@ class TileXRUDMATransport { std::map> peerRemoteEids_; std::map> peerQpRouteEids_; std::map> peerQpRouteWeights_; + std::map> logicalPeerQpRouteEids_; + std::map> logicalPeerQpRouteWeights_; std::map states_; std::map localEidByEid_; MemoryRegionMap registeredMem_; - std::map> remoteMemHandlesByPeer_; - std::map localMemInfoByEid_; + std::map>> remoteMemHandlesByPeer_; + std::vector> localMemInfoByRegion_; + std::vector registeredRegions_; GM_ADDR udmaInfoDev_ = nullptr; GM_ADDR eidTableDev_ = nullptr; uint32_t udmaInfoSize_ = 0; diff --git a/src/include/tilexr_api.h b/src/include/tilexr_api.h index 20b75a9f..fcdb61ef 100644 --- a/src/include/tilexr_api.h +++ b/src/include/tilexr_api.h @@ -13,6 +13,7 @@ #include #include #include "comm_args.h" +#include "tilexr_udma_reg.h" #ifdef __cplusplus extern "C" { #endif // __cplusplus @@ -40,6 +41,12 @@ int TileXRCommNextMagic(TileXRCommPtr comm, int64_t *magic); int TileXRUDMARegister(TileXRCommPtr comm, GM_ADDR localPtr, size_t bytes, TileXRUDMAMemHandle *handle); +// Registers up to four independent MRs as one logical, offset-addressed space. +// Adjacent VMM mappings remain contiguous to kernels; each MR gets a dedicated +// physical QP lane because this UDMA data plane cannot switch MRs on one QP. +int TileXRUDMARegisterRegions(TileXRCommPtr comm, const TileXR::TileXRUDMARegionDesc *regions, + uint32_t regionCount, TileXRUDMAMemHandle *handle); + int TileXRUDMAUnregister(TileXRCommPtr comm, TileXRUDMAMemHandle handle); int TileXRGetUDMARegistryDev(TileXRCommPtr comm, GM_ADDR ®istryPtr); diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index 94584c08..dfee0472 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -30,11 +30,29 @@ constexpr bool TILEXR_UDMA_ARCH_SUPPORTED = true; constexpr bool TILEXR_UDMA_ARCH_SUPPORTED = false; #endif +constexpr uint32_t TILEXR_UDMA_MAX_WQE_BYTES = 64U * 1024U * 1024U; + struct UDMASignalParams { __gm__ uint64_t* sigAddr; uint64_t signal; + uint32_t tid; + uint32_t tokenValue; +}; + +struct UDMARegionLocation { + uint32_t regionIndex; + uint64_t regionOffset; + uint64_t bytesAvailable; + __gm__ uint8_t* addr; }; +__aicore__ inline uint32_t UDMAChunkBytes(uint64_t remaining, uint64_t bytesAvailable) +{ + uint64_t chunk = remaining < bytesAvailable ? remaining : bytesAvailable; + if (chunk > TILEXR_UDMA_MAX_WQE_BYTES) chunk = TILEXR_UDMA_MAX_WQE_BYTES; + return static_cast(chunk); +} + __aicore__ inline bool UDMAEnabled(const __gm__ CommArgs* args) { return args != nullptr && ((args->extraFlag & ExtraFlag::UDMA) != 0) && args->udmaInfoPtr != nullptr; @@ -64,17 +82,41 @@ __aicore__ inline bool UDMARegisteredRangeValid( static_cast(targetRank) >= registry->rankSize) { return false; } - const auto& region = registry->regions[targetRank]; - if (region.base == nullptr || byteOffset > region.bytes) { - return false; + uint64_t totalBytes = 0; + for (uint32_t regionIndex = 0; regionIndex < registry->regionCount; ++regionIndex) { + const auto& region = registry->regions[targetRank][regionIndex]; + if (region.base == nullptr || region.bytes == 0 || totalBytes > UINT64_MAX - region.bytes) { + return false; + } + totalBytes += region.bytes; } - return byteCount <= region.bytes - byteOffset; + return byteOffset <= totalBytes && byteCount <= totalBytes - byteOffset; +} + +__aicore__ inline bool UDMAResolveRegisteredOffset( + const __gm__ TileXRUDMARegistry* registry, int targetRank, uint64_t byteOffset, + UDMARegionLocation& location) +{ + uint64_t cursor = 0; + for (uint32_t regionIndex = 0; regionIndex < registry->regionCount; ++regionIndex) { + const auto& region = registry->regions[targetRank][regionIndex]; + if (byteOffset >= cursor && byteOffset - cursor < region.bytes) { + location.regionIndex = regionIndex; + location.regionOffset = byteOffset - cursor; + location.bytesAvailable = region.bytes - location.regionOffset; + location.addr = reinterpret_cast<__gm__ uint8_t*>(region.base + location.regionOffset); + return true; + } + cursor += region.bytes; + } + return false; } __aicore__ inline __gm__ uint8_t* UDMARegisteredRemoteAddr( const __gm__ TileXRUDMARegistry* registry, int targetRank, uint64_t byteOffset) { - return reinterpret_cast<__gm__ uint8_t*>(registry->regions[targetRank].base + byteOffset); + UDMARegionLocation location {}; + return UDMAResolveRegisteredOffset(registry, targetRank, byteOffset, location) ? location.addr : nullptr; } __aicore__ inline void UDMACleanCacheLines(__gm__ uint8_t* addr, uint64_t length) @@ -109,10 +151,13 @@ __aicore__ inline __gm__ UDMACQCtx* UDMAGetSCQCtx(__gm__ UDMAInfo* udmaInfo, uin return reinterpret_cast<__gm__ UDMACQCtx*>(udmaInfo->scqPtr + (pe * qpNum + qpIdx) * sizeof(UDMACQCtx)); } -__aicore__ inline __gm__ UDMAMemInfo* UDMAGetRemoteMemInfo(__gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx) +__aicore__ inline __gm__ UDMAMemInfo* UDMAGetRemoteMemInfo( + __gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx, uint32_t regionIndex = 0) { uint32_t qpNum = udmaInfo->qpNum; - return reinterpret_cast<__gm__ UDMAMemInfo*>(udmaInfo->memPtr + sizeof(UDMAMemInfo) * (pe * qpNum + qpIdx)); + uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; + return reinterpret_cast<__gm__ UDMAMemInfo*>(udmaInfo->memPtr + sizeof(UDMAMemInfo) * + ((pe * qpNum + qpIdx) * regionCount + regionIndex)); } __aicore__ inline uint32_t UDMAGetQpWeight(__gm__ UDMAInfo* udmaInfo, uint32_t pe, uint32_t qpIdx) @@ -122,10 +167,25 @@ __aicore__ inline uint32_t UDMAGetQpWeight(__gm__ UDMAInfo* udmaInfo, uint32_t p return 1; } auto weights = reinterpret_cast<__gm__ uint32_t*>(udmaInfo->qpWeightPtr); - uint32_t weight = weights[pe * qpNum + qpIdx]; + uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; + uint32_t weight = weights[pe * qpNum + qpIdx * regionCount]; return weight == 0 ? 1 : weight; } +__aicore__ inline uint32_t UDMAGetLogicalQpNum(const __gm__ UDMAInfo* udmaInfo) +{ + uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; + uint32_t logicalQpNum = udmaInfo->qpNum / regionCount; + return logicalQpNum == 0 ? 1 : logicalQpNum; +} + +__aicore__ inline uint32_t UDMAGetRegionQpIndex( + const __gm__ UDMAInfo* udmaInfo, uint32_t logicalQpIdx, uint32_t regionIndex) +{ + uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; + return logicalQpIdx * regionCount + regionIndex; +} + __aicore__ inline void UDMAPollCQUpdateInfo( uint32_t curTail, __gm__ UDMACQCtx* cqCtxEntry, __gm__ UDMAWQCtx* wqCtxEntry) { @@ -226,7 +286,10 @@ __aicore__ inline void UDMAFillSqeCtx( __gm__ uint64_t* rmtEid = reinterpret_cast<__gm__ uint64_t*>(remoteMemInfo->eidAddr); sqeCtx->rmtEidL = rmtEid[0]; sqeCtx->rmtEidH = rmtEid[1]; - UDMAFillNotifyData(sqeCtx, remoteMemInfo->tid, remoteMemInfo->rmtTokenValue, signalParams); + UDMAFillNotifyData(sqeCtx, + signalParams == nullptr ? remoteMemInfo->tid : signalParams->tid, + signalParams == nullptr ? remoteMemInfo->rmtTokenValue : signalParams->tokenValue, + signalParams); } __aicore__ inline void UDMAFillSgeCtx( @@ -258,16 +321,19 @@ __aicore__ inline void UDMAPostSendUpdateInfo(uint32_t curHead, __gm__ UDMAWQCtx __aicore__ inline void UDMAPostSend( __gm__ UDMAInfo* udmaInfo, __gm__ uint8_t* remoteAddr, __gm__ uint8_t* localAddr, - uint32_t pe, uint32_t qpIdx, uint64_t messageLen, UDMAOpcode opcode, const UDMASignalParams* signalParams) + uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen, + UDMAOpcode opcode, const UDMASignalParams* signalParams) { - __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, pe, qpIdx); + uint32_t physicalQpIdx = UDMAGetRegionQpIndex(udmaInfo, qpIdx, regionIndex); + if (physicalQpIdx >= udmaInfo->qpNum) return; + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, pe, physicalQpIdx); uint32_t wqeSize = 1U << qpCtxEntry->baseBkShift; uint32_t depth = qpCtxEntry->depth; uint32_t curHead = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->headAddr), 0); uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); - UDMAPollCQWhenSQOverflow(udmaInfo, qpCtxEntry, wqeCnt, pe, qpIdx); + UDMAPollCQWhenSQOverflow(udmaInfo, qpCtxEntry, wqeCnt, pe, physicalQpIdx); - __gm__ UDMAMemInfo* remoteMemInfo = UDMAGetRemoteMemInfo(udmaInfo, pe, qpIdx); + __gm__ UDMAMemInfo* remoteMemInfo = UDMAGetRemoteMemInfo(udmaInfo, pe, physicalQpIdx, regionIndex); __gm__ uint8_t* wqeAddr = reinterpret_cast<__gm__ uint8_t*>(qpCtxEntry->bufAddr + wqeSize * (curHead % depth)); __gm__ UDMASqeCtx* sqeCtx = reinterpret_cast<__gm__ UDMASqeCtx*>(wqeAddr); @@ -285,32 +351,45 @@ __aicore__ inline void UDMAPostSend( __aicore__ inline void UDMAWrite( const __gm__ CommArgs* args, __gm__ uint8_t* remoteAddr, __gm__ uint8_t* localAddr, - uint32_t pe, uint32_t qpIdx, uint64_t messageLen) + uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen) { if constexpr (TILEXR_UDMA_ARCH_SUPPORTED) { - UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, messageLen, UDMAOpcode::WRITE, nullptr); + UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, regionIndex, + messageLen, UDMAOpcode::WRITE, nullptr); } } __aicore__ inline void UDMARead( const __gm__ CommArgs* args, __gm__ uint8_t* localAddr, __gm__ uint8_t* remoteAddr, - uint32_t pe, uint32_t qpIdx, uint64_t messageLen) + uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen) { if constexpr (TILEXR_UDMA_ARCH_SUPPORTED) { - UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, messageLen, UDMAOpcode::READ, nullptr); + UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, regionIndex, + messageLen, UDMAOpcode::READ, nullptr); } } __aicore__ inline void UDMAWriteNotify( const __gm__ CommArgs* args, __gm__ uint8_t* remoteAddr, __gm__ uint8_t* localAddr, - uint32_t pe, uint32_t qpIdx, uint64_t messageLen, const UDMASignalParams* signalParams) + uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen, + const UDMASignalParams* signalParams) { if constexpr (TILEXR_UDMA_ARCH_SUPPORTED) { - UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, messageLen, + UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, regionIndex, messageLen, UDMAOpcode::WRITE_WITH_NOTIFY, signalParams); } } +__aicore__ inline void UDMAQuietInternalOnQp( + const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx, uint32_t regionIndex) +{ + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + uint32_t physicalQpIdx = UDMAGetRegionQpIndex(udmaInfo, qpIdx, regionIndex); + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, physicalQpIdx); + uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + (void)UDMAPollCQ(udmaInfo, targetRank, physicalQpIdx, wqeCnt); +} + template __aicore__ inline void UDMAPutNbiOnQp( const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx, @@ -321,9 +400,21 @@ __aicore__ inline void UDMAPutNbiOnQp( auto registry = GetUDMARegistry(args); if (!UDMARegisteredRangeValid(registry, targetRank, byteOffset, byteCount)) return; - auto remoteAddr = UDMARegisteredRemoteAddr(registry, targetRank, byteOffset); - UDMAWrite(args, remoteAddr, reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)), - targetRank, qpIdx, byteCount); + auto localAddr = reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)); + uint64_t currentOffset = byteOffset; + uint64_t remaining = byteCount; + while (remaining != 0) { + UDMARegionLocation location {}; + if (!UDMAResolveRegisteredOffset(registry, targetRank, currentOffset, location)) return; + uint32_t chunk = UDMAChunkBytes(remaining, location.bytesAvailable); + UDMAWrite(args, location.addr, localAddr, targetRank, qpIdx, location.regionIndex, chunk); + currentOffset += chunk; + localAddr += chunk; + remaining -= chunk; + if (remaining != 0 && chunk == location.bytesAvailable) { + UDMAQuietInternalOnQp(args, targetRank, qpIdx, location.regionIndex); + } + } } template @@ -350,8 +441,21 @@ __aicore__ inline void UDMAGetNbiOnQp( auto registry = GetUDMARegistry(args); if (!UDMARegisteredRangeValid(registry, sourceRank, byteOffset, byteCount)) return; - auto remoteAddr = UDMARegisteredRemoteAddr(registry, sourceRank, byteOffset); - UDMARead(args, reinterpret_cast<__gm__ uint8_t*>(localDst), remoteAddr, sourceRank, qpIdx, byteCount); + auto localAddr = reinterpret_cast<__gm__ uint8_t*>(localDst); + uint64_t currentOffset = byteOffset; + uint64_t remaining = byteCount; + while (remaining != 0) { + UDMARegionLocation location {}; + if (!UDMAResolveRegisteredOffset(registry, sourceRank, currentOffset, location)) return; + uint32_t chunk = UDMAChunkBytes(remaining, location.bytesAvailable); + UDMARead(args, localAddr, location.addr, sourceRank, qpIdx, location.regionIndex, chunk); + currentOffset += chunk; + localAddr += chunk; + remaining -= chunk; + if (remaining != 0 && chunk == location.bytesAvailable) { + UDMAQuietInternalOnQp(args, sourceRank, qpIdx, location.regionIndex); + } + } } template @@ -382,13 +486,38 @@ __aicore__ inline void UDMAPutSignalNbiOnQp( return; } + UDMARegionLocation signalLocation {}; + if (!UDMAResolveRegisteredOffset(registry, targetRank, signalByteOffset, signalLocation)) return; + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + uint32_t signalPhysicalQp = UDMAGetRegionQpIndex(udmaInfo, qpIdx, signalLocation.regionIndex); + __gm__ UDMAMemInfo* signalMemInfo = + UDMAGetRemoteMemInfo(udmaInfo, targetRank, signalPhysicalQp, signalLocation.regionIndex); UDMASignalParams signalParams = {}; - signalParams.sigAddr = reinterpret_cast<__gm__ uint64_t*>( - UDMARegisteredRemoteAddr(registry, targetRank, signalByteOffset)); + signalParams.sigAddr = reinterpret_cast<__gm__ uint64_t*>(signalLocation.addr); signalParams.signal = signal; - auto remoteAddr = UDMARegisteredRemoteAddr(registry, targetRank, byteOffset); - UDMAWriteNotify(args, remoteAddr, reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)), - targetRank, qpIdx, byteCount, &signalParams); + signalParams.tid = signalMemInfo->tid; + signalParams.tokenValue = signalMemInfo->rmtTokenValue; + + auto localAddr = reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)); + uint64_t currentOffset = byteOffset; + uint64_t remaining = byteCount; + while (remaining != 0) { + UDMARegionLocation location {}; + if (!UDMAResolveRegisteredOffset(registry, targetRank, currentOffset, location)) return; + uint32_t chunk = UDMAChunkBytes(remaining, location.bytesAvailable); + if (remaining == chunk) { + UDMAWriteNotify(args, location.addr, localAddr, targetRank, qpIdx, + location.regionIndex, chunk, &signalParams); + } else { + UDMAWrite(args, location.addr, localAddr, targetRank, qpIdx, location.regionIndex, chunk); + } + currentOffset += chunk; + localAddr += chunk; + remaining -= chunk; + if (remaining != 0 && chunk == location.bytesAvailable) { + UDMAQuietInternalOnQp(args, targetRank, qpIdx, location.regionIndex); + } + } } template @@ -407,40 +536,45 @@ __aicore__ inline void UDMAPutRegisteredSignalNbi( UDMAPutSignalNbi(args, targetRank, localSrc, byteOffset, byteCount, signalByteOffset, signal); } +__aicore__ inline uint32_t UDMAPollLogicalQpRegions( + const __gm__ CommArgs* args, int targetRank, uint32_t logicalQpIdx) +{ + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; + uint32_t status = 0; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + uint32_t physicalQpIdx = UDMAGetRegionQpIndex(udmaInfo, logicalQpIdx, regionIndex); + if (physicalQpIdx >= udmaInfo->qpNum) return 0xFFFFFFFFU; + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, physicalQpIdx); + uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + uint32_t current = UDMAPollCQ(udmaInfo, targetRank, physicalQpIdx, wqeCnt); + if (status == 0 && current != 0) status = current; + } + return status; +} + __aicore__ inline void UDMAQuiet(const __gm__ CommArgs* args, int targetRank) { if (!UDMAEnabled(args)) return; - __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); - __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, 0); - uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); - (void)UDMAPollCQ(udmaInfo, targetRank, 0, wqeCnt); + (void)UDMAPollLogicalQpRegions(args, targetRank, 0); } __aicore__ inline void UDMAQuietOnQp(const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx) { if (!UDMAEnabled(args)) return; - __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); - __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, qpIdx); - uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); - (void)UDMAPollCQ(udmaInfo, targetRank, qpIdx, wqeCnt); + (void)UDMAPollLogicalQpRegions(args, targetRank, qpIdx); } __aicore__ inline uint32_t UDMAQuietStatus(const __gm__ CommArgs* args, int targetRank) { if (!UDMAEnabled(args)) return 0xFFFFFFFFU; - __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); - __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, 0); - uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); - return UDMAPollCQ(udmaInfo, targetRank, 0, wqeCnt); + return UDMAPollLogicalQpRegions(args, targetRank, 0); } __aicore__ inline uint32_t UDMAQuietStatusOnQp(const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx) { if (!UDMAEnabled(args)) return 0xFFFFFFFFU; - __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); - __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, qpIdx); - uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); - return UDMAPollCQ(udmaInfo, targetRank, qpIdx, wqeCnt); + return UDMAPollLogicalQpRegions(args, targetRank, qpIdx); } } // namespace TileXR diff --git a/src/include/tilexr_udma_reg.h b/src/include/tilexr_udma_reg.h index ccdca377..6e47363c 100644 --- a/src/include/tilexr_udma_reg.h +++ b/src/include/tilexr_udma_reg.h @@ -14,20 +14,24 @@ 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_REGISTRY_VERSION = 2U; +constexpr uint32_t TILEXR_UDMA_MAX_REGIONS = 4U; struct TileXRUDMARegionDesc { GM_ADDR base = nullptr; uint64_t bytes = 0; }; +// Region descriptors are concatenated in array order for device-side offsets. +// A 4 GiB ping/pong workspace can use four 1 GiB MRs while exposing ping at +// offset 0 and pong at offset 2 GiB in one continuous VMM reservation. + struct TileXRUDMARegistry { uint32_t magic = TILEXR_UDMA_REGISTRY_MAGIC; uint32_t version = TILEXR_UDMA_REGISTRY_VERSION; uint32_t rankSize = 0; uint32_t regionCount = 0; - TileXRUDMARegionDesc regions[TILEXR_MAX_RANK_SIZE] = {}; + TileXRUDMARegionDesc regions[TILEXR_MAX_RANK_SIZE][TILEXR_UDMA_MAX_REGIONS] = {}; }; inline bool UDMARegistryValid(const TileXRUDMARegistry *registry, int expectedRankSize) @@ -50,11 +54,15 @@ inline bool UDMARegionContains(const TileXRUDMARegistry *registry, int rank, uin if (rank < 0 || static_cast(rank) >= registry->rankSize) { return false; } - const auto ®ion = registry->regions[rank]; - if (region.base == nullptr || byteOffset > region.bytes) { - return false; + uint64_t totalBytes = 0; + for (uint32_t regionIndex = 0; regionIndex < registry->regionCount; ++regionIndex) { + const auto ®ion = registry->regions[rank][regionIndex]; + if (region.base == nullptr || region.bytes == 0 || totalBytes > UINT64_MAX - region.bytes) { + return false; + } + totalBytes += region.bytes; } - return byteCount <= region.bytes - byteOffset; + return byteOffset <= totalBytes && byteCount <= totalBytes - byteOffset; } inline GM_ADDR UDMARemoteAddr(const TileXRUDMARegistry *registry, int rank, uint64_t byteOffset) @@ -62,7 +70,18 @@ inline GM_ADDR UDMARemoteAddr(const TileXRUDMARegistry *registry, int rank, uint if (registry == nullptr || rank < 0 || static_cast(rank) >= registry->rankSize) { return nullptr; } - return registry->regions[rank].base + byteOffset; + uint64_t cursor = 0; + for (uint32_t regionIndex = 0; regionIndex < registry->regionCount; ++regionIndex) { + const auto ®ion = registry->regions[rank][regionIndex]; + if (region.base == nullptr || region.bytes == 0 || cursor > UINT64_MAX - region.bytes) { + return nullptr; + } + if (byteOffset >= cursor && byteOffset - cursor < region.bytes) { + return region.base + (byteOffset - cursor); + } + cursor += region.bytes; + } + return nullptr; } } // namespace TileXR diff --git a/src/include/tilexr_udma_types.h b/src/include/tilexr_udma_types.h index 9b344854..7a25a61e 100644 --- a/src/include/tilexr_udma_types.h +++ b/src/include/tilexr_udma_types.h @@ -70,6 +70,7 @@ struct UDMACQCtx { struct UDMAInfo { uint32_t qpNum; + uint32_t regionCount; uint64_t sqPtr; uint64_t rqPtr; uint64_t scqPtr; diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index a3b34915..5f307e57 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -273,6 +273,7 @@ if(BUILD_TILEXR_UDMA_DEMO) -o "${TILEXR_UDMA_DEMO_KERNEL_SO}" DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_demo_kernel.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_constants.h" "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_fullmesh_trace.h" "${TILEXR_ROOT}/src/include/tilexr_udma.h" "${TILEXR_ROOT}/src/include/tilexr_data_as_flag.h" diff --git a/tests/udma/demo/run_tilexr_udma_demo.sh b/tests/udma/demo/run_tilexr_udma_demo.sh index c6781137..d526e5ad 100755 --- a/tests/udma/demo/run_tilexr_udma_demo.sh +++ b/tests/udma/demo/run_tilexr_udma_demo.sh @@ -36,7 +36,7 @@ echo "==========================================" echo " TileXR UDMA Communication Demo" echo "==========================================" echo "Binary: ${bin}" -echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all, 3=all-reduce, 7=all-to-all bigdata UDMA)" +echo "Test type: ${test_type} (0=all-gather put, 1=put-signal, 2=all-to-all, 3=all-reduce, 7=all-to-all bigdata UDMA, 9=4GiB VMM ping+pong via 4x1GiB MR)" echo "Rank size: ${rank_size}" echo "Elements/rank: ${elements_per_rank}" echo "NPU count: ${npu_count}" diff --git a/tests/udma/demo/tilexr_udma_alltoall_constants.h b/tests/udma/demo/tilexr_udma_alltoall_constants.h new file mode 100644 index 00000000..ffbc56fe --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_constants.h @@ -0,0 +1,19 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_UDMA_ALLTOALL_CONSTANTS_H +#define TILEXR_UDMA_ALLTOALL_CONSTANTS_H + +#include + +namespace TileXR { +namespace Demo { + +constexpr uint64_t kAllToAllBigDataMultiNodePeerSlotBytes = 64ULL * 1024ULL * 1024ULL; + +} // namespace Demo +} // namespace TileXR + +#endif // TILEXR_UDMA_ALLTOALL_CONSTANTS_H diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 600fcbd3..ac394520 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -171,7 +171,7 @@ __aicore__ inline void AllToAllGroupSelectRouteQps( uint32_t& primaryWeight, uint32_t& secondaryWeight) { auto udmaInfo = TileXR::GetUDMAInfo(args); - const uint32_t qpCount = udmaInfo->qpNum == 0U ? 1U : udmaInfo->qpNum; + const uint32_t qpCount = TileXR::UDMAGetLogicalQpNum(udmaInfo); primaryQp = 0U; primaryWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); for (uint32_t qp = 1U; qp < qpCount; ++qp) { diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 2867b465..a7f7f9bb 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -30,7 +30,7 @@ constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; -constexpr size_t kAllToAllGroupMaxRegisteredBytes = 1ULL << 30; +constexpr size_t kAllToAllGroupMaxRegisteredBytes = 4ULL << 30; struct AllToAllGroupPlan { bool valid = false; diff --git a/tests/udma/demo/tilexr_udma_alltoall_layout.h b/tests/udma/demo/tilexr_udma_alltoall_layout.h index 60a698d4..644efff1 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_layout.h @@ -12,6 +12,8 @@ #include #include +#include "tilexr_udma_alltoall_constants.h" + namespace TileXR { namespace Demo { @@ -26,8 +28,7 @@ struct AllToAllChunkPlan { }; constexpr size_t kAllToAllBigDataMaxRegisteredBytes = 128ULL * 1024ULL * 1024ULL; -constexpr size_t kAllToAllBigDataMultiNodeRegisteredBytes = 1024ULL * 1024ULL * 1024ULL; -constexpr size_t kAllToAllBigDataMultiNodePeerSlotBytes = 16ULL * 1024ULL * 1024ULL; +constexpr size_t kAllToAllBigDataMultiNodeRegisteredBytes = 4ULL * 1024ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllBigDataControlSlotBytes = 128ULL; constexpr uint32_t kAllToAllBigDataCoresPerPeer = 5U; constexpr uint32_t kAllToAllBigDataSingleNodeShards = 2U; diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 1d5a5017..396f211f 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -47,6 +47,9 @@ extern void launch_tilexr_udma_p2p_latency( extern void launch_tilexr_datacopy_latency( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR debug, int32_t elementsPerPeer, int32_t chunkElements); +extern void launch_tilexr_udma_vmm_regions_probe( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR localBase, + uint64_t regionBytes, uint32_t regionCount); extern void launch_tilexr_udma_all_to_all_fused( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR signal, GM_ADDR debug, int32_t elementsPerPeer, @@ -109,6 +112,11 @@ constexpr int kDemoBarrierPortOffset = 97; constexpr int kConnectRetryCount = 500; constexpr int kConnectRetrySleepMs = 10; constexpr int kBigDataProfileStageFull = 8; +constexpr uint32_t kVmmProbeRegionCount = 4; +constexpr uint64_t kVmmProbeRegionBytes = 1ULL << 30; +constexpr uint64_t kVmmProbeMarkerDstOffset = 64; +constexpr uint64_t kVmmProbeMarkerSrcOffset = 128; +constexpr uint64_t kVmmProbeBoundaryBytes = 64; struct BarrierEndpoint { std::string host; @@ -443,6 +451,230 @@ bool DemoBarrierAll(int rank, int rankSize, const std::string& step) return true; } +struct VmmMultiRegionAllocation { + void* base = nullptr; + size_t regionBytes = 0; + uint32_t regionCount = 0; + uint32_t mappedCount = 0; + std::array physical {}; + std::array regions {}; +}; + +VmmMultiRegionAllocation gRegisteredVmm; + +void ReleaseVmmMultiRegion(int rank, VmmMultiRegionAllocation& allocation) +{ + for (uint32_t i = allocation.mappedCount; i > 0; --i) { + void* regionBase = static_cast(allocation.base) + + static_cast(i - 1) * allocation.regionBytes; + CheckAcl(rank, "aclrtUnmapMem region " + std::to_string(i - 1), aclrtUnmapMem(regionBase)); + } + for (uint32_t i = allocation.regionCount; i > 0; --i) { + if (allocation.physical[i - 1] != nullptr) { + CheckAcl(rank, "aclrtFreePhysical region " + std::to_string(i - 1), + aclrtFreePhysical(allocation.physical[i - 1])); + } + } + if (allocation.base != nullptr) { + CheckAcl(rank, "aclrtReleaseMemAddress", aclrtReleaseMemAddress(allocation.base)); + } + allocation = VmmMultiRegionAllocation {}; +} + +bool AllocateVmmMultiRegion( + int rank, int deviceId, size_t regionBytes, uint32_t regionCount, + VmmMultiRegionAllocation& allocation) +{ + if (regionCount == 0 || regionCount > allocation.regions.size()) { + return false; + } + aclrtPhysicalMemProp prop {}; + prop.handleType = ACL_MEM_HANDLE_TYPE_NONE; + prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED; + prop.memAttr = ACL_HBM_MEM_HUGE; + prop.location.id = deviceId; + prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE; + + size_t granularity = 0; + if (!CheckAcl(rank, "aclrtMemGetAllocationGranularity", + aclrtMemGetAllocationGranularity( + &prop, ACL_RT_MEM_ALLOC_GRANULARITY_RECOMMENDED, &granularity)) || + granularity == 0 || regionBytes % granularity != 0) { + std::cerr << "[rank " << rank << "] ERROR: invalid VMM granularity=" << granularity + << " for regionBytes=" << regionBytes << std::endl; + return false; + } + + allocation.regionBytes = regionBytes; + allocation.regionCount = regionCount; + const size_t totalBytes = regionBytes * regionCount; + if (!CheckAcl(rank, "aclrtReserveMemAddress multi-region", + aclrtReserveMemAddress(&allocation.base, totalBytes, 0, nullptr, 1))) { + allocation = VmmMultiRegionAllocation {}; + return false; + } + for (uint32_t i = 0; i < regionCount; ++i) { + void* regionBase = static_cast(allocation.base) + static_cast(i) * regionBytes; + if (!CheckAcl(rank, "aclrtMallocPhysical region " + std::to_string(i), + aclrtMallocPhysical(&allocation.physical[i], regionBytes, &prop, 0)) || + !CheckAcl(rank, "aclrtMapMem region " + std::to_string(i), + aclrtMapMem(regionBase, regionBytes, 0, allocation.physical[i], 0))) { + ReleaseVmmMultiRegion(rank, allocation); + return false; + } + ++allocation.mappedCount; + allocation.regions[i].base = static_cast(regionBase); + allocation.regions[i].bytes = regionBytes; + } + return true; +} + +bool RunVmmMultiRegionProbe( + int rank, int rankSize, int deviceId, TileXRCommPtr comm, aclrtStream stream, GM_ADDR commArgsDev) +{ + const size_t regionBytes = static_cast(kVmmProbeRegionBytes); + const size_t totalBytes = regionBytes * kVmmProbeRegionCount; + aclrtPhysicalMemProp prop {}; + prop.handleType = ACL_MEM_HANDLE_TYPE_NONE; + prop.allocationType = ACL_MEM_ALLOCATION_TYPE_PINNED; + prop.memAttr = ACL_HBM_MEM_HUGE; + prop.location.id = deviceId; + prop.location.type = ACL_MEM_LOCATION_TYPE_DEVICE; + + size_t granularity = 0; + if (!CheckAcl(rank, "aclrtMemGetAllocationGranularity", + aclrtMemGetAllocationGranularity( + &prop, ACL_RT_MEM_ALLOC_GRANULARITY_RECOMMENDED, &granularity)) || + granularity == 0 || regionBytes % granularity != 0) { + std::cerr << "[rank " << rank << "] ERROR: invalid VMM granularity=" << granularity + << " for regionBytes=" << regionBytes << std::endl; + return false; + } + + void* base = nullptr; + std::array physical {}; + uint32_t mappedCount = 0; + auto releaseVmm = [&]() { + for (uint32_t i = mappedCount; i > 0; --i) { + void* regionBase = static_cast(base) + static_cast(i - 1) * regionBytes; + CheckAcl(rank, "aclrtUnmapMem region " + std::to_string(i - 1), aclrtUnmapMem(regionBase)); + } + for (uint32_t i = kVmmProbeRegionCount; i > 0; --i) { + if (physical[i - 1] != nullptr) { + CheckAcl(rank, "aclrtFreePhysical region " + std::to_string(i - 1), + aclrtFreePhysical(physical[i - 1])); + } + } + if (base != nullptr) { + CheckAcl(rank, "aclrtReleaseMemAddress", aclrtReleaseMemAddress(base)); + } + }; + + if (!CheckAcl(rank, "aclrtReserveMemAddress 4GiB", + aclrtReserveMemAddress(&base, totalBytes, 0, nullptr, 1))) { + return false; + } + std::array regions {}; + for (uint32_t i = 0; i < kVmmProbeRegionCount; ++i) { + void* regionBase = static_cast(base) + static_cast(i) * regionBytes; + if (!CheckAcl(rank, "aclrtMallocPhysical region " + std::to_string(i), + aclrtMallocPhysical(&physical[i], regionBytes, &prop, 0)) || + !CheckAcl(rank, "aclrtMapMem region " + std::to_string(i), + aclrtMapMem(regionBase, regionBytes, 0, physical[i], 0))) { + releaseVmm(); + return false; + } + ++mappedCount; + regions[i].base = static_cast(regionBase); + regions[i].bytes = regionBytes; + std::cout << "[rank " << rank << "] VMM region " << i + << " base=" << regionBase << " bytes=" << regionBytes << std::endl; + } + const uint64_t expectedEnd = reinterpret_cast(base) + totalBytes; + const uint64_t actualEnd = reinterpret_cast(regions.back().base) + regions.back().bytes; + if (actualEnd != expectedEnd) { + std::cerr << "[rank " << rank << "] ERROR: VMM VA range is not contiguous" << std::endl; + releaseVmm(); + return false; + } + + TileXRUDMAMemHandle handle = 0; + int ret = TileXRUDMARegisterRegions(comm, regions.data(), regions.size(), &handle); + if (!CheckTileXR(rank, "TileXRUDMARegisterRegions 4x1GiB", ret)) { + releaseVmm(); + return false; + } + + const int predecessor = (rank - 1 + rankSize) % rankSize; + bool ok = true; + for (uint32_t i = 0; i < kVmmProbeRegionCount; ++i) { + const uint64_t marker = (static_cast(rank) << 32) | (0xA5000000ULL + i); + const uint64_t zero = 0; + ok = CheckAcl(rank, "init marker source " + std::to_string(i), + aclrtMemcpy(regions[0].base + kVmmProbeMarkerSrcOffset + i * sizeof(marker), sizeof(marker), + &marker, sizeof(marker), ACL_MEMCPY_HOST_TO_DEVICE)) && ok; + ok = CheckAcl(rank, "clear marker destination " + std::to_string(i), + aclrtMemcpy(regions[i].base + kVmmProbeMarkerDstOffset, sizeof(zero), + &zero, sizeof(zero), ACL_MEMCPY_HOST_TO_DEVICE)) && ok; + } + std::array boundarySource {}; + for (uint32_t i = 0; i < boundarySource.size(); ++i) { + boundarySource[i] = static_cast((rank * 17 + i) & 0xFF); + } + GM_ADDR boundarySourceAddr = static_cast(base) + 4096; + GM_ADDR boundaryDestinationAddr = + static_cast(base) + regionBytes - kVmmProbeBoundaryBytes / 2; + ok = CheckAcl(rank, "init cross-region boundary source", + aclrtMemcpy(boundarySourceAddr, boundarySource.size(), boundarySource.data(), boundarySource.size(), + ACL_MEMCPY_HOST_TO_DEVICE)) && ok; + if (!ok || !DemoBarrierAll(rank, rankSize, "VMM regions initialized")) { + TileXRUDMAUnregister(comm, handle); + releaseVmm(); + return false; + } + + launch_tilexr_udma_vmm_regions_probe( + 1, stream, commArgsDev, static_cast(base), regionBytes, kVmmProbeRegionCount); + ok = CheckAcl(rank, "aclrtSynchronizeStream VMM region probe", aclrtSynchronizeStream(stream)); + ok = DemoBarrierAll(rank, rankSize, "VMM region UDMA writes complete") && ok; + + for (uint32_t i = 0; i < kVmmProbeRegionCount; ++i) { + uint64_t actual = 0; + const uint64_t expected = (static_cast(predecessor) << 32) | (0xA5000000ULL + i); + ok = CheckAcl(rank, "read marker destination " + std::to_string(i), + aclrtMemcpy(&actual, sizeof(actual), regions[i].base + kVmmProbeMarkerDstOffset, + sizeof(actual), ACL_MEMCPY_DEVICE_TO_HOST)) && ok; + if (actual != expected) { + std::cerr << "[rank " << rank << "] ERROR: region " << i + << " marker=" << actual << " expected=" << expected << std::endl; + ok = false; + } + } + std::array boundaryActual {}; + ok = CheckAcl(rank, "read cross-region boundary", + aclrtMemcpy(boundaryActual.data(), boundaryActual.size(), boundaryDestinationAddr, boundaryActual.size(), + ACL_MEMCPY_DEVICE_TO_HOST)) && ok; + for (uint32_t i = 0; i < boundaryActual.size(); ++i) { + const uint8_t expected = static_cast((predecessor * 17 + i) & 0xFF); + if (boundaryActual[i] != expected) { + std::cerr << "[rank " << rank << "] ERROR: boundary byte " << i + << "=" << static_cast(boundaryActual[i]) + << " expected=" << static_cast(expected) << std::endl; + ok = false; + break; + } + } + + ok = CheckTileXR(rank, "TileXRUDMAUnregister VMM regions", + TileXRUDMAUnregister(comm, handle)) && ok; + releaseVmm(); + if (ok) { + std::cout << "[rank " << rank + << "] VMM 4GiB / ping+pong 2GiB each / 4x1GiB MR probe success" << std::endl; + } + return ok; +} + bool ValidateData(int rank, int rankSize, const std::vector& data, int32_t elementsPerRank) { bool ok = true; @@ -640,7 +872,7 @@ bool CopyChunkDeviceToHost( bool RunGroupedAllToAll( int rank, int rankSize, int32_t elementsPerPeer, - TileXRCommPtr comm, aclrtStream stream, GM_ADDR commArgsDev) + int deviceId, TileXRCommPtr comm, aclrtStream stream, GM_ADDR commArgsDev) { constexpr uint32_t kErrorWordsPerCore = 12U; constexpr uint32_t kErrorCoreCount = TileXR::Demo::kAllToAllGroupBlockDim; @@ -771,6 +1003,7 @@ bool RunGroupedAllToAll( int32_t* input = nullptr; int32_t* output = nullptr; void* registeredMemory = nullptr; + VmmMultiRegionAllocation registeredVmm; std::array groupTraceDevices {}; std::array, kRouteStageCount> hostGroupTraces; aclrtEvent stageStartEvent = nullptr; @@ -783,7 +1016,10 @@ bool RunGroupedAllToAll( TileXRUDMAUnregister(comm, handle)); registered = false; } - if (registeredMemory != nullptr) { + if (registeredVmm.base != nullptr) { + ReleaseVmmMultiRegion(rank, registeredVmm); + registeredMemory = nullptr; + } else if (registeredMemory != nullptr) { aclrtFree(registeredMemory); registeredMemory = nullptr; } @@ -811,13 +1047,30 @@ bool RunGroupedAllToAll( } }; + constexpr size_t kGroupedRegionBytes = 1ULL << 30; + const uint32_t groupedRegionCount = static_cast( + (plan.registeredBytes + kGroupedRegionBytes - 1U) / kGroupedRegionBytes); + const bool useMultiRegion = plan.registeredBytes > kGroupedRegionBytes; if (!CheckAcl(rank, "aclrtMalloc grouped input", aclrtMalloc(reinterpret_cast(&input), dataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || !CheckAcl(rank, "aclrtMalloc grouped output", - aclrtMalloc(reinterpret_cast(&output), dataBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || - !CheckAcl(rank, "aclrtMalloc grouped registered memory", - aclrtMalloc(®isteredMemory, plan.registeredBytes, ACL_MEM_MALLOC_HUGE_FIRST)) || - !CopyHostToDevice(rank, input, dataBytes, hostInput.data(), dataBytes, "grouped input") || + aclrtMalloc(reinterpret_cast(&output), dataBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + release(); + return false; + } + if (useMultiRegion) { + if (!AllocateVmmMultiRegion( + rank, deviceId, kGroupedRegionBytes, groupedRegionCount, registeredVmm)) { + release(); + return false; + } + registeredMemory = registeredVmm.base; + } else if (!CheckAcl(rank, "aclrtMalloc grouped registered memory", + aclrtMalloc(®isteredMemory, plan.registeredBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + release(); + return false; + } + if (!CopyHostToDevice(rank, input, dataBytes, hostInput.data(), dataBytes, "grouped input") || !CopyHostToDevice(rank, output, dataBytes, hostOutput.data(), dataBytes, "grouped output init") || !CheckAcl(rank, "aclrtMemset grouped registered memory", aclrtMemset(registeredMemory, plan.registeredBytes, 0, plan.registeredBytes))) { @@ -857,9 +1110,14 @@ bool RunGroupedAllToAll( } } - const int registerRet = TileXRUDMARegister( - comm, static_cast(registeredMemory), plan.registeredBytes, &handle); - if (!CheckTileXR(rank, "TileXRUDMARegister grouped alltoall", registerRet)) { + const int registerRet = useMultiRegion ? + TileXRUDMARegisterRegions( + comm, registeredVmm.regions.data(), registeredVmm.regionCount, &handle) : + TileXRUDMARegister( + comm, static_cast(registeredMemory), plan.registeredBytes, &handle); + if (!CheckTileXR(rank, useMultiRegion ? + "TileXRUDMARegisterRegions grouped alltoall" : + "TileXRUDMARegister grouped alltoall", registerRet)) { release(); return false; } @@ -884,6 +1142,7 @@ bool RunGroupedAllToAll( " signalOffset0=" + std::to_string(plan.signalOffset[0]) + " signalOffset1=" + std::to_string(plan.signalOffset[1]) + " controlOffset=" + std::to_string(plan.controlOffset) + + " regionCount=" + std::to_string(useMultiRegion ? groupedRegionCount : 1U) + " groupWidth=" + std::to_string(plan.groupWidth) + " groups=" + std::to_string(plan.groupCount) + " passes=" + std::to_string(plan.passCount)); @@ -1055,8 +1314,13 @@ void Cleanup( TileXRCommPtr comm, aclrtStream stream, void* registeredMemory, int32_t* debug, int rank, int deviceId) { if (registeredMemory != nullptr) { - PrintStatus(rank, "aclrtFree registered memory"); - aclrtFree(registeredMemory); + if (registeredMemory == gRegisteredVmm.base) { + PrintStatus(rank, "release VMM registered memory"); + ReleaseVmmMultiRegion(rank, gRegisteredVmm); + } else { + PrintStatus(rank, "aclrtFree registered memory"); + aclrtFree(registeredMemory); + } } if (debug != nullptr) { PrintStatus(rank, "aclrtFree debug"); @@ -1138,9 +1402,16 @@ int main(int argc, char** argv) return 1; } + if (testType == 9) { + const bool ok = RunVmmMultiRegionProbe( + rank, rankSize, deviceId, comm, stream, commArgsDev); + Cleanup(comm, stream, nullptr, nullptr, rank, deviceId); + return ok ? 0 : 1; + } + if (testType == 8) { const bool ok = RunGroupedAllToAll( - rank, rankSize, elementsPerRank, comm, stream, commArgsDev); + rank, rankSize, elementsPerRank, deviceId, comm, stream, commArgsDev); Cleanup(comm, stream, nullptr, nullptr, rank, deviceId); if (!ok) { std::cerr << "[rank " << rank << "] TileXR grouped alltoall demo failed" << std::endl; @@ -1292,13 +1563,22 @@ int main(int argc, char** argv) if (allocBytes < registeredBytes) { allocBytes = registeredBytes; } + const bool useBigDataMultiRegionVmm = testType == 7 && + TileXR::Demo::AllToAllBigDataIsMultiNode(rankSize, bigDataRanksPerNode); if (!CheckAcl(rank, "aclrtMalloc debug", aclrtMalloc(reinterpret_cast(&debug), kDebugWords * sizeof(int32_t), ACL_MEM_MALLOC_HUGE_FIRST)) || - !CheckAcl(rank, "aclrtMalloc registered memory", aclrtMalloc(®isteredMemory, - allocBytes, ACL_MEM_MALLOC_HUGE_FIRST))) { + (useBigDataMultiRegionVmm && + !AllocateVmmMultiRegion(rank, deviceId, static_cast(kVmmProbeRegionBytes), + kVmmProbeRegionCount, gRegisteredVmm)) || + (!useBigDataMultiRegionVmm && + !CheckAcl(rank, "aclrtMalloc registered memory", aclrtMalloc(®isteredMemory, + allocBytes, ACL_MEM_MALLOC_HUGE_FIRST)))) { Cleanup(comm, stream, registeredMemory, debug, rank, deviceId); return 1; } + if (useBigDataMultiRegionVmm) { + registeredMemory = gRegisteredVmm.base; + } auto data = static_cast(registeredMemory); auto input = reinterpret_cast(static_cast(registeredMemory) + inputOffset); auto output = reinterpret_cast(static_cast(registeredMemory) + outputOffset); @@ -1435,8 +1715,11 @@ int main(int argc, char** argv) return 1; } if (!udmaRegistered) { - int registerRet = TileXRUDMARegister(comm, static_cast(registeredMemory), - registeredBytes, &udmaHandle); + int registerRet = useBigDataMultiRegionVmm ? + TileXRUDMARegisterRegions(comm, gRegisteredVmm.regions.data(), + gRegisteredVmm.regionCount, &udmaHandle) : + TileXRUDMARegister(comm, static_cast(registeredMemory), + registeredBytes, &udmaHandle); if (registerRet != TileXR::TILEXR_SUCCESS) { std::cerr << "[rank " << rank << "] ERROR: bigdata alltoall UDMA registration failed" << " ret=" << registerRet << " regBytes=" << registeredBytes << std::endl; diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index 9db21051..b9a0ae77 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -6,6 +6,7 @@ #include "kernel_operator.h" #include "tilexr_data_as_flag.h" #include "tilexr_udma.h" +#include "tilexr_udma_alltoall_constants.h" #include "tilexr_udma_fullmesh_trace.h" constexpr int32_t TILEXR_UDMA_DEMO_MAGIC = 0x5444554d; // "TDUM" @@ -56,7 +57,8 @@ constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES = 64 * 1024; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_PINGPONG_BYTES = TILEXR_UDMA_DEMO_BIGDATA_RELAY_UB_BYTES * 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U; -constexpr uint64_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 16ULL * 1024ULL * 1024ULL; +constexpr uint64_t TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = + TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U; constexpr uint32_t TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS = TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS; @@ -660,7 +662,7 @@ __aicore__ inline uint32_t BigDataSelectWeightedQp( const __gm__ TileXR::CommArgs* args, int32_t peer, bool selectMax) { auto udmaInfo = TileXR::GetUDMAInfo(args); - const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; + const uint32_t qpCount = TileXR::UDMAGetLogicalQpNum(udmaInfo); uint32_t selected = 0U; uint32_t selectedWeight = TileXR::UDMAGetQpWeight(udmaInfo, peer, 0U); for (uint32_t qpIdx = 1U; qpIdx < qpCount; ++qpIdx) { @@ -677,7 +679,7 @@ __aicore__ inline uint32_t BigDataSelectDistinctWeightedQp( const __gm__ TileXR::CommArgs* args, int32_t peer, uint32_t avoidQp, bool selectMax) { auto udmaInfo = TileXR::GetUDMAInfo(args); - const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; + const uint32_t qpCount = TileXR::UDMAGetLogicalQpNum(udmaInfo); if (qpCount <= 1U) { return 0U; } @@ -1072,7 +1074,7 @@ __aicore__ inline void BigDataSendPeerWorker( debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = static_cast(wqCtx->localTokenId); debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = - static_cast(reinterpret_cast(registry->regions[peer].base) & + static_cast(reinterpret_cast(registry->regions[peer][0].base) & 0xFFFFFFFFU); debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); @@ -2181,7 +2183,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_kernel( debug[TILEXR_UDMA_DEMO_DEBUG_WQE_BEFORE_BASE + peer] = static_cast(wqeBefore); debug[TILEXR_UDMA_DEMO_DEBUG_LOCAL_TOKEN_BASE + peer] = static_cast(wqCtx->localTokenId); debug[TILEXR_UDMA_DEMO_DEBUG_REMOTE_BASE_LOW_BASE + peer] = - static_cast(reinterpret_cast(registry->regions[peer].base) & 0xFFFFFFFFU); + static_cast(reinterpret_cast(registry->regions[peer][0].base) & 0xFFFFFFFFU); debug[TILEXR_UDMA_DEMO_DEBUG_MEM_ADDR_LOW_BASE + peer] = static_cast(remoteMemInfo->addr & 0xFFFFFFFFU); debug[TILEXR_UDMA_DEMO_DEBUG_TPN_BASE + peer] = static_cast(remoteMemInfo->tpn); @@ -2557,6 +2559,31 @@ extern "C" __global__ __aicore__ void tilexr_udma_registered_smoke_kernel( } } +extern "C" __global__ __aicore__ void tilexr_udma_vmm_regions_probe_kernel( + GM_ADDR commArgsGM, GM_ADDR localBaseGM, uint64_t regionBytes, uint32_t regionCount) +{ + auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + auto localBase = reinterpret_cast<__gm__ uint8_t*>(localBaseGM); + if (!TileXR::UDMARegistryEnabled(args) || args->rankSize < 2 || regionCount < 4) { + return; + } + const int targetRank = (args->rank + 1) % args->rankSize; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + auto source = localBase + 128 + static_cast(regionIndex) * sizeof(uint64_t); + const uint64_t destinationOffset = static_cast(regionIndex) * regionBytes + 64; + TileXR::UDMAPutNbiOnQp( + args, targetRank, 0, source, destinationOffset, sizeof(uint64_t)); + TileXR::UDMAQuiet(args, targetRank); + } + + constexpr uint32_t boundaryBytes = 64; + auto boundarySource = localBase + 4096; + const uint64_t boundaryDestinationOffset = regionBytes - boundaryBytes / 2; + TileXR::UDMAPutNbiOnQp( + args, targetRank, 0, boundarySource, boundaryDestinationOffset, boundaryBytes); + TileXR::UDMAQuiet(args, targetRank); +} + // DMA-based scatter: write data to all peers' IPC staging area via DataCopyPad. // Split from the fused kernel to allow host-side sync between scatter and gather, // which guarantees P2P write visibility without fragile in-kernel flag polling. @@ -2776,6 +2803,14 @@ void launch_tilexr_udma_registered_smoke( commArgs, local, debug, bytes, signal); } +void launch_tilexr_udma_vmm_regions_probe( + uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR localBase, + uint64_t regionBytes, uint32_t regionCount) +{ + tilexr_udma_vmm_regions_probe_kernel<<>>( + commArgs, localBase, regionBytes, regionCount); +} + // --------------------------------------------------------------------------- // Latency micro-kernels (testType 4 = P2P-only, testType 5 = DataCopy-only). // These mirror the two halves of tilexr_udma_all_to_all_kernel so the P2P @@ -2819,7 +2854,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_p2p_latency_kernel( const uint64_t payloadBytes = AllToAllPayloadBytes(effectiveChunkElements); const uint32_t bytes = static_cast(payloadBytes); auto udmaInfo = TileXR::GetUDMAInfo(args); - const uint32_t qpCount = udmaInfo->qpNum == 0 ? 1U : udmaInfo->qpNum; + const uint32_t qpCount = TileXR::UDMAGetLogicalQpNum(udmaInfo); // One block per peer, skip self (peer == rank): no local copy here. for (int32_t peer = blockIdx; peer < rankSize; peer += blockNum) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index e42115ff..b1e9ae87 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -166,7 +166,22 @@ void TestPlan() constexpr int32_t thirtyTwoMiBElements = 8 * 1024 * 1024; CHECK_EQ(TileXR::Demo::PlanAllToAllGroup( - rankSize, thirtyTwoMiBElements, thirtyTwoMiBElements).valid, false); + rankSize, thirtyTwoMiBElements, thirtyTwoMiBElements).valid, true); + + constexpr int32_t oneGiBPerRankElementsPerPeer = 16 * 1024 * 1024; + const auto oneGiBPerRank = TileXR::Demo::PlanAllToAllGroup( + rankSize, oneGiBPerRankElementsPerPeer, oneGiBPerRankElementsPerPeer); + CHECK_EQ(oneGiBPerRank.valid, true); + CHECK_EQ(oneGiBPerRank.passCount, 1U); + CHECK_EQ(oneGiBPerRank.payloadPlaneBytes, 1ULL << 30); + CHECK_EQ(oneGiBPerRank.registeredBytes > (2ULL << 30), true); + CHECK_EQ(oneGiBPerRank.registeredBytes <= + TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); + + constexpr int32_t twoGiBPerRankElementsPerPeer = 32 * 1024 * 1024; + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup( + rankSize, twoGiBPerRankElementsPerPeer, + twoGiBPerRankElementsPerPeer).valid, false); } void TestChannelPolicy() diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp index 97785c13..00dde09b 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_layout.cpp @@ -179,8 +179,8 @@ void TestAllToAllBigDataPlan() const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); CHECK_EQ(TileXR::Demo::kAllToAllBigDataMaxRegisteredBytes, 128ULL * 1024ULL * 1024ULL); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, 1024ULL * 1024ULL * 1024ULL); - CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 16ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodeRegisteredBytes, 4ULL * 1024ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes, 64ULL * 1024ULL * 1024ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataControlSlotBytes, 128ULL); CHECK_EQ(TileXR::Demo::kAllToAllBigDataCoresPerPeer, 5U); CHECK_EQ(TileXR::Demo::kAllToAllBigDataSingleNodeShards, 2U); @@ -246,6 +246,20 @@ void TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone() static_cast(elementsPerPeer) * sizeof(int32_t)); } +void TestAllToAllBigData16PRankPayload1GiBPlan() +{ + constexpr int rankSize = 16; + constexpr int32_t elementsPerPeer = 16 * 1024 * 1024; + const auto plan = TileXR::Demo::PlanAllToAllBigDataUdma(rankSize, elementsPerPeer); + + CHECK_EQ(plan.registeredBytes, 4ULL * 1024ULL * 1024ULL * 1024ULL); + CHECK_EQ(plan.passCount, 1U); + CHECK_EQ(plan.chunkElements, elementsPerPeer); + CHECK_EQ(plan.chunkBytesPerPeer, 64ULL * 1024ULL * 1024ULL); + CHECK_EQ(plan.dataBytes > 0, true); + CHECK_EQ(plan.dataBytes + plan.controlBytes + plan.signalBytes <= plan.registeredBytes, true); +} + void TestAllToAllBigDataMultiNodeSmallPayloadUsesPayloadSlot() { constexpr int rankSize = 64; @@ -462,6 +476,9 @@ void TestAllToAllBigDataSource() ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo.cpp"); const std::string kernel = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_demo_kernel.cpp"); + const std::string constants = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + + "/tests/udma/demo/tilexr_udma_alltoall_constants.h"); const std::string udma = ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); @@ -523,7 +540,10 @@ void TestAllToAllBigDataSource() CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_SINGLE_NODE_SHARDS = 2U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_LOCAL_COPY_SHARDS"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_PINGPONG_SLOTS = 2U"); - CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_PEER_SLOT_BYTES = 16ULL * 1024ULL * 1024ULL"); + CHECK_CONTAINS(kernel, "#include \"tilexr_udma_alltoall_constants.h\""); + CHECK_CONTAINS(kernel, "TileXR::Demo::kAllToAllBigDataMultiNodePeerSlotBytes"); + CHECK_CONTAINS(constants, + "kAllToAllBigDataMultiNodePeerSlotBytes = 64ULL * 1024ULL * 1024ULL"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_RANKS_PER_NODE = 8"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_CONTROL_SHARDS = 32U"); CHECK_CONTAINS(kernel, "TILEXR_UDMA_DEMO_BIGDATA_MULTINODE_BLOCK_DIM ="); @@ -769,6 +789,7 @@ int main() TestAllToAllMaxRank256With64MiBPerRank(); TestAllToAllBigDataPlan(); TestAllToAllBigDataMultiNodePlanUses16ShardsAndRemoteSendDone(); + TestAllToAllBigData16PRankPayload1GiBPlan(); TestAllToAllBigDataMultiNodeSmallPayloadUsesPayloadSlot(); TestAllToAllBigDataForce35CorePlanFor8P(); TestAllToAllBigDataBlockDim(); diff --git a/tests/udma/unit/test_tilexr_udma_registry.cpp b/tests/udma/unit/test_tilexr_udma_registry.cpp index c053d629..34b06fd7 100644 --- a/tests/udma/unit/test_tilexr_udma_registry.cpp +++ b/tests/udma/unit/test_tilexr_udma_registry.cpp @@ -33,10 +33,10 @@ void TestRemoteAddressCalculation() registry.version = TileXR::TILEXR_UDMA_REGISTRY_VERSION; registry.regionCount = 1; registry.rankSize = 2; - registry.regions[0].base = reinterpret_cast(0x100000); - registry.regions[0].bytes = 4096; - registry.regions[1].base = reinterpret_cast(0x200000); - registry.regions[1].bytes = 2048; + registry.regions[0][0].base = reinterpret_cast(0x100000); + registry.regions[0][0].bytes = 4096; + registry.regions[1][0].base = reinterpret_cast(0x200000); + registry.regions[1][0].bytes = 2048; CHECK_TRUE(TileXR::UDMARegistryValid(®istry, 2)); CHECK_TRUE(TileXR::UDMARegionContains(®istry, 0, 128, 256)); @@ -56,21 +56,38 @@ void TestRankScaleLimit() registry.version = TileXR::TILEXR_UDMA_REGISTRY_VERSION; registry.regionCount = 1; registry.rankSize = 1024; - registry.regions[TileXR::TILEXR_MAX_RANK_SIZE - 1].base = + registry.regions[TileXR::TILEXR_MAX_RANK_SIZE - 1][0].base = reinterpret_cast(0x300000); - registry.regions[TileXR::TILEXR_MAX_RANK_SIZE - 1].bytes = 4096; + registry.regions[TileXR::TILEXR_MAX_RANK_SIZE - 1][0].bytes = 4096; CHECK_TRUE(TileXR::UDMARegistryValid(®istry, 1024)); CHECK_TRUE(TileXR::UDMARegionContains(®istry, 1023, 2048, 2048)); CHECK_TRUE(!TileXR::UDMARegionContains(®istry, 1024, 0, 1)); } +void TestLogicalOffsetsAcrossRegions() +{ + TileXR::TileXRUDMARegistry registry = {}; + registry.rankSize = 1; + registry.regionCount = 2; + registry.regions[0][0] = {reinterpret_cast(0x100000), 1024}; + registry.regions[0][1] = {reinterpret_cast(0x200000), 2048}; + + CHECK_TRUE(TileXR::UDMARegistryValid(®istry, 1)); + CHECK_TRUE(TileXR::UDMARegionContains(®istry, 0, 1000, 80)); + CHECK_TRUE(TileXR::UDMARegionContains(®istry, 0, 1024 - 8, 16)); + CHECK_TRUE(!TileXR::UDMARegionContains(®istry, 0, 3064, 9)); + CHECK_EQ(reinterpret_cast(TileXR::UDMARemoteAddr(®istry, 0, 1024)), + static_cast(0x200000)); +} + } // namespace int main() { TestRemoteAddressCalculation(); TestRankScaleLimit(); + TestLogicalOffsetsAcrossRegions(); if (g_failures != 0) { std::cerr << g_failures << " registry checks failed" << std::endl; return 1; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index e095b912..d79fe494 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -83,6 +83,7 @@ void TestHostLayoutUsesDeviceRelativePointers() CHECK_EQ(ret, TileXR::TILEXR_UDMA_LAYOUT_SUCCESS); CHECK_EQ(info.qpNum, 1U); + CHECK_EQ(info.regionCount, 1U); CHECK_TRUE(info.sqPtr >= deviceBase); CHECK_TRUE(info.rqPtr > info.sqPtr); CHECK_TRUE(info.scqPtr > info.rqPtr); @@ -111,6 +112,29 @@ void TestHostLayoutUsesDeviceRelativePointers() CHECK_EQ(imageWeights[1], 1U); } +void TestMultiRegionMemoryLayout() +{ + std::vector sq(2), rq(2); + std::vector scq(2), rcq(2); + std::vector mem(4); + std::vector weights(2, 1); + for (size_t i = 0; i < mem.size(); ++i) { + mem[i].addr = 0x100000 + i * 0x1000; + mem[i].tid = static_cast(10 + i); + } + TileXR::UDMAInfo info = {}; + std::vector bytes; + const uintptr_t base = 0x80000000; + const int ret = TileXR::BuildUDMAInfoImage( + base, 1, 2, sq, rq, scq, rcq, mem, weights, info, bytes); + CHECK_EQ(ret, TileXR::TILEXR_UDMA_LAYOUT_SUCCESS); + CHECK_EQ(info.regionCount, 2U); + const auto* imageMem = reinterpret_cast( + bytes.data() + (info.memPtr - base)); + CHECK_EQ(imageMem[3].addr, static_cast(0x103000)); + CHECK_EQ(imageMem[3].tid, 13U); +} + void TestRejectsMismatchedArrays() { std::vector sq(2); @@ -382,6 +406,7 @@ int main() { TestHostLayoutUsesDeviceRelativePointers(); TestRejectsMismatchedArrays(); + TestMultiRegionMemoryLayout(); TestMultiRouteQpMappingRepeatsEachRoute(); TestMultiRouteQpMappingRejectsEmptyInputs(); TestMultiRouteQpWeightsUseRouteBandwidth(); From d30ad52b8aa558c4a2b01751eefd088e578fb054 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 30 Jul 2026 15:55:22 +0800 Subject: [PATCH 129/163] feat(udma): support shared QPs across memory regions --- src/comm/udma/tilexr_udma_layout.cpp | 22 +- src/comm/udma/tilexr_udma_layout.h | 8 +- src/comm/udma/tilexr_udma_transport.cpp | 204 +++++++++++------- src/comm/udma/tilexr_udma_transport.h | 2 +- .../test_tilexr_udma_transport_layout.cpp | 15 +- 5 files changed, 170 insertions(+), 81 deletions(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 968fb3cd..ab78e38b 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -10,6 +10,7 @@ #include #include #include +#include namespace TileXR { namespace { @@ -166,9 +167,26 @@ uint32_t UDMASharedQpLane( return lanesPerDirection + (backward - 1) % lanesPerDirection; } -size_t UDMASharedQpPoolSize(uint32_t laneCount, uint32_t eidCount) +uint32_t UDMASharedQpIndex(uint32_t regionIndex, uint32_t lane, uint32_t laneCount) { - return static_cast(laneCount) * eidCount; + if (laneCount == 0 || lane >= laneCount || + regionIndex > (std::numeric_limits::max() - lane) / laneCount) { + return std::numeric_limits::max(); + } + return regionIndex * laneCount + lane; +} + +size_t UDMASharedQpPoolSize(uint32_t laneCount, uint32_t eidCount, uint32_t regionCount) +{ + if (laneCount == 0 || eidCount == 0 || regionCount == 0 || + static_cast(eidCount) > std::numeric_limits::max() / laneCount) { + return 0; + } + const size_t lanesPerRegion = static_cast(laneCount) * eidCount; + if (static_cast(regionCount) > std::numeric_limits::max() / lanesPerRegion) { + return 0; + } + return lanesPerRegion * regionCount; } std::vector SelectExplicitUDMARouteEids( diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index 691164c6..16192703 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -68,9 +68,15 @@ uint32_t UDMASharedQpLane( int rankSize, uint32_t laneCount); +uint32_t UDMASharedQpIndex( + uint32_t regionIndex, + uint32_t lane, + uint32_t laneCount); + size_t UDMASharedQpPoolSize( uint32_t laneCount, - uint32_t eidCount); + uint32_t eidCount, + uint32_t regionCount = 1U); std::vector SelectExplicitUDMARouteEids( const char* routeList, diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 24813b3e..8abf3661 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -883,11 +883,18 @@ int TileXRUDMATransport::CreateQueues() states_[state.eidIndex] = state; auto& activeState = states_[state.eidIndex]; if (sharedQpPool_) { - for (uint32_t lane = 0; lane < sharedQpLaneCount_; ++lane) { - int ret = CreateSharedQueue(activeState, lane); - if (ret != TILEXR_SUCCESS) { - CleanupQueues(); - return ret; + const uint32_t regionCount = static_cast(registeredRegions_.size()); + if (regionCount == 0) { + CleanupQueues(); + return TILEXR_ERROR_INTERNAL; + } + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + for (uint32_t lane = 0; lane < sharedQpLaneCount_; ++lane) { + int ret = CreateSharedQueue(activeState, regionIndex, lane); + if (ret != TILEXR_SUCCESS) { + CleanupQueues(); + return ret; + } } } continue; @@ -930,9 +937,14 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t return CreateQueue(state, peer, qpIdx, localQpIdx, false); } -int TileXRUDMATransport::CreateSharedQueue(PerEidState& state, uint32_t lane) +int TileXRUDMATransport::CreateSharedQueue( + PerEidState& state, uint32_t regionIndex, uint32_t lane) { - return CreateQueue(state, -1, lane, lane, true); + const uint32_t queueIndex = UDMASharedQpIndex(regionIndex, lane, sharedQpLaneCount_); + if (queueIndex == UINT32_MAX) { + return TILEXR_ERROR_PARA_CHECK_FAIL; + } + return CreateQueue(state, -1, queueIndex, queueIndex, true); } int TileXRUDMATransport::CreateQueue( @@ -1163,7 +1175,7 @@ int TileXRUDMATransport::ImportSharedQueues() uint32_t valid; uint32_t eidIndex; uint32_t lane; - uint32_t reserved; + uint32_t regionIndex; QpKeyT key; }; @@ -1182,7 +1194,12 @@ int TileXRUDMATransport::ImportSharedQueues() return TILEXR_ERROR_INTERNAL; } - const size_t recordsPerRank = static_cast(options_.rankSize) * maxRoutes; + const uint32_t regionCount = static_cast(registeredRegions_.size()); + if (regionCount == 0) { + return TILEXR_ERROR_INTERNAL; + } + const size_t recordsPerRank = + static_cast(options_.rankSize) * maxRoutes * regionCount; std::vector sendRecords(recordsPerRank); std::vector recvRecords(recordsPerRank); for (int peer = 0; peer < options_.rankSize; ++peer) { @@ -1204,20 +1221,27 @@ int TileXRUDMATransport::ImportSharedQueues() if (stateIt == states_.end()) { return TILEXR_ERROR_INTERNAL; } - const auto queueIt = stateIt->second.sharedQueues.find(lane); - if (queueIt == stateIt->second.sharedQueues.end()) { - return TILEXR_ERROR_INTERNAL; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + const uint32_t sharedQpIndex = + UDMASharedQpIndex(regionIndex, lane, sharedQpLaneCount_); + const auto queueIt = stateIt->second.sharedQueues.find(sharedQpIndex); + if (queueIt == stateIt->second.sharedQueues.end()) { + return TILEXR_ERROR_INTERNAL; + } + const size_t recordIndex = + (static_cast(peer) * maxRoutes + route) * regionCount + regionIndex; + auto& record = sendRecords[recordIndex]; + record.valid = 1; + record.eidIndex = localEid; + record.lane = lane; + record.regionIndex = regionIndex; + record.key = queueIt->second.qpInfo.key; } - auto& record = sendRecords[static_cast(peer) * maxRoutes + route]; - record.valid = 1; - record.eidIndex = localEid; - record.lane = lane; - record.key = queueIt->second.qpInfo.key; } } ret = options_.exchange->AllToAll( - sendRecords.data(), maxRoutes, recvRecords.data()); + sendRecords.data(), static_cast(maxRoutes) * regionCount, recvRecords.data()); if (ret != TILEXR_SUCCESS) { return ret; } @@ -1246,52 +1270,65 @@ int TileXRUDMATransport::ImportSharedQueues() const uint32_t localEid = localRoutes[route]; const uint32_t remoteEid = remoteRoutes[route]; auto stateIt = states_.find(localEid); - if (stateIt == states_.end() || stateIt->second.sharedQueues.count(localLane) == 0 || - remoteEid >= eidCount_) { - return TILEXR_ERROR_INTERNAL; - } - const size_t remoteSlot = static_cast(peer) * maxRoutes + route; - if (remoteSlot >= recvRecords.size()) { - return TILEXR_ERROR_INTERNAL; - } - const auto& record = recvRecords[remoteSlot]; - if (record.valid == 0 || record.eidIndex != remoteEid || record.lane != remoteLane) { - TILEXR_LOG(ERROR) << "Shared QP exchange mismatch rank " << options_.rank - << " peer=" << peer - << " route=" << route - << " expectedEid=" << remoteEid - << " receivedEid=" << record.eidIndex - << " expectedLane=" << remoteLane - << " receivedLane=" << record.lane - << " valid=" << record.valid; - return TILEXR_ERROR_INTERNAL; - } - QpImportInfoT importInfo {}; - importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; - importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; - importInfo.in.ub.type = TARGET_TYPE_JETTY; - importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; - importInfo.in.ub.tpType = 1; - importInfo.in.key = record.key; - PerEidState::SharedRemoteQueueState remote {}; - ret = loader_.RaCtxQpImport(stateIt->second.ctxHandle, &importInfo, &remote.remoteQpHandle); - if (ret != 0 || remote.remoteQpHandle == nullptr) { + if (stateIt == states_.end() || remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - remote.tpn = importInfo.out.ub.tpn; - const uint32_t qpIdx = static_cast(route) * qpsPerRoute_; - stateIt->second.sharedRemoteQueues[QueueKey(peer, qpIdx, qpNum_)] = remote; - if (diag) { - TILEXR_LOG(INFO) << "UDMA diag import shared qp rank " << options_.rank - << " peer=" << peer - << " qpIdx=" << qpIdx - << " localLane=" << localLane - << " remoteLane=" << remoteLane - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " remoteQp=" << remote.remoteQpHandle - << " tpn=" << remote.tpn; + for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { + const uint32_t localSharedQpIndex = + UDMASharedQpIndex(regionIndex, localLane, sharedQpLaneCount_); + if (stateIt->second.sharedQueues.count(localSharedQpIndex) == 0) { + return TILEXR_ERROR_INTERNAL; + } + const size_t remoteSlot = + (static_cast(peer) * maxRoutes + route) * regionCount + regionIndex; + if (remoteSlot >= recvRecords.size()) { + return TILEXR_ERROR_INTERNAL; + } + const auto& record = recvRecords[remoteSlot]; + if (record.valid == 0 || record.eidIndex != remoteEid || + record.lane != remoteLane || record.regionIndex != regionIndex) { + TILEXR_LOG(ERROR) << "Shared QP exchange mismatch rank " << options_.rank + << " peer=" << peer + << " route=" << route + << " region=" << regionIndex + << " expectedEid=" << remoteEid + << " receivedEid=" << record.eidIndex + << " expectedLane=" << remoteLane + << " receivedLane=" << record.lane + << " receivedRegion=" << record.regionIndex + << " valid=" << record.valid; + return TILEXR_ERROR_INTERNAL; + } + QpImportInfoT importInfo {}; + importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; + importInfo.in.ub.type = TARGET_TYPE_JETTY; + importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + importInfo.in.ub.tpType = 1; + importInfo.in.key = record.key; + PerEidState::SharedRemoteQueueState remote {}; + ret = loader_.RaCtxQpImport( + stateIt->second.ctxHandle, &importInfo, &remote.remoteQpHandle); + if (ret != 0 || remote.remoteQpHandle == nullptr) { + return TILEXR_ERROR_INTERNAL; + } + remote.tpn = importInfo.out.ub.tpn; + const uint32_t qpIdx = + static_cast(route) * qpsPerRoute_ + regionIndex; + stateIt->second.sharedRemoteQueues[QueueKey(peer, qpIdx, qpNum_)] = remote; + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag import shared qp rank " << options_.rank + << " peer=" << peer + << " qpIdx=" << qpIdx + << " region=" << regionIndex + << " localLane=" << localLane + << " remoteLane=" << remoteLane + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " remoteQp=" << remote.remoteQpHandle + << " tpn=" << remote.tpn; + } } } } @@ -1372,6 +1409,7 @@ int TileXRUDMATransport::RefreshUDMAInfo() for (int rank = 0; rank < options_.rankSize; ++rank) { for (uint32_t qpIdx = 0; qpIdx < qpNum_; ++qpIdx) { + const uint32_t regionIndex = qpIdx % regionCount; uint32_t localEid = fallbackEid; uint32_t remoteEid = fallbackEid; uint32_t routeQpIdx = qpIdx; @@ -1410,9 +1448,19 @@ int TileXRUDMATransport::RefreshUDMAInfo() PerEidState::SharedRemoteQueueState* sharedRemotePtr = nullptr; if (rank == options_.rank) { if (sharedQpPool_ && !state.sharedQueues.empty()) { - queuePtr = &state.sharedQueues.begin()->second; + const uint32_t sharedQpIndex = + UDMASharedQpIndex(regionIndex, 0, sharedQpLaneCount_); + auto queueIt = state.sharedQueues.find(sharedQpIndex); + if (queueIt != state.sharedQueues.end()) { + queuePtr = &queueIt->second; + } } else if (sharedQpPool_ && !fallbackIt->second.sharedQueues.empty()) { - queuePtr = &fallbackIt->second.sharedQueues.begin()->second; + const uint32_t sharedQpIndex = + UDMASharedQpIndex(regionIndex, 0, sharedQpLaneCount_); + auto queueIt = fallbackIt->second.sharedQueues.find(sharedQpIndex); + if (queueIt != fallbackIt->second.sharedQueues.end()) { + queuePtr = &queueIt->second; + } } else if (!state.peerQueues.empty()) { queuePtr = &state.peerQueues.begin()->second; } else if (!fallbackIt->second.peerQueues.empty()) { @@ -1421,7 +1469,9 @@ int TileXRUDMATransport::RefreshUDMAInfo() } else if (sharedQpPool_) { const uint32_t lane = UDMASharedQpLane( options_.rank, rank, options_.rankSize, sharedQpLaneCount_); - auto queueIt = state.sharedQueues.find(lane); + const uint32_t sharedQpIndex = + UDMASharedQpIndex(regionIndex, lane, sharedQpLaneCount_); + auto queueIt = state.sharedQueues.find(sharedQpIndex); if (queueIt == state.sharedQueues.end()) { return TILEXR_ERROR_INTERNAL; } @@ -1452,13 +1502,14 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_ERROR_INTERNAL; } auto& queue = *queuePtr; - if (!registeredRegions_.empty()) { + if (regionIndex < localMemInfoByRegion_.size()) { const auto localMrsIt = registeredMem_.find( - reinterpret_cast(registeredRegions_.front().base)); - const auto& localMrs = localMrsIt->second; - const auto localMrIt = localMrs.find(localEid); - if (localMrIt != localMrs.end()) { - queue.localWq.localTokenId = localMrIt->second.tokenId; + reinterpret_cast(registeredRegions_[regionIndex].base)); + if (localMrsIt != registeredMem_.end()) { + const auto localMrIt = localMrsIt->second.find(localEid); + if (localMrIt != localMrsIt->second.end()) { + queue.localWq.localTokenId = localMrIt->second.tokenId; + } } } const size_t entryIndex = static_cast(rank) * qpNum_ + qpIdx; @@ -1579,10 +1630,6 @@ int TileXRUDMATransport::ConfigureRegionQueues(uint32_t regionCount) if (regionCount == 0 || logicalQpNum_ == 0 || logicalQpsPerRoute_ == 0) { return TILEXR_ERROR_PARA_CHECK_FAIL; } - if (sharedQpPool_ && regionCount > 1) { - TILEXR_LOG(WARN) << "UDMA multi-region registration is incompatible with shared QP pool"; - return TILEXR_ERROR_NOT_FOUND; - } constexpr uint32_t maxQpsPerRoute = 64; if (logicalQpsPerRoute_ > maxQpsPerRoute / regionCount || logicalQpNum_ > UINT32_MAX / regionCount) { @@ -1590,6 +1637,13 @@ int TileXRUDMATransport::ConfigureRegionQueues(uint32_t regionCount) << logicalQpsPerRoute_ << " regionCount=" << regionCount; return TILEXR_ERROR_PARA_CHECK_FAIL; } + if (sharedQpPool_ && + UDMASharedQpPoolSize(sharedQpLaneCount_, static_cast(ctxHandleByEid_.size()), regionCount) == 0) { + TILEXR_LOG(WARN) << "UDMA region-aware shared QP pool size is invalid: lanes=" + << sharedQpLaneCount_ << " eids=" << ctxHandleByEid_.size() + << " regionCount=" << regionCount; + return TILEXR_ERROR_PARA_CHECK_FAIL; + } qpsPerRoute_ = logicalQpsPerRoute_ * regionCount; qpNum_ = logicalQpNum_ * regionCount; diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index 32227140..3e0ed44e 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -55,7 +55,7 @@ class TileXRUDMATransport { int CreateContexts(); int CreateQueues(); int CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx); - int CreateSharedQueue(PerEidState& state, uint32_t lane); + int CreateSharedQueue(PerEidState& state, uint32_t regionIndex, uint32_t lane); int CreateQueue(PerEidState& state, int peer, uint32_t qpIdx, uint32_t queueIndex, bool shared); int ImportQueues(); int ImportSharedQueues(); diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index d79fe494..0ff6be94 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -216,11 +216,21 @@ void TestSharedQpLanesAreUniqueWithinEveryGroup() } } -void TestSharedQpPoolScalesWithLanesAndEidsOnly() +void TestSharedQpPoolScalesWithRegions() { CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2), static_cast(32)); CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 1), static_cast(16)); + CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2, 4), static_cast(128)); CHECK_EQ(TileXR::UDMASharedQpPoolSize(0, 2), static_cast(0)); + CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2, 0), static_cast(0)); +} + +void TestSharedQpIndexAddsRegionDimension() +{ + CHECK_EQ(TileXR::UDMASharedQpIndex(0, 3, 16), 3U); + CHECK_EQ(TileXR::UDMASharedQpIndex(1, 3, 16), 19U); + CHECK_EQ(TileXR::UDMASharedQpIndex(3, 15, 16), 63U); + CHECK_EQ(TileXR::UDMASharedQpIndex(0, 16, 16), UINT32_MAX); } void TestSharedQpLaneRejectsInvalidInputs() @@ -412,7 +422,8 @@ int main() TestMultiRouteQpWeightsUseRouteBandwidth(); TestSharedQpLaneMatchesGroupedPeerOrder(); TestSharedQpLanesAreUniqueWithinEveryGroup(); - TestSharedQpPoolScalesWithLanesAndEidsOnly(); + TestSharedQpPoolScalesWithRegions(); + TestSharedQpIndexAddsRegionDimension(); TestSharedQpLaneRejectsInvalidInputs(); TestSocketExchangeSupportsPersonalizedAllToAll(); TestExplicitRouteSelectionKeepsRequestedCandidateOrder(); From 0e7696d06c77f8e9bffb55e18b56784269e26e14 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 30 Jul 2026 17:42:45 +0800 Subject: [PATCH 130/163] feat(udma): reuse shared QPs across memory regions --- src/comm/udma/tilexr_udma_layout.cpp | 17 +- src/comm/udma/tilexr_udma_layout.h | 4 +- src/comm/udma/tilexr_udma_transport.cpp | 224 ++++++++---------- src/comm/udma/tilexr_udma_transport.h | 2 +- src/include/tilexr_udma.h | 34 +-- tests/udma/demo/tilexr_udma_demo.cpp | 38 ++- tests/udma/demo/tilexr_udma_demo_kernel.cpp | 11 + .../test_tilexr_udma_transport_layout.cpp | 34 ++- 8 files changed, 189 insertions(+), 175 deletions(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index ab78e38b..d4a64cb5 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -167,26 +167,21 @@ uint32_t UDMASharedQpLane( return lanesPerDirection + (backward - 1) % lanesPerDirection; } -uint32_t UDMASharedQpIndex(uint32_t regionIndex, uint32_t lane, uint32_t laneCount) +uint32_t UDMASharedQpIndex(uint32_t lane, uint32_t laneCount) { - if (laneCount == 0 || lane >= laneCount || - regionIndex > (std::numeric_limits::max() - lane) / laneCount) { + if (laneCount == 0 || lane >= laneCount) { return std::numeric_limits::max(); } - return regionIndex * laneCount + lane; + return lane; } -size_t UDMASharedQpPoolSize(uint32_t laneCount, uint32_t eidCount, uint32_t regionCount) +size_t UDMASharedQpPoolSize(uint32_t laneCount, uint32_t eidCount) { - if (laneCount == 0 || eidCount == 0 || regionCount == 0 || + if (laneCount == 0 || eidCount == 0 || static_cast(eidCount) > std::numeric_limits::max() / laneCount) { return 0; } - const size_t lanesPerRegion = static_cast(laneCount) * eidCount; - if (static_cast(regionCount) > std::numeric_limits::max() / lanesPerRegion) { - return 0; - } - return lanesPerRegion * regionCount; + return static_cast(laneCount) * eidCount; } std::vector SelectExplicitUDMARouteEids( diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index 16192703..5343bb2f 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -69,14 +69,12 @@ uint32_t UDMASharedQpLane( uint32_t laneCount); uint32_t UDMASharedQpIndex( - uint32_t regionIndex, uint32_t lane, uint32_t laneCount); size_t UDMASharedQpPoolSize( uint32_t laneCount, - uint32_t eidCount, - uint32_t regionCount = 1U); + uint32_t eidCount); std::vector SelectExplicitUDMARouteEids( const char* routeList, diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 8abf3661..69278ab8 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -883,18 +883,11 @@ int TileXRUDMATransport::CreateQueues() states_[state.eidIndex] = state; auto& activeState = states_[state.eidIndex]; if (sharedQpPool_) { - const uint32_t regionCount = static_cast(registeredRegions_.size()); - if (regionCount == 0) { - CleanupQueues(); - return TILEXR_ERROR_INTERNAL; - } - for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { - for (uint32_t lane = 0; lane < sharedQpLaneCount_; ++lane) { - int ret = CreateSharedQueue(activeState, regionIndex, lane); - if (ret != TILEXR_SUCCESS) { - CleanupQueues(); - return ret; - } + for (uint32_t lane = 0; lane < sharedQpLaneCount_; ++lane) { + int ret = CreateSharedQueue(activeState, lane); + if (ret != TILEXR_SUCCESS) { + CleanupQueues(); + return ret; } } continue; @@ -928,7 +921,22 @@ int TileXRUDMATransport::CreateQueues() } } } - return states_.empty() ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; + if (states_.empty()) { + return TILEXR_ERROR_INTERNAL; + } + if (UDMADiagEnabled()) { + size_t physicalQpCount = 0; + for (const auto& stateEntry : states_) { + physicalQpCount += sharedQpPool_ ? stateEntry.second.sharedQueues.size() : + stateEntry.second.peerQueues.size(); + } + TILEXR_LOG(INFO) << "UDMA diag queue topology rank " << options_.rank + << " shared=" << sharedQpPool_ + << " physicalQps=" << physicalQpCount + << " logicalQpNum=" << qpNum_ + << " regionCount=" << registeredRegions_.size(); + } + return TILEXR_SUCCESS; } int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx) @@ -937,10 +945,9 @@ int TileXRUDMATransport::CreatePeerQueue(PerEidState& state, int peer, uint32_t return CreateQueue(state, peer, qpIdx, localQpIdx, false); } -int TileXRUDMATransport::CreateSharedQueue( - PerEidState& state, uint32_t regionIndex, uint32_t lane) +int TileXRUDMATransport::CreateSharedQueue(PerEidState& state, uint32_t lane) { - const uint32_t queueIndex = UDMASharedQpIndex(regionIndex, lane, sharedQpLaneCount_); + const uint32_t queueIndex = UDMASharedQpIndex(lane, sharedQpLaneCount_); if (queueIndex == UINT32_MAX) { return TILEXR_ERROR_PARA_CHECK_FAIL; } @@ -1175,7 +1182,7 @@ int TileXRUDMATransport::ImportSharedQueues() uint32_t valid; uint32_t eidIndex; uint32_t lane; - uint32_t regionIndex; + uint32_t reserved; QpKeyT key; }; @@ -1194,12 +1201,7 @@ int TileXRUDMATransport::ImportSharedQueues() return TILEXR_ERROR_INTERNAL; } - const uint32_t regionCount = static_cast(registeredRegions_.size()); - if (regionCount == 0) { - return TILEXR_ERROR_INTERNAL; - } - const size_t recordsPerRank = - static_cast(options_.rankSize) * maxRoutes * regionCount; + const size_t recordsPerRank = static_cast(options_.rankSize) * maxRoutes; std::vector sendRecords(recordsPerRank); std::vector recvRecords(recordsPerRank); for (int peer = 0; peer < options_.rankSize; ++peer) { @@ -1221,27 +1223,20 @@ int TileXRUDMATransport::ImportSharedQueues() if (stateIt == states_.end()) { return TILEXR_ERROR_INTERNAL; } - for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { - const uint32_t sharedQpIndex = - UDMASharedQpIndex(regionIndex, lane, sharedQpLaneCount_); - const auto queueIt = stateIt->second.sharedQueues.find(sharedQpIndex); - if (queueIt == stateIt->second.sharedQueues.end()) { - return TILEXR_ERROR_INTERNAL; - } - const size_t recordIndex = - (static_cast(peer) * maxRoutes + route) * regionCount + regionIndex; - auto& record = sendRecords[recordIndex]; - record.valid = 1; - record.eidIndex = localEid; - record.lane = lane; - record.regionIndex = regionIndex; - record.key = queueIt->second.qpInfo.key; + const uint32_t sharedQpIndex = UDMASharedQpIndex(lane, sharedQpLaneCount_); + const auto queueIt = stateIt->second.sharedQueues.find(sharedQpIndex); + if (queueIt == stateIt->second.sharedQueues.end()) { + return TILEXR_ERROR_INTERNAL; } + auto& record = sendRecords[static_cast(peer) * maxRoutes + route]; + record.valid = 1; + record.eidIndex = localEid; + record.lane = lane; + record.key = queueIt->second.qpInfo.key; } } - ret = options_.exchange->AllToAll( - sendRecords.data(), static_cast(maxRoutes) * regionCount, recvRecords.data()); + ret = options_.exchange->AllToAll(sendRecords.data(), maxRoutes, recvRecords.data()); if (ret != TILEXR_SUCCESS) { return ret; } @@ -1273,62 +1268,54 @@ int TileXRUDMATransport::ImportSharedQueues() if (stateIt == states_.end() || remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { - const uint32_t localSharedQpIndex = - UDMASharedQpIndex(regionIndex, localLane, sharedQpLaneCount_); - if (stateIt->second.sharedQueues.count(localSharedQpIndex) == 0) { - return TILEXR_ERROR_INTERNAL; - } - const size_t remoteSlot = - (static_cast(peer) * maxRoutes + route) * regionCount + regionIndex; - if (remoteSlot >= recvRecords.size()) { - return TILEXR_ERROR_INTERNAL; - } - const auto& record = recvRecords[remoteSlot]; - if (record.valid == 0 || record.eidIndex != remoteEid || - record.lane != remoteLane || record.regionIndex != regionIndex) { - TILEXR_LOG(ERROR) << "Shared QP exchange mismatch rank " << options_.rank - << " peer=" << peer - << " route=" << route - << " region=" << regionIndex - << " expectedEid=" << remoteEid - << " receivedEid=" << record.eidIndex - << " expectedLane=" << remoteLane - << " receivedLane=" << record.lane - << " receivedRegion=" << record.regionIndex - << " valid=" << record.valid; - return TILEXR_ERROR_INTERNAL; - } - QpImportInfoT importInfo {}; - importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; - importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; - importInfo.in.ub.type = TARGET_TYPE_JETTY; - importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; - importInfo.in.ub.tpType = 1; - importInfo.in.key = record.key; - PerEidState::SharedRemoteQueueState remote {}; - ret = loader_.RaCtxQpImport( - stateIt->second.ctxHandle, &importInfo, &remote.remoteQpHandle); - if (ret != 0 || remote.remoteQpHandle == nullptr) { - return TILEXR_ERROR_INTERNAL; - } - remote.tpn = importInfo.out.ub.tpn; - const uint32_t qpIdx = - static_cast(route) * qpsPerRoute_ + regionIndex; - stateIt->second.sharedRemoteQueues[QueueKey(peer, qpIdx, qpNum_)] = remote; - if (diag) { - TILEXR_LOG(INFO) << "UDMA diag import shared qp rank " << options_.rank - << " peer=" << peer - << " qpIdx=" << qpIdx - << " region=" << regionIndex - << " localLane=" << localLane - << " remoteLane=" << remoteLane - << " localEid=" << localEid - << " remoteEid=" << remoteEid - << " remoteQp=" << remote.remoteQpHandle - << " tpn=" << remote.tpn; - } + const uint32_t localSharedQpIndex = + UDMASharedQpIndex(localLane, sharedQpLaneCount_); + if (stateIt->second.sharedQueues.count(localSharedQpIndex) == 0) { + return TILEXR_ERROR_INTERNAL; + } + const size_t remoteSlot = static_cast(peer) * maxRoutes + route; + if (remoteSlot >= recvRecords.size()) { + return TILEXR_ERROR_INTERNAL; + } + const auto& record = recvRecords[remoteSlot]; + if (record.valid == 0 || record.eidIndex != remoteEid || record.lane != remoteLane) { + TILEXR_LOG(ERROR) << "Shared QP exchange mismatch rank " << options_.rank + << " peer=" << peer + << " route=" << route + << " expectedEid=" << remoteEid + << " receivedEid=" << record.eidIndex + << " expectedLane=" << remoteLane + << " receivedLane=" << record.lane + << " valid=" << record.valid; + return TILEXR_ERROR_INTERNAL; + } + QpImportInfoT importInfo {}; + importInfo.in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + importInfo.in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + importInfo.in.ub.policy = JETTY_GRP_POLICY_RR; + importInfo.in.ub.type = TARGET_TYPE_JETTY; + importInfo.in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + importInfo.in.ub.tpType = 1; + importInfo.in.key = record.key; + PerEidState::SharedRemoteQueueState remote {}; + ret = loader_.RaCtxQpImport( + stateIt->second.ctxHandle, &importInfo, &remote.remoteQpHandle); + if (ret != 0 || remote.remoteQpHandle == nullptr) { + return TILEXR_ERROR_INTERNAL; + } + remote.tpn = importInfo.out.ub.tpn; + const uint32_t qpIdx = static_cast(route) * qpsPerRoute_; + stateIt->second.sharedRemoteQueues[QueueKey(peer, qpIdx, qpNum_)] = remote; + if (diag) { + TILEXR_LOG(INFO) << "UDMA diag import shared qp rank " << options_.rank + << " peer=" << peer + << " qpIdx=" << qpIdx + << " localLane=" << localLane + << " remoteLane=" << remoteLane + << " localEid=" << localEid + << " remoteEid=" << remoteEid + << " remoteQp=" << remote.remoteQpHandle + << " tpn=" << remote.tpn; } } } @@ -1409,7 +1396,6 @@ int TileXRUDMATransport::RefreshUDMAInfo() for (int rank = 0; rank < options_.rankSize; ++rank) { for (uint32_t qpIdx = 0; qpIdx < qpNum_; ++qpIdx) { - const uint32_t regionIndex = qpIdx % regionCount; uint32_t localEid = fallbackEid; uint32_t remoteEid = fallbackEid; uint32_t routeQpIdx = qpIdx; @@ -1448,15 +1434,13 @@ int TileXRUDMATransport::RefreshUDMAInfo() PerEidState::SharedRemoteQueueState* sharedRemotePtr = nullptr; if (rank == options_.rank) { if (sharedQpPool_ && !state.sharedQueues.empty()) { - const uint32_t sharedQpIndex = - UDMASharedQpIndex(regionIndex, 0, sharedQpLaneCount_); + const uint32_t sharedQpIndex = UDMASharedQpIndex(0, sharedQpLaneCount_); auto queueIt = state.sharedQueues.find(sharedQpIndex); if (queueIt != state.sharedQueues.end()) { queuePtr = &queueIt->second; } } else if (sharedQpPool_ && !fallbackIt->second.sharedQueues.empty()) { - const uint32_t sharedQpIndex = - UDMASharedQpIndex(regionIndex, 0, sharedQpLaneCount_); + const uint32_t sharedQpIndex = UDMASharedQpIndex(0, sharedQpLaneCount_); auto queueIt = fallbackIt->second.sharedQueues.find(sharedQpIndex); if (queueIt != fallbackIt->second.sharedQueues.end()) { queuePtr = &queueIt->second; @@ -1469,8 +1453,7 @@ int TileXRUDMATransport::RefreshUDMAInfo() } else if (sharedQpPool_) { const uint32_t lane = UDMASharedQpLane( options_.rank, rank, options_.rankSize, sharedQpLaneCount_); - const uint32_t sharedQpIndex = - UDMASharedQpIndex(regionIndex, lane, sharedQpLaneCount_); + const uint32_t sharedQpIndex = UDMASharedQpIndex(lane, sharedQpLaneCount_); auto queueIt = state.sharedQueues.find(sharedQpIndex); if (queueIt == state.sharedQueues.end()) { return TILEXR_ERROR_INTERNAL; @@ -1502,9 +1485,9 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_ERROR_INTERNAL; } auto& queue = *queuePtr; - if (regionIndex < localMemInfoByRegion_.size()) { + if (!registeredRegions_.empty()) { const auto localMrsIt = registeredMem_.find( - reinterpret_cast(registeredRegions_[regionIndex].base)); + reinterpret_cast(registeredRegions_.front().base)); if (localMrsIt != registeredMem_.end()) { const auto localMrIt = localMrsIt->second.find(localEid); if (localMrIt != localMrsIt->second.end()) { @@ -1631,36 +1614,23 @@ int TileXRUDMATransport::ConfigureRegionQueues(uint32_t regionCount) return TILEXR_ERROR_PARA_CHECK_FAIL; } constexpr uint32_t maxQpsPerRoute = 64; - if (logicalQpsPerRoute_ > maxQpsPerRoute / regionCount || - logicalQpNum_ > UINT32_MAX / regionCount) { - TILEXR_LOG(WARN) << "UDMA region QP expansion exceeds limit: logicalQpsPerRoute=" - << logicalQpsPerRoute_ << " regionCount=" << regionCount; + if (logicalQpsPerRoute_ > maxQpsPerRoute) { + TILEXR_LOG(WARN) << "UDMA logical QP count exceeds limit: logicalQpsPerRoute=" + << logicalQpsPerRoute_; return TILEXR_ERROR_PARA_CHECK_FAIL; } if (sharedQpPool_ && - UDMASharedQpPoolSize(sharedQpLaneCount_, static_cast(ctxHandleByEid_.size()), regionCount) == 0) { - TILEXR_LOG(WARN) << "UDMA region-aware shared QP pool size is invalid: lanes=" + UDMASharedQpPoolSize(sharedQpLaneCount_, static_cast(ctxHandleByEid_.size())) == 0) { + TILEXR_LOG(WARN) << "UDMA shared QP pool size is invalid: lanes=" << sharedQpLaneCount_ << " eids=" << ctxHandleByEid_.size() << " regionCount=" << regionCount; return TILEXR_ERROR_PARA_CHECK_FAIL; } - qpsPerRoute_ = logicalQpsPerRoute_ * regionCount; - qpNum_ = logicalQpNum_ * regionCount; - peerQpRouteEids_.clear(); - peerQpRouteWeights_.clear(); - for (const auto& entry : logicalPeerQpRouteEids_) { - auto& expanded = peerQpRouteEids_[entry.first]; - for (uint32_t eid : entry.second) { - expanded.insert(expanded.end(), regionCount, eid); - } - } - for (const auto& entry : logicalPeerQpRouteWeights_) { - auto& expanded = peerQpRouteWeights_[entry.first]; - for (uint32_t weight : entry.second) { - expanded.insert(expanded.end(), regionCount, weight); - } - } + qpsPerRoute_ = logicalQpsPerRoute_; + qpNum_ = logicalQpNum_; + peerQpRouteEids_ = logicalPeerQpRouteEids_; + peerQpRouteWeights_ = logicalPeerQpRouteWeights_; return TILEXR_SUCCESS; } diff --git a/src/comm/udma/tilexr_udma_transport.h b/src/comm/udma/tilexr_udma_transport.h index 3e0ed44e..32227140 100644 --- a/src/comm/udma/tilexr_udma_transport.h +++ b/src/comm/udma/tilexr_udma_transport.h @@ -55,7 +55,7 @@ class TileXRUDMATransport { int CreateContexts(); int CreateQueues(); int CreatePeerQueue(PerEidState& state, int peer, uint32_t qpIdx); - int CreateSharedQueue(PerEidState& state, uint32_t regionIndex, uint32_t lane); + int CreateSharedQueue(PerEidState& state, uint32_t lane); int CreateQueue(PerEidState& state, int peer, uint32_t qpIdx, uint32_t queueIndex, bool shared); int ImportQueues(); int ImportSharedQueues(); diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index dfee0472..3170d71f 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -167,23 +167,20 @@ __aicore__ inline uint32_t UDMAGetQpWeight(__gm__ UDMAInfo* udmaInfo, uint32_t p return 1; } auto weights = reinterpret_cast<__gm__ uint32_t*>(udmaInfo->qpWeightPtr); - uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; - uint32_t weight = weights[pe * qpNum + qpIdx * regionCount]; + uint32_t weight = weights[pe * qpNum + qpIdx]; return weight == 0 ? 1 : weight; } __aicore__ inline uint32_t UDMAGetLogicalQpNum(const __gm__ UDMAInfo* udmaInfo) { - uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; - uint32_t logicalQpNum = udmaInfo->qpNum / regionCount; - return logicalQpNum == 0 ? 1 : logicalQpNum; + return udmaInfo->qpNum == 0 ? 1 : udmaInfo->qpNum; } __aicore__ inline uint32_t UDMAGetRegionQpIndex( const __gm__ UDMAInfo* udmaInfo, uint32_t logicalQpIdx, uint32_t regionIndex) { uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; - return logicalQpIdx * regionCount + regionIndex; + return regionIndex < regionCount ? logicalQpIdx : udmaInfo->qpNum; } __aicore__ inline void UDMAPollCQUpdateInfo( @@ -536,45 +533,38 @@ __aicore__ inline void UDMAPutRegisteredSignalNbi( UDMAPutSignalNbi(args, targetRank, localSrc, byteOffset, byteCount, signalByteOffset, signal); } -__aicore__ inline uint32_t UDMAPollLogicalQpRegions( +__aicore__ inline uint32_t UDMAPollLogicalQp( const __gm__ CommArgs* args, int targetRank, uint32_t logicalQpIdx) { __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); - uint32_t regionCount = udmaInfo->regionCount == 0 ? 1 : udmaInfo->regionCount; - uint32_t status = 0; - for (uint32_t regionIndex = 0; regionIndex < regionCount; ++regionIndex) { - uint32_t physicalQpIdx = UDMAGetRegionQpIndex(udmaInfo, logicalQpIdx, regionIndex); - if (physicalQpIdx >= udmaInfo->qpNum) return 0xFFFFFFFFU; - __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, physicalQpIdx); - uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); - uint32_t current = UDMAPollCQ(udmaInfo, targetRank, physicalQpIdx, wqeCnt); - if (status == 0 && current != 0) status = current; - } - return status; + if (logicalQpIdx >= udmaInfo->qpNum) return 0xFFFFFFFFU; + __gm__ UDMAWQCtx* qpCtxEntry = UDMAGetWQCtx(udmaInfo, targetRank, logicalQpIdx); + uint32_t wqeCnt = ld_dev(reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); + return UDMAPollCQ(udmaInfo, targetRank, logicalQpIdx, wqeCnt); } __aicore__ inline void UDMAQuiet(const __gm__ CommArgs* args, int targetRank) { if (!UDMAEnabled(args)) return; - (void)UDMAPollLogicalQpRegions(args, targetRank, 0); + (void)UDMAPollLogicalQp(args, targetRank, 0); } __aicore__ inline void UDMAQuietOnQp(const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx) { if (!UDMAEnabled(args)) return; - (void)UDMAPollLogicalQpRegions(args, targetRank, qpIdx); + (void)UDMAPollLogicalQp(args, targetRank, qpIdx); } __aicore__ inline uint32_t UDMAQuietStatus(const __gm__ CommArgs* args, int targetRank) { if (!UDMAEnabled(args)) return 0xFFFFFFFFU; - return UDMAPollLogicalQpRegions(args, targetRank, 0); + return UDMAPollLogicalQp(args, targetRank, 0); } __aicore__ inline uint32_t UDMAQuietStatusOnQp(const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx) { if (!UDMAEnabled(args)) return 0xFFFFFFFFU; - return UDMAPollLogicalQpRegions(args, targetRank, qpIdx); + return UDMAPollLogicalQp(args, targetRank, qpIdx); } } // namespace TileXR diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 396f211f..faccdbe7 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -117,6 +117,8 @@ constexpr uint64_t kVmmProbeRegionBytes = 1ULL << 30; constexpr uint64_t kVmmProbeMarkerDstOffset = 64; constexpr uint64_t kVmmProbeMarkerSrcOffset = 128; constexpr uint64_t kVmmProbeBoundaryBytes = 64; +constexpr uint64_t kVmmProbeNotifySrcOffset = 8192; +constexpr uint64_t kVmmProbeNotifyDstOffset = 16384; struct BarrierEndpoint { std::string host; @@ -617,6 +619,18 @@ bool RunVmmMultiRegionProbe( aclrtMemcpy(regions[i].base + kVmmProbeMarkerDstOffset, sizeof(zero), &zero, sizeof(zero), ACL_MEMCPY_HOST_TO_DEVICE)) && ok; } + const uint64_t notifyPayload = (static_cast(rank) << 32) | 0xB6000000ULL; + const uint64_t zero = 0; + ok = CheckAcl(rank, "init cross-MR notify source", + aclrtMemcpy(static_cast(base) + kVmmProbeNotifySrcOffset, + sizeof(notifyPayload), ¬ifyPayload, sizeof(notifyPayload), + ACL_MEMCPY_HOST_TO_DEVICE)) && ok; + ok = CheckAcl(rank, "clear cross-MR notify payload", + aclrtMemcpy(static_cast(base) + kVmmProbeNotifyDstOffset, + sizeof(zero), &zero, sizeof(zero), ACL_MEMCPY_HOST_TO_DEVICE)) && ok; + ok = CheckAcl(rank, "clear cross-MR notify signal", + aclrtMemcpy(static_cast(base) + regionBytes + kVmmProbeNotifyDstOffset, + sizeof(zero), &zero, sizeof(zero), ACL_MEMCPY_HOST_TO_DEVICE)) && ok; std::array boundarySource {}; for (uint32_t i = 0; i < boundarySource.size(); ++i) { boundarySource[i] = static_cast((rank * 17 + i) & 0xFF); @@ -650,6 +664,27 @@ bool RunVmmMultiRegionProbe( ok = false; } } + uint64_t notifyPayloadActual = 0; + uint64_t notifySignalActual = 0; + const uint64_t notifyPayloadExpected = + (static_cast(predecessor) << 32) | 0xB6000000ULL; + const uint64_t notifySignalExpected = + (static_cast(predecessor) << 32) | 0xC7000000ULL; + ok = CheckAcl(rank, "read cross-MR notify payload", + aclrtMemcpy(¬ifyPayloadActual, sizeof(notifyPayloadActual), + static_cast(base) + kVmmProbeNotifyDstOffset, + sizeof(notifyPayloadActual), ACL_MEMCPY_DEVICE_TO_HOST)) && ok; + ok = CheckAcl(rank, "read cross-MR notify signal", + aclrtMemcpy(¬ifySignalActual, sizeof(notifySignalActual), + static_cast(base) + regionBytes + kVmmProbeNotifyDstOffset, + sizeof(notifySignalActual), ACL_MEMCPY_DEVICE_TO_HOST)) && ok; + if (notifyPayloadActual != notifyPayloadExpected || notifySignalActual != notifySignalExpected) { + std::cerr << "[rank " << rank << "] ERROR: cross-MR notify payload=" + << notifyPayloadActual << " expectedPayload=" << notifyPayloadExpected + << " signal=" << notifySignalActual + << " expectedSignal=" << notifySignalExpected << std::endl; + ok = false; + } std::array boundaryActual {}; ok = CheckAcl(rank, "read cross-region boundary", aclrtMemcpy(boundaryActual.data(), boundaryActual.size(), boundaryDestinationAddr, boundaryActual.size(), @@ -670,7 +705,8 @@ bool RunVmmMultiRegionProbe( releaseVmm(); if (ok) { std::cout << "[rank " << rank - << "] VMM 4GiB / ping+pong 2GiB each / 4x1GiB MR probe success" << std::endl; + << "] VMM 4GiB / shared QP cross-MR payload+notify / 4x1GiB MR probe success" + << std::endl; } return ok; } diff --git a/tests/udma/demo/tilexr_udma_demo_kernel.cpp b/tests/udma/demo/tilexr_udma_demo_kernel.cpp index b9a0ae77..e515059f 100644 --- a/tests/udma/demo/tilexr_udma_demo_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_demo_kernel.cpp @@ -2576,6 +2576,17 @@ extern "C" __global__ __aicore__ void tilexr_udma_vmm_regions_probe_kernel( TileXR::UDMAQuiet(args, targetRank); } + constexpr uint64_t notifySourceOffset = 8192; + constexpr uint64_t notifyDestinationOffset = 16384; + auto notifySource = localBase + notifySourceOffset; + const uint64_t notifySignalOffset = regionBytes + notifyDestinationOffset; + const uint64_t notifySignal = + (static_cast(args->rank) << 32) | 0xC7000000ULL; + TileXR::UDMAPutSignalNbiOnQp( + args, targetRank, 0, notifySource, notifyDestinationOffset, + sizeof(uint64_t), notifySignalOffset, notifySignal); + TileXR::UDMAQuiet(args, targetRank); + constexpr uint32_t boundaryBytes = 64; auto boundarySource = localBase + 4096; const uint64_t boundaryDestinationOffset = regionBytes - boundaryBytes / 2; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 0ff6be94..080f1c39 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -216,21 +216,19 @@ void TestSharedQpLanesAreUniqueWithinEveryGroup() } } -void TestSharedQpPoolScalesWithRegions() +void TestSharedQpPoolIsIndependentOfRegions() { CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2), static_cast(32)); CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 1), static_cast(16)); - CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2, 4), static_cast(128)); CHECK_EQ(TileXR::UDMASharedQpPoolSize(0, 2), static_cast(0)); - CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 2, 0), static_cast(0)); + CHECK_EQ(TileXR::UDMASharedQpPoolSize(16, 0), static_cast(0)); } -void TestSharedQpIndexAddsRegionDimension() +void TestSharedQpIndexUsesLaneOnly() { - CHECK_EQ(TileXR::UDMASharedQpIndex(0, 3, 16), 3U); - CHECK_EQ(TileXR::UDMASharedQpIndex(1, 3, 16), 19U); - CHECK_EQ(TileXR::UDMASharedQpIndex(3, 15, 16), 63U); - CHECK_EQ(TileXR::UDMASharedQpIndex(0, 16, 16), UINT32_MAX); + CHECK_EQ(TileXR::UDMASharedQpIndex(3, 16), 3U); + CHECK_EQ(TileXR::UDMASharedQpIndex(15, 16), 15U); + CHECK_EQ(TileXR::UDMASharedQpIndex(16, 16), UINT32_MAX); } void TestSharedQpLaneRejectsInvalidInputs() @@ -336,6 +334,21 @@ void TestTransportHasOptInSharedQpPool() CHECK_NOT_CONTAINS(transport, "allSharedKeys"); CHECK_CONTAINS(transport, "state.sharedRemoteQueues.clear()"); CHECK_CONTAINS(transport, "auto cleanupQueue = [&]()"); + CHECK_CONTAINS(transport, "qpsPerRoute_ = logicalQpsPerRoute_"); + CHECK_CONTAINS(transport, "qpNum_ = logicalQpNum_"); + CHECK_NOT_CONTAINS(transport, "logicalQpsPerRoute_ * regionCount"); + CHECK_NOT_CONTAINS(transport, "logicalQpNum_ * regionCount"); +} + +void TestDeviceReusesLogicalQpAcrossRegions() +{ + const std::string device = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); + + CHECK_CONTAINS(device, "return regionIndex < regionCount ? logicalQpIdx : udmaInfo->qpNum"); + CHECK_CONTAINS(device, "return udmaInfo->qpNum == 0 ? 1 : udmaInfo->qpNum"); + CHECK_NOT_CONTAINS(device, "logicalQpIdx * regionCount + regionIndex"); + CHECK_NOT_CONTAINS(device, "qpIdx * regionCount"); } void TestSocketExchangeSupportsPersonalizedAllToAll() @@ -422,8 +435,8 @@ int main() TestMultiRouteQpWeightsUseRouteBandwidth(); TestSharedQpLaneMatchesGroupedPeerOrder(); TestSharedQpLanesAreUniqueWithinEveryGroup(); - TestSharedQpPoolScalesWithRegions(); - TestSharedQpIndexAddsRegionDimension(); + TestSharedQpPoolIsIndependentOfRegions(); + TestSharedQpIndexUsesLaneOnly(); TestSharedQpLaneRejectsInvalidInputs(); TestSocketExchangeSupportsPersonalizedAllToAll(); TestExplicitRouteSelectionKeepsRequestedCandidateOrder(); @@ -434,6 +447,7 @@ int main() TestExplicitNodeIdentityOverridesMachineId(); TestTransportUsesPerPeerQueues(); TestTransportHasOptInSharedQpPool(); + TestDeviceReusesLogicalQpAcrossRegions(); TestRootInfoEidBytesSelectRuntimeContexts(); TestRootInfoDeviceOffsetDoesNotDependOnEntryOrder(); TestMemoryRegistrationUsesOfficialUbFlags(); From 4221476d19145a7ed7a920089196f81ddd9aa1d1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 30 Jul 2026 19:04:37 +0800 Subject: [PATCH 131/163] feat(udma): allow 16 GiB grouped alltoall payloads --- .../demo/tilexr_udma_alltoall_group_layout.h | 10 ++++++++- ...test_tilexr_udma_alltoall_group_layout.cpp | 22 ++++++++++++++++--- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index a7f7f9bb..055edcee 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -30,7 +30,12 @@ constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupControlBytes = 4096U; -constexpr size_t kAllToAllGroupMaxRegisteredBytes = 4ULL << 30; +constexpr size_t kAllToAllGroupMaxPayloadBytes = 16ULL << 30; +constexpr size_t kAllToAllGroupMaxRegisteredBytes = + 2U * kAllToAllGroupMaxPayloadBytes + + 2U * static_cast(kAllToAllGroupMaxRankSize) * + kAllToAllGroupSignalSlotBytes + + kAllToAllGroupControlBytes; struct AllToAllGroupPlan { bool valid = false; @@ -227,6 +232,9 @@ inline AllToAllGroupPlan PlanAllToAllGroup( static_cast(rankSize), kAllToAllGroupSignalSlotBytes, plan.signalPlaneBytes)) { return AllToAllGroupPlan {}; } + if (plan.payloadPlaneBytes > kAllToAllGroupMaxPayloadBytes) { + return AllToAllGroupPlan {}; + } plan.payloadOffset[0] = 0U; if (!AllToAllGroupAlignUp(plan.payloadPlaneBytes, plan.payloadOffset[1])) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index b1e9ae87..92b6e35c 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -178,10 +178,26 @@ void TestPlan() CHECK_EQ(oneGiBPerRank.registeredBytes <= TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); - constexpr int32_t twoGiBPerRankElementsPerPeer = 32 * 1024 * 1024; + constexpr size_t largeRankBytes[] = { + 2ULL << 30, 4ULL << 30, 8ULL << 30, 16ULL << 30}; + for (const size_t rankBytes : largeRankBytes) { + const int32_t largeElementsPerPeer = static_cast( + rankBytes / (static_cast(rankSize) * sizeof(int32_t))); + const auto largePlan = TileXR::Demo::PlanAllToAllGroup( + rankSize, largeElementsPerPeer, largeElementsPerPeer); + CHECK_EQ(largePlan.valid, true); + CHECK_EQ(largePlan.passCount, 1U); + CHECK_EQ(largePlan.payloadPlaneBytes, rankBytes); + CHECK_EQ(largePlan.registeredBytes <= + TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); + } + + constexpr size_t tooLargeRankBytes = 32ULL << 30; + const int32_t tooLargeElementsPerPeer = static_cast( + tooLargeRankBytes / (static_cast(rankSize) * sizeof(int32_t))); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup( - rankSize, twoGiBPerRankElementsPerPeer, - twoGiBPerRankElementsPerPeer).valid, false); + rankSize, tooLargeElementsPerPeer, + tooLargeElementsPerPeer).valid, false); } void TestChannelPolicy() From 846193bd2ca46bba0530d3a4e471e606d549a3e1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 30 Jul 2026 19:12:24 +0800 Subject: [PATCH 132/163] feat(udma): expand registered memory regions --- src/include/tilexr_udma.h | 1 + src/include/tilexr_udma_reg.h | 8 +++---- tests/udma/demo/tilexr_udma_demo.cpp | 8 ++++--- tests/udma/unit/test_tilexr_udma_registry.cpp | 23 +++++++++++++++++++ 4 files changed, 33 insertions(+), 7 deletions(-) diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index 3170d71f..7cafcb3c 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -78,6 +78,7 @@ __aicore__ inline bool UDMARegisteredRangeValid( { if (registry == nullptr || registry->magic != TILEXR_UDMA_REGISTRY_MAGIC || registry->version != TILEXR_UDMA_REGISTRY_VERSION || registry->regionCount == 0 || + registry->regionCount > TILEXR_UDMA_MAX_REGIONS || registry->rankSize > TILEXR_MAX_RANK_SIZE || targetRank < 0 || static_cast(targetRank) >= registry->rankSize) { return false; diff --git a/src/include/tilexr_udma_reg.h b/src/include/tilexr_udma_reg.h index 6e47363c..2825cad9 100644 --- a/src/include/tilexr_udma_reg.h +++ b/src/include/tilexr_udma_reg.h @@ -14,8 +14,8 @@ namespace TileXR { constexpr uint32_t TILEXR_UDMA_REGISTRY_MAGIC = 0x54585255U; // TXRU -constexpr uint32_t TILEXR_UDMA_REGISTRY_VERSION = 2U; -constexpr uint32_t TILEXR_UDMA_MAX_REGIONS = 4U; +constexpr uint32_t TILEXR_UDMA_REGISTRY_VERSION = 3U; +constexpr uint32_t TILEXR_UDMA_MAX_REGIONS = 33U; struct TileXRUDMARegionDesc { GM_ADDR base = nullptr; @@ -23,8 +23,8 @@ struct TileXRUDMARegionDesc { }; // Region descriptors are concatenated in array order for device-side offsets. -// A 4 GiB ping/pong workspace can use four 1 GiB MRs while exposing ping at -// offset 0 and pong at offset 2 GiB in one continuous VMM reservation. +// A 16 GiB ping/pong workspace plus control data can use 33 one-GiB MRs while +// exposing them as one continuous logical registered-memory reservation. struct TileXRUDMARegistry { uint32_t magic = TILEXR_UDMA_REGISTRY_MAGIC; diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index faccdbe7..ab68497a 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -458,8 +458,8 @@ struct VmmMultiRegionAllocation { size_t regionBytes = 0; uint32_t regionCount = 0; uint32_t mappedCount = 0; - std::array physical {}; - std::array regions {}; + std::vector physical; + std::vector regions; }; VmmMultiRegionAllocation gRegisteredVmm; @@ -487,7 +487,7 @@ bool AllocateVmmMultiRegion( int rank, int deviceId, size_t regionBytes, uint32_t regionCount, VmmMultiRegionAllocation& allocation) { - if (regionCount == 0 || regionCount > allocation.regions.size()) { + if (regionCount == 0 || regionCount > TileXR::TILEXR_UDMA_MAX_REGIONS) { return false; } aclrtPhysicalMemProp prop {}; @@ -509,6 +509,8 @@ bool AllocateVmmMultiRegion( allocation.regionBytes = regionBytes; allocation.regionCount = regionCount; + allocation.physical.assign(regionCount, nullptr); + allocation.regions.resize(regionCount); const size_t totalBytes = regionBytes * regionCount; if (!CheckAcl(rank, "aclrtReserveMemAddress multi-region", aclrtReserveMemAddress(&allocation.base, totalBytes, 0, nullptr, 1))) { diff --git a/tests/udma/unit/test_tilexr_udma_registry.cpp b/tests/udma/unit/test_tilexr_udma_registry.cpp index 34b06fd7..5acee300 100644 --- a/tests/udma/unit/test_tilexr_udma_registry.cpp +++ b/tests/udma/unit/test_tilexr_udma_registry.cpp @@ -81,6 +81,28 @@ void TestLogicalOffsetsAcrossRegions() static_cast(0x200000)); } +void TestMaximumRegionCount() +{ + CHECK_EQ(TileXR::TILEXR_UDMA_MAX_REGIONS, 33U); + + TileXR::TileXRUDMARegistry registry = {}; + registry.rankSize = 1; + registry.regionCount = TileXR::TILEXR_UDMA_MAX_REGIONS; + for (uint32_t region = 0; region < registry.regionCount; ++region) { + registry.regions[0][region] = { + reinterpret_cast(0x100000 + static_cast(region) * 0x1000), + 1024}; + } + + const uint64_t lastOffset = + static_cast(TileXR::TILEXR_UDMA_MAX_REGIONS - 1U) * 1024U; + CHECK_TRUE(TileXR::UDMARegistryValid(®istry, 1)); + CHECK_TRUE(TileXR::UDMARegionContains(®istry, 0, lastOffset, 1024)); + CHECK_EQ(reinterpret_cast( + TileXR::UDMARemoteAddr(®istry, 0, lastOffset)), + static_cast(0x120000)); +} + } // namespace int main() @@ -88,6 +110,7 @@ int main() TestRemoteAddressCalculation(); TestRankScaleLimit(); TestLogicalOffsetsAcrossRegions(); + TestMaximumRegionCount(); if (g_failures != 0) { std::cerr << g_failures << " registry checks failed" << std::endl; return 1; From 0e76a37190cf5475e8b42d2692bf981310f149cc Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 00:39:15 +0800 Subject: [PATCH 133/163] perf(udma): restore single-region signal fast path --- src/include/tilexr_udma.h | 46 +++++++++++++++++++ .../test_tilexr_udma_transport_layout.cpp | 11 +++++ 2 files changed, 57 insertions(+) diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index 7cafcb3c..ecce0e51 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -83,6 +83,11 @@ __aicore__ inline bool UDMARegisteredRangeValid( static_cast(targetRank) >= registry->rankSize) { return false; } + if (registry->regionCount == 1U) { + const auto& region = registry->regions[targetRank][0]; + return region.base != nullptr && region.bytes != 0 && + byteOffset <= region.bytes && byteCount <= region.bytes - byteOffset; + } uint64_t totalBytes = 0; for (uint32_t regionIndex = 0; regionIndex < registry->regionCount; ++regionIndex) { const auto& region = registry->regions[targetRank][regionIndex]; @@ -98,6 +103,17 @@ __aicore__ inline bool UDMAResolveRegisteredOffset( const __gm__ TileXRUDMARegistry* registry, int targetRank, uint64_t byteOffset, UDMARegionLocation& location) { + if (registry->regionCount == 1U) { + const auto& region = registry->regions[targetRank][0]; + if (byteOffset >= region.bytes) { + return false; + } + location.regionIndex = 0U; + location.regionOffset = byteOffset; + location.bytesAvailable = region.bytes - byteOffset; + location.addr = reinterpret_cast<__gm__ uint8_t*>(region.base + byteOffset); + return true; + } uint64_t cursor = 0; for (uint32_t regionIndex = 0; regionIndex < registry->regionCount; ++regionIndex) { const auto& region = registry->regions[targetRank][regionIndex]; @@ -479,6 +495,36 @@ __aicore__ inline void UDMAPutSignalNbiOnQp( if (!UDMARegistryEnabled(args)) return; auto registry = GetUDMARegistry(args); + if (registry != nullptr && registry->regionCount == 1U) { + if (registry->magic != TILEXR_UDMA_REGISTRY_MAGIC || + registry->version != TILEXR_UDMA_REGISTRY_VERSION || + registry->rankSize > TILEXR_MAX_RANK_SIZE || targetRank < 0 || + static_cast(targetRank) >= registry->rankSize) { + return; + } + const auto& region = registry->regions[targetRank][0]; + if (region.base == nullptr || region.bytes == 0 || + byteOffset > region.bytes || byteCount > region.bytes - byteOffset || + signalByteOffset > region.bytes || + sizeof(uint64_t) > region.bytes - signalByteOffset) { + return; + } + + __gm__ UDMAInfo* udmaInfo = GetUDMAInfo(args); + __gm__ UDMAMemInfo* signalMemInfo = + UDMAGetRemoteMemInfo(udmaInfo, targetRank, qpIdx, 0U); + UDMASignalParams signalParams = {}; + signalParams.sigAddr = reinterpret_cast<__gm__ uint64_t*>( + region.base + signalByteOffset); + signalParams.signal = signal; + signalParams.tid = signalMemInfo->tid; + signalParams.tokenValue = signalMemInfo->rmtTokenValue; + UDMAWriteNotify( + args, reinterpret_cast<__gm__ uint8_t*>(region.base + byteOffset), + reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)), + targetRank, qpIdx, 0U, byteCount, &signalParams); + return; + } if (!UDMARegisteredRangeValid(registry, targetRank, byteOffset, byteCount) || !UDMARegisteredRangeValid(registry, targetRank, signalByteOffset, sizeof(uint64_t))) { return; diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 080f1c39..13fcc075 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -351,6 +351,16 @@ void TestDeviceReusesLogicalQpAcrossRegions() CHECK_NOT_CONTAINS(device, "qpIdx * regionCount"); } +void TestDeviceHasSingleRegionFastPath() +{ + const std::string device = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma.h"); + + CHECK_CONTAINS(device, "if (registry->regionCount == 1U)"); + CHECK_CONTAINS(device, "const auto& region = registry->regions[targetRank][0]"); + CHECK_CONTAINS(device, "targetRank, qpIdx, 0U, byteCount, &signalParams"); +} + void TestSocketExchangeSupportsPersonalizedAllToAll() { const std::string exchange = @@ -448,6 +458,7 @@ int main() TestTransportUsesPerPeerQueues(); TestTransportHasOptInSharedQpPool(); TestDeviceReusesLogicalQpAcrossRegions(); + TestDeviceHasSingleRegionFastPath(); TestRootInfoEidBytesSelectRuntimeContexts(); TestRootInfoDeviceOffsetDoesNotDependOnEntryOrder(); TestMemoryRegistrationUsesOfficialUbFlags(); From 1e942b810ef7f0d4a4e2dbe8debe4eb1983d1920 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 02:19:04 +0800 Subject: [PATCH 134/163] fix(udma): align grouped measured launches --- tests/udma/demo/tilexr_udma_demo.cpp | 4 ++++ .../udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 9 +++++++++ 2 files changed, 13 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index ab68497a..1cb39f16 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -1230,6 +1230,10 @@ bool RunGroupedAllToAll( release(); return false; } + if (!DemoBarrierAll(rank, rankSize, "grouped measured ready")) { + release(); + return false; + } const auto begin = std::chrono::steady_clock::now(); for (int iter = 0; iter < repeat; ++iter, ++invocationId) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 92b6e35c..9a9fe1b0 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -564,6 +564,15 @@ void TestHostStructure() CHECK_CONTAINS(grouped, "\" complete\""); CHECK_CONTAINS(grouped, "DemoBarrierAll(rank, rankSize, \"grouped route stages ready\")"); + CHECK_CONTAINS(grouped, + "DemoBarrierAll(rank, rankSize, \"grouped measured ready\")"); + const size_t warmupSync = grouped.find( + "aclrtSynchronizeStream grouped warmup"); + const size_t measuredBarrier = grouped.find( + "DemoBarrierAll(rank, rankSize, \"grouped measured ready\")"); + const size_t measuredBegin = grouped.find( + "const auto begin = std::chrono::steady_clock::now()"); + CHECK_TRUE(warmupSync < measuredBarrier && measuredBarrier < measuredBegin); CHECK_CONTAINS(demo, "\"/tilexr_group_trace_\" + stageName + \"_rank_\""); } From 35726b45b4c1d87c5914005b1d791316bbe5c49b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 02:20:05 +0800 Subject: [PATCH 135/163] test(udma): use grouped layout assertion helper --- tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 9a9fe1b0..6ad8f26e 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -572,7 +572,7 @@ void TestHostStructure() "DemoBarrierAll(rank, rankSize, \"grouped measured ready\")"); const size_t measuredBegin = grouped.find( "const auto begin = std::chrono::steady_clock::now()"); - CHECK_TRUE(warmupSync < measuredBarrier && measuredBarrier < measuredBegin); + CHECK_EQ(warmupSync < measuredBarrier && measuredBarrier < measuredBegin, true); CHECK_CONTAINS(demo, "\"/tilexr_group_trace_\" + stageName + \"_rank_\""); } From 079fbccf1374dda79bdd52ee855afe5554c4eff3 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 09:46:43 +0800 Subject: [PATCH 136/163] fix(udma): extend CQ polling for large batches --- src/include/tilexr_udma_types.h | 2 +- tests/udma/unit/test_tilexr_udma_transport_layout.cpp | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/include/tilexr_udma_types.h b/src/include/tilexr_udma_types.h index 7a25a61e..a82a0f5e 100644 --- a/src/include/tilexr_udma_types.h +++ b/src/include/tilexr_udma_types.h @@ -16,7 +16,7 @@ constexpr uint32_t TILEXR_UDMA_MAX_SQE_BB_NUM = 4; constexpr uint32_t TILEXR_UDMA_SQ_BB_COUNT = TILEXR_UDMA_SQ_DEPTH * TILEXR_UDMA_MAX_SQE_BB_NUM; constexpr uint32_t TILEXR_UDMA_NUM_CQE_PER_POLL = 100; constexpr uint32_t TILEXR_UDMA_CACHE_LINE_SIZE = 64; -constexpr uint32_t TILEXR_UDMA_MAX_RETRY_TIMES = 1000000; +constexpr uint32_t TILEXR_UDMA_MAX_RETRY_TIMES = 100000000; enum class UDMAOpcode : uint32_t { WRITE = 3, diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 13fcc075..48fdd29d 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -433,6 +433,14 @@ void TestDeviceSqeInitializesOfficialFields() CHECK_CONTAINS(device, "UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth"); } +void TestCompletionPollingAllowsLargeOutstandingBatches() +{ + const std::string types = + ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma_types.h"); + + CHECK_CONTAINS(types, "TILEXR_UDMA_MAX_RETRY_TIMES = 100000000"); +} + } // namespace int main() @@ -464,6 +472,7 @@ int main() TestMemoryRegistrationUsesOfficialUbFlags(); TestDeviceSgeUsesPerPeerLocalTokenId(); TestDeviceSqeInitializesOfficialFields(); + TestCompletionPollingAllowsLargeOutstandingBatches(); if (g_failures != 0) { std::cerr << g_failures << " UDMA transport layout checks failed" << std::endl; return 1; From 9150bcd09576ce5bd6b551b0f953d4c20bcd7745 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 10:02:50 +0800 Subject: [PATCH 137/163] Revert "fix(udma): extend CQ polling for large batches" This reverts commit 079fbccf1374dda79bdd52ee855afe5554c4eff3. --- src/include/tilexr_udma_types.h | 2 +- tests/udma/unit/test_tilexr_udma_transport_layout.cpp | 9 --------- 2 files changed, 1 insertion(+), 10 deletions(-) diff --git a/src/include/tilexr_udma_types.h b/src/include/tilexr_udma_types.h index a82a0f5e..7a25a61e 100644 --- a/src/include/tilexr_udma_types.h +++ b/src/include/tilexr_udma_types.h @@ -16,7 +16,7 @@ constexpr uint32_t TILEXR_UDMA_MAX_SQE_BB_NUM = 4; constexpr uint32_t TILEXR_UDMA_SQ_BB_COUNT = TILEXR_UDMA_SQ_DEPTH * TILEXR_UDMA_MAX_SQE_BB_NUM; constexpr uint32_t TILEXR_UDMA_NUM_CQE_PER_POLL = 100; constexpr uint32_t TILEXR_UDMA_CACHE_LINE_SIZE = 64; -constexpr uint32_t TILEXR_UDMA_MAX_RETRY_TIMES = 100000000; +constexpr uint32_t TILEXR_UDMA_MAX_RETRY_TIMES = 1000000; enum class UDMAOpcode : uint32_t { WRITE = 3, diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 48fdd29d..13fcc075 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -433,14 +433,6 @@ void TestDeviceSqeInitializesOfficialFields() CHECK_CONTAINS(device, "UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth"); } -void TestCompletionPollingAllowsLargeOutstandingBatches() -{ - const std::string types = - ReadFile(std::string(TILEXR_SOURCE_ROOT) + "/src/include/tilexr_udma_types.h"); - - CHECK_CONTAINS(types, "TILEXR_UDMA_MAX_RETRY_TIMES = 100000000"); -} - } // namespace int main() @@ -472,7 +464,6 @@ int main() TestMemoryRegistrationUsesOfficialUbFlags(); TestDeviceSgeUsesPerPeerLocalTokenId(); TestDeviceSqeInitializesOfficialFields(); - TestCompletionPollingAllowsLargeOutstandingBatches(); if (g_failures != 0) { std::cerr << g_failures << " UDMA transport layout checks failed" << std::endl; return 1; From 63c42e2048490b40662302fef95cb2ae75ac4c66 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 11:03:19 +0800 Subject: [PATCH 138/163] perf(udma): split grouped data and signal writes --- src/include/tilexr_udma.h | 57 +++++++++++++++++-- .../tilexr_udma_alltoall_group_kernel.cpp | 35 +++++++++++- .../demo/tilexr_udma_alltoall_group_layout.h | 16 +++++- ...test_tilexr_udma_alltoall_group_layout.cpp | 15 ++++- .../test_tilexr_udma_transport_layout.cpp | 6 +- 5 files changed, 120 insertions(+), 9 deletions(-) diff --git a/src/include/tilexr_udma.h b/src/include/tilexr_udma.h index ecce0e51..9424b646 100644 --- a/src/include/tilexr_udma.h +++ b/src/include/tilexr_udma.h @@ -13,6 +13,11 @@ namespace TileXR { +constexpr uint32_t TILEXR_UDMA_SQE_FLAG_COMPLETION = 0x20U; +constexpr uint32_t TILEXR_UDMA_SQE_FLAG_STRONG_ORDER = 0x02U; +constexpr uint32_t TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION = + TILEXR_UDMA_SQE_FLAG_COMPLETION | TILEXR_UDMA_SQE_FLAG_STRONG_ORDER; + /** * @file tilexr_udma.h * @brief Device-side UDMA wrapper for TileXR-registered memory. @@ -272,11 +277,12 @@ __aicore__ inline void UDMAFillNotifyData( __aicore__ inline void UDMAFillSqeCtx( __gm__ UDMASqeCtx* sqeCtx, __gm__ uint8_t* remoteAddr, __gm__ UDMAMemInfo* remoteMemInfo, - uint32_t curHead, uint32_t depth, UDMAOpcode opcode, const UDMASignalParams* signalParams) + uint32_t curHead, uint32_t depth, UDMAOpcode opcode, const UDMASignalParams* signalParams, + uint32_t sqeFlag = TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION) { sqeCtx->sqeBbIdx = curHead % depth; sqeCtx->opcode = static_cast(opcode); - sqeCtx->flag = 0b00100010; + sqeCtx->flag = sqeFlag; sqeCtx->rsv0 = 0; sqeCtx->nf = 0; sqeCtx->tokenEn = remoteMemInfo->tokenValueValid; @@ -336,7 +342,8 @@ __aicore__ inline void UDMAPostSendUpdateInfo(uint32_t curHead, __gm__ UDMAWQCtx __aicore__ inline void UDMAPostSend( __gm__ UDMAInfo* udmaInfo, __gm__ uint8_t* remoteAddr, __gm__ uint8_t* localAddr, uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen, - UDMAOpcode opcode, const UDMASignalParams* signalParams) + UDMAOpcode opcode, const UDMASignalParams* signalParams, + uint32_t sqeFlag = TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION) { uint32_t physicalQpIdx = UDMAGetRegionQpIndex(udmaInfo, qpIdx, regionIndex); if (physicalQpIdx >= udmaInfo->qpNum) return; @@ -351,7 +358,9 @@ __aicore__ inline void UDMAPostSend( __gm__ uint8_t* wqeAddr = reinterpret_cast<__gm__ uint8_t*>(qpCtxEntry->bufAddr + wqeSize * (curHead % depth)); __gm__ UDMASqeCtx* sqeCtx = reinterpret_cast<__gm__ UDMASqeCtx*>(wqeAddr); - UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth, opcode, signalParams); + UDMAFillSqeCtx( + sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth, + opcode, signalParams, sqeFlag); __gm__ UDMASgeCtx* sgeCtx = reinterpret_cast<__gm__ UDMASgeCtx*>(UDMAGetSgeCtxAddr(wqeAddr, opcode)); UDMAFillSgeCtx(sgeCtx, messageLen, localAddr, qpCtxEntry); @@ -363,6 +372,17 @@ __aicore__ inline void UDMAPostSend( st_dev(wqeCnt, reinterpret_cast<__gm__ uint32_t*>(qpCtxEntry->wqeCntAddr), 0); } +__aicore__ inline void UDMAWriteWithFlag( + const __gm__ CommArgs* args, __gm__ uint8_t* remoteAddr, __gm__ uint8_t* localAddr, + uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen, + uint32_t sqeFlag) +{ + if constexpr (TILEXR_UDMA_ARCH_SUPPORTED) { + UDMAPostSend(GetUDMAInfo(args), remoteAddr, localAddr, pe, qpIdx, regionIndex, + messageLen, UDMAOpcode::WRITE, nullptr, sqeFlag); + } +} + __aicore__ inline void UDMAWrite( const __gm__ CommArgs* args, __gm__ uint8_t* remoteAddr, __gm__ uint8_t* localAddr, uint32_t pe, uint32_t qpIdx, uint32_t regionIndex, uint64_t messageLen) @@ -438,6 +458,35 @@ __aicore__ inline void UDMAPutNbi( UDMAPutNbiOnQp(args, targetRank, 0, localSrc, byteOffset, byteCount); } +template +__aicore__ inline void UDMAPutNbiOnQpWithFlag( + const __gm__ CommArgs* args, int targetRank, uint32_t qpIdx, + const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount, + uint32_t sqeFlag) +{ + if (!UDMARegistryEnabled(args)) return; + + auto registry = GetUDMARegistry(args); + if (!UDMARegisteredRangeValid(registry, targetRank, byteOffset, byteCount)) return; + + auto localAddr = reinterpret_cast<__gm__ uint8_t*>(const_cast<__gm__ T*>(localSrc)); + uint64_t currentOffset = byteOffset; + uint64_t remaining = byteCount; + while (remaining != 0) { + UDMARegionLocation location {}; + if (!UDMAResolveRegisteredOffset(registry, targetRank, currentOffset, location)) return; + uint32_t chunk = UDMAChunkBytes(remaining, location.bytesAvailable); + UDMAWriteWithFlag(args, location.addr, localAddr, targetRank, qpIdx, + location.regionIndex, chunk, sqeFlag); + currentOffset += chunk; + localAddr += chunk; + remaining -= chunk; + if (remaining != 0 && chunk == location.bytesAvailable) { + UDMAQuietInternalOnQp(args, targetRank, qpIdx, location.regionIndex); + } + } +} + template __aicore__ inline void UDMAPutRegisteredNbi( const __gm__ CommArgs* args, int targetRank, const __gm__ T* localSrc, uint64_t byteOffset, uint32_t byteCount) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index ac394520..7cac7127 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -21,6 +21,9 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 1024U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; constexpr uint64_t TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES = 10000000000ULL; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_WORDS = 12U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_BYTES = + TILEXR_ALLTOALL_GROUP_ERROR_WORDS * TILEXR_ALLTOALL_GROUP_BLOCK_DIM * + sizeof(uint32_t); constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; @@ -428,6 +431,18 @@ struct AllToAllGroupQuietState { uint32_t pendingCount = 0U; }; +__aicore__ inline uint32_t AllToAllGroupSignalSourceSlot( + const AllToAllGroupQuietState&) +{ + return 0U; +} + +__aicore__ inline uint32_t AllToAllGroupSignalSourceSlot( + const AllToAllGroupQuietState& state) +{ + return state.pendingCount; +} + __aicore__ inline bool AllToAllGroupFlushQuiet( const __gm__ TileXR::CommArgs* args, AllToAllGroupPendingQuiet* pending, uint32_t& pendingCount, @@ -853,10 +868,26 @@ __aicore__ inline void AllToAllGroupKernelImpl( static_cast(rank) * TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE + static_cast(route) * TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE; + const uint64_t signalSourceIndex = + static_cast(blockIdx) * + TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH + + AllToAllGroupSignalSourceSlot(quietState); + auto signalLocal = reinterpret_cast<__gm__ uint64_t*>( + reinterpret_cast<__gm__ uint8_t*>(debug) + + TILEXR_ALLTOALL_GROUP_ERROR_BYTES + + signalSourceIndex * sizeof(uint64_t)); + *signalLocal = expectedToken; + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t*>(signalLocal), sizeof(uint64_t)); const uint64_t putBegin = AllToAllGroupTraceCycle(groupTrace); - TileXR::UDMAPutSignalNbiOnQp( + TileXR::UDMAPutNbiOnQpWithFlag( args, peer, selectedQp, localSrc, remotePayloadOffset, - segmentElements * sizeof(int32_t), remoteSignalOffset, expectedToken); + segmentElements * sizeof(int32_t), + TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); + TileXR::UDMAPutNbiOnQpWithFlag( + args, peer, selectedQp, signalLocal, remoteSignalOffset, + sizeof(uint64_t), + TileXR::TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION); AllToAllGroupTraceRecordTask( groupTrace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendPutSignal, groupCount, passCount, diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 055edcee..062f77a2 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -25,11 +25,22 @@ constexpr uint32_t kAllToAllGroupRouteSignalStride = 512U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 1024U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; constexpr uint32_t kAllToAllGroupSendWorkerCount = 32U; +constexpr uint32_t kAllToAllGroupErrorWordsPerCore = 12U; +constexpr uint32_t kAllToAllGroupErrorCoreCount = 64U; +constexpr size_t kAllToAllGroupErrorBytes = + static_cast(kAllToAllGroupErrorWordsPerCore) * + kAllToAllGroupErrorCoreCount * sizeof(uint32_t); +constexpr size_t kAllToAllGroupSignalSourceSlots = + static_cast(kAllToAllGroupSendWorkerCount) * + kAllToAllGroupMaxQuietBatch; +constexpr size_t kAllToAllGroupSignalSourceBytes = + kAllToAllGroupSignalSourceSlots * sizeof(uint64_t); constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; -constexpr size_t kAllToAllGroupControlBytes = 4096U; +constexpr size_t kAllToAllGroupControlBytes = + kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes; constexpr size_t kAllToAllGroupMaxPayloadBytes = 16ULL << 30; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 2U * kAllToAllGroupMaxPayloadBytes + @@ -50,6 +61,8 @@ struct AllToAllGroupPlan { size_t signalOffset[kAllToAllGroupPingPongSlots] = {0, 0}; size_t controlOffset = 0; size_t controlBytes = kAllToAllGroupControlBytes; + size_t signalSourceOffset = 0; + size_t signalSourceBytes = kAllToAllGroupSignalSourceBytes; size_t registeredBytes = 0; }; @@ -252,6 +265,7 @@ inline AllToAllGroupPlan PlanAllToAllGroup( !AllToAllGroupAlignUp(cursor, plan.registeredBytes)) { return AllToAllGroupPlan {}; } + plan.signalSourceOffset = plan.controlOffset + kAllToAllGroupErrorBytes; if (plan.registeredBytes > kAllToAllGroupMaxRegisteredBytes) { return AllToAllGroupPlan {}; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 6ad8f26e..0489fb50 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -143,7 +143,14 @@ void TestPlan() CHECK_EQ(plan.signalOffset[0] >= plan.payloadOffset[1] + plan.payloadPlaneBytes, true); CHECK_EQ(plan.signalOffset[1] >= plan.signalOffset[0] + plan.signalPlaneBytes, true); CHECK_EQ(plan.controlOffset >= plan.signalOffset[1] + plan.signalPlaneBytes, true); + CHECK_EQ(plan.signalSourceOffset, + plan.controlOffset + TileXR::Demo::kAllToAllGroupErrorBytes); + CHECK_EQ(plan.signalSourceBytes, + TileXR::Demo::kAllToAllGroupSignalSourceBytes); + CHECK_EQ(plan.signalSourceOffset + plan.signalSourceBytes <= + plan.controlOffset + plan.controlBytes, true); CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); + CHECK_EQ(TileXR::Demo::kAllToAllGroupMaxPayloadBytes, 16ULL << 30); const auto chunked = TileXR::Demo::PlanAllToAllGroup( rankSize, elementsPerPeer, elementsPerPeer / 4); @@ -488,7 +495,13 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers"); CHECK_CONTAINS(kernel, "const uint32_t traceCore = copyoutWorkers == 8U"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx"); - CHECK_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); + CHECK_CONTAINS(kernel, "UDMAPutNbiOnQpWithFlag"); + CHECK_CONTAINS(kernel, "UDMAPutNbiOnQpWithFlag"); + CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbiOnQp"); + CHECK_CONTAINS(kernel, "AllToAllGroupSignalSourceSlot(quietState)"); + CHECK_CONTAINS(kernel, + "AllToAllGroupSignalSourceSlot(\n const AllToAllGroupQuietState&)"); + CHECK_CONTAINS(kernel, "return state.pendingCount;"); CHECK_CONTAINS(kernel, "AllToAllGroupPendingQuiet"); CHECK_CONTAINS(kernel, "AllToAllGroupFlushQuiet"); CHECK_CONTAINS(kernel, diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index 13fcc075..a9f9a5af 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -430,7 +430,11 @@ void TestDeviceSqeInitializesOfficialFields() CHECK_CONTAINS(device, "sqeCtx->rsv1 = 0"); CHECK_CONTAINS(device, "sqeCtx->rsv2 = 0"); CHECK_CONTAINS(device, "sqeCtx->rsv3 = 0"); - CHECK_CONTAINS(device, "UDMAFillSqeCtx(sqeCtx, remoteAddr, remoteMemInfo, curHead, qpCtxEntry->depth"); + CHECK_CONTAINS(device, "UDMAFillSqeCtx("); + CHECK_CONTAINS(device, "opcode, signalParams, sqeFlag"); + CHECK_CONTAINS(device, "sqeCtx->flag = sqeFlag"); + CHECK_CONTAINS(device, "UDMAPutNbiOnQpWithFlag"); + CHECK_CONTAINS(device, "location.regionIndex, chunk, sqeFlag"); } } // namespace From 62ba1f7eb270b51cc04f3e0657b419f189c4cd05 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 15:57:32 +0800 Subject: [PATCH 139/163] fix(trace): label grouped send workers correctly --- ...lexr_udma_alltoall_group_trace_to_chrome.py | 6 +++--- ...lexr_udma_alltoall_group_trace_to_chrome.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index e68ad61f..40209c41 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -178,14 +178,14 @@ def build_chrome_trace(rank_traces): sources.append(rank_trace["path"]) events.append(_metadata("process_name", rank, 0, f"rank {rank}")) for core in range(MAX_CORES): - role = "send" if core < 16 else "receive" + role = "send" if core < 32 else "receive" events.append(_metadata("thread_name", rank, core, f"core{core} {role}")) for iteration in range(header["iteration_count"]): base = bases[(rank, iteration)] offset_us = iteration_offsets[iteration] for core in range(MAX_CORES): - role = "send" if core < 16 else "receive" + role = "send" if core < 32 else "receive" kernel = _read_span( data, kernel_span_offset(iteration, core), f"kernel rank={rank} iter={iteration} core={core}") @@ -217,7 +217,7 @@ def build_chrome_trace(rank_traces): "iteration": iteration, "group": group, "pass": pass_index, - "lane": core if core < 16 else core - 16, + "lane": core % 16 if core < 32 else core - 32, "peer": peer, "qp": None if qp == NO_QP else qp, "role": role, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index 4b8a9be4..f805a7f7 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -105,6 +105,24 @@ def test_converts_all_grouped_pipeline_phases(self): self.assertEqual(trace["otherData"]["displayTimeUnit"], "ns") json.loads(json.dumps(trace)) + def test_labels_32_send_and_32_receive_cores(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilexr_group_trace_rank_0.bin" + self.make_trace(path) + + trace = MODULE.build_chrome_trace([MODULE.read_rank_trace(path)]) + thread_names = { + event["tid"]: event["args"]["name"] + for event in trace["traceEvents"] + if event.get("name") == "thread_name" + } + + self.assertEqual(thread_names[0], "core0 send") + self.assertEqual(thread_names[16], "core16 send") + self.assertEqual(thread_names[31], "core31 send") + self.assertEqual(thread_names[32], "core32 receive") + self.assertEqual(thread_names[63], "core63 receive") + def test_reads_suffixed_stage_trace(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "tilexr_group_trace_primary_rank_0.bin" From aa0eae4225efbf0367752d48730dd99e1646c91b Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 18:35:31 +0800 Subject: [PATCH 140/163] feat(udma): add grouped alltoall ingress credits --- ...6-08-02-grouped-alltoall-ingress-credit.md | 120 ++++++ ...-grouped-alltoall-ingress-credit-design.md | 165 ++++++++ .../tilexr_udma_alltoall_group_kernel.cpp | 354 +++++++++++++++++- .../demo/tilexr_udma_alltoall_group_layout.h | 108 +++++- .../demo/tilexr_udma_alltoall_group_trace.h | 5 +- ...exr_udma_alltoall_group_trace_to_chrome.py | 35 +- tests/udma/demo/tilexr_udma_demo.cpp | 39 +- ...test_tilexr_udma_alltoall_group_layout.cpp | 164 +++++++- ...exr_udma_alltoall_group_trace_to_chrome.py | 33 +- 9 files changed, 989 insertions(+), 34 deletions(-) create mode 100644 docs/plans/2026-08-02-grouped-alltoall-ingress-credit.md create mode 100644 docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md diff --git a/docs/plans/2026-08-02-grouped-alltoall-ingress-credit.md b/docs/plans/2026-08-02-grouped-alltoall-ingress-credit.md new file mode 100644 index 00000000..c92103da --- /dev/null +++ b/docs/plans/2026-08-02-grouped-alltoall-ingress-credit.md @@ -0,0 +1,120 @@ +# Grouped AllToAll Ingress Credit Implementation Plan + +## Goal And Scope + +Implement the default-disabled `window=1` ingress-credit protocol approved in +`docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md`. The change +must bound payload admission to one source per destination lane while +preserving the existing path exactly when credits are disabled. + +This plan does not change route weights, cabinet topology interpretation, +quiet batching, copyout assignment, or multi-pass credit semantics. Initial +hardware acceptance is single-pass only. + +## Task 1: Pure Layout And Schedule Contracts + +**Objective and role:** Define host-testable helpers for credit enablement, +token construction, receive ownership, and next-group peer mapping. + +**Background and prerequisites:** The approved spec is authoritative. Existing +peer mapping and plan helpers live in +`tests/udma/demo/tilexr_udma_alltoall_group_layout.h`. + +**Modification scope:** The grouped layout header and +`tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp`. + +**Constraints and non-goals:** Preserve existing peer mapping and default +layout behavior. Do not add device-only dependencies to the host-testable +header. + +**Acceptance and verification:** Unit tests cover 256 ranks, group zero, +inactive final lanes, owner selection, shared-route credit identity, invocation +separation, and default-disabled behavior. Run the grouped layout unit binary. + +**Artifacts and interfaces:** Stable helper contracts consumed by Host and +Kernel work. + +## Task 2: Registered Credit Plane And Host Interface + +**Objective and role:** Reserve ping-pong credit storage and pass a validated +ingress-window value and credit offsets into kernel launch. + +**Background and prerequisites:** Task 1 token and layout contracts are fixed. +The grouped plan already owns payload, signal, and control regions. + +**Modification scope:** Grouped layout plan, demo environment parsing and +allocation/launch wiring, grouped kernel declaration and launch wrapper, plus +focused source-contract tests. + +**Constraints and non-goals:** `INGRESS_WINDOW=0` remains the default. Only 0 +and 1 are accepted. Existing payload and signal offsets remain stable where +practical; registered-size overflow checks remain mandatory. + +**Acceptance and verification:** Unit tests prove valid/invalid environment +and launch wiring contracts, credit planes do not overlap, and max registered +size accounts for both planes. + +**Artifacts and interfaces:** Kernel receives two credit offsets and the +ingress-window value. + +## Task 3: Kernel Credit Wait And Publication + +**Objective and role:** Gate each destination send once per group, publish a +local request after both route-ready tokens arrive, and let the primary send +core issue the next lane credit. + +**Background and prerequisites:** Tasks 1 and 2 provide mapping, storage, and +arguments. Existing MTE token wait and UDMA NBI primitives are reused. + +**Modification scope:** +`tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp`, existing UDMA helper +interfaces only if required, and grouped source-contract tests. + +**Constraints and non-goals:** Do not add per-credit quiet. Only cores 32..47 +publish local requests; only primary send cores publish UDMA credits, preserving +the shared-SQ single-producer rule. Local request slots are 64-byte isolated so +independent receive-core cache maintenance cannot lose another lane's request. +Primary and secondary workers wait on the same destination credit. The +default-disabled branch must not execute a credit wait or write. If source +lifetime cannot be made safe without per-credit quiet, stop and revise the +design instead of weakening this constraint. + +**Acceptance and verification:** Static/source tests prove placement, +single-producer ownership, request isolation, and unique shared-SQ completion +reclamation. Compile checks validate Ascend C overloads and launch signatures. +The exact final source passes a clean b131 Bisheng build and 4x8 hardware +validation; 256P admission-bound validation remains pending. + +**Artifacts and interfaces:** Trace-visible or debug-visible credit wait and +publish phases where trace capacity permits. + +## Task 4: Verification And Delivery Readiness + +**Objective and role:** Establish local correctness and document remaining +hardware evidence. + +**Background and prerequisites:** Tasks 1 through 3 are complete. + +**Modification scope:** Tests and documentation required to keep validation +commands and limitations accurate. + +**Constraints and non-goals:** Run only the explicitly authorized NPU matrix. +Do not modify or stage unrelated existing worktree files. + +**Acceptance and verification:** Run focused unit tests, relevant Python +parser/source tests, formatting/build checks available on Windows, inspect the +final diff, and record the exact 2x8 then 4x8 then 256P hardware matrix. The +implemented path has completed 4x8 functional and warmup-5/repeat-50 +performance validation; 256P admission-bound validation remains pending. + +**Artifacts and interfaces:** A reviewable implementation, verification +summary, and scoped Git commit after the final checks pass. + +## Key Risks + +- Credit NBI source storage may be reused before hardware consumes it. +- Two receive copy slices may publish duplicate credits. +- Primary and secondary send workers may wait on different tokens by mistake. +- A final inactive lane may wait for or publish a nonexistent peer. +- Registered-memory growth may cross a region boundary or maximum-size check. +- Trace storage may need expansion if new phases are recorded. diff --git a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md new file mode 100644 index 00000000..088895f3 --- /dev/null +++ b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md @@ -0,0 +1,165 @@ +# Grouped AllToAll Ingress Credit Design + +## Status + +Implemented and validated on 4x8 hardware on 2026-08-02. The experimental +feature remains disabled by default. The strict global admission bound still +requires 256P all-rank trace or admission-counter validation. + +## Goal + +Bound each destination rank to at most 16 concurrent payload source ranks in +the grouped AllToAll data path. Preserve independent lane progress so one slow +source does not impose a full-rank or full-group barrier. + +The first implementation is experimental and disabled by default. Existing +behavior and performance remain unchanged unless ingress credits are enabled. + +## Non-goals + +- Limiting the number of physical port operations to 16. A source may still + use primary and secondary routes after it receives one destination credit. +- Adding a global barrier between peer groups. +- Waiting for receive-copy completion before releasing network capacity. +- Changing topology-derived primary/secondary route weights. +- Enabling credit control for multi-pass payloads in the first hardware test. + +## Existing Behavior + +At `rankSize=256` and `groupWidth=16`, every rank has 16 logical peer groups. +Each group contains up to eight forward and eight backward peers. The 32 send +workers are split into 16 primary-route workers and 16 secondary-route workers. + +Each worker advances to its next group immediately after its own quiet +completes. Workers and ranks do not share group progress. Consequently, a +destination may receive payloads from more than one logical group at the same +time even though each group contains no more than 16 source ranks. + +## Credit Protocol + +There are 16 independent destination lane chains. Group zero is initially +enabled. For group `g > 0`, a source must observe a credit from its destination +for `(invocation, g)` before either route posts payload data. + +For destination rank `D`, lane `L`, and group `g`: + +1. The designated receive owner waits until all payload routes expected from + `peer(D, g, L)` have published their data-ready tokens. +2. Before receive-copy, the owner publishes a local request for lane `L`. +3. Primary send core `L` observes that request and posts one credit token to + `peer(D, g + 1, L)`, when that peer exists. +4. The next source observes the credit in its local registered credit plane and + may post both primary and secondary payload routes to `D`. + +At most one source per destination lane can therefore be admitted. With 16 +lanes, the strict payload-source bound is 16 ranks per destination. + +Credit dependencies are monotonic from group `g` to `g + 1`. Group zero has no +dependency, so the protocol does not introduce a cyclic startup wait. + +## Credit Storage And Tokens + +Credits use a separate registered-memory plane. A credit received by source +rank `S` is indexed by destination rank `D`, because `S` communicates with `D` +only once per invocation: + +```text +credit[pingPongSlot][destinationRank] +``` + +The expected value encodes at least the invocation and destination group. The +receiver accepts a value greater than or equal to the expected token, matching +the existing data-ready stale-token policy. Ping-pong slots prevent adjacent +invocations from reusing a live credit location. + +Primary and secondary send workers for the same peer wait on the same credit. +Credits control source-rank admission, not route admission. + +## Receive Ownership + +With 32 copyout workers, workers 0 through 15 and 16 through 31 can wait on the +same peer signal while copying different payload slices. Only the first slice, +kernel cores 32 through 47, owns the local credit request. Cores 48 through 63 +never request credits. + +The owner publishes its request after both expected primary and secondary data +tokens arrive and before MTE receive-copy begins. This avoids placing copyout +latency on the send admission path. + +Each `(pingPongSlot, lane)` request occupies its own 64-byte cache line. Packing +multiple lanes into one line is incorrect: independent receive cores clean and +invalidate the line concurrently, so one core can overwrite another lane's +request with stale cache-line contents. + +## Credit Submission + +Credit publication is an NBI 8-byte UDMA write issued only by the corresponding +primary send core. Receive cores do not issue UDMA operations because they can +share a payload SQ with a send core, while SQ head and WQE count updates require +a single producer. + +The credit path must not perform a quiet for every credit. Token source storage +remains immutable until completion through distinct lane/group slots. At the +end of the invocation, the primary send core reclaims credit completions once +per unique underlying shared SQ, identified by its WQE-count address. This +preserves single-producer SQ ownership and avoids duplicate quiet operations +when multiple peer/QP views refer to the same shared queue. + +If the available UDMA interface cannot safely keep credit source data alive +without per-credit quiet, implementation pauses for a revised design rather +than adding a per-credit quiet to the critical path. + +## Configuration + +Add an experimental ingress window setting: + +```text +TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW=0|1 +``` + +- `0`: current behavior and default. +- `1`: one admitted source per lane, at most 16 payload source ranks per + destination. + +A later `window=2` extension may trade a 32-source bound for more tolerance of +credit latency, but it is outside the first implementation. + +## Validation + +Host/unit tests must cover: + +- 256-rank group and lane predecessor/successor mapping. +- No credit wait for group zero. +- One shared credit for primary and secondary routes. +- Single credit publisher when 32 copyout workers duplicate receive lanes. +- Correct inactive-lane behavior at the diameter and final partial group. +- Credit token separation across invocations and ping-pong slots. +- Default-disabled compatibility. + +Hardware validation starts with 2x8 functional coverage, then 4x8, then 256P. +The 256P comparison uses 1 GiB/rank, multi channel, single pass, warmup 5 and +repeat 50. Evidence includes P50 and range plus per-phase `credit-wait`, +`send-quiet`, and `receive-wait` trace durations. A 1 KiB single-channel case +checks that the default-disabled path has no latency regression. + +The first trace-enabled run is diagnostic only and is not compared directly +with trace-disabled performance. + +The final 4x8, warmup-5/repeat-50, trace-disabled comparison passed data +validation on all 32 ranks: + +```text +1 KiB/rank single: window0 P50 38.224 us, window1 P50 52.763 us +1 GiB/rank multi: window0 P50 5329.210 us, window1 P50 5317.925 us +``` + +This proves the request/credit protocol progresses across repeated invocations +and does not regress 1 GiB throughput in the 4x8 environment. It does not prove +the 16-source global bound because 4x8 has only two peer groups. + +## Residual Risk + +Rank0-only trace cannot prove the global source-card bound. Hardware validation +needs either all-rank trace or lightweight per-destination admission counters. +The physical four-port capacity and cabinet-aware route weighting are separate +topology problems and are not solved by this credit protocol. diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 7cac7127..f07f5573 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -16,6 +16,7 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_BLOCK_DIM = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH = 16U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH = 4U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH = 64U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 1024U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -24,9 +25,25 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_WORDS = 12U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_BYTES = TILEXR_ALLTOALL_GROUP_ERROR_WORDS * TILEXR_ALLTOALL_GROUP_BLOCK_DIM * sizeof(uint32_t); +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES = + TILEXR_ALLTOALL_GROUP_SEND_WORKERS * + TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH * sizeof(uint64_t); +constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_OFFSET = + TILEXR_ALLTOALL_GROUP_ERROR_BYTES + + TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_BYTES = + TILEXR_ALLTOALL_GROUP_SEND_CORES * + TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT * sizeof(uint64_t); +constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_OFFSET = + TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_OFFSET + + TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_BYTES; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_STRIDE = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT = 4U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_QUIET = 5U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_REQUEST_WAIT = 6U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED = 0U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY = 2U; @@ -168,6 +185,11 @@ __aicore__ inline bool AllToAllGroupRemoteAssistDevice( return copyoutWorkers >= 32U && worker >= TILEXR_ALLTOALL_GROUP_SEND_CORES; } +__aicore__ inline bool AllToAllGroupCreditOwnerDevice(uint32_t worker) +{ + return worker < TILEXR_ALLTOALL_GROUP_SEND_CORES; +} + __aicore__ inline void AllToAllGroupSelectRouteQps( const __gm__ TileXR::CommArgs* args, int32_t peer, uint32_t& primaryQp, uint32_t& secondaryQp, @@ -333,6 +355,130 @@ __aicore__ inline bool AllToAllGroupWaitRouteTokensMte( return true; } +__aicore__ inline void AllToAllGroupRecordError( + __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, + uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, + uint32_t quietStatus, uint64_t expectedToken, uint64_t observedToken); + +__aicore__ inline int32_t AllToAllGroupNextCreditPeerDevice( + int32_t rank, int32_t rankSize, uint32_t completedGroup, + uint32_t lane, uint32_t groupCount, uint32_t groupWidth) +{ + if (completedGroup + 1U >= groupCount) { + return -1; + } + return AllToAllGroupDevicePeer( + rank, rankSize, completedGroup + 1U, lane, groupWidth); +} + +__aicore__ inline uint32_t AllToAllGroupCreditQp( + const __gm__ TileXR::CommArgs* args, int32_t peer) +{ + uint32_t primaryQp = 0U; + uint32_t secondaryQp = 0U; + uint32_t primaryWeight = 0U; + uint32_t secondaryWeight = 0U; + AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp, + primaryWeight, secondaryWeight); + return primaryQp; +} + +__aicore__ inline void AllToAllGroupPublishNextCredit( + const __gm__ TileXR::CommArgs* args, __gm__ int32_t* debug, + int32_t rank, int32_t rankSize, uint32_t invocationId, + uint32_t completedGroup, uint32_t lane, uint32_t groupCount, + uint32_t groupWidth, uint64_t remoteCreditOffset) +{ + const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( + rank, rankSize, completedGroup, lane, groupCount, groupWidth); + if (nextPeer < 0) { + return; + } + const uint64_t creditToken = + AllToAllGroupDeviceToken(invocationId, completedGroup + 1U, 0U); + const uint64_t sourceIndex = + static_cast(lane) * TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT + + completedGroup; + auto creditLocal = reinterpret_cast<__gm__ uint64_t*>( + reinterpret_cast<__gm__ uint8_t*>(debug) + + TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_OFFSET + + sourceIndex * sizeof(uint64_t)); + *creditLocal = creditToken; + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t*>(creditLocal), sizeof(uint64_t)); + TileXR::UDMAPutNbiOnQpWithFlag( + args, nextPeer, AllToAllGroupCreditQp(args, nextPeer), creditLocal, + remoteCreditOffset + static_cast(rank) * sizeof(uint64_t), + sizeof(uint64_t), TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); +} + +__aicore__ inline __gm__ uint64_t* AllToAllGroupCreditRequest( + __gm__ int32_t* debug, uint32_t slot, uint32_t lane) +{ + const uint64_t requestIndex = + static_cast(slot) * TILEXR_ALLTOALL_GROUP_SEND_CORES + lane; + return reinterpret_cast<__gm__ uint64_t*>( + reinterpret_cast<__gm__ uint8_t*>(debug) + + TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_OFFSET + + requestIndex * TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_STRIDE); +} + +__aicore__ inline void AllToAllGroupPublishCreditRequest( + __gm__ int32_t* debug, uint32_t invocationId, uint32_t completedGroup, + uint32_t slot, uint32_t lane, uint32_t groupCount) +{ + if (completedGroup + 1U >= groupCount) { + return; + } + auto request = AllToAllGroupCreditRequest(debug, slot, lane); + *request = AllToAllGroupDeviceToken(invocationId, completedGroup + 1U, 0U); + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t*>(request), sizeof(uint64_t)); +} + +__aicore__ inline bool AllToAllGroupFinishCredits( + const __gm__ TileXR::CommArgs* args, __gm__ int32_t* debug, + uint32_t blockIdx, int32_t rank, int32_t rankSize, + uint32_t invocationId, uint32_t lane, uint32_t groupCount, + uint32_t groupWidth, uint32_t publishedCreditCount) +{ + uint64_t completedQueueKeys[TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT] = {}; + uint32_t completedQueueCount = 0U; + bool success = true; + for (uint32_t group = 0U; group < publishedCreditCount; ++group) { + const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( + rank, rankSize, group, lane, groupCount, groupWidth); + if (nextPeer < 0) { + continue; + } + const uint32_t qpIdx = AllToAllGroupCreditQp(args, nextPeer); + auto udmaInfo = TileXR::GetUDMAInfo(args); + auto wq = TileXR::UDMAGetWQCtx(udmaInfo, nextPeer, qpIdx); + const uint64_t queueKey = wq->wqeCntAddr; + bool alreadyCompleted = false; + for (uint32_t index = 0U; index < completedQueueCount; ++index) { + if (completedQueueKeys[index] == queueKey) { + alreadyCompleted = true; + break; + } + } + if (alreadyCompleted) { + continue; + } + completedQueueKeys[completedQueueCount++] = queueKey; + const uint32_t quietStatus = + TileXR::UDMAQuietStatusOnQp(args, nextPeer, qpIdx); + if (quietStatus != 0U) { + AllToAllGroupRecordError(debug, blockIdx, + TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_QUIET, group, 0U, + nextPeer, qpIdx, quietStatus, + AllToAllGroupDeviceToken(invocationId, group + 1U, 0U), 0ULL); + success = false; + } + } + return success; +} + __aicore__ inline void AllToAllGroupRecordError( __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, @@ -533,7 +679,7 @@ __aicore__ inline bool AllToAllGroupFinishQuiet( } // namespace -template +template __aicore__ inline void AllToAllGroupKernelImpl( GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, @@ -541,6 +687,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, uint32_t groupWidth, uint32_t quietBatch) @@ -577,7 +724,12 @@ __aicore__ inline void AllToAllGroupKernelImpl( elementsPerPeer <= 0 || chunkElements <= 0 || passCount == 0U || groupCount == 0U || groupCount != static_cast( (rankSize - 1 + static_cast(groupWidth) - 1) / - static_cast(groupWidth))) { + static_cast(groupWidth)) || + (IngressCredit && + (debug == nullptr || + groupWidth != TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH || + passCount != 1U || + routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED))) { AllToAllGroupRecordError(debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CONFIG, 0U, 0U, -1, 0U, 0U, 0ULL, 0ULL); AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, @@ -585,6 +737,17 @@ __aicore__ inline void AllToAllGroupKernelImpl( return; } + if constexpr (IngressCredit) { + auto priorError = reinterpret_cast<__gm__ AllToAllGroupDeviceError*>( + debug + static_cast(blockIdx) * + TILEXR_ALLTOALL_GROUP_ERROR_WORDS); + if (priorError->valid != 0U) { + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } + } + AscendC::TPipe pipe; AscendC::TBuf relayTBuf; pipe.InitBuffer(relayTBuf, TILEXR_ALLTOALL_GROUP_RELAY_BYTES); @@ -593,6 +756,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( const uint32_t slot = invocationId & 1U; const uint64_t payloadOffsets[2] = {payloadOffset0, payloadOffset1}; const uint64_t signalOffsets[2] = {signalOffset0, signalOffset1}; + const uint64_t creditOffsets[2] = {creditOffset0, creditOffset1}; const uint64_t bytesPerPeer = static_cast(elementsPerPeer) * sizeof(int32_t); @@ -741,6 +905,13 @@ __aicore__ inline void AllToAllGroupKernelImpl( peer, TileXR::Demo::kAllToAllGroupTraceNoQp, waitBegin, AllToAllGroupTraceCycle(groupTrace)); } + if constexpr (IngressCredit) { + if (AllToAllGroupCreditOwnerDevice(worker) && + pass + 1U == passCount) { + AllToAllGroupPublishCreditRequest( + debug, invocationId, group, slot, lane, groupCount); + } + } if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { continue; } @@ -794,6 +965,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( const uint32_t lane = blockIdx % TILEXR_ALLTOALL_GROUP_SEND_CORES; const uint32_t workerRoute = blockIdx / TILEXR_ALLTOALL_GROUP_SEND_CORES; AllToAllGroupQuietState quietState; + uint32_t publishedCreditCount = 0U; for (uint32_t group = 0U; group < groupCount; ++group) { const int32_t peer = AllToAllGroupDevicePeer( rank, rankSize, group, lane, groupWidth); @@ -821,6 +993,89 @@ __aicore__ inline void AllToAllGroupKernelImpl( static_cast(elementsPerPeer), primaryWeight, secondaryWeight, primaryRouteParts, primaryTotalElements, secondaryTotalElements); + if constexpr (IngressCredit) { + const uint32_t routeElements = workerRoute == 0U ? + primaryTotalElements : secondaryTotalElements; + if (group != 0U) { + const uint64_t expectedCredit = + AllToAllGroupDeviceToken(invocationId, group, 0U); + const uint64_t creditWaitBegin = + AllToAllGroupTraceCycle(groupTrace); + if (workerRoute == 0U) { + auto request = AllToAllGroupCreditRequest(debug, slot, lane); + uint64_t observedRequest = 0ULL; + if (!AllToAllGroupWaitTokenMte( + request, expectedCredit, + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, + relayLocal, observedRequest)) { + AllToAllGroupRecordError(debug, blockIdx, + TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_REQUEST_WAIT, + group, 0U, peer, workerRoute, 0U, + expectedCredit, observedRequest); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, 0U, + TileXR::Demo::kAllToAllGroupTraceCreditWait, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + creditWaitBegin, + AllToAllGroupTraceCycle(groupTrace)); + (void)AllToAllGroupFinishCredits( + args, debug, blockIdx, rank, rankSize, + invocationId, lane, groupCount, groupWidth, + publishedCreditCount); + AllToAllGroupTraceRecordKernel( + groupTrace, traceIteration, blockIdx, kernelBegin, + AllToAllGroupTraceCycle(groupTrace)); + return; + } + AllToAllGroupPublishNextCredit( + args, debug, rank, rankSize, invocationId, group - 1U, + lane, groupCount, groupWidth, creditOffsets[slot]); + ++publishedCreditCount; + } + if (routeElements != 0U) { + auto creditSignal = reinterpret_cast<__gm__ uint64_t*>( + registeredMemory + creditOffsets[slot] + + static_cast(peer) * sizeof(uint64_t)); + uint64_t observedCredit = 0ULL; + if (!AllToAllGroupWaitTokenMte( + creditSignal, expectedCredit, + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, + relayLocal, observedCredit)) { + AllToAllGroupRecordError(debug, blockIdx, + TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT, group, 0U, + peer, workerRoute, 0U, expectedCredit, + observedCredit); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, 0U, + TileXR::Demo::kAllToAllGroupTraceCreditWait, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + creditWaitBegin, + AllToAllGroupTraceCycle(groupTrace)); + if (workerRoute == 0U) { + (void)AllToAllGroupFinishCredits( + args, debug, blockIdx, rank, rankSize, + invocationId, lane, groupCount, groupWidth, + publishedCreditCount); + } + AllToAllGroupTraceRecordKernel( + groupTrace, traceIteration, blockIdx, kernelBegin, + AllToAllGroupTraceCycle(groupTrace)); + return; + } + } + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, group, 0U, + TileXR::Demo::kAllToAllGroupTraceCreditWait, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + creditWaitBegin, AllToAllGroupTraceCycle(groupTrace)); + } + if (routeElements == 0U) { + continue; + } + } for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; @@ -896,6 +1151,14 @@ __aicore__ inline void AllToAllGroupKernelImpl( args, quietState, quietBatch, peer, selectedQp, group, pass, expectedToken, debug, blockIdx, groupTrace, traceIteration, groupCount, passCount)) { + if constexpr (IngressCredit) { + if (workerRoute == 0U) { + (void)AllToAllGroupFinishCredits( + args, debug, blockIdx, rank, rankSize, + invocationId, lane, groupCount, groupWidth, + publishedCreditCount); + } + } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; @@ -906,10 +1169,26 @@ __aicore__ inline void AllToAllGroupKernelImpl( if (!AllToAllGroupFinishQuiet( args, quietState, debug, blockIdx, groupTrace, traceIteration, groupCount, passCount)) { + if constexpr (IngressCredit) { + if (workerRoute == 0U) { + (void)AllToAllGroupFinishCredits( + args, debug, blockIdx, rank, rankSize, invocationId, lane, + groupCount, groupWidth, publishedCreditCount); + } + } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } + if constexpr (IngressCredit) { + if (workerRoute == 0U && !AllToAllGroupFinishCredits( + args, debug, blockIdx, rank, rankSize, invocationId, lane, + groupCount, groupWidth, publishedCreditCount)) { + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } + } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); } @@ -925,10 +1204,11 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, uint32_t groupWidth, uint32_t quietBatch) { - AllToAllGroupKernelImpl( + AllToAllGroupKernelImpl( commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + 0ULL, 0ULL, groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, groupWidth, quietBatch); } @@ -944,10 +1224,53 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_kernel( uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, uint32_t groupWidth, uint32_t quietBatch) { - AllToAllGroupKernelImpl( + AllToAllGroupKernelImpl( commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + 0ULL, 0ULL, + groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, + groupWidth, quietBatch); +} + +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_credit_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, + GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, + GM_ADDR groupTraceGM, uint32_t traceIteration, + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) +{ + AllToAllGroupKernelImpl( + commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + creditOffset0, creditOffset1, + groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, + groupWidth, quietBatch); +} + +extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_credit_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, + GM_ADDR registeredMemoryGM, GM_ADDR debugGM, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, + GM_ADDR groupTraceGM, uint32_t traceIteration, + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) +{ + AllToAllGroupKernelImpl( + commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + creditOffset0, creditOffset1, groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, groupWidth, quietBatch); } @@ -959,23 +1282,40 @@ void launch_tilexr_udma_all_to_all_group( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTrace, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch) + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) { - if (quietBatch == 1U) { + if (ingressWindow == 0U && quietBatch == 1U) { tilexr_udma_all_to_all_group_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, groupTrace, traceIteration, routeStage, multiChannel, primaryRouteParts, groupWidth, quietBatch); - } else { + } else if (ingressWindow == 0U) { tilexr_udma_all_to_all_group_batch_kernel<<>>( commArgs, input, output, registeredMemory, debug, invocationId, elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, groupTrace, traceIteration, routeStage, multiChannel, primaryRouteParts, groupWidth, quietBatch); + } else if (quietBatch == 1U) { + tilexr_udma_all_to_all_group_credit_kernel<<>>( + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + creditOffset0, creditOffset1, + groupTrace, traceIteration, routeStage, + multiChannel, primaryRouteParts, groupWidth, quietBatch, ingressWindow); + } else { + tilexr_udma_all_to_all_group_batch_credit_kernel<<>>( + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + creditOffset0, creditOffset1, + groupTrace, traceIteration, routeStage, + multiChannel, primaryRouteParts, groupWidth, quietBatch, ingressWindow); } } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 062f77a2..750db316 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -20,11 +20,13 @@ constexpr uint32_t kAllToAllGroupWidth = 16U; constexpr uint32_t kAllToAllGroupHalfWidth = 8U; constexpr uint32_t kAllToAllGroupExperimentalWidth = 4U; constexpr uint32_t kAllToAllGroupMaxQuietBatch = 64U; +constexpr uint32_t kAllToAllGroupMaxIngressWindow = 1U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupRouteSignalStride = 512U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 1024U; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; constexpr uint32_t kAllToAllGroupSendWorkerCount = 32U; +constexpr uint32_t kAllToAllGroupMaxGroupCount = 64U; constexpr uint32_t kAllToAllGroupErrorWordsPerCore = 12U; constexpr uint32_t kAllToAllGroupErrorCoreCount = 64U; constexpr size_t kAllToAllGroupErrorBytes = @@ -35,17 +37,32 @@ constexpr size_t kAllToAllGroupSignalSourceSlots = kAllToAllGroupMaxQuietBatch; constexpr size_t kAllToAllGroupSignalSourceBytes = kAllToAllGroupSignalSourceSlots * sizeof(uint64_t); +constexpr size_t kAllToAllGroupCreditSourceSlots = + static_cast(kAllToAllGroupSendCoreCount) * + kAllToAllGroupMaxGroupCount; +constexpr size_t kAllToAllGroupCreditSourceBytes = + kAllToAllGroupCreditSourceSlots * sizeof(uint64_t); +constexpr size_t kAllToAllGroupCreditRequestSlots = + static_cast(kAllToAllGroupPingPongSlots) * + kAllToAllGroupSendCoreCount; +constexpr size_t kAllToAllGroupCreditRequestStride = 64U; +constexpr size_t kAllToAllGroupCreditRequestBytes = + kAllToAllGroupCreditRequestSlots * kAllToAllGroupCreditRequestStride; constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; -constexpr size_t kAllToAllGroupControlBytes = +constexpr size_t kAllToAllGroupBaseControlBytes = kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes; +constexpr size_t kAllToAllGroupControlBytes = + kAllToAllGroupBaseControlBytes + kAllToAllGroupCreditSourceBytes + + kAllToAllGroupCreditRequestBytes; constexpr size_t kAllToAllGroupMaxPayloadBytes = 16ULL << 30; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 2U * kAllToAllGroupMaxPayloadBytes + 2U * static_cast(kAllToAllGroupMaxRankSize) * kAllToAllGroupSignalSlotBytes + + 2U * static_cast(kAllToAllGroupMaxRankSize) * sizeof(uint64_t) + kAllToAllGroupControlBytes; struct AllToAllGroupPlan { @@ -59,10 +76,16 @@ struct AllToAllGroupPlan { size_t payloadOffset[kAllToAllGroupPingPongSlots] = {0, 0}; size_t signalPlaneBytes = 0; size_t signalOffset[kAllToAllGroupPingPongSlots] = {0, 0}; + size_t creditPlaneBytes = 0; + size_t creditOffset[kAllToAllGroupPingPongSlots] = {0, 0}; size_t controlOffset = 0; - size_t controlBytes = kAllToAllGroupControlBytes; + size_t controlBytes = kAllToAllGroupBaseControlBytes; size_t signalSourceOffset = 0; size_t signalSourceBytes = kAllToAllGroupSignalSourceBytes; + size_t creditSourceOffset = 0; + size_t creditSourceBytes = 0; + size_t creditRequestOffset = 0; + size_t creditRequestBytes = 0; size_t registeredBytes = 0; }; @@ -78,6 +101,11 @@ inline bool AllToAllGroupValidQuietBatch(uint32_t quietBatch) (quietBatch & (quietBatch - 1U)) == 0U; } +inline bool AllToAllGroupValidIngressWindow(uint32_t ingressWindow) +{ + return ingressWindow <= kAllToAllGroupMaxIngressWindow; +} + enum class AllToAllGroupChannelMode : uint32_t { kAuto = 0U, kSingle = 1U, @@ -107,6 +135,25 @@ inline size_t AllToAllGroupSignalByteOffset(uint32_t sourceRank, uint32_t route) static_cast(route) * kAllToAllGroupRouteSignalStride; } +inline size_t AllToAllGroupCreditByteOffset(uint32_t destinationRank) +{ + return static_cast(destinationRank) * sizeof(uint64_t); +} + +inline size_t AllToAllGroupCreditSourceByteOffset( + uint32_t lane, uint32_t completedGroup) +{ + return (static_cast(lane) * kAllToAllGroupMaxGroupCount + + completedGroup) * sizeof(uint64_t); +} + +inline size_t AllToAllGroupCreditRequestByteOffset( + uint32_t slot, uint32_t lane) +{ + return (static_cast(slot) * kAllToAllGroupSendCoreCount + lane) * + kAllToAllGroupCreditRequestStride; +} + inline bool AllToAllGroupValidRankSize(int rankSize) { return rankSize >= kAllToAllGroupMinRankSize && @@ -177,6 +224,23 @@ inline int32_t AllToAllGroupPeer( (rank - distance + rankSize) % rankSize; } +inline int32_t AllToAllGroupNextCreditPeer( + int rank, int rankSize, uint32_t completedGroup, uint32_t lane, + uint32_t groupWidth = kAllToAllGroupWidth) +{ + const uint32_t groupCount = AllToAllGroupCount(rankSize, groupWidth); + if (groupCount == 0U || completedGroup + 1U >= groupCount) { + return -1; + } + return AllToAllGroupPeer( + rank, rankSize, completedGroup + 1U, lane, groupWidth); +} + +inline bool AllToAllGroupCreditOwner(uint32_t copyoutWorker) +{ + return copyoutWorker < kAllToAllGroupSendCoreCount; +} + inline uint64_t AllToAllGroupToken( uint32_t invocationId, uint32_t group, uint32_t pass) { @@ -187,6 +251,12 @@ inline uint64_t AllToAllGroupToken( (static_cast(pass) + 1ULL); } +inline uint64_t AllToAllGroupCreditToken( + uint32_t invocationId, uint32_t group) +{ + return AllToAllGroupToken(invocationId, group, 0U); +} + inline bool AllToAllGroupCheckedAdd(size_t lhs, size_t rhs, size_t& result) { if (rhs > std::numeric_limits::max() - lhs) { @@ -217,11 +287,14 @@ inline bool AllToAllGroupAlignUp(size_t value, size_t& result) inline AllToAllGroupPlan PlanAllToAllGroup( int rankSize, int32_t elementsPerPeer, int32_t chunkElements, - uint32_t groupWidth = kAllToAllGroupWidth) + uint32_t groupWidth = kAllToAllGroupWidth, + uint32_t ingressWindow = 0U) { AllToAllGroupPlan plan {}; if (!AllToAllGroupValidRankSize(rankSize) || !AllToAllGroupValidWidth(groupWidth) || + !AllToAllGroupValidIngressWindow(ingressWindow) || + (ingressWindow != 0U && groupWidth != kAllToAllGroupWidth) || elementsPerPeer <= 0 || chunkElements <= 0) { return plan; } @@ -245,6 +318,10 @@ inline AllToAllGroupPlan PlanAllToAllGroup( static_cast(rankSize), kAllToAllGroupSignalSlotBytes, plan.signalPlaneBytes)) { return AllToAllGroupPlan {}; } + if (ingressWindow != 0U && !AllToAllGroupCheckedMul( + static_cast(rankSize), sizeof(uint64_t), plan.creditPlaneBytes)) { + return AllToAllGroupPlan {}; + } if (plan.payloadPlaneBytes > kAllToAllGroupMaxPayloadBytes) { return AllToAllGroupPlan {}; } @@ -259,13 +336,34 @@ inline AllToAllGroupPlan PlanAllToAllGroup( !AllToAllGroupAlignUp(cursor, plan.signalOffset[0]) || !AllToAllGroupCheckedAdd(plan.signalOffset[0], plan.signalPlaneBytes, cursor) || !AllToAllGroupAlignUp(cursor, plan.signalOffset[1]) || - !AllToAllGroupCheckedAdd(plan.signalOffset[1], plan.signalPlaneBytes, cursor) || - !AllToAllGroupAlignUp(cursor, plan.controlOffset) || + !AllToAllGroupCheckedAdd(plan.signalOffset[1], plan.signalPlaneBytes, cursor)) { + return AllToAllGroupPlan {}; + } + if (ingressWindow != 0U) { + plan.creditSourceBytes = kAllToAllGroupCreditSourceBytes; + plan.creditRequestBytes = kAllToAllGroupCreditRequestBytes; + plan.controlBytes = kAllToAllGroupControlBytes; + if (!AllToAllGroupAlignUp(cursor, plan.creditOffset[0]) || + !AllToAllGroupCheckedAdd( + plan.creditOffset[0], plan.creditPlaneBytes, cursor) || + !AllToAllGroupAlignUp(cursor, plan.creditOffset[1]) || + !AllToAllGroupCheckedAdd( + plan.creditOffset[1], plan.creditPlaneBytes, cursor)) { + return AllToAllGroupPlan {}; + } + } + if (!AllToAllGroupAlignUp(cursor, plan.controlOffset) || !AllToAllGroupCheckedAdd(plan.controlOffset, plan.controlBytes, cursor) || !AllToAllGroupAlignUp(cursor, plan.registeredBytes)) { return AllToAllGroupPlan {}; } plan.signalSourceOffset = plan.controlOffset + kAllToAllGroupErrorBytes; + if (ingressWindow != 0U) { + plan.creditSourceOffset = + plan.signalSourceOffset + kAllToAllGroupSignalSourceBytes; + plan.creditRequestOffset = + plan.creditSourceOffset + kAllToAllGroupCreditSourceBytes; + } if (plan.registeredBytes > kAllToAllGroupMaxRegisteredBytes) { return AllToAllGroupPlan {}; } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index 98884a62..ac1b5691 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -14,12 +14,12 @@ namespace TileXR { namespace Demo { constexpr uint32_t kAllToAllGroupTraceMagic = 0x47545243U; // "GTRC" -constexpr uint32_t kAllToAllGroupTraceVersion = 1U; +constexpr uint32_t kAllToAllGroupTraceVersion = 2U; constexpr size_t kAllToAllGroupTraceBytes = 128ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; -constexpr uint32_t kAllToAllGroupTracePhaseCount = 5U; +constexpr uint32_t kAllToAllGroupTracePhaseCount = 6U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; constexpr size_t kAllToAllGroupTraceCacheLineBytes = 128U; @@ -30,6 +30,7 @@ enum AllToAllGroupTracePhase : uint32_t { kAllToAllGroupTraceSendQuiet = 2U, kAllToAllGroupTraceReceiveWait = 3U, kAllToAllGroupTraceReceiveCopy = 4U, + kAllToAllGroupTraceCreditWait = 5U, }; struct AllToAllGroupTraceSpan { diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index 40209c41..ce2c2960 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -6,12 +6,14 @@ TRACE_MAGIC = 0x47545243 -TRACE_VERSION = 1 +TRACE_VERSION = 2 +LEGACY_TRACE_VERSION = 1 TRACE_BYTES = 128 * 1024 * 1024 HEADER_BYTES = 4096 MAX_ITERATIONS = 50 MAX_CORES = 64 PHASE_COUNT = 5 +CURRENT_PHASE_COUNT = 6 SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 TASK_FORMAT = " TRACE_BYTES: raise ValueError(f"trace capacity exceeded in {path}: required={required}") with path.open("rb") as stream: @@ -196,7 +208,7 @@ def build_chrome_trace(rank_traces): {"iteration": iteration, "role": role}, offset_us)) for group in range(header["group_count"]): for pass_index in range(header["pass_count"]): - for phase in range(PHASE_COUNT): + for phase in range(header["phase_count"]): label = ( f"task rank={rank} iter={iteration} core={core} " f"group={group} pass={pass_index} phase={phase}") @@ -204,7 +216,8 @@ def build_chrome_trace(rank_traces): data, task_span_offset( iteration, core, group, pass_index, phase, - header["group_count"], header["pass_count"]), + header["group_count"], header["pass_count"], + header["phase_count"]), label, ) if task is None: diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 1cb39f16..ccb27fda 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -62,9 +62,10 @@ extern void launch_tilexr_udma_all_to_all_group( uint32_t passCount, uint32_t groupCount, uint64_t payloadOffset0, uint64_t payloadOffset1, uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTrace, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch); + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -936,10 +937,27 @@ bool RunGroupedAllToAll( return false; } const uint32_t quietBatch = static_cast(quietBatchValue); + const int ingressWindowValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW", 0); + if (ingressWindowValue < 0 || !TileXR::Demo::AllToAllGroupValidIngressWindow( + static_cast(ingressWindowValue))) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW" + << " must be 0 or 1, got " << ingressWindowValue << std::endl; + return false; + } + const uint32_t ingressWindow = static_cast(ingressWindowValue); + if (ingressWindow != 0U && groupWidth != TileXR::Demo::kAllToAllGroupWidth) { + std::cerr << "[rank " << rank + << "] ERROR: grouped ingress credit currently requires groupWidth=16" + << " groupWidth=" << groupWidth << std::endl; + return false; + } const int32_t requestedChunkElements = std::max( 1, GetEnvInt("TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS", elementsPerPeer)); const auto plan = TileXR::Demo::PlanAllToAllGroup( - rankSize, elementsPerPeer, requestedChunkElements, groupWidth); + rankSize, elementsPerPeer, requestedChunkElements, groupWidth, + ingressWindow); if (!plan.valid) { std::cerr << "[rank " << rank << "] ERROR: invalid grouped alltoall plan" << " rankSize=" << rankSize @@ -947,6 +965,12 @@ bool RunGroupedAllToAll( << " chunkElements=" << requestedChunkElements << std::endl; return false; } + if (ingressWindow != 0U && plan.passCount != 1U) { + std::cerr << "[rank " << rank + << "] ERROR: grouped ingress credit currently requires single pass" + << " passCount=" << plan.passCount << std::endl; + return false; + } const int channelModeValue = GetEnvInt( "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE", 0); if (channelModeValue < 0 || !TileXR::Demo::AllToAllGroupValidChannelMode( @@ -1179,6 +1203,11 @@ bool RunGroupedAllToAll( " signalPlaneBytes=" + std::to_string(plan.signalPlaneBytes) + " signalOffset0=" + std::to_string(plan.signalOffset[0]) + " signalOffset1=" + std::to_string(plan.signalOffset[1]) + + " creditPlaneBytes=" + std::to_string(plan.creditPlaneBytes) + + " creditOffset0=" + std::to_string(plan.creditOffset[0]) + + " creditOffset1=" + std::to_string(plan.creditOffset[1]) + + " creditRequestOffset=" + std::to_string(plan.creditRequestOffset) + + " creditRequestBytes=" + std::to_string(plan.creditRequestBytes) + " controlOffset=" + std::to_string(plan.controlOffset) + " regionCount=" + std::to_string(useMultiRegion ? groupedRegionCount : 1U) + " groupWidth=" + std::to_string(plan.groupWidth) + @@ -1194,6 +1223,7 @@ bool RunGroupedAllToAll( " blockDim=" + std::to_string(groupBlockDim) + " useSecondaryRoute=" + std::to_string(useSecondaryRouteValue) + " quietBatch=" + std::to_string(quietBatch) + + " ingressWindow=" + std::to_string(ingressWindow) + " routeStages=" + std::to_string(routeStagesValue)); if (routeStages && @@ -1213,10 +1243,13 @@ bool RunGroupedAllToAll( plan.passCount, plan.groupCount, plan.payloadOffset[0], plan.payloadOffset[1], plan.signalOffset[0], plan.signalOffset[1], + plan.creditOffset[0], plan.creditOffset[1], reinterpret_cast(trace), traceIteration, static_cast(routeStage), multiChannel ? 1U : 0U, primaryRouteParts, - groupWidth, quietBatch); + groupWidth, quietBatch, + routeStage == TileXR::Demo::AllToAllGroupRouteStage::kCombined ? + ingressWindow : 0U); }; double totalUs = 0.0; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 0489fb50..168238ff 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -142,6 +142,9 @@ void TestPlan() CHECK_EQ(TileXR::Demo::AllToAllGroupSignalByteOffset(3U, 1U), 3584ULL); CHECK_EQ(plan.signalOffset[0] >= plan.payloadOffset[1] + plan.payloadPlaneBytes, true); CHECK_EQ(plan.signalOffset[1] >= plan.signalOffset[0] + plan.signalPlaneBytes, true); + CHECK_EQ(plan.creditPlaneBytes, 0ULL); + CHECK_EQ(plan.creditOffset[0], 0ULL); + CHECK_EQ(plan.creditOffset[1], 0ULL); CHECK_EQ(plan.controlOffset >= plan.signalOffset[1] + plan.signalPlaneBytes, true); CHECK_EQ(plan.signalSourceOffset, plan.controlOffset + TileXR::Demo::kAllToAllGroupErrorBytes); @@ -149,9 +152,66 @@ void TestPlan() TileXR::Demo::kAllToAllGroupSignalSourceBytes); CHECK_EQ(plan.signalSourceOffset + plan.signalSourceBytes <= plan.controlOffset + plan.controlBytes, true); + CHECK_EQ(plan.creditSourceOffset, 0ULL); + CHECK_EQ(plan.creditSourceBytes, 0ULL); + CHECK_EQ(plan.creditRequestOffset, 0ULL); + CHECK_EQ(plan.creditRequestBytes, 0ULL); + const auto legacyAlign = [](size_t value) { + return (value + TileXR::Demo::kAllToAllGroupAlignment - 1U) & + ~(TileXR::Demo::kAllToAllGroupAlignment - 1U); + }; + const size_t legacyPayloadOffset1 = legacyAlign(plan.payloadPlaneBytes); + const size_t legacySignalOffset0 = legacyAlign( + legacyPayloadOffset1 + plan.payloadPlaneBytes); + const size_t legacySignalOffset1 = legacyAlign( + legacySignalOffset0 + plan.signalPlaneBytes); + const size_t legacyControlOffset = legacyAlign( + legacySignalOffset1 + plan.signalPlaneBytes); + const size_t legacyRegisteredBytes = legacyAlign( + legacyControlOffset + TileXR::Demo::kAllToAllGroupBaseControlBytes); + CHECK_EQ(plan.payloadOffset[1], legacyPayloadOffset1); + CHECK_EQ(plan.signalOffset[0], legacySignalOffset0); + CHECK_EQ(plan.signalOffset[1], legacySignalOffset1); + CHECK_EQ(plan.controlOffset, legacyControlOffset); + CHECK_EQ(plan.registeredBytes, legacyRegisteredBytes); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditByteOffset(3U), 24ULL); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditSourceByteOffset(1U, 2U), + (TileXR::Demo::kAllToAllGroupMaxGroupCount + 2ULL) * sizeof(uint64_t)); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditRequestByteOffset(1U, 2U), + (TileXR::Demo::kAllToAllGroupSendCoreCount + 2ULL) * + TileXR::Demo::kAllToAllGroupCreditRequestStride); CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); CHECK_EQ(TileXR::Demo::kAllToAllGroupMaxPayloadBytes, 16ULL << 30); + const auto ingressPlan = TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, elementsPerPeer, + TileXR::Demo::kAllToAllGroupWidth, 1U); + CHECK_EQ(ingressPlan.valid, true); + CHECK_EQ(ingressPlan.creditPlaneBytes, + static_cast(rankSize) * sizeof(uint64_t)); + CHECK_EQ(ingressPlan.creditOffset[0] >= + ingressPlan.signalOffset[1] + ingressPlan.signalPlaneBytes, true); + CHECK_EQ(ingressPlan.creditOffset[1] >= + ingressPlan.creditOffset[0] + ingressPlan.creditPlaneBytes, true); + CHECK_EQ(ingressPlan.controlOffset >= + ingressPlan.creditOffset[1] + ingressPlan.creditPlaneBytes, true); + CHECK_EQ(ingressPlan.creditSourceOffset, + ingressPlan.signalSourceOffset + + TileXR::Demo::kAllToAllGroupSignalSourceBytes); + CHECK_EQ(ingressPlan.creditSourceOffset + ingressPlan.creditSourceBytes <= + ingressPlan.controlOffset + ingressPlan.controlBytes, true); + CHECK_EQ(ingressPlan.creditRequestOffset, + ingressPlan.creditSourceOffset + + TileXR::Demo::kAllToAllGroupCreditSourceBytes); + CHECK_EQ(ingressPlan.creditRequestBytes, + TileXR::Demo::kAllToAllGroupCreditRequestBytes); + CHECK_EQ(ingressPlan.creditRequestOffset + ingressPlan.creditRequestBytes <= + ingressPlan.controlOffset + ingressPlan.controlBytes, true); + CHECK_EQ(ingressPlan.registeredBytes > plan.registeredBytes, true); + CHECK_EQ(TileXR::Demo::PlanAllToAllGroup( + rankSize, elementsPerPeer, elementsPerPeer, + TileXR::Demo::kAllToAllGroupExperimentalWidth, 1U).valid, false); + const auto chunked = TileXR::Demo::PlanAllToAllGroup( rankSize, elementsPerPeer, elementsPerPeer / 4); CHECK_EQ(chunked.valid, true); @@ -255,8 +315,9 @@ void TestScalePlanAndTraceCapacity() CHECK_EQ(TileXR::Demo::kAllToAllGroupTraceBytes, 128ULL * 1024ULL * 1024ULL); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 3U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 4U), true); - CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 6U), false); } @@ -271,6 +332,69 @@ void TestTokens() CHECK_EQ(TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U) > token49, true); CHECK_EQ(TileXR::Demo::AllToAllGroupToken(49U, 1U, 1U) > TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditToken(49U, 1U), + TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U)); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditToken(50U, 1U) > + TileXR::Demo::AllToAllGroupCreditToken(49U, 1U), true); +} + +void TestIngressCreditPolicy() +{ + CHECK_EQ(TileXR::Demo::AllToAllGroupValidIngressWindow(0U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidIngressWindow(1U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidIngressWindow(2U), false); + + constexpr int rankSize = 256; + constexpr uint32_t groupCount = 16U; + for (int rank = 0; rank < rankSize; ++rank) { + for (uint32_t group = 0U; group < groupCount; ++group) { + for (uint32_t lane = 0U; + lane < TileXR::Demo::kAllToAllGroupWidth; ++lane) { + const int32_t nextPeer = TileXR::Demo::AllToAllGroupNextCreditPeer( + rank, rankSize, group, lane); + const int32_t expected = group + 1U < groupCount ? + TileXR::Demo::AllToAllGroupPeer( + rank, rankSize, group + 1U, lane) : -1; + CHECK_EQ(nextPeer, expected); + if (nextPeer >= 0) { + uint32_t senderLane = lane ^ + TileXR::Demo::kAllToAllGroupHalfWidth; + if (TileXR::Demo::AllToAllGroupPeer( + nextPeer, rankSize, group + 1U, + senderLane) < 0) { + senderLane = lane; + } + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer( + nextPeer, rankSize, group + 1U, senderLane), rank); + } + } + } + } + + CHECK_EQ(TileXR::Demo::AllToAllGroupNextCreditPeer(0, 64, 2U, 0U), 25); + CHECK_EQ(TileXR::Demo::AllToAllGroupNextCreditPeer(0, 64, 3U, 0U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupNextCreditPeer(0, 64, 2U, 15U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(0U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(15U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(16U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(31U), false); + for (uint32_t lane = 0U; + lane < TileXR::Demo::kAllToAllGroupSendCoreCount; ++lane) { + const uint32_t primaryWorker = lane; + const uint32_t secondaryWorker = + lane + TileXR::Demo::kAllToAllGroupSendCoreCount; + CHECK_EQ(primaryWorker % TileXR::Demo::kAllToAllGroupSendCoreCount, + secondaryWorker % TileXR::Demo::kAllToAllGroupSendCoreCount); + const int32_t primaryPeer = TileXR::Demo::AllToAllGroupPeer( + 0, rankSize, 3U, + primaryWorker % TileXR::Demo::kAllToAllGroupSendCoreCount); + const int32_t secondaryPeer = TileXR::Demo::AllToAllGroupPeer( + 0, rankSize, 3U, + secondaryWorker % TileXR::Demo::kAllToAllGroupSendCoreCount); + CHECK_EQ(primaryPeer, secondaryPeer); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditByteOffset(primaryPeer), + TileXR::Demo::AllToAllGroupCreditByteOffset(secondaryPeer)); + } } void TestDualRoutePeerPolicy() @@ -507,7 +631,7 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, request.peer, request.qpIdx)"); CHECK_CONTAINS(kernel, "state.pendingCount != quietBatch"); - CHECK_CONTAINS(kernel, "template "); + CHECK_CONTAINS(kernel, "template "); CHECK_CONTAINS(kernel, "struct AllToAllGroupQuietState"); CHECK_CONTAINS(kernel, "AllToAllGroupQuietState quietState"); CHECK_CONTAINS(kernel, "AllToAllGroupCompleteQuiet("); @@ -515,8 +639,40 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, peer, selectedQp)"); CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_batch_kernel"); + CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_credit_kernel"); + CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_batch_credit_kernel"); + CHECK_CONTAINS(kernel, "AllToAllGroupKernelImpl"); + CHECK_CONTAINS(kernel, "AllToAllGroupKernelImpl"); + CHECK_CONTAINS(kernel, "AllToAllGroupKernelImpl"); + CHECK_CONTAINS(kernel, "AllToAllGroupKernelImpl"); + CHECK_CONTAINS(kernel, "if constexpr (IngressCredit)"); CHECK_CONTAINS(kernel, "if (quietBatch == 1U)"); CHECK_CONTAINS(kernel, "uint32_t groupWidth, uint32_t quietBatch"); + CHECK_CONTAINS(kernel, "uint64_t creditOffset0, uint64_t creditOffset1"); + CHECK_CONTAINS(kernel, "uint32_t ingressWindow"); + CHECK_CONTAINS(kernel, "AllToAllGroupPublishNextCredit"); + CHECK_CONTAINS(kernel, "AllToAllGroupPublishCreditRequest"); + CHECK_CONTAINS(kernel, "AllToAllGroupCreditRequest(debug, slot, lane)"); + CHECK_CONTAINS(kernel, "AllToAllGroupFinishCredits"); + CHECK_CONTAINS(kernel, "AllToAllGroupCreditOwnerDevice(worker)"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_QUIET"); + CHECK_CONTAINS(kernel, + "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_REQUEST_WAIT"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceCreditWait"); + CHECK_CONTAINS(kernel, "completedQueueKeys"); + CHECK_CONTAINS(kernel, "publishedCreditCount"); + CHECK_CONTAINS(kernel, "TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION"); + const size_t publishCreditBegin = kernel.find( + "__aicore__ inline void AllToAllGroupPublishNextCredit"); + const size_t publishCreditEnd = kernel.find( + "__aicore__ inline bool AllToAllGroupFinishCredits", publishCreditBegin); + const std::string publishCredit = publishCreditBegin == std::string::npos ? + std::string() : kernel.substr(publishCreditBegin, + publishCreditEnd == std::string::npos ? std::string::npos : + publishCreditEnd - publishCreditBegin); + CHECK_CONTAINS(publishCredit, "UDMAPutNbiOnQpWithFlag"); + CHECK_NOT_CONTAINS(publishCredit, "UDMAQuiet"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsReceiveDevice(routeStage)"); @@ -545,6 +701,9 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_QUIET_BATCH"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW"); + CHECK_CONTAINS(demo, "grouped ingress credit currently requires single pass"); + CHECK_CONTAINS(demo, "grouped ingress credit currently requires groupWidth=16"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS"); CHECK_CONTAINS(demo, "kAllToAllGroupSendWorkerCount"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); @@ -598,6 +757,7 @@ int main() TestChannelPolicy(); TestScalePlanAndTraceCapacity(); TestTokens(); + TestIngressCreditPolicy(); TestDualRoutePeerPolicy(); TestDualRouteQpWeights(); TestRouteStages(); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index f805a7f7..94017f3e 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -31,17 +31,21 @@ def test_assigns_each_core_disjoint_cache_lines(self): def make_trace( self, path, *, rank=0, magic=None, iteration_count=1, group_count=1, pass_count=1, core_count=64, + version=None, phase_count=None, ): + version = MODULE.TRACE_VERSION if version is None else version + phase_count = ( + MODULE.CURRENT_PHASE_COUNT if phase_count is None else phase_count) header = struct.pack( MODULE.HEADER_FORMAT, MODULE.TRACE_MAGIC if magic is None else magic, - MODULE.TRACE_VERSION, + version, rank, iteration_count, group_count, pass_count, core_count, - MODULE.PHASE_COUNT, + phase_count, 1000, MODULE.TRACE_BYTES, MODULE.HEADER_BYTES, @@ -53,6 +57,7 @@ def make_trace( (0, 0, 0, 2, 1300, 1400, 1, 3), (16, 0, 0, 3, 1400, 1500, 1, MODULE.NO_QP), (16, 0, 0, 4, 1500, 1600, 1, MODULE.NO_QP), + (0, 0, 0, 5, 1050, 1100, 1, MODULE.NO_QP), ) with path.open("wb") as stream: stream.truncate(MODULE.TRACE_BYTES) @@ -62,8 +67,11 @@ def make_trace( stream.seek(MODULE.kernel_span_offset(0, core)) stream.write(struct.pack("= phase_count: + continue stream.seek(MODULE.task_span_offset( - 0, core, group, pass_index, phase, group_count, pass_count)) + 0, core, group, pass_index, phase, group_count, pass_count, + phase_count)) stream.write(struct.pack( MODULE.TASK_FORMAT, begin, end, peer, qp)) @@ -95,7 +103,7 @@ def test_converts_all_grouped_pipeline_phases(self): {event["name"] for event in complete}, { "kernel", "self-copy", "send-put-signal", "send-quiet", - "receive-wait", "receive-copy", + "receive-wait", "receive-copy", "credit-wait", }, ) send = next(event for event in complete if event["name"] == "send-put-signal") @@ -187,6 +195,23 @@ def test_rejects_half_written_span(self): with self.assertRaisesRegex(ValueError, "incomplete"): MODULE.build_chrome_trace([MODULE.read_rank_trace(path)]) + def test_reads_legacy_five_phase_trace(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "legacy.bin" + self.make_trace( + path, version=MODULE.LEGACY_TRACE_VERSION, + phase_count=MODULE.PHASE_COUNT) + + rank_trace = MODULE.read_rank_trace(path) + trace = MODULE.build_chrome_trace([rank_trace]) + + self.assertEqual(rank_trace["header"]["phase_count"], 5) + names = { + event["name"] for event in trace["traceEvents"] + if event.get("ph") == "X" + } + self.assertNotIn("credit-wait", names) + def test_main_writes_json_without_dumps(self): with tempfile.TemporaryDirectory() as directory: source = Path(directory) / "rank0.bin" From db811f6ded7173bffb0a789ba851e30cde5bd31f Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 18:58:13 +0800 Subject: [PATCH 141/163] docs(udma): record 64-rank ingress credit results --- ...02-grouped-alltoall-ingress-credit-design.md | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md index 088895f3..4eb64bc2 100644 --- a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md +++ b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md @@ -2,7 +2,7 @@ ## Status -Implemented and validated on 4x8 hardware on 2026-08-02. The experimental +Implemented and validated through 8x8 hardware on 2026-08-02. The experimental feature remains disabled by default. The strict global admission bound still requires 256P all-rank trace or admission-counter validation. @@ -157,6 +157,21 @@ This proves the request/credit protocol progresses across repeated invocations and does not regress 1 GiB throughput in the 4x8 environment. It does not prove the 16-source global bound because 4x8 has only two peer groups. +The 8x8 A10 run used the fixed host order `226, 223, 220, 217, 198, 195, +192, 189`, single pass, shared QP, trace disabled, and warmup-5/repeat-50. All +64 ranks passed data validation: + +```text +1 KiB/rank single: window0 P50 73.067 us, window1 P50 105.930 us +1 GiB/rank multi: window0 P50 6516.335 us, window1 P50 6336.370 us +``` + +At 64 ranks, ingress credit reduced the 1 GiB P50 by 179.965 us (2.76%), from +approximately 153.46 GiB/s to 157.82 GiB/s. The 1 KiB case paid 32.863 us of +additional fixed control latency. This run exercises four peer groups and +therefore provides stronger repeated credit-chain evidence than 4x8, but it +still does not prove the 256P global admission bound. + ## Residual Risk Rank0-only trace cannot prove the global source-card bound. Hardware validation From f61b36b56c8aaccd3207fac7e00986810f5d4945 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 20:07:35 +0800 Subject: [PATCH 142/163] perf(udma): publish ingress credits through dedicated IPC --- ...-grouped-alltoall-ingress-credit-design.md | 64 +++--- src/comm/tilexr_comm.cpp | 107 ++++++++++ src/comm/tilexr_comm.h | 6 + src/include/comm_args.h | 4 + tests/comm/unit/test_tilexr_source_guards.cpp | 17 ++ .../tilexr_udma_alltoall_group_kernel.cpp | 187 ++---------------- .../demo/tilexr_udma_alltoall_group_layout.h | 67 +------ tests/udma/demo/tilexr_udma_demo.cpp | 26 ++- ...test_tilexr_udma_alltoall_group_layout.cpp | 53 ++--- 9 files changed, 221 insertions(+), 310 deletions(-) diff --git a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md index 4eb64bc2..b3286cb9 100644 --- a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md +++ b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md @@ -2,9 +2,10 @@ ## Status -Implemented and validated through 8x8 hardware on 2026-08-02. The experimental -feature remains disabled by default. The strict global admission bound still -requires 256P all-rank trace or admission-counter validation. +The original request-plus-UDMA-credit implementation was validated through 8x8 +hardware on 2026-08-02. It has since been replaced by direct receive-core +publication through a dedicated credit IPC allocation and requires renewed +hardware validation. The experimental feature remains disabled by default. ## Goal @@ -45,10 +46,9 @@ For destination rank `D`, lane `L`, and group `g`: 1. The designated receive owner waits until all payload routes expected from `peer(D, g, L)` have published their data-ready tokens. -2. Before receive-copy, the owner publishes a local request for lane `L`. -3. Primary send core `L` observes that request and posts one credit token to - `peer(D, g + 1, L)`, when that peer exists. -4. The next source observes the credit in its local registered credit plane and +2. Before receive-copy, the owner directly writes one credit token through the + dedicated IPC mapping of `peer(D, g + 1, L)`, when that peer exists. +3. The next source observes the credit in its local dedicated credit buffer and may post both primary and secondary payload routes to `D`. At most one source per destination lane can therefore be admitted. With 16 @@ -59,12 +59,13 @@ dependency, so the protocol does not introduce a cyclic startup wait. ## Credit Storage And Tokens -Credits use a separate registered-memory plane. A credit received by source -rank `S` is indexed by destination rank `D`, because `S` communicates with `D` -only once per invocation: +Each rank allocates a dedicated 1 MiB IPC buffer when +`TILEXR_ENABLE_CREDIT_IPC=1`. A credit received by source rank `S` is indexed by +destination rank `D`; each entry occupies 512 bytes and two fixed 512 KiB +ping-pong planes support 1024 ranks: ```text -credit[pingPongSlot][destinationRank] +credit[pingPongSlot][destinationRank * 512] ``` The expected value encodes at least the invocation and destination group. The @@ -79,35 +80,24 @@ Credits control source-rank admission, not route admission. With 32 copyout workers, workers 0 through 15 and 16 through 31 can wait on the same peer signal while copying different payload slices. Only the first slice, -kernel cores 32 through 47, owns the local credit request. Cores 48 through 63 -never request credits. +kernel cores 32 through 47, owns direct credit publication. Cores 48 through 63 +never publish credits. -The owner publishes its request after both expected primary and secondary data +The owner publishes its credit after both expected primary and secondary data tokens arrive and before MTE receive-copy begins. This avoids placing copyout latency on the send admission path. -Each `(pingPongSlot, lane)` request occupies its own 64-byte cache line. Packing -multiple lanes into one line is incorrect: independent receive cores clean and -invalidate the line concurrently, so one core can overwrite another lane's -request with stale cache-line contents. - ## Credit Submission -Credit publication is an NBI 8-byte UDMA write issued only by the corresponding -primary send core. Receive cores do not issue UDMA operations because they can -share a payload SQ with a send core, while SQ head and WQE count updates require -a single producer. - -The credit path must not perform a quiet for every credit. Token source storage -remains immutable until completion through distinct lane/group slots. At the -end of the invocation, the primary send core reclaims credit completions once -per unique underlying shared SQ, identified by its WQE-count address. This -preserves single-producer SQ ownership and avoids duplicate quiet operations -when multiple peer/QP views refer to the same shared queue. +Credit publication is a direct 8-byte GM store by the receive owner through +`creditMems[nextSource]`. It does not consume a UDMA WQE, QP, completion, or +quiet, and it does not add another producer to a shared UDMA SQ. Both send +routes poll the same local dedicated credit address through MTE. -If the available UDMA interface cannot safely keep credit source data alive -without per-credit quiet, implementation pauses for a revised design rather -than adding a per-credit quiet to the critical path. +The dedicated allocation is independent of the existing optional communication +IPC buffer, so grouped ingress credit can run with `TILEXR_ENABLE_IPC=0`. It +adds one IPC mapping per peer and process, but only 1 MiB of device memory per +rank. Host code rejects ingress-credit execution if any mapping is missing. ## Configuration @@ -121,6 +111,9 @@ TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW=0|1 - `1`: one admitted source per lane, at most 16 payload source ranks per destination. +`window=1` additionally requires `TILEXR_ENABLE_CREDIT_IPC=1` before +communicator initialization. + A later `window=2` extension may trade a 32-source bound for more tolerance of credit latency, but it is outside the first implementation. @@ -153,7 +146,7 @@ validation on all 32 ranks: 1 GiB/rank multi: window0 P50 5329.210 us, window1 P50 5317.925 us ``` -This proves the request/credit protocol progresses across repeated invocations +This proved the previous request/UDMA-credit protocol progressed across repeated invocations and does not regress 1 GiB throughput in the 4x8 environment. It does not prove the 16-source global bound because 4x8 has only two peer groups. @@ -166,7 +159,8 @@ The 8x8 A10 run used the fixed host order `226, 223, 220, 217, 198, 195, 1 GiB/rank multi: window0 P50 6516.335 us, window1 P50 6336.370 us ``` -At 64 ranks, ingress credit reduced the 1 GiB P50 by 179.965 us (2.76%), from +Under the previous request/UDMA-credit implementation, ingress credit reduced +the 1 GiB P50 by 179.965 us (2.76%), from approximately 153.46 GiB/s to 157.82 GiB/s. The 1 KiB case paid 32.863 us of additional fixed control latency. This run exercises four peer groups and therefore provides stronger repeated credit-chain evidence than 4x8, but it diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 20288c53..8c9d0dd5 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -277,6 +277,7 @@ int TileXRComm::SyncCommArgs() commArgs_.localRankSize = localRankSize_; for (int i = 0; i < rankSize_; ++i) { commArgs_.peerMems[i] = peerMem_[i]; // 这里不会越界，之前有逻辑校验过越界了 + commArgs_.creditMems[i] = creditIpcMem_[i]; } if (isEnableMsprofOp_) { @@ -534,6 +535,20 @@ void TileXRComm::CloseIpcMem() } } +void TileXRComm::CloseCreditIpcMem() +{ + for (int i = 0; i < rankSize_; ++i) { + if (i == rank_ || creditIpcMem_[i] == nullptr) { + continue; + } + int ret = rtIpcCloseMemory(static_cast(creditIpcMem_[i])); + if (ret != RT_ERROR_NONE) { + TILEXR_LOG(WARN) << "Close credit ipc[" << i << "] memory failed! ret: " << ret; + } + creditIpcMem_[i] = nullptr; + } +} + void TileXRComm::FreePeerMem(GM_ADDR &mem) const { if (mem != nullptr) { @@ -588,6 +603,13 @@ int TileXRComm::Init() } else { TILEXR_LOG(INFO) << "TileXR IPC memory disabled by environment"; } + if (IsEnvEnabled("TILEXR_ENABLE_CREDIT_IPC", false)) { + if (InitCreditCommMem() != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "InitCreditCommMem failed!"; + return TILEXR_ERROR_INTERNAL; + } + TILEXR_LOG(INFO) << "Dedicated credit IPC memory initialized"; + } // 新增：初始化 UDMA ret = InitUDMA(); @@ -956,6 +978,86 @@ int TileXRComm::InitCommMem() return TILEXR_SUCCESS; } +int TileXRComm::InitCreditCommMem() +{ + uint32_t pids[TILEXR_MAX_RANK_SIZE] = {0}; + int ret = GetPid(pids); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "GetPid for credit IPC failed! ret: " << ret; + return ret; + } + int64_t sdids[TILEXR_MAX_RANK_SIZE] = {0}; + ret = GetSidId(sdids, rankSize_); + if (ret != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "GetSidId for credit IPC failed! ret: " << ret; + return ret; + } + return InitCreditIpcMem(pids, sdids); +} + +int TileXRComm::InitCreditIpcMem(const uint32_t *pids, const int64_t *sdids) +{ + aclError aclRet = aclrtMalloc( + reinterpret_cast(&creditIpcMem_[rank_]), CREDIT_IPC_BYTES, + (GetChipName() == ChipName::CHIP_310P3) ? + ACL_MEM_MALLOC_HUGE_FIRST_P2P : ACL_MEM_MALLOC_HUGE_FIRST); + if (aclRet != ACL_SUCCESS) { + TILEXR_LOG(ERROR) << "allocate credit IPC memory failed: " << aclRet; + return TILEXR_ERROR_INTERNAL; + } + aclRet = aclrtMemset( + creditIpcMem_[rank_], CREDIT_IPC_BYTES, 0, CREDIT_IPC_BYTES); + if (aclRet != ACL_SUCCESS) { + TILEXR_LOG(ERROR) << "initialize credit IPC memory failed: " << aclRet; + return TILEXR_ERROR_INTERNAL; + } + + char nameBuffer[IPC_NAME_SIZE] = {}; + if (rtIpcSetMemoryName( + creditIpcMem_[rank_], CREDIT_IPC_BYTES, + nameBuffer, IPC_NAME_SIZE) != RT_ERROR_NONE) { + TILEXR_LOG(ERROR) << "set credit IPC memory name failed"; + return TILEXR_ERROR_INTERNAL; + } + string name(nameBuffer); + if (SetIpcPidSdid(name, pids, sdids) != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "set credit IPC pid/sdid failed"; + return TILEXR_ERROR_INTERNAL; + } + char names[TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE] = {}; + name.resize(IPC_NAME_SIZE); + if (GetName(name, names) != TILEXR_SUCCESS) { + TILEXR_LOG(ERROR) << "gather credit IPC memory names failed"; + return TILEXR_ERROR_INTERNAL; + } + return OpenCreditIpcMem(names); +} + +int TileXRComm::OpenCreditIpcMem( + const char names[TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]) +{ + static mutex mut; + lock_guard lock(mut); + for (int i = 0; i < rankSize_; ++i) { + if (i == rank_) { + continue; + } + if (SkipUnusedChannel910B2C(rank_, i, GetChipName())) { + continue; + } + int ret = rtIpcOpenMemory( + reinterpret_cast(&creditIpcMem_[i]), names[i]); + if (ret != RT_ERROR_NONE) { + CloseCreditIpcMem(); + TILEXR_LOG(ERROR) << "rank: " << rank_ << " creditIpcMem: " << i + << " IpcOpenMemory err " << ret; + return TILEXR_ERROR_INTERNAL; + } + } + creditIpcMemInited_ = true; + return TILEXR_SUCCESS; +} + int TileXRComm::OpenIpcMem(const char names[TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]) { static mutex mut; @@ -1042,6 +1144,10 @@ TileXRComm::~TileXRComm() CloseIpcMem(); ipcMemInited_ = false; } + if (creditIpcMemInited_) { + CloseCreditIpcMem(); + creditIpcMemInited_ = false; + } if (socketExchange_) { delete socketExchange_; socketExchange_ = nullptr; @@ -1056,6 +1162,7 @@ TileXRComm::~TileXRComm() } FreeUDMARegistry(); FreePeerMem(peerMem_[rank_]); + FreePeerMem(creditIpcMem_[rank_]); FreePeerMem(commArgsPtr_); if (udmaTransport_ != nullptr) { diff --git a/src/comm/tilexr_comm.h b/src/comm/tilexr_comm.h index 83afd890..066980e7 100644 --- a/src/comm/tilexr_comm.h +++ b/src/comm/tilexr_comm.h @@ -60,12 +60,16 @@ class TileXRComm { int SetMemoryName(std::string &name); int SetIpcPidSdid(std::string &name, const uint32_t *pids, const int64_t *sdids) const; int OpenIpcMem(const char names[TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]); + int InitCreditIpcMem(const uint32_t *pids, const int64_t *sdids); + int OpenCreditIpcMem(const char names[TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]); int GetDev(); int GetDevThread(const std::string &uid = ""); int EnablePeerAccess(); int InitCommMem(); + int InitCreditCommMem(); int InitCommon(); void CloseIpcMem(); + void CloseCreditIpcMem(); void FreePeerMem(GM_ADDR &mem) const; int InitMem(); int GetSidId(int64_t sdids[TILEXR_MAX_RANK_SIZE], int rankSize); @@ -89,6 +93,7 @@ class TileXRComm { std::atomic magic_ {1}; bool inited_ = false; bool ipcMemInited_ = false; + bool creditIpcMemInited_ = false; std::string uid_ = {}; std::vector devList_ = {}; int commDomain_ = {}; @@ -96,6 +101,7 @@ class TileXRComm { // shared ping pong buff，这个地址就是一开始申请在HBM上的，所以host上可以取到，但不能直接修改。 GM_ADDR peerMem_[TILEXR_MAX_RANK_SIZE] = {}; + GM_ADDR creditIpcMem_[TILEXR_MAX_RANK_SIZE] = {}; PhysicalInfo physicalInfo_ = {}; CommArgs commArgs_ = {}; // host侧 GM_ADDR commArgsPtr_ = nullptr; // device侧 diff --git a/src/include/comm_args.h b/src/include/comm_args.h index 861c9b0c..754d0ad0 100644 --- a/src/include/comm_args.h +++ b/src/include/comm_args.h @@ -38,6 +38,9 @@ constexpr int TILEXR_MAX_RANK_SIZE = 1024; // 最大支持的npu卡数 constexpr int RANK_SIZE_TWO = 2; // 可用SIO的规模，以及是否需要跨卡搬运数据核的分界规模 constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // 前2MB作为flag标志位，之后100MB作为数据存储 +constexpr int64_t CREDIT_IPC_STRIDE = 512; +constexpr int64_t CREDIT_IPC_SLOT_BYTES = TILEXR_MAX_RANK_SIZE * CREDIT_IPC_STRIDE; +constexpr int64_t CREDIT_IPC_BYTES = 2 * CREDIT_IPC_SLOT_BYTES; constexpr int64_t SYNC_FLAG_BIT_NUM = 10; // cce 算子在用 constexpr int64_t MEM_DMA_UNIT_INT_NUM = 4; constexpr int64_t EVENT_ID_MASK = 0xFFFFFFFF; @@ -105,6 +108,7 @@ struct CommArgs { int localRankSize = -1; // 此参数是指fullmesh互联的卡数 uint32_t extraFlag = 0; // 32 bit map，具体每一位的含义就在此文件正上方 GM_ADDR peerMems[TILEXR_MAX_RANK_SIZE] = {}; // 传入初始化获得的buff，所有allreduce都是同一个参数 + GM_ADDR creditMems[TILEXR_MAX_RANK_SIZE] = {}; // optional dedicated grouped ingress-credit IPC buffers /** * @param sendCountMatrix 大小是rankSize*rankSize的一维数组 * eg: sendCountMatrix[1] 的数值，对应二维数组的[0][1]，表示 卡0 要给 卡1 发送的数据个数 diff --git a/tests/comm/unit/test_tilexr_source_guards.cpp b/tests/comm/unit/test_tilexr_source_guards.cpp index 8ad23f64..f8408ac8 100644 --- a/tests/comm/unit/test_tilexr_source_guards.cpp +++ b/tests/comm/unit/test_tilexr_source_guards.cpp @@ -102,6 +102,22 @@ void TestCommInitChecksDeviceCommArgsSync() CheckContains(path, text, "ret = InitMem();"); } +void TestDedicatedCreditIpcLifecycle() +{ + const std::string cppPath = "src/comm/tilexr_comm.cpp"; + const std::string headerPath = "src/include/comm_args.h"; + const auto cppText = ReadFile(cppPath); + const auto headerText = ReadFile(headerPath); + + CheckContains(headerPath, headerText, "GM_ADDR creditMems[TILEXR_MAX_RANK_SIZE]"); + CheckContains(headerPath, headerText, "CREDIT_IPC_BYTES"); + CheckContains(cppPath, cppText, "TILEXR_ENABLE_CREDIT_IPC"); + CheckContains(cppPath, cppText, "InitCreditIpcMem(pids, sdids)"); + CheckContains(cppPath, cppText, "rtIpcOpenMemory("); + CheckContains(cppPath, cppText, "CloseCreditIpcMem();"); + CheckContains(cppPath, cppText, "FreePeerMem(creditIpcMem_[rank_]);"); +} + void TestCWrappersDoNotPublishFailedCommunicators() { const std::string path = "src/comm/comm_wrap.cpp"; @@ -195,6 +211,7 @@ int main() { TestOpenSourceTarballsAreNotTracked(); TestCommInitChecksDeviceCommArgsSync(); + TestDedicatedCreditIpcLifecycle(); TestCWrappersDoNotPublishFailedCommunicators(); TestDumpInitCleansFailedAllocations(); TestAscend950UsesHccsTopologyForPeerLinks(); diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index f07f5573..33a3835c 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -28,22 +28,11 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_BYTES = constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES = TILEXR_ALLTOALL_GROUP_SEND_WORKERS * TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH * sizeof(uint64_t); -constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_OFFSET = - TILEXR_ALLTOALL_GROUP_ERROR_BYTES + - TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_BYTES = - TILEXR_ALLTOALL_GROUP_SEND_CORES * - TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT * sizeof(uint64_t); -constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_OFFSET = - TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_OFFSET + - TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_BYTES; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_STRIDE = 64U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT = 4U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_QUIET = 5U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_REQUEST_WAIT = 6U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED = 0U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY = 2U; @@ -371,23 +360,11 @@ __aicore__ inline int32_t AllToAllGroupNextCreditPeerDevice( rank, rankSize, completedGroup + 1U, lane, groupWidth); } -__aicore__ inline uint32_t AllToAllGroupCreditQp( - const __gm__ TileXR::CommArgs* args, int32_t peer) -{ - uint32_t primaryQp = 0U; - uint32_t secondaryQp = 0U; - uint32_t primaryWeight = 0U; - uint32_t secondaryWeight = 0U; - AllToAllGroupSelectRouteQps(args, peer, primaryQp, secondaryQp, - primaryWeight, secondaryWeight); - return primaryQp; -} - __aicore__ inline void AllToAllGroupPublishNextCredit( - const __gm__ TileXR::CommArgs* args, __gm__ int32_t* debug, + const __gm__ TileXR::CommArgs* args, int32_t rank, int32_t rankSize, uint32_t invocationId, uint32_t completedGroup, uint32_t lane, uint32_t groupCount, - uint32_t groupWidth, uint64_t remoteCreditOffset) + uint32_t groupWidth, uint64_t creditOffset) { const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( rank, rankSize, completedGroup, lane, groupCount, groupWidth); @@ -396,87 +373,10 @@ __aicore__ inline void AllToAllGroupPublishNextCredit( } const uint64_t creditToken = AllToAllGroupDeviceToken(invocationId, completedGroup + 1U, 0U); - const uint64_t sourceIndex = - static_cast(lane) * TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT + - completedGroup; - auto creditLocal = reinterpret_cast<__gm__ uint64_t*>( - reinterpret_cast<__gm__ uint8_t*>(debug) + - TILEXR_ALLTOALL_GROUP_CREDIT_SOURCE_OFFSET + - sourceIndex * sizeof(uint64_t)); - *creditLocal = creditToken; - TileXR::UDMACleanCacheLines( - reinterpret_cast<__gm__ uint8_t*>(creditLocal), sizeof(uint64_t)); - TileXR::UDMAPutNbiOnQpWithFlag( - args, nextPeer, AllToAllGroupCreditQp(args, nextPeer), creditLocal, - remoteCreditOffset + static_cast(rank) * sizeof(uint64_t), - sizeof(uint64_t), TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); -} - -__aicore__ inline __gm__ uint64_t* AllToAllGroupCreditRequest( - __gm__ int32_t* debug, uint32_t slot, uint32_t lane) -{ - const uint64_t requestIndex = - static_cast(slot) * TILEXR_ALLTOALL_GROUP_SEND_CORES + lane; - return reinterpret_cast<__gm__ uint64_t*>( - reinterpret_cast<__gm__ uint8_t*>(debug) + - TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_OFFSET + - requestIndex * TILEXR_ALLTOALL_GROUP_CREDIT_REQUEST_STRIDE); -} - -__aicore__ inline void AllToAllGroupPublishCreditRequest( - __gm__ int32_t* debug, uint32_t invocationId, uint32_t completedGroup, - uint32_t slot, uint32_t lane, uint32_t groupCount) -{ - if (completedGroup + 1U >= groupCount) { - return; - } - auto request = AllToAllGroupCreditRequest(debug, slot, lane); - *request = AllToAllGroupDeviceToken(invocationId, completedGroup + 1U, 0U); - TileXR::UDMACleanCacheLines( - reinterpret_cast<__gm__ uint8_t*>(request), sizeof(uint64_t)); -} - -__aicore__ inline bool AllToAllGroupFinishCredits( - const __gm__ TileXR::CommArgs* args, __gm__ int32_t* debug, - uint32_t blockIdx, int32_t rank, int32_t rankSize, - uint32_t invocationId, uint32_t lane, uint32_t groupCount, - uint32_t groupWidth, uint32_t publishedCreditCount) -{ - uint64_t completedQueueKeys[TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT] = {}; - uint32_t completedQueueCount = 0U; - bool success = true; - for (uint32_t group = 0U; group < publishedCreditCount; ++group) { - const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( - rank, rankSize, group, lane, groupCount, groupWidth); - if (nextPeer < 0) { - continue; - } - const uint32_t qpIdx = AllToAllGroupCreditQp(args, nextPeer); - auto udmaInfo = TileXR::GetUDMAInfo(args); - auto wq = TileXR::UDMAGetWQCtx(udmaInfo, nextPeer, qpIdx); - const uint64_t queueKey = wq->wqeCntAddr; - bool alreadyCompleted = false; - for (uint32_t index = 0U; index < completedQueueCount; ++index) { - if (completedQueueKeys[index] == queueKey) { - alreadyCompleted = true; - break; - } - } - if (alreadyCompleted) { - continue; - } - completedQueueKeys[completedQueueCount++] = queueKey; - const uint32_t quietStatus = - TileXR::UDMAQuietStatusOnQp(args, nextPeer, qpIdx); - if (quietStatus != 0U) { - AllToAllGroupRecordError(debug, blockIdx, - TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_QUIET, group, 0U, - nextPeer, qpIdx, quietStatus, - AllToAllGroupDeviceToken(invocationId, group + 1U, 0U), 0ULL); - success = false; - } - } - return success; + auto remoteCredit = reinterpret_cast<__gm__ uint64_t*>( + args->creditMems[nextPeer] + creditOffset + + static_cast(rank) * TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE); + *remoteCredit = creditToken; } __aicore__ inline void AllToAllGroupRecordError( @@ -908,8 +808,9 @@ __aicore__ inline void AllToAllGroupKernelImpl( if constexpr (IngressCredit) { if (AllToAllGroupCreditOwnerDevice(worker) && pass + 1U == passCount) { - AllToAllGroupPublishCreditRequest( - debug, invocationId, group, slot, lane, groupCount); + AllToAllGroupPublishNextCredit( + args, rank, rankSize, invocationId, group, lane, + groupCount, groupWidth, creditOffsets[slot]); } } if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { @@ -965,7 +866,6 @@ __aicore__ inline void AllToAllGroupKernelImpl( const uint32_t lane = blockIdx % TILEXR_ALLTOALL_GROUP_SEND_CORES; const uint32_t workerRoute = blockIdx / TILEXR_ALLTOALL_GROUP_SEND_CORES; AllToAllGroupQuietState quietState; - uint32_t publishedCreditCount = 0U; for (uint32_t group = 0U; group < groupCount; ++group) { const int32_t peer = AllToAllGroupDevicePeer( rank, rankSize, group, lane, groupWidth); @@ -1001,42 +901,11 @@ __aicore__ inline void AllToAllGroupKernelImpl( AllToAllGroupDeviceToken(invocationId, group, 0U); const uint64_t creditWaitBegin = AllToAllGroupTraceCycle(groupTrace); - if (workerRoute == 0U) { - auto request = AllToAllGroupCreditRequest(debug, slot, lane); - uint64_t observedRequest = 0ULL; - if (!AllToAllGroupWaitTokenMte( - request, expectedCredit, - TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, - relayLocal, observedRequest)) { - AllToAllGroupRecordError(debug, blockIdx, - TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_REQUEST_WAIT, - group, 0U, peer, workerRoute, 0U, - expectedCredit, observedRequest); - AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, blockIdx, group, 0U, - TileXR::Demo::kAllToAllGroupTraceCreditWait, - groupCount, passCount, peer, - TileXR::Demo::kAllToAllGroupTraceNoQp, - creditWaitBegin, - AllToAllGroupTraceCycle(groupTrace)); - (void)AllToAllGroupFinishCredits( - args, debug, blockIdx, rank, rankSize, - invocationId, lane, groupCount, groupWidth, - publishedCreditCount); - AllToAllGroupTraceRecordKernel( - groupTrace, traceIteration, blockIdx, kernelBegin, - AllToAllGroupTraceCycle(groupTrace)); - return; - } - AllToAllGroupPublishNextCredit( - args, debug, rank, rankSize, invocationId, group - 1U, - lane, groupCount, groupWidth, creditOffsets[slot]); - ++publishedCreditCount; - } if (routeElements != 0U) { auto creditSignal = reinterpret_cast<__gm__ uint64_t*>( - registeredMemory + creditOffsets[slot] + - static_cast(peer) * sizeof(uint64_t)); + args->creditMems[rank] + creditOffsets[slot] + + static_cast(peer) * + TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE); uint64_t observedCredit = 0ULL; if (!AllToAllGroupWaitTokenMte( creditSignal, expectedCredit, @@ -1053,12 +922,6 @@ __aicore__ inline void AllToAllGroupKernelImpl( TileXR::Demo::kAllToAllGroupTraceNoQp, creditWaitBegin, AllToAllGroupTraceCycle(groupTrace)); - if (workerRoute == 0U) { - (void)AllToAllGroupFinishCredits( - args, debug, blockIdx, rank, rankSize, - invocationId, lane, groupCount, groupWidth, - publishedCreditCount); - } AllToAllGroupTraceRecordKernel( groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -1151,14 +1014,6 @@ __aicore__ inline void AllToAllGroupKernelImpl( args, quietState, quietBatch, peer, selectedQp, group, pass, expectedToken, debug, blockIdx, groupTrace, traceIteration, groupCount, passCount)) { - if constexpr (IngressCredit) { - if (workerRoute == 0U) { - (void)AllToAllGroupFinishCredits( - args, debug, blockIdx, rank, rankSize, - invocationId, lane, groupCount, groupWidth, - publishedCreditCount); - } - } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; @@ -1169,26 +1024,10 @@ __aicore__ inline void AllToAllGroupKernelImpl( if (!AllToAllGroupFinishQuiet( args, quietState, debug, blockIdx, groupTrace, traceIteration, groupCount, passCount)) { - if constexpr (IngressCredit) { - if (workerRoute == 0U) { - (void)AllToAllGroupFinishCredits( - args, debug, blockIdx, rank, rankSize, invocationId, lane, - groupCount, groupWidth, publishedCreditCount); - } - } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } - if constexpr (IngressCredit) { - if (workerRoute == 0U && !AllToAllGroupFinishCredits( - args, debug, blockIdx, rank, rankSize, invocationId, lane, - groupCount, groupWidth, publishedCreditCount)) { - AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, - kernelBegin, AllToAllGroupTraceCycle(groupTrace)); - return; - } - } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index 750db316..b02b9472 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -24,6 +24,9 @@ constexpr uint32_t kAllToAllGroupMaxIngressWindow = 1U; constexpr uint32_t kAllToAllGroupPingPongSlots = 2U; constexpr uint32_t kAllToAllGroupRouteSignalStride = 512U; constexpr uint32_t kAllToAllGroupSignalSlotBytes = 1024U; +constexpr size_t kAllToAllGroupCreditStride = 512U; +constexpr size_t kAllToAllGroupCreditSlotBytes = + static_cast(kAllToAllGroupMaxRankSize) * kAllToAllGroupCreditStride; constexpr uint32_t kAllToAllGroupSendCoreCount = 16U; constexpr uint32_t kAllToAllGroupSendWorkerCount = 32U; constexpr uint32_t kAllToAllGroupMaxGroupCount = 64U; @@ -37,33 +40,18 @@ constexpr size_t kAllToAllGroupSignalSourceSlots = kAllToAllGroupMaxQuietBatch; constexpr size_t kAllToAllGroupSignalSourceBytes = kAllToAllGroupSignalSourceSlots * sizeof(uint64_t); -constexpr size_t kAllToAllGroupCreditSourceSlots = - static_cast(kAllToAllGroupSendCoreCount) * - kAllToAllGroupMaxGroupCount; -constexpr size_t kAllToAllGroupCreditSourceBytes = - kAllToAllGroupCreditSourceSlots * sizeof(uint64_t); -constexpr size_t kAllToAllGroupCreditRequestSlots = - static_cast(kAllToAllGroupPingPongSlots) * - kAllToAllGroupSendCoreCount; -constexpr size_t kAllToAllGroupCreditRequestStride = 64U; -constexpr size_t kAllToAllGroupCreditRequestBytes = - kAllToAllGroupCreditRequestSlots * kAllToAllGroupCreditRequestStride; constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupBaseControlBytes = kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes; -constexpr size_t kAllToAllGroupControlBytes = - kAllToAllGroupBaseControlBytes + kAllToAllGroupCreditSourceBytes + - kAllToAllGroupCreditRequestBytes; constexpr size_t kAllToAllGroupMaxPayloadBytes = 16ULL << 30; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 2U * kAllToAllGroupMaxPayloadBytes + 2U * static_cast(kAllToAllGroupMaxRankSize) * kAllToAllGroupSignalSlotBytes + - 2U * static_cast(kAllToAllGroupMaxRankSize) * sizeof(uint64_t) + - kAllToAllGroupControlBytes; + kAllToAllGroupBaseControlBytes; struct AllToAllGroupPlan { bool valid = false; @@ -82,10 +70,6 @@ struct AllToAllGroupPlan { size_t controlBytes = kAllToAllGroupBaseControlBytes; size_t signalSourceOffset = 0; size_t signalSourceBytes = kAllToAllGroupSignalSourceBytes; - size_t creditSourceOffset = 0; - size_t creditSourceBytes = 0; - size_t creditRequestOffset = 0; - size_t creditRequestBytes = 0; size_t registeredBytes = 0; }; @@ -137,21 +121,7 @@ inline size_t AllToAllGroupSignalByteOffset(uint32_t sourceRank, uint32_t route) inline size_t AllToAllGroupCreditByteOffset(uint32_t destinationRank) { - return static_cast(destinationRank) * sizeof(uint64_t); -} - -inline size_t AllToAllGroupCreditSourceByteOffset( - uint32_t lane, uint32_t completedGroup) -{ - return (static_cast(lane) * kAllToAllGroupMaxGroupCount + - completedGroup) * sizeof(uint64_t); -} - -inline size_t AllToAllGroupCreditRequestByteOffset( - uint32_t slot, uint32_t lane) -{ - return (static_cast(slot) * kAllToAllGroupSendCoreCount + lane) * - kAllToAllGroupCreditRequestStride; + return static_cast(destinationRank) * kAllToAllGroupCreditStride; } inline bool AllToAllGroupValidRankSize(int rankSize) @@ -318,9 +288,11 @@ inline AllToAllGroupPlan PlanAllToAllGroup( static_cast(rankSize), kAllToAllGroupSignalSlotBytes, plan.signalPlaneBytes)) { return AllToAllGroupPlan {}; } - if (ingressWindow != 0U && !AllToAllGroupCheckedMul( - static_cast(rankSize), sizeof(uint64_t), plan.creditPlaneBytes)) { - return AllToAllGroupPlan {}; + if (ingressWindow != 0U) { + plan.creditPlaneBytes = static_cast(rankSize) * + kAllToAllGroupCreditStride; + plan.creditOffset[0] = 0U; + plan.creditOffset[1] = kAllToAllGroupCreditSlotBytes; } if (plan.payloadPlaneBytes > kAllToAllGroupMaxPayloadBytes) { return AllToAllGroupPlan {}; @@ -339,31 +311,12 @@ inline AllToAllGroupPlan PlanAllToAllGroup( !AllToAllGroupCheckedAdd(plan.signalOffset[1], plan.signalPlaneBytes, cursor)) { return AllToAllGroupPlan {}; } - if (ingressWindow != 0U) { - plan.creditSourceBytes = kAllToAllGroupCreditSourceBytes; - plan.creditRequestBytes = kAllToAllGroupCreditRequestBytes; - plan.controlBytes = kAllToAllGroupControlBytes; - if (!AllToAllGroupAlignUp(cursor, plan.creditOffset[0]) || - !AllToAllGroupCheckedAdd( - plan.creditOffset[0], plan.creditPlaneBytes, cursor) || - !AllToAllGroupAlignUp(cursor, plan.creditOffset[1]) || - !AllToAllGroupCheckedAdd( - plan.creditOffset[1], plan.creditPlaneBytes, cursor)) { - return AllToAllGroupPlan {}; - } - } if (!AllToAllGroupAlignUp(cursor, plan.controlOffset) || !AllToAllGroupCheckedAdd(plan.controlOffset, plan.controlBytes, cursor) || !AllToAllGroupAlignUp(cursor, plan.registeredBytes)) { return AllToAllGroupPlan {}; } plan.signalSourceOffset = plan.controlOffset + kAllToAllGroupErrorBytes; - if (ingressWindow != 0U) { - plan.creditSourceOffset = - plan.signalSourceOffset + kAllToAllGroupSignalSourceBytes; - plan.creditRequestOffset = - plan.creditSourceOffset + kAllToAllGroupCreditSourceBytes; - } if (plan.registeredBytes > kAllToAllGroupMaxRegisteredBytes) { return AllToAllGroupPlan {}; } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index ccb27fda..c7711ba7 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -911,8 +911,15 @@ bool CopyChunkDeviceToHost( bool RunGroupedAllToAll( int rank, int rankSize, int32_t elementsPerPeer, - int deviceId, TileXRCommPtr comm, aclrtStream stream, GM_ADDR commArgsDev) + int deviceId, TileXRCommPtr comm, aclrtStream stream, + const TileXR::CommArgs& commArgsHost, GM_ADDR commArgsDev) { + static_assert(TileXR::Demo::kAllToAllGroupCreditStride == + static_cast(TileXR::CREDIT_IPC_STRIDE), + "grouped credit stride must match the communicator IPC layout"); + static_assert(TileXR::Demo::kAllToAllGroupCreditSlotBytes == + static_cast(TileXR::CREDIT_IPC_SLOT_BYTES), + "grouped credit slot must match the communicator IPC layout"); constexpr uint32_t kErrorWordsPerCore = 12U; constexpr uint32_t kErrorCoreCount = TileXR::Demo::kAllToAllGroupBlockDim; const int groupWidthValue = GetEnvInt( @@ -953,6 +960,18 @@ bool RunGroupedAllToAll( << " groupWidth=" << groupWidth << std::endl; return false; } + if (ingressWindow != 0U) { + for (int peer = 0; peer < rankSize; ++peer) { + if (commArgsHost.creditMems[peer] == nullptr) { + std::cerr << "[rank " << rank + << "] ERROR: grouped ingress credit requires dedicated" + << " credit IPC mappings for every rank; missing peer=" + << peer << ". Set TILEXR_ENABLE_CREDIT_IPC=1." + << std::endl; + return false; + } + } + } const int32_t requestedChunkElements = std::max( 1, GetEnvInt("TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS", elementsPerPeer)); const auto plan = TileXR::Demo::PlanAllToAllGroup( @@ -1206,8 +1225,6 @@ bool RunGroupedAllToAll( " creditPlaneBytes=" + std::to_string(plan.creditPlaneBytes) + " creditOffset0=" + std::to_string(plan.creditOffset[0]) + " creditOffset1=" + std::to_string(plan.creditOffset[1]) + - " creditRequestOffset=" + std::to_string(plan.creditRequestOffset) + - " creditRequestBytes=" + std::to_string(plan.creditRequestBytes) + " controlOffset=" + std::to_string(plan.controlOffset) + " regionCount=" + std::to_string(useMultiRegion ? groupedRegionCount : 1U) + " groupWidth=" + std::to_string(plan.groupWidth) + @@ -1486,7 +1503,8 @@ int main(int argc, char** argv) if (testType == 8) { const bool ok = RunGroupedAllToAll( - rank, rankSize, elementsPerRank, deviceId, comm, stream, commArgsDev); + rank, rankSize, elementsPerRank, deviceId, comm, stream, + *commArgsHost, commArgsDev); Cleanup(comm, stream, nullptr, nullptr, rank, deviceId); if (!ok) { std::cerr << "[rank " << rank << "] TileXR grouped alltoall demo failed" << std::endl; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 168238ff..757bd9e9 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -152,10 +152,6 @@ void TestPlan() TileXR::Demo::kAllToAllGroupSignalSourceBytes); CHECK_EQ(plan.signalSourceOffset + plan.signalSourceBytes <= plan.controlOffset + plan.controlBytes, true); - CHECK_EQ(plan.creditSourceOffset, 0ULL); - CHECK_EQ(plan.creditSourceBytes, 0ULL); - CHECK_EQ(plan.creditRequestOffset, 0ULL); - CHECK_EQ(plan.creditRequestBytes, 0ULL); const auto legacyAlign = [](size_t value) { return (value + TileXR::Demo::kAllToAllGroupAlignment - 1U) & ~(TileXR::Demo::kAllToAllGroupAlignment - 1U); @@ -174,12 +170,8 @@ void TestPlan() CHECK_EQ(plan.signalOffset[1], legacySignalOffset1); CHECK_EQ(plan.controlOffset, legacyControlOffset); CHECK_EQ(plan.registeredBytes, legacyRegisteredBytes); - CHECK_EQ(TileXR::Demo::AllToAllGroupCreditByteOffset(3U), 24ULL); - CHECK_EQ(TileXR::Demo::AllToAllGroupCreditSourceByteOffset(1U, 2U), - (TileXR::Demo::kAllToAllGroupMaxGroupCount + 2ULL) * sizeof(uint64_t)); - CHECK_EQ(TileXR::Demo::AllToAllGroupCreditRequestByteOffset(1U, 2U), - (TileXR::Demo::kAllToAllGroupSendCoreCount + 2ULL) * - TileXR::Demo::kAllToAllGroupCreditRequestStride); + CHECK_EQ(TileXR::Demo::AllToAllGroupCreditByteOffset(3U), + 3ULL * TileXR::Demo::kAllToAllGroupCreditStride); CHECK_EQ(plan.registeredBytes <= TileXR::Demo::kAllToAllGroupMaxRegisteredBytes, true); CHECK_EQ(TileXR::Demo::kAllToAllGroupMaxPayloadBytes, 16ULL << 30); @@ -188,26 +180,11 @@ void TestPlan() TileXR::Demo::kAllToAllGroupWidth, 1U); CHECK_EQ(ingressPlan.valid, true); CHECK_EQ(ingressPlan.creditPlaneBytes, - static_cast(rankSize) * sizeof(uint64_t)); - CHECK_EQ(ingressPlan.creditOffset[0] >= - ingressPlan.signalOffset[1] + ingressPlan.signalPlaneBytes, true); - CHECK_EQ(ingressPlan.creditOffset[1] >= - ingressPlan.creditOffset[0] + ingressPlan.creditPlaneBytes, true); - CHECK_EQ(ingressPlan.controlOffset >= - ingressPlan.creditOffset[1] + ingressPlan.creditPlaneBytes, true); - CHECK_EQ(ingressPlan.creditSourceOffset, - ingressPlan.signalSourceOffset + - TileXR::Demo::kAllToAllGroupSignalSourceBytes); - CHECK_EQ(ingressPlan.creditSourceOffset + ingressPlan.creditSourceBytes <= - ingressPlan.controlOffset + ingressPlan.controlBytes, true); - CHECK_EQ(ingressPlan.creditRequestOffset, - ingressPlan.creditSourceOffset + - TileXR::Demo::kAllToAllGroupCreditSourceBytes); - CHECK_EQ(ingressPlan.creditRequestBytes, - TileXR::Demo::kAllToAllGroupCreditRequestBytes); - CHECK_EQ(ingressPlan.creditRequestOffset + ingressPlan.creditRequestBytes <= - ingressPlan.controlOffset + ingressPlan.controlBytes, true); - CHECK_EQ(ingressPlan.registeredBytes > plan.registeredBytes, true); + static_cast(rankSize) * TileXR::Demo::kAllToAllGroupCreditStride); + CHECK_EQ(ingressPlan.creditOffset[0], 0ULL); + CHECK_EQ(ingressPlan.creditOffset[1], + TileXR::Demo::kAllToAllGroupCreditSlotBytes); + CHECK_EQ(ingressPlan.registeredBytes, plan.registeredBytes); CHECK_EQ(TileXR::Demo::PlanAllToAllGroup( rankSize, elementsPerPeer, elementsPerPeer, TileXR::Demo::kAllToAllGroupExperimentalWidth, 1U).valid, false); @@ -651,27 +628,21 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "uint64_t creditOffset0, uint64_t creditOffset1"); CHECK_CONTAINS(kernel, "uint32_t ingressWindow"); CHECK_CONTAINS(kernel, "AllToAllGroupPublishNextCredit"); - CHECK_CONTAINS(kernel, "AllToAllGroupPublishCreditRequest"); - CHECK_CONTAINS(kernel, "AllToAllGroupCreditRequest(debug, slot, lane)"); - CHECK_CONTAINS(kernel, "AllToAllGroupFinishCredits"); CHECK_CONTAINS(kernel, "AllToAllGroupCreditOwnerDevice(worker)"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT"); - CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_QUIET"); - CHECK_CONTAINS(kernel, - "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_REQUEST_WAIT"); CHECK_CONTAINS(kernel, "kAllToAllGroupTraceCreditWait"); - CHECK_CONTAINS(kernel, "completedQueueKeys"); - CHECK_CONTAINS(kernel, "publishedCreditCount"); CHECK_CONTAINS(kernel, "TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION"); const size_t publishCreditBegin = kernel.find( "__aicore__ inline void AllToAllGroupPublishNextCredit"); const size_t publishCreditEnd = kernel.find( - "__aicore__ inline bool AllToAllGroupFinishCredits", publishCreditBegin); + "__aicore__ inline void AllToAllGroupRecordError", publishCreditBegin); const std::string publishCredit = publishCreditBegin == std::string::npos ? std::string() : kernel.substr(publishCreditBegin, publishCreditEnd == std::string::npos ? std::string::npos : publishCreditEnd - publishCreditBegin); - CHECK_CONTAINS(publishCredit, "UDMAPutNbiOnQpWithFlag"); + CHECK_CONTAINS(publishCredit, "args->creditMems[nextPeer]"); + CHECK_CONTAINS(publishCredit, "*remoteCredit = creditToken"); + CHECK_NOT_CONTAINS(publishCredit, "UDMAPutNbiOnQpWithFlag"); CHECK_NOT_CONTAINS(publishCredit, "UDMAQuiet"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); @@ -702,6 +673,8 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_QUIET_BATCH"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW"); + CHECK_CONTAINS(demo, "TILEXR_ENABLE_CREDIT_IPC=1"); + CHECK_CONTAINS(demo, "*commArgsHost, commArgsDev"); CHECK_CONTAINS(demo, "grouped ingress credit currently requires single pass"); CHECK_CONTAINS(demo, "grouped ingress credit currently requires groupWidth=16"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS"); From 10f17d74dc7a8df3cf5755b32e28d33ebb64892d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 20:30:18 +0800 Subject: [PATCH 143/163] fix(udma): transfer grouped credits through MTE --- ...-grouped-alltoall-ingress-credit-design.md | 11 ++-- .../tilexr_udma_alltoall_group_kernel.cpp | 54 +++++++++++++++++-- ...test_tilexr_udma_alltoall_group_layout.cpp | 7 ++- 3 files changed, 63 insertions(+), 9 deletions(-) diff --git a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md index b3286cb9..20aab01a 100644 --- a/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md +++ b/docs/specs/2026-08-02-grouped-alltoall-ingress-credit-design.md @@ -89,10 +89,13 @@ latency on the send admission path. ## Credit Submission -Credit publication is a direct 8-byte GM store by the receive owner through -`creditMems[nextSource]`. It does not consume a UDMA WQE, QP, completion, or -quiet, and it does not add another producer to a shared UDMA SQ. Both send -routes poll the same local dedicated credit address through MTE. +Credit publication is one 512-byte MTE copy by the receive owner through +`creditMems[nextSource]`. The token occupies the first 8 bytes of the slot; +copying the complete slot gives the remote IPC write an explicit MTE3 +completion and keeps adjacent credits on separate transfer units. It does not +consume a UDMA WQE, QP, completion, or quiet, and it does not add another +producer to a shared UDMA SQ. Both send routes poll the same local dedicated +credit slot through a 512-byte MTE2 copy. The dedicated allocation is independent of the existing optional communication IPC buffer, so grouped ingress credit can run with `TILEXR_ENABLE_IPC=0`. It diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 33a3835c..fcbf8448 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -29,6 +29,10 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES = TILEXR_ALLTOALL_GROUP_SEND_WORKERS * TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH * sizeof(uint64_t); constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE = 512U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_WORDS = + TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE / sizeof(uint64_t); +static_assert(TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE == TileXR::CREDIT_IPC_STRIDE, + "grouped credit slot must match the runtime IPC layout"); constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; @@ -316,6 +320,35 @@ __aicore__ inline bool AllToAllGroupWaitTokenMte( return observed >= expectedToken; } +__aicore__ inline uint64_t AllToAllGroupLoadCreditMte( + __gm__ uint64_t* credit, AscendC::LocalTensor relayLocal) +{ + AscendC::GlobalTensor creditGlobal; + creditGlobal.SetGlobalBuffer( + credit, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); + auto creditLocal = relayLocal.ReinterpretCast(); + AscendC::DataCopy( + creditLocal, creditGlobal, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + return creditLocal.GetValue(0); +} + +__aicore__ inline bool AllToAllGroupWaitCreditMte( + __gm__ uint64_t* credit, uint64_t expectedToken, uint64_t timeoutCycles, + AscendC::LocalTensor relayLocal, uint64_t& observed) +{ + const uint64_t begin = static_cast(AscendC::GetSystemCycle()); + observed = AllToAllGroupLoadCreditMte(credit, relayLocal); + while (observed < expectedToken) { + if (static_cast(AscendC::GetSystemCycle()) - begin >= timeoutCycles) { + return false; + } + observed = AllToAllGroupLoadCreditMte(credit, relayLocal); + } + return observed >= expectedToken; +} + __aicore__ inline bool AllToAllGroupWaitRouteTokensMte( __gm__ uint64_t* primarySignal, __gm__ uint64_t* secondarySignal, bool waitPrimary, bool waitSecondary, uint64_t expectedToken, @@ -364,7 +397,8 @@ __aicore__ inline void AllToAllGroupPublishNextCredit( const __gm__ TileXR::CommArgs* args, int32_t rank, int32_t rankSize, uint32_t invocationId, uint32_t completedGroup, uint32_t lane, uint32_t groupCount, - uint32_t groupWidth, uint64_t creditOffset) + uint32_t groupWidth, uint64_t creditOffset, + AscendC::LocalTensor relayLocal) { const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( rank, rankSize, completedGroup, lane, groupCount, groupWidth); @@ -376,7 +410,19 @@ __aicore__ inline void AllToAllGroupPublishNextCredit( auto remoteCredit = reinterpret_cast<__gm__ uint64_t*>( args->creditMems[nextPeer] + creditOffset + static_cast(rank) * TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE); - *remoteCredit = creditToken; + auto creditLocal = relayLocal.ReinterpretCast(); + creditLocal.SetValue(0, creditToken); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor remoteCreditGlobal; + remoteCreditGlobal.SetGlobalBuffer( + remoteCredit, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); + AscendC::DataCopy( + remoteCreditGlobal, creditLocal, + TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); } __aicore__ inline void AllToAllGroupRecordError( @@ -810,7 +856,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( pass + 1U == passCount) { AllToAllGroupPublishNextCredit( args, rank, rankSize, invocationId, group, lane, - groupCount, groupWidth, creditOffsets[slot]); + groupCount, groupWidth, creditOffsets[slot], relayLocal); } } if (!AllToAllGroupStageRunsCopyDevice(routeStage)) { @@ -907,7 +953,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( static_cast(peer) * TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE); uint64_t observedCredit = 0ULL; - if (!AllToAllGroupWaitTokenMte( + if (!AllToAllGroupWaitCreditMte( creditSignal, expectedCredit, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, relayLocal, observedCredit)) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 757bd9e9..57a5572d 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -641,10 +641,15 @@ void TestKernelStructure() publishCreditEnd == std::string::npos ? std::string::npos : publishCreditEnd - publishCreditBegin); CHECK_CONTAINS(publishCredit, "args->creditMems[nextPeer]"); - CHECK_CONTAINS(publishCredit, "*remoteCredit = creditToken"); + CHECK_CONTAINS(publishCredit, "TILEXR_ALLTOALL_GROUP_CREDIT_WORDS"); + CHECK_CONTAINS(publishCredit, "AscendC::HardEvent::S_MTE3"); + CHECK_CONTAINS(publishCredit, "AscendC::HardEvent::MTE3_S"); + CHECK_CONTAINS(publishCredit, "AscendC::DataCopy("); + CHECK_NOT_CONTAINS(publishCredit, "*remoteCredit = creditToken"); CHECK_NOT_CONTAINS(publishCredit, "UDMAPutNbiOnQpWithFlag"); CHECK_NOT_CONTAINS(publishCredit, "UDMAQuiet"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTokenMte"); + CHECK_CONTAINS(kernel, "AllToAllGroupWaitCreditMte"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsSendDevice(routeStage)"); CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsReceiveDevice(routeStage)"); CHECK_CONTAINS(kernel, "AllToAllGroupReceivePeerInRouteStageDevice"); From 869c75fbc880658d05fa630ec6cc1f783c3efadd Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Sun, 2 Aug 2026 22:35:17 +0800 Subject: [PATCH 144/163] feat(udma): reduce grouped copyout cores with SDMA --- src/comm/CMakeLists.txt | 5 + src/comm/comm_wrap.cpp | 5 +- src/comm/sdma/tilexr_sdma_a5_backend.cpp | 766 ++++++++++++++++++ src/comm/sdma/tilexr_sdma_a5_backend.h | 138 ++++ src/comm/sdma/tilexr_sdma_a5_cleanup.h | 173 ++++ src/comm/sdma/tilexr_sdma_transport.cpp | 75 +- src/comm/sdma/tilexr_sdma_transport.h | 32 +- src/comm/tilexr_comm.cpp | 37 +- src/comm/tilexr_comm.h | 3 +- src/include/tilexr_sdma.h | 36 +- src/include/tilexr_sdma_a5.h | 203 +++++ src/include/tilexr_sdma_a5_types.h | 200 +++++ tests/sdma/CMakeLists.txt | 19 + tests/sdma/run_tests.sh | 1 + .../unit/test_tilexr_sdma_a5_validation.cpp | 271 +++++++ .../test_tilexr_sdma_transport_disabled.cpp | 32 +- tests/udma/CMakeLists.txt | 124 ++- .../cmake/embed_grouped_alltoall_kernel.cmake | 37 + .../tilexr_udma_alltoall_group_kernel.cpp | 91 +-- .../tilexr_udma_alltoall_group_launcher.cpp | 227 ++++++ .../demo/tilexr_udma_alltoall_group_layout.h | 3 +- ...exr_udma_alltoall_group_trace_to_chrome.py | 10 +- tests/udma/demo/tilexr_udma_demo.cpp | 38 +- ...test_tilexr_udma_alltoall_group_layout.cpp | 45 +- ...exr_udma_alltoall_group_trace_to_chrome.py | 28 +- 25 files changed, 2452 insertions(+), 147 deletions(-) create mode 100644 src/comm/sdma/tilexr_sdma_a5_backend.cpp create mode 100644 src/comm/sdma/tilexr_sdma_a5_backend.h create mode 100644 src/comm/sdma/tilexr_sdma_a5_cleanup.h create mode 100644 src/include/tilexr_sdma_a5.h create mode 100644 src/include/tilexr_sdma_a5_types.h create mode 100644 tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp create mode 100644 tests/udma/cmake/embed_grouped_alltoall_kernel.cmake create mode 100644 tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp diff --git a/src/comm/CMakeLists.txt b/src/comm/CMakeLists.txt index fe13647a..22af9cef 100644 --- a/src/comm/CMakeLists.txt +++ b/src/comm/CMakeLists.txt @@ -116,6 +116,9 @@ set(TILEXR_SOURCE_FILE tilexr_comm.cpp udma/tilexr_udma_transport.cpp sdma/tilexr_sdma_transport.h sdma/tilexr_sdma_transport.cpp + sdma/tilexr_sdma_a5_cleanup.h + sdma/tilexr_sdma_a5_backend.h + sdma/tilexr_sdma_a5_backend.cpp ) add_library(tile-comm SHARED ${TILEXR_SOURCE_FILE}) @@ -176,6 +179,8 @@ install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_udma_reg.h ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_udma_types.h ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma_types.h + ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma_a5_types.h + ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma_a5.h ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma.h ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_sdma_compat.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/src/comm/comm_wrap.cpp b/src/comm/comm_wrap.cpp index bb5c1465..39db3580 100644 --- a/src/comm/comm_wrap.cpp +++ b/src/comm/comm_wrap.cpp @@ -329,7 +329,10 @@ int TileXRCommDestroy(TileXRCommPtr comm) return TILEXR_INVALID_VALUE; } auto *c = static_cast(comm); - + if (!c->PrepareDestroy()) { + TILEXR_LOG(ERROR) << "TileXR SDMA cleanup failed; communicator retained for retry"; + return TILEXR_ERROR_INTERNAL; + } delete c; return TILEXR_SUCCESS; } diff --git a/src/comm/sdma/tilexr_sdma_a5_backend.cpp b/src/comm/sdma/tilexr_sdma_a5_backend.cpp new file mode 100644 index 00000000..aa5320fa --- /dev/null +++ b/src/comm/sdma/tilexr_sdma_a5_backend.cpp @@ -0,0 +1,766 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#include "sdma/tilexr_sdma_a5_backend.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "driver/ascend_hal.h" +#include "sdma/tilexr_sdma_a5_cleanup.h" +#include "tilexr_log.h" + +#ifndef ACL_STREAM_DEVICE_USE_ONLY +#define ACL_STREAM_DEVICE_USE_ONLY 0x00000020U +#endif + +namespace TileXR { +namespace { + +constexpr size_t kBuiltinWorkspaceBytes = 16U * 1024U; +constexpr aclError kExpectedAicpuQueryFailure = + static_cast(detail::TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS); +constexpr int32_t kDeviceInfoModuleType = 0; +constexpr int32_t kPhysicalDieInfoType = 19; + +using RtGetDevicePhyIdByIndexFn = int32_t (*)(uint32_t, uint32_t*); +using RtStreamGetSqidFn = int32_t (*)(const void*, uint32_t*); +using RtStreamGetCqidFn = int32_t (*)(const void*, uint32_t*, uint32_t*); +using RtGetDeviceInfoFn = int32_t (*)(uint32_t, int32_t, int32_t, int64_t*); +using HalResAddrMapFn = drvError_t (*)(unsigned int, res_addr_info*, + unsigned long*, unsigned int*); +using HalResAddrUnmapFn = drvError_t (*)(unsigned int, res_addr_info*); +using AclCreateTensorFn = aclTensor* (*)(const int64_t*, uint64_t, aclDataType, + const int64_t*, int64_t, aclFormat, + const int64_t*, uint64_t, void*); +using AclDestroyTensorFn = int32_t (*)(const aclTensor*); +using AclnnQueryWorkspaceFn = aclnnStatus (*)(const aclTensor*, aclTensor*, + uint64_t*, aclOpExecutor**); +using AclnnQueryFn = aclnnStatus (*)(void*, uint64_t, aclOpExecutor*, aclrtStream); + +template +bool LoadSymbol(void* handle, const char* name, T& symbol) +{ + symbol = reinterpret_cast(dlsym(handle, name)); + if (symbol != nullptr) { + return true; + } + TILEXR_LOG(WARN) << "TileXR A5 SDMA missing runtime symbol " << name; + return false; +} + +class A5RuntimeApi { +public: + ~A5RuntimeApi() + { + Close(); + } + + bool Load() + { + Close(); + runtimeHandle_ = dlopen("libruntime.so", RTLD_NOW | RTLD_LOCAL); + if (runtimeHandle_ == nullptr) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not load libruntime.so: " << dlerror(); + return false; + } + if (!LoadSymbol(runtimeHandle_, "rtGetDevicePhyIdByIndex", getPhysicalDevice) || + !LoadSymbol(runtimeHandle_, "rtStreamGetSqid", getSqId) || + !LoadSymbol(runtimeHandle_, "rtStreamGetCqid", getCqId) || + !LoadSymbol(runtimeHandle_, "rtGetDeviceInfo", getDeviceInfo)) { + Close(); + return false; + } + + opapiHandle_ = dlopen("libopapi.so", RTLD_NOW | RTLD_LOCAL); + if (opapiHandle_ == nullptr) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not load libopapi.so: " << dlerror(); + Close(); + return false; + } + if (!LoadSymbol(opapiHandle_, "aclCreateTensor", createTensor) || + !LoadSymbol(opapiHandle_, "aclDestroyTensor", destroyTensor) || + !LoadSymbol(opapiHandle_, "aclnnShmemSdmaStarsQueryGetWorkspaceSize", prepareQuery) || + !LoadSymbol(opapiHandle_, "aclnnShmemSdmaStarsQuery", executeQuery)) { + Close(); + return false; + } + + if (!LoadSymbol(RTLD_DEFAULT, "halResAddrMap", mapResource) || + !LoadSymbol(RTLD_DEFAULT, "halResAddrUnmap", unmapResource)) { + Close(); + return false; + } + return true; + } + + void Close() + { + getPhysicalDevice = nullptr; + getSqId = nullptr; + getCqId = nullptr; + getDeviceInfo = nullptr; + createTensor = nullptr; + destroyTensor = nullptr; + prepareQuery = nullptr; + executeQuery = nullptr; + mapResource = nullptr; + unmapResource = nullptr; + if (opapiHandle_ != nullptr) { + (void)dlclose(opapiHandle_); + opapiHandle_ = nullptr; + } + if (runtimeHandle_ != nullptr) { + (void)dlclose(runtimeHandle_); + runtimeHandle_ = nullptr; + } + } + + RtGetDevicePhyIdByIndexFn getPhysicalDevice = nullptr; + RtStreamGetSqidFn getSqId = nullptr; + RtStreamGetCqidFn getCqId = nullptr; + RtGetDeviceInfoFn getDeviceInfo = nullptr; + AclCreateTensorFn createTensor = nullptr; + AclDestroyTensorFn destroyTensor = nullptr; + AclnnQueryWorkspaceFn prepareQuery = nullptr; + AclnnQueryFn executeQuery = nullptr; + HalResAddrMapFn mapResource = nullptr; + HalResAddrUnmapFn unmapResource = nullptr; + +private: + void* runtimeHandle_ = nullptr; + void* opapiHandle_ = nullptr; +}; + +struct QuerySnapshot { + uint32_t flag = 0U; + uint32_t totalQueueCount = 0U; + aclError syncStatus = ACL_SUCCESS; + std::vector channels; +}; + +int CleanupSetCurrentContext(void*, void* context) +{ + return static_cast( + aclrtSetCurrentContext(static_cast(context))); +} + +int CleanupDestroyStream(void*, void* stream) +{ + return static_cast(aclrtDestroyStream(static_cast(stream))); +} + +int CleanupDestroyContext(void*, void* context) +{ + return static_cast(aclrtDestroyContext(static_cast(context))); +} + +int CleanupFreeDevice(void*, void* address) +{ + return static_cast(aclrtFree(address)); +} + +int CleanupDestroyTensor(void* opaque, const void* tensor) +{ + A5RuntimeApi* api = static_cast(opaque); + return api == nullptr || api->destroyTensor == nullptr + ? -1 + : api->destroyTensor(static_cast(tensor)); +} + +detail::A5QueryCleanupOps MakeQueryCleanupOps(A5RuntimeApi& api) +{ + detail::A5QueryCleanupOps ops; + ops.opaque = &api; + ops.setCurrentContext = CleanupSetCurrentContext; + ops.destroyStream = CleanupDestroyStream; + ops.destroyContext = CleanupDestroyContext; + ops.freeDevice = CleanupFreeDevice; + ops.destroyTensor = CleanupDestroyTensor; + return ops; +} + +bool AllocateTrackedBuffer(std::vector& buffers, size_t bytes, + bool zero, void*& address) +{ + address = nullptr; + if (aclrtMalloc(&address, bytes, ACL_MEM_MALLOC_HUGE_FIRST) != ACL_SUCCESS) { + return false; + } + buffers.push_back(address); + return !zero || aclrtMemset(address, bytes, 0, bytes) == ACL_SUCCESS; +} + +bool CreateTrackedUint64Tensor(A5RuntimeApi& api, + detail::A5PendingQueryCleanup& cleanup, + void* address, int64_t elements, + aclTensor*& tensor) +{ + const int64_t shape[] = {elements}; + const int64_t strides[] = {1}; + tensor = api.createTensor(shape, 1U, ACL_UINT64, strides, 0, + ACL_FORMAT_ND, shape, 1U, address); + if (tensor == nullptr) { + return false; + } + cleanup.tensors.push_back(tensor); + return true; +} + +bool CheckRuntimeHealth(detail::A5PendingQueryCleanup& cleanup, void* scratch) +{ + aclrtStream stream = nullptr; + const aclError createStatus = aclrtCreateStream(&stream); + cleanup.healthStream = stream; + if (createStatus != ACL_SUCCESS || stream == nullptr) { + return false; + } + const aclError memsetStatus = aclrtMemsetAsync( + scratch, sizeof(uint64_t), 0xA5, sizeof(uint64_t), stream); + const aclError syncStatus = memsetStatus == ACL_SUCCESS + ? aclrtSynchronizeStream(stream) + : memsetStatus; + uint64_t value = 0U; + const aclError copyStatus = syncStatus == ACL_SUCCESS + ? aclrtMemcpy(&value, sizeof(value), scratch, sizeof(value), + ACL_MEMCPY_DEVICE_TO_HOST) + : syncStatus; + return memsetStatus == ACL_SUCCESS && syncStatus == ACL_SUCCESS && + copyStatus == ACL_SUCCESS && value == 0xA5A5A5A5A5A5A5A5ULL; +} + +bool FinishQuery(A5RuntimeApi& api, + detail::A5PendingQueryCleanup& cleanup, + std::vector& pending, + bool result) +{ + const bool released = detail::CleanupA5QueryResources( + cleanup, MakeQueryCleanupOps(api), cleanup.ownerContext); + if (!cleanup.Empty()) { + pending.push_back(std::move(cleanup)); + } + return result && released; +} + +bool RunBuiltinQuery(A5RuntimeApi& api, + int32_t logicalDevice, + const std::vector& streams, + std::vector& pending, + QuerySnapshot& snapshot) +{ + if (streams.empty() || streams.size() > detail::TILEXR_SDMA_A5_CHANNEL_COUNT) { + return false; + } + + detail::A5PendingQueryCleanup cleanup; + aclrtContext ownerContext = nullptr; + if (aclrtGetCurrentContext(&ownerContext) != ACL_SUCCESS || + ownerContext == nullptr) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA query requires an active owner context"; + return false; + } + cleanup.ownerContext = ownerContext; + + const size_t streamsBytes = streams.size() * sizeof(detail::A5BuiltinStreamInfo); + void* streamsDev = nullptr; + void* resourceDev = nullptr; + void* builtinWorkspaceDev = nullptr; + void* inputDev = nullptr; + void* outputDev = nullptr; + void* opWorkspaceDev = nullptr; + if (!AllocateTrackedBuffer(cleanup.ownerBuffers, streamsBytes, false, streamsDev) || + !AllocateTrackedBuffer(cleanup.ownerBuffers, + sizeof(detail::A5BuiltinOpResource), true, resourceDev) || + !AllocateTrackedBuffer(cleanup.ownerBuffers, + kBuiltinWorkspaceBytes, true, builtinWorkspaceDev) || + !AllocateTrackedBuffer(cleanup.ownerBuffers, + 2U * sizeof(uint64_t), false, inputDev) || + !AllocateTrackedBuffer(cleanup.ownerBuffers, + sizeof(uint64_t), true, outputDev) || + aclrtMemcpy(streamsDev, streamsBytes, streams.data(), streamsBytes, + ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA query buffer setup failed"; + return FinishQuery(api, cleanup, pending, false); + } + + detail::A5BuiltinOpResource resource {}; + resource.size = streams.size(); + resource.streamsAddress = reinterpret_cast(streamsDev); + resource.workspaceAddress = reinterpret_cast(builtinWorkspaceDev); + const uint64_t inputs[] = { + reinterpret_cast(resourceDev), + reinterpret_cast(builtinWorkspaceDev), + }; + if (aclrtMemcpy(resourceDev, sizeof(resource), &resource, sizeof(resource), + ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS || + aclrtMemcpy(inputDev, sizeof(inputs), inputs, sizeof(inputs), + ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA query resource upload failed"; + return FinishQuery(api, cleanup, pending, false); + } + + aclrtContext isolated = nullptr; + const aclError createContextStatus = aclrtCreateContext(&isolated, logicalDevice); + cleanup.isolatedContext = isolated; + if (createContextStatus != ACL_SUCCESS || isolated == nullptr) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA isolated query context creation failed"; + return FinishQuery(api, cleanup, pending, false); + } + + aclTensor* inputTensor = nullptr; + aclTensor* outputTensor = nullptr; + aclrtStream queryStream = nullptr; + bool queryLaunched = false; + aclError syncStatus = ACL_SUCCESS; + uint64_t opWorkspaceBytes = 0U; + aclOpExecutor* executor = nullptr; + bool ok = CreateTrackedUint64Tensor( + api, cleanup, inputDev, 2, inputTensor) && + CreateTrackedUint64Tensor(api, cleanup, outputDev, 1, outputTensor) && + api.prepareQuery(inputTensor, outputTensor, + &opWorkspaceBytes, &executor) == ACL_SUCCESS; + if (ok && opWorkspaceBytes != 0U) { + ok = AllocateTrackedBuffer(cleanup.isolatedBuffers, + static_cast(opWorkspaceBytes), false, + opWorkspaceDev); + } + if (ok) { + const aclError createStreamStatus = aclrtCreateStreamWithConfig( + &queryStream, 0, ACL_STREAM_FAST_LAUNCH | ACL_STREAM_FAST_SYNC); + cleanup.queryStream = queryStream; + ok = createStreamStatus == ACL_SUCCESS && queryStream != nullptr; + } + if (ok) { + aclrtStreamAttrValue failureMode {}; + failureMode.failureMode = 0; + ok = aclrtSetStreamAttribute( + queryStream, ACL_STREAM_ATTR_FAILURE_MODE, &failureMode) == ACL_SUCCESS; + } + if (ok) { + const aclnnStatus launchStatus = api.executeQuery( + opWorkspaceDev, opWorkspaceBytes, executor, queryStream); + queryLaunched = launchStatus == ACL_SUCCESS; + ok = queryLaunched; + if (queryLaunched) { + syncStatus = aclrtSynchronizeStream(queryStream); + } + } + + if (aclrtSetCurrentContext(ownerContext) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not restore query owner context"; + return FinishQuery(api, cleanup, pending, false); + } + if (!ok || !queryLaunched) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA built-in query launch failed"; + return FinishQuery(api, cleanup, pending, false); + } + + const size_t snapshotBytes = sizeof(detail::A5BuiltinWorkspaceHeader) + + streams.size() * sizeof(detail::A5BuiltinChannelInfo); + std::vector bytes(snapshotBytes, 0U); + if (aclrtMemcpy(bytes.data(), bytes.size(), builtinWorkspaceDev, bytes.size(), + ACL_MEMCPY_DEVICE_TO_HOST) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA query workspace download failed"; + return FinishQuery(api, cleanup, pending, false); + } + detail::A5BuiltinWorkspaceHeader header {}; + std::memcpy(&header, bytes.data(), sizeof(header)); + snapshot.flag = header.flag; + snapshot.totalQueueCount = header.totalQueueCount; + snapshot.syncStatus = syncStatus; + snapshot.channels.resize(streams.size()); + std::memcpy(snapshot.channels.data(), bytes.data() + sizeof(header), + snapshot.channels.size() * sizeof(snapshot.channels[0])); + + if (syncStatus == kExpectedAicpuQueryFailure && + !CheckRuntimeHealth(cleanup, outputDev)) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA context health check failed after expected AICPU error"; + return FinishQuery(api, cleanup, pending, false); + } + return FinishQuery(api, cleanup, pending, true); +} + +int32_t QueryHostSq(uint32_t physicalDevice, uint32_t sqId, + drvSqCqPropType_t property, uint32_t (&values)[3]) +{ + halSqCqQueryInfo query {}; + query.type = DRV_NORMAL_TYPE; + query.tsId = 0U; + query.sqId = sqId; + query.cqId = 0U; + query.prop = property; + const drvError_t status = halSqCqQuery(physicalDevice, &query); + values[0] = query.value[0]; + values[1] = query.value[1]; + values[2] = query.value[2]; + return static_cast(status); +} + +} // namespace + +struct TileXRA5SDMABackend::Impl { + struct OwnedChannel { + aclrtStream stream = nullptr; + res_addr_info mapInfo {}; + bool mapped = false; + uint64_t rtsqAddress = 0U; + uint32_t rtsqLength = 0U; + detail::A5HostChannelIdentity identity {}; + detail::A5BuiltinChannelInfo query {}; + }; + + int32_t logicalDevice = -1; + uint32_t physicalDevice = 0U; + uint32_t physicalDieId = 0U; + A5RuntimeApi api; + aclrtContext ownerContext = nullptr; + aclrtContext restoreContext = nullptr; + bool restorePending = false; + void* workspaceDev = nullptr; + std::vector pendingQueries; + std::array channels {}; + + bool HasOwnedResources() const + { + if (workspaceDev != nullptr || !pendingQueries.empty()) { + return true; + } + for (const OwnedChannel& channel : channels) { + if (channel.mapped || channel.stream != nullptr) { + return true; + } + } + return false; + } + + void EraseCompletedQueries() + { + auto query = pendingQueries.begin(); + while (query != pendingQueries.end()) { + if (query->Empty()) { + query = pendingQueries.erase(query); + } else { + ++query; + } + } + } +}; + +TileXRA5SDMABackend::TileXRA5SDMABackend() = default; + +TileXRA5SDMABackend::~TileXRA5SDMABackend() +{ + (void)Shutdown(); +} + +bool TileXRA5SDMABackend::Init(int32_t deviceId) +{ + if (impl_ != nullptr) { + TILEXR_LOG(ERROR) << "TileXR A5 SDMA backend contains state before initialization"; + return false; + } + std::unique_ptr state(new (std::nothrow) Impl()); + if (state == nullptr) { + return false; + } + state->logicalDevice = deviceId; + if (aclrtGetCurrentContext(&state->ownerContext) != ACL_SUCCESS || + state->ownerContext == nullptr) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA requires an active owner context"; + return false; + } + + if (!state->api.Load() || + state->api.getPhysicalDevice( + static_cast(deviceId), &state->physicalDevice) != 0) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA runtime discovery failed"; + return false; + } + int64_t physicalDie = -1; + if (state->api.getDeviceInfo( + static_cast(deviceId), kDeviceInfoModuleType, + kPhysicalDieInfoType, &physicalDie) != 0 || + physicalDie < 0 || static_cast(physicalDie) > + static_cast(std::numeric_limits::max())) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA physical die discovery failed"; + return false; + } + state->physicalDieId = static_cast(physicalDie); + impl_ = std::move(state); + A5RuntimeApi& api = impl_->api; + + std::vector streamInfos; + streamInfos.reserve(detail::TILEXR_SDMA_A5_CHANNEL_COUNT); + for (uint32_t index = 0U; index < detail::TILEXR_SDMA_A5_CHANNEL_COUNT; ++index) { + Impl::OwnedChannel& owned = impl_->channels[index]; + if (aclrtCreateStreamWithConfig( + &owned.stream, 0, ACL_STREAM_DEVICE_USE_ONLY) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA STARS stream creation failed at channel " << index; + Shutdown(); + return false; + } + int32_t streamId = -1; + uint32_t sqId = 0U; + uint32_t cqId = 0U; + uint32_t logicalCqId = 0U; + if (aclrtStreamGetId(owned.stream, &streamId) != ACL_SUCCESS || streamId < 0 || + api.getSqId(owned.stream, &sqId) != 0 || + api.getCqId(owned.stream, &cqId, &logicalCqId) != 0) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA stream identifier query failed at channel " << index; + Shutdown(); + return false; + } + owned.identity = { + static_cast(streamId), sqId, cqId, logicalCqId, impl_->physicalDieId, + }; + owned.mapInfo.id = 0U; + owned.mapInfo.target_proc_type = PROCESS_CP1; + owned.mapInfo.res_type = RES_ADDR_TYPE_STARS_RTSQ; + owned.mapInfo.res_id = sqId; + unsigned long mappedAddress = 0UL; + unsigned int mappedLength = 0U; + const drvError_t mapStatus = api.mapResource( + impl_->physicalDevice, &owned.mapInfo, &mappedAddress, &mappedLength); + owned.mapped = mapStatus == DRV_ERROR_NONE; + if (mapStatus != DRV_ERROR_NONE || mappedAddress == 0UL || + mappedLength < sizeof(uint32_t)) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA RTSQ map failed at channel " << index + << ", status " << mapStatus; + Shutdown(); + return false; + } + owned.rtsqAddress = static_cast(mappedAddress); + owned.rtsqLength = mappedLength; + + detail::A5BuiltinStreamInfo info {}; + info.stream = reinterpret_cast(owned.stream); + info.context = reinterpret_cast(impl_->ownerContext); + info.streamId = streamId; + info.sqId = sqId; + info.cqId = cqId; + info.logicalCqId = logicalCqId; + info.deviceId = static_cast(impl_->physicalDieId); + streamInfos.push_back(info); + } + + QuerySnapshot batch; + if (!RunBuiltinQuery( + api, deviceId, streamInfos, impl_->pendingQueries, batch)) { + Shutdown(); + return false; + } + const detail::A5QueryResultKind batchKind = batch.channels.empty() + ? detail::A5QueryResultKind::INVALID + : detail::ClassifyA5QueryResult( + static_cast(batch.syncStatus), batch.flag, batch.totalQueueCount, + batch.channels[0], impl_->channels[0].identity); + if (batchKind == detail::A5QueryResultKind::COMPLETE) { + if (batch.channels.size() != detail::TILEXR_SDMA_A5_CHANNEL_COUNT) { + Shutdown(); + return false; + } + for (uint32_t index = 0U; index < detail::TILEXR_SDMA_A5_CHANNEL_COUNT; ++index) { + if (!detail::ValidateA5BuiltinChannel( + batch.channels[index], impl_->channels[index].identity, true)) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA complete query validation failed at channel " << index; + Shutdown(); + return false; + } + impl_->channels[index].query = batch.channels[index]; + } + } else if (batchKind == detail::A5QueryResultKind::EXPECTED_PARTIAL) { + impl_->channels[0].query = batch.channels[0]; + for (uint32_t index = 1U; index < detail::TILEXR_SDMA_A5_CHANNEL_COUNT; ++index) { + std::vector oneStream(1U, streamInfos[index]); + QuerySnapshot isolated; + if (!RunBuiltinQuery( + api, deviceId, oneStream, impl_->pendingQueries, isolated) || + isolated.channels.size() != 1U || + detail::ClassifyA5QueryResult( + static_cast(isolated.syncStatus), isolated.flag, + isolated.totalQueueCount, isolated.channels[0], + impl_->channels[index].identity) != + detail::A5QueryResultKind::EXPECTED_PARTIAL) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA isolated query validation failed at channel " << index; + Shutdown(); + return false; + } + impl_->channels[index].query = isolated.channels[0]; + } + } else { + TILEXR_LOG(WARN) << "TileXR A5 SDMA query returned unsupported status " << batch.syncStatus; + Shutdown(); + return false; + } + + detail::A5SdmaWorkspace hostWorkspace {}; + hostWorkspace.header.magic = detail::TILEXR_SDMA_A5_WORKSPACE_MAGIC; + hostWorkspace.header.abiVersion = detail::TILEXR_SDMA_A5_ABI_VERSION; + hostWorkspace.header.backendKind = detail::TILEXR_SDMA_A5_BACKEND_KIND; + hostWorkspace.header.channelCount = detail::TILEXR_SDMA_A5_CHANNEL_COUNT; + hostWorkspace.header.sqeSize = detail::TILEXR_SDMA_A5_SQE_BYTES; + hostWorkspace.header.channelStride = sizeof(detail::A5SdmaChannel); + hostWorkspace.header.workspaceSize = sizeof(hostWorkspace); + hostWorkspace.header.maxTransferBytes = + static_cast(detail::TILEXR_SDMA_A5_MAX_TRANSFER_BYTES); + + if (aclrtMalloc(&impl_->workspaceDev, sizeof(hostWorkspace), + ACL_MEM_MALLOC_HUGE_FIRST) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA workspace allocation failed"; + Shutdown(); + return false; + } + const uint64_t workspaceBase = reinterpret_cast(impl_->workspaceDev); + for (uint32_t index = 0U; index < detail::TILEXR_SDMA_A5_CHANNEL_COUNT; ++index) { + const Impl::OwnedChannel& owned = impl_->channels[index]; + detail::A5SdmaChannel& channel = hostWorkspace.channels[index]; + channel.sqBase = owned.query.sqBase; + channel.rtsqAddress = owned.rtsqAddress; + channel.completionPayloadAddress = workspaceBase + + offsetof(detail::A5SdmaWorkspace, completionPayloads) + + index * sizeof(detail::A5SdmaCompletionLine); + channel.completionRecordAddress = workspaceBase + + offsetof(detail::A5SdmaWorkspace, completionRecords) + + index * sizeof(detail::A5SdmaCompletionLine); + channel.depth = owned.query.sqDepth; + channel.head = owned.query.sqHead; + channel.tail = owned.query.sqTail; + channel.taskId = detail::A5SdmaQueueDistance( + owned.query.sqHead, owned.query.sqTail, owned.query.sqDepth); + channel.rtsqLength = owned.rtsqLength; + channel.streamId = owned.identity.streamId; + channel.sqId = owned.identity.sqId; + channel.cqId = owned.identity.cqId; + channel.logicalCqId = owned.identity.logicalCqId; + channel.physicalDieId = owned.identity.physicalDieId; + + uint32_t tailValues[3] = {0U, 0U, 0U}; + uint32_t sqeSizeValues[3] = {0U, 0U, 0U}; + if (QueryHostSq(impl_->physicalDevice, channel.sqId, + DRV_SQCQ_PROP_SQ_TAIL, tailValues) != 0 || + QueryHostSq(impl_->physicalDevice, channel.sqId, + DRV_SQCQ_PROP_SQE_SIZE, sqeSizeValues) != 0 || + tailValues[0] != channel.tail || + sqeSizeValues[0] != detail::TILEXR_SDMA_A5_SQE_BYTES) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA Host SQ cross-check failed at channel " << index; + Shutdown(); + return false; + } + } + if (aclrtMemcpy(impl_->workspaceDev, sizeof(hostWorkspace), &hostWorkspace, + sizeof(hostWorkspace), ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA workspace upload failed"; + Shutdown(); + return false; + } + + TILEXR_LOG(INFO) << "TileXR A5 direct SDMA initialized on device " << deviceId + << " with " << detail::TILEXR_SDMA_A5_CHANNEL_COUNT << " channels"; + return true; +} + +bool TileXRA5SDMABackend::Shutdown() +{ + if (impl_ == nullptr) { + return true; + } + if (impl_->restorePending) { + if (aclrtSetCurrentContext(impl_->restoreContext) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not complete pending context restore"; + return false; + } + impl_->restoreContext = nullptr; + impl_->restorePending = false; + if (!impl_->HasOwnedResources()) { + impl_.reset(); + return true; + } + } + + const detail::A5QueryCleanupOps queryCleanupOps = + MakeQueryCleanupOps(impl_->api); + for (size_t reverse = impl_->pendingQueries.size(); reverse > 0U; --reverse) { + detail::A5PendingQueryCleanup& query = impl_->pendingQueries[reverse - 1U]; + if (query.restorePending && + !detail::CleanupA5QueryResources( + query, queryCleanupOps, query.restoreContext)) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA query context restore remains pending"; + return false; + } + } + impl_->EraseCompletedQueries(); + + aclrtContext previous = nullptr; + if (aclrtGetCurrentContext(&previous) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not capture current context for cleanup"; + return false; + } + for (size_t reverse = impl_->pendingQueries.size(); reverse > 0U; --reverse) { + detail::A5PendingQueryCleanup& query = impl_->pendingQueries[reverse - 1U]; + if (!detail::CleanupA5QueryResources(query, queryCleanupOps, previous)) { + impl_->EraseCompletedQueries(); + TILEXR_LOG(WARN) << "TileXR A5 SDMA query cleanup incomplete; retained for retry"; + return false; + } + } + impl_->EraseCompletedQueries(); + + if (impl_->ownerContext == nullptr || + aclrtSetCurrentContext(impl_->ownerContext) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not switch to owner context for cleanup"; + return false; + } + + if (impl_->workspaceDev != nullptr) { + if (aclrtFree(impl_->workspaceDev) == ACL_SUCCESS) { + impl_->workspaceDev = nullptr; + } + } + if (impl_->workspaceDev == nullptr) { + for (size_t reverse = detail::TILEXR_SDMA_A5_CHANNEL_COUNT; reverse > 0U; --reverse) { + Impl::OwnedChannel& owned = impl_->channels[reverse - 1U]; + bool mappingReleased = true; + if (owned.mapped) { + if (impl_->api.unmapResource( + impl_->physicalDevice, &owned.mapInfo) == DRV_ERROR_NONE) { + owned.mapped = false; + } else { + mappingReleased = false; + } + } + if (mappingReleased && owned.stream != nullptr && + aclrtDestroyStream(owned.stream) == ACL_SUCCESS) { + owned.stream = nullptr; + } + } + } + if (previous != impl_->ownerContext && + aclrtSetCurrentContext(previous) != ACL_SUCCESS) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA could not restore previous context after cleanup"; + impl_->restoreContext = previous; + impl_->restorePending = true; + return false; + } + if (impl_->HasOwnedResources()) { + TILEXR_LOG(WARN) << "TileXR A5 SDMA cleanup incomplete; retained resources for retry"; + return false; + } + impl_.reset(); + return true; +} + +GM_ADDR TileXRA5SDMABackend::GetWorkspaceDev() const +{ + return impl_ == nullptr ? nullptr : static_cast(impl_->workspaceDev); +} + +} // namespace TileXR diff --git a/src/comm/sdma/tilexr_sdma_a5_backend.h b/src/comm/sdma/tilexr_sdma_a5_backend.h new file mode 100644 index 00000000..d1b4a592 --- /dev/null +++ b/src/comm/sdma/tilexr_sdma_a5_backend.h @@ -0,0 +1,138 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_SDMA_A5_BACKEND_H +#define TILEXR_SDMA_A5_BACKEND_H + +#include +#include +#include + +#include "comm_args.h" +#include "tilexr_sdma_a5_types.h" + +namespace TileXR { + +namespace detail { + +struct alignas(64) A5BuiltinStreamInfo { + uint64_t stream; + uint64_t context; + int32_t streamId; + uint32_t sqId; + uint32_t cqId; + uint32_t logicalCqId; + uint64_t cqeAddress; + int32_t deviceId; + uint8_t reserved[20]; +}; + +struct alignas(64) A5BuiltinOpResource { + uint64_t size; + uint64_t streamsAddress; + uint64_t workspaceAddress; + uint8_t reserved[40]; +}; + +struct alignas(64) A5BuiltinChannelInfo { + uint32_t sqHead; + uint32_t sqTail; + uint64_t sqBase; + uint64_t sqRegisterBase; + uint32_t sqDepth; + uint32_t sqId; + uint32_t cqId; + uint32_t logicalCqId; + uint64_t cqeAddress; + uint32_t reportCqeCount; + uint32_t streamId; + uint32_t deviceId; + uint8_t reserved[4]; +}; + +struct alignas(64) A5BuiltinWorkspaceHeader { + uint32_t flag; + uint32_t totalQueueCount; + uint8_t reserved[56]; +}; + +struct A5HostChannelIdentity { + uint32_t streamId; + uint32_t sqId; + uint32_t cqId; + uint32_t logicalCqId; + uint32_t physicalDieId; +}; + +constexpr int32_t TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS = 507018; + +enum class A5QueryResultKind : uint32_t { + INVALID = 0U, + COMPLETE = 1U, + EXPECTED_PARTIAL = 2U, +}; + +inline bool ValidateA5BuiltinChannel(const A5BuiltinChannelInfo& channel, + const A5HostChannelIdentity& expected, + bool requireRegisterBase) +{ + return channel.sqBase != 0U && + (!requireRegisterBase || channel.sqRegisterBase != 0U) && + A5SdmaQueueHasCapacity(channel.sqHead, channel.sqTail, channel.sqDepth) && + channel.streamId == expected.streamId && channel.sqId == expected.sqId && + channel.cqId == expected.cqId && channel.logicalCqId == expected.logicalCqId && + channel.deviceId == expected.physicalDieId; +} + +inline A5QueryResultKind ClassifyA5QueryResult(int32_t syncStatus, + uint32_t flag, + uint32_t totalQueueCount, + const A5BuiltinChannelInfo& firstChannel, + const A5HostChannelIdentity& expected) +{ + if (syncStatus == 0 && flag == 1U && + totalQueueCount == TILEXR_SDMA_A5_CHANNEL_COUNT && + ValidateA5BuiltinChannel(firstChannel, expected, true)) { + return A5QueryResultKind::COMPLETE; + } + if (syncStatus == TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS && flag == 0U && + totalQueueCount == 0U && firstChannel.sqRegisterBase == 0U && + ValidateA5BuiltinChannel(firstChannel, expected, false)) { + return A5QueryResultKind::EXPECTED_PARTIAL; + } + return A5QueryResultKind::INVALID; +} + +static_assert(sizeof(A5BuiltinStreamInfo) == 64U, "unexpected built-in stream ABI"); +static_assert(sizeof(A5BuiltinOpResource) == 64U, "unexpected built-in resource ABI"); +static_assert(sizeof(A5BuiltinChannelInfo) == 64U, "unexpected built-in channel ABI"); +static_assert(sizeof(A5BuiltinWorkspaceHeader) == 64U, "unexpected built-in header ABI"); +static_assert(offsetof(A5BuiltinChannelInfo, sqBase) == 8U, "unexpected built-in SQ base offset"); +static_assert(offsetof(A5BuiltinChannelInfo, sqRegisterBase) == 16U, + "unexpected built-in register offset"); +static_assert(offsetof(A5BuiltinChannelInfo, streamId) == 52U, + "unexpected built-in stream ID offset"); + +} // namespace detail + +class TileXRA5SDMABackend { +public: + TileXRA5SDMABackend(); + ~TileXRA5SDMABackend(); + TileXRA5SDMABackend(const TileXRA5SDMABackend&) = delete; + TileXRA5SDMABackend& operator=(const TileXRA5SDMABackend&) = delete; + + bool Init(int32_t deviceId); + bool Shutdown(); + GM_ADDR GetWorkspaceDev() const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace TileXR + +#endif // TILEXR_SDMA_A5_BACKEND_H diff --git a/src/comm/sdma/tilexr_sdma_a5_cleanup.h b/src/comm/sdma/tilexr_sdma_a5_cleanup.h new file mode 100644 index 00000000..b5098021 --- /dev/null +++ b/src/comm/sdma/tilexr_sdma_a5_cleanup.h @@ -0,0 +1,173 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_SDMA_A5_CLEANUP_H +#define TILEXR_SDMA_A5_CLEANUP_H + +#include +#include + +namespace TileXR { +namespace detail { + +using A5CleanupFn = int (*)(void*, void*); +using A5ConstCleanupFn = int (*)(void*, const void*); + +struct A5QueryCleanupOps { + void* opaque = nullptr; + A5CleanupFn setCurrentContext = nullptr; + A5CleanupFn destroyStream = nullptr; + A5CleanupFn destroyContext = nullptr; + A5CleanupFn freeDevice = nullptr; + A5ConstCleanupFn destroyTensor = nullptr; +}; + +struct A5PendingQueryCleanup { + void* ownerContext = nullptr; + void* isolatedContext = nullptr; + void* queryStream = nullptr; + void* healthStream = nullptr; + std::vector ownerBuffers; + std::vector isolatedBuffers; + std::vector tensors; + void* restoreContext = nullptr; + bool restorePending = false; + + bool Empty() const + { + return isolatedContext == nullptr && queryStream == nullptr && + healthStream == nullptr && ownerBuffers.empty() && + isolatedBuffers.empty() && tensors.empty() && !restorePending; + } +}; + +inline bool A5SetCleanupContext(const A5QueryCleanupOps& ops, void* context) +{ + return ops.setCurrentContext != nullptr && + ops.setCurrentContext(ops.opaque, context) == 0; +} + +inline bool A5ReleaseCleanupHandle(void*& handle, A5CleanupFn cleanup, + const A5QueryCleanupOps& ops) +{ + if (handle == nullptr) { + return true; + } + if (cleanup != nullptr && cleanup(ops.opaque, handle) == 0) { + handle = nullptr; + return true; + } + return false; +} + +inline bool A5ReleaseCleanupHandles(std::vector& handles, + A5CleanupFn cleanup, + const A5QueryCleanupOps& ops) +{ + bool released = true; + size_t reverse = handles.size(); + while (reverse > 0U) { + --reverse; + if (handles[reverse] != nullptr && cleanup != nullptr && + cleanup(ops.opaque, handles[reverse]) == 0) { + handles.erase(handles.begin() + static_cast(reverse)); + } else { + released = false; + } + } + return released; +} + +inline bool A5ReleaseCleanupTensors(std::vector& tensors, + const A5QueryCleanupOps& ops) +{ + bool released = true; + size_t reverse = tensors.size(); + while (reverse > 0U) { + --reverse; + if (tensors[reverse] != nullptr && ops.destroyTensor != nullptr && + ops.destroyTensor(ops.opaque, tensors[reverse]) == 0) { + tensors.erase(tensors.begin() + static_cast(reverse)); + } else { + released = false; + } + } + return released; +} + +inline bool CleanupA5QueryResources(A5PendingQueryCleanup& state, + const A5QueryCleanupOps& ops, + void* callerContext) +{ + bool released = true; + if (state.restorePending) { + if (!A5SetCleanupContext(ops, state.restoreContext)) { + return false; + } + state.restoreContext = nullptr; + state.restorePending = false; + } + + if (state.healthStream != nullptr) { + if (state.ownerContext == nullptr || + !A5SetCleanupContext(ops, state.ownerContext)) { + released = false; + } else { + released = A5ReleaseCleanupHandle( + state.healthStream, ops.destroyStream, ops) && released; + } + } + + const bool hasIsolatedResources = state.queryStream != nullptr || + !state.isolatedBuffers.empty() || !state.tensors.empty(); + if (state.isolatedContext != nullptr) { + if (!A5SetCleanupContext(ops, state.isolatedContext)) { + released = false; + } else { + released = A5ReleaseCleanupHandle( + state.queryStream, ops.destroyStream, ops) && released; + if (state.queryStream == nullptr) { + released = A5ReleaseCleanupHandles( + state.isolatedBuffers, ops.freeDevice, ops) && released; + released = A5ReleaseCleanupTensors(state.tensors, ops) && released; + } + if (state.queryStream == nullptr && state.tensors.empty() && + state.isolatedBuffers.empty()) { + released = A5ReleaseCleanupHandle( + state.isolatedContext, ops.destroyContext, ops) && released; + } + } + } else if (hasIsolatedResources) { + released = false; + } + + const bool isolatedReleased = state.isolatedContext == nullptr && + state.queryStream == nullptr && state.isolatedBuffers.empty() && + state.tensors.empty(); + if (!state.ownerBuffers.empty() && state.healthStream == nullptr && + isolatedReleased) { + if (state.ownerContext == nullptr || + !A5SetCleanupContext(ops, state.ownerContext)) { + released = false; + } else { + released = A5ReleaseCleanupHandles( + state.ownerBuffers, ops.freeDevice, ops) && released; + } + } else if (!state.ownerBuffers.empty()) { + released = false; + } + + if (!A5SetCleanupContext(ops, callerContext)) { + state.restoreContext = callerContext; + state.restorePending = true; + released = false; + } + return released && state.Empty(); +} + +} // namespace detail +} // namespace TileXR + +#endif // TILEXR_SDMA_A5_CLEANUP_H diff --git a/src/comm/sdma/tilexr_sdma_transport.cpp b/src/comm/sdma/tilexr_sdma_transport.cpp index 08ee0b6a..9b770757 100644 --- a/src/comm/sdma/tilexr_sdma_transport.cpp +++ b/src/comm/sdma/tilexr_sdma_transport.cpp @@ -9,6 +9,8 @@ #include #include +#include "acl/acl_rt.h" +#include "sdma/tilexr_sdma_a5_backend.h" #include "tilexr_log.h" #include "tilexr_types.h" @@ -19,8 +21,10 @@ namespace TileXR { struct TileXRSDMATransport::Impl { + std::unique_ptr a5Backend; #if TILEXR_HAVE_PTO_SDMA pto::comm::sdma::SdmaWorkspaceManager workspaceManager; + bool ptoInitialized = false; #endif }; @@ -28,7 +32,7 @@ TileXRSDMATransport::TileXRSDMATransport() = default; TileXRSDMATransport::~TileXRSDMATransport() { - Shutdown(); + (void)Shutdown(); } bool TileXRSDMATransport::EnvEnabled() @@ -43,7 +47,11 @@ bool TileXRSDMATransport::EnvEnabled() int TileXRSDMATransport::Init(const TileXRSDMATransportOptions& options) { - Shutdown(); + if (impl_ != nullptr) { + TILEXR_LOG(ERROR) << "TileXR SDMA transport contains state before initialization"; + lastStatus_ = SDMAInitStatus::INIT_FAILED; + return TILEXR_ERROR_INTERNAL; + } options_ = options; available_ = false; workspaceDev_ = nullptr; @@ -54,19 +62,53 @@ int TileXRSDMATransport::Init(const TileXRSDMATransportOptions& options) return TILEXR_SUCCESS; } -#if TILEXR_HAVE_PTO_SDMA + const char* socName = aclrtGetSocName(); + const detail::SDMABackendKind backend = detail::ClassifySDMABackend(socName); + if (backend == detail::SDMABackendKind::UNSUPPORTED) { + lastStatus_ = SDMAInitStatus::PTO_UNAVAILABLE; + TILEXR_LOG(WARN) << "TileXR SDMA unsupported on SoC " + << (socName == nullptr ? "unknown" : socName); + return TILEXR_SUCCESS; + } + impl_.reset(new (std::nothrow) Impl()); if (impl_ == nullptr) { lastStatus_ = SDMAInitStatus::INIT_FAILED; - TILEXR_LOG(WARN) << "TileXR SDMA workspace manager allocation failed"; + TILEXR_LOG(WARN) << "TileXR SDMA implementation allocation failed"; + return TILEXR_SUCCESS; + } + + if (backend == detail::SDMABackendKind::A5_DIRECT) { + impl_->a5Backend.reset(new (std::nothrow) TileXRA5SDMABackend()); + if (impl_->a5Backend == nullptr) { + impl_.reset(); + lastStatus_ = SDMAInitStatus::INIT_FAILED; + return TILEXR_SUCCESS; + } + if (!impl_->a5Backend->Init(options_.devId)) { + lastStatus_ = SDMAInitStatus::INIT_FAILED; + TILEXR_LOG(WARN) << "TileXR A5 direct SDMA unavailable; communicator will continue without SDMA"; + return TILEXR_SUCCESS; + } + workspaceDev_ = impl_->a5Backend->GetWorkspaceDev(); + if (workspaceDev_ == nullptr) { + (void)Shutdown(); + lastStatus_ = SDMAInitStatus::NULL_WORKSPACE; + return TILEXR_SUCCESS; + } + available_ = true; + lastStatus_ = SDMAInitStatus::INITIALIZED; return TILEXR_SUCCESS; } + +#if TILEXR_HAVE_PTO_SDMA if (!impl_->workspaceManager.Init()) { lastStatus_ = SDMAInitStatus::INIT_FAILED; TILEXR_LOG(WARN) << "TileXR SDMA workspace manager init failed"; impl_.reset(); return TILEXR_SUCCESS; } + impl_->ptoInitialized = true; workspaceDev_ = static_cast(impl_->workspaceManager.GetWorkspaceAddr()); if (workspaceDev_ == nullptr) { lastStatus_ = SDMAInitStatus::NULL_WORKSPACE; @@ -83,20 +125,35 @@ int TileXRSDMATransport::Init(const TileXRSDMATransportOptions& options) #else lastStatus_ = SDMAInitStatus::PTO_UNAVAILABLE; TILEXR_LOG(WARN) << "TileXR SDMA PTO headers unavailable at build time"; + impl_.reset(); return TILEXR_SUCCESS; #endif } -void TileXRSDMATransport::Shutdown() +bool TileXRSDMATransport::Shutdown() { -#if TILEXR_HAVE_PTO_SDMA + bool cleanupComplete = true; if (impl_ != nullptr) { - impl_->workspaceManager.Finalize(); - impl_.reset(); - } + if (impl_->a5Backend != nullptr) { + if (impl_->a5Backend->Shutdown()) { + impl_->a5Backend.reset(); + } else { + cleanupComplete = false; + } + } +#if TILEXR_HAVE_PTO_SDMA + if (impl_->ptoInitialized) { + impl_->workspaceManager.Finalize(); + impl_->ptoInitialized = false; + } #endif + if (cleanupComplete) { + impl_.reset(); + } + } available_ = false; workspaceDev_ = nullptr; + return cleanupComplete; } bool TileXRSDMATransport::IsAvailable() const diff --git a/src/comm/sdma/tilexr_sdma_transport.h b/src/comm/sdma/tilexr_sdma_transport.h index 2f8ed9f9..93f778c3 100644 --- a/src/comm/sdma/tilexr_sdma_transport.h +++ b/src/comm/sdma/tilexr_sdma_transport.h @@ -6,6 +6,7 @@ #ifndef TILEXR_SDMA_TRANSPORT_H #define TILEXR_SDMA_TRANSPORT_H +#include #include #include "comm_args.h" @@ -13,6 +14,35 @@ namespace TileXR { +namespace detail { + +enum class SDMABackendKind : uint32_t { + UNSUPPORTED = 0U, + PTO = 1U, + A5_DIRECT = 2U, +}; + +inline bool SDMASocHasPrefix(const char* socName, const char* prefix) +{ + return socName != nullptr && + std::strncmp(socName, prefix, std::strlen(prefix)) == 0; +} + +inline SDMABackendKind ClassifySDMABackend(const char* socName) +{ + if (SDMASocHasPrefix(socName, "Ascend950")) { + return SDMABackendKind::A5_DIRECT; + } + if (SDMASocHasPrefix(socName, "Ascend910B") || + SDMASocHasPrefix(socName, "Ascend910A") || + SDMASocHasPrefix(socName, "Ascend910_93")) { + return SDMABackendKind::PTO; + } + return SDMABackendKind::UNSUPPORTED; +} + +} // namespace detail + struct TileXRSDMATransportOptions { int devId = 0; }; @@ -25,7 +55,7 @@ class TileXRSDMATransport { TileXRSDMATransport& operator=(const TileXRSDMATransport&) = delete; int Init(const TileXRSDMATransportOptions& options); - void Shutdown(); + bool Shutdown(); bool IsAvailable() const; GM_ADDR GetWorkspaceDev() const; diff --git a/src/comm/tilexr_comm.cpp b/src/comm/tilexr_comm.cpp index 8c9d0dd5..a8fba19f 100644 --- a/src/comm/tilexr_comm.cpp +++ b/src/comm/tilexr_comm.cpp @@ -218,7 +218,11 @@ int TileXRComm::InitSDMA() lock_guard lock(g_sdmaMtx); g_sdmaUnavailable = true; } - sdmaTransport_.reset(); + if (sdmaTransport_->Shutdown()) { + sdmaTransport_.reset(); + } else { + TILEXR_LOG(ERROR) << "TileXR SDMA cleanup incomplete after initialization failure"; + } sdmaWorkspaceDev_ = nullptr; commArgs_.sdmaWorkspacePtr = nullptr; return TILEXR_SUCCESS; @@ -231,7 +235,11 @@ int TileXRComm::InitSDMA() commArgs_.extraFlag &= ~ExtraFlag::SDMA; commArgs_.sdmaWorkspacePtr = nullptr; sdmaWorkspaceDev_ = nullptr; - sdmaTransport_.reset(); + if (sdmaTransport_->Shutdown()) { + sdmaTransport_.reset(); + } else { + TILEXR_LOG(ERROR) << "TileXR SDMA cleanup incomplete after null workspace"; + } return TILEXR_SUCCESS; } @@ -242,16 +250,31 @@ int TileXRComm::InitSDMA() return TILEXR_SUCCESS; } -void TileXRComm::ResetSDMAState() +bool TileXRComm::ResetSDMAState() { + const uint32_t oldExtraFlag = commArgs_.extraFlag; + const GM_ADDR oldWorkspace = commArgs_.sdmaWorkspacePtr; commArgs_.extraFlag &= ~ExtraFlag::SDMA; commArgs_.sdmaWorkspacePtr = nullptr; - sdmaWorkspaceDev_ = nullptr; - sdmaInitStatus_ = SDMAInitStatus::DISABLED_BY_ENV; + if (commArgsPtr_ != nullptr && UpdateCommArgsDev() != TILEXR_SUCCESS) { + commArgs_.extraFlag = oldExtraFlag; + commArgs_.sdmaWorkspacePtr = oldWorkspace; + return false; + } if (sdmaTransport_ != nullptr) { - sdmaTransport_->Shutdown(); + if (!sdmaTransport_->Shutdown()) { + return false; + } sdmaTransport_.reset(); } + sdmaWorkspaceDev_ = nullptr; + sdmaInitStatus_ = SDMAInitStatus::DISABLED_BY_ENV; + return true; +} + +bool TileXRComm::PrepareDestroy() +{ + return ResetSDMAState(); } bool TileXRComm::IsSDMAAvailable() const @@ -1170,7 +1193,7 @@ TileXRComm::~TileXRComm() udmaTransport_.reset(); } udmaInfoDev_ = nullptr; - ResetSDMAState(); + (void)ResetSDMAState(); } TileXRComm::TileXRComm(int rank, int rankSize) : rank_(rank), rankSize_(rankSize) diff --git a/src/comm/tilexr_comm.h b/src/comm/tilexr_comm.h index 066980e7..55bfcf94 100644 --- a/src/comm/tilexr_comm.h +++ b/src/comm/tilexr_comm.h @@ -51,6 +51,7 @@ class TileXRComm { bool IsSDMAAvailable() const; GM_ADDR GetSDMAWorkspacePtr() const; SDMAInitStatus GetSDMAInitStatus() const; + bool PrepareDestroy(); std::string PrintDFX(); friend class Lccl; friend class Lcoc; @@ -81,7 +82,7 @@ class TileXRComm { int InitSDMA(); int UpdateCommArgsDev(); void FreeUDMARegistry(); - void ResetSDMAState(); + bool ResetSDMAState(); private: int rank_ = 0; // global rank id diff --git a/src/include/tilexr_sdma.h b/src/include/tilexr_sdma.h index db9fdf79..fd5b6035 100644 --- a/src/include/tilexr_sdma.h +++ b/src/include/tilexr_sdma.h @@ -19,7 +19,18 @@ #define TILEXR_HAVE_PTO_SDMA 0 #endif -#if TILEXR_ASCENDC_AICORE_COMPILE && defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA +#if TILEXR_ASCENDC_AICORE_COMPILE && \ + (((defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)) || \ + (defined(CATLASS_ARCH) && (CATLASS_ARCH == 3510))) || \ + defined(TILEXR_SDMA_FORCE_A5)) +#define TILEXR_SDMA_A5_AICORE_COMPILE 1 +#include "tilexr_sdma_a5.h" +#else +#define TILEXR_SDMA_A5_AICORE_COMPILE 0 +#endif + +#if TILEXR_ASCENDC_AICORE_COMPILE && !TILEXR_SDMA_A5_AICORE_COMPILE && \ + defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA #include "tilexr_sdma_compat.h" #endif @@ -47,7 +58,15 @@ __aicore__ inline uint64_t SDMACopyNbi( uint64_t bytes, uint32_t channelGroupIdx = TILEXR_SDMA_AUTO_CHANNEL_GROUP) { -#if defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA +#if TILEXR_SDMA_A5_AICORE_COMPILE + if (!SDMAEnabled(args) || dst == nullptr || src == nullptr || bytes == 0) { + return 0; + } + const uint32_t resolvedGroup = SDMAResolveChannelGroup(channelGroupIdx); + return detail::A5SdmaCopyNbi( + reinterpret_cast<__gm__ uint8_t*>(args->sdmaWorkspacePtr), + dst, src, bytes, resolvedGroup); +#elif defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA if (!SDMAEnabled(args) || dst == nullptr || src == nullptr || bytes == 0) { return 0; } @@ -75,7 +94,18 @@ __aicore__ inline bool SDMAWait( uint64_t eventHandle, uint32_t channelGroupIdx = TILEXR_SDMA_AUTO_CHANNEL_GROUP) { -#if defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA +#if TILEXR_SDMA_A5_AICORE_COMPILE + if (eventHandle == 0) { + return true; + } + if (!SDMAEnabled(args)) { + return false; + } + const uint32_t resolvedGroup = SDMAResolveChannelGroup(channelGroupIdx); + return detail::A5SdmaWaitEvent( + reinterpret_cast<__gm__ uint8_t*>(args->sdmaWorkspacePtr), + eventHandle, resolvedGroup); +#elif defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA if (eventHandle == 0) { return true; } diff --git a/src/include/tilexr_sdma_a5.h b/src/include/tilexr_sdma_a5.h new file mode 100644 index 00000000..cce2c087 --- /dev/null +++ b/src/include/tilexr_sdma_a5.h @@ -0,0 +1,203 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_SDMA_A5_H +#define TILEXR_SDMA_A5_H + +#include "kernel_operator.h" +#include "tilexr_sdma_a5_types.h" + +namespace TileXR { +namespace detail { + +__aicore__ inline void A5SdmaCleanCacheLine(__gm__ uint8_t* address) +{ + AscendC::GlobalTensor line; + line.SetGlobalBuffer(address); + __asm__ __volatile__(""); + AscendC::DataCacheCleanAndInvalid< + uint8_t, AscendC::CacheLine::SINGLE_CACHE_LINE, + AscendC::DcciDst::CACHELINE_OUT>(line); + __asm__ __volatile__(""); +} + +__aicore__ inline void A5SdmaZeroSqe(__gm__ A5SdmaSqe* sqe) +{ + __gm__ uint32_t* words = reinterpret_cast<__gm__ uint32_t*>(sqe); + for (uint32_t index = 0U; index < TILEXR_SDMA_A5_SQE_BYTES / sizeof(uint32_t); ++index) { + words[index] = 0U; + } +} + +__aicore__ inline void A5SdmaBuildSqe(__gm__ A5SdmaSqe* sqe, + uint32_t streamId, + uint32_t taskId, + uint64_t source, + uint64_t destination, + uint32_t bytes) +{ + A5SdmaZeroSqe(sqe); + sqe->typeLockUnlock = static_cast(TILEXR_SDMA_A5_SQE_TYPE); + sqe->rtStreamId = static_cast(streamId); + sqe->taskId = static_cast(taskId); + sqe->kernelCredit = static_cast(TILEXR_SDMA_A5_KERNEL_CREDIT); + sqe->transferFlags = (1U << 8U) | (1U << 9U) | (1U << 10U) | (1U << 11U); + sqe->mpamFlags = static_cast(TILEXR_SDMA_A5_QOS << 3U); + sqe->srcAddressLow = static_cast(source & 0xFFFFFFFFULL); + sqe->srcAddressHigh = static_cast(source >> 32U); + sqe->dstAddressLow = static_cast(destination & 0xFFFFFFFFULL); + sqe->dstAddressHigh = static_cast(destination >> 32U); + sqe->length = bytes; +} + +__aicore__ inline void A5SdmaRingDoorbell(uint64_t address, uint32_t tail) +{ + __ubuf__ uint32_t value[8]; + value[0] = tail; + pipe_barrier(PIPE_ALL); + copy_ubuf_to_gm_align_v2(reinterpret_cast<__gm__ uint32_t*>(address), value, + 0, 1, sizeof(uint32_t), 0, 0, 0); + set_flag(PIPE_MTE3, PIPE_MTE2, static_cast(0)); + wait_flag(PIPE_MTE3, PIPE_MTE2, static_cast(0)); +} + +__aicore__ inline uint32_t A5SdmaReadCompletion(__gm__ A5SdmaCompletionLine* completion) +{ + __ubuf__ uint32_t line[TILEXR_SDMA_A5_COMPLETION_BYTES / sizeof(uint32_t)]; + copy_gm_to_ubuf_align_v2( + line, reinterpret_cast<__gm__ uint32_t*>(completion), + 0, 1, TILEXR_SDMA_A5_COMPLETION_BYTES, 0, 0, false, 0, 0, 0); + set_flag(PIPE_MTE2, PIPE_S, static_cast(0)); + wait_flag(PIPE_MTE2, PIPE_S, static_cast(0)); + return line[0]; +} + +__aicore__ inline bool A5SdmaWorkspaceValid(const __gm__ A5SdmaWorkspace* workspace) +{ + return workspace != nullptr && + workspace->header.magic == TILEXR_SDMA_A5_WORKSPACE_MAGIC && + workspace->header.abiVersion == TILEXR_SDMA_A5_ABI_VERSION && + workspace->header.backendKind == TILEXR_SDMA_A5_BACKEND_KIND && + workspace->header.channelCount == TILEXR_SDMA_A5_CHANNEL_COUNT && + workspace->header.sqeSize == TILEXR_SDMA_A5_SQE_BYTES && + workspace->header.channelStride == sizeof(A5SdmaChannel) && + workspace->header.workspaceSize == sizeof(A5SdmaWorkspace); +} + +__aicore__ inline uint64_t A5SdmaCopyNbi(__gm__ uint8_t* workspaceAddress, + __gm__ uint8_t* destination, + __gm__ uint8_t* source, + uint64_t bytes, + uint32_t channelIndex) +{ + __gm__ A5SdmaWorkspace* workspace = + reinterpret_cast<__gm__ A5SdmaWorkspace*>(workspaceAddress); + if (!A5SdmaWorkspaceValid(workspace) || destination == nullptr || source == nullptr || + !A5SdmaTransferLengthValid(bytes) || !A5SdmaChannelValid(channelIndex)) { + return 0ULL; + } + + __gm__ A5SdmaChannel* channel = &workspace->channels[channelIndex]; + if (channel->sqBase == 0U || channel->rtsqAddress == 0U || + channel->completionPayloadAddress == 0U || channel->completionRecordAddress == 0U || + channel->rtsqLength < sizeof(uint32_t) || + !A5SdmaQueueHasCapacity(channel->head, channel->tail, channel->depth) || + channel->streamId > 0xFFFFU) { + return 0ULL; + } + if (AscendC::AtomicCas(&channel->outstanding, 0U, 1U) != 0U) { + return 0ULL; + } + + const uint32_t generation = A5SdmaNextGeneration(channel->generation); + const uint32_t dataIndex = channel->tail; + const uint32_t completionIndex = (dataIndex + 1U) % channel->depth; + const uint32_t newTail = A5SdmaAdvanceTail(dataIndex, channel->depth); + __gm__ A5SdmaCompletionLine* payload = reinterpret_cast<__gm__ A5SdmaCompletionLine*>( + channel->completionPayloadAddress); + __gm__ A5SdmaCompletionLine* completion = reinterpret_cast<__gm__ A5SdmaCompletionLine*>( + channel->completionRecordAddress); + __gm__ uint32_t* payloadWords = reinterpret_cast<__gm__ uint32_t*>(payload); + __gm__ uint32_t* completionWords = reinterpret_cast<__gm__ uint32_t*>(completion); + for (uint32_t index = 0U; index < TILEXR_SDMA_A5_COMPLETION_BYTES / sizeof(uint32_t); ++index) { + payloadWords[index] = 0U; + completionWords[index] = 0U; + } + payload->generation = generation; + + __gm__ A5SdmaSqe* sqBase = reinterpret_cast<__gm__ A5SdmaSqe*>(channel->sqBase); + A5SdmaBuildSqe(sqBase + dataIndex, channel->streamId, channel->taskId, + reinterpret_cast(source), + reinterpret_cast(destination), + static_cast(bytes)); + A5SdmaBuildSqe(sqBase + completionIndex, channel->streamId, channel->taskId + 1U, + channel->completionPayloadAddress, + channel->completionRecordAddress, + TILEXR_SDMA_A5_COMPLETION_BYTES); + + channel->generation = generation; + channel->tail = newTail; + channel->taskId = A5SdmaAdvanceTaskId(channel->taskId); + pipe_barrier(PIPE_ALL); + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(payload)); + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(completion)); + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(sqBase + dataIndex)); + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(sqBase + completionIndex)); + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(channel)); + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(channel) + 64U); + pipe_barrier(PIPE_ALL); + dsb(DSB_DDR); + A5SdmaRingDoorbell(channel->rtsqAddress, newTail); + return A5SdmaEncodeEvent(channelIndex, generation); +} + +__aicore__ inline bool A5SdmaWaitEvent(__gm__ uint8_t* workspaceAddress, + uint64_t event, + uint32_t expectedChannel) +{ + uint32_t channelIndex = 0U; + uint32_t generation = 0U; + if (!A5SdmaDecodeEvent(event, channelIndex, generation) || channelIndex != expectedChannel) { + return false; + } + __gm__ A5SdmaWorkspace* workspace = + reinterpret_cast<__gm__ A5SdmaWorkspace*>(workspaceAddress); + if (!A5SdmaWorkspaceValid(workspace)) { + return false; + } + __gm__ A5SdmaChannel* channel = &workspace->channels[channelIndex]; + if (channel->generation != generation || + AscendC::AtomicCas(&channel->outstanding, 1U, 1U) != 1U || + channel->completionRecordAddress == 0U) { + return false; + } + __gm__ A5SdmaCompletionLine* completion = reinterpret_cast<__gm__ A5SdmaCompletionLine*>( + channel->completionRecordAddress); + const uint64_t waitBegin = static_cast(AscendC::GetSystemCycle()); + while (A5SdmaReadCompletion(completion) != generation) { + if (static_cast(AscendC::GetSystemCycle()) - waitBegin >= + TILEXR_SDMA_A5_WAIT_TIMEOUT_CYCLES) { + return false; + } + } + dsb(DSB_DDR); + if (channel->generation != generation) { + return false; + } + channel->head = channel->tail; + A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(channel)); + pipe_barrier(PIPE_ALL); + dsb(DSB_DDR); + if (channel->generation != generation || + AscendC::AtomicCas(&channel->outstanding, 1U, 0U) != 1U) { + return false; + } + return true; +} + +} // namespace detail +} // namespace TileXR + +#endif // TILEXR_SDMA_A5_H diff --git a/src/include/tilexr_sdma_a5_types.h b/src/include/tilexr_sdma_a5_types.h new file mode 100644 index 00000000..89098112 --- /dev/null +++ b/src/include/tilexr_sdma_a5_types.h @@ -0,0 +1,200 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#ifndef TILEXR_SDMA_A5_TYPES_H +#define TILEXR_SDMA_A5_TYPES_H + +#include +#include + +namespace TileXR { +namespace detail { + +#if defined(__CCE__) && defined(__CCE_IS_AICORE__) +#define TILEXR_SDMA_A5_HOST_DEVICE_INLINE __aicore__ inline +#else +#define TILEXR_SDMA_A5_HOST_DEVICE_INLINE inline +#endif + +constexpr uint32_t TILEXR_SDMA_A5_WORKSPACE_MAGIC = 0x41355344U; // "A5SD" +constexpr uint16_t TILEXR_SDMA_A5_ABI_VERSION = 1U; +constexpr uint16_t TILEXR_SDMA_A5_BACKEND_KIND = 2U; +constexpr uint32_t TILEXR_SDMA_A5_CHANNEL_COUNT = 48U; +constexpr uint32_t TILEXR_SDMA_A5_SQE_BYTES = 64U; +constexpr uint32_t TILEXR_SDMA_A5_COMPLETION_BYTES = 64U; +constexpr uint64_t TILEXR_SDMA_A5_MAX_TRANSFER_BYTES = 0xFFFFFFFFULL; +constexpr uint32_t TILEXR_SDMA_A5_SQE_TYPE = 11U; +constexpr uint32_t TILEXR_SDMA_A5_KERNEL_CREDIT = 254U; +constexpr uint32_t TILEXR_SDMA_A5_QOS = 6U; +constexpr uint64_t TILEXR_SDMA_A5_WAIT_TIMEOUT_CYCLES = 10000000000ULL; + +constexpr uint64_t TILEXR_SDMA_A5_EVENT_MAGIC = 0xA5D5ULL; +constexpr uint32_t TILEXR_SDMA_A5_EVENT_MAGIC_SHIFT = 48U; +constexpr uint32_t TILEXR_SDMA_A5_EVENT_CHANNEL_SHIFT = 40U; +constexpr uint64_t TILEXR_SDMA_A5_EVENT_RESERVED_MASK = 0x000000FF00000000ULL; +constexpr uint64_t TILEXR_SDMA_A5_EVENT_GENERATION_MASK = 0xFFFFFFFFULL; + +struct alignas(64) A5SdmaWorkspaceHeader { + uint32_t magic; + uint16_t abiVersion; + uint16_t backendKind; + uint32_t channelCount; + uint32_t sqeSize; + uint32_t channelStride; + uint32_t workspaceSize; + uint32_t maxTransferBytes; + uint32_t reserved[9]; +}; + +struct alignas(64) A5SdmaChannel { + uint64_t sqBase; + uint64_t rtsqAddress; + uint64_t completionPayloadAddress; + uint64_t completionRecordAddress; + uint32_t depth; + uint32_t head; + uint32_t tail; + uint32_t taskId; + uint32_t rtsqLength; + uint32_t streamId; + uint32_t sqId; + uint32_t cqId; + uint32_t logicalCqId; + uint32_t physicalDieId; + uint32_t generation; + uint32_t reserved0[13]; + uint32_t outstanding; + uint32_t reserved1[15]; +}; + +struct alignas(64) A5SdmaCompletionLine { + uint32_t generation; + uint32_t reserved[15]; +}; + +struct A5SdmaSqe { + uint8_t typeLockUnlock; + uint8_t controlFlags; + uint16_t numBlocks; + uint16_t rtStreamId; + uint16_t taskId; + uint32_t reserved0; + uint16_t reserved1; + uint8_t kernelCredit; + uint8_t reserved2; + uint32_t transferFlags; + uint16_t sqeId; + uint8_t mpamPartId; + uint8_t mpamFlags; + uint16_t srcStreamId; + uint16_t srcSubStreamId; + uint16_t dstStreamId; + uint16_t dstSubStreamId; + uint32_t srcAddressLow; + uint32_t srcAddressHigh; + uint32_t dstAddressLow; + uint32_t dstAddressHigh; + uint32_t length; + uint32_t srcOffsetLow; + uint32_t dstOffsetLow; + uint16_t srcOffsetHigh; + uint16_t dstOffsetHigh; +}; + +struct alignas(64) A5SdmaWorkspace { + A5SdmaWorkspaceHeader header; + A5SdmaChannel channels[TILEXR_SDMA_A5_CHANNEL_COUNT]; + A5SdmaCompletionLine completionPayloads[TILEXR_SDMA_A5_CHANNEL_COUNT]; + A5SdmaCompletionLine completionRecords[TILEXR_SDMA_A5_CHANNEL_COUNT]; +}; + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE bool A5SdmaChannelValid(uint32_t channel) +{ + return channel < TILEXR_SDMA_A5_CHANNEL_COUNT; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE bool A5SdmaTransferLengthValid(uint64_t bytes) +{ + return bytes != 0U && bytes <= TILEXR_SDMA_A5_MAX_TRANSFER_BYTES; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE bool A5SdmaQueueStateValid(uint32_t tail, uint32_t depth) +{ + return depth >= 3U && tail < depth; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint32_t A5SdmaAdvanceTail(uint32_t tail, uint32_t depth) +{ + return (tail + 2U) % depth; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint32_t A5SdmaQueueDistance( + uint32_t head, uint32_t tail, uint32_t depth) +{ + return (tail + depth - head) % depth; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE bool A5SdmaQueueHasCapacity( + uint32_t head, uint32_t tail, uint32_t depth) +{ + if (depth < 3U || head >= depth || tail >= depth) { + return false; + } + return A5SdmaQueueDistance(head, tail, depth) <= depth - 3U; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint32_t A5SdmaAdvanceTaskId(uint32_t taskId) +{ + return (taskId + 2U) & 0xFFFFU; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint32_t A5SdmaNextGeneration(uint32_t generation) +{ + return generation == 0xFFFFFFFFU ? 1U : generation + 1U; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint64_t A5SdmaEncodeEvent( + uint32_t channel, uint32_t generation) +{ + return !A5SdmaChannelValid(channel) || generation == 0U + ? 0ULL + : (TILEXR_SDMA_A5_EVENT_MAGIC << TILEXR_SDMA_A5_EVENT_MAGIC_SHIFT) | + (static_cast(channel) << TILEXR_SDMA_A5_EVENT_CHANNEL_SHIFT) | + static_cast(generation); +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE bool A5SdmaDecodeEvent( + uint64_t event, uint32_t& channel, uint32_t& generation) +{ + if ((event >> TILEXR_SDMA_A5_EVENT_MAGIC_SHIFT) != TILEXR_SDMA_A5_EVENT_MAGIC || + (event & TILEXR_SDMA_A5_EVENT_RESERVED_MASK) != 0U) { + return false; + } + channel = static_cast((event >> TILEXR_SDMA_A5_EVENT_CHANNEL_SHIFT) & 0xFFULL); + generation = static_cast(event & TILEXR_SDMA_A5_EVENT_GENERATION_MASK); + return A5SdmaChannelValid(channel) && generation != 0U; +} + +static_assert(sizeof(A5SdmaWorkspaceHeader) == 64U, "A5 workspace header must be 64 bytes"); +static_assert(sizeof(A5SdmaChannel) == 192U, "A5 channel ABI must be 192 bytes"); +static_assert(sizeof(A5SdmaCompletionLine) == 64U, "A5 completion line must be 64 bytes"); +static_assert(sizeof(A5SdmaSqe) == TILEXR_SDMA_A5_SQE_BYTES, "A5 SQE must be 64 bytes"); +static_assert(alignof(A5SdmaWorkspace) == 64U, "A5 workspace must be cache-line aligned"); +static_assert(sizeof(A5SdmaWorkspace) % 64U == 0U, "A5 workspace size must be cache-line aligned"); +static_assert(offsetof(A5SdmaChannel, sqBase) == 0U, "unexpected A5 SQ base offset"); +static_assert(offsetof(A5SdmaChannel, completionPayloadAddress) == 16U, + "unexpected A5 completion payload offset"); +static_assert(offsetof(A5SdmaChannel, generation) == 72U, "unexpected A5 generation offset"); +static_assert(offsetof(A5SdmaChannel, outstanding) == 128U, "unexpected A5 outstanding offset"); +static_assert(offsetof(A5SdmaSqe, srcAddressLow) == 32U, "unexpected A5 source address offset"); +static_assert(offsetof(A5SdmaSqe, length) == 48U, "unexpected A5 length offset"); +static_assert(offsetof(A5SdmaWorkspace, channels) == 64U, "unexpected A5 channel array offset"); + +#undef TILEXR_SDMA_A5_HOST_DEVICE_INLINE + +} // namespace detail +} // namespace TileXR + +#endif // TILEXR_SDMA_A5_TYPES_H diff --git a/tests/sdma/CMakeLists.txt b/tests/sdma/CMakeLists.txt index 17c85a30..8742331d 100644 --- a/tests/sdma/CMakeLists.txt +++ b/tests/sdma/CMakeLists.txt @@ -77,6 +77,18 @@ target_link_libraries(test_tilexr_sdma_api_invalid add_executable(test_tilexr_sdma_transport_disabled unit/test_tilexr_sdma_transport_disabled.cpp ${TILEXR_ROOT}/src/comm/sdma/tilexr_sdma_transport.cpp + ${TILEXR_ROOT}/src/comm/sdma/tilexr_sdma_a5_backend.cpp +) + +add_executable(test_tilexr_sdma_a5_validation + unit/test_tilexr_sdma_a5_validation.cpp +) +target_include_directories(test_tilexr_sdma_a5_validation PRIVATE + ${TILEXR_ROOT}/src/comm + ${TILEXR_ROOT}/src/include + ${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/ + ${ASCEND_HOME_PATH}/${ARCH}-linux/include/ + ${ASCEND_DRIVER_PATH}/kernel/inc ) target_include_directories(test_tilexr_sdma_transport_disabled PRIVATE ${TILEXR_ROOT}/src/comm @@ -86,6 +98,12 @@ target_include_directories(test_tilexr_sdma_transport_disabled PRIVATE ${ASCEND_HOME_PATH}/${ARCH}-linux/include/ ${ASCEND_DRIVER_PATH}/kernel/inc ) +target_link_libraries(test_tilexr_sdma_transport_disabled + ascendcl + runtime + ascend_hal + dl +) add_executable(test_tilexr_sdma_comm_wiring unit/test_tilexr_sdma_comm_wiring.cpp @@ -132,6 +150,7 @@ set(INSTALL_TARGETS test_tilexr_sdma_metadata test_tilexr_sdma_api_invalid test_tilexr_sdma_transport_disabled + test_tilexr_sdma_a5_validation test_tilexr_sdma_comm_wiring test_tilexr_sdma_source_guard test_tilexr_sdma_header_compile diff --git a/tests/sdma/run_tests.sh b/tests/sdma/run_tests.sh index 793cb9bb..3591e4ba 100755 --- a/tests/sdma/run_tests.sh +++ b/tests/sdma/run_tests.sh @@ -63,6 +63,7 @@ fi "${INSTALL_DIR}/bin/test_tilexr_sdma_metadata" "${INSTALL_DIR}/bin/test_tilexr_sdma_transport_disabled" +"${INSTALL_DIR}/bin/test_tilexr_sdma_a5_validation" "${INSTALL_DIR}/bin/test_tilexr_sdma_comm_wiring" "${INSTALL_DIR}/bin/test_tilexr_sdma_source_guard" "${INSTALL_DIR}/bin/test_tilexr_sdma_header_compile" diff --git a/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp b/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp new file mode 100644 index 00000000..fd33e7b6 --- /dev/null +++ b/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp @@ -0,0 +1,271 @@ +#include +#include +#include +#include + +#include "sdma/tilexr_sdma_a5_cleanup.h" +#include "sdma/tilexr_sdma_a5_backend.h" + +namespace { + +int g_failures = 0; + +#define CHECK_TRUE(expr) \ + do { \ + if (!(expr)) { \ + std::cerr << "CHECK_TRUE failed at line " << __LINE__ << ": " #expr << std::endl; \ + ++g_failures; \ + } \ + } while (0) + +TileXR::detail::A5BuiltinChannelInfo ValidChannel() +{ + TileXR::detail::A5BuiltinChannelInfo channel {}; + channel.sqHead = 1U; + channel.sqTail = 2U; + channel.sqBase = 0x1000U; + channel.sqRegisterBase = 0x2000U; + channel.sqDepth = 8U; + channel.sqId = 3U; + channel.cqId = 4U; + channel.logicalCqId = 5U; + channel.streamId = 6U; + channel.deviceId = 7U; + return channel; +} + +TileXR::detail::A5HostChannelIdentity ValidIdentity() +{ + return {6U, 3U, 4U, 5U, 7U}; +} + +void* Handle(uintptr_t value) +{ + return reinterpret_cast(value); +} + +struct FakeCleanupRuntime { + std::string failedOperation; + void* failedHandle = nullptr; + bool failureConsumed = false; + std::vector calls; + + int Call(const char* operation, void* handle) + { + calls.push_back(operation); + if (!failureConsumed && failedOperation == operation && + failedHandle == handle) { + failureConsumed = true; + return 1; + } + return 0; + } +}; + +int FakeSetContext(void* opaque, void* handle) +{ + return static_cast(opaque)->Call("set", handle); +} + +int FakeDestroyStream(void* opaque, void* handle) +{ + return static_cast(opaque)->Call("stream", handle); +} + +int FakeDestroyContext(void* opaque, void* handle) +{ + return static_cast(opaque)->Call("context", handle); +} + +int FakeFreeDevice(void* opaque, void* handle) +{ + return static_cast(opaque)->Call("free", handle); +} + +int FakeDestroyTensor(void* opaque, const void* handle) +{ + return static_cast(opaque)->Call( + "tensor", const_cast(handle)); +} + +TileXR::detail::A5QueryCleanupOps FakeCleanupOps(FakeCleanupRuntime& runtime) +{ + TileXR::detail::A5QueryCleanupOps ops; + ops.opaque = &runtime; + ops.setCurrentContext = FakeSetContext; + ops.destroyStream = FakeDestroyStream; + ops.destroyContext = FakeDestroyContext; + ops.freeDevice = FakeFreeDevice; + ops.destroyTensor = FakeDestroyTensor; + return ops; +} + +TileXR::detail::A5PendingQueryCleanup FullCleanupState() +{ + TileXR::detail::A5PendingQueryCleanup state; + state.ownerContext = Handle(1U); + state.isolatedContext = Handle(2U); + state.queryStream = Handle(3U); + state.healthStream = Handle(4U); + state.ownerBuffers = {Handle(5U), Handle(6U)}; + state.isolatedBuffers = {Handle(7U)}; + state.tensors = {Handle(8U), Handle(9U)}; + return state; +} + +bool ContainsHandle(const TileXR::detail::A5PendingQueryCleanup& state, + void* handle) +{ + if (state.isolatedContext == handle || state.queryStream == handle || + state.healthStream == handle || state.restoreContext == handle) { + return true; + } + for (void* buffer : state.ownerBuffers) { + if (buffer == handle) { + return true; + } + } + for (void* buffer : state.isolatedBuffers) { + if (buffer == handle) { + return true; + } + } + for (const void* tensor : state.tensors) { + if (tensor == handle) { + return true; + } + } + return false; +} + +void TestCompleteAndExpectedPartialClassification() +{ + using namespace TileXR::detail; + A5BuiltinChannelInfo channel = ValidChannel(); + const A5HostChannelIdentity identity = ValidIdentity(); + CHECK_TRUE(ClassifyA5QueryResult(0, 1U, 48U, channel, identity) == + A5QueryResultKind::COMPLETE); + channel.sqRegisterBase = 0U; + CHECK_TRUE(ClassifyA5QueryResult(TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS, + 0U, 0U, channel, identity) == + A5QueryResultKind::EXPECTED_PARTIAL); +} + +void TestPartialClassificationFailsClosed() +{ + using namespace TileXR::detail; + A5BuiltinChannelInfo channel = ValidChannel(); + channel.sqRegisterBase = 0U; + const A5HostChannelIdentity identity = ValidIdentity(); + CHECK_TRUE(ClassifyA5QueryResult(507019, 0U, 0U, channel, identity) == + A5QueryResultKind::INVALID); + CHECK_TRUE(ClassifyA5QueryResult(TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS, + 1U, 0U, channel, identity) == + A5QueryResultKind::INVALID); + channel.sqTail = channel.sqDepth; + CHECK_TRUE(ClassifyA5QueryResult(TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS, + 0U, 0U, channel, identity) == + A5QueryResultKind::INVALID); + channel = ValidChannel(); + channel.sqRegisterBase = 0U; + channel.streamId += 1U; + CHECK_TRUE(ClassifyA5QueryResult(TILEXR_SDMA_A5_EXPECTED_QUERY_STATUS, + 0U, 0U, channel, identity) == + A5QueryResultKind::INVALID); + channel = ValidChannel(); + channel.sqHead = 0U; + channel.sqTail = channel.sqDepth - 1U; + CHECK_TRUE(ClassifyA5QueryResult(0, 1U, TILEXR_SDMA_A5_CHANNEL_COUNT, + channel, identity) == A5QueryResultKind::INVALID); +} + +void TestCompleteClassificationRequiresFinishedHeader() +{ + using namespace TileXR::detail; + const A5BuiltinChannelInfo channel = ValidChannel(); + const A5HostChannelIdentity identity = ValidIdentity(); + CHECK_TRUE(ClassifyA5QueryResult(0, 0U, TILEXR_SDMA_A5_CHANNEL_COUNT, + channel, identity) == A5QueryResultKind::INVALID); + CHECK_TRUE(ClassifyA5QueryResult(0, 1U, TILEXR_SDMA_A5_CHANNEL_COUNT - 1U, + channel, identity) == A5QueryResultKind::INVALID); + CHECK_TRUE(ClassifyA5QueryResult(0, 1U, 0U, channel, identity) == + A5QueryResultKind::INVALID); +} + +void TestGroupedCopyWaitIsBounded() +{ + CHECK_TRUE(TileXR::detail::TILEXR_SDMA_A5_WAIT_TIMEOUT_CYCLES != 0ULL); +} + +void TestCleanupFailuresRetainHandlesForRetry() +{ + struct FailureCase { + const char* operation; + void* handle; + }; + const std::vector failures = { + {"set", Handle(2U)}, + {"stream", Handle(4U)}, + {"stream", Handle(3U)}, + {"free", Handle(7U)}, + {"tensor", Handle(9U)}, + {"context", Handle(2U)}, + {"free", Handle(6U)}, + }; + for (const FailureCase& failure : failures) { + FakeCleanupRuntime runtime; + runtime.failedOperation = failure.operation; + runtime.failedHandle = failure.handle; + TileXR::detail::A5PendingQueryCleanup state = FullCleanupState(); + const TileXR::detail::A5QueryCleanupOps ops = FakeCleanupOps(runtime); + CHECK_TRUE(!TileXR::detail::CleanupA5QueryResources( + state, ops, state.ownerContext)); + CHECK_TRUE(runtime.failureConsumed); + CHECK_TRUE(ContainsHandle(state, failure.handle)); + CHECK_TRUE(!state.Empty()); + + CHECK_TRUE(TileXR::detail::CleanupA5QueryResources( + state, ops, state.ownerContext)); + CHECK_TRUE(state.Empty()); + } +} + +void TestCleanupRestoreFailureIsRetryable() +{ + FakeCleanupRuntime runtime; + runtime.failedOperation = "set"; + runtime.failedHandle = Handle(1U); + TileXR::detail::A5PendingQueryCleanup state; + state.ownerContext = Handle(1U); + state.isolatedContext = Handle(2U); + state.queryStream = Handle(3U); + const TileXR::detail::A5QueryCleanupOps ops = FakeCleanupOps(runtime); + + CHECK_TRUE(!TileXR::detail::CleanupA5QueryResources( + state, ops, state.ownerContext)); + CHECK_TRUE(state.restorePending); + CHECK_TRUE(state.restoreContext == state.ownerContext); + CHECK_TRUE(!state.Empty()); + + CHECK_TRUE(TileXR::detail::CleanupA5QueryResources( + state, ops, state.restoreContext)); + CHECK_TRUE(state.Empty()); +} + +} // namespace + +int main() +{ + TestCompleteAndExpectedPartialClassification(); + TestPartialClassificationFailsClosed(); + TestCompleteClassificationRequiresFinishedHeader(); + TestGroupedCopyWaitIsBounded(); + TestCleanupFailuresRetainHandlesForRetry(); + TestCleanupRestoreFailureIsRetryable(); + if (g_failures != 0) { + std::cerr << g_failures << " A5 SDMA validation checks failed" << std::endl; + return 1; + } + std::cout << "TileXR A5 SDMA validation checks passed" << std::endl; + return 0; +} diff --git a/tests/sdma/unit/test_tilexr_sdma_transport_disabled.cpp b/tests/sdma/unit/test_tilexr_sdma_transport_disabled.cpp index 7e85967d..2504b22b 100644 --- a/tests/sdma/unit/test_tilexr_sdma_transport_disabled.cpp +++ b/tests/sdma/unit/test_tilexr_sdma_transport_disabled.cpp @@ -2,6 +2,7 @@ #include #include +#include "acl/acl_rt.h" #include "sdma/tilexr_sdma_transport.h" #include "tilexr_sdma_types.h" #include "tilexr_types.h" @@ -84,7 +85,27 @@ void TestEnvZeroSkipsInitialization() CHECK_EQ(transport.GetLastStatus(), TileXR::SDMAInitStatus::DISABLED_BY_ENV); } -void TestEnvOneReportsPTOUnavailable() +TileXR::SDMAInitStatus ExpectedUnavailableStatus() +{ + return TileXR::detail::ClassifySDMABackend(aclrtGetSocName()) == + TileXR::detail::SDMABackendKind::A5_DIRECT + ? TileXR::SDMAInitStatus::INIT_FAILED + : TileXR::SDMAInitStatus::PTO_UNAVAILABLE; +} + +void TestBackendClassification() +{ + using TileXR::detail::ClassifySDMABackend; + using TileXR::detail::SDMABackendKind; + CHECK_EQ(ClassifySDMABackend("Ascend950PR_9589"), SDMABackendKind::A5_DIRECT); + CHECK_EQ(ClassifySDMABackend("Ascend910B3"), SDMABackendKind::PTO); + CHECK_EQ(ClassifySDMABackend("Ascend910A"), SDMABackendKind::PTO); + CHECK_EQ(ClassifySDMABackend("Ascend910_9391"), SDMABackendKind::PTO); + CHECK_EQ(ClassifySDMABackend("Ascend310P"), SDMABackendKind::UNSUPPORTED); + CHECK_EQ(ClassifySDMABackend(nullptr), SDMABackendKind::UNSUPPORTED); +} + +void TestEnvOneReportsBackendUnavailable() { EnvGuard env; setenv("TILEXR_ENABLE_SDMA", "1", 1); @@ -94,7 +115,7 @@ void TestEnvOneReportsPTOUnavailable() CHECK_EQ(transport.Init(options), TileXR::TILEXR_SUCCESS); CHECK_TRUE(!transport.IsAvailable()); CHECK_TRUE(transport.GetWorkspaceDev() == nullptr); - CHECK_EQ(transport.GetLastStatus(), TileXR::SDMAInitStatus::PTO_UNAVAILABLE); + CHECK_EQ(transport.GetLastStatus(), ExpectedUnavailableStatus()); } void TestSameInstanceTransitionsResetState() @@ -114,12 +135,12 @@ void TestSameInstanceTransitionsResetState() CHECK_EQ(transport.Init(options), TileXR::TILEXR_SUCCESS); CHECK_TRUE(!transport.IsAvailable()); CHECK_TRUE(transport.GetWorkspaceDev() == nullptr); - CHECK_EQ(transport.GetLastStatus(), TileXR::SDMAInitStatus::PTO_UNAVAILABLE); + CHECK_EQ(transport.GetLastStatus(), ExpectedUnavailableStatus()); transport.Shutdown(); CHECK_TRUE(!transport.IsAvailable()); CHECK_TRUE(transport.GetWorkspaceDev() == nullptr); - CHECK_EQ(transport.GetLastStatus(), TileXR::SDMAInitStatus::PTO_UNAVAILABLE); + CHECK_EQ(transport.GetLastStatus(), ExpectedUnavailableStatus()); unsetenv("TILEXR_ENABLE_SDMA"); CHECK_EQ(transport.Init(options), TileXR::TILEXR_SUCCESS); @@ -134,7 +155,8 @@ int main() { TestEnvDisabledSkipsInitialization(); TestEnvZeroSkipsInitialization(); - TestEnvOneReportsPTOUnavailable(); + TestBackendClassification(); + TestEnvOneReportsBackendUnavailable(); TestSameInstanceTransitionsResetState(); if (g_failures != 0) { std::cerr << g_failures << " SDMA transport disabled checks failed" << std::endl; diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index 5f307e57..ec0edc5c 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -226,8 +226,6 @@ if(BUILD_TILEXR_UDMA_DEMO) endif() set(TILEXR_UDMA_DEMO_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_udma_demo_kernel.so") - set(TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO - "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_udma_alltoall_group_kernel.so") set(TILEXR_UDMA_DEMO_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/compiler/tikcpp -I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw @@ -282,62 +280,100 @@ if(BUILD_TILEXR_UDMA_DEMO) ) add_custom_target(tilexr_udma_demo_kernel ALL DEPENDS "${TILEXR_UDMA_DEMO_KERNEL_SO}") + list(APPEND CMAKE_MODULE_PATH "${TILEXR_ROOT}/src/collectives/kernels/cmake") + enable_language(CCE) + set(TILEXR_GROUPED_ALLTOALL_KERNEL_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp") + set_source_files_properties( + "${TILEXR_GROUPED_ALLTOALL_KERNEL_SOURCE}" PROPERTIES LANGUAGE CCE) + add_library(tilexr_udma_alltoall_group_device OBJECT + "${TILEXR_GROUPED_ALLTOALL_KERNEL_SOURCE}") + target_compile_options(tilexr_udma_alltoall_group_device PRIVATE + -O2 + -std=gnu++17 + --cce-aicore-only + -Wno-deprecated-declarations + "SHELL:-mllvm -cce-aicore-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-function-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-record-overflow=true" + "SHELL:-mllvm -cce-aicore-addr-transform" + "SHELL:-mllvm -cce-aicore-dcci-insert-for-scalar=false" + ${TILEXR_UDMA_AICORE_ARCH}) + target_compile_definitions(tilexr_udma_alltoall_group_device PRIVATE + CATLASS_ARCH=${TILEXR_UDMA_CATLASS_ARCH} + TILEXR_UDMA_FORCE_ENABLE=1) + target_include_directories(tilexr_udma_alltoall_group_device PRIVATE + ${ASCEND_HOME_PATH}/compiler/tikcpp + ${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw + ${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw/impl + ${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw/interface + ${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp + ${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw + ${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw/impl + ${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw/interface + ${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/ + ${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/runtime/ + ${ASCEND_HOME_PATH}/${ARCH}-linux/include/ + ${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include/ + ${ASCEND_DRIVER_PATH}/kernel/inc + ${TILEXR_ROOT}/3rdparty + ${TILEXR_ROOT}/src/include + ${CMAKE_CURRENT_SOURCE_DIR}/demo) + + set(TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT + "${CMAKE_CURRENT_BINARY_DIR}/tilexr_udma_alltoall_group_kernel.o") add_custom_command( - OUTPUT "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" - COMMAND ${BISHENG_EXECUTABLE} - ${TILEXR_UDMA_KERNEL_COMPILE_OPTIONS} - -std=gnu++17 - -fPIC - -shared - ${TILEXR_UDMA_KERNEL_LINK_OPTIONS} - -DCATLASS_ARCH=${TILEXR_UDMA_CATLASS_ARCH} - -DTILEXR_UDMA_FORCE_ENABLE=1 - ${TILEXR_UDMA_DEMO_KERNEL_INCLUDES} - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" - -L${ASCEND_DRIVER_PATH}/lib64/driver - -L${ASCEND_HOME_PATH}/${ARCH}-linux/lib64 - -L${ASCEND_HOME_PATH}/${ARCH}-linux/devlib - -lruntime - -lascendcl - -lstdc++ - -lm - -ltiling_api - -lplatform - -lc_sec - -ldl - -lnnopbase - -lpthread - -o "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" - DEPENDS - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_kernel.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_layout.h" - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_route.h" - "${CMAKE_CURRENT_SOURCE_DIR}/demo/tilexr_udma_alltoall_group_trace.h" - "${TILEXR_ROOT}/src/include/tilexr_udma.h" - VERBATIM - COMMENT "Building TileXR grouped AllToAll kernel with bisheng" - ) - add_custom_target(tilexr_udma_alltoall_group_kernel ALL - DEPENDS "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}") + OUTPUT "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}" + COMMAND ${CMAKE_CCE_LINKER} -m aicorelinux -Ttext=0 + $ + --static -o "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}" + --allow-multiple-definition + COMMAND truncate -c -s 8388608 + "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}" + DEPENDS tilexr_udma_alltoall_group_device + COMMAND_EXPAND_LISTS + VERBATIM) + add_custom_target(tilexr_udma_alltoall_group_device_binary + DEPENDS "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}") + + set(TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP + "${CMAKE_CURRENT_BINARY_DIR}/tilexr_udma_alltoall_group_kernel_embed.cpp") + add_custom_command( + OUTPUT "${TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP}" + COMMAND ${CMAKE_COMMAND} + -DTILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT=${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT} + -DTILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP=${TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP} + -P "${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_grouped_alltoall_kernel.cmake" + DEPENDS tilexr_udma_alltoall_group_device_binary + "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}" + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_grouped_alltoall_kernel.cmake" + VERBATIM) + + add_library(tilexr_udma_alltoall_group_kernel SHARED + demo/tilexr_udma_alltoall_group_launcher.cpp + "${TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP}") + add_dependencies(tilexr_udma_alltoall_group_kernel + tilexr_udma_alltoall_group_device_binary) + target_link_libraries(tilexr_udma_alltoall_group_kernel PRIVATE runtime) + set_target_properties(tilexr_udma_alltoall_group_kernel PROPERTIES + OUTPUT_NAME tilexr_udma_alltoall_group_kernel) add_executable(tilexr_udma_demo demo/tilexr_udma_demo.cpp ) - add_dependencies(tilexr_udma_demo - tilexr_udma_demo_kernel tilexr_udma_alltoall_group_kernel) + add_dependencies(tilexr_udma_demo tilexr_udma_demo_kernel) target_link_directories(tilexr_udma_demo PRIVATE ${CMAKE_CURRENT_BINARY_DIR}) target_link_libraries(tilexr_udma_demo "${TILEXR_UDMA_DEMO_KERNEL_SO}" - "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" + tilexr_udma_alltoall_group_kernel ${TILEXR_LIB} ascendcl runtime ascend_hal ) - list(APPEND INSTALL_TARGETS tilexr_udma_demo) + list(APPEND INSTALL_TARGETS + tilexr_udma_demo tilexr_udma_alltoall_group_kernel) install(FILES "${TILEXR_UDMA_DEMO_KERNEL_SO}" DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) - install(FILES "${TILEXR_UDMA_ALLTOALL_GROUP_KERNEL_SO}" - DESTINATION ${CMAKE_INSTALL_PREFIX}/lib) message(STATUS "TileXR UDMA demo enabled with ${BISHENG_EXECUTABLE}, SOC=${TILEXR_UDMA_DEMO_SOC_TYPE}") endif() endif() diff --git a/tests/udma/cmake/embed_grouped_alltoall_kernel.cmake b/tests/udma/cmake/embed_grouped_alltoall_kernel.cmake new file mode 100644 index 00000000..8afb7262 --- /dev/null +++ b/tests/udma/cmake/embed_grouped_alltoall_kernel.cmake @@ -0,0 +1,37 @@ +# Copyright (c) 2024-2026 TileXR Project +# Licensed under the Apache License, Version 2.0 + +if(NOT DEFINED TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT) + message(FATAL_ERROR "TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT is required") +endif() +if(NOT DEFINED TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP) + message(FATAL_ERROR "TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP is required") +endif() +if(NOT EXISTS "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}") + message(FATAL_ERROR + "Grouped AllToAll CCE object not found: ${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}") +endif() + +file(SIZE "${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}" _tilexr_grouped_kernel_size) +file(WRITE "${TILEXR_GROUPED_ALLTOALL_KERNEL_EMBED_CPP}" "/* + * Generated by tests/udma/cmake/embed_grouped_alltoall_kernel.cmake. + */ +#include +#include + +extern \"C\" { +extern const std::size_t TileXRGroupedAllToAllKernelBinarySize = ${_tilexr_grouped_kernel_size}; +} + +asm(R\"( +.section .rodata, \"a\", @progbits +.balign 16 +.global TileXRGroupedAllToAllKernelBinaryData +.type TileXRGroupedAllToAllKernelBinaryData, @object +TileXRGroupedAllToAllKernelBinaryData: +.incbin \"${TILEXR_GROUPED_ALLTOALL_KERNEL_OBJECT}\" +TileXRGroupedAllToAllKernelBinaryDataEnd: +.size TileXRGroupedAllToAllKernelBinaryData, TileXRGroupedAllToAllKernelBinaryDataEnd - TileXRGroupedAllToAllKernelBinaryData +.previous +)\"); +") diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index fcbf8448..45f28c36 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -4,6 +4,7 @@ */ #include "kernel_operator.h" +#include "tilexr_sdma.h" #include "tilexr_udma.h" #include "tilexr_udma_alltoall_group_route.h" #include "tilexr_udma_alltoall_group_trace.h" @@ -37,6 +38,10 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CONFIG = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT = 4U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_SDMA = 5U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK = 0U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SDMA_COMPLETE = 1U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_SDMA_FAILED = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED = 0U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_LOCAL = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY = 2U; @@ -291,6 +296,24 @@ __aicore__ inline void AllToAllGroupCopyMte( AscendC::PipeBarrier(); } +__aicore__ inline uint32_t AllToAllGroupCopySdma( + const __gm__ TileXR::CommArgs* args, __gm__ uint8_t* dst, + __gm__ uint8_t* src, uint32_t bytes, uint32_t channel, + uint64_t& event) +{ + event = 0ULL; + if (!TileXR::SDMAEnabled(args)) { + return TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK; + } + event = TileXR::SDMACopyNbi(args, dst, src, bytes, channel); + if (event == 0ULL) { + return TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK; + } + return TileXR::SDMAWait(args, event, channel) ? + TILEXR_ALLTOALL_GROUP_SDMA_COMPLETE : + TILEXR_ALLTOALL_GROUP_SDMA_FAILED; +} + __aicore__ inline uint64_t AllToAllGroupLoadTokenMte( __gm__ uint64_t* signal, AscendC::LocalTensor relayLocal) { @@ -638,7 +661,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, uint32_t groupWidth, uint32_t quietBatch) { - constexpr uint32_t copyoutWorkers = 32U; + constexpr uint32_t copyoutWorkers = 1U; const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? reinterpret_cast<__gm__ uint8_t*>(groupTraceGM) : nullptr; @@ -651,7 +674,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( const int32_t rank = args->rank; const int32_t rankSize = args->rankSize; - if ((copyoutWorkers != 8U && copyoutWorkers != 16U && + if ((copyoutWorkers != 1U && copyoutWorkers != 8U && copyoutWorkers != 16U && copyoutWorkers != 32U && copyoutWorkers != 48U) || routeStage > TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_NO_COPY || multiChannel > 1U || @@ -776,7 +799,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( static_cast(elementsPerPeer), primaryWeight, secondaryWeight, primaryRouteParts, primaryTotalElements, secondaryTotalElements); - const uint32_t traceCore = copyoutWorkers == 8U ? + const uint32_t traceCore = copyoutWorkers < TILEXR_ALLTOALL_GROUP_SEND_CORES ? TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx; for (uint32_t pass = 0U; pass < passCount; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; @@ -890,12 +913,27 @@ __aicore__ inline void AllToAllGroupKernelImpl( output + static_cast(peer) * elementsPerPeer + chunkElementOffset + copyElementBegin); const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); - AllToAllGroupCopyMte(relayDst, relaySrc, copyBytes, relayLocal); + uint64_t sdmaEvent = 0ULL; + const uint32_t sdmaStatus = AllToAllGroupCopySdma( + args, relayDst, relaySrc, copyBytes, worker, sdmaEvent); + if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK) { + AllToAllGroupCopyMte(relayDst, relaySrc, copyBytes, relayLocal); + } AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveCopy, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); + if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FAILED) { + AllToAllGroupRecordError( + debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_SDMA, + group, pass, peer, worker, 0U, + static_cast(worker), sdmaEvent); + AllToAllGroupTraceRecordKernel( + groupTrace, traceIteration, blockIdx, kernelBegin, + AllToAllGroupTraceCycle(groupTrace)); + return; + } } } } @@ -1159,48 +1197,3 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_credit_ groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, groupWidth, quietBatch); } - -void launch_tilexr_udma_all_to_all_group( - uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, - GM_ADDR registeredMemory, GM_ADDR debug, uint32_t invocationId, - int32_t elementsPerPeer, int32_t chunkElements, - uint32_t passCount, uint32_t groupCount, - uint64_t payloadOffset0, uint64_t payloadOffset1, - uint64_t signalOffset0, uint64_t signalOffset1, - uint64_t creditOffset0, uint64_t creditOffset1, - GM_ADDR groupTrace, uint32_t traceIteration, - uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) -{ - if (ingressWindow == 0U && quietBatch == 1U) { - tilexr_udma_all_to_all_group_kernel<<>>( - commArgs, input, output, registeredMemory, debug, invocationId, - elementsPerPeer, chunkElements, passCount, groupCount, - payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, routeStage, - multiChannel, primaryRouteParts, groupWidth, quietBatch); - } else if (ingressWindow == 0U) { - tilexr_udma_all_to_all_group_batch_kernel<<>>( - commArgs, input, output, registeredMemory, debug, invocationId, - elementsPerPeer, chunkElements, passCount, groupCount, - payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - groupTrace, traceIteration, routeStage, - multiChannel, primaryRouteParts, groupWidth, quietBatch); - } else if (quietBatch == 1U) { - tilexr_udma_all_to_all_group_credit_kernel<<>>( - commArgs, input, output, registeredMemory, debug, invocationId, - elementsPerPeer, chunkElements, passCount, groupCount, - payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - creditOffset0, creditOffset1, - groupTrace, traceIteration, routeStage, - multiChannel, primaryRouteParts, groupWidth, quietBatch, ingressWindow); - } else { - tilexr_udma_all_to_all_group_batch_credit_kernel<<>>( - commArgs, input, output, registeredMemory, debug, invocationId, - elementsPerPeer, chunkElements, passCount, groupCount, - payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, - creditOffset0, creditOffset1, - groupTrace, traceIteration, routeStage, - multiChannel, primaryRouteParts, groupWidth, quietBatch, ingressWindow); - } -} diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp new file mode 100644 index 00000000..0c72b247 --- /dev/null +++ b/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp @@ -0,0 +1,227 @@ +/* + * Copyright (c) 2024-2026 TileXR Project + * Licensed under the Apache License, Version 2.0 + */ + +#include +#include +#include + +#include "runtime/kernel.h" + +extern "C" { +extern const unsigned char TileXRGroupedAllToAllKernelBinaryData[]; +extern const std::size_t TileXRGroupedAllToAllKernelBinarySize; +} + +namespace { + +constexpr uint32_t TILEXR_RT_DEV_BINARY_MAGIC_ELF_AIVEC = 0x41415246U; +constexpr char TILEXR_GROUPED_ALLTOALL_KERNEL_NAME[] = + "tilexr_udma_all_to_all_group_kernel"; +constexpr char TILEXR_GROUPED_ALLTOALL_BATCH_KERNEL_NAME[] = + "tilexr_udma_all_to_all_group_batch_kernel"; +constexpr char TILEXR_GROUPED_ALLTOALL_CREDIT_KERNEL_NAME[] = + "tilexr_udma_all_to_all_group_credit_kernel"; +constexpr char TILEXR_GROUPED_ALLTOALL_BATCH_CREDIT_KERNEL_NAME[] = + "tilexr_udma_all_to_all_group_batch_credit_kernel"; + +struct alignas(8) GroupedAllToAllKernelArgs { + uint8_t* commArgs; + uint8_t* input; + uint8_t* output; + uint8_t* registeredMemory; + uint8_t* debug; + uint32_t invocationId; + int32_t elementsPerPeer; + int32_t chunkElements; + uint32_t passCount; + uint32_t groupCount; + uint64_t payloadOffset0; + uint64_t payloadOffset1; + uint64_t signalOffset0; + uint64_t signalOffset1; + uint8_t* groupTrace; + uint32_t traceIteration; + uint32_t routeStage; + uint32_t multiChannel; + uint32_t primaryRouteParts; + uint32_t groupWidth; + uint32_t quietBatch; +}; + +struct alignas(8) GroupedAllToAllCreditKernelArgs { + uint8_t* commArgs; + uint8_t* input; + uint8_t* output; + uint8_t* registeredMemory; + uint8_t* debug; + uint32_t invocationId; + int32_t elementsPerPeer; + int32_t chunkElements; + uint32_t passCount; + uint32_t groupCount; + uint64_t payloadOffset0; + uint64_t payloadOffset1; + uint64_t signalOffset0; + uint64_t signalOffset1; + uint64_t creditOffset0; + uint64_t creditOffset1; + uint8_t* groupTrace; + uint32_t traceIteration; + uint32_t routeStage; + uint32_t multiChannel; + uint32_t primaryRouteParts; + uint32_t groupWidth; + uint32_t quietBatch; + uint32_t ingressWindow; +}; + +static_assert(sizeof(GroupedAllToAllKernelArgs) == 128U, + "grouped alltoall kernel argument ABI changed"); +static_assert(offsetof(GroupedAllToAllKernelArgs, payloadOffset0) == 64U, + "grouped alltoall payload offset ABI changed"); +static_assert(offsetof(GroupedAllToAllKernelArgs, groupTrace) == 96U, + "grouped alltoall trace argument ABI changed"); +static_assert(sizeof(GroupedAllToAllCreditKernelArgs) == 152U, + "grouped alltoall credit kernel argument ABI changed"); +static_assert(offsetof(GroupedAllToAllCreditKernelArgs, creditOffset0) == 96U, + "grouped alltoall credit offset ABI changed"); +static_assert(offsetof(GroupedAllToAllCreditKernelArgs, groupTrace) == 112U, + "grouped alltoall credit trace argument ABI changed"); + +std::mutex g_groupedKernelRegistrationMutex; +bool g_groupedKernelRegistered = false; +rtError_t g_groupedKernelRegistrationStatus = RT_ERROR_NONE; +void* g_groupedKernelBinaryHandle = nullptr; +uint8_t g_groupedKernelFunctionToken = 0U; +uint8_t g_groupedBatchKernelFunctionToken = 0U; +uint8_t g_groupedCreditKernelFunctionToken = 0U; +uint8_t g_groupedBatchCreditKernelFunctionToken = 0U; + +int8_t* GroupedKernelFunctionSignature() +{ + return reinterpret_cast(&g_groupedKernelFunctionToken); +} + +int8_t* GroupedBatchKernelFunctionSignature() +{ + return reinterpret_cast(&g_groupedBatchKernelFunctionToken); +} + +int8_t* GroupedCreditKernelFunctionSignature() +{ + return reinterpret_cast(&g_groupedCreditKernelFunctionToken); +} + +int8_t* GroupedBatchCreditKernelFunctionSignature() +{ + return reinterpret_cast(&g_groupedBatchCreditKernelFunctionToken); +} + +rtError_t RegisterGroupedKernel(int8_t* signature, const char* name) +{ + return rtFunctionRegister( + g_groupedKernelBinaryHandle, signature, name, name, 0U); +} + +rtError_t EnsureGroupedKernelRegistered() +{ + std::lock_guard guard(g_groupedKernelRegistrationMutex); + if (g_groupedKernelRegistered) { + return g_groupedKernelRegistrationStatus; + } + + rtDevBinary_t binary {}; + binary.data = TileXRGroupedAllToAllKernelBinaryData; + binary.length = static_cast(TileXRGroupedAllToAllKernelBinarySize); + binary.magic = TILEXR_RT_DEV_BINARY_MAGIC_ELF_AIVEC; + binary.version = 0U; + + g_groupedKernelRegistrationStatus = + rtDevBinaryRegister(&binary, &g_groupedKernelBinaryHandle); + if (g_groupedKernelRegistrationStatus != RT_ERROR_NONE) { + return g_groupedKernelRegistrationStatus; + } + g_groupedKernelRegistrationStatus = RegisterGroupedKernel( + GroupedKernelFunctionSignature(), TILEXR_GROUPED_ALLTOALL_KERNEL_NAME); + if (g_groupedKernelRegistrationStatus != RT_ERROR_NONE) { + return g_groupedKernelRegistrationStatus; + } + g_groupedKernelRegistrationStatus = RegisterGroupedKernel( + GroupedBatchKernelFunctionSignature(), + TILEXR_GROUPED_ALLTOALL_BATCH_KERNEL_NAME); + if (g_groupedKernelRegistrationStatus != RT_ERROR_NONE) { + return g_groupedKernelRegistrationStatus; + } + g_groupedKernelRegistrationStatus = RegisterGroupedKernel( + GroupedCreditKernelFunctionSignature(), + TILEXR_GROUPED_ALLTOALL_CREDIT_KERNEL_NAME); + if (g_groupedKernelRegistrationStatus != RT_ERROR_NONE) { + return g_groupedKernelRegistrationStatus; + } + g_groupedKernelRegistrationStatus = RegisterGroupedKernel( + GroupedBatchCreditKernelFunctionSignature(), + TILEXR_GROUPED_ALLTOALL_BATCH_CREDIT_KERNEL_NAME); + if (g_groupedKernelRegistrationStatus == RT_ERROR_NONE) { + g_groupedKernelRegistered = true; + } + return g_groupedKernelRegistrationStatus; +} + +} // namespace + +int launch_tilexr_udma_all_to_all_group( + uint32_t blockDim, void* stream, uint8_t* commArgs, uint8_t* input, uint8_t* output, + uint8_t* registeredMemory, uint8_t* debug, uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, + uint8_t* groupTrace, uint32_t traceIteration, + uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) +{ + const rtError_t registrationStatus = EnsureGroupedKernelRegistered(); + if (registrationStatus != RT_ERROR_NONE) { + return static_cast(registrationStatus); + } + + const bool useCredit = ingressWindow != 0U; + const bool useBatch = quietBatch != 1U; + GroupedAllToAllKernelArgs args {}; + GroupedAllToAllCreditKernelArgs creditArgs {}; + int8_t* functionSignature = nullptr; + if (useCredit) { + creditArgs = { + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + creditOffset0, creditOffset1, groupTrace, traceIteration, + routeStage, multiChannel, primaryRouteParts, groupWidth, + quietBatch, ingressWindow, + }; + functionSignature = useBatch ? GroupedBatchCreditKernelFunctionSignature() : + GroupedCreditKernelFunctionSignature(); + } else { + args = { + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + groupTrace, traceIteration, routeStage, multiChannel, + primaryRouteParts, groupWidth, quietBatch, + }; + functionSignature = useBatch ? GroupedBatchKernelFunctionSignature() : + GroupedKernelFunctionSignature(); + } + rtArgsEx_t argsInfo {}; + argsInfo.args = useCredit ? static_cast(&creditArgs) : static_cast(&args); + argsInfo.argsSize = useCredit ? sizeof(creditArgs) : sizeof(args); + + rtTaskCfgInfo_t cfgInfo {}; + cfgInfo.schemMode = RT_SCHEM_MODE_NORMAL; + return static_cast(rtKernelLaunchWithFlagV2( + functionSignature, blockDim, &argsInfo, nullptr, + static_cast(stream), 0U, &cfgInfo)); +} diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index b02b9472..e8a3ab25 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -132,7 +132,8 @@ inline bool AllToAllGroupValidRankSize(int rankSize) inline bool AllToAllGroupValidCopyoutWorkers(uint32_t workers) { - return workers == 8U || workers == 16U || workers == 32U || workers == 48U; + return workers == 1U || workers == 8U || workers == 16U || + workers == 32U || workers == 48U; } inline uint32_t AllToAllGroupBlockDim(uint32_t sendWorkers, uint32_t copyoutWorkers) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index ce2c2960..a6d8f935 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -12,6 +12,8 @@ HEADER_BYTES = 4096 MAX_ITERATIONS = 50 MAX_CORES = 64 +SEND_WORKER_COUNT = 32 +LANE_COUNT = 16 PHASE_COUNT = 5 CURRENT_PHASE_COUNT = 6 SPAN_BYTES = 16 @@ -190,14 +192,14 @@ def build_chrome_trace(rank_traces): sources.append(rank_trace["path"]) events.append(_metadata("process_name", rank, 0, f"rank {rank}")) for core in range(MAX_CORES): - role = "send" if core < 32 else "receive" - events.append(_metadata("thread_name", rank, core, f"core{core} {role}")) + role = "send" if core < SEND_WORKER_COUNT else "receive" + events.append(_metadata("thread_name", rank, core, f"core{core:02d} {role}")) for iteration in range(header["iteration_count"]): base = bases[(rank, iteration)] offset_us = iteration_offsets[iteration] for core in range(MAX_CORES): - role = "send" if core < 32 else "receive" + role = "send" if core < SEND_WORKER_COUNT else "receive" kernel = _read_span( data, kernel_span_offset(iteration, core), f"kernel rank={rank} iter={iteration} core={core}") @@ -230,7 +232,7 @@ def build_chrome_trace(rank_traces): "iteration": iteration, "group": group, "pass": pass_index, - "lane": core % 16 if core < 32 else core - 32, + "lane": core % LANE_COUNT, "peer": peer, "qp": None if qp == NO_QP else qp, "role": role, diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index c7711ba7..498f6096 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -55,7 +55,7 @@ extern void launch_tilexr_udma_all_to_all_fused( GM_ADDR udmaMem, GM_ADDR signal, GM_ADDR debug, int32_t elementsPerPeer, uint64_t udmaMemByteOffset, uint64_t signalByteOffsetBase, int32_t chunkElements, uint32_t passCount, uint32_t loopCount); -extern void launch_tilexr_udma_all_to_all_group( +extern int launch_tilexr_udma_all_to_all_group( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR registeredMemory, GM_ADDR debug, uint32_t invocationId, int32_t elementsPerPeer, int32_t chunkElements, @@ -1027,7 +1027,7 @@ bool RunGroupedAllToAll( static_cast(primaryRoutePartsValue); constexpr uint32_t sendWorkers = TileXR::Demo::kAllToAllGroupSendWorkerCount; - constexpr uint32_t copyoutWorkers = 32U; + constexpr uint32_t copyoutWorkers = 1U; const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim( sendWorkers, copyoutWorkers); const int routeStagesValue = GetEnvInt( @@ -1251,8 +1251,8 @@ bool RunGroupedAllToAll( uint32_t invocationId = 0U; auto launchGroupStage = [&](TileXR::Demo::AllToAllGroupRouteStage routeStage, - void* trace, uint32_t traceIteration) { - launch_tilexr_udma_all_to_all_group( + void* trace, uint32_t traceIteration) -> bool { + const int launchRet = launch_tilexr_udma_all_to_all_group( groupBlockDim, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), reinterpret_cast(registeredMemory), reinterpret_cast(debug), @@ -1267,13 +1267,24 @@ bool RunGroupedAllToAll( groupWidth, quietBatch, routeStage == TileXR::Demo::AllToAllGroupRouteStage::kCombined ? ingressWindow : 0U); + if (launchRet != 0) { + std::cerr << "[rank " << rank + << "] rtKernelLaunchWithFlagV2 grouped failed: " + << launchRet << std::endl; + return false; + } + return true; }; double totalUs = 0.0; std::array stageTotalUs {}; if (!routeStages) { for (int iter = 0; iter < warmup; ++iter, ++invocationId) { - launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, nullptr, 0U); + if (!launchGroupStage( + TileXR::Demo::AllToAllGroupRouteStage::kCombined, nullptr, 0U)) { + release(); + return false; + } } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped warmup", aclrtSynchronizeStream(stream))) { @@ -1287,8 +1298,11 @@ bool RunGroupedAllToAll( const auto begin = std::chrono::steady_clock::now(); for (int iter = 0; iter < repeat; ++iter, ++invocationId) { - launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, - groupTraceDevices[0], static_cast(iter)); + if (!launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, + groupTraceDevices[0], static_cast(iter))) { + release(); + return false; + } } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped measured", aclrtSynchronizeStream(stream))) { @@ -1300,7 +1314,9 @@ bool RunGroupedAllToAll( } else { auto runStageBatch = [&](size_t stageIndex) -> bool { for (int iter = 0; iter < warmup; ++iter, ++invocationId) { - launchGroupStage(stagedRouteStages[stageIndex], nullptr, 0U); + if (!launchGroupStage(stagedRouteStages[stageIndex], nullptr, 0U)) { + return false; + } } if (!CheckAcl(rank, "aclrtSynchronizeStream grouped stage warmup", aclrtSynchronizeStream(stream)) || @@ -1309,8 +1325,10 @@ bool RunGroupedAllToAll( return false; } for (int iter = 0; iter < repeat; ++iter, ++invocationId) { - launchGroupStage(stagedRouteStages[stageIndex], - groupTraceDevices[stageIndex], static_cast(iter)); + if (!launchGroupStage(stagedRouteStages[stageIndex], + groupTraceDevices[stageIndex], static_cast(iter))) { + return false; + } } if (!CheckAcl(rank, "aclrtRecordEvent grouped stage end", aclrtRecordEvent(stageEndEvent, stream)) || diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 57a5572d..8414d667 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -534,6 +534,7 @@ void TestRouteStages() void TestCopyoutWorkerPolicy() { CHECK_EQ(TileXR::Demo::kAllToAllGroupBlockDim, 64U); + CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(1U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(8U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(16U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(32U), true); @@ -541,6 +542,7 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupValidCopyoutWorkers(12U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U, 32U), 64U); + CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U, 1U), 33U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(16U, 48U), 0U); CHECK_EQ(TileXR::Demo::AllToAllGroupBlockDim(32U, 48U), 0U); @@ -567,12 +569,25 @@ void TestCopyoutWorkerPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 0U, 48U), 0); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(47U, 0U, 48U), 15); CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(32U, 1U, 48U), -1); + + lanes.clear(); + for (uint32_t assignment = 0U; assignment < 16U; ++assignment) { + lanes.insert(TileXR::Demo::AllToAllGroupCopyoutLane( + 0U, assignment, 1U)); + } + CHECK_EQ(lanes.size(), 16U); + CHECK_EQ(*lanes.begin(), 0); + CHECK_EQ(*lanes.rbegin(), 15); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(0U, 16U, 1U), -1); + CHECK_EQ(TileXR::Demo::AllToAllGroupCopyoutLane(1U, 0U, 1U), -1); } void TestKernelStructure() { const std::string kernel = ReadFile( std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp"); + const std::string launcher = ReadFile( + std::string(TILEXR_SOURCE_ROOT) + "/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp"); CHECK_CONTAINS(kernel, "tilexr_udma_all_to_all_group_kernel"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_CORES"); CHECK_CONTAINS(kernel, "#include \"tilexr_udma_alltoall_group_route.h\""); @@ -583,7 +598,7 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); - CHECK_CONTAINS(kernel, "constexpr uint32_t copyoutWorkers = 32U"); + CHECK_CONTAINS(kernel, "constexpr uint32_t copyoutWorkers = 1U"); CHECK_CONTAINS(kernel, "uint32_t multiChannel, uint32_t primaryRouteParts"); CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); CHECK_CONTAINS(kernel, @@ -593,8 +608,17 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupRemoteAssistDevice"); CHECK_CONTAINS(kernel, "copySliceCount"); CHECK_CONTAINS(kernel, "copySliceIndex"); + CHECK_CONTAINS(kernel, "#include \"tilexr_sdma.h\""); + CHECK_CONTAINS(kernel, "AllToAllGroupCopySdma"); + CHECK_CONTAINS(kernel, + "args, relayDst, relaySrc, copyBytes, worker, sdmaEvent"); + CHECK_CONTAINS(kernel, "event == 0ULL"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SDMA_FAILED"); + CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_SDMA"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers"); - CHECK_CONTAINS(kernel, "const uint32_t traceCore = copyoutWorkers == 8U"); + CHECK_CONTAINS(kernel, + "const uint32_t traceCore = copyoutWorkers < TILEXR_ALLTOALL_GROUP_SEND_CORES"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx"); CHECK_CONTAINS(kernel, "UDMAPutNbiOnQpWithFlag"); CHECK_CONTAINS(kernel, "UDMAPutNbiOnQpWithFlag"); @@ -623,7 +647,6 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupKernelImpl"); CHECK_CONTAINS(kernel, "AllToAllGroupKernelImpl"); CHECK_CONTAINS(kernel, "if constexpr (IngressCredit)"); - CHECK_CONTAINS(kernel, "if (quietBatch == 1U)"); CHECK_CONTAINS(kernel, "uint32_t groupWidth, uint32_t quietBatch"); CHECK_CONTAINS(kernel, "uint64_t creditOffset0, uint64_t creditOffset1"); CHECK_CONTAINS(kernel, "uint32_t ingressWindow"); @@ -656,7 +679,21 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupStageRunsCopyDevice(routeStage)"); CHECK_CONTAINS(kernel, "AllToAllGroupStageWaitsForSignalDevice(routeStage)"); CHECK_CONTAINS(kernel, "observed >= expectedToken"); - CHECK_CONTAINS(kernel, "launch_tilexr_udma_all_to_all_group"); + CHECK_NOT_CONTAINS(kernel, "<<<"); + CHECK_CONTAINS(launcher, "rtDevBinaryRegister"); + CHECK_CONTAINS(launcher, "rtFunctionRegister"); + CHECK_CONTAINS(launcher, "rtKernelLaunchWithFlagV2"); + CHECK_CONTAINS(launcher, "GroupedAllToAllKernelArgs"); + CHECK_CONTAINS(launcher, "sizeof(GroupedAllToAllKernelArgs) == 128U"); + CHECK_CONTAINS(launcher, "GroupedAllToAllCreditKernelArgs"); + CHECK_CONTAINS(launcher, "sizeof(GroupedAllToAllCreditKernelArgs) == 152U"); + CHECK_CONTAINS(launcher, "TILEXR_GROUPED_ALLTOALL_BATCH_KERNEL_NAME"); + CHECK_CONTAINS(launcher, "TILEXR_GROUPED_ALLTOALL_CREDIT_KERNEL_NAME"); + CHECK_CONTAINS(launcher, "TILEXR_GROUPED_ALLTOALL_BATCH_CREDIT_KERNEL_NAME"); + CHECK_CONTAINS(launcher, "const bool useCredit = ingressWindow != 0U"); + CHECK_CONTAINS(launcher, "const bool useBatch = quietBatch != 1U"); + CHECK_CONTAINS(launcher, "cfgInfo.schemMode = RT_SCHEM_MODE_NORMAL"); + CHECK_NOT_CONTAINS(launcher, "<<<"); CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); CHECK_NOT_CONTAINS(kernel, "SyncAll"); CHECK_NOT_CONTAINS(kernel, "elementsPerPeer) * lane /"); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index 94017f3e..4658fe1d 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -52,18 +52,19 @@ def make_trace( MODULE.TASK_BASE_OFFSET, ) spans = ( - (16, 0, 0, 0, 1100, 1200, rank, MODULE.NO_QP), + (32, 0, 0, 0, 1100, 1200, rank, MODULE.NO_QP), (0, 0, 0, 1, 1200, 1300, 1, 3), + (16, 0, 0, 1, 1250, 1350, 2, 4), (0, 0, 0, 2, 1300, 1400, 1, 3), - (16, 0, 0, 3, 1400, 1500, 1, MODULE.NO_QP), - (16, 0, 0, 4, 1500, 1600, 1, MODULE.NO_QP), + (32, 0, 0, 3, 1400, 1500, 1, MODULE.NO_QP), + (32, 0, 0, 4, 1500, 1600, 1, MODULE.NO_QP), (0, 0, 0, 5, 1050, 1100, 1, MODULE.NO_QP), ) with path.open("wb") as stream: stream.truncate(MODULE.TRACE_BYTES) stream.seek(0) stream.write(header) - for core in (0, 16): + for core in (0, 16, 32): stream.seek(MODULE.kernel_span_offset(0, core)) stream.write(struct.pack(" Date: Tue, 4 Aug 2026 12:06:05 +0800 Subject: [PATCH 145/163] fix(udma): select grouped copyout resources by SDMA mode --- tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp | 4 ++-- tests/udma/demo/tilexr_udma_demo.cpp | 8 +++++++- .../unit/test_tilexr_udma_alltoall_group_layout.cpp | 10 +++++++--- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 45f28c36..e24285dd 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -661,12 +661,12 @@ __aicore__ inline void AllToAllGroupKernelImpl( uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, uint32_t groupWidth, uint32_t quietBatch) { - constexpr uint32_t copyoutWorkers = 1U; const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? reinterpret_cast<__gm__ uint8_t*>(groupTraceGM) : nullptr; const uint64_t kernelBegin = AllToAllGroupTraceCycle(groupTrace); auto args = reinterpret_cast<__gm__ TileXR::CommArgs*>(commArgsGM); + const uint32_t copyoutWorkers = TileXR::SDMAEnabled(args) ? 1U : 32U; auto input = reinterpret_cast<__gm__ int32_t*>(inputGM); auto output = reinterpret_cast<__gm__ int32_t*>(outputGM); auto registeredMemory = reinterpret_cast<__gm__ uint8_t*>(registeredMemoryGM); @@ -915,7 +915,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); uint64_t sdmaEvent = 0ULL; const uint32_t sdmaStatus = AllToAllGroupCopySdma( - args, relayDst, relaySrc, copyBytes, worker, sdmaEvent); + args, relayDst, relaySrc, copyBytes, lane, sdmaEvent); if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK) { AllToAllGroupCopyMte(relayDst, relaySrc, copyBytes, relayLocal); } diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 498f6096..7ec62065 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -1026,8 +1026,13 @@ bool RunGroupedAllToAll( TileXR::Demo::kAllToAllGroupAutoPrimaryParts : static_cast(primaryRoutePartsValue); + bool sdmaAvailable = false; + if (!CheckTileXR(rank, "TileXRSDMAAvailable grouped alltoall", + TileXRSDMAAvailable(comm, &sdmaAvailable))) { + return false; + } constexpr uint32_t sendWorkers = TileXR::Demo::kAllToAllGroupSendWorkerCount; - constexpr uint32_t copyoutWorkers = 1U; + const uint32_t copyoutWorkers = sdmaAvailable ? 1U : 32U; const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim( sendWorkers, copyoutWorkers); const int routeStagesValue = GetEnvInt( @@ -1235,6 +1240,7 @@ bool RunGroupedAllToAll( " channelMode=" + std::to_string(channelModeValue) + " multiChannel=" + std::to_string(multiChannel ? 1 : 0) + " primaryRouteParts=" + std::to_string(primaryRoutePartsValue) + + " sdmaAvailable=" + std::to_string(sdmaAvailable ? 1 : 0) + " sendWorkers=" + std::to_string(sendWorkers) + " copyoutWorkers=" + std::to_string(copyoutWorkers) + " blockDim=" + std::to_string(groupBlockDim) + diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 8414d667..ef1b02f1 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -598,7 +598,10 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); - CHECK_CONTAINS(kernel, "constexpr uint32_t copyoutWorkers = 1U"); + CHECK_CONTAINS(kernel, + "const uint32_t copyoutWorkers = TileXR::SDMAEnabled(args) ? 1U : 32U"); + CHECK_CONTAINS(kernel, + "args, relayDst, relaySrc, copyBytes, lane, sdmaEvent"); CHECK_CONTAINS(kernel, "uint32_t multiChannel, uint32_t primaryRouteParts"); CHECK_CONTAINS(kernel, "AllToAllGroupPeerInRouteStageDevice"); CHECK_CONTAINS(kernel, @@ -610,8 +613,6 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "copySliceIndex"); CHECK_CONTAINS(kernel, "#include \"tilexr_sdma.h\""); CHECK_CONTAINS(kernel, "AllToAllGroupCopySdma"); - CHECK_CONTAINS(kernel, - "args, relayDst, relaySrc, copyBytes, worker, sdmaEvent"); CHECK_CONTAINS(kernel, "event == 0ULL"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SDMA_FAILED"); @@ -721,6 +722,9 @@ void TestHostStructure() CHECK_CONTAINS(demo, "grouped ingress credit currently requires groupWidth=16"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_PRIMARY_ROUTE_PARTS"); CHECK_CONTAINS(demo, "kAllToAllGroupSendWorkerCount"); + CHECK_CONTAINS(demo, "TileXRSDMAAvailable(comm, &sdmaAvailable)"); + CHECK_CONTAINS(demo, + "const uint32_t copyoutWorkers = sdmaAvailable ? 1U : 32U"); CHECK_CONTAINS(demo, "grouped alltoall registeredBytes="); CHECK_CONTAINS(demo, "grouped alltoall warmup="); const size_t begin = demo.find("bool RunGroupedAllToAll("); From 2d4c65c0297e5c9bd8b56a1869778a7c1543746f Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 13:06:32 +0800 Subject: [PATCH 146/163] docs(udma): document grouped credit IPC MTE flow --- docs/grouped-alltoall-credit-ipc-mte.md | 423 ++++++++++++++++++++++++ 1 file changed, 423 insertions(+) create mode 100644 docs/grouped-alltoall-credit-ipc-mte.md diff --git a/docs/grouped-alltoall-credit-ipc-mte.md b/docs/grouped-alltoall-credit-ipc-mte.md new file mode 100644 index 00000000..6cf01791 --- /dev/null +++ b/docs/grouped-alltoall-credit-ipc-mte.md @@ -0,0 +1,423 @@ +# Grouped AllToAll Credit IPC 与 MTE Set/Wait 实现 + +本文整理 TileXR grouped all-to-all ingress credit 在当前版本 `b51dc2e` 中的完整实现，覆盖通信域初始化、IPC 内存布局、Host 到 Device 参数传递，以及 Device 侧通过 MTE 发布和轮询 credit 的流程。其它项目可以复用该设计，但必须先确认目标硬件支持对已映射 IPC GM 地址执行 MTE 访问。 + +## 1. 功能语义 + +Credit 用于限制同一张目标卡在相邻 group 中承受的并发 ingress。基本时序是： + +```text +发送端 S 接收端 R +group g: UDMA put + signal --------> 等待 payload signal 成功 + MTE3 写 credit(g+1) + 到 S 拥有的 IPC credit buffer +group g+1: MTE2 轮询本地 credit <---- credit 对 S 可见 +group g+1: UDMA put + signal -------> +``` + +当前发布点位于接收端确认 UDMA payload signal 到达之后、receive-copy 之前。因此当前 credit 的准确含义是： + +> 上一 group 的远端数据已经到达，接收端允许下一 group ingress。 + +它不表示上一批数据已被 receive-copy 或应用层消费。如果复用同一块 payload 存储要求“消费完成后才能覆盖”，必须把 credit 发布点后移到 copyout/消费完成之后。 + +Credit 通道只传控制 token，使用 IPC GM 加 MTE，不使用 UDMA QP。 + +## 2. 代码位置 + +| 模块 | 文件 | 职责 | +|---|---|---| +| 公共 ABI 和常量 | `src/include/comm_args.h` | IPC 大小、stride、`CommArgs::creditMems` | +| 通信域 Host 初始化 | `src/comm/tilexr_comm.cpp` | 申请、授权、交换、映射和释放 IPC memory | +| 布局与 Host 校验 | `tests/udma/demo/tilexr_udma_alltoall_group_layout.h` | ping-pong slot、offset、参数约束 | +| Demo Host 流程 | `tests/udma/demo/tilexr_udma_demo.cpp` | 环境变量、映射完整性和 single-pass 校验 | +| Kernel launch ABI | `tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp` | credit offset 和 ingressWindow 下发 | +| Device set/wait | `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` | MTE3 publish、MTE2 poll 和调度 | + +## 3. IPC 内存参数与布局 + +公共常量定义如下： + +```cpp +constexpr int TILEXR_MAX_RANK_SIZE = 1024; +constexpr int64_t CREDIT_IPC_STRIDE = 512; +constexpr int64_t CREDIT_IPC_SLOT_BYTES = + TILEXR_MAX_RANK_SIZE * CREDIT_IPC_STRIDE; // 512 KiB +constexpr int64_t CREDIT_IPC_BYTES = + 2 * CREDIT_IPC_SLOT_BYTES; // 1 MiB +``` + +每个 rank 独立申请一块固定 1 MiB 的 credit IPC memory。它按 invocation 奇偶分成两个 slot： + +```text +rank R 拥有的 credit buffer，共 1 MiB + +slot 0: [0, 512 KiB) +slot 1: [512 KiB, 1 MiB) + +slot 内： +rank index 0: [0 * 512, 1 * 512) +rank index 1: [1 * 512, 2 * 512) +... +rank index 1023 +``` + +每个 credit 占位 512 B。token 只使用第一个 `uint64_t`，但 set 和 wait 都搬运完整 512 B，即 64 个 `uint64_t`。固定 stride 同时满足数据搬运对齐，并隔离不同 peer 的 cache line/传输单元，避免伪共享。 + +内存按最大 1024 rank 固定布局，而不是按当前 `rankSize` 缩小。因此不同进程和 Host/Kernel 无需协商动态 slot 大小。 + +## 4. 通信域 Host 初始化 + +### 4.1 开关和初始化顺序 + +通信域通过以下变量启用专用 credit IPC memory： + +```bash +export TILEXR_ENABLE_CREDIT_IPC=1 +``` + +普通多进程通信域中的相关初始化顺序为： + +```text +GetDev +InitCommon +InitCommMem（TILEXR_ENABLE_IPC 开启时） +InitCreditCommMem（TILEXR_ENABLE_CREDIT_IPC=1 时） +InitUDMA +InitSDMA +SyncCommArgs +``` + +Credit IPC 必须在 `SyncCommArgs()` 前完成，因为映射后的 Device 地址要写入 `CommArgs`。 + +### 4.2 收集 IPC 授权信息 + +`InitCreditCommMem()` 首先收集所有进程的 PID 和 SDID： + +```text +rtDeviceGetBareTgid(localPid) +socketExchange_->AllGather(localPid, allPids) + +rtGetDeviceInfo(..., infoTypeSdid=26, localSdid) +socketExchange_->AllGather(localSdid, allSdids) +``` + +这里使用 bare TGID，目的是在容器环境中取得 Runtime IPC 授权需要的宿主侧进程标识。 + +IPC 授权模式由下列变量控制： + +```bash +export TILEXR_IPC_PID_MODE=pid +# 或 +export TILEXR_IPC_PID_MODE=sdid +``` + +当前代码对 `CHIP_910_9391 <= chip < CHIP_950` 默认选择 SDID，其它情况默认 PID；显式环境变量优先。授权由 `SetIpcPidSdid()` 对本 rank 的 IPC name 添加所有远端访问者。 + +对应 Runtime API 为： + +```cpp +// PID 模式 +rtSetIpcMemPid(ipcName, &peerPid, HCCL_IPC_PID_ARRAY_SIZE); + +// SDID 模式；失败时当前实现回退到 PID 模式 +rtSetIpcMemorySuperPodPid( + ipcName, peerSdid, &peerPid, HCCL_IPC_PID_ARRAY_SIZE); +``` + +### 4.3 申请、命名和映射 + +每个 rank 的 Host 流程如下： + +```cpp +// 1. 本 rank 申请 owner buffer。 +aclrtMalloc(&creditIpcMem_[rank], CREDIT_IPC_BYTES, + ACL_MEM_MALLOC_HUGE_FIRST); + +// 310P3 使用 ACL_MEM_MALLOC_HUGE_FIRST_P2P。 + +// 2. 只初始化一次，后续依靠单调 token 复用。 +aclrtMemset(creditIpcMem_[rank], CREDIT_IPC_BYTES, 0, CREDIT_IPC_BYTES); + +// 3. 导出 IPC name，并授权所有 peer PID/SDID。 +rtIpcSetMemoryName(creditIpcMem_[rank], CREDIT_IPC_BYTES, + localName, IPC_NAME_SIZE); +SetIpcPidSdid(localName, allPids, allSdids); + +// 4. AllGather 每个 rank 的 IPC name。 +GetName(localName, allNames); + +// 5. 打开其它 rank 的 owner buffer。 +rtIpcOpenMemory(&creditIpcMem_[peer], allNames[peer]); +``` + +完成后，在任意 rank `R` 的进程中： + +```text +creditIpcMem_[R] = R 自己申请的本地 owner 地址 +creditIpcMem_[peer] = peer buffer 映射到 R 地址空间后的 Device GM 地址 +``` + +`OpenCreditIpcMem()` 跳过 self，也沿用 `SkipUnusedChannel910B2C()` 对不使用的 910B2C 链路做过滤。使用 credit 的 Host 代码随后要求当前通信范围内每个 `creditMems[peer]` 都非空，因此迁移时必须让“跳过映射”和“实际参与 peer”保持一致。 + +### 4.4 Host 到 Device 参数传递 + +`CommArgs` 提供固定长度的映射表： + +```cpp +struct CommArgs { + // ... + GM_ADDR creditMems[TILEXR_MAX_RANK_SIZE] = {}; +}; +``` + +`SyncCommArgs()` 将 Host 通信域中的地址逐项复制进去： + +```cpp +for (int i = 0; i < rankSize_; ++i) { + commArgs_.creditMems[i] = creditIpcMem_[i]; +} + +aclrtMalloc(&commArgsPtr_, sizeof(commArgs_), ACL_MEM_MALLOC_HUGE_FIRST); +aclrtMemcpy(commArgsPtr_, sizeof(commArgs_), + &commArgs_, sizeof(commArgs_), + ACL_MEMCPY_HOST_TO_DEVICE); +``` + +Kernel 获得 `commArgsPtr_` 后，可以直接把 `args->creditMems[peer]` 作为 `__gm__` 地址用于 MTE。这里传递的是 Device 可访问的 owner/IPC mapping 地址，不是 Host 虚拟地址或 IPC name。 + +Kernel launch ABI 另外传递： + +```cpp +uint64_t creditOffset0; +uint64_t creditOffset1; +uint32_t ingressWindow; +``` + +当前 credit kernel 参数结构为 152 B。Host 和 Kernel 的参数结构、字段顺序及对齐必须同步修改，不能只改一侧。 + +## 5. Ping-pong 和 token + +当前 invocation 使用两个 slot 交替复用： + +```cpp +slot = invocationId & 1U; +creditOffset[0] = 0; +creditOffset[1] = CREDIT_IPC_SLOT_BYTES; +``` + +Token 格式为： + +```cpp +uint64_t token = + (uint64_t(invocationId + 1) << 32) | + (uint64_t(slot) << 31) | + (uint64_t(group) << 16) | + uint64_t(pass + 1); +``` + +Credit 当前只支持 single pass，因此 credit token 的 `pass` 固定为 0。消费者判断： + +```cpp +observed >= expectedToken +``` + +Buffer 只在通信域初始化时清零一次。之后 token 随 invocation/group 单调推进，旧值无需再次清零。迁移时必须确保 invocation 不会在 buffer 生命周期内回绕，否则 `>=` 比较将不再安全。 + +## 6. Device 侧 MTE set + +### 6.1 地址所有权 + +接收 rank `R` 完成 group `g` 的 payload signal wait 后，计算下一 group 同 lane 的发送者 `S`，然后写入： + +```cpp +args->creditMems[S] + + creditOffset[slot] + + R * CREDIT_IPC_STRIDE +``` + +也就是： + +```text +buffer owner = 下一 group 的发送者 S +slot index = 当前 invocation 的 ping-pong slot +entry index = 授权者/目标接收 rank R +``` + +这个方向非常关键。R 不是写自己的 buffer，而是通过 IPC mapping 写 S 拥有的 buffer；这样 S 在发送下一 group 前只需轮询自己的本地 owner buffer。 + +### 6.2 MTE3 发布代码 + +当前实现先在 UB 中构造 512 B credit block，再用 MTE3 写远端 IPC GM： + +```cpp +auto creditLocal = relayLocal.ReinterpretCast(); +creditLocal.SetValue(0, creditToken); + +SetFlag(EVENT_ID0); +WaitFlag(EVENT_ID0); + +GlobalTensor remoteCreditGlobal; +remoteCreditGlobal.SetGlobalBuffer(remoteCredit, 64); +DataCopy(remoteCreditGlobal, creditLocal, 64); // 64 * 8 B = 512 B + +SetFlag(EVENT_ID0); +WaitFlag(EVENT_ID0); +``` + +两个同步方向分别解决： + +- `S_MTE3`：保证 MTE3 读取 UB 前，scalar 对 `creditLocal` 的 token 写入已经可见。 +- `MTE3_S`：保证 scalar 路径继续执行前，本次 UB→GM/IPC 写已经完成。 + +不要用 `GlobalTensor::SetValue()` 代替生产路径中的 MTE copy。当前实现专门经 UB 和 MTE3 写入远端 IPC memory，以获得明确的数据搬运与同步语义。 + +### 6.3 发布者唯一性 + +Receive-copy 可能有 32 个 worker，但 credit 只允许指定 owner 发布： + +```text +worker < 16 +``` + +这保证每个 lane/entry 只有一个 core 写 token，避免两个 core 并发写同一 512 B entry。当前代码还要求在最后一个 pass 发布；由于 ingress credit 只允许 single pass，实际就是 pass 0。 + +## 7. Device 侧 MTE wait + +发送 rank `S` 在发送 group `g > 0` 到目标 `R` 前，轮询： + +```cpp +args->creditMems[S] + + creditOffset[slot] + + R * CREDIT_IPC_STRIDE +``` + +从 S 的视角，`args->creditMems[S]` 是自己的本地 owner buffer。entry `R` 表示目标接收端 R 已经授予本次发送 credit。 + +每次 load 都将完整 512 B 从 GM 搬到 UB： + +```cpp +GlobalTensor creditGlobal; +creditGlobal.SetGlobalBuffer(credit, 64); +auto creditLocal = relayLocal.ReinterpretCast(); + +DataCopy(creditLocal, creditGlobal, 64); // GM -> UB, 512 B +SetFlag(EVENT_ID0); +WaitFlag(EVENT_ID0); + +uint64_t observed = creditLocal.GetValue(0); +``` + +`MTE2_S` 保证 scalar 读取 UB token 前，GM→UB 的 MTE2 搬运已经完成。 + +轮询逻辑为： + +```cpp +const uint64_t begin = GetSystemCycle(); +do { + observed = LoadCreditMte(creditAddress, relayLocal); + if (GetSystemCycle() - begin >= timeoutCycles) { + return false; + } +} while (observed < expectedToken); +``` + +当前 timeout 是 `10000000000` cycles。超时会记录 credit-wait stage、group、peer、expected 和 observed token，并提前退出该 kernel core。 + +Group 0 不等待，因为还没有前一 group 可以授予 credit。对当前 route 数据量为 0 的 peer，也跳过 wait 和对应 send。 + +## 8. Host 侧运行约束 + +Demo 中还需要启用 ingress window： + +```bash +export TILEXR_ENABLE_CREDIT_IPC=1 +export TILEXR_DEMO_ALLTOALL_GROUP_INGRESS_WINDOW=1 +``` + +当前实现限制如下： + +```text +ingressWindow = 0 或 1 +credit 模式要求 groupWidth = 16 +credit 模式要求 CHUNK_ELEMENTS = ELEMENTS，即 single pass +creditMems[0..rankSize-1] 必须全部非空 +``` + +Host 用静态断言保护通信域和 demo 的布局常量： + +```cpp +static_assert(kAllToAllGroupCreditStride == CREDIT_IPC_STRIDE); +static_assert(kAllToAllGroupCreditSlotBytes == CREDIT_IPC_SLOT_BYTES); +``` + +其它项目也应在 ABI 两侧增加同类断言，避免 stride 或最大 rank 数不一致造成静默越界。 + +## 9. 释放顺序 + +通信域销毁时必须区分远端 mapping 和本地 owner allocation： + +```text +1. kernel/stream 已停止访问 CommArgs 和 credit memory +2. rtIpcCloseMemory() 关闭所有 peer 的远端 mapping +3. aclrtFree() 释放本 rank 的 owner credit allocation +4. aclrtFree() 释放 Device CommArgs +``` + +当前代码由 `CloseCreditIpcMem()` 关闭 `creditIpcMem_[peer]`，再通过 `FreePeerMem(creditIpcMem_[rank_])` 释放本地 allocation。不能对远端 mapping 调用 `aclrtFree()`，也不能在 peer 尚可能访问时提前释放 owner buffer。 + +## 10. 迁移到其它项目的最小步骤 + +1. 定义 Host/Device 共用的 `MAX_RANK_SIZE`、512 B stride、两个 slot 和总字节数。 +2. 每 rank 申请一块独立的 P2P/IPC Device GM buffer，并清零一次。 +3. 收集所有进程的 bare PID 和设备 SDID。 +4. 导出本 rank IPC name/handle，并按目标芯片要求授权 PID 或 SDID。 +5. AllGather IPC name/handle，在每个 rank 打开所有参与 peer 的映射。 +6. 将 owner 地址和所有 IPC mapping 填入 Device 可见的 `creditMems[]`。 +7. 把 `CommArgs` 复制到 Device，并在 kernel launch ABI 中传入 slot offsets 和 enable flag。 +8. Producer 在 UB 写 token，经 `S_MTE3` 后用 MTE3 搬 512 B 到远端 IPC GM，再执行 `MTE3_S`。 +9. Consumer 用 MTE2 从本地 owner GM 搬 512 B 到 UB，经 `MTE2_S` 后读取 token 并轮询。 +10. 使用随 epoch/invocation 单调增加的 token，并用 ping-pong slot 隔离相邻 invocation。 +11. 明确定义 credit 发布点是“数据到达”“copyout 完成”还是“应用消费完成”。 +12. 销毁时先停止 kernel，再关闭远端 mapping，最后释放本地 owner buffer。 + +## 11. 常见错误与检查方法 + +| 问题 | 表现或风险 | 检查方法 | +|---|---|---| +| `creditMems[peer]` 为空 | Device 访问异常或 Host 拒绝启动 | 启动前逐 peer 校验映射 | +| owner/index 写反 | 永久 credit-wait 超时 | 确认 R 写 `creditMems[S] + R*stride`，S 读 `creditMems[S] + R*stride` | +| Host/Kernel stride 不一致 | 读到其它 peer token或越界 | 共享头文件并加 `static_assert` | +| 只搬 8 B 或不满足平台对齐 | DataCopy 行为不稳定或性能异常 | 保持 512 B 对齐和 512 B copy | +| 缺少 `S_MTE3` | MTE3 可能读到 UB 旧值 | scalar 写 UB 后建立 S→MTE3 依赖 | +| 缺少 `MTE3_S` | credit 未完成便复用 UB/继续调度 | publish 后建立 MTE3→S 依赖 | +| 缺少 `MTE2_S` | scalar 可能在 load 完成前读取 UB | load 后建立 MTE2→S 依赖 | +| 多个 core 发布同一 entry | 重复或竞争写 token | 为每个 lane/peer 指定唯一 publisher | +| token 不单调或 invocation 回绕 | 旧 token 被误判为新 credit | 评估 token 位宽和通信域最长生命周期 | +| 发布点过早 | payload 被覆盖而尚未消费 | 根据 buffer 复用语义移动 publish 点 | +| IPC 授权模式不匹配 | `rtIpcOpenMemory` 失败 | 核对 PID/SDID、容器 bare PID 和芯片规则 | +| 超出 MTE 可访问拓扑 | 远端 IPC 地址不能被 MTE 访问 | 先在目标超节点/P2P 范围做最小 set/wait 验证 | + +调试超时时，至少记录： + +```text +local rank +peer rank +buffer owner rank +slot / byte offset +group / pass / invocation +expected token +observed token +publisher core +waiter core +``` + +其中 expected/observed token 可以直接拆出 invocation、slot、group 和 pass，用于区分“对端没有发布”“读错 entry”以及“读到了上一轮 token”。 + +## 12. 设计边界 + +- 该 credit 方案依赖超节点或目标 P2P 拓扑内，MTE 能访问 `rtIpcOpenMemory` 返回的 GM mapping。它不是任意跨节点的通用控制通道。 +- 512 B stride 是当前实现的 Host/Device ABI，不只是性能参数；修改时必须同步布局、申请大小、DataCopy 长度、地址计算和静态断言。 +- 两个 ping-pong slot 解决相邻 invocation 干扰，但不能替代正确的 epoch/token 设计。 +- Busy polling 会占用 AIV/MTE2 资源。若快慢卡严重，credit-wait 本身可能成为关键路径，需要结合调度顺序、窗口大小或硬件通知机制进一步优化。 +- 当前 `ingressWindow=1` 是严格的相邻 group credit，不等同于可配置的多 credit 滑动窗口。扩展到 window > 1 时，需要重新设计 entry 状态或 token/ack 关系，不能只放宽 Host 参数校验。 From 748ca74c356209cdf603b043dcd98933d93e87b3 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 13:14:32 +0800 Subject: [PATCH 147/163] docs(udma): add reusable credit IPC implementation --- docs/grouped-alltoall-credit-ipc-mte.md | 637 +++++++++++++++++++++++- 1 file changed, 631 insertions(+), 6 deletions(-) diff --git a/docs/grouped-alltoall-credit-ipc-mte.md b/docs/grouped-alltoall-credit-ipc-mte.md index 6cf01791..a97b50d7 100644 --- a/docs/grouped-alltoall-credit-ipc-mte.md +++ b/docs/grouped-alltoall-credit-ipc-mte.md @@ -27,12 +27,12 @@ Credit 通道只传控制 token，使用 IPC GM 加 MTE，不使用 UDMA QP。 | 模块 | 文件 | 职责 | |---|---|---| -| 公共 ABI 和常量 | `src/include/comm_args.h` | IPC 大小、stride、`CommArgs::creditMems` | -| 通信域 Host 初始化 | `src/comm/tilexr_comm.cpp` | 申请、授权、交换、映射和释放 IPC memory | -| 布局与 Host 校验 | `tests/udma/demo/tilexr_udma_alltoall_group_layout.h` | ping-pong slot、offset、参数约束 | -| Demo Host 流程 | `tests/udma/demo/tilexr_udma_demo.cpp` | 环境变量、映射完整性和 single-pass 校验 | -| Kernel launch ABI | `tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp` | credit offset 和 ingressWindow 下发 | -| Device set/wait | `tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` | MTE3 publish、MTE2 poll 和调度 | +| 公共 ABI 和常量 | `D:/workspace/TileXR/.kilo/worktrees/serene-beak/src/include/comm_args.h` | IPC 大小、stride、`CommArgs::creditMems` | +| 通信域 Host 初始化 | `D:/workspace/TileXR/.kilo/worktrees/serene-beak/src/comm/tilexr_comm.cpp` | 申请、授权、交换、映射和释放 IPC memory | +| 布局与 Host 校验 | `D:/workspace/TileXR/.kilo/worktrees/serene-beak/tests/udma/demo/tilexr_udma_alltoall_group_layout.h` | ping-pong slot、offset、参数约束 | +| Demo Host 流程 | `D:/workspace/TileXR/.kilo/worktrees/serene-beak/tests/udma/demo/tilexr_udma_demo.cpp` | 环境变量、映射完整性和 single-pass 校验 | +| Kernel launch ABI | `D:/workspace/TileXR/.kilo/worktrees/serene-beak/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp` | credit offset 和 ingressWindow 下发 | +| Device set/wait | `D:/workspace/TileXR/.kilo/worktrees/serene-beak/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp` | MTE3 publish、MTE2 poll 和调度 | ## 3. IPC 内存参数与布局 @@ -421,3 +421,628 @@ waiter core - 两个 ping-pong slot 解决相邻 invocation 干扰，但不能替代正确的 epoch/token 设计。 - Busy polling 会占用 AIV/MTE2 资源。若快慢卡严重，credit-wait 本身可能成为关键路径，需要结合调度顺序、窗口大小或硬件通知机制进一步优化。 - 当前 `ingressWindow=1` 是严格的相邻 group credit，不等同于可配置的多 credit 滑动窗口。扩展到 window > 1 时，需要重新设计 entry 状态或 token/ack 关系，不能只放宽 Host 参数校验。 + +## 13. 可直接提取的实际代码 + +本节不是伪代码，来自当前 `b51dc2e` 实现。迁移时可以直接复制，再替换目标项目的日志、错误码、通信域类名和 all-gather 接口。 + +### 13.1 依赖和唯一适配点 + +Host 代码依赖 CANN ACL/Runtime IPC API： + +```cpp +#include +#include "runtime/dev.h" +#include "runtime/mem.h" + +#include +#include +#include +#include +#include +``` + +目标项目需要提供以下已有上下文： + +```text +rank_ / rankSize_ 当前 rank 和通信域大小 +devList_[rank_] 当前 rank 对应的逻辑 device id +physicalInfo_.chipName 芯片型号，用于选择 PID/SDID 模式 +socketExchange_->AllGather(...) Host 侧全通信域 all-gather +SkipUnusedChannel910B2C(...) 可选；没有特殊拓扑时删除该判断 +TILEXR_SUCCESS / TILEXR_ERROR_* 替换成目标项目错误码 +TILEXR_LOG 替换成目标项目日志 +``` + +其中 `AllGather` 必须在所有 rank 上以完全相同的顺序调用。支持 SDID 的芯片依次执行 PID、SDID 和 IPC name 三次 collective；其它芯片执行 PID 和 IPC name 两次 collective。通信域内不能让部分 rank 单独跳过其中一次。 + +### 13.2 Host/Device 共用头文件 + +这部分必须放在 Host 和 Ascend C kernel 都包含的公共头文件中： + +```cpp +#pragma once + +#include + +#ifndef GM_ADDR +using GM_ADDR = uint8_t*; +#endif + +namespace TileXR { + +constexpr int TILEXR_MAX_RANK_SIZE = 1024; +constexpr int IPC_NAME_SIZE = 65; +constexpr int HCCL_IPC_PID_ARRAY_SIZE = 1; +constexpr int64_t CREDIT_IPC_STRIDE = 512; +constexpr int64_t CREDIT_IPC_SLOT_BYTES = + TILEXR_MAX_RANK_SIZE * CREDIT_IPC_STRIDE; +constexpr int64_t CREDIT_IPC_BYTES = 2 * CREDIT_IPC_SLOT_BYTES; + +struct CommArgs { + int rank = 0; + int localRank = -1; + int rankSize = 0; + int localRankSize = -1; + // 其它通信参数放在这里。 + GM_ADDR creditMems[TILEXR_MAX_RANK_SIZE] = {}; +}; + +} // namespace TileXR +``` + +如果目标项目已经有 `CommArgs`，只添加 `creditMems[]` 字段。Host 和 kernel 必须使用同一份结构体定义，不要维护两份相似 ABI。 + +### 13.3 通信域类成员 + +当前通信域需要的成员和方法声明如下： + +```cpp +private: + int GetPid(uint32_t *pids); + int GetSidId(int64_t sdids[TileXR::TILEXR_MAX_RANK_SIZE], int rankSize); + int GetName( + std::string &name, + char names[TileXR::TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]) const; + int SetIpcPidSdid( + std::string &name, const uint32_t *pids, + const int64_t *sdids) const; + int InitCreditCommMem(); + int InitCreditIpcMem(const uint32_t *pids, const int64_t *sdids); + int OpenCreditIpcMem( + const char names[TileXR::TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]); + void CloseCreditIpcMem(); + + GM_ADDR creditIpcMem_[TileXR::TILEXR_MAX_RANK_SIZE] = {}; + bool creditIpcMemInited_ = false; +``` + +`IPC_NAME_SIZE` 使用目标 CANN/Runtime 头文件中的 IPC name 容量定义；所有 rank 必须交换固定的 `IPC_NAME_SIZE` 字节，而不是 `strlen(name)` 字节。 + +### 13.4 PID、SDID 和 IPC name all-gather + +以下是当前实际实现。`socketExchange_->AllGather(local, count, output)` 的 `count` 是元素数： + +```cpp +int TileXRComm::GetPid(uint32_t *pids) +{ + if (rtDeviceGetBareTgid(&pids[rank_]) != RT_ERROR_NONE) { + return TILEXR_ERROR_INTERNAL; + } + return socketExchange_->AllGather(&pids[rank_], 1, pids); +} + +int TileXRComm::GetSidId( + int64_t sdids[TileXR::TILEXR_MAX_RANK_SIZE], int rankSize) +{ + if (rank_ >= rankSize) { + return TILEXR_ERROR_INTERNAL; + } + + if (physicalInfo_.chipName >= ChipName::CHIP_910_9391 && + physicalInfo_.chipName < ChipName::RESERVED) { + constexpr int kRtModuleTypeSystem = 0; + constexpr int kInfoTypeSdid = 26; + if (rtGetDeviceInfo( + devList_[rank_], kRtModuleTypeSystem, kInfoTypeSdid, + &sdids[rank_]) != RT_ERROR_NONE) { + return TILEXR_ERROR_INTERNAL; + } + return socketExchange_->AllGather(&sdids[rank_], 1, sdids); + } + return TILEXR_SUCCESS; +} + +int TileXRComm::GetName( + std::string &name, + char names[TileXR::TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]) const +{ + return socketExchange_->AllGather( + name.c_str(), IPC_NAME_SIZE, names[0]); +} +``` + +如果目标项目的 all-gather 以字节数为单位，PID/SDID 的发送长度应分别改为 `sizeof(uint32_t)` 和 `sizeof(int64_t)`；不要机械保留这里的 `1`。 + +### 13.5 IPC PID/SDID 授权 + +这是当前实际授权函数： + +```cpp +int TileXRComm::SetIpcPidSdid( + std::string &name, const uint32_t *pids, + const int64_t *sdids) const +{ + const char *modeEnv = std::getenv("TILEXR_IPC_PID_MODE"); + const bool forcePid = + modeEnv != nullptr && std::string(modeEnv) == "pid"; + const bool forceSdid = + modeEnv != nullptr && std::string(modeEnv) == "sdid"; + const bool defaultSdid = + physicalInfo_.chipName >= ChipName::CHIP_910_9391 && + physicalInfo_.chipName < ChipName::CHIP_950; + const bool useSdid = forceSdid || (!forcePid && defaultSdid); + + for (int i = 0; i < rankSize_; ++i) { + if (i == rank_) { + continue; + } + + int32_t pid = static_cast(pids[i]); + if (!useSdid) { + if (rtSetIpcMemPid( + name.c_str(), &pid, + HCCL_IPC_PID_ARRAY_SIZE) != RT_ERROR_NONE) { + return TILEXR_ERROR_INTERNAL; + } + continue; + } + + int ret = rtSetIpcMemorySuperPodPid( + name.c_str(), sdids[i], &pid, HCCL_IPC_PID_ARRAY_SIZE); + if (ret != RT_ERROR_NONE) { + ret = rtSetIpcMemPid( + name.c_str(), &pid, HCCL_IPC_PID_ARRAY_SIZE); + if (ret != RT_ERROR_NONE) { + return TILEXR_ERROR_INTERNAL; + } + } + } + return TILEXR_SUCCESS; +} +``` + +### 13.6 IPC memory 申请、交换和打开 + +以下函数构成通信域 credit 初始化主体： + +```cpp +int TileXRComm::InitCreditCommMem() +{ + uint32_t pids[TileXR::TILEXR_MAX_RANK_SIZE] = {}; + int ret = GetPid(pids); + if (ret != TILEXR_SUCCESS) { + return ret; + } + + int64_t sdids[TileXR::TILEXR_MAX_RANK_SIZE] = {}; + ret = GetSidId(sdids, rankSize_); + if (ret != TILEXR_SUCCESS) { + return ret; + } + return InitCreditIpcMem(pids, sdids); +} + +int TileXRComm::InitCreditIpcMem( + const uint32_t *pids, const int64_t *sdids) +{ + const auto policy = + GetChipName() == ChipName::CHIP_310P3 + ? ACL_MEM_MALLOC_HUGE_FIRST_P2P + : ACL_MEM_MALLOC_HUGE_FIRST; + + aclError aclRet = aclrtMalloc( + reinterpret_cast(&creditIpcMem_[rank_]), + TileXR::CREDIT_IPC_BYTES, policy); + if (aclRet != ACL_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + + aclRet = aclrtMemset( + creditIpcMem_[rank_], TileXR::CREDIT_IPC_BYTES, + 0, TileXR::CREDIT_IPC_BYTES); + if (aclRet != ACL_SUCCESS) { + aclrtFree(creditIpcMem_[rank_]); + creditIpcMem_[rank_] = nullptr; + return TILEXR_ERROR_INTERNAL; + } + + char nameBuffer[IPC_NAME_SIZE] = {}; + if (rtIpcSetMemoryName( + creditIpcMem_[rank_], TileXR::CREDIT_IPC_BYTES, + nameBuffer, IPC_NAME_SIZE) != RT_ERROR_NONE) { + aclrtFree(creditIpcMem_[rank_]); + creditIpcMem_[rank_] = nullptr; + return TILEXR_ERROR_INTERNAL; + } + + std::string name(nameBuffer); + if (SetIpcPidSdid(name, pids, sdids) != TILEXR_SUCCESS) { + aclrtFree(creditIpcMem_[rank_]); + creditIpcMem_[rank_] = nullptr; + return TILEXR_ERROR_INTERNAL; + } + + char names[TileXR::TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE] = {}; + name.resize(IPC_NAME_SIZE); + if (GetName(name, names) != TILEXR_SUCCESS) { + aclrtFree(creditIpcMem_[rank_]); + creditIpcMem_[rank_] = nullptr; + return TILEXR_ERROR_INTERNAL; + } + return OpenCreditIpcMem(names); +} + +int TileXRComm::OpenCreditIpcMem( + const char names[TileXR::TILEXR_MAX_RANK_SIZE][IPC_NAME_SIZE]) +{ + static std::mutex mutex; + std::lock_guard lock(mutex); + + for (int peer = 0; peer < rankSize_; ++peer) { + if (peer == rank_) { + continue; + } + if (SkipUnusedChannel910B2C(rank_, peer, GetChipName())) { + continue; + } + if (rtIpcOpenMemory( + reinterpret_cast(&creditIpcMem_[peer]), + names[peer]) != RT_ERROR_NONE) { + CloseCreditIpcMem(); + aclrtFree(creditIpcMem_[rank_]); + creditIpcMem_[rank_] = nullptr; + return TILEXR_ERROR_INTERNAL; + } + } + creditIpcMemInited_ = true; + return TILEXR_SUCCESS; +} +``` + +上面在当前代码基础上补齐了初始化失败时本 rank owner allocation 的释放，适合作为其它项目直接采用的版本。没有 910B2C 特殊拓扑时，应删除 `SkipUnusedChannel910B2C()` 判断，确保每个通信 peer 都完成映射。 + +通信域主初始化中的调用位置必须在 `SyncCommArgs()` 之前： + +```cpp +if (IsEnvEnabled("TILEXR_ENABLE_CREDIT_IPC", false)) { + const int ret = InitCreditCommMem(); + if (ret != TILEXR_SUCCESS) { + return ret; + } +} + +// InitUDMA / InitSDMA ... + +const int ret = SyncCommArgs(); +if (ret != TILEXR_SUCCESS) { + return ret; +} +``` + +### 13.7 `CommArgs` Host 到 Device + +实际同步代码的 credit 部分如下： + +```cpp +int TileXRComm::SyncCommArgs() +{ + commArgs_.rank = rank_; + commArgs_.localRank = localRank_; + commArgs_.rankSize = rankSize_; + commArgs_.localRankSize = localRankSize_; + + for (int peer = 0; peer < rankSize_; ++peer) { + commArgs_.creditMems[peer] = creditIpcMem_[peer]; + } + + int ret = aclrtMalloc( + reinterpret_cast(&commArgsPtr_), sizeof(commArgs_), + ACL_MEM_MALLOC_HUGE_FIRST); + if (ret != ACL_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + + ret = aclrtMemcpy( + commArgsPtr_, sizeof(commArgs_), + &commArgs_, sizeof(commArgs_), + ACL_MEMCPY_HOST_TO_DEVICE); + if (ret != ACL_SUCCESS) { + aclrtFree(commArgsPtr_); + commArgsPtr_ = nullptr; + return TILEXR_ERROR_INTERNAL; + } + return TILEXR_SUCCESS; +} +``` + +启用 credit 的 Host 在 launch 前必须验证： + +```cpp +for (int peer = 0; peer < rankSize; ++peer) { + if (commArgsHost.creditMems[peer] == nullptr) { + std::cerr << "missing credit IPC mapping, peer=" << peer << '\n'; + return false; + } +} + +static_assert( + kAllToAllGroupCreditStride == + static_cast(TileXR::CREDIT_IPC_STRIDE)); +static_assert( + kAllToAllGroupCreditSlotBytes == + static_cast(TileXR::CREDIT_IPC_SLOT_BYTES)); +``` + +### 13.8 Device 侧完整 MTE set/wait helper + +以下代码可直接放入 Ascend C kernel。调用者提供至少 512 B 的 `relayLocal`，并保证同一个 core 不会同时把它用于其它未完成的 MTE 操作： + +```cpp +#include "kernel_operator.h" + +namespace { + +constexpr uint32_t kCreditStride = 512U; +constexpr uint32_t kCreditWords = kCreditStride / sizeof(uint64_t); +constexpr uint64_t kCreditWaitTimeoutCycles = 10000000000ULL; + +static_assert( + kCreditStride == TileXR::CREDIT_IPC_STRIDE, + "credit stride must match Host IPC layout"); + +__aicore__ inline uint64_t MakeCreditToken( + uint32_t invocationId, uint32_t group, uint32_t pass) +{ + const uint64_t invocation = static_cast(invocationId) + 1ULL; + const uint64_t slot = static_cast(invocationId & 1U); + return (invocation << 32U) | + (slot << 31U) | + (static_cast(group) << 16U) | + (static_cast(pass) + 1ULL); +} + +__aicore__ inline uint64_t LoadCreditMte( + __gm__ uint64_t *credit, + AscendC::LocalTensor relayLocal) +{ + AscendC::GlobalTensor creditGlobal; + creditGlobal.SetGlobalBuffer(credit, kCreditWords); + auto creditLocal = relayLocal.ReinterpretCast(); + + AscendC::DataCopy(creditLocal, creditGlobal, kCreditWords); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + return creditLocal.GetValue(0); +} + +__aicore__ inline bool WaitCreditMte( + __gm__ uint64_t *credit, uint64_t expectedToken, + uint64_t timeoutCycles, + AscendC::LocalTensor relayLocal, + uint64_t &observed) +{ + const uint64_t begin = + static_cast(AscendC::GetSystemCycle()); + observed = LoadCreditMte(credit, relayLocal); + while (observed < expectedToken) { + if (static_cast(AscendC::GetSystemCycle()) - begin >= + timeoutCycles) { + return false; + } + observed = LoadCreditMte(credit, relayLocal); + } + return true; +} + +__aicore__ inline void PublishCreditMte( + __gm__ uint8_t *remoteOwnerBuffer, + uint64_t slotOffset, int32_t grantingRank, + uint64_t creditToken, + AscendC::LocalTensor relayLocal) +{ + auto remoteCredit = reinterpret_cast<__gm__ uint64_t *>( + remoteOwnerBuffer + slotOffset + + static_cast(grantingRank) * kCreditStride); + + auto creditLocal = relayLocal.ReinterpretCast(); + creditLocal.SetValue(0, creditToken); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor remoteCreditGlobal; + remoteCreditGlobal.SetGlobalBuffer(remoteCredit, kCreditWords); + AscendC::DataCopy(remoteCreditGlobal, creditLocal, kCreditWords); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +} // namespace +``` + +`PublishCreditMte()` 是从当前 `AllToAllGroupPublishNextCredit()` 提取出的通用形式，数据搬运和同步指令完全相同，只把 grouped peer 计算移到调用侧，方便其它项目复用。 + +### 13.9 Device 调用位置和地址计算 + +接收端 `R` 确认 group `g` 的 payload 到达后，为下一 group 的发送端 `S` 发布 credit： + +```cpp +const uint32_t slot = invocationId & 1U; +const uint64_t slotOffsets[2] = { + 0ULL, + static_cast(TileXR::CREDIT_IPC_SLOT_BYTES) +}; +const uint64_t token = MakeCreditToken(invocationId, group + 1U, 0U); + +// args->creditMems[S] 是 S 拥有的 buffer 在 R 地址空间中的 IPC mapping。 +PublishCreditMte( + args->creditMems[S], slotOffsets[slot], R, token, relayLocal); +``` + +发送端 `S` 在向 `R` 发送 group `g > 0` 前等待： + +```cpp +const uint32_t slot = invocationId & 1U; +const uint64_t slotOffset = slot == 0U + ? 0ULL + : static_cast(TileXR::CREDIT_IPC_SLOT_BYTES); +const uint64_t expected = MakeCreditToken(invocationId, group, 0U); + +// args->creditMems[S] 是 S 自己的 owner buffer；entry R 由 R 写入。 +auto credit = reinterpret_cast<__gm__ uint64_t *>( + args->creditMems[S] + slotOffset + + static_cast(R) * kCreditStride); + +uint64_t observed = 0ULL; +if (!WaitCreditMte( + credit, expected, kCreditWaitTimeoutCycles, + relayLocal, observed)) { + // 必须记录 S、R、group、expected 和 observed，并让整个算子报错。 + return; +} +``` + +Grouped all-to-all 当前实际变量代入为： + +```text +S = 当前 send core 所在 rank +R = 当前 group/lane 算出的 peer +发布侧 next S = 接收 rank 在下一 group/同 lane 算出的 nextPeer +group 0 不 wait +最后一个 group 不再发布 next credit +routeElements == 0 时不 wait，也不发送该 route +``` + +### 13.10 Kernel launch 的实际 Host/Device ABI + +当前代码不是通过 `<<<>>>` 启动，而是注册 kernel binary 后使用 `rtKernelLaunchWithFlagV2()`。Credit kernel 的实际参数结构如下： + +```cpp +struct alignas(8) GroupedAllToAllCreditKernelArgs { + uint8_t* commArgs; + uint8_t* input; + uint8_t* output; + uint8_t* registeredMemory; + uint8_t* debug; + uint32_t invocationId; + int32_t elementsPerPeer; + int32_t chunkElements; + uint32_t passCount; + uint32_t groupCount; + uint64_t payloadOffset0; + uint64_t payloadOffset1; + uint64_t signalOffset0; + uint64_t signalOffset1; + uint64_t creditOffset0; + uint64_t creditOffset1; + uint8_t* groupTrace; + uint32_t traceIteration; + uint32_t routeStage; + uint32_t multiChannel; + uint32_t primaryRouteParts; + uint32_t groupWidth; + uint32_t quietBatch; + uint32_t ingressWindow; +}; + +static_assert(sizeof(GroupedAllToAllCreditKernelArgs) == 152U); +static_assert(offsetof(GroupedAllToAllCreditKernelArgs, creditOffset0) == 96U); +static_assert(offsetof(GroupedAllToAllCreditKernelArgs, groupTrace) == 112U); +``` + +实际组参和 launch 方式： + +```cpp +GroupedAllToAllCreditKernelArgs args { + commArgs, input, output, registeredMemory, debug, invocationId, + elementsPerPeer, chunkElements, passCount, groupCount, + payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, + creditOffset0, creditOffset1, groupTrace, traceIteration, + routeStage, multiChannel, primaryRouteParts, groupWidth, + quietBatch, ingressWindow, +}; + +rtArgsEx_t argsInfo {}; +argsInfo.args = static_cast(&args); +argsInfo.argsSize = sizeof(args); + +rtTaskCfgInfo_t cfgInfo {}; +cfgInfo.schemMode = RT_SCHEM_MODE_NORMAL; +const rtError_t ret = rtKernelLaunchWithFlagV2( + groupedCreditKernelFunctionSignature, + blockDim, + &argsInfo, + nullptr, + static_cast(stream), + 0U, + &cfgInfo); +``` + +Device kernel 声明必须保持完全相同的参数顺序： + +```cpp +extern "C" __global__ __aicore__ void grouped_alltoall_credit_kernel( + GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR outputGM, + GM_ADDR registeredMemoryGM, GM_ADDR debugGM, + uint32_t invocationId, + int32_t elementsPerPeer, int32_t chunkElements, + uint32_t passCount, uint32_t groupCount, + uint64_t payloadOffset0, uint64_t payloadOffset1, + uint64_t signalOffset0, uint64_t signalOffset1, + uint64_t creditOffset0, uint64_t creditOffset1, + GM_ADDR groupTraceGM, uint32_t traceIteration, + uint32_t routeStage, uint32_t multiChannel, + uint32_t primaryRouteParts, uint32_t groupWidth, + uint32_t quietBatch, uint32_t ingressWindow) +{ + // 调用包含 PublishCreditMte/WaitCreditMte 的 kernel 实现。 +} +``` + +如果目标项目使用直接 kernel launch，也仍应保留参数结构的 `sizeof/offsetof` 校验；Host 和 Device 参数错位通常不会在编译期报错，而会表现为错误地址或 kernel 卡死。 + +### 13.11 完整释放代码 + +```cpp +void TileXRComm::CloseCreditIpcMem() +{ + for (int peer = 0; peer < rankSize_; ++peer) { + if (peer == rank_ || creditIpcMem_[peer] == nullptr) { + continue; + } + rtIpcCloseMemory(static_cast(creditIpcMem_[peer])); + creditIpcMem_[peer] = nullptr; + } +} + +TileXRComm::~TileXRComm() +{ + // 析构前必须保证所有使用 creditMems 的 kernel/stream 已结束。 + if (creditIpcMemInited_) { + CloseCreditIpcMem(); + creditIpcMemInited_ = false; + } + + if (creditIpcMem_[rank_] != nullptr) { + aclrtFree(creditIpcMem_[rank_]); + creditIpcMem_[rank_] = nullptr; + } + + if (commArgsPtr_ != nullptr) { + aclrtFree(commArgsPtr_); + commArgsPtr_ = nullptr; + } +} +``` + +如果通信域析构前没有隐式 stream synchronize，目标项目必须显式同步；仅 Host barrier 不能证明 Device 已停止访问 IPC mapping。 From 2719f4140a2d9e3953145faafb5131c5de9b9bb1 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 13:15:41 +0800 Subject: [PATCH 148/163] feat(udma): trace SDMA submit and wait phases --- .../tilexr_udma_alltoall_group_kernel.cpp | 51 ++++++++++++++----- .../demo/tilexr_udma_alltoall_group_trace.h | 6 ++- ...exr_udma_alltoall_group_trace_to_chrome.py | 18 ++++--- ...test_tilexr_udma_alltoall_group_layout.cpp | 2 +- ...exr_udma_alltoall_group_trace_to_chrome.py | 16 ++++++ 5 files changed, 71 insertions(+), 22 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index e24285dd..844c925c 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -296,7 +296,7 @@ __aicore__ inline void AllToAllGroupCopyMte( AscendC::PipeBarrier(); } -__aicore__ inline uint32_t AllToAllGroupCopySdma( +__aicore__ inline uint32_t AllToAllGroupCopySdmaSubmit( const __gm__ TileXR::CommArgs* args, __gm__ uint8_t* dst, __gm__ uint8_t* src, uint32_t bytes, uint32_t channel, uint64_t& event) @@ -309,6 +309,12 @@ __aicore__ inline uint32_t AllToAllGroupCopySdma( if (event == 0ULL) { return TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK; } + return TILEXR_ALLTOALL_GROUP_SDMA_COMPLETE; +} + +__aicore__ inline uint32_t AllToAllGroupCopySdmaWait( + const __gm__ TileXR::CommArgs* args, uint64_t event, uint32_t channel) +{ return TileXR::SDMAWait(args, event, channel) ? TILEXR_ALLTOALL_GROUP_SDMA_COMPLETE : TILEXR_ALLTOALL_GROUP_SDMA_FAILED; @@ -914,9 +920,38 @@ __aicore__ inline void AllToAllGroupKernelImpl( chunkElementOffset + copyElementBegin); const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); uint64_t sdmaEvent = 0ULL; - const uint32_t sdmaStatus = AllToAllGroupCopySdma( + const uint64_t sdmaSubmitBegin = AllToAllGroupTraceCycle(groupTrace); + const uint32_t sdmaSubmitStatus = AllToAllGroupCopySdmaSubmit( args, relayDst, relaySrc, copyBytes, lane, sdmaEvent); - if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK) { + const uint64_t sdmaSubmitEnd = AllToAllGroupTraceCycle(groupTrace); + if (sdmaSubmitStatus != TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK) { + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceSdmaSubmit, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + sdmaSubmitBegin, sdmaSubmitEnd); + const uint64_t sdmaWaitBegin = AllToAllGroupTraceCycle(groupTrace); + const uint32_t sdmaStatus = AllToAllGroupCopySdmaWait( + args, sdmaEvent, lane); + const uint64_t sdmaWaitEnd = AllToAllGroupTraceCycle(groupTrace); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceSdmaWait, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + sdmaWaitBegin, sdmaWaitEnd); + if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FAILED) { + AllToAllGroupTraceRecordError( + debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_SDMA, + group, pass, peer, worker, 0U, + static_cast(worker), sdmaEvent); + AllToAllGroupTraceRecordKernel( + groupTrace, traceIteration, blockIdx, kernelBegin, + AllToAllGroupTraceCycle(groupTrace)); + return; + } + } else { AllToAllGroupCopyMte(relayDst, relaySrc, copyBytes, relayLocal); } AllToAllGroupTraceRecordTask( @@ -924,16 +959,6 @@ __aicore__ inline void AllToAllGroupKernelImpl( TileXR::Demo::kAllToAllGroupTraceReceiveCopy, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); - if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FAILED) { - AllToAllGroupRecordError( - debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_SDMA, - group, pass, peer, worker, 0U, - static_cast(worker), sdmaEvent); - AllToAllGroupTraceRecordKernel( - groupTrace, traceIteration, blockIdx, kernelBegin, - AllToAllGroupTraceCycle(groupTrace)); - return; - } } } } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index ac1b5691..2a2bf6e6 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -14,12 +14,12 @@ namespace TileXR { namespace Demo { constexpr uint32_t kAllToAllGroupTraceMagic = 0x47545243U; // "GTRC" -constexpr uint32_t kAllToAllGroupTraceVersion = 2U; +constexpr uint32_t kAllToAllGroupTraceVersion = 3U; constexpr size_t kAllToAllGroupTraceBytes = 128ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; -constexpr uint32_t kAllToAllGroupTracePhaseCount = 6U; +constexpr uint32_t kAllToAllGroupTracePhaseCount = 8U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; constexpr size_t kAllToAllGroupTraceCacheLineBytes = 128U; @@ -31,6 +31,8 @@ enum AllToAllGroupTracePhase : uint32_t { kAllToAllGroupTraceReceiveWait = 3U, kAllToAllGroupTraceReceiveCopy = 4U, kAllToAllGroupTraceCreditWait = 5U, + kAllToAllGroupTraceSdmaSubmit = 6U, + kAllToAllGroupTraceSdmaWait = 7U, }; struct AllToAllGroupTraceSpan { diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index a6d8f935..0247e1c9 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -6,7 +6,7 @@ TRACE_MAGIC = 0x47545243 -TRACE_VERSION = 2 +TRACE_VERSION = 3 LEGACY_TRACE_VERSION = 1 TRACE_BYTES = 128 * 1024 * 1024 HEADER_BYTES = 4096 @@ -15,7 +15,9 @@ SEND_WORKER_COUNT = 32 LANE_COUNT = 16 PHASE_COUNT = 5 -CURRENT_PHASE_COUNT = 6 +TRACE_V2_PHASE_COUNT = 6 +TRACE_V3_PHASE_COUNT = 8 +CURRENT_PHASE_COUNT = TRACE_V3_PHASE_COUNT SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 TASK_FORMAT = " Date: Tue, 4 Aug 2026 13:29:09 +0800 Subject: [PATCH 149/163] fix(udma): use grouped error recorder in SDMA trace --- tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp | 2 +- tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 844c925c..084992c7 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -942,7 +942,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( TileXR::Demo::kAllToAllGroupTraceNoQp, sdmaWaitBegin, sdmaWaitEnd); if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FAILED) { - AllToAllGroupTraceRecordError( + AllToAllGroupRecordError( debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_SDMA, group, pass, peer, worker, 0U, static_cast(worker), sdmaEvent); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index a3985a14..8d85ee05 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -617,6 +617,9 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SDMA_FAILED"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_SDMA"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaSubmit"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaWait"); + CHECK_NOT_CONTAINS(kernel, "AllToAllGroupTraceRecordError"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers"); CHECK_CONTAINS(kernel, "const uint32_t traceCore = copyoutWorkers < TILEXR_ALLTOALL_GROUP_SEND_CORES"); From bea490a9a580446785840ff75d5acf3c1b765dba Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 13:53:40 +0800 Subject: [PATCH 150/163] feat(udma): trace SDMA submit pipeline stages --- src/include/tilexr_sdma.h | 7 +- src/include/tilexr_sdma_a5.h | 29 ++++++- src/include/tilexr_sdma_types.h | 16 ++++ .../tilexr_udma_alltoall_group_kernel.cpp | 80 +++++++++++++++---- .../demo/tilexr_udma_alltoall_group_trace.h | 15 +++- ...exr_udma_alltoall_group_trace_to_chrome.py | 78 ++++++++++++------ ...test_tilexr_udma_alltoall_group_layout.cpp | 8 +- ...exr_udma_alltoall_group_trace_to_chrome.py | 42 ++++++++-- 8 files changed, 223 insertions(+), 52 deletions(-) diff --git a/src/include/tilexr_sdma.h b/src/include/tilexr_sdma.h index fd5b6035..4733a0a3 100644 --- a/src/include/tilexr_sdma.h +++ b/src/include/tilexr_sdma.h @@ -56,7 +56,8 @@ __aicore__ inline uint64_t SDMACopyNbi( __gm__ uint8_t* dst, __gm__ uint8_t* src, uint64_t bytes, - uint32_t channelGroupIdx = TILEXR_SDMA_AUTO_CHANNEL_GROUP) + uint32_t channelGroupIdx = TILEXR_SDMA_AUTO_CHANNEL_GROUP, + SDMASubmitTrace* submitTrace = nullptr) { #if TILEXR_SDMA_A5_AICORE_COMPILE if (!SDMAEnabled(args) || dst == nullptr || src == nullptr || bytes == 0) { @@ -65,8 +66,9 @@ __aicore__ inline uint64_t SDMACopyNbi( const uint32_t resolvedGroup = SDMAResolveChannelGroup(channelGroupIdx); return detail::A5SdmaCopyNbi( reinterpret_cast<__gm__ uint8_t*>(args->sdmaWorkspacePtr), - dst, src, bytes, resolvedGroup); + dst, src, bytes, resolvedGroup, submitTrace); #elif defined(TILEXR_HAVE_PTO_SDMA) && TILEXR_HAVE_PTO_SDMA + (void)submitTrace; if (!SDMAEnabled(args) || dst == nullptr || src == nullptr || bytes == 0) { return 0; } @@ -85,6 +87,7 @@ __aicore__ inline uint64_t SDMACopyNbi( (void)src; (void)bytes; (void)channelGroupIdx; + (void)submitTrace; return 0; #endif } diff --git a/src/include/tilexr_sdma_a5.h b/src/include/tilexr_sdma_a5.h index cce2c087..6e96aa41 100644 --- a/src/include/tilexr_sdma_a5.h +++ b/src/include/tilexr_sdma_a5.h @@ -8,6 +8,7 @@ #include "kernel_operator.h" #include "tilexr_sdma_a5_types.h" +#include "tilexr_sdma_types.h" namespace TileXR { namespace detail { @@ -90,8 +91,12 @@ __aicore__ inline uint64_t A5SdmaCopyNbi(__gm__ uint8_t* workspaceAddress, __gm__ uint8_t* destination, __gm__ uint8_t* source, uint64_t bytes, - uint32_t channelIndex) + uint32_t channelIndex, + SDMASubmitTrace* submitTrace) { + if (submitTrace != nullptr) { + submitTrace->prepareBegin = static_cast(AscendC::GetSystemCycle()); + } __gm__ A5SdmaWorkspace* workspace = reinterpret_cast<__gm__ A5SdmaWorkspace*>(workspaceAddress); if (!A5SdmaWorkspaceValid(workspace) || destination == nullptr || source == nullptr || @@ -115,6 +120,13 @@ __aicore__ inline uint64_t A5SdmaCopyNbi(__gm__ uint8_t* workspaceAddress, const uint32_t dataIndex = channel->tail; const uint32_t completionIndex = (dataIndex + 1U) % channel->depth; const uint32_t newTail = A5SdmaAdvanceTail(dataIndex, channel->depth); + if (submitTrace != nullptr) { + submitTrace->head = channel->head; + submitTrace->tail = dataIndex; + submitTrace->newTail = newTail; + submitTrace->depth = channel->depth; + submitTrace->generation = generation; + } __gm__ A5SdmaCompletionLine* payload = reinterpret_cast<__gm__ A5SdmaCompletionLine*>( channel->completionPayloadAddress); __gm__ A5SdmaCompletionLine* completion = reinterpret_cast<__gm__ A5SdmaCompletionLine*>( @@ -140,6 +152,10 @@ __aicore__ inline uint64_t A5SdmaCopyNbi(__gm__ uint8_t* workspaceAddress, channel->generation = generation; channel->tail = newTail; channel->taskId = A5SdmaAdvanceTaskId(channel->taskId); + if (submitTrace != nullptr) { + submitTrace->prepareEnd = static_cast(AscendC::GetSystemCycle()); + submitTrace->cacheCleanBegin = submitTrace->prepareEnd; + } pipe_barrier(PIPE_ALL); A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(payload)); A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(completion)); @@ -147,9 +163,20 @@ __aicore__ inline uint64_t A5SdmaCopyNbi(__gm__ uint8_t* workspaceAddress, A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(sqBase + completionIndex)); A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(channel)); A5SdmaCleanCacheLine(reinterpret_cast<__gm__ uint8_t*>(channel) + 64U); + if (submitTrace != nullptr) { + submitTrace->cacheCleanEnd = static_cast(AscendC::GetSystemCycle()); + submitTrace->dsbBegin = submitTrace->cacheCleanEnd; + } pipe_barrier(PIPE_ALL); dsb(DSB_DDR); + if (submitTrace != nullptr) { + submitTrace->dsbEnd = static_cast(AscendC::GetSystemCycle()); + submitTrace->doorbellBegin = submitTrace->dsbEnd; + } A5SdmaRingDoorbell(channel->rtsqAddress, newTail); + if (submitTrace != nullptr) { + submitTrace->doorbellEnd = static_cast(AscendC::GetSystemCycle()); + } return A5SdmaEncodeEvent(channelIndex, generation); } diff --git a/src/include/tilexr_sdma_types.h b/src/include/tilexr_sdma_types.h index 62a3d40d..4039b011 100644 --- a/src/include/tilexr_sdma_types.h +++ b/src/include/tilexr_sdma_types.h @@ -25,6 +25,22 @@ enum class SDMAInitStatus : int32_t { NULL_WORKSPACE = 4, }; +struct SDMASubmitTrace { + uint64_t prepareBegin; + uint64_t prepareEnd; + uint64_t cacheCleanBegin; + uint64_t cacheCleanEnd; + uint64_t dsbBegin; + uint64_t dsbEnd; + uint64_t doorbellBegin; + uint64_t doorbellEnd; + uint32_t head; + uint32_t tail; + uint32_t newTail; + uint32_t depth; + uint32_t generation; +}; + } // namespace TileXR #endif // TILEXR_SDMA_TYPES_H diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 084992c7..daef1b7c 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -299,13 +299,14 @@ __aicore__ inline void AllToAllGroupCopyMte( __aicore__ inline uint32_t AllToAllGroupCopySdmaSubmit( const __gm__ TileXR::CommArgs* args, __gm__ uint8_t* dst, __gm__ uint8_t* src, uint32_t bytes, uint32_t channel, - uint64_t& event) + uint64_t& event, TileXR::SDMASubmitTrace& submitTrace) { event = 0ULL; if (!TileXR::SDMAEnabled(args)) { return TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK; } - event = TileXR::SDMACopyNbi(args, dst, src, bytes, channel); + event = TileXR::SDMACopyNbi( + args, dst, src, bytes, channel, &submitTrace); if (event == 0ULL) { return TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK; } @@ -504,7 +505,10 @@ __aicore__ inline void AllToAllGroupTraceRecordTask( __gm__ uint8_t* trace, uint32_t iteration, uint32_t core, uint32_t group, uint32_t pass, uint32_t phase, uint32_t groupCount, uint32_t passCount, int32_t peer, uint32_t qpIdx, - uint64_t beginCycle, uint64_t endCycle) + uint64_t beginCycle, uint64_t endCycle, + uint32_t sdmaHead = 0U, uint32_t sdmaTail = 0U, + uint32_t sdmaNewTail = 0U, uint32_t sdmaDepth = 0U, + uint32_t sdmaGeneration = 0U) { if (trace == nullptr || iteration >= TileXR::Demo::kAllToAllGroupTraceMaxIterations || core >= TileXR::Demo::kAllToAllGroupTraceCoreCount || group >= groupCount || @@ -531,6 +535,11 @@ __aicore__ inline void AllToAllGroupTraceRecordTask( trace + offset); span->peer = peer; span->qpIdx = qpIdx; + span->sdmaHead = sdmaHead; + span->sdmaTail = sdmaTail; + span->sdmaNewTail = sdmaNewTail; + span->sdmaDepth = sdmaDepth; + span->sdmaGeneration = sdmaGeneration; span->beginCycle = beginCycle; span->endCycle = endCycle; } @@ -920,27 +929,69 @@ __aicore__ inline void AllToAllGroupKernelImpl( chunkElementOffset + copyElementBegin); const uint64_t receiveCopyBegin = AllToAllGroupTraceCycle(groupTrace); uint64_t sdmaEvent = 0ULL; + TileXR::SDMASubmitTrace sdmaTrace {}; const uint64_t sdmaSubmitBegin = AllToAllGroupTraceCycle(groupTrace); const uint32_t sdmaSubmitStatus = AllToAllGroupCopySdmaSubmit( - args, relayDst, relaySrc, copyBytes, lane, sdmaEvent); + args, relayDst, relaySrc, copyBytes, lane, sdmaEvent, sdmaTrace); const uint64_t sdmaSubmitEnd = AllToAllGroupTraceCycle(groupTrace); if (sdmaSubmitStatus != TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK) { + const uint64_t sdmaWaitBegin = AllToAllGroupTraceCycle(groupTrace); + const uint32_t sdmaStatus = AllToAllGroupCopySdmaWait( + args, sdmaEvent, lane); + const uint64_t sdmaWaitEnd = AllToAllGroupTraceCycle(groupTrace); + const uint64_t receiveCopyEnd = AllToAllGroupTraceCycle(groupTrace); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceReceiveCopy, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + receiveCopyBegin, receiveCopyEnd); AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceSdmaSubmit, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, - sdmaSubmitBegin, sdmaSubmitEnd); - const uint64_t sdmaWaitBegin = AllToAllGroupTraceCycle(groupTrace); - const uint32_t sdmaStatus = AllToAllGroupCopySdmaWait( - args, sdmaEvent, lane); - const uint64_t sdmaWaitEnd = AllToAllGroupTraceCycle(groupTrace); + sdmaSubmitBegin, sdmaSubmitEnd, + sdmaTrace.head, sdmaTrace.tail, sdmaTrace.newTail, + sdmaTrace.depth, sdmaTrace.generation); AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceSdmaWait, groupCount, passCount, peer, TileXR::Demo::kAllToAllGroupTraceNoQp, sdmaWaitBegin, sdmaWaitEnd); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceSdmaPrepare, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + sdmaTrace.prepareBegin, sdmaTrace.prepareEnd, + sdmaTrace.head, sdmaTrace.tail, sdmaTrace.newTail, + sdmaTrace.depth, sdmaTrace.generation); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceSdmaCacheClean, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + sdmaTrace.cacheCleanBegin, sdmaTrace.cacheCleanEnd, + sdmaTrace.head, sdmaTrace.tail, sdmaTrace.newTail, + sdmaTrace.depth, sdmaTrace.generation); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceSdmaDsb, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + sdmaTrace.dsbBegin, sdmaTrace.dsbEnd, + sdmaTrace.head, sdmaTrace.tail, sdmaTrace.newTail, + sdmaTrace.depth, sdmaTrace.generation); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceSdmaDoorbell, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + sdmaTrace.doorbellBegin, sdmaTrace.doorbellEnd, + sdmaTrace.head, sdmaTrace.tail, sdmaTrace.newTail, + sdmaTrace.depth, sdmaTrace.generation); if (sdmaStatus == TILEXR_ALLTOALL_GROUP_SDMA_FAILED) { AllToAllGroupRecordError( debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_SDMA, @@ -953,12 +1004,13 @@ __aicore__ inline void AllToAllGroupKernelImpl( } } else { AllToAllGroupCopyMte(relayDst, relaySrc, copyBytes, relayLocal); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, traceCore, group, pass, + TileXR::Demo::kAllToAllGroupTraceReceiveCopy, + groupCount, passCount, peer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } - AllToAllGroupTraceRecordTask( - groupTrace, traceIteration, traceCore, group, pass, - TileXR::Demo::kAllToAllGroupTraceReceiveCopy, groupCount, passCount, - peer, TileXR::Demo::kAllToAllGroupTraceNoQp, - receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } } } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h index 2a2bf6e6..59d2a655 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace.h @@ -14,12 +14,12 @@ namespace TileXR { namespace Demo { constexpr uint32_t kAllToAllGroupTraceMagic = 0x47545243U; // "GTRC" -constexpr uint32_t kAllToAllGroupTraceVersion = 3U; +constexpr uint32_t kAllToAllGroupTraceVersion = 4U; constexpr size_t kAllToAllGroupTraceBytes = 128ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupTraceHeaderBytes = 4096ULL; constexpr uint32_t kAllToAllGroupTraceMaxIterations = 50U; constexpr uint32_t kAllToAllGroupTraceCoreCount = 64U; -constexpr uint32_t kAllToAllGroupTracePhaseCount = 8U; +constexpr uint32_t kAllToAllGroupTracePhaseCount = 12U; constexpr uint32_t kAllToAllGroupTraceNoQp = 0xFFFFFFFFU; constexpr uint64_t kAllToAllGroupTraceCyclesPerUs = 1000ULL; constexpr size_t kAllToAllGroupTraceCacheLineBytes = 128U; @@ -33,6 +33,10 @@ enum AllToAllGroupTracePhase : uint32_t { kAllToAllGroupTraceCreditWait = 5U, kAllToAllGroupTraceSdmaSubmit = 6U, kAllToAllGroupTraceSdmaWait = 7U, + kAllToAllGroupTraceSdmaPrepare = 8U, + kAllToAllGroupTraceSdmaCacheClean = 9U, + kAllToAllGroupTraceSdmaDsb = 10U, + kAllToAllGroupTraceSdmaDoorbell = 11U, }; struct AllToAllGroupTraceSpan { @@ -45,6 +49,11 @@ struct AllToAllGroupTraceTaskSpan { uint64_t endCycle; int32_t peer; uint32_t qpIdx; + uint32_t sdmaHead; + uint32_t sdmaTail; + uint32_t sdmaNewTail; + uint32_t sdmaDepth; + uint32_t sdmaGeneration; }; struct AllToAllGroupTraceHeader { @@ -146,7 +155,7 @@ inline size_t AllToAllGroupTraceTaskSpanOffset( static_assert(sizeof(AllToAllGroupTraceSpan) == 16U, "group trace kernel span must contain two uint64 timestamps"); -static_assert(sizeof(AllToAllGroupTraceTaskSpan) == 24U, +static_assert(sizeof(AllToAllGroupTraceTaskSpan) == 48U, "group trace task span layout changed"); static_assert(sizeof(AllToAllGroupTraceHeader) <= kAllToAllGroupTraceHeaderBytes, "group trace header must fit its region"); diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py index 0247e1c9..e5294640 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/demo/tilexr_udma_alltoall_group_trace_to_chrome.py @@ -6,7 +6,7 @@ TRACE_MAGIC = 0x47545243 -TRACE_VERSION = 3 +TRACE_VERSION = 4 LEGACY_TRACE_VERSION = 1 TRACE_BYTES = 128 * 1024 * 1024 HEADER_BYTES = 4096 @@ -17,10 +17,13 @@ PHASE_COUNT = 5 TRACE_V2_PHASE_COUNT = 6 TRACE_V3_PHASE_COUNT = 8 -CURRENT_PHASE_COUNT = TRACE_V3_PHASE_COUNT +TRACE_V4_PHASE_COUNT = 12 +CURRENT_PHASE_COUNT = TRACE_V4_PHASE_COUNT SPAN_BYTES = 16 CACHE_LINE_BYTES = 128 -TASK_FORMAT = " len(data): +def _read_task(data, offset, label, task_format, task_bytes): + if offset < 0 or offset + task_bytes > len(data): raise ValueError(f"task offset out of range for {label}: {offset}") - begin, end, peer, qp = struct.unpack_from(TASK_FORMAT, data, offset) + fields = struct.unpack_from(task_format, data, offset) + begin, end, peer, qp = fields[:4] if begin == 0 and end == 0: return None if begin == 0 or end == 0: raise ValueError(f"incomplete {label}: begin={begin} end={end}") if end < begin: raise ValueError(f"invalid {label}: begin={begin} end={end}") - return begin, end, peer, qp + metadata = fields[4:] if len(fields) > 4 else (0, 0, 0, 0, 0) + return begin, end, peer, qp, *metadata def read_rank_trace(path): @@ -111,7 +120,7 @@ def read_rank_trace(path): } if header["magic"] != TRACE_MAGIC: raise ValueError(f"invalid trace magic in {path}") - if header["version"] not in (LEGACY_TRACE_VERSION, 2, TRACE_VERSION): + if header["version"] not in (LEGACY_TRACE_VERSION, 2, 3, TRACE_VERSION): raise ValueError(f"unsupported trace version {header['version']} in {path}") if header["trace_bytes"] != TRACE_BYTES: raise ValueError(f"trace byte dimension mismatch in {path}") @@ -122,8 +131,11 @@ def read_rank_trace(path): expected_phase_count = { LEGACY_TRACE_VERSION: PHASE_COUNT, 2: TRACE_V2_PHASE_COUNT, - TRACE_VERSION: TRACE_V3_PHASE_COUNT, + 3: TRACE_V3_PHASE_COUNT, + TRACE_VERSION: TRACE_V4_PHASE_COUNT, }[header["version"]] + task_format = TASK_FORMAT if header["version"] == TRACE_VERSION else LEGACY_TASK_FORMAT + task_bytes = TASK_BYTES if header["version"] == TRACE_VERSION else LEGACY_TASK_BYTES if (header["core_count"] != MAX_CORES or header["phase_count"] != expected_phase_count): raise ValueError(f"core/phase dimension mismatch in {path}") @@ -134,13 +146,15 @@ def read_rank_trace(path): raise ValueError(f"invalid cycle frequency in {path}") required = layout_bytes( header["iteration_count"], header["group_count"], - header["pass_count"], header["phase_count"]) + header["pass_count"], header["phase_count"], task_bytes) if required > TRACE_BYTES: raise ValueError(f"trace capacity exceeded in {path}: required={required}") with path.open("rb") as stream: data = stream.read(required) if len(data) != required: raise ValueError(f"short trace read in {path}: read={len(data)} required={required}") + header["task_format"] = task_format + header["task_bytes"] = task_bytes return {"path": str(path), "header": header, "data": data} @@ -225,24 +239,36 @@ def build_chrome_trace(rank_traces): task_span_offset( iteration, core, group, pass_index, phase, header["group_count"], header["pass_count"], - header["phase_count"]), + header["phase_count"], header["task_bytes"]), label, + header["task_format"], header["task_bytes"], ) if task is None: continue - begin, end, peer, qp = task + (begin, end, peer, qp, sdma_head, sdma_tail, + sdma_new_tail, sdma_depth, sdma_generation) = task + task_args = { + "iteration": iteration, + "group": group, + "pass": pass_index, + "lane": core % LANE_COUNT, + "peer": peer, + "qp": None if qp == NO_QP else qp, + "role": role, + } + if sdma_depth != 0: + task_args.update({ + "sdmaHead": sdma_head, + "sdmaTail": sdma_tail, + "sdmaNewTail": sdma_new_tail, + "sdmaDepth": sdma_depth, + "sdmaGeneration": sdma_generation, + "sdmaWrapped": sdma_new_tail < sdma_tail, + }) events.append(_event( PHASE_NAMES[phase], role, rank, core, begin, end, base, header["cycles_per_us"], - { - "iteration": iteration, - "group": group, - "pass": pass_index, - "lane": core % LANE_COUNT, - "peer": peer, - "qp": None if qp == NO_QP else qp, - "role": role, - }, + task_args, offset_us, )) return { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 8d85ee05..178dd8fb 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -292,7 +292,9 @@ void TestScalePlanAndTraceCapacity() CHECK_EQ(TileXR::Demo::kAllToAllGroupTraceBytes, 128ULL * 1024ULL * 1024ULL); - CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 3U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 1U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 2U), false); + CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 3U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 4U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 5U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupTraceLayoutFits(50U, 64U, 6U), false); @@ -619,6 +621,10 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_SDMA"); CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaSubmit"); CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaWait"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaPrepare"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaCacheClean"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaDsb"); + CHECK_CONTAINS(kernel, "kAllToAllGroupTraceSdmaDoorbell"); CHECK_NOT_CONTAINS(kernel, "AllToAllGroupTraceRecordError"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers"); CHECK_CONTAINS(kernel, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py index c417e3e7..13f3fca9 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_trace_to_chrome.py @@ -36,6 +36,10 @@ def make_trace( version = MODULE.TRACE_VERSION if version is None else version phase_count = ( MODULE.CURRENT_PHASE_COUNT if phase_count is None else phase_count) + task_format = ( + MODULE.TASK_FORMAT if version == MODULE.TRACE_VERSION + else MODULE.LEGACY_TASK_FORMAT) + task_bytes = struct.calcsize(task_format) header = struct.pack( MODULE.HEADER_FORMAT, MODULE.TRACE_MAGIC if magic is None else magic, @@ -61,6 +65,10 @@ def make_trace( (0, 0, 0, 5, 1050, 1100, 1, MODULE.NO_QP), (32, 0, 0, 6, 1600, 1650, 1, MODULE.NO_QP), (32, 0, 0, 7, 1650, 1700, 1, MODULE.NO_QP), + (32, 0, 0, 8, 1605, 1615, 1, MODULE.NO_QP), + (32, 0, 0, 9, 1615, 1625, 1, MODULE.NO_QP), + (32, 0, 0, 10, 1625, 1635, 1, MODULE.NO_QP), + (32, 0, 0, 11, 1635, 1645, 1, MODULE.NO_QP), ) with path.open("wb") as stream: stream.truncate(MODULE.TRACE_BYTES) @@ -74,9 +82,11 @@ def make_trace( continue stream.seek(MODULE.task_span_offset( 0, core, group, pass_index, phase, group_count, pass_count, - phase_count)) - stream.write(struct.pack( - MODULE.TASK_FORMAT, begin, end, peer, qp)) + phase_count, task_bytes)) + values = [begin, end, peer, qp] + if version == MODULE.TRACE_VERSION: + values.extend([0, 62, 0, 64, 33]) + stream.write(struct.pack(task_format, *values)) def write_at(self, path, offset, payload): with path.open("r+b") as stream: @@ -108,6 +118,8 @@ def test_converts_all_grouped_pipeline_phases(self): "kernel", "self-copy", "send-put-signal", "send-quiet", "receive-wait", "receive-copy", "credit-wait", "sdma-submit", "sdma-wait", + "sdma-prepare", "sdma-cache-clean", "sdma-dsb", + "sdma-doorbell", }, ) send = next(event for event in complete if event["name"] == "send-put-signal") @@ -120,6 +132,11 @@ def test_converts_all_grouped_pipeline_phases(self): self.assertEqual(second_send["args"]["role"], "send") self.assertEqual(second_send["args"]["lane"], 0) self.assertEqual(trace["otherData"]["displayTimeUnit"], "ns") + submit = next(event for event in complete if event["name"] == "sdma-submit") + self.assertEqual(submit["args"]["sdmaTail"], 62) + self.assertEqual(submit["args"]["sdmaNewTail"], 0) + self.assertEqual(submit["args"]["sdmaDepth"], 64) + self.assertTrue(submit["args"]["sdmaWrapped"]) json.loads(json.dumps(trace)) def test_labels_32_send_and_32_receive_cores(self): @@ -170,7 +187,8 @@ def test_normalizes_ranks_independently(self): offset = MODULE.task_span_offset(0, 0, 0, 0, 1, 1, 1) self.write_at( second, offset, - struct.pack(MODULE.TASK_FORMAT, 5200, 5300, 2, 3)) + struct.pack(MODULE.TASK_FORMAT, 5200, 5300, 2, 3, + 0, 0, 0, 0, 0)) trace = MODULE.build_chrome_trace([ MODULE.read_rank_trace(first), MODULE.read_rank_trace(second)]) @@ -206,7 +224,8 @@ def test_rejects_half_written_span(self): offset = MODULE.task_span_offset(0, 0, 0, 0, 1, 1, 1) self.write_at( path, offset, - struct.pack(MODULE.TASK_FORMAT, 0, 1300, 1, 3)) + struct.pack(MODULE.TASK_FORMAT, 0, 1300, 1, 3, + 0, 0, 0, 0, 0)) with self.assertRaisesRegex(ValueError, "incomplete"): MODULE.build_chrome_trace([MODULE.read_rank_trace(path)]) @@ -240,6 +259,19 @@ def test_reads_version_two_six_phase_trace(self): self.assertIn("credit-wait", names) self.assertNotIn("sdma-submit", names) + def test_reads_version_three_eight_phase_trace(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "version3.bin" + self.make_trace(path, version=3, phase_count=MODULE.TRACE_V3_PHASE_COUNT) + rank_trace = MODULE.read_rank_trace(path) + self.assertEqual(rank_trace["header"]["phase_count"], 8) + names = { + event["name"] for event in MODULE.build_chrome_trace([rank_trace])["traceEvents"] + if event.get("ph") == "X" + } + self.assertIn("sdma-submit", names) + self.assertNotIn("sdma-prepare", names) + def test_main_writes_json_without_dumps(self): with tempfile.TemporaryDirectory() as directory: source = Path(directory) / "rank0.bin" From 24a95d0d05baf035d1345eabede44952617d866c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 16:04:04 +0800 Subject: [PATCH 151/163] perf(udma): close grouped ingress credit window --- .../tilexr_udma_alltoall_group_kernel.cpp | 122 ++++++++++++++++-- .../demo/tilexr_udma_alltoall_group_layout.h | 13 ++ ...test_tilexr_udma_alltoall_group_layout.cpp | 53 +++++++- 3 files changed, 172 insertions(+), 16 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index daef1b7c..1710aef6 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -102,6 +102,12 @@ __aicore__ inline uint64_t AllToAllGroupDeviceToken( (static_cast(pass) + 1ULL); } +__aicore__ inline uint64_t AllToAllGroupDeviceTerminalCreditToken( + uint32_t invocationId, uint32_t groupCount) +{ + return AllToAllGroupDeviceToken(invocationId, groupCount, 0U); +} + __aicore__ inline bool AllToAllGroupPeerInRouteStageDevice( int32_t rank, int32_t peer, uint32_t routeStage) { @@ -423,22 +429,13 @@ __aicore__ inline int32_t AllToAllGroupNextCreditPeerDevice( rank, rankSize, completedGroup + 1U, lane, groupWidth); } -__aicore__ inline void AllToAllGroupPublishNextCredit( +__aicore__ inline void AllToAllGroupPublishCredit( const __gm__ TileXR::CommArgs* args, - int32_t rank, int32_t rankSize, uint32_t invocationId, - uint32_t completedGroup, uint32_t lane, uint32_t groupCount, - uint32_t groupWidth, uint64_t creditOffset, + int32_t rank, int32_t peer, uint64_t creditToken, uint64_t creditOffset, AscendC::LocalTensor relayLocal) { - const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( - rank, rankSize, completedGroup, lane, groupCount, groupWidth); - if (nextPeer < 0) { - return; - } - const uint64_t creditToken = - AllToAllGroupDeviceToken(invocationId, completedGroup + 1U, 0U); auto remoteCredit = reinterpret_cast<__gm__ uint64_t*>( - args->creditMems[nextPeer] + creditOffset + + args->creditMems[peer] + creditOffset + static_cast(rank) * TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE); auto creditLocal = relayLocal.ReinterpretCast(); creditLocal.SetValue(0, creditToken); @@ -455,6 +452,51 @@ __aicore__ inline void AllToAllGroupPublishNextCredit( AscendC::WaitFlag(EVENT_ID0); } +__aicore__ inline void AllToAllGroupPublishNextCredit( + const __gm__ TileXR::CommArgs* args, + int32_t rank, int32_t rankSize, uint32_t invocationId, + uint32_t completedGroup, uint32_t lane, uint32_t groupCount, + uint32_t groupWidth, uint64_t creditOffset, + AscendC::LocalTensor relayLocal) +{ + const int32_t nextPeer = AllToAllGroupNextCreditPeerDevice( + rank, rankSize, completedGroup, lane, groupCount, groupWidth); + if (nextPeer < 0) { + return; + } + AllToAllGroupPublishCredit( + args, rank, nextPeer, + AllToAllGroupDeviceToken(invocationId, completedGroup + 1U, 0U), + creditOffset, relayLocal); +} + +__aicore__ inline void AllToAllGroupPublishTerminalCredits( + const __gm__ TileXR::CommArgs* args, + int32_t rank, int32_t rankSize, uint32_t invocationId, + uint32_t worker, uint32_t copyoutWorkers, uint32_t groupCount, + uint32_t groupWidth, uint64_t creditOffset, + AscendC::LocalTensor relayLocal) +{ + if (!AllToAllGroupCreditOwnerDevice(worker)) { + return; + } + const uint64_t terminalToken = + AllToAllGroupDeviceTerminalCreditToken(invocationId, groupCount); + for (uint32_t assignment = 0U; ; ++assignment) { + const int32_t laneValue = AllToAllGroupCopyoutLaneDevice( + worker, assignment, copyoutWorkers); + if (laneValue < 0) { + return; + } + const int32_t firstPeer = AllToAllGroupDevicePeer( + rank, rankSize, 0U, static_cast(laneValue), groupWidth); + if (firstPeer >= 0) { + AllToAllGroupPublishCredit( + args, rank, firstPeer, terminalToken, creditOffset, relayLocal); + } + } +} + __aicore__ inline void AllToAllGroupRecordError( __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, @@ -544,6 +586,45 @@ __aicore__ inline void AllToAllGroupTraceRecordTask( span->endCycle = endCycle; } +__aicore__ inline bool AllToAllGroupWaitTerminalCredit( + const __gm__ TileXR::CommArgs* args, + int32_t rank, int32_t rankSize, uint32_t invocationId, + uint32_t lane, uint32_t groupCount, uint32_t passCount, + uint32_t groupWidth, uint64_t creditOffset, + AscendC::LocalTensor relayLocal, + __gm__ int32_t* debug, uint32_t blockIdx, + __gm__ uint8_t* groupTrace, uint32_t traceIteration) +{ + const int32_t firstPeer = AllToAllGroupDevicePeer( + rank, rankSize, 0U, lane, groupWidth); + if (firstPeer < 0) { + return true; + } + const uint64_t expectedCredit = + AllToAllGroupDeviceTerminalCreditToken(invocationId, groupCount); + auto creditSignal = reinterpret_cast<__gm__ uint64_t*>( + args->creditMems[rank] + creditOffset + + static_cast(firstPeer) * TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE); + const uint64_t creditWaitBegin = AllToAllGroupTraceCycle(groupTrace); + uint64_t observedCredit = 0ULL; + const bool ok = AllToAllGroupWaitCreditMte( + creditSignal, expectedCredit, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, + relayLocal, observedCredit); + AllToAllGroupTraceRecordTask( + groupTrace, traceIteration, blockIdx, 0U, 0U, + TileXR::Demo::kAllToAllGroupTraceCreditWait, + groupCount, passCount, firstPeer, + TileXR::Demo::kAllToAllGroupTraceNoQp, + creditWaitBegin, AllToAllGroupTraceCycle(groupTrace)); + if (!ok) { + AllToAllGroupRecordError( + debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT, + groupCount, 0U, firstPeer, 0U, 0U, + expectedCredit, observedCredit); + } + return ok; +} + struct AllToAllGroupPendingQuiet { int32_t peer; uint32_t qpIdx; @@ -1014,6 +1095,12 @@ __aicore__ inline void AllToAllGroupKernelImpl( } } } + if constexpr (IngressCredit) { + AscendC::SyncAll(); + AllToAllGroupPublishTerminalCredits( + args, rank, rankSize, invocationId, worker, copyoutWorkers, + groupCount, groupWidth, creditOffsets[slot], relayLocal); + } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; @@ -1189,6 +1276,17 @@ __aicore__ inline void AllToAllGroupKernelImpl( kernelBegin, AllToAllGroupTraceCycle(groupTrace)); return; } + if constexpr (IngressCredit) { + AscendC::SyncAll(); + if (workerRoute == 0U && !AllToAllGroupWaitTerminalCredit( + args, rank, rankSize, invocationId, lane, groupCount, passCount, + groupWidth, creditOffsets[slot], relayLocal, + debug, blockIdx, groupTrace, traceIteration)) { + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } + } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index e8a3ab25..a1cf52d4 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -207,6 +207,13 @@ inline int32_t AllToAllGroupNextCreditPeer( rank, rankSize, completedGroup + 1U, lane, groupWidth); } +inline int32_t AllToAllGroupTerminalCreditPeer( + int rank, int rankSize, uint32_t lane, + uint32_t groupWidth = kAllToAllGroupWidth) +{ + return AllToAllGroupPeer(rank, rankSize, 0U, lane, groupWidth); +} + inline bool AllToAllGroupCreditOwner(uint32_t copyoutWorker) { return copyoutWorker < kAllToAllGroupSendCoreCount; @@ -228,6 +235,12 @@ inline uint64_t AllToAllGroupCreditToken( return AllToAllGroupToken(invocationId, group, 0U); } +inline uint64_t AllToAllGroupTerminalCreditToken( + uint32_t invocationId, uint32_t groupCount) +{ + return AllToAllGroupToken(invocationId, groupCount, 0U); +} + inline bool AllToAllGroupCheckedAdd(size_t lhs, size_t rhs, size_t& result) { if (rhs > std::numeric_limits::max() - lhs) { diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 178dd8fb..bdeb7eae 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -315,6 +315,10 @@ void TestTokens() TileXR::Demo::AllToAllGroupToken(49U, 1U, 0U)); CHECK_EQ(TileXR::Demo::AllToAllGroupCreditToken(50U, 1U) > TileXR::Demo::AllToAllGroupCreditToken(49U, 1U), true); + CHECK_EQ(TileXR::Demo::AllToAllGroupTerminalCreditToken(49U, 16U), + TileXR::Demo::AllToAllGroupToken(49U, 16U, 0U)); + CHECK_EQ(TileXR::Demo::AllToAllGroupTerminalCreditToken(49U, 16U) > + TileXR::Demo::AllToAllGroupCreditToken(49U, 15U), true); } void TestIngressCreditPolicy() @@ -357,6 +361,44 @@ void TestIngressCreditPolicy() CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(15U), true); CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(16U), false); CHECK_EQ(TileXR::Demo::AllToAllGroupCreditOwner(31U), false); + + for (const int terminalRankSize : {16, 32, 256, 512}) { + const uint32_t terminalGroupCount = + TileXR::Demo::AllToAllGroupCount(terminalRankSize); + for (int rank = 0; rank < terminalRankSize; ++rank) { + uint32_t chainedCredits = 0U; + uint32_t terminalCredits = 0U; + for (uint32_t group = 0U; group < terminalGroupCount; ++group) { + for (uint32_t lane = 0U; + lane < TileXR::Demo::kAllToAllGroupWidth; ++lane) { + if (TileXR::Demo::AllToAllGroupNextCreditPeer( + rank, terminalRankSize, group, lane) >= 0) { + ++chainedCredits; + } + } + } + for (uint32_t lane = 0U; + lane < TileXR::Demo::kAllToAllGroupWidth; ++lane) { + const int32_t terminalPeer = + TileXR::Demo::AllToAllGroupTerminalCreditPeer( + rank, terminalRankSize, lane); + if (terminalPeer < 0) { + continue; + } + ++terminalCredits; + uint32_t senderLane = lane ^ + TileXR::Demo::kAllToAllGroupHalfWidth; + if (TileXR::Demo::AllToAllGroupPeer( + terminalPeer, terminalRankSize, 0U, senderLane) < 0) { + senderLane = lane; + } + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer( + terminalPeer, terminalRankSize, 0U, senderLane), rank); + } + CHECK_EQ(chainedCredits + terminalCredits, + static_cast(terminalRankSize - 1)); + } + } for (uint32_t lane = 0U; lane < TileXR::Demo::kAllToAllGroupSendCoreCount; ++lane) { const uint32_t primaryWorker = lane; @@ -661,19 +703,23 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "uint64_t creditOffset0, uint64_t creditOffset1"); CHECK_CONTAINS(kernel, "uint32_t ingressWindow"); CHECK_CONTAINS(kernel, "AllToAllGroupPublishNextCredit"); + CHECK_CONTAINS(kernel, "AllToAllGroupPublishTerminalCredits"); + CHECK_CONTAINS(kernel, "AllToAllGroupWaitTerminalCredit"); + CHECK_CONTAINS(kernel, "AllToAllGroupDeviceTerminalCreditToken"); + CHECK_CONTAINS(kernel, "AscendC::SyncAll()"); CHECK_CONTAINS(kernel, "AllToAllGroupCreditOwnerDevice(worker)"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT"); CHECK_CONTAINS(kernel, "kAllToAllGroupTraceCreditWait"); CHECK_CONTAINS(kernel, "TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION"); const size_t publishCreditBegin = kernel.find( - "__aicore__ inline void AllToAllGroupPublishNextCredit"); + "__aicore__ inline void AllToAllGroupPublishCredit"); const size_t publishCreditEnd = kernel.find( - "__aicore__ inline void AllToAllGroupRecordError", publishCreditBegin); + "__aicore__ inline void AllToAllGroupPublishNextCredit", publishCreditBegin); const std::string publishCredit = publishCreditBegin == std::string::npos ? std::string() : kernel.substr(publishCreditBegin, publishCreditEnd == std::string::npos ? std::string::npos : publishCreditEnd - publishCreditBegin); - CHECK_CONTAINS(publishCredit, "args->creditMems[nextPeer]"); + CHECK_CONTAINS(publishCredit, "args->creditMems[peer]"); CHECK_CONTAINS(publishCredit, "TILEXR_ALLTOALL_GROUP_CREDIT_WORDS"); CHECK_CONTAINS(publishCredit, "AscendC::HardEvent::S_MTE3"); CHECK_CONTAINS(publishCredit, "AscendC::HardEvent::MTE3_S"); @@ -705,7 +751,6 @@ void TestKernelStructure() CHECK_CONTAINS(launcher, "cfgInfo.schemMode = RT_SCHEM_MODE_NORMAL"); CHECK_NOT_CONTAINS(launcher, "<<<"); CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); - CHECK_NOT_CONTAINS(kernel, "SyncAll"); CHECK_NOT_CONTAINS(kernel, "elementsPerPeer) * lane /"); } From 0b31864bf7b2131861c0e5b887e4a082fea30006 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 16:12:16 +0800 Subject: [PATCH 152/163] fix(udma): avoid grouped terminal credit barrier --- .../tilexr_udma_alltoall_group_kernel.cpp | 98 ++++++++++++++++++- .../demo/tilexr_udma_alltoall_group_layout.h | 14 ++- ...test_tilexr_udma_alltoall_group_layout.cpp | 10 +- 3 files changed, 114 insertions(+), 8 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 1710aef6..6f23a7d0 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -29,6 +29,12 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_BYTES = constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES = TILEXR_ALLTOALL_GROUP_SEND_WORKERS * TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH * sizeof(uint64_t); +constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_COPYOUT_SLICES = 3U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_SLICES = + TILEXR_ALLTOALL_GROUP_MAX_COPYOUT_SLICES - 1U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_STRIDE = 512U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_OFFSET = + TILEXR_ALLTOALL_GROUP_ERROR_BYTES + TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES; constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_WORDS = TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE / sizeof(uint64_t); @@ -497,6 +503,37 @@ __aicore__ inline void AllToAllGroupPublishTerminalCredits( } } +__aicore__ inline __gm__ uint64_t* AllToAllGroupTerminalAssistSignal( + __gm__ int32_t* debug, uint32_t lane, uint32_t copySliceIndex) +{ + const uint32_t assistIndex = + lane * TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_SLICES + copySliceIndex - 1U; + return reinterpret_cast<__gm__ uint64_t*>( + reinterpret_cast<__gm__ uint8_t*>(debug) + + TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_OFFSET + + static_cast(assistIndex) * + TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_STRIDE); +} + +__aicore__ inline void AllToAllGroupPublishTerminalAssist( + __gm__ int32_t* debug, uint32_t lane, uint32_t copySliceIndex, + uint64_t terminalToken, AscendC::LocalTensor relayLocal) +{ + auto signal = AllToAllGroupTerminalAssistSignal( + debug, lane, copySliceIndex); + auto signalLocal = relayLocal.ReinterpretCast(); + signalLocal.SetValue(0, terminalToken); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::GlobalTensor signalGlobal; + signalGlobal.SetGlobalBuffer( + signal, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); + AscendC::DataCopy( + signalGlobal, signalLocal, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + __aicore__ inline void AllToAllGroupRecordError( __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, @@ -625,6 +662,43 @@ __aicore__ inline bool AllToAllGroupWaitTerminalCredit( return ok; } +__aicore__ inline bool AllToAllGroupWaitTerminalAssistCredits( + __gm__ int32_t* debug, uint32_t blockIdx, + uint32_t worker, uint32_t copyoutWorkers, uint32_t groupCount, + uint64_t terminalToken, + AscendC::LocalTensor relayLocal) +{ + if (!AllToAllGroupCreditOwnerDevice(worker) || copyoutWorkers < 32U) { + return true; + } + const uint32_t copySliceCount = + copyoutWorkers / TILEXR_ALLTOALL_GROUP_SEND_CORES; + for (uint32_t assignment = 0U; ; ++assignment) { + const int32_t laneValue = AllToAllGroupCopyoutLaneDevice( + worker, assignment, copyoutWorkers); + if (laneValue < 0) { + return true; + } + const uint32_t lane = static_cast(laneValue); + for (uint32_t copySliceIndex = 1U; + copySliceIndex < copySliceCount; ++copySliceIndex) { + uint64_t observedToken = 0ULL; + if (!AllToAllGroupWaitCreditMte( + AllToAllGroupTerminalAssistSignal( + debug, lane, copySliceIndex), + terminalToken, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, + relayLocal, observedToken)) { + AllToAllGroupRecordError( + debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT, + groupCount, 0U, + static_cast(lane), copySliceIndex, 0U, + terminalToken, observedToken); + return false; + } + } + } +} + struct AllToAllGroupPendingQuiet { int32_t peer; uint32_t qpIdx; @@ -1096,10 +1170,25 @@ __aicore__ inline void AllToAllGroupKernelImpl( } } if constexpr (IngressCredit) { - AscendC::SyncAll(); - AllToAllGroupPublishTerminalCredits( - args, rank, rankSize, invocationId, worker, copyoutWorkers, - groupCount, groupWidth, creditOffsets[slot], relayLocal); + const uint64_t terminalToken = + AllToAllGroupDeviceTerminalCreditToken(invocationId, groupCount); + if (AllToAllGroupRemoteAssistDevice(worker, copyoutWorkers)) { + AllToAllGroupPublishTerminalAssist( + debug, worker % TILEXR_ALLTOALL_GROUP_SEND_CORES, + worker / TILEXR_ALLTOALL_GROUP_SEND_CORES, + terminalToken, relayLocal); + } else if (AllToAllGroupWaitTerminalAssistCredits( + debug, blockIdx, worker, copyoutWorkers, groupCount, + terminalToken, relayLocal)) { + AllToAllGroupPublishTerminalCredits( + args, rank, rankSize, invocationId, worker, copyoutWorkers, + groupCount, groupWidth, creditOffsets[slot], relayLocal); + } else { + AllToAllGroupTraceRecordKernel( + groupTrace, traceIteration, blockIdx, kernelBegin, + AllToAllGroupTraceCycle(groupTrace)); + return; + } } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -1277,7 +1366,6 @@ __aicore__ inline void AllToAllGroupKernelImpl( return; } if constexpr (IngressCredit) { - AscendC::SyncAll(); if (workerRoute == 0U && !AllToAllGroupWaitTerminalCredit( args, rank, rankSize, invocationId, lane, groupCount, passCount, groupWidth, creditOffsets[slot], relayLocal, diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index a1cf52d4..a45a6076 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -40,12 +40,20 @@ constexpr size_t kAllToAllGroupSignalSourceSlots = kAllToAllGroupMaxQuietBatch; constexpr size_t kAllToAllGroupSignalSourceBytes = kAllToAllGroupSignalSourceSlots * sizeof(uint64_t); +constexpr uint32_t kAllToAllGroupMaxCopyoutSlices = 3U; +constexpr uint32_t kAllToAllGroupTerminalAssistSlices = + kAllToAllGroupMaxCopyoutSlices - 1U; +constexpr size_t kAllToAllGroupTerminalAssistStride = 512U; +constexpr size_t kAllToAllGroupTerminalAssistBytes = + static_cast(kAllToAllGroupSendCoreCount) * + kAllToAllGroupTerminalAssistSlices * kAllToAllGroupTerminalAssistStride; constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupBaseControlBytes = - kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes; + kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes + + kAllToAllGroupTerminalAssistBytes; constexpr size_t kAllToAllGroupMaxPayloadBytes = 16ULL << 30; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 2U * kAllToAllGroupMaxPayloadBytes + @@ -70,6 +78,8 @@ struct AllToAllGroupPlan { size_t controlBytes = kAllToAllGroupBaseControlBytes; size_t signalSourceOffset = 0; size_t signalSourceBytes = kAllToAllGroupSignalSourceBytes; + size_t terminalAssistOffset = 0; + size_t terminalAssistBytes = kAllToAllGroupTerminalAssistBytes; size_t registeredBytes = 0; }; @@ -331,6 +341,8 @@ inline AllToAllGroupPlan PlanAllToAllGroup( return AllToAllGroupPlan {}; } plan.signalSourceOffset = plan.controlOffset + kAllToAllGroupErrorBytes; + plan.terminalAssistOffset = + plan.signalSourceOffset + plan.signalSourceBytes; if (plan.registeredBytes > kAllToAllGroupMaxRegisteredBytes) { return AllToAllGroupPlan {}; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index bdeb7eae..48b86f13 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -150,7 +150,11 @@ void TestPlan() plan.controlOffset + TileXR::Demo::kAllToAllGroupErrorBytes); CHECK_EQ(plan.signalSourceBytes, TileXR::Demo::kAllToAllGroupSignalSourceBytes); - CHECK_EQ(plan.signalSourceOffset + plan.signalSourceBytes <= + CHECK_EQ(plan.terminalAssistOffset, + plan.signalSourceOffset + plan.signalSourceBytes); + CHECK_EQ(plan.terminalAssistBytes, + TileXR::Demo::kAllToAllGroupTerminalAssistBytes); + CHECK_EQ(plan.terminalAssistOffset + plan.terminalAssistBytes <= plan.controlOffset + plan.controlBytes, true); const auto legacyAlign = [](size_t value) { return (value + TileXR::Demo::kAllToAllGroupAlignment - 1U) & @@ -704,9 +708,10 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "uint32_t ingressWindow"); CHECK_CONTAINS(kernel, "AllToAllGroupPublishNextCredit"); CHECK_CONTAINS(kernel, "AllToAllGroupPublishTerminalCredits"); + CHECK_CONTAINS(kernel, "AllToAllGroupPublishTerminalAssist"); + CHECK_CONTAINS(kernel, "AllToAllGroupWaitTerminalAssistCredits"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTerminalCredit"); CHECK_CONTAINS(kernel, "AllToAllGroupDeviceTerminalCreditToken"); - CHECK_CONTAINS(kernel, "AscendC::SyncAll()"); CHECK_CONTAINS(kernel, "AllToAllGroupCreditOwnerDevice(worker)"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT"); CHECK_CONTAINS(kernel, "kAllToAllGroupTraceCreditWait"); @@ -751,6 +756,7 @@ void TestKernelStructure() CHECK_CONTAINS(launcher, "cfgInfo.schemMode = RT_SCHEM_MODE_NORMAL"); CHECK_NOT_CONTAINS(launcher, "<<<"); CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); + CHECK_NOT_CONTAINS(kernel, "SyncAll"); CHECK_NOT_CONTAINS(kernel, "elementsPerPeer) * lane /"); } From beeb7ab0a000e2c401f9aeac3c788541c9c1cef3 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 16:20:25 +0800 Subject: [PATCH 153/163] Revert "fix(udma): avoid grouped terminal credit barrier" This reverts commit 0b31864bf7b2131861c0e5b887e4a082fea30006. --- .../tilexr_udma_alltoall_group_kernel.cpp | 98 +------------------ .../demo/tilexr_udma_alltoall_group_layout.h | 14 +-- ...test_tilexr_udma_alltoall_group_layout.cpp | 10 +- 3 files changed, 8 insertions(+), 114 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 6f23a7d0..1710aef6 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -29,12 +29,6 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_ERROR_BYTES = constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES = TILEXR_ALLTOALL_GROUP_SEND_WORKERS * TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH * sizeof(uint64_t); -constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_COPYOUT_SLICES = 3U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_SLICES = - TILEXR_ALLTOALL_GROUP_MAX_COPYOUT_SLICES - 1U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_STRIDE = 512U; -constexpr uint32_t TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_OFFSET = - TILEXR_ALLTOALL_GROUP_ERROR_BYTES + TILEXR_ALLTOALL_GROUP_SIGNAL_SOURCE_BYTES; constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_CREDIT_WORDS = TILEXR_ALLTOALL_GROUP_CREDIT_STRIDE / sizeof(uint64_t); @@ -503,37 +497,6 @@ __aicore__ inline void AllToAllGroupPublishTerminalCredits( } } -__aicore__ inline __gm__ uint64_t* AllToAllGroupTerminalAssistSignal( - __gm__ int32_t* debug, uint32_t lane, uint32_t copySliceIndex) -{ - const uint32_t assistIndex = - lane * TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_SLICES + copySliceIndex - 1U; - return reinterpret_cast<__gm__ uint64_t*>( - reinterpret_cast<__gm__ uint8_t*>(debug) + - TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_OFFSET + - static_cast(assistIndex) * - TILEXR_ALLTOALL_GROUP_TERMINAL_ASSIST_STRIDE); -} - -__aicore__ inline void AllToAllGroupPublishTerminalAssist( - __gm__ int32_t* debug, uint32_t lane, uint32_t copySliceIndex, - uint64_t terminalToken, AscendC::LocalTensor relayLocal) -{ - auto signal = AllToAllGroupTerminalAssistSignal( - debug, lane, copySliceIndex); - auto signalLocal = relayLocal.ReinterpretCast(); - signalLocal.SetValue(0, terminalToken); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); - AscendC::GlobalTensor signalGlobal; - signalGlobal.SetGlobalBuffer( - signal, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); - AscendC::DataCopy( - signalGlobal, signalLocal, TILEXR_ALLTOALL_GROUP_CREDIT_WORDS); - AscendC::SetFlag(EVENT_ID0); - AscendC::WaitFlag(EVENT_ID0); -} - __aicore__ inline void AllToAllGroupRecordError( __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, @@ -662,43 +625,6 @@ __aicore__ inline bool AllToAllGroupWaitTerminalCredit( return ok; } -__aicore__ inline bool AllToAllGroupWaitTerminalAssistCredits( - __gm__ int32_t* debug, uint32_t blockIdx, - uint32_t worker, uint32_t copyoutWorkers, uint32_t groupCount, - uint64_t terminalToken, - AscendC::LocalTensor relayLocal) -{ - if (!AllToAllGroupCreditOwnerDevice(worker) || copyoutWorkers < 32U) { - return true; - } - const uint32_t copySliceCount = - copyoutWorkers / TILEXR_ALLTOALL_GROUP_SEND_CORES; - for (uint32_t assignment = 0U; ; ++assignment) { - const int32_t laneValue = AllToAllGroupCopyoutLaneDevice( - worker, assignment, copyoutWorkers); - if (laneValue < 0) { - return true; - } - const uint32_t lane = static_cast(laneValue); - for (uint32_t copySliceIndex = 1U; - copySliceIndex < copySliceCount; ++copySliceIndex) { - uint64_t observedToken = 0ULL; - if (!AllToAllGroupWaitCreditMte( - AllToAllGroupTerminalAssistSignal( - debug, lane, copySliceIndex), - terminalToken, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, - relayLocal, observedToken)) { - AllToAllGroupRecordError( - debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT, - groupCount, 0U, - static_cast(lane), copySliceIndex, 0U, - terminalToken, observedToken); - return false; - } - } - } -} - struct AllToAllGroupPendingQuiet { int32_t peer; uint32_t qpIdx; @@ -1170,25 +1096,10 @@ __aicore__ inline void AllToAllGroupKernelImpl( } } if constexpr (IngressCredit) { - const uint64_t terminalToken = - AllToAllGroupDeviceTerminalCreditToken(invocationId, groupCount); - if (AllToAllGroupRemoteAssistDevice(worker, copyoutWorkers)) { - AllToAllGroupPublishTerminalAssist( - debug, worker % TILEXR_ALLTOALL_GROUP_SEND_CORES, - worker / TILEXR_ALLTOALL_GROUP_SEND_CORES, - terminalToken, relayLocal); - } else if (AllToAllGroupWaitTerminalAssistCredits( - debug, blockIdx, worker, copyoutWorkers, groupCount, - terminalToken, relayLocal)) { - AllToAllGroupPublishTerminalCredits( - args, rank, rankSize, invocationId, worker, copyoutWorkers, - groupCount, groupWidth, creditOffsets[slot], relayLocal); - } else { - AllToAllGroupTraceRecordKernel( - groupTrace, traceIteration, blockIdx, kernelBegin, - AllToAllGroupTraceCycle(groupTrace)); - return; - } + AscendC::SyncAll(); + AllToAllGroupPublishTerminalCredits( + args, rank, rankSize, invocationId, worker, copyoutWorkers, + groupCount, groupWidth, creditOffsets[slot], relayLocal); } AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, kernelBegin, AllToAllGroupTraceCycle(groupTrace)); @@ -1366,6 +1277,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( return; } if constexpr (IngressCredit) { + AscendC::SyncAll(); if (workerRoute == 0U && !AllToAllGroupWaitTerminalCredit( args, rank, rankSize, invocationId, lane, groupCount, passCount, groupWidth, creditOffsets[slot], relayLocal, diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index a45a6076..a1cf52d4 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -40,20 +40,12 @@ constexpr size_t kAllToAllGroupSignalSourceSlots = kAllToAllGroupMaxQuietBatch; constexpr size_t kAllToAllGroupSignalSourceBytes = kAllToAllGroupSignalSourceSlots * sizeof(uint64_t); -constexpr uint32_t kAllToAllGroupMaxCopyoutSlices = 3U; -constexpr uint32_t kAllToAllGroupTerminalAssistSlices = - kAllToAllGroupMaxCopyoutSlices - 1U; -constexpr size_t kAllToAllGroupTerminalAssistStride = 512U; -constexpr size_t kAllToAllGroupTerminalAssistBytes = - static_cast(kAllToAllGroupSendCoreCount) * - kAllToAllGroupTerminalAssistSlices * kAllToAllGroupTerminalAssistStride; constexpr uint32_t kAllToAllGroupBlockDim = 64U; constexpr size_t kAllToAllGroupMultiChannelThresholdBytes = 150ULL * 1024ULL * 1024ULL; constexpr size_t kAllToAllGroupAlignment = 512U; constexpr size_t kAllToAllGroupBaseControlBytes = - kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes + - kAllToAllGroupTerminalAssistBytes; + kAllToAllGroupErrorBytes + kAllToAllGroupSignalSourceBytes; constexpr size_t kAllToAllGroupMaxPayloadBytes = 16ULL << 30; constexpr size_t kAllToAllGroupMaxRegisteredBytes = 2U * kAllToAllGroupMaxPayloadBytes + @@ -78,8 +70,6 @@ struct AllToAllGroupPlan { size_t controlBytes = kAllToAllGroupBaseControlBytes; size_t signalSourceOffset = 0; size_t signalSourceBytes = kAllToAllGroupSignalSourceBytes; - size_t terminalAssistOffset = 0; - size_t terminalAssistBytes = kAllToAllGroupTerminalAssistBytes; size_t registeredBytes = 0; }; @@ -341,8 +331,6 @@ inline AllToAllGroupPlan PlanAllToAllGroup( return AllToAllGroupPlan {}; } plan.signalSourceOffset = plan.controlOffset + kAllToAllGroupErrorBytes; - plan.terminalAssistOffset = - plan.signalSourceOffset + plan.signalSourceBytes; if (plan.registeredBytes > kAllToAllGroupMaxRegisteredBytes) { return AllToAllGroupPlan {}; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 48b86f13..bdeb7eae 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -150,11 +150,7 @@ void TestPlan() plan.controlOffset + TileXR::Demo::kAllToAllGroupErrorBytes); CHECK_EQ(plan.signalSourceBytes, TileXR::Demo::kAllToAllGroupSignalSourceBytes); - CHECK_EQ(plan.terminalAssistOffset, - plan.signalSourceOffset + plan.signalSourceBytes); - CHECK_EQ(plan.terminalAssistBytes, - TileXR::Demo::kAllToAllGroupTerminalAssistBytes); - CHECK_EQ(plan.terminalAssistOffset + plan.terminalAssistBytes <= + CHECK_EQ(plan.signalSourceOffset + plan.signalSourceBytes <= plan.controlOffset + plan.controlBytes, true); const auto legacyAlign = [](size_t value) { return (value + TileXR::Demo::kAllToAllGroupAlignment - 1U) & @@ -708,10 +704,9 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "uint32_t ingressWindow"); CHECK_CONTAINS(kernel, "AllToAllGroupPublishNextCredit"); CHECK_CONTAINS(kernel, "AllToAllGroupPublishTerminalCredits"); - CHECK_CONTAINS(kernel, "AllToAllGroupPublishTerminalAssist"); - CHECK_CONTAINS(kernel, "AllToAllGroupWaitTerminalAssistCredits"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitTerminalCredit"); CHECK_CONTAINS(kernel, "AllToAllGroupDeviceTerminalCreditToken"); + CHECK_CONTAINS(kernel, "AscendC::SyncAll()"); CHECK_CONTAINS(kernel, "AllToAllGroupCreditOwnerDevice(worker)"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT"); CHECK_CONTAINS(kernel, "kAllToAllGroupTraceCreditWait"); @@ -756,7 +751,6 @@ void TestKernelStructure() CHECK_CONTAINS(launcher, "cfgInfo.schemMode = RT_SCHEM_MODE_NORMAL"); CHECK_NOT_CONTAINS(launcher, "<<<"); CHECK_NOT_CONTAINS(kernel, "UDMAPutSignalNbi"); - CHECK_NOT_CONTAINS(kernel, "SyncAll"); CHECK_NOT_CONTAINS(kernel, "elementsPerPeer) * lane /"); } From 7ab4b105c93f94a87b38de8b8cd0ea62aa616dc8 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 17:17:41 +0800 Subject: [PATCH 154/163] perf(udma): schedule ready grouped copyouts --- .../tilexr_udma_alltoall_group_kernel.cpp | 96 ++++++++++++++++++- ...test_tilexr_udma_alltoall_group_layout.cpp | 6 ++ 2 files changed, 99 insertions(+), 3 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 1710aef6..afcb1e5a 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -18,6 +18,10 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH = 16U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_EXPERIMENTAL_WIDTH = 4U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH = 64U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_MAX_GROUP_COUNT = 64U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_READY_BITMAP_WORDS = + (TileXR::TILEXR_MAX_RANK_SIZE + 63U) / 64U; +static_assert(TILEXR_ALLTOALL_GROUP_READY_BITMAP_WORDS * 64U >= + TileXR::TILEXR_MAX_RANK_SIZE, "ready bitmap must cover every peer task"); constexpr uint32_t TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE = 512U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SIGNAL_STRIDE = 1024U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_RELAY_BYTES = 64U * 1024U; @@ -413,6 +417,20 @@ __aicore__ inline bool AllToAllGroupWaitRouteTokensMte( return true; } +__aicore__ inline bool AllToAllGroupRouteTokensReadyMte( + __gm__ uint64_t* primarySignal, __gm__ uint64_t* secondarySignal, + bool waitPrimary, bool waitSecondary, uint64_t expectedToken, + AscendC::LocalTensor relayLocal, uint64_t& observed) +{ + const uint64_t primaryObserved = waitPrimary ? + AllToAllGroupLoadTokenMte(primarySignal, relayLocal) : expectedToken; + const uint64_t secondaryObserved = waitSecondary ? + AllToAllGroupLoadTokenMte(secondarySignal, relayLocal) : expectedToken; + observed = primaryObserved < secondaryObserved ? + primaryObserved : secondaryObserved; + return observed >= expectedToken; +} + __aicore__ inline void AllToAllGroupRecordError( __gm__ int32_t* debug, uint32_t blockIdx, uint32_t stage, uint32_t group, uint32_t pass, int32_t peer, uint32_t qpIdx, @@ -856,6 +874,28 @@ __aicore__ inline void AllToAllGroupKernelImpl( selfCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } + const bool readyDrivenCopyout = copyoutWorkers == 1U && + TileXR::SDMAEnabled(args) && + routeStage == TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_COMBINED; + const uint32_t schedulerPassCount = readyDrivenCopyout ? passCount : 1U; + for (uint32_t schedulerPass = 0U; + schedulerPass < schedulerPassCount; ++schedulerPass) { + uint64_t completedTasks[TILEXR_ALLTOALL_GROUP_READY_BITMAP_WORDS]; + for (uint32_t word = 0U; + word < TILEXR_ALLTOALL_GROUP_READY_BITMAP_WORDS; ++word) { + completedTasks[word] = 0ULL; + } + uint64_t noProgressBegin = + static_cast(AscendC::GetSystemCycle()); + bool schedulerDone = false; + while (!schedulerDone) { + bool pendingTask = false; + bool madeProgress = false; + uint32_t lastPendingGroup = 0U; + uint32_t lastPendingPass = schedulerPass; + int32_t lastPendingPeer = -1; + uint64_t lastExpectedToken = 0ULL; + uint64_t lastObservedToken = 0ULL; for (uint32_t group = 0U; group < groupCount; ++group) { for (uint32_t assignment = 0U; ; ++assignment) { const int32_t laneValue = AllToAllGroupCopyoutLaneDevice( @@ -864,6 +904,12 @@ __aicore__ inline void AllToAllGroupKernelImpl( break; } const uint32_t lane = static_cast(laneValue); + const uint32_t taskIndex = group * groupWidth + lane; + const uint32_t taskWord = taskIndex / 64U; + const uint64_t taskMask = 1ULL << (taskIndex % 64U); + if (readyDrivenCopyout && (completedTasks[taskWord] & taskMask) != 0ULL) { + continue; + } const int32_t peer = AllToAllGroupDevicePeer( rank, rankSize, group, lane, groupWidth); if (peer < 0) { @@ -897,7 +943,9 @@ __aicore__ inline void AllToAllGroupKernelImpl( primaryTotalElements, secondaryTotalElements); const uint32_t traceCore = copyoutWorkers < TILEXR_ALLTOALL_GROUP_SEND_CORES ? TILEXR_ALLTOALL_GROUP_SEND_WORKERS + lane : blockIdx; - for (uint32_t pass = 0U; pass < passCount; ++pass) { + const uint32_t passBegin = readyDrivenCopyout ? schedulerPass : 0U; + const uint32_t passEnd = readyDrivenCopyout ? schedulerPass + 1U : passCount; + for (uint32_t pass = passBegin; pass < passEnd; ++pass) { const int32_t chunkElementOffset = static_cast(pass) * chunkElements; const int32_t remaining = elementsPerPeer - chunkElementOffset; if (remaining <= 0) { @@ -934,6 +982,9 @@ __aicore__ inline void AllToAllGroupKernelImpl( secondaryElements == 0U)) { continue; } + if (readyDrivenCopyout) { + pendingTask = true; + } const uint64_t expectedToken = AllToAllGroupDeviceToken(invocationId, group, pass); auto primarySignal = reinterpret_cast<__gm__ uint64_t*>( @@ -949,10 +1000,23 @@ __aicore__ inline void AllToAllGroupKernelImpl( routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_SECONDARY; const bool waitSecondary = secondaryElements != 0U && routeStage != TILEXR_ALLTOALL_GROUP_ROUTE_STAGE_PRIMARY; - if (!AllToAllGroupWaitRouteTokensMte( + const bool ready = readyDrivenCopyout ? + AllToAllGroupRouteTokensReadyMte( + primarySignal, secondarySignal, waitPrimary, waitSecondary, + expectedToken, relayLocal, observed) : + AllToAllGroupWaitRouteTokensMte( primarySignal, secondarySignal, waitPrimary, waitSecondary, expectedToken, TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES, - relayLocal, observed)) { + relayLocal, observed); + if (!ready && readyDrivenCopyout) { + lastPendingGroup = group; + lastPendingPass = pass; + lastPendingPeer = peer; + lastExpectedToken = expectedToken; + lastObservedToken = observed; + continue; + } + if (!ready) { AllToAllGroupTraceRecordTask( groupTrace, traceIteration, traceCore, group, pass, TileXR::Demo::kAllToAllGroupTraceReceiveWait, groupCount, passCount, @@ -1092,8 +1156,34 @@ __aicore__ inline void AllToAllGroupKernelImpl( TileXR::Demo::kAllToAllGroupTraceNoQp, receiveCopyBegin, AllToAllGroupTraceCycle(groupTrace)); } + if (readyDrivenCopyout) { + completedTasks[taskWord] |= taskMask; + madeProgress = true; + } } } + } + if (!readyDrivenCopyout || !pendingTask) { + schedulerDone = true; + continue; + } + if (madeProgress) { + noProgressBegin = + static_cast(AscendC::GetSystemCycle()); + continue; + } + if (static_cast(AscendC::GetSystemCycle()) - + noProgressBegin >= TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES) { + AllToAllGroupRecordError( + debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_WAIT, + lastPendingGroup, lastPendingPass, lastPendingPeer, + 0U, 0U, lastExpectedToken, lastObservedToken); + AllToAllGroupTraceRecordKernel( + groupTrace, traceIteration, blockIdx, kernelBegin, + AllToAllGroupTraceCycle(groupTrace)); + return; + } + } } if constexpr (IngressCredit) { AscendC::SyncAll(); diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index bdeb7eae..6a66e4b8 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -639,6 +639,12 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "AllToAllGroupSplitByRouteDevice"); CHECK_CONTAINS(kernel, "TILEXR_ALLTOALL_GROUP_ROUTE_SIGNAL_STRIDE"); CHECK_CONTAINS(kernel, "AllToAllGroupWaitRouteTokensMte"); + CHECK_CONTAINS(kernel, "AllToAllGroupRouteTokensReadyMte"); + CHECK_CONTAINS(kernel, "readyDrivenCopyout"); + CHECK_CONTAINS(kernel, "taskIndex = group * groupWidth + lane"); + CHECK_CONTAINS(kernel, "completedTasks[taskWord] |= taskMask"); + CHECK_CONTAINS(kernel, "while (!schedulerDone)"); + CHECK_CONTAINS(kernel, "noProgressBegin"); CHECK_CONTAINS(kernel, "secondaryQp"); CHECK_CONTAINS(kernel, "selectedQp"); CHECK_CONTAINS(kernel, "copyoutWorkers"); From a36e91de11a5068adcf7696a416a95962f2f3638 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Tue, 4 Aug 2026 19:59:22 +0800 Subject: [PATCH 155/163] perf(udma): prewarm grouped SDMA SQ pages --- src/include/tilexr_sdma.h | 16 ++++++ src/include/tilexr_sdma_a5.h | 56 +++++++++++++++++++ src/include/tilexr_sdma_a5_types.h | 24 ++++++++ .../unit/test_tilexr_sdma_a5_validation.cpp | 12 ++++ .../tilexr_udma_alltoall_group_kernel.cpp | 32 ++++++++--- .../tilexr_udma_alltoall_group_launcher.cpp | 11 ++-- tests/udma/demo/tilexr_udma_demo.cpp | 27 ++++++++- ...test_tilexr_udma_alltoall_group_layout.cpp | 3 +- 8 files changed, 164 insertions(+), 17 deletions(-) diff --git a/src/include/tilexr_sdma.h b/src/include/tilexr_sdma.h index 4733a0a3..de58c1f0 100644 --- a/src/include/tilexr_sdma.h +++ b/src/include/tilexr_sdma.h @@ -51,6 +51,22 @@ __aicore__ inline uint32_t SDMAResolveChannelGroup(uint32_t channelGroupIdx) return channelGroupIdx; } +__aicore__ inline bool SDMAPrewarmSqPages( + const __gm__ CommArgs* args, uint32_t channelCount, + AscendC::LocalTensor scratch) +{ +#if TILEXR_SDMA_A5_AICORE_COMPILE + return SDMAEnabled(args) && detail::A5SdmaPrewarmSqPages( + reinterpret_cast<__gm__ uint8_t*>(args->sdmaWorkspacePtr), + channelCount, scratch); +#else + (void)args; + (void)channelCount; + (void)scratch; + return false; +#endif +} + __aicore__ inline uint64_t SDMACopyNbi( const __gm__ CommArgs* args, __gm__ uint8_t* dst, diff --git a/src/include/tilexr_sdma_a5.h b/src/include/tilexr_sdma_a5.h index 6e96aa41..fa437f2d 100644 --- a/src/include/tilexr_sdma_a5.h +++ b/src/include/tilexr_sdma_a5.h @@ -87,6 +87,62 @@ __aicore__ inline bool A5SdmaWorkspaceValid(const __gm__ A5SdmaWorkspace* worksp workspace->header.workspaceSize == sizeof(A5SdmaWorkspace); } +__aicore__ inline bool A5SdmaPrewarmSqPages( + __gm__ uint8_t* workspaceAddress, uint32_t channelCount, + AscendC::LocalTensor scratch) +{ + __gm__ A5SdmaWorkspace* workspace = + reinterpret_cast<__gm__ A5SdmaWorkspace*>(workspaceAddress); + if (!A5SdmaWorkspaceValid(workspace) || channelCount == 0U || + channelCount > TILEXR_SDMA_A5_CHANNEL_COUNT) { + return false; + } + + for (uint32_t channelIndex = 0U; channelIndex < channelCount; ++channelIndex) { + __gm__ A5SdmaChannel* channel = &workspace->channels[channelIndex]; + const uint32_t originalHead = channel->head; + const uint32_t originalTail = channel->tail; + const uint32_t originalGeneration = channel->generation; + const uint32_t originalTaskId = channel->taskId; + if (channel->sqBase == 0U || channel->depth == 0U || + originalHead != originalTail || channel->outstanding != 0U) { + return false; + } + + const uint32_t pageCount = A5SdmaSqPageCount(channel->depth); + for (uint32_t page = 0U; page < pageCount; ++page) { + const uint64_t offset = A5SdmaSqPageOffset(page); + if (offset + TILEXR_SDMA_A5_PREWARM_BYTES > + A5SdmaSqBytes(channel->depth)) { + return false; + } + AscendC::GlobalTensor pageGlobal; + pageGlobal.SetGlobalBuffer( + reinterpret_cast<__gm__ uint8_t*>(channel->sqBase) + offset); + AscendC::DataCopyPadExtParams padIn {false, 0U, 0U, 0}; + AscendC::DataCopyExtParams copyIn { + 1U, TILEXR_SDMA_A5_PREWARM_BYTES, 0U, 0U, 0U}; + AscendC::DataCopyPad(scratch, pageGlobal, copyIn, padIn); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::DataCopyExtParams copyOut { + 1U, TILEXR_SDMA_A5_PREWARM_BYTES, 0U, 0U, 0U}; + AscendC::DataCopyPad(pageGlobal, scratch, copyOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::PipeBarrier(); + dsb(DSB_DDR); + if (channel->head != originalHead || channel->tail != originalTail || + channel->generation != originalGeneration || + channel->taskId != originalTaskId || channel->outstanding != 0U) { + return false; + } + } + return true; +} + __aicore__ inline uint64_t A5SdmaCopyNbi(__gm__ uint8_t* workspaceAddress, __gm__ uint8_t* destination, __gm__ uint8_t* source, diff --git a/src/include/tilexr_sdma_a5_types.h b/src/include/tilexr_sdma_a5_types.h index 89098112..edbc3cbd 100644 --- a/src/include/tilexr_sdma_a5_types.h +++ b/src/include/tilexr_sdma_a5_types.h @@ -24,6 +24,8 @@ constexpr uint16_t TILEXR_SDMA_A5_BACKEND_KIND = 2U; constexpr uint32_t TILEXR_SDMA_A5_CHANNEL_COUNT = 48U; constexpr uint32_t TILEXR_SDMA_A5_SQE_BYTES = 64U; constexpr uint32_t TILEXR_SDMA_A5_COMPLETION_BYTES = 64U; +constexpr uint32_t TILEXR_SDMA_A5_SQ_PAGE_BYTES = 4096U; +constexpr uint32_t TILEXR_SDMA_A5_PREWARM_BYTES = 64U; constexpr uint64_t TILEXR_SDMA_A5_MAX_TRANSFER_BYTES = 0xFFFFFFFFULL; constexpr uint32_t TILEXR_SDMA_A5_SQE_TYPE = 11U; constexpr uint32_t TILEXR_SDMA_A5_KERNEL_CREDIT = 254U; @@ -125,6 +127,24 @@ TILEXR_SDMA_A5_HOST_DEVICE_INLINE bool A5SdmaQueueStateValid(uint32_t tail, uint return depth >= 3U && tail < depth; } +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint64_t A5SdmaSqBytes(uint32_t depth) +{ + return static_cast(depth) * TILEXR_SDMA_A5_SQE_BYTES; +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint32_t A5SdmaSqPageCount(uint32_t depth) +{ + const uint64_t bytes = A5SdmaSqBytes(depth); + return static_cast( + (bytes + TILEXR_SDMA_A5_SQ_PAGE_BYTES - 1U) / + TILEXR_SDMA_A5_SQ_PAGE_BYTES); +} + +TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint64_t A5SdmaSqPageOffset(uint32_t page) +{ + return static_cast(page) * TILEXR_SDMA_A5_SQ_PAGE_BYTES; +} + TILEXR_SDMA_A5_HOST_DEVICE_INLINE uint32_t A5SdmaAdvanceTail(uint32_t tail, uint32_t depth) { return (tail + 2U) % depth; @@ -181,6 +201,10 @@ static_assert(sizeof(A5SdmaWorkspaceHeader) == 64U, "A5 workspace header must be static_assert(sizeof(A5SdmaChannel) == 192U, "A5 channel ABI must be 192 bytes"); static_assert(sizeof(A5SdmaCompletionLine) == 64U, "A5 completion line must be 64 bytes"); static_assert(sizeof(A5SdmaSqe) == TILEXR_SDMA_A5_SQE_BYTES, "A5 SQE must be 64 bytes"); +static_assert(TILEXR_SDMA_A5_SQ_PAGE_BYTES % TILEXR_SDMA_A5_SQE_BYTES == 0U, + "A5 SQ pages must contain whole SQEs"); +static_assert(TILEXR_SDMA_A5_PREWARM_BYTES == TILEXR_SDMA_A5_SQE_BYTES, + "A5 SQ prewarm must preserve one complete SQE"); static_assert(alignof(A5SdmaWorkspace) == 64U, "A5 workspace must be cache-line aligned"); static_assert(sizeof(A5SdmaWorkspace) % 64U == 0U, "A5 workspace size must be cache-line aligned"); static_assert(offsetof(A5SdmaChannel, sqBase) == 0U, "unexpected A5 SQ base offset"); diff --git a/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp b/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp index fd33e7b6..add28810 100644 --- a/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp +++ b/tests/sdma/unit/test_tilexr_sdma_a5_validation.cpp @@ -197,6 +197,17 @@ void TestGroupedCopyWaitIsBounded() CHECK_TRUE(TileXR::detail::TILEXR_SDMA_A5_WAIT_TIMEOUT_CYCLES != 0ULL); } +void TestSqPrewarmPageLayout() +{ + using namespace TileXR::detail; + CHECK_TRUE(A5SdmaSqBytes(2049U) == 131136ULL); + CHECK_TRUE(A5SdmaSqPageCount(2049U) == 33U); + CHECK_TRUE(A5SdmaSqPageOffset(0U) == 0ULL); + CHECK_TRUE(A5SdmaSqPageOffset(7U) == 28672ULL); + CHECK_TRUE(A5SdmaSqPageOffset(32U) + TILEXR_SDMA_A5_PREWARM_BYTES == + A5SdmaSqBytes(2049U)); +} + void TestCleanupFailuresRetainHandlesForRetry() { struct FailureCase { @@ -260,6 +271,7 @@ int main() TestPartialClassificationFailsClosed(); TestCompleteClassificationRequiresFinishedHeader(); TestGroupedCopyWaitIsBounded(); + TestSqPrewarmPageLayout(); TestCleanupFailuresRetainHandlesForRetry(); TestCleanupRestoreFailureIsRetryable(); if (g_failures != 0) { diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index afcb1e5a..2210f020 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -43,6 +43,7 @@ constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_QUIET = 2U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_WAIT = 3U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_CREDIT_WAIT = 4U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_SDMA = 5U; +constexpr uint32_t TILEXR_ALLTOALL_GROUP_STAGE_SDMA_PREWARM = 6U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SDMA_FALLBACK = 0U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SDMA_COMPLETE = 1U; constexpr uint32_t TILEXR_ALLTOALL_GROUP_SDMA_FAILED = 2U; @@ -773,7 +774,7 @@ __aicore__ inline void AllToAllGroupKernelImpl( uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch) + uint32_t groupWidth, uint32_t quietBatch, uint32_t prewarmSq) { const uint32_t blockIdx = static_cast(AscendC::GetBlockIdx()); auto groupTrace = blockIdx < TileXR::Demo::kAllToAllGroupTraceCoreCount ? @@ -797,6 +798,8 @@ __aicore__ inline void AllToAllGroupKernelImpl( quietBatch == 0U || quietBatch > TILEXR_ALLTOALL_GROUP_MAX_QUIET_BATCH || (quietBatch & (quietBatch - 1U)) != 0U || + prewarmSq > 1U || + (prewarmSq != 0U && copyoutWorkers != 1U) || (primaryRouteParts > TileXR::Demo::kAllToAllGroupRouteParts && primaryRouteParts != TileXR::Demo::kAllToAllGroupAutoPrimaryParts) || TILEXR_ALLTOALL_GROUP_SEND_WORKERS + copyoutWorkers > @@ -850,6 +853,15 @@ __aicore__ inline void AllToAllGroupKernelImpl( return; } const uint32_t worker = blockIdx - TILEXR_ALLTOALL_GROUP_SEND_WORKERS; + if (prewarmSq != 0U && worker == 0U && + !TileXR::SDMAPrewarmSqPages(args, groupWidth, relayLocal)) { + AllToAllGroupRecordError( + debug, blockIdx, TILEXR_ALLTOALL_GROUP_STAGE_SDMA_PREWARM, + 0U, 0U, -1, 0U, 0U, groupWidth, 0ULL); + AllToAllGroupTraceRecordKernel(groupTrace, traceIteration, blockIdx, + kernelBegin, AllToAllGroupTraceCycle(groupTrace)); + return; + } const uint32_t selfCopyWorkers = copyoutWorkers >= 32U ? 16U : copyoutWorkers; const int32_t selfBegin = worker < selfCopyWorkers ? static_cast( static_cast(elementsPerPeer) * worker / selfCopyWorkers) : 0; @@ -1390,7 +1402,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch) + uint32_t groupWidth, uint32_t quietBatch, uint32_t prewarmSq) { AllToAllGroupKernelImpl( commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, @@ -1398,7 +1410,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_kernel( payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, 0ULL, 0ULL, groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, - groupWidth, quietBatch); + groupWidth, quietBatch, prewarmSq); } extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_kernel( @@ -1410,7 +1422,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_kernel( uint64_t signalOffset0, uint64_t signalOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch) + uint32_t groupWidth, uint32_t quietBatch, uint32_t prewarmSq) { AllToAllGroupKernelImpl( commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, @@ -1418,7 +1430,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_kernel( payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, 0ULL, 0ULL, groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, - groupWidth, quietBatch); + groupWidth, quietBatch, prewarmSq); } extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_credit_kernel( @@ -1431,7 +1443,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_credit_kernel uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow, + uint32_t prewarmSq) { AllToAllGroupKernelImpl( commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, @@ -1439,7 +1452,7 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_credit_kernel payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, creditOffset0, creditOffset1, groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, - groupWidth, quietBatch); + groupWidth, quietBatch, prewarmSq); } extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_credit_kernel( @@ -1452,7 +1465,8 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_credit_ uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTraceGM, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow, + uint32_t prewarmSq) { AllToAllGroupKernelImpl( commArgsGM, inputGM, outputGM, registeredMemoryGM, debugGM, invocationId, @@ -1460,5 +1474,5 @@ extern "C" __global__ __aicore__ void tilexr_udma_all_to_all_group_batch_credit_ payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, creditOffset0, creditOffset1, groupTraceGM, traceIteration, routeStage, multiChannel, primaryRouteParts, - groupWidth, quietBatch); + groupWidth, quietBatch, prewarmSq); } diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp index 0c72b247..4b9b84e0 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_launcher.cpp @@ -48,6 +48,7 @@ struct alignas(8) GroupedAllToAllKernelArgs { uint32_t primaryRouteParts; uint32_t groupWidth; uint32_t quietBatch; + uint32_t prewarmSq; }; struct alignas(8) GroupedAllToAllCreditKernelArgs { @@ -75,9 +76,10 @@ struct alignas(8) GroupedAllToAllCreditKernelArgs { uint32_t groupWidth; uint32_t quietBatch; uint32_t ingressWindow; + uint32_t prewarmSq; }; -static_assert(sizeof(GroupedAllToAllKernelArgs) == 128U, +static_assert(sizeof(GroupedAllToAllKernelArgs) == 136U, "grouped alltoall kernel argument ABI changed"); static_assert(offsetof(GroupedAllToAllKernelArgs, payloadOffset0) == 64U, "grouped alltoall payload offset ABI changed"); @@ -181,7 +183,8 @@ int launch_tilexr_udma_all_to_all_group( uint64_t creditOffset0, uint64_t creditOffset1, uint8_t* groupTrace, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow) + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow, + uint32_t prewarmSq) { const rtError_t registrationStatus = EnsureGroupedKernelRegistered(); if (registrationStatus != RT_ERROR_NONE) { @@ -200,7 +203,7 @@ int launch_tilexr_udma_all_to_all_group( payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, creditOffset0, creditOffset1, groupTrace, traceIteration, routeStage, multiChannel, primaryRouteParts, groupWidth, - quietBatch, ingressWindow, + quietBatch, ingressWindow, prewarmSq, }; functionSignature = useBatch ? GroupedBatchCreditKernelFunctionSignature() : GroupedCreditKernelFunctionSignature(); @@ -210,7 +213,7 @@ int launch_tilexr_udma_all_to_all_group( elementsPerPeer, chunkElements, passCount, groupCount, payloadOffset0, payloadOffset1, signalOffset0, signalOffset1, groupTrace, traceIteration, routeStage, multiChannel, - primaryRouteParts, groupWidth, quietBatch, + primaryRouteParts, groupWidth, quietBatch, prewarmSq, }; functionSignature = useBatch ? GroupedBatchKernelFunctionSignature() : GroupedKernelFunctionSignature(); diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 7ec62065..1824771f 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -65,7 +65,8 @@ extern int launch_tilexr_udma_all_to_all_group( uint64_t creditOffset0, uint64_t creditOffset1, GM_ADDR groupTrace, uint32_t traceIteration, uint32_t routeStage, uint32_t multiChannel, uint32_t primaryRouteParts, - uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow); + uint32_t groupWidth, uint32_t quietBatch, uint32_t ingressWindow, + uint32_t prewarmSq); extern void launch_tilexr_udma_all_to_all_bigdata( uint32_t blockDim, void* stream, GM_ADDR commArgs, GM_ADDR input, GM_ADDR output, GM_ADDR udmaMem, GM_ADDR debug, GM_ADDR fullmeshTrace, uint32_t fullmeshTraceIteration, @@ -1033,6 +1034,17 @@ bool RunGroupedAllToAll( } constexpr uint32_t sendWorkers = TileXR::Demo::kAllToAllGroupSendWorkerCount; const uint32_t copyoutWorkers = sdmaAvailable ? 1U : 32U; + const int prewarmSqValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_PREWARM_SQ", 0); + if ((prewarmSqValue != 0 && prewarmSqValue != 1) || + (prewarmSqValue != 0 && !sdmaAvailable)) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_PREWARM_SQ" + << " must be 0 or 1 and requires SDMA, got " + << prewarmSqValue << std::endl; + return false; + } + const uint32_t prewarmSq = static_cast(prewarmSqValue); const uint32_t groupBlockDim = TileXR::Demo::AllToAllGroupBlockDim( sendWorkers, copyoutWorkers); const int routeStagesValue = GetEnvInt( @@ -1247,6 +1259,7 @@ bool RunGroupedAllToAll( " useSecondaryRoute=" + std::to_string(useSecondaryRouteValue) + " quietBatch=" + std::to_string(quietBatch) + " ingressWindow=" + std::to_string(ingressWindow) + + " prewarmSq=" + std::to_string(prewarmSq) + " routeStages=" + std::to_string(routeStagesValue)); if (routeStages && @@ -1256,8 +1269,12 @@ bool RunGroupedAllToAll( } uint32_t invocationId = 0U; + bool prewarmSqPending = prewarmSq != 0U; auto launchGroupStage = [&](TileXR::Demo::AllToAllGroupRouteStage routeStage, - void* trace, uint32_t traceIteration) -> bool { + void* trace, uint32_t traceIteration) -> bool { + const uint32_t prewarmThisLaunch = + prewarmSqPending && + routeStage == TileXR::Demo::AllToAllGroupRouteStage::kCombined ? 1U : 0U; const int launchRet = launch_tilexr_udma_all_to_all_group( groupBlockDim, stream, commArgsDev, reinterpret_cast(input), reinterpret_cast(output), @@ -1272,13 +1289,17 @@ bool RunGroupedAllToAll( multiChannel ? 1U : 0U, primaryRouteParts, groupWidth, quietBatch, routeStage == TileXR::Demo::AllToAllGroupRouteStage::kCombined ? - ingressWindow : 0U); + ingressWindow : 0U, + prewarmThisLaunch); if (launchRet != 0) { std::cerr << "[rank " << rank << "] rtKernelLaunchWithFlagV2 grouped failed: " << launchRet << std::endl; return false; } + if (prewarmThisLaunch != 0U) { + prewarmSqPending = false; + } return true; }; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 6a66e4b8..fda5b900 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -746,9 +746,10 @@ void TestKernelStructure() CHECK_CONTAINS(launcher, "rtFunctionRegister"); CHECK_CONTAINS(launcher, "rtKernelLaunchWithFlagV2"); CHECK_CONTAINS(launcher, "GroupedAllToAllKernelArgs"); - CHECK_CONTAINS(launcher, "sizeof(GroupedAllToAllKernelArgs) == 128U"); + CHECK_CONTAINS(launcher, "sizeof(GroupedAllToAllKernelArgs) == 136U"); CHECK_CONTAINS(launcher, "GroupedAllToAllCreditKernelArgs"); CHECK_CONTAINS(launcher, "sizeof(GroupedAllToAllCreditKernelArgs) == 152U"); + CHECK_CONTAINS(launcher, "prewarmSq"); CHECK_CONTAINS(launcher, "TILEXR_GROUPED_ALLTOALL_BATCH_KERNEL_NAME"); CHECK_CONTAINS(launcher, "TILEXR_GROUPED_ALLTOALL_CREDIT_KERNEL_NAME"); CHECK_CONTAINS(launcher, "TILEXR_GROUPED_ALLTOALL_BATCH_CREDIT_KERNEL_NAME"); From abdd99c1f981861008470073b254aa315c6708ba Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 5 Aug 2026 20:54:25 +0800 Subject: [PATCH 156/163] perf(udma): balance grouped 128p peer schedule --- .../tilexr_udma_alltoall_group_kernel.cpp | 42 +++++++++++ .../demo/tilexr_udma_alltoall_group_layout.h | 40 +++++++++++ ...test_tilexr_udma_alltoall_group_layout.cpp | 70 +++++++++++++++++++ 3 files changed, 152 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 2210f020..7e1351ce 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -83,6 +83,48 @@ __aicore__ inline int32_t AllToAllGroupDevicePeer( lane >= groupWidth) { return -1; } + if (rankSize == 128 && groupWidth == TILEXR_ALLTOALL_GROUP_DEFAULT_WIDTH) { + constexpr int32_t partitionSize = 64; + const int32_t partitionBase = rank / partitionSize * partitionSize; + const int32_t oppositeBase = partitionSize - partitionBase; + const int32_t partitionRank = rank % partitionSize; + if (lane < TILEXR_ALLTOALL_GROUP_SEND_CORES / 2U) { + constexpr uint32_t sideWidth = + TILEXR_ALLTOALL_GROUP_SEND_CORES / 4U; + const uint32_t index = lane < sideWidth ? lane : lane - sideWidth; + const int32_t distance = static_cast( + group * sideWidth + index + 1U); + if (distance > partitionSize / 2 || + (lane >= sideWidth && distance == partitionSize / 2)) { + return -1; + } + const int32_t peer = lane < sideWidth ? + (partitionRank + distance) % partitionSize : + (partitionRank - distance + partitionSize) % partitionSize; + return partitionBase + peer; + } + + const uint32_t remoteLane = + lane - TILEXR_ALLTOALL_GROUP_SEND_CORES / 2U; + int32_t offset = 0; + if (group == 0U) { + // Zero and the diameter are self-inverse, so they share one group. + if (remoteLane < 4U) { + offset = static_cast(remoteLane); + } else if (remoteLane == 4U) { + offset = partitionSize / 2; + } else { + offset = -static_cast(remoteLane - 4U); + } + } else { + const uint32_t index = remoteLane < 4U ? + remoteLane : remoteLane - 4U; + const int32_t distance = static_cast(group * 4U + index); + offset = remoteLane < 4U ? distance : -distance; + } + return oppositeBase + + (partitionRank + offset + partitionSize) % partitionSize; + } const uint32_t halfWidth = groupWidth / 2U; const uint32_t index = lane < halfWidth ? lane : lane - halfWidth; const int32_t distance = static_cast( diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h index a1cf52d4..57c307cc 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_layout.h +++ b/tests/udma/demo/tilexr_udma_alltoall_group_layout.h @@ -181,6 +181,46 @@ inline int32_t AllToAllGroupPeer( group >= AllToAllGroupCount(rankSize, groupWidth)) { return -1; } + if (rankSize == 128 && groupWidth == kAllToAllGroupWidth) { + constexpr int32_t partitionSize = 64; + const int32_t partitionBase = rank / partitionSize * partitionSize; + const int32_t oppositeBase = partitionSize - partitionBase; + const int32_t partitionRank = rank % partitionSize; + if (lane < kAllToAllGroupHalfWidth) { + constexpr uint32_t sideWidth = kAllToAllGroupHalfWidth / 2U; + const uint32_t index = lane < sideWidth ? lane : lane - sideWidth; + const int32_t distance = static_cast( + group * sideWidth + index + 1U); + if (distance > partitionSize / 2 || + (lane >= sideWidth && distance == partitionSize / 2)) { + return -1; + } + const int32_t peer = lane < sideWidth ? + (partitionRank + distance) % partitionSize : + (partitionRank - distance + partitionSize) % partitionSize; + return partitionBase + peer; + } + + const uint32_t remoteLane = lane - kAllToAllGroupHalfWidth; + int32_t offset = 0; + if (group == 0U) { + // Zero and the diameter are self-inverse, so they share one group. + if (remoteLane < 4U) { + offset = static_cast(remoteLane); + } else if (remoteLane == 4U) { + offset = partitionSize / 2; + } else { + offset = -static_cast(remoteLane - 4U); + } + } else { + const uint32_t index = remoteLane < 4U ? + remoteLane : remoteLane - 4U; + const int32_t distance = static_cast(group * 4U + index); + offset = remoteLane < 4U ? distance : -distance; + } + return oppositeBase + + (partitionRank + offset + partitionSize) % partitionSize; + } const uint32_t halfWidth = groupWidth / 2U; const uint32_t index = lane < halfWidth ? lane : lane - halfWidth; const int32_t distance = static_cast( diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index fda5b900..7866e9b5 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -120,6 +120,75 @@ void TestSchedules() CHECK_EQ(TileXR::Demo::AllToAllGroupPeer(0, 16, 0, 4, width), -1); } +void Test128PartitionBalancedSchedule() +{ + constexpr int rankSize = 128; + constexpr uint32_t groupWidth = TileXR::Demo::kAllToAllGroupWidth; + constexpr uint32_t groupCount = 8U; + const int expectedRank0[groupCount][groupWidth] = { + {1, 2, 3, 4, 63, 62, 61, 60, + 64, 65, 66, 67, 96, 127, 126, 125}, + {5, 6, 7, 8, 59, 58, 57, 56, + 68, 69, 70, 71, 124, 123, 122, 121}, + {9, 10, 11, 12, 55, 54, 53, 52, + 72, 73, 74, 75, 120, 119, 118, 117}, + {13, 14, 15, 16, 51, 50, 49, 48, + 76, 77, 78, 79, 116, 115, 114, 113}, + {17, 18, 19, 20, 47, 46, 45, 44, + 80, 81, 82, 83, 112, 111, 110, 109}, + {21, 22, 23, 24, 43, 42, 41, 40, + 84, 85, 86, 87, 108, 107, 106, 105}, + {25, 26, 27, 28, 39, 38, 37, 36, + 88, 89, 90, 91, 104, 103, 102, 101}, + {29, 30, 31, 32, 35, 34, 33, -1, + 92, 93, 94, 95, 100, 99, 98, 97}, + }; + + for (uint32_t group = 0U; group < groupCount; ++group) { + for (uint32_t lane = 0U; lane < groupWidth; ++lane) { + CHECK_EQ(TileXR::Demo::AllToAllGroupPeer( + 0, rankSize, group, lane), expectedRank0[group][lane]); + } + } + + for (int rank = 0; rank < rankSize; ++rank) { + const int partition = rank / 64; + std::set peers; + for (uint32_t group = 0U; group < groupCount; ++group) { + uint32_t localPeers = 0U; + uint32_t remotePeers = 0U; + for (uint32_t lane = 0U; lane < groupWidth; ++lane) { + const int peer = TileXR::Demo::AllToAllGroupPeer( + rank, rankSize, group, lane); + if (peer < 0) { + CHECK_EQ(group, groupCount - 1U); + CHECK_EQ(lane, 7U); + continue; + } + CHECK_EQ(peers.insert(peer).second, true); + if (peer / 64 == partition) { + ++localPeers; + CHECK_EQ(lane < 8U, true); + } else { + ++remotePeers; + CHECK_EQ(lane >= 8U, true); + } + bool reciprocal = false; + for (uint32_t remoteLane = 0U; + remoteLane < groupWidth; ++remoteLane) { + reciprocal = reciprocal || + TileXR::Demo::AllToAllGroupPeer( + peer, rankSize, group, remoteLane) == rank; + } + CHECK_EQ(reciprocal, true); + } + CHECK_EQ(localPeers, group + 1U == groupCount ? 7U : 8U); + CHECK_EQ(remotePeers, 8U); + } + CHECK_EQ(peers.size(), static_cast(rankSize - 1)); + } +} + void TestPlan() { constexpr int rankSize = 16; @@ -833,6 +902,7 @@ void TestHostStructure() int main() { TestSchedules(); + Test128PartitionBalancedSchedule(); TestPlan(); TestChannelPolicy(); TestScalePlanAndTraceCapacity(); From 0dc483b72c3b68107375a5735789863cdcd8b89c Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 5 Aug 2026 21:54:38 +0800 Subject: [PATCH 157/163] test(udma): add grouped single-stage isolation --- tests/udma/demo/tilexr_udma_demo.cpp | 38 +++++++++++++++---- ...test_tilexr_udma_alltoall_group_layout.cpp | 1 + 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 1824771f..213f697e 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -1056,6 +1056,27 @@ bool RunGroupedAllToAll( return false; } const bool routeStages = routeStagesValue == 1; + const int routeStageValue = GetEnvInt( + "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE", -1); + if (routeStageValue < -1 || + (routeStageValue >= 0 && !TileXR::Demo::AllToAllGroupValidRouteStage( + static_cast(routeStageValue)))) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE must be -1 or 0..10, got " + << routeStageValue << std::endl; + return false; + } + if (routeStages && routeStageValue >= 0) { + std::cerr << "[rank " << rank + << "] ERROR: TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES and " + "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE cannot both be enabled" + << std::endl; + return false; + } + const bool singleRouteStage = routeStageValue >= 0; + const auto selectedRouteStage = singleRouteStage ? + static_cast(routeStageValue) : + TileXR::Demo::AllToAllGroupRouteStage::kCombined; constexpr size_t kRouteStageCount = 10U; const std::array stagedRouteStages {{ @@ -1260,7 +1281,8 @@ bool RunGroupedAllToAll( " quietBatch=" + std::to_string(quietBatch) + " ingressWindow=" + std::to_string(ingressWindow) + " prewarmSq=" + std::to_string(prewarmSq) + - " routeStages=" + std::to_string(routeStagesValue)); + " routeStages=" + std::to_string(routeStagesValue) + + " routeStage=" + std::to_string(routeStageValue)); if (routeStages && !DemoBarrierAll(rank, rankSize, "grouped route stages ready")) { @@ -1307,8 +1329,7 @@ bool RunGroupedAllToAll( std::array stageTotalUs {}; if (!routeStages) { for (int iter = 0; iter < warmup; ++iter, ++invocationId) { - if (!launchGroupStage( - TileXR::Demo::AllToAllGroupRouteStage::kCombined, nullptr, 0U)) { + if (!launchGroupStage(selectedRouteStage, nullptr, 0U)) { release(); return false; } @@ -1325,7 +1346,7 @@ bool RunGroupedAllToAll( const auto begin = std::chrono::steady_clock::now(); for (int iter = 0; iter < repeat; ++iter, ++invocationId) { - if (!launchGroupStage(TileXR::Demo::AllToAllGroupRouteStage::kCombined, + if (!launchGroupStage(selectedRouteStage, groupTraceDevices[0], static_cast(iter))) { release(); return false; @@ -1387,8 +1408,10 @@ bool RunGroupedAllToAll( std::vector hostDebug(kErrorWordsPerCore * kErrorCoreCount, 0); const size_t debugBytes = hostDebug.size() * sizeof(int32_t); - bool copyOk = CopyDeviceToHost( - rank, hostOutput.data(), dataBytes, output, dataBytes, "grouped alltoall output") && + const bool validateOutput = !singleRouteStage || + selectedRouteStage == TileXR::Demo::AllToAllGroupRouteStage::kCombined; + bool copyOk = (!validateOutput || CopyDeviceToHost( + rank, hostOutput.data(), dataBytes, output, dataBytes, "grouped alltoall output")) && CopyDeviceToHost(rank, hostDebug.data(), debugBytes, debug, debugBytes, "grouped alltoall debug"); if (traceEnabled) { @@ -1442,7 +1465,8 @@ bool RunGroupedAllToAll( } const bool valid = copyOk && debugOk && - ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerPeer); + (!validateOutput || + ValidateAllToAllData(rank, rankSize, hostOutput, elementsPerPeer)); release(); return valid; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 7866e9b5..47f29057 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -841,6 +841,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "launch_tilexr_udma_all_to_all_group"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); + CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); From e36784ee873f36677c5dd7f27144999290f50925 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 5 Aug 2026 22:08:07 +0800 Subject: [PATCH 158/163] fix(udma): retain regions through grouped stage runs --- tests/udma/demo/tilexr_udma_demo.cpp | 6 ++++++ tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp | 1 + 2 files changed, 7 insertions(+) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 213f697e..5f08e426 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -1406,6 +1406,12 @@ bool RunGroupedAllToAll( } } + if (singleRouteStage && + !DemoBarrierAll(rank, rankSize, "grouped single route stage complete")) { + release(); + return false; + } + std::vector hostDebug(kErrorWordsPerCore * kErrorCoreCount, 0); const size_t debugBytes = hostDebug.size() * sizeof(int32_t); const bool validateOutput = !singleRouteStage || diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 47f29057..2325e525 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -842,6 +842,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHUNK_ELEMENTS"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE"); + CHECK_CONTAINS(demo, "grouped single route stage complete"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); From 425dcbecdbbbf38583aade1747b8351a3fca7015 Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 5 Aug 2026 22:22:02 +0800 Subject: [PATCH 159/163] fix(udma): isolate repeated demo barriers --- tests/udma/demo/tilexr_udma_demo.cpp | 10 ++++++---- .../unit/test_tilexr_udma_alltoall_group_layout.cpp | 1 + 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index 5f08e426..ee69f7f6 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -278,7 +278,7 @@ bool WriteGroupTraceBinary( return true; } -BarrierEndpoint GetBarrierEndpoint() +BarrierEndpoint GetBarrierEndpoint(uint32_t sequence) { std::string host = "127.0.0.1"; int basePort = kDefaultCommPort; @@ -297,9 +297,10 @@ BarrierEndpoint GetBarrierEndpoint() basePort = std::atoi(value.c_str() + colon + 1); } } - int barrierPort = basePort + kDemoBarrierPortOffset; + int barrierPort = basePort + kDemoBarrierPortOffset + static_cast(sequence); if (barrierPort <= 0 || barrierPort > 65535) { - barrierPort = kDefaultCommPort + kDemoBarrierPortOffset; + barrierPort = kDefaultCommPort + kDemoBarrierPortOffset + + static_cast(sequence); } return BarrierEndpoint{host, static_cast(barrierPort)}; } @@ -396,7 +397,8 @@ bool DemoBarrierAll(int rank, int rankSize, const std::string& step) return true; } - BarrierEndpoint endpoint = GetBarrierEndpoint(); + static uint32_t barrierSequence = 0U; + BarrierEndpoint endpoint = GetBarrierEndpoint(barrierSequence++); PrintStatus(rank, "demo tcp barrier begin: " + step + " host=" + endpoint.host + " port=" + std::to_string(endpoint.port)); constexpr uint8_t kArrive = 1; diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 2325e525..4343ce8d 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -843,6 +843,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGES"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE"); CHECK_CONTAINS(demo, "grouped single route stage complete"); + CHECK_CONTAINS(demo, "GetBarrierEndpoint(barrierSequence++)"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); From cd712ca75d31fe92da03918c86ee62f49434ca3f Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 5 Aug 2026 22:34:56 +0800 Subject: [PATCH 160/163] fix(udma): wait for grouped terminal barrier --- tests/udma/demo/tilexr_udma_demo.cpp | 19 +++++++++++++------ ...test_tilexr_udma_alltoall_group_layout.cpp | 1 + 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_demo.cpp b/tests/udma/demo/tilexr_udma_demo.cpp index ee69f7f6..f28cf891 100644 --- a/tests/udma/demo/tilexr_udma_demo.cpp +++ b/tests/udma/demo/tilexr_udma_demo.cpp @@ -368,7 +368,8 @@ int CreateBarrierServer(uint16_t port) return fd; } -int ConnectBarrierServer(const std::string& host, uint16_t port) +int ConnectBarrierServer(const std::string& host, uint16_t port, + int retryCount = kConnectRetryCount) { sockaddr_in addr{}; addr.sin_family = AF_INET; @@ -377,7 +378,7 @@ int ConnectBarrierServer(const std::string& host, uint16_t port) } addr.sin_port = htons(port); - for (int attempt = 0; attempt < kConnectRetryCount; ++attempt) { + for (int attempt = 0; attempt < retryCount; ++attempt) { int fd = socket(AF_INET, SOCK_STREAM, 0); if (fd < 0) { return -1; @@ -391,7 +392,8 @@ int ConnectBarrierServer(const std::string& host, uint16_t port) return -1; } -bool DemoBarrierAll(int rank, int rankSize, const std::string& step) +bool DemoBarrierAll(int rank, int rankSize, const std::string& step, + int connectRetryCount = kConnectRetryCount) { if (rankSize <= 1) { return true; @@ -438,7 +440,8 @@ bool DemoBarrierAll(int rank, int rankSize, const std::string& step) return false; } } else { - int fd = ConnectBarrierServer(endpoint.host, endpoint.port); + int fd = ConnectBarrierServer( + endpoint.host, endpoint.port, connectRetryCount); if (fd < 0) { std::cerr << "[rank " << rank << "] ERROR: failed to connect demo barrier on " << endpoint.host << ":" << endpoint.port << std::endl; @@ -1408,8 +1411,12 @@ bool RunGroupedAllToAll( } } - if (singleRouteStage && - !DemoBarrierAll(rank, rankSize, "grouped single route stage complete")) { + const int terminalBarrierRetryCount = std::max(kConnectRetryCount, + GetEnvInt("TILEXR_DEMO_TIMEOUT_SECONDS", 1800) * + 1000 / kConnectRetrySleepMs); + if (singleRouteStage && !DemoBarrierAll( + rank, rankSize, "grouped single route stage complete", + terminalBarrierRetryCount)) { release(); return false; } diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 4343ce8d..6e6fa1a3 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -844,6 +844,7 @@ void TestHostStructure() CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_ROUTE_STAGE"); CHECK_CONTAINS(demo, "grouped single route stage complete"); CHECK_CONTAINS(demo, "GetBarrierEndpoint(barrierSequence++)"); + CHECK_CONTAINS(demo, "terminalBarrierRetryCount"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_USE_SECONDARY_ROUTE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_CHANNEL_MODE"); CHECK_CONTAINS(demo, "TILEXR_DEMO_ALLTOALL_GROUP_WIDTH"); From 2a2ae0122964564838dbfde5552fdbcb27c3bbbf Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Wed, 5 Aug 2026 23:59:24 +0800 Subject: [PATCH 161/163] fix(udma): extend grouped quiet polling --- tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp | 10 ++++++++-- .../unit/test_tilexr_udma_alltoall_group_layout.cpp | 1 + 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 7e1351ce..2155274f 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -754,8 +754,14 @@ __aicore__ inline bool AllToAllGroupCompleteQuiet( uint32_t groupCount, uint32_t passCount) { const uint64_t quietBegin = AllToAllGroupTraceCycle(trace); - const uint32_t quietStatus = - TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); + const uint64_t quietDeadlineBegin = + static_cast(AscendC::GetSystemCycle()); + uint32_t quietStatus = 0U; + do { + quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); + } while (quietStatus == 0xFFU && + static_cast(AscendC::GetSystemCycle()) - quietDeadlineBegin < + TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES); AllToAllGroupTraceRecordTask( trace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 6e6fa1a3..627e33b6 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -759,6 +759,7 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, request.peer, request.qpIdx)"); CHECK_CONTAINS(kernel, "state.pendingCount != quietBatch"); + CHECK_CONTAINS(kernel, "quietStatus == 0xFFU"); CHECK_CONTAINS(kernel, "template "); CHECK_CONTAINS(kernel, "struct AllToAllGroupQuietState"); CHECK_CONTAINS(kernel, "AllToAllGroupQuietState quietState"); From 47dd5e7e8e3b1a0ab6aad75c7e0662132bcffb5d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 6 Aug 2026 00:30:56 +0800 Subject: [PATCH 162/163] fix(udma): align 128p grouped shared QP lanes --- src/comm/udma/tilexr_udma_layout.cpp | 28 +++++++++++++++++++ tests/udma/CMakeLists.txt | 1 + .../tilexr_udma_alltoall_group_kernel.cpp | 10 ++----- ...test_tilexr_udma_alltoall_group_layout.cpp | 1 - .../test_tilexr_udma_transport_layout.cpp | 28 +++++++++++++++++++ 5 files changed, 59 insertions(+), 9 deletions(-) diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index d4a64cb5..2d95572f 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -158,6 +158,34 @@ uint32_t UDMASharedQpLane( return laneCount; } + if (rankSize == 128 && laneCount == 16) { + constexpr int partitionSize = 64; + const int rankPartition = rank / partitionSize; + const int peerPartition = peer / partitionSize; + const uint32_t rankInPartition = static_cast(rank % partitionSize); + const uint32_t peerInPartition = static_cast(peer % partitionSize); + const uint32_t forward = + (peerInPartition + partitionSize - rankInPartition) % partitionSize; + if (rankPartition == peerPartition) { + if (forward <= partitionSize / 2) { + return (forward - 1U) % 4U; + } + const uint32_t backward = partitionSize - forward; + return 4U + (backward - 1U) % 4U; + } + if (forward == 0U) { + return 8U; + } + if (forward < partitionSize / 2) { + return 8U + forward % 4U; + } + if (forward == partitionSize / 2) { + return 12U; + } + const uint32_t backward = partitionSize - forward; + return 12U + backward % 4U; + } + const uint32_t lanesPerDirection = laneCount / 2; const uint32_t forward = static_cast((peer - rank + rankSize) % rankSize); const uint32_t backward = static_cast(rankSize) - forward; diff --git a/tests/udma/CMakeLists.txt b/tests/udma/CMakeLists.txt index ec0edc5c..15370ec9 100644 --- a/tests/udma/CMakeLists.txt +++ b/tests/udma/CMakeLists.txt @@ -144,6 +144,7 @@ target_include_directories(test_tilexr_sock_exchange_layout PRIVATE target_include_directories(test_tilexr_udma_transport_layout PRIVATE ${TILEXR_ROOT}/src/include ${TILEXR_ROOT}/src/comm + ${CMAKE_CURRENT_SOURCE_DIR}/demo ) target_compile_definitions(test_tilexr_udma_transport_layout PRIVATE diff --git a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp index 2155274f..7e1351ce 100644 --- a/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp +++ b/tests/udma/demo/tilexr_udma_alltoall_group_kernel.cpp @@ -754,14 +754,8 @@ __aicore__ inline bool AllToAllGroupCompleteQuiet( uint32_t groupCount, uint32_t passCount) { const uint64_t quietBegin = AllToAllGroupTraceCycle(trace); - const uint64_t quietDeadlineBegin = - static_cast(AscendC::GetSystemCycle()); - uint32_t quietStatus = 0U; - do { - quietStatus = TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); - } while (quietStatus == 0xFFU && - static_cast(AscendC::GetSystemCycle()) - quietDeadlineBegin < - TILEXR_ALLTOALL_GROUP_WAIT_TIMEOUT_CYCLES); + const uint32_t quietStatus = + TileXR::UDMAQuietStatusOnQp(args, peer, selectedQp); AllToAllGroupTraceRecordTask( trace, traceIteration, blockIdx, group, pass, TileXR::Demo::kAllToAllGroupTraceSendQuiet, groupCount, passCount, diff --git a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp index 627e33b6..6e6fa1a3 100644 --- a/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_alltoall_group_layout.cpp @@ -759,7 +759,6 @@ void TestKernelStructure() CHECK_CONTAINS(kernel, "UDMAQuietStatusOnQp(args, request.peer, request.qpIdx)"); CHECK_CONTAINS(kernel, "state.pendingCount != quietBatch"); - CHECK_CONTAINS(kernel, "quietStatus == 0xFFU"); CHECK_CONTAINS(kernel, "template "); CHECK_CONTAINS(kernel, "struct AllToAllGroupQuietState"); CHECK_CONTAINS(kernel, "AllToAllGroupQuietState quietState"); diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index a9f9a5af..c6410cd8 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -8,6 +8,7 @@ #include #include "udma/tilexr_udma_layout.h" +#include "tilexr_udma_alltoall_group_layout.h" #ifndef TILEXR_SOURCE_ROOT #define TILEXR_SOURCE_ROOT "." @@ -189,6 +190,32 @@ void TestSharedQpLaneMatchesGroupedPeerOrder() CHECK_EQ(TileXR::UDMASharedQpLane(0, 0, 64, 16), 16U); } +void TestSharedQpLaneMatchesBalanced128Schedule() +{ + constexpr int rankSize = 128; + constexpr uint32_t groupWidth = TileXR::Demo::kAllToAllGroupWidth; + const uint32_t groupCount = + TileXR::Demo::AllToAllGroupCount(rankSize, groupWidth); + for (int rank = 0; rank < rankSize; ++rank) { + for (uint32_t group = 0U; group < groupCount; ++group) { + bool used[16] = {}; + for (uint32_t lane = 0U; lane < groupWidth; ++lane) { + const int peer = TileXR::Demo::AllToAllGroupPeer( + rank, rankSize, group, lane, groupWidth); + if (peer < 0) { + continue; + } + const uint32_t sharedLane = + TileXR::UDMASharedQpLane(rank, peer, rankSize, groupWidth); + CHECK_EQ(sharedLane, lane); + CHECK_TRUE(sharedLane < groupWidth); + CHECK_TRUE(!used[sharedLane]); + used[sharedLane] = true; + } + } + } +} + void TestSharedQpLanesAreUniqueWithinEveryGroup() { for (int rankSize = 8; rankSize <= 1024; rankSize += 8) { @@ -448,6 +475,7 @@ int main() TestMultiRouteQpMappingRejectsEmptyInputs(); TestMultiRouteQpWeightsUseRouteBandwidth(); TestSharedQpLaneMatchesGroupedPeerOrder(); + TestSharedQpLaneMatchesBalanced128Schedule(); TestSharedQpLanesAreUniqueWithinEveryGroup(); TestSharedQpPoolIsIndependentOfRegions(); TestSharedQpIndexUsesLaneOnly(); From b4f03a5862b6e05cfb824817f946d2a34548df6d Mon Sep 17 00:00:00 2001 From: "hechangcheng (A)" Date: Thu, 6 Aug 2026 00:35:28 +0800 Subject: [PATCH 163/163] test(udma): use balanced lane contract at 128p --- tests/udma/unit/test_tilexr_udma_transport_layout.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp index c6410cd8..cdb6178a 100644 --- a/tests/udma/unit/test_tilexr_udma_transport_layout.cpp +++ b/tests/udma/unit/test_tilexr_udma_transport_layout.cpp @@ -219,6 +219,9 @@ void TestSharedQpLaneMatchesBalanced128Schedule() void TestSharedQpLanesAreUniqueWithinEveryGroup() { for (int rankSize = 8; rankSize <= 1024; rankSize += 8) { + if (rankSize == 128) { + continue; + } const uint32_t groupCount = static_cast((rankSize - 1 + 15) / 16); for (int rank = 0; rank < rankSize; ++rank) { for (uint32_t group = 0; group < groupCount; ++group) {