PTO 逐元素倒数平方根指令 TRSQRT 详解从数学定义到 A5 向量内核实现【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa导读TRSQRT 是 Ascend CANN 开源仓库 pto-isa 中基于 Parallel Tile OperationPTO虚拟指令集定义的逐元素倒数平方根Reciprocal Square Roottile 级运算指令计算dst 1 / sqrt(src)。本文以 TRSQRT 指令规范文档 为骨架结合仓库内 A5、A2A3、CPU 多平台实现源码与 ST 测试用例完整讲解其数学语义、两级汇编语法、C Intrinsic 接口、约束校验规则、精度路径选择与临时空间tmp的真实行为。读完本文你将能够使用 C Intrinsic 或汇编形式编写 TRSQRT 算子理解默认精度与高精度实现的内在差异并知晓如何通过仓库测试用例验证结果正确性。TRSQRT 的语义与数学定义TRSQRT 是 PTO 指令集中的一种逐元素Elementwise一元运算作用于向量Vector位置的 tile。它对有效区域valid region内的每个元素独立计算其倒数平方根。对于有效区域内的每个元素(i, j)运算满足$$\mathrm{dst}{i,j} \frac{1}{\sqrt{\mathrm{src}{i,j}}}$$即先对输入取平方根再求其倒数。该运算在归一化如 RMSNorm、LayerNorm 中的1/sqrt(x)缩放项、信号处理、数值算法中都是高频基础算子。TRSQRT 与仓库中的平方根指令 TSQRT 不同TSQRT 只计算sqrt(src)而 TRSQRT 在其基础上增加一次倒数运算文档明确其默认精度实现由vsqrt与vdiv两条底层向量指令组合完成。需要特别说明的是指令文档指出其定义域/NaN 行为是目标平台相关的target-defined例如当src 0除零或输入为负数负数的平方根无实数结果时不同硬件的具体表现可能不同编写算子时应在上层做好输入约束。汇编语法两级抽象形式TRSQRT 在 PTO 指令集中以三种形式出现分别对应不同开发模式与汇编层级同步形式PTO Assembly%dst trsqrt %src : !pto.tile...AS Level 1SSA 形式用于 Auto 模式%dst pto.trsqrt %src : !pto.tile... - !pto.tile...AS Level 2DPS 形式显式绑定 tile bufferpto.trsqrt ins(%src : !pto.tile_buf...) outs(%dst : !pto.tile_buf...)其中!pto.tile...是 PTO 虚拟 ISA 中表示 tile 类型的方言类型...处可展开为 tile 的形状、数据类型与布局等属性。Level 1SSA适合 Auto 模式——编译器/运行时负责 tile 的放置与调度Level 2DPS则显式声明ins/outs的tile_buf适用于需要精确控制资源绑定的手动开发场景。仓库 PTO-Virtual-ISA-Manual.md 对这两级汇编抽象有系统性描述。C Intrinsic 接口在应用开发中更常用的是 C Intrinsic。TRSQRT 在 include/pto/common/pto_instr.hpp 中声明了两个重载均以RsqrtAlgorithm模板参数控制精度策略默认RsqrtAlgorithm::DEFAULTtemplate auto PrecisionType RsqrtAlgorithm::DEFAULT, typename TileDataDst, typename TileDataSrc, typename... WaitEvents, std::enable_if_tall_events_vWaitEvents..., int 0 PTO_INST RecordEvent TRSQRT(TileDataDst dst, TileDataSrc src, WaitEvents ... events); template auto PrecisionType RsqrtAlgorithm::DEFAULT, typename TileDataDst, typename TileDataSrc, typename TileDataTmp, typename... WaitEvents, std::enable_if_tis_tile_data_vTileDataTmp all_events_vWaitEvents..., int 0 PTO_INST RecordEvent TRSQRT(TileDataDst dst, TileDataSrc src, TileDataTmp tmp, WaitEvents ... events);两个重载的使用要点2 参数重载TRSQRT(dst, src, events...)基本形式通过TRSQRT_IMPLPrecisionType(dst, src)下发计算3 参数重载TRSQRT(dst, src, tmp, events...)多接收一个临时 tile内部调用TRSQRT_IMPLPrecisionType(dst, src, tmp)为高精度路径预留接口两者都返回RecordEvent可通过EventOp::TRSQRT, Op::...与前后指令如 TLOAD、TSTORE构成事件依赖链实现异步流水调度变参WaitEvents受all_events_vWaitEvents...约束确保只接受合法事件类型。精度枚举RsqrtAlgorithm定义于 include/pto/common/type.hppenum class RsqrtAlgorithm : uint8_t { DEFAULT, HIGH_PRECISION };约束与编译期校验TRSQRT 对 tile 的形态有严格约束这些约束在 NPU 实现中以static_assert编译期与PTO_ASSERT运行期双重方式执行。核心规则如下约束类别规则检查时机数据类型DType必须为float或half编译期Tile 位置TileData::Loc TileType::Vec必须在向量单元编译期布局必须为行主序TileData::isRowMajor编译期静态有效边界ValidRow Rows且ValidCol Cols编译期运行时形状匹配src.GetValidRow() dst.GetValidRow()且src.GetValidCol() dst.GetValidCol()运行期迭代域以dst.GetValidRow()/dst.GetValidCol()为迭代范围运行期以 include/pto/npu/a5/TRsqrt.hpp 中的TRSQRT_IMPL为例可以看到一组完整的静态断言任何一个不满足都会在编译阶段直接报错static_assert(DstTile::isRowMajor SrcTile::isRowMajor, TRSQRT: Not supported Layout type); static_assert(DstTile::Loc TileType::Vec SrcTile::Loc TileType::Vec, TRSQRT: TileType of src and dst tiles must be TileType::Vec.); static_assert(DstTile::ValidCol DstTile::Cols, TRSQRT: Number of dsts valid columns must not be greater than number of tile columns.); static_assert(DstTile::ValidRow DstTile::Rows, TRSQRT: Number of dsts valid rows must not be greater than number of tile rows.); static_assert(std::is_same_vtypename DstTile::DType, typename SrcTile::DType, TRSQRT: The data type of dst must be consistent with of src); static_assert(std::is_same_vtypename DstTile::DType, float32_t || std::is_same_vtypename DstTile::DType, float || std::is_same_vtypename DstTile::DType, float16_t || std::is_same_vtypename DstTile::DType, half, TRSQRT: Invalid data type.);运行期则通过PTO_ASSERT检查 src/dst 的有效行列是否一致随后读取dst.GetValidRow()与dst.GetValidCol()作为实际计算域。注意迭代域以dst的 valid 区域为准文档明确说明因此保证 src 与 dst valid 区域一致是正确性的前提。临时空间tmp的真实行为A5 上的接口兼容设计这是 TRSQRT 一个容易被误用的设计点。指令文档明确说明了tmp参数在不同平台上的真实行为2 参数重载无 tmp不需要临时空间默认精度实现直接用vsqrtvdiv两条向量指令完成先开方、再取倒数。3 参数重载带 tmp接口接收tmp但当前 A5 实现并未使用它。从 include/pto/npu/a5/TRsqrt.hpp 可以看到3 参数版本直接委托给 2 参数实现template auto PrecisionType RsqrtAlgorithm::DEFAULT, typename DstTile, typename SrcTile, typename TmpTile PTO_INTERNAL void TRSQRT_IMPL(DstTile dst, SrcTile src, TmpTile tmp) { TRSQRT_IMPLPrecisionType(dst, src); }tmp之所以保留在 C Intrinsic 签名中是为 API 兼容性以及未来潜在的高精度路径预留。CPU 实现include/pto/cpu/TRSqrt.hpp同样采取委托策略。值得注意的平台差异A2A3 平台的实现并不相同。在 include/pto/npu/a2a3/TUnaryOp.hpp 中3 参数重载会真正调用TRsqrtHighPrecision——先用vector_dup把tmp初始化为 1.0再逐行执行vsqrt经pipe_barrier(PIPE_V)同步后执行vdiv完成1/sqrt(x)。因此如果为 A2A3 平台使用高精度路径tmp是实际会被使用的临时缓冲而 A5 上当前无论是否传入tmp计算路径完全一致。精度路径DEFAULT 与 HIGH_PRECISION 的实现差异从 include/pto/npu/a5/TRsqrt.hpp 的 1D 内核可以看出 A5 上两种精度策略的底层差异DEFAULT默认精度对每个向量块执行vsqrt(tmpReg, srcReg, pReg, MODE_ZEROING)后再执行vdiv(dstReg, oneReg, tmpReg, pReg)其中oneReg通过vdup预加载为常量 1.0HIGH_PRECISION高精度对float走SqrtFloatImplDivIEEE754FloatImpl对half走SqrtPrecisionImplDivIEEE754HalfImpl即使用符合 IEEE 754 语义的平方根与除法实现相关实现见 include/pto/npu/a5/custom/TSqrtHp.hpp 与 include/pto/npu/a5/custom/Div754.hpp以获得更精确的结果。指令的循环次数由有效元素总数与单次 repeat 元素数决定nRepeatElem CCE_VL / sizeof(T)向量处理单元一个 repeat 能处理的元素个数repeat 次数为CeilDivision(validRow * validCol, nRepeatElem)。A2A3 平台的默认路径则直接使用硬件近似倒数平方根指令vrsqrt见 include/pto/npu/a2a3/TUnaryOp.hpp 中的RsqrtOp高精度路径才改用vsqrtvdiv组合。这体现了不同硬件代际在指令级实现上的差异。使用示例Auto 与 Manual 模式Auto 模式Auto 模式下 tile 的放置与调度由编译器/运行时管理只需声明TileTileType::Vec, T, Rows, Cols并直接调用 Intrinsic#include pto/pto-inst.hpp using namespace pto; void example_auto() { using TileT TileTileType::Vec, float, 16, 16; TileT src, dst; TRSQRT(dst, src); }Manual 模式Manual 模式需要先用TASSIGN为 tile 显式绑定地址UB 缓冲再发起计算#include pto/pto-inst.hpp using namespace pto; void example_manual() { using TileT TileTileType::Vec, float, 16, 16; TileT src, dst; TASSIGN(src, 0x1000); TASSIGN(dst, 0x2000); TRSQRT(dst, src); }对应的汇编形式中Manual 模式在指令前通过pto.tassign完成资源绑定tile 操作数可省略但显式绑定更清晰# Manual mode: resources must be bound explicitly before issuing the instruction. # Optional for tile operands: # pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) %dst pto.trsqrt %src : !pto.tile... - !pto.tile...Auto 模式则由编译器管理放置与调度直接发射 SSA 形式的pto.trsqrt。源码级内核实现解析A5include/pto/npu/a5/TRsqrt.hpp 根据 tile 形态自动选择不同的内核变体。顶层分发逻辑TRsqrt先根据编译期信息判断若ValidCol Cols整行有效或行数为 1纯一维数据则走 1D 路径否则走 2D 行分片路径。1D 路径TRsqrt_1D_Switch按VFImplKind再细分VFIMPL_1D_NO_POST_UPDATEvlds/vsts使用绝对偏移i * nRepeatElem不更新基址VFIMPL_2D_POST_UPDATE/VFIMPL_2D_NO_POST_UPDATE映射到 2D 内核按行 stride 寻址默认TRsqrt_1D_PostUpdatevlds/vsts带POST_UPDATE标志每次 repeat 后自动累加基址。2D 路径TRsqrt_2D对外层行循环、内层列分块循环每次加载src i * SrcRowStride j * nRepeatElem处的数据块计算后存储到dst i * DstRowStride j * nRepeatElem行与行之间通过编译期常量RowStride跳转天然适配行主序 tile 的 strided 布局。VFImplKind枚举及后续向量指令下发机制可在 include/pto/npu/a5/vf/vf_defs.hpp 与 include/pto/npu/a5/vf/vf_common.hpp 中进一步追溯。CPU 侧参考实现include/pto/cpu/TRSqrt.hpp则把每个元素提升为double计算1.0 / std::sqrt(x)后再窄化回原类型并通过cpu::parallel_for_rows按行并行是跨平台正确性验证如 tests/cpu/st/testcase/trsqrt的黄金参照。测试验证A5 ST 用例与精度阈值仓库为 TRSQRT 提供了完整的 STSystem Test验证A5 平台用例位于 tests/npu/a5/src/st/testcase/trsqrt包含三个文件trsqrt_kernel.cpp内核定义通过EventOp::TLOAD, Op::TRSQRT与EventOp::TRSQRT, Op::TSTORE_VEC构建TLOAD → TRSQRT → TSTORE事件链并以模板参数highPrecision、isInPlace控制精度与原地/非原地模式main.cppgtest 驱动读取input.bin上板执行后与golden.bin比对gen_data.py生成输入与黄金数据。测试覆盖了 8 个代表性用例tests/npu/a5/src/st/testcase/trsqrt/main.cpp用例类型形态dst/srcvalid模式case1float64x64 / 64x6464x64高精度 原地case2float64x64 / 64x6464x64高精度 非原地case3half64x64 / 64x6464x64高精度 原地case4half64x64 / 64x6464x64高精度 非原地case5float128x128 / 64x6464x64默认 dst 大于 srccase6float64x64 / 128x12832x32默认 src 大于 dstcase7half128x256 / 64x6464x64默认 dst 大于 srccase8half64x64 / 128x25632x32默认 src 大于 dstcase5~case8 特意让 dst 与 src 的物理 tile 尺寸不同、valid 区域小于物理尺寸用以验证有效边界与 stride 寻址的正确性。精度容差tests/npu/a5/src/st/testcase/trsqrt/main.cpp按类型与精度分级float默认eps 0.00005fhalf默认eps 0.0005f高精度模式下收紧到eps 0.0000001f——这从测试侧印证了 HIGH_PRECISION 路径IEEE 754 除法 高精度开方的精度优势。除 A5 外TRSQRT 在 A2A3tests/npu/a2a3/src/st/testcase/trsqrt、kirin9030tests/npu/kirin9030/src/st/testcase/trsqrt、kirinDev0000tests/npu/kirinDev0000/src/st/testcase/trsqrt以及 CPUtests/cpu/st/testcase/trsqrt平台均有对应的内核与测试用例体现了该指令跨平台cross-platform的一致性设计代价模型侧include/pto/costmodel/a5/vf_costmodel.hpp 也将PtoOpcode::TRSQRT映射到TRSQRT名称供性能仿真与开销估算使用参见 docs/costmodel/perf-sim-user-guide.md。总结TRSQRT 是 PTO 指令集中实现1/sqrt(x)的基础逐元素向量指令文档规范、C Intrinsic、多平台内核与测试用例共同构成了完整的技术闭环。理解它的关键在于三点一是其两级汇编语法SSA/DPS与 Auto/Manual 两种开发模式的对应关系二是tmp参数在不同平台A5 预留兼容 vs A2A3 高精度实际使用的行为差异三是 DEFAULT 与 HIGH_PRECISION 两条精度路径vsqrtvdiv组合 vs IEEE 754 高精度实现在精度与开销上的权衡。当你需要编写涉及归一化、数值缩放等场景的 PTO 算子时可直接参考本文示例与 TRSQRT 指令规范并借助仓库测试用例完成端到端验证。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/19 10:53:57 · 阅读 1.2k PTO 逐元素倒数平方根指令 TRSQRT 详解：从数学定义到 A5 向量内核实现
PTO 逐元素倒数平方根指令 TRSQRT 详解从数学定义到 A5 向量内核实现【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa导读TRSQRT 是 Ascend CANN 开源仓库 pto-isa 中基于 Parallel Tile OperationPTO虚拟指令集定义的逐元素倒数平方根Reciprocal Square Roottile 级运算指令计算dst 1 / sqrt(src)。本文以 TRSQRT 指令规范文档 为骨架结合仓库内 A5、A2A3、CPU 多平台实现源码与 ST 测试用例完整讲解其数学语义、两级汇编语法、C Intrinsic 接口、约束校验规则、精度路径选择与临时空间tmp的真实行为。读完本文你将能够使用 C Intrinsic 或汇编形式编写 TRSQRT 算子理解默认精度与高精度实现的内在差异并知晓如何通过仓库测试用例验证结果正确性。TRSQRT 的语义与数学定义TRSQRT 是 PTO 指令集中的一种逐元素Elementwise一元运算作用于向量Vector位置的 tile。它对有效区域valid region内的每个元素独立计算其倒数平方根。对于有效区域内的每个元素(i, j)运算满足$$\mathrm{dst}{i,j} \frac{1}{\sqrt{\mathrm{src}{i,j}}}$$即先对输入取平方根再求其倒数。该运算在归一化如 RMSNorm、LayerNorm 中的1/sqrt(x)缩放项、信号处理、数值算法中都是高频基础算子。TRSQRT 与仓库中的平方根指令 TSQRT 不同TSQRT 只计算sqrt(src)而 TRSQRT 在其基础上增加一次倒数运算文档明确其默认精度实现由vsqrt与vdiv两条底层向量指令组合完成。需要特别说明的是指令文档指出其定义域/NaN 行为是目标平台相关的target-defined例如当src 0除零或输入为负数负数的平方根无实数结果时不同硬件的具体表现可能不同编写算子时应在上层做好输入约束。汇编语法两级抽象形式TRSQRT 在 PTO 指令集中以三种形式出现分别对应不同开发模式与汇编层级同步形式PTO Assembly%dst trsqrt %src : !pto.tile...AS Level 1SSA 形式用于 Auto 模式%dst pto.trsqrt %src : !pto.tile... - !pto.tile...AS Level 2DPS 形式显式绑定 tile bufferpto.trsqrt ins(%src : !pto.tile_buf...) outs(%dst : !pto.tile_buf...)其中!pto.tile...是 PTO 虚拟 ISA 中表示 tile 类型的方言类型...处可展开为 tile 的形状、数据类型与布局等属性。Level 1SSA适合 Auto 模式——编译器/运行时负责 tile 的放置与调度Level 2DPS则显式声明ins/outs的tile_buf适用于需要精确控制资源绑定的手动开发场景。仓库 PTO-Virtual-ISA-Manual.md 对这两级汇编抽象有系统性描述。C Intrinsic 接口在应用开发中更常用的是 C Intrinsic。TRSQRT 在 include/pto/common/pto_instr.hpp 中声明了两个重载均以RsqrtAlgorithm模板参数控制精度策略默认RsqrtAlgorithm::DEFAULTtemplate auto PrecisionType RsqrtAlgorithm::DEFAULT, typename TileDataDst, typename TileDataSrc, typename... WaitEvents, std::enable_if_tall_events_vWaitEvents..., int 0 PTO_INST RecordEvent TRSQRT(TileDataDst dst, TileDataSrc src, WaitEvents ... events); template auto PrecisionType RsqrtAlgorithm::DEFAULT, typename TileDataDst, typename TileDataSrc, typename TileDataTmp, typename... WaitEvents, std::enable_if_tis_tile_data_vTileDataTmp all_events_vWaitEvents..., int 0 PTO_INST RecordEvent TRSQRT(TileDataDst dst, TileDataSrc src, TileDataTmp tmp, WaitEvents ... events);两个重载的使用要点2 参数重载TRSQRT(dst, src, events...)基本形式通过TRSQRT_IMPLPrecisionType(dst, src)下发计算3 参数重载TRSQRT(dst, src, tmp, events...)多接收一个临时 tile内部调用TRSQRT_IMPLPrecisionType(dst, src, tmp)为高精度路径预留接口两者都返回RecordEvent可通过EventOp::TRSQRT, Op::...与前后指令如 TLOAD、TSTORE构成事件依赖链实现异步流水调度变参WaitEvents受all_events_vWaitEvents...约束确保只接受合法事件类型。精度枚举RsqrtAlgorithm定义于 include/pto/common/type.hppenum class RsqrtAlgorithm : uint8_t { DEFAULT, HIGH_PRECISION };约束与编译期校验TRSQRT 对 tile 的形态有严格约束这些约束在 NPU 实现中以static_assert编译期与PTO_ASSERT运行期双重方式执行。核心规则如下约束类别规则检查时机数据类型DType必须为float或half编译期Tile 位置TileData::Loc TileType::Vec必须在向量单元编译期布局必须为行主序TileData::isRowMajor编译期静态有效边界ValidRow Rows且ValidCol Cols编译期运行时形状匹配src.GetValidRow() dst.GetValidRow()且src.GetValidCol() dst.GetValidCol()运行期迭代域以dst.GetValidRow()/dst.GetValidCol()为迭代范围运行期以 include/pto/npu/a5/TRsqrt.hpp 中的TRSQRT_IMPL为例可以看到一组完整的静态断言任何一个不满足都会在编译阶段直接报错static_assert(DstTile::isRowMajor SrcTile::isRowMajor, TRSQRT: Not supported Layout type); static_assert(DstTile::Loc TileType::Vec SrcTile::Loc TileType::Vec, TRSQRT: TileType of src and dst tiles must be TileType::Vec.); static_assert(DstTile::ValidCol DstTile::Cols, TRSQRT: Number of dsts valid columns must not be greater than number of tile columns.); static_assert(DstTile::ValidRow DstTile::Rows, TRSQRT: Number of dsts valid rows must not be greater than number of tile rows.); static_assert(std::is_same_vtypename DstTile::DType, typename SrcTile::DType, TRSQRT: The data type of dst must be consistent with of src); static_assert(std::is_same_vtypename DstTile::DType, float32_t || std::is_same_vtypename DstTile::DType, float || std::is_same_vtypename DstTile::DType, float16_t || std::is_same_vtypename DstTile::DType, half, TRSQRT: Invalid data type.);运行期则通过PTO_ASSERT检查 src/dst 的有效行列是否一致随后读取dst.GetValidRow()与dst.GetValidCol()作为实际计算域。注意迭代域以dst的 valid 区域为准文档明确说明因此保证 src 与 dst valid 区域一致是正确性的前提。临时空间tmp的真实行为A5 上的接口兼容设计这是 TRSQRT 一个容易被误用的设计点。指令文档明确说明了tmp参数在不同平台上的真实行为2 参数重载无 tmp不需要临时空间默认精度实现直接用vsqrtvdiv两条向量指令完成先开方、再取倒数。3 参数重载带 tmp接口接收tmp但当前 A5 实现并未使用它。从 include/pto/npu/a5/TRsqrt.hpp 可以看到3 参数版本直接委托给 2 参数实现template auto PrecisionType RsqrtAlgorithm::DEFAULT, typename DstTile, typename SrcTile, typename TmpTile PTO_INTERNAL void TRSQRT_IMPL(DstTile dst, SrcTile src, TmpTile tmp) { TRSQRT_IMPLPrecisionType(dst, src); }tmp之所以保留在 C Intrinsic 签名中是为 API 兼容性以及未来潜在的高精度路径预留。CPU 实现include/pto/cpu/TRSqrt.hpp同样采取委托策略。值得注意的平台差异A2A3 平台的实现并不相同。在 include/pto/npu/a2a3/TUnaryOp.hpp 中3 参数重载会真正调用TRsqrtHighPrecision——先用vector_dup把tmp初始化为 1.0再逐行执行vsqrt经pipe_barrier(PIPE_V)同步后执行vdiv完成1/sqrt(x)。因此如果为 A2A3 平台使用高精度路径tmp是实际会被使用的临时缓冲而 A5 上当前无论是否传入tmp计算路径完全一致。精度路径DEFAULT 与 HIGH_PRECISION 的实现差异从 include/pto/npu/a5/TRsqrt.hpp 的 1D 内核可以看出 A5 上两种精度策略的底层差异DEFAULT默认精度对每个向量块执行vsqrt(tmpReg, srcReg, pReg, MODE_ZEROING)后再执行vdiv(dstReg, oneReg, tmpReg, pReg)其中oneReg通过vdup预加载为常量 1.0HIGH_PRECISION高精度对float走SqrtFloatImplDivIEEE754FloatImpl对half走SqrtPrecisionImplDivIEEE754HalfImpl即使用符合 IEEE 754 语义的平方根与除法实现相关实现见 include/pto/npu/a5/custom/TSqrtHp.hpp 与 include/pto/npu/a5/custom/Div754.hpp以获得更精确的结果。指令的循环次数由有效元素总数与单次 repeat 元素数决定nRepeatElem CCE_VL / sizeof(T)向量处理单元一个 repeat 能处理的元素个数repeat 次数为CeilDivision(validRow * validCol, nRepeatElem)。A2A3 平台的默认路径则直接使用硬件近似倒数平方根指令vrsqrt见 include/pto/npu/a2a3/TUnaryOp.hpp 中的RsqrtOp高精度路径才改用vsqrtvdiv组合。这体现了不同硬件代际在指令级实现上的差异。使用示例Auto 与 Manual 模式Auto 模式Auto 模式下 tile 的放置与调度由编译器/运行时管理只需声明TileTileType::Vec, T, Rows, Cols并直接调用 Intrinsic#include pto/pto-inst.hpp using namespace pto; void example_auto() { using TileT TileTileType::Vec, float, 16, 16; TileT src, dst; TRSQRT(dst, src); }Manual 模式Manual 模式需要先用TASSIGN为 tile 显式绑定地址UB 缓冲再发起计算#include pto/pto-inst.hpp using namespace pto; void example_manual() { using TileT TileTileType::Vec, float, 16, 16; TileT src, dst; TASSIGN(src, 0x1000); TASSIGN(dst, 0x2000); TRSQRT(dst, src); }对应的汇编形式中Manual 模式在指令前通过pto.tassign完成资源绑定tile 操作数可省略但显式绑定更清晰# Manual mode: resources must be bound explicitly before issuing the instruction. # Optional for tile operands: # pto.tassign %arg0, tile(0x1000) # pto.tassign %arg1, tile(0x2000) %dst pto.trsqrt %src : !pto.tile... - !pto.tile...Auto 模式则由编译器管理放置与调度直接发射 SSA 形式的pto.trsqrt。源码级内核实现解析A5include/pto/npu/a5/TRsqrt.hpp 根据 tile 形态自动选择不同的内核变体。顶层分发逻辑TRsqrt先根据编译期信息判断若ValidCol Cols整行有效或行数为 1纯一维数据则走 1D 路径否则走 2D 行分片路径。1D 路径TRsqrt_1D_Switch按VFImplKind再细分VFIMPL_1D_NO_POST_UPDATEvlds/vsts使用绝对偏移i * nRepeatElem不更新基址VFIMPL_2D_POST_UPDATE/VFIMPL_2D_NO_POST_UPDATE映射到 2D 内核按行 stride 寻址默认TRsqrt_1D_PostUpdatevlds/vsts带POST_UPDATE标志每次 repeat 后自动累加基址。2D 路径TRsqrt_2D对外层行循环、内层列分块循环每次加载src i * SrcRowStride j * nRepeatElem处的数据块计算后存储到dst i * DstRowStride j * nRepeatElem行与行之间通过编译期常量RowStride跳转天然适配行主序 tile 的 strided 布局。VFImplKind枚举及后续向量指令下发机制可在 include/pto/npu/a5/vf/vf_defs.hpp 与 include/pto/npu/a5/vf/vf_common.hpp 中进一步追溯。CPU 侧参考实现include/pto/cpu/TRSqrt.hpp则把每个元素提升为double计算1.0 / std::sqrt(x)后再窄化回原类型并通过cpu::parallel_for_rows按行并行是跨平台正确性验证如 tests/cpu/st/testcase/trsqrt的黄金参照。测试验证A5 ST 用例与精度阈值仓库为 TRSQRT 提供了完整的 STSystem Test验证A5 平台用例位于 tests/npu/a5/src/st/testcase/trsqrt包含三个文件trsqrt_kernel.cpp内核定义通过EventOp::TLOAD, Op::TRSQRT与EventOp::TRSQRT, Op::TSTORE_VEC构建TLOAD → TRSQRT → TSTORE事件链并以模板参数highPrecision、isInPlace控制精度与原地/非原地模式main.cppgtest 驱动读取input.bin上板执行后与golden.bin比对gen_data.py生成输入与黄金数据。测试覆盖了 8 个代表性用例tests/npu/a5/src/st/testcase/trsqrt/main.cpp用例类型形态dst/srcvalid模式case1float64x64 / 64x6464x64高精度 原地case2float64x64 / 64x6464x64高精度 非原地case3half64x64 / 64x6464x64高精度 原地case4half64x64 / 64x6464x64高精度 非原地case5float128x128 / 64x6464x64默认 dst 大于 srccase6float64x64 / 128x12832x32默认 src 大于 dstcase7half128x256 / 64x6464x64默认 dst 大于 srccase8half64x64 / 128x25632x32默认 src 大于 dstcase5~case8 特意让 dst 与 src 的物理 tile 尺寸不同、valid 区域小于物理尺寸用以验证有效边界与 stride 寻址的正确性。精度容差tests/npu/a5/src/st/testcase/trsqrt/main.cpp按类型与精度分级float默认eps 0.00005fhalf默认eps 0.0005f高精度模式下收紧到eps 0.0000001f——这从测试侧印证了 HIGH_PRECISION 路径IEEE 754 除法 高精度开方的精度优势。除 A5 外TRSQRT 在 A2A3tests/npu/a2a3/src/st/testcase/trsqrt、kirin9030tests/npu/kirin9030/src/st/testcase/trsqrt、kirinDev0000tests/npu/kirinDev0000/src/st/testcase/trsqrt以及 CPUtests/cpu/st/testcase/trsqrt平台均有对应的内核与测试用例体现了该指令跨平台cross-platform的一致性设计代价模型侧include/pto/costmodel/a5/vf_costmodel.hpp 也将PtoOpcode::TRSQRT映射到TRSQRT名称供性能仿真与开销估算使用参见 docs/costmodel/perf-sim-user-guide.md。总结TRSQRT 是 PTO 指令集中实现1/sqrt(x)的基础逐元素向量指令文档规范、C Intrinsic、多平台内核与测试用例共同构成了完整的技术闭环。理解它的关键在于三点一是其两级汇编语法SSA/DPS与 Auto/Manual 两种开发模式的对应关系二是tmp参数在不同平台A5 预留兼容 vs A2A3 高精度实际使用的行为差异三是 DEFAULT 与 HIGH_PRECISION 两条精度路径vsqrtvdiv组合 vs IEEE 754 高精度实现在精度与开销上的权衡。当你需要编写涉及归一化、数值缩放等场景的 PTO 算子时可直接参考本文示例与 TRSQRT 指令规范并借助仓库测试用例完成端到端验证。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考