oil-gas-ops-prospect

🚀 概述

oil-gas-ops-prospect 是面向油气勘探（oil & gas exploration）领域的昇腾自定义算子库。仓名即 oil-gas-ops（油气算子）+ prospect（勘探目标），面向地震成像、全波形反演等勘探计算场景。

本仓按交付物形态分为三类算子后端；Python 对外接口按算子名暴露，不区分实现后端：

后端	目录	交付物	适用场景
AscendC	ascendc/	CANN OPP .run 算子包，暴露 aclnn 接口	需要精细控制核间切分与流水、追求极致性能的算子
Triton-Ascend	triton-ascend/	Python whl 内的 Triton kernel	用 Python 快速表达、依赖 autograd、形态多变的算子
Composite-Torch	composite-torch/	Python whl 内的纯 PyTorch 组合模块	组合现有 torch 算子即可表达、无需自定义 kernel 的算子

📌 版本配套

当前仓库已验证通过的 CANN 版本如下：

CANN 版本	发布时间
9.0.0 官方版本	2026/04/30

本仓库需要安装 Toolkit 包、对应 SoC 的 Ops 包，以及与 CANN 配套的 Triton-Ascend（已验证 triton-ascend==3.2.1）。当前面向 Ascend 910B（Ascend910B1）。完整的安装步骤、SoC 支持矩阵与环境变量配置见环境部署。

⬇️ 源码下载

git clone https://gitcode.com/cann/oil-gas-ops-prospect.git
cd oil-gas-ops-prospect

⚡️ 快速入门

编译与安装

仓库根目录提供三种构建入口。

单独编译 OPP .run 包

bash build.sh --pkg

只编译 AscendC OPP .run，写入 output/<SOC>/，不编 Python whl，也不安装。

只编译不安装

bash build.sh

同时编译 OPP .run 与 Python whl，产物写入 output/，不安装到 CANN / Python，也不跑冒烟测试。

编译并安装

bash build.sh --install

编译 OPP 与 whl，安装到 CANN / Python，并跑 whl 冒烟测试。

产物统一落到仓库根目录 output/：

output/
├── <SOC>/oil_gas_ops_prospect_opp_<os>_<arch>.run   # OPP 算子包
└── wheels/oil_gas_ops_prospect-*.whl                # Python whl

安装完成后，后续 shell 中需要保留自定义 opapi 动态库路径：

export LD_LIBRARY_PATH=${ASCEND_HOME_PATH}/opp/vendors/customize/op_api/lib:${ASCEND_HOME_PATH}/lib64:$LD_LIBRARY_PATH

常用参数：

bash build.sh --help                                # 查看全部选项
bash build.sh --pkg -h                              # Package Options
bash build.sh -u -h                                 # Test Options
bash build.sh --clean                               # 先清理再构建（可与上面组合）
bash build.sh --ops=complex_mul,complex_mul_grad    # 只构建指定算子
bash build.sh --install --skip-tests                # 安装后跳过 whl 冒烟测试
bash build.sh --soc=Ascend910B1 -j16                # 指定 SOC 与并行度

完整的构建入口与参数说明见编译参数说明。

接口使用

import torch
import torch_npu  # noqa: F401
from oil_gas_ops_prospect.real_fft import real_irfft, real_rfft

x = torch.randn(8, 256, device="npu")
xr, xi = real_rfft(x, dim=-1)          # 实数正向 FFT，频率轴长度 256 // 2 + 1 = 129
y = real_irfft(xr, xi, n=256, dim=-1)  # 从半谱重建实信号

可直接运行的示例见 examples/，更多调用方式见算子调用。

测试

分类见 tests/README.md。

bash build.sh -u                                 # 全部套件（需 910B）
bash build.sh -u --ut                            # UT：切分 + 上板小 launch
bash build.sh -u --st                            # ST：多 shape aclnn
bash build.sh -u --precision                     # 精度
bash build.sh -u --perf                          # 性能

bash tests/run_whl_tests.sh                      # 只验证已安装的 whl

算子在 NPU 上的精度、性能、显存必须在真实 910B 上验证通过。问题定位见调试与性能分析。

📖 接口支持清单

当前交付 4 个 AscendC 算子、2 个 Triton-Ascend 算子和 1 个 Composite-Torch 算子：

