Chunked pto.comm.tput (staging tile smaller than transfer) does not drain its stores before a following TNOTIFY → peer reads partial data (a2a3)

[YunjiQin]

Summary

When the VEC staging tile passed to pto::comm::TPUT is smaller than the transfer extent (so TPUT internally 2-D-slides the transfer through the tile in multiple chunks), the generated kernel does not drain the chunked MTE3 stores before a following cross-rank pto::comm::TNOTIFY. The signal therefore races ahead of the last in-flight chunk stores; the peer's TWAIT unblocks and reads partial data.

Single-shot TPUT (staging tile == transfer extent) does not exhibit this — it is self-draining.

This is the same insert-sync / pto.comm.tput hazard family as #706 and #730, but specific to the chunked (sub-tile staging) TPUT path.

Repro

Frontend (PyPTO) — a 2-rank ring-shuffle pld.tensor.put of an [8, 64] FP32 window, with the staging tile forced down to [2, 64] (or [2, 32]) so TPUT must auto-chunk:

• rank r pushes its [8,64] slice into peer (r+1)%2's slice, then TNOTIFYs the peer; peer TWAITs then reads its slice back.
• Golden: out[r] == in[(r-1)%2].

Observed on device (a2a3 / Ascend 910B, 2 ranks):

• chunked staging ([2,32] or [2,64] tile, full [8,64] transfer) → wrong: one direction fully correct, the other has the top rows correct and the bottom rows zero (a partial transfer; max diff = 511). Reproduced with both column-split and row-only chunking.
• single-shot staging (full [8,64] tile) → correct.

Generated C++ (PTOAS output)

The chunked TPUT region (next_levels/chip_orch/kernels/aiv/ring_step.cpp). Note the GlobalTensor is the full Shape<1,1,1,8,64> transfer while the staging tile is Tile<...,2,32,...,-1,-1,...> (allocated 2x32, ValidRow=ValidCol=-1 = DYNAMIC):

GlobalTensor<float, Shape<1,1,1,8,64>, Stride<512,512,512,64,1>, Layout::ND> v36 = ...;  // full transfer
// v28 = Tile<TileType::Vec, float, 2, 32, RowMajor, -1, -1, NoneBox, ...>               // sub-tile staging

pipe_barrier(PIPE_ALL);                          // orders the preceding source TSTORE (ok)
wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1);
pto::comm::TPUT(v36, v26, v28);                  // chunked: 4x2 internal slide, many MTE3 stores
set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2);
// >>> nothing drains the chunked stores here <<<
pto::comm::TNOTIFY(v41, v12, v11);               // cross-rank signal races the in-flight stores
...
pto::comm::TWAIT(v44, v12, v10);                 // peer unblocks and reads partial data

PTOAS emits only the local set_flag(PIPE_MTE3, PIPE_MTE2, ...) after TPUT; there is no pipe_barrier(PIPE_ALL) (or equivalent MTE3 drain) between the chunked TPUT and the following TNOTIFY.

Expected

A chunked pto::comm::TPUT followed by a visibility/signal op (TNOTIFY) must guarantee all chunk stores have landed before the signal fires — same effective completion guarantee as single-shot TPUT. Either:

1. insert-sync should emit a full pipe_barrier(PIPE_ALL) (or MTE3 drain) between a chunked pto.comm.tput and a following pto.system.notify / pto.comm.tnotify; or
2. the chunked TPUT implementation should be self-completing (drain its last chunk before returning), matching single-shot semantics.

Workaround (verified)

Inserting a pipe_barrier(PIPE_ALL) immediately after TPUT (before TNOTIFY) makes the chunked case correct on device. We currently inject this from the PyPTO frontend codegen as a stopgap, but it should be PTOAS's responsibility (single-shot TPUT needs no such caller-inserted barrier).

Environment

• Platform: a2a3 (Ascend 910B), 2 ranks
• pto-isa commit: 016396b57e2c17093f1194e6acd89bb112b0ab24
• Frontend: PyPTO pld.tensor.put with sub-tile VEC staging (auto-chunk)

Related

• [Pass Bug] insert-sync: missing MTE3->MTE2 hazard between pto.tstore and pto.comm.tput (TPUT reads src on MTE2) #706 — [Pass Bug] insert-sync: missing MTE3->MTE2 hazard between pto.tstore and pto.comm.tput
• [Bug] insert-sync: MTE3->MTE2 pipe flag for same-address GM store->load round-trip doesn't guarantee GM visibility — intermittent stale read / MTE hang (a2a3, vector spmd) #730 — [Bug] insert-sync: MTE3->MTE2 pipe flag for same-address GM store->load round-trip doesn't guarantee GM visibility

