ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

PyPTO Pro 向量函数 vf.abs 详解:实数取绝对值与复数取模的双寄存器实现

PyPTO Pro 向量函数 vf.abs 详解:实数取绝对值与复数取模的双寄存器实现 PyPTO Pro 向量函数 vf.abs 详解实数取绝对值与复数取模的双寄存器实现【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptovf.abs是 PyPTO ProParallel Tensor/Tile Operation 编程范式向量函数Vector FunctionVF体系中的基础算术运算接口用于对寄存器RegTensor中的有效元素逐个取绝对值实数或取模复数是编写昇腾 NPU 向量计算内核时最常用的单目运算之一。本文以 abs.md 为核心结合仓库源码与配套文档完整讲解其功能语义、参数规则、双寄存器复数模式原理并给出可直接运行的真机验证示例帮助开发者快速在 PyPTO Pro 内核中落地绝对值计算。产品支持情况vf.abs属于 Ascend 950 系列专属的 VF 寄存器级计算能力目前的产品支持矩阵如下产品形态支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持即vf.abs仅在 Ascend 950 系列上可用。在使用前请先通过pypto_pro.platform或运行环境确认目标设备型号A2/A3 系列上应改用 Tile 级的pl.abs等高层接口。该支持矩阵与 reg_tensor.md、mask_reg.md 等 VF 基础容器文档保持一致说明整个 VF 指令族在 950 系列统一开放。功能说明vf.abs的计算语义分为实数与复数两大类二者映射到不同的硬件指令路径。实数类型逐个取绝对值对实数类型DT_INT8、DT_INT16、DT_INT32、DT_FP16、DT_FP32的源寄存器src对其中每个有效元素逐个取绝对值结果写入目的寄存器dst的对应位置$$dst_i |src_i|$$其中i为寄存器中的有效 lane元素下标有效由掩码寄存器preg的对应比特位决定详见下文 mask_reg 掩码机制。复数类型双寄存器模式取模对复数类型vf.abs通过双寄存器模式dual-register mode实现。该模式下复数由两个同类型浮点寄存器组合表示DT_FP16 双寄存器模式由两个 DT_FP16 组成一个复数实部、虚部各 16 位共 32 位DT_FP32 双寄存器模式由两个 DT_FP32 组成一个复数实部、虚部各 32 位共 64 位。其中reg[0]存放实部reg[1]存放虚部。对src中的有效元素逐个取模即复数绝对值公式如下$$dst_i |src_i| (\alpha^2 \beta^2)^{1/2}$$其中 $src_i \alpha \beta i$α 为复数实部β 为复数虚部。注意复数模式下vf.abs输入的两个寄存器分别承载实部与虚部输出dst即为实数的模值不再需要实部/虚部分别存储这与其他支持双寄存器的复数运算如加减乘的语义不同编写内核时需特别注意寄存器布局。双寄存器模式的存储结构双寄存器模式的数据从 UBUnified BufferTile 搬运到寄存器时采用交错搬入distributed interleave方式读取2 * VL的数据量偶数索引元素存入reg[0]奇数索引元素存入reg[1]。以 DT_FP16 为例两个寄存器共存储 512B 数据reg[0]与reg[1]各存 128 个 DT_FP16 元素RegTensor 搬运 DT_FP16 双寄存器交错存储结构DT_FP32 双寄存器模式同理从 UB 以 DIST_DINTLV_B32 模式读取2 * VL数据量交错搬运两个寄存器各存 64 个 DT_FP32 元素共 512BRegTensor 搬运 DT_FP32 双寄存器交错存储结构关于单/双寄存器模式下各类 dtype 的搬运细节含 b64 类型的低位/高位拆分规则可进一步阅读 reg_tensor.md。函数原型abs(src, preg, mode: Optional[MergeMode] None) - dst该接口在仓库中的声明位于 python/pypto_pro/language/_vf_api.pyvf.absdocstring 中给出的数学语义为$$dstReg_i |srcReg_i|$$与文档描述一致对每个 mask 生效的 lane 计算绝对值并写入目的寄存器。参数说明参数输入/输出说明src输入源操作数reg_tensor。源操作数src与目的操作数dst的数据类型保持一致。支持的数据类型DT_INT8、DT_INT16、DT_INT32、DT_FP16、DT_FP32、DT_INT64。其中 DT_FP16 和 DT_FP32 支持双寄存器模式用于复数取模运算。preg输入mask_reg元素级有效性控制掩码。mode输入可选对应 MergeMode 类型-pypto_pro.language.MergeMode.ZEROING默认preg未筛选的元素在dst中置 0。-pypto_pro.language.MergeMode.MERGING当前不支持。参数细节深入src数据类型说明接口在语法层面支持 6 种数据类型DT_INT8/DT_INT16/DT_INT32/DT_FP16/DT_FP32/DT_INT64其中实数取绝对值路径覆盖 DT_INT8、DT_INT16、DT_INT32、DT_FP16、DT_FP32复数取模路径覆盖 DT_FP16 与 DT_FP32 的双寄存器模式DT_INT64 走独立的整型取绝对值路径见下文 INT64 调用示例。从 reg_tensor.md 的寄存器容量表可知RegTensor 总大小固定为 256 字节因此不同 dtype 对应不同元素个数DT_INT8/DT_FP16/DT_INT32/DT_INT64 分别对应 256/128/64/32 个元素。这一点决定了单次vf.abs处理的 lane 数也解释了为何 INT64 示例中 Tile 形状取[1, 32]恰好填满一个寄存器。preg掩码语义根据 mask_reg.mdmask_reg 总位宽固定为 256 bit粒度由 dtype 决定如 DT_FP32 为 b32 粒度256 bit 覆盖 64 个元素。vf.abs执行时比特位为 1有效该元素参与运算结果写入dst对应位置比特位为 0无效该元素不参与运算在 ZEROING 模式下dst对应位置置 0。mode行为MergeMode.ZEROING是默认且当前唯一支持的取值。MergeMode.MERGING保留目标寄存器原值在当前设备上不支持使用时会报错。该限制同样适用于vf.add、vf.sub、vf.exp等绝大多数 VF 计算接口见 MergeMode.md。约束说明原文档对本接口的约束说明为无。需要说明的是所谓无约束是相对其他 VF 接口如数据类型范围限制而言实际使用时仍需遵循 VF 体系的基础规则src与dst数据类型必须一致寄存器在pypto_pro.language.vector_function函数内创建和使用函数结束后自动释放RegTensor 寄存器数量上限为 32超出上限的数据会写入预留的 8K UB 内存并可能引起性能劣化详见 reg_tensor.mdmask_reg 数量上限为 16详见 mask_reg.md。返回值说明返回dst目的操作数类型为 reg_tensor支持的数据类型与src一致。需要特别关注整型数据的非饱和截断行为整型计算结果若超出数据类型的表示范围将采用非饱和截断wrap-around而非饱和钳位。例如 DT_INT8 类型当src为 -128 时其绝对值 128 超出 DT_INT8 的表示范围-128~127结果会被截断为 -128。这在编写涉及边界值如 INT8_MIN/INT64_MIN的算法时是必须规避的陷阱建议在算法层面提前处理或使用更高位宽的整型中间结果。调用链与底层实现结合源码为帮助理解vf.abs在内核中的完整调用路径下面梳理其前后端接线VF 函数入口pl.vector_function装饰器定义于 python/pypto_pro/language/parser/decorator.py将用户函数标记为可被编译器解析的 SIMD/SIMT 向量函数数据加载vf.load_align(tile, offset)python/pypto_pro/language/_vf_api.py以 vlds 指令将 UB Tile 中的对齐数据加载进寄存器支持整数 offset、AddrReg 或[row, col]列表形式赋值形式隐式声明目的寄存器掩码创建vf.create_mask(pattern, dtype)python/pypto_pro/language/_vf_api.py创建并初始化掩码寄存器pattern默认MaskPattern.ALLdtype默认 FP32也可显式指定与运算寄存器一致的 dtype如 DT_INT64 场景核心运算vf.abs(src, preg, mode)完成取绝对值/取模经 IR 下降映射到硬件向量指令结果写回vf.store_align(tile, src, *args)python/pypto_pro/language/_vf_api.py以 vsts 指令将寄存器数据存回 UB Tile。由此可看出vf.abs是典型的加载-计算-存储三段式 VF 流水中的计算节点其前后均与 UB 内存打交道中间全程在寄存器上完成避免频繁访存。调用示例下面两个示例均来自 abs.md 原文可直接作为独立脚本运行需 Ascend 950 环境、pypto_pro与torch_npu均已安装。基本调用示例DT_FP32import 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_a vf.load_align(src_tile, 0) reg_out vf.abs(reg_a, preg) vf.store_align(dst_tile, reg_out, 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, torch.abs(a), rtol1e-5, atol1e-5) if __name__ __main__: test_example() print(PASSED)示例要点逐段解析pl.vector_function装饰的example_vf是 VF 计算单元vf.create_mask(...)生成全 1 掩码所有元素有效vf.load_align(src_tile, 0)从 Tile 偏移 0 处加载数据vf.abs完成取绝对值vf.store_align将结果写回目标 Tilepl.jit()装饰的example_kernel是内核入口通过pl.TileType(shape[1, 64], ...)声明 UB 上的 Tile 布局[1, 64]恰好对应一个 DT_FP32 寄存器的 64 个元素pl.make_tile_group(addrs0x0/0x100, mutex_ids[0]/[1])为输入输出分配互斥的地址区间pl.section_vector()限定向量计算区内部依次完成全局内存到 UB 的pl.load、VF 计算、UB 到全局的pl.store测试侧用torch.randn构造含正负值的输入并以torch.abs作为参考实现做精度比对rtol1e-5, atol1e-5这是验证 VF 运算正确性的标准做法。INT64 数据类型示例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(patternpl.MaskPattern.ALL, dtypepl.DT_INT64) reg_a vf.load_align(src_tile, 0) reg_out vf.abs(reg_a, preg) 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], dtypepl.DT_INT64, target_memorypl.MemorySpace.Vec) in_a_grp pl.make_tile_group(typetf, addrs0, mutex_ids[0]) in_a in_a_grp.current() t_out_grp pl.make_tile_group(typetf, addrs256, 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 fnpu:{device_id} core_nums 1 torch.npu.set_device(device) a torch.randint(-100, 100, [1, 32], devicedevice, dtypetorch.int64) out torch.empty([1, 32], devicedevice, dtypetorch.int64) example_kernel_int64None, core_nums torch.npu.synchronize() torch.testing.assert_close(out, torch.abs(a), rtol0, atol0) if __name__ __main__: test_example_int64() print(PASSED)INT64 示例与基本示例的差异掩码 dtype 显式指定为pl.DT_INT64以匹配 b64 掩码粒度256 bit 覆盖 32 个元素Tile 形状改为[1, 32]因为 DT_INT64 每个寄存器仅容纳 32 个元素寄存器总容量 256B / 8B输入改用torch.randint(-100, 100, ...)构造整型数据参考结果仍为torch.abs由于整型逐位精确断言精度收紧为rtol0, atol0即要求完全一致。需要注意INT64 取绝对值同样受非饱和截断约束若输入包含 INT64_MIN-2^63其绝对值 2^63 无法用 INT64 表示结果会被截断为 INT64_MIN。示例中数据范围-100~100刻意避开了该边界。使用建议与注意事项先确认设备vf.abs仅在 Ascend 950 系列支持A2/A3 系列请改用 Tile 级接口寄存器与掩码 dtype 对齐建议vf.create_mask的 dtype 与运算寄存器保持一致混用不同粒度时需自行判断结果行为警惕整型边界溢出对 INT8_MIN/INT64_MIN 等边界值取绝对值会发生非饱和截断业务上应先做防溢出处理复数场景的寄存器布局双寄存器模式下reg[0]存实部、reg[1]存虚部数据加载时偶数/奇数索引元素被分别交错存入两个寄存器编写 load/store 时应使用与reg_tensor双寄存器搬入模式匹配的 Tile 布局数据量按2 * VL组织控制寄存器用量VF 函数内 RegTensor 上限 32、MaskReg 上限 16超出会落入预留 UB 空间并影响性能循环内运算时注意及时释放或复用寄存器。总结vf.abs是 PyPTO Pro VF 体系中兼具基础与特殊双重属性的运算接口对 DT_INT8/DT_INT16/DT_INT32/DT_FP16/DT_FP32/DT_INT64 执行逐元素取绝对值对 DT_FP16/DT_FP32 则通过双寄存器交错存储实现复数取模。理解其掩码语义、非饱和截断行为与双寄存器数据布局是正确编写昇腾 950 向量内核的前提。本文提供的两个可运行示例覆盖了浮点精度比对与整型逐位精确校验两条验证路径开发者可直接将其改造为自身算子的 VF 计算单元并配合 reg_tensor.md、mask_reg.md 与 MergeMode.md 深入掌握 VF 编程的完整知识。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表