PyPTO vf.neg 向量取反算子深度解析:基于 mask 的寄存器级取反运算实战
发布时间:2026/9/20 3:33:29 作者:尧图编辑部 阅读量:1,286

PyPTO vf.neg 向量取反算子深度解析基于 mask 的寄存器级取反运算实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptovf.neg是 PyPTO 向量函数Vector FunctionVF体系中基础算术basic_arithmetic模块提供的逐元素取反接口用于在向量寄存器层面根据掩码对源数据执行dst_i -src_i运算。本文围绕该算子的产品支持范围、函数原型、mask 驱动语义、MergeMode 行为以及 FP32/INT64 两种完整可运行示例展开讲解并辅以源码与测试佐证帮助读者在 Ascend 950 系列平台上快速编写、验证基于vf.neg的向量算子。产品支持情况vf.neg属于 SIMD-API 中寄存器计算reg_computation类接口其支持范围与所在的 PyPTO 向量函数体系保持一致具体如下产品形态支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持这意味着当前版本的vf.neg仅面向 Ascend 950 系列硬件提供向量寄存器取反能力在使用前应确认目标设备型号。功能说明mask 驱动的逐元素取反vf.neg对输入寄存器src中的每个元素执行算术取反并将结果写入目的寄存器dst计算公式为$$dst_i -src_i$$运算的逐元素有效性由掩码寄存器preg控制mask 位为 1有效对应元素参与取反运算结果写入dst对应位置mask 位为 0无效对应元素不参与运算dst对应位置根据mode参数决定处理方式默认置 0。这一语义与 PyPTO 向量函数的通用 mask 机制一致。mask 是 VF 计算中元素级有效性的核心控制容器详见 vf.mask_regmask 寄存器总位宽固定为 256 bit其粒度由 dtype 决定例如 FP32 场景下每 4 bit 对应一个元素即一个 mask 可覆盖 64 个 FP32 元素INT64 场景下每 8 bit 对应一个元素覆盖 32 个 INT64 元素。实际使用时mask 的 dtype 一般与配套的数据寄存器一致典型创建方式为vf.create_mask(patternpl.MaskPattern.ALL, dtype...)表示所有元素均参与运算。函数原型neg(src, preg, mode: Optional[MergeMode] None) - dst该接口在 PyPTO Python 语言层的声明位于 python/pypto_pro/language/_vf_api.py其文档字符串明确描述对于每个 mask 位激活的通道lanei计算src[i]的算术取反并写入dst[i]。它是pypto_pro.language中vf命名空间下注册的向量函数指令之一只能在pl.vector_function修饰的向量函数内调用。参数说明参数输入/输出说明src输入源操作数类型为 reg_tensor向量寄存器VF 计算的基本数据容器源操作数 src 与目的操作数 dst 的数据类型保持一致。支持的数据类型为DT_INT8、DT_INT16、DT_FP16、DT_INT32、DT_FP32、DT_INT64。preg输入mask_reg掩码寄存器控制哪些元素参与取反运算。mode输入可选对应 MergeMode 类型。-pypto_pro.language.MergeMode.ZEROING默认preg 未筛选的元素在 dst 中置 0。-pypto_pro.language.MergeMode.MERGING当前不支持。关于数据类型可以从 vf.reg_tensor 的寄存器模型进一步理解寄存器总大小固定为 256 字节不同 dtype 决定单个寄存器的元素个数例如 DT_FP32 为 64 个元素、DT_INT64 为 32 个元素。因此vf.neg在一次调用中处理的元素数量与 dtype 强相关编写算子时 Tile 的 shape 应与寄存器容量匹配如 FP32 取[1, 64]INT64 取[1, 32]。MergeMode 语义详解mode参数的类型 MergeMode 在源码中定义为一个枚举类class MergeMode(enum.Enum): ZEROING ... # mask未选中位置置零默认 MERGING ... # mask未选中位置保留目标寄存器原值ZEROING默认mask 未选中位为 0的元素其目的寄存器对应位置被置 0。这是vf.neg的默认行为也是当前唯一支持的模式MERGINGmask 未选中位置保留目标寄存器原值。vf.neg当前不支持该模式。在vf.neg中mode参数可省略传None即采用默认 ZEROING 行为。由于当前仅支持 ZEROING实际编写代码时可省略该参数语义上等价于全量元素参与运算mask 未覆盖位置输出 0。约束说明vf.neg本身无额外约束。需要留意的是其依赖的数据容器约束这些约束在使用时会一并生效寄存器在pypto_pro.language.vector_function函数内创建和使用函数结束后自动释放创建寄存器后必须通过vf.load_align或vf.full初始化数据否则内容未定义数据寄存器RegTensor数量上限为 32mask 寄存器上限为 16编译器会自动复用生命周期结束的寄存器FP8/FP4 类型为存储类型不支持直接参与算术运算因此不在vf.neg的支持类型列表中需先通过vf.astype转换为 FP32/BF16/FP16 等计算类型。返回值说明返回 dst 目的操作数类型为 reg_tensor支持的数据类型与 src 中的说明一致DT_INT8、DT_INT16、DT_FP16、DT_INT32、DT_FP32、DT_INT64。该寄存器由编译器在赋值形式中自动声明例如示例中的reg_out vf.neg(reg_a, preg)。调用示例基本调用示例DT_FP32以下示例展示了一个完整的取反算子在向量函数中创建全量 mask、从 UB Tile 对齐加载数据到寄存器、执行vf.neg取反再将对齐存储回目的 Tile最后在 kernel 中编排 Tile 的搬运与计算并用torch.testing.assert_close与-a的结果进行对比验证。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_a vf.load_align(src_tile, 0) reg_out vf.neg(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, -a, rtol1e-5, atol1e-5) if __name__ __main__: test_example() print(PASSED)示例要点说明tf pl.TileType(shape[1, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec)声明 Vector 内存空间的 64 元素 FP32 Tile与单个 FP32 寄存器容量64 元素对齐pl.make_tile_group(typetf, addrs0x0, mutex_ids[0])以固定地址创建 Tile 组mutex_ids用于约束共享内存互斥pl.section_vector()将向量搬移与 VF 计算放入向量指令段执行输入数据采用torch.randn验证逻辑为out -aFP32 场景使用 rtol/atol 容差比较。INT64 数据类型示例vf.neg同样支持 DT_INT64 数据示例中寄存器数量随 dtype 变化INT64 元素宽度 64 bit单个寄存器仅容纳 32 个元素因此 Tile shape 调整为[1, 32]验证时使用整数精确相等比较rtol0, atol0。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.neg(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, -a, rtol0, atol0) if __name__ __main__: test_example_int64() print(PASSED)该示例在仓库的向量函数测试体系中也有同类用法可供参考例如 python/tests/st/pypto_pro/frontend/vf_api/test_vf_basic_ops.py 中通过vf.neg(reg_a, preg)构造寄存器级取反运算验证vf.neg在 VF 前端下的行为此外 python/tests/st/pypto_pro/frontend/element_wise/test_abs_addc_and.py 与 python/tests/st/pypto_pro/frontend/element_wise/test_vector_operations.py 展示了 Tile 级pl.neg的调用与验证方式可作为对比参考。运行前提两个示例的宿主代码test_example/test_example_int64展示了 PyPTO VF 算子的标准验证流程通过环境变量TILE_FWK_DEVICE_ID默认 0选择 NPU 设备调用torch.npu.set_device(device)设置当前设备以example_kernelNone, core_nums形式启动 kernelNone表示同步启动core_nums1指定使用 1 个 AI Core使用torch.npu.synchronize()等待设备侧计算完成用torch.testing.assert_close与 PyTorch 的-a结果对比校验通过后打印PASSED。运行前需确保环境已安装torch、torch_npu以及本仓库的pypto_pro语言运行时且设备为支持 Ascend 950 系列 NPU 的推理/训练环境。整体向量函数编程范式与 Tile 组、Vector section 等概念的入门介绍可参考 PyPTO 的快速上手文档 docs/zh/guide/quick_start/pro 与编程指南 docs/zh/guide/programming_guide/pro。小结vf.neg是 PyPTO 向量函数体系中实现元素级取反的最小算子单元其价值体现在与 mask 机制、寄存器模型和 Tile 编排的深度协同通过create_mask控制参与运算的元素范围通过load_align/store_align完成 UB 与寄存器间的数据流转再配合section_vector将整个取反流水编排进向量指令段。无论是 FP32 浮点取反还是 INT64 整数取反读者都可以基于本文的完整示例直接修改 dtype、Tile shape 与地址参数快速移植到自己的算子中并借助torch.testing.assert_close完成正确性验证。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考