Skip to content

feat: 支持 tput/tget 动态分区 shape #1079

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhangstevenunity merged 1 commit into from
Aug 3, 2026
Merged

feat: 支持 tput/tget 动态分区 shape #1079

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 61 additions & 7 deletions docs/PTO_IR_manual.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10201,7 +10201,7 @@ This section documents PTO communication primitives. PTOAS currently exposes:

##### `pto.comm.tput` - Synchronous Remote Write

**Summary:** Lowers to `pto::comm::TPUT(...)` and copies data from local GM to remote GM through a VEC staging tile.
**Summary:** Copies data synchronously from local GM to remote GM through a VEC staging tile.

**Arguments:**

Expand All @@ -10214,9 +10214,29 @@ This section documents PTO communication primitives. PTOAS currently exposes:

**Constraints & Verification:**

- `dst` / `src` must be GM-shaped values with positive static shapes.
- `dst` and `src` must have the same element type and static shape.
- `ping` / `pong` must be local VEC tile-like values whose element type matches `src`.
- `dst` / `src` must be GM-shaped values. Static dimensions must be positive.
- Dynamic dimensions are supported only when both operands are
`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

@zhangstevenunity zhangstevenunity Aug 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

within its backing tensor view. A zero extent denotes an empty transfer.
Comment on lines +10222 to +10223

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Permit constant-zero partition extents

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 👍 / 👎.

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

@zhangstevenunity zhangstevenunity Aug 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

dynamic. If a transfer enters the chunked path, a static `v_row` / `v_col`
must exactly divide the corresponding logical row / column extent. Use
dynamic valid dimensions when a partial final chunk is possible.
- When `pong` is present, it must have the same type as `ping`.

**Semantics:**

For every logical index in the common `src` / `dst` shape, the operation reads
the local `src` element and writes the corresponding remote `dst` element.
`atomic_none` performs a normal write; an atomic mode such as `atomic_add`
combines the source value with the destination according to that mode. The
staging bundle may divide the logical range into chunks but does not change the
logical transfer extent.

**Examples:**

Expand All @@ -10228,11 +10248,36 @@ pto.comm.tput(%dst, %src, buf(%ping) : !pto.partition_tensor_view<128xf32>, !pto
pto.comm.tput(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>, !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {atomicType = #pto<atomic_type atomic_add>}
```

For a variable-size transfer, construct both partitions with the same runtime
extent and use dynamic staging valid dimensions when the extent may require a
partial final chunk:

```mlir
%dst_part = pto.partition_view %dst_view,
offsets = [%c0, %c0], sizes = [%rows, %c4096]
: !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
%src_part = pto.partition_view %src_view,
offsets = [%c0, %c0], sizes = [%rows, %c4096]
: !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
%stage = pto.alloc_tile addr = %c0_i64
valid_row = %c1 valid_col = %c4096
: !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096,
v_row=?, v_col=?, blayout=row_major,
slayout=none_box, fractal=512, pad=0>
pto.comm.tput(%dst_part, %src_part, buf(%stage)
: !pto.partition_tensor_view<?x4096xi8>,
!pto.partition_tensor_view<?x4096xi8>,
!pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096,
v_row=?, v_col=?, blayout=row_major,
slayout=none_box, fractal=512, pad=0>)
{atomicType = #pto<atomic_type atomic_none>}
```

---

##### `pto.comm.tget` - Synchronous Remote Read

**Summary:** Lowers to `pto::comm::TGET(...)` and copies data from remote GM to local GM through a VEC staging tile.
**Summary:** Copies data synchronously from remote GM to local GM through a VEC staging tile.

**Arguments:**

