已关闭
[Requirement|需求建议]: 新增 besseli1e算子 #2216
kangjiaming创建于  7月10日关闭于  7月13日
kangjiaming
7月10日 创建

[Requirement|需求建议]: 新增 besseli1e算子 #2216

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

一、背景信息 (必填)

BesselI1e 算子用于计算指数缩放的第一阶修正贝塞尔函数(Exponentially Scaled Modified Bessel Function of the First Kind, Order 1)。

计算公式:

outi=einputiI1(inputi)out_i = e^{-|input_i|} \cdot I_1(input_i)

其中 I1(x)I_1(x) 是第一阶修正贝塞尔函数,指数缩放因子 exe^{-|x|} 使得函数在大参数值时数值稳定,避免普通贝塞尔函数在参数较大时出现的数值溢出问题。

该算子在物理学(热传导、波动问题)、信号处理、统计学(von Mises 分布)等领域有广泛应用。TensorFlow 和 SciPy 均已提供该算子的实现(tf.math.bessel_i1e / scipy.special.i1e)。

二、价值/作用 (必填)

  1. 补齐数学函数能力:BesselI1e 是基础数学函数,与已有的 BesselI0e 算子配套,完善特殊函数算子库。
  2. 数值稳定性:相比直接计算 I1(x)I_1(x) 再乘以指数衰减因子,BesselI1e 在大参数值时具有更好的数值稳定性,避免浮点溢出。
  3. 框架兼容性:对标 TensorFlow 的 tf.math.bessel_i1e 算子,提升 Ascend 平台对主流框架模型的兼容性。
  4. 科学计算场景:满足物理仿真、电磁场计算、统计分布等科学计算场景对高阶贝塞尔函数的需求。

三、设计方案 (必填)

3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)

使能方式 是否支持
图模式(GE IR 构图)
Aclnn 直调 ×
PyTorch 训练 ×

当前仅支持图模式调用,通过 GE 算子 IR(bessel_i1e_proto.h)构图方式使用。

3.2 总体设计

BesselI1e 为 Elementwise 算子,对输入 Tensor 逐元素计算 exI1(x)e^{-|x|} \cdot I_1(x)

整体实现分为三层:

  • op_graph 层:GE 算子原型注册(bessel_i1e_proto.h),定义算子输入输出及属性。
  • op_host 层:算子定义(bessel_i1e_def.cpp)、形状推导(bessel_i1e_infershape.cpp)、Tiling 参数计算(按架构分目录 arch32/arch35/)。
  • op_kernel 层:AI Core 核函数实现(bessel_i1e_apt.cpp),采用分段多项式逼近算法。

3.2.1 算子支持的数据类型

参数名 输入/输出 数据类型 数据格式
x 输入 FLOAT、FLOAT16、BF16 ND
y 输出 FLOAT、FLOAT16、BF16 ND

输入与输出的数据类型和形状必须一致。

3.2.2 host侧设计

算子定义(bessel_i1e_def.cpp)

  • 注册算子 BesselI1e,定义 1 个输入 x 和 1 个输出 y
  • 输入输出均支持 BF16、FLOAT16、FLOAT 数据类型,ND 数据格式。
  • 配置 AI Core 参数:DynamicCompileStaticFlag(true)DynamicRankSupportFlag(true)DynamicShapeSupportFlag(true)PrecisionReduceFlag(true)
  • 支持硬件配置:ascend910bascend910_93ascend950

形状推导(bessel_i1e_infershape.cpp)

  • Elementwise 算子,输出形状与输入形状相同:*output_shape = *input_shape

Tiling 参数计算(arch35/bessel_i1e_tiling_arch35.cpp)

  • 根据输入数据量和硬件规格计算 Tiling 参数,包括 block 数量、每次迭代处理的数据量等。
  • arch32 架构复用 arch35 的 Tiling 逻辑。

3.2.3 kernel侧设计

核函数入口(bessel_i1e_apt.cpp)

  • 注册 AICore 核函数 bessel_i1e,采用 AIV_ONLY 模式(纯向量计算,无需 Cube 单元)。
  • 通过 REGISTER_TILING_DEFAULT 注册 Tiling 数据结构,从 GM 地址读取 Tiling 参数。

计算实现(arch35/bessel_i1e.h)

  • 采用分段多项式逼近算法计算 I1(x)I_1(x)
    • 分段点 SEGMENT_POINT = 3.75,将计算区间分为 x3.75|x| \leq 3.75x>3.75|x| > 3.75 两段。
    • 小参数区间使用多项式逼近 I1(x)I_1(x),再乘以 exe^{-|x|}
    • 大参数区间使用渐近展开式直接计算 exI1(x)e^{-|x|} \cdot I_1(x),避免中间结果溢出。
  • 利用 AscendC 编程框架实现向量化计算,支持多核并行。

3.3 支持硬件

产品 是否支持
Ascend 950PR / Ascend 950DT
Atlas A3 训练系列产品 / Atlas A3 推理系列产品
Atlas A2 训练系列产品 / Atlas A2 推理系列产品
Atlas 200I/500 A2 推理产品 ×
Atlas 推理系列产品 ×
Atlas 训练系列产品 ×

对应硬件配置标识:ascend910bascend910_93ascend950

3.4 算子约束限制

  1. 输入与输出的数据类型和形状必须一致。
  2. 仅支持图模式调用,不支持 aclnn 接口直调。
  3. 仅支持 ND 数据格式。
  4. 不支持 DOUBLE 数据类型(与 BesselI0e 不同,BesselI1e 当前未注册 DOUBLE 类型)。

💡 备注(选填)

  • 对标 TensorFlow 算子:tf.math.bessel_i1e
  • 对标 SciPy 函数:scipy.special.i1e
  • 已有同系列算子 BesselI0e 可供参考,两者实现架构一致。
  • 代码已由 CANNBot 生成,详见 math/bessel_i1e/ 目录。
likedislike
Kkangjiaming
7月10日 修改了issue 的描述
Kkangjiaming
7月10日 关联了pull request:besseli1e算子开发
陈思
陈思成员
7月12日 评论:

/assign @kangjiaming

likedislike
CANN-robotCANN-robot成员
7月12日 将 kangjiaming 设为负责人
CANN-robotCANN-robot成员
7月13日 关闭了 issue
CANN-robotCANN-robot成员
7月13日 添加了label:resolved