cann_ops_nn

面向 Ascend NPU 的高性能算子扩展库，通过 JIT 编译将 PyTorch 接口与 ACLNN 算子库桥接。

构建与安装

前置条件

• 操作系统：Linux
• Python：3.8+
• 编译器：GCC 9.4.0+
• 框架：PyTorch>=2.6.0、torch_npu（需匹配 PyTorch 版本）
• 工具包：Ascend CANN Toolkit

构建 Wheel 包

支持两种构建方式：通过 build.sh 构建和直接使用 python3 -m build 构建。

方式一：通过 build.sh 构建（推荐）

# 构建整包（包含所有算子）
bash build.sh --torch_extension

# 构建单算子包（仅包含指定算子）
bash build.sh --torch_extension --ops=swiglu_group --vendor_name=custom

# 构建多算子包
bash build.sh --torch_extension --ops=swiglu_group,swiglu_group_quant --vendor_name=custom

# 构建实验性算子包（仅包含 experimental 目录下的算子）
bash build.sh --torch_extension --experimental

构建完成后，wheel 包会自动复制到 build_out/ 目录。

参数说明：

参数	必选	说明
--torch_extension	是	仅构建 torch_extension wheel 包，不执行 cmake 编译
--ops=op1,op2,...	否	指定编译的算子名（逗号分隔），不指定则编译所有算子
--vendor_name=name	否	指定子包名后缀，用于子包命名和隔离。不指定 --ops 时此参数无效，默认为 custom
--experimental	否	仅编译 experimental 目录下的算子。不指定则跳过 experimental 目录

包命名规则：

场景	条件	包名	安装目录
整包	不指定 --ops	cann_ops_nn	cann_ops_nn/
单算子/多算子包	指定 --ops，--vendor_name 可选	cann_ops_nn_<vendor>	cann_ops_nn_<vendor>/

命名逻辑： 不指定 --ops 时构建整包，包名固定为 cann_ops_nn；指定 --ops 时构建子包，包名为 cann_ops_nn_ 拼接 --vendor_name 的值（未指定则默认 custom）。整包与子包安装目录物理隔离，可共存。

方式二：直接使用 python3 -m build 构建

# 安装依赖
python3 -m pip install -r requirements.txt

# 构建整包
python3 -m build --wheel -n

# 构建单算子包
TORCH_EXTENSION_OPS=swiglu_group TORCH_EXTENSION_VENDOR=custom python3 -m build --wheel -n

安装

# 安装整包
python3 -m pip install dist/cann_ops_nn-*.whl

# 安装单算子包
python3 -m pip install dist/cann_ops_nn_custom-*.whl

整包与子包共存机制

整包和单算子包可以同时安装，互不冲突：

• 整包安装到 cann_ops_nn/ 目录，包含所有算子。
• 单算子包安装到 cann_ops_nn_<vendor>/ 目录，与整包物理隔离。
• 单算子包通过 entry point 机制注册算子，优先级高于整包。用户调用 cann_ops_nn.swiglu_group(x) 时，若子包已安装则使用子包的算子实现。
• 卸载单算子包后，整包的同名算子自动接管。

# 安装整包
pip install cann_ops_nn-*.whl

# 安装单算子包（覆盖整包中的同名算子）
pip install cann_ops_nn_custom-*.whl

# 卸载单算子包（整包算子自动恢复）
pip uninstall cann-ops-nn-custom

目录结构

├── torch_extension
│   ├── build.sh                              # 构建脚本（支持 --torch_extension 参数）
│   ├── setup.py                              # wheel 打包配置（自动收集分布式算子）
│   ├── requirements.txt
│   └── cann_ops_nn                           # 整包安装目录（框架公共部分）
│       ├── __init__.py                       # 包入口，from .ops import *
│       ├── op_builder/
│       │   └── builder.py                    # OpBuilder 基类，JIT 编译管理
│       ├── common/
│       │   └── aclnn_common.h                # ACLNN_CMD 宏等公共能力
│       ├── csrc/
│       │   └── extension.cpp                 # 占位模块（满足 PyTorch 扩展要求）
│       ├── ops/
│       │   └── __init__.py                   # 自动发现算子 + entry point 加载（构建时填充）
│       └── docs/
│           └── torch_extension_guidelines.md # 开发规范
├── <category>/<op>/torch_extension/          # 算子开发目录（仓库根目录下）
│   ├── <op>.py                               # Python 前端（OpBuilder + schema + meta）
│   ├── __init__.py
│   ├── graph_convert_<op>.py                 # （可选）图回退转换
│   └── csrc/<op>.cpp                         # C++ kernel wrapper
└── build_out/                                # build.sh 输出目录

构建时收集： setup.py 在打包阶段自动遍历仓库中所有 <category>/<op>/torch_extension/，把 <op>.py、csrc/<op>.cpp 汇集到 wheel 包目录 cann_ops_nn/ops/<category>/<op>/ 与 cann_ops_nn/csrc/ 下。开发者无需手动在 cann_ops_nn/ 下创建算子文件。

快速入门

import torch
import torch_npu
import cann_ops_nn

# 初始化 NPU 张量
x = torch.randn(16, 32, dtype=torch.float16).npu()

# 调用算子（两种方式等价）
result = cann_ops_nn.swiglu_group(x)
result = cann_ops_nn.ops.swiglu_group(x)

print(result.shape)  # torch.Size([16, 16])

开发者指南：新增算子

以添加新算子 new_operator（类别 my_cat）为例，在算子目录 <my_cat>/new_operator/ 下创建 torch_extension/ 子目录，包含 C++ 绑定和 Python 构建器。

1. C++ 后端 (<my_cat>/new_operator/torch_extension/csrc/new_operator.cpp)

该文件将 PyTorch 张量桥接到 ACLNN C-API。

#include <torch/extension.h>
#include "aclnnop/aclnn_new_operator.h"
#include "aclnn_common.h"  // 通过 -I cann_ops_nn/common 解析

namespace cann_ops_nn {
namespace my_cat {

at::Tensor new_operator(
    const at::Tensor& input1,
    const at::Tensor& input2,
    int64_t param1,
    const std::string& param2)
{
    // 设备检查
    TORCH_CHECK(input1.device().type() == at::kPrivateUse1, "input1 must be on NPU device");
    TORCH_CHECK(input2.device().type() == at::kPrivateUse1, "input2 must be on NPU device");

    // 输出创建
    at::Tensor out = at::empty(output_shape,
        at::TensorOptions().dtype(output_dtype).device(at::kPrivateUse1));

    // ACLNN 调用（宏内部自动处理类型转换和 workspace）
    ACLNN_CMD(aclnnNewOperator, input1, input2, out);

    return out;
}

}  // namespace my_cat
}  // namespace cann_ops_nn

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("new_operator", &cann_ops_nn::my_cat::new_operator,
          "NewOperator on NPU");
}

2. Python 前端 (<my_cat>/new_operator/torch_extension/new_operator.py)

该文件管理 JIT 编译逻辑并将算子注册到 PyTorch Dispatcher。

import torch
import torch_npu
from torch.library import impl
from cann_ops_nn.op_builder import OpBuilder, get_as_library

class NewOperatorOpBuilder(OpBuilder):
    def __init__(self):
        super().__init__("new_operator")

    def sources(self):
        """C++ 源码路径。"""
        return [self.resolve_source("new_operator.cpp")]

    def schema(self) -> str:
        """PyTorch 算子签名。"""
        return "new_operator(Tensor input1, Tensor input2, int param1=0, str param2=\"\") -> Tensor"

    def register_meta(self):
        """
        注册 Meta 实现（形状/类型推导）。
        对 Autograd 和 FakeTensor 支持至关重要。
        """
        @impl(get_as_library(), self.name, "Meta")
        def new_operator_meta(input1, input2, param1=0, param2=""):
            return torch.empty_like(input1)

# 实例化构建器
builder = NewOperatorOpBuilder()
builder._ensure_initialized()

@impl(get_as_library(), builder.name, "PrivateUse1")
def new_operator(input1, input2, param1=0, param2=""):
    """
    Dispatcher 的 NPU 实现。
    'PrivateUse1' 是自定义 NPU 后端的分发键。
    """
    op_module = builder.load()  # 编译/加载 .so 文件
    return op_module.new_operator(input1, input2, param1, param2)

技术说明

组件	职责
OpBuilder	使用 ninja 处理 C++ 源码的 JIT 编译
Meta 分发	允许 PyTorch 在不运行 NPU 代码的情况下推导输出形状/类型
PrivateUse1	PyTorch 路由 NPU 特定操作使用的后端分发键
ACLNN_CMD 宏	自动处理 at::Tensor → aclTensor* 类型转换、workspace 申请/释放、流调度

确定性计算

通过 torch.use_deterministic_algorithms(True) 开启确定性模式后，所有经由 ACLNN_CMD 调用的算子会自动将确定性标志传递给 CANN 运行时，确保同一输入产生相同输出：

import torch
import cann_ops_nn

torch.use_deterministic_algorithms(True)
# 后续算子调用将使用确定性实现

实现原理：ACLNN_CMD 宏在每次算子调用前读取 at::globalContext().deterministicAlgorithms()，通过 aclrtCtxSetSysParamOpt(ACL_OPT_DETERMINISTIC, ...) 将标志写入当前 device context，算子根据该标志选择确定性实现路径。

算子库查找机制

JIT 编译的 .so 在运行时按以下顺序查找算子符号（dlopen/dlsym）：

1. libcust_opapi.so（vendor 单算子包，可选）：遍历 ASCEND_CUSTOM_OPP_PATH 和 LD_LIBRARY_PATH 中的路径，加载 libcust_opapi.so
2. libopapi_nn.so（仓库构建安装的算子库）

vendor 包为可选项，未安装时直接使用 libopapi_nn.so。

# 设置 vendor 包路径（可选，安装 vendor 包后自动生成）
source <vendor_path>/set_env.bash
# 或手动设置
export LD_LIBRARY_PATH=<vendor_path>/op_api/lib:$LD_LIBRARY_PATH
export ASCEND_CUSTOM_OPP_PATH=<vendor_path>:$ASCEND_CUSTOM_OPP_PATH