Expand All @@ -10245,8 +10290,17 @@ pto.comm.tput(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view<128xf32

**Constraints & Verification:**

- Same GM/global-like and staging constraints as `pto.comm.tput`.
- `dst` and `src` must have the same element type and static shape.
- The GM/global-like, dynamic partition, runtime extent, zero-length, and
staging constraints are the same as for `pto.comm.tput`.
- `dst` and `src` must have the same element type and identical static/dynamic
shape signatures.

**Semantics:**

For every logical index in the common `src` / `dst` shape, the operation reads
the remote `src` element and writes the corresponding local `dst` element. The
staging bundle may divide the logical range into chunks but does not change the
logical transfer extent.

**Examples:**

Expand Down
51 changes: 40 additions & 11 deletions lib/PTO/IR/PTO.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3934,8 +3934,14 @@ static bool isCommGlobalLikeType(Type ty) {
return isa<pto::TensorViewType, pto::PartitionTensorViewType>(ty);
}

static LogicalResult verifyCommGlobalLike(Operation *op, Value value,
StringRef name) {
enum class CommGlobalShapePolicy {
StaticOnly,
AllowDynamicPartitionView,
};

static LogicalResult verifyCommGlobalLike(
Operation *op, Value value, StringRef name,
CommGlobalShapePolicy policy = CommGlobalShapePolicy::StaticOnly) {
Type ty = value.getType();
if (!isCommGlobalLikeType(ty))
return op->emitOpError()
Expand All @@ -3944,10 +3950,23 @@ static LogicalResult verifyCommGlobalLike(Operation *op, Value value,
SmallVector<int64_t, 4> shape = getShapeVec(ty);
if (shape.empty())
return op->emitOpError() << "expects " << name << " to have rank >= 1";

bool opAllowsDynamic =
policy == CommGlobalShapePolicy::AllowDynamicPartitionView;
bool isAllowedDynamicType = isa<pto::PartitionTensorViewType>(ty);
for (int64_t dim : shape) {
if (dim == ShapedType::kDynamic || dim <= 0)
return op->emitOpError() << "expects " << name
<< " to have a positive static shape";
if (dim == ShapedType::kDynamic) {

@zhangstevenunity zhangstevenunity Aug 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Dynamic is allowed on any dimension, but only the non-innermost ones are safe

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 TILE

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

if (!opAllowsDynamic)
return op->emitOpError()
<< "does not support dynamic dimensions on " << name;
if (!isAllowedDynamicType)
return op->emitOpError() << "allows dynamic dimensions on " << name
<< " only for partition_tensor_view";
continue;
}
if (dim <= 0)
return op->emitOpError() << "expects every static dimension of " << name
<< " to be positive";
}
return success();
}
Expand Down Expand Up @@ -16299,32 +16318,42 @@ LogicalResult TGetAsyncOp::verify() {
}

LogicalResult TPutOp::verify() {
if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) ||
failed(verifyCommGlobalLike(*this, getSrc(), "src")) ||
if (failed(verifyCommGlobalLike(
*this, getDst(), "dst",
CommGlobalShapePolicy::AllowDynamicPartitionView)) ||
failed(verifyCommGlobalLike(
*this, getSrc(), "src",
CommGlobalShapePolicy::AllowDynamicPartitionView)) ||
failed(verifyCommStagingTileLike(*this, getPing(), "ping")) ||
failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping",
"pong")))
return failure();
if (getElemTy(getDst().getType()) != getElemTy(getSrc().getType()))
return emitOpError("expects src and dst to have the same element type");
if (getShapeVec(getDst().getType()) != getShapeVec(getSrc().getType()))
return emitOpError("expects src and dst to have the same static shape");
return emitOpError(
"expects src and dst to have the same static/dynamic shape signature");
if (getElemTy(getPing().getType()) != getElemTy(getSrc().getType()))
return emitOpError("expects staging tile element type to match src/dst");
return success();
}

