3. 算子定义与注册：接口规范、前向反向注册与属性绑定

好，咱们今天聊聊算子。说白了，算子就是深度学习框架里的“积木块”。你搭建一个网络，本质上就是在拼这些积木。我刚开始接触框架开发时，觉得算子不就是个函数嘛，前向算一下，反向求个导，完事了。后来才发现，这里面的门道可不少。

一个算子要想被框架“认领”，得走一套完整的流程：定义接口、实现逻辑、注册到框架、绑定属性和输入输出。今天我就带你走一遍这个流程。

3.1 算子接口规范：框架的“通用语言”

每个框架都有自己的算子接口规范。但万变不离其宗，核心就三样东西：输入（Inputs）、输出（Outputs）、属性（Attrs）。

我拿一个简单的 Conv2D（二维卷积）算子举例。它的接口大概长这样：

// 伪代码，展示核心结构
OpDef Conv2D {
    // 输入
    Input: input_tensor   // 形状 [N, C, H, W]
    Input: weight_tensor  // 形状 [K, C, R, S]
    Input: bias_tensor    // 可选，形状 [K]

    // 输出
    Output: output_tensor // 形状 [N, K, H_out, W_out]

    // 属性
    Attr: strides         // 类型: list of ints, 默认值: [1, 1]
    Attr: padding         // 类型: string, 可选值: "valid", "same"
    Attr: dilation        // 类型: list of ints, 默认值: [1, 1]
}

你看，接口把“这个算子需要什么、产出什么、怎么配置”说得清清楚楚。框架拿到这个定义，就知道怎么去调度它。

3.2 前向与反向算子注册：让框架“认识”你

接口定义好了，接下来就是注册。注册，就是告诉框架：“嘿，我这儿有个叫 Conv2D 的算子，你记一下，以后遇到它，就调用我的代码。”

注册通常分两步：前向注册 和 反向注册。

3.2.1 前向注册

前向注册，就是把算子的计算逻辑挂到框架的算子库里。我见过很多框架的做法，比如用宏或者注册函数：

// 以 C++ 为例，伪代码
REGISTER_OP_KERNEL("Conv2D", CPU, Conv2DKernel)
    .Input("input_tensor")
    .Input("weight_tensor")
    .Output("output_tensor")
    .Attr("strides", {1, 1})
    .Attr("padding", "valid");

这段代码干了三件事：

• 告诉框架，算子名叫 "Conv2D"
• 指定了它的输入输出和属性
• 绑定了实际的执行函数 Conv2DKernel

这样，当框架的计算图里出现 "Conv2D" 节点时，它就能找到对应的代码去执行。

3.2.2 反向注册

反向注册，就是注册梯度计算逻辑。这是自动求导的关键。我记得有一次，我写了一个自定义算子，前向跑得飞快，但一训练就报错。查了半天，发现是忘了注册反向算子。

反向注册的套路和前向类似：

REGISTER_OP_GRADIENT("Conv2D", Conv2DGrad)
    .Input("input_tensor")
    .Input("weight_tensor")
    .Input("grad_output_tensor")  // 来自上一层的梯度
    .Output("grad_input_tensor")
    .Output("grad_weight_tensor");

这里要注意，反向算子的输入，除了前向的输入，还要加上“上游传下来的梯度”。输出则是“对每个输入的梯度”。

3.3 算子属性与输入输出的绑定：把“积木”拼起来

注册完了，框架只是知道了“有这么一个算子”。但具体怎么用，还得靠属性（Attr）和输入输出（Tensor）的绑定。

属性，是算子的“配置参数”。比如卷积的步长、填充方式。这些参数在计算图构建时就确定了，不会随着数据变化。

输入输出，是算子的“数据流”。数据从输入 Tensor 流进来，经过计算，从输出 Tensor 流出去。

绑定的过程，其实就是把属性和 Tensor 的信息，传递给算子的执行函数。我画个图，你一看就明白：

从图里能看出来，输入 Tensor 和属性是“并行”进入算子的。算子根据属性配置，对输入数据进行计算，最后产出输出 Tensor。

在实际的框架代码里，绑定通常发生在算子注册时，或者计算图构建时。比如在 PyTorch 里，你定义一个 nn.Conv2d 模块，其实就是在绑定属性和输入输出。

3.4 小结：算子注册的“三步走”

好了，咱们把算子定义和注册的流程捋一遍：

1. 定义接口：明确输入、输出、属性，形成“契约”。
2. 注册前向和反向：把算子的计算逻辑和梯度逻辑，挂到框架的算子库里。
3. 绑定属性和 Tensor：把配置和数据，传递给算子的执行函数。

这三步走完，一个算子才算真正被框架“接纳”。以后你写自定义算子，照着这个流程来，基本不会出错。

嗯，今天就先聊到这儿。算子这块内容不少，但核心就是这些。你回去可以拿一个简单的算子，比如加法或者 ReLU，试着走一遍这个流程，感受会更深。