ExpandMixedKernel Pass

Expands mixed InCore functions into separate AIC (Cube) + AIV (Vector) kernels wrapped in a Group function. Non-mixed InCore functions get their FunctionType converted to AIC or AIV.

Overview

After OutlineIncoreScopes and ConvertTensorToTileOps, InCore functions may contain both Cube ops (tile.matmul, tile.gemv, etc.) and Vector ops (tile.add, tile.exp, etc.). Some ops like tile.load, tile.store, tile.move, and tile.reshape are classified as Cube or Vector based on the MemorySpace of their tile operands. Functions containing ops from both sides are mixed InCore functions. Hardware requires Cube and Vector operations to run on separate core types, so this pass splits them into:

• AIC function (FunctionType::AIC) — contains only Cube + shared ops
• AIV function (FunctionType::AIV) — contains only Vector + shared ops
• Group function (FunctionType::Group) — calls AIC then AIV, replaces the original

When an existing Group function already calls the InCore function (e.g. from OutlineClusterScopes), the pass rewrites that Group in-place to call AIC + AIV directly, avoiding a redundant Group wrapper. When a standalone Spmd wrapper calls the InCore function, the pass preserves the Spmd wrapper and creates a Group callee under it so launch semantics stay on FunctionType::Spmd.

For non-mixed InCore functions (pure Cube or pure Vector), the pass converts FunctionType::InCore to the corresponding type without splitting:

• Pure Cube → FunctionType::AIC
• Pure Vector or shared-only → FunctionType::AIV

After this pass, no FunctionType::InCore functions remain in the program.

Shared argument ABI for hand-written Groups

One runtime MixedKernels submission creates a co-scheduled AIC/AIV task whose active lanes share one CoreTaskArgs payload. Consequently, the two compiled kernel wrappers must unpack the same tensors-first argument layout. This is automatic for Groups created by the mixed-kernel splitter, because both member calls forward the complete Group signature.

A hand-written pl.cluster() may instead submit members with different source signatures, for example cube(q, k, v, sink, task) and vec(out, task). The outlined Group captures the union of those values. This pass normalizes both DSL members to that Group signature and rewrites both inner Call or Submit nodes to forward the complete list. Unused parameters remain valid wrapper arguments; each cloned body still references only the values from its original call.

Normalization also binds dynamic dimensions in each member's parameter types to the rewritten Group arguments, using the same type-binding rules as inline expansion. The bindings apply to local aliases, loop-carried values, expression types, and declared return types. For example, a member workspace Tensor[[8, N]] passed a Group workspace Tensor[[8, t + 32]] retains t + 32 throughout its cloned body; an unbound member-local N must not survive in a type. Replacements are applied once so caller-owned extent expressions keep their original meaning. The widened Call/Submit results and their receiving aliases use the normalized member return types, including Submit's trailing TaskId.

The pass first records every producer return position, its callee-scoped SSA expression, type, and optional output-parameter binding. GroupOutputBindings then follows the consumer's SSA aliases and tuple projections to a particular producer call and return position. Only after identifying that output does it resolve the actual storage passed at the call site and bind it to the shared Group parameter list. Producer-local SSA variables are never copied into a consumer body.

Computed Tensor results have already been materialized by ConvertTensorToTileOps as explicit Out parameters and stores. That pass exposes internal output storage through Group parameters and allocates it in the outer caller, even when the Group does not return it. Existing complete-parameter writebacks reuse their original storage. This also covers reordered Call/Submit tuple results, direct pair[i] arguments, and projections through local tuple aliases or nested tuples. pl.no_dep uses the same mapping. Submit's omitted-output gap and trailing communication contexts are respected. Callee summaries are shared for one normalization run and then discarded. Views and slices are not treated as complete-parameter aliases; a result with no caller-visible carrier is diagnosed instead of producing a free variable. Device-computed scalar returns require an explicit Tensor output under the existing kernel-return contract.

This binding does not insert a cross-core handshake. Hand-written Group members run concurrently; their existing synchronization protocol must order a producer store before a consumer load.

When both member kernels are exclusive to the Group, their names are preserved. If either kernel has another call site, the pass creates a Group-private adapter pair and rewrites import_peer_buffer(peer_func=...) to the adapter peer names, leaving the original ABI available to the other callers. External-source and runtime-bound members cannot be body-cloned; when their layouts differ, the pass rejects the Group and requires both declarations to use the same signature.

This normalization happens before InjectGMPipeBuffer, so a backend-injected __gm_pipe_buffer parameter is appended to the already-identical member signatures. It does not require sync_start: AIC and AIV are subslots of one mixed task, while sync_start controls multi-block SPMD launch admission. When the Group is called through an enclosing Spmd wrapper, orchestration codegen traces complete-buffer alias chains back to the wrapper parameters before ordering the shared argument payload.

Unsplittable transpose: checked earlier, not here

A tile.transpose that swaps the split axis cannot be split correctly: the per-lane data migrates to the other dimension while the split still halves the original axis, so the result would be mis-shaped and miscompute. This pass no longer reports it. The check moved to where the context it needs still exists:

Shape	Reported by
A pl.split_aiv region	AivSplitValid check (l), with LowerAutoVectorSplit as the backstop
A function-level pl.split (no region)	LowerAutoVectorSplit

