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
|
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