LogicalResult TGetOp::verify() {
if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) ||
failed(verifyCommGlobalLike(*this, getSrc(), "src")) ||
if (failed(verifyCommGlobalLike(
*this, getDst(), "dst",
CommGlobalShapePolicy::AllowDynamicPartitionView)) ||
failed(verifyCommGlobalLike(
*this, getSrc(), "src",
CommGlobalShapePolicy::AllowDynamicPartitionView)) ||
failed(verifyCommStagingTileLike(*this, getPing(), "ping")) ||
failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping",
"pong")))
return failure();
if (getElemTy(getDst().getType()) != getElemTy(getSrc().getType()))
return emitOpError("expects src and dst to have the same element type");
if (getShapeVec(getDst().getType()) != getShapeVec(getSrc().getType()))
return emitOpError("expects src and dst to have the same static shape");
return emitOpError(
"expects src and dst to have the same static/dynamic shape signature");
if (getElemTy(getPing().getType()) != getElemTy(getSrc().getType()))
return emitOpError("expects staging tile element type to match src/dst");
return success();
Expand Down
21 changes: 21 additions & 0 deletions test/lit/pto/comm_dynamic_async_scope_invalid.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
// CANN Open Software License Agreement Version 2.0 (the "License").
// Please refer to the License for details. You may not use this file except in compliance with the License.
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
// See LICENSE in the root of the software repository for the full text of the License.

// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s

module {
func.func @dynamic_async_p2p(
%dst: !pto.partition_tensor_view<?xf32>,
%src: !pto.partition_tensor_view<?xf32>,
%session: !pto.async_session) {
%event = pto.comm.tput_async(%dst, %src, %session : !pto.partition_tensor_view<?xf32>, !pto.partition_tensor_view<?xf32>, !pto.async_session) -> !pto.async_event
return
}
}

// CHECK: error: 'pto.comm.tput_async' op expects dst to have a static shape
21 changes: 21 additions & 0 deletions test/lit/pto/comm_dynamic_collective_scope_invalid.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
// CANN Open Software License Agreement Version 2.0 (the "License").
// Please refer to the License for details. You may not use this file except in compliance with the License.
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
// See LICENSE in the root of the software repository for the full text of the License.

// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s

