PyPTO 非对齐存储接口 vf.store_unalign 全解析：变长向量与掩码数据的高效 Tile 搬出

PyPTO 非对齐存储接口 vf.store_unalign 全解析：变长向量与掩码数据的高效 Tile 搬出

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

vf.store_unalign是 CANN PyPTO 向量编程（Vector Function，VF）中用于非对齐数据搬出的核心接口：它将变长向量数据（reg_tensor）或掩码数据（mask_reg）写入 UB（Unified Buffer）上的 Tile 地址，并配合vf.unalign_reg_for_store()与vf.store_unalign_post()构成完整的非对齐存储序列。本文以 store_unalign.md 为骨架，结合仓库中 VF API 声明 与 CCE 后端代码生成实现，讲解 ureg 状态机原理、参数语义、指令映射与四类实战用法（基本搬出、AddrReg 变长步长、mask_reg 打包、INT64 宽数据类型），帮助你在非对齐数据流场景中写出正确且高性能的搬出代码。

产品支持情况

vf.store_unalign的可用性随昇腾硬件平台而异，具体如下：

在 A2/A3 系列产品上，非对齐存储需要改用其他搬出方案（如对齐存储配合标量搬移），使用前请先确认目标运行平台的架构能力。

功能说明：非对齐存储解决什么问题

对齐存储（vf.store_align）要求寄存器数据与 Tile 地址按固定粒度对齐；而实际算子往往需要把变长向量或mask 比较结果写入到地址不满足对齐约束的 Tile 位置。vf.store_unalign正是为此设计：

• 将变长向量数据（reg_tensor）或 mask_reg 数据写入 Tile 地址 dstAddr；
• 当src为 mask_reg 时，后端自动分派 mask_reg 非对齐存储路径，无需使用者手工区分指令形态；
• 必须与vf.unalign_reg_for_store()（分配对齐跟踪寄存器）和vf.store_unalign_post()（搬出后处理）配对使用。

从仓库中的 API 声明可以看到，该接口属于语句形式（statement form）的副作用操作：_vf_api.py中store_unalign与store_unalign_post均无返回值，目的操作数是 UB Tile 指针而非寄存器。在 CCE 代码生成阶段，cce_codegen.cpp将这些接口列入skip_ops集合——即vf.store_unalign、vf.store_unalign_post、vf.unalign_reg_for_store等操作的首参是 Tile 指针而非寄存器，编译器不会为其自动声明 RegTensor，这也与“存储类接口无返回值、以显式 tile 为目的地”的语义保持一致。

非对齐搬出原理：ureg 状态机与两阶段搬出

非对齐搬出依赖一个被称为ureg（unalign register）的寄存器状态来缓存边界数据。根据 ureg 当前是否为空，搬出分为两种场景。以下均以“源 reg_tensor 256B、目标 Tile 地址 dstAddr ~ 304”为例说明。

场景一：ureg 为空（第一次迭代）

处理流程如下：

1. 调用vf.store_unalign：此时 ureg 内无有效数据，表示连续非对齐搬出的起始状态。系统将源 reg_tensor 中对应 Tile 地址48 ~ 288的数据写入 dstAddr；同时，源 reg_tensor 中对应 Tile 地址288 ~ 304的数据会被写入 ureg 缓存。
2. 调用vf.store_unalign_post：进行非对齐搬出后处理，将 ureg 中缓存的 288 ~ 304 数据写入 Tile。

第一次迭代的尾部 16B 数据并不立即落盘，而是暂存于 ureg，等待后处理阶段统一刷出，从而保证目标 Tile 上 [dstAddr, 304) 区间的数据完整、连续。

场景二：ureg 不为空（除第一次迭代）

处理流程如下：

1. 调用vf.store_unalign：此时 ureg 内有有效数据，系统将ureg 中 Tile 地址 32 ~ dstAddr 对应的数据与源 reg_tensor 中 Tile 地址 dstAddr ~ 288 对应的数据进行拼接，将拼接结果写入 Tile 地址 dstAddr；同时，源 reg_tensor 中对应 Tile 地址288 ~ 304的数据仍会被写入 ureg，供下一次迭代使用。
2. 调用vf.store_unalign_post：将 ureg 中缓存的 288 ~ 304 数据写入 Tile。

