pypto_pro.language.TensorLayout#

产品支持情况#

Ascend 950PR/Ascend 950DT：支持

Atlas A3 训练系列产品/Atlas A3 推理系列产品：不支持

Atlas A2 训练系列产品/Atlas A2 推理系列产品：不支持

功能说明#

数据布局枚举，用于描述GM Tensor的存储形式和Tile的数据排列形式。

GM Tensor支持ND、NZ，默认ND。
Tile支持ND、DN、NZ、ZN、NN、ZZ，默认值由内存空间和芯片架构决定。

取值#

取值	数据排列	适用对象	典型用途
`pl.ND`	非分形行主序，最后一维连续	Tensor / Tile	普通GM Tensor（默认）；UB Tile（默认）；Scaling buffer
`pl.DN`	非分形列主序	Tile	UB上`[ROWS, 1]`列向量（归约结果、histogram索引）
`pl.NZ`	NZ分形排列	Tensor / Tile	GM Tensor；L1 Mat（默认）；A5的L0A（默认）；L0C Acc（默认）
`pl.ZN`	ZN分形排列	Tile	L0B Right（默认）；转置搬入时的L1 Mat
`pl.ZZ`	ZZ分形排列	Tile	A3的L0A（默认）；MX矩阵计算中，A矩阵的E8M0分组缩放因子在L1和ScaleLeft中的布局
`pl.NN`	NN分形排列	Tile	MX矩阵计算中，B矩阵的E8M0分组缩放因子在L1和ScaleRight中的布局

以上短名称分别等价于pypto_pro.language.TensorLayout.ND、DN、NZ、ZN、NN、ZZ。

布局声明#

Tensor布局#

GM Tensor支持ND（行主序，默认）和NZ（分形布局）：

x: pl.Tensor[[64, 128], pl.DT_FP16]                   # 默认 ND
x_nz: pl.Tensor[[64, 128], pl.DT_FP16, pl.NZ]         # NZ 分形布局

pl.NZ只声明GM内存的布局，不会把普通ND buffer自动转换成NZ。调用kernel前，输入buffer必须已经按NZ物理顺序完成packing；NZ输出也必须使用按NZ格式分配的buffer。高维NZ Tensor的最后两轴固定解释为[M, N]，所有前导轴均作为batch轴，不支持在layout中指定任意分形轴。

NZ将逻辑[..., M, N]存储为[..., ceil(N/C0), ceil(M/16), 16, C0]。M/N无需分形对齐，但底层storage必须按align(M, 16) * align(N, C0)的容量分配并使用上述NZ物理排布；仅分配M * N元素的紧凑buffer不受支持。补齐区不属于逻辑Tensor内容，框架不会为传入的buffer自动扩容或完成packing。INT8、FP8E4M3FN、FP8E5M2、FP8E8M0和HF8的C0为32，FP4E2M1和FP4E1M2为64，FP16和BF16为16，FP32和INT32为8。FP4的M/N同样按逻辑元素计数，不使用packed字节数作为Tensor shape。

MX矩阵计算使用的E8M0分组缩放因子在GM中仍声明为普通ND Tensor；物理shape和搬运约束见matmul_mx和load。

[!IMPORTANT]重要普通ND GM Tensor的转置搬运由load/load_tile的order参数决定（order=[1,0]即is_transpose=True），需与L1 Tile布局ZN配合。GM NZ只支持与NZ Tile同布局正序搬运，不支持通过order转置。详见下文转置搬入。

Tile布局#

Tile通过pl.TileType的layout参数指定。不指定时，默认值由内存空间和芯片架构决定：

内存空间	A3默认	A5默认	额外允许
`Vec`（UB）	无默认值	无默认值	`ND`；`DN`（仅特定API要求的列主序场景）；`NZ`
`Mat`（L1）	`NZ`	`NZ`	`ZN`（转置搬入）；`DT_FP8E8M0`还允许`ZZ`、`NN`；`UINT64`/`INT64`还允许`ND`
`Left`（L0A）	`ZZ`	`NZ`	`ZZ`、`NZ`
`Right`（L0B）	`ZN`	`ZN`	—
`Acc`（L0C）	`NZ`	`NZ`	—
`Scaling`	`ND`	`ND`	—
`ScaleLeft`	—	`ZZ`	—
`ScaleRight`	—	`NN`	—

