CANN ops-math Polar 算子 NPU 实战指南：基于 pybind 与 msprof 的精度-性能一体化测试框架

Polar 是 CANN ops-math 数学算子库中由极坐标模长（abs）与幅角（angle）构造复数张量的基础算子，对应 PyTorch 的 torch.polar(abs, angle)。本文以 experimental/math/polar/tests/pybind/README.md 为核心，完整讲解该算子配套的 pybind 测试框架：从 wheel 构建、msprof 计时、复数精度校验到性能基线门禁的完整闭环，并深入到 pytorch_npu_helper.hpp 的动态 dlopen 机制与 test_op.py 的 16 个测试用例设计，帮助读者掌握在 NPU 上对自定义算子做正确性与性能联合验收的完整方法。

一、框架定位：Polar 算子的 S8 同款验收流水线

Polar 算子的计算公式为：

$$out_i = input_i \times (\cos(angle_i) + i \cdot \sin(angle_i))$$

其中 input（模长）与 angle（幅角，弧度）为 fp32 张量且支持 NumPy 广播规则，out 为 complex64 复数张量，shape 等于两者广播后的结果。对应接口为 aclnnPolar(input, angle, out)，其对齐基准为 ops-math 仓 math/complex 下的 l0 拼接版参考实现。

围绕该算子，tests/pybind/ 提供了一套与 S8 测试习惯完全一致的验收框架（"S8 同款"），整体流程为：重建 pybind wheel → msprof 计时 → get_time 提取耗时 → 性能基线校验，同时通过 test_op.py 完成多场景正确性验证。该框架只在 NPU 环境可运行（依赖 torch_npu），本地无法自检。

二、文件构成与职责分工

文件	来源	说明
common/pytorch_npu_helper.hpp	逐字复用 S8	EXEC_NPU_CMD 通用胶水（算子无关，运行时动态 dlopen aclnnPolar 及其 GetWorkspaceSize）
setup.py / get_time.py	逐字复用 S8	pybind wheel 构建 / 解析 msprof 生成的 op_summary*.csv 并取第 [20:40] 条记录的均值
extension/custom_op.cpp	Polar 适配	EXEC_NPU_CMD(aclnnPolar, input, angle, result)，result 为 complex64、shape 由 at::infer_size 按广播规则推导，循环 50 轮供 msprof 取平均
test_op.py	Polar 适配	16 个 case + 复数精度 verify（实部/虚部各自绝对误差判据，TOL/RATE=1e-4）
run.sh	Polar 适配	重建 wheel → msprof 计时 → get_time → 基线校验；LD_LIBRARY_PATH 指向 custom_math vendor

这套结构与 S8 的 case_910b/<Op>/ 布局完全一致，任何算子都可以按同样的五个文件模板快速搭建验收环境。

三、前置条件：编译并安装自定义 Polar 算子

测试框架本身不负责算子编译，算子需要先在 ops-math 仓编译成自定义算子包并安装到 NPU 环境的 OPP vendor 目录。该步骤只需做一次，或在每次修改 kernel 后重做：

# 在 ops-math 仓编译并安装自定义 Polar 算子
bash build.sh --pkg --soc=ascend910b --ops=polar -j16     # A3 产品用 --soc=ascend910_93
./build_out/cann-ops-math-*linux*.run

安装后算子包位于 ${ASCEND_HOME_PATH}/opp/vendors/custom_math/，这也是 run.sh 中 LD_LIBRARY_PATH 指向的 vendor 名称。值得说明的是，这套构建走的是 ops-math 仓级构建（--pkg --ops=polar），与 S8 每算子自带 build.sh 的构建方式解耦，属于本框架与 S8 在构建层唯一的差异。

从 op_api/aclnn_polar.cpp 的源码看，L2 层的 aclnnPolar 是"Contiguous + numpy 广播（BroadcastTo）+ L0 Polar + ViewCopy"的组合：kernel 本身按同 shape elementwise 实现，所有广播与非连续处理在 L2 层完成，因此任意 shape、任意广播组合（≤8 维）都能正确处理。这一设计也解释了为什么测试框架在 host 侧用 at::infer_size 推导广播输出 shape，而正确性验证完全聚焦在广播语义与复数数值精度上。

四、运行方法：一条命令完成精度与性能验收

在 NPU 环境进入测试目录后，直接执行：

cd case_910b/Polar
bash run.sh 2          # 跑 case2（广播主战场）

run.sh 的参数是 case 编号。单次运行的预期输出包含两部分：

• 正确性结果：xxx verify result pass!（来自 test_op.py 的 verify_result）；
• 性能结果：time_base = ... time_use = ...（来自 get_time.py 提取的 msprof 耗时均值）。

run.sh 内部的关键动作依次为：

