Skip to content

BlockNzTensorViews Pass

Overview

BlockNzTensorViews turns a logical pl.NZ tensor into the blocked form pto-isa's Layout::NZ GlobalTensor requires, and retargets the tile.load that reads it or the supported tile.store that writes it.

A pl.Tensor[[R, C], dtype, pl.NZ] annotation — or [[B, R, C], ...] — is a declaration of the GM byte order: a reader expects PTO-native NZ fractal bytes, and a supported accumulator store produces them. Adding the annotation does not convert existing bytes. The DSL keeps the logical shape and logical slicing; this pass supplies the physical description the backend needs.

The payoff is that a matmul B-operand TLOAD becomes NZ→NZ instead of ND→NZ, which removes the online fractal conversion from every weight load.

The blocked form

With c0 = the number of elements in a 32-byte C0 line (256 / dtype bits; 32 for INT8) and a 16-row fractal, pto-isa describes an NZ buffer as (pto/common/pto_tile.hpp, TileShape2D / BaseShape2D specialisations for Layout::NZ):

shape   = [B, C/c0, R/16, 16, c0]
strides = [C*R, R*c0, 16*c0, c0, 1]

Reading the shape from the inside out: c0 contiguous elements form one 32-byte C0 line, 16 rows form one 16 x c0 fractal (512 bytes), R/16 fractals walk down the row axis, and C/c0 steps between column blocks. That is "column blocks outside, row fractals inside" — the same byte order the tile side expresses as blayout=col_major, slayout=row_major, fractal=512.

The rank is fixed at 5, not logical rank + 2. The leading batch slot B exists whether or not the logical tensor has a leading axis, so a logical rank-2 [N, K] weight blocks to [1, K/c0, N/16, 16, c0] with the batch materialised as 1. PTOAS checks this arity directly — a rank-4 view is refused with 'pto.make_tensor_view' op user-specified layout=nz requires a rank-5 view, no matter how self-consistent the surrounding IR is.

Why NZ needs no stride rule of its own

Row-major strides over the blocked shape are pto-isa's NZ strides:

slot row-major derivation BaseShape2D<T, R, C, NZ>
c0 1 1
16 c0 C0Size
R/16 16*c0 FRACTAL_NZ_ROW*C0Size
C/c0 (R/16)*16*c0 = R*c0 rows*C0Size
leading (C/c0)*R*c0 = C*R cols*rows

So once the shape is blocked, NZ is an ordinary member of the row-major family and BuildLogicalStridesFromLayout handles it through the same BuildRowMajorStrides path as ND. MaterializeTensorStrides (pass 33) fills the stride later; this pass only rewrites the shape.

This amends RFC #1300's claim that NZ has "no logical-stride representation" — true of a logical 2-D shape, false of the blocked rank-5 one.

Position in the pipeline

... -> LowerCompositeOps -> FlattenTileNdTo2D -> BlockNzTensorViews -> LegalizeTileCast -> ...

Three constraints fix this slot:

  • After ConvertTensorToTileOps / LowerCompositeOps — the tile.load and tile.store ops phase 2 rewrites must already exist.
  • After FlattenTileNdTo2D — the destination tile must already be the logical 2-D operand. Blocking a still-ND-rank tile would leave a tile.load whose type annotation and argument ranks cannot both be printed, breaking the printer round-trip.
  • Before MaterializeTensorStrides — which asserts every NZ view is blocked and then fills its row-major stride.

FlattenTileNdTo2D skips its ND2NZ source-window collapse for an NZ source (that collapse exists because ND→NZ needs a 2-D GlobalTensor; NZ→NZ does not), so the logical window is still intact when this pass runs.

Behavior

Phase 1 — block every NZ TensorType shape.

# before
w: pl.Tensor[[32, 2048, 4096], pl.INT8, pl.NZ]

# after  (c0 = 32:  4096/32 = 128,  2048/16 = 128)
w: pl.Tensor[[32, 128, 128, 16, 32], pl.INT8, pl.NZ]

A logical rank-2 weight lands on the same rank, with the batch synthesised:

# before
w: pl.Tensor[[256, 512], pl.INT8, pl.NZ]

