人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载StoreDist 是 PyPTOParallel Tensor/Tile Operation 编程范式中用于描述 vf.store_align 数据存储分布模式的枚举类型它决定了数据从寄存器reg_tensor / mask_reg搬运到统一缓冲区UBTile 的方式。本文以 StoreDist 官方文档 为主体结合 vf.store_align 接口文档 与 前端 API 声明源码 展开帮助你在编写 vector function 时精准选择存储分布模式实现普通对齐存储、首元素存储、压缩存储、交错存储与 mask 压缩等能力。产品支持情况StoreDist 的使用与硬件平台强相关当前仓库文档明确给出的支持矩阵如下Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持也就是说StoreDist 相关存储分布能力目前仅在 Ascend 950 系列产品上可用编写代码前请先确认目标运行环境。功能定位StoreDist 是什么StoreDist 定义了 vf.store_align 的数据存储分布模式用于控制从寄存器到 UB 的数据搬运方式。它是vf.store_align的可选关键字参数dist的取值类型完整函数原型见 store_align 文档store_align(tile, src, *args, dist: Optional[StoreDist] None, data_copy_mode: Optional[DataCopyMode] None, block_strideNone, repeat_strideNone, post_update: bool False)从 前端 API 声明 中可以看到store_align对应底层vsts指令语句形式dst是 UB Tile 而非寄存器而dist参数的语义在源码 docstring 中被描述为 selecting the store distribution pattern并注明与 AscendCReg::StoreDist对应。概括来说粗粒度模式名NORM、FIRST_ELEMENT、PACK、PACK4、INTLV会根据源数据类型自动选择位宽粒度变体带位宽限定的模式名如NORM_B8、INTLV_B32则显式指定按 8/16/32/64 位宽粒度执行。下图展示了vf.store_align连续对齐搬出时部分分布模式的搬出示意原型定义StoreDist 完整枚举StoreDist 是一个继承自enum.Enum的类型完整定义如下与官方文档一致class StoreDist(enum.Enum): # reg_tensor单搬出模式 NORM ... # 普通对齐存储默认根据dtype自动选择位宽粒度64位宽数据类型DT_INT64/DT_UINT64只支持此模式 NORM_B8 ... # 按8位宽类型普通存储 NORM_B16 ... # 按16位宽类型普通存储 NORM_B32 ... # 按32位宽类型普通存储 FIRST_ELEMENT ... # 仅存储lane 0首个元素根据dtype自动选择位宽粒度 FIRST_ELEMENT_B8 ... # 按8位宽类型仅存储首个元素 FIRST_ELEMENT_B16 ... # 按16位宽类型仅存储首个元素 FIRST_ELEMENT_B32 ... # 按32位宽类型仅存储首个元素 PACK ... # 压缩存储根据mask将src中有效元素的低半部分bit数据连续存储于dst中根据dtype自动选择位宽粒度 PACK_B16 ... # 按16位宽类型压缩存储 PACK_B32 ... # 按32位宽类型压缩存储 PACK_B64 ... # 按B64粒度压缩存储 PACK4 ... # 4元素压缩存储将有效元素的低8bit数据连续存储32位宽类型 PACK4_B32 ... # 按32位宽类型4元素压缩存储 # reg_tensor双搬出模式 INTLV ... # 交错存储将src0、src1中的元素交错存储根据dtype自动选择位宽类型 INTLV_B8 ... # 按8位宽类型交错存储 INTLV_B16 ... # 按16位宽类型交错存储 INTLV_B32 ... # 按32位宽类型交错存储 # mask_reg模式 # NORM 同上搬运VL/8数据 # PACK 同上每间隔1bit舍弃数据将VL/8的数据压缩为VL/16搬出在 PyPTO 中通过pypto_pro.language命名空间访问即pl.StoreDist.XXXStoreDist已在 language/init.py 中导出。枚举成员分类解读StoreDist 的成员按用途分为三大类理解这三大类有助于正确选型reg_tensor 单搬出模式将一个 reg_tensor 中的 VL 数据量搬出到 Tile包括NORM系列普通存储、FIRST_ELEMENT系列仅存首个元素、PACK系列压缩存储、PACK4系列4 元素压缩。reg_tensor 双搬出模式interleave将两个 reg_tensor 中的元素交错搬出到 Tiledst长度为 2*VL包括INTLV系列。该模式需要两个源寄存器。mask_reg 模式当源操作数已通过vf.create_mask预声明为 mask_reg 时后端自动分派 mask_reg 存储路径仅支持NORM与PACK不支持显式粒度形式。约束说明各模式支持的 StoreDist 取值StoreDist 的取值并非在任意场景下都可用需要根据目标模式区分。以下是约束汇总reg_tensor 单搬出模式支持NORM含NORM_B8/NORM_B16/NORM_B32、FIRST_ELEMENT含FIRST_ELEMENT_B8/B16/B32、PACK含PACK_B16/B32/B64、PACK4含PACK4_B32。reg_tensor 双搬出模式支持INTLV含INTLV_B8/INTLV_B16/INTLV_B32需要两个源寄存器。mask_reg 模式支持NORM、PACK不支持显式粒度形式。单搬出模式 dist 取值与对齐约束结合 vf.store_align 文档 中表1reg_tensor 单搬出模式下各取值的行为与对齐约束如下dist 取值含义对齐约束ByteStoreDist.NORM正常模式通用根据数据类型自动选择位宽粒度搬运 VL 数据。64 位宽数据类型 DT_INT64、DT_UINT64 只支持此模式。32StoreDist.NORM_B8按 8 位宽类型普通存储搬运 VL 数据。32StoreDist.NORM_B16按 16 位宽类型普通存储搬运 VL 数据。32StoreDist.NORM_B32按 32 位宽类型普通存储搬运 VL 数据。32StoreDist.FIRST_ELEMENT忽略 mask仅向 dst 搬出 src 第一个元素通用根据数据类型自动选择位宽粒度。按 dtype 宽度StoreDist.FIRST_ELEMENT_B8忽略 mask仅向 dst 搬出 src 第一个元素数据类型为 8 位宽类型。1StoreDist.FIRST_ELEMENT_B16忽略 mask仅向 dst 搬出 src 第一个元素数据类型为 16 位宽类型。2StoreDist.FIRST_ELEMENT_B32忽略 mask仅向 dst 搬出 src 第一个元素数据类型为 32 位宽类型。4StoreDist.PACK压缩模式通用根据 mask 将 src 中有效元素的低半部分 bit 数据连续存储于 dst 中根据数据类型自动选择位宽粒度。min(32, VL/2)StoreDist.PACK_B16压缩模式按 16 位宽类型粒度将有效元素的低半部分 bit 数据连续存储于 dst。min(32, VL/2)StoreDist.PACK_B32压缩模式按 32 位宽类型粒度连续存储。min(32, VL/2)StoreDist.PACK_B64压缩模式按 64 位宽类型粒度连续存储。min(32, VL/2)StoreDist.PACK44 元素压缩模式通用根据 mask 将 src 中有效元素的低 8bit四分之一数据连续存储于 dst。min(32, VL/4)StoreDist.PACK4_B324 元素压缩模式按 32 位宽类型粒度将有效元素的低 8bit四分之一数据连续存储于 dst。min(32, VL/4)要点解读NORM 是唯一默认模式也是最通用的选择。当数据类型为 64 位宽DT_INT64/DT_UINT64时只允许 NORM 模式这一点在使用 64 位数据类型时必须注意。FIRST_ELEMENT 系列忽略 mask只搬出 lane 0 的元素可用于取首元素这类数据抽取场景其对齐约束随显式粒度B8/B16/B32降为 1/2/4 字节。PACK 系列压缩的是低半部分 bit 数据即每个有效元素只保留低半位宽PACK4 则只保留低 8bit四分之一两者对齐约束分别为 min(32, VL/2) 与 min(32, VL/4)。双搬出交错模式 dist 取值与对齐约束对应 store_align 文档表2双搬出模式的取值如下dist 取值含义对齐约束ByteStoreDist.INTLV交错存储通用将 src0、src1 中的元素交错存储于 dst 中根据数据类型自动选择位宽粒度。32StoreDist.INTLV_B8按 8 位宽类型交错存储。32StoreDist.INTLV_B16按 16 位宽类型交错存储。32StoreDist.INTLV_B32按 32 位宽粒度交错存储。32INTLV 系列需要两个源寄存器偶数/奇数源典型用法是与vf.load_align(..., distpl.LoadDist.DINTLV_B32)配对加载时按交错反解出偶/奇两个寄存器存储时再用INTLV_B32交错写回。mask_reg 模式 dist 取值对应 store_align 文档表3当源为 mask_reg 时dist 取值含义对齐约束ByteStoreDist.NORM正常模式搬运 VL/8 数据。VL/8StoreDist.PACK压缩模式每间隔 1bit 舍弃数据将 VL/8 的数据压缩为 VL/16 搬出。VL/16注意 mask_reg 模式下不支持显式粒度形式即不能用NORM_B8、PACK_B16等只能使用NORM或PACK。调用示例最基础的调用方式如下来自官方文档示例import pypto_pro.language as pl pl.vector_function def vf_kernel(): vf.store_align(ub_tile, reg, distpl.StoreDist.NORM)其中reg为 reg_tensor 源操作数ub_tile为 32 字节对齐的 Tile 目的地址distpl.StoreDist.NORM使用默认的普通对齐存储模式。实战示例一普通对齐存储完整可运行的 FP32 示例含 jit kernel 与测试验证参考 store_align 文档import os import pypto_pro.language as pl import torch import torch_npu pl.vector_function def example_vf(src_tile, dst_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32) reg vf.load_align(src_tile, 0) vf.store_align(dst_tile, reg, preg) 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], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) in_a_grp pl.make_tile_group(typetf, addrs0x0, mutex_ids[0]) in_a in_a_grp.current() t_out_grp pl.make_tile_group(typetf, addrs0x100, 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 fnpu:{device_id} core_nums 1 torch.npu.set_device(device) a torch.randn([1, 64], devicedevice, dtypetorch.float32) out torch.empty([1, 64], devicedevice, dtypetorch.float32) example_kernelNone, core_nums torch.npu.synchronize() torch.testing.assert_close(out, a, rtol1e-5, atol1e-5) if __name__ __main__: test_example() print(PASSED)实战示例二interleaved 交错存储使用distpypto_pro.language.StoreDist.INTLV_B32将偶数/奇数寄存器交错写入 Tile与vf.load_align的DINTLV_B32加载模式配对pl.vector_function def example_vf(src_tile, dst_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32) dst_even, dst_odd vf.load_align(src_tile, 0, distpl.LoadDist.DINTLV_B32) vf.store_align(dst_tile, dst_even, dst_odd, preg, distpl.StoreDist.INTLV_B32)调用时 Tile 形状取[1, 128]两倍于单寄存器 VL 的元素数因为交错存储后dst长度为 2*VL。测试中torch.testing.assert_close(out, a)验证了加载-交错-存储的往返一致性。实战示例三PACK 压缩存储含 mask_reg 场景仓库测试 test_vf_basic_ops.py 中给出了多种实际用法例如reg_tensor 压缩存储vf.store_align(t_f0, reg_a, preg_f32, distpl.StoreDist.PACK)4 元素压缩vf.store_align(t_f0, reg_a, preg, distpl.StoreDist.PACK4)mask_reg 压缩存储vf.store_align(t_u0, mask_a, distpl.StoreDist.PACK)源为 mask_reg 时无需传谓词 mask。mask_reg 的完整链路参考 store_align 文档的 mask_reg 存储示例先用vf.ge比较产生 mask_reg以PACK模式存入 mask 缓冲 Tilevf.mem_bar同步后以LoadDist.US读回 mask再参与后续计算。该场景下PACK模式每间隔 1bit 舍弃数据把 VL/8 的 mask 数据压缩为 VL/16 搬出与普通元素压缩的语义不同。实战示例四NORM_B16 按位宽存储与 INT64 限制量化类算子如 test_quant_lightning_indexer_vf.py中常使用distpl.StoreDist.NORM_B16将 FP32 寄存器按 16 位宽粒度普通存储以及INTLV_B32完成两路寄存器交错写盘。对于 64 位宽数据DT_INT64/DT_UINT64只能使用StoreDist.NORM根据 dtype 自动选择位宽粒度使用vf.store_align直接存储即可pl.vector_function def example_vf_int64(src_tile, dst_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_INT64) reg_a vf.load_align(src_tile, 0) vf.store_align(dst_tile, reg_a, preg)与 vf.store_align 其他参数的配合StoreDist 并非孤立使用它与vf.store_align的offset、post_update、data_copy_mode等参数协同工作完整参数说明见 store_align 文档offset 偏移整数、[row, col]列表线性偏移 row * shape[1] col或vf.create_addr_reg创建的 AddrReg。dist分布模式与偏移机制正交可任意组合。post_updateTrue搬运后目标地址自动累进适用于循环内连续存储。测试test_vf_basic_ops.py中即有distpl.StoreDist.PACK, post_updateTrue的组合用法。data_copy_modeDataCopyMode.DATA_BLOCK_COPY启用非连续 DataBlock32B 单位搬运此时dist仍可指定分布模式。三类搬运接口普通搬运 / PostUpdate 扩展搬运 / AddrReg 存储偏移量均支持配合 StoreDist 使用。总结StoreDist 是 PyPTO 向量编程中控制寄存器 → UB Tile数据分布的核心枚举按模式选型单搬出用NORM/FIRST_ELEMENT/PACK/PACK4双搬出用INTLV需双源寄存器mask_reg 源只能用NORM/PACK按位宽选型粗粒度名按 dtype 自动选择位宽_B8/_B16/_B32/_B64后缀显式指定粒度64 位数据类型仅支持NORM按场景选型普通搬运用NORM取首元素用FIRST_ELEMENT数据压缩用PACK/PACK4交错排布用INTLV。更多细节可继续阅读 vf.store_align 接口文档 及其完整调用示例并结合 StoreDist 官方定义 与仓库测试用例test_vf_basic_ops.py验证各类模式的实际行为。赞分享人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐PyPTO vf.store_align 深度解析寄存器数据对齐搬出至 Tile 的存储指令全指南PyPTO vf.store_align 深度解析寄存器数据对齐搬出至 Tile 的存储指令全指南 vf.store_align 是 PyPTO 向量函数人工智能编译器模型编译高性能计算深度学习CANNPyPTO LoadDist 枚举完全指南从 UB Tile 到寄存器的数据分布搬运模式PyPTO LoadDist 枚举完全指南从 UB Tile 到寄存器的数据分布搬运模式 导读 本文系统讲解 CANN PyPTOParallel Tens人工智能编译器模型编译高性能计算深度学习CANNCANN PyPTO TensorLayout 数据布局枚举详解GM Tensor 与 Tile 的存储排列与分形布局CANN PyPTO TensorLayout 数据布局枚举详解GM Tensor 与 Tile 的存储排列与分形布局 本文聚焦 CANN PyPTO 编程范人工智能编译器模型编译高性能计算深度学习CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考