
MulNoNan【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×本算子目前仅在 Ascend 950 (dav-3510 / Atlas 350 加速卡) 架构上实现。功能说明算子功能二元 element-wise 安全乘法把Mul中0 · inf NaN与0 · NaN NaN两类异常屏蔽为 0避免在反向传播 / loss 计算等场景把 NaN 污染传递下去。计算公式$$ y_i \begin{cases} 0, x_{2,i} 0 \ x_{1,i} \cdot x_{2,i}, x_{2,i} \ne 0 \end{cases} $$仅判x2一侧。x2 -0同样进入零臂IEEE-754-0 0。x2 \ne 0时所有 IEEE 行为按普通乘法保留NaN / Inf 正常传播。第三方对标等价于 TensorFlowtf.math.multiply_no_nans。参数说明参数名输入/输出描述数据类型数据格式x1输入乘法第一个输入张量。FLOAT16, FLOAT, INT32, BFLOAT16NDx2输入乘法第二个输入张量判 0 主体与x1须可广播。FLOAT16, FLOAT, INT32, BFLOAT16NDy输出x1 * x2的结果x2 0处强制为 0。shape 为x1、x2广播后的统一形状。FLOAT16, FLOAT, INT32, BFLOAT16ND约束说明x1、x2、y必须为同一种 dtype不支持 mix-dtype。支持任意 NumPy 广播形态含标量[1]、单维 broadcast、跨 rank broadcast。支持动态 shape 与动态 rank。实现方案层文件说明计算图原型op_graph/mul_no_nan_proto.hREG_OP(MulNoNan)二输入一输出算子定义op_host/mul_no_nan_def.cppOpDef::AddConfig(ascend950, ...)InferShapeop_host/mul_no_nan_infershape.cpp复用Ops::Base::InferShape4BroadcastTilingop_host/arch35/mul_no_nan_tiling_arch35.{h,cpp}按 dtype 分支调用Ops::Base::BroadcastBaseTilingOpDagDAGop_kernel/arch35/mul_no_nan_dag.hfp32/int32 通路原生计算fp16/bf16 通路提升 fp32 中间精度Structop_kernel/arch35/mul_no_nan_struct.hBRC_TEMP_SCH_MODE_KEY_DECL/SELKernel 入口op_kernel/mul_no_nan_apt.cppKERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)BroadcastSchschMode, OpDagfp32 / int32 通路In0/In1 -- CopyInBrc -- -- Mul(x1,x2) -- -- Select(mask, MulRes, Zero) -- CopyOut -- Out0 Const(0) -- Duplicate -- Zero -- | -- Compare(NE, x2, Zero) -- maskVec::Compareu8, T, NE输出位掩码Vec::Selectu8, T, TENSOR_TENSOR按mask ? MulRes : Zero逐元素选择。x2 0时即使MulRes是NaN0 · inf/0 · NaN也被 Select 丢弃输出 0。fp16 / bf16 通路In0/In1 -- CopyInBrc -- Cast(-fp32) -- -- Mul -- -- Select -- Cast(-T,RINT) -- CopyOut -- Out0 Const(0,fp32) -- Duplicate -- Zero -- | -- Compare(NE) -- mask把 fp16 / bf16 整体提升到 fp32 做 cmp/sel/mul末端用CAST_MODE_RINTround-to-nearest-even回退。原因与 DSLmul_no_nan.py在 fp16vcmpsel不可用时 fallback 到 fp32 的行为一致避免 fp16 中间溢出如3e4 · 3e4在 fp16 中先 inf 再被 select 处理会 引入额外的 saturation 不确定性fp32 中间有充足动态范围。调用说明调用方式样例代码说明图模式test_geir_mul_no_nan通过算子IR构图方式调用 MulNoNan 算子覆盖 fp32/fp16/bf16/int32 基础用例 0·inf、0·NaN、-0、广播等关键特殊值用例。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考