feat(distributed): Complete arbitrary-length allreduce support - #2161
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe change expands distributed allreduce support to FP16/FP32 and four reduction operators, rewrites ring segmentation and tail handling, inlines remote address translation with fencing, aligns communication buffers, and updates runtime templates, validation, tests, and documentation. ChangesDistributed collective support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant UserProgram
participant LowerCompositeOps
participant CommContext
participant RuntimeKernel
UserProgram->>LowerCompositeOps: invoke allreduce with dtype, operator, and mode
LowerCompositeOps->>CommContext: derive peer offsets and aligned segment shapes
LowerCompositeOps->>RuntimeKernel: emit reduction instruction and transfer metadata
RuntimeKernel->>RuntimeKernel: load, reduce, restore valid shape, and store
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9c3fe8ac13
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/pypto/language/distributed/op/tile_ops.py (1)
49-101: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winKeep physical-tail reads lowering-only. The public DSL/API currently exposes
allow_physical_tail_padding, whichremote_loadcan translate into reading an additional FP16 tail element and marking it valid. Move this variant behind a lowering-only IR constructor/operation so user callers cannot read logical tail padding.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pypto/language/distributed/op/tile_ops.py` around lines 49 - 101, Keep physical-tail padding lowering-only: remove allow_physical_tail_padding from the public distributed tile API around remote_load in python/pypto/language/distributed/op/tile_ops.py (lines 49-101), and expose it only through a dedicated lowering-only IR constructor/operation in python/pypto/ir/op/distributed/tile_ops.py (lines 32-70). Update the remote_load lowering and verification paths in src/ir/op/distributed/remote_load.cpp (lines 146-165 and 206-213) to accept and process the variant only through that internal operation, preventing user-facing calls from reading logical tail padding.
🧹 Nitpick comments (8)
tests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.py (1)
94-101: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFP16
stage_colsis uncapped, unlike the FP32 path.FP32 clamps the staging tile to
STAGE_CHUNK(8192), but the FP16 branch setsstage_colsto the full 16-alignedsize. Atsize=8193that is already a ~16 KB tile, and any future parametrization with a larger FP16 length (the FP32 test already uses 65537) would request a >128 KB staging tile. Clamping toSTAGE_CHUNKkeeps both branches bounded.♻️ Suggested clamp
elif dtype_bytes == 2: alignment = 32 // dtype_bytes - stage_cols = ((size + alignment - 1) // alignment) * alignment + stage_cols = min(STAGE_CHUNK, ((size + alignment - 1) // alignment) * alignment)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.py` around lines 94 - 101, Update the FP16 branch in the stage_cols calculation to cap the 16-aligned size at STAGE_CHUNK, matching the existing FP32 path. Preserve the alignment behavior while ensuring stage_cols never exceeds the staging limit.tests/ut/ir/transforms/test_lower_composite_ops.py (2)
1800-1813: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHoist the ad-hoc
CallCollectorto a module-level helper.This exact
IRVisitorsubclass is now redefined inline in at least four tests in this file (868-881, 1800-1807, 1875-1882, plus the mesh partial-valid test). A single module-level_collect_calls(program)helper would remove the duplication and keep future tests from copying it again.♻️ Suggested helper
class _CallCollector(ir.IRVisitor): def __init__(self) -> None: super().__init__() self.calls: list[ir.Call] = [] def visit_call(self, op: ir.Call) -> None: self.calls.append(op) super().visit_call(op) def _collect_calls(program) -> list[ir.Call]: collector = _CallCollector() collector.visit_program(program) return collector.callsAlso applies to: 1875-1885
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/ir/transforms/test_lower_composite_ops.py` around lines 1800 - 1813, Introduce module-level _CallCollector and _collect_calls(program) helpers in test_lower_composite_ops.py, then replace each inline CallCollector definition and traversal—including the tests around lines 868, 1800, 1875, and the mesh partial-valid test—with _collect_calls(program). Preserve each test’s existing call filtering and assertions.
1838-1839: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueName the ring chunk-sizing constants instead of inlining lowering policy.
4096,8192, the+ 7 // 8/+ 15 // 16alignment rounding, and especially the+ 15FP16 tail-slack term are the lowering's internal segmentation policy replicated verbatim in two tests. If the pass retunes any of them, the failure surfaces as an opaque numeric mismatch. Promoting them to named module constants next to_ALLREDUCE_FP32_CHUNK(e.g._RING_FP32_CHUNK,_RING_FP16_CHUNK,_RING_FP16_TAIL_SLACK) documents the coupling and gives one place to update.Also applies to: 1901-1902
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/ir/transforms/test_lower_composite_ops.py` around lines 1838 - 1839, Replace the inlined ring chunk-sizing literals in both test cases around the expected chunk calculations with named module-level constants alongside _ALLREDUCE_FP32_CHUNK, including separate FP32/FP16 chunk sizes and the FP16 tail-slack value. Use those constants for the alignment rounding and tail-slack calculations so the tests document and centralize the lowering policy.tests/ut/codegen/distributed/test_distributed_pto_codegen.py (1)
575-613: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThese constant assertions are weaker than their test names imply.
arith.constant 2/arith.constant 4(bothi64andindex) are ubiquitous in generated MLIR — strides, tuple offsets, shapes. A regression that indexed the wrongCommContextslot or divided by the wrong element width would very likely still find a matching constant elsewhere in the kernel and pass. Since these tests exist specifically to pin the ABI contract, consider anchoring each constant to its consumer, e.g. capture the SSA name defined by thearith.constant {win_idx_unit} : indexline and assert it feeds thepto.load_scalarthat readswindowsIn, and that thedivsiright operand is the element-width constant. Theset(funcs) == {"kernel"}anddivsipresence checks are solid as-is.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/codegen/distributed/test_distributed_pto_codegen.py` around lines 575 - 613, Strengthen test_remote_load_uses_comm_layout_constants by anchoring the rank and windowsIn index constants to their specific SSA consumers instead of matching ubiquitous constant text. Capture the SSA value for the windowsIn constant and verify it is used by the pto.load_scalar reading windowsIn; likewise validate the rank-slot constant’s corresponding load. Also assert that the arith.divsi right operand is the FP16 element-width constant, while preserving the existing function-set and divsi presence checks.tests/ut/codegen/distributed/test_host_orch_distributed.py (1)
275-281: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStale/truncated comment above the aligned-size assertions.
Line 278 ends mid-sentence ("(the inner parens come from the") and Line 279 repeats it; the quoted
window_size=((64 * 4)),also no longer matches the aligned expression now asserted on Line 281.📝 Proposed comment cleanup
- # Single slot → `window_size=((64 * 4)),` (the inner parens come from the - # The inner parentheses come from the Mul expression. + # Single slot → `window_size=((((64 * 4) + 31) // 32) * 32),`. The inner + # parentheses come from the Mul expression.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/codegen/distributed/test_host_orch_distributed.py` around lines 275 - 281, Clean up the comments above the aligned_size assertion: remove the truncated and duplicated sentence, and replace the stale single-slot example with a description matching the aligned expression asserted by the window_size regex. Keep the assertion and its expected formatting unchanged.src/ir/transforms/lower_composite_ops_pass.cpp (2)
1325-1339: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winInactive AG lanes issue a real remote read.
The
elsebranch performs a fullpld.tile.remote_loadfromleft_peerat[0, 0]purely to satisfy theIfExprtype agreement, unlike the reduce-scatter path (Line 1217) which uses a localtile.loadfor the same purpose. On ragged sizes this adds a per-subchunk cross-rank transfer that is immediately discarded. A localtile.loadplaceholder would produce the same TileType without the interconnect traffic.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ir/transforms/lower_composite_ops_pass.cpp` around lines 1325 - 1339, The inactive-AG else branch in the lowering around recv_ag_placeholder should avoid performing a remote read solely for type agreement. Replace the pld.tile.remote_load bound to placeholder_loaded with the local tile.load placeholder pattern used by the reduce-scatter path, preserving the subsequent restore_remote_valid_shape and tile.fillpad_inplace operations.
1276-1301: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
ag_send_idxandag_recv_idxreduce to the same index — collapse or document.
send_idx = (left − step + 1 + NR) % NRwithleft = (my_rank − 1 + NR) % NRis algebraically(my_rank − step) % NR, which is exactlyrecv_idx. Sosend_begin/send_colsandrecv_begin/recv_colsalways coincide, and the four extraBinds plus the separatesend_offsets/recv_offsetstuples are redundant.This matters beyond tidiness: the store at Line 1347 is guarded by
active = subcol < send_colswhile its extent comes fromrecv_valid_cols. That is only safe because the two segment lengths are identical; if the AG index formula is ever intended to differ (segments are ragged, so lengths differ by up to one chunk boundary), the guard would either drop the recv tail or emit a non-positivevalid_colsstore. Either collapse to a single segment computation, or guard the store onrecv_colsand add a comment recording the invariant.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ir/transforms/lower_composite_ops_pass.cpp` around lines 1276 - 1301, The send and receive segment calculations in the AG lowering loop duplicate the same index and create a guard/extent mismatch risk. Collapse ag_send_idx/ag_recv_idx and their associated begin/end/column bindings and offset tuples into one shared segment computation, then use that shared values consistently for the load and store paths; alternatively, guard the store using recv_cols and document the required equality invariant.tests/st/distributed/test_l3_host_tensor_allreduce.py (1)
25-58: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_expected_allreduce/_make_rank_inputsare now duplicated across ST files with diverging signatures.The same two helpers exist in
tests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.py(Lines 42-77), but the dtype kwarg is spelleddtype=here andtorch_dtype=there, and theprod/FP16 input-pattern branch differs. Consider hoisting one copy into a shared ST helper module so the golden reference cannot drift between the mesh, ring, and host suites.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/st/distributed/test_l3_host_tensor_allreduce.py` around lines 25 - 58, Consolidate _expected_allreduce and _make_rank_inputs into a shared ST distributed helper module, then update the mesh, ring, and host allreduce suites to import and use them. Standardize the dtype keyword name and the prod/FP16 input-generation behavior across callers, removing the duplicated local implementations while preserving existing test expectations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/en/dev/distributed_ops.md`:
- Around line 381-382: Update the helper implementation references in both
language versions: docs/en/dev/distributed_ops.md lines 381-382 and
docs/zh-cn/dev/distributed_ops.md lines 330-331 must point to
src/backend/common/pto_ops_distributed.cpp. Keep the EmitCommRemoteOffset and
EmitCommRemoteView descriptions unchanged.
In `@docs/en/dev/passes/39-lower_host_tensor_collectives.md`:
- Around line 60-63: Qualify the 32-byte physical tail-transfer rounding as
applying only to FP16 data in both host allreduce descriptions: update
docs/en/dev/passes/39-lower_host_tensor_collectives.md lines 60-63 and
docs/zh-cn/dev/passes/39-lower_host_tensor_collectives.md lines 60-62 with
equivalent wording, while preserving the existing logical tensor-shape behavior.
In `@python/pypto/language/distributed/op/tensor_ops.py`:
- Line 572: Replace the Unicode minus sign in the signal expression within the
tensor operation documentation with an ASCII hyphen, changing the visible range
to use 2*(NR-1) while preserving the surrounding text.
In `@src/backend/common/pto_ops_distributed.cpp`:
- Around line 209-236: Add distributed rank-3 test coverage for an implicit
column-vector tensor with a constant trailing dimension of 1 and no explicit
TensorView, exercising use_column_vector_convention in the distributed tensor
tests. Verify its generated strides match the existing general DN-layout
behavior.
In `@tests/ut/ir/transforms/test_lower_host_tensor_collectives.py`:
- Around line 459-461: Update
test_host_allreduce_rejects_unsupported_dtype_before_lowering so its
pytest.raises match pattern reflects the current allreduce error wording,
matching “target dtype must be FP16 or FP32” rather than the obsolete FP32-only
text.
---
Outside diff comments:
In `@python/pypto/language/distributed/op/tile_ops.py`:
- Around line 49-101: Keep physical-tail padding lowering-only: remove
allow_physical_tail_padding from the public distributed tile API around
remote_load in python/pypto/language/distributed/op/tile_ops.py (lines 49-101),
and expose it only through a dedicated lowering-only IR constructor/operation in
python/pypto/ir/op/distributed/tile_ops.py (lines 32-70). Update the remote_load
lowering and verification paths in src/ir/op/distributed/remote_load.cpp (lines
146-165 and 206-213) to accept and process the variant only through that
internal operation, preventing user-facing calls from reading logical tail
padding.
---
Nitpick comments:
In `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 1325-1339: The inactive-AG else branch in the lowering around
recv_ag_placeholder should avoid performing a remote read solely for type
agreement. Replace the pld.tile.remote_load bound to placeholder_loaded with the
local tile.load placeholder pattern used by the reduce-scatter path, preserving
the subsequent restore_remote_valid_shape and tile.fillpad_inplace operations.
- Around line 1276-1301: The send and receive segment calculations in the AG
lowering loop duplicate the same index and create a guard/extent mismatch risk.
Collapse ag_send_idx/ag_recv_idx and their associated begin/end/column bindings
and offset tuples into one shared segment computation, then use that shared
values consistently for the load and store paths; alternatively, guard the store
using recv_cols and document the required equality invariant.
In `@tests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.py`:
- Around line 94-101: Update the FP16 branch in the stage_cols calculation to
cap the 16-aligned size at STAGE_CHUNK, matching the existing FP32 path.
Preserve the alignment behavior while ensuring stage_cols never exceeds the
staging limit.
In `@tests/st/distributed/test_l3_host_tensor_allreduce.py`:
- Around line 25-58: Consolidate _expected_allreduce and _make_rank_inputs into
a shared ST distributed helper module, then update the mesh, ring, and host
allreduce suites to import and use them. Standardize the dtype keyword name and
the prod/FP16 input-generation behavior across callers, removing the duplicated
local implementations while preserving existing test expectations.
In `@tests/ut/codegen/distributed/test_distributed_pto_codegen.py`:
- Around line 575-613: Strengthen test_remote_load_uses_comm_layout_constants by
anchoring the rank and windowsIn index constants to their specific SSA consumers
instead of matching ubiquitous constant text. Capture the SSA value for the
windowsIn constant and verify it is used by the pto.load_scalar reading
windowsIn; likewise validate the rank-slot constant’s corresponding load. Also
assert that the arith.divsi right operand is the FP16 element-width constant,
while preserving the existing function-set and divsi presence checks.
In `@tests/ut/codegen/distributed/test_host_orch_distributed.py`:
- Around line 275-281: Clean up the comments above the aligned_size assertion:
remove the truncated and duplicated sentence, and replace the stale single-slot
example with a description matching the aligned expression asserted by the
window_size regex. Keep the assertion and its expected formatting unchanged.
In `@tests/ut/ir/transforms/test_lower_composite_ops.py`:
- Around line 1800-1813: Introduce module-level _CallCollector and
_collect_calls(program) helpers in test_lower_composite_ops.py, then replace
each inline CallCollector definition and traversal—including the tests around
lines 868, 1800, 1875, and the mesh partial-valid test—with
_collect_calls(program). Preserve each test’s existing call filtering and
assertions.
- Around line 1838-1839: Replace the inlined ring chunk-sizing literals in both
test cases around the expected chunk calculations with named module-level
constants alongside _ALLREDUCE_FP32_CHUNK, including separate FP32/FP16 chunk
sizes and the FP16 tail-slack value. Use those constants for the alignment
rounding and tail-slack calculations so the tests document and centralize the
lowering policy.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c11bea55-50cf-4adb-bd5c-08a99bada9ed
📒 Files selected for processing (40)
docs/en/dev/distributed_ops.mddocs/en/dev/passes/12-lower_composite_ops.mddocs/en/dev/passes/39-lower_host_tensor_collectives.mddocs/zh-cn/dev/distributed_ops.mddocs/zh-cn/dev/passes/12-lower_composite_ops.mddocs/zh-cn/dev/passes/39-lower_host_tensor_collectives.mdinclude/pypto/codegen/distributed/comm_layout.hinclude/pypto/codegen/pto/pto_codegen.hinclude/pypto/ir/comm.hpython/bindings/modules/ir.cpppython/pypto/ir/op/distributed/tile_ops.pypython/pypto/language/distributed/op/system_ops.pypython/pypto/language/distributed/op/tensor_ops.pypython/pypto/language/distributed/op/tile_ops.pypython/pypto/pypto_core/ir.pyipython/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.inpython/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.insrc/backend/common/pto_ops_distributed.cppsrc/codegen/distributed/distributed_codegen.cppsrc/codegen/distributed/distributed_ops_codegen.cppsrc/codegen/pto/pto_codegen.cppsrc/ir/op/distributed/allreduce.cppsrc/ir/op/distributed/collective.cppsrc/ir/op/distributed/get.cppsrc/ir/op/distributed/remote_load.cppsrc/ir/op/distributed/remote_store.cppsrc/ir/op/distributed/system.cppsrc/ir/transforms/flatten_tile_nd_to_2d/rewrite.cppsrc/ir/transforms/lower_composite_ops_pass.cpptests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_barrier_intrinsic.pytests/st/distributed/test_l3_host_tensor_allreduce.pytests/st/runtime/ops/test_fillpad.pytests/ut/codegen/distributed/test_comm_layout.pytests/ut/codegen/distributed/test_distributed_pto_codegen.pytests/ut/codegen/distributed/test_host_orch_distributed.pytests/ut/ir/test_distributed_ops.pytests/ut/ir/transforms/test_lower_composite_ops.pytests/ut/ir/transforms/test_lower_host_tensor_collectives.py
💤 Files with no reviewable changes (3)
- tests/st/distributed/collectives/test_l3_tensor_barrier_intrinsic.py
- src/codegen/pto/pto_codegen.cpp
- include/pypto/codegen/pto/pto_codegen.h
5e35f0d to
02599b0
Compare
|
Addressed the actionable automated-review feedback in 02599b0. Implemented:
One suggestion was adjusted rather than applied literally: the host builtin rounds ragged load spans for both FP16 and FP32, while only the InCore mesh/ring padded remote-tail path is FP16-specific. The docs now state that distinction. I did not apply the low-value CallCollector/chunk-constant/shared-ST-helper consolidation suggestions because they are unrelated refactors with no correctness impact and would unnecessarily broaden this PR. Two independent full-diff reviewers rechecked the final revision against upstream/main and both returned PASS. Remote CI validation is intentionally left to the author. |
24fa312 to
36177ab
Compare
44c8f3f to
fcc90a9
Compare
fcc90a9 to
08ff10c
Compare
2 participants
Summary
SIZE % NR == 0contract to arbitrary positive element counts, includingSIZE < NR, ragged segments, and inputs larger than UB.Sum,Max,Min, andProdfor FP16/FP32 across mesh, ring, and generated host-allreduce variants.Implementation
[1, N]stream and process them in physical chunks of at most 16 KiB.valid_shape, so short or non-aligned lengths do not over-read or over-store.Validation
git diff --check upstream/main...HEADgit range-diffreports both rebased commits as patch-equivalent to the reviewed and validated revisions.507018was observed.Current constraints
for/whileloops remain rejected.Follow-up to #2084.