对比两个场景可以发现：ureg 的作用是跨迭代缓存上一轮搬出的尾部残片，并在本轮搬出时与当前源数据的头部拼接，从而把“非对齐搬出”分解为“对齐主体搬出 + 边界残片拼接”，这正是非对齐搬出性能优于逐元素搬移的关键所在。

连续非对齐搬入搬出：ureg 缓存复用与性能优化

在循环中连续执行非对齐搬运时，搬入（vf.load_unalign）与搬出（vf.store_unalign）可充分利用 ureg 缓存避免重复调用前后处理接口：

• 连续非对齐搬入：vf.load_unalign会将后续未对齐的数据缓存至 ureg，所以下一次搬入不需要再次调用vf.load_unalign_pre，只需在迭代开始前调用一次vf.load_unalign_pre，即可实现非对齐搬入的性能优化。
• 连续非对齐搬出：下一次迭代的vf.store_unalign会把本次迭代vf.store_unalign缓存至 ureg 中的数据写入 Tile，因此本次迭代不需要调用vf.store_unalign_post来刷出 ureg 数据，只需在迭代结束后调用一次vf.store_unalign_post，即可实现非对齐搬出的性能优化。

简言之：搬入侧把load_unalign_pre提到循环外只做一次，搬出侧把store_unalign_post放到循环后只做一次，中间每轮迭代的 ureg 数据由相邻迭代自动接力。

函数原型与参数说明

store_unalign(tile, src, align_reg, stride=None, post_update: bool = False)

各参数语义如下表所示：

关于stride的两种形态，源码声明给出了更底层的指令映射（见 _vf_api.py）：

• 整型标量 stride→ 发射vstus（strided 模式），地址按固定步长推进；
• AddrReg stride→ 发射vstu（AddrReg 模式），AddrReg 提供一组元素偏移地址，用于 scatter 形式的非对齐搬出，且vstu恒为 POST_UPDATE 语义。

配套的vf.store_unalign_post(tile, align_reg, stride, post_update=False)完成搬出序列的收尾（见 _vf_api.py）：

• stride 为整型标量时发射vstas（strided 模式）；
• stride 为 AddrReg 时发射vsta（AddrReg 模式），必须与store_unalign中的vstu成对使用，且vsta恒为 POST_UPDATE；
• 用于 mask_reg 搬出的后处理时，该接口的 stride 参数不生效。

约束说明与返回值

• vf.store_unalign必须与vf.store_unalign_post()配对使用，并在vf.store_unalign_post()之前调用；
• 调用vf.store_unalign/vf.store_unalign_post之前，需先通过vf.unalign_reg_for_store()分配对齐状态寄存器（该接口返回非对齐寄存器句柄）；
• 返回值：无。

调用示例

基本非对齐存储

非对齐搬出通常与非对齐搬入配套使用，形成完整的非对齐数据搬运闭环：vf.load_unalign_init分配非对齐搬入寄存器，vf.load_unalign_pre初始化缓存，vf.load_unalign执行搬入，vf.unalign_reg_for_store分配搬出对齐寄存器，vf.store_unalign执行搬出，vf.store_unalign_post刷出剩余数据。

import os import pypto_pro.language as pl import torch import torch_npu @pl.vector_function def example_vf(src_tile, dst_tile): ureg = vf.load_unalign_init() vf.load_unalign_pre(ureg, src_tile) src_reg = vf.load_unalign(ureg, src_tile, post_update=True) store_ureg = vf.unalign_reg_for_store() vf.store_unalign(dst_tile, src_reg, store_ureg, 64, post_update=True) vf.store_unalign_post(dst_tile, store_ureg, 0, post_update=True) @pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], ): tf = pl.TileType(shape=[1, 64], dtype=pl.DT_FP32, target_memory=pl.MemorySpace.Vec) in_a_grp = pl.make_tile_group(type=tf, addrs=0x0, mutex_ids=[0]) in_a = in_a_grp.current() t_out_grp = pl.make_tile_group(type=tf, addrs=0x100, mutex_ids=[1]) t_out = t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) example_vf(in_a, t_out) pl.store(out, t_out, [0, 0]) def test_example(): device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0)) device = f"npu:{device_id}" core_nums = 1 torch.npu.set_device(device) a = torch.randn([1, 64], device=device, dtype=torch.float32) out = torch.empty([1, 64], device=device, dtype=torch.float32) example_kernelNone, core_nums torch.npu.synchronize() torch.testing.assert_close(out, a, rtol=1e-5, atol=1e-5) if __name__ == "__main__": test_example() print("PASSED")

