ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

CANN ops-math BesselI1e 算子解析:指数缩放第一阶修正贝塞尔函数在 NPU 上的实现与调用

CANN ops-math BesselI1e 算子解析:指数缩放第一阶修正贝塞尔函数在 NPU 上的实现与调用 CANN ops-math BesselI1e 算子解析指数缩放第一阶修正贝塞尔函数在 NPU 上的实现与调用【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathBesselI1e 是 CANN ops-math 数学算子库中用于计算指数缩放第一阶修正贝塞尔函数Exponentially Scaled Modified Bessel Function of Order 1的基础算子。本文以 math/bessel_i1e/README.md 为核心结合算子 IR、Host 侧 shape 推导与 tiling、Kernel 侧分段多项式逼近算法等源码系统讲解该算子的数学定义、产品支持情况、参数约束、图模式调用方式以及底层实现原理帮助开发者掌握在昇腾 NPU 上正确使用与理解 BesselI1e 算子的完整方法。功能说明与数学背景BesselI1e 算子以张量x为输入逐元素element-wise计算指数缩放的第一阶修正贝塞尔函数其计算公式为$$ out_i e^{-|input_i|} \cdot I_1(input_i) $$其中 $I_1(x)$ 是第一阶修正贝塞尔函数Modified Bessel Function of the First Kind, Order 1。指数缩放因子 $e^{-|x|}$ 的作用是数值稳定在 $|x|$ 较大时$I_1(x)$ 本身会以 $\sim e^{x}/\sqrt{2\pi x}$ 量级快速增长直接计算容易溢出或损失精度乘以 $e^{-|x|}$ 后函数值被压缩到 $O(1/\sqrt{|x|})$ 量级从而在宽泛的参数区间内保持数值稳定。从算子 IR 注释可以看出该算子与 TensorFlow 框架的BesselI1e算子兼容见 math/bessel_i1e/op_graph/bessel_i1e_proto.h因此在使用 CANN 的图模式对接 TensorFlow 模型时可直接承载对应的框架算子。产品支持情况README 明确给出了各产品的支持矩阵产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×对照算子 Host 侧注册配置math/bessel_i1e/op_host/bessel_i1e_def.cpp该算子为AICore()添加了ascend910b、ascend910_93、ascend950三档芯片配置同时 Kernel 侧存在arch35对应 Atlas A3/Ascend 950 等新架构与arch22对应 Atlas A2 等两套 tiling 实现与上表的产品支持范围相互印证。其中arch22的 tiling 实现直接复用arch35的逻辑见 math/bessel_i1e/op_host/arch22/bessel_i1e_tiling_arch22.cpp。参数说明算子包含一个输入和一个输出具体参数如下参数名输入/输出/属性描述数据类型数据格式x输入待进行 BesselI1e 计算的入参FLOAT、FLOAT16、BF16NDy输出BesselI1e 计算的出参FLOAT、FLOAT16、BF16ND几点补充说明依据源码确认数据类型README 列出 FLOAT、FLOAT16、BF16 三种。算子 IR 的TensorType({FloatingDataType, DT_BF16})定义以及 tiling 源码中supportedDtype {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}的校验集合均与此一致即支持浮点类型与 BF16。数据格式仅支持 ND 格式算子定义中Format({ge::FORMAT_ND, ...})与UnknownShapeFormat({ge::FORMAT_ND, ...})均固定为 ND。输出与输入同 shape、同 dtypeshape 推导将输出 shape 直接拷贝为输入 shape见下节数据类型推导将输出类型设置为输入类型xDataType。约束说明README 中该算子的约束说明为无。从实现上看算子本身没有复杂的属性无 attr、无可选输入唯一需要注意的是输入x与输出y的 shape 大小必须一致tiling 阶段会校验inputShapeX.GetShapeSize() ! outShapeY.GetShapeSize()并报错见 math/bessel_i1e/op_host/arch35/bessel_i1e_tiling_arch35.cpp输入维度上限为 8 维超出会报dim num of x must be less than or equal to 8标量0 维输入会被视为 shape 为{1}的向量参与计算。调用方式图模式调用README 提供的调用方式为图模式调用即通过算子 IR 构图GE 图的方式调用 BesselI1e 算子对应样例为 math/bessel_i1e/examples/test_geir_bessel_i1e.cpp算子 IR 定义在 math/bessel_i1e/op_graph/bessel_i1e_proto.h。1. 算子 IR原型定义算子原型使用REG_OP宏注册接口极简——仅一个输入x与一个输出y无任何属性REG_OP(BesselI1e) .INPUT(x, TensorType({FloatingDataType, DT_BF16})) .OUTPUT(y, TensorType({FloatingDataType, DT_BF16})) .OP_END_FACTORY_REG(BesselI1e)2. 构图与执行流程样例程序的核心流程分为五步可归纳为初始化 GE通过ge::GEInitialize初始化global_options中设置了ge.exec.deviceId0与ge.graphRunMode1。创建算子节点并构图在CreateOpInGraph中创建op::BesselI1e(bessel_i1e_0)算子对象构造op::Data(placeholder0)作为输入占位节点设置输入 shape 为{32, 4, 4, 4}默认DT_FLOAT类型通过bessel_op.set_input_x(placeholder0)建立数据依赖并将节点加入Graph。设置图输入输出graph.SetInputs(inputs).SetOutputs(outputs)将 placeholder 与 BesselI1e 算子分别指定为整图的输入与输出。创建 Session 并运行ge::Session创建会话session-AddGraph(graph_id, graph)添加图session-RunGraph(graph_id, input, output)触发整图编译与执行。落盘输出将输出张量按bessel_i1e_geir_output_i.bin命名写入二进制文件便于后续比对验证。样例中一个值得注意的细节测试数据由GenTestData生成取值为(i % 10) * 0.5f即落在[0, 4.5]区间——该区间横跨了 Kernel 算法中分段逼近的切换点 $|x| 3.75$详见下文能够同时覆盖大小参数两条计算路径。3. 配套示例的辅助设施整个 BesselI1e 算子目录中还包含构建与测试脚本math/bessel_i1e/CMakeLists.txt算子各子模块op_host、op_kernel、op_graph、examples、tests的构建入口math/bessel_i1e/tests/ut/run.sh 与 math/bessel_i1e/tests/ut/test_bessel_i1e_ut.cpp单测运行脚本与测试用例可用于验证算子正确性math/bessel_i1e/framework/bessel_i1e_tf_plugin.cppTensorFlow 框架侧的算子插件适配承载框架算子到 CANN 算子的对接。源码级实现剖析1. Host 侧shape 与数据类型推导shape 推导实现于 math/bessel_i1e/op_host/bessel_i1e_infershape.cpp逻辑为输出 shape 直接等于输入 shape逐元素算子的标准行为gert::Shape* output_shape context-GetOutputShape(0); *output_shape *input_shape;数据类型推导实现于 math/bessel_i1e/op_graph/bessel_i1e_graph_infer.cpp将输出数据类型设置为输入数据类型。2. Host 侧tiling 计算tiling 是昇腾算子将任意 shape 切分为 NPU 可并行执行粒度核数 × UB 分块的关键步骤。arch35的 tiling 逻辑math/bessel_i1e/op_host/arch35/bessel_i1e_tiling_arch35.cpp产出一个仅含三个字段的BesselI1eTilingData结构见 math/bessel_i1e/op_kernel/arch35/bessel_i1e_tiling_data.hstruct BesselI1eTilingData { int64_t totalNum 0; // 输入总元素个数 int64_t blockFactor 0; // 每个核处理的基本元素块大小 int64_t ubFactor 0; // 单次 UB 搬运/计算的元素块大小 };计算要点核数分配通过platform_ascendc::PlatformAscendC查询可用 AIV 核数coreNumblockFactor CeilAlign(CeilDiv(totalIdx, coreNum), ubBlockSize / dtypeSize)保证块大小按 UB 对齐粒度对齐usedCoreNum CeilDiv(totalIdx, blockFactor)为实际启用核数UB 分块ubFactor依据 UB 总容量、内部队列数量3 个队列 × 双缓冲因子 2 6以及对齐粒度计算得到Workspace大小为 0该算子不需要额外 workspace模板参数下发通过ASCENDC_TPL_SEL_PARAM(context, dTypeX)将输入数据类型C_DT_FLOAT / C_DT_FLOAT16 / C_DT_BF16作为模板参数下发驱动 Kernel 侧实例化对应数据类型版本见 math/bessel_i1e/op_kernel/arch35/bessel_i1e_tiling_key.h。值得强调的是bessel_i1e_def.cpp中注册的DynamicCompileStaticFlag(true)、DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)表明该算子支持动态 shape 与动态 ranktiling 在运行时根据实际输入 shape 现场计算这正是上述 tiling 数据逐字段推导的原因。3. Kernel 侧分段多项式逼近算法Kernel 入口位于 math/bessel_i1e/op_kernel/bessel_i1e_apt.cpp通过REGISTER_TILING_DEFAULT读取 tiling 数据后调用NsBesselI1e::BesselI1eT完成计算算法主体在 math/bessel_i1e/op_kernel/arch35/bessel_i1e.h。1数据流与流水线。Process()中按ubFactor将每核数据切成多轮循环每轮依次执行CopyIn → Compute → CopyOut输入、输出、临时缓冲均使用双缓冲队列BUFFER_NUM 2通过流水线隐藏搬运与计算延迟。CopyIn/CopyOut对half/bfloat16_t输入做了类型适配低精度数据先转成float参与计算计算完成后再转回原类型写回BF16 回写使用CAST_RINT舍入模式保证中间精度。2分段逼近策略。由于 $I_1(x)$ 在奇点与无穷远处的渐进行为不同算法以SEGMENT_POINT 3.75f为界分为两段这也是常见修正贝塞尔函数多项式近似的经典分界点如 Cephes 库的划分方式小参数区间$|x| 3.75$令 $t |x|/3.75$利用itrBefore[7]七项系数做 Horner 循环求值得到多项式 $p(t^2)$同时用EXP_COEFF[13]共 13 项泰勒系数在 $q |x|/4$ 处计算 $e^{-|x|/4}$再通过两次平方e e * e; e e * e恢复出 $e^{-|x|}$最终result e * |x| * poly。大参数区间$|x| \ge 3.75$令 $t 3.75/|x|$利用itrAfter[9]九项系数做 Horner 求值得到 $p(t)$再除以 $\sqrt{|x|}$对应 $I_1(x) \sim e^{x}/\sqrt{2\pi x}$ 的渐近主项得到result poly / sqrt(|x|)。奇函数性质恢复符号$I_1(x)$ 为奇函数即 $I_1(-x) -I_1(x)$因此算法先对 $|x|$ 计算、最后再乘上sign恢复正负号。3精度说明。源码头注释指出exp(-|x|/4)的泰勒展开由 9 项增加到 13 项后最大误差从约2.6e-6下降到约6.1e-10约 4000 倍提升且所有多项式求值均采用循环式 Horner 评估for (int k ...; k 0; k--) { poly poly * t2 itrBefore[k]; }刻意避免深层嵌套表达式——这是针对 Ascend C 编译器无法处理过深嵌套的工程化取舍。这两点均属于代码注释与实现本身传达的信息可作为理解其精度设计与编码风格的依据。总结BesselI1e 算子是 CANN ops-math 中实现指数缩放第一阶修正贝塞尔函数的基础数学算子功能上以e^{-|x|}·I_1(x)的数值稳定形式逐元素计算支持 FLOAT/FLOAT16/BF16 三种数据类型与 ND 格式覆盖 Ascend 950 与 Atlas A2/A3 系列产品实现上由算子 IRop_graph、shape/类型推导与动态 tilingop_host、分段多项式逼近 Kernelop_kernel三层构成并通过图模式样例与 UT 测试闭环验证。对开发者而言既可以直接在 GE 图中通过op::BesselI1e构图调用也可以将本文的算法剖析作为在 NPU 上实现类似特殊函数算子的参考范本。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表