Lower PTOAS no-tmp frontend ops to PTO-ISA tmp-parameterized interfaces with implicit-tmp materialization

[FangRui0]

Background

PTOAS ODS 提供了一批 tile op 的 无 tmp 前端接口（tmp operand 为 optional，可省略）。此前当用户不写 tmp 时，PTOAS 会继续 lowering 到后端不带 tmp 的 C++ 接口。

本 PR 改为：将这些无 tmp 前端接口 lowering 到 PTO-ISA 的有 tmp 参数的接口，由 PTOAS 构造 tmp，并交由 memplan 统一规划 tmp 与所有变量的地址。

What this PR does

新增 PTOMaterializeImplicitTmp pass（lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp），在 memplan 之前为省略了 tmp 的 op 构造一个隐式 tmp buffer，并改写到带 tmp 参数的 PTO-ISA 形态。覆盖以下 op：

• 标量初始化：TCI
• 行展开二元（mode-1 col-major src1 需 tmp 展开）：TRowExpandAdd/Sub/Mul/Div/Max/Min
• 可选 tmp 的算子：TColSum、TGather、TQuant、TPow/TPowS、TSort32、TMrgSort、TXor/TXorS、TCvt
• 行归约/行 ArgMax/ArgMin、TRowSum/RowProd/RowMax/RowMin、TRem/RemS、TPRelu、TSel/SelS、TTrans 等（通过 replaceFixedDpsOpWithTmp / replaceRowReductionWithTmp / materializeFixedMandatoryTmp 统一处理）

同时在各 op 的 verifier 中加入隐式 tmp 契约校验：当 tmp 无法被隐式构造（如 level3/native 路径跳过 PlanMemory，或 src 为动态 valid-shape 无法推导 tmp 类型）时，给出明确报错。

Key design points

• requireExplicitTmp = (level == Level3)：level3/native 路径跳过 PlanMemory，无法构造隐式 tmp，此时 mode-1 行展开 / 动态形状 TQuant 等 op 必须显式给 tmp，否则 verifier 拒绝。
• level2（含 PlanMemory）路径自动构造隐式 tmp，并把它纳入 memplan 的统一地址规划。

示例（trowexpand_implicit_tmp_materialization 实际生成 IR）

mode-1（src1 为 8x1 col-major，需要 tmp）：

%0 = pto.alloc_tile addr = 0    : !pto.tile_buf<vec, 8x16xf32>
%1 = pto.alloc_tile addr = 512  : !pto.tile_buf<vec, 8x1xf32, blayout=col_major>
%2 = pto.alloc_tile addr = 768  : !pto.tile_buf<vec, 8x16xf32>
%3 = pto.alloc_tile addr = 1280 : !pto.tile_buf<vec, 1x2048xf32>   ← ptoas 构造的 tmp
pto.trowexpandadd ins(%0, %1, %3, ...) outs(%2, ...)              ← lowering 到带 tmp 接口

mode-2（src1 为 8x8 足够大，可复用，无需 tmp）：

pto.trowexpandadd ins(%0, %1, ...) outs(%2, ...)                  ← 保持无 tmp

memplan 给变量与 tmp 分配了互不重叠的地址（0/512/768/1280）。

Validation

远端 A3 环境（LLVM 21.1.8 工具链，aarch64）：

• 构建：433/433 target，ptoas 0.56
• 全套 check-pto：1565 passed / 0 failed / 1 unsupported（unsupported 为环境正常项）

实现逻辑经实际生成 IR 验证正确（见上文 trowexpand 示例：无 tmp 前端接口 → lowering 到带 tmp 的 PTO-ISA 接口，PTOAS 构造 tmp，memplan 统一规划地址）。

[zhangstevenunity]

Review: implicit-tmp materialization (deep correctness pass)

