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.

Shape expressions may retain tensor.dim(x, axis). Codegen emits the selected declared extent as an INDEX operand, including constant, symbolic and composite extents; it does not require an existing tensor view for x or a prior IR fold. Only scalar symbols in that extent become dynamic-dimension ABI arguments, not the tensor operand itself. Printed type annotations preserve these reads when parsed again, without evaluating parameter references as Python values. Shape-only assignment annotations replace an inferred shape only when direct dimension-read temporaries prove structural equivalence; otherwise the inferred runtime shape is retained. Constant extents are emitted directly as INDEX, including unsigned constants.

Dynamic shape arguments across outlined boundaries¶

An outlined scope captures scalar symbols needed by its tensor parameter types, even when the kernel body never reads them directly. A bare axis such as [B, 16] binds B through its tensor descriptor. Supported invertible axes such as [B * 2, 16] also retain descriptor extraction, so the caller need not have a separate scalar binding for B. A composite axis such as [(B + 15) // 16 * 16 + T * 4, 16] does not bind either symbol: the outliner conservatively adds explicit B and T parameters and passes their caller values. These symbols must have caller-side runtime bindings: scalar parameters, body value definitions, or bare dimensions of tensor parameters. A symbol appearing only inside a non-invertible type expression is rejected during outlining, before it can become an undeclared orchestration argument. It keeps the expression intact; it does not replace it with a new ROWS symbol.

IR scalar parameters precede annotations that reference them. Runtime dispatch and the native kernel signature instead use tensor pointers first, then explicit scalar payloads, then any remaining descriptor-derived dimensions. PTO codegen and the wrapper share collect_tensor_shape_dyn_vars for this last group. An explicit scalar wins over descriptor extraction, deduplicated by Var identity, so a symbol used by both a shape and the body is forwarded exactly once.

For authored kernels, descriptor extraction still supports bare dimensions and single-symbol addition, subtraction, and nonzero constant multiplication. Floor division has no general inverse: B // 5 == 2 does not prove B == 10. Pass such symbols explicitly or expose them through another supported dimension; otherwise wrapper generation reports an error. Existing runtime descriptors and scalar packing are unchanged.

Architecture¶

Explicit Buffer input¶

GenerateBufferFunction validates explicitly constructed Buffer IR before emission. Its GM path supports static, packed ND rank-2 FP16/BF16/FP32/INT32/INT8 Tensor parameters and normalized Tensor parameter returns. It shares the existing GM tensor-view prologue and native tensors-first/scalars-last ABI. Tensor returns remain in IR for orchestration aliasing; they do not create native return values.

buffer.load and buffer.store each render an explicit window as pto.partition_view followed by pto.tload or pto.tstore. Their destination, offsets, and valid extents come from ordinary operands. buffer.add, mul, and copy render their explicit buffer destinations. Each pto.alloc_tile comes from a buffer.alloc in the same scope; an address is emitted exactly when its operand is present, independently of the legacy emission flag. GM transfers never create a buffer, infer valid-state updates, or reconstruct a logical Tile.

With the default enable_buffer_ir=True, the pipeline runs LowerTileToBuffer for selected whole kernels before entering this direct path. Unsupported kernels retain Functional Tile IR; False keeps all kernels on that path. See Buffer contracts for descriptor, direction, dynamic-window, and ABI limits. Native compilation tests establish syntax and operand dataflow; numerical execution is a separate integration requirement.

Storage aliases arrive as explicit buffer.subview and buffer.reshape operations; a subview with a valid tuple renders a valid [..] clause whose constant or SSA operands type each native result dimension. Codegen maps them to pto.subview and pto.treshape, preserving the source handle and result descriptor without allocating, copying, or inferring storage windows. Descriptor bounds and alias legality are verified before emission.

Mat, Left, Right and Acc descriptors render as loc=mat|left|right|acc with their fractal blayout/slayout/fractal fields, and must be whole fractal boxes: a partial box is rejected with the same extent-to-allocate diagnostic as the legacy path. buffer.matmul and buffer.matmul_acc render as pto.tmatmul and pto.tmatmul.acc; the accumulating form passes its destination as the first input, as PTOAS requires. buffer.extract renders as pto.textract with index offsets, and a Mat to Left/Right buffer.copy as pto.tmov.

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.Ascend910,
)

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`, then `pto.tinsert` for supported ND Vec copies / FIXPIPE writebacks; otherwise `pto.subview` + `pto.tmov`
`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. tile.slice and the fallback tile.assemble path use pto.subview, a pure 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. ND (row_major, none_box) Vec-to-Vec copies of FP32, FP16, and BF16 use pto.tinsert with the original destination and explicit offsets. This preserves the destination row pitch and copies only the source's valid region, leaving the rest of the target intact. A subview + tmov is insufficient for column insertion on A2/A3: PTOAS resolves the view to its parent's physical dimensions, then rejects the unequal source/destination shapes (#2717). Other copies retain the pto.subview + pto.tmov path. A2/A3 PTOAS rejects INT16/INT32 TINSERT, and the pinned ISA rejects INT8/UINT8 TINSERT with dynamic valid columns; existing whole-tile integer TMOV copies are preserved. On A2/A3, the insertion column offset must be 32-byte aligned (8 FP32 or 16 FP16/BF16 elements). Codegen diagnoses a constant offset violating this restriction; callers must preserve it for dynamic offsets too.

FIXPIPE epilogue on the Acc writeback. A converting Acc → Mat assemble and every Acc → GM store are the cube's fix-pipe draining L0C, which can multiply by an FP32 scale and apply ReLU on the way out; both ops carry that as pre_quant / pre_relu. The order is ReLU, then the multiply, then the destination clamp — the activation is a pre-quant stage (pto-isa's ReluPreMode) that sees the raw accumulator, so the pair computes maximum(tile, 0) * scale, not maximum(tile * scale, 0). Those agree for every scale > 0, which is why only a negative scale separates them; acc_to_gm_negative_scale in tests/st/runtime/ops/test_fixpipe_epilogue.py measures it on a2a3.

