[Feature] Support GM-side NZ tensor layout so implicit matmul weight loads become NZ-to-NZ

[lyfne123]

Summary

Allow a GM tensor to declare that its bytes are already stored in PTO-native
NZ fractal order, keep its logical shape and slicing at the DSL level, and
have the compiler synthesize the blocked NZ GlobalTensor descriptor that
implicit matmul weight loads consume — so a TLOAD of a matmul B operand becomes
NZ→NZ instead of ND→NZ.

Today TensorLayout::NZ is rejected on any TensorType ("NZ is tile-only"), so
there is no way to express this.

Motivation / Use Case

The DeepSeek-V4 Flash MTP decode routed W1/W3 INT8 matmuls take logical ND
weights [N_LOCAL, MOE_INTER, D]. Every generated cube load therefore performs
ND→NZ conversion while streaming each weight tile: the GM descriptor is
Shape<1,1,1,256,512> / Stride<...,4096,1>, so each fractal row segment is only
32 contiguous bytes at a 4096-byte stride.

Two generated-.cpp experiments (pypto-lib#1039) were validated on A2/A3 with
exact INT32 comparison:

1. Replace the weight GM descriptors with PTO-native NZ descriptors and upload
   correspondingly packed weights.
2. Additionally read those buffers through the L2-cache-disable address alias
   (tracked separately — see Related).

Measured task_slot_0 means over five real decode expert distributions, against
CANN's aclnnGroupedMatmulWeightNz:

Decode distribution	CANN WeightNZ	PyPTO ND	PyPTO direct NZ
learned P10	73.333 us	84.134 us	77.134 us
learned P50	130.057 us	152.573 us	138.688 us
learned P90	178.753 us	199.144 us	182.748 us
hash P50	271.523 us	306.437 us	278.829 us
MTP P50	178.455 us	200.301 us	182.630 us

Aggregate gap vs CANN: ND +13.28%, direct NZ +3.35%. Direct NZ alone
recovers roughly three quarters of the gap.

The experiment is unusable as a supported path because it requires editing
generated C++ after compilation. The two patches change only the GM
descriptor — tile buffers, TLOAD, TEXTRACT, matmul, synchronization and the
L0 loops are byte-identical to the ND build. So the missing capability is
descriptor synthesis in the frontend, not codegen or ISA work.

Proposed API / Behavior

Expected use case

The logical shape must stay [E, N, K] so the frontend, golden reference and
slicing are unchanged; only the physical storage order of the bytes differs.

import pypto.language as pl

E, N, K = 32, 2048, 4096          # local experts, moe_intermediate, hidden
N_TILE, K_TILE = 256, 512


@pl.jit
def w13_gmm(
    x: pl.Tensor[[512, K], pl.INT8],
    # Logical shape stays [E, N, K]. The layout annotation asserts only that the
    # bytes in GM are already in PTO-native NZ fractal order.
    w1: pl.Tensor[[E, N, K], pl.INT8, pl.NZ_GM],
    out: pl.Out[pl.Tensor[[512, N], pl.INT32]],
):
    for e in pl.parallel(E):
        for nb in pl.spmd(N // N_TILE, name_hint="w1_mm"):
            n0 = nb * N_TILE
            xt = pl.slice(x, [64, K_TILE], [0, 0])
            # Sliced with *logical* (n, k) coordinates, exactly as the ND form is.
            w1_k0 = w1[e : e + 1, n0 : n0 + N_TILE, 0 : K_TILE]
            acc = pl.matmul(xt, w1_k0, b_trans=True, out_dtype=pl.INT32)
            for k0 in pl.pipeline(K_TILE, K, K_TILE, stage=2):
                x_k = pl.slice(x, [64, K_TILE], [0, k0])
                w1_k = w1[e : e + 1, n0 : n0 + N_TILE, k0 : k0 + K_TILE]
                acc = pl.matmul_acc(acc, x_k, w1_k, b_trans=True)
            out[0:64, n0 : n0 + N_TILE] = pl.reshape(acc, [64, N_TILE])
    return out

Expected lowering

With c0 = 32 / sizeof(dtype) (32 for INT8) and a 16-row fractal, the compiler
should synthesize the blocked rank-(r+2) view:

make_tensor_view %w1,
    shape   = [E, K/c0, N/16, 16, c0],
    strides = [K/c0 * N * c0, N * c0, 16 * c0, c0, 1]
    {layout = #pto.layout<nz>}

partition_view  offsets = [e, k0/c0, n0/16, 0, 0],
                sizes   = [1, K_TILE/c0, N_TILE/16, 16, c0]

which lowers to exactly the descriptor pto-isa documents as canonical NZ
(tests/npu/a5/src/st/testcase/tmov_nd2nz/, tinsert/) and which the validated
experiment hand-wrote:

GlobalTensor<int8_t, pto::Shape<1, 16, 16, 16, 32>,
             pto::Stride<8388608, 65536, 512, 32, 1>,
             pto::Layout::NZ>

Everything downstream is unchanged: the destination tile is already
blayout=col_major, slayout=row_major, fractal=512, and pto-isa already supports
NZ→NZ TLOAD.

Scope for a first milestone

Read path only, aligned only, matmul B operand only:

• N % 16 == 0, K % c0 == 0, and slice offsets n0 % 16 == 0, k0 % c0 == 0.
  Anything else must be rejected with a clear diagnostic, never silently
  mis-addressed.
• No NZ stores, no arbitrary NZ slicing, no padding, no dynamic blocked shapes.

Suggested design: a new layout value rather than relaxing NZ

TensorLayout already has the precedent: MX_A_ZZ / MX_B_NN are GM packed
layouts that are legal on a TensorType, with their own stride rule
(tensor_view_semantics.h:98), their own tile view (fractal = kMXScaleFractal),
their own op guards (tile_ops/memory.cpp:152) and their own predicate
IsMxTensorLayout().

Adding a new value (e.g. TensorLayout::NZ_GM) rather than relaxing NZ keeps
every existing "NZ is tile-only" check honest and makes each site an explicit
opt-in, instead of requiring a re-audit of all eight of them.

Alternatives Considered

• Keep online ND→NZ conversion. Correct, but retains most of the measured gap.
• Require the caller to declare the blocked physical shape ([E, K/c0, N/16, 16, c0]) as an ordinary ND tensor. This compiles today, but it forces every
  caller and the golden reference onto the physical shape, and the logical
  slice w[e, n0:n0+256, k0:k0+512] can no longer be written at all.
• Relax TensorLayout::NZ on TensorType. Rejected in favour of a new enum
  value, per above.
• Keep patching generated C++. Proves the optimization; unsafe under
  recompilation and unusable as a supported path.

Additional Context

Current blocker

pl.NZ on a tensor parses and reaches IR (pinned by
tests/ut/jit/test_jit_compile_extraction.py::TestNzOnTensorIsNotJitSpecific),
but dies at the first layout pass:

ValueError: MaterializeTensorStrides: NZ layout is tile-only and not allowed on
a tensor type. Annotate the tensor as pl.ND or pl.DN, and produce NZ on a Tile
instead.
Check failed: view.layout != TensorLayout::NZ
  at src/ir/transforms/materialize_tensor_strides_pass.cpp:76

Minimal repro — nz_probe.py (not committed; save anywhere and run)

"""Minimal probe: does a GM tensor annotated pl.NZ survive the compile pipeline?"""
import pypto.language as pl


@pl.jit
def nz_mm(
    x: pl.Tensor[[64, 512], pl.INT8],
    w: pl.Tensor[[256, 512], pl.INT8, pl.NZ],
    out: pl.Out[pl.Tensor[[64, 256], pl.INT32]],
):
    for _ in pl.spmd(1, name_hint="nz_mm"):
        xt = pl.slice(x, [64, 512], [0, 0])
        wt = w[0:256, 0:512]
        acc = pl.matmul(xt, wt, b_trans=True, out_dtype=pl.INT32)
        out[0:64, 0:256] = pl.reshape(acc, [64, 256])
    return out


if __name__ == "__main__":
    from pypto.runtime.runner import RunConfig

    try:
        nz_mm.lower(config=RunConfig(platform="a2a3"))
        print("PROBE: compiled OK")
    except Exception as e:
        print("PROBE FAILED:", type(e).__name__)
        print(str(e)[:2000])

Sites that reject NZ on a TensorType

Location	Role
src/ir/transforms/materialize_tensor_strides_pass.cpp:76	CheckNoNzOnTensorType — hit first
include/pypto/ir/transforms/utils/tensor_view_semantics.h:123	BuildLogicalStridesFromLayout CHECK(false) for NZ
include/pypto/ir/transforms/utils/tensor_view_semantics.h:209	CheckCanonicalView rejects NZ
src/ir/verifier/verify_tensor_view_canonical.cpp:88	verifier restates it
src/ir/op/tensor_ops/transform.cpp:268	tensor.reinterpret_view rejects NZ
src/ir/op/tensor_ops/transform.cpp:485,487	tensor.view rejects NZ src and dst
src/ir/transforms/window_externalization/common.cpp:272	MakeWindowTensorView returns nullopt for NZ
src/backend/common/pto_ops_distributed.cpp:641	distributed defer_wait (unrelated to this request)

Two of these need attention beyond a simple opt-in:

• flatten_tile_nd_to_2d collapses the weight to a 2-D ND view
  (rewrite_utils.cpp:171 calls CanonicalizeView(flat_shape, layout); this is
  what produces the [65536, 4096] view in the ND .pto). For an NZ-family
  tensor it must produce the blocked form instead.
• MakeWindowTensorView returning nullopt means an NZ weight loses window
  externalization. Worth evaluating so a resident weight does not silently lose
  scope hoisting.

RFC #1300 premise

#1300 established that NZ has "no logical-stride representation". That holds for
a 2-D logical shape, but not for the blocked rank-(r+2) shape: the strides
above are ordinary row-major over [E, K/c0, N/16, 16, c0]. This request is
effectively an amendment to that clause, which is why an RFC revision may be
worth doing before the code change.

Downstream dependency: PTOAS rejects the descriptor today

Even with the frontend fixed, PTOAS 0.60 currently refuses the descriptor. Taking
the ND .pto from the pypto-lib#1039 evidence pack and rewriting the four weight
views into the blocked NZ form above:

error: layout mismatch: user-specified layout=nz but inferred=nd    (x4)
Error: Pass execution failed.

Changing only layout<nz> to layout<nd> in the same file compiles, and PTOAS
emits:

GlobalTensor<int8_t, pto::Shape<1, 16, 16, 16, 32>,
             pto::Stride<8388608, 65536, 512, 32, 1>,
             pto::Layout::ND>   // <- the validated experiment has NZ here

i.e. the descriptor round-trips correctly and only the layout tag is wrong. PTOAS
has an Infer GlobalTensor layout (ND/DN/NZ) for make_tensor_view pass that
derives the layout structurally; blocked NZ and ND are structurally identical
(both row-major), so it infers nd and overrides the explicit annotation
rather than trusting it. Setting the batch dim to 1 does not help.

A PTOAS issue for this will be filed separately; this pypto feature cannot land
end-to-end until it is fixed.

Environment used for the checks above

Component	Version
pypto-lib	48de053
pypto	135ddcf2
simpler	1f27a157 (detached)
ptoas	0.60
pto-isa	f51c92f6
CANN	9.0.0

Related: hw-native-sys/pypto-lib#1039 (umbrella), #1300 (TensorType layout RFC)