ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

PyPTO vf.unalign_reg_for_store 详解:非对齐存储的 alignment tracker 寄存器分配与使用

PyPTO vf.unalign_reg_for_store 详解:非对齐存储的 alignment tracker 寄存器分配与使用 人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载本篇文章以 PyPTOParallel Tensor/Tile Operation 编程范式中向量函数pl.vector_function非对齐存储调用链为核心系统讲解vf.unalign_reg_for_store()的定位、函数原型、返回值、配对约束与底层原理。你将掌握如何为一个完整的非对齐搬出流程分配 alignment tracker 寄存器并学会在连续迭代、AddrReg 变长步长、mask_reg 打包等典型场景中正确组织vf.store_unalign/vf.store_unalign_post调用序列。一、功能定位非对齐存储调用链的状态寄存器分配入口在 PyPTO 的向量编程模型中reg_tensor支持对非 32 字节对齐地址的连续访问。非对齐存储unaligned store指将reg_tensor/mask_reg中变长的向量数据写入 Tile 地址时数据可能跨越对齐边界、存在多出来的尾部字节无法一次性写出的场景。vf.unalign_reg_for_store()正是为这类非对齐存储分配alignment tracker 寄存器即 align_reg的接口。该寄存器贯穿后续的vf.store_unalign/vf.store_unalign_post调用链用于追踪未对齐字节的累积状态vf.store_unalign把超出对齐边界的数据暂存其中vf.store_unalign_post再将其刷出从而保证非对齐搬出的数据完整性与连续性。从 Python API 声明 的 docstring 可以印证其定位Must be called before store_unalign/store_unalign_post to allocate the alignment state register——即它必须在vf.store_unalign/vf.store_unalign_post之前调用用于分配对齐状态寄存器返回非对齐寄存器句柄。二、产品支持情况该接口在 接口文档 中声明的产品支持情况如下Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持说明vf.unalign_reg_for_store与整个非对齐搬入搬出接口族vf.load_unalign_init、vf.load_unalign_pre、vf.load_unalign、vf.store_unalign、vf.store_unalign_post等的产品支持范围一致均为 Ascend 950 系列专属能力。三、函数原型与返回说明unalign_reg_for_store() - align_reg参数说明无。返回值说明返回align_reg即 alignment tracker 寄存器其类型为目标reg_tensor。约束说明接口本身无独立约束但其分配出的align_reg必须贯穿后续非对齐存储调用链使用即必须在vf.store_unalign/vf.store_unalign_post之前调用见 _vf_api.pyvf.store_unalign与vf.store_unalign_post必须配对使用且后处理必须在搬出之后调用见 store_unalign_post 约束。四、工作原理align_reg 如何追踪未对齐字节vf.unalign_reg_for_store()分配的align_reg本质上是跨对齐边界的尾部数据缓存 偏移状态追踪二合一寄存器。结合 store_unalign 非对齐搬出原理可将非对齐搬出分为两种场景理解场景一align_regureg为空第一次迭代从源 reg_tensor256B读取数据搬运至目标 Tile 地址dstAddr ~ 304调用vf.store_unalign此时 ureg 内无有效数据表示连续非对齐搬出的起始状态。将源 reg_tensor 中对应 Tile 地址 48 ~ 288 的数据写入 dstAddr同时源 reg_tensor 中对应 Tile 地址 288 ~ 304 的尾部数据被写入 uregalign_reg调用vf.store_unalign_post做非对齐搬出后处理将 ureg 中缓存的数据写入 Tile 地址 288 ~ 304。场景二align_regureg不为空除第一次迭代调用vf.store_unalign此时 ureg 内有有效数据系统将 ureg 中 Tile 地址 32 ~ dstAddr 对应的缓存数据与源 reg_tensor 中 Tile 地址 dstAddr ~ 288 对应的数据进行拼接结果写入 Tile 地址 dstAddr同时源 reg_tensor 中对应 Tile 地址 288 ~ 304 的数据再次写入 ureg为下一次搬出做好准备调用vf.store_unalign_post将 ureg 中缓存的数据写入 Tile 地址 288 ~ 304。从 Python API 声明 可以看到指令层面的映射vf.store_unalign(tile, src, align_reg, strideNone, post_updateFalse)stride为整型标量时发射vstusstrided 模式stride为AddrReg时发射vstuAddrReg 模式散播式非对齐存储恒为 POST_UPDATEvf.store_unalign_post(tile, align_reg, stride, post_updateFalse)stride为整型标量时发射vstas为AddrReg时发射vsta且必须与vf.store_unalign的 AddrReg 模式vstu配对。align_reg正是这些硬件指令之间共享的对齐状态载体。连续非对齐搬出优化把 post 移出循环连续非对齐搬出时下一次迭代的vf.store_unalign会自动将本次迭代暂存在 ureg 中的数据写入 Tile因此本次迭代无需立即调用vf.store_unalign_post刷出 ureg 数据只需在整个迭代结束后调用一次vf.store_unalign_post即可实现非对齐搬出的性能优化。这与非对齐搬入侧vf.load_unalign自动缓存尾部数据、只需迭代开始前调用一次vf.load_unalign_pre形成对称设计。五、典型调用模式与完整示例5.1 基本用法完整非对齐搬入搬出流程以下示例完整继承自 unalign_reg_for_store 调用示例展示了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_updateTrue) store_ureg vf.unalign_reg_for_store() vf.store_unalign(dst_tile, src_reg, store_ureg, 64, post_updateTrue) vf.store_unalign_post(dst_tile, store_ureg, 0, post_updateTrue) 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)关键点解析vf.unalign_reg_for_store()无参数、无副作用依赖只需在搬出调用前执行一次即可获得store_uregvf.store_unalign(dst_tile, src_reg, store_ureg, 64, post_updateTrue)中stride64为整型标量表示地址更新步长元素个数配合post_updateTrue使 tracker 自动累进到下一段vf.store_unalign_post(dst_tile, store_ureg, 0, post_updateTrue)将store_ureg中缓存的尾部数据刷出到 Tile上层 kernel 通过pl.TileType、pl.make_tile_group以显式地址0x0/0x100与mutex_ids声明 Tile 组并在pl.section_vector()向量段内完成加载、向量函数调用与存储。5.2 AddrReg 模式变长步长的非对齐搬出当stride传入由vf.create_addr_reg创建的AddrReg时vf.store_unalign发射 vstu 指令使用一组向量偏移地址替代标量步长适用于变长步长的散播式非对齐搬出场景vf.store_unalign_post必须同样以该AddrReg配对发射 vsta 指令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_updateTrue) store_ureg vf.unalign_reg_for_store() for i in pl.range(0, 1, 1): addr_reg vf.create_addr_reg(64, dtypepl.DT_FP32) vf.store_unalign(dst_tile, src_reg, store_ureg, addr_reg, post_updateTrue) vf.store_unalign_post(dst_tile, store_ureg, addr_reg)5.3 mask_reg 模式自动分派的打包存储路径当src为mask_reg时vf.store_unalign自动分派 mask_reg 非对齐存储路径无需传 stridemask_reg 32 字节数据按 16 位宽DT_INT16、DT_UINT16、DT_FP16、DT_BF16打包为 16 字节或按 32 位宽DT_INT32、DT_UINT32、DT_FP32打包为 8 字节写入 Tile硬件从每 2bit16 位宽/ 4bit32 位宽中提取最低有效位LSB。此时vf.unalign_reg_for_store()分配的store_ureg同样作为 alignment tracker 传入pl.vector_function def example_vf(src_tile, mask_buf_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.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_updateTrue)5.4 INT64 等宽数据类型场景非对齐存储支持的数据类型包括 DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_FP16、DT_BF16、DT_INT32、DT_UINT32、DT_FP32、DT_INT64、DT_UINT64、DT_FP8E4M3FN、DT_FP8E5M2、DT_FP8E8M0、DT_HF8、DT_FP4E2M1、DT_FP4E1M2。以 INT64 为例vf.unalign_reg_for_store()的使用方式与基本模式一致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) ureg vf.unalign_reg_for_store() vf.store_unalign(dst_tile, reg_a, ureg, 64, post_updateTrue) vf.store_unalign_post(dst_tile, ureg, 0, post_updateTrue) reg_out reg_a vf.store_align(dst_tile, reg_out, preg)六、源码级印证与使用要点在 python/pypto_pro/language/_vf_api.py 中vf命名空间下的相关 API 全部以_api_decl静态方法形式声明构成完整的非对齐存储接口族接口底层指令/行为与 align_reg 的关系unalign_reg_for_store()分配 alignment tracker 状态寄存器本接口返回的align_reg是以下所有调用的共享状态store_unalign(tile, src, align_reg, stride, post_update)vstus标量 stride/ vstuAddrReg消费并更新align_reg暂存跨边界尾部数据store_unalign_post(tile, align_reg, stride, post_update)vstas标量 stride/ vstaAddrReg读取align_reg中的剩余字节并刷出到 Tilesqueeze_store_unalign/squeeze_store_unalign_postvstur / vstar同样接收align_reg但基于vf.squeeze(gather_modeSTORE_REG)写入 AR 寄存器的有效字节数作为隐式步长实际使用中的核心要点归纳先分配、后使用vf.unalign_reg_for_store()必须在vf.store_unalign/vf.store_unalign_post之前调用且一个完整的搬出流程通常只分配一次严格配对vf.store_unalign与vf.store_unalign_post必须成对出现后处理在搬出之后调用AddrReg 模式下vstu/vsta必须使用同一个AddrReg循环优化连续多次非对齐搬出时将vf.store_unalign_post移出循环、仅在整个迭代结束后调用一次由vf.store_unalign在迭代间自动接力 ureg 中的缓存数据可显著降低后处理指令开销对称理解搬出侧的unalign_reg_for_store→store_unalign→store_unalign_post与搬入侧的load_unalign_init→load_unalign_pre→load_unalign详见 load_unalign 文档一一对应是 PyPTO 处理非 32 字节对齐内存访问的统一机制。七、总结vf.unalign_reg_for_store()是 PyPTO 非对齐存储调用链的起点它分配并返回 alignment tracker 寄存器align_reg该寄存器在vf.store_unalign暂存跨边界尾部数据与vf.store_unalign_post刷出剩余字节之间传递未对齐字节的累积状态支撑了变长数据在非 32 字节对齐地址上的高效连续搬出。无论是标量 stride 的 vstus/vstas 模式、AddrReg 变长步长的 vstu/vsta 模式还是 mask_reg 自动分派路径vf.unalign_reg_for_store()都是必须先执行的分配入口。该能力目前仅由 Ascend 950 系列产品支持在 A2/A3 平台上不可用使用时需注意平台前提。赞分享人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐PyPTO vf.load_unalign_init 详解为非对齐搬入分配 UnalignRegForLoad 寄存器PyPTO vf.load_unalign_init 详解为非对齐搬入分配 UnalignRegForLoad 寄存器 导读 vf.load_unalign_人工智能编译器模型编译高性能计算深度学习CANNvf.squeeze_store_unalign_post 详解PyPTO 向量编程范式下非对齐压缩存储的收尾与 AR 寄存器残量处理vf.squeeze_store_unalign_post 详解PyPTO 向量编程范式下非对齐压缩存储的收尾与 AR 寄存器残量处理 导读 vf.squee人工智能编译器模型编译高性能计算深度学习CANNCANN PyPTO vf.squeeze_store_unalign 详解基于 vstur 的无偏移非对齐存储CANN PyPTO vf.squeeze_store_unalign 详解基于 vstur 的无偏移非对齐存储 导读 vf.squeeze_store_un人工智能编译器模型编译高性能计算深度学习CANN创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表