Reporting it here could only ever produce a worse diagnostic: the walk that would raise the region case is AppendConsumed, which erases the wrapper it would have to name, so it could state neither the region's mode nor the offending result nor a fix. And the function-level case sat here only because this pass sees every function, including the pure-vector pl.split kernels LowerAutoVectorSplit declines to lower; that pass now validates its pass-through branch, which closes the gap without the late diagnostic.

A whole-function backstop remains here, as an INTERNAL_CHECK_SPAN. This pass is documented below as invocable on its own after InferTileMemorySpace, and a bare pass call does not enforce required properties — only PassPipeline does. So a caller who skips LowerAutoVectorSplit arrives with AivSplitLoweredValid unmet, and without the guard the kernel expands silently and SplitVectorKernel mis-shapes it. It is internal rather than user-facing because the failure is an unmet pass prerequisite, not a fact the author can read off their source, so the message names the skipped step instead of giving authoring advice.

split_axis::FindTransposeSplitHazard is still the shared detector: it flags the first tile.transpose whose source is non-singleton on the split axis (a singleton source carries no split data — the no-op broadcast case — and is left split; a dynamic, non-ConstInt extent is treated as non-singleton and flagged conservatively).

SplitRegionConsumer validates each retained or synthesized region using its own split axis, erases the wrapper, and records placement on the same consumed body used by both mixed-function classification and expansion. Pure-AIV functions also consume their regions here. Functions with no regions skip region consumption without rebuilding their bodies. Erased-region comments are prepended to the first emitted statement without mutating input metadata, including when a boundary expands into transport operations.

The pass requires AivSplitLoweredValid. Source boundaries use strict operand memory rules; lowered verification retains those checks for locally defined operands. Expansion checks operand availability, using lexical placement for inline calls as well as bound values, on the producing lane, so parameters shared by both lanes remain valid regardless of their memory annotation. Boundary result memory still describes the consuming lane.

Cross-core data transfer at CV boundaries is handled by splitting explicit tile.move ops into tpush/tpop pairs:

Direction	AIC side	AIV side
Cube→Vector (e.g. Acc→Vec)	tpush_to_aiv(source_tile)	dest_var = tpop_from_aic()
Vector→Cube (e.g. Vec→Mat/Left/Right)	dest_var = tpop_from_aiv()	tile.move to adapt fractal layout, then tpush_to_aic(adapted_tile)

Fractal TileView layout: Cross-core transfer tile views are computed by BuildCrossCoreTransferView based on the destination memory space. The mapping differs between backends:

Ascend950 (a5) — hardware cross-core pipe carries data in fractal layout:

Direction	Push/Pop TileView (blayout, slayout)	Name
Vec→Left	col_major, row_major	NZ
Vec→Right	col_major, row_major	NZ
Vec→Mat	must be explicitly set in move	—
Mat/Acc→Vec	must be explicitly set in move	—

Vec→Right crosses as NZ, not as the ZN a Right operand ends up in: a5 inserts the Vec tile into the Mat FIFO through TINSERT_IMPL<TInsertMode::NZ>, so the bridge tile has to stay NZ. The ZN view belongs to the Mat → Right tile.move the cube side makes after the pop, which is one step later than this table describes — the same split the a2a3 table below spells out.

Ascend910 (a2a3) — cross-core transfer goes through GM → Mat, and Mat only supports the NZ layout. Both Left and Right destinations use NZ at the transfer boundary; the final Left/Right layout is resolved by the subsequent Mat → Left/Right tile.move (MTE1):

Direction	Push/Pop TileView (blayout, slayout)	Name
Vec→Left	col_major, row_major	NZ
Vec→Right	col_major, row_major	NZ
Vec→Mat	preserve original	—
Mat/Acc→Vec	preserve original	—

On both backends, the ordinary-data AIV push side (V→C) inserts a tile.move before tpush_to_aic to convert the source tile into the required fractal layout. The tile.move helper (CreateMove) propagates blayout/slayout kwargs when the result type carries a TileView. Ascend950 MX scales use the dedicated row/row carrier described below.

Hand-written pipes get the same adapter

The rule above describes the boundary-move path, which only sees the pipes this pass builds while expanding an InCore function. A pipe authored directly (pl.reserve_buffer, pl.{aic,aiv}_initialize_pipe and pl.tpush_to_aic, written out by hand) never reaches it. On a backend where RequiresVtoCFractalAdapt() holds, such a push would ship a bare ND tile into a FIFO the cube reads as fractal, scattering every element of the popped tile.

AdaptManualVtoCPush closes that gap. It runs as the pass's final phase over every AIV function the pass emits, including pure-vector InCore bodies converted to AIV and the AIV half of a mixed body. Ordinary data keeps the fixed cube-side Mat transfer view. For MX, the pass resolves the AIV import to its AIC reserve and pairs V2C operations by pipe id and occurrence order, so each push is planned from its actual consumer tpop contract. Missing peers, ambiguous initializers, mismatched push/pop counts or splits, and undersized slots are rejected instead of guessing a layout.

The rewrite applies to FP8E8M0 MX scales as well as ordinary data tiles, preserves the original push kwargs and attrs (including id), and is idempotent: a push already staged into the boundary view is left unchanged. The backend capability is consulted only after a manual V→C push is encountered.

GM-mediated cross-lane dependencies

A tile.move is not the only way data crosses the CV boundary. When one lane writes a GM tensor with tile.store and the other lane reads the same tensor with tile.load, the data crosses through GM. Because neither op is a tile.move, CV-boundary detection misses the dependency, and without a fence the two split kernels race on the shared GM region (issue #1433).

CollectGmCrossLaneSyncs detects these store/load pairs and emits a pure-synchronisation handshake: the data still flows through GM unchanged, while a tpush (producer, right after the store) and a matching tpop (consumer, right before the load) establish the missing happens-before edge. The popped tile is a fence token, freed immediately by FinalizeTpopTfrees; BuildAutomaticPipeSetup then injects the same pipe setup as for tile.move boundaries.

Producer (writes GM)	AIC side	AIV side
Cube store → Vector load	tile.store ...; tpush_to_aiv(stored_tile)	tok = tpop_from_aic(); tfree_to_aic(tok); tile.load ...

To stay deadlock-free, a handshake is emitted only when (1) the GM origin tensor has exactly one producer-lane store, (2) the opposite-lane load lives in the same structural body (same loop/branch, so the tpush/tpop execute the same number of times), and (3) the store precedes the load. Pairs split across different loops or branches are left untouched.

Only the Cube→Vector direction (cube tile.store → vector tile.load) is fenced. The AIC tpush sends the stored tile raw, exactly as the normal boundary C2V push does on both backends, and the AIV tpop lands in Vec. The reverse Vector→Cube direction would require the V→C fractal-layout adaptation that the tile.move boundary path applies before tpush_to_aic; emitting a raw-tile sync there would violate the cross-core transport contract, so V2C GM exchanges are left unfenced.

When split kernels contain cross-core tpush/tpop, the pass also prepends the required frontend pipe setup automatically:

• system.reserve_buffer(...) on the consumer side
• system.import_peer_buffer(...) on the producer side
• system.aic_initialize_pipe(...) / system.aiv_initialize_pipe(...) on both sides

In addition, the pass inserts system.tfree_to_aic(...) / system.tfree_to_aiv(...) after every consumer-side tpop chain.

Setup is derived from the split bodies:

• dir_mask: C2V=1, V2C=2, bidirectional=3
• id: omitted for automatic setup, so PTOAS uses the default frontend pipe id 0
• slot_size: max tile byte size across all directions (shape * dtype bits / 8)
• slot_num: ring depth — 2 by default, in every direction; always emitted explicitly; overridable per scope (see below)
• local_slot_num: optional A2/A3 local consumer depth; omission uses slot_num
• buffer_size: slot_size * local_slot_num when supplied, otherwise slot_size * slot_num
• buffer names: <func>_c2v_slot_buffer / <func>_v2c_slot_buffer
• reserve-buffer base: AUTO on insertion, then resolved to an explicit address by AllocateMemoryAddr

When cross-core directions use different tile sizes, the pass picks max(all observed tile byte sizes) as the common slot_size for initialize_pipe. Smaller tiles leave unused bytes in each slot but hardware correctness is preserved. Explicit user-authored programs can still create multiple independent pipes by supplying different id values to initialize_pipe and matching tpush / tpop / tfree ops.

MX scale V2C transport

On Ascend950, a mixed quant_mx → matmul_mx path transports both results over V2C. FP8E8M0 fractal-32 scales never use the NZ TINSERT path. A full-valid logical row/row scale is pushed directly. A logical col/col scale takes the official, zero-copy tile.transpose_view after any carrier preparation, producing a row/row carrier with the physically transposed shape. If producer and consumer logical layouts differ, a real Vec→Vec tile.move converts to the consumer layout before that carrier step.

The AIC tpop keeps the consumer's public logical shape, layout, and valid_shape. For a partial logical scale, the producer transports the full physical box so the ND insert is contiguous and 32-byte aligned. The pass first materializes a private Vec carrier, then widens only that carrier with tile.set_validshape; this keeps slice/reshape views legal and leaves the original logical scale metadata unchanged. The consumer still ignores padding outside its original logical valid region. The carrier's full byte size must fit the paired pipe's slot_size.

Explicit pl.aic_gather uses the same MX carrier path. Its type deduction keeps FP8E8M0 row/row/32 and col/col/32 scale layouts (and tensor-level MX_A_ZZ / MX_B_NN), so the following Mat→LeftScale/RightScale move sees the scale format. Ordinary shard/gather tiles still use the destination's normal transfer layout.

SplitMode.NONE preserves the logical shape. For fully valid scale boxes, UP_DOWN gathers MX_A_ZZ row bands and LEFT_RIGHT gathers MX_B_NN column bands. The latter's transpose view changes both endpoints' physical split to UP_DOWN; the Mat pop retains the full logical col/col shape. Carrier validation compares each producer with its half of that logical shape and sizes the slot for both lanes. Scale-group-axis splits and partially valid split scale boxes are rejected because they need additional boxed-layout transport handling.

Gathered Mat tiles retain FIFO slots until their consumers release them. Gathering all four MX matmul operands before consuming any requires pl.cross_core_slot(slot_num=4); two slots are insufficient even with correct scale layouts.

Overriding the slot count (slot_num)

The default slot count (cross_core_pipe::kDefaultAutoPipeSlotNum = 2) can be tuned per scope with the pl.cross_core_slot(slot_num=N) optimization entry:

# Deepen the auto-inserted ring so the cube can run further ahead of the vector
# core, at the cost of consumer-side buffer space.
with pl.spmd(4, optimizations=[pl.cross_core_slot(slot_num=4)]):
    ...

# Or on the pl.at form, alongside an independent split mode:
with pl.at(level=pl.Level.CORE_GROUP,
           optimizations=[pl.split(pl.SplitMode.UP_DOWN),
                          pl.cross_core_slot(slot_num=16)]):
    ...

The value is carried as the slot_num scope attr, propagated by OutlineIncoreScopes to the outlined function's slot_num attr, and read here. It drives the transport ring and the emitted initialize_pipe slot_num attribute. Without an explicit local depth, it also sizes the reserved buffer (slot_size * slot_num). Automatic setup always emits the attribute, whether or not the scope overrides it: PTOAS derives its own depth from dir_mask (8 unidirectional, 4 bidirectional) when the clause is absent, which would index past a buffer reserved for a shallower ring. Only hand-written pl.system.{aic,aiv}_initialize_pipe still reaches that fallback. slot_num must be positive and is orthogonal to splitting — it sizes a data channel, it does not partition work. It applies in whichever directions the scope actually uses (cube->vector, vector->cube, or both), and to any scope that drives a pipe: a kernel with no split at all still drives one (on a2a3 via dual-AIV dispatch — see below). It is ignored only when the outlined scope ends up with no cross-core ops.

Deprecated: pl.split(MODE, slot_num=N) still accepts the slot count and emits a DeprecationWarning. That spelling forced authors to name a split mode they did not want — the pl.split(pl.SplitMode.NONE, slot_num=N) idiom — which OutlineIncoreScopes now rejects on any scope carrying pl.split_aiv regions. Move the value to pl.cross_core_slot(slot_num=N); the printer already normalises to that form.

On A2/A3, set local_slot_num to keep a smaller on-chip consumer FIFO while retaining the GM transport depth:

with pl.spmd(4, optimizations=[pl.cross_core_slot(slot_num=4, local_slot_num=1)]):
    ...

Both counts must be positive integers and local_slot_num <= slot_num. Omitting local_slot_num (or passing None) preserves the existing footprint. The outliner propagates both attributes; both peer initializers emit the same explicit local depth, and each on-chip reserve uses slot_size * local_slot_num. GM allocation still uses slot_num. The common maximum slot_size is unchanged. The pass checks the finalized FIFO schedule and rejects a local depth smaller than the number of simultaneously retained consumer pops. It does not reschedule consumers to make a smaller depth fit.

Explicit local depth is rejected on A5, even when equal to slot_num, because PTOAS does not support the attribute for that frontend pipe lowering. Omit it for portable source. Like slot_num, it has no effect on scopes without an automatic cross-core pipe.

For consumer-side cross-core tiles, the pass also ensures each tile.tpop_* has a matching system.tfree_*. When an existing free is obviously too early, the pass delays it to a later statement in the same block without reordering independent tpop chains. When AIC must post-process a tile.tpop_from_aiv result with a same-side tile.move (for example Mat -> Left/Right/Bias), the generated system.tfree_to_aiv(...) is rewritten to free the canonical popped tile value and can be delayed past that carrier chain.

For Ascend910 (a2a3), mixed kernels with no function split mode (split unset or SplitMode.None) are also supported. In that case the pass keeps a single AIV kernel body, marks it for dual AIV dispatch, and later lowering emits a runtime subblock_idx branch: AIV lane 0 executes the original body, while AIV lane 1 replays the cross-core handshakes with tile-producing replay work forced to valid_shape=[0, 0], and visible tile.store writes suppressed. This keeps the AIC/AIV handshakes balanced for pl.at(level=CORE_GROUP) no-split mixed kernels while the secondary sync lane avoids real DMA/compute work.

This replay path applies only to non-split_aiv mixed kernels. A function carrying a pl.split_aiv region — any mode, including the task-parallel pl.SplitMode.NONE — is stamped split_aiv by LowerAutoVectorSplit (pass 23), so SplitVectorKernel routes it through the split path (both AIV lanes dispatch via dual_aiv_dispatch) and never the lane-0-only replay above. For a NONE region this is exactly right: both lanes run the full body for disjoint, aiv_id-dispatched work — dropping lane-1 stores would be a miscompile.

Requirements:

• Input IR must have tile ops (run ConvertTensorToTileOps first)
• Input IR must have InCore scopes outlined (run OutlineIncoreScopes first)
• Tile ops must be flattened to 2D (run FlattenTileNdTo2D first)
• Tile memory space must be inferred (run InferTileMemorySpace first)
• Cross-core fractal TileView assignment is supported on Ascend950 and Ascend910 backends

When to use: Run after InferTileMemorySpace when InCore functions may contain both Cube and Vector tile operations.

Note: This pass is part of the default tile optimization pipeline. You can also invoke it explicitly via passes.expand_mixed_kernel()(program) when debugging or building a custom pass pipeline.

API

C++	Python	Level
pass::ExpandMixedKernel()	passes.expand_mixed_kernel()	Program-level

Python usage:

from pypto.pypto_core import passes

expand_pass = passes.expand_mixed_kernel()
program_expanded = expand_pass(program)

Algorithm

Phase 1 — Pre-scan:
  Identify InCore functions that have existing Group callers and
  callers that still need the original function name to remain callable.

Phase 2 — Expand each InCore function F:
  1. Recursively classify affinity of all statements (including inside loops/conditionals)
  2. Detect CV boundary moves: tile.move ops crossing cube↔vector memory spaces
     (recorded in a separate boundary_moves map, not as a distinct affinity)
  2a. Detect GM-mediated cross-lane store/load pairs (CollectGmCrossLaneSyncs):
      a tile.store on one lane + a tile.load on the other from the same GM
      tensor origin, scheduled as a tpush/tpop sync fence (see "GM-mediated
      cross-lane dependencies" above)
  3. If not mixed (no CUBE ops or no VECTOR ops, and no boundary moves):
     convert FunctionType to AIC (pure Cube) or AIV (pure Vector / shared-only)
  4. Build AIC body: keep CUBE + SHARED stmts, prune VECTOR, recurse into MIXED loops
     - For boundary move (Cube→Vector): emit tpush_to_aiv(source_tile)
     - For boundary move (Vector→Cube): emit dest_var = tpop_from_aiv() with fractal TileView
  4a. For an op-driven boundary (tile.aiv_shard / tile.aic_gather), derive the
      pto-isa split CODE for the tpush/tpop pair from the op's authored MODE plus
      the full-width tile's extents (BoundaryTransportSplitCode ->
      split_axis::ShardSplitCode / GatherSplitCode) — see "The split code" below
  5. Build AIV body: symmetric (keep VECTOR + SHARED, prune CUBE)
     - For boundary move (Cube→Vector): emit dest_var = tpop_from_aic() with fractal TileView
     - For boundary move (Vector→Cube): emit tile.move to adapt fractal layout, then tpush_to_aic(adapted_tile)
  5a. Assign fractal TileView to all boundary tpop result types and pre-tpush tile.move ops
      via BuildCrossCoreTransferView:
      - Ascend950: Left→NZ, Right→ZN, Mat/Vec→preserve
      - Ascend910: Left→NZ, Right→NZ (Mat only supports NZ), Mat/Vec→preserve
  6. Repair loop-carried state on both bodies
     - Strip dead iter_args whose carried values are unused on this side —
       uses include the while condition, and close over the loop's own yield,
       so a carry read only to feed another live slot (a multi-entry FIFO
       rotation) stays alive
     - Pull back missing init-value definitions for surviving iter_args
     - Rewrite dangling yields to identity yields when a branch-local value was stripped
     - Seed branch-result cleanup with function parameters and pending remap keys,
       preserving live if/yield results that directly reference a parameter or
       an expression such as `n - 1` on either lane
     - Remap dangling tile.store result vars (SSA versions stripped by AIC-side splitting) to the corresponding output parameter
  7. Run dead code elimination on both bodies (recursive into loops)
  8. Normalize loop-carried state again, because DCE may remove a SHARED-only
     post-loop use that temporarily kept an iter_arg alive
  9. Run dead code elimination again to clean up init-value chains exposed
     by the second strip
 10. Ensure each consumer-side `tpop` chain has a matching
     `system.tfree_to_aic` / `system.tfree_to_aiv`, delaying obviously early frees within the same block when needed
 11. If the split bodies use cross-core tile ops and do not already contain setup,
      derive reserve/import/initialize_pipe prologues and prepend them
 12. Create AIC function (no return) and AIV function (original return)
     - On Ascend910 no-split mixed kernels, tag the generated AIV with dual-dispatch metadata
       so later lowering launches the same AIV kernel on both vector lanes and rewrites the secondary lane's tile
       replay path to `valid_shape=[0, 0]`
 13. If a non-Group caller still needs the original function name
     (for example a standalone Spmd wrapper): also create a Group
     function (calls AIC then AIV)
 14. If there are only existing Group callers: skip the extra Group wrapper

Phase 3 — Rewrite Group callers:
  For each Group function that calls a split InCore, replace the InCore call
  with an AIC call + AIV call sequence (EvalStmt for AIC, AssignStmt for AIV).

Phase 4 — Normalize hand-written mixed Group ABIs:
  For each Group whose AIC/AIV calls do not both forward the complete Group
  signature, rebuild the DSL member bodies against one canonical parameter list
  and rewrite both Call/Submit nodes to pass it. Preserve exclusive member names;
  otherwise create Group-private adapters and retarget peer_func references.

Affinity classification:

Affinity	Ops	Classification Rule
CUBE	tile.matmul, tile.matmul_acc, tile.matmul_bias, tile.gemv, tile.gemv_acc, tile.gemv_bias, tile.batch_matmul, tile.batch_matmul_acc	Always CUBE (op name)
CUBE or VECTOR	tile.load	By target_memory kwarg: cube memory (Mat, Left, Right, Acc, Bias) → CUBE; Vec → VECTOR
CUBE or VECTOR	tile.store, tile.reshape	By source tile's memory_space: cube memory → CUBE; Vec → VECTOR
MIXED	tile.move crossing cube↔vector memory	Leaf cross-side move — also recorded in the boundary_moves map (see below)
CUBE or VECTOR	tile.move (same-side)	By source tile's memory_space
VECTOR	All other tile.* ops (tile.add, tile.exp, tile.sub, etc.)	Always VECTOR (op name)
VECTOR	pld.tile.put, pld.tile.get, pld.tensor.put, pld.tensor.get	Declared set_core_affinity(VECTOR) — TPUT/TGET stream through a VEC staging tile, which ptoas enforces
SHARED	Non-tile ops, function calls, control flow, scalar ops	—
SHARED	pld.system.notify, pld.system.wait	Core-agnostic by ISA (pure scalar/GM), so no affinity is declared. notify also declares set_no_duplicate() (both NotifyOp forms) — a cube-lane copy can release the peer before the vector lane's TPUT lands the data; wait does not, since it blocks and its cube-lane copy is load-bearing
MIXED	Compound statements containing both CUBE and VECTOR children	—
VECTOR	Region-local, intrinsically SHARED no-duplicate call without a stated lane	Pass-local lexical placement

Region placement keeps pld.system.notify off AIC. It preserves duplicate-safe wait, stated lane affinity, and both boundary endpoints. The map is internal to this pass; no call attributes carry placement. See LowerAutoVectorSplit for the shared contract.

CV boundary detection: A tile.move is a CV boundary when its source tile memory and target memory are on different core sides. Cube-side memory: Mat, Left, Right, Acc, Bias. Vector-side memory: Vec. Same-side moves (e.g. Mat→Left) are classified by their source memory as usual. Boundary leaf moves are tagged MIXED for affinity purposes and are also recorded in a separate boundary_moves map; the cross-core direction (Cube→Vector vs Vector→Cube) is recovered via ClassifyMoveDirection at the call sites that need it (CollectCVBoundaryMoves, BuildCoreBody).

Nested structure handling: ForStmt, IfStmt, and WhileStmt containing mixed ops are duplicated into both AIC and AIV bodies with recursively pruned contents.

Loop-state repair after splitting: mixed-loop control flow is intentionally preserved during body construction, which can leave one side with extra iter_args, missing init-value definitions, or yields that reference stripped branch-local values. The pass repairs those cases before DCE, then normalizes loop-carried state once more after DCE because dead shared aliases can disappear and make an iter_arg removable only at that stage. A final DCE pass cleans up any init-value chains that become dead after the second strip.

Which carries count as live: uses in the loop body, the while condition, or post-loop uses of the corresponding result seed liveness. A carry used only by the while condition stays live even when its final result is unused. An iter_arg read only by the loop's own trailing yield is still live when the slot it feeds is live — that is exactly the shape of a multi-entry FIFO, whose entry N is read nowhere but the yield that rotates it into entry N-1. Liveness therefore closes over the yield instead of being read off the body alone, and the strip asserts that neither the while condition nor a surviving yield value names a carry it just dropped. Dropping such a carry would leave a free Var that no later pass rejects and PTO codegen reports as "no MLIR mapping for MemRef base ...".

Group wrapper param-returns: a newly created Group wrapper returns its own parameters when every return position traces to a param writeback (via return_lineage::ReturnedParamIndices) — the AIV call is emitted as a bare EvalStmt and the wrapper returns the matching Group params directly (e.g. return out_0), keeping the ReturnParamsExplicit invariant. Only when some position is not a param writeback does it fall back to the legacy result = aiv_call(); return result form.

The split code

tile.aiv_shard / tile.aic_gather carry the authored mode (0 / 1 / 2); the transport pair this pass mints carries pto-isa's TileSplitAxis code (0..4), which additionally says how the two AIV lanes' runtime extents relate — that is how the consumer finds lane 1's band inside the FIFO slot:

Code	pto-isa	Lane 1's band starts at	Requires
1 / 2	TILE_UP_DOWN / TILE_LEFT_RIGHT	e1 * pitch	e0 == e1
3 / 4	TILE_UP_DOWN_ODD / TILE_LEFT_RIGHT_ODD	(e1 + 1) * pitch	e0 == e1 + 1

BoundaryTransportSplitCode reads the lane extents eL = clamp(V - L*S, 0, S) off the FULL-width tile — the shard's operand (Cube→Vector) or the gather's result (Vector→Cube). S is the partition stride: the tile's physical half, or the balanced lane_stride=S attr when LowerAutoVectorSplit rebalanced a ragged boundary. So an odd split axis (an odd physical box, an odd valid extent inside an even one, or a rebalanced odd valid region) takes the odd code, an empty lane 1 keeps the even one, and a ragged extent pto-isa cannot place — only reachable on the box partition — is rejected with the extents that would work. Both lane bodies run the same derivation on the same inputs, so the AIC and AIV sides always agree. The Vector→Cube gather has no odd form (pto-isa's vector-side producer offsets lane 1 by lane 1's own extent), so it is even-only. Full derivation: LowerAutoVectorSplit.

Example 1: InCore without existing Group caller

When Orchestration calls InCore directly, a new Group wrapper is created.

Before (after InferTileMemorySpace):

@pl.program
class Before:
    @pl.function(type=pl.FunctionType.InCore)
    def compute_incore_0(self, x: pl.Tensor[[16, 128], pl.BF16],
                         y: pl.Tensor[[128, 128], pl.BF16],
                         out_0: pl.Out[pl.Tensor[[16, 128], pl.FP32]]
                         ) -> pl.Tensor[[16, 128], pl.FP32]:
        x_mat: pl.Tile[[16, 128], pl.BF16] = pl.load(x, [0, 0], [16, 128], target_memory=pl.Mem.Mat)
        y_mat: pl.Tile[[128, 128], pl.BF16] = pl.load(y, [0, 0], [128, 128], target_memory=pl.Mem.Mat)
        x_left: pl.Tile[[16, 128], pl.BF16] = pl.move(x_mat, target_memory=pl.Mem.Left)
        y_right: pl.Tile[[128, 128], pl.BF16] = pl.move(y_mat, target_memory=pl.Mem.Right)
        z_tile: pl.Tile[[16, 128], pl.FP32] = pl.matmul(x_left, y_right)
        z_vec: pl.Tile[[16, 128], pl.FP32] = pl.move(z_tile, target_memory=pl.Mem.Vec)
        out_0 = pl.store(z_vec, [0, 0], out_0)
        return out_0

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x, y) -> pl.Tensor[[16, 128], pl.FP32]:
        out_0 = pl.create_tensor([16, 128], dtype=pl.FP32)
        return self.compute_incore_0(x, y, out_0)

After (conceptual — actual IR includes type annotations on all variables; auto-generated pipe setup is omitted here for brevity):

@pl.program
class After:
    @pl.function(type=pl.FunctionType.AIC)
    def compute_incore_0_aic(self, x, y, out_0):
        x_mat = pl.load(x, [0, 0], [16, 128], target_memory=pl.Mem.Mat)  # CUBE: load to Mat
        y_mat = pl.load(y, [0, 0], [128, 128], target_memory=pl.Mem.Mat) # CUBE: load to Mat
        x_left = pl.move(x_mat, target_memory=pl.Mem.Left)   # CUBE: Mat→Left (same-side)
        y_right = pl.move(y_mat, target_memory=pl.Mem.Right)  # CUBE: Mat→Right (same-side)
        z_tile = pl.matmul(x_left, y_right)              # CUBE op
        pl.tile.tpush_to_aiv(z_tile, split=0)        # boundary move: push Acc tile to AIV

    @pl.function(type=pl.FunctionType.AIV)
    def compute_incore_0_aiv(self, x, y, out_0) -> pl.Tensor[[16, 128], pl.FP32]:
        # tpop result carries fractal TileView (no layout for Vec destination → preserved)
        z_vec: pl.Tile[[16, 128], pl.FP32, pl.Mem.Vec, pl.TileView()] = pl.tile.tpop_from_aic(split=0)
        out_0 = pl.store(z_vec, [0, 0], out_0)           # VECTOR op
        pl.system.tfree_to_aic(z_vec)                    # release the consumed cross-core slot
        return out_0

    @pl.function(type=pl.FunctionType.Group)
    def compute_incore_0(self, x, y, out_0) -> pl.Tensor[[16, 128], pl.FP32]:
        self.compute_incore_0_aic(x, y, out_0)
        self.compute_incore_0_aiv(x, y, out_0)
        return out_0  # param-return: the AIV result writes through out_0

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x, y) -> pl.Tensor[[16, 128], pl.FP32]:
        out_0 = pl.create_tensor([16, 128], dtype=pl.FP32)
        return self.compute_incore_0(x, y, out_0)  # calls Group (same name)

Example 2: InCore with existing Group caller

When OutlineClusterScopes has already created a Group function calling the InCore, the pass rewrites the existing Group instead of creating a new wrapper. When OutlineClusterScopes has created a standalone Spmd wrapper, the pass keeps that Spmd function and retargets it to a new Group callee with the original InCore name.

Before:

@pl.program
class Before:
    @pl.function(type=pl.FunctionType.InCore)
    def compute_incore_0(self, x, y, out_0) -> pl.Tensor[[16, 128], pl.FP32]:
        # ... mixed Cube + Vector ops ...

    @pl.function(type=pl.FunctionType.Group)
    def compute_group(self, x, y, out_0) -> pl.Tensor[[16, 128], pl.FP32]:
        result = self.compute_incore_0(x, y, out_0)
        return result

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x, y) -> pl.Tensor[[16, 128], pl.FP32]:
        out_0 = pl.create_tensor([16, 128], dtype=pl.FP32)
        return self.compute_group(x, y, out_0)

After — existing Group is rewritten, no compute_incore_0 Group wrapper:

@pl.program
class After:
    @pl.function(type=pl.FunctionType.AIC)
    def compute_incore_0_aic(self, x, y, out_0):
        # ... Cube ops + tpush ...

    @pl.function(type=pl.FunctionType.AIV)
    def compute_incore_0_aiv(self, x, y, out_0) -> pl.Tensor[[16, 128], pl.FP32]:
        # ... tpop + Vector ops ...

    @pl.function(type=pl.FunctionType.Group)
    def compute_group(self, x, y, out_0) -> pl.Tensor[[16, 128], pl.FP32]:
        self.compute_incore_0_aic(x, y, out_0)       # rewritten: AIC call
        result = self.compute_incore_0_aiv(x, y, out_0)  # rewritten: AIV call
        return result

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x, y) -> pl.Tensor[[16, 128], pl.FP32]:
        out_0 = pl.create_tensor([16, 128], dtype=pl.FP32)
        return self.compute_group(x, y, out_0)  # unchanged

Loop initializer repair tracks lexical visibility: enclosing definitions, parameters, loop variables, carries and pending cross-core remaps remain available in nested bodies. It restores only missing producers; compound statements export only their result variables. Branch-local bindings never become visible to sibling branches.

Implementation

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

Implementation:

• src/ir/transforms/expand_mixed_kernel_pass.cpp: mixed-kernel expansion, cross-core transport, and pass orchestration.
• src/ir/transforms/expand_mixed_kernel/group_abi.cpp: hand-written Group ABI normalization, output binding analysis, and member adapter construction. The private group_abi.h declares its result and entry point.

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

Tests: tests/ut/ir/transforms/test_expand_mixed_kernel_a2a3.py, test_expand_mixed_kernel_a5.py, test_expand_mixed_kernel_split_aiv.py, and test_group_output_bindings.py.

Pass Properties

Property	Value
Required	SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D, TileMemoryInferred, NormalizedStmtStructure, AivSplitLoweredValid
Produced	SSAForm, MixedKernelExpanded, NormalizedStmtStructure, HardSyncallOccupancyValid, AccCompactValid
Invalidated	AccCompactValid, AivSplitLoweredValid

HardSyncallOccupancyValid is produced here not by anything this pass rewrites, but because resolving each kernel's FunctionType to AIV/AIC/Group is the precondition the hard-syncall occupancy verifier depends on. That verifier fires once, right after this pass.