I built this branch (LLVM 21, clang++, check-pto = 1648 passed / 0 failed / 1 unsupported) and cross-checked every synthesized tmp size against pto-isa (include/pto/npu/{a2a3,a5} + docs/isa/*.md).

What is right

The tmp specs are accurate -- this is the strongest part of the PR:

• TCI 768 B (b32) / 1792 B (b16) matches docs/isa/TCI.md exactly.
• TCVT head/tail/half-to-i8 formulas match docs/isa/TCVT.md term for term.
• TROWEXPAND* ceil(R/8)*256 and 30*256 match docs/isa/TROWEXPANDADD.md, and rejecting mode 2 with an explicit tmp is correct ("This overload only supports Mode 1") -- that closes a real latent miscompile, since master accepted mode-2 IR and still emitted the 4-arg overload.
• TMRGSORT tmp.Cols >= sum(src.Cols) matches docs/isa/TMRGSORT.md; the a5 smoke test really was undersized (1280 < 1536), so this fixes a latent OOB.
• TSEL 8/16 B, TREM 2 rows, TREMS 1 row, TXOR full-dst all match the headers.
• The A5 getEffects gating is right: I confirmed a5/TPow.hpp, a5/TCvt.hpp, a5/TGather.hpp, a5/TXor.hpp, a5/TSel.hpp, a5/TRowReduce.hpp all take tmp and never touch it.
• Adding Read alongside Write on tmp is correct (scratch is read-modify-write), and the PlanMemoryModern forbid-alias extension to non-DPS inputs is the right companion.

Blocking

P1 -- A5 + level3 silently emits an unaddressed tmp tile. Row reductions, txor/txors and the tsel/tsels/tprelu/trem/trems group use requireExplicitTmp && !isA5, so on A5 at level3 they skip the "requires explicit tmp" error and still create pto.alloc_tile with no addr. PlanMemory is skipped at level3, and ptoas's own "alloc_tile requires a base addr at level3" guard runs before this pass, so the synthesized alloc escapes both. Reproduced -- details inline.

Non-blocking but user-visible

• P2 pto.tprelu with an explicit tmp and a dynamic dst valid_shape is now a hard error. tmp was mandatory pre-PR, so this rejects IR that exists today.
• P2 dynamic-valid-shape pto.tquant on A2/A3 loses its no-tmp fallback and becomes a compile error (tquant_no_implicit_tmp_a3.pto flipped from a passing test to not ptoas).
• P2 the mode-1 trowexpand tmp is a fixed 8 KB where the PR's own verifier computes 256 B for validRow = 8 -- new UB pressure at level2 where master needed zero planned bytes.
• P3 the tsel tmp dtype tightening rejects previously valid IR for no backend reason.
• P3 a stray operandSegmentSizes discardable attribute is attached to 15 ops that do not carry AttrSizedOperandSegments.
• P3 two unrelated level3 sync regression guards were edited without need (verified: their pre-PR versions still pass against this build).
• Nits: two new lit files have no license header; tquant_no_implicit_tmp_a3.pto had its header text changed from "THIS SOFTWARE" to "THIS PROGRAM".
• Heads-up: the raw generated builder order changed (TRowMaxOp(src, tmp, dst) -> TRowMaxOp(src, dst, tmp=...)); for tsel/txor/ttrans/row-reductions an old positional call now binds dst<-tmp and tmp<-dst with both being tile_bufs, so it can bind silently. The ptodsl surface wrappers are preserved, so this only affects raw-binding users like test/samples/*.py -- worth a release note.

Requesting changes for the P1 only; everything else is a judgement call for you.

[zhangstevenunity]

P1 (blocking): on A5 at level3 this materialises a tmp alloc_tile that never gets an address.

requireExplicitTmp && !isA5 here (and at lines 450 / 474, and via isA5 ? false : requireExplicitTmp in materializeFixedMandatoryTmp at lines 561 / 577 / 592 / 610 / 625) means A5 + level3 skips the diagnostic and falls through to createAllocTmp, which emits pto.alloc_tile with no addr. level3 skips PlanMemory, and the "pto.alloc_tile requires a base 'addr' operand when --pto-level=level3" guard in tools/ptoas/ptoas.cpp runs before this pass, so the synthesized alloc is checked by neither.

Reproduced on this branch (ptoas --pto-arch=a5 --pto-level=level3) with a pto.trowmax and a pto.tsel that have no tmp:

%0 = pto.alloc_tile addr = %c0_i64    : !pto.tile_buf<vec, 2x16xf32>
%1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf<vec, 2x8xf32, valid=2x1>
%2 = pto.alloc_tile                   : !pto.tile_buf<vec, 1x8xf32>   <-- no addr
pto.trowmax ins(%0, %2 : ...) outs(%1 : ...)

and the emitted C++ (rc=0, no diagnostic):

Tile<TileType::Vec, float, 2, 16, ...> v7 = v6;
TASSIGN(v7, v8);                      // every user tile gets TASSIGN
...
Tile<TileType::Vec, float, 1, 8, ...> v12;
Tile<TileType::Vec, float, 1, 8, ...> v13 = v12;   // no TASSIGN at all
TROWMAX(v10, v7, v13);

pto/common/pto_tile.hpp leaves data_ uninitialised in the default Tile() ctor outside __PTO_AUTO__, so v13 carries an undefined UB address into the call. It is latent today only because A5's TROWMAX/TSEL/TXOR/TPRELU/TREM all ignore tmp -- but it is exactly the situation the design doc rules out ("level3 ... does not automatically create an unaddressed tmp"), and it is inconsistent with the rest of the pass: tprelu/tcvt/tmrgsort/tpow all raise "requires explicit tmp when PlanMemory is skipped" on A5+level3, and TTransOp (line 655) passes requireExplicitTmp straight through, so ttrans errors while tsel next to it silently materialises.

Note that A5 still needs a placeholder at level2 (a5 TROWMAX_IMPL has no 3-arg overload, and PTORowMaxToEmitC calls peelUnrealized(adaptor.getTmp()) unconditionally, which would null-deref on an absent tmp) -- so the fix is not "skip A5", it is "use requireExplicitTmp uniformly" so A5+level3 raises the same error as A2/A3.

Test gap that hid this: implicit_tmp_a5_skip_no_tmp.pto is the only A5+level3 test and it only covers tci/tcvt/tquant/trowexpandadd, i.e. exactly the four ops that return early on A5.

[zhangstevenunity]

P2: this rejects IR that compiles on main today.

pto.tprelu's tmp was mandatory before this PR, so every existing A2/A3 tprelu already carries one -- and any of them built on a dynamic-valid-shape tile now fails to verify:

error: 'pto.tprelu' op expects A2/A3 tprelu dst valid_shape to be static when tmp is provided

(reproduced on this branch with pto.alloc_tile valid_row = %vr valid_col = %vc : !pto.tile_buf<..., v_row=?, v_col=?, ...>; on main this verifies fine, since the two checks immediately above are both guarded with != ShapedType::kDynamic).

Same pattern in two more places added by this PR:

• lib/PTO/IR/PTO.cpp:6322 / :6364 -- "expects static src valid_shape and element size to verify tcolsum tmp" (binary tcolsum).
• lib/PTO/IR/PTO.cpp:6441 -- "expects static src shape and dst valid_shape to verify tcvt tmp".

Suggestion: make the capacity check conditional the same way the neighbouring checks are -- skip it when the valid shape is dynamic instead of turning the whole op into an error. A capacity bound that cannot be computed should not invalidate the op.

[zhangstevenunity]

P2: dynamic-valid-shape TQUANT on A2/A3 loses a working path.

tquant_no_implicit_tmp_a3.pto went from a passing test to RUN: not ptoas in this PR. On main that kernel lowered through the 4-arg no-tmp overload; now it is a hard compile error.

makeSameShapeTmpType bails because the valid shape is dynamic, even though the allocated shape (32x32) is perfectly static. Two ways out, both already in the codebase:

1. pto.alloc_tile accepts valid_row / valid_col operands -- forward the src's dynamic valid dims into the synthesized tmp (the A2/A3 tquant verifier wants verifyTileBufSameValidShape(src, tmp), which that satisfies), or
2. fall back to preserving the no-tmp form when the tmp cannot be constructed, exactly as materializeTRowExpandTmp already does at level3 (PTOMaterializeImplicitTmp.cpp:257-263).

Turning a previously compiling kernel into an error is the one outcome that is strictly worse than either.

[zhangstevenunity]

P2: 8192 bytes is ~32x the documented minimum, and it is new UB pressure that master did not have.

getTRowExpandTmpMinBytes (lib/PTO/IR/PTO.cpp:909) already computes the exact bound from docs/isa/TROWEXPANDADD.md -- ceil(R/8)*256 for R < 256, 7680 otherwise. makeTRowExpandTmpType ignores it and always allocates 8 KB.

Measured on this branch, a level2 A3 kernel with 8x704xf32 tiles and one validRow = 8 pto.trowexpandmul:

%8 = pto.alloc_tile addr = %c180224_i64 : !pto.tile_buf<vec, 8x1xf32, blayout=col_major>
%9 = pto.alloc_tile addr = %c180480_i64 : !pto.tile_buf<vec, 1x2048xf32>   <-- 8192 B

188,672 of the 192 KB UB. The documented requirement for validRow = 8 is 256 bytes. On master this kernel needed zero planned bytes for the tmp (mode 1 used pto-isa's internal TMP_UB_OFFSET scratch), so any existing level2 kernel sitting within 8 KB of the UB ceiling now fails to plan, and the diagnostic will point at memory planning rather than at this pass.

Since the dst valid_row is known statically here, cols = getTRowExpandTmpMinBytes(dstValid[0]) / elemBytes would give the same guarantee at 1/32 the cost. Keeping 8192 only for the dynamic-valid-row case would be reasonable.

[zhangstevenunity]

P3: this dtype constraint rejects previously valid IR without a backend justification.

a2a3/TSel.hpp does __ubuf__ uint32_t* cmpMaskPtr = (__ubuf__ uint32_t*)__cce_get_tile_ptr(tmp); -- it reinterprets the tmp buffer regardless of TmpTile::DType and writes only cmpmaskLen uint32s (8 B for b32 data, 16 B for b16). A 2x128xf16 tmp is 512 bytes and 32-byte aligned, so it satisfies the real contract completely; only the declared element type differs.

The byte-capacity check you added two lines below is the genuine constraint, and it is a good addition. The dtype check on top of it is what forced tsel_bf16.pto, select_tile_native.pto and plan_memory_inplace_forbid_alias.pto to be rewritten in this PR -- which is direct evidence that downstream kernels use a same-dtype tsel tmp and will now stop compiling:

error: 'pto.tsel' op expects A2/A3 tsel tmp element type to be 4 bytes wide

Either drop it and keep only the capacity bound, or keep it and call it out as an intentional breaking change in the PR description / release notes. Note tsels right below takes the opposite rule (tmp elem must equal the data elem), so the two ops now disagree on what a tmp dtype means.

[zhangstevenunity]

P3: this writes an operandSegmentSizes attribute onto ops that do not have AttrSizedOperandSegments.

None of the 15 ops switched to hasCustomAssemblyFormat in this PR carry the AttrSizedOperandSegments trait (they each have exactly one optional operand, so MLIR infers the operand index from the operand count). Confirmed with pto-test-opt --mlir-print-op-generic:

%0 = "pto.alloc_tile"() <{operandSegmentSizes = array<i32: 0, 0, 0>}> : ...   // inherent, real
"pto.trowmax"(%0, %1) {operandSegmentSizes = array<i32: 1, 0, 1>} : ...       // discardable, dead

Note the {...} vs <{...}>: on pto.trowmax it lands in the discardable dictionary, where nothing validates or reads it. Consequences:

• it is added only on the parse path -- ops built by the C++ / Python builders do not get it, so two structurally identical pto.trowmax ops now have different attribute dictionaries and will not be merged by CSE / OperationEquivalence;
• rebuildWithOperands in the new pass strips it (copyAttrsExceptOperandSegments), so it silently disappears after materialisation;
• if the trait is ever added to these ops, a stale 1, 0, 1 left behind by a rewrite that inserted a tmp would mis-index the operands rather than fail loudly.

Same in parseOptionalTmpFixedDpsOp (line 12495), TCvtOp::parse (line 9642), TXorOp::parse and TXorSOp::parse. Simplest fix is to drop the result.addAttribute(...) calls for the ops without the trait (and then the elidedAttrs entries in the printers are unnecessary too).

[zhangstevenunity]

P3: this edit is not needed, and it removes the only coverage of the path it was guarding.

materializeTRowExpandTmp deliberately preserves the no-tmp form at level3 (PTOMaterializeImplicitTmp.cpp:257-263, and rowexpandadd_level3_no_tmp_preserved.pto asserts exactly that), so these level3 kernels are unaffected by the pass.

Verified by running the pre-PR versions of both files against this branch's ptoas:

issue646_pipev_repeat_prune.pto  : ptoas rc=0, FileCheck rc=0
issue533_loop_zero_trip_sync...  : ptoas rc=0, FileCheck rc=0

Both still pass unchanged. Since issue #646 is specifically about barrier pruning between TADD and a no-tmp TROWEXPANDDIV -- which is what the level3 frontend actually emits -- adding a tmp here weakens the regression guard. Suggest reverting both files (4 sites in this one, 3 in issue533_loop_zero_trip_sync_regression.pto and its _gss twin).

[zhangstevenunity]

Nit: this new file and test/lit/pto/tci_implicit_tmp_level3_invalid.pto have no license header. .claude/CLAUDE.md requires the PR386 OAT.3 header on new files, and the other 16 files added by this PR all have it.

Related, in test/lit/pto/tquant_no_implicit_tmp_a3.pto the PR changed line 5 from "THIS SOFTWARE IS PROVIDED" to "THIS PROGRAM IS PROVIDED" -- that looks unintentional; "THIS SOFTWARE" is the canonical wording used everywhere else including PTOMaterializeImplicitTmp.cpp. test/lit/pto/tci_implicit_tmp_materialization.pto has the same typo.

[zhangstevenunity]

Heads-up rather than a defect: making tmp optional reorders the raw generated builder from (src, tmp, dst) to (src, dst, tmp=None), which is why these samples had to change.

For TSelOp, TXorOp, TTransOp, TColArgMax/MinOp and the row reductions, the old and new positional slots are all tile_buf, so an existing caller written as pto.TRowMaxOp(src, tmp, dst) will still bind -- with dst and tmp swapped. pto.TSelSOp is the loud case (an i16 scalar lands where a tile is expected), but the others can fail silently.

The ptodsl surface wrappers keep their (src, tmp, dst) signature, so the blast radius is limited to code using the raw _pto / pto.*Op bindings directly -- which is what test/samples/*.py do. Worth a line in the PR description or release notes so downstream raw-binding users know to audit their call sites.