ops-gnn 开发指导
本文档提供 ops-gnn 库的开发环境构建、算子开发指南和测试说明。
一、环境构建
1.1 安装 CANN 软件
安装前准备
离线安装时，请单击获取链接下载 CANN 软件包，并上传到安装环境任意路径。
安装 CANN
请使用 CANN 9.1.0-beta.1 及以上版本，其他版本暂不支持。
chmod +x Ascend-cann-toolkit_${VERSION}_linux-$(arch).run
./Ascend-cann-toolkit_${VERSION}_linux-$(arch).run --install
其中 ${VERSION} 表示对应的 CANN 版本（如 9.1.0），$(arch) 表示 CPU 架构。
安装后配置
source /usr/local/Ascend/cann-9.1.0-beta.1/bin/setenv.bash
若 CANN 安装在其他路径，请替换为实际路径。
1.2 CANN 详细安装指南
开发者可访问昇腾文档-昇腾社区 → CANN 社区版 → 软件安装，查看 CANN 软件安装引导，根据机器环境、操作系统和业务场景选择后阅读详细安装步骤。
1.3 依赖安装
ops-gnn 依赖以下组件：
| 组件 | 版本要求 | 说明 |
|---|---|---|
| Python | 3.9+ | |
| CMake | 3.18+ | C++ 构建系统 |
| GCC | 7.0+ | C++17 编译器 |
| PyTorch | 2.7+ | 需从昇腾社区获取 NPU 适配版本 |
| torch_npu | 26.0.0 | PyTorch NPU 设备后端 |
| CANN Toolkit | 9.1.0-beta.1+ | AscendC 编译器、Bisheng 编译器、运行时库 |
torch 和 torch_npu 需从昇腾社区下载页面获取适配 Ascend 硬件的版本。
# 系统工具链
sudo apt install build-essential cmake
# 测试依赖（可选）
pip3 install pytest pytest-cov
1.4 安装 ops-gnn
# 激活 CANN 环境
source /usr/local/Ascend/cann-9.1.0-beta.1/bin/setenv.bash
# 方式一：pip 开发模式安装
cd /path/to/ops-gnn
pip install -e .
# 方式二：构建脚本安装
cd scripts
./build.sh python
# 方式三：CMake 手动编译
mkdir build_cmake && cd build_cmake
cmake ..
cmake --build .
二、项目架构
2.1 分层架构设计
ops-gnn 采用 PyTorch 扩展 + AscendC 内核的分层架构：
┌─────────────────────────────────────────────┐
│ Python 接口层 │
│ python/ops_gnn/<op>.py │
│ (类型注解、参数预处理、文档) │
├─────────────────────────────────────────────┤
│ PyTorch 绑定层 │
│ csrc/pybind.cpp │
│ (PYBIND11_MODULE 注册、参数映射) │
├─────────────────────────────────────────────┤
│ Host 端算子层 │
│ csrc/npu/host/<op>/<op>.cpp/.h │
│ (参数解析、Tiling 计算、Kernel 启动、流管理) │
├─────────────────────────────────────────────┤
│ AscendC Kernel 层 │
│ csrc/npu/kernel/<op>/<op>_kernel.cpp/.h │
│ (AscendC SIMT/VF 编程、向量计算、数据搬移) │
└─────────────────────────────────────────────┘
架构说明：
- Python 接口层：提供用户友好的 Python API，包含类型注解、docstring 和使用示例
- PyTorch 绑定层：通过 pybind11 将 C++ 函数暴露给 Python，处理 Tensor 参数传递
- Host 端算子层：运行在主机 CPU，负责参数校验、Tiling 数据计算、AscendC Kernel 启动和 Stream 同步
- AscendC Kernel 层：运行在 NPU 设备端（AIV 核），实现核心计算逻辑
2.2 目录结构
ops-gnn
├── csrc/ # C++/AscendC 源码目录
│ ├── pybind.cpp # PyTorch 绑定代码
│ └── npu/ # NPU 相关代码
│ ├── host/ # Host 端代码（按算子分类）
│ │ ├── add_sample/
│ │ │ ├── add_sample.h # Host 接口声明
│ │ │ └── add_sample.cpp # Host 实现（Tiling + Launch）
│ │ └── segment_max_csr/
│ │ ├── segment_max_csr.h
│ │ └── segment_max_csr.cpp
│ └── kernel/ # AscendC 内核实现（按算子分类）
│ ├── add_sample/
│ │ ├── add_sample_kernel.h # Kernel Launch 接口
│ │ └── add_sample_kernel.cpp # Kernel 实现（AscendC SIMT）
│ └── segment_max_csr/
│ ├── segment_max_csr_kernel.h # Kernel Launch 接口
│ ├── segment_max_csr_kernel.cpp # Kernel 入口 + 模板实例化
│ ├── segment_max_csr_kernel_impl.h # Kernel 核心实现类
│ └── segment_max_csr_tiling.h # Tiling 数据结构
├── docs/ # 文档目录
├── python/ # Python 源码目录
│ └── ops_gnn/ # Python 包目录
│ ├── __init__.py # 包初始化、导出列表
│ ├── add_sample.py # add_sample Python 接口
│ ├── segment_max_csr.py # segment_max_csr Python 接口
│ └── typing.py # 类型别名定义
├── test/ # 测试目录
│ ├── test_import.py # 导入验证测试
│ ├── test_example.py # add_sample 算子测试
│ └── test_segment_max_csr.py # segment_max_csr 算子测试
├── scripts/ # 构建脚本目录
│ └── build.sh # 统一构建脚本（支持 python/cpp/all）
├── cmake/ # CMake 配置
│ └── OpsGNNConfig.cmake.in # CMake 包配置模板
├── CMakeLists.txt # CMake 构建配置
├── setup.py # Python 安装脚本（setuptools + cmake）
├── setup.cfg # setuptools 配置
├── pyproject.toml # 现代 Python 项目配置
├── MANIFEST.in # 打包清单
└── LICENSE # CANN 许可证
2.3 核心文件说明
| 文件 | 功能说明 |
|---|---|
csrc/pybind.cpp |
PyTorch 绑定入口，通过 PYBIND11_MODULE 注册所有 C++ 算子到 Python |
csrc/npu/host/<op>/<op>.h |
Host 端算子接口声明，定义函数签名 |
csrc/npu/host/<op>/<op>.cpp |
Host 端算子实现：Tensor 维度解析、Tiling 参数计算、dtype 分发、Stream 管理 |
csrc/npu/kernel/<op>/<op>_kernel.h |
Kernel Launch 函数声明（Host 端调用入口） |
csrc/npu/kernel/<op>/<op>_kernel.cpp |
Kernel Launch 实现 + 模板显式实例化 + <<<>>> 启动语法 |
csrc/npu/kernel/<op>/<op>_kernel_impl.h |
AscendC Kernel 核心类实现（Init → Process → Compute 流水线） |
csrc/npu/kernel/<op>/<op>_tiling.h |
Tiling 数据结构定义（传递给 Device 端的参数结构体） |
python/ops_gnn/<op>.py |
Python 接口封装：类型注解、参数默认值处理、调用 _pybind.<op> |
python/ops_gnn/__init__.py |
包入口，从各模块导入并注册到 __all__ |
python/ops_gnn/typing.py |
类型别名（Tensor, OptTensor） |
三、代码规范
3.1 命名规范
C++ 层：
- 文件名：小写下划线分隔，如
segment_max_csr_kernel.h - 函数名：大驼峰，如
LaunchSegmentMaxCsrKernel - 结构体名：大驼峰，如
SegmentMaxCsrTilingData - 类名：大驼峰，如
SegmentMaxCsrKernel - 模板参数：大驼峰或单字母大写，如
typename T
Python 层：
- 文件名：小写下划线分隔，如
segment_max_csr.py - 函数名：小写下划线分隔，如
segment_max_csr - 与 PyTorch 风格保持一致
3.2 代码风格
- C++ 遵循 C++17 标准，使用
#pragma once头文件保护 - AscendC Kernel 使用
__aicore__、__global__、__gm__等修饰符 - 使用 namespace
AscendC下的 API - Python 接口使用类型注解（
Tensor,OptTensor,Optional[Tensor]） - 每个文件顶部包含 CANN Open Software License 版权声明
3.3 文件组织规范
每个算子严格遵循以下文件拆分：
csrc/npu/
├── host/<op>/
│ ├── <op>.h # Host 接口声明
│ └── <op>.cpp # Host 实现
└── kernel/<op>/
├── <op>_kernel.h # Kernel Launch 声明
├── <op>_kernel.cpp # Kernel Launch 实现 + 模板实例化
├── <op>_kernel_impl.h # Kernel 核心类（简单算子可合并到 kernel.cpp）
└── <op>_tiling.h # Tiling 结构体
四、构建系统
4.1 CMake 构建
CMakeLists.txt 负责编译 AscendC Kernel 和 PyTorch 绑定库：
- 使用 bisheng 编译器将 Kernel 源文件编译为
libopsgnn_npu_kernel.so - 将
csrc/pybind.cpp+ host 端代码编译为_pybind.so - 产物安装到
output/kernel/
关键 CMake 变量：
| 变量 | 说明 | 默认值 |
|---|---|---|
NPU_ARCH |
NPU 架构 | dav-3510（950）/ dav-2201（910B） |
WITH_PYTHON |
是否编译 Python 绑定 | ON |
ASCEND_HOME_PATH |
CANN 安装路径 | 从环境变量读取 |
Python3_ROOT_DIR |
Python 根目录 | build.sh 自动导出 |
4.2 setuptools 构建
setup.py 调用 CMake 编译并生成 wheel 包。build_with_cmake() 构建前会自动清理旧 CMakeCache.txt，避免 pip 临时目录导致路径不一致。
五、测试指南
5.1 运行测试
pytest test/ -v # 运行所有测试
pytest test/test_segment_max_csr.py -v # 单个算子测试
pytest test/test_segment_max_csr.py::test_func -v # 单个测试用例
5.2 测试编写模式
torch.npu.set_device(4)指定 NPU 设备torch.manual_seed(42)保证可复现- 用
.npu()方法创建 NPU Tensor（如torch.randint(...).npu()） - 调用
ops_gnn.<op>(...) - 验证结果：设备类型、形状、数值正确性
5.3 测试覆盖要求
| 场景 | 说明 |
|---|---|
| 基本功能 | 典型输入 |
| 不同 dtype | float32, float16, int32, int16 |
| 边界情况 | 空 segment、单元素 |
| 可选参数 | optional_out |
| 多维输入 | 2D/3D Tensor |
| 广播场景 | indptr 广播 |
六、算子开发实例
本节以 segment_max_csr 为例，说明开发新算子的流程。
6.1 调用链分析
用户代码
└─→ ops_gnn.segment_max_csr(src, indptr, optional_out)
│ python/ops_gnn/segment_max_csr.py
│ (参数默认值处理：None → 空 Tensor)
└─→ _pybind.segment_max_csr(src, indptr, optional_out)
│ csrc/pybind.cpp
│ (PYBIND11_MODULE 注册)
└─→ segment_max_csr() [Host]
│ csrc/npu/host/segment_max_csr/segment_max_csr.cpp
│ (维度解析、Tiling填充、dtype分发、Stream管理)
└─→ LaunchSegmentMaxCsrKernel<T>()
│ csrc/npu/kernel/segment_max_csr/segment_max_csr_kernel.cpp
│ (获取AIV核数、<<<>>>启动)
└─→ segment_max_csr_kernel<T> [Device]
│ segment_max_csr_kernel_impl.h
│ (Init → Process → Compute 流水线)
每层职责：
- Python（
python/ops_gnn/<op>.py）：类型注解、可选参数默认值、调用_pybind.<op> - PyBind（
csrc/pybind.cpp）：m.def("<op>", &func, py::arg(...), ...)注册 C++ 函数 - Host（
csrc/npu/host/<op>/）：Tensor 维度解析 → Tiling 参数计算 →torch::empty创建输出 →aclrtCreateStream→ 按scalar_type分发模板 Launch →aclrtSynchronizeStream→ 销毁 Stream - Kernel Launch（
csrc/npu/kernel/<op>/<op>_kernel.cpp）：GetCoreNumAiv()获取核数 →<<<coreNum, nullptr, stream>>>启动 → 显式模板实例化 - Kernel 实现（
<op>_kernel_impl.h）：Init()解析 Tiling + 分配 Buffer/Event →Process()按 Block 分配工作范围 →Compute()双缓冲流水线（DataCopy → Max/Add → DataCopy）
6.2 新增算子文件清单
以 segment_max_csr 为模板，每个新算子需创建：
csrc/npu/host/<op>/
├── <op>.h # torch::Tensor <op>(torch::Tensor ...);
└── <op>.cpp # Host 实现
csrc/npu/kernel/<op>/
├── <op>_kernel.h # template<typename T> void Launch<Op>Kernel(...);
├── <op>_kernel.cpp # Launch 实现 + 模板实例化
├── <op>_kernel_impl.h # Kernel 核心类 (Init/Process/Compute)
└── <op>_tiling.h # struct <Op>TilingData { uint32_t ... };
python/ops_gnn/
└── <op>.py # Python 接口
test/
└── test_<op>.py # 单元测试
此外需修改两个文件：
csrc/pybind.cpp：#include "host/<op>/<op>.h"+m.def("<op>", ...)python/ops_gnn/__init__.py：from .<op> import <op>+ 加入__all__
简单算子可合并文件：无 Tiling 时省略 _tiling.h，逻辑简单时 _kernel_impl.h 可合并到 _kernel.cpp（参考 add_sample）。
6.3 两种 Kernel 模式
| 特性 | add_sample | segment_max_csr |
|---|---|---|
| 文件数 | 1个 kernel.cpp | kernel + impl + tiling（4文件） |
| 编程模型 | SIMT (__simt_vf__ + VF_CALL) |
Kernel类 (TPipe + Buffer + Event) |
| 数据搬移 | 直接 GM 读写 | DataCopy + 双缓冲流水线 |
| Tiling | 无（标量参数直传） | Tiling 结构体 |
| 适用 | 逐元素操作 | 规约、分段、多级流水线 |
6.4 开发流程总结
- 参照
segment_max_csr创建目录结构和文件 - 定义 Tiling 结构体（仅
uint32_t类型，不含指针） - 实现 Kernel 核心类（Init → Process → Compute），用 TPipe + Event 做双缓冲流水线
- 实现 Launch 函数，
<<<>>>用#pragma GCC diagnostic包围 - Host 端：维度解析 → Tiling 填充 →
torch::empty→aclrtCreateStream→ dtype 分发 Launch → 同步 → 销毁 Stream - 在
pybind.cpp中m.def(...)注册 - 编写 Python 接口（类型注解 + docstring + 默认值处理）
- 更新
__init__.py导出 - 编写测试（至少覆盖基本功能、多个 dtype、边界情况）