[Bug] pto.comm.tput/tget verifier rejects dynamic partition-view shapes that the rest of the pipeline already lowers

[YunjiQin]

Component

Verifier / IR semantics (lib/PTO/IR)

Description

verifyCommGlobalLike() rejects any ? (dynamic) dimension on the dst/src partition views of pto.comm.tput and pto.comm.tget, so a variable-size transfer cannot be expressed — even though the rest of the pipeline already handles it end to end.

The verifier is the only blocker. I confirmed this by bypassing it (see "Proof the lowering already works" below): with the check out of the way, a fully ?-dimensioned pto.comm.tput lowers cleanly to a runtime-extent GlobalTensor and a valid pto::comm::TPUT(...) call, on both v0.48 and v0.54.

Why the rejection looks unintentional rather than a real constraint:

1. pto.partition_view with a ? result type + dynamic size SSA already lowers fine — it just isn't allowed to reach a comm op. The same partition view feeding pto.tload compiles without complaint (case C below).
2. The live lowering path is dynamic-capable. pto.partition_view → (PTOViewToMemref.cpp:1330 lowerPartitionViewOps) → memref.subview with dynamic sizes → SubviewToEmitCPattern (PTOToEmitC.cpp:3766), which passes adaptor.getSizes() straight into the pto::Shape constructor (PTOToEmitC.cpp:4172-4184). The static-only early-return in buildGlobalTensorFromMemref (PTOToEmitC.cpp:4406) sits on the memref-direct path, which comm ops built from partition views never take.
3. GlobalTensor can express runtime extents. pto::DYNAMIC = -1 and pto::Shape / pto::Stride have runtime constructors (pto-isa include/pto/common/pto_tile.hpp), which is exactly what SubviewToEmitCPattern emits.
4. pto-isa reads the shape at runtime. TPUT_IMPL (pto-isa include/pto/comm/a2a3/TPut.hpp) obtains extents via srcGlobalData.GetShape(GlobalTensorDim::DIM_n), with no shape-related static_assert; the 2-D slide auto-chunks a transfer larger than the staging tile.

Impact: push-based collectives in pypto must transfer the full compile-time capacity instead of the actual runtime count. In all_to_all_v / MoE dispatch, a peer that receives 1 of 32 possible rows still costs 32 rows of interconnect bandwidth. pypto codegen already emits the dynamic form (!pto.partition_tensor_view<?x64xf16> with a static chunked staging tile), so it fails at the ptoas verifier today — see pypto issue #2213.

Reproduction (minimal)

Case A — dynamic pto.comm.tput (fails). Identical to the static version except sizes uses a runtime %rows and the partition-view type is <?x4096xi8>:

module attributes {pto.target_arch = "a2a3"} {
  func.func @tput_dynamic(%arg0: !pto.ptr<i8>, %arg1: !pto.ptr<i8>, %arg2: i32,
                          %__pypto_spmd_block_idx: i32, %__pypto_spmd_block_num: i32)
      attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
    %c0_i64 = arith.constant 0 : i64
    %c0_index = arith.constant 0 : index
    %c1_index = arith.constant 1 : index
    %c32_index = arith.constant 32 : index
    %c4096_index = arith.constant 4096 : index
    %rows = arith.index_cast %arg2 : i32 to index
    %dst_view = pto.make_tensor_view %arg0, shape = [%c32_index, %c4096_index], strides = [%c4096_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xi8>
    %src_view = pto.make_tensor_view %arg1, shape = [%c32_index, %c4096_index], strides = [%c4096_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xi8>
    %stage = pto.alloc_tile addr = %c0_i64 valid_row = %c1_index valid_col = %c4096_index : !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>
    %dst_pview = pto.partition_view %dst_view, offsets = [%c0_index, %c0_index], sizes = [%rows, %c4096_index] : !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
    %src_pview = pto.partition_view %src_view, offsets = [%c0_index, %c0_index], sizes = [%rows, %c4096_index] : !pto.tensor_view<?x?xi8> -> !pto.partition_tensor_view<?x4096xi8>
    pto.barrier <PIPE_ALL>
    pto.comm.tput(%dst_pview, %src_pview, 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.barrier <PIPE_ALL>
    return
  }
}

ptoas dynamic.pto -o dynamic.cpp --pto-level=level3
# same file with sizes = [%c32_index, ...] and <32x4096xi8> succeeds

Case B — same file with pto.comm.tget: same error ('pto.comm.tget' op expects dst to have a positive static shape).

Case C — the same dynamic partition view feeding pto.tload (succeeds). Shows the ? partition view itself is fine; only the comm ops reject it:

    %pv = pto.partition_view %view, offsets = [%c0_index, %c0_index], sizes = [%rows, %c32_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<?x32xf32>
    pto.tload ins(%pv : !pto.partition_tensor_view<?x32xf32>) outs(%tb : !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>)

emits the runtime extent and TLOAD without error:

pto::Shape<1, 1, 1, -1, 32> v10 = pto::Shape<1, 1, 1, -1, 32>(v8);
GlobalTensor<float, pto::Shape<1, 1, 1, -1, 32>, pto::Stride<-1, -1, -1, 32, 1>, pto::Layout::ND> v12 = ...;
TLOAD(v13, v12);

Proof the lowering already works. Taking Case A and only routing the staging tile through pto.bind_tile (so shouldBypassDecodedMemrefVerifier, PTO.cpp:514, skips the check) makes the same ?x4096 tput lower successfully on both v0.48 and v0.54:

    %ubuf = memref.alloc() : memref<1x4096xi8, #pto.address_space<vec>>
    %stage = pto.bind_tile %ubuf, %c1_index, %c4096_index
      {config = #pto.tile_buf_config<blayout=0 : i32, slayout=0 : i32, s_fractal_size=512, pad=0 : i32>}
      : memref<1x4096xi8, #pto.address_space<vec>> -> memref<1x4096xi8, #pto.address_space<vec>>
    pto.comm.tput(%dst_pview, %src_pview, buf(%stage) : !pto.partition_tensor_view<?x4096xi8>, !pto.partition_tensor_view<?x4096xi8>, memref<1x4096xi8, #pto.address_space<vec>>) {atomicType = #pto<atomic_type atomic_none>}

$ ptoas dyn_bypass.pto -o dyn_bypass.cpp --pto-level=level2   # exit 0, no diagnostics

pto::Shape<1, 1, 1, -1, 4096> v11 = pto::Shape<1, 1, 1, -1, 4096>(v9);            // v9 = runtime rows
pto::Stride<-1, -1, -1, 4096, 1> v12 = pto::Stride<-1, -1, -1, 4096, 1>(v10, v10, v10);
GlobalTensor<int8_t, pto::Shape<1, 1, 1, -1, 4096>, pto::Stride<-1, -1, -1, 4096, 1>, pto::Layout::ND> v13 = ...;
pto::comm::TPUT(v13, v17, v19);

--emit-pto-ir on that run confirms the intermediate form the EmitC patterns consume — a dynamic memref.subview, with the comm op already memref-typed at that point:

%subview = memref.subview %reinterpret_cast[0, 0] [%0, 4096] [1, 1] {layout = #pto.layout<nd>} : memref<?x?xi8, strided<[?, ?], offset: ?>, #pto.address_space<gm>> to memref<?x4096xi8, strided<[?, ?], offset: ?>, #pto.address_space<gm>>
pto.comm.tput(%subview, %subview_1, buf(%2) : memref<?x4096xi8, strided<[?, ?], offset: ?>, #pto.address_space<gm>>, ...)

Expected behavior

pto.comm.tput / pto.comm.tget should accept ? dimensions on their dst/src partition views and lower them to a runtime-extent GlobalTensor, exactly as Case C and the bypass run already do — the number of rows on the wire then equals the runtime payload instead of the compile-time capacity.

Suggested fix in verifyCommGlobalLike (lib/PTO/IR/PTO.cpp:4187-4191): keep rejecting non-positive static dims, allow kDynamic:

for (int64_t dim : shape) {
  if (dim != ShapedType::kDynamic && dim <= 0)
    return op->emitOpError() << "expects " << name
                             << " to have a positive or dynamic shape";
}

Two related spots to consider in the same review:

• TPutOp::verify() / TGetOp::verify() (PTO.cpp:16844 / :16862) compare getShapeVec(dst) != getShapeVec(src) under the message "same static shape". With ? allowed, two dynamic extents comparing equal as kDynamic is the right structural check, but the diagnostic wording should be updated.
• verifyCommSignalLike and verifyCommStagingTileLike (PTO.cpp:4195 / :4206) reuse the same positive-static loop. Signals and VEC staging tiles genuinely want static shapes, so please keep those strict — only the GM dst/src views of tput/tget need relaxing.

Actual behavior / error logs

$ ptoas dynamic.pto -o dynamic.cpp --pto-level=level3
loc("dynamic.pto":17:5): error: 'pto.comm.tput' op expects dst to have a positive static shape
Error: Failed to parse MLIR.

$ ptoas dyn_tget.pto -o dyn_tget.cpp --pto-level=level3
loc("dyn_tget.pto":17:5): error: 'pto.comm.tget' op expects dst to have a positive static shape
Error: Failed to parse MLIR.

Reproduced identically on released binaries v0.48, v0.50 and v0.54, and with real pypto codegen output (pld.tensor.put(..., shape=[n, 64], chunk_rows=4, chunk_cols=32)):

loc("pypto_dyn_put.pto":22:3): error: 'pto.comm.tput' op expects dst to have a positive static shape
Error: Failed to parse MLIR.

Origin (v0.54, lib/PTO/IR/PTO.cpp:4177-4192):

static LogicalResult verifyCommGlobalLike(Operation *op, Value value,
                                          StringRef name) {
  ...
  for (int64_t dim : shape) {
    if (dim == ShapedType::kDynamic || dim <= 0)
      return op->emitOpError() << "expects " << name
                               << " to have a positive static shape";
  }
  return success();
}

Git commit

195fdb8

Host platform

Linux (aarch64)

Target Ascend arch (if relevant)

a3

PTOAS build level (if relevant)

level3

[georgebisbas]

Thanks for this @YunjiQin