【CANN社区任务2026】用 Ascend C 把 Polar 算子做到基线 1.78×：六算子融合、complex64 交织与广播反超实录
CANN 训练营第二季社区任务 · 序号 04-5 Polar 目标：基于 Ascend C 实现 aclnnPolar，对齐开源仓 l0 拼接参考实现并新增广播支持，验收通过后贡献至 cann/ops-math。
一、任务是什么
torch.polar(abs, angle) 用极坐标构造复数：
输入 input（模长）、angle（弧度）都是 float32，输出是 complex64。规格看着很简单——逐元素、无 attr、各元素独立，天然适合多核并行。但真正动手时，三个"看着简单、做起来扎手"的点几乎占了我全部的调试时间：怎么把一串小算子融成一个 kernel、complex64 在 910B 上怎么摆、广播怎么在 kernel 里算而不退化成标量循环。这篇就把这三块的踩坑和破局过程记下来。
对齐基准很特殊：这个算子没有 TBE 历史实现，任务书明确"和 aclnn_polar.cpp 对齐即可"。所以我的性能基线就是开源仓里那份 l0 拼接参考实现（系统 libopapi.so 里已导出的 aclnnPolar）。
二、为什么能快：从"六个小算子拼接"到"一个融合 kernel"
先看参考实现干了什么。它是纯 host 端的 l0 API 拼接：
大约 6 个独立的设备算子 串起来。每一个都有自己的 kernel 启动 dispatch 开销，而且每一个中间结果都要在 GM（HBM）和 UB 之间搬一趟——粗算下来 GM↔UB 往返 ≥10 次。小 shape 下启动延迟就把计算时间盖住了，大 shape 下则是被这些中间结果的 HBM 带宽拖死。
Ascend C 的单 kernel 融合正好对症：
kernel 里的流水线就是标准三段式：CopyIn 一次把 input/angle 搬进 UB，Compute 在 UB 里链式跑完 Cos → Sin → Mul×2 → 交织，CopyOut 一次写出。中间结果一次都不落 HBM。
结果：16M（[4096,4096]）这个"所有核参与"的核心验收场景，l0 基线 1665µs，融合版 937.98µs，快 1.78×。融合本身就是最大的一块性能来源。
三、硬骨头一：complex64 在 910B 上怎么交织？
这是最坑的一块。910B（dav_c220）的 vector 单元不原生支持 complex64。complex64 在内存里是实部虚部交替存的 [r0, i0, r1, i1, …]，而我在 UB 里算出来的是分离的两段：一段全是实部、一段全是虚部。问题就是：怎么把分离的 real/imag 高效地交织成 [r,i,r,i,…] 写出去。
我一口气试了五种方案，前四种全部撞墙：
最终方案的关键是：交织本质上是一个固定的地址置换——第 j 个输出元素该从哪个源地址取，是一个只跟 tileLen 有关的静态函数：
偶数位取实部区、奇数位取虚部区。这个偏移表在 Init() 里用纯 int32 向量指令一次性构造好，之后每个 tile 都复用同一张表，Gather 一条指令就把分离的 (real区, imag区) 交织成 (r,i,r,i,…)。全程没有一次逐元素的标量操作。
四、硬骨头二：广播搬进 kernel —— 全矢量化 unravel，拒绝标量
任务的新增功能点是支持 numpy 广播（input 和 angle 维度可以不一致，如 [3,1,4] + [2,4] → [3,2,4]）。参考实现里广播是 l0 的 Mul 内部处理的；我要在单 kernel 里自己实现。
难点在于：广播意味着 output 的某个线性下标 idx，要反查回 input/angle 里对应的元素地址。最朴素的写法是逐元素做 unravel（把线性 idx 拆成多维坐标，再点乘各维 stride）——但逐元素标量在昇腾 vector 上是性能毒药。
我的做法是把整个 unravel 全程矢量化：host 侧先算好每个操作数在各输出轴上的元素 stride（被广播的轴 stride 记为 0），kernel 侧对一整个 tile 的输出线性下标，逐输出轴做：
整条链没有一个标量循环。这里有个小陷阱值得记：直接用浮点除法求坐标商会在边界上抖（floor(x/D) 偶尔差 1），所以加了两次 CompareScalar + Select 做精确修正——这比事后 debug"某些位置结果偏一格"要省心得多。
另外为了不让全形（未广播）场景背上广播的开销，tiling 里用 inFull/anFull 两个标志做运行期判断：操作数 numel == 输出 numel 就是没广播，走连续 DataCopyPad 快路径；否则走上面的矢量化 unravel。全路径统一一个 tilingKey（0），不为广播单开 funcEntry——轻量分支没有性能损失，却省掉了多套 kernel 的编译和维护成本。
五、硬骨头三：唯一两个"输给基线"的 case —— 80MB 的 HBM 陷阱
同 shape 和一般广播都做完、50 个验收 case 里 48 个都稳赢基线 之后，剩下两个死活过不去：
（这一段的基线是 torch.polar。）
这俩的共同点：angle 是 input 末维的周期性重复。我一开始在 op_api 层用 BroadcastTo 把 angle 物理展开到输出 shape（因为同 shape kernel 要求两个操作数一样大）。于是 Test_036 的 HBM 流量变成了：
而基线用的是 view（stride=0） 读 angle，压根不展开——省掉了那 16MB 的展开写，所以它反而更快。我是被自己的 BroadcastTo 拖累了。
破局的关键观察：angle 的周期长度 K（= angle.numel），每个 tile 取整数个周期时，cos/sin 的 tile pattern 对所有 tile 完全相同。那就没必要每个 tile 都重算，更没必要把 angle 物理展开。方案（我叫它"方案 A / period pre-compute"）：
- op_api：满足资格时 angle 不做 BroadcastTo，保留原 shape
[K] 直接传 kernel； - op_host tiling：识别这种 inner-broadcast 模式，按 K-block 分核，让每个核的负载是 K 的整数倍（tile 起点对齐周期）；
- op_kernel：
Process() 里加一个运行期分支（还是纯 data 字段，不用多注册 kernel）： PrecomputeBcastCosSin()：读一次 angle[0:K]，算出 cos/sin 的一个周期，循环填满一个常驻 tile buffer——整个核只做一次；ProcessTileBcast()：主循环里只读 input，Mul(cosTile, input) 得实部、Mul(sinTile, input) 得虚部，Gather 交织写出。没有 angle 搬运、没有每 tile 的 Sin/Cos 重算。
HBM 从 80MB 降到 48MB（省掉 angle 的 16MB 展开写 + 16MB 读 → 只剩 256B），Sin/Cos 从"每 tile 重算"降到"整核一次"。
结果：
从输基线直接反超到 3.5×，而且顺带让所有 angle 周期广播的 case 普遍再快约 2×。
六、验收结果
精度（官方 AscendOpTest，任务权威判据）：6/6 全 PASS，含 16M [4096,4096] 和全部广播用例（低维→高维、标量 input、双向广播、高维非 32B 对齐）。判据是实部、虚部各自纯绝对误差 ≤ 1e-4。全程 fp32 计算，complex64 经 ReinterpretCast 视作交织的 fp32 分量，无 Cast 中间精度损失。
性能（vs l0 参考基线）：
全场景优于基线，核心场景达基线 178%，远超"所有核参与场景 ≥95%"的门禁。泛化交叉集 16 个 case（8D 满秩、中间轴广播、双向多轴、标量×高维、负 abs、大角度归约、非对齐 inner）在严格复数判据下全 PASS。
七、几条可复用的心得
融合是访存受限小算子的第一性收益。参考 6 个小算子拼接的瓶颈不在算力，而在启动 dispatch + 中间结果的 HBM 往返。融合成单 kernel（启动 6→1、搬运 ≥10→2）就是最大的一块提速，其它优化都是锦上添花。
**数据重排优先想"静态地址置换 + Gather"**。别急着找专用交织指令（很可能是新架构独有），也别用 Transpose 硬凑。把置换表在 Init 里构造一次、全 tile 复用，又通用又快。
广播要全矢量化，逐元素标量是毒药。unravel、坐标修正、取数全用向量指令（Muls/Floor/CompareScalar/Select/Gather），host 把 stride 算好下发。fp 除法边界抖动用 CompareScalar+Select 精确修正。
小心自己引入的 HBM 流量。BroadcastTo 物理展开一时爽，但对周期性广播来说，那次展开写就是你输给"用 view/stride=0 的基线"的原因。能 pre-compute 周期 pattern 的，就别物理展开。
运行期分支 > 多 tilingKey。全形/广播/inner-broadcast 三条路都靠 tiling 里的纯 data 标志做运行期判断，统一一个 tilingKey，省掉多套 kernel 的编译与维护，轻量分支没有可测的性能损失。
"实部等于金标准的虚部"这类串位现象，先查 stride/连续读。DataCopy(srcStride=0) 不是"重复同段广播"，它是连续读——这个坑我调了很久才定位。
八、结语
Polar 这个算子规格上很朴素，但把它从"6 个小算子拼接"做成"1 个融合 kernel"，中间要跨过 complex64 交织、kernel 侧广播、周期广播的 HBM 陷阱三道坎。最终精度 6/6 全 PASS、核心场景 1.78× 基线、50 个 case 无一输给基线。希望这几条踩坑记录能给做昇腾算子的同学省点时间。
代码已贡献并合入 cann/ops-math：https://gitcode.com/cann/ops-math/pull/2923
感谢 CANN 训练营和社区评审的耐心 review。
点赞
收藏
回复
分享
举报
【CANN社区任务2026】用 Ascend C 把 Polar 算子做到基线 1.78×：六算子融合、complex64 交织与广播反超实录
一、任务是什么
torch.polar(abs, angle)用极坐标构造复数：out = abs · (cos(angle) + i·sin(angle)) 即 out.real = abs · cos(angle) out.imag = abs · sin(angle)输入
input（模长）、angle（弧度）都是 float32，输出是 complex64。规格看着很简单——逐元素、无 attr、各元素独立，天然适合多核并行。但真正动手时，三个"看着简单、做起来扎手"的点几乎占了我全部的调试时间：怎么把一串小算子融成一个 kernel、complex64 在 910B 上怎么摆、广播怎么在 kernel 里算而不退化成标量循环。这篇就把这三块的踩坑和破局过程记下来。对齐基准很特殊：这个算子没有 TBE 历史实现，任务书明确"和
aclnn_polar.cpp对齐即可"。所以我的性能基线就是开源仓里那份 l0 拼接参考实现（系统libopapi.so里已导出的aclnnPolar）。二、为什么能快：从"六个小算子拼接"到"一个融合 kernel"
先看参考实现干了什么。它是纯 host 端的 l0 API 拼接：
大约 6 个独立的设备算子 串起来。每一个都有自己的 kernel 启动 dispatch 开销，而且每一个中间结果都要在 GM（HBM）和 UB 之间搬一趟——粗算下来 GM↔UB 往返 ≥10 次。小 shape 下启动延迟就把计算时间盖住了，大 shape 下则是被这些中间结果的 HBM 带宽拖死。
Ascend C 的单 kernel 融合正好对症：
kernel 里的流水线就是标准三段式：
CopyIn一次把 input/angle 搬进 UB，Compute在 UB 里链式跑完Cos → Sin → Mul×2 → 交织，CopyOut一次写出。中间结果一次都不落 HBM。结果：16M（
[4096,4096]）这个"所有核参与"的核心验收场景，l0 基线 1665µs，融合版 937.98µs，快 1.78×。融合本身就是最大的一块性能来源。三、硬骨头一：complex64 在 910B 上怎么交织？
这是最坑的一块。910B（dav_c220）的 vector 单元不原生支持 complex64。complex64 在内存里是实部虚部交替存的
[r0, i0, r1, i1, …]，而我在 UB 里算出来的是分离的两段：一段全是实部、一段全是虚部。问题就是：怎么把分离的 real/imag 高效地交织成[r,i,r,i,…]写出去。我一口气试了五种方案，前四种全部撞墙：
Copy<float,true>+ maskTranspose（NCHW2NHWC）vintlv/Interleave指令DataCopyPadGather+ 静态偏移表最终方案的关键是：交织本质上是一个固定的地址置换——第 j 个输出元素该从哪个源地址取，是一个只跟 tileLen 有关的静态函数：
偶数位取实部区、奇数位取虚部区。这个偏移表在
Init()里用纯 int32 向量指令一次性构造好，之后每个 tile 都复用同一张表，Gather一条指令就把分离的(real区, imag区)交织成(r,i,r,i,…)。全程没有一次逐元素的标量操作。四、硬骨头二：广播搬进 kernel —— 全矢量化 unravel，拒绝标量
任务的新增功能点是支持 numpy 广播（input 和 angle 维度可以不一致，如
[3,1,4]+[2,4]→[3,2,4]）。参考实现里广播是 l0 的Mul内部处理的；我要在单 kernel 里自己实现。难点在于：广播意味着 output 的某个线性下标 idx，要反查回 input/angle 里对应的元素地址。最朴素的写法是逐元素做 unravel（把线性 idx 拆成多维坐标，再点乘各维 stride）——但逐元素标量在昇腾 vector 上是性能毒药。
我的做法是把整个 unravel 全程矢量化：host 侧先算好每个操作数在各输出轴上的元素 stride（被广播的轴 stride 记为 0），kernel 侧对一整个 tile 的输出线性下标，逐输出轴做：
整条链没有一个标量循环。这里有个小陷阱值得记：直接用浮点除法求坐标商会在边界上抖（
floor(x/D)偶尔差 1），所以加了两次CompareScalar + Select做精确修正——这比事后 debug"某些位置结果偏一格"要省心得多。另外为了不让全形（未广播）场景背上广播的开销，tiling 里用
inFull/anFull两个标志做运行期判断：操作数 numel == 输出 numel 就是没广播，走连续DataCopyPad快路径；否则走上面的矢量化 unravel。全路径统一一个 tilingKey（0），不为广播单开 funcEntry——轻量分支没有性能损失，却省掉了多套 kernel 的编译和维护成本。五、硬骨头三：唯一两个"输给基线"的 case —— 80MB 的 HBM 陷阱
同 shape 和一般广播都做完、50 个验收 case 里 48 个都稳赢基线 之后，剩下两个死活过不去：
[256,256,64]×[64][128,128,512]×[512]（这一段的基线是
torch.polar。）这俩的共同点：angle 是 input 末维的周期性重复。我一开始在 op_api 层用
BroadcastTo把 angle 物理展开到输出 shape（因为同 shape kernel 要求两个操作数一样大）。于是 Test_036 的 HBM 流量变成了：而基线用的是 view（stride=0） 读 angle，压根不展开——省掉了那 16MB 的展开写，所以它反而更快。我是被自己的 BroadcastTo 拖累了。
破局的关键观察：angle 的周期长度 K（= angle.numel），每个 tile 取整数个周期时，cos/sin 的 tile pattern 对所有 tile 完全相同。那就没必要每个 tile 都重算，更没必要把 angle 物理展开。方案（我叫它"方案 A / period pre-compute"）：
[K]直接传 kernel；Process()里加一个运行期分支（还是纯 data 字段，不用多注册 kernel）：PrecomputeBcastCosSin()：读一次angle[0:K]，算出 cos/sin 的一个周期，循环填满一个常驻 tile buffer——整个核只做一次；ProcessTileBcast()：主循环里只读 input，Mul(cosTile, input)得实部、Mul(sinTile, input)得虚部，Gather 交织写出。没有 angle 搬运、没有每 tile 的 Sin/Cos 重算。HBM 从 80MB 降到 48MB（省掉 angle 的 16MB 展开写 + 16MB 读 → 只剩 256B），Sin/Cos 从"每 tile 重算"降到"整核一次"。
结果：
从输基线直接反超到 3.5×，而且顺带让所有 angle 周期广播的 case 普遍再快约 2×。
六、验收结果
精度（官方 AscendOpTest，任务权威判据）：6/6 全 PASS，含 16M
[4096,4096]和全部广播用例（低维→高维、标量 input、双向广播、高维非 32B 对齐）。判据是实部、虚部各自纯绝对误差 ≤ 1e-4。全程 fp32 计算，complex64 经ReinterpretCast视作交织的 fp32 分量，无 Cast 中间精度损失。性能（vs l0 参考基线）：
[2,6,10][4096,4096]（核心验收）[4,1,8]×[4,5,8]全场景优于基线，核心场景达基线 178%，远超"所有核参与场景 ≥95%"的门禁。泛化交叉集 16 个 case（8D 满秩、中间轴广播、双向多轴、标量×高维、负 abs、大角度归约、非对齐 inner）在严格复数判据下全 PASS。
七、几条可复用的心得
融合是访存受限小算子的第一性收益。参考 6 个小算子拼接的瓶颈不在算力，而在启动 dispatch + 中间结果的 HBM 往返。融合成单 kernel（启动 6→1、搬运 ≥10→2）就是最大的一块提速，其它优化都是锦上添花。
**数据重排优先想"静态地址置换 + Gather"**。别急着找专用交织指令（很可能是新架构独有），也别用 Transpose 硬凑。把置换表在 Init 里构造一次、全 tile 复用，又通用又快。
广播要全矢量化，逐元素标量是毒药。unravel、坐标修正、取数全用向量指令（
Muls/Floor/CompareScalar/Select/Gather），host 把 stride 算好下发。fp 除法边界抖动用CompareScalar+Select精确修正。小心自己引入的 HBM 流量。
BroadcastTo物理展开一时爽，但对周期性广播来说，那次展开写就是你输给"用 view/stride=0 的基线"的原因。能 pre-compute 周期 pattern 的，就别物理展开。运行期分支 > 多 tilingKey。全形/广播/inner-broadcast 三条路都靠 tiling 里的纯 data 标志做运行期判断，统一一个 tilingKey，省掉多套 kernel 的编译与维护，轻量分支没有可测的性能损失。
"实部等于金标准的虚部"这类串位现象，先查 stride/连续读。
DataCopy(srcStride=0)不是"重复同段广播"，它是连续读——这个坑我调了很久才定位。八、结语
Polar 这个算子规格上很朴素，但把它从"6 个小算子拼接"做成"1 个融合 kernel"，中间要跨过 complex64 交织、kernel 侧广播、周期广播的 HBM 陷阱三道坎。最终精度 6/6 全 PASS、核心场景 1.78× 基线、50 个 case 无一输给基线。希望这几条踩坑记录能给做昇腾算子的同学省点时间。
代码已贡献并合入
cann/ops-math：https://gitcode.com/cann/ops-math/pull/2923感谢 CANN 训练营和社区评审的耐心 review。