oil-gas-ops-prospect三路线架构揭秘：AscendC、Triton-Ascend、Composite-Torch该如何选型

【免费下载链接】oil-gas-ops-prospect面向油气勘探（oil & gas exploration）领域的昇腾自定义算子库。仓名即 oil-gas-ops（油气算子）+ prospect（勘探目标），面向地震成像、全波形反演等勘探计算场景项目地址: https://gitcode.com/cann/oil-gas-ops-prospect

oil-gas-ops-prospect是一个面向油气勘探领域的昇腾自定义算子库，覆盖地震成像、全波形反演等勘探计算场景。这个仓库最有特色的地方，是把自定义算子按交付物形态组织成AscendC、Triton-Ascend、Composite-Torch 三条路线。本文以项目已交付的 7 个算子为例，帮新手快速搞清楚：自己的算子到底该走哪条路线。

一句话看懂：为什么需要三条算子路线 🧭

昇腾上做自定义算子，常面临"开发成本 vs 性能上限"的取舍：写 C++ AscendC 内核性能上限最高，但开发周期也最长；反过来，纯 PyTorch 组合几行代码就能写完，但只能用现有算子拼。

oil-gas-ops-prospect 用"交付物形态"把三条路线切得很清楚：

各目录的完整职责划分见目录结构。

路线一：AscendC——追求极致性能的"精调"路线 🔥

交付物：安装 CANN OPP.run算子包后，可通过aclnn两段式接口调用；Python 侧同样封装成模块。

仓库里的 4 个 AscendC 算子，都是典型的计算密集场景：

• ComplexMul / ComplexMulGrad——复数逐点乘及其反向，源码在 ascendc/operators/complex_mul/
• FusedBiasSoftmax——把 bias 累加与 softmax 三次发射融合进单个 kernel，原地写回
• FusedSoftmaxGrad——注意力 softmax 反向融合（JVP），还可走 AIC 矩乘 + AIV 向量的 MIX 路径

以 ComplexMul 为例，每个算子目录分两层组织：

• op_host/——算子定义（OpDef）、形状推断、tiling 切分算法。比如 complex_mul_tiling_core.h 负责把数据在 AIV 核间按 32B 对齐均分，并按 UB 缓冲预算分拍
• op_kernel/——kernel 实现，真正跑在 NPU AICore 上

这条路线的成本在于切分算法要自己写、自己测。仓库为每个算子都准备了 googletest 单元测试：test_complex_mul_tiling.cpp，覆盖"每个元素恰好被一个核处理、无空洞无重叠"这类切分不变量。

路线二：Triton-Ascend——用 Python 快速表达复杂算子 ⚡

交付物：打包在 Python whl 里的 Triton kernel，支持 autograd，Python 直接 import 就能用。

仓库里的 2 个 Triton-Ascend 算子是实数 FFT——real_rfft/real_irfft（源码在 triton-ascend/real_fft/），是地震成像与全波形反演里的高频变换。实现思路很巧妙：把 DFT 基矩阵构造出来，当作tl.dot矩阵乘跑在昇腾的cube 单元上——让"FFT"变成高吞吐的矩阵乘，而不是低吞吐的逐元素向量操作。

两个值得借鉴的工程细节：

1. CPU 回退：不在 NPU 上、或 Triton 不可用时自动回退torch.fft，接口语义一致，开发时可以先在 CPU 上自检，示例见 real_fft_example.py
2. AOT 编译检查：aot_compile_check.py 保证 kernel 的每个分支都能通过 TTIR → Linalg → bishengir → npubin 编译链路，避免"分支只在上板运行时才暴露编译错误"

路线三：Composite-Torch——零内核的纯 PyTorch 组合 🧩

交付物：Python whl 内的纯 PyTorch 组合模块，前向 / 反向用torch.autograd.Function组织，没有任何自定义 kernel。

当前代表算子是Conv3d（composite-torch/conv3d_custom/）：作为nn.Conv3d的等价替换，自定义前向 / 反向，state_dict 与原生 Conv3d 完全互通。它的反向用分步构造（grad_weight + 零膨胀 + 翻转转置卷积）替代原生路径，在 NPU 上处理大 kernel depthwise 卷积时比原生aclnnConvolutionBackward更稳定。

这条路线最大的好处是CPU / NPU 行为一致：不需要 NPU 开发环境就能开发测试，tests/composite/conv3d_custom/ 里的 11 种实际训练调用模式在 CPU 上就能全部跑通。

如何选型：算子路线选择速查表 ✅

要新增算子时，按顺序问自己三个问题即可：

官方算子开发指南开头就给了这张选型表，并写明了每条路线的完整开发流程。顺带提醒：新增 AscendC 算子还要在 ascendc/operators/manifest.tsv 登记一行注册信息。

统一 Python 接口：调用时不用关心是哪条路线 🎯

一个容易让人困惑的问题：路线不同，调用方式会不会不一样？答案是——对外完全一样。

三条路线最终都打包进同一个 Python 包oil_gas_ops_prospect，按算子名暴露，不区分实现后端：

这个包名由 setup.py 在打包时组装，源码树里并没有这个目录；而 scripts/local_pkg.py 会在裸 checkout 时自动把各算子目录挂到同一导入名下，装不装 whl，调用方式都一样：

import torch import torch_npu # noqa: F401 from oil_gas_ops_prospect.real_fft import real_rfft x = torch.randn(8, 256, device="npu") xr, xi = real_rfft(x, dim=-1) # 实数正向 FFT，频率轴长度 129

完整接口签名与环境变量见接口列表和算子调用。

快速上手：3 条命令从克隆到跑通 🚀

当前已验证配套CANN 9.0.0 + Triton-Ascend 3.2.1，面向Ascend 910B。完整安装步骤见环境部署：

git clone https://gitcode.com/cann/oil-gas-ops-prospect.git cd oil-gas-ops-prospect bash build.sh --install # 编译 OPP .run 与 Python whl、安装并跑冒烟测试

产物统一落在output/目录；如果只想产出 AscendC 算子包，用bash build.sh --pkg即可；跑全量测试套件则用bash build.sh -u（套件清单见 tests/suites.json）。

总结

• AscendC：性能上限路线，精细控制切分与流水，适合复数乘法、融合注意力这类计算密集算子
• Triton-Ascend：快速迭代路线，Python 表达 + autograd，适合实数 FFT 这类形态多变的算子
• Composite-Torch：零内核路线，纯 PyTorch 组合，适合现有算子可表达的场景
• 无论走哪条路线，对外接口都统一为oil_gas_ops_prospect.<算子名>，切换与维护成本极低

先花十分钟浏览算子列表，把 7 个已交付算子的数学定义与约束过一遍，再动手写你的第一个勘探算子吧！

【免费下载链接】oil-gas-ops-prospect面向油气勘探（oil & gas exploration）领域的昇腾自定义算子库。仓名即 oil-gas-ops（油气算子）+ prospect（勘探目标），面向地震成像、全波形反演等勘探计算场景项目地址: https://gitcode.com/cann/oil-gas-ops-prospect