引言：ATB与自定义算子的关系

ATB（Ascend Tensor Boost）是昇腾CANN生态中的高性能算子库，提供了丰富的神经网络算子和优化实现。在实际开发中，开发者经常会遇到ATB现有算子无法满足需求的情况，这时就需要开发自定义算子并注册到ATB中。本文将详细讲解自定义算子的完整开发流程，从算子定义、实现、注册到验证，帮助开发者将自定义算子无缝集成到ATB生态中。

自定义算子的开发流程

在昇腾CANN中，开发一个能被ATB使用的自定义算子，需要遵循以下标准流程：

算子需求分析 → 算子原型定义 → Ascend C实现 → 算子注册 → 编译部署 → 验证测试

这个流程确保了自定义算子的正确性、性能和可维护性。

Step 1：算子原型定义

算子原型定义了算子的输入输出接口、属性参数等。在昇腾CANN中，算子原型通过opbase框架定义：

# custom_op_proto.py
from opbase import OperatorBase, RegOp

# WHY: 使用RegOp装饰器注册算子原型，这样GE才能识别这个算子
@RegOp("CustomSoftmax")
class CustomSoftmax(OperatorBase):
    """自定义Softmax算子原型定义"""
    
    def __init__(self):
        super().__init__()
        # WHY: 定义输入tensor，名称必须与实现中的输入名一致
        self.add_input("input_x", "N *", "float16")
        # WHY: 定义输出tensor，dtype需要与输入匹配或通过type推导
        self.add_output("output_y", "N *", "float16")
        # WHY: 定义属性参数，axis表示softmax计算的维度
        self.add_attr("axis", "int", default=-1)
        # WHY: 设置算子类型，影响算子在图中的位置和优化策略
        self.set_op_type("CustomSoftmax")

Step 2：Ascend C算子实现

算子原型定义完成后，需要使用Ascend C编程语言实现算子的计算逻辑。Ascend C是昇腾CANN专为NPU设计的算子开发语言，提供了高效的向量计算和矩阵计算接口。

// custom_softmax.cpp
#include "lib/ascendc/softmax.h"

// WHY: 使用extern "C"确保函数名不被C++编译器修饰，便于Python端调用
extern "C" __global__ __aicore__ void custom_softmax_kernel(
    __gm__ float16_t* input,
    __gm__ float16_t* output,
    int32_t axis,
    uint32_t total_elements,
    uint32_t dim_size
) {
    // WHY: 获取当前AI Core的ID，用于数据分片，实现多核并行
    int32_t core_id = GetBlockIdx();
    int32_t core_num = GetBlockNum();
    
    // WHY: 计算每个AI Core处理的数据范围，实现负载均衡
    uint32_t elements_per_core = (total_elements + core_num - 1) / core_num;
    uint32_t start = core_id * elements_per_core;
    uint32_t end = min(start + elements_per_core, total_elements);
    
    // WHY: 使用LocalTensor在Local Buffer中操作，避免反复访问Global Memory
    LocalTensor<float16_t> input_local = LocalTensor<float16_t>(input + start);
    LocalTensor<float16_t> output_local = LocalTensor<float16_t>(output + start);
    
    // WHY: 分步计算softmax：max → exp → sum → div，提升数值稳定性
    // Step 1: 计算最大值（防止exp溢出）
    float16_t max_val = AscendC::VectorMax(input_local, dim_size);
    
    // Step 2: 计算exp(x - max)
    AscendC::VectorSubScalar(input_local, max_val, dim_size);
    AscendC::VectorExp(input_local, dim_size);
    
    // Step 3: 计算sum
    float16_t sum_val = AscendC::VectorSum(input_local, dim_size);
    
    // Step 4: 归一化
    AscendC::VectorDivScalar(input_local, sum_val, dim_size);
    
    // WHY: 将结果写回Global Memory，完成计算
    AscendC::DataCopy(output_local, input_local, dim_size);
}

Step 3：算子注册到ATB

算子实现完成后，需要将其注册到ATB中，这样才能在模型中调用。昇腾CANN提供了多种注册方式，这里介绍最常用的Python接口注册：

# register_custom_op.py
import torch
import torch_npu
from opbase import OperatorRegistry

# WHY: 定义Python端的算子接口，封装底层C++实现，提供用户友好的API
class CustomSoftmaxOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, axis=-1):
        # WHY: 保存axis用于反向传播（如果需要）
        ctx.axis = axis
        
        # WHY: 调用底层NPU算子，input_npu已经是NPU上的tensor
        output = torch_npu.npu_custom_softmax(input, axis)
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        # WHY: 实现反向传播，确保算子可用于训练
        axis = ctx.axis
        # ... 反向传播计算 ...
        return grad_input, None