调用示例#

转置搬入#

GM ND Tensor搬入L1 Mat Tile时，两轴顺序一致可省略order；需要交换两轴时，load设置order=[1, 0]，目标Tile布局使用ZN。框架根据order生成对应的TLOAD指令。

以C[M, N] = A[M, K] @ B[K, N]为例：

操作数	Tensor shape	是否转置	`load`的`order`	L1 Mat Tile layout
左矩阵A	`[M, K]`	否	`[0, 1]`（默认）	`NZ`（默认）
左矩阵A	`[K, M]`	是	`[1, 0]`	`ZN`
右矩阵B	`[K, N]`	否	`[0, 1]`（默认）	`NZ`（默认）
右矩阵B	`[N, K]`	是	`[1, 0]`	`ZN`

左矩阵转置搬入#

Tensor shape为[K, M]（与L1 Tile的[M, K]轴序相反），load设置order=[1, 0]，L1 Tile配ZN：

@pl.jit(auto_mutex=True)
def kernel_left_transpose(
    a: pl.Tensor[[K, M], pl.DT_FP16],               # [K, M]，需转置
    b: pl.Tensor[[K, N], pl.DT_FP16],               # [K, N]，不转置
    out: pl.Tensor[[M, N], pl.DT_FP32],
):
    a_l1 = pl.make_tile_group(
        type=pl.TileType(shape=[M, K], dtype=pl.DT_FP16,
                         target_memory=pl.MemorySpace.Mat, layout=pl.ZN),  # ZN
        addrs=0x00000, mutex_ids=[0])
    b_l1 = pl.make_tile_group(
        type=pl.TileType(shape=[K, N], dtype=pl.DT_FP16,
                         target_memory=pl.MemorySpace.Mat),                # NZ（默认）
        addrs=0x10000, mutex_ids=[1])
    ...
    with pl.section_cube():
        cur_a = a_l1.current()
        pl.load(cur_a, a, [0, 0], order=[1, 0])    # 转置搬入
        cur_b = b_l1.current()
        pl.load(cur_b, b, [0, 0])                   # 不转置
        ...

右矩阵转置搬入#

Tensor shape为[N, K]（与L1 Tile的[K, N]轴序相反），load设置order=[1, 0]，L1 Tile配ZN：

@pl.jit(auto_mutex=True)
def kernel_right_transpose(
    a: pl.Tensor[[M, K], pl.DT_FP16],               # [M, K]，不转置
    b: pl.Tensor[[N, K], pl.DT_FP16],               # [N, K]，需转置
    out: pl.Tensor[[M, N], pl.DT_FP32],
):
    a_l1 = pl.make_tile_group(
        type=pl.TileType(shape=[M, K], dtype=pl.DT_FP16,
                         target_memory=pl.MemorySpace.Mat),                # NZ（默认）
        addrs=0x00000, mutex_ids=[0])
    b_l1 = pl.make_tile_group(
        type=pl.TileType(shape=[K, N], dtype=pl.DT_FP16,
                         target_memory=pl.MemorySpace.Mat, layout=pl.ZN),  # ZN
        addrs=0x10000, mutex_ids=[1])
    ...
    with pl.section_cube():
        cur_a = a_l1.current()
        pl.load(cur_a, a, [0, 0])                   # 不转置
        cur_b = b_l1.current()
        pl.load(cur_b, b, [0, 0], order=[1, 0])    # 转置搬入
        ...

UB Tile的ND与DN#

UB Tile大部分情况使用ND（行主序）。DN（列主序）仅在特定API要求时使用，典型场景是归约操作产生[ROWS, 1]列向量：

# 普通数据 Tile：ND（行主序）
tile_src = pl.TileType(shape=[32, 128], dtype=pl.DT_UINT16,
                       target_memory=pl.MemorySpace.Vec, layout=pl.ND)