代码要点：

• 标量 stride 取64且post_update=True，表示每段搬出后 tracker 自动前进 64 个元素；store_unalign_post使用stride=0收尾；
• 测试用例通过torch.testing.assert_close校验搬入搬出往返后的数据一致性，可作为该接口正确性的可复现验证手段。

AddrReg 非对齐存储示例

当stride参数传入 AddrReg（由vf.create_addr_reg创建）时，AddrReg 提供一组向量偏移地址，适用于变长步长的非对齐搬出场景——每条 lane 的目的地址由向量偏移独立指定，这正是vstu/vsta指令的 scatter 语义。

import os import pypto_pro.language as pl import torch import torch_npu @pl.vector_function def example_vf(src_tile, dst_tile): ureg = vf.load_unalign_init() vf.load_unalign_pre(ureg, src_tile) src_reg = vf.load_unalign(ureg, src_tile, post_update=True) store_ureg = vf.unalign_reg_for_store() for i in pl.range(0, 1, 1): addr_reg = vf.create_addr_reg(64, dtype=pl.DT_FP32) vf.store_unalign(dst_tile, src_reg, store_ureg, addr_reg, post_update=True) vf.store_unalign_post(dst_tile, store_ureg, addr_reg) @pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], ): tf = pl.TileType(shape=[1, 64], dtype=pl.DT_FP32, target_memory=pl.MemorySpace.Vec) in_a_grp = pl.make_tile_group(type=tf, addrs=0x0, mutex_ids=[0]) in_a = in_a_grp.current() t_out_grp = pl.make_tile_group(type=tf, addrs=0x100, mutex_ids=[1]) t_out = t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) example_vf(in_a, t_out) pl.store(out, t_out, [0, 0]) def test_example(): device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0)) device = f"npu:{device_id}" core_nums = 1 torch.npu.set_device(device) a = torch.randn([1, 64], device=device, dtype=torch.float32) out = torch.empty([1, 64], device=device, dtype=torch.float32) example_kernelNone, core_nums torch.npu.synchronize() assert out.shape == torch.Size([1, 64]) if __name__ == "__main__": test_example() print("PASSED")

使用 AddrReg 时注意：

• vf.create_addr_reg(64, dtype=pl.DT_FP32)创建包含 64 个 FP32 元素偏移的地址寄存器，作为每次搬出的向量步长；
• 主搬出与后处理必须同时使用同一个 AddrReg（store_unalign与store_unalign_post的 stride 均为addr_reg），否则会破坏vstu/vsta的配对约束。

mask_reg 非对齐存储示例

当src为 mask_reg 时，vf.store_unalign自动分派 mask_reg 非对齐存储路径，硬件会将 mask_reg 的 32 字节数据按位宽打包后写入 Tile：

• 按 16 位宽（DT_INT16、DT_UINT16、DT_FP16、DT_BF16）打包为 16 字节；
• 按 32 位宽（DT_INT32、DT_UINT32、DT_FP32）打包为 8 字节；
• 硬件从每 2bit（16 位宽）/ 4bit（32 位宽）中提取最低有效位（LSB）。