1. 设置 LD_LIBRARY_PATH，优先指向 ${ASCEND_OPP_PATH:-${ASCEND_HOME_PATH}/opp}/vendors/custom_math/op_api/lib/，并兼容 customize 兜底；
2. 每次运行前删除 ./dist ./build ./custom_ops.egg-info 并重建 wheel（避免上次运行其他算子残留的 wheel signature 不匹配），pip3 install dist/custom_ops*.whl --force-reinstall 强制重装；
3. 清掉旧的 PROF* 目录，用 timeout 180 msprof --application="python3 test_op.py $1" 在 msprof 采样下执行目标 case，超时（退出码 124）直接判失败；
4. 用 get_time.py 解析 msprof 产出的 op_summary*.csv 得到 time_use；
5. 与哨兵基线 time_base 比较，time_use == 0 或 time_use >= time_base 均判定失败。

五、测试用例设计：从广播主战场到极端边界

test_op.py 的 case_data 字典是全部测试数据的入口，新增 case 只需在 case_data 中加一个键（与 S8 习惯一致），框架会自动完成 golden 计算、NPU 执行与精度比对。README 中列出的 8 个基础 case 覆盖了广播、性能、对齐与数值归约四类考点：

case	input shape	angle shape	考点
1	[2,6,10]	[2,6,10]	同 shape 基础
2	[4,1,8]	[4,5,8]	广播：低维→高维（新增功能主战场）
3	[1]	[3,4,5]	广播：标量 input
4	[8,1]	[1,7]	广播：双向 → [8,7]
5	[4096,4096]	[4096,4096]	大 shape（性能）
6	[3,5,17,269]	同	高维 + inner 非 32B 对齐（1076B/行）
7	[1]	[1]	1 元素边界
8	[64,1024]	[64,1024]	大角度（Sin/Cos 范围归约压力）

结合 test_op.py 实际代码，case_data 中已进一步扩充到 16 个 case，可视为该表的实战延伸，值得逐一对照：

• case9（[2,2,2,2,2,2,2,2] 同 shape）：8D 满秩，对应 OpDef 的 MAX_DIM=8 上限；
• case10（input [3,1,5,1,7] × angle [3,4,5,6,7]）：5D 中间轴广播，覆盖多轴同时扩展的中间轴场景；
• case11（input [8,1,1] × angle [1,4,5] → [8,4,5]）：双向多轴广播；
• case12（标量 input [1] × 6D angle [2,3,2,3,2,3]）：标量与高维张量广播；
• case13（input 负值 [-50,50]，shape [128,257]）：验证负模长语义——torch.polar 中负 abs 合法，会翻转相位（翻号）；
• case14（input [5,1] × angle [5,300]，角度 [-50000,-1]）：大负角度 + rank 不一致广播；
• case15（input [7,1,13] × angle [7,11,13]）：非 32B 对齐 inner（13 个 fp32=52B）+ 广播的组合；
• case16（input [20000] × angle [1]）：一侧 rank-1 大向量 × 标量的极端形态。

这套 case 组合实际上把广播的三种形态（低维→高维、标量、双向）、维度上限（8D）、非对齐 inner 维度、负模长与大角度数值范围全部纳入了回归范围。

六、复数精度验证：实部/虚部分别判绝对误差

test_op.py 中的 verify_result 是精度门禁的核心。由于 complex64 在 AscendOpTest 的 accuracy_config 中没有内置默认阈值，框架采用与 AscendOpTest compare_complex 逐行等价的 fp32 分量判据：默认 err_threshold = [1e-4, 1e-4]（[绝对偏差, 错误率]）。

具体判据为：

• 实部、虚部各自纯绝对误差 ≤ TOL（TOL = 1e-4），无相对回退；
• 二者同时满足才算该元素正确；
• 错误元素数 > total × RATE（RATE = 1e-4）判失败；
• 失败时打印前 10 个越界元素的下标、NPU 结果、golden 结果及 dRe / dIm 偏差，便于定位。

验证前还通过 golden_polar 构造基准：先用 torch.broadcast_tensors 显式广播 input 与 angle，再调用 torch.polar 生成 complex64 golden，确保 golden 与算子的广播语义一致，避免"参考实现本身未广播"造成的假失败。NPU 侧则通过 custom_ops_lib.custom_op(input_npu, angle_npu) 触发算子执行。

七、性能计时：msprof 采样 + op_summary 均值提取

性能计时的数据链路由 run.sh 与 get_time.py 配合完成：

1. 采样：run.sh 用 msprof --application="python3 test_op.py $1" 包裹被测脚本，msprof 会为每次算子执行产出 op_summary*.csv；
2. 重复执行：extension/custom_op.cpp 中 round = 50，即在循环中重复执行 aclnnPolar 50 次，为统计均值提供足量样本；
3. 提取：get_time.py 递归扫描 ./ 下所有 op_summary*.csv，读取每行的 Task Duration(us) 字段，换算为纳秒（int(float(x) * 1000000)），并取第 [20:40] 条记录的均值——跳过前 20 条预热样本，规避首轮调度/显存分配抖动；
4. 门禁：run.sh 将均值与 time_base 比较。当前 time_base=9999999999999 是哨兵值，仅承担正确性门禁（任何非零耗时都通过），真实性能基线需待硬件实测出 l0 参考实现耗时后回填（任务要求 ≥ 基线 95%）。

