Repository navigation
feat(rccl): TDM async-to-LDS path for LL128 user buffers on gfx1250 - #12375
mustafabar wants to merge 24 commits into
Conversation
✅ All Policy Checks Passed
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
07a6b41 to
0e2b2c0
Compare
github-actions
Bot
added
the
documentation
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unaligned fallback operations can race pending async stores, destination alignment violates the copy contract, and registered-buffer behavior remains untested.
Review effort: Balanced
Findings: 1
Open (3)
What changed in this PR
Adds an opt-in gfx1250 async-to-LDS path for LL128 user-buffer transfers.
Changes:
- Implements asynchronous LL128 source loads and destination stores.
- Adds build-system and installer gating.
- Adds MI455 validation configuration and documentation.
| File | Description |
|---|---|
projects/rccl/src/device/tdm/ll128Tdm.h |
Implements TDM staging operations. |
projects/rccl/src/device/prims_ll128.h |
Integrates TDM into LL128 primitives. |
projects/rccl/CMakeLists.txt |
Adds the build option. |
projects/rccl/src/CMakeLists.txt |
Registers the new header. |
projects/rccl/install.sh |
Adds the installer flag. |
projects/rccl/docs/install/building-installing.rst |
Documents the flag. |
projects/rccl/tools/scripts/test_runner/configs/mi455_tdm_ll128_func.json |
Adds functional validation suites. |
projects/rccl/CHANGELOG.md |
Records the feature. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
![amd-argus[bot]](https://avatars.githubusercontent.com/in/4740020?s=60&v=4){kind=small}
There was a problem hiding this comment.
ARGUS
Bottom line: Worth settling before approval: the staged window at ll128Tdm.h:84 reaches the async store engine with no __threadfence_block(), and RCCL's __syncwarp() is a no-op.
Comment @argus review to re-run, or @argus re-review after pushing.
| int ix = g * WARP_SIZE - LineSkip * (g / 2) + wid - (g % 2) * (wid / (LineElems / 2)); | ||
| if (!flagThread || g % 2 == 0) storeShmem128(shm8 + 2 * ix, regs[2 * g + 0], regs[2 * g + 1]); | ||
| } | ||
| __syncwarp(); |
There was a problem hiding this comment.
High
tdmStoreRegs publishes the staged window to asyncStoreFromLDS behind only __syncwarp(), which common.h:19 defines to nothing, so nothing orders the lanes' storeShmem128 writes against the DMA engine about to read them. That no-op is sound for the existing shmem fallback, where producer and consumer are both the same wave's DS pipeline, but the engine is not. The in-tree tile store fences first for exactly this reason: ncclSymkTileFenceSmem() at symmetric/primitives.cuh:169 emits __threadfence_block() and is called right before every ncclSymkTileStore.
| __syncwarp(); | |
| } | |
| __threadfence_block(); | |
| __syncwarp(); |
There was a problem hiding this comment.
Resolved by 2f15b7cdd8.
| { | ||
| "name": "TDM_LL128_Ring_ReduceScatter_UserBufferRegistration", | ||
| "description": "Registered user buffers for ReduceScatter under forced Ring + LL128", | ||
| "test_filter": "ReduceScatter.UserBufferRegistration:ReduceScatter.ManagedMemUserBufferRegistration" |
There was a problem hiding this comment.
Medium
ReduceScatter.UserBufferRegistration and ReduceScatter.ManagedMemUserBufferRegistration do not exist anywhere: ReduceScatterTests.cpp registers only OutOfPlace, OutOfPlaceGraph, InPlace, InPlaceGraph, ManagedMem and ManagedMemGraph. A filter matching nothing is recorded as SKIPPED (lib/test_executor.py:219-238) and only FAILED and TIMEOUT are counted, so this entry is a permanent green covering no ReduceScatter registration path at all. The same two names also appear as exclusions at line 71, where test_categories.yaml:160 has plain ReduceScatter.* with none.
There was a problem hiding this comment.
Confirmed and fixed in a43f574. ReduceScatter registers only OutOfPlace, OutOfPlaceGraph, InPlace, InPlaceGraph, ManagedMem and ManagedMemGraph, so both filters matched nothing. Removed the entry from the registration suite and the matching dead exclusions from the Ring suite.
There was a problem hiding this comment.
Resolved by a43f574d61.
| ] | ||
| } | ||
| }, | ||
| "test_suites": [ |
There was a problem hiding this comment.
Medium
These three suites would pass unchanged with ll128Tdm.h deleted and prims_ll128.h reverted, because every test they name is a pre-existing LL128 correctness test that observes nothing about which copy path the kernel took. The missing observable is a differential one: the same collective and datatype matrix built with and without ENABLE_TDM_PRIM_LL128, asserting the output buffers match, which fails if the TDM shuffle indices or the drain ordering are wrong. Could we add that arm?
There was a problem hiding this comment.
Agreed that these suites observe correctness under a TDM build rather than proving the TDM path ran. A differential arm needs two builds and a harness to compare outputs, which a single test-runner config cannot express, so it belongs in a follow-up rather than here. Tracking it alongside the accumulator path, which is also unexercised.
| "NCCL_PROTO": "LL128", | ||
| "NCCL_ALGO": "Ring" | ||
| }, | ||
| "tests": [ |
There was a problem hiding this comment.
Medium
The Ring suite runs 4 of the 12 patterns in pow2_list (test_categories.yaml:154-166) and the Tree suite 2, though the description at line 4 says both mirror full_pow2_ring_ll128 and full_pow2_tree_ll128, which alias that list verbatim. In-place Scatter is the omission that matters: CollectiveArgs.cpp:79 offsets its destination by globalRank * numOutputBytes, making it the one case that hands tdmStoreRegs an arbitrarily aligned dst, which is the input the new store leg is least exercised on. Could we add Scatter.* to the Ring list?
There was a problem hiding this comment.
Added Scatter to the Ring suite in a43f574. The in-place destination offset is the input the store leg is least exercised on, so this is the most useful of the missing patterns.
There was a problem hiding this comment.
Resolved by a43f574d61.
|
|
||
| static constexpr int TdmAlign = 128; | ||
| // Mirrors loadUser128/storeUser128: system scope only when the buffer is registered. | ||
| static constexpr CachePolicy TdmLoadDevPolicy = createCachePolicy(TemporalHint::NT, MemScope::DEV); |
There was a problem hiding this comment.
Low
TdmLoadDevPolicy uses MemScope::DEV, but the loadUser128 branch it mirrors resolves to load128NT, which on gfx1250 is a system-scope global_load_b128 because RCCL_LL_FIFO_SYS_SCOPE is 1 there (nccl_device/rccl_ptr.h:82). TDM_LL128_ON only holds on gfx1250, so both loadUser128 branches are system scope on the one arch this file compiles for, and the comment above is inverted for the load leg. The other three policies do mirror their originals. Please check.
There was a problem hiding this comment.
Confirmed and fixed in a43f574. RCCL_LL_FIFO_SYS_SCOPE is 1 on gfx1250 and TDM_LL128_ON only holds there, so both loadUser128 branches are system scope. Collapsed the two load policies into one system-scope policy and dropped the now redundant userBypass branch on the load. Measured at 539 and 566 GB/s against 557 and 569 before, which sits inside the baseline spread.
There was a problem hiding this comment.
Resolved by a43f574d61.
| asyncWait<0>(); | ||
| #pragma unroll | ||
| for (int g = 0; g < WordPerThread / 2; g++) { | ||
| int ix = g * WARP_SIZE - LineSkip * (g / 2) + wid - (g % 2) * (wid / (LineElems / 2)); |
There was a problem hiding this comment.
Low
The LL128 slot-index expression now sits in four byte-identical copies, at prims_ll128.h:267, prims_ll128.h:344, and lines 65 and 81 here, each with its own local LineElems and LineSkip pair; two of the four arrive with this PR. algorithms/dda/device/CollCommon_ll128.h:99 already factors the same formula as ddaLL128ChunkIx(g, wid) over class-scope constants. Could we pull one ll128SlotIx(g) member into the class and have all four sites call it?
There was a problem hiding this comment.
Factored a tdmSlotIx member in a43f574 and both copies in this file now call it. I left the two pre-existing copies in prims_ll128.h alone on purpose. That file tracks upstream NCCL closely and this change is structured to keep its footprint to guarded one-line call sites, so editing stock lines there would widen the conflict surface on the next sync.
There was a problem hiding this comment.
Resolved by a43f574d61.
| * Added `ncclGinFenceLevel` barrier semantics on the Anvil SDMA GIN backend (`NCCL_GIN_TYPE=7`). `gin.get`, `flushAsync` and `wait` are implemented instead of trapping, and a barrier's signal is now ordered behind the payload on the peer's SDMA queue, so `Put`, `Get` and the default `Put | Get` fences drain what they promise. `Put` covers a rank's puts to itself, and `ncclGinAllContexts` fences every GIN context. Single-node validated on gfx950. | ||
| * Added gfx1250 (MI450) Tensor Data Mover staging to the symmetric AllGather, AllReduce and ReduceScatter kernels, where NCCL uses NVIDIA's TMA. Their deep loops move each tile through the warp's shared-memory scratch with the DMA engine instead of per-lane vector loads; on ROCm that path previously compiled out entirely, so the `Tma` kernels were never generated. Launch tuning was also corrected so these kernels can use up to a quarter of a GPU's WGPs, where only 14 were reached before, which is where most of the large-message gain comes from. Opt-in via `NCCL_SYM_TMA_ENABLE=1`, and reached only when the application opts into symmetric memory. `NCCL_SYM_TMA_ENABLE=2` additionally forces these kernels wherever the collective has one, skipping both the tuner's cost comparison and the deep-loop size bar, which is intended for A/B measurement rather than production. The new kernels are generated only when gfx1250 is in `GPU_TARGETS`, and every other architecture keeps the vector path unchanged. | ||
| * Removed GIN rocSHMEM GDA dependency on rocSHMEM GDA bitcode: GIN rocSHMEM GDA is now header-only. The CMake `roc::rccl` target provides the required include directories in its `INSTALL_INTERFACE` property. | ||
| * Added an experimental gfx1250 (MI450) Tensor Data Mover path for the LL128 protocol's user source and destination buffers, staging each warp slice through shared memory with the async-to-LDS engine instead of per-lane vector accesses. The receive and send FIFO legs are not eligible, because both interleave protocol flag words into the payload. Excluded from the default build: it requires `--enable-tdm-prim-ll128` at build time. There is no runtime variable, because a runtime check leaves both store paths in the kernel and costs roughly 3x. Registered user buffers and the accumulator path are not yet exercised. |
There was a problem hiding this comment.
Low
The entry closes with "Registered user buffers and the accumulator path are not yet exercised", but the same PR adds the TDM LL128 - Registered user buffers suite with "enabled": true at mi455_tdm_ll128_func.json:167, running the UserBufferRegistration and ManagedMemUserBufferRegistration cases and described as covering the system-scope cache policy. The accumulator half of the sentence is accurate. Could we settle the registered-buffer half, so a reader knows whether that path has been validated?
There was a problem hiding this comment.
Fixed in a43f574. The entry now states only the accumulator path as unexercised.
There was a problem hiding this comment.
Resolved by a43f574d61.
| "-t", | ||
| "--enable-tdm-prim-ll128" | ||
| ], | ||
| "cmake_options": "-DBITCODE_LIB_ARCH=gfx1250", |
There was a problem hiding this comment.
Low
-DBITCODE_LIB_ARCH=gfx1250 selects nothing here: it is read only by bindings/ir, and that subdirectory is added only under if(EMIT_LLVM_IR) at CMakeLists.txt:817, which this config never sets. The arch actually comes from -l, which is BUILD_LOCAL_GPU_TARGET_ONLY and resolves to whatever the runner has, so the line reads as though the target under test is pinned when nothing pins it. Both sibling gfx1250 configs pass -DEMIT_LLVM_IR=ON alongside it (mi455_ainic_roce_func.json:33).
There was a problem hiding this comment.
Confirmed and removed in a43f574. bindings/ir is added only under EMIT_LLVM_IR, which this config does not set, so the line selected nothing. The arch comes from -l as you say.
There was a problem hiding this comment.
Resolved by a43f574d61.
| --enable_backtrace) build_bfd=true; shift ;; | ||
| --enable-mpi-tests) enable_mpi_tests=true; shift ;; | ||
| --enable-rccl-ep-tests) enable_rccl_ep_tests=true; shift ;; | ||
| --enable-tdm-prim-ll128) enable_tdm_prim_ll128=true; shift ;; |
There was a problem hiding this comment.
Nit
| --enable-tdm-prim-ll128) enable_tdm_prim_ll128=true; shift ;; | |
| --enable-tdm-prim-ll128) enable_tdm_prim_ll128=true; shift ;; |
There was a problem hiding this comment.
Resolved by a43f574d61.
|
@argus re-review |
There was a problem hiding this comment.
ARGUS
Bottom line: Worth settling before approval: the staged window at ll128Tdm.h:82 still reaches the async store engine with no __threadfence_block(), and RCCL's __syncwarp() compiles to nothing.
Comment @argus review to re-run, or @argus re-review after pushing.
| } | ||
| } | ||
|
|
||
| #if TDM_LL128_ON |
There was a problem hiding this comment.
Low
The 8 #if TDM_LL128_ON regions in this file (247, 253, 315, 333, 493, 508, 513, 530) wrap 8 one-line payloads in 16 preprocessor lines, in a file you want to keep close to upstream. Moving one #if/#else into ll128Tdm.h, with tdmLoadBegin/tdmStoreRegs returning false, an empty tdmDrain() and a tdmSetLoadAllowed(bool), turns all 8 call sites into plain calls the compiler folds away and drops about 14 of the 33 lines added here. The include at line 9 is already unconditional, so ASYNC_COPY_SUPPORTED stays defined either way.
| "NCCL_DEBUG": "INFO" | ||
| } | ||
| }, | ||
| "tdm_ll128_ring": { |
There was a problem hiding this comment.
Low
tdm_ll128_ring, tdm_ll128_tree and tdm_ll128_ubr repeat the same six scalars and the same nine env_variables; tree differs only at NCCL_ALGO and ubr in nothing but its tests. A childless tdm_ll128_base extended by all three removes about 28 of this file's 169 lines, which is the shape net_ib_base and cast_base already use in mi455_ainic_roce_func.json. The base has to omit tests, since _merge_tests at lib/test_config.py:275 unions parent and child lists rather than overriding.
There was a problem hiding this comment.
ARGUS
Bottom line: Worth settling before approval: the staged window at ll128Tdm.h:82 still reaches the async store engine with no __threadfence_block(), and RCCL's __syncwarp() compiles to nothing.
Comment @argus review to re-run, or @argus re-review after pushing.
Source load overlaps the receive-FIFO spin; destination store is staged in LDS and pushed out as one bulk transfer. Implementation in tdm/ll128Tdm.h, gated by TDM_LL128.
Adds RCCL_TDM_LL128_ENABLE, mirroring RCCL_TDM_SIMPLE_ENABLE. The build flag TDM_LL128 now defaults off, since leaving the path compiled in and predicated off costs about 5% on AllGather at 1 GiB.
The path is compiled in wherever the async-to-LDS builtins exist and is selected at runtime with RCCL_TDM_LL128_ENABLE=1.
Replaces the RCCL_TDM_LL128_ENABLE runtime check. Off by default. tdmWindow() aligns by offset instead of casting through uintptr_t.
Adds the install flag to the build options and a CHANGELOG entry.
Builds with --enable-tdm-prim-ll128 and runs the unit suite under forced
LL128 on Ring and Tree, mirroring the full_pow2_{ring,tree}_ll128
categories. Runnable by hand until gfx1250 CI capacity exists.
loadRegsBegin's misaligned fallback also stages through that window, so returning early left it racing the previous slice's async store. Also adds a registered-buffer suite to the test-runner config.
Load policy is system scope, matching loadUser128 on gfx1250 where RCCL_LL_FIFO_SYS_SCOPE is 1 on both branches. Factors the repeated LL128 slot-index expression. Test config drops ReduceScatter registration filters that match no test, adds Scatter for its unaligned destination, and drops BITCODE_LIB_ARCH, which is only read under EMIT_LLVM_IR.
Co-authored-by: mustafabar <2463168+mustafabar@users.noreply.github.com>
RCCL_TDM_LL128_ENABLE=1, captured in the primitive's constructor so the hot path tests a register instead of reading ncclShmem per slice.
Replaces the async-to-LDS builtins with tdm::. Keeping both the staged and the vector path live spilled the LL128 ring body under the former, which cost 5x on a gated RING/LL128 AllReduce; the tensor mover does not.
56c87b8 to
b5dba58
Compare
|
No undocumented RCCL or NCCL environment-variable changes. |
There was a problem hiding this comment.
ARGUS
Findings
1 finding(s) that could not be posted as inline comments
projects/rccl/tools/scripts/test_runner/configs/mi455_tdm_ll128_func.json:14(High) The new TDM LL128 validation config never sets RCCL_TDM_LL128_ENABLE, so all three suites run with the path offmi455_tdm_ll128_func.jsonnever setsRCCL_TDM_LL128_ENABLE, so none of the three suites reach the new code. The build flag is only half the gate: the param defaults to 0 at init.cc:225, feedscomm->tdmLl128Enableat init.cc:2048, becomestdmEnableat prims_ll128.h:619, and short-circuitstdmLoadBegin,tdmStoreRegsandtdmDrainat ll128Tdm.h:45, :79 and :100. This top-level block merges into every test, so one line covers all three suites.
Bottom line: Worth settling before approval: mi455_tdm_ll128_func.json never sets RCCL_TDM_LL128_ENABLE, so all three suites run with the TDM path off, and the staged window at ll128Tdm.h:87 still reaches the store engine with no __threadfence_block().
Comment @argus review to re-run, or @argus re-review after pushing.
The build flag is half the gate: RCCL_TDM_LL128_ENABLE defaults to 0, so all three suites ran with the path off. Also add a capability arm, since an SDK without tensor-mover support compiles the path away and passes green.
…taging Staging the source through the same per-warp window corrupted exactly one slice in ~2% of 2-node ReduceScatter runs at 4 GiB, and bought nothing: destination staging alone matches both legs within 1% at every size. Store leg only: 0 of 100 runs wrong, against 2 of 100 with the load leg.
Removed comments regarding source buffer staging issues and FIFO eligibility.
Three byte-identical copies of the flag-hole compensation collapse into one __host__ __device__ ll128SlotIx(). Host-callable so tests can check the layout without a GPU. Addresses the ARGUS note on the duplicated formula.
Three tests, landed together because they share one fixture:
Ll128TdmBuildFlagMatchesExpectation - fails when --enable-tdm-prim-ll128 did
not reach the device pass, gated on RCCL_TEST_EXPECT_TDM_LL128 so default
builds skip rather than fail.
Ll128SlotIndicesPackDensely - host-only; the slot indices must tile the
staged window exactly once.
Ll128StoreLegMatchesScalarPath - the staged store byte-for-byte against the
scalar storeUser128 path it replaces, 512 warps x 4 seeds.
Needs 1 GPU, so it runs on 4-GPU gfx1250 nodes where the 8-GPU suites cannot.
… config Sets RCCL_TEST_EXPECT_TDM_LL128=1 and adds two arms to the capability suite, so a build with the flag silently compiled out now goes red instead of green.
There was a problem hiding this comment.
ARGUS
Bottom line: Worth settling before approval: slice N's async TDM store can still be reading the LDS window when slice N+1's unaligned loadRegsBegin overwrites the same per-warp scratch, with no tdmWait in between (ll128Tdm.h:38).
Comment @argus review to re-run, or @argus re-review after pushing.
loadRegsBegin's unaligned fallback stages through the same per-warp scratch the TDM store leg hands the mover, and the only drain was in the next tdmStoreRegs, which runs after the clobber. Restores the wait that 1f1ba9b dropped along with the load leg. Not reproduced: 9 collective/offset configs on the unfixed build are clean, since the fallback's global loads give the in-flight store time to drain.
Adds the variable to the gfx1250 dispatch table in env-variables.rst, which is what the env-var changelog gate asks for, and refreshes the CHANGELOG entry: it still described the source-load leg that 1f1ba9b removed, and now carries the measured numbers and the new tests.
Routes the two LL128 user-buffer legs through the gfx1250 async-to-LDS engine. The path is off by default and is opted into at build time with
./install.sh --enable-tdm-prim-ll128.Motivation
LL128 reads the user source buffer and writes the user destination buffer with per-lane vector accesses. Both legs can be handed to the async-to-LDS engine on gfx1250. The source read then overlaps the receive-FIFO spin, and the destination write goes out as a single bulk transfer instead of one store per lane.
The receive and send FIFO legs are not eligible. The receive path re-evaluates a predicate on the values it has just loaded, and both FIFO legs interleave protocol flag words into the payload, so neither survives a bulk copy.
Technical Details
Async state is held in class members instead of being threaded through parameters, which leaves
loadRegsBegin,loadRegsFinish,storeRegsandrecvReduceSendCopywith their original signatures.load and store share one per-warp staging window. Every write to that window first drains whatever async operation is still reading it, and the last store is drained before the barrier that publishes completion.
Gating is build time only.
ENABLE_TDM_PRIM_LL128defaults to off and is set by the new install flag, mirroring howENABLE_TDM_SIMPLEis wired. There is deliberately no runtime switch. An earlier revision gated the path on an environment variable, which cost 3.3x whether the variable was set or not, because a runtime condition leaves both the TDM bulk store and the original per-lane store in the compiled kernel. Caching the flag in a member, specialisingstoreRegson a compile-time bool, and hoisting the decision out of the slice loop all failed to recover that. Every variant carrying one store path ran at full speed and every variant carrying two did not.Issue Tracking
JIRA ID : AICOMRCCL-1637
Test Plan
Single node and two-node run with 4/8 gfx1250 GPUs on ROCm 10.2. All arms were built in one session with identical
ONLY_FUNCS. The baseline is pristine develop at9c02ecdc30, which is the develop-side parent of this branch's merge.Runs force
NCCL_ALGO=RINGandNCCL_PROTO=LL128so the conventional primitives are exercised, usingall_reduce_perf -b 128M -e 512M -f 4 -g 1 -A 1 -n 20over two repetitions. Three configurations were measured: develop, this branch built without the flag, and this branch built with the flag.Test coverage is still pending. LL128 has existing coverage if the TDM + LL128 build is enabled
Test Result
Results on AR, RS and AG are included in the Jira AICOMRCCL-1637
Submission Checklist