module {
func.func @dynamic_collective(
%src: !pto.partition_tensor_view<?x128xf32>,
%peer: !pto.partition_tensor_view<?x128xf32>,
%stage: !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {
pto.comm.tbroadcast(%src, recv(%stage), group(%peer) : !pto.partition_tensor_view<?x128xf32>, !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=128, v_row=1, v_col=128, blayout=row_major, slayout=none_box, fractal=512, pad=0>, !pto.partition_tensor_view<?x128xf32>) {root = 0 : i32}
return
}
}

// CHECK: error: 'pto.comm.tbroadcast' op does not support dynamic dimensions on src
19 changes: 19 additions & 0 deletions test/lit/pto/comm_dynamic_signal_scope_invalid.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
// CANN Open Software License Agreement Version 2.0 (the "License").
// Please refer to the License for details. You may not use this file except in compliance with the License.
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
// See LICENSE in the root of the software repository for the full text of the License.

// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s

module {
func.func @dynamic_signal(
%signal: !pto.partition_tensor_view<?xi32>, %value: i32) {
pto.comm.tnotify(%signal, %value : !pto.partition_tensor_view<?xi32>, i32) {notifyOp = #pto<notify_op set>}
return
}
}

// CHECK: error: 'pto.comm.tnotify' op does not support dynamic dimensions on signal
68 changes: 68 additions & 0 deletions test/lit/pto/comm_p2p_dynamic_partition_emitc.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
// CANN Open Software License Agreement Version 2.0 (the "License").
// Please refer to the License for details. You may not use this file except in compliance with the License.
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
// See LICENSE in the root of the software repository for the full text of the License.

// RUN: ptoas --pto-arch=a3 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC
// RUN: ptoas --pto-arch=a3 --pto-level=level3 --mlir-print-ir-after=pto-resolve-buffer-select %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE

module {
func.func @comm_p2p_dynamic_partition(
%dst_ptr: !pto.ptr<i8>, %src_ptr: !pto.ptr<i8>, %runtime_rows: i32) {
%c0_i64 = arith.constant 0 : i64
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
%c4096 = arith.constant 4096 : index
%rows = arith.index_cast %runtime_rows : i32 to index
%dst_view = pto.make_tensor_view %dst_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi8>
%src_view = pto.make_tensor_view %src_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi8>
%dst = pto.partition_view %dst_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
%src = pto.partition_view %src_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
%stage = pto.alloc_tile addr = %c0_i64 valid_row = %c1 valid_col = %c4096 : !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>
pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view<?x4096xi8>, !pto.partition_tensor_view<?x4096xi8>, !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {atomicType = #pto<atomic_type atomic_none>}
pto.comm.tget(%dst, %src, buf(%stage) : !pto.partition_tensor_view<?x4096xi8>, !pto.partition_tensor_view<?x4096xi8>, !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>)
return
}

// A static valid shape remains legal. At runtime its row/column values must
// divide the transfer shape whenever the transfer enters the chunked path.
func.func @comm_p2p_dynamic_partition_static_valid(
%dst_ptr: !pto.ptr<i8>, %src_ptr: !pto.ptr<i8>, %runtime_rows: i32) {
%c0_i64 = arith.constant 0 : i64
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
%c4096 = arith.constant 4096 : index
%rows = arith.index_cast %runtime_rows : i32 to index
%dst_view = pto.make_tensor_view %dst_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi8>
%src_view = pto.make_tensor_view %src_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi8>
%dst = pto.partition_view %dst_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
%src = pto.partition_view %src_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
%stage = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=1, v_col=4096, blayout=row_major, slayout=none_box, fractal=512, pad=0>
pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view<?x4096xi8>, !pto.partition_tensor_view<?x4096xi8>, !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=1, v_col=4096, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {atomicType = #pto<atomic_type atomic_none>}
return
}
}

// EMITC: pto::Shape<1, 1, 1, -1, 4096>
// EMITC-SAME: ({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^)]+}})

@zhangstevenunity zhangstevenunity Aug 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge These wildcards accept any five constructor arguments, so the test cannot fail

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.

// EMITC: pto::Stride<-1, -1, -1, -1, -1>
// EMITC-SAME: ({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^)]+}})
// EMITC: GlobalTensor<int8_t, pto::Shape<1, 1, 1, -1, 4096>, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND>
// EMITC: pto::comm::TPUT(
// EMITC: pto::comm::TGET(

// NATIVE: IR Dump After PTOResolveBufferSelect
// NATIVE-LABEL: func.func @comm_p2p_dynamic_partition
// NATIVE: %[[ROWS:.*]] = arith.index_cast
// NATIVE: %[[DST:.*]] = pto.partition_view {{.*}} sizes = [%[[ROWS]], %c4096]
// NATIVE: %[[SRC:.*]] = pto.partition_view {{.*}} sizes = [%[[ROWS]], %c4096]
// NATIVE: pto.comm.tput(%[[DST]], %[[SRC]], buf(
// NATIVE-SAME: !pto.partition_tensor_view<?x4096xi8>
// NATIVE: pto.comm.tget(%[[DST]], %[[SRC]], buf(
// NATIVE-SAME: !pto.partition_tensor_view<?x4096xi8>
// NATIVE-NOT: memref.subview
21 changes: 21 additions & 0 deletions test/lit/pto/comm_p2p_dynamic_partition_verify_invalid.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
// Copyright (c) 2026 Huawei Technologies Co., Ltd.
// This program is free software, you can redistribute it and/or modify it under the terms and conditions of
// CANN Open Software License Agreement Version 2.0 (the "License").
// Please refer to the License for details. You may not use this file except in compliance with the License.
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
// See LICENSE in the root of the software repository for the full text of the License.

// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s

module {
func.func @shape_signature_mismatch(
%dst: !pto.partition_tensor_view<?x4096xi8>,
%src: !pto.partition_tensor_view<?x2048xi8>,
%stage: !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=1, v_col=4096, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {
pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view<?x4096xi8>, !pto.partition_tensor_view<?x2048xi8>, !pto.tile_buf<loc=vec, dtype=i8, rows=1, cols=4096, v_row=1, v_col=4096, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {atomicType = #pto<atomic_type atomic_none>}
return
}
}

// CHECK: error: 'pto.comm.tput' op expects src and dst to have the same static/dynamic shape signature
Loading
Loading