import os import pypto_pro.language as pl import torch import torch_npu @pl.vector_function def example_vf(src_tile, mask_buf_tile): preg = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_FP32) reg_a = vf.load_align(src_tile, 0) cmp_mask = vf.ge(reg_a, 0.0, preg) ureg = vf.unalign_reg_for_store() vf.store_unalign(mask_buf_tile, cmp_mask, ureg) vf.store_unalign_post(mask_buf_tile, ureg, 0, post_update=True) @pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_UINT32], ): tf = pl.TileType(shape=[1, 64], dtype=pl.DT_FP32, target_memory=pl.MemorySpace.Vec) tu = pl.TileType(shape=[1, 64], dtype=pl.DT_UINT32, target_memory=pl.MemorySpace.Vec) in_a_grp = pl.make_tile_group(type=tf, addrs=0x0, mutex_ids=[0]) in_a = in_a_grp.current() t_mask_grp = pl.make_tile_group(type=tu, addrs=0x100, mutex_ids=[1]) t_mask = t_mask_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) pl.load(t_mask, out, [0, 0]) example_vf(in_a, t_mask) pl.store(out, t_mask, [0, 0]) def test_example_2(): device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0)) device = f"npu:{device_id}" core_nums = 1 torch.npu.set_device(device) a = torch.ones([1, 64], device=device, dtype=torch.float32) out = torch.zeros([1, 64], device=device, dtype=torch.int32) example_kernelNone, core_nums torch.npu.synchronize() assert (out != 0).any(), "全正输入应产生非零打包掩码" a = -torch.ones([1, 64], device=device, dtype=torch.float32) out = torch.zeros([1, 64], device=device, dtype=torch.int32) example_kernelNone, core_nums torch.npu.synchronize() assert (out == 0).all(), "全负输入应产生零打包掩码" if __name__ == "__main__": test_example_2() print("PASSED")

要点说明：

• 先通过vf.ge对输入张量逐元素与 0.0 比较得到cmp_mask（mask_reg），再整体非对齐搬出到 mask 缓冲 Tile；
• src为 mask_reg 时不传 stride，后处理store_unalign_post的 stride 参数对 mask 路径不生效；
• 测试用例利用“全正输入产生非零打包掩码、全负输入产生零打包掩码”的语义做双向断言，是验证 mask 打包路径的正确性范式。

INT64 数据类型示例

宽数据类型（8 字节元素）同样支持非对齐搬出，示例中以DT_INT64演示了非对齐存储与对齐存储（vf.store_align）的混用：

import os import pypto_pro.language as pl import torch import torch_npu @pl.vector_function def example_vf_int64(src_tile, dst_tile): preg = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_INT64) reg_a = vf.load_align(src_tile, 0) ureg = vf.unalign_reg_for_store() vf.store_unalign(dst_tile, reg_a, ureg, 64, post_update=True) vf.store_unalign_post(dst_tile, ureg, 0, post_update=True) reg_out = reg_a vf.store_align(dst_tile, reg_out, preg) @pl.jit() def example_kernel_int64( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_INT64], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_INT64], ): tf = pl.TileType(shape=[1, 32], dtype=pl.DT_INT64, target_memory=pl.MemorySpace.Vec) in_a_grp = pl.make_tile_group(type=tf, addrs=0, mutex_ids=[0]) in_a = in_a_grp.current() t_out_grp = pl.make_tile_group(type=tf, addrs=256, mutex_ids=[1]) t_out = t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) example_vf_int64(in_a, t_out) pl.store(out, t_out, [0, 0]) def test_example_int64(): device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0)) device = f"npu:{device_id}" core_nums = 1 torch.npu.set_device(device) a = torch.randint(-100, 100, [1, 32], device=device, dtype=torch.int64) out = torch.empty([1, 32], device=device, dtype=torch.int64) example_kernel_int64None, core_nums torch.npu.synchronize() torch.testing.assert_close(out, a, rtol=0, atol=0) if __name__ == "__main__": test_example_int64() print("PASSED")

实战要点小结

• 配对使用：vf.store_unalign必须与vf.store_unalign_post成对出现，且前者在前；两者共用的align_reg必须先由vf.unalign_reg_for_store()分配。
• 循环内免后处理：连续迭代场景下，把vf.store_unalign_post移到循环外只调用一次，由 ureg 在相邻迭代间接力，可显著减少指令开销；搬入侧的vf.load_unalign_pre同理前移。
• stride 选择：固定步长用整型标量（vstus/vstas），变长步长或 scatter 型搬出用vf.create_addr_reg创建的 AddrReg（vstu/vsta），且主搬出与后处理必须使用同一个 AddrReg。
• mask 路径：src为 mask_reg 时无需传 stride，后端自动分派打包路径（16 位宽打包为 16 字节、32 位宽打包为 8 字节，取 LSB）。
• 数据与类型一致性：tile与src数据类型必须一致，支持从 8 位整型到 FP8/FP4 再到 64 位整型的全谱系数据类型。
• 平台约束：当前仅在 Ascend 950PR/Ascend 950DT 上受支持，A2/A3 系列产品不可用，移植代码前请先核对目标平台。

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto