PyPTO vf.abs_sub 接口深度解析寄存器级绝对差计算的实现与实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读vf.abs_sub是 PyPTOParallel Tensor/Tile Operation 编程范式SIMD 编程模型Vector Function简称 VF中用于逐元素计算绝对差Absolute Difference的寄存器级向量计算接口。它适用于需要计算两路数据逐元素相减再取绝对值的场景例如图像差异度量、误差统计、模板匹配等。读完本文你将掌握vf.abs_sub的产品支持范围、函数原型、参数语义、Mask 合并模式MergeMode行为并能直接复用在 Ascend 950 系列产品上经过验证的 FP32 与 INT64 两种完整调用示例同时了解其底层如何被编译为vabsdif硬件指令INT64 场景下如何通过借位链减法序列模拟。产品支持情况vf.abs_sub属于 PyPTO 为 Ascend 950 系列推出的 SIMD 向量指令集能力当前仓库文档明确的产品支持矩阵如下产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持在使用该接口前请确认目标硬件属于 Ascend 950 系列在 Atlas A2/A3 产品上该接口不可用需改用其他等价的逐元素计算组合。功能说明该接口根据谓词寄存器preg的筛选结果将src0与src1逐元素相减后求绝对值计算结果写入目标寄存器dst。其逐 lane通道计算公式为$$ dst_i \lvert src0_i - src1_i \rvert $$即对于preg中处于活跃选中状态的每个元素位执行|src0 - src1|并写入dst的对应位置。这一语义在源码层的 API 声明中也有明确描述参见 python/pypto_pro/language/_vf_api.py其中注明该接口映射到硬件vabsdif指令。函数原型abs_sub(src0, src1, preg, mode: Optional[MergeMode] None) - dst该接口在 IR 层注册为vf.abs_sub属于VFOp操作类别共包含 4 个操作数dst、src0、src1、mask并携带mode属性注册定义见 framework/src/interface/ir/op/vf_ops.cppREGISTER_OP(vf.abs_sub) .set_op_category(VFOp) .set_description(Absolute difference |src0 - src1|) .add_argument(dst, Destination register) .add_argument(src0, Source register 0) .add_argument(src1, Source register 1) .add_argument(mask, Mask register) .set_attrint(mode) .f_deduce_type(DeduceVFFromDstArg);其中f_deduce_type(DeduceVFFromDstArg)表示结果类型由目标寄存器参数推断这与返回值数据类型与src0一致的接口约定相呼应。参数说明参数输入/输出说明src0输入源操作数 0reg_tensor支持的数据类型为DT_FP16、DT_FP32、DT_INT64。src1输入源操作数 1reg_tensor支持的数据类型和 src0 中的说明一致。preg输入mask_reg。mode输入可选对应 MergeMode 类型。- pypto_pro.language.MergeMode.ZEROING默认preg 未筛选的元素在 dst 中置 0。- pypto_pro.language.MergeMode.MERGING 当前不支持。关于 MergeMode 的补充说明MergeMode 定义了 VF 计算指令中 mask 未选中元素非活跃元素在目标寄存器中的处理方式class MergeMode(enum.Enum): ZEROING ... # mask未选中位置置零默认 MERGING ... # mask未选中位置保留目标寄存器原值对于vf.abs_sub当前仅支持ZEROING模式默认值即未选中的元素在dst中被清零MERGING模式尚未开放。这一点在代码生成后端也有强制校验见 framework/src/interface/pypto_pro/backend/backend_cce_vf_ops.cpp 中的VFZeroingOnly(op, vf.abs_sub)调用。约束说明无。vf.abs_sub不要求额外对齐、边界或特殊内存约束仅需满足上文数据类型与 MergeMode 的限制。返回值说明返回dst目标 reg_tensor支持的数据类型和src0中的说明一致DT_FP16、DT_FP32、DT_INT64。从 IR 层实现看dst、src0、src1三者类型必须完全一致代码生成阶段会对此进行强校验若类型不一致会抛出形如vf.abs_sub requires dst, src0, src1 to have the same type的错误见 backend_cce_vf_ops.cpp。调用示例基本调用示例FP32以下示例在pl.vector_function中创建全活跃 mask加载两路 FP32 数据执行vf.abs_sub后将结果写回并在 Host 侧与torch.abs(a - b)对齐校验import os import pypto_pro.language as pl import torch import torch_npu pl.vector_function def example_vf(src_a, src_b, dst_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32) reg_a vf.load_align(src_a, 0) reg_b vf.load_align(src_b, 0) reg_out vf.abs_sub(reg_a, reg_b, preg) vf.store_align(dst_tile, reg_out, preg) pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], b: 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() in_b_grp pl.make_tile_group(typetf, addrs0x100, mutex_ids[1]) in_b in_b_grp.current() t_out_grp pl.make_tile_group(typetf, addrs0x200, mutex_ids[2]) t_out t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) pl.load(in_b, b, [0, 0]) example_vf(in_a, in_b, 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) b 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 - b), rtol1e-5, atol1e-5) if __name__ __main__: test_example() print(PASSED)要点拆解vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32)生成全活跃谓词 mask其数据类型与参与计算的寄存器保持一致pl.load/pl.store在全局张量与 Vector 内存空间的 Tile 之间搬运数据vf.load_align/vf.store_align完成对齐加载与存储Host 侧通过torch.npu.set_device切换到目标 NPUexample_kernel[None, core_nums]指定 1 核启动使用torch.testing.assert_close以rtol1e-5, atol1e-5的容差做数值对齐验证。INT64 数据类型示例vf.abs_sub支持 INT64 整型运算适用于误差取整、整数像素差等无需浮点精度的场景import os import pypto_pro.language as pl import torch import torch_npu pl.vector_function def example_vf_int64(src_a, src_b, dst_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_INT64) reg_a vf.load_align(src_a, 0) reg_b vf.load_align(src_b, 0) reg_out vf.abs_sub(reg_a, reg_b, 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], b: 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() in_b_grp pl.make_tile_group(typetf, addrs256, mutex_ids[1]) in_b in_b_grp.current() t_out_grp pl.make_tile_group(typetf, addrs512, mutex_ids[2]) t_out t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) pl.load(in_b, b, [0, 0]) example_vf_int64(in_a, in_b, 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) b 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 - b), rtol0, atol0) if __name__ __main__: test_example_int64() print(PASSED)要点拆解使用torch.randint(-100, 100, ...)构造包含正负数的随机整数输入可覆盖符号处理路径Tile 形状调整为[1, 32]寄存器地址按 INT64 位宽换算0 / 256 / 512整型结果使用rtol0, atol0的严格相等校验与浮点示例的容差校验形成对比。底层实现原理从 vf.abs_sub 到 vabsdifFP16 / FP32直接映射硬件指令对于 FP16 与 FP32vf.abs_sub在 CCE 代码生成阶段直接发射硬件vabsdif指令。核心实现位于 framework/src/interface/pypto_pro/backend/backend_cce_vf_ops.cpp其关键逻辑如下static std::string EmitVFAbsSub(const ir::CallPtr op, codegen::CodegenBase codegen_base) { auto codegen dynamic_castcodegen::CCECodegen(codegen_base); CHECK(op-args_.size() 4) vf.abs_sub requires 4 args (dst, src0, src1, mask); DataType s0_dt GetExprDtype(op-args_[1]); CHECK((s0_dt DataType::FP16 || s0_dt DataType::FP32 || s0_dt DataType::INT64)) vf.abs_sub src0 only supports FP16/FP32/INT64, got DTypeStr(s0_dt); ... std::string mode VFZeroingOnly(op, vf.abs_sub); ... }生成逻辑同时承担了两类校验职责操作数个数校验vf.abs_sub必须恰好携带 4 个参数dst、src0、src1、mask否则直接报错数据类型白名单校验src0 / src1 仅允许 FP16、FP32、INT64 三种类型且 dst 必须与两个源操作数同类型。这两层校验与 vf_ops.cpp 中注册的 4 参数签名共同构成了接口的类型安全边界。该行为由单元测试锁定在 framework/tests/ut/interface/src/pypto_pro/backend/test_backend_cce_vf_ops.cpp 中测试通过expect_binary(vf.abs_sub, vabsdif()断言 FP 场景下生成的指令前缀为vabsdif(。INT64借位链减法序列模拟INT64b64场景下由于硬件vabsdif没有原生的 b64 单寄存器重载编译器将其展开为一段多指令模拟序列思路为|src0 - src1| abs(vsub(src0, src1))具体步骤见 backend_cce_vf_ops.cpp使用vdintlv将两个 INT64 源寄存器分别拆分为 b32 低半/高半通过vsubc/vsubcs借位链完成 64 位减法borrow-chain sub得到diff {lo, hi}用vcmp_lt对高半做有符号符号位检查hi 0判定负数对负数通过借位链完成0 - diff的取反再用vsel按符号选择原差或取反结果最后vintlv将高低半重新交织回 b64 目标寄存器并按 ZEROING 语义处理 mask 未选中位。这段序列在单元测试 test_backend_cce_vf_ops.cpp 中通过ExpectInvoke(codegen, vf.abs_sub, ...)被完整断言包括ppack、vdintlv、vsubc、vsel、vintlv等指令的发射顺序。因此INT64 的绝对差计算并非单条硬件指令而是一段受控的指令序列其计算结果在语义上与 FP 路径完全一致。小结vf.abs_sub是 PyPTO SIMD 向量函数库中语义简洁、类型约束明确的逐元素绝对差接口支持 FP16、FP32 与 INT64 三种数据类型当前仅在 Ascend 950PR/950DT 上可用。它对外呈现统一的dst_i |src0_i - src1_i|语义对内则依据数据类型走两条不同的代码生成路径FP 类型直接发射vabsdifINT64 通过借位链减法与符号取反序列模拟。配合 reg_tensor、mask_reg 与 MergeMode 三份关联文档读者可以在 PyPTO 的 Tile 编程框架中将其组合进更复杂的向量算子如误差统计、绝对差求和等流水线。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考