PTO Codegen¶

The PTO Codegen (PTOCodegen) generates MLIR code in PTO-ISA dialect from PyPTO IR. It transforms high-level PyPTO programs into low-level PTO instructions suitable for accelerator execution.

Design Principle: Strict 1-to-1 Mapping¶

Codegen must be a strict 1-to-1 translation from IR to generated code. Each IR node maps directly to its corresponding output construct — no optimization, analysis, or indirection transformation should occur in the codegen layer.

Belongs in codegen	Belongs in an earlier pass
IR node → output code mapping	Data-flow analysis (e.g., tracing return values to parameters)
Type/format conversion (DataType → MLIR type)	IR restructuring or canonicalization
Name mangling and SSA bookkeeping	Optimization or simplification

Why: Codegen that embeds analysis becomes fragile — it duplicates logic that passes already handle, and it's harder to test in isolation. Keeping codegen a straightforward translation ensures it stays predictable and maintainable.

When analysis is found in codegen: File a tracking issue and refactor it into a dedicated pass when bandwidth allows. #814 was an example: return-to-parameter tracing in orchestration codegen has been refactored into the NormalizeReturnOrder pass.

Overview¶

Key Features¶

Automatic MLIR Generation: Converts PyPTO IR to PTO-ISA MLIR dialect
Structured Code Generation: Outputs constants, tensor views, allocations in order
Implicit Lowering: Automatically generates pto.partition_view from tile.load/tile.store
MemRef-based Allocation: Maps IR MemRef objects to pto.alloc_tile operations
Type-aware Conversion: Derives tile_buf/tensor_view types from TileType metadata
PTOAS Type Annotations: Emits typed ins/outs clauses for all operations

Generation Order¶

The codegen generates MLIR in the following fixed order:

Constants: arith.constant for index and float values
Tensor Views: pto.make_tensor_view for all tensor parameters
Allocations: pto.alloc_tile for all tile buffers (based on MemRef)
Operations: Function body with load, compute, store operations

The tensor-view and allocation prologue is rendered into a buffer before the constants block is finalized, so a constant that appears only inside a shape or stride expression — e.g. the 2 in a composite parameter dim M * 2 — is still declared in the constants block before its use.

Architecture¶

Class Structure¶

Header: include/pypto/codegen/pto/pto_codegen.h

namespace pypto::codegen {

class PTOCodegen : public CodegenBase {
 public:
  PTOCodegen();
  explicit PTOCodegen(const backend::Backend* backend);

  std::string Generate(const ir::ProgramPtr& program);

  // CodegenBase interface
  std::string GetCurrentResultTarget() const override;
  void Emit(const std::string& line) override;
  std::string GetExprAsCode(const ir::ExprPtr& expr) override;
  std::string GetTypeString(const DataType& dtype) const override;

  // PTO-specific helpers for operator codegen
  std::string NewTemp();
  std::string GetOrCreateTensorView(const ir::VarPtr& tensor);
  std::string GetOrEmitConstant(int64_t value, DataType dt);   // int/index overload
  std::string GetOrEmitConstant(double value, DataType dt);    // float overload
  std::string GetTensorViewTypeString(const ir::TensorType* tensor_type) const;
  std::string GetTileBufTypeString(const ir::MemRef* memref) const;
  std::string GetExprTypeAnnotation(const ir::ExprPtr& expr);
  std::string GetCurrentResultTileBufTypeString() const;
};

}  // namespace codegen

Implementation Components¶

File: src/codegen/pto/pto_codegen.cpp

Component	Purpose
`PTOCodegen`	Main visitor class (inherits `CodegenBase`) for IR traversal
`MemRefCollectorVisitor`	Collects MemRef objects and their associated TileType for allocation
Helper functions	`DataTypeToMLIRImpl()`, `MemorySpaceToMLIR()`

Python API¶

Basic Usage¶

from pypto.ir import compile, OptimizationStrategy
from pypto.backend import BackendType
import pypto.language as pl

@pl.program
class MyKernel:
    @pl.function
    def vector_add(self,
                   a: pl.Tensor[[32, 32], pl.FP32],
                   b: pl.Tensor[[32, 32], pl.FP32]):
        tile_a = pl.load(a, [0, 0], [32, 32])
        tile_b = pl.load(b, [0, 0], [32, 32])
        tile_c = pl.add(tile_a, tile_b)
        pl.store(tile_c, [0, 0], a)

# Compile with PTO backend
output_dir = compile(
    MyKernel,
    strategy=OptimizationStrategy.Default,
    backend_type=BackendType.Ascend910B,
)

The compile() function automatically applies the selected optimization strategy and invokes the appropriate codegen based on backend_type. Default is the only optimization strategy.

Direct Codegen Access¶

from pypto.pypto_core import codegen

# After pass transformations
pto_codegen = codegen.PTOCodegen()
pto_code = pto_codegen.generate(transformed_program)
print(pto_code)

Operator Mappings¶

Tile Operations → PTO Instructions¶

PyPTO Operation	Generated PTO-ISA
`tile.load(tensor, [row, col], [h, w])`	`pto.partition_view` + `pto.tload`
`tile.store(tile, [row, col], tensor)`	`pto.partition_view` + `pto.tstore`
`tile.slice(tile, [h, w], [row, col][, valid_shape=...])`	`pto.subview` (zero-copy view; `valid [...]` clause emitted only when `valid_shape` is supplied)
`tile.assemble(target, source, [row, col])`	(optional) `pto.tmov target -> dst` + `pto.subview dst[row, col] sizes [src.rows, src.cols]` + `pto.tmov src -> dst_view`
`tile.set_validshape(tile, vr, vc)`	`pto.set_validshape`; a view operand is rejected (see below)
`tile.mul(lhs, rhs)`	`pto.tmul`
`tile.addc(src0, src1, carry)`	`pto.taddc` (`src0 + src1 + carry`)
`tile.subc(src0, src1, carry)`	`pto.tsubc` (`src0 - src1 + carry`)
`tile.addsc(src0, scalar, carry)`	`pto.taddsc` (`src0 + scalar + carry`)
`tile.subsc(src0, scalar, carry)`	`pto.tsubsc` (`src0 - scalar + carry`)
`tile.adds(tile, scalar)`	`pto.tadds` (tile + scalar)
`tile.and_(lhs, rhs)` / `tile.ands(lhs, scalar)`	`pto.tand` / `pto.tands`; scalar is same-width signless `iN`
`tile.or_(lhs, rhs)` / `tile.ors(lhs, scalar)`	`pto.tor` / `pto.tors`; scalar is same-width signless `iN`
`tile.xor(lhs, rhs, tmp)` / `tile.xors(lhs, scalar, tmp)`	`pto.txor` / `pto.txors`; scalar is same-width signless `iN`
`tile.fillpad_expand(src, shape)`	`pto.tfillpad_expand ins(%src) outs(%dst)` (the `shape` tuple is type-deduction only; the larger `dst` and its pad come from the result type)

