Polar inner-broadcast 性能优化方案 A【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math2026-05-21。本轮针对 acceptance 50-case 中唯一性能不达 baseline 的两个 caseTest_036 / Test_044 做 angle inner-broadcast 专用优化从输基线 0.8x 反超到 1.7-3.5x并使所有 angle 周期广播 case 普遍 ~2x 提速。1. 问题定位测试口径pybind EXEC_NPU_CMD(aclnnPolar)msprof op_summary Task Duration每调用设备时50 次取中位。 baseline torch.polarNPU 内置Sin/Cos/Mul/Complex 基础算子拼接。优化前全 50 case 仅这两个输基线Caseinput × angleout 元素优化前 oursbaselinevs baseTest_036[256,256,64] ×[64]4M168.8 µs147.2 µs0.87x❌Test_044[128,128,512] ×[512]8M331.4 µs276.6 µs0.83x❌根因这两个是 angle 为 input trailing-dim 的周期重复 广播。原实现在 op_api 层用BroadcastTo把 angle 物理展开到 out shape同 shape kernel 要求于是Test_036先写 16 MBangle 展开→ kernel 再读 input 16 MB 读展开后 angle 16 MB 写 out 32 MB 80 MB HBMbaseline 用 viewstride0读 angle省掉展开的一次 HBM 写故反而更快。2. 方案 Ainner-broadcast 周期 pre-computekernel 内运行时分支无 tilingKey核心观察angle 周期长度 K angle.numel每个 tile 取整数个 period →cos/sin tile pattern 对所有 tile 完全相同。实现三处资格判定严格一致op_api / op_host tiling / pybind资格 inNoutN anNinN inN%anN0 anN2048 anN%80op_api (aclnn_polar.cpp)资格命中时 angle不 BroadcastTo保留原 shape [K] 传 kernelop_host (polar_tiling.cpp)识别bcastMode1按K-block 分核每核负载 K 整数倍 → tile 起点 period 对齐op_kernel (polar.h)Process()内if(bcastMode1)分支纯 data 字段无需多 kernel 注册PrecomputeBcastCosSin()读 angle[0:K] → Cos/Sin 得 cosT/sinT[K] → 循环 mT/K 次 broadcast 填满 常驻 buffercosTile[0:T] / sinTile[T:2T]只做一次ProcessTileBcast()main loop只读 input→Mul(cosTile,input)→real /Mul(sinTile,input)→imag → Gather 交织 → 写 out。无 angle 搬运、无每 tile Sin/Cos 重算。HBM/算力收益Test_036, 4M80 MB →48 MB省 angle 展开 16 MB angle 读 16 MB→256 B Sin/Cos 从每 tile 重算降到一次 precompute。⚠️ 实现陷阱DataCopy(srcStride0)不是重复同段 broadcast——它连续读 src会把 cosT 之后的 sinT 读进第二 period症状npu.realgolden.imag。正解是把 broadcast 提到 Precompute 一次性循环填满。3. 优化结果关注的两个 case从输基线反超CaseK优化前优化后kernel 提速baselinevs base 前→后Test_03664168.842.14.0x147.20.87x →3.50xTest_044512331.4161.92.0x276.60.83x →1.71x其它 angle 周期广播 case顺带普遍 ~2xCaseout 元素K优化前优化后baselinevs baseTest_035262K6417.615.754.73.47xTest_03716M646543217652.38xTest_03867M642595126133402.65xTest_039134M325181253167712.67xTest_040268M16103545144136332.65xTest_041537M82070110287273612.66xTest_047134M85181258167842.63xTest_04833M32130164116392.56x未启用优化K1/2/4 不满足 32B 对齐走 same-shape 兜底的 Test_042/045/046因规模大268-537M baseline 同为 HBM boundsame-shape 已 1.31-1.47x 胜基线无需再优化。全 50-case 结论无一输基线类别case加速比小同 shape (1-12)same-shape2.5-3.4x大同 shape (13-29)same-shape2.5-3.1x小广播 (30-34, K8)same-shape 兜底3.2-5.3x大广播 (35-50)bcast 资格 ~2x kernel 提速非资格规模大也胜1.3-3.5x4. 官方 AscendOpTest 验证ops-math 工程customize_math vendorpython run_test.py -i .../acceptance/op.json -c Test_036,Test_044 \ --op-type custom --op-path $ASCEND_OPP_PATH/vendors/customize_math/op_api --build结果result.csvcasecompare_resultTest_036passTest_044pass精度判据 err_threshold[1e-3, 1e-3]官方工具确认数值正确。5. 改动文件ops-mathop_kernel/polar_tiling_data.hinN/anN/bcastModeop_host/polar_tiling.cppinner-bcast 检测 K-block 分核op_kernel/polar.hProcess()分支 PrecomputeBcastCosSin()ProcessTileBcast()op_api/aclnn_polar.cpp资格命中时 angle 不 BroadcastToop_api/polar.cpp注释更正同 shape 与非 inner-bcast 广播路径中间轴 / input 侧广播如 Test_043/049/050行为不变由 same-shape kernel op_api BroadcastTo 兜底。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考