ConvertTensorToTileOps Pass¶

Converts tensor operations to tile operations in InCore functions and updates orchestration call sites.

Overview¶

After OutlineIncoreScopes extracts InCore scopes into separate functions, those functions still operate on TensorType variables using tensor.* operations. This pass lowers them to TileType variables with tile.* operations that map directly to PTO-ISA instructions.

The pass also updates call sites in orchestration/opaque functions: for each new output parameter added to an InCore function, a tensor.create is inserted at the call site.

Requirements:

Input IR must be in SSA form
InCore scopes must be outlined (run OutlineIncoreScopes first)
Statement structure must be normalized

When to use: Run after OutlineClusterScopes and before OptimizeOrchTensors.

API¶

C++	Python	Level
`pass::ConvertTensorToTileOps()`	`passes.convert_tensor_to_tile_ops()`	Program-level

Python usage:

from pypto.pypto_core import passes

convert_pass = passes.convert_tensor_to_tile_ops()
program_tiled = convert_pass(program)

Algorithm¶

The pass operates in three program-level phases:

Phase 1: Transform InCore Functions¶

For each FunctionType::InCore function:

Pre-scan consumer memory demand (ConsumerSpaceCollector): Collect, for every variable, the memory space its downstream consumers need, driven by the input_reqs declared in OpConversionRegistry. A tensor.slice feeding tensor.matmul is the motivating case — it needs a Mat tile.load (natural, plus a zero-copy tile.transpose_view when transposed) rather than landing in Vec and being moved out again — but the mechanism is general, not matmul-specific.
Analyze writes and share read-only loads: Before conversion, trace writes through GM aliases and control flow using operator argument effects. A parameter used by a default tile-input converter may share an entry tile.load only when it is not written in the function. Writes to local values derived from a parameter also conservatively disable sharing for that parameter. Calls with opaque effects disable entry-load sharing. The loaded tile lives in a separate operand cache; the GM parameter retains its identity and type. Load space follows consumer demand, or remains unset for InferTileMemorySpace to determine.
Convert body via TensorToTileMutator: Convert each registered call using OpConversionRegistry. BridgeInputSpaces supplies tile operands at the consuming statement, reusing read-only entry loads for default inputs or loading mutable GM operands at each use. Memory converters retain their GM operands and perform their own loading/storing. Local computed results still map to tiles. A value-flow analysis identifies tile-valued loop/branch results: GM handles carried through writes remain GM, while computed carries receive tile seeds and compatible yields.
Insert tile.store (exit stores): For each return value converted from TensorType to TileType, add an Out parameter and insert tile.store(tile, zeros, out_param). If the return value comes from a tile.assemble loop, the loop is rewritten to use tile.store directly (conversion-time assemble-loop rewrite; distinct from OptimizeOrchTensors Pattern 3 which handles cross-function optimization).
Upgrade written param directions: An alias-origin analysis (AnalyzeCallAccess) attributes every read/write back to the parameter it originates from, then upgrades each In param that is written to Out (write-only) or InOut (read and written). Which argument an operator writes is not decided here — it is read from that operator's registry declaration (set_arg_effect, see Operators), so tile.store, tile.mscatter, tensor.write, tensor.assemble, tensor.expand_clone, the pld.tile.* / pld.tensor.* push and pull family, pld.system.notify, system.syncall and the composite collectives all reach the same analysis through one table. An argument declared Write is not counted as a read: a store landing on a sub-region never reads the untouched remainder. Kwarg-dependent effects resolve per call, so an atomic store or an AtomicAdd notify marks its destination read+write while the plain forms do not. Params the user already declared Out / InOut are left as-is.

An operator that never declared its effects still counts as reading every argument. That default is now confined to operators the registry has nothing to say about — most of them functional — rather than being the fallback for a hand-maintained list that a new write operator silently escaped.

GM Identity and Read/Write Ordering¶

Loading a GM tensor for computation creates a tile value, not an alias that can replace every use of the tensor. For example, this kernel computes from x and then changes one GM element:

@pl.jit.incore
def kernel(
    x: pl.InOut[pl.Tensor[[16, 32], pl.FP32]],
    out: pl.Out[pl.Tensor[[16, 1], pl.FP32]],
):
    r = pl.row_max(x)
    out[0:16, 0:1] = r
    pl.tensor.write(x, [0, 0], pl.const(1.0, pl.FP32))
    return out

The reduction receives a tile.load(x, ...) result. The final write remains tensor.write(x, ...), emitted as a GM pto.store_scalar; it does not become a tile.write into the reduction's input tile. The same rule applies to pld.DistributedTensor parameters, scalar reads, and returned GM aliases.

When a later computation reads x after a GM write, it loads the updated value at that statement. Mutable GM loads are not shared across writes, branches, or loop iterations. Returning another tensor, or returning nothing, does not discard the GM write. InOut describes argument effects; it does not request an unconditional whole-tensor store at function exit.

GM Store Coherence Restriction¶

An InCore function may not combine a bulk store (tensor.assemble, lowered to tile.store) with tensor.write on the same underlying GM tensor. The bulk store uses MTE3 while the scalar store uses the D-cache path; barriers alone do not guarantee cache-line coherence between those paths on A2/A3. The pass therefore rejects the combination instead of emitting code that can silently lose either the scalar overrides or neighbouring bulk-written bytes.