八、源码剖析：EXEC_NPU_CMD 如何动态拉起 aclnnPolar

框架的算子无关性来自 pytorch_npu_helper.hpp 中的 EXEC_NPU_CMD 宏。它不依赖编译期链接 op_api 头文件，而是在运行时通过 dlopen/dlsym 查找符号：

• GetOpApiFuncAddr 依次尝试 libcust_opapi.so（自定义算子库）与 libopapi.so（标准算子库），找到首个可用符号即返回，这正是 run.sh 需要把 custom_math vendor 的 op_api/lib 提前放进 LD_LIBRARY_PATH 的原因；
• 对 aclnnPolar，宏按约定查找 aclnnPolarGetWorkspaceSize 与 aclnnPolar 两个符号，先调用 GetWorkspaceSize 获取 workspace 大小，需要时分配 NPU workspace tensor，再通过 OpCommand 的自定义 handler 在 NPU 流上执行；
• 参数侧通过 ConvertType 将 at::Tensor 转为 aclTensor（内置了 at::ScalarType → aclDataType 的完整映射表，如 Float → ACL_FLOAT、ComplexFloat → ACL_COMPLEX64），执行完经 ReleaseConvertTypes 释放，并接入 InitHugeMemThreadLocal/UnInitHugeMemThreadLocal 管理线程级内存。

extension/custom_op.cpp 在其上只做 Polar 适配：用 at::infer_size(input.sizes(), angle.sizes()) 按 numpy 广播规则推导输出 shape，创建 c10::kComplexFloat 空张量作为 result，然后 EXEC_NPU_CMD(aclnnPolar, input, angle, result)，最后通过 TORCH_LIBRARY(myops, ...) + PrivateUse1 后端分发注册为 torch 自定义算子，并导出 custom_op pybind 入口供 Python 侧调用。

九、与 S8 的差异及基线对照版本

与 S8 的差异（仅构建层）：

• S8：每算子自带 S8/<Op>/build.sh，auto_submit.sh 可直接 build；
• Polar：构建走 ops-math 仓的 build.sh --pkg --ops=polar（仓级构建），与本测试框架解耦；
• auto_submit.sh -o Polar 仍可复用其 upload/run 阶段；build 阶段需替换为 ops-math 命令（待算子工程落地后再接线）。

基线对照版本：仓库同时保留了 tests/pybind_baseline/ 目录，其 README.md 与 pybind 版描述一致（同样"Polar 测试框架（S8 同款）"），test_op.py 的 case_data 也包含 case1–16 全量用例，run.sh 的流程（重建 wheel → msprof → get_time → 基线校验、custom_math vendor 的 LD 指向、哨兵基线）与 pybind 版完全相同。两个目录可互为对照：pybind 作为主测试框架，pybind_baseline 作为基线参考实现（l0 拼接版）的验收载体，后续回填真实性能基线时可在两套框架间做横向比对。

十、注意事项与当前限制

使用该框架时需明确以下三点当前状态：

1. 精度阈值待回填：当前 verify_result 使用 fp32 分量默认阈值 rtol/atol=1e-4；任务要求对齐"AscendOpTest 默认阈值"，待 AscendOpTest 工具实际阈值确认后需在 test_op.py 的 verify_result 回填；
2. 性能基线为哨兵值：run.sh:time_base 目前是 9999999999999，仅作正确性门禁；硬件实测出 l0 参考实现耗时后需回填真实基线（验收要求 ≥ 基线 95%）；
3. 环境依赖：框架只能在 NPU 环境运行（依赖 torch_npu 及 custom_math vendor 算子包），本地无法自检——这一点与 S8 相同。

十一、进一步阅读

• 算子功能、参数与约束的完整定义见 experimental/math/polar/README.md（input/angle 为 FLOAT、ND 格式，out 为 COMPLEX64，维度不超过 8 维）；
• aclnn 接口调用样例见 examples/test_aclnn_polar.cpp，接口说明见 docs/aclnnPolar.md；
• 算子设计与自测报告见 docs/design.md 与 tests/自测报告.md；
• 广播场景的性能优化思路见 docs/性能优化_inner_broadcast.md。

综上，tests/pybind 框架以"S8 同款"的五文件结构，为 Polar 算子提供了一条从 NPU 算子包安装、pybind wheel 构建、16 类广播/边界 case 精度验证到 msprof 性能采样的完整验收闭环；EXEC_NPU_CMD 的动态 dlopen 机制使其可被任意算子零改动复用，而哨兵基线、待回填阈值等标注也清楚地划出了当前框架的完成度边界。