接上一篇文章聊的 metadef，这篇说一个完整的话题：自定义算子怎么从零注册到昇腾 CANN 里，让 PyTorch 或者 TensorFlow 能直接调用。

这个过程涉及到几个环节：算子实现（核函数）、元数据定义（metadef）、算子注册（OpLoader）、框架适配（PyTorch 的 autograd 或者 TensorFlow 的 op 包装）。哪一步出了问题，调用都会失败。

整体链路

先建立一个整体概念。自定义算子在昇腾上跑通，分四步：

1. 写核函数（C++/Ascend C）
   ↓ 生成 .o 文件或者 .so
2. 写 metadef（JSON/proto）
   ↓ 描述算子的接口
3. 注册算子（OpLoader/ACL）
   ↓ 把 .so 和 metadef 关联起来
4. 框架适配（torch.autograd.function / tf.raw_ops）
   ↓ 让 PyTorch/TF 能调用

每一步都有人踩过坑。

第一步：核函数实现

核函数是真正跑在 NPU 上的代码。昇腾的核函数可以用两种语言写：C++ 和 Ascend C。Ascend C 是昇腾提供的 DSL，适合写复杂的算子逻辑；简单场景直接用 C++ 写也可以。

一个最简单的例子：两个 tensor 逐元素相加的算子。

// add_kernel.cl (OpenCL 风格，也可以用 Ascend C)
__kernel void element_add(
    __global const float* a,
    __global const float* b,
    __global float* c,
    const int size
) {
    int gid = get_global_id(0);
    if (gid < size) {
        c[gid] = a[gid] + b[gid];  // 这里没有做 softmax，单纯相加
    }
}

Ascend C 的写法会不太一样，用的是昇腾特有的编程模型：

// add_kernel.cpp（Ascend C）
#include "acl/acl.h"

class ElementAddKernel : public OpKernelBase {
public:
    ElementAddKernel() = default;
    ~ElementAddKernel() = default;

    // Compute 实现
    Status Compute(const OpKernelInput& input, OpKernelOutput* output) override {
        // 获取输入tensor
        const auto& x = input.GetTensor(0);
        const auto& y = input.GetTensor(1);
        auto* z = output->GetTensor(0);

        // 获取数据指针和shape
        float* x_ptr = x.Data<float>();
        float* y_ptr = y.Data<float>();
        float* z_ptr = z.Data<float>();
        int64_t size = x.Size();

        // 逐元素相加，注意这里要处理向量化
        // 昇腾 NPU 喜欢 32 或 64 字节对齐的数据
        for (int64_t i = 0; i < size; i++) {
            z_ptr[i] = x_ptr[i] + y_ptr[i];
        }

        return SUCCESS;
    }
};

编译成 .so：

# 编译脚本
aoc -kernel add_kernel.cpp -o libelement_add.so \
    -I${ACL_ROOT}/include \
    -L${ACL_ROOT}/lib64 \
    -lacl

第二步：写 metadef

核函数写好了，接下来用 metadef 描述它的接口。metadef 里最重要的几个字段：算子名字、输入输出描述、属性描述。

{
  "op_name": "element_add",
  "op_type": "Custom",
  "input_desc": [
    {
      "name": "x",
      "dtype": ["float32"],
      "format": ["ND"],
      "shape": [-1]
    },
    {
      "name": "y",
      "dtype": ["float32"],
      "format": ["ND"],
      "shape": [-1]
    }
  ],
  "output_desc": [
    {
      "name": "z",
      "dtype": ["float32"],
      "format": ["ND"],
      "shape": [-1]
    }
  ]
}

这里 dtype 和 shape 都用列表表示，表示支持多种组合。比如 dtype: ["float32", "float16"] 表示这个算子可以接受 float32 和 float16 两种输入。

第三步：注册算子

有了核函数和 metadef，接下来要把它们注册到 ACL 里，这样 ACL 才能根据名字找到对应的实现。

import acl

# 初始化 ACL
acl.init()
acl.rt.set_device(0)

# 加载核函数 .so
ret = acl.ops.load_operator_library("/path/to/libelement_add.so")
if ret != 0:
    raise RuntimeError(f"Failed to load operator library: {ret}")

# 注册算子（把名字和实际实现关联起来）
ret = acl.op.register_operator("element_add")
if ret != 0:
    raise RuntimeError(f"Failed to register operator: {ret}")

# 注册算子模型（关联 metadef 和核函数）
ret = acl.op.register_operator_model(
    "element_add",
    "/path/to/element_add.json",  # metadef 文件
    "element_add"                  # 核函数里的实际名字
)
if ret != 0:
    raise RuntimeError(f"Failed to register operator model: {ret}")

print("算子注册成功")

注册成功之后，理论上 ACL 就知道 element_add 是什么、怎么调用了。但这时候还不能在 PyTorch 里直接用，需要做第四步的框架适配。

第四步：PyTorch 适配

PyTorch 昇腾适配自定义算子，主要靠 torch.autograd.Function 和 torch.autograd.function:

import torch
from torch.autograd import Function
import acl
import numpy as np

class ElementAdd(Function):
    @staticmethod
    def forward(ctx, x, y):
        # 这里调用昇腾 ACL 的单算子执行接口
        # 注意要先确保算子已经注册过了
        z = acl.ops.element_add(x, y)
        return z

    @staticmethod
    def backward(ctx, grad_output):
        # 反向也要注册对应的反向算子
        grad_x = grad_output
        grad_y = grad_output
        return grad_x, grad_y

# 包装成 nn.Module，方便在模型里用
class ElementAddModule(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        return ElementAdd.apply(x, y)

一个常见的坑：反向算子。

自定义算子如果要在训练里用，必须注册反向传播的实现。没有反向的话，模型只能做推理，不能做训练。很多新手只注册了前向算子，训练的时候才发现梯度回不来。

# 反向算子的注册（以 ACL 接口为例）
ret = acl.op.register_operator_gradient(
    "element_add",  # 前向算子名字
    "/path/to/element_add_grad.json"  # 反向算子的 metadef
)

调试注册问题

注册流程里最容易出错的地方：

1. .so 路径问题

路径必须是绝对路径，相对路径在 ACL 里行为不一致。注册之前先确认文件存在：

import os
lib_path = "/path/to/libelement_add.so"
assert os.path.exists(lib_path), f"Library not found: {lib_path}"

2. metadef 格式错误

用官方提供的 validator 先过一遍：

python -m metadef.validator /path/to/element_add.json

3. 注册顺序问题

必须先 load_operator_library，再 register_operator，再 register_operator_model。顺序搞反会报奇怪的链接错误。

4. 多进程重复注册

推理服务如果是多 worker 模式，要确保算子只注册一次。可以用 torch.distributed 的 barrier 或者单例模式控制。

# 用环境变量控制只注册一次
import os
if os.environ.get("RANK", "0") == "0":
    register_custom_ops()

性能相关的补充

自定义算子的性能往往不如昇腾原生算子，原因很直接：原生算子是昇腾工程师手写的，深度优化过的。自定义算子如果不做特殊处理，就是最朴素的实现。

几个提升性能的方向：

向量化：昇腾 NPU 的向量单元一次能处理 16/32/64 个元素，如果循环里一次只处理一个元素，利用率会很低。

内存对齐：输入 tensor 的地址最好 32 字节对齐，否则向量化指令可能触发 memory misaligned 的异常处理逻辑，拖慢速度。

融合：如果一个计算图里有多个自定义算子逐个执行，考虑把它们合并成一个，减少中间结果的显存写入。

仓库在 https://atomgit.com/cann/metadef，可以参考官方自定义算子的注册示例。