# 归约结果列向量：DN（列主序）
tile_red = pl.TileType(shape=[TILE_ROWS, 1], dtype=pl.DT_FP32,
                       target_memory=pl.MemorySpace.Vec, layout=pl.DN)

Cube分形布局#

Cube计算的L1/L0A/L0B/L0C各级Buffer使用分形布局，默认值由内存空间决定：

# L1 Mat：默认 NZ
mat_type = pl.TileType(shape=[64, 64], dtype=pl.DT_FP32,
                       target_memory=pl.MemorySpace.Mat, layout=pl.NZ)

# L0B Right：默认 ZN
right_type = pl.TileType(shape=[64, 64], dtype=pl.DT_FP32,
                         target_memory=pl.MemorySpace.Right, layout=pl.ZN)

# L0C Acc：默认 NZ，fp32 需指定 fractal=1024
acc_type = pl.TileType(shape=[64, 64], dtype=pl.DT_FP32,
                       target_memory=pl.MemorySpace.Acc, layout=pl.NZ, fractal=1024)

A3架构L0A的ZZ布局#

A3架构下L0A（左矩阵）默认ZZ；A5架构下默认NZ。以下为显式指定ZZ的用法：

tile_type = pl.TileType(
    shape=[128, 128], dtype=pl.DT_FP16,
    target_memory=pl.MemorySpace.Left, layout=pl.ZZ,
)

NZ Tensor输入输出#

Tensor标注为NZ时，pl.load/pl.store按NZ物理布局访问GM，目标/源Tile也必须使用pl.NZ。以下示例展示二维GM NZ在UB中的原始搬入和写回：

@pl.jit()
def copy_nz_kernel(
    nz_in: pl.Tensor[[64, 64], pl.DT_FP16, pl.NZ],
    nz_out: pl.Tensor[[64, 64], pl.DT_FP16, pl.NZ],
):
    tile_type = pl.TileType(
        shape=[64, 64], dtype=pl.DT_FP16,
        target_memory=pl.MemorySpace.Vec, layout=pl.NZ,
    )
    tile = pl.make_tile(tile_type, addr=0x0000)
    with pl.section_vector():
        pl.load(tile, nz_in, [0, 0])
        pl.store(nz_out, tile, [0, 0])

高维NZ沿用相同写法，例如pl.Tensor[[B, H, M, N], dtype, pl.NZ]固定以最后两轴M/N作为分形轴，B/H为batch轴。GM NZ不支持通过order=[1, 0]转置，也不支持直接搬入ND/ZN Tile。完整分形、offset以及L0C直接写回限制见load和store。

MX scale的ZZ与NN布局#

ZZ和NN分别用于存放MX矩阵计算中A矩阵和B矩阵的E8M0 scale。L1 Mat Tile默认使用NZ，因此A矩阵的scale需要显式指定ZZ，B矩阵的scale需要显式指定NN；ScaleLeft和ScaleRight Tile则分别默认使用ZZ和NN。

# A/B矩阵的E8M0 scale逻辑shape分别为[M,G]和[G,N]，其中G=K/32。
M, G, N = 64, 4, 64

# L1 Mat的默认布局是NZ，因此需要显式指定ZZ或NN。
scale_a_l1_type = pl.TileType(
    shape=[M, G],
    dtype=pl.DT_FP8E8M0,
    target_memory=pl.MemorySpace.Mat,
    layout=pl.ZZ,
)
scale_b_l1_type = pl.TileType(
    shape=[G, N],
    dtype=pl.DT_FP8E8M0,
    target_memory=pl.MemorySpace.Mat,
    layout=pl.NN,
)

# ScaleLeft/ScaleRight分别默认使用ZZ/NN，无需再次指定layout。
scale_a_type = pl.TileType(
    shape=[M, G],
    dtype=pl.DT_FP8E8M0,
    target_memory=pl.MemorySpace.ScaleLeft,
)
scale_b_type = pl.TileType(
    shape=[G, N],
    dtype=pl.DT_FP8E8M0,
    target_memory=pl.MemorySpace.ScaleRight,
)