Before applying the restriction, the pass canonicalizes a simple constant scalar-fill loop into tensor.full plus tensor.assemble, which then lowers to tile.full plus tile.store. The loop must be sequential, zero-based, and unit-stride; its body may contain one or more tensor.write operations and nothing else. Each write must cover a contiguous region with one loop-invariant constant value, and the flattened region must satisfy MTE3's 32-byte row alignment. This keeps full-block fallback fills on the MTE3 path instead of rejecting them as mixed stores. Dynamic values, non-canonicalizable partial updates, unaligned regions, and strided scalar loops remain on the D-cache path.

The pass also stages a full-tensor tensor.full initialization followed only by scalar updates to that tensor. The scalar updates are redirected to the local tile, and one tile.store writes the completed value to GM at function exit. This supports dynamic sparse-map construction without mixing MTE3 and D-cache stores. The rewrite requires a zero-offset, full-shape initialization backed by a private tensor.full result; partial initialization, aliases, additional DMA stores, or other uses of the GM target remain rejected.

The check follows assignment aliases and loop/branch carries. It is deliberately conservative: it rejects mixed store paths to one GM tensor even when source offsets appear disjoint, because the compiler does not yet prove cache-line separation across symbolic views and control flow. Mixed paths to distinct GM tensors remain valid.

# Rejected: a partial assignment becomes MTE3 TSTORE, then pl.write uses D-cache.
output[0:1, 0:16] = pl.full([1, 16], dtype=pl.INT32, value=-1)
for i in pl.range(4):
    pl.write(output, [0, i], pl.cast(i, pl.INT32))

# Supported: update the local on-chip value, then issue one GM store.
staged = pl.full([1, 32], dtype=pl.INT32, value=-1)
for i in pl.range(4):
    pl.write(staged, [0, i], pl.cast(i, pl.INT32))
output[0:1, 0:32] = staged

Using tensor.write for every element is also supported when a single bulk store is not practical.

Cache-Policy Declarations → `tile.load` `cache` Kwarg¶

This pass is where a declared GM cache policy stops being metadata and becomes part of the access. OutlineIncoreScopes left the declarations on the InCore function as the cache_policy attr — std::vector<std::pair<int32_t, int>> (param index, CachePolicy as int). Phase 1 turns those indices back into param Var identities once per function, then adds {"cache", <policy>} to every tile.load whose source arg is a listed param: the entry loads it synthesises, the consumer-driven Mat loads, the input-space bridge loads, and any tile.load already in the body (user-written or produced by an earlier pass). The attr is erased when the transformed function is rebuilt — nothing downstream may see it, because its param indices go stale the moment a later pass grows the param list.

Precedence is per access: an explicit pl.load(..., cache=...) kwarg already on the load always wins over the scope declaration, in both directions, so cache=pl.CachePolicy.DEFAULT opts one read back into the cache inside a bypassing scope. From here the kwarg simply rides the op through the remaining passes to codegen, the way target_memory does. See GM Cache-Access Policy.

Phase 2a: Propagate Added Outputs Through Spmd/Group Wrappers¶

OutlineClusterScopes produces Spmd/Group wrappers that are transparent 1:1 forwarders of their params to a single inner InCore call. When Phase 1 appends Out params to that InCore callee, the wrapper must mirror the appended params on its own signature and forward them through the inner call — otherwise orchestration codegen's BuildWrapperReorderedParams invariant (every inner-call Var arg resolves to a wrapper param) breaks.

For each FunctionType::Spmd / FunctionType::Group function:

ForwardedCallFinder locates the first call to a transformed InCore (one whose Phase 1 added at least one Out param).
If found, the wrapper signature is extended with matching Out params (same type as the InCore's appended params, reusing the name_hint_), and WrapperForwardMutator rewrites the inner call to append the new vars as forward args and adopt the callee's new return type. tensor.create is not synthesised in the wrapper — allocation remains the caller's responsibility.
If no forwarded transformed-InCore call is found, the wrapper is left unchanged.

Phase 2b: Update Orchestration Call Sites¶

For each orchestration / opaque function that calls a transformed InCore function or a wrapper that absorbed output params in Phase 2a:

Insert tensor.create for each added output parameter
Append created tensors as extra arguments to the call

InCore, Spmd, and Group functions are skipped from this phase — they were already rewritten in Phase 1 / 2a.

Where a Bridged Load Lands¶

An operand listed in a conversion's input_reqs reaches the converter as a tile: if it is still a TensorType, BridgeInputSpaces synthesizes the tile.load. That load has to name a memory space, and the space is derived from the consuming operator's own declaration, never written a second time here.

Each entry names the operand rather than the space:

// tensor.matmul: lhs and rhs are bridged as the cube operands of whichever
// matmul the rank dispatch emits.
{{0, {BridgeSpaceOf({"tile.matmul", "tile.batch_matmul"}, 0), "a_trans", /*cube_m_axis=*/true}},
 {1, {BridgeSpaceOf({"tile.matmul", "tile.batch_matmul"}, 1), "b_trans"}}}

BridgeSpaceOf reads that argument's set_input_memory constraint out of OpRegistry and stages it through StagingSpaceForLoad. OpRegistry is therefore the single place an operand's memory space is stated, for this pass and for InferTileMemorySpace alike.

A registration cannot write a space of its own. InputSpaceReq::demanded_space is an OperandSpace, a type with no public constructor and no conversion from MemorySpace, so {MemorySpace::Vec, ...} does not compile. The only way to obtain one is BridgeSpaceOf / OperandSpaceOf, which is what makes "derived, never authored" an invariant of the type rather than a convention the next registration can quietly ignore.

A requirement names every operator that can consume the operand. The list is braced even when there is one (BridgeSpaceOf({"tile.sel"}, 2)). It matters when a converter dispatches: tensor.matmul emits tile.matmul for a 2-D operand and tile.batch_matmul for an ND one, so deriving from the flat op alone would read an operator an ND call never becomes. All listed operators must declare the same constraint. That agreement is structural rather than lucky — FlattenTileNdTo2D unrolls the batched op into the flat one, so the same operand lands in the same buffer either way — and a divergence throws instead of silently bridging against the wrong name.

The constraint and the load target are different spaces, and that is the point. tile.matmul constrains its operands to Left / Right — L0A and L0B — but MTE2 fills no L0 buffer, so a load can never satisfy the constraint directly:

Operand constraint	Load lands in	Last hop
`Vec`	`Vec`	none
`Mat`	`Mat`	none
`Left` / `Right` / `Bias` / `LeftScale` / `RightScale`	`Mat`	`tile.move` added by `InferTileMemorySpace` Phase 2
`Acc`	—	not loadable; an accumulator must be created in L0C

StagingSpaceForLoad (src/ir/memref.cpp) holds that table, and InferTileMemorySpace applies the same function when it retargets a producer it placed itself — so a bridged load and an inferred load put the same operand in the same buffer.

Failures are loud and early. Naming an unregistered operator, or one that declares no set_input_memory for that argument, throws while the conversion registry is being built rather than leaving the bridge with a plausible-looking default. A new operator cannot silently fall into the gap between the two registries.

Not every operand goes through this. Two other mechanisms place a tile, and neither is a missing declaration:

Self-loading ops (tensor.slice, tensor.gather, tensor.paged_gather, tensor.gather_row, tensor.assemble, …) synthesize their own loads inside the converter, because the load's shape or offset is part of the lowering.
Inherit-input ops (tile.assemble, tile.gather_row) declare set_output_memory_inherit_input() or set_output_reuses_input(idx). Their space follows the tile they are given — Mat for an L1 paged-gather accumulator, Vec for a UB one — so pinning it with set_input_memory would force a wrong tile.move instead of describing the hardware.

MatmulSlice Pattern¶

This is the motivating instance of the general input_reqs flow above, not a mechanism of its own. When tensor.slice feeds into tensor.matmul or tensor.matmul_acc, the slice must produce a Mat-space tile instead of a Vec-space tile. ConsumerSpaceCollector reads that demand off the matmul's input_reqs like any other consumer's, and the producer answers it with a natural Mat tile.load; a transposed operand (a_trans for LHS, b_trans for RHS) gets a zero-copy tile.transpose_view at the matmul site. Nothing here matches on the operator being a matmul — an operand of any op declaring a non-Vec requirement reaches its producer the same way.

The demand is propagated through zero-copy metadata ops that declare set_output_memory_inherit_input() — tensor.slice, tensor.view, tensor.reshape, tensor.reinterpret_view, tensor.set_validshape. So an operand written as pl.matmul(pl.set_validshape(a[:, :K], rows, K), b) still loads straight to Mat. An op that aliases its input's storage but omits that declaration breaks the chain: the operand materializes in Vec and needs a tile.move to Mat, which is a vector→cube boundary that flips an otherwise pure-CUBE InCore scope to MIXED and makes ExpandMixedKernel split it into an AIC/AIV pair.

Cube Operand M-Axis Boxing¶

A cube operand has two independent extents, and only one of them is constrained.

The logical extent is essentially free. pto.tmatmul derives M, K, and N from the operands' valid region, bounded at [1, 4095] with no divisibility rule; pto.mad carries a disable_gemv clause whose whole purpose is selecting the L0A organization at %m == 1.
The physical extent must be a whole number of NZ fractal boxes. A box is 16 rows tall on every generation and dtype; it is 32 bytes / sizeof(dtype) wide (16 for FP16, 32 for INT8). ptoas enforces it directly — 'pto.alloc_tile' op expects result boxed tile rows to be a multiple of innerRows (16) — and pto-isa's TExtract repeats it as a static assertion on the Mat source it reads.

So every cube tile the matmul's M axis runs through is allocated with that extent rounded up to the box, and the tensor's true extent declared as valid_shape:

# Before  (M = 100)
y = tensor.matmul(a, b)          # a: Tensor[[100, 256]]

# After
a_mat = pl.tile.load(a, [0, 0], [112, 256], [100, 256], target_memory=pl.Mem.Mat)
y_tile = pl.tile.matmul(a_mat, b_mat)   # Tile[[112, 512]] valid [100, 512]

The padding is free on both axes of the cost model. A tile.load moves only the valid extent, so the DMA is unchanged; and the MAD cost is ceil(M/16) passes, which rounding M up to a multiple of 16 leaves unchanged. The hardware addresses the narrower valid region inside the full box through compact mode, and tile.store writes only the valid rows, so the observable result is identical.

Making M a multiple of 16 here is also what keeps AutoTileMatmulL0 legal at the boundary: that pass picks a 16-aligned tile and peels the remainder, and a multiple of 16 can only be split into 16-aligned pieces, tail included. It therefore needs no boundary special case.

Which axis carries M¶

M does not always land on the row axis. An a_trans operand is loaded naturally and reinterpreted by a zero-copy tile.transpose_view, so its loaded tile has K on rows and M on columns — the box rule follows M to whichever axis holds it (InputSpaceReq::cube_m_axis resolves the axis from the operand's own transpose flag):

# a: Tensor[[128, 100]], a_trans=True — M is the load's column extent
a_mat = pl.tile.load(a, [0, 0], [128, 112], [128, 100], target_memory=pl.Mem.Mat)
a_t   = pl.tile.transpose_view(a_mat)   # Tile[[112, 128]] valid [100, 128]

That also means the granularity differs per tile even though the padded extent must not: an Acc box is 16 rows for every dtype, while a transposed operand's column box is 32 / sizeof(dtype). A transposed INT8 operand therefore needs 32, and the accumulator paired with it has to adopt that same 32 — each taking its own would give a 128-row product against a 112-row accumulator, which tile.matmul_acc rejects. InputSpaceReq::m_align_from_arg names the left operand as the single decider, and the alignment is the lcm of its box and the accumulator's 16 rows (every granularity involved is a power of two, so the lcm is the max).

The rule rides on the demand, not on one code path. Four sites can answer a matmul's operand requirement, and all four box: BridgeInputSpaces for an operand that is still a tensor at the call; HandleConsumerDrivenLoad when a tensor.slice (or any set_output_memory_inherit_input() chain) answers it at the producer; the Phase-1 entry loop when the producer is a parameter, which is what a pl.matmul(pl.set_validshape(a, ...), b) reaches; and HandleBoxedAccCreate for an accumulator, which is allocated rather than loaded. ConsumerSpaceReq therefore carries the box flag alongside the memory space. When several consumers share one producer, the rows are boxed only if every one of them asks for it, so a consumer that reads the tile at its declared physical shape is never handed a padded one.

`tensor.matmul_acc` boxes its accumulator with its operand¶

tensor.matmul_acc takes exactly tensor.matmul's M constraint, because the two cube tiles the M axis runs through must be boxed together: tile.matmul_acc requires the accumulator and the matrix product to agree on physical M, so boxing the left operand alone would trade a ptoas rejection for an operand-mismatch one.

The accumulator is never loaded — nothing but the matrix unit writes L0C — so its allocation is the site that answers the demand. HandleBoxedAccCreate rewrites the tensor.create that seeds it, and the narrowing rides on a separate tile.set_validshape because tile.create takes no valid extent:

# Before  (M = 100)
acc = pl.create_tensor([100, 64], pl.FP32)
c = pl.matmul_acc(acc, a, b)

# After
acc_storage = pl.tile.create([112, 64], dtype=pl.FP32, target_memory=pl.Mem.Acc, compact=True)
acc_tile = pl.tile.set_validshape(acc_storage, 100, 64)   # Tile[[112, 64]] valid [100, 64]
a_mat = pl.tile.load(a, [0, 0], [112, 128], [100, 128], target_memory=pl.Mem.Mat)
c_tile = pl.tile.matmul_acc(acc_tile, a_mat, b_mat)

Two details this path settles that the operand path does not:

The space is stated, not left to InferTileMemorySpace. The plain tensor.create conversion deliberately leaves target_memory unset, having no consumer context to derive it from; here there is one, and it is the same demand that asked for the boxing. Stating it also matters for correctness: an Acc tile's implicit view is boxed NZ, so a seed left unresolved would carry the raw row-major view and disagree with the tile.matmul_acc result it is carried against across a split-K loop.
The demand crosses the loop carry. A split-K accumulator is allocated before the loop and reaches its tile.matmul_acc as an IterArg, so ConsumerSpaceCollector matches operands with AsVarLike (an IterArg carries its own ObjectKind) and records a propagation edge from each IterArg to the value that seeds it.

The boxed accumulator is declared compact. mad lays the product out at a pitch of ceil(validRow/16)*16 — 112 for a 100-row product — while a non-compact reader derives its stride from the physical row count, which the box rounded to 112 or, at a 32-row alignment, to 128. Compact makes every reader recompute the pitch mad actually used; without it AccCompactValid rejects the program (issue #2470).

The `N` granularity¶

M's alignment reconciles the tiles it runs through — a Mat operand and its accumulator (see ResolveCubeMAlignment). N's has to reconcile memory spaces instead, which M's does not:

the right operand is loaded into Mat and then promoted on to Right by InferTileMemorySpace, and Right boxes the same 512-byte fractal under the opposite slayout, which swaps the row and column granularities:

Space slayout fp32 box fp16 box

Mat row_major 16 x 8 16 x 16

Right col_major 8 x 16 16 x 16
the product lands in an Acc, whose N box is 16 for every dtype.

Both numbers are real, and ptoas asks for each against the tile that carries it, so rounding to Mat's 8 alone yields a box that clears the load and is rejected one op later at Right. ResolveCubeNAlignment therefore takes the lcm of all three; every granularity is a power of two, so the lcm is the maximum.

Scope, and what is deliberately left out:

Case	Boxed?	Why
2-D `tensor.matmul` left operand	Yes, on rows	Its row axis is M, whose box height is 16 for every dtype
2-D `tensor.matmul_acc` left operand and accumulator	Yes, on rows	Same axis, same rule — and the op requires the two to agree on physical M, so they move together
`a_trans` left operand, and the accumulator paired with it	Yes, on columns	The natural load's row axis is K, so M is the column extent; the accumulator adopts that operand's column granularity via `m_align_from_arg`
2-D `tensor.matmul` right operand	Yes, on the axis carrying `N`	Columns normally; rows under `b_trans`, whose natural load puts `K` on the columns. `N`'s padded cells land outside the result's valid region, so nothing reads them
`tensor.matmul_acc` right operand	No	Its accumulator has to agree with the product on physical `N` just as it does on `M`, which needs the same cross-tile decider reaching the `tile.create` that allocates it
Either operand's `K` axis	No	Padding the reduction axis would feed uninitialised L1 into the sum unless the hardware masks by valid col, which is unverified; `AutoTileMatmulL0`'s `PH-AT-007` declines it for the same reason
Rank >= 3 operand	No	It lowers to `tile.batch_matmul`, whose rows `FlattenTileNdTo2D` row-packs into one `[B*M, N]` tile — the box rule binds that packed extent, not this dimension
Dynamic extent	No	No compile-time box to round to

Every case this pass declines still reaches a PyPTO-level error rather than a ptoas one: PTO codegen validates the box grid of every pto.alloc_tile it emits.

Transpose Lowering¶

tensor.transpose lowers to a plain 3-arg tile.transpose(input, axis1, axis2). The PTO pto.ttrans instruction needs a scratch workspace tile (same shape/dtype as the source), but that scratch is a pure codegen detail — not a semantic operand. FlattenTileNdTo2D is the sole owner of scratch materialization: it emits the codegen-ready 4-arg form (tile.create + tile.transpose(..., tmp)) for both 2D and per-page >2D transposes, still before the memory allocator runs (so the scratch gets a real UB address). Keeping scratch out of the high-level op means tensor.transpose and the DSL pl.tile.transpose(tile, axis1, axis2) stay 1:1 with the semantic operation.

# Before
y = tensor.transpose(x, 0, 1)

# After (this pass)
y_tile = pl.tile.transpose(x_tile, 0, 1)   # 3-arg, no scratch

# After FlattenTileNdTo2D (scratch materialized there)
transpose_tmp = pl.tile.create(x.shape, x.dtype, target_memory=x.memory_space)
y_tile = pl.tile.transpose(x_tile, 0, 1, transpose_tmp)

Scatter Update Lowering¶

tensor.scatter_update / tile.scatter_update (whole-row scatter, dim=-2 only) lower to a per-element tile.scatter (pto.tscatter) plus a tile.sel preserve-blend. The hardware pto.tscatter writes per element using a flattened destination index (dst.flat[idx[k, c]] = src[k, c]) and treats its dst operand as write-only (unwritten slots are not preserved), so the pass reconstructs the "keep input on unwritten rows" semantics itself.

The whole-row update input[index.flat[k], :] = src[k, :] is expressed as a flat index:

flat_idx[k, c] = index.flat[k] * d + c          # d = feature width (= src cols)

The flat-index arithmetic is built entirely in i32, and only the finished row-major [n, d] index is narrowed to the pto.tscatter-required width (i16 for 2-byte data, i32 for 4-byte) via a single trailing tile.cast. Computing in i32 keeps every intermediate tile in a canonical, 32-byte-aligned, row-major layout — narrowing earlier would either cast a col_major [n, 1] view (which tile.cast mis-orders) or produce an unaligned 2-byte [b, s] tile (cols * 2 bytes is not 32-byte aligned).

Generated PTO op sequence (FP32 [32, 32] input, [2, 8] index, [16, 32] src):

#	PTO op	Produces
1–3	`pto.tload` ×3	`input_tile`, `index_tile`, `src_tile`
4	`pto.tci`	column arange `[1, d]` = `0..d-1`
5	`pto.texpands`	zero template `[n, d]`
6	`pto.tcolexpand`	`col_nd[k, c] = c`
7	`pto.tmuls`	`row_base[k] = index.flat[k] * d` (index reshaped to `[n, 1]`)
8	`pto.trowexpandadd`	`flat_idx = col_nd + row_base` → `[n, d]`
8a	`pto.tcvt`	narrow `flat_idx` i32→i16 (2-byte dtypes only)
9	`pto.texpands`	zeroed scatter base `[m, d]`
10	`pto.tscatter`	`scattered` = src into zeroed base (written = src, unwritten = 0)
11–12	`pto.texpands` ×2	mask zero base `[m, d]`, ones src `[n, d]`
13	`pto.tscatter`	`mask` = ones into zeroed base (written = 1, unwritten = 0)
14	`pto.tcmps`	`pred = (mask != 0)`
15	`pto.tsel`	`out = sel(pred, scattered, input_tile)`
16	`pto.tstore`	write `out` to the output tensor

tile.sel (not input * mask) reconstructs the preserve blend so the lowering emits no pto.tmul, which A2/A3 reject for bf16/i8. The index reshape [b, s] → [n, 1] is a buffer-view realias, not a separate PTO op.

Flat Gather Lowering¶

pl.gather(src, index=idx) and pl.gather(src, idx) emit the existing tensor.gather without a dim attribute. This means flat element indexing, out = src.reshape(-1)[idx]; an explicit dim keeps the existing axis semantics.

The converter chooses the operation from the source's type after earlier producers have been lowered:

Source at conversion	Lowering
GM `TensorType` / local `DistributedTensorType` window	Keep the source in GM; load only the indices if necessary; emit `tile.mgather(..., coalesce="elem")` into Vec
On-chip `TileType`	Reuse a proven packed source, otherwise pack it; allocate index-shaped INT32 scratch; emit `tile.gather`

Flat gather is self-loading: generic operand bridging must not load the whole GM source. Vec indices are reused; their memory space and unboxed row-major layout are rechecked after producer lowering. Scratch/result valid shapes are restored explicitly for tile.gather. See the operator contract for dtype, shape, alignment, and bounds requirements.

Strided sources are packed using tile.extract for floating point or exact integer tile.adds(..., 0) for INT16/INT32, which A2/A3 TEXTRACT cannot copy. The converter receives a read-only ConversionContext containing packed-storage facts collected from earlier lowered producers in one SSA traversal. Registered functional tile results with fresh storage are proven packed; plain aliases and tile.set_validshape preserve that proof. Other views, parameters, and control-flow results remain unknown and keep the copy. An empty TileView.stride alone is not proof: tile.slice can retain its parent's physical pitch without storing it there.

Paged Gather Lowering¶

tensor.paged_gather(src, indices, block_table, ...) gathers scattered rows of a paged KV pool directly into an on-chip buffer (L1 / Mem.Mat by default, or UB / Mem.Vec). The hardware pto.tgather instruction can only write UB, so paged-gather-to-L1 is not an indexed gather instruction — it is a fully-scalar per-row GM → on-chip DMA loop on the Cube core (AIC). src, indices, and block_table are kept as GM tensors (the op is registered self-loading, so the framework does not preload them into Vec tiles).

The pass materializes the loop directly:

rows = tensor.dim(indices, last_axis)                  # runtime gathered-row count
acc  = tile.create([max_indices, size], target_memory=space)   # static on-chip buffer
for i in [0, rows):                                    # ForStmt, iter_arg = acc
    idx   = tensor.read(indices, [i])                  # scalar GM read (pto.load_scalar)
    phys  = block_table[idx // block_size] * block_size + idx % block_size   # scalar
    acc   = tile.gather_row(acc, src, [i, 0], [phys, col_off], [1, size])    # GM->on-chip
    yield acc

tile.gather_row is a DPS op that writes one physical GM row straight into a sub-region of the accumulator: pto.subview of acc + pto.partition_view of src + pto.tload (GM → on-chip) — no pto.tmov. An L1→L1 tmov is unsupported on a2a3 (L1 can only be filled via GM → L1 tload), so the row is loaded straight into the accumulator's sub-region rather than assembled.

Only the small index / page-table metadata is scalar-read from GM; the bulk KV data goes straight GM → L1 via pto.tload and never touches UB — eliminating the GM round-trip that a gather_kv → qk_pv pipeline pays today. is_trans=True (Mat only) loads each row transposed into column offset [0, i], giving the matmul B-operand layout. max_indices sizes the L1 buffer statically; the runtime rows count drives the loop bound, so dynamic gather counts are supported.

Boxed (NZ) sub-region alignment. An L1 (Mem.Mat) accumulator carries the matmul-operand NZ fractal layout, where pto.subview sizes must be whole multiples of the inner box (M0 = 16 rows; C0 = fractal_bytes / dtype_bytes / 16 cols). A per-row gather writes a single row, so tile.gather_row codegen emits a box-aligned physical size (phys_rows = round_up(1, 16), phys_cols = round_up(size, C0)) while marking only the real extent valid (valid = [1, size]); the tload then fills just that row. UB (Mem.Vec, slayout = none_box) tiles have no inner box and use the exact [1, size] size. The gathered L1 tile is consumed by tensor.matmul directly (its natural use as a matmul operand).

Kernel-Driven Gather (`tensor.create_l1` + `tensor.gather_row`)¶

tensor.paged_gather hardcodes its per-row source address (block_table[idx // bs] * bs + idx % bs). When the kernel needs arbitrary gather logic — multi-source selection, invalid-row clamping, overlay pools — it builds the same L1 accumulator itself from two tensor-level primitives, the flexible counterpart of paged_gather:

Op	Lowers to	Role
`tensor.create_l1(shape, dtype, transpose=...)`	`tile.create(target_memory=Mat, transpose=...)`	seed the loop-carried L1 accumulator
`tensor.gather_row(acc, src, dst_off, src_off, shapes, valid_shape=..., transpose=...)`	`tile.gather_row` (DPS)	DMA one caller-addressed GM row into `acc`

Both deduce a TensorType, so the gathered result composes with tensor-level tensor.matmul / softmax; both are registered self-loading (src stays GM). The caller computes src_off and the dst_off slot, then fills the accumulator row by row in its own loop.

Dynamic transfer length (valid_shape). shapes must stay compile-time constant: it becomes pto.subview's sizes, which the PTO dialect types as a static I64ArrayAttr (SubViewOp in PTOOps.td). The optional valid_shape carries the runtime extent instead — it feeds the subview's valid_row / valid_col, declared Optional<Index> SSA operands, and the GM pto.partition_view sizes, which accept a dynamic ? dim. So a dynamic row count changes neither the allocation nor the box alignment below: the sub-region stays statically shapes-sized and only the copy length varies. Omitting valid_shape transfers the whole window, which is the pre-existing behaviour.

This turns a run of consecutive rows whose length is only known at runtime into one call instead of a guarded per-row loop:

kv = pl.create_l1([128, HEAD_DIM], pl.BF16)
# r1 is a runtime Scalar[INDEX] — e.g. a page-boundary split point
kv = pl.gather_row(kv, pool, [0, 0],  [b0, 0], [128, HEAD_DIM], valid_shape=[r1, HEAD_DIM])
kv = pl.gather_row(kv, pool, [r1, 0], [b1, 0], [128, HEAD_DIM], valid_shape=[128 - r1, HEAD_DIM])
oi = pl.matmul(q, kv, b_trans=True)

Bounds are on the written region, not the window. shapes sizes the static pto.subview, so with a runtime dst_offset the declared window may reach past the destination — in the example above run 2 declares rows [r1, r1 + 128) on a 128-row tile. That is sound only because the transfer is bounded by valid_shape, not by shapes: the rows actually written are [r1, r1 + (128 - r1)), which stay inside. The caller's obligation is therefore dst_offset + valid_shape <= dst.shape per dimension; dst_offset + shapes need not fit. The two-run form above is covered on device by test_gather_row_two_run_split.

valid_shape is keyword-only in the DSL — transpose already owned the sixth positional slot, and taking it would silently reinterpret an existing gather_row(..., shapes, True) call. At the IR level it is a positional operand rather than an attr precisely because it may be a runtime value: it has to stay in the use-def chain so SSA/liveness keep the scalar alive. It is rejected together with transpose=True (see below) — that path would need a runtime column extent on a boxed NZ tile, which is unverified on device. The deducer rejects any provable violation of 0 <= valid_shape[i] <= shapes[i]; a symbolic extent it cannot decide is accepted, which is the case the operand exists for.

Transpose (ZN) for a b_trans matmul operand. transpose=True makes the gathered tile a transposed matmul B-operand without a GM round-trip:

tensor.create_l1(..., transpose=True) allocates the transposed Mat (ZN) fractal (blayout = row_major, slayout = col_major) — the layout a b_trans operand carries.
tensor.gather_row(..., transpose=True) places the GM row [r, c] as the L1 column [c, r]. pto.tload does not transpose, so codegen presents src as a DN-strided view (pto.make_tensor_view ... {layout = #pto.layout<dn>}, shape/strides swapped, same base ptr) and partitions the row as a column — the tload then runs DN → NZ, which is the transpose. (paged_gather's is_trans=True shares this tile.gather_row path.) A straight ND → NZ tload would scramble the fractal layout.

AIV-Split Boundary Lowering (`tensor.aiv_shard` / `tensor.aic_gather`)¶

tensor.aiv_shard / tensor.aic_gather are the @pl.jit / pl.spmd author-facing form of the cube↔vector AIV-split boundary. They are emitted by pl.aiv_shard(x) / pl.aic_gather(x) when the operand x is a high-level Tensor (e.g. a pl.matmul result), inside a for aiv_id in pl.split_aiv(...) region:

raw = pl.matmul(q, k, b_trans=True, out_dtype=pl.FP32)   # Tensor, OUTSIDE the region
for aiv_id in pl.split_aiv(2, mode=pl.SplitMode.UP_DOWN):
    h = pl.aiv_shard(raw)     # C->V: tensor.aiv_shard — full [M, N] -> lane half [M/2, N]
    s = pl.softmax(h)         # AIV vector work on the half
    full = pl.aic_gather(s)   # V->C: tensor.aic_gather — halves -> full [M, N]
oi = pl.matmul(full, v, out_dtype=pl.FP32)               # Tensor, OUTSIDE the region

This pass lowers each 1:1 to its tile op (tensor.aiv_shard → tile.aiv_shard, tensor.aic_gather → tile.aic_gather), so from here on the IR is byte-identical to what the AUTO pl.split path produces via LowerAutoVectorSplit (pass 23). ExpandMixedKernel (pass 24) then folds both into the cross-core tpush/tpop machinery.

Constraints (enforced by the tensor-level deducer and the DSL parser, not this pass):

2D-only — UP_DOWN / LEFT_RIGHT are only well-defined on the 2D physical tile view; an N-D operand is rejected with a pl.reshape-to-2D hint (an N-D tensor would flatten to [product(leading), last], so a pre-flatten row split would not match the contiguous half the lowering physically takes).
Region-only — the tensor.* form is reachable solely through the pl.split_aiv region (which supplies the split mode). The outlined low-level pl.tile.aiv_shard(t, split=N) form stays tile-only; a Tensor operand there is rejected.
Distributed rejected — a DistributedTensorType operand is out of scope (AIV/AIC split only) and is rejected upstream.

Conversion details:

Split-kwarg forwarding. The split int attr (1 = UP_DOWN/axis0, 2 = LEFT_RIGHT/axis1, the tpush/tpop encoding) is passed through verbatim to the tile op, which halves (shard) or doubles (gather) the split-axis extent.
Boundary memory. The tile-level split deducer intentionally leaves the boundary memory space null (the deduction fixpoint must not inherit an input-side layout); OpRegistry::Create then fills it from the tile op's set_output_memory declaration, so this converter needs no re-attachment of its own. LowerAutoVectorSplit builds its aiv_shard / aic_gather through the same Create, which is what keeps the two paths byte-identical — one declaration, read once. That space is the consuming lane's: tile.aiv_shard → Vec (AIV pops the half into UB), tile.aic_gather → Mat (AIC pops the full tile into L1, the space ExpandMixedKernel builds its V→C tpop in). The operand side is the mirror — Acc for the shard, Vec for the gather — and is enforced by the AivSplitValid verifier rather than declared as an input constraint, because a violated input constraint would make InferTileMemorySpace insert a move to the required space instead of reporting the authoring error.
No synthesized load. The realistic (region-only) operand is already an on-chip tile by the time the converter runs (its producer — a cube matmul for aiv_shard, a Vec vector op for aic_gather — lowered earlier in this same pass), so no tile.load is injected; aiv_shard / aic_gather are the cross-core transfer.

Recognized before this pass. Because the tensor.* form survives from OutlineIncoreScopes until this pass runs, earlier stages already treat it as the AIV-split boundary: ClassifyCallAffinity rolls both tensor.* and tile.* shard/gather up as MIXED (so cube/vector outlining splits correctly), and SplitAivStructuralVerifier requires both forms to be region-scoped.

Example¶

Before:

@pl.program
class Before:
    @pl.function(type=pl.FunctionType.InCore)
    def main_incore_0(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
        y: pl.Tensor[[64], pl.FP32] = pl.add(x, x)
        return y

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
        y: pl.Tensor[[64], pl.FP32] = self.main_incore_0(x)
        return y

After:

@pl.program
class After:
    @pl.function(type=pl.FunctionType.InCore)
    def main_incore_0(
        self, x: pl.Tensor[[64], pl.FP32],
        ret0_out: pl.Out[pl.Tensor[[64], pl.FP32]]
    ) -> pl.Tensor[[64], pl.FP32]:
        x_tile: pl.Tile[[64], pl.FP32] = pl.load(x, (0,), (64,))
        y_tile: pl.Tile[[64], pl.FP32] = pl.tile.add(x_tile, x_tile)
        ret0_store: pl.Tensor[[64], pl.FP32] = pl.store(y_tile, (0,), ret0_out)
        return ret0_store

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x: pl.Tensor[[64], pl.FP32]) -> pl.Tensor[[64], pl.FP32]:
        ret0_out: pl.Tensor[[64], pl.FP32] = pl.tensor.create((64,), dtype=pl.FP32)
        y: pl.Tensor[[64], pl.FP32] = self.main_incore_0(x, ret0_out)
        return y

Key changes:

pl.add(x, x) → pl.tile.add(x_tile, x_tile) (op conversion)
tile.load inserted at entry, tile.store at exit
Out parameter ret0_out added to InCore function
tensor.create inserted at orchestration call site

Loop-Carry Valid-Shape Repair¶

tensor.matmul drops its operands' valid_shape, so an accumulator only becomes narrower than the seed it is carried from once this pass produces a tile.matmul over a row-narrowed left operand:

acc = pl.create_tensor([M, N], dtype=pl.INT32)          # full box
for k0 in pl.pipeline(0, K, K_TILE, stage=2):
    xk = pl.slice(x, [M, K_TILE], [m0, k0], valid_shape=[v, K_TILE])   # runtime v
    acc = pl.matmul_acc(acc, xk, wk, b_trans=True)      # narrowed, compact result

The carry is typed from its init value alone — ConvertToSSA mints the IterArg from the seed, this pass re-mints it from the converted seed, and both force the loop's return_var back to that type — so the narrowing dies at the loop boundary. mad writes L0C at an N-fractal stride of ceil(v/16)*16 while a reader that believes the full box height walks it at the physical row pitch, corrupting every N-fractal above the first (issue #2470).

Before returning, this pass therefore calls narrow_loop_carry::NarrowAccCarries on each function: an Acc carry seeded by tile.create is re-declared at the extent its yields prove — tile.create(compact=True) plus tile.set_validshape — and the body's def-use closure is re-typed through the operators' own deducers. Repairing it in the pass that creates it keeps the pipeline verifiable; leaving it would publish a carry the TypeCheck diagnostic and the AccCompactValid property verifier reject. FlattenTileNdTo2D calls the same helper, for an ND seed whose narrowing only appears when tile.batch_matmul is unrolled into 2D matmuls.

A carry is left exactly as it is when the two readings of its buffer cannot disagree — a single-fractal-block [16, N] accumulator packs to its physical rows whatever its valid rows — or when the narrowed extent is only computed inside the loop body, where the re-declared seed could not name it.

Implementation¶

Header: include/pypto/ir/transforms/passes.h

Implementation: src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp

Python binding: python/bindings/modules/passes.cpp

Tests: tests/ut/ir/transforms/test_convert_tensor_to_tile_ops.py, tests/ut/ir/transforms/test_narrow_loop_carry_valid_shape.py (the carry repair)

Pass Properties¶

Property	Value
Required	SSAForm, SplitIncoreOrch, NormalizedStmtStructure
Produced	SSAForm, IncoreTileOps, NormalizedStmtStructure, AivSplitValid
Invalidated	AivSplitValid

AivSplitValid is both invalidated and re-produced, which forces a second verification of the split regions here. OutlineIncoreScopes establishes the property while the AIV-split boundary is still tensor.aiv_shard / tensor.aic_gather; a TensorType carries no memory space, so the verifier's boundary memory-contract check is necessarily skipped there. This pass rewrites those ops to their tile form and attaches the declared boundary memory, which is exactly what that check inspects.

Key Components¶

Component	Role
`TensorArgsInConvertedOpsCollector`	IRVisitor — identifies tensor params needing entry loads
`ConsumerSpaceCollector`	IRVisitor — collects each variable's consumer-demanded memory space from `input_reqs`
`TypePropagatingMutator`	Base IRMutator — propagates type changes through control flow
`TensorToTileMutator`	IRMutator — converts tensor ops to tile ops via OpConversionRegistry
`ForwardedCallFinder`	IRVisitor — locates the wrapper's call into a transformed InCore (Phase 2a)
`WrapperForwardMutator`	IRMutator — appends new Out args to the wrapper's inner call (Phase 2a)
`CallSiteUpdateMutator`	IRMutator — inserts tensor.create at orchestration call sites (Phase 2b)
`IncoreTileOpsVerifier`	IRVisitor — verifies no TensorType ops remain in InCore functions

Scope¶

Function type	Action
InCore	Converted (tensor ops → tile ops); Phase 1 may append `Out` params
Spmd / Group (forwarding to a transformed InCore)	Signature mirrors the InCore's new `Out` params; inner call forwards them (Phase 2a)
Spmd / Group (no transformed-InCore forwarding)	Unchanged
Orchestration / Opaque	Call sites updated — `tensor.create` inserted for each new `Out` param (Phase 2b)

Space	`slayout`	fp32 box	fp16 box
`Mat`	`row_major`	16 x 8	16 x 16
`Right`	`col_major`	8 x 16	16 x 16