feat(distributed): add multicore HOST AllReduce - #2160
Conversation
📝 WalkthroughWalkthroughThe PR adds a keyword-only ChangesHOST AllReduce multicore support
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Suggested labels: 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: cc80601c8b
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| CHECK(pypto::backend::BackendConfig::IsConfigured()) | ||
| << "builtin.tensor.allreduce codegen requires a configured backend"; | ||
| const auto* backend = pypto::backend::GetBackend(); | ||
| const int max_block_dim = backend->GetCoreCount(ir::CoreType::CUBE); |
There was a problem hiding this comment.
The generated builtin is submitted with rt_submit_aiv_task, so its maximum launch width must come from CoreType::VECTOR, not CoreType::CUBE. The configured 910B model has 48 vector cores versus 24 cube cores, and 950 has 36 versus 18, so valid requests such as core_num=32 are incorrectly rejected even though the target device has enough AIV cores.
Useful? React with 👍 / 👎.
| for loop in _collect_for_stmts(host.body) | ||
| if isinstance(loop.body, ir.EvalStmt) | ||
| and isinstance(loop.body.expr, ir.Call) | ||
| and loop.body.expr.op.name == "builtin.tensor.allreduce" |
There was a problem hiding this comment.
Replace this raw .op.name literal comparison, and the equivalent new comparisons in the signal-synthesis test, with ir.get_op("...").name. A typo or operator rename currently makes the comprehension silently select no nodes and fail later with an ambiguous assertion instead of reporting the invalid operator at the comparison site, contrary to the repository's required operator-identity convention. .claude/rules/operator-identity-checks.mdL5-L13
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/proposals/host_allreduce_spmd_multicore.md`:
- Around line 5-6: Update the proposal status to reflect that multicore HOST
AllReduce is implemented and validated, replacing the outdated unimplemented
wording in docs/en/dev/proposals/host_allreduce_spmd_multicore.md lines 5-6 and
the matching Chinese wording in
docs/zh-cn/dev/proposals/host_allreduce_spmd_multicore.md line 5; mark both
documents as implemented/historical or describe only remaining limitations.
In `@src/codegen/distributed/distributed_ops_codegen.cpp`:
- Around line 227-230: Update the core_num capacity check in the
builtin.tensor.allreduce generation path to use
backend->GetCoreCount(ir::CoreType::VECTOR), matching the AIV launch capacity
used by launch_core_count_method (set_block_num/set_core_num), while preserving
the existing validation and error message behavior.
🪄 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: 6df44fe7-0c7f-4507-adf5-fb27864e075d
📒 Files selected for processing (24)
docs/en/dev/distributed_ops.mddocs/en/dev/passes/37-synthesize_allreduce_signals.mddocs/en/dev/passes/39-lower_host_tensor_collectives.mddocs/en/dev/proposals/host_allreduce_spmd_multicore.mddocs/zh-cn/dev/distributed_ops.mddocs/zh-cn/dev/passes/37-synthesize_allreduce_signals.mddocs/zh-cn/dev/passes/39-lower_host_tensor_collectives.mddocs/zh-cn/dev/proposals/host_allreduce_spmd_multicore.mdpython/pypto/ir/op/distributed/tensor_ops.pypython/pypto/language/distributed/op/tensor_ops.pypython/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.inpython/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.insrc/codegen/distributed/distributed_ops_codegen.cppsrc/ir/op/distributed/allreduce.cppsrc/ir/op/distributed/collective.cppsrc/ir/transforms/lower_composite_ops_pass.cppsrc/ir/transforms/lower_host_tensor_collectives_pass.cppsrc/ir/transforms/synthesize_allreduce_signals_pass.cpptests/st/distributed/test_l3_host_tensor_allreduce_multicore.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.pytests/ut/ir/transforms/test_materialize_comm_domain_scopes.py
| Proposed design for the first hardware-validation implementation. This document | ||
| does not describe an implemented feature yet. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the proposal status to reflect the implemented feature.
These statements say multicore HOST AllReduce is not implemented, while this PR implements and validates it. Mark the proposals as implemented/historical, or revise the wording to identify only remaining limitations.
docs/en/dev/proposals/host_allreduce_spmd_multicore.md#L5-L6: replace the “does not describe an implemented feature yet” status.docs/zh-cn/dev/proposals/host_allreduce_spmd_multicore.md#L5-L5: replace “目前还没有实现对应功能” with matching implemented/historical status.
📍 Affects 2 files
docs/en/dev/proposals/host_allreduce_spmd_multicore.md#L5-L6(this comment)docs/zh-cn/dev/proposals/host_allreduce_spmd_multicore.md#L5-L5
🤖 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 `@docs/en/dev/proposals/host_allreduce_spmd_multicore.md` around lines 5 - 6,
Update the proposal status to reflect that multicore HOST AllReduce is
implemented and validated, replacing the outdated unimplemented wording in
docs/en/dev/proposals/host_allreduce_spmd_multicore.md lines 5-6 and the
matching Chinese wording in
docs/zh-cn/dev/proposals/host_allreduce_spmd_multicore.md line 5; mark both
documents as implemented/historical or describe only remaining limitations.
| const int max_block_dim = backend->GetCoreCount(ir::CoreType::CUBE); | ||
| CHECK(core_num <= max_block_dim) << "builtin.tensor.allreduce core_num (" << core_num | ||
| << ") exceeds the backend logical block capacity (" << max_block_dim | ||
| << ")"; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C3 'GetCoreCount\(|GetLaunchSpecCoreCountMethod|set_block_num|set_core_num|CoreType::(CUBE|VECTOR)' include src testsRepository: hw-native-sys/pypto
Length of output: 29309
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== distributed allreduce relevant section =="
sed -n '190,245p' src/codegen/distributed/distributed_ops_codegen.cpp
echo "== backend core count implementation =="
sed -n '365,405p' src/backend/common/backend.cpp
sed -n '110,190p' src/backend/common/soc.cpp
echo "== handler counts and set methods =="
rg -n -C4 'GetLaunchSpecCoreCountMethod|set_block_num|set_core_num|block_num|core_num|available_cluster_count|available_aiv_count' include/pypto backend templates src tests/ut/codegen/distributed src/codegen/distributed | head -n 220
echo "== locate allreduce templates =="
fd -a 'allreduce|tensor.allreduce|builtin_tensor_allreduce|collective' | sed 's#^\./##' | head -n 200
echo "== allreduce template relevant files =="
python3 - <<'PY'
from pathlib import Path
for p in Path('.').rglob('*'):
if p.is_file() and 'allreduce' in p.name.lower().replace('_', '').replace('-', ''):
print(p)
PYRepository: hw-native-sys/pypto
Length of output: 29702
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate allreduce templates =="
git ls-files | rg -n -i 'allreduce|collectives|builtins' | head -n 200
echo "== allreduce template files =="
python3 - <<'PY'
from pathlib import Path
matches = []
base = '.'
for p in Path(base).rglob('*'):
if not p.is_file():
continue
s = str(p)
if 'allreduce' in s.lower() or ('collectives' in s.lower() and 'allreduce' in p.name.lower()):
matches.append(s)
matches.sort()
for m in matches:
print(m)
PY
echo "== allreduce template content (if any) =="
python3 - <<'PY'
from pathlib import Path
for p in Path('.').rglob('*'):
if not p.is_file():
continue
if 'allreduce' in p.name.lower() or ('collectives' in p.parent.name.lower() and p.exists()):
try:
text = p.read_text(errors='ignore')
except Exception as e:
print(f"UNREADABLE {p}: {e}")
continue
if 'allreduce' in text.lower() or 'submit_allreduce_kernel' in text or 'core_count' in text:
print(f"\n== {p} ==")
print(text[:6000])
PYRepository: hw-native-sys/pypto
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
for f in \
python/pypto/runtime/builtins/collectives/allreduce/templates/entry.cpp.in \
python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in \
python/pypto/runtime/builtins/collectives/allreduce/templates/kernel_config.py.in \
tests/ut/codegen/distributed/test_host_orch_distributed.py \
tests/st/distributed/test_l3_host_tensor_allreduce_multicore.py
do
echo "== ${f##*/} =="
sed -n '1,260p' "$f"
done
echo "== backend handlers =="
sed -n '1,70p' include/pypto/backend/910B/backend_910b_handler.h
sed -n '1,70p' include/pypto/backend/950/backend_950_handler.hRepository: hw-native-sys/pypto
Length of output: 37292
Use VECTOR core count for AIV launch capacity.
builtin.tensor.allreduce renders AIV dispatch with launch_core_count_method (set_block_num / set_core_num), while this check caps core_num against CoreType::CUBE. 950 has fewer AIV/available cluster cores than AIC cores in this config, so 950 can accept values it cannot launch and may reject valid valid 910B AIV launch widths; base the check on the same VECTOR runtime capacity used for the AIV allreduce grid.
🤖 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/codegen/distributed/distributed_ops_codegen.cpp` around lines 227 - 230,
Update the core_num capacity check in the builtin.tensor.allreduce generation
path to use backend->GetCoreCount(ir::CoreType::VECTOR), matching the AIV launch
capacity used by launch_core_count_method (set_block_num/set_core_num), while
preserving the existing validation and error message behavior.
1 participant
Summary
core_numoption to HOST tensor AllReduce while preserving the single-core defaultTesting
ruff check .ruff format --check .pyright262 passed7801 passed, 2 skipped3 passedon A2/A3 (2x2,2x4, and4x2rank/core cases)