尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

CANN ops-math 之 Ndtri 算子深度解析:标准正态分位数函数的 Cephes 有理逼近实现与 aclnn 调用指南

CANN ops-math 之 Ndtri 算子深度解析:标准正态分位数函数的 Cephes 有理逼近实现与 aclnn 调用指南 CANN ops-math 之 Ndtri 算子深度解析标准正态分位数函数的 Cephes 有理逼近实现与 aclnn 调用指南【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 仓库中 experimental/math/ndtri/README.md 为核心骨架结合 op_kernel、op_host、op_api 与示例测试源码系统讲解 Ndtri 算子的数学定义、分区间计算策略、参数约束、精度标准与 aclnn 调用流程帮助开发者理解其底层实现原理并快速完成算子编译、验证与集成。产品支持情况Ndtri 算子当前仅适配 Ascend 950 系列产品其余主流昇腾系列产品暂不支持。README 中给出的产品支持矩阵如下产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×Kirin xxx×该结论在源码中得到印证op_api/ndtri.cpp 中的IsAiCoreSupport明确要求当前 NPU 架构必须为NpuArch::DAV_3510即 arch35 / Ascend950否则直接返回不支持op_host/ndtri_def.cpp 中也仅通过AICore().AddConfig(ascend950, aiCoreConfig)注册了 ascend950 这一 AI Core 配置。功能说明算子功能与计算公式Ndtri 算子逐元素计算标准正态分布累积分布函数CDF的反函数又称 probit 函数或分位数函数作用是将概率张量映射为对应的标准正态分位点。其数学定义为$$ y \mathrm{ndtri}(p) \Phi^{-1}(p) \sqrt{2} \cdot \mathrm{erf}^{-1}(2p - 1), \quad p \in (0, 1) $$其中 $\Phi$ 为标准正态分布的累积分布函数$\mathrm{erf}^{-1}$ 为逆误差函数。该定义与 PyTorchtorch.special.ndtri、SciPyscipy.special.ndtri、TensorFlowtf.math.ndtri保持一致README 参考资源中列出均为对标基线算法本源来自 Cephes 数学库的ndtri.c。分区间有理逼近策略直接计算 $\mathrm{erf}^{-1}$ 代价高昂工程实现普遍采用分段有理逼近。本算子基于 Cephes 数学库的分区间有理逼近实现设边界常量val_sub e⁻² ≈ 0.1353、res_exp 1 − e⁻² ≈ 0.8647按下表选择计算路径区间条件计算方法中心区$e^{-2} \le p \le 1 - e^{-2}$cal_p0$y \approx \sqrt{2\pi}\cdot((p-0.5) (p-0.5)^2 \cdot P_0/Q_0)$左尾$p e^{-2}$cal_sub cal_p12$x\sqrt{-2\ln p}$$y \approx -(x - \ln x / x - P_{1/2}/Q_{1/2})$右尾$p 1 - e^{-2}$cal_sub cal_p12$x\sqrt{-2\ln(1-p)}$$y \approx (x - \ln x / x - P_{1/2}/Q_{1/2})$从源码 op_kernel/ndtri_coeffs.h 可以看到这套逼近实际使用三组有理分式系数各自覆盖不同的自变量区间中心区P0/Q0以z (p − 0.5)²为变量LIST_P0共 5 项、LIST_Q0共 8 项逼近式为 $y \approx \sqrt{2\pi} \cdot (p-0.5) \cdot (1 z \cdot P_0(z)/Q_0(z))$尾部区P1/Q1以x sqrt(−2·ln(q))为变量覆盖x ∈ [2, 8]q 为 p 或 1−p 中较小的一侧概率LIST_P1共 9 项、LIST_Q1共 8 项极尾区P2/Q2覆盖x 8的极端小概率情形对应p e⁻³²LIST_P2共 9 项、LIST_Q2共 8 项。系数文件中还定义了关键边界常量NDTRI_VAL_SUB 0.1353352832366127fe⁻²、NDTRI_RES_EXP 0.8646647167633873f1 − e⁻²、NDTRI_SQRT_2PI 2.50662827463100050242f√2π、NDTRI_X_BOUNDARY 8.0fP1/Q1 与 P2/Q2 的分界、NDTRI_SAFE_LO 1.1754944e-38fFLT_MIN概率钳制下限。系数文件头部注释说明这些系数源自 Cephes 数学库的 double 精度版本被转换为单精度 float 以适配 NPU arch35 向量流水线精度损失远在 FP32 阈值2⁻¹³ ≈ 1.22e-4之内。特殊值处理特殊值行为对齐 PyTorchtorch.special.ndtri与 Cephes 实现输入 p输出 y0$-\infty$1$\infty$$p 0$ 或 $p 1$$\mathrm{NaN}$$\mathrm{NaN}$ / $\pm\infty$$\mathrm{NaN}$参数说明Ndtri 算子共两个参数均为 ND 格式参数名输入/输出/属性描述数据类型数据格式self输入公式中的输入概率张量p期望元素取值范围p∈ (0, 1)越界或非法值按特殊值规则处理。支持 0-8 维支持空 Tensor。FLOAT、FLOAT16、BFLOAT16NDout输出公式中的输出y标准正态分位点。shape 与 dtype 均与 self 一致。FLOAT、FLOAT16、BFLOAT16ND该参数定义与 op_host/ndtri_def.cpp 完全对应输入self与输出out均声明为REQUIRED支持ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16三种数据类型与ge::FORMAT_ND格式并带AutoContiguous()属性对非连续输入由框架自动做连续化处理。约束说明同构约束out 与 self 必须同 shape、同 dtype、同 format。形状与数据类型推导分别由 op_host/ndtri_infershape.cpp*outShape *selfShape与 op_host/ndtri_infer_datatype.cpp输出 dtype 直接取输入 dtype在编译期保证。维度与空 Tensor支持 0~8 维支持空 Tensornumel0 时 short-circuit 返回Kernel 的Process()在blockLength_ 0时直接 return。连续性兜底非连续 Tensor 由 aclnn L2 层通过 Contiguous 兜底处理Kernel 仅处理连续内存self的AutoContiguous()声明正是这一机制在算子定义层的体现。低精度内部提精度fp16 / bf16 输入在 Kernel 内部统一 Cast 到 fp32 参与多项式计算结果再 Cast 回原 dtype详见下文计算流程 Step 0/Step 7。精度标准浮点计算类社区标准并发判定 MERE Threshold 且 MARE 10×Threshold数据类型ThresholdFLOAT322⁻¹³FLOAT162⁻¹⁰BFLOAT162⁻⁷调用说明aclnn 调用与编译验证调用方式调用样例说明aclnn调用test_aclnn_ndtri参见仓库 docs/zh 目录下的算子调用指南完成算子编译和验证。ACLNN L0 API 实现在 op_api/ndtri.cpp 中Ndtri(const aclTensor* self, aclOpExecutor* executor)先经IsAiCoreSupport校验架构仅DAV_3510与 dtype仅 FLOAT/FLOAT16/BF16随后通过executor-AllocTensor按输入 view shape 与 dtype 分配输出张量最终以ADD_TO_LAUNCHER_LIST_AICORE(Ndtri, OP_INPUT(self), OP_OUTPUT(out))将算子下发到 AI Core 执行。综合测试样例examples/arch35/test_aclnn_ndtri.cpp 提供了可直接参考的完整验证程序其核心验证思路是回环验证利用标准正态 CDF $\Phi(z) 0.5 \cdot (1 \mathrm{erf}(z/\sqrt{2}))$ 检验Φ(ndtri(p)) ≈ p。测试覆盖shape 覆盖小 shape7 元素、对齐 shape1024、2D shape32×32、大 shape100000、非对齐 shape3×5dtype 覆盖float32、float16概率值覆盖边界附近0.001 / 0.999、小/中/大概率区间固定 15 值0.001、0.01、0.05、0.1、0.2、0.3、0.4、0.5、0.6、0.7、0.8、0.9、0.95、0.99、0.999以及 (0.001, 0.999) 区间均匀随机采样避开 0/1 奇点。调用模式为标准的两阶段 aclnn 流程aclnnNdtriGetWorkspaceSize(xT, outT, wsSize, executor)查询 workspace 并创建执行器再调用aclnnNdtri(wsAddr, wsSize, executor, stream)执行最后aclrtSynchronizeStream同步并回读结果比对。float32 用例的 roundtrip 绝对误差容限为 1e-4float16 用例为 5e-2。源码级实现原理Kernel 计算流水7 步op_kernel/ndtri_kernel.h 中的NdtriT, K_ALIGN类实现了完整的单 tile 计算流采用CopyIn → Compute → CopyOut的双缓冲BUFFER_NUM 2流水结构Step 0 输入 Cast 到 fp32fp32 输入通过Adds(p, inFp32, 0.0f)等价拷贝到独立 fp32 bufferfp16/bf16 通过Cast(CAST_NONE)无损提升精度Step 1 构造 maskBuildMasks生成三个 uint8 掩码——maskTail尾区|p−0.5| ≥ 0.5 − e⁻²、maskNegp ≥ 0.5决定尾部符号、maskSpecialp ≤ 0或p ≥ 1或 NaN其中 NaN 检测使用Compare(p, scratch, NE)技巧对 p 依次做Maxs(p, −INF)、Mins(·, INF)将 NaN 回填为非 NaN 值后再比较规避了 AscendCCompareScalar不支持同张量比较的限制Step 2 概率钳制pSafe clamp(p, FLT_MIN, 1 − FLT_MIN)保证后续ln(q)不出现负无穷参数Step 3 尾区计算cal_tail(pSafe, maskNeg)内部组合cal_sub求x sqrt(−2·ln(q))与x0 x − ln(x)/x和cal_p12按x 8掩码在 P1/Q1 与 P2/Q2 两组逼近之间 select 修正项最后按 maskNeg 决定正负号Step 4 中心区计算cal_p0以pm p − 0.5、z pm²展开有理逼近并乘√2πStep 5/6 结果合成先select(maskTail, yTail, yCenter)再select(maskSpecial, ySpecial, y)覆盖 0、1、NaN、±∞ 等特殊值BuildSpecialY构造p0 → −inf、p1 → inf、其余 → NaNStep 7 输出 Cast 回原 dtypefp16/bf16 使用Cast(CAST_RINT)四舍五入。值得注意的实现细节由于Compare/CompareScalarAPI 对 count 所占空间有 256B 对齐硬约束FP32 下即 64 元素倍数尾块长度会向上对齐到 64 的倍数并在 padding 区域填入中性值 0.5f位于中心区既不触发尾部分支也不触发特殊分支最终由DataCopyPad的blockLen currentNum * sizeof(T)保证 padding 位置的输出不会写回 GM。有理逼近的 Tensor 化实现op_kernel/ndtri_compute.h 将 Cephes 的标量算法张量化PolEvl按 Cephespolevl约定coefs[0]为最高次项系数用 Horner 法则计算 $P(x)$PlEvl按p1evl约定计算首项系数为 1 的 $Q(x) x^n \cdots$PolEvlPlEvlDiv(tmpP, tmpQ)得到有理函数 $P(x)/Q(x)$CalSubLn → Muls(−2) → Sqrt求 x再Ln → Div → Sub求x0CalP12z 1/x后分别计算两路有理逼近用CompareScalar(x, 8.0, LT)生成掩码后Select合并CalTail组合以上步骤并以Select(maskNeg, base, −base)完成左右尾的符号统一。实现注释还记录了性能优化结论Horner 每步的Mul Adds串联在 arch35 编译层会自动融合为单条 FusedMulAdd 指令手工改写为Duplicate FusedMulAdd的等价形式经 msprof 对比大 shape 耗时差 1.5%处于测量噪声带内故保留语义最清晰的原始写法。Kernel 入口与 TilingKernel 入口 op_kernel/ndtri.cpp 以DTYPE_SELF由 def 驱动实例化 T 模板实参和K_ALIGN32B 对齐标记0非对齐、1对齐为模板参数通过REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT获取 Tiling 数据后调用op.Init → op.Process。Tiling 逻辑位于 op_host/ndtri_tiling.cpp要点包括dtype 校验仅接受DT_FLOAT / DT_FLOAT16 / DT_BF16多核切分perCoreRaw CeilDiv(totalNum, coreNum)再按 32B/64 元素对齐得到blockFactorSetBlockDim(usedCoreNum)设置实际使用的核数UB 切分以availableUb ubSize − 48KB 保留为可用空间按 64 字节/元素估算 tile 大小并对 256 对齐得到ubFactorTilingKey 派发DispatchTilingKey根据totalNum % alignElem 0设置isAlign对应 {fp32, fp16, bf16} × {对齐, 非对齐} 共 6 个编译实例kernel 头文件注释证实 6 个 TilingKey 均为真实实现。算子定义与工程接入op_host/ndtri_def.cpp 中AiCoreConfig的关键配置包括DynamicCompileStaticFlag(true)动态编译静态化、DynamicRankSupportFlag(true)与DynamicShapeSupportFlag(true)动态 rank 与动态 shape 支持、PrecisionReduceFlag(true)并通过ExtendCfgInfo(opFile.value, ndtri)将算子定义与 Kernel 入口文件 op_kernel/ndtri.cpp 关联。仓库级构建接入由 experimental/math/ndtri/CMakeLists.txt 中的add_all_modules_sources(OPTYPE ndtri ACLNNTYPE aclnn_exclude)完成与 ops-math 仓库统一的算子模块构建体系保持一致。参考资源Ndtri 算子的对标基线与算法来源README 参考资源此处不展开外部链接仅保留可检索线索TensorFlowtf.math.ndtri对标基线Cephes Mathematical Library 的ndtri算法本源ndtri_coeffs.h头文件头部保留了原始出处声明SciPyscipy.special.ndtriPyTorchtorch.special.ndtri仓库内的核心实现文件索引README、Kernel 主体、系数表、计算核心、Tiling 实现、算子定义、ACLNN L0 API、综合测试样例。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表