tile.slice / tile.assemble lowering details. Both ops are lowered through pto.subview, which is a pure view alias of the source tile (no data movement, no extra pto.alloc_tile). pto.subview requires the result tile_buf to share dtype, memory_space, blayout, slayout, fractal, pad, and compact with the source — DeduceTileSliceType propagates those five TileView fields from the source so the produced TileType satisfies the constraints by construction. Backend codegen also runs a CheckSubviewTileCompat guard at lowering time:

Source and result must both carry an explicit TileView.
dtype, blayout, slayout, fractal, pad, and compact must match exactly.
pad must be PadValue::null — pto.subview is a view, not a fillpad, so use tile.fillpad on the slice result if zero/min/max padding is required.

For tile.assemble, the leading pto.tmov target → dst is only emitted when buffer reuse did not collapse target and the destination buffer; in that case it preserves any data outside the insertion window. The trailing pto.tmov src → dst_view is the actual data write into the sub-window carved out by pto.subview.

tile.set_validshape lowering details. pto.set_validshape mutates the operand's valid_row / valid_col operands, so the operand must be a handle that has them: an alloc, an scf.if result, a cross-core pop slot. A view — the pto.subview a tile.slice lowers to, or a pto.treshape — carries its valid extent in its own type instead, so ptoas rejects the op against one; PyPTO therefore rejects it first, with a message pointing at the slice. View-ness is tracked at those emission sites rather than inferred from the rendered dims: a slice given a runtime valid_shape renders v_row=?, v_col=? exactly as an alloc-backed handle does. To narrow a view, pass valid_shape= to the slice (it lands in pto.subview's valid [...] clause and accepts runtime extents), or call set_validshape on the source tile before taking the view.

Cross-Core Operations → PTO Instructions¶

PyPTO Operation	Generated PTO-ISA	Description
`tile.tpush_to_aiv(tile, split=N[, id=I])`	`pto.tpush_to_aiv ins(%tile : type) {[id = I, ]split = N}`	Cube → Vector push
`tile.tpush_to_aic(tile, split=N[, id=I])`	`pto.tpush_to_aic ins(%tile : type) {[id = I, ]split = N}`	Vector → Cube push
`tile.tpop_from_aic(split=N[, id=I])`	`%buf = pto.tpop_from_aic {[id = I, ]split = N} -> type`	Pop from Cube pipe
`tile.tpop_from_aiv(split=N[, id=I])`	`%buf = pto.tpop_from_aiv {[id = I, ]split = N} -> type`	Pop from Vector pipe
`system.tfree_to_aic(tile_from_tpop[, id=I])`	`pto.tfree_from_aic {[id = I, ]split = N}`	Release a consumer slot back to Cube
`system.tfree_to_aiv(tile_from_tpop[, id=I])`	`pto.tfree_from_aiv {[id = I, ]split = N}`	Release a consumer slot back to Vector
`system.aic_initialize_pipe(...)`	`pto.aic_initialize_pipe {[id = I, ]dir_mask = D, slot_size = S[, slot_num = N][, local_slot_num = L]} (c2v_consumer_buf = %ssa : i32, v2c_consumer_buf = %ssa : i32)`	Cube pipe init (`slot_num`/`local_slot_num` emitted only when set; otherwise PTOAS uses its defaults)
`system.aiv_initialize_pipe(...)`	`pto.aiv_initialize_pipe {[id = I, ]dir_mask = D, slot_size = S[, slot_num = N][, local_slot_num = L]} (c2v_consumer_buf = %ssa : i32, v2c_consumer_buf = %ssa : i32)`	Vector pipe init (`slot_num`/`local_slot_num` emitted only when set; otherwise PTOAS uses its defaults)
`system.reserve_buffer(...)`	`%name = pto.reserve_buffer {name = "N", size = S, location = #pto.address_space<loc>, auto = false, base = B} -> i32`	Reserve buffer (`auto = true`, `base` omitted under `memory_planner=PTOAS`)
`system.import_peer_buffer(...)`	`%name = pto.import_reserved_buffer {name = "N", peer_func = @F} -> i32`	Import peer buffer
`system.syncall(core_type=C)`	`pto.syncall() mode = #pto.sync_all_mode<hard>, core_type = #pto.sync_core_type<C>`	Cross-core all-participant barrier (hard/FFTS form)
`system.syncall(mode="soft", core_type=C, gm_workspace=ws, used_cores=N)`	`pto.syncall(%gm_ptr[, %used] : !pto.ptr<i32>[, i32]) mode = #pto.sync_all_mode<soft>, core_type = #pto.sync_core_type<C>`	Current PTO-ISA soft/GM-polling barrier (partial occupancy; at least 64-byte GM workspace; explicit `N=0` derives the count from device launch registers and omits `%used`)

Notes:

Push ops use an ins() clause with a typed tile buffer; frontend pop ops produce an SSA result with a -> !pto.tile_buf<...> result type
id is optional. When omitted, PTOAS defaults to frontend pipe id 0. Use explicit ids only when authoring multiple independent frontend pipes; automatic bidirectional mixed-kernel setup keeps a single dir_mask = 3 pipe.
If the pushed tile was allocated with dynamic valid_row / valid_col operands or updated by tile.set_validshape, tpush emits the same tile handle after its runtime valid shape has been updated. For split tpush, codegen temporarily uses the full physical transport box, then restores the producer tile's logical valid shape.
split is pto-isa's TileSplitAxis, printed verbatim. 0 = no split, 1 / 2 = up-down / left-right, and 3 / 4 = the same two axes over an odd extent:

Code	pto-isa	Lane 0	Lane 1	Lane 1's band starts at
0	`TILE_NO_SPLIT`	whole tile	(single reader)	—
1 / 2	`TILE_UP_DOWN` / `TILE_LEFT_RIGHT`	`e0`	`e1`	`e1 * pitch`
3 / 4	`TILE_UP_DOWN_ODD` / `TILE_LEFT_RIGHT_ODD`	`e0`	`e1`	`(e1 + 1) * pitch`

eL is lane L's runtime valid extent on the split axis — the ISA reads it off the popped tile (popVecTileFromGMFiFo), so the even codes require e0 == e1 and the odd ones e0 == e1 + 1. LowerAutoVectorSplit materializes those extents and ExpandMixedKernel picks the matching code. - The Cube-to-Vector FIFO carries a compacted rectangle: the producer stores its valid_row x valid_col block at a valid_col row pitch, and each consumer lane reads its band back with the same pitch (gmStrideR = valid_col, doubled for the left-right codes). A partial valid shape on one side of the transport and not the other therefore mis-strides the pop — silent corruption of the valid region, since the ISA's matching assertion is compiled out in release builds. So a partial Acc-to-Vec transfer uses the full physical box for TPOP and for the TPUSH columns, whether the transfer is no-split or split, and restores the logical valid shape immediately on each side of the transport — on the consumer side through a metadata-only pto.treshape (a frontend tpop result is not a locally bound PTOAS tile, so pto.set_validshape cannot restore it in place). The split-axis extent is one exception: it stays per-lane on the TPOP operands, because that is what tells the ISA where lane 1's band begins — which is also why a per-lane extent the compiler could not verify must never get there. A pl.split_aiv boundary whose split-axis extent is a runtime value keeps the FULL box on the popped tile (split_axis::WithFullSplitAxisValid) so the even code's band lands on the box half, and carries the lane's own extent on the consumers instead. - The row extent of a no-split Acc-to-Vec TPUSH is the other exception: it stays exactly as the producer wrote it. TPUSH runs TStoreAccNz2nd out of L0C, whose source pitch is ceil(validRow/16)*16 for a compact tile and TileData::Rows otherwise, while mad laid the product out at the pitch implied by the L0A operand's valid rows. Widening validRow before the push re-derives that pitch from the physical box, so the fix-pipe walks L0C at a stride mad never wrote at — with a 64-row box valid to 16 the push picks up N-fractal 4j for every fractal j (issue #2510). The rows past validRow stay stale in the slot, which is what a narrowed valid_shape already promises about its invalid region, and the transport moves validRow rows instead of the whole box. - A split Acc-to-Vec transport cannot take that route: lane 1 reads the band starting at the box half, which exists only if the producer wrote the full box — and writing it means reading L0C at the physical pitch, which is not the pitch mad used. The two requirements are mutually exclusive, so a row-narrowed compact accumulator crossing a pl.split / pl.split_aiv boundary is rejected with a message naming both DSL alternatives (narrow the result instead of the operand, or stage the accumulator through GM), rather than lowered into silently skewed data — measured on device at 1808 of 8192 elements wrong before the refusal. The refusal is gated on the pitches actually differing, so a single-fractal-block accumulator (ceil(validRow/16)*16 == Rows) keeps crossing as before. - When a tpop result TileView.valid_shape differs from the physical tile shape, PTO codegen emits PTOAS frontend operands as %buf = pto.tpop_from_*(%valid_row, %valid_col) {[id = I, ]split = N} -> !pto.tile_buf<..., v_row=?, v_col=?, ...>. This covers dynamic expressions and static non-full shapes such as [0, 0]; the operands carry the logical extents used by compute and store. The full-box Cube-to-Vector transport above overrides this for a statically-shaped, non-empty partial pop, because pto.treshape carries no valid-row/valid-col operands and so can only restore static logical extents. - For split consumers of a hand-written pop, SplitVectorKernel localizes those dynamic tpop valid-shape operands per subblock (for example global [8, 16] becomes [8, 16] then [0, 16] under up/down split of a [16, 16] tile). An odd split axis reaches the operands the same way — a [17, 128] tile pops [9, 128] on lane 0 and [8, 128] on lane 1 under split = 3. - system.tfree_* derives split from its tile argument, so the frontend must free the exact SSA value produced by tile.tpop_*, even though the PTO instruction itself does not take the tile as an explicit operand - ExpandMixedKernel now auto-generates consumer-side system.tfree_* after split-generated tile.tpop_*, preserving tpop -> direct users -> tfree -> next tpop - reserve_buffer and import_reserved_buffer return i32 SSA values; initialize_pipe references them as operands - Under memory_planner=PYPTO or DSA_RP, AllocateMemoryAddr resolves reserve_buffer(base=AUTO) before PTO emission, so PTO emits auto = false, base = <value>. Under memory_planner=PTOAS that pass is skipped, so PTO emits auto = true with base omitted (ptoas rejects both attributes together) and ptoas PlanMemory places the reserved region - reserve_buffer location is mat for AIC functions, vec for AIV/InCore functions - import_reserved_buffer uses MLIR symbol syntax (@func_name) for peer_func - Buffer name and peer_func strings are validated by CheckSafeIdentifier (alphanumeric + underscore only)

Parameter Type Handling¶

PyPTO Type	MLIR Parameter Type	Post-processing
`TensorType`	`!pto.ptr<dtype>`	Generate `pto.make_tensor_view`
`ScalarType`	`dtype` (e.g., `f32`)	Direct usage as `%argN`
`TileType`	Not allowed as parameter	Must be computed internally

Code Generation Details¶

Tensor View Generation¶

For each TensorType parameter, the codegen generates:

%0 = pto.make_tensor_view %arg0,
     shape = [%c32_index, %c32_index]
     strides = [%c32_index, %c1_index]
     {layout = #pto.layout<nd>}
     : !pto.tensor_view<?x?xf32>

Key aspects:

Shape from TensorType.shape_
Strides computed as row-major: [dim1, 1] for 2D tensors
Constants (%c32_index, %c1_index) auto-generated, including any that appear only inside a composite shape/stride expression
Composite dims lower to arithmetic, e.g. M * 2 → arith.muli %M, %c2_index
Tensor view type uses ? for each dimension (e.g., ?x?xf32 for 2D)

Layout Handling for 2D Tensors¶

The layout attribute on make_tensor_view tells PTOAS the memory layout convention. The codegen determines shape, strides, and layout based on the tensor's IR type and shape:

Case	Shape emitted	Strides emitted	Layout	Notes
ND `[R, C]`	`[R, C]`	`[C, 1]`	`nd`	Standard row-major
DN `[R, C]` (both > 1)	`[C, R]`	`[1, C]`	`dn`	Shape swapped for PTOAS column-major convention
Column vector `[M, 1]`	`[M, 1]`	`[1, M]`	`dn`	Auto-detected, no DN annotation needed

Column vector auto-DN: Any 2D tensor whose last dimension is a compile-time constant 1 (i.e., shape [M, 1]) is automatically emitted with layout = dn and strides [1, M]. This is required because PTOAS always infers DN for the shape/stride pattern [M, 1] / [1, 1], making the degenerate ND representation ambiguous. The codegen resolves this by always using unambiguous DN strides. Users do not need to annotate [M, 1] tensors with pl.DN in the DSL.

Example for a [16, 1] column vector (no DN annotation in DSL):

%col_view = pto.make_tensor_view %arg1,
    shape = [%c16_index, %c1_index], strides = [%c1_index, %c16_index]
    {layout = #pto.layout<dn>}
    : !pto.tensor_view<?x?xf32>

Allocation Generation¶

Based on TileType variables collected from the function body. Each tile variable gets its own pto.alloc_tile instruction with an explicit addr attribute derived from the variable's MemRef. Variables sharing the same MemRef share the same address:

%mi_tile = pto.alloc_tile addr = %c8320_i64 : !pto.tile_buf<loc=vec, dtype=f32, rows=16, cols=1,
                      v_row=16, v_col=1, blayout=col_major,
                      slayout=none_box, fractal=512, pad=0>
%mi_tile_nd = pto.alloc_tile addr = %c8320_i64 : !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=16,
                      v_row=1, v_col=16, blayout=row_major,
                      slayout=none_box, fractal=512, pad=0>

Tile variable → alloc_tile mapping:

Memory space (TileType.memory_space_) → loc attribute (using PTO address space names)
Tile dtype and dimensions derived from each variable's own TileType metadata
One allocation per tile variable (not per unique MemRef)
addr attribute from MemRef.addr_, emitted as arith.constant ... : i64
Variables sharing the same MemRef produce the same addr SSA value

Who plans memory: `compile(memory_planner=...)`¶

Who assigns the physical addr is selected by the memory_planner option (ir.compile(..., memory_planner=passes.MemoryPlanner.PYPTO | DSA_RP | PTOAS), default PYPTO). It threads to both the pass pipeline (via PassContext) and codegen:

Mode	Pipeline	`pto.alloc_tile`	`pto.reserve_buffer`	ptoas
`PYPTO` (default)	runs `MaterializeSemanticAliases` + `MemoryReuse` + `AllocateMemoryAddr`	emits `addr = <const>` (from `MemRef.byte_offset_`)	`auto = false, base = <const>`	`--pto-level=level3` (trusts baked addresses)
`DSA_RP`	runs `MaterializeSemanticAliases` + `AllocateMemoryAddr`; skips `MemoryReuse`	emits the in-process canonical-greedy DSA-RP `addr = <const>`	`auto = false, base = <const>`	`--pto-level=level3` (trusts baked addresses)
`PTOAS`	runs `MaterializeSemanticAliases`; skips `MemoryReuse` + `AllocateMemoryAddr`	omits `addr` (`PTOCodegen.generate(emit_tile_addr=False)`)	`auto = true` (no `base`)	`--pto-level=level2` (ptoas `PlanMemory` does reuse + addresses)

Memory planning is split into two passes: MaterializeSemanticAliases forces semantics-required aliasing (loop-carried accumulators, in-place ops) to share one MemRef, while MemoryReuse does opportunistic lifetime-based coalescing of independent buffers for PYPTO. DSA_RP skips that coalescing and places the independent identities under capacity and reuse penalties in AllocateMemoryAddr. InitMemRef + MaterializeSemanticAliases run in all three modes, so must-alias buffers survive. In PTOAS mode, ptoas PlanMemory (which level2 requires, rejecting any addr operand) performs lifetime reuse and address assignment.

Caveat: PTOAS mode skips the Ascend910B load + tpop_from_aic in-place hazard guard (part of MemoryReuse) and reserve-buffer base resolution (AllocateMemoryAddr); those are deferred to ptoas. compile() emits a warning — verify affected kernels on-device.

Multi-slot declarations become one ptoas region (`PTOAS` mode)¶

A declared multi-slot allocation (pl.MemRef(slots=N), see Python syntax) is not lowered to N alloc_tiles. It maps onto ptoas's own multi-buffer pair, one region declared in the function head and one slot selection per use:

%l0c_mb = pto.alloc_multi_tile valid_row = %c64_index valid_col = %c64_index
        : !pto.multi_tile_buf<!pto.tile_buf<loc=vec, dtype=f32, rows=64, cols=64, ...>, count=2>
scf.for %i = %c0_index to %c4_index step %c1_index {
  %0 = arith.remsi %i, %c2_index : index
  %t = pto.multi_tile_get %l0c_mb[%0]
     : !pto.multi_tile_buf<..., count=2> -> !pto.tile_buf<loc=vec, ...>
  ...
}

Two properties matter:

No addr. ptoas PlanMemory places the region and is forbidden to merge its slots, which is what carries the author's separation into level2.
The operand is the slot index, not the byte offset InitMemRef derived from it. ptoas matches the index's affine form (i % 2) to decide which accesses can share a slot, and that is what earns the rotation per-slot (dynamic) event ids — iteration i's load overlapping iteration i-1's compute.

PlanMultiBufferRegions decides eligibility before the body walk; a shape ptoas cannot describe (slots holding differently shaped tiles, slots declaring different valid shapes, two slots live at once inside a loop, a space other than Vec / Mat / Acc, a runtime valid shape, a slot carried out of an if or loop as a phi, a count outside ptoas's [2, 16]) is a ValueError naming the shape, because falling back to per-slot alloc_tile would let ptoas plan the slots on top of each other.

One slot per iteration. The co-live rejection is not a shape ptoas fails to type — it is one it fails to synchronize. ptoas 0.54 derives the per-slot WAR guard only for the first multi_tile_get of an iteration; given two, the second load is emitted with no wait_flag, so the next iteration overwrites that slot while the current one still reads it. Measured wrong on device, so codegen refuses the shape and points at the PyPTO planner, whose baked addresses and PyPTO-emitted sync handle it. Straight-line code is unaffected — with no loop there is no cross-iteration reuse to guard. Filed as PTOAS#1118; lifting the restriction is one condition in PlanMultiBufferRegions.

Under PYPTO no region is emitted at all: at --pto-level=level3 ptoas does not fold its per-slot address fan-out, so the region form would lose the slot analysis it exists for (PTOAS#1106).

Load Operation Transformation¶

PyPTO IR:

tile_a = pl.load(tensor_a, [0, 0], [32, 32])

Generated MLIR (two operations):

# 1. Create partition view
%3 = pto.partition_view %tensor_view, offsets = [%c0_index, %c0_index],
                 sizes = [%c32_index, %c32_index]
                 : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<32x32xf32>

# 2. Load into tile buffer
pto.tload ins(%3 : !pto.partition_tensor_view<32x32xf32>)
          outs(%tile_buf : !pto.tile_buf<loc=vec, ...>)

Key transformations:

Tensor parameter → tensor_view lookup
Offsets/sizes from tile.load arguments
Output tile_buf from variable's MemRef with type derived from TileType

Store Operation Transformation¶

PyPTO IR:

pl.store(tile_c, [0, 0], tensor_out)

Generated MLIR:

# 1. Create partition view for output
%5 = pto.partition_view %output_view, offsets = [%c0_index, %c0_index],
                 sizes = [%c32_index, %c32_index]
                 : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<32x32xf32>

# 2. Store from tile buffer
pto.tstore ins(%tile_buf : !pto.tile_buf<loc=vec, ...>)
           outs(%5 : !pto.partition_tensor_view<32x32xf32>)

Compute Operations¶

Example: Tile Multiplication¶

PyPTO:

tile_c = pl.mul(tile_a, tile_b)

MLIR:

pto.tmul ins(%tile_a_buf : !pto.tile_buf<...>,
             %tile_b_buf : !pto.tile_buf<...>)
         outs(%tile_c_buf : !pto.tile_buf<...>)

Result handling:

Result variable's MemRef determines output tile_buf
Input operands resolved through variable name lookup
All ins/outs clauses include type annotations

Source Locations (`loc`)¶

Every emitted operation carries a trailing MLIR location built from the IR Span, e.g. pto.tadd ins(...) outs(...) loc("kernels/attn.py":41:9). ptoas propagates loc() verbatim into its diagnostics, so a verifier rejection names the user's .py line instead of a line in the generated .pto — a file that, under @pl.jit, the user never sees (spans there are already remapped from the synthesized <jit:name> text back to the real source).

Which span is used — bound at two levels, the second refining the first:

Level	Bound in	Source
Statement (primary)	`PTOCodegen::VisitStmt`	`Stmt::span_`
Call (refinement)	`PTOCodegen::VisitExpr_(CallPtr)`	`Call::span_`, only when nested inside the statement span

The containment test is what makes this correct. Call::span_ is column-accurate when preserved, but passes that synthesize tile ops (ConvertTensorToTileOps) rebuild the Call carrying the enclosing function's span while leaving the AssignStmt's own span intact. Such a span begins before the statement, fails containment, and is discarded in favour of the statement span — otherwise most operations would report the def line.

No location is emitted for: region braces, separators and block labels (loc(...) is legal only at the end of a complete operation, so these use EmitStructural() rather than Emit()); arith.constant in the constants section (deduplicated across uses, so no single span fits); and nodes whose span is unknown or has no filename.

Disabling — Generate(program, emit_tile_addr, emit_source_loc), compile(..., emit_source_loc=...), or PYPTO_EMIT_PTO_LOC=0. The output is then byte-identical to the location-free form; this is the escape hatch for a ptoas build whose parser rejects a trailing location, since ptoas ships independently of PyPTO.

Boxed Tile Extents¶

Every pto.alloc_tile PyPTO emits is validated against the box grid PTOAS will check it against. PTO addresses a boxed tile one box at a time, so a tile whose physical extent is not a whole number of boxes has no address at all.

The rule mirrors PTOAS' verifyBoxedTileLayout exactly:

Layout	Box (rows x cols)
fractal 1024 (`Acc`)	16 x 16
fractal 512, `slayout = row_major` (`Mat` / `Left`)	16 x (32 / sizeof(dtype))
fractal 512, `slayout = col_major` (`Right`, the transposed dual)	(32 / sizeof(dtype)) x 16
`slayout = none_box`	not boxed — no rule

with PTOAS' own exemptions kept: the row rule is skipped for Vec and for a single-row tile (the NZ map degenerates there), while the column rule always applies. The MX-scale fractal and sub-byte carriers are left to PTOAS, which diagnoses them itself.

Why here rather than in PTOAS. PTOAS rejects the same shape, but its message names its own internals and offers no remedy:

'pto.alloc_tile' op expects result boxed tile rows to be a multiple of innerRows (16), but got 100

Raising at the emission site instead reports the tile, the axis, the extent to reach, and how to reach it:

a Mat tile of physical shape [100, 128] and dtype fp16 must be a whole number of
16x16 fractal boxes, but its row extent 100 is not a multiple of 16. PTO addresses
a boxed tile one box at a time, so a partial box has no address. The *logical*
extent is free -- allocate 112 on that axis and declare 100 as the tile's
valid_shape (`valid_shape=` on pl.load / pl.tile.create + pl.set_validshape),
which moves and computes only the real data. A tensor-level pl.matmul /
pl.matmul_acc does this for its M axis automatically.

ComputeAllocTileFields is the single choke point every allocation passes through — the per-variable declaration, the hoisted extra_alloc_tiles, and the control-flow paths alike — so the check sees exactly what is emitted and cannot drift from it. A tensor-level pl.matmul / pl.matmul_acc never trips it on its M axis, which is boxed for the user in ConvertTensorToTileOps. The axes that remain the user's responsibility are K and N.

Complete Example¶

Input: PyPTO Program¶

import pypto.language as pl

@pl.program
class MulKernel:
    @pl.function
    def mul_kernel_2d(self,
                     a: pl.Tensor[[32, 32], pl.FP32],
                     b: pl.Tensor[[32, 32], pl.FP32],
                     c: pl.Tensor[[32, 32], pl.FP32]):
        # Load tiles
        tile_a = pl.load(a, [0, 0], [32, 32])
        tile_b = pl.load(b, [0, 0], [32, 32])

        # Multiply
        tile_c = pl.mul(tile_a, tile_b)

        # Store result
        pl.store(tile_c, [0, 0], c)

Output: PTO-ISA MLIR¶

module {
  func.func @mul_kernel_2d(%arg0: !pto.ptr<f32>,
                          %arg1: !pto.ptr<f32>,
                          %arg2: !pto.ptr<f32>) {
    // Constants
    %c32_index = arith.constant 32 : index
    %c1_index = arith.constant 1 : index
    %c0_index = arith.constant 0 : index

    // Tensor views
    %3 = pto.make_tensor_view %arg0, shape = [%c32_index, %c32_index]
         strides = [%c32_index, %c1_index] : !pto.tensor_view<?x?xf32>
    %4 = pto.make_tensor_view %arg1, shape = [%c32_index, %c32_index]
         strides = [%c32_index, %c1_index] : !pto.tensor_view<?x?xf32>
    %5 = pto.make_tensor_view %arg2, shape = [%c32_index, %c32_index]
         strides = [%c32_index, %c1_index] : !pto.tensor_view<?x?xf32>

    // Allocations
    %0 = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, ...>
    %1 = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, ...>
    %2 = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, ...>

    // Load tile_a
    %6 = pto.partition_view %3, offsets = [%c0_index, %c0_index], sizes = [%c32_index, %c32_index]
         : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<32x32xf32>
    pto.tload ins(%6 : !pto.partition_tensor_view<32x32xf32>)
              outs(%0 : !pto.tile_buf<...>)

    // Load tile_b
    %7 = pto.partition_view %4, offsets = [%c0_index, %c0_index], sizes = [%c32_index, %c32_index]
         : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<32x32xf32>
    pto.tload ins(%7 : !pto.partition_tensor_view<32x32xf32>)
              outs(%1 : !pto.tile_buf<...>)

    // Multiply
    pto.tmul ins(%0 : !pto.tile_buf<...>, %1 : !pto.tile_buf<...>)
             outs(%2 : !pto.tile_buf<...>)

    // Store tile_c
    %8 = pto.partition_view %5, offsets = [%c0_index, %c0_index], sizes = [%c32_index, %c32_index]
         : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<32x32xf32>
    pto.tstore ins(%2 : !pto.tile_buf<...>)
               outs(%8 : !pto.partition_tensor_view<32x32xf32>)

    return
  }
}

Variable Mapping¶

Internal Tracking¶

The codegen maintains several mappings to track MLIR variable names:

Mapping	Purpose	Example
`var_to_mlir_`	IR variable → MLIR SSA name	`"tile_a"` → `"%0"`
`tensor_to_view_`	Parameter → tensor_view	`"a"` → `"%3"`
`memref_to_mlir_`	MemRef pointer → tile_buf	`memref.get()` → `"%0"`
`memref_to_tile_type_`	MemRef pointer → TileType	Used for deriving tile_buf types

SSA value naming:

Parameters: %arg0, %arg1, %arg2, ...
Constants: %c0_index, %c1_index, %c32_index, %c0_i64, %cst, ...
Results: %0, %1, %2, ...

MemRef-based Resolution¶

For operations like tile.mul:

tile_c = pl.mul(tile_a, tile_b)

The codegen:

Resolves tile_a → %0 via var_to_mlir_
Resolves tile_b → %1 via var_to_mlir_
Gets tile_c's MemRef from its TileType
Maps MemRef → %2 via memref_to_mlir_
Gets tile_buf type from memref_to_tile_type_
Generates: pto.tmul ins(%0 : !pto.tile_buf<...>, %1 : !pto.tile_buf<...>) outs(%2 : !pto.tile_buf<...>)

Type Conversions¶

DataType Mapping¶

PyPTO DataType	MLIR Type
`DataType::FP32`	`f32`
`DataType::FP16`	`f16`
`DataType::BF16`	`bf16`
`DataType::INT32`	`i32`
`DataType::INT64`	`i64`
`DataType::INT8`	`i8`
`DataType::UINT8`	`ui8`

Memory Space Mapping¶

PyPTO MemorySpace	PTO Address Space
`MemorySpace::DDR`	`gm` (global memory)
`MemorySpace::Vec`	`vec` (vector buffer)
`MemorySpace::Mat`	`mat` (matrix buffer)
`MemorySpace::Left`	`left`
`MemorySpace::Right`	`right`
`MemorySpace::Acc`	`acc` (accumulator)
`MemorySpace::Bias`	`bias` (bias buffer)

Tile Buffer Attributes¶

Generated alloc_tile operations derive dtype and dimensions from TileType metadata, and layout/fractal/pad/compact mode from the associated TileView (when available):

!pto.tile_buf<
  loc=vec,             // PTO address space (from MemorySpace)
  dtype=f32,           // Element data type (from TileType)
  rows=32,             // Tile height (from TileType shape)
  cols=32,             // Tile width (from TileType shape)
  v_row=32,            // Virtual row size (= rows)
  v_col=32,            // Virtual column size (= cols)
  blayout=row_major,   // Block layout (from TileView, default: row_major)
  slayout=none_box,    // Scatter layout (from TileView, default: none_box)
  fractal=512,         // Fractal size in bytes, not elements (from TileView, default: 512)
  pad=0,               // Pad mode as int (from TileView, default: 0/null)
  compact=1            // Optional compact mode (normal=1; omitted for null=0)
>

TileView-derived attributes:

Attribute	Source	Enum Values	Default
`blayout`	`TileView::blayout`	`none_box`, `row_major`, `col_major`	`row_major`
`slayout`	`TileView::slayout`	`none_box`, `row_major`, `col_major`	`none_box`
`fractal`	`TileView::fractal`	uint64	`512`
`pad`	`TileView::pad`	`null(0)`, `zero(1)`, `max(2)`, `min(3)`	`null(0)`
`compact`	`TileView::compact`	`null(0)`, `normal(1)`	`null(0)`

When no TileView is associated with the MemRef, the codegen falls back to the default values listed above. The compact attribute is omitted for its null default. Two paths set normal(1) automatically:

A partial tile.extract into L0A/L0B, so TEXTRACT transfers only the logical valid_shape instead of treating box-alignment padding as data.
An Acc (L0C) tile whose valid rows are not provably equal to its physical rows, as produced by the tile.matmul, tile.matmul_bias, and tile.matmul_mx deducers. mad always lays the product out with an N-fractal stride of ceil(validRow/16)*16 taken from the lhs valid rows, while every Acc reader derives its stride from the compile-time physical Rows unless the tile is compact. Without the flag a runtime-narrowed accumulator is read back at a different pitch than it was written at. Only the row extent decides this — every Acc stride the ISA derives is a function of validRow alone, so a narrowed column extent keeps the non-compact form.

Compact is stamped only where the accumulator's layout is established, never re-derived on an alias of it. tile.matmul_acc (and matmul_mx_acc) inherit the accumulator operand's mode, because the op reuses that operand's buffer in place and codegen aliases the two only when their full tile config matches. tile.set_validshape likewise inherits: it is metadata-only and may run after the buffer was written, so the pitch its readers must use is the one mad already wrote at — deriving a new one from the narrowed rows would re-interpret bytes that were never repacked.

A buffer can also declare the mode at creation: tile.create(..., target_memory=Acc, compact=True). A fresh L0C buffer has no prior bytes to re-interpret, so a declaration is not an alias re-derivation — and it is what AutoTileMatmulL0 puts on the accumulator seed it synthesizes when it splits K, since tile.matmul_acc inherits from that seed and a non-compact one drags the whole chain, and the reader after the loop, back to the physical pitch. A declaration is also the only form that survives: a type a pass stamps onto a call is discarded as soon as any later pass re-deduces it (InferTileMemorySpace does), whereas a kwarg is re-read every time.

AccCompactValid (see Verifier) checks both halves of the contract: every tile.matmul_acc accumulates into a compact buffer when mad's pitch differs from the accumulator's physical row count, and no tile outside Left/Right/Acc carries a compact mode at all.

Note that the Acc → L1 readers (TExtractAccToMat, TMovCcToCb) have no CompactMode branch in PTO-ISA on either a2a3 or a5, so a runtime-narrowed accumulator consumed by tile.extract / tile.move into L1 still reads at the physical Rows pitch. That gap needs a matching PTO-ISA change.

Kernel Wrapper Generation (PTO Backend)¶

When compiling with the PTO backend via ir.compile(), a kernel wrapper is automatically generated for each InCore function to bridge the ptoas output to the orchestration calling convention.

Pipeline¶

InCore Function → PTOCodegen → .pto → ptoas → .cpp → kernel_wrapper → kernels/aiv/<name>.cpp

Each InCore function is compiled independently through ptoas. The final wrapper file combines:

Preprocessed ptoas code (with __global__ AICORE → static)
kernel_entry(__gm__ int64_t* args) wrapper that unpacks the args array and forwards to the ptoas function

Output Structure¶

When the program contains an Orchestration function, the PTO backend generates the following output structure:

output_dir/
├── passes_dump/                     # IR after each pass
├── ptoas_passes/                    # Optional ptoas IR after each pass
│   └── <kernel-or-group>/            # ptoas/MLIR-managed dump tree
├── ptoas/                           # Intermediates
│   ├── <func_name>.pto              # MLIR from PTOCodegen
│   └── <func_name>.cpp              # C++ from ptoas
├── kernels/aiv/
│   └── <func_name>.cpp              # Final wrapper
├── orchestration/
│   └── <orch_func_name>.cpp         # simpler runtime orchestration code
└── kernel_config.py                 # Runtime/orchestration/kernel config

ptoas_passes/ is emitted only when ir.compile(..., dump_ptoas_passes=True) or RunConfig(dump_ptoas_passes=True) is used.

The orchestration codegen generates identical orchestration C++ code using the simpler runtime API (rt_submit_task, make_tensor_external, etc.).

Runtime configuration (`kernel_config.py`)¶

kernel_config.py exposes a RUNTIME_CONFIG dict that the dispatch path reads to launch the program. Stable keys:

Key	When emitted	Notes
`runtime`	Always	`"tensormap_and_ringbuffer"` (default) or `"host_build_graph"` — the wire name of the `RuntimeKind` selected by `ir.compile(runtime=...)`, or by wrapping the call in `PassContext([], runtime=...)`.
`aicpu_thread_num`	Always (`0`)	`0` selects the runtime's architecture default (a2a3: 4; a5: 5); callers may explicitly override it.

The runtime is carried by PassContext as an ir::RuntimeKind, not by a codegen-only argument, so passes that must legalize IR for a specific runtime switch on PassContext::GetRuntime() rather than comparing strings. It is an enum rather than a name because the set is closed — one enumerator per implementation under runtime/src/<arch>/runtime/ — so a typo is a compile error instead of a value that surfaces much later as an opaque CCEC error about a nonexistent include directory.

The wire name crosses the ABI boundary in exactly one place each way: ir::RuntimeKindToName when writing kernel_config.py, and ir::RuntimeKindFromName when reading one back. Both are re-exported to Python as passes.runtime_kind_to_name / passes.runtime_kind_from_name.

The runtime is also part of the @pl.jit cache key, so a host_build_graph call cannot reuse an artifact compiled for tensormap_and_ringbuffer, whose kernel_config.py names a runtime no matching worker would bind.

Argument Unpacking¶

The wrapper unpacks int64_t* args following the standard convention:

Parameter Type	Unpacking Pattern
`TensorType`	`ChipTensor*` → `buffer.addr` → typed pointer
`ScalarType`	`uint64_t` → union decode → typed value

SPMD Identity Parameters¶

tile.get_block_idx(), tile.get_block_num(), and tile.get_subblock_idx() lower to synthetic i32 parameters that PTOCodegen appends at the end of the func.func signature using named SSAs (%__pypto_spmd_block_idx, %__pypto_spmd_block_num, %__pypto_spmd_subblock_idx). The IR contract for these ops is unchanged — the synthetic params live only in the generated MLIR / C++ and never appear in Function.params. They are appended in the canonical order block_idx, block_num, subblock_idx, each gated independently on the ops the function actually uses.

func.func @spmd_kernel(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f32>,
                       %__pypto_spmd_block_idx: i32,
                       %__pypto_spmd_block_num: i32)
                       attributes { ... } {
  %0 = arith.index_cast %__pypto_spmd_block_idx : i32 to index
  // ... use %0 as the block index ...
}

The kernel wrapper resolves the runtime values once from intrinsic.h::get_block_idx(args) / get_block_num(args) and forwards them as the trailing two call args:

extern "C" __aicore__ __attribute__((always_inline))
void kernel_entry(__gm__ int64_t* args) {
    // Read logical SPMD block identity from runtime dispatch payload
    int32_t __pypto_spmd_block_idx = get_block_idx(args);
    int32_t __pypto_spmd_block_num = get_block_num(args);

    // ... tensor / scalar / dyn-dim unpacking ...

    // Forward to ptoas-generated function (block args at the end)
    spmd_kernel(a, out, __pypto_spmd_block_idx, __pypto_spmd_block_num);
}

subblock_idx (AIV lane). tile.get_subblock_idx() uses the same synthetic-param channel: the wrapper resolves it from intrinsic.h::get_sub_block_id(args) (the runtime per-core lane id the scheduler stores in GlobalContext.sub_block_id) and appends __pypto_spmd_subblock_idx after any block-identity args. It deliberately reads the runtime lane id rather than the ccec get_subblockid() register, which returns a stale value under the tensormap_and_ringbuffer dispatch.

Split AIV FIFO endpoints use that runtime value even when the tensor program does not call tile.get_subblock_idx(): after PTOAS lowers a split endpoint, the wrapper backend forwards the lane as the third argument of PTO-ISA's explicit TPUSH(pipe, tile, subblock_id) / TPOP(...) overload. If the function has no synthetic subblock parameter already, the backend adds a private trailing parameter to the generated function. The explicit overload derives the byte offset from each call's actual tile type, so one automatic pipe can safely carry differently sized sequential transfers. Like block identity, the wrapper resolves the runtime lane unconditionally because GlobalContext.sub_block_id is populated by the scheduler on every platform. The endpoint argument is build-guarded: device builds pass the lane to the explicit overload, while CPU simulation and in-core cost-model builds retain PTO-ISA's normal two-argument endpoint because those implementations already model their lane context and do not expose the device-only explicit-lane overload.

Detection scope. Both layers detect SPMD usage on a per-function basis:

MemRefCollectorVisitor::UsesSpmdBlockOps / UsesSubblockOp (C++, src/codegen/pto/pto_codegen.cpp) drive whether PTOCodegen appends the block / subblock params to a given function's signature.
_uses_spmd_block_ops / _uses_dynamic_subblock_id (Python, python/pypto/backend/pto_backend.py) drive whether the wrapper appends the matching locals to the inner call site. Split TPUSH / TPOP endpoints are additionally detected by _runtime_split_fifo_endpoint_counts; they reuse that subblock argument or request the private one described above.

For non-SPMD sibling functions in an SPMD group (group_uses_spmd=True but the function itself does not call tile.get_block_*), the wrapper still declares the two block locals because the __gm_pipe_buffer sharding logic in _generate_arg_unpacking consumes them — but it does not append them to the inner call, matching the function's MLIR signature.

This replaces the earlier macro-shadow + [[block_local]] static / static thread_local bridge plus #pragma push_macro / #undef / pop_macro dance. Block and lane identity now flow through the call graph like every other per-launch value (tensor pointers, scalar args, dynamic dims).

Implementation¶

Module: python/pypto/backend/pto_backend.py

Key functions:

generate() — entry point: produces all PTO backend files (kernels + orchestration + config)
_preprocess_ptoas_output() — strips duplicate includes, makes functions static
_generate_arg_unpacking() — generates C++ unpacking code from IR parameter types
_generate_kernel_wrapper() — assembles the complete wrapper file