名称 / OpType	后端 / 交付物	功能	dtype / format
ComplexMul	AscendC / OPP .run（AIV-only）	复数逐点乘：(ur, ui), (vr, vi) -> (wr, wi)	float32 / ND，全部张量同 shape
ComplexMulGrad	AscendC / OPP .run（AIV-only）	ComplexMul 的反向	float32 / ND，全部张量同 shape
FusedBiasSoftmax	AscendC / OPP .run	softmax(x + bias1 + bias2, -1)，原地写回	fp32 / fp16 / bf16，ND
FusedSoftmaxGrad	AscendC / OPP .run（可走 MIX）	注意力 softmax JVP：g = go @ V^T 后原地覆写 probs	fp32 / fp16 / bf16，ND
real_rfft	Triton-Ascend / Python whl	实数正向 FFT：x -> (xr, xi)，频率轴长度 n // 2 + 1	NPU kernel 为 float32，非 NPU 环境回退到 PyTorch
real_irfft	Triton-Ascend / Python whl	从实部、虚部半谱重建实数：(xr, xi), n -> y	NPU kernel 为 float32，非 NPU 环境回退到 PyTorch
Conv3d	Composite-Torch / Python whl	nn.Conv3d 等价替换（自定义前/反向，state_dict 互通）	任意（同 nn.Conv3d）

数学定义、约束与支持范围见算子列表；完整接口签名、参数与环境变量见接口列表。

🔍 目录结构

oil-gas-ops-prospect/
├── ascendc/                # AscendC 算子（交付物：CANN OPP .run 包）
│   ├── operators/          #   算子源码 + manifest.tsv 注册表
│   ├── pybind/             #   Python API（打包为 oil_gas_ops_prospect）+ C++ 绑定
│   ├── scripts/            #   msopgen 构建流程内部模块
│   └── build.sh            #   AscendC 侧入口（--pkg 打包 / -u --ophost 跑 UT）
├── triton-ascend/          # Triton-Ascend 算子（real_fft 打包为 oil_gas_ops_prospect.real_fft）
│   ├── real_fft/           #   real_rfft / real_irfft
│   └── tools/              #   AOT 编译检查
├── composite-torch/        # Composite-Torch 算子（纯 PyTorch 组合，打包为 oil_gas_ops_prospect.<op>）
│   └── conv3d_custom/      #   Conv3d（nn.Conv3d 等价替换）
├── tests/                   # UT / ST / Python 精度性能；operators/ 不再放测试
│   ├── ut/op_host/         #   tiling googletest
│   ├── st/aclnn<OpType>/   #   两段式 aclnn 系统测试
│   ├── ascendc/ophost/     #   汇总编译 ophost_ut（googletest）
│   ├── ascendc/complex_mul/#   NPU 精度/性能对比
│   ├── ascendc/fused_bias_softmax/
│   └── triton/real_fft/    #   rFFT / iRFFT 精度（含 CPU 回退路径）
├── examples/               # Python 示例 + examples/aclnn/<op>/ C++ 样例
├── scripts/                # 辅助脚本
├── docs/                   # 文档
├── cmake/                  # CMake 模块（三方件解析、安全编译选项）
├── .gitcode/workflows/     # GitCode Action 门禁流水线
├── build.sh                # 顶层构建入口（默认只编包；--install 安装；-u 跑测试）
├── CMakeLists.txt          # 顶层 CMake（ophost tiling UT）
└── version.cmake           # 版本与依赖定义

逐目录职责说明见目录结构。

🛠️ 开发指南

新增 AscendC / Triton-Ascend / Composite 算子的完整流程、目录约定、命名约定与自检方式见算子开发指南。

提交前自检：

bash build.sh -u                        # 全部测试套件（需 910B）
bash scripts/check_license_header.sh    # 许可证头检查

💬 相关信息

• 快速入门
• 文档导航
• 目录结构
• 算子列表
• 接口列表
• 编译参数说明
• 算子开发指南
• 贡献指南
• 安全声明
• 变更记录
• 许可证（Apache License 2.0）
• 所属 SIG：oil-gas-engineering

---

PS：本项目功能和文档正在持续更新和完善中，欢迎您关注最新版本。

• 问题反馈：通过 GitCode Issues 提交问题。
• 社区互动：通过 GitCode 讨论 参与交流。