# after  (c0 = 32:  512/32 = 16,  256/16 = 16;  batch materialised as 1)
w: pl.Tensor[[1, 16, 16, 16, 32], pl.INT8, pl.NZ]

Phase 2 — retarget supported loads and stores.

# before  (slicing w[1:2, 256:512, 512:1024] out of an [E, N, K] weight)
wt: pl.Tile[[256, 512], pl.INT8, pl.Mem.Mat] =
    pl.tile.load(w, [1, 256, 512], [1, 256, 512], target_memory=pl.Mem.Mat)

# after   (offsets -> [.., k0/c0, n0/16, 0, 0];  sizes -> blocked)
wt: pl.Tile[[256, 512], pl.INT8, pl.Mem.Mat] =
    pl.tile.load(w, [1, 16, 16, 0, 0], [1, 16, 16, 16, 32], target_memory=pl.Mem.Mat)

Accumulator stores

On A2/A3, pl.store(acc, offsets, out) can write a fully valid, static 2-D FP32 Acc tile directly to FP32 pl.NZ GM. It preserves the logical tile and mathematical values; FP32 NZ uses c0 = 8, while the Acc tile keeps its own on-chip layout. No extra GM layout-conversion kernel is inserted.

For a 64 x 16 tile and a logical [64, 16] destination, the store's offsets become [0, 0, 0, 0, 0], and its fourth IR argument becomes the explicit partition shape [1, 2, 4, 16, 8]. The existing rank-N store emitter consumes that shape; it must not derive a rank-2 GM partition from the Acc tile. The descriptor represents:

using NzShape = pto::Shape<1, 2, 4, 16, 8>;
using NzStride = pto::Stride<1024, 512, 128, 8, 1>;
pto::GlobalTensor<float, NzShape, NzStride, pto::Layout::NZ> dst(
    output_address, NzShape(), NzStride());
pto::TSTORE(dst, accumulator);

The store call's result remains an alias of the destination tensor. Rebuilding the call preserves its kwargs and attrs, including st_phase=Final; changing the GM coordinates must not change the tile type or unit-flag behavior.

The supported write scope is narrower than the read scope:

Condition Store behavior
Backend and dtypes A2/A3, FP32 Acc to FP32 GM only
Tile shape and valid region Static, positive, 2-D, fully valid; rows and columns are multiples of 16
Emitted-store extent limits Rows at most 65535, columns at most 4095; alignment makes the effective maxima 65520 and 4080
Emitted-source allocation Each physical Acc tile must fit the backend's L0C capacity (currently 128 KiB on A2/A3)
Destination row window Entire parent row axis, at row offset 0
Leading axes Exactly one panel per store; a nonzero batch offset is allowed
Column offset Non-negative and provably a multiple of c0 = 8
Unknown alignment or sign Rejected; existing offset facts are reused rather than assuming validity
Atomic, pre_quant, pre_relu, or dtype conversion Rejected
A5, Vec/Mat sources, dynamic or partially valid tiles Rejected

The initial logical source may exceed L0C or the store instruction's column field. BlockNzTensorViews validates its static, fully valid shape, attributes, and complete-parent-row geometry without requiring that one logical value already be a physical allocation. It still runs before AutoTileMatmulL0.

For a supported Mat-resident matmul consumed by one NZ store, AutoTileMatmulL0 understands the blocked rank-5 partition and splits the output along N only, keeping m == M and permitting the existing K reduction tiling. A column chunk beginning at ni adds ni / 8 to the blocked column offset; its partition is [1, n / 8, M / 16, 16, 8]. The logical result is computed and drained as smaller Acc tiles, without materializing the oversized full Acc. For example, a 64 x 1024 FP32 output is 256 KiB in total, but its individual output chunks can fit the 128 KiB L0C. The same applies when the logical N exceeds the instruction's 4095 limit: each emitted store must satisfy that limit after tiling.

At final emission, every actual source is checked against both the ISA field limits and StaticPhysicalAllocationBytes / the backend handler's GetCapacityBytes(MemorySpace::Acc). These checks apply to each emitted Acc tile, not the GM tensor's total size. If full M leaves no legal N/K tile, AutoTileMatmulL0 reports a dedicated compile-time diagnostic. A producer/use pattern that cannot be retiled remains subject to allocation checks and the final NZ-store guard, so it cannot silently bypass the physical limits. Explicit column tiling remains available; every chunk must cover the destination's complete row axis.

