算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载导读本文以 CANN ops-math 仓库中 Polar 算子设计文档 为核心完整讲解由极坐标模长与幅角构造复数张量的aclnnPolar算子从torch.polar语义、无 TBE 历史实现的背景到基于开源仓 l0 拼接参考实现、最终以 Ascend C 单 kernel 融合方案落地的全过程。读者将掌握该算子的功能规格与约束、host 侧 Tiling 与 kernel 侧三段式流水设计、复数交织的关键选型以及最终的精度与性能验收结果并能结合仓库源码experimental/math/polar目录深入理解每一处实现细节。一、需求背景极坐标构造复数的 Polar 算子1.1 算子来源与语义Polar 算子源自昇腾 CANN 训练营第二季社区任务序号04-5 Polar基于 Ascend C 实现aclnnPolar对齐开源仓 l0 拼接参考实现并新增广播支持验收通过后贡献至cann/ops-math。torch.polar(abs, angle)由极坐标构造复数out input·(cos(angle)i·sin(angle))即实部out.real input·cos(angle)、虚部out.imag input·sin(angle)。本算子即为该语义的 NPU 实现提供与 PyTorchtorch.polar等效的功能支持模型迁移。1.2 无 TBE 历史实现以 l0 拼接为对齐基准任务书明确说明本算子无 TBE 历史实现暂无 tbe 实现和aclnn_polar.cpp对齐即可因此以开源仓的 l0 拼接参考实现为对齐基准与性能基线。参考源码位于 math/polar/op_api/aclnn_polar.cpp依赖 l0 APIsin.h/cos.h/mul.h/complex.h其中l0op::Complex为已注册设备算子系统libopapi.so已导出aclnnPolar/aclnnPolarGetWorkspaceSize接口。1.3 l0 参考实现现状分析参考实现的参数规格如下参数含义支持 dtype约束形状input模长 absfloat32与 angle 同 dtype≤8维angle角度(弧度)float32与 input 同 dtype≤8维可与 input 广播out复数结果complex64恒 complex64broadcast(input,angle)参考实现逻辑与 math/polar/op_api/aclnn_polar.cpp 源码一致aclnnPolarGetWorkspaceSize为纯 host 端 l0 拼接 ——CheckParams空指针/dtype/广播 shape 校验→Contiguous×2 →Sin(angle)、Cos(angle)→Mul(·, input)×2 得虚/实部 →Complex构造复数 →ViewCopy写回输出约束为DTYPE{DT_FLOAT}、OUTPUT{DT_COMPLEX64}、dtype 须匹配、MAX_DIM8、广播经OP_CHECK_BROADCAST_AND_INFER_SHAPE推导。核心问题该实现将计算拆成Contiguous → Sin → Cos → Mul → Mul → Complex约 6 个独立设备算子每个算子都有 kernel 启动 dispatch 开销中间结果必经 GM↔UB 搬运≥10 次。在小 shape 场景下启动延迟远超计算本身带宽也被白白浪费。这正是引入 Ascend C 单 kernel 融合优化的动机。参考l0实现流程图与源码逻辑一致二、算子功能规格与原型2.1 参数规格名称类别dtypeformat介绍input输入float32ND极坐标模长 absangle输入float32ND极坐标幅角(弧度)dtype 须与 input 一致out输出complex64ND复数 tensorreal/imag 交替存储数学公式out input·cos(angle) i·input·sin(angle)无属性(attr)。支持广播input 与 angle 维度可不一致遵循 numpy 广播规则这是本任务的新增功能点。仓库中的算子原型定义在 polar_def.cppinput、angle均声明为ge::DT_FLOATFORMAT_ND输出out恒为ge::DT_COMPLEX64FORMAT_ND三个 tensor 均带AutoContiguous()并通过AICore().AddConfig(ascend910b)/AddConfig(ascend910_93)声明双平台支持。该定义与设计文档的规格完全一致。2.2 相关约束平台Atlas A2 训练系列产品ascend910b/ Atlas A3 训练系列产品ascend910_93双平台支持dtypeinput 与 angle 必须一致float32out 恒 complex64维度≤8 维angle 取值为弧度、无范围限制三角函数为周期函数不支持fp16/bf16/fp64/complex128任务范围外参考实现亦不支持。算子测试描述文件 tests/Polar.json 与上述规格一致输入float/ND、输出complex64/ND。三、需求分析与总体设计思路3.1 需求拆解设计文档将任务拆解为四点实现Cos Sin Mul×2 复数交织全融合 kernel消除小算子拼接瓶颈广播语义的支持input 与 angle 维度可不一致numpy 规则fp32(input/angle) → complex64(out)与参考一致≤8 维性能 ≥ 参考 l0 的 95%所有核参与场景为核心验收精度满足 AscendOpTest 默认阈值。3.2 总体思路单 kernel 融合参考实现的痛点在于多次 kernel dispatch 与中间结果的 GM↔UB 往返。Ascend C 单 kernel 融合方案可以做到CopyIn 一次搬入、Compute 在 UB 内链式完成Cos → Sin → Mul → 交织、中间结果驻留 UB、CopyOut 一次写出 ——1 次启动、2 次 GM↔UB 搬运。算子特性分析来自设计文档数学公式out.real input·cos(angle)out.imag input·sin(angle)算子特性逐元素、各元素独立天然多核并行确定性计算无 Reduce/Scatter每输出元素 2 次三角函数 2 次乘法、读 2 输入写 1 复数输出元素间无依赖无需跨核同步单核流水线可满排数据类型input、angle float32须同 dtypeout complex64。910B vector 单元不原生支持 complex64经ReinterpretCast视作 2N 个交织 fp32 处理支持形状ND≤8 维out.shape broadcast(input, angle)。算子分 Host 侧参数校验、广播 shape/stride 计算、Tiling 下发与 Kernel 侧Init Process 三段式 CopyIn→Compute→CopyOut两部分实现。四、host 侧设计详解4.1 InferShape / InferDataTypeInferShape按 numpy 广播规则推导 out shape右对齐各轴取 max对齐参考实现的OP_CHECK_BROADCAST_AND_INFER_SHAPEInferDataTypeout 恒置DT_COMPLEX64与 input dtype 无关。源码实现见 polar_infershape.cppInferShape4Polar取两输入 rank 的较大者为输出 rank逐轴对右对齐的维度取 max缺失轴视为 1InferDataType4Polar直接SetOutputDataType(0, ge::DT_COMPLEX64)。4.2 广播 shape 计算广播计算遵循右对齐、缺失轴首部补 1逐维相等或其一为 1 否则失败每维取 max。示例[3,1,4][2,4]→[3,2,4]。host 侧对 input/angle 分别计算其在各输出轴的元素 stride广播轴 0并判定全形快路径标志操作数 numel out numel ⟺ 未广播。4.3 分核策略优先满核取平台 AIV 核数小 case 少核每核 ≥MIN_PER_CORE 2048元素降低 launch/同步开销强制偶数核vector core 两两绑定2 时下取偶。能均分则无大小核之分不能均分时余块分到前几核前rem核各per1其余核各per。4.4 数据分块与内存优化充分利用 UB910B 192KB。初始设计采用BUFFER_NUM 1尾块同路径处理DataCopyPad 任意字节无碎片。tileLen计算固定占用 ≈52×tileLen 字节分布于 8 个 UB buffer ——qAbs/qAng各tileLen·4BVECIN、qOut/bufPacked/bufOff各2·tileLen·4B复数/实虚 packed/交织偏移表、bufBC5·tileLen·4B矢量化 unravel scratch、bufInFull/bufAnFull广播操作数 numel×4B全形时仅 32B 占位。取tileLen 2048固定 ≈104KB ≤ 192KB余量留给广播操作数整块载入。tile 循环为while(done coreLen) { n min(coreLen - done, tileLen); …; done n; }。注TPipe 单 kernel InitBuffer 张量上限为 8BUFFER_NUM2 时每个计 2初版设计 BUFFER_NUM1、8 个 buffer 恰在上限超限会直接 run failed 且无报错这是 UB buffer 规划时必须严格看守的约束。4.5 tilingkey 策略与 TilingData 字段tilingkey 策略单 kernel 全路径全形/广播按 tiling 的inFull/anFull标志运行期判断不设多 tilingkey统一 0—— 避免多 funcEntry 编译/维护成本轻量分支无性能损失。仅 fp32→complex64OpDef 已限定 dtype无 bf16 兜底。TilingData 字段totalLen, tileLen, bigCoreNum/coreNum(偶数), bigCoreLen/smallCoreLen, outRank, inFull/anFull(全形快路径标志), inNumel/anNumel(广播整块载入用), outDims[8], inStr[8]/anStr8其中outDims/inStr/anStr经TILING_DATA_FIELD_DEF_ARR宏声明。五、kernel 侧设计详解单 kernelInit ProcessCopyIn/Compute/CopyOut全程 fp32 计算。kernel 入口见 polar.cpp核心实现见 polar.h。5.1 Init 阶段按blockIdx计算本核区间大核/小核分块BuildOffsetTable()纯 int32 矢量构造交织静态偏移表off[j] 4·(j1) (j1)·4T一次性构造、全 tile 复用、无逐元素标量计算广播操作数经LoadFull整块载入 UB。5.2 CopyInFetchOperand全形操作数→ 连续DataCopyPad快路径广播操作数→ 矢量化 unravel对 tile 输出线性 idx逐输出轴Muls(1/D) Floor求商加两次CompareScalar Select精确 ±1 修正算坐标累加偏移 →Castint32 →Gather从 UB 整块操作数取数全矢量无逐元素标量。5.3 ComputeCos(angle)→ packed 实部区、Sin(angle)→ packed 虚部区 → 原地Mul × abs→Gather用静态偏移表交织为(r, i, r, i, …)。5.4 CopyOutDataCopyPad写 complex64 GM视作 2N 个 fp32 大块对齐。Ascend C 实现流程图5.5 复数交织方案选型910Bdav_c220的输出需要将分离的 real/imag 交织为[r0, i0, r1, i1, …]设计文档对比了多种候选方案方案结果原因Copyfloat, truemask❌mask 同时作用于 src/dst无法错位TransposeNCHW2NHWC✅但慢受 16×16 分形约束占绝大部分耗时vintlv/Interleave❌dav_c310 独有c220 不支持stridedDataCopyPad❌VECOUT 32B 粒度逐 4B 写出对 MTE 是灾难Gather 静态偏移表✅ 选用纯向量指令偏移表 Init 一次构造、全 tile 复用最终选型为Gather 静态偏移表。在落地代码 polar.h 的BuildOffsetTable中可看到具体实现CreateVecIndex生成 0..2T-1 索引经ShiftRight/Muls/Add组合出字节偏移off[j] 4T·j (4-8T)·(j1)InterleaveAndStore中Gather(outL, packed, off32, 0, 2n)一步完成交织。5.6 与参考流程图的差异点#参考 l0本 Ascend C原因1约 6 个独立设备算子串接各 HBM 往返 dispatch单 kernel 融合省 ≥5 次中间 HBM 多次 dispatch核心场景实测快 1.78×2l0op::Complex构造复数Gather 静态偏移表交织910B vector 不支持 complex64规避 strided/Transpose 约束3广播由 l0 Mul 内部处理host 算 stride kernel 矢量化 unravel原生融合需自实现广播全矢量避免标量4Contiguous ViewCopy 处理非连续框架传连续 GM 无需aclnn/测试输入连续省拷贝融合优势汇总kernel 启动约 6 次 → 1 次GM↔UB 搬运 ≥10 次 → 2 次中间结果留在片上GM→UB 消除angle 一次读入后 cos/sin 同 tile 计算参考实现无法复用。六、落地实现设计文档到最终代码的演进设计文档描述的初版方案是广播完全迁移到 kernel 侧host 算 stride、kernel 内矢量化 unravel。而从仓库最终代码看落地时采取了aclnn L2 层 同 shape 单 kernel inner-broadcast 快路径的优化分工这是一处值得关注的实现演进6.1 L2 层Contiguous BroadcastTo 单 kernel 直调experimental/math/polar/op_api/aclnn_polar.cpp 的aclnnPolarGetWorkspaceSize流程为CheckParamsCheckNotNull空指针→CheckDtypeValidinput/angle 为DT_FLOAT且一致、out 为DT_COMPLEX64→CheckBroadcastShape≤8 维、OP_CHECK_BROADCASTBroadcastInferShape与 out shape 比对空张量短路input-IsEmpty() || angle-IsEmpty()时*workspaceSize 0直接返回Contiguous ×2解决 view/transpose 等非连续输入广播处理input 恒BroadcastTo到 out shapeangle 满足inner-broadcast 资格时保留原 shape[K]kernelbcastMode1周期复用省一次 BroadcastTo 的 HBM 写、每 tile 的 angle 搬运及 Sin/Cos 重算否则也BroadcastTo兜底。资格判定与 polar_tiling.cpp 严格一致inN outN, anN inN, inN % anN 0, anN 2048, anN % 8 032B 对齐l0op::Polar调用同 shape 单 kernelkernel 假设 input.shape angle.shape out.shapeViewCopy把 kernel 输出拷到调用者提供的 out。6.2 Tiling 侧bcastMode 双路径分核polar_tiling.cpp 中IsAngleInnerBcast判定bcastModebcastMode 1angle 周期 K 远小于 inputSplitCoresBcast按K-block 分核——每核负载是 K 的整数倍coreStart % K 0保证 tile 起点与 period 对齐前remBlk核各多分 1 个 K 块bcastMode 0同 shape elementwiseSplitCoresElementwise按MIN_PER_CORE 2048收敛核数 ClampWantEven强制偶数核。CalcSinCosTmpSize通过GetCosMaxMinTmpSize取maxValue而非 minValue 计算 Sin/Cos 显式 sharedTmpBufferminValue 仅勉强够用Sin/Cos 会按小 buffer 分块多趟计算大 same-shape case 性能会暴跌文档记录实测 16M same-shapecosMin 约 5800µs → cosMax 约 650µs约 9 倍差距再 clamp 到 UB 上限防越界。6.3 Kernel 侧双 buffer 周期复用快路径最终 kernelpolar.h采用BUFFER_NUM2 双 buffer资源布局qInVECIN2 buffer × (2×T fp32)合并存放[abs(0:n) | ang(T:Tn)]CopyIn 阶段DataCopyPad input → qIn[0:n]、DataCopyPad angle → qIn[T:Tn]Compute 阶段 Sin/Cos 取qIn[T:Tn]、Mul 取qIn[0:n]qOutVECOUT2 buffer × (2×T fp32)存放 complex64 交织结果bufPackedVECCALC1 buffer × (2×T fp32)[cos | sin]中间结果bufOffVECCALC1 buffer × (2×T u32)Gather 静态 offset 表bufSinCosTmpVECCALCtilingData.tmpBufferSizeSin/Cos 显式 sharedTmpBuffer。TPipe Tensor 数 qIn(2) qOut(2) bufPacked(1) bufOff(1) bufSinCosTmp(1) 7≤8 上限UB 用量 32 32 16 16 24 120 KB / 192 KB。bcastMode 1时走PrecomputeBcastCosSin快路径Init 阶段对angle[0:K]只做一次 Cos/Sin并把 cosTile/sinTile 周期复制填满bufSinCosTmp_的[0:T]/[T:2T]常驻区主循环ProcessTileBcast只读 input HBM每 tile 仅Mul × 2 → Gather彻底省去每 tile 的 angle 搬运与三角函数重算。off必 K 对齐、n是 K 整数倍由 host K-block 分核与主循环共同保证。该分工的收益是任意 shape / 任意广播组合≤8 维都能正确处理不会因 outNumel 过大或中间轴广播触发 kernel 内 fp32 unravel 失精或 UB 容量越界见 aclnn_polar.cpp 文件头注释。七、支持硬件与使能方式芯片版本勾选Atlas A2 训练系列产品 / Atlas A2 推理系列产品 (ascend910b)√Atlas A3 训练系列产品 / Atlas A3 推理系列产品 (ascend910_93)√上层框架勾选TF 训练/推理Pytorch 训练/推理√pybind 调 aclnn 测试ATC 推理Aclnn 直调√OPAT 调优SGAT 子图切分主要适配ACLNN调用框架。仓库 README.md 的产品支持表与之呼应并额外说明 Ascend 950PR/950DT、Atlas 200I/500 A2 推理产品等不支持。aclnn 接口遵循 CANN 标准的两段式规范先aclnnPolarGetWorkspaceSize获取 workspace 与 executor再aclnnPolar执行完整接口说明与调用示例见 docs/aclnnPolar.md可直接编译运行的 C 样例见 examples/test_aclnn_polar.cpp。调用样例中的核心流程详见 aclnnPolar.md// 第一段入参校验 workspace 计算 ret aclnnPolarGetWorkspaceSize(input, angle, out, workspaceSize, executor); CHECK_RET(ret ACLNN_SUCCESS, ...); // 申请 workspace void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); ... } // 第二段执行计算 ret aclnnPolar(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACLNN_SUCCESS, ...); ret aclrtSynchronizeStream(stream);注意 out 为 complex64每元素占 real/imag 两个 float申请与读取输出内存时按numel × 2 × sizeof(float)计算字节数。返回值与错误码第一段接口完成入参校验出错场景返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001input、angle、out 是空指针ACLNN_ERR_PARAM_INVALID161002数据类型/格式不在支持范围、input 与 angle dtype 不一致、shape 超过 8 维、不满足广播推导、out shape 与广播结果不一致或非 COMPLEX64八、算子约束限制dtypeinput/angle 仅 float32须同 dtypeout 仅 complex64。不支持 fp16/bf16/fp64/complex128≤8 维input/angle 须满足 numpy 广播广播操作数 numel 需可整体载入 UB验收 shape 量级满足超大广播待分块优化——落地代码以 L2 层 BroadcastTo 兜底规避了该限制无 attr确定性计算aclnnPolar 默认确定性实现。九、特性交叉分析本算子为逐元素计算输出仅依赖对应位置输入不涉及跨元素聚合Reduce/Scatter、数据搬移Reshape/Transpose 语义、量化Quant与现有特性无冲突。广播与各类边界条件已交叉验证8D 满秩同形、5D 中间轴广播、双向多轴广播、标量 input×高维 angle、负 abs、大角度归约、非 32B 对齐 inner广播、大向量×标量等用例case9–16均严格判据 PASS无异常。十、可维可测分析精度与性能验收10.1 精度标准与实测验收标准描述来源精度满足 AscendOpTest 默认阈值不低于参考fp32→complex64 实/虚部各 max_abs1e-4、rate1e-4AscendOpTestHIT1920/AscendOpTest性能所有核参与场景 ≥ 参考 l0 的 95%小 shape 提供仿真分析《算子任务书》精度细化与实测complex64 在 AscendOpTestaccuracy_config无内置默认值因此用例 JSON 显式配置err_threshold[1e-4,1e-4]fp32 分量默认compare_complex为实/虚部各自纯绝对误差判定。官方 AscendOpTest 实跑 6 个用例同 shape 小 shape/16M 广播低→高/标量/双向 高维非对齐全部 PASS含 16M[4096,4096]。全程 fp32 计算不引入 Cast 中间精度损失complex64 经 ReinterpretCast 视作交织 fp32 分量精度等价 fp32。10.2 性能实测性能对比vs 系统 l0 参考基线官方 AscendOpTest msprof / 每调用设备时场景l0 基线本算子结论小 [2,6,10]14.25 µs12.52 µs更快 ✓16M [4096,4096]所有核1665 µs937.98 µs快 1.78×核心验收达标✓✓广播 [4,1,8]×[4,5,8]15.17 µs14.92 µs更快 ✓全场景优于 l0 参考实现远超 ≥95% 的验收要求。详细自测过程与数据可参考 自测报告 与 测试步骤指导。10.3 兼容性分析本算子为新增算子昇腾仓原无 Ascend C 原生 Polar亦无 TBE 历史实现不涉及历史兼容性问题。aclnn 接口遵循 CANN 标准规范提供与 PyTorchtorch.polar等效功能可支撑模型迁移。总结Polar 算子是一个典型的小算子拼接 → 单 kernel 融合优化案例从约 6 个 l0 设备算子串接≥10 次 GM↔UB 搬运、6 次 dispatch到 Ascend C 单 kernel 内完成Cos/Sin Mul×2 复数交织2 次搬运、1 次启动核心场景 16M 性能提升 1.78×同时通过 aclnn L2 层BroadcastTo Contiguous与 kernel 侧bcastModeinner-broadcast 周期复用快路径的合理分工完整支持 numpy 广播语义。其设计文档、接口文档、可运行样例与测试报告均随算子沉淀在仓库experimental/math/polar目录中可作为昇腾 Ascend C 算子融合开发与广播语义实现的可参考范例。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math AcosGradV2 算子全解析Ascend A2 上反余弦梯度计算与单 Kernel 融合实现CANN ops math AcosGradV2 算子全解析Ascend A2 上反余弦梯度计算与单 Kernel 融合实现 导读 本文基于 AcosGrad算子库人工智能CANNCANN ops-math Lerp 算子全解析从 aclnnLerp 接口到 Ascend 内核实现CANN ops math Lerp 算子全解析从 aclnnLerp 接口到 Ascend 内核实现 LerpLinear Interpolation线算子库人工智能CANN推荐项目ShadowImageView - 简单易用的阴影效果图像视图组件推荐项目ShadowImageView 简单易用的阴影效果图像视图组件 项目简介 在Android开发中为UI元素添加阴影效果可以提升界面的视觉吸引力和层次算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN ops-math Polar 算子设计解读：从 l0 算子拼接到 Ascend C 单 Kernel 融合实现
算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载导读本文以 CANN ops-math 仓库中 Polar 算子设计文档 为核心完整讲解由极坐标模长与幅角构造复数张量的aclnnPolar算子从torch.polar语义、无 TBE 历史实现的背景到基于开源仓 l0 拼接参考实现、最终以 Ascend C 单 kernel 融合方案落地的全过程。读者将掌握该算子的功能规格与约束、host 侧 Tiling 与 kernel 侧三段式流水设计、复数交织的关键选型以及最终的精度与性能验收结果并能结合仓库源码experimental/math/polar目录深入理解每一处实现细节。一、需求背景极坐标构造复数的 Polar 算子1.1 算子来源与语义Polar 算子源自昇腾 CANN 训练营第二季社区任务序号04-5 Polar基于 Ascend C 实现aclnnPolar对齐开源仓 l0 拼接参考实现并新增广播支持验收通过后贡献至cann/ops-math。torch.polar(abs, angle)由极坐标构造复数out input·(cos(angle)i·sin(angle))即实部out.real input·cos(angle)、虚部out.imag input·sin(angle)。本算子即为该语义的 NPU 实现提供与 PyTorchtorch.polar等效的功能支持模型迁移。1.2 无 TBE 历史实现以 l0 拼接为对齐基准任务书明确说明本算子无 TBE 历史实现暂无 tbe 实现和aclnn_polar.cpp对齐即可因此以开源仓的 l0 拼接参考实现为对齐基准与性能基线。参考源码位于 math/polar/op_api/aclnn_polar.cpp依赖 l0 APIsin.h/cos.h/mul.h/complex.h其中l0op::Complex为已注册设备算子系统libopapi.so已导出aclnnPolar/aclnnPolarGetWorkspaceSize接口。1.3 l0 参考实现现状分析参考实现的参数规格如下参数含义支持 dtype约束形状input模长 absfloat32与 angle 同 dtype≤8维angle角度(弧度)float32与 input 同 dtype≤8维可与 input 广播out复数结果complex64恒 complex64broadcast(input,angle)参考实现逻辑与 math/polar/op_api/aclnn_polar.cpp 源码一致aclnnPolarGetWorkspaceSize为纯 host 端 l0 拼接 ——CheckParams空指针/dtype/广播 shape 校验→Contiguous×2 →Sin(angle)、Cos(angle)→Mul(·, input)×2 得虚/实部 →Complex构造复数 →ViewCopy写回输出约束为DTYPE{DT_FLOAT}、OUTPUT{DT_COMPLEX64}、dtype 须匹配、MAX_DIM8、广播经OP_CHECK_BROADCAST_AND_INFER_SHAPE推导。核心问题该实现将计算拆成Contiguous → Sin → Cos → Mul → Mul → Complex约 6 个独立设备算子每个算子都有 kernel 启动 dispatch 开销中间结果必经 GM↔UB 搬运≥10 次。在小 shape 场景下启动延迟远超计算本身带宽也被白白浪费。这正是引入 Ascend C 单 kernel 融合优化的动机。参考l0实现流程图与源码逻辑一致二、算子功能规格与原型2.1 参数规格名称类别dtypeformat介绍input输入float32ND极坐标模长 absangle输入float32ND极坐标幅角(弧度)dtype 须与 input 一致out输出complex64ND复数 tensorreal/imag 交替存储数学公式out input·cos(angle) i·input·sin(angle)无属性(attr)。支持广播input 与 angle 维度可不一致遵循 numpy 广播规则这是本任务的新增功能点。仓库中的算子原型定义在 polar_def.cppinput、angle均声明为ge::DT_FLOATFORMAT_ND输出out恒为ge::DT_COMPLEX64FORMAT_ND三个 tensor 均带AutoContiguous()并通过AICore().AddConfig(ascend910b)/AddConfig(ascend910_93)声明双平台支持。该定义与设计文档的规格完全一致。2.2 相关约束平台Atlas A2 训练系列产品ascend910b/ Atlas A3 训练系列产品ascend910_93双平台支持dtypeinput 与 angle 必须一致float32out 恒 complex64维度≤8 维angle 取值为弧度、无范围限制三角函数为周期函数不支持fp16/bf16/fp64/complex128任务范围外参考实现亦不支持。算子测试描述文件 tests/Polar.json 与上述规格一致输入float/ND、输出complex64/ND。三、需求分析与总体设计思路3.1 需求拆解设计文档将任务拆解为四点实现Cos Sin Mul×2 复数交织全融合 kernel消除小算子拼接瓶颈广播语义的支持input 与 angle 维度可不一致numpy 规则fp32(input/angle) → complex64(out)与参考一致≤8 维性能 ≥ 参考 l0 的 95%所有核参与场景为核心验收精度满足 AscendOpTest 默认阈值。3.2 总体思路单 kernel 融合参考实现的痛点在于多次 kernel dispatch 与中间结果的 GM↔UB 往返。Ascend C 单 kernel 融合方案可以做到CopyIn 一次搬入、Compute 在 UB 内链式完成Cos → Sin → Mul → 交织、中间结果驻留 UB、CopyOut 一次写出 ——1 次启动、2 次 GM↔UB 搬运。算子特性分析来自设计文档数学公式out.real input·cos(angle)out.imag input·sin(angle)算子特性逐元素、各元素独立天然多核并行确定性计算无 Reduce/Scatter每输出元素 2 次三角函数 2 次乘法、读 2 输入写 1 复数输出元素间无依赖无需跨核同步单核流水线可满排数据类型input、angle float32须同 dtypeout complex64。910B vector 单元不原生支持 complex64经ReinterpretCast视作 2N 个交织 fp32 处理支持形状ND≤8 维out.shape broadcast(input, angle)。算子分 Host 侧参数校验、广播 shape/stride 计算、Tiling 下发与 Kernel 侧Init Process 三段式 CopyIn→Compute→CopyOut两部分实现。四、host 侧设计详解4.1 InferShape / InferDataTypeInferShape按 numpy 广播规则推导 out shape右对齐各轴取 max对齐参考实现的OP_CHECK_BROADCAST_AND_INFER_SHAPEInferDataTypeout 恒置DT_COMPLEX64与 input dtype 无关。源码实现见 polar_infershape.cppInferShape4Polar取两输入 rank 的较大者为输出 rank逐轴对右对齐的维度取 max缺失轴视为 1InferDataType4Polar直接SetOutputDataType(0, ge::DT_COMPLEX64)。4.2 广播 shape 计算广播计算遵循右对齐、缺失轴首部补 1逐维相等或其一为 1 否则失败每维取 max。示例[3,1,4][2,4]→[3,2,4]。host 侧对 input/angle 分别计算其在各输出轴的元素 stride广播轴 0并判定全形快路径标志操作数 numel out numel ⟺ 未广播。4.3 分核策略优先满核取平台 AIV 核数小 case 少核每核 ≥MIN_PER_CORE 2048元素降低 launch/同步开销强制偶数核vector core 两两绑定2 时下取偶。能均分则无大小核之分不能均分时余块分到前几核前rem核各per1其余核各per。4.4 数据分块与内存优化充分利用 UB910B 192KB。初始设计采用BUFFER_NUM 1尾块同路径处理DataCopyPad 任意字节无碎片。tileLen计算固定占用 ≈52×tileLen 字节分布于 8 个 UB buffer ——qAbs/qAng各tileLen·4BVECIN、qOut/bufPacked/bufOff各2·tileLen·4B复数/实虚 packed/交织偏移表、bufBC5·tileLen·4B矢量化 unravel scratch、bufInFull/bufAnFull广播操作数 numel×4B全形时仅 32B 占位。取tileLen 2048固定 ≈104KB ≤ 192KB余量留给广播操作数整块载入。tile 循环为while(done coreLen) { n min(coreLen - done, tileLen); …; done n; }。注TPipe 单 kernel InitBuffer 张量上限为 8BUFFER_NUM2 时每个计 2初版设计 BUFFER_NUM1、8 个 buffer 恰在上限超限会直接 run failed 且无报错这是 UB buffer 规划时必须严格看守的约束。4.5 tilingkey 策略与 TilingData 字段tilingkey 策略单 kernel 全路径全形/广播按 tiling 的inFull/anFull标志运行期判断不设多 tilingkey统一 0—— 避免多 funcEntry 编译/维护成本轻量分支无性能损失。仅 fp32→complex64OpDef 已限定 dtype无 bf16 兜底。TilingData 字段totalLen, tileLen, bigCoreNum/coreNum(偶数), bigCoreLen/smallCoreLen, outRank, inFull/anFull(全形快路径标志), inNumel/anNumel(广播整块载入用), outDims[8], inStr[8]/anStr8其中outDims/inStr/anStr经TILING_DATA_FIELD_DEF_ARR宏声明。五、kernel 侧设计详解单 kernelInit ProcessCopyIn/Compute/CopyOut全程 fp32 计算。kernel 入口见 polar.cpp核心实现见 polar.h。5.1 Init 阶段按blockIdx计算本核区间大核/小核分块BuildOffsetTable()纯 int32 矢量构造交织静态偏移表off[j] 4·(j1) (j1)·4T一次性构造、全 tile 复用、无逐元素标量计算广播操作数经LoadFull整块载入 UB。5.2 CopyInFetchOperand全形操作数→ 连续DataCopyPad快路径广播操作数→ 矢量化 unravel对 tile 输出线性 idx逐输出轴Muls(1/D) Floor求商加两次CompareScalar Select精确 ±1 修正算坐标累加偏移 →Castint32 →Gather从 UB 整块操作数取数全矢量无逐元素标量。5.3 ComputeCos(angle)→ packed 实部区、Sin(angle)→ packed 虚部区 → 原地Mul × abs→Gather用静态偏移表交织为(r, i, r, i, …)。5.4 CopyOutDataCopyPad写 complex64 GM视作 2N 个 fp32 大块对齐。Ascend C 实现流程图5.5 复数交织方案选型910Bdav_c220的输出需要将分离的 real/imag 交织为[r0, i0, r1, i1, …]设计文档对比了多种候选方案方案结果原因Copyfloat, truemask❌mask 同时作用于 src/dst无法错位TransposeNCHW2NHWC✅但慢受 16×16 分形约束占绝大部分耗时vintlv/Interleave❌dav_c310 独有c220 不支持stridedDataCopyPad❌VECOUT 32B 粒度逐 4B 写出对 MTE 是灾难Gather 静态偏移表✅ 选用纯向量指令偏移表 Init 一次构造、全 tile 复用最终选型为Gather 静态偏移表。在落地代码 polar.h 的BuildOffsetTable中可看到具体实现CreateVecIndex生成 0..2T-1 索引经ShiftRight/Muls/Add组合出字节偏移off[j] 4T·j (4-8T)·(j1)InterleaveAndStore中Gather(outL, packed, off32, 0, 2n)一步完成交织。5.6 与参考流程图的差异点#参考 l0本 Ascend C原因1约 6 个独立设备算子串接各 HBM 往返 dispatch单 kernel 融合省 ≥5 次中间 HBM 多次 dispatch核心场景实测快 1.78×2l0op::Complex构造复数Gather 静态偏移表交织910B vector 不支持 complex64规避 strided/Transpose 约束3广播由 l0 Mul 内部处理host 算 stride kernel 矢量化 unravel原生融合需自实现广播全矢量避免标量4Contiguous ViewCopy 处理非连续框架传连续 GM 无需aclnn/测试输入连续省拷贝融合优势汇总kernel 启动约 6 次 → 1 次GM↔UB 搬运 ≥10 次 → 2 次中间结果留在片上GM→UB 消除angle 一次读入后 cos/sin 同 tile 计算参考实现无法复用。六、落地实现设计文档到最终代码的演进设计文档描述的初版方案是广播完全迁移到 kernel 侧host 算 stride、kernel 内矢量化 unravel。而从仓库最终代码看落地时采取了aclnn L2 层 同 shape 单 kernel inner-broadcast 快路径的优化分工这是一处值得关注的实现演进6.1 L2 层Contiguous BroadcastTo 单 kernel 直调experimental/math/polar/op_api/aclnn_polar.cpp 的aclnnPolarGetWorkspaceSize流程为CheckParamsCheckNotNull空指针→CheckDtypeValidinput/angle 为DT_FLOAT且一致、out 为DT_COMPLEX64→CheckBroadcastShape≤8 维、OP_CHECK_BROADCASTBroadcastInferShape与 out shape 比对空张量短路input-IsEmpty() || angle-IsEmpty()时*workspaceSize 0直接返回Contiguous ×2解决 view/transpose 等非连续输入广播处理input 恒BroadcastTo到 out shapeangle 满足inner-broadcast 资格时保留原 shape[K]kernelbcastMode1周期复用省一次 BroadcastTo 的 HBM 写、每 tile 的 angle 搬运及 Sin/Cos 重算否则也BroadcastTo兜底。资格判定与 polar_tiling.cpp 严格一致inN outN, anN inN, inN % anN 0, anN 2048, anN % 8 032B 对齐l0op::Polar调用同 shape 单 kernelkernel 假设 input.shape angle.shape out.shapeViewCopy把 kernel 输出拷到调用者提供的 out。6.2 Tiling 侧bcastMode 双路径分核polar_tiling.cpp 中IsAngleInnerBcast判定bcastModebcastMode 1angle 周期 K 远小于 inputSplitCoresBcast按K-block 分核——每核负载是 K 的整数倍coreStart % K 0保证 tile 起点与 period 对齐前remBlk核各多分 1 个 K 块bcastMode 0同 shape elementwiseSplitCoresElementwise按MIN_PER_CORE 2048收敛核数 ClampWantEven强制偶数核。CalcSinCosTmpSize通过GetCosMaxMinTmpSize取maxValue而非 minValue 计算 Sin/Cos 显式 sharedTmpBufferminValue 仅勉强够用Sin/Cos 会按小 buffer 分块多趟计算大 same-shape case 性能会暴跌文档记录实测 16M same-shapecosMin 约 5800µs → cosMax 约 650µs约 9 倍差距再 clamp 到 UB 上限防越界。6.3 Kernel 侧双 buffer 周期复用快路径最终 kernelpolar.h采用BUFFER_NUM2 双 buffer资源布局qInVECIN2 buffer × (2×T fp32)合并存放[abs(0:n) | ang(T:Tn)]CopyIn 阶段DataCopyPad input → qIn[0:n]、DataCopyPad angle → qIn[T:Tn]Compute 阶段 Sin/Cos 取qIn[T:Tn]、Mul 取qIn[0:n]qOutVECOUT2 buffer × (2×T fp32)存放 complex64 交织结果bufPackedVECCALC1 buffer × (2×T fp32)[cos | sin]中间结果bufOffVECCALC1 buffer × (2×T u32)Gather 静态 offset 表bufSinCosTmpVECCALCtilingData.tmpBufferSizeSin/Cos 显式 sharedTmpBuffer。TPipe Tensor 数 qIn(2) qOut(2) bufPacked(1) bufOff(1) bufSinCosTmp(1) 7≤8 上限UB 用量 32 32 16 16 24 120 KB / 192 KB。bcastMode 1时走PrecomputeBcastCosSin快路径Init 阶段对angle[0:K]只做一次 Cos/Sin并把 cosTile/sinTile 周期复制填满bufSinCosTmp_的[0:T]/[T:2T]常驻区主循环ProcessTileBcast只读 input HBM每 tile 仅Mul × 2 → Gather彻底省去每 tile 的 angle 搬运与三角函数重算。off必 K 对齐、n是 K 整数倍由 host K-block 分核与主循环共同保证。该分工的收益是任意 shape / 任意广播组合≤8 维都能正确处理不会因 outNumel 过大或中间轴广播触发 kernel 内 fp32 unravel 失精或 UB 容量越界见 aclnn_polar.cpp 文件头注释。七、支持硬件与使能方式芯片版本勾选Atlas A2 训练系列产品 / Atlas A2 推理系列产品 (ascend910b)√Atlas A3 训练系列产品 / Atlas A3 推理系列产品 (ascend910_93)√上层框架勾选TF 训练/推理Pytorch 训练/推理√pybind 调 aclnn 测试ATC 推理Aclnn 直调√OPAT 调优SGAT 子图切分主要适配ACLNN调用框架。仓库 README.md 的产品支持表与之呼应并额外说明 Ascend 950PR/950DT、Atlas 200I/500 A2 推理产品等不支持。aclnn 接口遵循 CANN 标准的两段式规范先aclnnPolarGetWorkspaceSize获取 workspace 与 executor再aclnnPolar执行完整接口说明与调用示例见 docs/aclnnPolar.md可直接编译运行的 C 样例见 examples/test_aclnn_polar.cpp。调用样例中的核心流程详见 aclnnPolar.md// 第一段入参校验 workspace 计算 ret aclnnPolarGetWorkspaceSize(input, angle, out, workspaceSize, executor); CHECK_RET(ret ACLNN_SUCCESS, ...); // 申请 workspace void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); ... } // 第二段执行计算 ret aclnnPolar(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACLNN_SUCCESS, ...); ret aclrtSynchronizeStream(stream);注意 out 为 complex64每元素占 real/imag 两个 float申请与读取输出内存时按numel × 2 × sizeof(float)计算字节数。返回值与错误码第一段接口完成入参校验出错场景返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001input、angle、out 是空指针ACLNN_ERR_PARAM_INVALID161002数据类型/格式不在支持范围、input 与 angle dtype 不一致、shape 超过 8 维、不满足广播推导、out shape 与广播结果不一致或非 COMPLEX64八、算子约束限制dtypeinput/angle 仅 float32须同 dtypeout 仅 complex64。不支持 fp16/bf16/fp64/complex128≤8 维input/angle 须满足 numpy 广播广播操作数 numel 需可整体载入 UB验收 shape 量级满足超大广播待分块优化——落地代码以 L2 层 BroadcastTo 兜底规避了该限制无 attr确定性计算aclnnPolar 默认确定性实现。九、特性交叉分析本算子为逐元素计算输出仅依赖对应位置输入不涉及跨元素聚合Reduce/Scatter、数据搬移Reshape/Transpose 语义、量化Quant与现有特性无冲突。广播与各类边界条件已交叉验证8D 满秩同形、5D 中间轴广播、双向多轴广播、标量 input×高维 angle、负 abs、大角度归约、非 32B 对齐 inner广播、大向量×标量等用例case9–16均严格判据 PASS无异常。十、可维可测分析精度与性能验收10.1 精度标准与实测验收标准描述来源精度满足 AscendOpTest 默认阈值不低于参考fp32→complex64 实/虚部各 max_abs1e-4、rate1e-4AscendOpTestHIT1920/AscendOpTest性能所有核参与场景 ≥ 参考 l0 的 95%小 shape 提供仿真分析《算子任务书》精度细化与实测complex64 在 AscendOpTestaccuracy_config无内置默认值因此用例 JSON 显式配置err_threshold[1e-4,1e-4]fp32 分量默认compare_complex为实/虚部各自纯绝对误差判定。官方 AscendOpTest 实跑 6 个用例同 shape 小 shape/16M 广播低→高/标量/双向 高维非对齐全部 PASS含 16M[4096,4096]。全程 fp32 计算不引入 Cast 中间精度损失complex64 经 ReinterpretCast 视作交织 fp32 分量精度等价 fp32。10.2 性能实测性能对比vs 系统 l0 参考基线官方 AscendOpTest msprof / 每调用设备时场景l0 基线本算子结论小 [2,6,10]14.25 µs12.52 µs更快 ✓16M [4096,4096]所有核1665 µs937.98 µs快 1.78×核心验收达标✓✓广播 [4,1,8]×[4,5,8]15.17 µs14.92 µs更快 ✓全场景优于 l0 参考实现远超 ≥95% 的验收要求。详细自测过程与数据可参考 自测报告 与 测试步骤指导。10.3 兼容性分析本算子为新增算子昇腾仓原无 Ascend C 原生 Polar亦无 TBE 历史实现不涉及历史兼容性问题。aclnn 接口遵循 CANN 标准规范提供与 PyTorchtorch.polar等效功能可支撑模型迁移。总结Polar 算子是一个典型的小算子拼接 → 单 kernel 融合优化案例从约 6 个 l0 设备算子串接≥10 次 GM↔UB 搬运、6 次 dispatch到 Ascend C 单 kernel 内完成Cos/Sin Mul×2 复数交织2 次搬运、1 次启动核心场景 16M 性能提升 1.78×同时通过 aclnn L2 层BroadcastTo Contiguous与 kernel 侧bcastModeinner-broadcast 周期复用快路径的合理分工完整支持 numpy 广播语义。其设计文档、接口文档、可运行样例与测试报告均随算子沉淀在仓库experimental/math/polar目录中可作为昇腾 Ascend C 算子融合开发与广播语义实现的可参考范例。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math AcosGradV2 算子全解析Ascend A2 上反余弦梯度计算与单 Kernel 融合实现CANN ops math AcosGradV2 算子全解析Ascend A2 上反余弦梯度计算与单 Kernel 融合实现 导读 本文基于 AcosGrad算子库人工智能CANNCANN ops-math Lerp 算子全解析从 aclnnLerp 接口到 Ascend 内核实现CANN ops math Lerp 算子全解析从 aclnnLerp 接口到 Ascend 内核实现 LerpLinear Interpolation线算子库人工智能CANN推荐项目ShadowImageView - 简单易用的阴影效果图像视图组件推荐项目ShadowImageView 简单易用的阴影效果图像视图组件 项目简介 在Android开发中为UI元素添加阴影效果可以提升界面的视觉吸引力和层次算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考