[YunjiQin]

Complete .pto for the repro (input PTOAS lowers)

Full chip kernel for the chunked ring-shuffle put ([8,64] FP32 transfer, [2,32] VEC staging tile). The relevant sequence is pto.comm.tput (line 30) → pto.comm.tnotify (line 38).

To reproduce the bug, delete line 31 (pto.barrier <PIPE_ALL>) — that barrier is our frontend stopgap. Without it the natural lowering goes straight from the chunked tput to tnotify, and the chunked MTE3 stores race the cross-rank signal. PTOAS should insert that drain itself (insert-sync), or the chunked TPUT should self-complete like single-shot.

Note: %tput_stage (line 21) is rows=2, cols=32, v_row=?, v_col=? (DYNAMIC valid) while both partition views are the full <8x64xf32> transfer — i.e. TPUT must auto-chunk (4×2 internal slide).

module attributes {pto.target_arch = "a2a3"} {
  func.func @ring_step(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f32>, %arg2: !pto.ptr<f32>, %arg3: !pto.ptr<f32>, %arg4: !pto.ptr<i32>, %arg5: i32, %arg6: !pto.ptr<i64>, %arg7: !pto.ptr<i64>, %arg8: !pto.ptr<i64>) attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
  %c0_i64 = arith.constant 0 : i64
  %c8_index = arith.constant 8 : index
  %c64_index = arith.constant 64 : index
  %c1_index = arith.constant 1 : index
  %c0_index = arith.constant 0 : index
  %c2_index = arith.constant 2 : index
  %c32_index = arith.constant 32 : index
  %c1_i32 = arith.constant 1 : i32
  %inp__ssa_v0_view = pto.make_tensor_view %arg0, shape = [%c8_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf32>
  %out__ssa_v0_view = pto.make_tensor_view %arg1, shape = [%c8_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf32>
  %src__ssa_v0_view = pto.make_tensor_view %arg2, shape = [%c8_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf32>
  %dst__ssa_v0_view = pto.make_tensor_view %arg3, shape = [%c8_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf32>
  %signal__ssa_v0_view = pto.make_tensor_view %arg4, shape = [%c1_index, %c1_index], strides = [%c1_index, %c1_index] {layout = #pto.layout<dn>}: !pto.tensor_view<?x?xi32>
  %local__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c8_index valid_col = %c64_index : !pto.tile_buf<loc=vec, dtype=f32, rows=8, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>
  %inp__ssa_v0_pview = pto.partition_view %inp__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c8_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<8x64xf32>
  pto.tload ins(%inp__ssa_v0_pview : !pto.partition_tensor_view<8x64xf32>) outs(%local__ssa_v0 : !pto.tile_buf<loc=vec, dtype=f32, rows=8, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>)
  %src__ssa_v0_pview = pto.partition_view %src__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c8_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<8x64xf32>
  pto.tstore ins(%local__ssa_v0 : !pto.tile_buf<loc=vec, dtype=f32, rows=8, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>) outs(%src__ssa_v0_pview : !pto.partition_tensor_view<8x64xf32>)
  %tput_stage = pto.alloc_tile addr = %c0_i64 valid_row = %c2_index valid_col = %c32_index : !pto.tile_buf<loc=vec, dtype=f32, rows=2, cols=32, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>
  %peer__ssa_v0_idx = arith.index_cast %arg5 : i32 to index
  %0 = func.call @CommRemoteOffset_f32(%arg7, %peer__ssa_v0_idx) : (!pto.ptr<i64>, index) -> index
  %1 = pto.addptr %arg3, %0 : !pto.ptr<f32> -> !pto.ptr<f32>
  %2 = arith.muli %c1_index, %c64_index : index
  %3 = pto.make_tensor_view %1, shape = [%c8_index, %c64_index], strides = [%2, %c1_index] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xf32>
  %dst__ssa_v0_peer_pview = pto.partition_view %3, offsets = [%c0_index, %c0_index], sizes = [%c8_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<8x64xf32>
  %src__ssa_v1_local_pview = pto.partition_view %src__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c8_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<8x64xf32>
  pto.barrier <PIPE_ALL>
  pto.comm.tput(%dst__ssa_v0_peer_pview, %src__ssa_v1_local_pview, buf(%tput_stage) : !pto.partition_tensor_view<8x64xf32>, !pto.partition_tensor_view<8x64xf32>, !pto.tile_buf<loc=vec, dtype=f32, rows=2, cols=32, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>) {atomicType = #pto<atomic_type atomic_none>}
  pto.barrier <PIPE_ALL>
  %4 = arith.index_cast %arg5 : i32 to index
  %5 = func.call @CommRemoteOffset_i32(%arg8, %4) : (!pto.ptr<i64>, index) -> index
  %6 = pto.addptr %arg4, %5 : !pto.ptr<i32> -> !pto.ptr<i32>
  %7 = arith.muli %c1_index, %c1_index : index
  %8 = pto.make_tensor_view %6, shape = [%c1_index, %c1_index], strides = [%7, %c1_index] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi32>
  %signal__ssa_v0_peer_pview = pto.partition_view %8, offsets = [%c0_index, %c0_index], sizes = [%c1_index, %c1_index] : !pto.tensor_view<?x?xi32> -> !pto.partition_tensor_view<1x1xi32>
  pto.comm.tnotify(%signal__ssa_v0_peer_pview, %c1_i32 : !pto.partition_tensor_view<1x1xi32>, i32) {notifyOp = #pto<notify_op atomic_add>}
  %signal__ssa_v0_local_pview = pto.partition_view %signal__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c1_index, %c1_index] : !pto.tensor_view<?x?xi32> -> !pto.partition_tensor_view<1x1xi32>
  pto.comm.twait(%signal__ssa_v0_local_pview, %c1_i32 : !pto.partition_tensor_view<1x1xi32>, i32) {cmp = #pto<wait_cmp ge>}
  %recv__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c8_index valid_col = %c64_index : !pto.tile_buf<loc=vec, dtype=f32, rows=8, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>
  %dst__ssa_v0_pview = pto.partition_view %dst__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c8_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<8x64xf32>
  pto.tload ins(%dst__ssa_v0_pview : !pto.partition_tensor_view<8x64xf32>) outs(%recv__ssa_v0 : !pto.tile_buf<loc=vec, dtype=f32, rows=8, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>)
  %out__ssa_v0_pview = pto.partition_view %out__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c8_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<8x64xf32>
  pto.tstore ins(%recv__ssa_v0 : !pto.tile_buf<loc=vec, dtype=f32, rows=8, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=none_box, fractal=512, pad=0>) outs(%out__ssa_v0_pview : !pto.partition_tensor_view<8x64xf32>)
  return
  }
  func.func private @CommRemoteOffset_i32(%ctx: !pto.ptr<i64>, %peer: index) -> index {
    %c_r = arith.constant 2 : index
    %c_w = arith.constant 4 : index
    %rk_pair = pto.load_scalar %ctx[%c_r] : !pto.ptr<i64> -> i64
    %rk_i32 = arith.trunci %rk_pair : i64 to i32
    %rk_idx = arith.index_cast %rk_i32 : i32 to index
    %lb_off = arith.addi %c_w, %rk_idx : index
    %lbase = pto.load_scalar %ctx[%lb_off] : !pto.ptr<i64> -> i64
    %pb_off = arith.addi %c_w, %peer : index
    %pbase = pto.load_scalar %ctx[%pb_off] : !pto.ptr<i64> -> i64
    %dbytes = arith.subi %pbase, %lbase : i64
    %esize = arith.constant 4 : i64
    %delems_i = arith.divsi %dbytes, %esize : i64
    %delems = arith.index_cast %delems_i : i64 to index
    return %delems : index
  }
  func.func private @CommRemoteOffset_f32(%ctx: !pto.ptr<i64>, %peer: index) -> index {
    %c_r = arith.constant 2 : index
    %c_w = arith.constant 4 : index
    %rk_pair = pto.load_scalar %ctx[%c_r] : !pto.ptr<i64> -> i64
    %rk_i32 = arith.trunci %rk_pair : i64 to i32
    %rk_idx = arith.index_cast %rk_i32 : i32 to index
    %lb_off = arith.addi %c_w, %rk_idx : index
    %lbase = pto.load_scalar %ctx[%lb_off] : !pto.ptr<i64> -> i64
    %pb_off = arith.addi %c_w, %peer : index
    %pbase = pto.load_scalar %ctx[%pb_off] : !pto.ptr<i64> -> i64
    %dbytes = arith.subi %pbase, %lbase : i64
    %esize = arith.constant 4 : i64
    %delems_i = arith.divsi %dbytes, %esize : i64
    %delems = arith.index_cast %delems_i : i64 to index
    return %delems : index
  }
}

[YunjiQin]

Likely the same root cause as #744 (cross-rank push not observable before the signal)

pto::comm::TPUT is implemented in pto-isa with the same TSTORE_IMPL that remote_store uses (comm/a2a3/TPut.hpp → npu/a2a3/TStore.hpp), writing to a peer-addressed GM view (base + CommRemoteOffset). That TSTORE_IMPL issues no dsb; the only dsb(DSB_DDR) is inside TNOTIFY_IMPL, after the signal store (comm/a2a3/TNotify.hpp). So a chunked tput → tnotify is missing the same producer-side fences #744 is about, and possibly more:

1. an MTE3 drain of the chunked stores before the signal — the chunked TPUT sliding path returns without draining its last chunk, so the data isn't even issued (this is the deterministic part of the failure); and
2. very likely the [Bug] pto.comm.tnotify: MTE drain doesn't cover DDR observability — cross-rank receivers still race after #711 #744 dsb(DSB_DDR) before the signal store — to force the pushed data to DDR before the signal becomes observable to the peer (the residual intermittent cross-rank race [Bug] pto.comm.tnotify: MTE drain doesn't cover DDR observability — cross-rank receivers still race after #711 #744 fixed for remote_store).

emitTNotifyMteDrain's drain-mask currently recognizes only a preceding pto.tstore as a peer-store source, so for a preceding pto.comm.tput it emits neither the MTE3 drain nor the dsb.

Suggested fix

Extend emitTNotifyMteDrain so a preceding peer-addressed pto.comm.tput is treated like pto.tstore — emit pipe_barrier(PIPE_MTE3); dsb(DSB_DDR) before TNOTIFY, i.e. the same fix as #744. (Single-shot tput doesn't surface this because TPUT_IMPL self-drains its single store, but the cross-rank dsb requirement is identical.)

Our PyPTO-side stopgap currently only inserts a pipe_barrier(PIPE_ALL) after the tput; that fixes the deterministic missing-drain but, per #744, a pipe barrier alone does not guarantee cross-rank DDR observability — the dsb likely still belongs here too.

[YunjiQin]

On-device validation of the fix (a2a3, 2 ranks, test_chunked_ring_shuffle)

Variant	Result
PyPTO no barrier · TNOTIFY_IMPL stock	30/30 FAIL (deterministic)
PyPTO pipe_barrier(PIPE_ALL) after tput · TNOTIFY_IMPL stock	140/140 pass (but a pipe barrier alone can't guarantee cross-rank DDR observability — likely a residual #744-style race)
PyPTO no barrier · TNOTIFY_IMPL + dsb(DSB_DDR) only at start	10/10 FAIL
PyPTO no barrier · TNOTIFY_IMPL + pipe_barrier(PIPE_ALL); dsb(DSB_DDR) at start	20/20 PASS

Conclusions:

• dsb alone is not enough (10/10 fail): the chunked TPUT's last store is still in the MTE3 pipe, and dsb (a memory barrier) doesn't drain the pipe — you need a pipe_barrier(PIPE_MTE3/ALL) first.
• The complete fix is pipe_barrier(PIPE_ALL); dsb(DSB_DDR) before the signal write, i.e. exactly [Bug] pto.comm.tnotify: MTE drain doesn't cover DDR observability — cross-rank receivers still race after #711 #744's pipe_barrier(PIPE_MTE3); dsb(DSB_DDR).

Recommended location

Putting pipe_barrier(PIPE_ALL); dsb(DSB_DDR) at the very start of TNOTIFY_IMPL (pto-isa comm/a2a3/TNotify.hpp, before the signal store) fixes this generically for any preceding cross-rank store — pto.comm.tput, remote_store (#744), etc. — and removes the need for any frontend/codegen-side per-op drain mask. (TNOTIFY_IMPL already has a trailing dsb(DSB_DDR); pipe_barrier(PIPE_ALL) after the signal; the missing piece is the same fence before it.) Validated 20/20 with no PyPTO-side barrier.

(The PyPTO pipe_barrier(PIPE_ALL) stopgap stays until this lands; note it only covers the deterministic pipe-drain half, not the DDR half.)