feat: 支持 tput/tget 动态分区 shape - #1079
Conversation
b55cadb to
eff6d19
Compare
HecreReed
changed the title
![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: 55285f0936
ℹ️ 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".
| - Runtime extents must be nonnegative, and each partition range must remain | ||
| within its backing tensor view. A zero extent denotes an empty transfer. |
There was a problem hiding this comment.
When an empty transfer is statically known, the documented zero-length behavior cannot be expressed through the supported pto.partition_view path: PartitionViewOp::verify() still rejects any constant size <= 0 as invalid. Thus a peer count supplied as a runtime SSA value may become zero and work, while the equivalent %c0 extent—or an extent constant-folded before verification—fails before reaching tput/tget. Either allow constant-zero partition sizes for this path or do not claim that zero extents are supported.
Useful? React with 👍 / 👎.
ba13c77 to
d2b1c70
Compare
zhangstevenunity
left a comment
zhangstevenunity
left a comment
There was a problem hiding this comment.
Review: dynamic partition shapes for pto.comm.tput / pto.comm.tget
Built this branch (d2b1c707) against LLVM 21 and ran the suite locally: check-pto = 1538 passed / 1 unsupported / 0 failed, all 11 comm lit tests green. I also read the emitted C++ by hand and the pinned pto-isa runtime (016396b5, the ci_sim.yml pin).
The core change is correctly scoped and I found no P1. Two P2s below.
What I verified as sound
- The policy default keeps every other comm op strict. Signals (
verifyCommSignalLike), collective group operands (verifyCommGlobalGroup) andtbroadcast/tgather/tscatter/treduceall still getStaticOnly. Async p2p never used this helper at all -- it goes throughverifyAsyncFlatContiguous1DGMViewLike, which is why the new async negative test correctly expects a different diagnostic string. The three new scope tests pin exactly that boundary. - Nothing outside the EmitC path reads these operands' shapes.
SyncMacroModelextracts only Values plus a lane count;PTOPlanMemoryuses them as liveness handles and never reachesGenerateBufferInfo;InsertSync'sgetStaticBufferSizeInBytesreturns0onkDynamicandMemoryDependentAnalyzertreats size 0 as may-alias;GraphSyncSolver'sMemInfoyields thekDynamicsentinel andcheckConflictfalls through toreturn true. All conservative -- no dropped hazard.ptobcinterns printed type strings and stores no dim immediates, so the round-trip is shape-agnostic by construction. - Lowering is right for the tested shape.
PTOPartitionViewStaticToEmitCstill wins atPatternBenefit(2)for static partitions; dynamic ones fall through toPTOPartitionViewToEmitC, which feedsadaptor.getSizes()intobuildRuntimeGlobalTensor. The emitted descriptor ispto::Shape<1, 1, 1, -1, 4096>(v6, v6, v6, v8, v4)withv8 = (int64_t) %runtime_rows-- correct.pto::Shape's 5-argument constructor carries nostatic_assert(unlike the 1..4-argument ones) and only stores the DYNAMIC slots, so passing all five is well formed. - Runtime side is right for
<?x4096xi8>.TPUT_IMPLreadsGetShape(DIM_3)at run time, early-returns on a zero extent, and the chunked path clampscurRowsand updatesRowMaskInternalbecausev_row=?.include/pto/comm/a5/TPut.hppis a one-line#includeof the a2a3 header, so there is no arch divergence here.
P2-1: the new EmitC test cannot fail on the property this PR exists to deliver
Mutation-tested. Changing PTOPartitionViewToEmitC (lib/PTO/Transforms/PTOToEmitC.cpp:8246) to pass adaptor.getOffsets() instead of adaptor.getSizes() makes every dynamic transfer describe zero rows instead of %rows:
correct: pto::Shape<1, 1, 1, -1, 4096>(v6, v6, v6, v8, v4) // v8 = (int64_t) %runtime_rows
mutated: pto::Shape<1, 1, 1, -1, 4096>(v6, v6, v6, v7, v7) // v7 = 0
After that rebuild comm_p2p_dynamic_partition_emitc.pto still passes. Fix suggested inline.
P2-2: ? is accepted on the innermost dimension, where pto-isa takes the length from the staging tile instead of the tensor
Details inline on PTO.cpp. Short version: buildRuntimeGlobalTensor right-aligns to rank 5, so the innermost dim always lands in DIM_4, and the single-shot path in TPUT_IMPL computes lenBurst from the tile's validCol rather than from gShape4. A rank-1 !pto.partition_tensor_view<?xf32> -- the obvious extension of this manual's own <128xf32> example -- over-writes the remote peer whenever the runtime extent is smaller than v_col.
Non-blocking notes
- Codex's zero-extent P2 reproduces; reproduction inline on the manual.
- Runtime extent equality between
dstandsrcis unverifiable and unverified.getShapeVec(dst) != getShapeVec(src)compares?against?structurally, andTPUT_IMPLdrives everything offsrc, so a shorterdstis a remote OOB write. The manual states the requirement; nothing enforces it. Same class as P2-2. - No coverage compiles or runs the generated kernel.
remote-npu-validationisskippingon this PR and the description confirms there is no CANN environment, so the whole verification story is FileCheck over emitted text plus a ptobc round-trip. For a change whose entire payload is a runtime descriptor, one executed case would be worth more than the six new text tests. The CPU simulator will not help with chunking --pto/cpu/comm/TPut.hppmodels TPUT as a shape-drivenCopy_Data(src, dst)and ignores the staging tile entirely -- but it would at least prove the emitted C++ compiles and movesnrows rather than 32.
| if (dim == ShapedType::kDynamic || dim <= 0) | ||
| return op->emitOpError() << "expects " << name | ||
| << " to have a positive static shape"; | ||
| if (dim == ShapedType::kDynamic) { |
There was a problem hiding this comment.
kDynamic is accepted in every position here. buildRuntimeGlobalTensor (lib/PTO/Transforms/PTOToEmitC.cpp:7970) right-aligns the shape into rank 5, so whichever dimension is innermost always lands in DIM_4. I checked what this branch actually emits at --pto-level=level3:
!pto.partition_tensor_view<?xf32> -> pto::Shape<1, 1, 1, 1, -1>
!pto.partition_tensor_view<4x?xf32> -> pto::Shape<1, 1, 1, 4, -1>
!pto.partition_tensor_view<?x4096xi8> -> pto::Shape<1, 1, 1, -1, 4096> <- the only shape this PR tests
Only the last one is safe, because pto-isa treats DIM_3 and DIM_4 very differently. In TPUT_IMPL (pinned 016396b5, include/pto/comm/a2a3/TPut.hpp) the single-shot path is taken when
if (totalLogicalRows <= ubChunkRows && logicalDims[4] <= ubChunkCols) { TLOAD; ...; TSTORE_IMPL; return; }and it hands the tensors straight to TLOAD/TSTORE. Both TLoadGm2ubNd2nd (include/pto/npu/a2a3/TLoad.hpp:66) and TStoreUb2gmNd2nd (include/pto/npu/a2a3/TStore.hpp:31) then do:
PTO_ASSERT(validCol == gShape4, "The validCol of TileData must be equal to the 5th dim(Shape4) of ND shape!");
uint16_t nBurst = gShape3; // rows: from the tensor
uint32_t lenBurst = validCol * sizeof(typename TileData::DType); // columns: from the TILEPTO_ASSERT is ((void)0) unless _DEBUG (include/pto/common/debug.h:33-35), so on a release device build that guard is gone.
Failure scenario. !pto.partition_tensor_view<?xf32> with a rows=1, cols=128, v_row=1, v_col=128 staging tile -- i.e. this manual's own <128xf32> tput example with the extent made dynamic. At runtime n = 100: logicalDims = (1,1,1,1,100), 100 <= 128 so the single-shot path runs, lenBurst = 128 * 4, and TSTORE writes 128 floats into the peer's 100-float window. 28 floats of silent remote out-of-bounds write. n > 128 is no better with a static v_col: TputProcessSlice clamps curCols for the tail chunk but only assigns ColMaskInternal when ValidCol == DYNAMIC, so the tail transfers v_col columns regardless.
The row dimension has none of this, because nBurst comes from gShape3.
I compiled both of those shapes against this branch; they lower clean and emit pto::comm::TPUT(...), so nothing downstream catches it.
To be fair this contract is pre-existing and unchecked for static shapes too -- TPutOp::verify never relates the staging tile's valid dims to the transfer extent. What changes is that for a ? dimension it becomes permanently uncheckable, by anyone.
Suggestion: restrict the relaxation to non-innermost dimensions. That covers every motivating case (pypto emits <?x64xf16>, #1069's repro is <?x4096xi8>) and it is a small addition here -- reject kDynamic when the index is shape.size() - 1. If you would rather keep it open, the manual should state plainly that a dynamic innermost extent requires v_col to equal it exactly at run time.
| } | ||
|
|
||
| // EMITC: pto::Shape<1, 1, 1, -1, 4096> | ||
| // EMITC-SAME: ({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^)]+}}) |
There was a problem hiding this comment.
Five {{[^,]+}} groups match whatever the compiler passes, so nothing here ties the descriptor's runtime extent to the partition size -- the one property this PR adds.
I proved that by mutation. In PTOPartitionViewToEmitC (lib/PTO/Transforms/PTOToEmitC.cpp:8246) I changed adaptor.getSizes() to adaptor.getOffsets(), so every dynamic transfer describes 0 rows instead of %rows:
correct: pto::Shape<1, 1, 1, -1, 4096> v22 = pto::Shape<1, 1, 1, -1, 4096>(v6, v6, v6, v8, v4); // v8 = (int64_t) v3
mutated: pto::Shape<1, 1, 1, -1, 4096> v21 = pto::Shape<1, 1, 1, -1, 4096>(v6, v6, v6, v7, v7); // v7 = 0
Rebuilt and re-ran this file: still passes. A compiler that silently moved nothing at all would ship green.
The NATIVE block does pin sizes = [%[[ROWS]], %c4096], but that only proves the MLIR still carries the value -- it says nothing about which value reaches the emitted pto::Shape. Capturing the cast closes the gap:
// EMITC: int64_t [[N:v[0-9]+]] = (int64_t) v3;
// EMITC: pto::Shape<1, 1, 1, -1, 4096>({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, [[N]], {{[^)]+}})
Separately: both the EMITC and NATIVE blocks are satisfied entirely by @comm_p2p_dynamic_partition. @comm_p2p_dynamic_partition_static_valid has no checks of its own, so today it only asserts "does not error". If the static-v_row case is meant to be pinned, it needs its own CHECK-LABEL-anchored block.
| within its backing tensor view. A zero extent denotes an empty transfer. | ||
| - `ping` / `pong` must be local VEC tile-like values whose element type matches | ||
| `src`. Their physical `rows` / `cols` must be positive static values. | ||
| - Staging `v_row` / `v_col` values must be positive and may be static or |
There was a problem hiding this comment.
This bullet documents only half of the runtime contract, and a program that follows it exactly can still produce the out-of-bounds write I describe in the note on lib/PTO/IR/PTO.cpp.
The divisibility rule stated here applies to the chunked path. But TPUT_IMPL only chunks when the transfer does not fit the tile (pinned pto-isa 016396b5, include/pto/comm/a2a3/TPut.hpp):
if (totalLogicalRows <= ubChunkRows && logicalDims[4] <= ubChunkCols) {
TLOAD(stagingTileData, srcGlobalData); ...; TSTORE_IMPL(dstGlobalData, stagingTileData); return;
}The single-shot branch requires something stricter that is written down nowhere: v_col must equal the column extent and v_row must equal the total row extent, because TLoadGm2ubNd2nd / TStoreUb2gmNd2nd derive lenBurst from validCol and assert validCol == gShape4 and validRow == gShape0 * gShape1 * gShape2 * gShape3. Those are PTO_ASSERTs, i.e. ((void)0) off _DEBUG.
So a transfer with a dynamic extent of 100 and v_col = 128 never enters the chunked path, satisfies every rule on this page, and still over-transfers. Worth spelling the single-shot requirement out here.
| `pto.partition_tensor_view` values. | ||
| - `dst` and `src` must have the same element type and identical static/dynamic | ||
| shape signatures. Corresponding dynamic extents must be equal at runtime. | ||
| - Runtime extents must be nonnegative, and each partition range must remain |
There was a problem hiding this comment.
Confirming Codex's P2 with a reproduction on this branch -- a constant zero extent is rejected before it can reach tput:
%dst = pto.partition_view %dst_view, offsets = [%c0, %c0], sizes = [%c0, %c4096]
: !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>$ ptoas --pto-arch=a3 --pto-level=level3 zero_extent.pto -o /dev/null
error: 'pto.partition_view' op size at dim 0 must be positive, got 0
Error: Failed to parse MLIR.
PartitionViewOp::verify (lib/PTO/IR/PTO.cpp:2531) rejects any constant size <= 0, so only a runtime SSA value that happens to be zero ever reaches TPUT_IMPL's totalLogicalRows == 0 early return.
The corollary worries me more than the wording: getConstIndexValue sees through folded constants, so a program that is legal at parse time hard-fails partition_view verification if a later pass constant-folds the size operand to 0 -- for example after inlining a caller that specializes the peer count to zero. Either relax the check to < 0 when the corresponding result dim is dynamic, or drop the zero-extent claim from the manual.
A5 板测成功
|
A3 板测失败
失败用例
|
A3 板测失败详情：PR #1079comm_p2p_binding_variants
comm_p2p
|
…nCore composite (#2280) ## Summary Replaces the **pull-model** engine of the ring allreduce on **both rails** with a **TPUT push model** (remote write), enabling O(1) `NeighborBarrier` on the HOST builtin and eliminating the pull-model NPU memory-ordering gap. - **HOST builtin** (`builtin.tensor.allreduce_ring`): reduce-scatter + allgather converted from `TLOAD`/`TSTORE` pull to `pto::comm::TPUT` push — `TPUT<AtomicAdd>` remote-accumulate for RS, non-atomic `TPUT` for AG. Ordering is `pipe_barrier(PIPE_ALL)` around every transfer + `dsb(DSB_DDR)` before `TNOTIFY` (mirrors the in-tree allgather/all_to_all host builtins; not a GM fence). The O(P²) `RoundBarrier` is replaced by the O(1) `NeighborBarrier` (notify/wait the two ring neighbours only), which is NPU-safe because the TPUT write pipeline orders the data ahead of the signal — the pull model could not provide that. - **InCore composite** (`LowerTensorRingAllReduceRule`): replaces `pld.tile.remote_load` pulls with `pld.tile.put` pushes (non-atomic TPUT + local reduce, **preserving Sum/Max/Min/Prod**). Race-free per-subchunk protocol: own-value read → ready barrier → push to right neighbour → push-done barrier → local read+reduce+store; barrier credits stay 2 per subchunk (signal shape unchanged). Ragged/arbitrary lengths and FP16 are preserved via balanced segments + valid shapes, with the shared VEC staging tile narrowed per transfer via `tile.set_validshape`. ## Requires PTOAS >= v0.55 (pypto pins v0.57) **This PR depends on [PTOAS v0.55](https://github.com/hw-native-sys/PTOAS/releases/tag/v0.55)** (release: [hw-native-sys/PTOAS#1069](hw-native-sys/PTOAS#1069), fixed in [PR #1079](hw-native-sys/PTOAS#1079)). The InCore composite's `pld.tile.put` transfers carry the **exact ragged `valid_cols`** as the partition-view extent. PTOAS ≤ v0.54 rejects dynamic partition-view shapes for `pto.comm.tput` (`'pto.comm.tput' op expects dst to have a positive static shape`), so the pure push model cannot compile below v0.55. The HOST builtin does not depend on this (its kernel is hand-written), but the composite rail does. The requirement is satisfied by the current pin: pypto now pins **PTOAS v0.57** (via #2291). The PR is rebased onto current `main` (2026-08-26, was 76 commits behind; re-rebased twice 2026-08-27 — first onto the #2530 runtime bump adopting the `ChipTensor`→`TaskTensor` kernel rename, then onto #2542 adopting the `42-lower_host_tensor_collectives` → `43-…` docs rename, with the PR's ring-doc edits re-homed) and merges cleanly. The UT tests pin the push structure (`pld.tile.put` + `tile.create` staging tile instead of `pld.tile.remote_load`). ## Rebased — merge-order with #2279 (self-clearing signals) resolved The rebase picked up #2279's self-clearing signal epilogue, which was written for the **pull-model `RoundBarrier`** (reset every peer's cell with `TNOTIFY(-1)` per round). That credit pattern does **not** match the push model's `NeighborBarrier`: - `NeighborBarrier` credits only the **two ring neighbours** per round — a single cell when `nranks == 2`, where both neighbours are the same peer and the cell carries two +1s. - The #2279 loop would corrupt the unused cells to −1 and, for `nranks == 2`, leave +1 stale credit in the one used cell — reintroducing the exact stale-credit barrier failure #2279 fixed. The epilogue now branches on `kUseNeighborBarrier`: it restores only the two neighbour cells per used row with `TNOTIFY(-1, AtomicAdd)` (twice on the shared cell when `nranks == 2`), keeping the all-peer reset for the `RoundBarrier` fallback. The ring builtin is therefore **self-clearing and signal-reuse-safe** across back-to-back calls, matching the other host builtins (#2279). The ring signal-reuse ST (`test_l3_host_tensor_allreduce_ring.py` reuse leg) is the NPU gate for the adapted epilogue. ## Issues this PR addresses - **#2242 (ring unaligned-data handling)**: the pull-model dcci-flush tail gap (item 1) is **moot** — the push model needs no cacheline flush (the receiver reads data the sender wrote remotely via TPUT, never a locally-TSTORE'd line). The 32-byte transfer-alignment concern (item 2) is handled by narrowing the staging tile's column mask (`ColMaskInternal` / `tile.set_validshape`) to the exact (possibly ragged) transfer extent, so partial tails transfer exactly and never over-read/overwrite adjacent slots. - **#2213 (PTOAS dynamic partition-view)**: closed as superseded by #2524; the `>= v0.55` dependency it describes is satisfied by the v0.57 pin. ## Verification (NPU silicon, 910B2, PTOAS v0.55) All on real NPUs (8x 910B2), P=2 and P=4: - `tests/st/distributed/test_l3_host_tensor_allreduce_ring.py` — HOST ring, P=2/4 ✅ (with `NeighborBarrier` enabled) - `tests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.py` — InCore ring, P=2/4, sizes {1, 17, 4097, 65537} (ragged + >UB), Sum/Max/Min/Prod, FP16 ✅ - `tests/st/distributed/collectives/test_l3_allreduce_ring.py` + `test_l3_ring_sizing_prewarm.py` — no regression ✅ - UTs: `test_lower_host_tensor_collectives.py`, `test_host_orch_distributed.py`, `test_lower_composite_ops.py` (+ numerical) all green ✅ **Total: 25/25 ST + 190 UT passed** (pre-rebase). The 2026-08-26 rebase + epilogue adaptation re-ran the ring UTs (173/174, the one failure is a pre-existing parser `TileView(pad=…)` roundtrip gap on main, unrelated to this PR); the 2026-08-27 re-rebases (onto the #2530 runtime bump and the #2542 docs rename) each re-ran the same 173/174. NPU ST should be re-confirmed for the signal-reuse leg. ## Trade-off note (ReduceOp) The HOST builtin is `ReduceOp::kSum` only by construction, so its `TPUT<AtomicAdd>` RS is fine. The composite keeps non-atomic push + local reduce to preserve Sum/Max/Min/Prod; only a remote-atomic `TPUT<AtomicAdd>` variant would be Sum-only (`AtomicType` has no `AtomicMax/Min`). ## Follow-ups (not in this PR) - **#2310** — lift the HOST-rail in-loop `pld.tensor.allreduce` restriction via shared-signal synthesis (the other half of the #2279 review). - **TPUT_ASYNC** (pto-isa) for an overlapped / IBing forward phase — optional perf follow-on (simpler #1383 / plan 50). ## Review notes - Addresses CodeRabbit feedback: the allgather ready-barrier rationale is corrected (counters are per-round; the real guarantee is the previous round's push-done barrier every rank passes before round k), and the `nranks == 2` `NeighborBarrier` behaviour is documented.
3 participants
当前状态
设计和实现已完成，并已合并最新
main（71894b5db8df35bc41436605890100380e26af8d）解决冲突。PR 已迁移到当前 tile-native 通信链路，不再依赖已删除的PTOViewToMemref/ decoded-memref bypass。实现内容
verifyCommGlobalLike()增加显式CommGlobalShapePolicy：默认继续要求静态 shape，仅同步pto.comm.tput/tget的dst/src使用AllowDynamicPartitionView。!pto.partition_tensor_view开放；直接动态tensor_view仍由 verifier 拒绝，静态维仍必须大于0。dst/src继续要求元素类型相同、静态/动态 shape 签名完全相同，并更新对应诊断。rows/cols均保持原有静态约束。PTOPartitionViewToEmitC -> buildRuntimeGlobalTensor -> PTOP2PCommToEmitC链路，不修改 P2P lowering、ptobc schema 或 pto-isa。合并 main 后的迁移
PTOCommType已采用 tile-nativetensor_view/partition_tensor_view/tile_bufoperand；旧的直接动态 memref 负例已移除。pto-view-to-memref改为pto-resolve-buffer-select，检查动态 size 保留在pto.partition_view，且中间 IR 不退回memref.subview。Shape<1, 1, 1, -1, 4096>配合Stride<-1, -1, -1, -1, -1>的五个运行时参数。tput/tgetoperand 顺序和可选 pong。测试覆盖
partition_view<?x4096xi8>的tput/tget正向 EmitC 和 tile-native pass seam。tensor_view、非正数静态维。atomic_add，并把解码结果重新交给 PTOAS。验证结果
CCACHE_DISABLE=1 ninja -C build PTOASPythonPackage ptobc：通过。ctest --test-dir build -R '^ptobc_comm_p2p_dynamic_v0_encode$' --output-on-failure：1/1 passed。ninja -C build check-pto：1539 tests，1538 passed，1 unsupported，0 failed。git diff --check：通过。main的 merge-tree：无冲突。当前没有匹配的 CANN 设备编译环境，因此生成 C++ 的设备侧仅编译验证仍作为有环境时的补充验证项；EmitC 的 runtime
Shape/Stride/GlobalTensor/TPUT/TGET形态已由 lit 固定。Closes #1069