AccCompactValid is invalidated and re-produced because the Cube->Vector boundary tile.move is rebuilt here as a tpush/tpop pair with a freshly built consumer type, so the Acc compact contract has to be re-checked on that new IR rather than trusted from InferTileMemorySpace.

Property Verifier

MixedKernelExpandedPropertyVerifier checks:

• no remaining FunctionType::InCore function contains both Cube and Vector tile ops
• tile.tpop_from_aiv in AIC lands in MemorySpace::Mat
• tile.tpop_from_aic in AIV lands in MemorySpace::Vec
• AIC/AIV functions with cross-core tpush/tpop also contain the required pipe setup
• every AIC/AIV tile.tpop_* has a matching system.tfree_*
• cross-core tile ops have statically known tile shapes (required for auto pipe setup)

This moves common failures (missing initialize_pipe, missing reserve_buffer / import_peer_buffer, missing tfree, non-static tile sizes, mismatched cross-core tpop/tfree pairs) from PTO codegen / ptoas time to immediately after ExpandMixedKernel. The verifier checks pairing by tile value/op; it does not prove that a tfree is placed after the true last use of that tile.

Design Decisions

Decision	Rationale
Move-based CV boundary detection	Explicit tile.move ops mark boundaries — no fragile variable data-flow analysis needed
boundary_moves map (not a separate enum value)	Boundary status is derivable from ClassifyMoveDirection at the call sites that need it; storing it as a side map keeps CoreAffinity to the four execution-side categories (CUBE / VECTOR / SHARED / MIXED)
MemorySpace-based classification for data-movement ops	tile.load/tile.store/tile.move/tile.reshape serve Cube or Vector depending on which memory they touch; InferTileMemorySpace sets this before the pass runs
Fractal TileView on tpop results	tpop result types carry the fractal TileView (NZ/ZN) directly rather than stripping layout — downstream passes and codegen see the correct layout without extra inference
Pre-tpush tile.move on AIV side	V→C transfers require data in fractal layout; inserting an explicit tile.move before tpush_to_aic makes the layout conversion visible in IR
CreateMove propagates layout kwargs	When the result type has a TileView, blayout/slayout are forwarded as kwargs so the generated tile.move call is self-describing
Group keeps original function name	When no existing Group caller: Orchestration call sites work unchanged — no call-site rewriting needed
Rewrite existing Group callers	When a Group already calls the InCore (e.g. from OutlineClusterScopes): rewrite it in-place to call AIC + AIV, avoiding redundant Group→Group nesting
Preserve standalone Spmd wrappers	When a standalone Spmd calls the InCore: keep FunctionType::Spmd, create a Group callee underneath, and keep core_num / sync_start on the Spmd wrapper
Parameters copied to all three functions	Simplifies wiring; DCE removes unused params in downstream passes
Recursive compound-stmt handling	Correctly splits mixed ops inside ForStmt, IfStmt, WhileStmt
Two-stage post-split loop-state repair	First makes loop-carried state valid, then re-strips iter_args after DCE removes dead shared aliases, with a final DCE to clean up exposed init-value chains
Auto-generated pipe setup	Tensor-level mixed kernels do not need handwritten reserve_buffer / import_peer_buffer / initialize_pipe; the pass derives them from cross-core tile ops
Auto-generated tfree chains	Consumer-side split kernels insert missing tfree calls, rewrite them to free the canonical popped tile value, and delay obviously early frees within the same block without reordering independent tpop chains
Max-slot-size policy	Uses max(all tile byte sizes) as the single initialize_pipe.slot_size, matching the backend assumption of one automatic reserve/import buffer per direction while preserving legacy bidirectional dir_mask=3 behavior