# WHY: 将自定义算子注册到torch_npu，这样用户可以直接调用torch.npu.custom_softmax
torch_npu.register_op("custom_softmax", CustomSoftmaxOp.apply)

# WHY: 同时注册到ATB的算子库，这样GE在图编译时也能识别这个算子
OperatorRegistry.register("CustomSoftmax", "ATB")

Step 4：编译与部署

自定义算子需要编译为NPU可加载的二进制文件。昇腾CANN提供了完整的编译工具链：

# WHY: 使用升腾CANN的算子编译工具，将Ascend C代码编译为NPU二进制
cann-opc --input custom_softmax.cpp \
         --output custom_softmax.o \
         --include-path /usr/local/Ascend/ascend-toolkit/latest/include

# WHY: 将编译好的算子打包为算子包，便于部署和分发
cann-op-package --input custom_softmax.o \
                --output custom_softmax.opp \
                --op-name CustomSoftmax

部署算子包：

# WHY: 将算子包安装到昇腾CANN的算子库目录，这样运行时可以自动加载
cann-op-install --package custom_softmax.opp \
                --install-path /usr/local/Ascend/ascend-toolkit/latest/opp

Step 5：验证测试

算子部署完成后，需要编写测试用例验证其功能正确性和性能：

# test_custom_softmax.py
import torch
import torch_npu

def test_custom_softmax():
    # WHY: 创建测试数据，同时在CPU和NPU上创建，便于结果比对
    input_cpu = torch.randn(32, 64, dtype=torch.float16)
    input_npu = input_cpu.npu()
    
    # WHY: 使用PyTorch原生softmax作为参考结果
    output_cpu = torch.softmax(input_cpu, dim=-1)
    
    # WHY: 调用自定义softmax算子
    output_npu = torch.npu.custom_softmax(input_npu, axis=-1)
    
    # WHY: 将NPU结果拷贝回CPU进行比对
    output_npu_cpu = output_npu.cpu()
    
    # WHY: 计算最大绝对误差，判断精度是否满足要求（通常1e-3对于float16可接受）
    max_abs_error = torch.max(torch.abs(output_cpu - output_npu_cpu))
    print(f"Max absolute error: {max_abs_error.item()}")
    
    # WHY: 验证概率性质：所有输出在(0,1)且和为1
    assert torch.all(output_npu_cpu > 0) and torch.all(output_npu_cpu < 1)
    assert torch.allclose(torch.sum(output_npu_cpu, dim=-1), torch.ones(32))
    
    print("CustomSoftmax test PASSED!")

if __name__ == "__main__":
    test_custom_softmax()

常见问题与调试技巧

问题1：算子注册后无法调用

可能原因：

• 算子名称不匹配（注册名与调用名不一致）
• 算子未正确编译或部署
• torch_npu版本与算子编译环境不匹配

解决方案：

# WHY: 列出所有已注册的NPU算子，确认自定义算子是否在列表中
print([op for op in dir(torch.npu) if "softmax" in op.lower()])

问题2：算子执行结果不正确

调试方法：

# WHY: 使用NPU的printf功能，在Ascend C代码中打印调试信息
AscendC::printf("input[0] = %f\n", input_local(0));

# WHY: 使用小规模数据，将中间结果拷贝回CPU逐元素比对

问题3：算子性能不如预期

优化方向：

// WHY: 检查数据搬运是否成为瓶颈，尽量使用Local Buffer
// WHY: 检查是否充分利用了向量计算单元，避免串行计算
// WHY: 使用Ascend C的性能分析工具找出热点

最佳实践总结

1. 先原型后实现：先在Python端定义好算子原型，再实现Ascend C代码
2. 充分测试：编写单元测试、性能测试、边界测试
3. 版本管理：自定义算子与torch_npu版本绑定，避免兼容性问题
4. 文档完善：为自定义算子编写详细文档，包括接口说明、使用示例、性能数据

结语

将自定义算子注册到ATB是扩展昇腾CANN能力的重要手段。通过本文介绍的标准流程，开发者可以高效地将自己的算法实现部署到昇腾NPU上，充分发挥硬件性能。随着昇腾CANN生态的不断完善，自定义算子的开发体验也将越来越好。

---

参考资源：

• 自定义算子开发指南：https://www.atomgit.com/ascend/cann/wikis/自定义算子开发
• Ascend C编程指南：https://www.atomgit.com/ascend/cann/wikis/AscendC
• ATB算子库文档：https://www.atomgit.com/ascend/atb/wikis/Home

相关仓库：

• ATB: https://www.atomgit.com/ascend/atb
• opbase: https://www.atomgit.com/ascend/opbase
• torch_npu: https://www.atomgit.com/ascend/torch_npu
• ops-nn: https://www.atomgit.com/ascend/ops-nn

---

本文档由 CANN 开源社区 AIGC 系统生成，遵循 昇腾CANN 开源协议。