[VMI] expose PLT POST_UPDATE remaining from create_mask for multi-tile vdhist

[peanutchan]

Summary

pto.vmi.create_mask(active_lanes, size=...) returns only a mask. For an 8×256 vdhist sweep over a dynamic remaining count (TileLang group_count tail / partial copy tile), that forces absolute per-tile lane counts:

for i in range(8):
    rem_from = n - i * 256          # independent SSA per tile
    mask = pto.vmi.create_mask(rem_from, size=256)
    acc = pto.vmi.vdhist(acc, src[i], mask)

The ISA-intended pattern is a carried remaining scalar with plt_b8(..., POST_UPDATE):

uint32_t rem = n;
for (uint16_t i = 0; i < 8; ++i) {
    auto mask = plt_b8(rem, POST_UPDATE);  // rem updated in place
    dhistv2(hist, src[i], mask, Bin_N0);
}

Today this ideal path is expressible in raw CCE / pto.plt_b8 / pto.make_mask, but not with pto.vmi.vdhist:

• pto.plt_b8 returns !pto.mask<b8>, which pto.vmi.vdhist rejects (operand #2 must be VMI logical mask type)
• pto.vmi.create_mask lowers each call independently via plt and discards scalar_out, so remaining cannot chain across tiles

Ask

Please add a VMI-surface API that returns remaining, e.g.:

mask, remained = pto.vmi.create_mask(remained, size=256)  # or create_mask_post_update(...)
# remained usable as next create_mask / loop-carried i32/index

Requirements:

1. Result mask must be a VMI logical mask accepted by vdhist / other VMI ops
2. remained must be the PLT POST_UPDATE scalar_out (or equivalent), so an 8-tile unroll can share one remaining counter without recomputing min(256, max(0, n - i*256)) / rematerializing 8 absolute counts
3. Documented interaction with TileLang T.vmi.create_mask so group_count-style kernels can lower to this path

Microbench (CCE, same 8×256 dhist slice)

Repro (in tilelang fork worktree):

examples/ascend/vmi/plt_post_update_microbench/

• cce/dhist_plt_compare_cce.cpp — absolute vs POST_UPDATE modes
• run_cce_launch.py — launch under scripts/sim_dsl.sh
• abs_create_mask_vmi.py — VMI absolute create_mask(n - i*256) IR reference

Setup: n=1800, outer=64 → 512 tile iterations (64×8). Both modes PASS (identical hist vs golden).

Opsim instr mix (Ascend950PR_9599)

instr	absolute calls	post_update calls	delta
RV_DHISTv2	512	512	0
RV_VLD	512	512	0
RV_PLT	512	512	0
RV_SSUB	512	0	+512
RV_SMIN	512	0	+512
RV_SAND	512	0	+512
RV_SADD	1024	0	+1024
RV_SCMP	576	64	+512
RV_SMOV	129	65	+64
total call_count	5208	2136	~2.4×

Absolute mode pays a full min/max/sub scalar chain per tile to rebuild the lane count that POST_UPDATE would have carried for free. Wall block time is still DHIST/VLD-dominated in this tiny kernel (~4.28 µs both), but the scalar tax matches what we see in TileLang group_count tail dumps (+PLT/+SMOV per unrolled hist tile).

VMI absolute IR shape

%rem = arith.subi %n, %c256 : index
%mask = pto.vmi.create_mask %rem : index -> !pto.vmi.mask<256xpred>
%acc = pto.vmi.vdhist %acc, %src, %mask ...

repeated per tile — no loop-carried remaining from PLT.

Context

Came up while optimizing TileLang VMI group_count (vdhist over copy tiles with optional tail). Tail-enabled aligned cases show extra per-tile RV_PLT/RV_SMOV vs the no-tail path; expectation is one remaining scalar + POST_UPDATE PLTs across the 8 hist tiles.

Related surface that already does the right thing for non-VMI ops: pto.make_mask / pto.plt_b* returning (mask, remained) (see tcmp_template.py).

[peanutchan]

Correction: reproduced the TileLang fully-unrolled SMOV→PLT pattern

The first microbench used a runtime loop with min/max/sub absolute counts, which showed an SSUB/SMIN tax — not the TileLang problem.

TileLang with_tail is compile-time unrolled: 8 independent tile_elems SSA values, and PTOAS create_mask lowering does SMOV te_i → S64; PLT for each tile (even when every te_i == 256).

Fair comparison (one 8-tile sweep, n=2048 full tiles)

Path	RV_PLT unique PCs	RV_SMOV unique PCs	Notes
PTODSL pto.vmi.create_mask(n - i*256) unrolled	8	8	Each SMOV: Sd[64]=0x100 from distinct Sn — matches TileLang with_tail
CCE plt_b8(rem, POST_UPDATE) unrolled	8	1	Only initial SMOV n → rem; POST_UPDATE carries remaining
CCE absolute (bisheng)	8	1	Does not reproduce TileLang; bisheng coalesces / uses SMOVI instead of per-tile SMOV→S64

PTODSL absolute SMOV detail (same shape as TileLang):

RV_SMOV Sd[64]=0x100 from Sn[4]
RV_SMOV Sd[64]=0x100 from Sn[6]
RV_SMOV Sd[64]=0x100 from Sn[8]
RV_SMOV Sd[64]=0x100 from Sn[10]
RV_SMOV Sd[64]=0x100 from Sn[12]
RV_SMOV Sd[64]=0x100 from Sn[14]
RV_SMOV Sd[64]=0x100 from Sn[16]
RV_SMOV Sd[64]=0x100 from Sn[18]

Repro paths:

• Absolute (problem): plt_post_update_microbench/abs_create_mask_vmi.py → opsim out/opsim_vmi_abs_create_mask
• Post-update (ideal): run_cce_launch.py --mode post_update_unroll → out/opsim_post_update_unroll_v2

So the ask stands: VMI needs (mask, remained) from PLT POST_UPDATE so an unrolled 8-tile vdhist sweep can keep one scalar in the PLT operand reg instead of rematerializing 8 absolute counts via SMOV.

[peanutchan]

Isolated performance: only 8×PLT + SMOV vs 8×PLT POST_UPDATE

Added plt_smov_only.py, with identical vector work in both kernels:

• 8 RV_PLT
• 8 dependent RV_VADD (keeps every mask live)
• same VBR/PSET/VST/MTE3

Only count delivery differs:

1. independent: eight runtime counts (256 each) → 8 RV_SMOV into fixed S64
2. post_update: one runtime remaining (2048) → 1 initial RV_SMOV, then carried PLT scalar_out

Both PASS with identical output.

Metric	8 independent counts	POST_UPDATE carry	Delta
RV_PLT	8 / 48 cycles	8 / 48 cycles	0
RV_SMOV	8 / 8 cycles	1 / 1 cycle	-7 calls / -7 RVECSU cycles
RV_VADD	8 / 56 cycles	8 / 56 cycles	0
RVECSU aggregate	11 / 55 cycles	4 / 48 cycles	-7 cycles
block cycles	2124	2123	-1 cycle (0.047%)
reported duration	0.79 us	0.79 us	no visible change

Conclusion: this optimization does remove exactly seven redundant SMOVs per 8-tile sweep, but SMOVs overlap with the unchanged PLT/VADD execution, so their isolated end-to-end impact is only 1 simulator cycle here. It should not be presented as explaining the full group_count tail slowdown; the 8 PLTs themselves remain.

[peanutchan]

PSET baseline: PLT latency is the dominant cost

Added a third mode to plt_smov_only.py: one reusable PSET ALL mask + the same 8 dependent VADDs.

All three kernels now have:

• identical 8-scalar entry ABI
• identical dynamic scalar loads and 7-add output-offset chain (forces all arguments live)
• identical 8 dependent VADDs, VST, MTE3 and output validation

All PASS.

Path	PSET	PLT	SMOV	RVECEX cycles	RVECSU cycles	block ticks	duration
PSET ALL baseline	1 / 6 cyc	0	0	74	53	1736	0.58 µs
PLT POST_UPDATE	1 / 6 cyc	8 / 48 cyc	1 / 1 cyc	122	48	2120	0.79 µs
8 independent counts	1 / 6 cyc	8 / 48 cyc	8 / 8 cyc	122	55	2125	0.79 µs

Deltas

• PSET → POST_UPDATE PLT: +48 RVECEX cycles; +384 block ticks; +0.21 µs / +36% in this tiny slice
• POST_UPDATE → independent SMOVs: +7 RVECSU issue cycles; only +5 block ticks here

The +384 block ticks exactly tracks 48 VF cycles × 8 simulator core ticks/VF cycle; SET_FLAG/WAIT_FLAG completion is delayed by the PLT chain.

Conclusion: the group_count tail latency increase is primarily the eight PLTs on the consumed-mask dependency path, not the seven redundant SMOVs. POST_UPDATE is still cleaner and removes SMOV pressure, but it retains the eight PLTs, so it cannot recover the no-tail/PSET performance by itself.

[peanutchan]

General inner histogram loop: PSET full tiles + at most one tail PLT

Simulated the proposed general VF form with real dhistv2:

full_mask = pset_b8(PAT_ALL);
full_tiles = n / 256;
tail = n % 256;
for (i = 0; i < full_tiles; ++i)
    dhistv2(hist, load(i), full_mask, Bin_N0);
if (tail != 0)
    dhistv2(hist, load(full_tiles), plt_b8(tail, POST_UPDATE), Bin_N0);

Compared with fully-unrolled 8× plt_b8(rem, POST_UPDATE); identical 8 DHIST operations and both outputs PASS.

n	path	PLT	PSET	DHIST	block ticks	duration
2048	8× post-update PLT	8 / 48 cyc	1	8 / 120 cyc	2418	0.96 µs
2048	full loop + tail	0	1	8 / 120 cyc	2338	0.93 µs
1800	8× post-update PLT	8 / 48 cyc	1	8 / 120 cyc	2418	0.96 µs
1800	full loop + tail	1 / 6 cyc	1	8 / 120 cyc	2338	0.93 µs

Result: -80 block ticks / ~-0.03 µs / 3.3% faster for both aligned and partial cases.

Although this removes 8 PLTs aligned or 7 PLTs for a partial tail, the dynamic VLOOP + dynamic VLD/address work overlaps/adds overhead, while DHIST remains dominant (120 RVECEX cycles). It is still the best general mask shape tested, but the measured full-dhist gain is modest rather than the 36% seen in the predicate-only slice.