The scale also selects the lowering — an INT32 accumulator reaches an FP16 Mat tile only because it carries one (DEQF16) — so one predicate, ir::CubeMatWritebackUsesFixpipe, decides it for the emitter and FixpipeEpilogueValid alike; a same-dtype Acc-to-Mat assemble converts nothing, stays a plain pto.tmov, and carries no epilogue. pto.tinsert and pto.tstore spell the packed scale operand differently, neither form guessable from the other — both pinned with their provenance at the emission sites. See codegen::EncodeFixpipePreQuant for the register layout, 99-verifier.md for the per-backend dtype tables.

PTOAS 0.65 resolves the scaled pto.tinsert overload ambiguity tracked by PTOAS#1570 by emitting a uint64_t scale and uint16_t row/column operands. Both A2/A3 and A5 therefore allow the validated INT32 Acc → FP16 Mat (DEQF16) form. Other scaled Acc-to-Mat dtype pairs remain rejected by FixpipeEpilogueValid; the wider per-backend tables continue to govern Acc-to-GM stores independently.

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 or 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 of a compact tile keeps the producer's rows too on a backend whose lanes locate their own band in a GM slot (A2/A3). No lane reads a slot row at or past validRow as data — a lane's band is clamp(V - lane * S, 0, S) rows starting at lane * S, with S the balanced ceil(V / 2) for a static extent and the box half for a deferred one — so the box rows were never needed, and widening them skews the compact L0C read exactly as above (1792 of 2048 valid elements wrong for a 64-row box valid to 16 across pl.split(UP_DOWN)). A non-compact tile reads L0C at Rows whatever its validRow, so its split transport keeps the full box. Where the push splits the rows between the lanes in hardware (BackendHandler::SplitsCubeToVectorTransportInHardware, A5's dual-mode TMOV), the transported rows are the lane partition, so a split push keeps the box and a compact tile whose two pitches differ is rejected there. A runtime row extent is transported the same way: both the pl.split and the pl.split_aiv forms then pop the full box (see the popped-tile rule above); lane 1 loads its whole half, and the rows past the valid extent are stale but never treated as data. - 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 Ascend910 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 {
  %slot = arith.remsi %i, %c2_index : index
  %t = pto.multi_tile_get %l0c_mb[%slot]
     : !pto.multi_tile_buf<..., count=2> -> !pto.tile_buf<loc=vec, ...>
  ...
}

Signed integer and INDEX slots use corrected floor remainders, preserving Python's % semantics. The correction is skipped — as it is above — when the dividend is provably non-negative and the divisor a positive constant, because arith.remsi already agrees with floor-modulo there.

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. The corrected remainder feeds multi_tile_get directly; it is never scaled by the slot's byte size. ptoas recognizes the index (i % 2) to decide which accesses can share a slot, and that 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. Two forms of it have gone wrong:

Form	ptoas behaviour	Upstream
Two slots filled and read in the same iteration	≤ 0.55 guards only the first `multi_tile_get`; the second load races the next iteration's write (measured wrong on device with 0.54). 0.56+ guards the body with one static event for the whole region — correct, but none of the per-slot overlap the region form exists for	PTOAS#1118, fixed in 0.56
Prefetch: slot 0 filled before the loop, then each iteration fills slot `(i+1) % 2` while reading slot `i % 2`	≤ 0.62 primes and drains both slots' events as for a one-slot rotation, off by one here: wrong data for an even trip count, a device hang for an odd one. 0.63 primes only the slot the loop writes first; 0.64 and 0.65 prime both again	PTOAS#1519, fixed in 0.63, regressed in 0.64

CoLiveSlotCollector counts slot selections per loop body, so it cannot tell the two forms apart, and the second bug is back after its fix. Measured on device, the prefetch form runs correctly under 0.63, but under 0.64 and 0.65 it returns a later block's data for an even trip count and hangs for an odd one; the same-iteration form runs correctly under 0.65. Codegen therefore refuses both and points at the PyPTO planner, whose baked-address alloc_tile path runs the same-iteration form correctly on device. Straight-line code is unaffected — with no loop there is no cross-iteration reuse to guard. Lifting the restriction is one condition in PlanMultiBufferRegions; it needs a pinned ptoas that runs the prefetch form correctly on device.

Under PYPTO no region is emitted at all: a region at --pto-level=level3 needs an explicit base addr, which codegen does not emit yet. ptoas is not the limit — given a constant addr it has derived the same per-slot sync at level3 as at level2 since 0.55 (PTOAS#1106, closed).

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`

Scalar Boolean Casts¶

pl.cast(flag, dtype) maps scalar False to numeric zero and True to numeric one, for both literals and runtime values. BOOL uses MLIR i1, but must not be interpreted as a signed one-bit integer during conversion.

Destination	Lowering
Signed integer	`arith.extui` from `i1` to the destination width
Unsigned integer	`arith.extui` to signless `iN`, then bridge to `uiN`
INDEX	`arith.index_castui`
Floating point	Zero-extend to `i32`, then `arith.sitofp`
BOOL	Identity

The floating-point path uses an integer carrier because direct (float)bool casts emitted by PTOAS are rejected by A3's device compiler. Ordinary signed integer widening still uses sign extension. These rules apply to scalar casts; tensor and tile conversion operators have their own lowering.

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)¶

The PTO backend wraps each ptoas-generated InCore function for runtime dispatch. See PTO Kernel Wrapper Generation for the output layout, runtime configuration, argument unpacking, and SPMD identity parameters.