The full-row restriction protects against a limitation of A2/A3 pto-isa's TStoreAccNz2nz. It computes the destination column-block pitch from the partition height rather than the parent tensor's stride:

uint32_t dstStride = gShape2 * gShape3 * c0Size >> SHIFT_BLOCK_BYTE;
// gStride1 is not used by this path.

For example, a 64 x 16 write into a [128, 32] parent needs column blocks 1024 FP32 elements apart, but this implementation uses 512. Aligned row offsets alone cannot make that write correct. Until the underlying stride handling supports row subwindows, they are rejected rather than silently miswritten. There is no single-column-block exception: even with FP32 C0=8, this Acc store requires its valid column count to be a multiple of 16.

Stack independent panels on a leading batch axis ([PANELS, R, C]), not the row axis ([PANELS * R, C]), and store one panel at a time. The blocked partition's batch extent is 1; its batch offset locates the panel. A2/A3's Acc store folds a larger batch extent into the column count rather than traversing independent panels using the parent batch stride.

This adds a supported hardware capability, not a performance guarantee. The combined CSA candidate reported in issue #3036 regressed in its measured workload; that experiment also changed the dataflow and the AIV reader, so its result does not isolate the cost of NZ store itself.

Symbolic trailing offsets

A trailing offset does not have to be a constant, but its alignment must be proven — never assumed. An offset arrives here as the SSA name it was bound to, not as the arithmetic that produced it, so IsProvableMultipleOf (tensor_view_semantics.h) walks two kinds of binding the pass collects from the enclosing function in one read-only sweep:

Binding Multiple of the axis factor Non-negative
AssignStmt (n0 = nb * 256) one factor of the product is a multiple both factors are non-negative
ForStmt (for k0 in pl.pipeline(512, 4096, 512)) start and step are both multiples start and step are both non-negative
ForStmt with a symbolic start (for ob in pl.range(core, TILES, CORES)) both recurse into start and step both recurse into start and step
Min / Max (min(800 - o0, 256)) both operands are multiples Min: both operands; Max: either one
tile.get_block_idx / tile.get_block_num — a lane number is never negative
FloorMod / FloorDiv by a positive constant ((blk % 2) * 512) the other factor carries it both recurse: the dividend must be non-negative too
ConstInt the value is a multiple the value is >= 0

Sums and products compose from those; a difference proves divisibility but never its sign, so it is refused. The current proof conservatively requires a non-negative dividend for both operations. Signed integer and INDEX FloorMod codegen corrects arith.remsi to floor semantics, so a positive divisor yields a non-negative remainder even for a negative dividend; this pass does not use that stronger fact. FloorDiv lowers to arith.floordivsi, which rounds the quotient toward negative infinity. Both columns must hold — see Why the sign is proven too. The grouped-matmul weight path that motivated the feature therefore compiles:

for nb in pl.spmd(N // N_TILE):
    n0 = nb * N_TILE                     # -> row fractal offset  n0 // 16
    for k0 in pl.pipeline(K_TILE, K, K_TILE, stage=2):
        wt = w[n0 : n0 + N_TILE, k0 : k0 + K_TILE]   # -> c0 block offset  k0 // c0

Anything the walk cannot prove is rejected with a diagnostic naming the provable forms. That refusal is the design: an NZ tensor addressed from a guessed coordinate reads the wrong fractal, and nothing downstream would notice.

Why the offset is divided, not re-associated

Only the result of the offset is divided (FloorDiv(offset, divisor)). The tempting rewrite — turning n0 = nb * 256 into nb * 16 and saving the runtime division — is unsound, because IR arithmetic wraps at its declared width while the rewritten form does not:

x : INT32 = 1 << 24
(x * 256) / 16   ==  0            # x * 256 wraps to 0 in i32
x * (256 / 16)   ==  268435456    # re-associated: no wrap, wrong fractal

Widening is not the culprit, and rebuilding at the original width does not help: for a * b = q * 2^W + r, the original yields r / d while any re-association yields r / d + q * 2^(W - log2 d). Dividing the offset as a whole keeps its own dtype and overflow behaviour intact.

The proof itself survives wraparound because every divisor here is a power of two and so divides 2^W: reducing a multiple of d modulo 2^W leaves a multiple of d. An INTERNAL_CHECK pins that precondition.

One further limit is deliberate: an IterArg is never resolved, because its value changes every iteration, so neither its initial value nor any binding recorded for it describes the value a given use sees.

Why the sign is proven too

Divisibility alone does not make a coordinate safe. n0 = -16 is a clean multiple of 16, and FloorDiv(n0, 16) is -1; codegen then clamps a negative pto.partition_view offset to 0 instead of failing, so the load reads fractal 0 and returns silently wrong data — the same shape #2543 fixed for row indices.

The constant path already refuses a negative literal, so proving only divisibility for a symbolic offset would mean the identical value is caught written inline and waved through once bound to a name. IsProvableNonNegative closes that gap.

The destination TileType is preserved verbatim: the GM partition becomes rank-5, the tile stays the logical 2-D operand. The load is therefore rebuilt with the explicit-type Call constructor rather than OpRegistry::Create, which would re-deduce a rank-5 tile from the blocked shapes argument.

After this pass no logical-shaped NZ TensorType survives, so nothing downstream needs to know NZ is special — including codegen, which derives the rank of pto.make_tensor_view, its !pto.tensor_view<> type and the pto.partition_view independently from TensorType::shape_ and must see them agree.

Generated code

For the logical rank-2 w: pl.Tensor[[256, 512], pl.INT8, pl.NZ] above — note the leading %c1 / %c131072 / 1x carrying the synthesised batch through all three sites:

%w_view = pto.make_tensor_view %arg1,
    shape = [%c1, %c16, %c16, %c16, %c32],
    strides = [%c131072, %c8192, %c512, %c32, %c1]
    {layout = #pto.layout<nz>} : !pto.tensor_view<?x?x?x?x?xi8>
%w_pview = pto.partition_view %w_view,
    offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c16, %c16, %c16, %c32]
    : !pto.tensor_view<?x?x?x?x?xi8> -> !pto.partition_tensor_view<1x16x16x16x32xi8>
pto.tload ins(%w_pview : !pto.partition_tensor_view<1x16x16x16x32xi8>)
          outs(%wt : !pto.tile_buf<loc=mat, dtype=i8, rows=256, cols=512,
                                   blayout=col_major, slayout=row_major, fractal=512, ...>)

which PTOAS 0.61 and later turns into a real NZ GlobalTensor:

GlobalTensor<int8_t, pto::Shape<1, 16, 16, 16, 32>,
             pto::Stride<131072, 8192, 512, 32, 1>, pto::Layout::NZ> ...;

On 0.60 and earlier the same IR does not assemble — see Assembler version below.

Scope and rejections

NZ loads and accumulator stores have deliberately bounded scopes. Everything outside them is rejected with a diagnostic naming the fix — an NZ tensor must never be silently mis-addressed.

Condition Outcome
shape[-2] % 16 != 0 rejected — a partial fractal has no representation
shape[-1] % c0 != 0 rejected — a partial C0 line has no representation
dynamic shape[-2] / shape[-1] rejected — divisibility cannot be proven
slice offset not fractal-aligned rejected — no blocked representation
symbolic trailing slice offset, alignment and sign both provable mapped — see Symbolic trailing offsets
symbolic trailing slice offset, alignment not provable rejected — never divided on the assumption that it is aligned
symbolic trailing slice offset, sign not provable rejected — a negative offset is clamped, not caught, at the partition view
logical rank 2 blocked to [1, C/c0, R/16, 16, c0] — batch materialised
logical rank 3 blocked to [B, C/c0, R/16, 16, c0] — leading axis is the batch
logical rank > 3 blocked with every leading axis folded into the batch (see below)
logical rank < 2 rejected — the trailing pair is the fractal plane
dynamic leading extent, rank > 3 rejected — the fold needs static extents to multiply
tile.load with target_memory != Mat (or absent) rejected — NZ→NZ is the cube operand path
tensor.slice narrowing the leading axes only blocked like the load that follows it (see below)
tensor.slice windowing the trailing [R, C] pair rejected — the window is not contiguous
tensor.reshape flattening the whole tensor to [N] kept as written — see Flattening an NZ tensor
tensor.reshape to any other shape rejected — it reinterprets coordinates the blocked form does not carry
tile.store to an NZ destination supported under Accumulator stores
consumer other than a supported load, store, slice, or whole-tensor flatten rejected — unsupported NZ coordinates must not pass through unchanged
explicit tensor stride or partial tensor valid_shape rejected
tile.load dynamic valid_shape[-2] provably a multiple of 16 (a ragged last tile) blocked — the row-fractal count becomes FloorDiv(rows, 16)
tile.load dynamic valid_shape[-2] that may end inside a fractal rejected — a partial fractal has no blocked form
distributed tensor rejected — remote_load has no NZ blocking
tensor.view / tensor.reinterpret_view of NZ rejected at op construction
GM row gap above 65535 blocks, on a multi-column-block load rejected — temporary, see GM row gap

GM row gap: a temporary guard

This row is not a design limit. It guards an upstream defect (hw-native-sys/pto-isa#317) and is deleted when that lands — along with CheckNzGmGapFitsBurstStride in the pass and its three unit tests.

TLoadGm2L1Nz2nz computes the GM gap between consecutive column blocks as a uint32_t and hands it to TLoadInstrGm2L1's uint16_t gmGap with no range test. Above 65535 32-byte blocks it wraps, and the load reads the wrong fractals — silently: PTOAS assembles the view, the CCE compiler accepts the kernel, and the matmul returns wrong numbers with no error at run time. This pass is the last layer that still knows which pl.NZ annotation is responsible, so the diagnostic belongs here.

For a blocked view the gap reduces to the row extent the load leaves behind:

gmGap = (gStride1 - gShape2*gShape3*gShape4) * sizeof(T) / 32
      = (R/16 - TR/16) * 16 * c0 * sizeof(T) / 32
      = R - TR                                       (32-byte blocks)

because c0 * sizeof(T) == 32 holds for every NZ view by construction — so the bound is dtype-independent, and a wider row tile is what lowers the gap, not a narrower one:

R row tile gap Result
65536 16 65520 accepted
65552 16 65536 rejected
65552 32 65520 accepted — same tensor, wider tile

The other workaround is to annotate a layer-stacked weight rank-3 with the stacked axis as the batch ([LAYERS, K, N]); its extent then rides gStride0 and a real for loop rather than the burst gap.

The window measured is the one codegen turns into the pto.partition_view — valid_shape when the load carries one, else shapes (src/backend/common/pto_ops_memory.cpp). That view is pto-isa's gShape, so a narrowed valid_shape loads fewer row fractals and leaves a larger gap than shapes alone would suggest. A [65552, 64] INT8 weight read with shapes=[32, 64] and valid_shape=[16, 64] emits partition_tensor_view<1x2x1x16x32> and so a gap of 65536, not the 65520 shapes implies. When the loaded row extent is dynamic (a ragged last tile), the check takes the worst case, nothing loaded, so a load that passes fits for every run-time width.

A single-column-block load is exempt. TLoadGm2L1Nz2nz passes the load's column-block extent as nBurst, and the DMA applies gmGap only when stepping from one burst to the next, so at one burst the truncated field is never read. Confirmed on device in pto-isa's own tload_gm2mat ST suite: an NZ int16 load with gShape1 = 1 and a 65536-block gap returns bit-exact data, while the same load at gShape1 = 2 corrupts 1837/4096 elements.

Leading axes fold into the one batch slot

pto-isa's NZ GlobalTensor has exactly one batch slot, so a logical [G, E, N, K] weight folds its two leading axes into it: the blocked batch is G*E, with stride C*R. The fold is exact because those axes are dense and row-major — it removes only the strides it multiplies back in — and an offset folds the same way, [g, e, 0, 0] addressing batch g*E + e (FoldNzLeadingOffsets).

That is not the re-association the trailing offsets refuse: nothing is divided and nothing is assumed about alignment, so the fold is exact for every coordinate, not only aligned ones. It is the same arithmetic the ND path performs at address computation, written once into the coordinate instead.

The extents being folded must be static — a dynamic one cannot be multiplied into the batch — which the diagnostic names.

A dynamic batch is supported for logical rank-3 [B, R, C] parameters. Their DeviceTensor.shape or StackedDeviceTensor.full_shape must also have rank 3: the entry reads B from the first logical extent. Rank-2 or rank-4+ arguments are rejected before dispatch, even when their blocked shape would otherwise match.

A rank-4 parameter is what a multi-card entry declares ([RANKS, E, R, C], sliced per rank before dispatch), so the fold is what lets a distributed program carry NZ weights at all.

Slicing an NZ tensor

A layer- or rank-stacked weight reaches its kernel through tensor.slice, so the slice blocks like the tile.load that follows it: the shapes and offsets become rank-5, and a rank-reducing scalar index (w[r]) needs no drop_dims afterwards because the fold already collapsed every leading axis.

Only the leading axes may be narrowed. A window inside the trailing [R, C] pair is rejected: in NZ order one layer's rows sit inside every fractal column block, so [layer*R, 0] selects C/c0 disjoint runs, and the blocked view has no stride of its own to describe them — MaterializeTensorStrides derives a row-major one from the blocked shape. Annotate the stacked axis as a leading axis ([LAYERS, R, C]) instead of stacking rows.

Flattening an NZ tensor

A rank-1 view of every element is layout-invariant: the blocked form permutes the index space, not the memory, so both spellings walk the same contiguous GM range in the same order. Such a tensor.reshape is therefore kept exactly as written — no coordinate rewrite, and the result is ND, which is what prefetch.async_prefetch wants of its source. Without it, annotating a weight pl.NZ would silently cost it its SDMA L2 warm.

Any other target shape does reinterpret coordinates — [256, 512] -> [128, 1024] pairs rows in logical row-major order, and in the blocked form those elements are scattered across fractal blocks — so it is rejected rather than addressed as if it were ND.

An NZ argument also arrives at the orchestration entry in its logical shape: the caller allocates the weight that way, and only the compiled parameter is blocked. The entry restates it in blocked terms once (a metadata-only reshape, same elements in the same order), so every Tensor::view derived from it clamps against the rank it is written in.

Sub-byte dtypes (INT4, UINT4, HF4, and BOOL) remain outside this lowering's supported scope. Packed FP4 x2 types have byte-sized carrier elements; their type definition does not enable NZ execution. The existing packed layout restrictions are documented in FP4.

The alignment diagnostics are user-facing (CHECK_SPAN → ValueError) and live in BlockNzShape. Downstream, an unblocked NZ view is a pass-ordering invariant instead (INTERNAL_CHECK_SPAN), enforced by CheckNzViewIsBlocked in MaterializeTensorStrides and by the TensorViewCanonical verifier.

Idempotence

Blocking is not idempotent — blocking a blocked shape would be wrong — and the structural IsBlockedNzShape test cannot distinguish a blocked shape from a logical one that merely ends in [16, c0]. The pass therefore stamps nz_tensor_views_blocked on each function it rewrites and returns early when it sees that attribute.

Assembler version

Whether this pass's output assembles at all depends on the PTOAS release.

PTOAS Behavior
≤ 0.60 Infers the layout structurally. Blocked NZ and ND are structurally identical (both row-major), so it infers nd, overrides the explicit nz annotation, and fails with layout mismatch: user-specified layout=nz but inferred=nd. No NZ view assembles, at any rank.
≥ 0.61 Treats an explicit ND / DN / NZ annotation as authoritative and validates it, so the descriptor above assembles. It also enforces NZ's arity directly: a view of any rank but 5 is refused with 'pto.make_tensor_view' op user-specified layout=nz requires a rank-5 view.

toolchain/versions.env pins a release at or above v0.61, so pl.NZ works end to end on the pinned toolchain. tests/st/runtime/ops/test_matmul_nz.py is what holds that: an ND activation against an NZ weight, with the host packer that produces the fractal bytes, plus sliced cases that pin both offset axes — the row fractal (n0 // 16) and the C0 column block (k0 // c0). A whole-tensor load leaves every offset zero, so the baseline case alone would exercise neither.

The 0.60 failure is safe rather than silent in both directions: it stops at the layout mismatch above, and even past that, pto-isa's ND→NZ TLOAD path requires staticShape[0..2] == 1, which the blocked dims violate — so the generated C++ fails a static_assert instead of computing wrong results.